WIP: persist pickles for TOAs loaded during a test run

conftest.py

@@ -1,8 +1,12 @@
 import os
+import shutil
+import tempfile
+from pathlib import Path
+
 import hypothesis
 import pytest
-from astropy.utils.data import check_download_cache
 from astropy.config import paths
+from astropy.utils.data import check_download_cache
 
 # This setup is drawn from Astropy and might not be entirely relevant to us;
 # in particular we don't have a cron run for slow tests.
@@ -42,3 +46,9 @@ def temp_cache(tmpdir):
     with paths.set_temp_cache(tmpdir):
         yield None
         check_download_cache()
+
+
+@pytest.fixture(scope="session")
+def pickle_dir():
+    with tempfile.TemporaryDirectory(prefix="pytest-pickles-") as d:
+        yield Path(d)

src/pint/models/binary_dd.py

@@ -84,16 +84,16 @@ def validate(self):
         # If any *DOT is set, we need T0
         for p in ("PBDOT", "OMDOT", "EDOT", "A1DOT"):
             if getattr(self, p).value is None:
-                getattr(self, p).set("0")
+                getattr(self, p).value = 0
                 getattr(self, p).frozen = True
 
         if self.GAMMA.value is None:
-            self.GAMMA.set("0")
+            self.GAMMA.value = 0
             self.GAMMA.frozen = True
 
         # If eccentricity is zero, freeze some parameters to 0
         # OM = 0 -> T0 = TASC
         if self.ECC.value == 0 or self.ECC.value is None:
             for p in ("ECC", "OM", "OMDOT", "EDOT"):
-                getattr(self, p).set("0")
+                getattr(self, p).value = 0
                 getattr(self, p).frozen = True

src/pint/models/noise_model.py

@@ -160,14 +160,18 @@ def scale_toa_sigma(self, toas):
             if np.any(mask):
                 sigma_scaled[mask] = np.hypot(sigma_scaled[mask], equad.quantity)
             else:
+                # Is this really a problem?
                 warnings.warn(f"EQUAD {equad} has no TOAs")
+                pass
         for efac_name in self.EFACs:
             efac = getattr(self, efac_name)
             mask = efac.select_toa_mask(toas)
             if np.any(mask):
                 sigma_scaled[mask] *= efac.quantity
             else:
+                # Is this really a problem?
                 warnings.warn(f"EFAC {efac} has no TOAs")
+                pass
         return sigma_scaled
 
     def sigma_scaled_cov_matrix(self, toas):

src/pint/pint_matrix.py

@@ -227,7 +227,7 @@ def get_label_along_axis(self, axis, label_name):
         DEPRECATED - use get_label(..., axis=axis)
         """
         warn(
-            "This method is deprecated. Use `parameter_covariance_matrix.get_label(..., axis=axis)[0]` instead of `parameter_covariance_matrix.get_label_along_axis(...)",
+            "This method is deprecated. Use `parameter_covariance_matrix.get_label(..., axis=axis)[0]` instead of `parameter_covariance_matrix.get_label_along_axis(axis, ...)",
             category=DeprecationWarning,
         )
         label_in_one_axis = self.axis_labels[axis]
@@ -615,7 +615,7 @@ def combine_design_matrices_by_param(matrix1, matrix2, padding=0.0):
     append_offset = matrix1.shape[0]
     base_matrix = matrix1.matrix
     for d_quantity in matrix2.derivative_quantity:
-        quantity_label = matrix2.get_label_along_axis(0, d_quantity)
+        quantity_label = matrix2.get_label(d_quantity, axis=0)[0]
         if d_quantity in base_quantity_index.keys():
             # Check quantity size
             d_quantity_size = quantity_label[3] - quantity_label[2]
@@ -658,7 +658,7 @@ def combine_design_matrices_by_param(matrix1, matrix2, padding=0.0):
     new_matrix.fill(padding)
     # Fill up the new_matrix with matrix2
     for quantity, new_idx in new_quantity_index.items():
-        old_idx = matrix2.get_label_along_axis(0, d_quantity)[2:4]
+        old_idx = matrix2.get_label(d_quantity, axis=0)[0][2:4]
         new_matrix[new_idx[0] : new_idx[1], :] = matrix2.matrix[
             old_idx[0] : old_idx[1], :
         ]

src/pint/polycos.py

@@ -502,46 +502,35 @@ def generate_polycos(
         method="TEMPO",
         numNodes=20,
     ):
-        """
-        Generate the polyco data.
+        """Generate the polyco data.
+
+        This function replaces anything in ``self`` with a freshly
+        generated polyco table.
 
         Parameters
         ----------
-        model : TimingModel
+        model : pint.models.TimingModel
             TimingModel to generate the Polycos with parameters
             setup.
-
-        mjdStart : float / numpy longdouble
+        mjdStart : float or np.longdouble
             Start time of polycos in mjd
-
-        mjdEnd : float / numpy longdouble
+        mjdEnd : float or np.longdouble
             Ending time of polycos in mjd
-
         obs : str
             Observatory code
-
         segLength : int
             Length of polyco segement [minutes]
-
         ncoeff : int
             Number of coefficents
-
         obsFreq : float
             Observing frequency [MHz]
-
         maxha : float optional. Default 12.0
             Maximum hour angle. Only 12.0 is supported for now.
-
         method : string optional ["TEMPO", "TEMPO2", ...] Default TEMPO
             Method to generate polycos. Only the TEMPO method is supported for now.
-
         numNodes : int optional. Default 20
             Number of nodes for fitting. It cannot be less then the number of
             coefficents.
-
-        Return
-        ---------
-        A polyco table.
         """
         mjdStart = data2longdouble(mjdStart)
         mjdEnd = data2longdouble(mjdEnd)

