Repository: PyFstat/PyFstat

Issue:
pycuda F-stat maps: add condition number checks?
LALSuite has had a check on the condition number of the antenna pattern matrix since 2018 ([code link](https://lscsoft.docs.ligo.org/lalsuite/lalpulsar/group___l_a_l_compute_a_m__h.html#gad7ba2b601622946ba930216477c832b4); bonus historical reference, though probably not public-readable: https://git.ligo.org/CW/software/lalsuite/-/issues/13 ) and returns a default value of 2F=4 if that check fails. My pycuda implementation does not have that check and hence can in rare cases (especially for very short-duration single-detector calculations, but even then only for specific input parameter combinations) produce larger values than expected.

Discussion:

Patch:
diff --git a/pyfstat/pyCUDAkernels/cudaTransientFstatExpWindow.cu b/pyfstat/pyCUDAkernels/cudaTransientFstatExpWindow.cu
index 85d9972..092771d 100644
--- a/pyfstat/pyCUDAkernels/cudaTransientFstatExpWindow.cu
+++ b/pyfstat/pyCUDAkernels/cudaTransientFstatExpWindow.cu
@@ -101,21 +101,32 @@ __global__ void cudaTransientFstatExpWindow ( float *input,
 
     }
 
-    /* get determinant */
-    float Dd = ( Ad * Bd - Cd * Cd );
-    float DdInv = 0.0f;
-    /* safety catch as in XLALWeightMultiAMCoeffs():
-     * make it so that in the end F=0 instead of -nan
+    /* get inverse antenna pattern determinant,
+     * following safety checks from
+     * XLALComputeAntennaPatternSqrtDeterminant()
+     * and estimateAntennaPatternConditionNumber()
      */
-    if ( Dd > 0.0 ) {
-      DdInv  = 1.0 / Dd;
+    float sumAB  = Ad + Bd;
+    float diffAB = Ad - Bd;
+    float disc   = sqrt ( diffAB*diffAB + 4.0 * Cd*Cd );
+    float denom = sumAB - disc;
+    float cond = (denom > 0) ? ((sumAB + disc) / denom) : INFINITY;
+    float DdInv = 0.0f;
+    if ( cond < 1e4 ) {
+      DdInv = 1.0 / ( Ad * Bd - Cd * Cd );
     }
 
-    /* from XLALComputeFstatFromFaFb */
-    float F  = DdInv * (  Bd * ( Fa_re*Fa_re + Fa_im*Fa_im )
-                        + Ad * ( Fb_re*Fb_re + Fb_im*Fb_im )
-                        - 2.0 * Cd * ( Fa_re * Fb_re + Fa_im * Fb_im )
-                       );
+    /* matching compute_fstat_from_fa_fb
+     * including default fallback = 0.5*E[2F] in noise
+     * when DdInv == 0 due to ill-conditionness of M_munu
+     */
+    float F = 2;
+    if ( DdInv > 0 ) {
+      F  = DdInv * (   Bd * ( Fa_re*Fa_re + Fa_im*Fa_im )
+                     + Ad * ( Fb_re*Fb_re + Fb_im*Fb_im )
+                     - 2.0 * Cd * ( Fa_re * Fb_re + Fa_im * Fb_im )
+                   );
+    }
 
     /* store result in Fstat-matrix
      * at unraveled index of element {m,n}
diff --git a/pyfstat/pyCUDAkernels/cudaTransientFstatRectWindow.cu b/pyfstat/pyCUDAkernels/cudaTransientFstatRectWindow.cu
index dd353bb..e451da0 100644
--- a/pyfstat/pyCUDAkernels/cudaTransientFstatRectWindow.cu
+++ b/pyfstat/pyCUDAkernels/cudaTransientFstatRectWindow.cu
@@ -90,21 +90,32 @@ __global__ void cudaTransientFstatRectWindow ( float *input,
         i_t1_last = i_t1 + 1;
       }
 
-      /* get determinant */
-      float Dd = ( Ad * Bd - Cd * Cd );
-      float DdInv = 0.0f;
-      /* safety catch as in XLALWeightMultiAMCoeffs():
-       * make it so that in the end F=0 instead of -nan
+      /* get inverse antenna pattern determinant,
+       * following safety checks from
+       * XLALComputeAntennaPatternSqrtDeterminant()
+       * and estimateAntennaPatternConditionNumber()
        */
-      if ( Dd > 0.0 ) {
-        DdInv  = 1.0 / Dd;
+      float sumAB  = Ad + Bd;
+      float diffAB = Ad - Bd;
+      float disc   = sqrt ( diffAB*diffAB + 4.0 * Cd*Cd );
+      float denom = sumAB - disc;
+      float cond = (denom > 0) ? ((sumAB + disc) / denom) : INFINITY;
+      float DdInv = 0.0f;
+      if ( cond < 1e4 ) {
+        DdInv = 1.0 / ( Ad * Bd - Cd * Cd );
       }
 
-      /* from XLALComputeFstatFromFaFb */
-      float F  = DdInv * (  Bd * ( Fa_re*Fa_re + Fa_im*Fa_im )
-                          + Ad * ( Fb_re*Fb_re + Fb_im*Fb_im )
-                          - 2.0 * Cd * ( Fa_re * Fb_re + Fa_im * Fb_im )
-                         );
+      /* matching compute_fstat_from_fa_fb
+       * including default fallback = 0.5*E[2F] in noise
+       * when DdInv == 0 due to ill-conditionness of M_munu
+       */
+      float F = 2;
+      if ( DdInv > 0 ) {
+        F  = DdInv * (   Bd * ( Fa_re*Fa_re + Fa_im*Fa_im )
+                       + Ad * ( Fb_re*Fb_re + Fb_im*Fb_im )
+                       - 2.0 * Cd * ( Fa_re * Fb_re + Fa_im * Fb_im )
+                     );
+      }
 
       /* store result in Fstat-matrix
        * at unraveled index of element {m,n}
diff --git a/pyfstat/tcw_fstat_map_funcs.py b/pyfstat/tcw_fstat_map_funcs.py
index 4e42325..36afcc0 100644
--- a/pyfstat/tcw_fstat_map_funcs.py
+++ b/pyfstat/tcw_fstat_map_funcs.py
@@ -266,6 +266,16 @@ def _optional_imports_pycuda():
     )
 
 
+def _get_transient_fstat_map_features():
+    """Helper function to check available features."""
+    features = {}
+    have_lal = _optional_import("lal")
+    have_lalpulsar = _optional_import("lalpulsar")
+    features["lal"] = have_lal and have_lalpulsar
+    features["pycuda"] = _optional_imports_pycuda()
+    return features
+
+
 def init_transient_fstat_map_features(feature="lal", cudaDeviceName=None):
     """Initialization of available modules (or 'features') for computing transient F-stat maps.
 
@@ -296,11 +306,7 @@ def init_transient_fstat_map_features(feature="lal", cudaDeviceName=None):
         A CUDA device context object, if assigned.
     """
 
-    features = {}
-    have_lal = _optional_import("lal")
-    have_lalpulsar = _optional_import("lalpulsar")
-    features["lal"] = have_lal and have_lalpulsar
-    features["pycuda"] = _optional_imports_pycuda()
+    features = _get_transient_fstat_map_features()
     logger.debug("Got the following features for transient F-stat maps:")
     logger.debug(features)
 


Test patch:
diff --git a/tests/test_tcw_fstat_map_funcs.py b/tests/test_tcw_fstat_map_funcs.py
new file mode 100644
index 0000000..68616cf
--- /dev/null
+++ b/tests/test_tcw_fstat_map_funcs.py
@@ -0,0 +1,116 @@
+import logging
+import os
+import shutil
+
+import lalpulsar
+import numpy as np
+import pytest
+
+import pyfstat
+
+
+@pytest.mark.parametrize("snr", [0, 10])
+@pytest.mark.parametrize("window", ["rect", "exp"])
+@pytest.mark.parametrize("tCWFstatMapVersion", ["lal", "pycuda"])
+def test_compute_transient_fstat_map(tCWFstatMapVersion, window, snr):
+    logging.info("Initialising transient FstatMap features...")
+    features = pyfstat.tcw_fstat_map_funcs._get_transient_fstat_map_features()
+    if tCWFstatMapVersion == "pycuda" and not features[tCWFstatMapVersion]:
+        pytest.skip(f"Feature {tCWFstatMapVersion} not available.")
+    (
+        tCWFstatMapFeatures,
+        gpu_context,
+    ) = pyfstat.tcw_fstat_map_funcs.init_transient_fstat_map_features(
+        tCWFstatMapVersion
+    )
+
+    outdir = "TestData"
+    # ensure a clean working directory
+    if os.path.isdir(outdir):
+        try:
+            shutil.rmtree(outdir)
+        except OSError:
+            logging.warning(f"{outdir} not removed prior to tests.")
+    os.makedirs(outdir, exist_ok=True)
+
+    Tstart = 700000000
+    Tsft = 1800
+    day = 86400
+    duration = day
+    windowRange = lalpulsar.transientWindowRange_t()
+    windowRange.type = 1
+    windowRange.t0 = Tstart
+    windowRange.t0Band = duration - 2 * Tsft
+    windowRange.dt0 = Tsft
+    windowRange.tau = 2 * Tsft
+    windowRange.tauBand = duration - 2 * Tsft
+    windowRange.dtau = Tsft
+
+    logging.info("Creating synthetic atoms...")
+    statsfile = os.path.join(outdir, "synthTS_H1L_stats1.dat")
+    atomsfile = os.path.join(outdir, "synthTS_H1L_atoms")
+    pyfstat.utils.run_commandline(
+        f"lalpulsar_synthesizeTransientStats --fixedSNR={snr} --IFOs=H1 --dataStartGPS {Tstart} --dataDuration {duration} --injectWindow-type={window} --injectWindow-tauDays={0.25} --injectWindow-tauDaysBand=0 --injectWindow-t0Days={0.25} --injectWindow-t0DaysBand=0 --searchWindow-type={window} --searchWindow-tauDays={windowRange.tau/day} --searchWindow-tauDaysBand={windowRange.tauBand/day} --searchWindow-t0Days={0} --searchWindow-t0DaysBand={windowRange.t0Band/day} --computeFtotal=TRUE --numDraws=1 --randSeed=1 --outputStats={statsfile} --outputAtoms={atomsfile}"
+    )
+    atomsfile += "_0001_of_0001.dat"
+
+    logging.info(f"Loading atoms from {atomsfile} ...")
+    atoms_in = pyfstat.utils.read_txt_file_with_header(
+        atomsfile,
+        comments="%%",
+    )
+    multiFatoms = lalpulsar.CreateMultiFstatAtomVector(1)
+    multiFatoms.data[0] = lalpulsar.CreateFstatAtomVector(len(atoms_in))
+    multiFatoms.data[0].TAtom = Tsft
+    for ts in range(0, len(atoms_in)):
+        multiFatoms.data[0].data[ts].timestamp = int(atoms_in[ts][0])
+        multiFatoms.data[0].data[ts].a2_alpha = float(atoms_in[ts][1])
+        multiFatoms.data[0].data[ts].b2_alpha = float(atoms_in[ts][2])
+        multiFatoms.data[0].data[ts].ab_alpha = float(atoms_in[ts][3])
+        multiFatoms.data[0].data[ts].Fa_alpha = float(atoms_in[ts][4]) + 1j * float(
+            atoms_in[ts][5]
+        )
+        multiFatoms.data[0].data[ts].Fb_alpha = float(atoms_in[ts][6]) + 1j * float(
+            atoms_in[ts][7]
+        )
+
+    logging.info("Computing transient FtatMap...")
+    (
+        FstatMap,
+        timingFstatMap,
+    ) = pyfstat.tcw_fstat_map_funcs.call_compute_transient_fstat_map(
+        version=tCWFstatMapVersion,
+        features=tCWFstatMapFeatures,
+        multiFstatAtoms=multiFatoms,
+        windowRange=windowRange,
+        BtSG=True,
+    )
+    assert not np.isnan(FstatMap.maxF)
+    assert not np.isnan(FstatMap.lnBtSG)
+
+    logging.info(f"Loading {statsfile} for comparison...")
+    stats = pyfstat.utils.read_txt_file_with_header(
+        statsfile,
+        comments="%%",
+    )
+    # Keitel&Ashton 2018: differences for exp can reach ~10% due to lalpulsar's lookup table
+    reltol = 0.15 if window == "exp" else 0.01
+    assert pytest.approx(2 * FstatMap.maxF, rel=reltol) == stats["maxTwoF"]
+    assert pytest.approx(FstatMap.lnBtSG, rel=reltol) == stats["logBstat"]
+    if snr > 0:
+        assert (
+            pytest.approx(FstatMap.t0_ML, abs=4 * Tsft)
+            == stats["t0_MLd"] * day + Tstart
+        )
+        assert pytest.approx(FstatMap.tau_ML, abs=4 * Tsft) == stats["tau_MLd"] * day
+    if window == "rect":
+        # first t0 / last tau entry of F_mn should correspond to total F-stat
+        # (which for historical reasons is hacked into "fkdot3" column of synth stats file)
+        assert pytest.approx(2 * FstatMap.F_mn[0, -1], rel=reltol) == stats["fkdot3"]
+
+    # cleanup
+    if gpu_context:
+        logging.info("Detaching GPU context...")
+        gpu_context.detach()
+    if os.path.isdir(outdir):
+        shutil.rmtree(outdir)