Merge branch 'nanograv:master' into cmx

CHANGELOG-unreleased.md

@@ -9,7 +9,10 @@ the released changes.
 
 ## Unreleased
 ### Changed
+- 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
+- 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`
 - When TCB->TDB conversion info is missing, will print parameter name
 - Piecewise-constant model for chromatic variations (CMX)
 ### Fixed

src/pint/plot_utils.py

@@ -283,7 +283,9 @@ def plot_priors(
         a, x = np.histogram(values[i], bins=bins, density=True)
         counts.append(a)
 
-    fig, axs = plt.subplots(len(keys), figsize=(8, 11), constrained_layout=True)
+    fig, axs = plt.subplots(
+        len(keys), figsize=(8, len(keys) * 1.5), constrained_layout=True
+    )
 
     for i, p in enumerate(keys):
         if i != len(keys[:-1]):

src/pint/scripts/event_optimize.py

@@ -414,6 +414,16 @@ def __init__(
             self.model, phs, phserr
         )
         self.n_fit_params = len(self.fitvals)
+        self.M, _, _ = self.model.designmatrix(self.toas)
+        self.M = self.M.transpose() * -self.model.F0.value
+        self.phases = self.get_event_phases()
+        self.linearize_model = False
+
+    def calc_phase_matrix(self, theta):
+        d_phs = np.zeros(len(self.toas))
+        for i in range(len(theta) - 1):
+            d_phs += self.M[i + 1] * (self.fitvals[i] - theta[i])
+        return (self.phases - d_phs) % 1
 
     def get_event_phases(self):
         """
@@ -446,7 +456,10 @@ def lnposterior(self, theta):
             return -np.inf, -np.inf, -np.inf
 
         # Call PINT to compute the phases
-        phases = self.get_event_phases()
+        if self.linearize_model:
+            phases = self.calc_phase_matrix(theta)
+        else:
+            phases = self.get_event_phases()
         lnlikelihood = profile_likelihood(
             theta[-1], self.xtemp, phases, self.template, self.weights
         )
@@ -686,6 +699,13 @@ def main(argv=None):
         action="store_true",
         dest="noautocorr",
     )
+    parser.add_argument(
+        "--linearize_model",
+        help="Calculates the phase at each MCMC step using the designmatrix",
+        default=False,
+        action="store_true",
+        dest="linearize_model",
+    )
 
     args = parser.parse_args(argv)
     pint.logging.setup(
@@ -862,6 +882,9 @@ def main(argv=None):
     # This way, one walker should always be in a good position
     pos[0] = ftr.fitvals
 
+    # How phase will be calculated at each step (either with the designmatrix orthe exact phase calculation)
+    ftr.linearize_model = args.linearize_model
+
     import emcee
 
     # Setting up a backend to save the chains into an h5 file
@@ -925,7 +948,7 @@ def chains_to_dict(names, sampler):
 
     def plot_chains(chain_dict, file=False):
         npts = len(chain_dict)
-        fig, axes = plt.subplots(npts, 1, sharex=True, figsize=(8, 9))
+        fig, axes = plt.subplots(npts, 1, sharex=True, figsize=(8, npts * 1.5))
         for ii, name in enumerate(chain_dict.keys()):
             axes[ii].plot(chain_dict[name], color="k", alpha=0.3)
             axes[ii].set_ylabel(name)
@@ -950,6 +973,7 @@ def plot_chains(chain_dict, file=False):
     lnprior_samps = blobs["lnprior"]
     lnlikelihood_samps = blobs["lnlikelihood"]
     lnpost_samps = lnprior_samps + lnlikelihood_samps
+    maxpost = lnpost_samps[:][burnin:].max()
     ind = np.unravel_index(
         np.argmax(lnpost_samps[:][burnin:]), lnpost_samps[:][burnin:].shape
     )
@@ -1000,8 +1024,15 @@ def plot_chains(chain_dict, file=False):
     ]
     ftr.set_param_uncertainties(dict(zip(ftr.fitkeys[:-1], errors[:-1])))
 
+    # Calculating the AIC and BIC
+    n_params = len(ftr.model.free_params)
+    AIC = 2 * (n_params - maxpost)
+    BIC = n_params * np.log(len(ts)) - 2 * maxpost
+    ftr.model.NTOA.value = ts.ntoas
     f = open(filename + "_post.par", "w")
     f.write(ftr.model.as_parfile())
+    f.write(f"\n#The AIC is {AIC}")
+    f.write(f"\n#The BIC is {BIC}")
     f.close()
 
     # Print the best MCMC values and ranges