Merge pull request #1867 from kerrm/glitch_tweak

Glitch tweak

CHANGELOG-unreleased.md

@@ -25,6 +25,7 @@ the released changes.
 - Changed WAVE_OM units from 1/d to rad/d.
 - When EQUAD is created from TNEQ, has proper TCB->TDB conversion info
 - TOA selection masks will work when only TOA is the first one
+- Condense code in Glitch model and add test coverage.
 ### Removed
 - macOS 12 CI 
 

src/pint/models/glitch.py

@@ -145,49 +145,48 @@ def setup(self):
         ]
         for idx in set(self.glitch_indices):
             for param in self.glitch_prop:
-                if not hasattr(self, param + "%d" % idx):
+                check = f"{param}{idx}"
+                if not hasattr(self, check):
                     param0 = getattr(self, f"{param}1")
                     self.add_param(param0.new_param(idx))
-                    getattr(self, param + "%d" % idx).value = 0.0
+                    getattr(self, check).value = 0.0
                 self.register_deriv_funcs(
-                    getattr(self, f"d_phase_d_{param[:-1]}"), param + "%d" % idx
+                    getattr(self, f"d_phase_d_{param[:-1]}"), check
                 )
 
     def validate(self):
         """Validate parameters input."""
         super().validate()
         for idx in set(self.glitch_indices):
-            if not hasattr(self, "GLEP_%d" % idx):
-                msg = "Glitch Epoch is needed for Glitch %d." % idx
-                raise MissingParameter("Glitch", "GLEP_%d" % idx, msg)
-            else:  # Check to see if both the epoch and phase are to be fit
-                if (
-                    hasattr(self, "GLPH_%d" % idx)
-                    and (not getattr(self, "GLEP_%d" % idx).frozen)
-                    and (not getattr(self, "GLPH_%d" % idx).frozen)
-                ):
-                    raise ValueError(
-                        "Both the glitch epoch and phase cannot be fit for Glitch %d."
-                        % idx
-                    )
+            glep = f"GLEP_{idx}"
+            glph = f"GLPH_{idx}"
+            if (not hasattr(self, glep)) or (getattr(self, glep).quantity is None):
+                msg = f"Glitch Epoch is needed for Glitch {idx}"
+                raise MissingParameter("Glitch", glep, msg)
+            # Check to see if both the epoch and phase are to be fit
+            if (
+                hasattr(self, glph)
+                and (not getattr(self, glep).frozen)
+                and (not getattr(self, glph).frozen)
+            ):
+                raise ValueError(
+                    f"Both the glitch epoch and phase cannot be fit for Glitch {idx}."
+                )
 
         # Check the Decay Term.
         glf0dparams = [x for x in self.params if x.startswith("GLF0D_")]
         for glf0dnm in glf0dparams:
             glf0d = getattr(self, glf0dnm)
             idx = glf0d.index
-            if glf0d.value != 0.0 and getattr(self, "GLTD_%d" % idx).value == 0.0:
-                msg = (
-                    "None zero GLF0D_%d parameter needs a none"
-                    " zero GLTD_%d parameter" % (idx, idx)
-                )
-                raise MissingParameter("Glitch", "GLTD_%d" % idx, msg)
+            if glf0d.value != 0.0 and getattr(self, f"GLTD_{idx}").value == 0.0:
+                msg = f"Non-zero GLF0D_{idx} parameter needs a non-zero GLTD_{idx} parameter"
+                raise MissingParameter("Glitch", f"GLTD_{idx}", msg)
 
     def print_par(self, format="pint"):
         result = ""
         for idx in set(self.glitch_indices):
             for param in self.glitch_prop:
-                par = getattr(self, param + "%d" % idx)
+                par = getattr(self, f"{param}{idx}")
                 result += par.as_parfile_line(format=format)
         return result
 
@@ -204,22 +203,21 @@ def glitch_phase(self, toas, delay):
             glep = getattr(self, glepnm)
             idx = glep.index
             eph = glep.value
-            dphs = getattr(self, "GLPH_%d" % idx).quantity
-            dF0 = getattr(self, "GLF0_%d" % idx).quantity
-            dF1 = getattr(self, "GLF1_%d" % idx).quantity
-            dF2 = getattr(self, "GLF2_%d" % idx).quantity
+            dphs = getattr(self, f"GLPH_{idx}").quantity
+            dF0 = getattr(self, f"GLF0_{idx}").quantity
+            dF1 = getattr(self, f"GLF1_{idx}").quantity
+            dF2 = getattr(self, f"GLF2_{idx}").quantity
             dt = (tbl["tdbld"] - eph) * u.day - delay
             dt = dt.to(u.second)
             affected = dt > 0.0  # TOAs affected by glitch
             # decay term
