woodbury

src/Vela.jl

@@ -18,8 +18,8 @@ include("toa/toa.jl")
 include("toa/wideband_toa.jl")
 include("parameter/parameter.jl")
 include("model/component.jl")
-include("kernel/kernel.jl")
-include("kernel/woodbury_kernel.jl")
+include("kernel/white_noise_kernel.jl")
+include("kernel/ecorr_kernel.jl")
 include("model/spindown.jl")
 include("model/phase_offset.jl")
 include("model/glitch.jl")
@@ -48,6 +48,7 @@ include("model/gp_noise.jl")
 include("model/dispersion_measurement_noise.jl")
 include("model/timing_model.jl")
 include("model/wideband_model.jl")
+include("kernel/woodbury_kernel.jl")
 include("prior/prior.jl")
 include("prior/prior_scaling.jl")
 include("prior/simple_prior.jl")

src/kernel/kernel.jl → src/kernel/ecorr_kernel.jl

@@ -1,19 +1,4 @@
-export Kernel, WhiteNoiseKernel, EcorrKernel, EcorrGroup
-
-"""
-    Kernel
-    
-Abstract base class of all likelihood kernels"""
-abstract type Kernel end
-
-"""A kernel representing only uncorrelated noise.
-The covariance matrix is diagonal.
-
-Reference:
-    [Hobbs+ 2006](http://doi.org/10.1111/j.1365-2966.2006.10302.x),
-    [Alam+ 2021](http://doi.org/10.3847/1538-4365/abc6a1)
-"""
-struct WhiteNoiseKernel <: Kernel end
+export EcorrKernel, EcorrGroup
 
 """Range of TOAs belonging to an ECORR block."""
 struct EcorrGroup
@@ -34,7 +19,6 @@ struct EcorrKernel <: Kernel
     ecorr_groups::Vector{EcorrGroup}
 end
 
-show(io::IO, ::MIME"text/plain", model::Kernel) = show(io, model)
 show(io::IO, ek::EcorrKernel) = print(
     io,
     "EcorrKernel($(length(unique(grp.index for grp in ek.ecorr_groups)) - 1) ECORRs, $(length(ek.ecorr_groups)) groups)",

src/kernel/white_noise_kernel.jl

@@ -0,0 +1,18 @@
+export Kernel, WhiteNoiseKernel
+
+"""
+    Kernel
+    
+Abstract base class of all likelihood kernels"""
+abstract type Kernel end
+
+show(io::IO, ::MIME"text/plain", kernel::Kernel) = show(io, kernel)
+
+"""A kernel representing only uncorrelated noise.
+The covariance matrix is diagonal.
+
+Reference:
+    [Hobbs+ 2006](http://doi.org/10.1111/j.1365-2966.2006.10302.x),
+    [Alam+ 2021](http://doi.org/10.3847/1538-4365/abc6a1)
+"""
+struct WhiteNoiseKernel <: Kernel end

src/kernel/woodbury_kernel.jl

@@ -1,7 +1,83 @@
-struct WoodburyKernel{InnerKernel<:Kernel,SectionsTuple<:Tuple} <: Kernel
+abstract type KernelSection end
+
+struct MarginalizedPowerlawRedNoiseGP{N} <: KernelSection
+    basis::Matrix{Float64}
+    ln_js::NTuple{N,Float32}
+end
+
+function weight(rn::MarginalizedPowerlawRedNoiseGP, params)
+    log10_A = params.TNREDAMP
+    γ = params.TNREDGAM
+    A = exp10(log10_A)
+    f1 = params.PLREDFREQ
+    w1 = powerlaw(A, γ, rn.f1, rn.f1)
+    return map(ln_j -> w1 * exp(-γ * ln_j), rn.ln_js)
+end
+
+struct WoodburyKernel{InnerKernel<:Kernel, SectionsTuple<:Tuple} <: Kernel
     inner_kernel::InnerKernel
     sections::SectionsTuple
     basis::Matrix{Float64}
+
+    function WoodburyKernel(inner_kernel::Kernel, sections::Tuple)
+        @assert all(map(s -> isa(s, KernelSection), sections))
+        basis = combine_bases(sections)
+        return new(inner_kernel, sections, basis)
+    end
 end
 
-abstract type KernelSection end
+function combine_bases(sections)
+    return hcat(map(section -> section.basis, sections)...)
+end
+
+function combine_weights(sections, params)
+
+end
+
+function apply_inner_kernel(
+    ::WhiteNoiseKernel,
+    model::TimingModel,
+    toas::Vector{TOA},
+    params::NamedTuple,
+    tzrphase::GQ,
+    basis::Matrix{Float64},
+    weights::Matrix{GQ{2,Float64}},
+)
+    ntoa = length(toas)
+    nbasis = size(basis)[2]
+    @assert size(basis)[1] == ntoa
+
+    logdet_N = 0.0
+    r_Ninv_r = 0.0
+    UT_Ninv_r = zeros(GQ{-1,Float64}, nbasis)
+    Σ = zeros(GQ{-2,Float64}, nbasis, nbasis)
+
+    for (ii, toa) in enumerate(toas)
+        ctoa = correct_toa(model, toa, params)
+
+        ς2 = scaled_toa_error_sqr(toa, ctoa)
+        logdet_N += log(value(ς2))
+
+        dϕ = GQ{Float64}(phase_residual(toa, ctoa) - tzrphase)
+        r = dϕ / doppler_shifted_spin_frequency(ctoa)
+        r_Ninv_r += value(r * r / ς2)
+
+        for jj = 1:nbasis
+            UT_Ninv_r[jj] += basis[ii, jj] * r / ς2
+
+            for kk = 1:jj
+                Σ[jj, kk] += basis[ii, jj] * basis[ii, kk] / ς2
+            end
+        end
+    end
+
+    for jj = 1:nbasis
+        Σ[jj, jj] += 1 / weights[jj]
+
+        for kk = (jj+1):nbasis
+            Σ[jj, kk] += Σ[kk, jj]
+        end
+    end
+
+    return logdet_N, r_Ninv_r, UT_Ninv_r, Σ
+end