update checks to pass when beam_coefs is provided

Author: tyler-a-cox

src/fftvis/__init__.py

@@ -10,6 +10,9 @@
 from .cpu.cpu_simulate import CPUSimulationEngine
 from .wrapper import create_simulation_engine, simulate_vis
 
+# Import beam basis decomposition utility
+from .core.beam_basis import compute_beam_basis
+
 # Import utility modules
 from . import utils, logutils
 
@@ -23,4 +26,6 @@
     "CPUSimulationEngine",
     "create_simulation_engine",
     "simulate_vis",
+    # Beam basis decomposition
+    "compute_beam_basis",
 ]

src/fftvis/core/__init__.py

@@ -2,3 +2,4 @@
 
 from .beams import BeamEvaluator
 from .simulate import SimulationEngine, default_accuracy_dict
+from .beam_basis import compute_beam_basis

src/fftvis/core/beam_basis.py

@@ -14,140 +14,56 @@
 
 
 def compute_beam_basis(
-    beam_list: list[UVBeam],
+    beam_list,
     freq: float,
     threshold: float = 1e-12,
-    polarized: bool = False,
-) -> tuple[list[UVBeam], np.ndarray]:
-    """Decompose a set of UVBeams into an SVD eigenbeam basis.
-
-    Each beam is first interpolated onto a common frequency grid to keep memory
-    usage predictable, then stacked into a matrix and SVD'd.  The basis is
-    truncated at the first component whose normalised singular value drops below
-    ``threshold``.
-
-    The decomposition satisfies::
-
-        beam_i(az, za, freq) ≈ sum_k  coeffs[i, k] * eigenbeams[k](az, za, freq)
-
-    so ``eigenbeams`` can be passed directly as ``beam_list`` to
-    ``CPUSimulationEngine.simulate`` and ``coeffs`` as ``beam_coeffs``.
-
-    Parameters
-    ----------
-    beam_list : list of UVBeam
-        Input beams to decompose.  All beams must share the same spatial grid
-        (``Nax1``, ``Nax2``) and ``beam_type`` (``efield`` or ``power``).
-    freq : float
-        Frequency in Hz at which every beam is evaluated before the SVD.
-    threshold : float, optional
-        Normalised singular value cutoff.  Basis components ``k`` with
-        ``s[k] / s[0] < threshold`` are discarded.  Default ``1e-3``.
-    polarized : bool, optional
-        If ``True``, the full Jones matrix (all ``Naxes_vec`` and ``Nfeeds``
-        components) is included in the SVD.  If ``False``, only the first
-        component ``data_array[0, 0, 0, ...]`` is used, matching the scalar
-        beam path in the simulator.  Default ``False``.
-
-    Returns
-    -------
-    eigenbeams : list of UVBeam
-        The ``K`` truncated eigenbeam objects.  Each has the same metadata and
-        spatial/frequency grid as the interpolated input beams.
-    coeffs : np.ndarray
-        Per-antenna coefficients of shape ``(N_beams, K)``.  These are the
-        values to pass as ``beam_coeffs`` to the simulator.
-
-    Notes
-    -----
-    The function allocates two large intermediate arrays simultaneously: the
-    full stacked beam matrix ``B`` of shape ``(N_beams, N_flat)`` and the ``Vh``
-    factor of the same shape.  Peak memory is therefore roughly
-    ``2 * N_beams * N_flat * itemsize``.  Choosing a sparse ``freq_grid``
-    directly reduces ``N_flat``.
-
-    Only ``efield`` beams have been tested; ``power`` beams should work but
-    the resulting eigenbeams will have real-valued ``data_array`` only if all
-    inputs are real.
+):
+    """
     """
     if len(beam_list) == 0:
         raise ValueError("beam_list must contain at least one beam.")
 
     n_beams = len(beam_list)
     freq_grid = np.atleast_1d(freq)
 
-    # ------------------------------------------------------------------
-    # Step 1: Interpolate every beam to the common frequency grid.
-    # ------------------------------------------------------------------
-    logger.info(
-        f"Interpolating {n_beams} beams to {len(freq_grid)} frequencies "
-        f"({freq_grid[0]/1e6:.1f}–{freq_grid[-1]/1e6:.1f} MHz)."
-    )
     interp_beams = []
     for idx, beam in enumerate(beam_list):
         interp_beams.append(
             beam.interp(freq_array=freq_grid, new_object=True)
         )
-        logger.debug(f"  Interpolated beam {idx + 1}/{n_beams}.")
 
-    # ------------------------------------------------------------------
-    # Step 2: Extract the relevant slice of data_array and flatten.
-    #
-    # data_array shape: (Naxes_vec, 1, Nfeeds, Nfreqs, Nax1, Nax2)
-    #
-    # Polarized  → use axes_vec and feeds: data_array[:, 0, :, ...]
-    #              flat shape per beam: Naxes_vec * Nfeeds * Nfreqs * Nax1 * Nax2
-    # Unpolarized → match evaluate_beam scalar path: data_array[0, 0, 0, ...]
-    #              flat shape per beam: Nfreqs * Nax1 * Nax2
-    # ------------------------------------------------------------------
     ref = interp_beams[0]
 
-    if polarized:
-        # Shape of the slice we care about: (Naxes_vec, Nfeeds, Nfreqs, Nax1, Nax2)
-        slice_shape = ref.data_array[:, 0, :, :, :, :].shape
-        flat_vecs = np.array(
-            [b.data_array[:, 0, :, :, :, :].ravel() for b in interp_beams]
-        )  # (N_beams, Naxes_vec * Nfeeds * Nfreqs * Nax1 * Nax2)
-    else:
-        # Shape: (Nfreqs, Nax1, Nax2)
-        slice_shape = ref.data_array[0, 0, 0, :, :, :].shape
-        flat_vecs = np.array(
-            [b.data_array[0, 0, 0, :, :, :].ravel() for b in interp_beams]
-        )  # (N_beams, Nfreqs * Nax1 * Nax2)
-
-    logger.info(
-        f"Beam matrix shape: {flat_vecs.shape}  "
-        f"({flat_vecs.nbytes / 1024**2:.1f} MB)."
-    )
+    # Shape of the slice we care about: 
+    # (Naxes_vec, Nfeeds, Nfreqs, Nax1) for healpix beams
+    # (Naxes_vec, Nfeeds, Nfreqs, Nax1, Nax2) for gridded beams
+    slice_shape = ref.data_array[:, :, 0].shape
+    flat_vecs = np.array(
+        [b.data_array[:, :, 0].ravel() for b in interp_beams]
+    )  # (N_beams, Naxes_vec * Nfeeds * Nfreqs * Nax1 * Nax2)
 
     # ------------------------------------------------------------------
     # Step 3: SVD.
     # B = U @ diag(s) @ Vh  →  beam_i ≈ sum_k (U[i,k]*s[k]) * Vh[k]
     # ------------------------------------------------------------------
-    logger.info("Computing SVD...")
     U, s, Vh = np.linalg.svd(flat_vecs, full_matrices=False)
 
     # ------------------------------------------------------------------
     # Step 4: Truncate at normalised singular value threshold.
     # ------------------------------------------------------------------
     s_norm = s / s[0]
     K = int(np.sum(s_norm >= threshold))
-    logger.info(
-        f"Retaining {K}/{len(s)} basis components "
-        f"(threshold={threshold}, "
-        f"min retained s_norm={s_norm[K-1]:.3e})."
-    )
 
     U_k  = U[:, :K]   # (N_beams, K)
     s_k  = s[:K]       # (K,)
     Vh_k = Vh[:K, :]  # (K, N_flat)
 
     # ------------------------------------------------------------------
     # Step 5: Compute per-antenna coefficients.
-    # coeffs[i, k] = U[i, k] * s[k]  so that  flat_vecs ≈ coeffs @ Vh_k
+    # beam_coefs[i, k] = U[i, k] * s[k]  so that  flat_vecs ≈ beam_coefs @ Vh_k
     # ------------------------------------------------------------------
-    coeffs = U_k * s_k[None, :]  # (N_beams, K)
-
+    beam_coefs = U_k * s_k[None, :]  # (N_beams, K)
+    
     # ------------------------------------------------------------------
     # Step 6: Build eigenbeam UVBeam objects by copying reference metadata
     # and replacing data_array with the reshaped Vh rows.
@@ -156,16 +72,7 @@ def compute_beam_basis(
     for k in range(K):
         eb = ref.copy()
         eigenbeam_slice = Vh_k[k].reshape(slice_shape)
-
-        if polarized:
-            # Restore the dropped size-1 axis: (Naxes_vec, 1, Nfeeds, Nfreqs, Nax1, Nax2)
-            eb.data_array = eigenbeam_slice[:, np.newaxis, :, :, :, :]
-        else:
-            # Restore to full data_array shape; fill non-primary components with zeros.
-            eb.data_array = np.zeros_like(ref.data_array)
-            eb.data_array[0, 0, 0, :, :, :] = eigenbeam_slice
-
+        eb.data_array = eigenbeam_slice[:, :, np.newaxis]
         eigenbeams.append(eb)
 
-    logger.info(f"Basis computation complete: {K} eigenbeams.")
-    return eigenbeams, coeffs
+    return eigenbeams, beam_coefs

src/fftvis/cpu/cpu_simulate.py

@@ -308,7 +308,7 @@ def _compute_basis_visibilities(
     flux_here: np.ndarray,
     ant1_idxs: np.ndarray,
     ant2_idxs: np.ndarray,
-    beam_coeffs: np.ndarray,
+    beam_coefs: np.ndarray,
     freqidx: int,
     topo: np.ndarray,
     uvw: np.ndarray,
@@ -358,7 +358,7 @@ def _compute_basis_visibilities(
         ``(nsrc, nfreqs, nfeeds, nfeeds)`` for polarized sky.
     ant1_idxs, ant2_idxs : np.ndarray
         Antenna indices for each baseline, each shape ``(nbls,)``.
-    beam_coeffs : np.ndarray
+    beam_coefs : np.ndarray
         Coefficients for each baseline, shape ``(nbls, nbasis)``. 
     freqidx : int
         Frequency index into flux_here.
@@ -416,8 +416,8 @@ def _compute_basis_visibilities(
     # Gather coefficients once, outside the loop.
     # The measurement equation is V_ij = A_i^H C A_j, so the left (ant1)
     # coefficients are conjugated and the right (ant2) are not.
-    ant1_c = beam_coeffs[ant1_idxs, :].conj()  # C_ik^*  (nbls, K)
-    ant2_c = beam_coeffs[ant2_idxs, :]          # C_jl    (nbls, K)
+    ant1_c = beam_coefs[ant1_idxs, :, freqidx].conj()  # C_ik^*  (nbls, K)
+    ant2_c = beam_coefs[ant2_idxs, :, freqidx]          # C_jl    (nbls, K)
 
     # Only iterate over the upper triangle (k <= l) and use the conjugate
     # symmetry V_tilde[l, k] = V_tilde[k, l]^* to handle the lower triangle
@@ -499,7 +499,7 @@ def _evaluate_vis_chunk_remote(
     type1_n_modes: int = None,
     trace_mem: bool = False,
     nchunks: int = 1,
-    beam_coeffs: np.ndarray = None,
+    beam_coefs: np.ndarray = None,
 ):
     """Ray-compatible remote version of _evaluate_vis_chunk."""
     engine = CPUSimulationEngine()  # pragma: no cover
@@ -529,7 +529,7 @@ def _evaluate_vis_chunk_remote(
         type1_n_modes=type1_n_modes,
         trace_mem=trace_mem,
         nchunks=nchunks,
-        beam_coeffs=beam_coeffs,
+        beam_coefs=beam_coefs,
     )
 
 
@@ -567,18 +567,18 @@ def simulate(
         enable_memory_monitor: bool = False,
         nchunks: int = 1,
         source_buffer=1.0,
-        beam_coeffs: np.ndarray = None,
+        beam_coefs: np.ndarray = None,
     ) -> np.ndarray:
         """
         Simulate visibilities using CPU implementation.
 
         Parameters
         ----------
