Merge pull request #1615 from abhisrkckl/fdjump-ugly

FDJUMP implementation (Hacky version)

CHANGELOG-unreleased.md

@@ -18,6 +18,7 @@ the released changes.
 - Options to add a TZR TOA (`AbsPhase`) during the creation of a `TimingModel` using `ModelBuilder.__call__`, `get_model`, and `get_model_and_toas`
 - `pint.print_info()` function for bug reporting
 - Added an autocorrelation function to check for chain convergence in `event_optimize`
+- A hacky implementation of system-dependent FD parameters (FDJUMP)
 - Minor doc updates to explain default NHARMS and missing derivative functions
 ### Fixed
 - Deleting JUMP1 from flag tables will not prevent fitting

src/pint/models/__init__.py

@@ -36,6 +36,7 @@
 from pint.models.solar_system_shapiro import SolarSystemShapiro
 from pint.models.solar_wind_dispersion import SolarWindDispersion, SolarWindDispersionX
 from pint.models.spindown import Spindown
+from pint.models.fdjump import FDJump
 
 # Import the main timing model classes
 from pint.models.timing_model import DEFAULT_ORDER, TimingModel

src/pint/models/fdjump.py

@@ -0,0 +1,205 @@
+"""System and frequency dependent delays to model profile evolution."""
+
+import re
+from warnings import warn
+
+import astropy.units as u
+import numpy as np
+
+from pint.models.parameter import boolParameter, maskParameter
+from pint.models.timing_model import DelayComponent
+
+fdjump_max_index = 20
+
+
+class FDJump(DelayComponent):
+    """A timing model for system-dependent frequency evolution of pulsar
+    profiles.
+
+    This model expresses the delay as a polynomial function of the
+    observing frequency/logarithm of observing frequency in the SSB frame.
+    This is intended to compensate for the delays introduced by frequency-dependent
+    profile structure when a different profiles are used for different systems.
+
+    The default behavior is to have FDJUMPs as polynomials of the observing
+    frequency (rather than log-frequency). This is different from the convention
+    used for global FD parameters. This choice is made to be compatible with tempo2.
+    This is controlled using the FDJUMPLOG parameter. "FDJUMPLOG Y" may not be
+    tempo2-compatible.
+
+    Note
+    ----
+    FDJUMPs have two indices: the polynomial/FD/prefix index and the system/mask
+    index. i.e., they have properties of both maskParameters such as JUMPs and
+    prefixParameters such as FDs. There is currently no elegant way in PINT to implement
+    such parameters due to the way parameter indexing is implemented; there is no way to
+    distinguish between mask and prefix indices.
+
+    Hence, they are implemented here as maskParameters as a stopgap measure.
+    This means that there must be an upper limit for the FD indices. This is controlled
+    using the `pint.models.fdjump.fdjump_max_index` variable, and is 20 by default.
+    Note that this is strictly a limitation of the implementation and not a property
+    of FDJUMPs themselves.
+
+    FDJUMPs appear in tempo2-format par files as "FDJUMPp", where p is the FD index.
+    The mask index is not explicitly mentioned in par files similar to JUMPs.
+    PINT understands both "FDJUMPp" and "FDpJUMP" as the same parameter in par files,
+    but the internal representation is always "FDpJUMPq", where q is the mask index.
+
+    PINT understands 'q' as the mask parameter just fine, but the identification of 'p'
+    as the prefix parameter is done in a hacky way.
+
+    This implementation may be overhauled in the future.
+
+    Parameters supported:
+
+    .. paramtable::
+        :class: pint.models.fdjump.FDJump
+    """
+
+    register = True
+    category = "fdjump"
+
+    def __init__(self):
+        super().__init__()
+
+        # Matches "FDpJUMPq" where p and q are integers.
+        self.param_regex = re.compile("^FD(\\d+)JUMP(\\d+)")
+
+        self.add_param(
+            boolParameter(
+                name="FDJUMPLOG",
+                value=False,
+                description="Whether to use log-frequency (Y) or linear-frequency (N) for computing FDJUMPs.",
+            )
+        )
+        for j in range(1, fdjump_max_index + 1):
+            self.add_param(
+                maskParameter(
+                    name=f"FD{j}JUMP",
+                    units="second",
+                    description=f"System-dependent FD parameter of polynomial index {j}",
+                )
+            )
+
+        self.delay_funcs_component += [self.fdjump_delay]
+
+    def setup(self):
+        super().setup()
+
+        self.fdjumps = [
+            mask_par
+            for mask_par in self.get_params_of_type("maskParameter")
+            if self.param_regex.match(mask_par)
+        ]
+
+        for fdj in self.fdjumps:
+            # prevents duplicates from being added to phase_deriv_funcs
+            if fdj in self.deriv_funcs.keys():
+                del self.deriv_funcs[fdj]
+            self.register_deriv_funcs(self.d_delay_d_FDJUMP, fdj)
+
+    def get_fd_index(self, par):
+        """Extract the FD index from an FDJUMP parameter name. In a parameter name
+        "FDpJUMPq", p is the FD/prefix index and q is the mask index.
+
+        Parameters
+        ----------
+        par: Parameter name (str)
+
+        Returns
+        -------
+        FD index (int)
+        """
+        if m := self.param_regex.match(par):
+            return int(m.groups()[0])
+        else:
+            raise ValueError(
+                f"The given parameter {par} does not correspond to an FDJUMP."
+            )
+
+    def get_freq_y(self, toas):
+        """Get frequency or log-frequency in GHz based on the FDJUMPLOG value.
+        Returns (freq/1_GHz) if FDJUMPLOG==N and log(freq/1_GHz) if FDJUMPLOG==Y.
+        Any non-finite values are replaced by zero.
+
+        Parameters
+        ----------
+        toas: pint.toa.TOAs
+
+        Returns
+        -------
+        (freq/1_GHz) or log(freq/1_GHz) depending on the value of FDJUMPLOG (float).
+        """
+        tbl = toas.table
+        try:
+            freq = self._parent.barycentric_radio_freq(toas)
+        except AttributeError:
+            warn("Using topocentric frequency for frequency dependent delay!")
+            freq = tbl["freq"]
+
+        y = (
+            np.log(freq.to(u.GHz).value)
+            if self.FDJUMPLOG.value
+            else freq.to(u.GHz).value
+        )
+        non_finite = np.invert(np.isfinite(y))
+        y[non_finite] = 0.0
+
+        return y
+
+    def fdjump_delay(self, toas, acc_delay=None):
+        """Calculate frequency dependent delay.
+
+        If FDJUMPLOG is Y, use the following expression (similar to global FD parameters):
+
+            FDJUMP_delay = sum_i(c_i * (log(obs_freq/1GHz))^i)
+
+        If FDJUMPLOG is N, use the following expression (same as in tempo2, default):
+
+            FDJUMP_delay = sum_i(c_i * (obs_freq/1GHz)^i)
+        """
+        y = self.get_freq_y(toas)
+
+        delay = np.zeros_like(y)
+        for fdjump in self.fdjumps:
+            fdj = getattr(self, fdjump)
+            if fdj.quantity is not None:
+                mask = fdj.select_toa_mask(toas)
+                ymask = y[mask]
+                fdidx = self.get_fd_index(fdjump)
+                fdcoeff = fdj.value
+                delay[mask] += fdcoeff * ymask**fdidx
+
+        return delay * u.s
+
+    def d_delay_d_FDJUMP(self, toas, param, acc_delay=None):
+        """Derivative of delay w.r.t. FDJUMP parameters."""
+        assert (
+            bool(self.param_regex.match(param))
+            and hasattr(self, param)
+            and getattr(self, param).quantity is not None
+        ), f"{param} is not present in the FDJUMP model."
+
+        y = self.get_freq_y(toas)
+        mask = getattr(self, param).select_toa_mask(toas)
+        ymask = y[mask]
+        fdidx = self.get_fd_index(param)
+
+        delay_derivative = np.zeros_like(y)
+        delay_derivative[mask] = ymask**fdidx
+
+        return delay_derivative * u.dimensionless_unscaled
+
+    def print_par(self, format="pint"):
+        par = super().print_par(format)
+
+        if format != "tempo2":
+            return par
+
+        for fdjump in self.fdjumps:
+            if getattr(self, fdjump).quantity is not None:
+                j = self.get_fd_index(fdjump)
+                par = par.replace(f"FD{j}JUMP", f"FDJUMP{j}")
+
+        return par

