diff --git a/CHANGELOG.md b/CHANGELOG.md index 0839ba15..2b317fe3 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,6 +1,7 @@ # Unreleased ## Added - `time_residuals`, `scaled_toa_uncertainties`, and `model_dm`, `from_pint` methods in `SPNTA` +- `Pulsar` class ## Changed ## Fixed ## Removed diff --git a/docs/src/pyvela.md b/docs/src/pyvela.md index 2e3a43d4..a894dd44 100644 --- a/docs/src/pyvela.md +++ b/docs/src/pyvela.md @@ -37,7 +37,14 @@ Here, `spnta.model` is a `Vela.TimingModel` object and `spnta.toas` is a `Vector function reads the `par` and `tim` files using the [`pint.models.get_model_and_toas`](https://nanograv-pint.readthedocs.io/en/latest/_autosummary/pint.models.model_builder.get_model_and_toas.html) function under the hood and converts the resulting `pint.models.TimingModel` and `pint.toa.TOAs` objects. Note that this conversion does -not conserve all the information, and the reverse is not possible. +not conserve all the information, and the reverse is not possible. + +The `SPNTA` class internally uses the `Pulsar` type to store the timing model and the TOAs together. +Currently it is assumed that all the TOAs are of the same paradigm (narrowband or wideband). Inhomogeneous +datasets are not supported. +```@docs +Pulsar +``` `cheat_prior_scale` defines the scale factor by which the frequentist uncertainties are multiplied to obtain the "cheat" prior widths. The `custom_priors` argument contains the user-defined prior distributions, diff --git a/pyvela/examples/run_example_emcee.py b/pyvela/examples/run_example_emcee.py index cc797ad7..4c095fbd 100755 --- a/pyvela/examples/run_example_emcee.py +++ b/pyvela/examples/run_example_emcee.py @@ -48,7 +48,7 @@ nwalkers, spnta.ndim, spnta.lnpost_vectorized, - moves=[emcee.moves.StretchMove(), emcee.moves.DESnookerMove()], + # moves=[emcee.moves.StretchMove(), emcee.moves.DESnookerMove()], vectorize=True, ) sampler.run_mcmc(p0, 6000, progress=True, progress_kwargs={"mininterval": 1}) diff --git a/pyvela/pyvela/priors.py b/pyvela/pyvela/priors.py index f714e7d5..fce1d87f 100644 --- a/pyvela/pyvela/priors.py +++ b/pyvela/pyvela/priors.py @@ -25,15 +25,15 @@ "DMEFAC": jl.LogNormal(0.0, 0.25), "PLREDCOS_": jl.Normal(), "PLREDSIN_": jl.Normal(), - "TNREDAMP": jl.Uniform(-16.0, -9.0), + "TNREDAMP": jl.Uniform(-18.0, -9.0), "TNREDGAM": jl.Uniform(0.0, 7.0), "PLDMCOS_": jl.Normal(), "PLDMSIN_": jl.Normal(), - "TNDMAMP": jl.Uniform(-16.0, -9.0), + "TNDMAMP": jl.Uniform(-18.0, -9.0), "TNDMGAM": jl.Uniform(0.0, 7.0), "PLCHROMCOS_": jl.Normal(), "PLCHROMSIN_": jl.Normal(), - "TNCHROMAMP": jl.Uniform(-16.0, -9.0), + "TNCHROMAMP": jl.Uniform(-18.0, -9.0), "TNCHROMGAM": jl.Uniform(0.0, 7.0), } diff --git a/pyvela/pyvela/spnta.py b/pyvela/pyvela/spnta.py index 9bd70e00..6dd74045 100644 --- a/pyvela/pyvela/spnta.py +++ b/pyvela/pyvela/spnta.py @@ -126,23 +126,7 @@ def __init__( custom_priors=custom_priors, ) - self._setup(model, toas) - - def _setup(self, model, toas): - self.model = model - self.toas = toas - - self.lnlike = vl.get_lnlike_func(self.model, self.toas) - self.lnprior = vl.get_lnprior_func(self.model) - self.prior_transform = vl.get_prior_transform_func(self.model) - self.lnpost = vl.get_lnpost_func(self.model, self.toas, False) - self.lnpost_vectorized = vl.get_lnpost_func(self.model, self.toas, True) - - self.param_names = list(vl.get_free_param_names(self.model)) - self.param_labels = list(vl.get_free_param_labels(self.model)) - self.scale_factors = np.array(vl.get_scale_factors(self.model)) - self.maxlike_params = np.array(vl.read_param_values_to_vector(self.model)) - self.ndim = len(self.param_names) + self.pulsar = vl.Pulsar(model, toas) self._check() @@ -156,46 +140,99 @@ def _check(self): assert all(np.isfinite([lnpr, lnl, lnp])) assert np.isclose(lnp, lnpv) + def lnlike(self, params): + return vl.calc_lnlike(self.pulsar, params) + + def lnprior(self, params): + return vl.calc_lnprior(self.pulsar, params) + + def prior_transform(self, cube): + return vl.prior_transform(self.pulsar, cube) + + def lnpost(self, params): + return vl.calc_lnpost(self.pulsar, params) + + def lnpost_vectorized(self, paramss): + return vl.calc_lnpost_vectorized(self.pulsar, paramss) + + @property + def model(self): + return self.pulsar.model + + @property + def toas(self): + return self.pulsar.toas + + @property + def param_names(self): + return list(vl.get_free_param_names(self.pulsar.model)) + + @property + def param_labels(self): + return list(vl.get_free_param_labels(self.pulsar.model)) + + @property + def scale_factors(self): + return np.array(vl.get_scale_factors(self.pulsar.model)) + + @property + def maxlike_params(self): + return np.array(vl.read_param_values_to_vector(self.pulsar.model)) + + @property + def ndim(self): + return len(self.param_names) + def rescale_samples(self, samples: np.ndarray) -> np.ndarray: """Rescale the samples from Vela's internal units to common units""" return samples / self.scale_factors def save_jlso(self, filename: str) -> None: """Write the model and TOAs as a JLSO file""" - vl.save_pulsar_data(filename, self.model, self.toas) + vl.save_pulsar_data(filename, self.pulsar.model, self.pulsar.toas) def is_wideband(self) -> bool: """Whether the TOAs are wideband.""" - return jl.isa(self.toas[0], vl.WidebandTOA) + return jl.isa(self.pulsar.toas[0], vl.WidebandTOA) def get_mjds(self) -> np.ndarray: """Get the MJDs of each TOA.""" if self.is_wideband(): return np.array( - [jl.Float64(vl.value(wtoa.toa.value)) / day_to_s for wtoa in self.toas] + [ + jl.Float64(vl.value(wtoa.toa.value)) / day_to_s + for wtoa in self.pulsar.toas + ] ) else: return np.array( - [jl.Float64(vl.value(toa.value)) / day_to_s for toa in self.toas] + [jl.Float64(vl.value(toa.value)) / day_to_s for toa in self.pulsar.toas] ) def time_residuals(self, params: np.ndarray) -> np.ndarray: """Get the timing residuals (s) for a given set of parameters.""" - params = vl.read_params(self.model, params) + params = vl.read_params(self.pulsar.model, params) return np.array( - list(map(vl.value, vl.form_residuals(self.model, self.toas, params))) + list( + map( + vl.value, + vl.form_residuals(self.pulsar.model, self.pulsar.toas, params), + ) + ) if not self.is_wideband() else [ vl.value(wr[0]) - for wr in vl.form_residuals(self.model, self.toas, params) + for wr in vl.form_residuals(self.pulsar.model, self.pulsar.toas, params) ] ) def scaled_toa_unceritainties(self, params: np.ndarray) -> np.ndarray: """Get the scaled TOA uncertainties (s) for a given set of parameters.""" - params = vl.read_params(self.model, params) + params = vl.read_params(self.pulsar.model, params) - ctoas = [vl.correct_toa(self.model, tvi, params) for tvi in self.toas] + ctoas = [ + vl.correct_toa(self.pulsar.model, tvi, params) for tvi in self.pulsar.toas + ] return np.sqrt( [ @@ -204,26 +241,29 @@ def scaled_toa_unceritainties(self, params: np.ndarray) -> np.ndarray: if not self.is_wideband() else vl.scaled_toa_error_sqr(tvi.toa, ctoa.toa_correction) ) - for (tvi, ctoa) in zip(self.toas, ctoas) + for (tvi, ctoa) in zip(self.pulsar.toas, ctoas) ] ) def model_dm(self, params: np.ndarray) -> np.ndarray: """Compute the model DM (dmu) for a given set of parameters.""" - params = vl.read_params(self.model, params) + params = vl.read_params(self.pulsar.model, params) if not self.is_wideband(): - dms = np.zeros(len(self.toas)) - for ii, toa in enumerate(self.toas): + dms = np.zeros(len(self.pulsar.toas)) + for ii, toa in enumerate(self.pulsar.toas): ctoa = vl.TOACorrection() - for component in self.model.components: + for component in self.pulsar.model.components: if vl.isa(component, vl.DispersionComponent): dms[ii] += vl.value( vl.dispersion_slope(component, toa, ctoa, params) ) ctoa = vl.correct_toa(component, toa, ctoa, params) else: - ctoas = [vl.correct_toa(self.model, tvi, params) for tvi in self.toas] + ctoas = [ + vl.correct_toa(self.pulsar.model, tvi, params) + for tvi in self.pulsar.toas + ] dms = np.array([vl.value(ctoa.dm_correction.model_dm) for ctoa in ctoas]) dmu_conversion_factor = 2.41e-16 # Hz / (DMconst * dmu) @@ -236,7 +276,8 @@ def load_jlso(cls, filename: str): spnta = cls.__new__(cls) model, toas = vl.load_pulsar_data(filename) spnta.jlsofile = filename - spnta._setup(model, toas) + spnta.pulsar = vl.Pulsar(model, toas) + spnta._check() return spnta @classmethod @@ -269,7 +310,8 @@ def from_pint( custom_priors=custom_priors, ) - spnta._setup(model_v, toas_v) + spnta.pulsar = vl.Pulsar(model_v, toas_v) + spnta._check() return spnta diff --git a/pyvela/tests/test_data_files.py b/pyvela/tests/test_data_files.py index 4a1046ae..0b4f4b1a 100644 --- a/pyvela/tests/test_data_files.py +++ b/pyvela/tests/test_data_files.py @@ -96,6 +96,11 @@ def test_data(model_and_toas): assert all(np.isfinite(spnta.model_dm(spnta.maxlike_params))) + assert ( + all([len(label) > 0 for label in spnta.param_labels]) + and len(spnta.param_labels) == spnta.ndim + ) + def test_chi2(model_and_toas): spnta: SPNTA diff --git a/src/Vela.jl b/src/Vela.jl index 47843ca8..7fb4e11f 100644 --- a/src/Vela.jl +++ b/src/Vela.jl @@ -59,6 +59,7 @@ include("likelihood/wls_likelihood.jl") include("likelihood/ecorr_likelihood.jl") include("likelihood/wideband_wls_likelihood.jl") include("likelihood/posterior.jl") +include("pulsar/pulsar.jl") include("readwrite/readwrite.jl") end diff --git a/src/likelihood/posterior.jl b/src/likelihood/posterior.jl index 5d6abf42..0d88d7b4 100644 --- a/src/likelihood/posterior.jl +++ b/src/likelihood/posterior.jl @@ -1,6 +1,29 @@ -export get_lnpost_func +export calc_lnpost, calc_lnpost_serial, calc_lnpost_vectorized, get_lnpost_func +function calc_lnpost(model::TimingModel, toas::Vector{T}, params) where {T<:TOABase} + lnpr = calc_lnprior(model, params) + return isnan(lnpr) ? lnpr : lnpr + calc_lnlike(model, toas, params) +end + +function calc_lnpost_serial(model::TimingModel, toas::Vector{T}, params) where {T<:TOABase} + lnpr = calc_lnprior(model, params) + return isnan(lnpr) ? lnpr : lnpr + calc_lnlike_serial(model, toas, params) +end + +function calc_lnpost_vectorized( + model::TimingModel, + toas::Vector{T}, + paramss, +) where {T<:TOABase} + nparamss = size(paramss)[1] + result = Vector{Float64}(undef, nparamss) + @threads :static for ii = 1:nparamss + result[ii] = calc_lnpost_serial(model, toas, paramss[ii, :]) + end + return result +end + """ get_lnpost_func(::TimingModel, toas::Vector{T}, vectorize::Bool = false) where {T<:TOABase} @@ -13,29 +36,11 @@ function get_lnpost_func( toas::Vector{T}, vectorize::Bool = false, ) where {T<:TOABase} - lnprior_func = get_lnprior_func(model) - lnlike_func = - vectorize ? get_lnlike_serial_func(model, toas) : get_lnlike_func(model, toas) - nfree = model.param_handler._nfree - - function lnpost(params) - lnpr = lnprior_func(params) - return isfinite(lnpr) ? lnpr + lnlike_func(params) : -Inf - end - - if vectorize - function _lnpost_vector(paramss) - @assert (length(size(paramss)) == 2 && size(paramss)[2] == nfree) "Please pass a 2D array to this function for vectorized evaluation, where each row has `Nfree` elements." - nparamss = size(paramss)[1] - result = Vector{Float64}(undef, nparamss) - @threads :static for ii = 1:nparamss - result[ii] = lnpost(paramss[ii, :]) - end - return result - end - - return _lnpost_vector + return if vectorize + paramss -> calc_lnpost_vectorized(model, toas, paramss) + elseif nthreads() == 1 + params -> calc_lnpost_serial(model, toas, params) else - return lnpost + params -> calc_lnpost(model, toas, params) end end diff --git a/src/prior/prior.jl b/src/prior/prior.jl index 572e42bb..cea4a656 100644 --- a/src/prior/prior.jl +++ b/src/prior/prior.jl @@ -1,4 +1,10 @@ -export Prior, lnprior, prior_transform, get_lnprior_func, get_prior_transform_func, distr +export Prior, + lnprior, + calc_lnprior, + prior_transform, + get_lnprior_func, + get_prior_transform_func, + distr """Abstract base class of all priors.""" abstract type Prior end @@ -8,15 +14,16 @@ lnprior(prior::Prior, params::NamedTuple) = logpdf(distr(prior, params), param_value(prior, params)) lnprior(priors, params::NamedTuple) = sum(prior -> lnprior(prior, params), priors) lnprior(priors, params) = mapreduce(lnprior, +, priors, params) -lnprior(model::TimingModel, params::NamedTuple) = lnprior(model.priors, params) -lnprior(model::TimingModel, params) = lnprior(model.priors, params) + +calc_lnprior(model::TimingModel, params::NamedTuple) = lnprior(model.priors, params) +calc_lnprior(model::TimingModel, params) = lnprior(model.priors, params) """ get_lnprior_func(::TimingModel)::Function Returns a callable that evaluates the log-prior given a collection of parameter values. """ -get_lnprior_func(model::TimingModel) = params -> lnprior(model, params) +get_lnprior_func(model::TimingModel) = params -> calc_lnprior(model, params) """Evaluate the prior transform function.""" prior_transform(priors, cube) = map(prior_transform, priors, cube) diff --git a/src/pulsar/pulsar.jl b/src/pulsar/pulsar.jl new file mode 100644 index 00000000..81fb2861 --- /dev/null +++ b/src/pulsar/pulsar.jl @@ -0,0 +1,23 @@ +export Pulsar + +""" + Pulsar{M<:TimingModel,T<:TOABase} + +Represents a pulsar dataset with a timing model and a set of TOAs. +All TOAs must be of the same paradigm (narrowband/wideband). +""" +struct Pulsar{M<:TimingModel,T<:TOABase} + model::M + toas::Vector{T} +end + +calc_lnlike(psr::Pulsar, params) = calc_lnlike(psr.model, psr.toas, params) +calc_lnlike_serial(psr::Pulsar, params) = calc_lnlike_serial(psr.model, psr.toas, params) + +calc_lnprior(psr::Pulsar, params) = calc_lnprior(psr.model, params) +prior_transform(psr::Pulsar, cube) = prior_transform(psr.model, cube) + +calc_lnpost(psr::Pulsar, params) = calc_lnpost(psr.model, psr.toas, params) +calc_lnpost_serial(psr::Pulsar, params) = calc_lnpost_serial(psr.model, psr.toas, params) +calc_lnpost_vectorized(psr::Pulsar, paramss) = + calc_lnpost_vectorized(psr.model, psr.toas, paramss) diff --git a/test/test_NGC6440E.jl b/test/test_NGC6440E.jl index ccaac8de..1b8395f5 100644 --- a/test/test_NGC6440E.jl +++ b/test/test_NGC6440E.jl @@ -146,12 +146,32 @@ params = model.param_handler._default_params_tuple parv = read_param_values_to_vector(model.param_handler, params) - calc_lnpost = get_lnpost_func(model, toas) - @test isfinite(calc_lnpost(params)) + calc_lnpost_ = get_lnpost_func(model, toas) + @test isfinite(calc_lnpost_(params)) calc_lnpost_vec = get_lnpost_func(model, toas, true) paramss = transpose([parv parv parv]) @test allequal(calc_lnpost_vec(paramss)) - @test calc_lnpost_vec(paramss)[1] ≈ calc_lnpost(params) + @test calc_lnpost_vec(paramss)[1] ≈ calc_lnpost_(params) + end + + @testset "pulsar" begin + psr = Pulsar(model, toas) + + params = model.param_handler._default_params_tuple + parv = read_param_values_to_vector(model.param_handler, params) + + @test calc_lnlike(psr, parv) ≈ calc_lnlike_serial(psr, parv) + + @test isfinite(calc_lnprior(psr, parv)) + + cube = 0.5 .* ones(length(parv)) + @test isfinite(calc_lnprior(psr, prior_transform(psr, cube))) + + @test calc_lnpost(psr, parv) == calc_lnpost(psr, params) + @test calc_lnpost(psr, parv) ≈ calc_lnpost_serial(psr, parv) + + paramss = transpose([parv parv parv]) + @test allequal(calc_lnpost_vectorized(psr, paramss)) end end diff --git a/test/test_pure_rotator.jl b/test/test_pure_rotator.jl index 1d053848..9e91e64f 100644 --- a/test/test_pure_rotator.jl +++ b/test/test_pure_rotator.jl @@ -117,7 +117,7 @@ calc_lnprior = get_lnprior_func(model) params = model.param_handler._default_params_tuple parv = read_param_values_to_vector(model.param_handler, params) - @test isfinite(calc_lnprior(model.param_handler._default_params_tuple)) + @test isfinite(calc_lnprior(params)) @test calc_lnprior(params) == calc_lnprior(parv) prior_transform = get_prior_transform_func(model)