Revise initial condition and shapes API (#335)

* Allow defining `InitialCondition` quantities as functions of position

* Refactor IC and shapes API

* Reformat

* Fix tests

* Update docs

* Add tests for initial condition

---------

Co-authored-by: Niklas Neher <[email protected]>

examples/fluid/moving_wall_2d.jl

@@ -36,7 +36,7 @@ n_wall_particles_y = size(tank.face_indices[2], 2)
 
 wall = RectangularShape(boundary_particle_spacing,
                         (boundary_layers, n_wall_particles_y),
-                        (wall_position, 0.0), fluid_density)
+                        (wall_position, 0.0), density=fluid_density)
 
 # Movement function
 f_y(t) = 0.0

examples/fluid/periodic_channel_2d.jl

@@ -26,7 +26,7 @@ state_equation = StateEquationCole(sound_speed, 7, fluid_density, atmospheric_pr
 
 tank = RectangularTank(fluid_particle_spacing, initial_fluid_size, tank_size, fluid_density,
                        n_layers=boundary_layers, spacing_ratio=spacing_ratio,
-                       faces=(false, false, true, true), init_velocity=initial_velocity)
+                       faces=(false, false, true, true), velocity=initial_velocity)
 
 # ==========================================================================================
 # ==== Fluid

examples/fsi/dam_break_2d.jl

@@ -56,11 +56,11 @@ n_particles_y = round(Int, length_beam / solid_particle_spacing) + 1
 # but not for solids. We therefore need to pass `tlsph=true`.
 plate = RectangularShape(solid_particle_spacing,
                          (n_particles_x, n_particles_y - 1),
-                         (2initial_fluid_size[1], solid_particle_spacing), solid_density,
-                         tlsph=true)
+                         (2initial_fluid_size[1], solid_particle_spacing),
+                         density=solid_density, tlsph=true)
 fixed_particles = RectangularShape(solid_particle_spacing,
                                    (n_particles_x, 1), (2initial_fluid_size[1], 0.0),
-                                   solid_density, tlsph=true)
+                                   density=solid_density, tlsph=true)
 
 solid = union(plate, fixed_particles)
 

examples/fsi/dam_break_gate_2d.jl

@@ -47,7 +47,7 @@ gate_height = initial_fluid_size[2] + 4 * fluid_particle_spacing
 gate = RectangularShape(boundary_particle_spacing,
                         (boundary_layers,
                          round(Int, gate_height / boundary_particle_spacing)),
-                        (initial_fluid_size[1], 0.0), fluid_density)
+                        (initial_fluid_size[1], 0.0), density=fluid_density)
 
 # Movement of the gate according to the paper
 f_x(t) = 0.0
@@ -79,11 +79,11 @@ n_particles_y = round(Int, length_beam / solid_particle_spacing) + 1
 plate_position = 0.6 - n_particles_x * solid_particle_spacing
 plate = RectangularShape(solid_particle_spacing,
                          (n_particles_x, n_particles_y - 1),
-                         (plate_position, solid_particle_spacing), solid_density,
-                         tlsph=true)
+                         (plate_position, solid_particle_spacing),
+                         density=solid_density, tlsph=true)
 fixed_particles = RectangularShape(solid_particle_spacing,
                                    (n_particles_x, 1), (plate_position, 0.0),
-                                   solid_density, tlsph=true)
+                                   density=solid_density, tlsph=true)
 
 solid = union(plate, fixed_particles)
 

examples/fsi/falling_water_column_2d.jl

@@ -30,7 +30,7 @@ state_equation = StateEquationCole(sound_speed, 7, fluid_density, atmospheric_pr
 
 fluid = RectangularShape(fluid_particle_spacing,
                          round.(Int, (initial_fluid_size ./ fluid_particle_spacing)),
-                         (0.1, 0.2), fluid_density)
+                         (0.1, 0.2), density=fluid_density)
 
 # ==========================================================================================
 # ==== Fluid

examples/n_body/n_body_benchmark_trixi.jl

@@ -55,20 +55,20 @@ coordinates = [0.0 4.84143144246472090e+0 8.34336671824457987e+0 1.2894369562139
                0.0 -1.16032004402742839e+0 4.12479856412430479e+0 -1.51111514016986312e+1 -2.59193146099879641e+1;
                0.0 -1.03622044471123109e-1 -4.03523417114321381e-1 -2.23307578892655734e-1 1.79258772950371181e-1]
 
