Zsearch results improvements

hendrics/base.py

@@ -5,6 +5,7 @@
 
 import copy
 import os
+import re
 import sys
 import tempfile
 import urllib
@@ -28,11 +29,17 @@
 )
 
 from astropy import log
+from astropy import units as u
 from astropy.io.registry import identify_format
 from astropy.table import Table
 
 try:
     from pint.models import get_model
+    from pint.models.parameter import prefixParameter
+    from pint.models.timing_model import (
+        Component,
+        TimingModel,
+    )  # Interface for timing model
 
     HAS_PINT = True
 except (ImportError, urllib.error.URLError):
@@ -219,6 +226,36 @@ def show_progress(a):
 )
 
 
+def create_empty_timing_model(source_name="source"):
+    if not HAS_PINT:
+        raise ImportError(
+            "You need the optional dependency PINT to use this "
+            "functionality: github.com/nanograv/pint"
+        )
+    all_components = Component.component_types
+    selected_components = ["AbsPhase", "AstrometryEquatorial", "Spindown"]
+    component_instances = []
+
+    # Initiate the component instances
+    for cp_name in selected_components:
+        component_class = all_components[cp_name]  # Get the component class
+        component_instance = component_class()  # Instantiate a component object
+        component_instances.append(component_instance)
+    tm = TimingModel(source_name, component_instances)
+
+    # Add, at the very least, a F2 to the parameter
+    f2 = prefixParameter(
+        parameter_type="float",
+        name="F2",
+        value=0.0,
+        units=u.Hz / (u.s) ** 2,
+        longdouble=True,
+        tcb2tdb_scale_factor=u.Quantity(1),
+    )
+    tm.components["Spindown"].add_param(f2, setup=True)
+    return tm
+
+
 def r_in(td, r_0):
     """Calculate incident countrate given dead time and detected countrate."""
     tau = 1 / r_0
