Woodbury Kernel

CHANGELOG.md

@@ -16,6 +16,7 @@
 - Save only the basename of input files in the summary file in in `pyvela` script
 ## Fixed
 - Correctly avoid likelihood computation when the parameter is outside prior range.
+- 2π factor in likelihood expression
 ## Removed
 - `compare_residuals.py` example script
 

src/Vela.jl

@@ -18,7 +18,8 @@ include("toa/toa.jl")
 include("toa/wideband_toa.jl")
 include("parameter/parameter.jl")
 include("model/component.jl")
-include("model/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")
@@ -47,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/model/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

@@ -0,0 +1,83 @@
+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
+
+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

src/likelihood/ecorr_likelihood.jl

@@ -58,7 +58,7 @@ function calc_lnlike(
     spawn_chunk(chunk) = @spawn _ecorr_lnlike_chunk(model, toas, params, tzrphase, chunk)
     tasks = map(spawn_chunk, chunks)
     result = sum(fetch, tasks)
-    return -value(result) / 2
+    return -value(result) / 2 - (data_length(toas) * ln_2π / 2)
 end
 
 function calc_lnlike_serial(
@@ -72,5 +72,5 @@ function calc_lnlike_serial(
             group -> _ecorr_lnlike_group(model, toas, params, tzrphase, group),
             model.kernel.ecorr_groups,
         ),
-    ) / 2
+    ) / 2 - (data_length(toas) * ln_2π / 2)
 end

src/likelihood/wls_likelihood.jl

@@ -1,6 +1,8 @@
 export calc_lnlike,
     calc_lnlike_serial, get_lnlike_serial_func, get_lnlike_parallel_func, get_lnlike_func
 
+const ln_2π = log(2π)
+
 """A single term in the pulsar timing log-likelihood expression (white noise-only).
 
 A factor of 1/2 is excluded here."""
@@ -21,6 +23,9 @@ _wls_lnlike_chunk(
     chunk,
 ) where {T<:TOABase} = sum(ii -> _wls_lnlike_term(model, toas[ii], params, tzrphase), chunk)
 
+data_length(toas::Vector{TOA}) = length(toas)
+data_length(toas::Vector{WidebandTOA}) = 2 * length(toas)
+
 function calc_lnlike(
     model::TimingModel{ComponentsTuple,WhiteNoiseKernel,PriorsTuple},
     toas::Vector{T},
@@ -31,7 +36,7 @@ function calc_lnlike(
     spawn_chunk(chunk) = @spawn _wls_lnlike_chunk(model, toas, params, tzrphase, chunk)
     tasks = map(spawn_chunk, chunks)
     result = sum(fetch, tasks)
-    return -result / 2
+    return -result / 2 - (data_length(toas) * ln_2π / 2)
 end
 
 """
@@ -52,7 +57,8 @@ function calc_lnlike_serial(
     params::NamedTuple,
 ) where {ComponentsTuple<:Tuple,PriorsTuple<:Tuple,T<:TOABase}
     tzrphase = calc_tzr_phase(model, params)
-    return -sum(toa -> _wls_lnlike_term(model, toa, params, tzrphase), toas) / 2
+    return -sum(toa -> _wls_lnlike_term(model, toa, params, tzrphase), toas) / 2 -
+           (data_length(toas) * ln_2π / 2)
 end
 
 """