src/pint/models/frequency_dependent.py

@@ -24,7 +24,7 @@ class FD(DelayComponent):
 
     @classmethod
     def _description_template(cls, x):
-        return "%d term of frequency dependent coefficients" % x
+        return f"{x} term of frequency dependent coefficients"
 
     register = True
     category = "frequency_dependent"
@@ -72,7 +72,7 @@ def FD_delay(self, toas, acc_delay=None):
         Time Measurements, and Analysis of 37 Millisecond Pulsars, The
         Astrophysical Journal, Volume 813, Issue 1, article id. 65, 31 pp.(2015).
         Eq.(2):
-        FDdelay = sum(c_i * (log(obs_freq/1GHz))^i)
+            FD_delay = sum_i(c_i * (log(obs_freq/1GHz))^i)
         """
         tbl = toas.table
         try:

src/pint/models/model_builder.py

@@ -7,6 +7,7 @@
 from pathlib import Path
 from astropy import units as u
 from loguru import logger as log
+import re
 
 from pint.models.astrometry import Astrometry
 from pint.models.parameter import maskParameter
@@ -64,6 +65,28 @@ def parse_parfile(parfile):
     return parfile_dict
 
 
+def _replace_fdjump_in_parfile_dict(pardict):
+    """Replace parameter names s of the form "FDJUMPp" by "FDpJUMP"
+    while reading the par file, where p is the prefix index.
+
+    Ideally, this should have been done using the parameter alias
+    mechanism, but there is no easy way to do this currently due to the
+    mask and prefix indices being treated in an identical manner.
+
+    See :class:`~pint.models.fdjump.FDJump` for more details."""
+    fdjumpn_regex = re.compile("^FDJUMP(\\d+)")
+    pardict_new = {}
+    for key, value in pardict.items():
+        if m := fdjumpn_regex.match(key):
+            j = int(m.groups()[0])
+            new_key = f"FD{j}JUMP"
+            pardict_new[new_key] = value
+        else:
+            pardict_new[key] = value
+
+    return pardict_new
+
+
 class ModelBuilder:
     """Class for building a `TimingModel` object from a parameter file.
 