-            dF0D = getattr(self, "GLF0D_%d" % idx).quantity
+            dF0D = getattr(self, f"GLF0D_{idx}").quantity
             if dF0D != 0.0:
-                tau = getattr(self, "GLTD_%d" % idx).quantity
+                tau = getattr(self, f"GLTD_{idx}").quantity
                 decayterm = dF0D * tau * (1.0 - np.exp(-(dt[affected] / tau)))
             else:
                 decayterm = u.Quantity(0)
 
-            log.debug(f"Glitch phase for glitch {idx}: {dphs} {dphs.unit}")
             phs[affected] += (
                 dphs
                 + dt[affected]
@@ -232,114 +230,91 @@ def glitch_phase(self, toas, delay):
             )
         return phs
 
-    def deriv_prep(self, toas, param, delay):
+    def deriv_prep(self, toas, param, delay, check_param):
         """Get the things we need for any of the derivative calcs"""
-        tbl = toas.table
         p, ids, idv = split_prefixed_name(param)
+        if p != f"{check_param}_":
+            raise ValueError(
+                f"Can not calculate d_phase_d_{check_param} with respect to {param}."
+            )
+        tbl = toas.table
         eph = getattr(self, f"GLEP_{ids}").value
         dt = (tbl["tdbld"] - eph) * u.day - delay
         dt = dt.to(u.second)
         affected = np.where(dt > 0.0)[0]
-        return tbl, p, ids, idv, dt, affected
+        par = getattr(self, param)
+        zeros = np.zeros(len(tbl), dtype=np.longdouble) << 1 / par.units
+        return tbl, p, ids, idv, dt, affected, par, zeros
 
     def d_phase_d_GLPH(self, toas, param, delay):
         """Calculate the derivative wrt GLPH"""
-        tbl, p, ids, idv, dt, affected = self.deriv_prep(toas, param, delay)
-        if p != "GLPH_":
-            raise ValueError(
-                f"Can not calculate d_phase_d_GLPH with respect to {param}."
-            )
-        par_GLPH = getattr(self, param)
-        dpdGLPH = np.zeros(len(tbl), dtype=np.longdouble) / par_GLPH.units
-        dpdGLPH[affected] += 1.0 / par_GLPH.units
+        tbl, p, ids, idv, dt, affected, par_GLPH, dpdGLPH = self.deriv_prep(
+            toas, param, delay, "GLPH"
+        )
+        dpdGLPH[affected] = 1.0 / par_GLPH.units
         return dpdGLPH
 
     def d_phase_d_GLF0(self, toas, param, delay):
         """Calculate the derivative wrt GLF0"""
-        tbl, p, ids, idv, dt, affected = self.deriv_prep(toas, param, delay)
-        if p != "GLF0_":
-            raise ValueError(
-                f"Can not calculate d_phase_d_GLF0 with respect to {param}."
-            )
-        par_GLF0 = getattr(self, param)
-        dpdGLF0 = np.zeros(len(tbl), dtype=np.longdouble) / par_GLF0.units
+        tbl, p, ids, idv, dt, affected, par_GLF0, dpdGLF0 = self.deriv_prep(
+            toas, param, delay, "GLF0"
+        )
         dpdGLF0[affected] = dt[affected]
         return dpdGLF0
 
     def d_phase_d_GLF1(self, toas, param, delay):
         """Calculate the derivative wrt GLF1"""
-        tbl, p, ids, idv, dt, affected = self.deriv_prep(toas, param, delay)
-        if p != "GLF1_":
-            raise ValueError(
-                f"Can not calculate d_phase_d_GLF1 with respect to {param}."
-            )
-        par_GLF1 = getattr(self, param)
-        dpdGLF1 = np.zeros(len(tbl), dtype=np.longdouble) / par_GLF1.units
-        dpdGLF1[affected] += np.longdouble(0.5) * dt[affected] * dt[affected]
+        tbl, p, ids, idv, dt, affected, par_GLF1, dpdGLF1 = self.deriv_prep(
+            toas, param, delay, "GLF1"
+        )
+        dpdGLF1[affected] = 0.5 * dt[affected] ** 2
         return dpdGLF1
 
     def d_phase_d_GLF2(self, toas, param, delay):
         """Calculate the derivative wrt GLF1"""
