Merge branch 'solarwind' of https://github.com/abhisrkckl/Vela.jl int…

…o solarwind

CHANGELOG.md

@@ -2,6 +2,7 @@
 ## Added
 - `time_residuals`, `scaled_toa_uncertainties`, and `model_dm`, `from_pint` methods in `SPNTA`
 - NE_SW derivatives in `SolarWind`
+- `Pulsar` class
 ## Changed
 ## Fixed
 ## Removed

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, 

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})

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),
 }
 

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
 

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

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

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

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)

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)