Merge pull request #1859 from abhisrkckl/ne_sw1

Time derivatives of NE_SW (NE_SW1, ...)

CHANGELOG-unreleased.md

@@ -12,6 +12,8 @@ the released changes.
 - Command line scripts now automatically do `allow_tcb` and `allow_T2` while reading par files.
 - Updated the `plot_chains` function in `event_optimize` so that the subplots are a fixed size to prevent the subplots from being condensed in the case of many fit parameters.
 ### Added
+- Time derivatives of NE_SW in `SolarWindDispersion`
+- New prefix pattern for `split_prefixed_name` to handle derivatives of NE_SW
 - Added an option `nbin` to `photonphase` to decide how many phase bins to use for the phaseogram
 - Added an option `linearize_model` to speed up the photon phases calculation within `event_optimize` through the designmatrix.
 - Added AIC and BIC calculation to be written in the post fit parfile from `event_optimize`

src/pint/models/dispersion_model.py

@@ -21,7 +21,7 @@
     taylor_horner,
     taylor_horner_deriv,
 )
-from pint import DMconst
+from pint import DMconst, dmu
 from pint.exceptions import MissingParameter, MissingTOAs
 
 # This value is cited from Duncan Lorimer, Michael Kramer, Handbook of Pulsar
@@ -76,7 +76,7 @@ def dm_value(self, toas):
             DM values at given TOAs in the unit of DM.
         """
         toas_table = toas if isinstance(toas, Table) else toas.table
-        dm = np.zeros(len(toas_table)) * self._parent.DM.units
+        dm = np.zeros(len(toas_table)) * dmu
 
         for dm_f in self.dm_value_funcs:
             dm += dm_f(toas)
@@ -183,8 +183,8 @@ def __init__(self):
 
     def setup(self):
         super().setup()
-        base_dms = list(self.get_prefix_mapping_component("DM").values())
-        base_dms += ["DM"]
+
+        base_dms = ["DM"] + list(self.get_prefix_mapping_component("DM").values())
 
         for dm_name in base_dms:
             self.register_deriv_funcs(self.d_delay_d_dmparam, dm_name)
@@ -206,10 +206,10 @@ def validate(self):
                 )
 
     def DM_derivative_unit(self, n):
-        return "pc cm^-3/yr^%d" % n if n else "pc cm^-3"
+        return f"pc cm^-3/yr^{n}" if n != 0 else "pc cm^-3"
 
     def DM_derivative_description(self, n):
-        return "%d'th time derivative of the dispersion measure" % n
+        return f"{n}'th time derivative of the dispersion measure"
 
     def get_DM_terms(self):
         """Return a list of the DM term values in the model: [DM, DM1, ..., DMn]"""
@@ -223,7 +223,7 @@ def base_dm(self, toas):
             if DMEPOCH is None:
                 # Should be ruled out by validate()
                 raise ValueError(
-                    f"DMEPOCH not set but some derivatives are not zero: {dm_terms}"
+                    f"DMEPOCH not set but some DM derivatives are not zero: {dm_terms}"
                 )
             else:
                 dt = (toas["tdbld"] - DMEPOCH) * u.day
@@ -673,7 +673,7 @@ def dmx_dm(self, toas):
             condition, tbl["mjd_float"]
         )
         # Get DMX delays
-        dm = np.zeros(len(tbl)) * self._parent.DM.units
+        dm = np.zeros(len(tbl)) * dmu
         for k, v in select_idx.items():
             dm[v] += getattr(self, k).quantity
         return dm

src/pint/models/solar_wind_dispersion.py

@@ -10,7 +10,12 @@
 
 from pint import DMconst
 from pint.models.dispersion_model import Dispersion
-from pint.models.parameter import floatParameter, intParameter, prefixParameter
+from pint.models.parameter import (
+    MJDParameter,
+    floatParameter,
+    intParameter,
+    prefixParameter,
+)
 import pint.utils
 from pint.models.timing_model import MissingTOAs
 from pint.toa_select import TOASelect
@@ -302,6 +307,26 @@ def __init__(self):
                 tcb2tdb_scale_factor=(const.c * DMconst),
             )
         )
+        self.add_param(
+            prefixParameter(
+                name="NE_SW1",
+                units="cm^-3 / yr",
+                value=0.0,
+                description="Derivative of Solar Wind density at 1 AU",
+                unit_template=self.NE_SW_derivative_unit,
+                description_template=self.NE_SW_derivative_description,
+                type_match="float",
+                tcb2tdb_scale_factor=(const.c * DMconst),
+            )
+        )
+        self.add_param(
+            MJDParameter(
+                name="SWEPOCH",
+                description="Epoch of NE_SW measurement",
+                time_scale="tdb",
+                tcb2tdb_scale_factor=u.Quantity(1),
+            )
+        )
         self.add_param(
             floatParameter(
                 name="SWP",
@@ -324,11 +349,24 @@ def __init__(self):
 
     def setup(self):
         super().setup()
-        self.register_dm_deriv_funcs(self.d_dm_d_ne_sw, "NE_SW")
-        self.register_deriv_funcs(self.d_delay_d_ne_sw, "NE_SW")
+
         self.register_dm_deriv_funcs(self.d_dm_d_swp, "SWP")
         self.register_deriv_funcs(self.d_delay_d_swp, "SWP")
 
+        base_nesws = ["NE_SW"] + list(
+            self.get_prefix_mapping_component("NE_SW").values()
+        )
+
+        for nesw_name in base_nesws:
+            self.register_deriv_funcs(self.d_delay_d_ne_sw, nesw_name)
+            self.register_dm_deriv_funcs(self.d_dm_d_ne_sw, nesw_name)
+
+    def NE_SW_derivative_unit(self, n):
+        return f"cm^-3/yr^{n}" if n != 0 else "cm^-3"
+
+    def NE_SW_derivative_description(self, n):
+        return f"{n}'th time derivative of the Solar Wind density at 1 AU"
+
     def solar_wind_geometry(self, toas):
         """Return the geometry of solar wind dispersion.
 