-velocities = [0.0 1.66007664274403694e-3 -2.76742510726862411e-3 2.96460137564761618e-3 2.68067772490389322e-3;
-              0.0 7.69901118419740425e-3 4.99852801234917238e-3 2.37847173959480950e-3 1.62824170038242295e-3;
-              0.0 -6.90460016972063023e-5 2.30417297573763929e-5 -2.96589568540237556e-5 -9.51592254519715870e-5] *
-             DAYS_PER_YEAR
+velocity = [0.0 1.66007664274403694e-3 -2.76742510726862411e-3 2.96460137564761618e-3 2.68067772490389322e-3;
+            0.0 7.69901118419740425e-3 4.99852801234917238e-3 2.37847173959480950e-3 1.62824170038242295e-3;
+            0.0 -6.90460016972063023e-5 2.30417297573763929e-5 -2.96589568540237556e-5 -9.51592254519715870e-5] *
+           DAYS_PER_YEAR
 
 masses = [
     1.0, 9.54791938424326609e-4, 2.85885980666130812e-4, 4.36624404335156298e-5,
     5.15138902046611451e-5,
 ] * SOLAR_MASS
 
 # Offset sun momentum
-velocities[:, 1] = -velocities[:, 2:end] * masses[2:end] / SOLAR_MASS
+velocity[:, 1] = -velocity[:, 2:end] * masses[2:end] / SOLAR_MASS
 
-initial_condition = InitialCondition(coordinates, velocities, masses, zeros(size(masses)))
+initial_condition = InitialCondition(; coordinates, velocity, density=0.0, mass=masses)
 
 G = 1.0
 particle_system = NBodySystem(initial_condition, G)
@@ -79,8 +79,8 @@ particle_system = NBodySystem(initial_condition, G)
 semi = Semidiscretization(particle_system)
 
 # This is significantly faster than using OrdinaryDiffEq.
-function symplectic_euler!(velocities, coordinates, semi)
-    v = vec(velocities)
+function symplectic_euler!(velocity, coordinates, semi)
+    v = vec(velocity)
     u = vec(coordinates)
     dv = copy(v)
     du = copy(u)
@@ -110,14 +110,14 @@ open(filename, "w") do f
     end
 end
 
-@printf("%.9f\n", energy(velocities, coordinates, particle_system, semi))
+@printf("%.9f\n", energy(velocity, coordinates, particle_system, semi))
 
 # Disable multithreading, since it adds a significant overhead for this small problem.
 disable_polyester_threads() do
-    symplectic_euler!(velocities, coordinates, semi)
+    symplectic_euler!(velocity, coordinates, semi)
 end
 
-@printf("%.9f\n", energy(velocities, coordinates, particle_system, semi))
+@printf("%.9f\n", energy(velocity, coordinates, particle_system, semi))
 
 # Enable timers again
 open(filename, "w") do f

examples/n_body/n_body_solar_system.jl

@@ -15,15 +15,15 @@ include("n_body_system.jl")
 coordinates = [5.5850e+08 5.1979e+10 -1.5041e+10 -1.1506e+09 -4.8883e+10 -8.1142e+11 -4.2780e+11 2.7878e+12 4.2097e+12;
                5.5850e+08 7.6928e+09 9.7080e+10 -1.3910e+11 -1.9686e+11 4.5462e+10 -1.3353e+12 9.9509e+11 -1.3834e+12;
                5.5850e+08 -1.2845e+09 4.4635e+10 -6.0330e+10 -8.8994e+10 3.9229e+10 -5.3311e+11 3.9639e+08 -6.7105e+11]
