Projection

def project_polarizations_to_network(  # noqa: PLR0913, PLR0915
    polarizations: Mapping[str, GWpyTimeSeries],
    detector_names: Sequence[DetectorSpec],
    *,
    right_ascension: float,
    declination: float,
    polarization_angle: float,
    earth_rotation: bool = True,
    precess_source_direction: bool = False,
    backend: str = "numpy",
    sinc_taps: int = DEFAULT_SINC_TAPS,
    kaiser_beta: float = DEFAULT_KAISER_BETA,
) -> dict[str, GWpyTimeSeries]:
    """Project tensor plus/cross strains onto detectors using detector geometry.

    Built-in and custom detector codes are resolved through the LAL cached
    detector registry. For ``earth_rotation=False``, the constant
    geocenter->detector delay is applied via an exact frequency-domain phase
    shift (``h(t-tau) <-> H(f)*exp(-2*pi*i*f*tau)``), which is lossless at all
    frequencies. For ``earth_rotation=True``, the polarizations are resampled at the
    time-dependent delayed times with the Kaiser-windowed sinc kernel in
    :mod:`gwmock_signal.projection.resampling`, gathered from a zero-padded copy so kernel
    taps reaching past either end read zero rather than repeating the endpoint.

    Args:
        polarizations: Mapping containing ``plus`` and ``cross`` GWpy time series
            on a common grid.
        detector_names: Sequence of IFO codes (e.g. ``H1``, ``L1``, ``V1``) or
            :class:`~gwmock_signal.detector.CustomDetector` instances, or a mix
            of both.
        right_ascension: Source right ascension in radians.
        declination: Source declination in radians.
        polarization_angle: Polarization angle psi in radians (tensor modes).
        earth_rotation: If ``True``, evaluate antenna patterns at time-dependent
            GPS times (recommended for longer signals). If ``False``, use a single
            reference time at the segment midpoint for patterns and delays.
        precess_source_direction: Whether to rotate ``right_ascension``/``declination`` from J2000
            into the mean equator and equinox *of date* before using them. **Which value is correct
            depends on the source type, because LAL itself uses two conventions and they disagree
            by 1.8e-04 s of geocentre-to-detector delay:**

            * ``False`` (the default) reproduces ``XLALTimeDelayFromEarthCenter`` and
              ``XLALComputeDetAMResponse`` -- ``gha = gmst - ra`` with no rotation -- which is what
              every compact-binary search and parameter-estimation code does (Bilby, PyCBC,
              LALInference, GstLAL). Strictly it mixes frames, since GMST is measured from the
              equinox of date; but a CBC injection is only useful if the pipeline that recovers it
              agrees about where the source was, and precessing here would shift the recovered
              right ascension by ~0.43 degrees by 2030.
            * ``True`` reproduces ``lalpulsar.XLALBarycenter``, which applies lunisolar precession.
              Required for continuous waves, where the SSB-to-geocentre part of the phase comes from
              a barycentering routine that precesses (both LAL's and ripple's do), so *not*
              rotating here would leave the site term inconsistent with the term it is added to.

            This is therefore a property of the generator, not of the projection, which is why it is
            an explicit argument with no clever default: see
            :func:`~gwmock_signal.projection.sidereal.precess_to_epoch`.
        backend: Which implementation evaluates the ``earth_rotation=True`` branch.
            ``"numpy"`` (the default) runs on the host. ``"jax"`` delegates to
            :func:`~gwmock_signal.projection.jax_projection.project_polarizations_td_rotating`,
            which runs on whatever JAX backend is configured and is several times faster even
            on a CPU because the whole per-sample kernel fuses into one compiled loop. The two
            agree to 1e-10 of peak, pinned by ``test_rotating_projection_matches_numpy_path``;
            the difference is floating-point reassociation, not a different model. Selected
            rather than automatic: an implicit switch on whether JAX imports would make the
            numerical output depend on what happens to be installed.
        sinc_taps: Taps in the band-limited resampling kernel used by the
            ``earth_rotation=True`` branch. More taps cost arithmetic and buy accuracy.
        kaiser_beta: Kaiser window shape parameter for that kernel.

    Returns:
            Mapping from each detector name to the projected strain as a GWpy time
            series (same length and sample rate as the inputs).

    Raises:
        TypeError: If ``polarizations`` is not a mapping of GWpy series as required.
        ValueError: If keys are missing, time grids disagree, a detector name is not
            recognized, ``backend`` is unknown, or ``backend="jax"`` is combined with
            ``earth_rotation=False``.
    """
    validate_projection_backend(backend)
    if backend == "jax" and not earth_rotation:
        # The constant-pattern branch is a frequency-domain phase shift, and no device
        # counterpart exists. Refused rather than silently served from the host path, which
        # would report a backend that did not run.
        raise ValueError(
            "backend='jax' is only available with earth_rotation=True. The constant-pattern "
            "branch applies one frequency-domain phase shift for the whole span and has no "
            "device implementation; it is also cheap, being the branch that skips the resampler."
        )
    hp, hc = _validate_polarizations(polarizations)
    normalized_names = [d if isinstance(d, str) else d.name for d in detector_names]
    if len(set(normalized_names)) != len(normalized_names):
        raise ValueError("detector_names must not contain duplicates.")
    detectors = _make_detectors(list(detector_names))

    time_array = cast(np.ndarray, hp.times.to_value())
    reference_time = float(0.5 * (time_array[0] + time_array[-1]))

    _warn_if_constant_pattern_is_stretched(time_array, earth_rotation=earth_rotation)

    # Resolved once, before anything else reads the sky position. Every branch below -- host scalar,
    # host array, and the device kernel, for the delay *and* the antenna pattern -- derives its
    # geometry from these numbers, so deciding here rather than at each site is what makes it
    # impossible for one path to precess and another not to.
    #
    # Anchored at the *first* sample, because that is the origin the device kernel's `sample_offsets`
    # counts from, so both backends read the same line without a second convention to keep straight.
    #
    # A position and a *rate*, not a position alone. Freezing it per segment looks harmless -- the
    # angles move 2306 arcseconds per century, so across 4096 s they drift 3e-6 arcseconds -- but that
    # argument is about drift *within* a segment and says nothing about the step *between* two of them.
    # The step broke continuous-wave phase coherence at 1.6e-08 of peak against a 1e-09 tolerance.
    # Linear in absolute time removes it: two abutting segments evaluate the same line at the same
    # absolute time, so they agree exactly where they meet whatever their anchors are.
    #
    # Zero rates when not precessing, rather than a separate code path: the same multiply-add then
    # holds the position fixed exactly, so the two conventions cannot diverge in anything but the
    # numbers they put in.
    if precess_source_direction:
        # Python floats on both branches, which `precessed_sky_anchor_and_rate` guarantees: 0-d
        # NumPy arrays here would be a distinct `jax.jit` signature from the other branch's weakly
        # typed floats, so a process using both conventions at one segment shape would compile the
        # same device kernel twice.
        (right_ascension, declination), (d_right_ascension, d_declination) = precessed_sky_anchor_and_rate(
            right_ascension, declination, float(time_array[0])
        )
    else:
        d_right_ascension = d_declination = 0.0

    # Dispatched here, before any of the host branch's preparation. Everything below -- two
    # rffts, the frequency grid, and per-sample Astropy GMST with its sines and cosines -- serves
    # only the NumPy branches, and the device path recomputes what it needs from an anchor and a
    # rate. Left after the dispatch it was pure waste, and not a small one: at 4096 s and 512 Hz
    # the Astropy call alone is 4.3 s on this machine and the unused arrays are ~120 MiB, so the
    # device route was neither as fast nor as independent of host scaling as it claimed.
    if backend == "jax":
        return _project_rotating_on_device(
            hp,
            hc,
            detectors,
            time_array=time_array,
            right_ascension=right_ascension,
            declination=declination,
            right_ascension_rate=d_right_ascension,
            declination_rate=d_declination,
            polarization_angle=polarization_angle,
            sinc_taps=sinc_taps,
            kaiser_beta=kaiser_beta,
        )

    hp_vals = hp.to_value()
    hc_vals = hc.to_value()

    # Precomputed once for the exact FD phase-shift (earth_rotation=False path).
    n_samples = len(hp_vals)
    dt = float(hp.dt.value)
    rfft_hp = np.fft.rfft(hp_vals)
    rfft_hc = np.fft.rfft(hc_vals)
    freqs_fd = np.fft.rfftfreq(n_samples, d=dt)

    strains: dict[str, GWpyTimeSeries] = {}

    # Per sample, not per segment: see the anchor-and-rate comment above.
    precession_offsets = time_array - time_array[0]
    right_ascension_array = right_ascension + d_right_ascension * precession_offsets
    declination_array = declination + d_declination * precession_offsets
    cosdec = np.cos(declination_array)
    sindec = np.sin(declination_array)
    cospsi = np.cos(polarization_angle)
    sinpsi = np.sin(polarization_angle)
    gmst_array = _gmst_accurate_array(time_array)
    gha_array = gmst_array - right_ascension_array
    cosgha = np.cos(gha_array)
    singha = np.sin(gha_array)

    for name, prefix in detectors:
        if earth_rotation:
            response, location = reconstructed_geometry(prefix)

            # Vectorized time delay: time_delay = -location · prop_dir / c
            prop_dir = np.stack([cosdec * cosgha, -cosdec * singha, sindec], axis=-1)
            time_delays = -np.dot(prop_dir, location) / constants.c.value

            # Antenna pattern at the detector-time sample, i.e. the same time coordinate
            # the output series is labelled with. Evaluating it at t + tau would mix the
            # detector and geocenter time coordinates; LALSuite and the bilby-x-g
            # frequency-domain implementation both use a single consistent coordinate.

            # Shape (N, 3) — polarization basis vectors
            x_vec = np.stack(
                [
                    -cospsi * singha - sinpsi * cosgha * sindec,
                    -cospsi * cosgha + sinpsi * singha * sindec,
                    sinpsi * cosdec,
                ],
                axis=-1,
            )
            y_vec = np.stack(
                [
                    sinpsi * singha - cospsi * cosgha * sindec,
                    sinpsi * cosgha + cospsi * singha * sindec,
                    cospsi * cosdec,
                ],
                axis=-1,
            )

            # dx[n] = response @ x_vec[n], using row-vector form: x_vec @ response.T
            dx = x_vec @ response.T
            dy = y_vec @ response.T
            fp_vals = np.sum(x_vec * dx - y_vec * dy, axis=-1)
            fc_vals = np.sum(x_vec * dy + y_vec * dx, axis=-1)

            # Resample at t - tau(t) with the shared band-limited kernel. Expressed as a
            # fractional sample index because the input grid is uniform. Gathered from a
            # zero-padded copy for the same reason as the device path -- taps reaching past
            # either end must read zero rather than clamping to and repeating the endpoint --
            # and with the same padding, or the two paths disagree at the edges by the
            # difference in padding alone.
            require_terrestrial_location(location, name=f"location of {prefix}")
            pad = edge_padding(float(hp.sample_rate.value), sinc_taps, kaiser_beta)
            index = pad + np.arange(len(time_array), dtype=float) - time_delays / dt
            hp_shifted = resample_uniform_sinc(np.pad(hp_vals, (pad, pad)), index, taps=sinc_taps, beta=kaiser_beta)
            hc_shifted = resample_uniform_sinc(np.pad(hc_vals, (pad, pad)), index, taps=sinc_taps, beta=kaiser_beta)
        else:
            time_delay = _time_delay_from_earth_center_lal(
                prefix,
                right_ascension=right_ascension,
                declination=declination,
                t_gps=reference_time,
            )
            fp_vals, fc_vals = _antenna_pattern_lal(
                prefix,
                right_ascension=right_ascension,
                declination=declination,
                polarization_angle=polarization_angle,
                t_gps=reference_time,
            )
            # Exact FD phase shift: h(t-τ) <-> H(f)·exp(-2πifτ)
            # Circular-wrap is negligible for tapered polarizations (tested in test suite).
            phase = np.exp(-2j * np.pi * freqs_fd * time_delay)
            hp_shifted = np.fft.irfft(rfft_hp * phase, n=n_samples)
            hc_shifted = np.fft.irfft(rfft_hc * phase, n=n_samples)

        response = fp_vals * hp_shifted + fc_vals * hc_shifted

        strains[name] = GWpyTimeSeries(
            response,
            t0=float(time_array[0]),
            sample_rate=hp.sample_rate,
            name=name,
        )

    return strains

def validate_projection_backend(value: str, *, parameter: str = "backend") -> str:
    """Return *value* if it names a projection implementation, else raise.

    Shared so the name is checked identically wherever it is accepted -- this function, and the
    simulators that take it at construction time and only reach the projection later. Checking it
    in each place separately is how one of them ends up accepting a name the projection does not.

    Args:
        value: The candidate backend name.
        parameter: What to call it in the error message, since callers spell the argument
            differently (``backend`` here, ``projection_backend`` on a simulator).

    Returns:
        *value* unchanged.

    Raises:
        ValueError: If *value* is not one of :data:`PROJECTION_BACKENDS`.
    """
    if value not in PROJECTION_BACKENDS:
        raise ValueError(f"{parameter} must be 'numpy' or 'jax', got {value!r}.")
    return value