@@ -375,22 +413,51 @@ def solar_wind_dm(self, toas):
         SWM==1:
             Hazboun et al. 2022
         """
-        if self.NE_SW.value == 0:
+        ne_sw_terms = self.get_NE_SW_terms()
+
+        if len(ne_sw_terms) == 1:
+            ne_sw = self.NE_SW.quantity * np.ones(len(toas))
+        else:
+            if any(t.value != 0 for t in ne_sw_terms[1:]):
+                SWEPOCH = self.SWEPOCH.value
+                if SWEPOCH is None:
+                    # Should be ruled out by validate()
+                    raise ValueError(
+                        f"SWEPOCH not set but some NE_SW derivatives are not zero: {ne_sw_terms}"
+                    )
+                else:
+                    dt = (toas["tdbld"] - SWEPOCH) * u.day
+                dt_value = dt.to_value(u.yr)
+            else:
+                dt_value = np.zeros(len(toas), dtype=np.longdouble)
+
+            ne_sw_terms_value = [d.value for d in ne_sw_terms]
+            ne_sw = (
+                pint.utils.taylor_horner(dt_value, ne_sw_terms_value) * self.NE_SW.units
+            )
+
+        if np.all(ne_sw.value == 0):
             return np.zeros(len(toas)) * u.pc / u.cm**3
+
         if self.SWM.value not in [0, 1]:
             raise NotImplementedError(
                 f"Solar Dispersion Delay not implemented for SWM {self.SWM.value}"
             )
+
         solar_wind_geometry = self.solar_wind_geometry(toas)
-        solar_wind_dm = self.NE_SW.quantity * solar_wind_geometry
+        solar_wind_dm = ne_sw * solar_wind_geometry
         return solar_wind_dm.to(u.pc / u.cm**3)
 
     def solar_wind_delay(self, toas, acc_delay=None):
         """This is a wrapper function to compute solar wind dispersion delay."""
-        if self.NE_SW.value == 0:
-            return np.zeros(len(toas)) * u.s
         return self.dispersion_type_delay(toas)
 
+    def get_NE_SW_terms(self):
+        """Return a list of the NE_SW term values in the model: [NE_SW, NE_SW1, ..., NE_SWn]"""
+        return [self.NE_SW.quantity] + self._parent.get_prefix_list(
+            "NE_SW", start_index=1
+        )
+
     def d_dm_d_ne_sw(self, toas, param_name, acc_delay=None):
         """Derivative of of DM wrt the solar wind ne amplitude."""
         if self.SWM.value in [0, 1]:
@@ -399,17 +466,31 @@ def d_dm_d_ne_sw(self, toas, param_name, acc_delay=None):
             raise NotImplementedError(
                 "Solar Dispersion Delay not implemented for SWM %d" % self.SWM.value
             )
-        return solar_wind_geometry
+
+        par = getattr(self, param_name)
+        order = 0 if param_name == "NE_SW" else par.index
+
+        ne_sws = self.get_NE_SW_terms()
+
+        ne_sw_terms = np.longdouble(np.zeros(len(ne_sws)))
+        ne_sw_terms[order] = np.longdouble(1.0)
+
+        if self.SWEPOCH.value is None:
+            if any(t.value != 0 for t in ne_sws[1:]):
+                # Should be ruled out by validate()
+                raise ValueError(f"SWEPOCH is not set but {param_name} is not zero")
+            SWEPOCH = 0
+        else:
+            SWEPOCH = self.SWEPOCH.value
+
+        dt = (toas["tdbld"] - SWEPOCH) * u.day
+
+        return (solar_wind_geometry * dt**order / scipy.special.factorial(order)).to(
+            u.pc / u.cm**3 / par.units
+        )
 
     def d_delay_d_ne_sw(self, toas, param_name, acc_delay=None):
-        try:
-            bfreq = self._parent.barycentric_radio_freq(toas)
-        except AttributeError:
-            warn("Using topocentric frequency for dedispersion!")
-            bfreq = toas.table["freq"]
-        deriv = self.d_delay_d_dmparam(toas, "NE_SW")
-        deriv[bfreq < 1.0 * u.MHz] = 0.0
-        return deriv
+        return self.d_delay_d_dmparam(toas, param_name)
 
     def d_solar_wind_geometry_d_swp(self, toas, param_name, acc_delay=None):
         """Derivative of solar_wind_geometry (path length) wrt power-law index p
