Simulator

def __init__(self, waveform_model: str, waveform_backend: WaveformBackend | None = None) -> None:
    """Initialise with the waveform model name.

    Args:
        waveform_model: Time-domain approximant name or any model
            registered with ``WaveformFactory``.
        waveform_backend: Optional waveform backend instance.
    """
    self._waveform_model = waveform_model
    self._waveform_factory = WaveformFactory(backend=waveform_backend)

def generate_polarizations(
    self,
    params: Mapping[str, Any],
    sampling_frequency: float,
    minimum_frequency: float,
) -> tuple[TimeSeries, TimeSeries]:
    """Delegate waveform generation to ``WaveformFactory``.

    Projection-specific keys are excluded because they are consumed by
    ``TransientSimulator.simulate`` rather than by waveform generation.

    Args:
        params: CBC source parameters using gwmock-pop canonical names
            (e.g. ``detector_frame_mass_1``, ``coa_time``).
        sampling_frequency: Sample rate in Hz.
        minimum_frequency: Low-frequency cutoff in Hz.

    Returns:
        Tuple of ``(hp, hc)`` GWpy ``TimeSeries`` objects.
    """
    waveform_params = {
        k: v
        for k, v in params.items()
        if k not in {"right_ascension", "declination", "polarization_angle", "coa_time"}
    }

    result = self._waveform_factory.generate(
        self._waveform_model,
        waveform_params,
        tc=params["coa_time"],
        sampling_frequency=sampling_frequency,
        minimum_frequency=minimum_frequency,
    )
    return result["plus"], result["cross"]

def register_waveform_model(
    self,
    name: str,
    factory: RegisteredWaveformFactory,
) -> None:
    """Register a one-shot waveform generator under ``name`` for this instance.

    Re-registering the same name with an equal factory is a no-op. Reusing
    the name for a different factory raises ``ValueError``.

    Args:
        name: Waveform model name used later as ``waveform_model``.
        factory: Callable returning either ``(plus, cross)`` or a
            ``{"plus": ..., "cross": ...}`` mapping, or a ``WaveformBackend``
            whose ``generate_td_waveform`` method should serve that name.
    """
    if not hasattr(self, "_waveform_factory"):
        raise AttributeError(
            "TransientSimulator subclasses must initialize _waveform_factory before registering waveform models."
        )

    registered_factories = self.__dict__.setdefault("_registered_waveform_models", {})
    if name in registered_factories:
        if registered_factories[name] == factory:
            return
        raise ValueError(f"Waveform model {name!r} is already registered for this simulator instance.")

    try:
        self._waveform_factory.get_model(name)
    except ValueError:
        pass
    else:
        raise ValueError(f"Waveform model {name!r} is already registered for this simulator instance.")

    if not isinstance(factory, WaveformBackend) and not callable(factory):
        raise TypeError(f"Waveform model {name!r} must be registered with a callable or WaveformBackend.")

    wrapped_factory = (
        _wrap_registered_waveform_backend(name, factory)
        if isinstance(factory, WaveformBackend)
        else _wrap_registered_waveform_callable(factory)
    )
    self._waveform_factory.register_model(name, wrapped_factory)
    registered_factories[name] = factory

def simulate(  # noqa: PLR0913
    self,
    params: Mapping[str, Any],
    detector_names: Sequence[str | CustomDetector],
    background: Mapping[str, TimeSeries] | None = None,
    *,
    sampling_frequency: float,
    minimum_frequency: float,
    earth_rotation: bool = True,
    interpolate_if_offset: bool = True,
) -> DetectorStrainStack:
    """Run the full injection pipeline and return a ``DetectorStrainStack``.

    Calls ``_validate_params``, ``generate_polarizations``,
    ``project_polarizations_to_network``, and ``inject_strain`` in order,
    then assembles the result into a ``DetectorStrainStack``.

    Args:
        params: Source parameters; must contain all ``required_params`` keys
            plus ``right_ascension``, ``declination``, and ``polarization_angle``
            for antenna-response projection.
        detector_names: IFO codes (e.g. ``'H1'``, ``'L1'``, ``'V1'``).
        background: Optional mapping of detector name to background
            ``TimeSeries`` (e.g. noise or zeros). If omitted, the projected
            detector response is returned directly on the waveform time grid.
        sampling_frequency: Sample rate in Hz.
        minimum_frequency: Low-frequency cutoff in Hz.
        earth_rotation: If ``True``, evaluate antenna patterns at
            time-dependent GPS times (recommended for longer signals).
            If ``False``, use a single reference time.
        interpolate_if_offset: If ``True``, interpolate the injection
            when its start is not on a target sample boundary.

    Returns:
        ``DetectorStrainStack`` containing the simulated strain for each
        detector in ``detector_names`` order.
    """
    self._validate_params(params)

    # Validate projection keys before direct params[...] access.
    # _validate_params only checks subclass-defined required_params. Non-CBC subclasses can pass validation and then fail mid-pipeline with KeyError.
    projection_keys = {"right_ascension", "declination", "polarization_angle"}
    missing_projection = projection_keys - set(params)
    if missing_projection:
        raise ValueError(f"Missing required parameters: {sorted(missing_projection)}")

    # Normalise detector entries so dict lookups use plain string keys.
    str_names: list[str] = [d.name if hasattr(d, "name") and not isinstance(d, str) else d for d in detector_names]

    hp, hc = self.generate_polarizations(params, sampling_frequency, minimum_frequency)
    projected = project_polarizations_to_network(
        {"plus": hp, "cross": hc},
        detector_names,
        right_ascension=params["right_ascension"],
        declination=params["declination"],
        polarization_angle=params["polarization_angle"],
        earth_rotation=earth_rotation,
    )
    if background is None:
        return DetectorStrainStack.from_mapping(str_names, projected)
    injected = {
        name: inject_strain(
            background[name],
            projected[name],
            interpolate_if_offset=interpolate_if_offset,
        )
        for name in str_names
    }
    return DetectorStrainStack.from_mapping(str_names, injected)