-velocities = [-1.4663 -1.5205e+04 -3.4770e+04 2.9288e+04 2.4533e+04 -1.0724e+03 8.7288e+03 -2.4913e+03 1.8271e+03;
-              11.1238 4.4189e+04 -5.5933e+03 -398.5759 -2.7622e+03 -1.1422e+04 -2.4369e+03 5.5197e+03 4.7731e+03;
-              4.8370 2.5180e+04 -316.8994 -172.5873 -1.9295e+03 -4.8696e+03 -1.3824e+03 2.4527e+03 1.9082e+03]
+velocity = [-1.4663 -1.5205e+04 -3.4770e+04 2.9288e+04 2.4533e+04 -1.0724e+03 8.7288e+03 -2.4913e+03 1.8271e+03;
+            11.1238 4.4189e+04 -5.5933e+03 -398.5759 -2.7622e+03 -1.1422e+04 -2.4369e+03 5.5197e+03 4.7731e+03;
+            4.8370 2.5180e+04 -316.8994 -172.5873 -1.9295e+03 -4.8696e+03 -1.3824e+03 2.4527e+03 1.9082e+03]
 
 masses = [
     1.99e30, 3.30e23, 4.87e24, 5.97e24, 6.42e23, 1.90e27, 5.68e26, 8.68e25, 1.02e26,
 ]
 
-initial_condition = InitialCondition(coordinates, velocities, masses, zeros(size(masses)))
+initial_condition = InitialCondition(; coordinates, velocity, density=0.0, mass=masses)
 
 G = 6.6743e-11
 particle_system = NBodySystem(initial_condition, G)

examples/solid/oscillating_beam_2d.jl

@@ -39,7 +39,7 @@ n_particles_per_dimension = (round(Int, length_beam / particle_spacing) +
 # from the boundary, which is correct for fluids, but not for solids.
 # We therefore need to pass `tlsph=true`.
 beam = RectangularShape(particle_spacing, n_particles_per_dimension,
-                        (0.0, 0.0), material_density, tlsph=true)
+                        (0.0, 0.0), density=material_density, tlsph=true)
 
 solid = union(beam, fixed_particles)
 

src/general/initial_condition.jl

@@ -1,6 +1,6 @@
 @doc raw"""
-    InitialCondition(coordinates, velocities, masses, densities; pressure=0.0,
-                     particle_spacing=-1.0)
+    InitialCondition(; coordinates, density, velocity=zeros(size(coordinates, 1)),
+                     mass=nothing, pressure=0.0, particle_spacing=-1.0)
 
 Struct to hold the initial configuration of the particles.
 
@@ -14,46 +14,58 @@ to build more complex geometries.
 
 # Arguments
 - `coordinates`: An array where the $i$-th column holds the coordinates of particle $i$.
-- `velocities`: An array where the $i$-th column holds the velocity of particle $i$.
-- `masses`: A vector holding the mass of each particle.
-- `densities`: A vector holding the density of each particle.
+- `density`:     Either a vector holding the density of each particle,
+                 or a function mapping each particle's coordinates to its density,
+                 or a scalar for a constant density over all particles.
 
 # Keywords
-- `pressure`: Either a vector of pressure values of each particle or a number for a constant
-              pressure over all particles. This is optional and only needed when using
-              the [`EntropicallyDampedSPHSystem`](@ref).
-- `particle_spacing`: The spacing between the particles. This is a number, as the spacing
+- `velocity`:   Either an array where the $i$-th column holds the velocity of particle $i$,
+                or a function mapping each particle's coordinates to its velocity,
+                or, for a constant fluid velocity, a vector holding this velocity.
+                Velocity is constant zero by default.
+- `mass`:       Either `nothing` (default) to automatically compute particle mass from particle
+                density and spacing, or a vector holding the mass of each particle,
+                or a function mapping each particle's coordinates to its mass,
+                or a scalar for a constant mass over all particles.
+- `pressure`:   Either a vector holding the pressure of each particle,
+                or a function mapping each particle's coordinates to its pressure,
+                or a scalar for a constant pressure over all particles. This is optional and
+                only needed when using the [`EntropicallyDampedSPHSystem`](@ref).
+- `particle_spacing`: The spacing between the particles. This is a scalar, as the spacing
                       is assumed to be uniform. This is only needed when using
-                      set operations on the `InitialCondition`.
+                      set operations on the `InitialCondition` or for automatic mass calculation.
 
 # Examples
 ```julia
 # Rectangle filled with particles
-initial_condition = RectangularShape(0.1, (3, 4), (-1.0, 1.0), 1.0)
+initial_condition = RectangularShape(0.1, (3, 4), (-1.0, 1.0), density=1.0)
 
 # Two spheres in one initial condition
 initial_condition = union(SphereShape(0.15, 0.5, (-1.0, 1.0), 1.0),
                           SphereShape(0.15, 0.2, (0.0, 1.0), 1.0))
 
 # Rectangle with a spherical hole
-shape1 = RectangularShape(0.1, (16, 13), (-0.8, 0.0), 1.0)
+shape1 = RectangularShape(0.1, (16, 13), (-0.8, 0.0), density=1.0)
 shape2 = SphereShape(0.1, 0.35, (0.0, 0.6), 1.0, sphere_type=RoundSphere())
 initial_condition = setdiff(shape1, shape2)
 
 # Intersect of a rectangle with a sphere. Note that this keeps the particles of the
 # rectangle that are in the intersect, while `intersect(shape2, shape1)` would consist of
 # the particles of the sphere that are in the intersect.
-shape1 = RectangularShape(0.1, (16, 13), (-0.8, 0.0), 1.0)
+shape1 = RectangularShape(0.1, (16, 13), (-0.8, 0.0), density=1.0)
 shape2 = SphereShape(0.1, 0.35, (0.0, 0.6), 1.0, sphere_type=RoundSphere())
 initial_condition = intersect(shape1, shape2)
 
 # Build `InitialCondition` manually
 coordinates = [0.0 1.0 1.0
                0.0 0.0 1.0]
-velocities = zero(coordinates)
-masses = ones(3)
-densities = 1000 * ones(3)
-initial_condition = InitialCondition(coordinates, velocities, masses, densities)
+velocity = zero(coordinates)
+mass = ones(3)
+density = 1000 * ones(3)
+initial_condition = InitialCondition(; coordinates, velocity, mass, density)
+
+# With functions
+initial_condition = InitialCondition(; coordinates, velocity=x -> 2x, density=1000.0)
 ```
 """
 struct InitialCondition{ELTYPE}