@@ -452,10 +533,12 @@ def d_delay_d_swp(self, toas, param_name, acc_delay=None):
 
     def print_par(self, format="pint"):
         result = ""
-        result += getattr(self, "NE_SW").as_parfile_line(format=format)
-        result += getattr(self, "SWM").as_parfile_line(format=format)
-        if self.SWM.value == 1:
-            result += getattr(self, "SWP").as_parfile_line(format=format)
+        for par in self.params:
+            if par == "SWP" and self.SWM.value != 1:
+                continue
+            else:
+                result += getattr(self, par).as_parfile_line(format=format)
+
         return result
 
     def get_max_dm(self):

src/pint/utils.py

@@ -357,7 +357,7 @@ def has_astropy_unit(x: Any) -> bool:
     re.compile(r"^([a-zA-Z]*\d+[a-zA-Z]+)(\d+)$"),  # For the prefix like T2EFAC2
     re.compile(r"^([a-zA-Z]+)0*(\d+)$"),  # For the prefix like F12
     re.compile(r"^([a-zA-Z0-9]+_)(\d+)$"),  # For the prefix like DMXR1_3
-    # re.compile(r'([a-zA-Z]\d[a-zA-Z]+)(\d+)'),  # for prefixes like PLANET_SHAPIRO2?
+    re.compile(r"([a-zA-Z]+_[a-zA-Z]+)(\d+)$"),  # for prefixes like NE_SW2?
 ]
 
 

tests/test_solar_wind.py

@@ -359,3 +359,83 @@ def test_overlapping_swx():
     dm3 = model3.components["SolarWindDispersionX"].swx_dm(toas)
     dm3[toas.get_mjds() >= 54180 * u.d] = 0
     assert np.allclose(dm2 - dm, dm3)
+
+
+def test_nesw_derivatives():
+    par = """
+        PSRJ            J1744-1134
+        ELONG           266.119458498                 6.000e-09
+        ELAT            11.80517508                   3.000e-08
+        DM              3.1379                        4.000e-04
+        PEPOCH          55000
+        F0              245.4261196898081             5.000e-13
+        F1              -5.38156E-16                  3.000e-21
+        POSEPOCH        59150
+        DMEPOCH         55000
+        CLK             TT(BIPM2018)
+        EPHEM           DE436
+        RM              1.62                          9.000e-02
+        PX              2.61                          9.000e-02
+        DM1             -7.0E-5                       3.000e-05
+        PMELONG         19.11                         2.000e-02
+        DM2             1.7E-5                        4.000e-06
+        PMELAT          -9.21                         1.300e-01
+        UNITS           TDB        
+        NE_SW           8           1
+        NE_SW1          0           1
+        NE_SW2          0           1
+        SWEPOCH         55000
+    """
+    m = get_model(StringIO(par))
+    t = make_fake_toas_uniform(54500, 55500, 1000, m, add_noise=True)
+    ftr = Fitter.auto(t, m)
+    ftr.fit_toas()
+
+    assert (
+        m.NE_SW.value - ftr.model.NE_SW.value
+    ) / ftr.model.NE_SW.uncertainty_value < 3
+    assert (
+        m.NE_SW1.value - ftr.model.NE_SW1.value
+    ) / ftr.model.NE_SW1.uncertainty_value < 3
+    assert (
+        m.NE_SW2.value - ftr.model.NE_SW2.value
+    ) / ftr.model.NE_SW2.uncertainty_value < 3
+
+    m.components["SolarWindDispersion"]
+
+
+def test_expression():
+    par = """
+        PSRJ            J1744-1134
+        ELONG           266.119458498                 6.000e-09
+        ELAT            11.80517508                   3.000e-08
+        DM              3.1379                        4.000e-04
+        PEPOCH          55000
+        F0              245.4261196898081             5.000e-13
+        F1              -5.38156E-16                  3.000e-21
+        POSEPOCH        59150
+        DMEPOCH         55000
+        CLK             TT(BIPM2018)
+        EPHEM           DE436
+        PX              2.61                          9.000e-02
+        PMELONG         19.11                         2.000e-02
+        PMELAT          -9.21                         1.300e-01
+        UNITS           TDB        
+        NE_SW           8           1
+        NE_SW1          1           1
+        SWEPOCH         55000
+    """
+    m = get_model(StringIO(par))
+    t = make_fake_toas_uniform(54500, 55500, 10, m, add_noise=True)
+
+    t0 = m.SWEPOCH.value * u.day
+    t1 = t["tdbld"][-1] * u.day
+    dt = t1 - t0
+
+    assert np.isclose(
+        (
+            m.components["SolarWindDispersion"].solar_wind_dm(t)
+            / m.components["SolarWindDispersion"].solar_wind_geometry(t)
+        )[-1],
+        (m.NE_SW.quantity + dt * m.NE_SW1.quantity),
+    )