Merge pull request #44 from hersle/minor_changes

Minor performance improvement and consistency improvements

Author: hersle

Project.toml

@@ -19,7 +19,10 @@ ModelingToolkit = "961ee093-0014-501f-94e3-6117800e7a78"
 NonlinearSolve = "8913a72c-1f9b-4ce2-8d82-65094dcecaec"
 NumericalIntegration = "e7bfaba1-d571-5449-8927-abc22e82249b"
 OhMyThreads = "67456a42-1dca-4109-a031-0a68de7e3ad5"
-OrdinaryDiffEq = "1dea7af3-3e70-54e6-95c3-0bf5283fa5ed"
+OrdinaryDiffEqNonlinearSolve = "127b3ac7-2247-4354-8eb6-78cf4e7c58e8"
+OrdinaryDiffEqRosenbrock = "43230ef6-c299-4910-a778-202eb28ce4ce"
+OrdinaryDiffEqSDIRK = "2d112036-d095-4a1e-ab9a-08536f3ecdbf"
+OrdinaryDiffEqTsit5 = "b1df2697-797e-41e3-8120-5422d3b24e4a"
 PeriodicTable = "7b2266bf-644c-5ea3-82d8-af4bbd25a884"
 PhysicalConstants = "5ad8b20f-a522-5ce9-bfc9-ddf1d5bda6ab"
 PreallocationTools = "d236fae5-4411-538c-8e31-a6e3d9e00b46"
@@ -53,7 +56,10 @@ ModelingToolkit = "~10.25.0"
 NonlinearSolve = "4.3"
 NumericalIntegration = "0.3.4"
 OhMyThreads = "0.8.0"
-OrdinaryDiffEq = "6.90.1"
+OrdinaryDiffEqNonlinearSolve = "1.12.0"
+OrdinaryDiffEqRosenbrock = "1.14.0"
+OrdinaryDiffEqSDIRK = "1.6.0"
+OrdinaryDiffEqTsit5 = "1.5.0"
 PeriodicTable = "1.2.1"
 PhysicalConstants = "0.2.4"
 PreallocationTools = "0.4.25"

docs/Project.toml

@@ -18,7 +18,10 @@ MKL = "33e6dc65-8f57-5167-99aa-e5a354878fb2"
 ModelingToolkit = "961ee093-0014-501f-94e3-6117800e7a78"
 Optim = "429524aa-4258-5aef-a3af-852621145aeb"
 Optimization = "7f7a1694-90dd-40f0-9382-eb1efda571ba"
-OrdinaryDiffEq = "1dea7af3-3e70-54e6-95c3-0bf5283fa5ed"
+OrdinaryDiffEqBDF = "6ad6398a-0878-4a85-9266-38940aa047c8"
+OrdinaryDiffEqRosenbrock = "43230ef6-c299-4910-a778-202eb28ce4ce"
+OrdinaryDiffEqSDIRK = "2d112036-d095-4a1e-ab9a-08536f3ecdbf"
+OrdinaryDiffEqTsit5 = "b1df2697-797e-41e3-8120-5422d3b24e4a"
 PDMats = "90014a1f-27ba-587c-ab20-58faa44d9150"
 PairPlots = "43a3c2be-4208-490b-832a-a21dcd55d7da"
 Plots = "91a5bcdd-55d7-5caf-9e0b-520d859cae80"

docs/src/benchmarks.md

@@ -4,7 +4,7 @@
     This page is a work in progress.
 
 ```@example bench
-using MKL, SymBoltz, OrdinaryDiffEq, BenchmarkTools, Plots, BenchmarkPlots, StatsPlots, InteractiveUtils
+using MKL, SymBoltz, OrdinaryDiffEqRosenbrock, OrdinaryDiffEqSDIRK, OrdinaryDiffEqBDF, BenchmarkTools, Plots, BenchmarkPlots, StatsPlots, InteractiveUtils
 M = SymBoltz.ΛCDM(ν = nothing, K = nothing, h = nothing)
 pars = SymBoltz.parameters_Planck18(M)
 prob = CosmologyProblem(M, pars)

docs/src/parameter_fitting.md

@@ -52,7 +52,7 @@ fig
 
 To predict [luminosity distances](@ref "Luminosity distance") theoretically, we solve the w0waCDM model:
 ```@example fit
-using SymBoltz
+using SymBoltz, OrdinaryDiffEqTsit5
 g = SymBoltz.metric()
 K = SymBoltz.curvature(g)
 X = SymBoltz.w0wa(g; analytical = true)
