Interpolating Qlm in time

MonteCarloMarginalizeCode/Code/RIFT/likelihood/Q_inner_product.py

@@ -52,3 +52,59 @@ def Q_inner_product_cupy(Q, A, start_indices, window_size):
     )
 
     return out
+
+
+def Q_inner_product_cubic_cupy(Q, A, start_indices, fractional_offsets, window_size):
+    """Cubic-interpolated Q inner product for fractional detector-time offsets.
+
+    ``start_indices`` are the integer floor indices of the first requested time
+    sample and ``fractional_offsets`` are the corresponding fractional parts in
+    [0, 1).  The CUDA kernel uses a four-point cubic Lagrange stencil along the
+    time axis and zero extension outside the precomputed Q buffer.
+    """
+    num_time_points, num_lms = Q.shape
+    num_extrinsic_samples, _ = A.shape
+
+    assert not cupy.isfortran(Q)
+    assert not cupy.isfortran(A)
+
+    out = cupy.empty(
+        (num_extrinsic_samples, window_size),
+        dtype=cupy.complex128,
+        order="C",
+    )
+
+    global _cuda_code
+    if _cuda_code is None:
+        path = os.path.join(os.path.dirname(__file__), 'cuda_Q_inner_product.cu')
+        if not (os.path.isfile(path)):
+            path = os.path.join(os.path.split(os.path.dirname(__file__))[0], 'cuda_Q_inner_product.cu')
+        with open(path, 'r') as f:
+            _cuda_code = f.read()
+            Q_prod_fn = cupy.RawKernel(_cuda_code, "Q_inner_cubic")
+    else:
+        Q_prod_fn = cupy.RawKernel(_cuda_code, "Q_inner_cubic")
+
+    float_prec = 16
+    # The cubic stencil uses more registers than the nearest-neighbor kernel, so
+    # use a conservative default block shape to keep launches portable across
+    # older GPUs.  Tune upward on newer cards with RIFT_Q_CUBIC_THREADS_X/Y.
+    num_threads_x = int(os.environ.get("RIFT_Q_CUBIC_THREADS_X", "4"))
+    num_threads_y = int(os.environ.get("RIFT_Q_CUBIC_THREADS_Y", "128"))
+    block_size = num_threads_x, num_threads_y, 0
+    grid_size = (
+        (num_extrinsic_samples+num_threads_x-1)//num_threads_x,
+        0,
+        0,
+    )
+    args = (
+        Q, A, start_indices, fractional_offsets, window_size,
+        num_time_points, num_extrinsic_samples, num_lms,
+        out,
+    )
+    Q_prod_fn(
+        grid_size, block_size, args,
+        shared_mem=cupy.int32(0),
+    )
+
+    return out

MonteCarloMarginalizeCode/Code/RIFT/likelihood/cuda_Q_inner_product.cu

@@ -57,4 +57,68 @@ extern "C" {
       }
     } // if
   } // Q_inner
