Support arbitrary spherically symmetric density profiles. Release v0.1.3

Author: islent

Project.toml

@@ -1,13 +1,14 @@
 name = "AstroIC"
 uuid = "5b3ca070-48c3-41ea-b469-89825dd01fde"
 authors = ["islent <[email protected]>"]
-version = "0.1.2"
+version = "0.1.3"
 
 [deps]
 AstroIO = "c85a633c-0c3f-44a2-bffe-7f9d0681b3e7"
 AstroLib = "c7932e45-9af1-51e7-9da9-f004cd3a462b"
 AstroSimBase = "c6a34d98-f626-4d8d-b450-a3f124eaada6"
 BangBang = "198e06fe-97b7-11e9-32a5-e1d131e6ad66"
+DataFrames = "a93c6f00-e57d-5684-b7b6-d8193f3e46c0"
 Dates = "ade2ca70-3891-5945-98fb-dc099432e06a"
 Dierckx = "39dd38d3-220a-591b-8e3c-4c3a8c710a94"
 Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f"
@@ -29,13 +30,14 @@ AstroIO = "0.1"
 AstroLib = "0.4"
 AstroSimBase = "0.1"
 BangBang = "0.3, 0.4, 0.5"
+DataFrames = "1"
 Dierckx = "0.5"
 Distributions = "0.25"
 DocStringExtensions = "0.8, 0.9"
-Optim = "1.9.4"
+Optim = "1"
 PhysicalConstants = "0.2"
 PhysicalParticles = "1"
-Plots = "1.40.8"
+Plots = "1"
 PrecompileTools = "1"
 QuadGK = "2"
 Reexport = "1"

src/AstroIC.jl

@@ -6,6 +6,7 @@ using PrecompileTools
 using DocStringExtensions
 using Reexport
 using Unitful, UnitfulAstro
+using DataFrames
 using Distributions
 using Random
 using BangBang
@@ -42,6 +43,7 @@ export
     GasCloud,
     ExponentialDisc,
     Bulge,
+    SphericalSystem,
 
     solarsystem,
 
@@ -64,6 +66,7 @@ include("distribution.jl")
 include("plummer.jl")
 include("disk.jl")
 include("bulge.jl")
+include("spherical.jl")
 include("gascloud.jl")
 include("solarsystem.jl")
 

src/distribution.jl

@@ -90,21 +90,49 @@ $(TYPEDSIGNATURES)
 
 Sample from the discrete pdf using rejection method
 