-        tbl, p, ids, idv, dt, affected = self.deriv_prep(toas, param, delay)
-        if p != "GLF2_":
-            raise ValueError(
-                f"Can not calculate d_phase_d_GLF2 with respect to {param}."
-            )
-        par_GLF2 = getattr(self, param)
-        dpdGLF2 = np.zeros(len(tbl), dtype=np.longdouble) / par_GLF2.units
-        dpdGLF2[affected] += (
-            np.longdouble(1.0) / 6.0 * dt[affected] * dt[affected] * dt[affected]
+        tbl, p, ids, idv, dt, affected, par_GLF2, dpdGLF2 = self.deriv_prep(
+            toas, param, delay, "GLF2"
         )
+        dpdGLF2[affected] = (1.0 / 6.0) * dt[affected] ** 3
         return dpdGLF2
 
     def d_phase_d_GLF0D(self, toas, param, delay):
         """Calculate the derivative wrt GLF0D"""
-        tbl, p, ids, idv, dt, affected = self.deriv_prep(toas, param, delay)
-        if p != "GLF0D_":
-            raise ValueError(
-                f"Can not calculate d_phase_d_GLF0D with respect to {param}."
-            )
-        par_GLF0D = getattr(self, param)
-        tau = getattr(self, "GLTD_%d" % idv).quantity
-        dpdGLF0D = np.zeros(len(tbl), dtype=np.longdouble) / par_GLF0D.units
-        dpdGLF0D[affected] += tau * (np.longdouble(1.0) - np.exp(-dt[affected] / tau))
+        tbl, p, ids, idv, dt, affected, par_GLF0D, dpdGLF0D = self.deriv_prep(
+            toas, param, delay, "GLF0D"
+        )
+        tau = getattr(self, f"GLTD_{ids}").quantity
+        dpdGLF0D[affected] = tau * (1.0 - np.exp(-dt[affected] / tau))
         return dpdGLF0D
 
     def d_phase_d_GLTD(self, toas, param, delay):
         """Calculate the derivative wrt GLTD"""
-        tbl, p, ids, idv, dt, affected = self.deriv_prep(toas, param, delay)
-        if p != "GLTD_":
-            raise ValueError(
-                f"Can not calculate d_phase_d_GLTD with respect to {param}."
-            )
-        par_GLTD = getattr(self, param)
+        tbl, p, ids, idv, dt, affected, par_GLTD, dpdGLTD = self.deriv_prep(
+            toas, param, delay, "GLTD"
+        )
         if par_GLTD.value == 0.0:
-            return np.zeros(len(tbl), dtype=np.longdouble) / par_GLTD.units
+            return dpdGLTD
         glf0d = getattr(self, f"GLF0D_{ids}").quantity
         tau = par_GLTD.quantity
-        dpdGLTD = np.zeros(len(tbl), dtype=np.longdouble) / par_GLTD.units
-        dpdGLTD[affected] += glf0d * (
-            np.longdouble(1.0) - np.exp(-dt[affected] / tau)
-        ) + glf0d * tau * (-np.exp(-dt[affected] / tau)) * dt[affected] / (tau * tau)
+        et = np.exp(-dt[affected] / tau)
+        dpdGLTD[affected] = glf0d * (
+            1.0 - np.exp(-dt[affected] / tau) * (1.0 + dt[affected] / tau)
+        )
         return dpdGLTD
 
     def d_phase_d_GLEP(self, toas, param, delay):
         """Calculate the derivative wrt GLEP"""
-        tbl, p, ids, idv, dt, affected = self.deriv_prep(toas, param, delay)
-        if p != "GLEP_":
-            raise ValueError(
-                f"Can not calculate d_phase_d_GLEP with respect to {param}."
-            )
-        par_GLEP = getattr(self, param)
+        tbl, p, ids, idv, dt, affected, par_GLEP, dpdGLEP = self.deriv_prep(
+            toas, param, delay, "GLEP"
+        )
         glf0 = getattr(self, f"GLF0_{ids}").quantity
         glf1 = getattr(self, f"GLF1_{ids}").quantity
         glf2 = getattr(self, f"GLF2_{ids}").quantity
         glf0d = getattr(self, f"GLF0D_{ids}").quantity
         tau = getattr(self, f"GLTD_{ids}").quantity
-        dpdGLEP = np.zeros(len(tbl), dtype=np.longdouble) / par_GLEP.units
         dpdGLEP[affected] += (
             -glf0 + -glf1 * dt[affected] + -0.5 * glf2 * dt[affected] ** 2
         )
         if tau.value != 0.0:
-            dpdGLEP[affected] += -glf0d / np.exp(dt[affected] / tau)
+            dpdGLEP[affected] -= glf0d * np.exp(-dt[affected] / tau)
         return dpdGLEP

tests/test_glitch.py

@@ -1,24 +1,69 @@
+from io import StringIO
 import os
 import pytest
 
 import astropy.units as u
 import numpy as np
 import pytest
 
+from pint.exceptions import MissingParameter
 import pint.fitter
-import pint.models
+from pint.models import get_model
 import pint.residuals
 import pint.toa
 from pinttestdata import datadir
 
 parfile = os.path.join(datadir, "prefixtest.par")
 timfile = os.path.join(datadir, "prefixtest.tim")
 
+basepar = """
+    PSRJ           J0835-4510
+    RAJ            08:35:20.61149
+    DECJ           -45:10:34.8751
+    F0             11.18965156782
+    PEPOCH         55305                       
+    DM             67.99
+    UNITS          TDB
+"""
+
+good = """
+    GLEP_1         55555
+    GLPH_1         0
+    GLF0_1         1.0e-6
+    GLF1_1         -1.0e-12
+    GLF0D_1        1.0e-6
+    GLTD_1         10.0
+"""
+
+# no exponential decay set
+bad1 = """
+    GLEP_1         55555
+    GLF0_1         1.0e-6
+    GLF1_1         -1.0e-12
+    GLF0D_1        1.0e-6
+    GLTD_1         0.0
+"""
+
+# no epoch set
+bad2 = """
+    GLPH_1         0
+    GLF0_1         1.0e-6
+    GLF1_1         -1.0e-12
+"""
+
+# fitting both epoch and glitch phase
+bad3 = """
+    GLEP_1         55555 1 0
+    GLPH_1         0 1 0
+    GLF0_1         1.0e-6
+    GLF1_1         -1.0e-12
+"""
+
 
 class TestGlitch:
     @classmethod
     def setup_class(cls):
-        cls.m = pint.models.get_model(parfile)
+        cls.m = get_model(parfile)
         cls.t = pint.toa.get_TOAs(timfile, ephem="DE405", include_bipm=False)
         cls.f = pint.fitter.WLSFitter(cls.t, cls.m)
 
@@ -37,11 +82,11 @@ def test_glitch_der(self):
         for pf in self.m.glitch_prop:
             for idx in set(self.m.glitch_indices):
                 if pf in ["GLF0D_", "GLTD_"]:
-                    getattr(self.m, "GLF0D_%d" % idx).value = 1.0
-                    getattr(self.m, "GLTD_%d" % idx).value = 100
+                    getattr(self.m, f"GLF0D_{idx}").value = 1.0
+                    getattr(self.m, f"GLTD_{idx}").value = 100
                 else:
-                    getattr(self.m, "GLF0D_%d" % idx).value = 0.0
-                    getattr(self.m, "GLTD_%d" % idx).value = 0.0
+                    getattr(self.m, f"GLF0D_{idx}").value = 0.0
+                    getattr(self.m, f"GLTD_{idx}").value = 0.0
                 param = pf + str(idx)
                 adf = self.m.d_phase_d_param(self.t, delay, param)
                 param_obj = getattr(self.m, param)
@@ -54,3 +99,13 @@ def test_glitch_der(self):
                 errormsg = f"Derivatives for {param} is not accurate, max relative difference is"
                 errormsg += " %lf" % np.nanmax(np.abs(r_diff.value))
                 assert np.nanmax(np.abs(r_diff.value)) < 1e-3, errormsg
+
+    def test_bad_input(self):
+        get_model(StringIO(basepar + good))
+        with pytest.raises(MissingParameter):
+            get_model(StringIO(basepar + bad1))
+        with pytest.raises(MissingParameter):
+            m = get_model(StringIO(basepar + bad2))
+            print(m)
+        with pytest.raises(ValueError):
+            get_model(StringIO(basepar + bad3))