+
+  __global__ void Q_inner_cubic(
+    const double2 * Q, const double2 * A,
+    const int * index_start,
+    const double * fractional_offset,
+    int window_size,
+    int num_time_points,
+    int num_extrinsic_samples,
+    int num_lms,
+    double2 * out
+  ){
+    size_t sample_idx = threadIdx.x + blockDim.x*blockIdx.x;
+    size_t t_idx = threadIdx.y + blockDim.y * blockIdx.y;
+
+    if (sample_idx < num_extrinsic_samples) {
+      int i_first_time = index_start[sample_idx];
+      double u = fractional_offset[sample_idx];
+      double w_m1 = -u * (u - 1.0) * (u - 2.0) / 6.0;
+      double w_0 = (u + 1.0) * (u - 1.0) * (u - 2.0) / 2.0;
+      double w_p1 = -(u + 1.0) * u * (u - 2.0) / 2.0;
+      double w_p2 = (u + 1.0) * u * (u - 1.0) / 6.0;
+
+      for (size_t i_time = t_idx; i_time < window_size; i_time+=blockDim.y) {
+        size_t i_output = sample_idx*window_size + i_time;
+        int q_time = i_first_time + (int)i_time;
+        double out_re = 0.0;
+        double out_im = 0.0;
+
+        for (size_t i_lm = 0; i_lm < num_lms; ++i_lm) {
+          double q_re = 0.0;
+          double q_im = 0.0;
+          int q_m1 = q_time - 1;
+          int q_0 = q_time;
+          int q_p1 = q_time + 1;
+          int q_p2 = q_time + 2;
+          if (q_m1 >= 0 && q_m1 < num_time_points) {
+            double2 q = Q[((size_t)q_m1)*num_lms + i_lm];
+            q_re += w_m1 * q.x;
+            q_im += w_m1 * q.y;
+          }
+          if (q_0 >= 0 && q_0 < num_time_points) {
+            double2 q = Q[((size_t)q_0)*num_lms + i_lm];
+            q_re += w_0 * q.x;
+            q_im += w_0 * q.y;
+          }
+          if (q_p1 >= 0 && q_p1 < num_time_points) {
+            double2 q = Q[((size_t)q_p1)*num_lms + i_lm];
+            q_re += w_p1 * q.x;
+            q_im += w_p1 * q.y;
+          }
+          if (q_p2 >= 0 && q_p2 < num_time_points) {
+            double2 q = Q[((size_t)q_p2)*num_lms + i_lm];
+            q_re += w_p2 * q.x;
+            q_im += w_p2 * q.y;
+          }
+          double2 a = A[sample_idx*num_lms + i_lm];
+          out_re += a.x*q_re - a.y*q_im;
+          out_im += a.x*q_im + a.y*q_re;
+        }
+
+        out[i_output] = make_double2(out_re, out_im);
+      }
+    }
+  } // Q_inner_cubic
 } // extern

MonteCarloMarginalizeCode/Code/RIFT/likelihood/factored_likelihood.py

@@ -1838,7 +1838,53 @@ def _factored_lnL_helper(kappa_sq, rho_sq):
     return kappa_sq - 0.5 * rho_sq
 
 
-def  DiscreteFactoredLogLikelihoodViaArrayVectorNoLoop(tvals, P_vec, lookupNKDict, rholmsArrayDict, ctUArrayDict,ctVArrayDict,epochDict,Lmax=2,array_output=False,xpy=np, loglikelihood=_factored_lnL_helper,return_lnLt=False,phase_marginalization=False,n_cal=1,cal_method='loop',cal_distmarg=None,cal_log_weights=None,return_cal_components=False):
+def _cubic_Q_window_numpy(Q_block, start_indices, fractional_offsets, npts):
+    """Return cubic-interpolated Q windows with zero extension.
+
+    Q_block has shape (n_time, n_lm).  The returned array has shape
+    (n_extrinsic, npts, n_lm), matching the CPU fallback layout used by NoLoop.
+    The interpolation is a local four-point cubic Lagrange stencil; at integer
+    offsets it reproduces the original samples exactly.
+    """
+    npts_extrinsic = len(start_indices)
+    n_lms_det = Q_block.shape[1]
+    Qlms = np.zeros((npts_extrinsic, npts, n_lms_det), dtype=np.complex128)
+    tgrid = np.arange(npts)
+    n_time = Q_block.shape[0]
+    for i in range(npts_extrinsic):
+        idxs = int(start_indices[i]) + tgrid
+        u = float(fractional_offsets[i])
+        u2 = u*u
+        u3 = u2*u
+        weights = (
+            -u * (u - 1.0) * (u - 2.0) / 6.0,
+            (u + 1.0) * (u - 1.0) * (u - 2.0) / 2.0,
+            -(u + 1.0) * u * (u - 2.0) / 2.0,
+            (u + 1.0) * u * (u - 1.0) / 6.0,
+        )
+        for offset, weight in zip((-1, 0, 1, 2), weights):
+            idxs_here = idxs + offset
+            valid = (idxs_here >= 0) & (idxs_here < n_time)
+            if np.any(valid):
+                Qlms[i, valid] += weight * Q_block[idxs_here[valid]]
+    return Qlms
+
+
+def _nearest_Q_window_numpy(Q_block, start_indices, npts, xpy=np):
+    """Return nearest-grid Q windows with zero extension."""
+    npts_extrinsic = len(start_indices)
+    n_lms_det = Q_block.shape[1]
+    Qlms = xpy.zeros((npts_extrinsic, npts, n_lms_det), dtype=np.complex128)
+    tgrid = np.arange(npts)
+    n_time = Q_block.shape[0]
+    for i in range(npts_extrinsic):
+        idxs = int(start_indices[i]) + tgrid
+        valid = (idxs >= 0) & (idxs < n_time)
+        Qlms[i, valid] = Q_block[idxs[valid]]
+    return Qlms
+
+
+def  DiscreteFactoredLogLikelihoodViaArrayVectorNoLoop(tvals, P_vec, lookupNKDict, rholmsArrayDict, ctUArrayDict,ctVArrayDict,epochDict,Lmax=2,array_output=False,xpy=np, loglikelihood=_factored_lnL_helper,return_lnLt=False,phase_marginalization=False,n_cal=1,cal_method='loop',cal_distmarg=None,cal_log_weights=None,return_cal_components=False,time_interp='nearest'):
     """
     DiscreteFactoredLogLikelihoodViaArray uses the array-ized data structures to compute the log likelihood,
     either as an array vs time *or* marginalized in time.
@@ -1878,9 +1924,21 @@ def  DiscreteFactoredLogLikelihoodViaArrayVectorNoLoop(tvals, P_vec, lookupNKDic
     cal_distmarg : dict or None
         Distance-marginalization table+params for the fused distmarg kernel; see
         RIFT.likelihood.Q_fused_calmarg.Q_fused_calmarg_distmarg_cupy.
+
+    time_interp : {'nearest', 'cubic'}
+        Detector-time sampling convention for the data term.  'nearest'
+        preserves the historical NoLoop integer-bin gather.  'cubic' evaluates
+        the precomputed Q_lm time series at the fractional detector arrival time
+        using a four-sample cubic Lagrange stencil, with zero extension outside
+        the precomputed buffer.
     """
     global distMpcRef
 