src/pint/toa.py

@@ -110,7 +110,7 @@ def get_TOAs(
     include_gps=None,
     planets=None,
     model=None,
-    usepickle=False,
+    usepickle=None,
     tdb_method="default",
     picklefilename=None,
     limits="warn",
@@ -166,12 +166,14 @@ def get_TOAs(
         If a valid timing model is passed, model commands (such as BIPM version,
         planet shapiro delay, and solar system ephemeris) that affect TOA loading
         are applied.
-    usepickle : bool
-        Whether to try to use pickle-based caching of loaded clock-corrected TOAs objects.
+    usepickle : bool or None
+        Whether to try to use pickle-based caching of loaded clock-corrected
+        TOAs objects. If None, setting picklefilename will trigger use of
+        pickles; to suppress this supply False.
     tdb_method : str
         Which method to use for the clock correction to TDB. See
         :func:`pint.observatory.Observatory.get_TDBs` for details.
-    picklefilename : str or None
+    picklefilename : Path or str or None
         Filename to use for caching loaded file. Defaults to adding ``.pickle.gz`` to the
         filename of the timfile, if there is one and only one. If no filename is available,
         or multiple filenames are provided, a specific filename must be provided.
@@ -222,6 +224,8 @@ def get_TOAs(
 
     updatepickle = False
     recalc = False
+    if usepickle is None:
+        usepickle = picklefilename is not None
     if usepickle:
         try:
             t = load_pickle(timfile, picklefilename=picklefilename)
@@ -322,7 +326,7 @@ def load_pickle(toafilename, picklefilename=None):
     toafilename : str
         Base filename of the TOAs; pickles will be searched for with
         ".pickle.gz", ".pickle", or just this filename.
-    picklefilename : str, optional
+    picklefilename : str or Path, optional
         Explicit filename to use.
 
     Returns
@@ -334,11 +338,17 @@ def load_pickle(toafilename, picklefilename=None):
     IOError
         If no pickle is found.
     """
-    picklefilenames = (
-        [toafilename + ext for ext in (".pickle.gz", ".pickle", "")]
-        if picklefilename is None
-        else [picklefilename]
-    )
+    if picklefilename is None:
+        picklefilenames = [toafilename + ext for ext in (".pickle.gz", ".pickle", "")]
+    else:
+        picklefilename = Path(picklefilename)
+        if picklefilename.is_dir():
+            picklefilenames = [
+                picklefilename / (Path(toafilename).name + ext)
+                for ext in (".pickle.gz", ".pickle", "")
+            ]
+        else:
+            picklefilenames = [picklefilename]
 
     lf = None
     for fn in picklefilenames:
@@ -372,6 +382,12 @@ def save_pickle(toas, picklefilename=None):
     """
     # Save the PINT version used to create this pickle file
     toas.pintversion = pint.__version__
+    pickle_dir = None
+    if picklefilename is not None:
+        picklefilename = Path(picklefilename)
+        if picklefilename.is_dir():
+            pickle_dir = picklefilename
+            picklefilename = None
     if picklefilename is not None:
         pass
     elif toas.merged:
@@ -382,9 +398,11 @@ def save_pickle(toas, picklefilename=None):
         if isinstance(toas.filename, (str, Path)):
             picklefilename = str(toas.filename) + ".pickle.gz"
         else:
-            picklefilename = toas.filename[0] + ".pickle.gz"
+            picklefilename = str(toas.filename[0]) + ".pickle.gz"
     else:
         raise ValueError("TOA pickle method needs a (single) filename.")
+    if pickle_dir is not None:
+        picklefilename = pickle_dir / Path(picklefilename).name
     with gzip.open(picklefilename, "wb") as f:
         pickle.dump(toas, f)
 

tests/datafile/get_tempo_result.py

@@ -1,6 +1,8 @@
+import argparse
+
 import tempo_utils as t1u
+
 from pint.utils import longdouble2str
-import argparse
 
 
 def get_tempo_result(parfile, timfile):

tests/datafile/make_J1713_libstempo.py

@@ -1,6 +1,7 @@
+import datetime
+
 import libstempo as T
 import numpy as np
-import datetime
 
 timJ1713 = "J1713+0747_NANOGrav_11yv0_short.tim"
 # one file in ICRS, one in ECL