PR #315: create combined EFAC + EQUAD noise object with tempo/tempo2 …

…convention

Looks good, harmonized on the e_e side.

enterprise/signals/signal_base.py

@@ -110,6 +110,9 @@ def param_names(self):
                 ret.append(p.name)
         return ret
 
+    def __repr__(self):
+        return "<Enterprise Signal object " + self.signal_id + "[" + ", ".join(p.name for p in self.params) + "]>"
+
     def get(self, parname, params={}):
         try:
             return params[self._params[parname].name]

enterprise/signals/white_signals.py

@@ -57,35 +57,55 @@ def efac_ndiag(toaerrs, efac=1.0):
     return efac ** 2 * toaerrs ** 2
 
 
-def MeasurementNoise(efac=parameter.Uniform(0.5, 1.5), selection=Selection(selections.no_selection), name=""):
-    """Class factory for EFAC type measurement noise."""
+@function
+def combined_ndiag(toaerrs, efac=1.0, log10_t2equad=-8):
+    return efac ** 2 * (toaerrs ** 2 + 10 ** (2 * log10_t2equad))
+
+
+def MeasurementNoise(
+    efac=parameter.Uniform(0.5, 1.5), log10_t2equad=None, selection=Selection(selections.no_selection), name="",
+):
 
-    varianceFunction = efac_ndiag(efac=efac)
+    """Class factory for EFAC+EQUAD measurement noise
+    (with tempo/tempo2/pint parameter convention, variance = efac^2 (toaerr^2 + t2equad^2)).
+    Leave out log10_t2equad to use EFAC noise only."""
+
+    varianceFunction = (
+        efac_ndiag(efac=efac) if log10_t2equad is None else combined_ndiag(efac=efac, log10_t2equad=log10_t2equad)
+    )
     BaseClass = WhiteNoise(varianceFunction, selection=selection, name=name)
 
     class MeasurementNoise(BaseClass):
-        signal_name = "efac"
-        signal_id = "efac_" + name if name else "efac"
+        signal_name = "measurement_noise"
+        signal_id = "measurement_noise_" + name if name else "measurement_noise"
 
     return MeasurementNoise
 
 
 @function
-def equad_ndiag(toas, log10_equad=-8):
-    return np.ones_like(toas) * 10 ** (2 * log10_equad)
+def tnequad_ndiag(toas, log10_tnequad=-8):
+    return np.ones_like(toas) * 10 ** (2 * log10_tnequad)
 
 
-def EquadNoise(log10_equad=parameter.Uniform(-10, -5), selection=Selection(selections.no_selection), name=""):
-    """Class factory for EQUAD type measurement noise."""
+def TNEquadNoise(log10_tnequad=parameter.Uniform(-10, -5), selection=Selection(selections.no_selection), name=""):
+    """Class factory for TNEQUAD type measurement noise (legacy, not multiplied by EFAC)."""
 
-    varianceFunction = equad_ndiag(log10_equad=log10_equad)
+    varianceFunction = tnequad_ndiag(log10_tnequad=log10_tnequad)
     BaseClass = WhiteNoise(varianceFunction, selection=selection, name=name)
 
-    class EquadNoise(BaseClass):
-        signal_name = "equad"
-        signal_id = "equad_" + name if name else "equad"
+    class TNEquadNoise(BaseClass):
+        signal_name = "tnequad"
+        signal_id = "tnequad_" + name if name else "tnequad"
+
+    return TNEquadNoise
+
 