+    if time_interp not in ('nearest', 'cubic'):
+        raise ValueError("time_interp must be 'nearest' or 'cubic'")
+    if time_interp != 'nearest' and cal_method == 'fused':
+        raise NotImplementedError("time_interp='{}' is not implemented for cal_method='fused'".format(time_interp))
+
     detectors = rholmsArrayDict.keys()
     npts = len(tvals)
     npts_extrinsic = len(P_vec.phi)
@@ -1980,7 +2038,13 @@ def  DiscreteFactoredLogLikelihoodViaArrayVectorNoLoop(tvals, P_vec, lookupNKDic
         )
         tfirst = t_det + tvals[0]
 
-        ifirst = (xpy.rint((tfirst) / deltaT) + 0.5).astype(np.int32)  # C uses 32 bit integers : be careful
+        sample_first = tfirst / deltaT
+        if time_interp == 'nearest':
+            ifirst = (xpy.rint(sample_first) + 0.5).astype(np.int32)  # C uses 32 bit integers : be careful
+            frac_first = None
+        else:
+            ifirst = xpy.floor(sample_first).astype(np.int32)
+            frac_first = (sample_first - xpy.floor(sample_first)).astype(np.float64)
 #        ilast = ifirst + npts
 
 
@@ -2055,15 +2119,23 @@ def  DiscreteFactoredLogLikelihoodViaArrayVectorNoLoop(tvals, P_vec, lookupNKDic
             # Shape Q = (npts_time_full, nlms)
             # Shape A=FY_conj = (npts_extrinsic, nlms)
             # shape result = (npts_extrinsic, npts_time_*window* = npts)
-            Q_prod_result = Q_inner_product.Q_inner_product_cupy(
-              Q, FY_conj,
-              ifirst, npts,
-              )
+            if time_interp == 'nearest':
+              Q_prod_result = Q_inner_product.Q_inner_product_cupy(
+                Q, FY_conj,
+                ifirst, npts,
+                )
+            else:
+              Q_prod_result = Q_inner_product.Q_inner_product_cubic_cupy(
+                Q, FY_conj,
+                ifirst, frac_first, npts,
+                )
           else:
             # Use old code completely unchanged ... very wasteful on memory management!
-            Qlms = xpy.empty((npts_extrinsic, npts, n_lms), dtype=np.complex128)
-            for i in range(npts_extrinsic):
-                Qlms[i] = rholmsArrayDict[det][...,ifirst[i]:(ifirst[i]+npts)].T
+            Q_block = rholmsArrayDict[det].T
+            if time_interp == 'nearest':
+                Qlms = _nearest_Q_window_numpy(Q_block, ifirst, npts, xpy=xpy)
+            else:
+                Qlms = _cubic_Q_window_numpy(Q_block, ifirst, frac_first, npts)
             if phase_marginalization:
                 Qlms[:, :, 1] = xpy.conj(Qlms[:, :, 1])
 
@@ -2087,7 +2159,7 @@ def  DiscreteFactoredLogLikelihoodViaArrayVectorNoLoop(tvals, P_vec, lookupNKDic
           npts_full_det = Q.shape[0]
           N_window_block = npts_full_det // n_cal
           FY_conj_cal = xpy.conj(F_vec_dummy_lm * Ylms_vec)
-          cal_cache[det] = (Q, FY_conj_cal, ifirst, N_window_block)
+          cal_cache[det] = (Q, FY_conj_cal, ifirst, N_window_block, frac_first)
         # lnL_t_accum += Q_prod_result * (distMpcRef/distMpc)[...,None]
 
         # lnL_t_accum += Q_inner_product.Q_inner_product_cupy(
@@ -2198,7 +2270,7 @@ def  DiscreteFactoredLogLikelihoodViaArrayVectorNoLoop(tvals, P_vec, lookupNKDic
     for c in range(n_cal):
         kappa_sq_c = xpy.zeros((npts_extrinsic, npts), dtype=np.complex128)
         for det in detectors:
-            Q_det, FY_conj_det, ifirst_det, N_window_block = cal_cache[det]
+            Q_det, FY_conj_det, ifirst_det, N_window_block, frac_first_det = cal_cache[det]
             # Restrict to THIS realization block and use the within-block offset, so
             # the window is confined to [0, N_window): an over-running window then
             # zeros (the shared kernel / CPU fill below guard against N_window_block)
@@ -2208,17 +2280,19 @@ def  DiscreteFactoredLogLikelihoodViaArrayVectorNoLoop(tvals, P_vec, lookupNKDic
             Q_block = Q_det[c*N_window_block:(c+1)*N_window_block]   # (N_window, n_lms)
             ifirst_within = ifirst_det.astype(np.int32)
             if not (xpy is np):
-                Q_prod_result = Q_inner_product.Q_inner_product_cupy(
-                    Q_block, FY_conj_det, ifirst_within, npts,
-                )
+                if time_interp == 'nearest':
+                    Q_prod_result = Q_inner_product.Q_inner_product_cupy(
+                        Q_block, FY_conj_det, ifirst_within, npts,
+                    )
+                else:
+                    Q_prod_result = Q_inner_product.Q_inner_product_cubic_cupy(
+                        Q_block, FY_conj_det, ifirst_within, frac_first_det, npts,
+                    )
             else:
-                n_lms_det = Q_block.shape[1]
-                Qlms = xpy.zeros((npts_extrinsic, npts, n_lms_det), dtype=np.complex128)
-                tgrid = np.arange(npts)
-                for i in range(npts_extrinsic):
-                    idxs = int(ifirst_within[i]) + tgrid
-                    valid = (idxs >= 0) & (idxs < N_window_block)
-                    Qlms[i, valid] = Q_block[idxs[valid]]
+                if time_interp == 'nearest':
+                    Qlms = _nearest_Q_window_numpy(Q_block, ifirst_within, npts, xpy=xpy)
+                else:
+                    Qlms = _cubic_Q_window_numpy(Q_block, ifirst_within, frac_first_det, npts)
                 # Q_det and FY_conj_det already encode any phase-marg conjugation
                 Q_prod_result = np.einsum("ej,etj->et", FY_conj_det, Qlms)
             kappa_sq_c += Q_prod_result * invDistMpc[..., np.newaxis]
@@ -2376,5 +2450,3 @@ def ComputeYlmsArrayVector(lookupNK, theta, phi):
 
             Ylms[indx] = lalylm(theta, phi, s, l, m)
     return Ylms
-
-

MonteCarloMarginalizeCode/Code/bin/integrate_likelihood_extrinsic

@@ -206,7 +206,7 @@ integration_params.add_option("--adapt-weight-exponent", type=float, default=1.0
 integration_params.add_option("--adapt-floor-level", type=float, default=0.1, help="Floor to use with weights (likelihood integrand) when doing adaptive sampling. This is necessary to ensure the *sampling* prior is non zero during adaptive sampling and to prevent overconvergence. Default is 0.1 (no floor)")
 integration_params.add_option("--adapt-adapt",action='store_true',help="Adapt the tempering exponent")
 integration_params.add_option("--adapt-log",action='store_true',help="Use a logarithmic tempering exponent")
-integration_params.add_option("--interpolate-time", default=False,help="If using time marginalization, compute using a continuously-interpolated array. (Default=false)")
+integration_params.add_option("--interpolate-time", default=False,help="If using the maintained NoLoop likelihood, evaluate Q_lm at fractional detector times using cubic interpolation instead of nearest sample bins. Accepts truthy values such as True/1/yes. (Default=false)")
 integration_params.add_option("--d-prior",default='Euclidean' ,type=str,help="Distance prior for dL.  Options are dL^2 (Euclidean) and 'pseudo-cosmo'  .")
 integration_params.add_option("--d-max", default=10000,type=float,help="Maximum distance in volume integral. Used to SET THE PRIOR; changing this value changes the numerical answer.")
 integration_params.add_option("--d-min", default=1,type=float,help="Minimum distance in volume integral. Used to SET THE PRIOR; changing this value changes the numerical answer.")
@@ -253,7 +253,26 @@ for pin_param in LIKELIHOOD_PINNABLE_PARAMS:
     pinnable.add_option(option, type=float, help="Pin the value of %s." % pin_param)
 optp.add_option_group(pinnable)
 
-opts, args = optp.parse_args()
+def _normalize_interpolate_time_argv(argv):
+    out = []
+    i = 0
+    while i < len(argv):
+        out.append(argv[i])
+        if argv[i] == "--interpolate-time" and (i + 1 == len(argv) or argv[i + 1].startswith("--")):
+            out.append("True")
+        i += 1
+    return out
+
+opts, args = optp.parse_args(_normalize_interpolate_time_argv(sys.argv[1:]))
+
+def _truthy_option(value):
+    if isinstance(value, bool):
+        return value
+    if value is None:
+        return False
+    return str(value).strip().lower() in ("1", "true", "t", "yes", "y", "on")
+
+opts._noloop_time_interp = "cubic" if _truthy_option(opts.interpolate_time) else "nearest"
 
 if opts.gpu and xpy_default is numpy:
     print( " Override --gpu  (not available);  use --force-xpy to require the identical code path is used (with xpy =[np|cupy]")
@@ -1181,7 +1200,8 @@ else: # Sum over time for every point in other extrinsic params
                   P.phi = xpy_default.mod(P.phi, 2*np.pi)
 
                 lnL = factored_likelihood.DiscreteFactoredLogLikelihoodViaArrayVectorNoLoop(tvals,
-                    P, lookupNKDict, rholmArrayDict, ctUArrayDict, ctVArrayDict,epochDict,Lmax=opts.l_max,xpy=xpy_default)
+                    P, lookupNKDict, rholmArrayDict, ctUArrayDict, ctVArrayDict,epochDict,Lmax=opts.l_max,xpy=xpy_default,
+                    time_interp=opts._noloop_time_interp)
                 nEvals +=len(right_ascension)
                 return identity_convert(xpy_default.exp(lnL-manual_avoid_overflow_logarithm))
         else:
@@ -1262,7 +1282,8 @@ else: # Sum over time for every point in other extrinsic params
                   P.phi = xpy_default.mod(P.phi, 2*np.pi)
 
                 lnL = factored_likelihood.DiscreteFactoredLogLikelihoodViaArrayVectorNoLoop(tvals,
-                    P, lookupNKDict, rholmArrayDict, ctUArrayDict, ctVArrayDict,epochDict,Lmax=opts.l_max,xpy=xpy_default, loglikelihood=distmarg_loglikelihood)
+                    P, lookupNKDict, rholmArrayDict, ctUArrayDict, ctVArrayDict,epochDict,Lmax=opts.l_max,xpy=xpy_default, loglikelihood=distmarg_loglikelihood,
+                    time_interp=opts._noloop_time_interp)
                 nEvals +=len(right_ascension)
                 return identity_convert(xpy_default.exp(lnL-manual_avoid_overflow_logarithm))