Merge pull request #200 from astro-group-bristol/fergus/lib-gradus-mo…

…dels Gradus Spectral Models
astro-group-bristol · Jun 23, 2024 · 090e845 · 090e845
2 parents 1d420d5 + e76b6e0
commit 090e845
Show file tree

Hide file tree

Showing 6 changed files with 93 additions and 4 deletions.
diff --git a/.github/workflows/auto-register.yml b/.github/workflows/auto-register.yml
@@ -6,6 +6,7 @@ on:
       - main
     paths:
       - 'Project.toml'
+      - 'lib/GradusSpectralModels/Project.toml'
 
 permissions:
   contents: read
@@ -45,6 +46,7 @@ jobs:
           using LocalRegistry
           Pkg.Registry.add(Pkg.RegistrySpec(url="https://github.com/astro-group-bristol/AstroRegistry/"))
           register(pwd(), registry="AstroRegistry", push=false)
+          register(joinpath(pwd(), "lib/GradusSpectralModels"), registry="AstroRegistry", push=false)
           '
       - name: Create new branch
         run: |

diff --git a/lib/GradusSpectralModels/Project.toml b/lib/GradusSpectralModels/Project.toml
@@ -0,0 +1,8 @@
+name = "GradusSpectralModels"
+uuid = "0bd1f26a-7f13-41ad-b364-0116c59c8ea1"
+authors = ["fjebaker <fergusbkr@gmail.com>"]
+version = "0.1.0"
+
+[deps]
+Gradus = "c5b7b928-ea15-451c-ad6f-a331a0f3e4b7"
+SpectralFitting = "f2c56810-742e-4b72-8bf4-27af3bb81a12"
diff --git a/lib/GradusSpectralModels/src/GradusSpectralModels.jl b/lib/GradusSpectralModels/src/GradusSpectralModels.jl
@@ -0,0 +1,54 @@
+module GradusSpectralModels
+
+using Gradus, SpectralFitting
+
+struct LineProfile{D,T} <: AbstractTableModel{T,Additive}
+    table::D
+    K::T
+    "Spin"
+    a::T
+    "Observer inclination (degrees off of the spin axis)."
+    θ::T
+    "Inner radius of the accretion disc."
+    rin::T
+    "Outer radius of the accretion disc."
+    rout::T
+    "Central emission line energy (keV)."
+    E₀::T
+end
+
+function LineProfile(
+    profile,
+    table::Gradus.CunninghamTransferTable;
+    K = FitParam(1.0),
+    a = FitParam(0.998),
+    θ = FitParam(45.0),
+    rin = FitParam(1.0),
+    rout = FitParam(100.0, upper_limit = 100.0),
+    E₀ = FitParam(1.0),
+)
+    setup = integration_setup(profile, table((get_value(θ), get_value(a))))
+    LineProfile((; setup = setup, table = table), K, a, θ, rin, rout, E₀)
+end
+
+function SpectralFitting.invoke!(output, domain, model::LineProfile)
+    grid = model.table.table((model.θ, model.a))
+    rmin = if model.rmin < grid.r_grid[1]
+        grid.r_grid[1]
+    else
+        model.rmin
+    end
+    Gradus.integrate_lineprofile!(
+        output,
+        model.table.setup,
+        grid,
+        domain;
+        rmin = rmin,
+        rmax = model.rout,
+    )
+    output
+end
+
+export LineProfile
+
+end # module GradusSpectralModels
diff --git a/src/transfer-functions/cunningham-transfer-functions.jl b/src/transfer-functions/cunningham-transfer-functions.jl
@@ -492,9 +492,34 @@ function transfer_function_grid(itfs::InterpolatingTransferBranches, Ng::Int)
     )
 end
 
+function make_transfer_function_table(
+    M::Type{<:KerrMetric},
+    d::AbstractAccretionDisc,
+    a_range::AbstractVector,
+    θ_range::AbstractVector;
+    verbose = true,
+    r_max = 500.0,
+    n_radii = 150,
+    kwargs...,
+)
+    function _mapper(a, θ)
+        if verbose
+            @info "a, θ = $((a, θ))"
+        end
+        m = M(1.0, a)
+        x = SVector(0.0, 10000.0, deg2rad(θ), 0.0)
+
+        radii = Gradus.Grids._inverse_grid(Gradus.isco(m) + 1e-2, r_max, n_radii)
+        Gradus.transfer_function_grid(m, x, d, radii; verbose = verbose, kwargs...)
+    end
+    grids = [_mapper(a, θ) for a in a_range, θ in θ_range]
+    Gradus.CunninghamTransferTable((θ_range, a_range), grids)
+end
+
 export CunninghamTransferData,
     TransferBranches,
     InterpolatingTransferBranches,
     splitbranches,
     interpolate_branches,
-    cunningham_transfer_function
+    cunningham_transfer_function,
+    make_transfer_function_table
diff --git a/src/transfer-functions/integration.jl b/src/transfer-functions/integration.jl
@@ -35,7 +35,7 @@ function IntegrationSetup(
     time = nothing,
     h = 1e-8,
     g_grid_upscale = 1,
-    n_radii = 500,
+    n_radii = 1000,
     quadrature_points = 23,
 )
     IntegrationSetup(

diff --git a/src/transfer-functions/types.jl b/src/transfer-functions/types.jl
@@ -64,13 +64,13 @@ end
 
 struct CunninghamTransferTable{N,T,CacheT}
     params::NTuple{N,Vector{T}}
-    grids::Vector{CunninghamTransferGrid{T}}
+    grids::Array{CunninghamTransferGrid{T},N}
     cache::CacheT
 end
 
 function CunninghamTransferTable(
     x::NTuple{N},
-    grids::AbstractVector{<:CunninghamTransferGrid},
+    grids::AbstractArray{<:CunninghamTransferGrid},
 ) where {N}
     _grids = reshape(grids, length.(x))
     cache = InterpolationCache{N}(_grids)