From 055b4a485bb87331f2b6fd3788f7c8b18432e563 Mon Sep 17 00:00:00 2001 From: Nathan Musoke Date: Fri, 26 Jan 2024 16:15:51 -0500 Subject: [PATCH] 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. --- examples/2d.ipynb | 13 +---- examples/init_soliton.jl | 53 ----------------- examples/soliton_velocity.jl | 4 +- src/UltraDark.jl | 20 ++++++- src/initialise.jl | 109 +++++++++++++++++++++++++++++++++++ test/sims/scale_factor.jl | 14 ++--- test/sims/soliton_static.jl | 14 ++--- 7 files changed, 142 insertions(+), 85 deletions(-) delete mode 100644 examples/init_soliton.jl create mode 100644 src/initialise.jl diff --git a/examples/2d.ipynb b/examples/2d.ipynb index b31c3ba..de406d9 100644 --- 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 deleted file mode 100644 index 15a2252..0000000 --- a/examples/init_soliton.jl +++ /dev/null @@ -1,53 +0,0 @@ -using UltraDark -using PencilFFTs -using NPZ - -function add_soliton(grids, mass, position, velocity, phase, t0) - add_soliton(grids, grids.ψx, mass, position, velocity, phase, t0) -end - -function add_soliton(grids, psi, mass, position, velocity, phase, t0) - - profile = npzread(joinpath(@__DIR__, "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 diff --git a/examples/soliton_velocity.jl b/examples/soliton_velocity.jl index 97fba6c..e2c47fc 100644 --- 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 index e6bb2bf..d60993b 100644 --- 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 new file mode 100644 index 0000000..99b1b80 --- /dev/null +++ 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 index bf77634..e8cfd4c 100644 --- 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 index 5ab85b1..6947da8 100644 --- 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