@@ -330,6 +353,11 @@ def _pintify_parfile(self, parfile, allow_name_mixing=False):
             parfile_dict = parse_parfile(parfile)
         else:
             parfile_dict = parfile
+
+        # This is a special-case-hack to deal with FDJUMP parameters.
+        # @TODO: Implement a general mechanism to deal with cases like this.
+        parfile_dict = _replace_fdjump_in_parfile_dict(parfile_dict)
+
         for k, v in parfile_dict.items():
             try:
                 pint_name, init0 = self.all_components.alias_to_pint_param(k)

tests/datafile/J1909-3744.NB.par

@@ -0,0 +1,53 @@
+# Created: 2023-08-04T12:11:27.791090
+# PINT_version: 0.9.6+53.g447cc061.dirty
+# User: Abhimanyu S (abhimanyu)
+# Host: abhimanyu-HP-Notebook
+# OS: Linux-5.15.0-71-generic-x86_64-with-glibc2.35
+# Python: 3.9.13 | packaged by conda-forge | (main, May 27 2022, 16:58:50) 
+# [GCC 10.3.0]
+# Format: pint
+PSRJ                           J1909-3744
+EPHEM                               DE440
+CLK                          TT(BIPM2019)
+UNITS                                 TDB
+START              58260.0274418715905556
+FINISH             59496.4896411149411921
+TIMEEPH                              FB90
+T2CMETHOD                        IAU2000B
+DILATEFREQ                              N
+DMDATA                                  N
+NTOA                                  446
+CHI2                   13166.811752135152
+RAJ                     19:09:47.42480160 0 0.00000036670717903709
+DECJ                   -37:44:14.90754000 0 0.00001603708497544348
+PMRA                    -9.51261566578421 0 0.001606518013582412
+PMDEC                 -35.776735309087286 0 0.0057765724888464666
+PX                      0.881638358140426 0 0.013674560477413309
+POSEPOCH           58999.9997541495914071
+F0                  339.31569192154273734 1 1.2497063444336193814e-12
+F1              -1.6123419575084687074e-15 1 4.3048206597047727197e-20
+F2                                    0.0
+PEPOCH             58999.9997541495914071
+CORRECT_TROPOSPHERE                         Y
+PLANET_SHAPIRO                          N
+NE_SW                                 0.0
+SWM                                   0.0
+DM                 10.3912220010111820715 1 5.91828415377580018e-06
+DM1             0.000118226481794639765015 1 5.4915807592283010527e-06
+DM2                                   0.0
+DMEPOCH            59629.9997443813168694
+BINARY                               ELL1
+PB                  1.5334494515017690281 0 1.3220826895008462359e-12
+PBDOT               5.053209847126312e-13 0 2.521391425976759e-15
+A1                     1.8979910746557984 0 2.117833848195567e-08
+XDOT              -3.9210500492434694e-16 0 7.037623119797259e-17
+M2                     0.2086645878980989 0 0.0012758962708215615
+SINI                   0.9980748600209708 0 5.844811546142236e-05
+TASC               55015.4279066849879101 0 9.73762114761625686e-10
+EPS1            2.6338029192577372046e-08 0 1.033583508265e-08
+EPS2            -1.0054034660955184656e-07 0 5.72187338209e-09
+# ECC             1.0393292586143181096e-07
+# OM                  165.32038917886206575
+TZRMJD             59000.9659140826392014
+TZRSITE                              gmrt
+TZRFRQ                        1310.097656