@@ -421,8 +458,6 @@ def simple_orbit_fun_from_parfile(
     correction_mjd : function
         Function that accepts times in MJDs and returns the deorbited times.
     """
-    from astropy import units
-
     if not HAS_PINT:
         raise ImportError(
             "You need the optional dependency PINT to use this "
@@ -438,7 +473,7 @@ def simple_orbit_fun_from_parfile(
 
     correction = interp1d(
         mjds,
-        (toalist.table["tdbld"] * units.d - delays).to(units.d).value,
+        (toalist.table["tdbld"] * u.d - delays).to(u.d).value,
         fill_value="extrapolate",
     )
 
@@ -490,6 +525,10 @@ def deorbit_events(events, parameter_file=None, invert=False, ephem=None):
     elif ephem is None:
         ephem = "DE421"
 
+    allnum_re = re.compile(r"\d{3}")
+    if allnum_re.match(ephem):
+        ephem = f"DE{ephem}"
+
     orbital_correction_fun = simple_orbit_fun_from_parfile(
         pepoch_mjd - 1,
         pepoch_mjd + length_d + 1,

hendrics/efsearch.py

@@ -29,11 +29,14 @@
 from astropy.utils.introspection import minversion
 
 from .base import (
+    HAS_PINT,
     HENDRICS_STAR_VALUE,
     adjust_dt_for_power_of_two,
+    create_empty_timing_model,
     deorbit_events,
     find_peaks_in_image,
     fold_detection_level,
+    get_model,
     hen_root,
     histogram,
     histogram2d,
@@ -45,10 +48,11 @@
 )
 from .fake import scramble
 from .ffa import _z_n_fast_cached, ffa_search, h_test
-from .fold import filter_energy
+from .fold import filter_energy, fold_events
 from .io import (
     HEN_FILE_EXTENSION,
     EFPeriodogram,
+    find_file_in_allowed_paths,
     get_file_type,
     load_events,
     load_folding,
@@ -1138,17 +1142,29 @@ def dyn_folding_search(
 def print_qffa_results(best_cand_table):
     newtable = copy.deepcopy(best_cand_table)
     good = ~np.isnan(newtable["pulse_amp"])
+    factor = newtable["ncounts"] / newtable["exposure"] / 100
     if len(newtable[good]) == 0:
         print("No pulsations found. Best candidate and upper limit:")
         good = 0
         newtable["Pulsed amplitude (%)"] = [f"<{a:g} (90%)" for a in newtable["pulse_amp_ul_0.9"]]
+        newtable["Pulsed amplitude (ct/s)"] = [
+            f"<{a * f:g} (90%)" for (a, f) in zip(newtable["pulse_amp_ul_0.9"], factor)
+        ]
     else:
         print("Best candidate(s):")
         newtable["Pulsed amplitude (%)"] = [
             f"{a:g} ± {e:g}" for (a, e) in zip(newtable["pulse_amp"], newtable["pulse_amp_err"])
         ]
+        newtable["Pulsed amplitude (ct/s)"] = [
+            f"{a * f:g} ± {e * f:g}"
+            for (a, e, f) in zip(newtable["pulse_amp"], newtable["pulse_amp_err"], factor)
+        ]
 
-    print(newtable["mjd", "f", "fdot", "fddot", "power", "Pulsed amplitude (%)"][good])
+    print(
+        newtable[
+            "mjd", "f", "fdot", "fddot", "power", "Pulsed amplitude (%)", "Pulsed amplitude (ct/s)"
+        ][good]
+    )
 
 
 def get_xy_boundaries_from_level(x, y, image, level, x0, y0):
@@ -1224,6 +1240,109 @@ def get_boundaries_from_level(x, y, level, x0):
     return min_x, max_x
 
 
+def get_best_solution_from_qffa(ef, best_cand_table, out_model_file=None, fold=False):
+    # Get these from the first row of the table
+    f, fdot, fddot, ul, max_stat = (
+        best_cand_table["f"][0],
+        best_cand_table["fdot"][0],
+        best_cand_table["fddot"][0],
+        best_cand_table["pulse_amp_ul_0.9"][0],
+        best_cand_table["power"][0],
+    )
+    nbin = best_cand_table.meta["nbin"]
+
+    if (filename := best_cand_table.meta["filename"]) is not None:
+        root = os.path.split(filename)[0]
+        filename = find_file_in_allowed_paths(filename, [".", root])
+        events = load_events(filename)
+        if ef.emin is not None or ef.emax is not None:
+            events, elabel = filter_energy(events, ef.emin, ef.emax)
+
+        if hasattr(ef, "parfile") and ef.parfile is not None:
+            root = os.path.split(filename)[0]
+            parfile = find_file_in_allowed_paths(ef.parfile, [".", root])
+            if not parfile:
+                warnings.warn(f"{ef.parfile} does not exist")
+            else:
+                ef.parfile = parfile
+
+            if parfile and os.path.exists(parfile):
+                events = deorbit_events(events, parfile)
+                base_model = get_model(parfile)
+        else:
+            if HAS_PINT:
+                base_model = create_empty_timing_model()
+
+        if hasattr(ef, "ref_time") and ef.ref_time is not None:
+            ref_time = ef.ref_time
+        elif hasattr(ef, "pepoch") and ef.pepoch is not None:
+            ref_time = (ef.pepoch - events.mjdref) * 86400
+        else:
+            ref_time = (events.time[0] + events.time[-1]) / 2
+
+        pepoch = ref_time / 86400 + events.mjdref
+        solution_dict = {
+            "F0": f,
+            "F1": fdot,
+            "F2": fddot,
+            "PEPOCH": pepoch,
+            "START": events.gti[0, 0] / 86400 + events.mjdref,
+            "FINISH": events.gti[-1, 1] / 86400 + events.mjdref,
+        }
+
+        if HAS_PINT and np.isnan(ul):
+            for k, v in solution_dict.items():
+                getattr(base_model, k).value = v
+            if out_model_file is None:
+                out_model_file = "best_model.par"
+            log.info(f"Writing par file to {out_model_file}")
+            base_model.write_parfile(out_model_file)
+
+        if fold:
+            phase, profile, profile_err = fold_events(
+                copy.deepcopy(events.time),
+                f,
+                fdot,
+                ref_time=ref_time,
+                # gti=copy.deepcopy(events.gti),
+                expocorr=False,
+                nbin=nbin,
+            )
+            table = Table(
+                {
+                    "phase": np.concatenate((phase, phase + 1)),
+                    "profile": np.concatenate((profile, profile)),
+                    "err": np.concatenate((profile_err, profile_err)),
+                }
+            )
+        else:
+            table = None
+        ntimes = max(8, np.rint(max_stat / 20).astype(int))
+        phascommand = (
+            f"HENphaseogram -f {solution_dict['F0']} "
+            f"--fdot {solution_dict['F1']} {ef.filename} -n {nbin} --ntimes {ntimes} --norm meansub"
+        )
+        if ef.parfile and os.path.exists(ef.parfile):
+            phascommand += f" --deorbit-par {parfile}"
+        if hasattr(ef, "emin") and ef.emin is not None:
+            phascommand += f" --emin {ef.emin}"
+        if hasattr(ef, "emin") and ef.emin is not None:
+            phascommand += f" --emax {ef.emax}"
+
+        if hasattr(events, "mjdref") and events.mjdref is not None:
+            phascommand += f" --pepoch {pepoch}"
+
+        log.info("To see the detailed phaseogram, " f"run {phascommand}")
+
+    elif not os.path.exists(ef.filename):
+        warnings.warn(ef.filename + " does not exist")
+        events = None
+        solution_dict = None
+        table = None
+
+    return events, solution_dict, table
+
+
 def _analyze_qffa_results(input_ef_periodogram, fname=None):
     """Search best candidates in a quasi-fast-folding search.
 
@@ -1276,6 +1395,8 @@ def _analyze_qffa_results(input_ef_periodogram, fname=None):
             "fname",
             "mjd",
             "power",
+            "exposure",
+            "ncounts",
             "f",
             "f_err_n",
             "f_err_p",
@@ -1298,6 +1419,8 @@ def _analyze_qffa_results(input_ef_periodogram, fname=None):
             float,
             float,
             float,
+            int,
+            float,
             float,
             float,
             float,
@@ -1316,9 +1439,10 @@ def _analyze_qffa_results(input_ef_periodogram, fname=None):
         ],
     )
     best_cand_table["power"].info.format = ".2f"