+!!! attentions
+    - Data must be unitless, and have no NaN
+    - `x` should start from 0
+    - the maximum of `pdf_array` would be the width of sampling box
+"""
+function rejection_sampling(x, pdf_array, rMax, NumSamples)
+    spl = Spline1D(ustrip.(x), ustrip.(pdf_array))
+    pdf_maximum = ustrip(maximum(pdf_array))
+
+    R = eltype(rMax)[]
+    sizehint!(R, NumSamples)
+
+    while length(R) < NumSamples
+        R_rand = rand() * rMax
+        if rand() < spl(ustrip(R_rand)) / pdf_maximum
+            push!(R, R_rand)
+        end
+    end
+    
+    return R
+end
+
+"""
+$(TYPEDSIGNATURES)
+
+Sample from the discrete pdf using rejection method
+
 !!! attentions
     - Data must be unitless, and have no NaN
     - `x` should start from 0
     - the maximum of `pdf` would be the width of sampling box
 """
-function rejection_sampling(x, pdf, rMax, NumSamples)
-    spl = Spline1D(x, pdf)
-    pdf_maximum = maximum(pdf)
+function rejection_sampling(pdf::Function, rMax, NumSamples)
+    x = collect(LinRange(rMax/10000, rMax, 10000))
+    pdf_array = pdf.(x)
+    pdf_maximum = maximum(pdf_array)
 
-    R = eltype(pdf)[]
+    R = eltype(rMax)[]
     sizehint!(R, NumSamples)
 
     while length(R) < NumSamples
         R_rand = rand() * rMax
-        if rand() < spl(R_rand) / pdf_maximum
+        if rand() < pdf(R_rand) / pdf_maximum
             push!(R, R_rand)
         end
     end

src/spherical.jl

@@ -0,0 +1,83 @@
+"""
+$(TYPEDEF)
+$(TYPEDFIELDS)
+
+`mass_shell` support two types:
+- `Dict` or `DataFrames`: discrete radii `:r` and shell masses `:m`
+- `function`: a radial function of shell mass
+"""
+struct SphericalSystem{I} <: InitialConditionConfig
+    collection::Collection
+    NumSamples::I
+
+    mass_shell
+end
+
+function Base.show(io::IO, config::SphericalSystem)
+    print(io,
+        "Config of spherical system Initial Conditions:",
+        "\n    Particle Collection: ", config.collection,
+        "\n      Number of Samples: ", config.NumSamples,
+    )
+end
+
+"""
+$(TYPEDSIGNATURES)
+
+"""
+function generate(config::SphericalSystem, units = uAstro;
+    MaxRadius = 0.0u"kpc",
+)
+    uLen = getuLength(units)
+    uMass = getuMass(units)
+    uVel = getuVel(units)
+    NumSamples = config.NumSamples
+
+    if iszero(MaxRadius)
+        error("Keyword `MaxRadius` should not be zero")
+    end
+
+    if config.mass_shell isa Function
+        R = rejection_sampling(config.mass_shell, MaxRadius, NumSamples)
+    elseif config.mass_shell isa Dict
+        if haskey(config.mass_shell, "r") && haskey(config.mass_shell, "m")
+            R = rejection_sampling(config.mass_shell["r"], config.mass_shell["m"], MaxRadius, NumSamples)
+        else
+            error("`mass_shell` of `SphericalSystem` must be `Dict` or `DataFrame` (with fields `:r`, `:m`) or `Function`")
+        end
+    elseif config.mass_shell isa DataFrame
+        if hasproperty(config.mass_shell, :r) && hasproperty(config.mass_shell, :m)
+            R = rejection_sampling(config.mass_shell.r, config.mass_shell.m, MaxRadius, NumSamples)
+        else
+            error("`mass_shell` of `SphericalSystem` must be `Dict` or `DataFrame` (with fields `:r`, `:m`) or `Function`")
+        end
+    else
+        error("`mass_shell` of `SphericalSystem` must be `Dict` or `DataFrame` (with fields `:r`, `:m`) or `Function`")
+    end
+
+    # Compute total mass
+    x = collect(LinRange(MaxRadius/10000, MaxRadius, 10000))
+    dx = x[2] - x[1]
+    if config.mass_shell isa Function
+        shell_masses = config.mass_shell.(x)
+        TotalMass = PhysicalParticles.NumericalIntegration.integrate(x, shell_masses) * ustrip(uLen, dx)
+    elseif config.mass_shell isa Dict || config.mass_shell isa DataFrame
+        spl = Spline1D(ustrip.(uLen, config.mass_shell["r"]), ustrip.(uMass/uLen, config.mass_shell["m"]))
+        TotalMass = PhysicalParticles.NumericalIntegration.integrate(x, spl(ustrip.(uLen, x))*uMass/uLen) * ustrip(uLen, dx)
+    end
+
+    pos = rand_pos_3d.(R)
+
+    #TODO: vel
+    vel = [PVector(uVel) for i in 1:NumSamples]
+
+    # Packing
+    particles = StructArray(Star(units, id = i) for i in 1:NumSamples)
+    assign_particles(particles, :Pos, pos)
+    assign_particles(particles, :Vel, vel)
+
+    Mmean = TotalMass / NumSamples
+    assign_particles(particles, :Mass, Mmean)
+
+    return particles
+end

test/runtests.jl

@@ -3,7 +3,8 @@ using Test
 using Unitful, UnitfulAstro
 
 using AstroSimBase
-using AstroIO, PhysicalParticles
+using PhysicalParticles
+using AstroIO
 
 using Dates
 using AstroIC
@@ -12,4 +13,5 @@ include("tools.jl")
 include("plummer.jl")
 include("disk.jl")
 include("bulge.jl")
+include("spherical.jl")
 include("solarsystem.jl")

test/spherical.jl

@@ -0,0 +1,23 @@
+# using AstroPlot
+
+@testset "SphericalSystem" begin
+    r_s = 2.0u"kpc"
+    MaxRadius = 10*r_s
+    NumSamples = 1000
+
+    r = collect(LinRange(0.001u"kpc", MaxRadius, 1000))
+    dr = r[2] - r[1]
+
+    shell_mass_func = x->1.0e8u"Msun/kpc^3"/(x/r_s)/(1+x/r_s)^2 * 4π * x^2
+    
+    config = SphericalSystem(STAR, NumSamples, shell_mass_func)
+    particles = generate(config; MaxRadius)
+    # plot_makie(particles)
+
+
+    shell_mass = shell_mass_func.(r)
+
+    config = SphericalSystem(STAR, NumSamples, Dict("r"=>r, "m"=>shell_mass))
+    particles = generate(config; MaxRadius)
+    # plot_makie(particles)
+end