-    return EquadNoise
+def EquadNoise(*args, **kwargs):
+    raise NotImplementedError(
+        "EquadNoise was removed in enterprise v3.3."
+        " Use MeasurementNoise to implement tempo/tempo2/pint definition [efac^2 (toaerr^2 + t2equad^2)]"
+        " or TNEquadNoise to obtain legacy enterprise definition for EQUAD only [tnequad^2]."
+    )
 
 
 def EcorrKernelNoise(

tests/test_hierarchical_parameter.py

@@ -125,15 +125,15 @@ def powerlaw(f, log10_A=-15):
         def log10(A=10 ** -16):
             return np.log10(A)
 
-        fquad = white_signals.EquadNoise(
-            log10_equad=parameter.Function(log10, A=parameter.Uniform(10 ** -17, 10 ** -14))
+        fquad = white_signals.TNEquadNoise(
+            log10_tnequad=parameter.Function(log10, A=parameter.Uniform(10 ** -17, 10 ** -14))
         )
 
         fquad1 = fquad(self.psr)
 
-        repr_A = "[B1855+09_log10_equad_A:Uniform(pmin=1e-17, pmax=1e-14)]"
-        repr_B = "[B1855+09_log10_equad_A:Uniform(pmax=1e-14, pmin=1e-17)]"
+        repr_A = "[B1855+09_log10_tnequad_A:Uniform(pmin=1e-17, pmax=1e-14)]"
+        repr_B = "[B1855+09_log10_tnequad_A:Uniform(pmax=1e-14, pmin=1e-17)]"
         assert str(fquad1.params) == repr_A or str(fquad1.params) == repr_B
         assert np.allclose(
-            fquad1.get_ndiag(params={"B1855+09_log10_equad_A": 10 ** -14})[:3], np.array([1e-28, 1e-28, 1e-28])
+            fquad1.get_ndiag(params={"B1855+09_log10_tnequad_A": 10 ** -14})[:3], np.array([1e-28, 1e-28, 1e-28])
         )

tests/test_likelihood.py

@@ -46,7 +46,7 @@ def get_noise_from_pal2(noisefile):
             par = "efac"
             flag = ln[0].split("efac-")[-1]
         elif "equad" in line:
-            par = "log10_equad"
+            par = "log10_t2equad"
             flag = ln[0].split("equad-")[-1]
         elif "jitter_q" in line:
             par = "log10_ecorr"
@@ -104,8 +104,7 @@ def compute_like(self, npsrs=1, inc_corr=False, inc_kernel=False, cholesky_spars
 
         selection = Selection(selections.by_backend)
 
-        ef = white_signals.MeasurementNoise(efac=efac, selection=selection)
-        eq = white_signals.EquadNoise(log10_equad=equad, selection=selection)
+        ms = white_signals.MeasurementNoise(efac=efac, log10_t2equad=equad, selection=selection)
         ec = white_signals.EcorrKernelNoise(log10_ecorr=ecorr, selection=selection)
 
         pl = utils.powerlaw(log10_A=log10_A, gamma=gamma)
@@ -130,9 +129,9 @@ def compute_like(self, npsrs=1, inc_corr=False, inc_kernel=False, cholesky_spars
             inc_kernel = [inc_kernel] * npsrs
 
         if inc_corr:
-            s = tm + ef + eq + ec + rn + crn
+            s = tm + ms + ec + rn + crn
         else:
-            s = tm + ef + eq + ec + rn
+            s = tm + ms + ec + rn
 
         models = []
         for ik, psr in zip(inc_kernel, psrs):
@@ -181,8 +180,9 @@ def compute_like(self, npsrs=1, inc_corr=False, inc_kernel=False, cholesky_spars
             nvec0 = np.zeros_like(psr.toas)
             for ct, flag in enumerate(flags):
                 ind = psr.backend_flags == flag
-                nvec0[ind] = efacs[ii][ct] ** 2 * psr.toaerrs[ind] ** 2
-                nvec0[ind] += 10 ** (2 * equads[ii][ct]) * np.ones(np.sum(ind))
+                nvec0[ind] = efacs[ii][ct] ** 2 * (
+                    psr.toaerrs[ind] ** 2 + 10 ** (2 * equads[ii][ct]) * np.ones(np.sum(ind))
+                )
 
             # get the basis
             bflags = psr.backend_flags

tests/test_set_parameter.py

@@ -31,7 +31,7 @@ def get_noise_from_pal2(noisefile):
             par = "efac"
             flag = ln[0].split("efac-")[-1]
         elif "equad" in line:
-            par = "log10_equad"
+            par = "log10_t2equad"
             flag = ln[0].split("equad-")[-1]
         elif "jitter_q" in line:
             par = "log10_ecorr"
@@ -78,14 +78,13 @@ def test_single_pulsar(self):
 
         selection = Selection(selections.by_backend)
 
-        ef = white_signals.MeasurementNoise(efac=efac, selection=selection)
-        eq = white_signals.EquadNoise(log10_equad=equad, selection=selection)
+        ms = white_signals.MeasurementNoise(efac=efac, log10_t2equad=equad, selection=selection)
         ec = white_signals.EcorrKernelNoise(log10_ecorr=ecorr, selection=selection)
 
         pl = utils.powerlaw(log10_A=log10_A, gamma=gamma)
         rn = gp_signals.FourierBasisGP(pl)
 
-        s = ef + eq + ec + rn
+        s = ms + ec + rn
         m = s(self.psrs[0])
 
         # set parameters
@@ -103,8 +102,9 @@ def test_single_pulsar(self):
         nvec0 = np.zeros_like(self.psrs[0].toas)
         for ct, flag in enumerate(np.unique(flags)):
             ind = flag == self.psrs[0].backend_flags
-            nvec0[ind] = efacs[ct] ** 2 * self.psrs[0].toaerrs[ind] ** 2
-            nvec0[ind] += 10 ** (2 * equads[ct]) * np.ones(np.sum(ind))
+            nvec0[ind] = efacs[ct] ** 2 * (
+                self.psrs[0].toaerrs[ind] ** 2 + 10 ** (2 * equads[ct]) * np.ones(np.sum(ind))
+            )
 
         # get the basis
         bflags = self.psrs[0].backend_flags
@@ -181,14 +181,13 @@ def test_pta(self):
 
         selection = Selection(selections.by_backend)
 
-        ef = white_signals.MeasurementNoise(efac=efac, selection=selection)
-        eq = white_signals.EquadNoise(log10_equad=equad, selection=selection)
+        ms = white_signals.MeasurementNoise(efac=efac, log10_t2equad=equad, selection=selection)
         ec = white_signals.EcorrKernelNoise(log10_ecorr=ecorr, selection=selection)
 
         pl = utils.powerlaw(log10_A=log10_A, gamma=gamma)
         rn = gp_signals.FourierBasisGP(pl)
 
-        s = ef + eq + ec + rn
+        s = ms + ec + rn
         pta = s(self.psrs[0]) + s(self.psrs[1])
 
         # set parameters
@@ -210,8 +209,9 @@ def test_pta(self):
             nvec0 = np.zeros_like(psr.toas)
             for ct, flag in enumerate(flags):
                 ind = psr.backend_flags == flag
-                nvec0[ind] = efacs[ii][ct] ** 2 * psr.toaerrs[ind] ** 2
-                nvec0[ind] += 10 ** (2 * equads[ii][ct]) * np.ones(np.sum(ind))
+                nvec0[ind] = efacs[ii][ct] ** 2 * (
+                    psr.toaerrs[ind] ** 2 + 10 ** (2 * equads[ii][ct]) * np.ones(np.sum(ind))
+                )
 
             # get the basis
             bflags = psr.backend_flags