def write(  # noqa: PLR0913
    self,
    path: str | Path,
    params: Mapping[str, Any],
    detector_names: Sequence[str | CustomDetector],
    background: Mapping[str, TimeSeries],
    *,
    sampling_frequency: float,
    minimum_frequency: float,
    format: Literal["gwf", "hdf5", "npy", "txt"] = "hdf5",  # noqa: A002
    earth_rotation: bool = True,
    interpolate_if_offset: bool = True,
) -> DetectorStrainStack:
    """Simulate a CBC injection, write the result, and save the parameter dict.

    Calls :meth:`simulate`, writes the returned ``DetectorStrainStack`` to
    *path* via :meth:`DetectorStrainStack.write`, and saves *params* as a
    JSON sidecar at ``<stem>_params.json`` next to *path*.

    Args:
        path: Output file path.
        params: CBC source parameters (same as :meth:`simulate`).
        detector_names: IFO codes for the target network.
        background: Mapping of detector name to background ``TimeSeries``.
        sampling_frequency: Sample rate in Hz.
        minimum_frequency: Low-frequency cutoff in Hz.
        format: Output format for the strain data — one of ``'gwf'``,
            ``'hdf5'``, ``'npy'``, or ``'txt'``.  Defaults to ``'hdf5'``.
        earth_rotation: Passed through to :meth:`simulate`.
        interpolate_if_offset: Passed through to :meth:`simulate`.

    Returns:
        The ``DetectorStrainStack`` that was written to disk.
    """
    result = self.simulate(
        params,
        detector_names,
        background,
        sampling_frequency=sampling_frequency,
        minimum_frequency=minimum_frequency,
        earth_rotation=earth_rotation,
        interpolate_if_offset=interpolate_if_offset,
    )
    result.write(path, format=format)

    params_path = Path(path).with_name(Path(path).stem + "_params.json")
    params_path.write_text(json.dumps(dict(params), default=_json_default, indent=4))

    return result

@abstractmethod
def simulate(  # noqa: PLR0913
    self,
    params: Mapping[str, Any],
    detector_names: Sequence[str],
    background: Mapping[str, TimeSeries] | None = None,
    *,
    sampling_frequency: float,
    minimum_frequency: float,
    earth_rotation: bool = True,
    interpolate_if_offset: bool = True,
) -> DetectorStrainStack:
    """Return detector strain for one source realization as a ``DetectorStrainStack``.

    Args:
        params: Source parameters.
        detector_names: IFO codes (e.g. ``'H1'``, ``'L1'``, ``'V1'``).
        background: Optional mapping of detector name to existing background
            strain. Non-transient subclasses may use this to inject into
            pre-existing data; transient subclasses may ignore it and return
            signal-only strain directly.
        sampling_frequency: Sample rate in Hz.
        minimum_frequency: Low-frequency cutoff in Hz.
        earth_rotation: Optional backend-specific flag controlling whether
            detector response should vary across the signal duration.
        interpolate_if_offset: Optional backend-specific flag controlling
            how off-grid injections should be handled.

    Returns:
        ``DetectorStrainStack`` containing one aligned strain channel per
        detector in ``detector_names`` order.
    """

@abstractmethod
def generate_polarizations(
    self,
    params: Mapping[str, Any],
    sampling_frequency: float,
    minimum_frequency: float,
) -> tuple[TimeSeries, TimeSeries]:
    """Generate plus and cross polarization time series.

    Args:
        params: Source parameters.
        sampling_frequency: Sample rate in Hz.
        minimum_frequency: Low-frequency cutoff in Hz.

    Returns:
        Tuple of ``(hp, hc)`` GWpy ``TimeSeries`` objects.
    """