+    best_cand_table["exposure"].info.format = ".2f"
     best_cand_table["power_cl_0.9"].info.format = ".2f"
     best_cand_table["fdot"].info.format = ".2e"
-    best_cand_table["fddot"].info.format = "g"
+    best_cand_table["fddot"].info.format = ".2e"
     best_cand_table["pulse_amp_cl_0.1"].info.format = ".2f"
     best_cand_table["pulse_amp_cl_0.9"].info.format = ".2f"
     best_cand_table["pulse_amp"].info.format = ".2f"
@@ -1381,6 +1505,8 @@ def _analyze_qffa_results(input_ef_periodogram, fname=None):
                 input_ef_periodogram.filename,
                 input_ef_periodogram.pepoch,
                 max_stat,
+                input_ef_periodogram.exposure,
+                input_ef_periodogram.ncounts,
                 f,
                 fmin - f,
                 fmax - f,
@@ -1430,6 +1556,7 @@ def _analyze_qffa_results(input_ef_periodogram, fname=None):
         and os.path.exists(input_ef_periodogram.filename)
     ):
         best_cand_table.meta["filename"] = input_ef_periodogram.filename
+
     return best_cand_table
 
 
@@ -1450,6 +1577,7 @@ def analyze_qffa_results(fname):
     best_cand_table = _analyze_qffa_results(ef, fname=fname)
 
     best_cand_table.write(fname + "_best_cands.csv", overwrite=True)
+
     return ef, best_cand_table
 
 
@@ -1815,6 +1943,7 @@ def _common_main(args, func):
             mjdref=mjdref,
             pepoch=mjdref + ref_time / 86400,
             oversample=args.oversample,
+            exposure=events.exposure,
         )
         efperiodogram.upperlim = pf_upper_limit(np.max(stats), events.time.size, n=args.N)
         efperiodogram.ncounts = events.time.size
@@ -2124,7 +2253,10 @@ def main_accelsearch(args=None):
     if Nbins > 10**8:
         log.info(f"The number of bins is more than 100 millions: {Nbins}. " "Using memmap.")
 
-    dt = adjust_dt_for_power_of_two(dt, max_length)
+    if hasattr(events, "dt") and events.dt is not None:
+        dt = events.suggest_compatible_dt(dt)
+    else:
+        dt = adjust_dt_for_power_of_two(dt, max_length)
 
     if args.pad_to_double:
         times = memmapped_arange(-0.5 * max_length, 1.5 * max_length, dt)

hendrics/io.py

@@ -57,7 +57,9 @@
     cpl256 = np.dtype([("real", np.longdouble), ("imag", np.longdouble)])
 
 
-class EFPeriodogram:
+class EFPeriodogram(StingrayObject):
+    main_array_attr = "freq"
+
     def __init__(
         self,
         freq=None,
@@ -81,6 +83,7 @@ def __init__(
         emax=None,
         ncounts=None,
         upperlim=None,
+        exposure=None,
     ):
         self.freq = freq
         self.stat = stat
@@ -103,6 +106,7 @@ def __init__(
         self.mjdref = mjdref
         self.upperlim = upperlim
         self.ncounts = ncounts
+        self.exposure = exposure
 
     def find_peaks(self, conflevel=99.0):
         from .base import fold_detection_level, z2_n_detection_level