Merge pull request #1634 from dlakaplan/addchi2

Include CHI2 in par output

CHANGELOG-unreleased.md

@@ -13,6 +13,7 @@ the released changes.
 - Moved design matrix normalization code from `pint.fitter` to the new `pint.utils.normalize_designmatrix()` function.
 - Made `Residuals` independent of `GLSFitter` (GLS chi2 is now computed using the new function `Residuals._calc_gls_chi2()`).
 ### Added
+- CHI2, CHI2R, TRES, DMRES now in postfit par files
 - Added `WaveX` model as a `DelayComponent` with Fourier amplitudes as fitted parameters
 - `Parameter.as_latex` method for latex representation of a parameter.
 - `pint.output.publish` module and `pintpublish` script for generating publication (LaTeX) output.

src/pint/fitter.py

@@ -415,24 +415,30 @@ def get_summary(self, nodmx=False):
                         pn, prefitpar.str_quantity(prefitpar.value), "", par.units
                     )
                 elif par.frozen:
-                    if (
-                        par.name == "START"
-                        and prefitpar.value is None
-                        or par.name != "START"
-                        and par.name == "FINISH"
+                    if par.name in ["START", "FINISH"] and prefitpar.value is None:
+                        s += ("{:" + spacingName + "s} {:20s} {:28g} {} \n").format(
+                            pn, " ", par.value, par.units
+                        )
+                    elif par.name in ["START", "FINISH"]:
+                        s += ("{:" + spacingName + "s} {:20g} {:28g} {} \n").format(
+                            pn, prefitpar.value, par.value, par.units
+                        )
+                    elif (
+                        par.name in ["CHI2", "CHI2R", "TRES", "DMRES"]
                         and prefitpar.value is None
                     ):
                         s += ("{:" + spacingName + "s} {:20s} {:28g} {} \n").format(
                             pn, " ", par.value, par.units
                         )
-                    elif par.name in ["START", "FINISH"]:
+                    elif par.name in ["CHI2", "CHI2R", "TRES", "DMRES"]:
                         s += ("{:" + spacingName + "s} {:20g} {:28g} {} \n").format(
                             pn, prefitpar.value, par.value, par.units
                         )
                     else:
                         s += ("{:" + spacingName + "s} {:20g} {:28s} {} \n").format(
                             pn, prefitpar.value, "", par.units
                         )
+
                 else:
                     # s += "{:14s} {:20g} {:20g} {:20.2g} {} \n".format(
                     #     pn,
@@ -698,6 +704,15 @@ def update_model(self, chi2=None):
             if self.toas.clock_corr_info["include_bipm"]
             else "TT(TAI)"
         )
+        if chi2 is not None:
+            # assume a fit has been done
+            self.model.CHI2.value = chi2
+            self.model.CHI2R.value = chi2 / self.resids.dof
+            if not self.is_wideband:
+                self.model.TRES.quantity = self.resids.rms_weighted()
+            else:
+                self.model.TRES.quantity = self.resids.rms_weighted()["toa"]
+                self.model.DMRES.quantity = self.resids.rms_weighted()["dm"]
 
     def reset_model(self):
         """Reset the current model to the initial model."""

src/pint/models/parameter.py

@@ -463,8 +463,8 @@ def as_parfile_line(self, format="pint"):
         name = self.name if self.use_alias is None else self.use_alias
 
         # special cases for parameter names that change depending on format
-        if self.name == "CHI2" and format.lower() != "pint":
-            # no CHI2 for TEMPO/TEMPO2
+        if self.name in ["DMRES"] and format.lower() not in ["pint"]:
+            # DMRES only for PINT
             return ""
         elif self.name == "SWM" and format.lower() != "pint":
             # no SWM for TEMPO/TEMPO2

src/pint/models/timing_model.py

@@ -89,13 +89,13 @@
 #
 # Comparisons with keywords in par file lines is done in a case insensitive way.
 ignore_params = {
-    "TRES",
+    #    "TRES",
     "TZRMJD",
     "TZRFRQ",
     "TZRSITE",
     "NITS",
     "IBOOT",
-    "CHI2R",
+    #    "CHI2R",
     "MODE",
     "PLANET_SHAPIRO2",
 }