@@ -68,7 +68,7 @@ dL = SymBoltz.distance_luminosity_function(M, pars_fixed, pars_varying, data.zcm
 
 # Show example predictions
 Mb = -19.3 # absolute supernova brightness (constant since SN-Ia are standard candles)
-bgopts = (alg = SymBoltz.Tsit5(), reltol = 1e-5, maxiters = 1e3)
+bgopts = (alg = Tsit5(), reltol = 1e-5, maxiters = 1e3)
 p0 = [9.3e-5, 0.3, 0.0, 0.7, -1.0, 0.0] # fiducial parameters
 μs = μ(p0)
 mbs = μs .+ Mb
@@ -92,12 +92,12 @@ using Turing
     wa ~ Uniform(-1.0, +1.0)
 
     p = [Ωr0, Ωm0, Ωk0, h, w0, wa]
-    μs_pred = μ_pred(p)
-    if isempty(μs_pred)
+    mbs_pred = μ_pred(p)
+    if isempty(mbs_pred)
         Turing.@addlogprob! -Inf
         return nothing
     end
-    mbs_pred = μs_pred .+ Mb
+    mbs_pred .+= Mb
     return mbs ~ MvNormal(mbs_pred, C) # read "measurements sampled from multivariate normal with predictions and covariance matrix"
 
     # equivalently:
@@ -238,7 +238,7 @@ pp_fc
 
 ```@setup fit
 #=
-using SymBoltz, OrdinaryDiffEq, Turing
+using SymBoltz, OrdinaryDiffEqTsit5, Turing
 
 function dL_fast(z, Ωm0, Ωk0, h; Ωr0 = 9.3e-5, aini = 1e-8, reltol = 1e-8, alg = Tsit5(), maxiters = 1e3)
     ΩΛ0 = 1 - Ωr0 - Ωm0 - Ωk0

notebooks/paper/paper.ipynb

@@ -903,12 +903,12 @@
     "    w0 ~ Uniform(-2.0, 0.0)\n",
     "\n",
     "    p = [Ωm0, h, w0]\n",
-    "    μs_pred = μ_pred(p)\n",
-    "    if isempty(μs_pred)\n",
+    "    mbs_pred = μ_pred(p)\n",
+    "    if isempty(mbs_pred)\n",
     "        Turing.@addlogprob! -Inf\n",
     "        return nothing\n",
     "    end\n",
-    "    mbs_pred = μs_pred .+ Mb\n",
+    "    mbs_pred .+= Mb\n",
     "    return mbs ~ MvNormal(mbs_pred, C) # read \"measurements sampled from multivariate normal with predictions and covariance matrix\"\n",
     "\n",
     "    # equivalently:\n",

src/SymBoltz.jl

@@ -3,7 +3,8 @@ module SymBoltz
 using Reexport
 import ModelingToolkit: parameters, unknowns # explicitly imported because they are extended
 @reexport using ModelingToolkit # re-export commonly used ModelingToolkit functions
-using OrdinaryDiffEq
+using OrdinaryDiffEqRosenbrock, OrdinaryDiffEqSDIRK, OrdinaryDiffEqTsit5
+using OrdinaryDiffEqNonlinearSolve: NLNewton
 using NonlinearSolve
 using OhMyThreads
 using Base.Threads

src/constants.jl

@@ -44,12 +44,7 @@ const EHet2p2s = EHet2p1s - EHet2s1s
 
 δkron(i, j) = (i == j ? 1 : 0) # Kronecker delta
 
-function k_dimensionless(k::Number, h)
-    if unit(k) == NoUnits
-        return k
-    else
-        H₀ = h * H100 # s⁻¹
-        k0 = (H₀ / c) / u"m"
-        return NoUnits(k / k0)
-    end
-end
+k_dimensionless(k::Number, h) = k
+k_dimensionless(k::Quantity, h) = NoUnits(k / (h*H100 / c / u"m"))
+k_dimensionless(k::Number, bgsol::ODESolution) = k
+k_dimensionless(k::Quantity, bgsol::ODESolution) = k_dimensionless(k, getsym(bgsol, :h)(bgsol))

src/models/inflation.jl

@@ -11,7 +11,7 @@ function harrison_zeldovich(g; kwargs...)
         P ~ 2*Num(π)^2 / k^3 * As * (k/kpivot)^(ns-1)
     ]
     defaults = [
-        kpivot => 0.05 / (H100*g.h/c) / Mpc # k = 0.05/Mpc ≠ 0.05/(Mpc/h)
+        kpivot => 0.05 / Mpc / (H100/c) / g.h # k = 0.05/Mpc ≠ 0.05/(Mpc/h)
         ln_As1e10 => NaN # keep uninitialized if not needed
         ns => NaN
     ]

src/observables/distances.jl

@@ -15,17 +15,17 @@ function distance_luminosity(sol::CosmologySolution, ivs = sol.bg.t, τ0 = sol[s
     r = sinc.(√(-Ωk0+0im)*χ/π) .* χ |> real # Julia's sinc(x) = sin(π*x) / (π*x)
     H0 = H100 * sol[M.g.h]
     a = sol(M.g.a, ivs)
-    return @. r / a * SymBoltz.c / H0 # to meters
+    return @. r / a * c / H0 # to meters
 end
 
 # TODO: test @inferred
-function distance_luminosity_function(M::System, pars_fixed, pars_varying, zs; bgopts = (alg = SymBoltz.Tsit5(), reltol = 1e-5, maxiters = 1e3))
+function distance_luminosity_function(M::System, pars_fixed, pars_varying, zs; bgopts = (alg = Tsit5(), reltol = 1e-5, maxiters = 1e3))
     isequal(ModelingToolkit.get_iv(M), M.g.a) || error("Independent variable must be $(M.g.a)")
 
     pars = merge(pars_fixed, Dict(pars_varying .=> NaN))
     as = @. 1 / (zs + 1)
     prob = CosmologyProblem(M, pars; pt = false, ivspan = (minimum(as), 1.0))
-    probgen = SymBoltz.parameter_updater(prob, pars_varying; build_initializeprob = Val{false})
+    probgen = parameter_updater(prob, pars_varying; build_initializeprob = Val{false})
 
     geta = getsym(prob, M.g.a)
     getτ = getsym(prob, M.τ)
@@ -42,8 +42,8 @@ function distance_luminosity_function(M::System, pars_fixed, pars_varying, zs; b
         τ0 = τ[end] # time today
         χ = τ0 .- τ
         r = @. real(sinc(√(-Ωk0+0im)*χ/π) * χ) # Julia's sinc(x) = sin(π*x) / (π*x)
-        H0 = SymBoltz.H100 * h
-        return @. r / a * SymBoltz.c / H0 # luminosity distance in meters
+        H0 = H100 * h
+        return @. r / a * c / H0 # luminosity distance in meters
     end
 end
 

src/observables/fourier.jl

@@ -3,20 +3,21 @@ using MatterPower
 using TwoFAST
 
 """
-    spectrum_primordial(k, h, As, ns=1.0; kp = k_dimensionless(0.05 / u"Mpc", h))
+    spectrum_primordial(k, h, As, ns=1.0; kp = 0.05/u"Mpc")
 
 Compute the primordial power spectrum
 ```math
 P₀(k) = 2π² Aₛ (k/kₚ)^{nₛ-1} / k³
 ```
 with spectral amplitude `As`, spectral index `ns` and pivot scale wavenumber `kp` at the wavenumber(s) `k`.
 """
-function spectrum_primordial(k, h, As, ns=1.0; kp = k_dimensionless(0.05 / u"Mpc", h))
-    P = @. 2*π^2 / k^3 * As
+function spectrum_primordial(k, h, As, ns=1.0; kp = 0.05/u"Mpc")
+    P = 2*π^2 * As ./ k.^3
 
-    # compute k / kpivot with equal wavenumber units
+    # ensure both k and kp to dimensionless wavenumbers k/(H₀/c) before taking ratio
     k = k_dimensionless.(k, h)
-    @. P *= (k/kp)^(ns-1)
+    kp = k_dimensionless(kp, h)
+    P .*= (k./kp).^(ns-1)
 
     return P
 end
@@ -121,7 +122,7 @@ function variance_matter(sol::CosmologySolution, R)
     P = spectrum_matter(sol, k)
     lgPspl = spline(log.(P), log.(k))
     Pf(k) = exp(lgPspl(log(k)))
-    R = 1 / k_dimensionless(1 / R, sol[M.g.h]) # make dimensionless
+    R = 1 / k_dimensionless(1 / R, sol.bg) # make dimensionless
     return MatterPower.sigma2(Pf, R)
 end
 """