-        beam_coeffs : np.ndarray, optional
+        beam_coefs : np.ndarray, optional
             Per-antenna SVD coefficients of shape (N_ant, K). When provided,
             beam_list is interpreted as K basis beams rather than per-antenna
             beams. Visibilities are computed over all K^2 basis pairs and then
-            contracted with beam_coeffs in post-processing.
+            contracted with beam_coefs in post-processing.
 
         See base class for all other parameter descriptions.
         """
@@ -743,8 +743,8 @@ def simulate(
             freqs = ray.put(freqs)
             beam_list = ray.put(beam_list)
             coord_mgr = ray.put(coord_mgr)
-            if beam_coeffs is not None:
-                beam_coeffs = ray.put(beam_coeffs)
+            if beam_coefs is not None:
+                beam_coefs = ray.put(beam_coefs)
             if trace_mem:
                 os.system("ray memory --units MB > after-puts.txt")
 
@@ -803,7 +803,7 @@ def simulate(
                     type1_n_modes=n_modes if is_gridded else None,
                     trace_mem=(nprocesses > 1 or force_use_ray) and trace_mem,
                     nchunks=nchunks,
-                    beam_coeffs=beam_coeffs,
+                    beam_coefs=beam_coefs,
                 )
             )
             if trace_mem:
@@ -856,14 +856,14 @@ def _evaluate_vis_chunk(
         type1_n_modes: int = None,
         trace_mem: bool = False,
         nchunks: int = 1,
-        beam_coeffs: np.ndarray = None,
+        beam_coefs: np.ndarray = None,
     ) -> np.ndarray:
         """
         Evaluate a chunk of visibility data using CPU.
 
         Parameters
         ----------