@@ -64,33 +76,105 @@ struct InitialCondition{ELTYPE}
     density          :: Array{ELTYPE, 1}
     pressure         :: Array{ELTYPE, 1}
 
-    function InitialCondition(coordinates, velocities, masses, densities; pressure=0.0,
-                              particle_spacing=-1.0)
+    function InitialCondition(; coordinates, density, velocity=zeros(size(coordinates, 1)),
+                              mass=nothing, pressure=0.0, particle_spacing=-1.0)
+        NDIMS = size(coordinates, 1)
+
+        return InitialCondition{NDIMS}(coordinates, velocity, mass, density,
+                                       pressure, particle_spacing)
+    end
+
+    # Function barrier to make `NDIMS` static and therefore SVectors type-stable
+    function InitialCondition{NDIMS}(coordinates, velocity, mass, density,
+                                     pressure, particle_spacing) where {NDIMS}
         ELTYPE = eltype(coordinates)
+        n_particles = size(coordinates, 2)
+
+        if n_particles == 0
+            return new{ELTYPE}(particle_spacing, coordinates, zeros(ELTYPE, NDIMS, 0),
+                               zeros(ELTYPE, 0), zeros(ELTYPE, 0), zeros(ELTYPE, 0))
+        end
+
+        # SVector of coordinates to pass to functions
+        coordinates_svector = reinterpret(reshape, SVector{NDIMS, ELTYPE}, coordinates)
+
+        if velocity isa AbstractMatrix
+            velocities = velocity
+            if size(coordinates) != size(velocities)
+                throw(ArgumentError("`coordinates` and `velocities` must be of the same size"))
+            end
+        else
+            # Assuming `velocity` is a scalar or a function
+            velocity_fun = wrap_function(velocity, Val(NDIMS))
+            if length(velocity_fun(coordinates_svector[1])) != NDIMS
+                throw(ArgumentError("`velocity` must be $NDIMS-dimensional " *
+                                    "for $NDIMS-dimensional `coordinates`"))
+            end
+            velocities_svector = velocity_fun.(coordinates_svector)
+            velocities = reinterpret(reshape, ELTYPE, velocities_svector)
+        end
 
-        if size(coordinates) != size(velocities)
-            throw(ArgumentError("`coordinates` and `velocities` must be of the same size"))
+        if density isa AbstractVector
+            if length(density) != n_particles
+                throw(ArgumentError("Expected: length(density) == size(coordinates, 2)\n" *
+                                    "Got: size(coordinates, 2) = $(size(coordinates, 2)), " *
+                                    "length(density) = $(length(density))"))
+            end
+            densities = density
+        else
+            density_fun = wrap_function(density, Val(NDIMS))
+            densities = density_fun.(coordinates_svector)
         end
 