def register_waveform_model(
    self,
    name: str,
    factory: RegisteredWaveformFactory,
) -> None:
    """Register a one-shot waveform generator under ``name`` for this instance.

    Re-registering the same name with an equal factory is a no-op. Reusing
    the name for a different factory raises ``ValueError``.

    Args:
        name: Waveform model name used later as ``waveform_model``.
        factory: Callable returning either ``(plus, cross)`` or a
            ``{"plus": ..., "cross": ...}`` mapping, or a ``WaveformBackend``
            whose ``generate_td_waveform`` method should serve that name.
    """
    if not hasattr(self, "_waveform_factory"):
        raise AttributeError(
            "TransientSimulator subclasses must initialize _waveform_factory before registering waveform models."
        )

    registered_factories = self.__dict__.setdefault("_registered_waveform_models", {})
    if name in registered_factories:
        if registered_factories[name] == factory:
            return
        raise ValueError(f"Waveform model {name!r} is already registered for this simulator instance.")

    try:
        self._waveform_factory.get_model(name)
    except ValueError:
        pass
    else:
        raise ValueError(f"Waveform model {name!r} is already registered for this simulator instance.")

    if not isinstance(factory, WaveformBackend) and not callable(factory):
        raise TypeError(f"Waveform model {name!r} must be registered with a callable or WaveformBackend.")

    wrapped_factory = (
        _wrap_registered_waveform_backend(name, factory)
        if isinstance(factory, WaveformBackend)
        else _wrap_registered_waveform_callable(factory)
    )
    self._waveform_factory.register_model(name, wrapped_factory)
    registered_factories[name] = factory

def simulate(  # noqa: PLR0913
    self,
    params: Mapping[str, Any],
    detector_names: Sequence[str | CustomDetector],
    background: Mapping[str, TimeSeries] | None = None,
    *,
    sampling_frequency: float,
    minimum_frequency: float,
    earth_rotation: bool = True,
    interpolate_if_offset: bool = True,
) -> DetectorStrainStack:
    """Run the full injection pipeline and return a ``DetectorStrainStack``.

    Calls ``_validate_params``, ``generate_polarizations``,
    ``project_polarizations_to_network``, and ``inject_strain`` in order,
    then assembles the result into a ``DetectorStrainStack``.

    Args:
        params: Source parameters; must contain all ``required_params`` keys
            plus ``right_ascension``, ``declination``, and ``polarization_angle``
            for antenna-response projection.
        detector_names: IFO codes (e.g. ``'H1'``, ``'L1'``, ``'V1'``).
        background: Optional mapping of detector name to background
            ``TimeSeries`` (e.g. noise or zeros). If omitted, the projected
            detector response is returned directly on the waveform time grid.
        sampling_frequency: Sample rate in Hz.
        minimum_frequency: Low-frequency cutoff in Hz.
        earth_rotation: If ``True``, evaluate antenna patterns at
            time-dependent GPS times (recommended for longer signals).
            If ``False``, use a single reference time.
        interpolate_if_offset: If ``True``, interpolate the injection
            when its start is not on a target sample boundary.

    Returns:
        ``DetectorStrainStack`` containing the simulated strain for each
        detector in ``detector_names`` order.
    """
    self._validate_params(params)

    # Validate projection keys before direct params[...] access.
    # _validate_params only checks subclass-defined required_params. Non-CBC subclasses can pass validation and then fail mid-pipeline with KeyError.
    projection_keys = {"right_ascension", "declination", "polarization_angle"}
    missing_projection = projection_keys - set(params)
    if missing_projection:
        raise ValueError(f"Missing required parameters: {sorted(missing_projection)}")

    # Normalise detector entries so dict lookups use plain string keys.
    str_names: list[str] = [d.name if hasattr(d, "name") and not isinstance(d, str) else d for d in detector_names]

    hp, hc = self.generate_polarizations(params, sampling_frequency, minimum_frequency)
    projected = project_polarizations_to_network(
        {"plus": hp, "cross": hc},
        detector_names,
        right_ascension=params["right_ascension"],
        declination=params["declination"],
        polarization_angle=params["polarization_angle"],
        earth_rotation=earth_rotation,
    )
    if background is None:
        return DetectorStrainStack.from_mapping(str_names, projected)
    injected = {
        name: inject_strain(
            background[name],
            projected[name],
            interpolate_if_offset=interpolate_if_offset,
        )
        for name in str_names
    }
    return DetectorStrainStack.from_mapping(str_names, injected)