-        beam_coeffs : np.ndarray, optional
+        beam_coefs : np.ndarray, optional
             Per-antenna SVD coefficients, shape (N_ant, K). When provided,
             beam_list contains the K basis beams and the standard beam-pair
             loop is replaced by the basis visibility path.
@@ -888,11 +888,11 @@ def _evaluate_vis_chunk(
 
         coord_mgr.setup()
 
-        use_basis = beam_coeffs is not None
+        use_basis = beam_coefs is not None
 
         if use_basis:
             # Pre-compute the per-baseline antenna index arrays once.
-            # These are used in post-processing to gather the right rows of beam_coeffs.
+            # These are used in post-processing to gather the right rows of beam_coefs.
             ant1_idxs = np.array([antnums.index(bl[0]) for bl in baselines])
             ant2_idxs = np.array([antnums.index(bl[1]) for bl in baselines])
         else:
@@ -972,7 +972,7 @@ def _evaluate_vis_chunk(
                                 flux_here=flux,
                                 ant1_idxs=ant1_idxs,
                                 ant2_idxs=ant2_idxs,
-                                beam_coeffs=beam_coeffs,
+                                beam_coefs=beam_coefs,
                                 freqidx=freqidx,
                                 topo=topo,
                                 uvw=uvw if not use_type1 else None,