feat: Soliton initialisation

Previously soliton initialisation was in a script which was imported in some examples. Move this into the package as a function so it is properly documented and tested. Add this function to precompilation.
musoke · Jan 26, 2024 · 055b4a4 · 055b4a4
1 parent 786dbc6
commit 055b4a4
Show file tree

Hide file tree

Showing 7 changed files with 142 additions and 85 deletions.
diff --git a/examples/2d.ipynb b/examples/2d.ipynb
@@ -79,15 +79,6 @@
     "Load a function that initialises solitons."
    ]
   },
-  {
-   "cell_type": "code",
-   "execution_count": 5,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "include(joinpath(@__DIR__, \"init_soliton.jl\"));"
-   ]
-  },
   {
    "cell_type": "markdown",
    "metadata": {},
@@ -128,8 +119,8 @@
     "phase_2 = π\n",
     "t0 = 0\n",
     "\n",
-    "add_soliton(grids, mass, position_1, velocity, phase_1, t0)\n",
-    "add_soliton(grids, mass, position_2, velocity, phase_2, t0)"
+    "UltraDark.Initialise.add_fdm_soliton!(grids, mass, position_1, velocity, phase_1, t0)\n",
+    "UltraDark.Initialise.add_fdm_soliton!(grids, mass, position_2, velocity, phase_2, t0)"
    ]
   },
   {

diff --git a/examples/init_soliton.jl b/examples/init_soliton.jl
diff --git a/examples/soliton_velocity.jl b/examples/soliton_velocity.jl
@@ -9,8 +9,6 @@ end
 comm = MPI.COMM_WORLD
 @info "Process ID $(MPI.Comm_rank(comm)+1) of $(MPI.Comm_size(comm)) with $(Threads.nthreads()) threads\n"
 
-include(joinpath(@__DIR__, "init_soliton.jl"))
-
 const resol = 64
 
 const num_snapshots = 20
@@ -37,7 +35,7 @@ function run_sim()
     npzwrite(joinpath(output_config.directory, "y.npy"), grids.y)
     npzwrite(joinpath(output_config.directory, "z.npy"), grids.z)
 
-    add_soliton(grids, mass, position0, velocity, phase, t0)
+    UltraDark.Initialise.add_fdm_soliton!(grids, mass, position0, velocity, phase, t0)
     @info "Initialized" MPI.Comm_rank(comm)
 
     simulate!(grids, options, output_config) == nothing

diff --git a/src/UltraDark.jl b/src/UltraDark.jl
@@ -29,6 +29,8 @@ import .Output: output_state, output_xyz, output_external_states_headers
 import .Output: output_summary_row, output_summary_header
 import .Config: SimulationConfig, constant_scale_factor, TimeStepOptions
 
+include("initialise.jl")
+
 """
     outer_step!(Δt, grids, constants; a=1.0)
     outer_step!(Δt, grids, constants, s; a=1.0)
@@ -277,6 +279,13 @@ end
 @setup_workload begin
 
     resol = 64
+
+    soliton_mass = 10.0
+    soliton_position = [0.0, 0.0, 0.0]
+    soliton_velocity = [0.0, 0.0, 0.0]
+    soliton_phase = 0.0
+    soliton_t0 = 0.0
+
     output_dir = mktempdir()
 
     @compile_workload begin
@@ -294,8 +303,15 @@ end
         options = Config.SimulationConfig()
 
         grids = Grids(10.0, resol)
-        # Set ψx to something non-zero
-        grids.ψx .= (grids.x .^ 2 .+ grids.y .^ 2 .+ grids.z .^ 2) .^ 0.5 ./ 1e9
+
+        Initialise.add_fdm_soliton!(
+            grids,
+            soliton_mass,
+            soliton_position,
+            soliton_velocity,
+            soliton_phase,
+            soliton_t0,
+        )
 
         simulate!(grids, options, output_config)
     end

diff --git a/src/initialise.jl b/src/initialise.jl
@@ -0,0 +1,109 @@
+module Initialise
+
+using PencilFFTs
+using NPZ
+
+using ..UltraDark: AbstractGrids
+
+export add_fdm_soliton!
+
+"""
+    add_fdm_soliton!(grids::AbstractGrids, mass, position, velocity, phase, t0)
+    add_fdm_soliton!(grids::AbstractGrids, psi, mass, position, velocity, phase, t0)
+
+Add a fuzzy dark matter soliton `grids`.  The
+density profile is rescaled to the desired mass and the phase is set to the desired
+velocity.
+
+If the argument `psi` is passed, the soliton is added to the array-like `psi` rather than
+than `grids.ψx`.
+
+Note that due to coarse grid effects, the mass of the added soliton may not match the input
+mass exactly.
+
+The included density profile from comes [from PyUltraLight](https://github.com/auckland-cosmo/PyUltraLight/blob/28d942cb3100b6df9976a6e6f392eaeeb2ce9689/Soliton%20Profile%20Files/initial_f.npy).
+
+# Examples
+
+The following  script adds a soliton to a grid and checks that the resulting total mass is
+correct.
+
+```jldoctest
+using UltraDark
+using UltraDark.Initialise
+
+resol = 64
+grids = Grids(10.0, resol)
+
+mass = 10.0
+position = [0.0, 0.0, 0.0]
+velocity = [0.0, 0.0, 0.0]
+phase = 0.0
+t0 = 0.0
+
+add_fdm_soliton!(grids, mass, position, velocity, phase, t0)
+
+# Ensure density is correct
+UltraDark.update_gravitational_potential!(grids, ())
+
+actual_mass = UltraDark.mass(grids)
+
+isapprox(mass, actual_mass, rtol = 1e-3)
+
+# output
+
+true
+```
+"""
+function add_fdm_soliton!(grids::AbstractGrids, mass, position, velocity, phase, t0)
+    add_fdm_soliton!(grids, grids.ψx, mass, position, velocity, phase, t0)
+end
+
+function add_fdm_soliton!(grids::AbstractGrids, psi, mass, position, velocity, phase, t0)
+
+    profile::Vector{Float64} =
+        npzread(joinpath(@__DIR__, "..", "examples", "initial_f.npy"))
+    delta_x = 0.00001
+    alpha = (mass / 3.883)^2
+    beta = 2.454
+
+    if typeof(psi) <: PencilArray
+        ψ_glob = global_view(psi)
+    elseif typeof(psi) <: Array
+        ψ_glob = psi
+    else
+        throw("Unrecognised array type")
+    end
+
+    for I in CartesianIndices(ψ_glob)
+
+        i, j, k = Tuple(I)  # unpack indices
+
+        x = grids.x[i]
+        y = grids.y[j]
+        z = grids.z[k]
+
+        vx = velocity[1]
+        vy = velocity[2]
+        vz = velocity[3]
+
+        distfromcentre =
+            ((x - position[1])^2 + (y - position[2])^2 + (z - position[3])^2)^0.5
+
+        if alpha^0.5 * distfromcentre <= 5.6
+            f = alpha * profile[trunc(Int, alpha^0.5 * (distfromcentre / delta_x + 1))]
+            f *= exp(
+                im * (
+                    alpha * beta * t0 + (x * vx + y * vy + z * vz) -
+                    0.5 * (vx^2 + vy^2 + vz^2) * t0
+                ),
+            )
+            f *= exp(im * phase)
+            ψ_glob[I] += f
+        else
+            ψ_glob[I] += 0
+        end
+    end
+end
+
+end # module
diff --git a/test/sims/scale_factor.jl b/test/sims/scale_factor.jl
@@ -8,8 +8,6 @@ using PencilFFTs
 
 resol = 64
 
-include("../../examples/init_soliton.jl")
-
 output_dir = mktempdir()
 
 for a in [UltraDark.constant_scale_factor, t -> t]
@@ -29,13 +27,13 @@ for a in [UltraDark.constant_scale_factor, t -> t]
 
     grids = Grids(10.0, resol)
 
-    mass = 10
-    position = [0, 0, 0]
-    velocity = [0, 0, 0]
-    phase = 0
-    t0 = 0
+    mass = 10.0
+    position = [0.0, 0.0, 0.0]
+    velocity = [0.0, 0.0, 0.0]
+    phase = 0.0
+    t0 = 0.0
 
-    add_soliton(grids, mass, position, velocity, phase, t0)
+    UltraDark.Initialise.add_fdm_soliton!(grids, mass, position, velocity, phase, t0)
 
     @test simulate!(grids, options, output_config) == nothing
 

diff --git a/test/sims/soliton_static.jl b/test/sims/soliton_static.jl
@@ -7,8 +7,6 @@ using PencilFFTs
 
 resol = 64
 
-include("../../examples/init_soliton.jl")
-
 output_dir = mktempdir()
 
 for grid_type in [Grids, PencilGrids]
@@ -28,13 +26,13 @@ for grid_type in [Grids, PencilGrids]
 
     grids = grid_type(10.0, resol)
 
-    mass = 10
-    position = [0, 0, 0]
-    velocity = [0, 0, 0]
-    phase = 0
-    t0 = 0
+    mass = 10.0
+    position = [0.0, 0.0, 0.0]
+    velocity = [0.0, 0.0, 0.0]
+    phase = 0.0
+    t0 = 0.0
 
-    add_soliton(grids, mass, position, velocity, phase, t0)
+    UltraDark.Initialise.add_fdm_soliton!(grids, mass, position, velocity, phase, t0)
 
     @test simulate!(grids, options, output_config) == nothing
 end