-        if !(size(coordinates, 2) == length(masses) == length(densities))
-            throw(ArgumentError("Expected: size(coordinates, 2) == length(masses) == length(densities)\n" *
-                                "Got: size(coordinates, 2) = $(size(coordinates, 2)), " *
-                                "length(masses) = $(length(masses)), " *
-                                "length(densities) = $(length(densities))"))
+        if pressure isa AbstractVector
+            if length(pressure) != n_particles
+                throw(ArgumentError("Expected: length(pressure) == size(coordinates, 2)\n" *
+                                    "Got: size(coordinates, 2) = $(size(coordinates, 2)), " *
+                                    "length(pressure) = $(length(pressure))"))
+            end
+            pressures = pressure
+        else
+            pressure_fun = wrap_function(pressure, Val(NDIMS))
+            pressures = pressure_fun.(coordinates_svector)
         end
 
-        if pressure isa Number
-            pressure = pressure * ones(ELTYPE, length(masses))
-        elseif length(pressure) != length(masses)
-            throw(ArgumentError("`pressure` must either be a scalar or a vector of the " *
-                                "same length as `masses`"))
+        if mass isa AbstractVector
+            if length(mass) != n_particles
+                throw(ArgumentError("Expected: length(mass) == size(coordinates, 2)\n" *
+                                    "Got: size(coordinates, 2) = $(size(coordinates, 2)), " *
+                                    "length(mass) = $(length(mass))"))
+            end
+            masses = mass
+        elseif mass === nothing
+            if particle_spacing < 0
+                throw(ArgumentError("`mass` must be specified when not using `particle_spacing`"))
+            end
+            particle_volume = particle_spacing^NDIMS
+            masses = particle_volume * densities
+        else
+            mass_fun = wrap_function(mass, Val(NDIMS))
+            masses = mass_fun.(coordinates_svector)
         end
 
         return new{ELTYPE}(particle_spacing, coordinates, velocities, masses,
-                           densities, pressure)
+                           densities, pressures)
     end
 end
 
+function wrap_function(function_::Function, ::Val)
+    # Already a function
+    return function_
+end
+
+function wrap_function(constant_scalar::Number, ::Val)
+    return coords -> constant_scalar
+end
+
+# For vectors and tuples
+function wrap_function(constant_vector, ::Val{NDIMS}) where {NDIMS}
+    return coords -> SVector{NDIMS}(constant_vector)
+end
+
 @inline function Base.ndims(initial_condition::InitialCondition)
     return size(initial_condition.coordinates, 1)
 end
@@ -127,8 +211,8 @@ function Base.union(initial_condition::InitialCondition, initial_conditions...)
     density = vcat(initial_condition.density, ic.density[valid_particles])
     pressure = vcat(initial_condition.pressure, ic.pressure[valid_particles])
 
-    result = InitialCondition(coordinates, velocity, mass, density, pressure=pressure,
-                              particle_spacing=particle_spacing)
+    result = InitialCondition{ndims(ic)}(coordinates, velocity, mass, density, pressure,
+                                         particle_spacing)
 
     return union(result, Base.tail(initial_conditions)...)
 end
@@ -157,8 +241,8 @@ function Base.setdiff(initial_condition::InitialCondition, initial_conditions...
     density = initial_condition.density[valid_particles]
     pressure = initial_condition.pressure[valid_particles]
 
-    result = InitialCondition(coordinates, velocity, mass, density, pressure=pressure,
-                              particle_spacing=particle_spacing)
+    result = InitialCondition{ndims(ic)}(coordinates, velocity, mass, density, pressure,
+                                         particle_spacing)
 
     return setdiff(result, Base.tail(initial_conditions)...)
 end
@@ -186,8 +270,8 @@ function Base.intersect(initial_condition::InitialCondition, initial_conditions.
     density = initial_condition.density[too_close]
     pressure = initial_condition.pressure[too_close]
 
-    result = InitialCondition(coordinates, velocity, mass, density, pressure=pressure,
-                              particle_spacing=particle_spacing)
+    result = InitialCondition{ndims(ic)}(coordinates, velocity, mass, density, pressure,
+                                         particle_spacing)
 
     return intersect(result, Base.tail(initial_conditions)...)
 end