@@ -339,13 +339,36 @@ def __init__(self, name="", components=[]):
         self.add_param_from_top(
             floatParameter(
                 name="CHI2",
-                value=0.0,
                 units="",
                 description="Chi-squared value obtained during fitting",
             ),
             "",
         )
+        self.add_param_from_top(
+            floatParameter(
+                name="CHI2R",
+                units="",
+                description="Reduced chi-squared value obtained during fitting",
+            ),
+            "",
+        )
 
+        self.add_param_from_top(
+            floatParameter(
+                name="TRES",
+                units=u.us,
+                description="TOA residual after fitting",
+            ),
+            "",
+        )
+        self.add_param_from_top(
+            floatParameter(
+                name="DMRES",
+                units=u.pc / u.cm**3,
+                description="DM residual after fitting (wideband only)",
+            ),
+            "",
+        )
         for cp in components:
             self.add_component(cp, setup=False, validate=False)
 
@@ -2286,7 +2309,7 @@ def compare(
                         else:
                             value2[pn] = str(otherpar.value)
                         if otherpar.value != par.value:
-                            if par.name in ["START", "FINISH", "CHI2", "NTOA"]:
+                            if par.name in ["START", "FINISH", "CHI2", "CHI2R", "NTOA"]:
                                 if verbosity == "max":
                                     log.info(
                                         "Parameter %s has changed between these models"

tests/test_chi2inpar.py

@@ -0,0 +1,75 @@
+from astropy import units as u
+import pytest
+from pint.models import get_model_and_toas, get_model
+import pint.fitter
+import pint.simulation
+from pinttestdata import datadir
+import os
+import io
+
+
+@pytest.fixture
+def wb():
+    m = get_model(os.path.join(datadir, "NGC6440E.par"))
+    t = pint.simulation.make_fake_toas_uniform(
+        55000, 58000, 20, model=m, freq=1400 * u.MHz, wideband=True, add_noise=True
+    )
+
+    return m, t
+
+
+@pytest.fixture
+def nb():
+    m = get_model(os.path.join(datadir, "NGC6440E.par"))
+    t = pint.simulation.make_fake_toas_uniform(
+        55000,
+        58000,
+        20,
+        model=m,
+        freq=1400 * u.MHz,
+        wideband=False,
+        add_noise=True,
+    )
+
+    return m, t
+
+
+@pytest.mark.parametrize(
+    "ft",
+    [
+        pint.fitter.WLSFitter,
+        pint.fitter.GLSFitter,
+        pint.fitter.DownhillWLSFitter,
+        pint.fitter.DownhillGLSFitter,
+    ],
+)
+def test_nbfitters(ft, nb):
+    m, t = nb
+    f = ft(t, m)
+    f.fit_toas()
+    for p in ["CHI2", "CHI2R", "TRES"]:
+        matched_line = [
+            line for line in f.model.as_parfile().split("\n") if line.startswith(p)
+        ]
+        assert matched_line
+        assert float(matched_line[0].split()[-1]) > 0
+    assert not [
+        line for line in f.model.as_parfile().split("\n") if line.startswith("DMRES")
+    ]
+    m2 = get_model(io.StringIO(f.model.as_parfile()))
+
+
+@pytest.mark.parametrize(
+    "ft", [pint.fitter.WidebandTOAFitter, pint.fitter.WidebandDownhillFitter]
+)
+def test_wbfitters(ft, wb):
+    m, t = wb
+    f = ft(t, m)
+    f.fit_toas()
+    for p in ["CHI2", "CHI2R", "TRES", "DMRES"]:
+        matched_line = [
+            line for line in f.model.as_parfile().split("\n") if line.startswith(p)
+        ]
+        assert matched_line
+        assert float(matched_line[0].split()[-1]) > 0
+    m2 = get_model(io.StringIO(f.model.as_parfile()))

tests/test_parfile_writing_format.py

@@ -29,9 +29,10 @@ def test_CHI2():
     )
 
     f = fitter.WLSFitter(toas=t, model=m)
+    f.fit_toas()
     assert "CHI2" in f.model.as_parfile()
-    assert "CHI2" not in f.model.as_parfile(format="tempo2")
-    assert "CHI2" not in f.model.as_parfile(format="tempo")
+    assert "CHI2" in f.model.as_parfile(format="tempo2")
+    assert "CHI2" in f.model.as_parfile(format="tempo")
 
 
 def test_T2CMETHOD():