Batched simulation

def assemble_segments(
    batch: BatchedDetectorStrain,
    *,
    segment_duration: float,
    segment_start_times: Sequence[float],
    backgrounds: Sequence[Mapping[str, TimeSeries]] | None = None,
    interpolate_if_offset: bool = True,
) -> list[DetectorStrainStack]:
    """Scatter the batched signals into fixed-duration data segments (in memory).

    Each output segment spans ``[start, start + segment_duration)``. A signal longer than
    ``segment_duration`` contributes its overlapping part to each of the consecutive segments
    it spans.

    When ``batch`` was generated against a :class:`~gwmock_signal.sampling_grid.SamplingGrid`
    -- see ``output_grid`` on :func:`simulate_cbc_batch` -- every signal already starts on the
    output lattice and superposition is an exact integer-offset add. The segment starts are
    then required to lie on that same grid, and are rejected rather than rounded if they do
    not. Otherwise signals fall between samples and
    :func:`~gwmock_signal.injection.inject_strains_sequential` resamples them, which is only
    accurate for heavily oversampled strain.

    Args:
        batch: Batched per-event/detector strain from :func:`simulate_cbc_batch`.
        segment_duration: Duration of every output segment, in seconds.
        segment_start_times: GPS start time of each output segment (typically a
            contiguous tiling, e.g. ``start + k * segment_duration``).
        backgrounds: Optional per-segment backgrounds, aligned with
            ``segment_start_times``; each maps detector name to a background
            ``TimeSeries`` to inject into. When ``None`` (default), zero-noise
            segments are created.
        interpolate_if_offset: Forwarded to ``inject_strains_sequential`` for
            signals whose start is not on a segment-sample boundary. Unused when the batch
            carries a sampling grid, because then nothing needs interpolating.

    Returns:
        One :class:`~gwmock_signal.multichannel.stack.DetectorStrainStack` per
        entry in ``segment_start_times`` (same order), with channels in
        ``batch.detector_names`` order.
    """
    if backgrounds is not None and len(backgrounds) != len(segment_start_times):
        raise ValueError("backgrounds must be aligned one-to-one with segment_start_times.")

    strain = np.asarray(batch.strain)
    _, _, n_samples = strain.shape
    sampling_frequency = batch.sampling_frequency
    dt = 1.0 / sampling_frequency
    aligned = batch.start_index is not None and batch.grid is not None
    if aligned:
        # Where the buffer actually begins, which differs from epoch + coa_time by the
        # fractional remainder the device absorbed. Using the requested time here would
        # misattribute overlap for events sitting within a fraction of a sample of a boundary.
        signal_start = np.asarray(batch.grid.time_of(batch.start_index), dtype=float)
    else:
        signal_start = batch.epoch + np.asarray(batch.coa_time, dtype=float)
    signal_end = signal_start + n_samples * dt
    n_segment_samples = round(segment_duration * sampling_frequency)
    detectors = batch.detector_names

    if aligned:
        # The catalogue wrapper checks this, but a direct caller reaches here too, and integer
        # overlap arithmetic on a rounded length would silently describe a different interval
        # from the one requested.
        exact_samples = segment_duration * sampling_frequency
        if abs(exact_samples - round(exact_samples)) > _WHOLE_SAMPLE_TOLERANCE:
            raise ValueError(
                f"segment_duration * sampling_frequency must be a whole number of samples for a "
                f"grid-aligned batch; {segment_duration} x {sampling_frequency} = {exact_samples}."
            )
        segment_index = batch.grid.require_on_lattice(
            np.asarray(segment_start_times, dtype=float), name="segment_start_times"
        )
        event_index = np.asarray(batch.start_index, dtype=np.int64)
        # Checked up front rather than per channel: a mismatched background is a caller error, and
        # discovering it on the last segment after assembling every earlier one wastes the work.
        _require_backgrounds_match_segments(
            backgrounds,
            detectors=detectors,
            segment_index=segment_index,
            n_segment_samples=n_segment_samples,
            grid=batch.grid,
            sampling_frequency=sampling_frequency,
        )

    segments: list[DetectorStrainStack] = []
    for k, raw_start in enumerate(segment_start_times):
        seg_start = float(raw_start)
        seg_end = seg_start + segment_duration
        if aligned:
            # Integer lattice arithmetic, not reconstructed GPS times. Both are on one lattice
            # by construction here, so comparing sample counts is exact and implements the
            # half-open convention [start, start + duration) without float round-off deciding
            # whether a signal ending exactly on a boundary belongs to the next segment.
            offset = event_index - int(segment_index[k])
            overlapping = np.nonzero((offset < n_segment_samples) & (offset + n_samples > 0))[0]
        else:
            overlapping = np.nonzero((signal_start < seg_end) & (signal_end > seg_start))[0]
        channels: dict[str, TimeSeries] = {}
        for d, name in enumerate(detectors):
            if aligned:
                channels[name] = _aligned_channel(
                    background=None if backgrounds is None else backgrounds[k][name],
                    strain=strain[:, d],
                    offsets=offset,
                    overlapping=overlapping,
                    n_segment_samples=n_segment_samples,
                    segment_start=seg_start,
                    sampling_frequency=sampling_frequency,
                )
                continue
            if backgrounds is not None:
                background = backgrounds[k][name]
            else:
                background = TimeSeries(np.zeros(n_segment_samples), t0=seg_start, sample_rate=sampling_frequency)
            injections = [TimeSeries(strain[i, d], t0=float(signal_start[i]), dt=dt) for i in overlapping]
            channels[name] = inject_strains_sequential(
                background, injections, interpolate_if_offset=interpolate_if_offset
            )
        segments.append(DetectorStrainStack.from_mapping(detectors, channels))
    return segments

def available_device_memory_bytes() -> int | None:
    """Return the memory limit of the default JAX device, or ``None`` if unknown.

    CPU devices do not report a limit, and neither do some older backends, so callers must
    treat ``None`` as "cannot check" rather than as "no limit".
    """
    try:
        import jax  # noqa: PLC0415 — optional [jax] dep, kept out of module import
    except ImportError:
        return None
    devices = jax.devices()
    if not devices:
        return None
    stats = getattr(devices[0], "memory_stats", lambda: None)()
    if not stats:
        return None
    limit = stats.get("bytes_limit")
    return int(limit) if limit else None

def estimate_batch_memory_bytes(
    n_events: int,
    n_detectors: int,
    n_samples: int,
    *,
    earth_rotation: bool = True,
) -> int:
    """Estimate peak device memory for one :func:`simulate_cbc_batch` call.

    A vmapped batch holds far more than the strain it returns: the measured peak for an
    IMRPhenomXPHM batch was about 28x its own output. The estimate is therefore
    ``n_events * n_samples * 8 * (generation + per_detector * n_detectors)``, with the
    coefficients above.

    !!! warning "One calibration point"

        The coefficients come from a single A100 measurement with IMRPhenomXPHM, and the
        split between detector-independent and per-detector buffers is assumed rather than
        measured. Treat this as an order-of-magnitude guard that produces a useful error
        message, not as an accurate predictor. Approximants with smaller graphs than
        IMRPhenomXPHM will be over-estimated, which only costs a smaller chunk.

    Args:
        n_events: Events in the batch.
        n_detectors: Detectors projected onto.
        n_samples: Samples per event segment.
        earth_rotation: Whether the rotating projection is used, which needs more
            simultaneous buffers per detector.

    Returns:
        Estimated peak bytes.
    """
    if min(n_events, n_detectors, n_samples) < 1:
        raise ValueError("n_events, n_detectors and n_samples must all be >= 1")
    per_detector = _PROJECTION_BUFFERS_PER_DETECTOR * (_ROTATION_BUFFER_MULTIPLIER if earth_rotation else 1.0)
    buffers = _GENERATION_BUFFERS + per_detector * n_detectors
    return int(n_events * n_samples * 8 * buffers)

def recommend_chunk_size(
    n_detectors: int,
    n_samples: int,
    *,
    earth_rotation: bool = True,
    memory_fraction: float = _DEFAULT_MEMORY_FRACTION,
    available_bytes: int | None = None,
) -> int | None:
    """Return the largest event count expected to fit, or ``None`` if unknown.

    Args:
        n_detectors: Detectors projected onto.
        n_samples: Samples per event segment.
        earth_rotation: Whether the rotating projection is used.
        memory_fraction: Fraction of device memory the batch may occupy; must be in ``(0, 1]``.
        available_bytes: Device memory limit; queried from JAX when omitted.

    Returns:
        A chunk size of at least 1, or ``None`` when the device limit is unknown.

    Raises:
        ValueError: If ``memory_fraction`` is outside ``(0, 1]``.
    """
    # A fraction above 1 would recommend a chunk larger than the device, i.e. it would hand
    # back exactly the out-of-memory abort this function exists to prevent.
    if not 0.0 < memory_fraction <= 1.0:
        raise ValueError(f"memory_fraction must be in (0, 1]; got {memory_fraction}.")
    limit = available_device_memory_bytes() if available_bytes is None else available_bytes
    if not limit:
        return None
    per_event = estimate_batch_memory_bytes(1, n_detectors, n_samples, earth_rotation=earth_rotation)
    return max(1, int(limit * memory_fraction // per_event))

def simulate_cbc_batch(  # noqa: PLR0913
    approximant: str,
    detector_names: Sequence[DetectorSpec],
    *,
    sampling_frequency: float,
    minimum_frequency: float,
    parameters: Mapping[str, object],
    backend: RippleBackend | None = None,
    earth_rotation: bool = True,
    output_grid: SamplingGrid | None = None,
) -> BatchedDetectorStrain:
    """Simulate a catalogue of CBC signals on device, one strain per event and detector.

    Evaluates ripple frequency-domain waveforms for the whole catalogue under
    ``jax.vmap`` (a single grid sized worst-case for the longest inspiral), then
    projects each event onto each detector with the JAX antenna pattern and geocenter
    delay and inverse-FFTs to strain. The antenna pattern and delay are evaluated per
    sample by default and once per event at the **segment midpoint** when
    ``earth_rotation=False``, matching the two branches of
    :func:`gwmock_signal.projection.network.project_polarizations_to_network`.

    Args:
        approximant: A supported ripple approximant name.
        detector_names: Built-in LAL interferometer codes (e.g. ``"H1"``, ``"L1"``) and/or
            :class:`~gwmock_signal.detector.CustomDetector` instances. A custom detector is
            resolved through the prefix it registers with LAL, but its output channel is keyed
            by its own ``name``.
        sampling_frequency: Sample rate in Hz.
        minimum_frequency: Low-frequency cutoff in Hz.
        parameters: Mapping of **canonical** gwmock-pop parameter names (no aliases)
            to equal-length 1-D arrays. In addition to the waveform parameters
            (masses, spins, distance, inclination, coa_phase) this must include
            ``right_ascension``, ``declination``, ``polarization_angle`` and
            ``coa_time``.
        backend: Optional configured :class:`RippleBackend` (e.g. with a fixed
            ``segment_duration`` or ``f_ref``). Defaults to ``RippleBackend()``.
        output_grid: Sample lattice the returned strain should start on. When given, each
            event's first sample is placed exactly on the grid and the sub-sample remainder is
            absorbed into the shift the projection already applies -- an exact resampling
            rather than a second, cruder one downstream. Superposition then becomes an
            integer-offset add. When omitted, buffers start at the arbitrary time
            ``epoch + coa_time`` and the consumer must resample; see
            :mod:`gwmock_signal.sampling_grid` for what that costs.
        earth_rotation: If ``True`` (default, matching
            :func:`~gwmock_signal.projection.network.project_polarizations_to_network`),
            evaluate the antenna pattern and geocenter delay per sample and resample the
            polarizations at the delayed times. If ``False``, evaluate both once at the
            segment midpoint and apply the delay as an exact frequency-domain phase
            shift, which is cheaper but only valid for signals short compared with an
            hour. A binary neutron star in the Einstein Telescope band occupies 2048 s
            at 10 Hz and 16384 s at 5 Hz, over which the detector sweeps tens of
            degrees, so ``False`` is not appropriate for that population.

    Returns:
        A :class:`BatchedDetectorStrain` with the ``(n_events, n_detectors,
        n_samples)`` strain and per-event timing metadata.
    """
    import jax.numpy as jnp  # noqa: PLC0415

    backend = backend or RippleBackend()
    # Resolved before anything expensive: an unknown detector should fail here, not after a
    # catalogue has been generated. The two halves are used for different things -- lookup_keys
    # index LAL's registry, output_names key the result -- and conflating them is what limited
    # this path to built-in interferometer codes.
    output_names, lookup_keys = _resolve_detector_specs(detector_names)

    # Before generating anything: the estimate includes the waveform-generation buffers, so a
    # check placed after generation could never fire for a batch that exhausts memory *during*
    # generation -- which is most of the estimate. The grid length does not need the waveform, only
    # the masses (or a pinned segment duration, which _segment_samples handles), so it can be sized
    # up front.
    _check_batch_fits(
        len(np.atleast_1d(np.asarray(_required(parameters, "coa_time")))),
        len(lookup_keys),
        _planned_n_samples(backend, parameters, minimum_frequency, sampling_frequency),
        earth_rotation=earth_rotation,
    )

    fd = backend.generate_fd_polarizations_batch(
        approximant,
        sampling_frequency=sampling_frequency,
        minimum_frequency=minimum_frequency,
        parameters=parameters,
    )
    n_samples = fd.n_samples
    dt = 1.0 / sampling_frequency
    merger_index, epoch = backend.coalescence_placement(n_samples, sampling_frequency)

    right_ascension = jnp.asarray(_required(parameters, "right_ascension"), dtype=jnp.float64)
    declination = jnp.asarray(_required(parameters, "declination"), dtype=jnp.float64)
    polarization_angle = jnp.asarray(_required(parameters, "polarization_angle"), dtype=jnp.float64)
    coa_time = np.asarray(_required(parameters, "coa_time"), dtype=float)

    # Split each event's desired start into a lattice index and the sub-sample remainder the
    # projection must absorb. Without a grid there is nothing to align to and the remainder is
    # zero, which leaves both branches exactly as they were.
    if output_grid is None:
        start_index = None
        alignment_shift = np.zeros_like(coa_time)
    else:
        if output_grid.sampling_frequency != sampling_frequency:
            raise ValueError(
                f"output_grid.sampling_frequency ({output_grid.sampling_frequency}) must equal "
                f"sampling_frequency ({sampling_frequency})."
            )
        start_index, alignment_shift = output_grid.split_index(coa_time + epoch)

    if earth_rotation:
        # The aligned buffer starts one fractional sample earlier than requested, so the
        # sidereal anchor must move with it or F(t) and tau(t) are evaluated up to a full
        # sample after the samples they multiply.
        segment_start_gps = coa_time + epoch - alignment_shift / sampling_frequency
        # No precession, and explicit zero rates rather than an omission. This is the batched
        # *compact-binary* path, so it follows the convention CBC searches use -- `gha = gmst - ra`
        # with the catalogue right ascension, matching `XLALTimeDelayFromEarthCenter` -- which is
        # `project_polarizations_to_network`'s default and what the equivalence tests compare it to.
        # `precess_source_direction` in that function documents why the continuous-wave path differs.
        strain = _project_rotating(
            fd,
            lookup_keys,
            n_samples=n_samples,
            sampling_frequency=sampling_frequency,
            merger_index=merger_index,
            segment_start_gps=segment_start_gps,
            right_ascension=jnp.asarray(right_ascension, dtype=jnp.float64),
            declination=jnp.asarray(declination, dtype=jnp.float64),
            right_ascension_rate=jnp.zeros_like(right_ascension),
            declination_rate=jnp.zeros_like(declination),
            polarization_angle=polarization_angle,
            alignment_shift=alignment_shift,
        )
        return BatchedDetectorStrain(
            strain=strain,
            detector_names=output_names,
            coa_time=coa_time,
            epoch=epoch,
            sampling_frequency=sampling_frequency,
            start_index=start_index,
            grid=output_grid,
        )

    # earth_rotation=False reference time: the midpoint of each event's placed segment.
    midpoint_offset = epoch + 0.5 * (n_samples - 1) * dt
    # Astropy is the single implementation of the sidereal model for both branches and
    # both projection paths; see gwmock_signal.projection.sidereal. The alignment shift moves
    # the buffer, so the midpoint reference time moves with it, as in the rotating branch.
    aligned_midpoint = coa_time + midpoint_offset - alignment_shift / sampling_frequency
    gmst = jnp.asarray(gmst_rad_astropy(aligned_midpoint), dtype=jnp.float64)

    project_batch = _static_projection_kernel(n_samples, sampling_frequency, merger_index)

    per_detector = []
    for key in lookup_keys:
        response, location = reconstructed_geometry(key)
        f_plus, f_cross = antenna_pattern(
            response,
            gmst,
            right_ascension=right_ascension,
            declination=declination,
            polarization_angle=polarization_angle,
        )
        time_delay = time_delay_from_geocenter(location, gmst, right_ascension=right_ascension, declination=declination)
        # The alignment offset is a pure time shift, so it rides on the same exact phase
        # factor as the geocenter delay: no interpolation is involved on this branch at all.
        per_detector.append(
            project_batch(
                fd.frequencies,
                fd.plus,
                fd.cross,
                f_plus,
                f_cross,
                time_delay + jnp.asarray(alignment_shift, dtype=jnp.float64) / sampling_frequency,
            )
        )

    strain = jnp.stack(per_detector, axis=1)  # (n_events, n_detectors, n_samples)
    return BatchedDetectorStrain(
        strain=strain,
        detector_names=output_names,
        coa_time=coa_time,
        epoch=epoch,
        sampling_frequency=sampling_frequency,
        start_index=start_index,
        grid=output_grid,
    )

def simulate_cbc_catalogue(  # noqa: PLR0913
    approximant: str,
    detector_names: Sequence[DetectorSpec],
    *,
    sampling_frequency: float,
    minimum_frequency: float,
    parameters: Mapping[str, object],
    segment_duration: float,
    start_time: float,
    end_time: float,
    backend: RippleBackend | None = None,
    earth_rotation: bool = True,
    n_chirp_mass_bins: int = 1,
    chunk_size: int | None = None,
    memory_fraction: float = _DEFAULT_MEMORY_FRACTION,
    align_to_output_grid: bool = True,
    interpolate_if_offset: bool = True,
) -> list[DetectorStrainStack]:
    """Generate a catalogue on device and assemble it into fixed-duration segments.

    Convenience wrapper that runs :func:`simulate_cbc_batch` and then
    :func:`assemble_segments`, tiling ``[start_time, end_time)`` into contiguous
    zero-noise segments of ``segment_duration``. Signals are placed at their
    ``coa_time`` and split across the segments they span; signals outside the span
    simply do not appear. For non-zero backgrounds use the two-step API
    (:func:`simulate_cbc_batch` then :func:`assemble_segments`) so you can supply a
    background per segment.

    Two independent memory controls (composable):

    - ``chunk_size`` bounds the *peak* generation memory by processing at most that
      many events per batched call. All chunks of a bin share that bin's grid, so
      chunking leaves the *model* untouched -- it agrees with processing the whole
      bin at once to a few times 1e-13 of peak, measured on both the aligned and
      unaligned assembly paths. It is **not** bit-for-bit: XLA picks different
      reduction orderings for different batch shapes, so the same event generated in
      a batch of four and a batch of two differs in the last few bits. Nothing
      physical changes; see :func:`_check_batch_fits`.
    - ``n_chirp_mass_bins`` bounds the *buffer length* by generating heavier events
      on shorter grids. Because each bin uses a different frequency resolution,
      binning is **not** bit-identical to a single grid (see below).

    Args:
        approximant: A supported ripple approximant name.
        detector_names: Built-in LAL interferometer codes (e.g. ``"H1"``, ``"L1"``) and/or
            :class:`~gwmock_signal.detector.CustomDetector` instances, as for
            :func:`simulate_cbc_batch`.
        sampling_frequency: Sample rate in Hz.
        minimum_frequency: Low-frequency cutoff in Hz.
        parameters: Canonical catalogue parameters as struct-of-arrays (see
            :func:`simulate_cbc_batch`).
        segment_duration: Duration of every output segment, in seconds.
        start_time: GPS start of the first segment.
        end_time: GPS time the tiling must cover up to; the final segment is the
            first one whose span reaches or passes ``end_time``.
        backend: Optional configured :class:`RippleBackend`. If it pins a
            ``segment_duration`` that grid is used for every bin (binning then
            saves no *buffer* memory, though it still splits the catalogue into
            per-bin calls, and the run keeps the single-grid model).
        earth_rotation: Forwarded to :func:`simulate_cbc_batch`. Defaults to ``True``, and
            also makes the automatic chunk size smaller, because the rotating path holds
            more buffers per detector.
        n_chirp_mass_bins: Number of chirp-mass groups generated separately, each on
            its own worst-case grid (lightest first), injected on top of the
            earlier bins. ``1`` (default) uses a single grid sized for the
            lowest-mass event. Binned output agrees with a single-grid run only at
            the per-event grid discretization level (a fraction of a percent in
            overlap) — the resolution the per-event path uses.
        chunk_size: Generate at most this many events per batched call (within each bin).
            Model-preserving whatever the value -- it only bounds peak memory, and agrees
            with an unchunked run to a few times 1e-13 of peak rather than bitwise. When omitted,
            a size is chosen from the device memory limit and the grid actually selected
            (see :func:`recommend_chunk_size`) — previously the default was no chunking at
            all, which is what made a large catalogue abort with a bare XLA
            out-of-memory error. Pass an explicit value to override the estimate.
        memory_fraction: Fraction of device memory an automatically chosen chunk may
            occupy. Ignored when ``chunk_size`` is given.
        align_to_output_grid: Generate every batch on the lattice defined by this function's
            own segment starts, so superposition is an exact integer-offset add. Defaults to
            ``True`` because the alternative resamples each signal with a cubic spline, which
            reaches 12% error at half Nyquist. Set ``False`` only to reproduce that older
            behaviour deliberately; it is a legacy mode, not a fallback.
        interpolate_if_offset: Forwarded to :func:`assemble_segments`.

    Returns:
        One :class:`~gwmock_signal.multichannel.stack.DetectorStrainStack` per
        segment, in time order.
    """
    if segment_duration <= 0:
        raise ValueError("segment_duration must be > 0")
    if end_time <= start_time:
        raise ValueError("end_time must be greater than start_time")
    if n_chirp_mass_bins < 1:
        raise ValueError("n_chirp_mass_bins must be >= 1")
    if chunk_size is not None and chunk_size < 1:
        raise ValueError("chunk_size must be >= 1")

    backend = backend or RippleBackend()
    # Resolved once here purely so an unusable detector is reported before any generation. Each
    # chunk resolves again inside simulate_cbc_batch, which is cheap: a CustomDetector caches its
    # LAL detector on first use, and a built-in code is a dictionary lookup.
    _resolve_detector_specs(detector_names)
    n_segments = int(np.ceil((end_time - start_time) / segment_duration))
    segment_start_times = start_time + np.arange(n_segments) * segment_duration

    output_grid: SamplingGrid | None = None
    if align_to_output_grid:
        samples_per_segment = segment_duration * sampling_frequency
        if abs(samples_per_segment - round(samples_per_segment)) > _WHOLE_SAMPLE_TOLERANCE:
            raise ValueError(
                f"segment_duration * sampling_frequency must be a whole number of samples for the "
                f"segments to share one lattice; {segment_duration} x {sampling_frequency} = "
                f"{samples_per_segment}. Adjust one of them, or pass align_to_output_grid=False to "
                f"accept resampled superposition."
            )
        # One grid for the whole catalogue, so every chunk and every chirp-mass bin lands on the
        # same lattice and can be superposed with integer offsets. The starts are then rebuilt
        # *from integer sample indices* rather than kept as start_time + k * segment_duration:
        # over a long span the repeated float multiplication accumulates representation error,
        # so a start intended to be on the lattice can drift off it even when the spacing is a
        # whole number of samples.
        output_grid = SamplingGrid(epoch=float(start_time), sampling_frequency=sampling_frequency)
        segment_start_times = output_grid.time_of(np.arange(n_segments, dtype=np.int64) * round(samples_per_segment))

    segments: list[DetectorStrainStack] | None = None
    for bin_indices in _chirp_mass_bins(parameters, n_chirp_mass_bins):
        # Pin the grid to this bin's worst case so every chunk of the bin shares it
        # (so chunking within the bin does not change the model). A user-pinned backend is left as-is.
        bin_backend = _bin_backend(backend, parameters, bin_indices, minimum_frequency, sampling_frequency)
        # Size the chunk from this bin's own grid: bins differ in n_samples by design, so a
        # single catalogue-wide chunk size would be wrong for all but one of them.
        effective_chunk = chunk_size
        if effective_chunk is None:
            bin_samples = _bin_n_samples(bin_backend, parameters, bin_indices, minimum_frequency, sampling_frequency)
            effective_chunk = recommend_chunk_size(
                len(tuple(detector_names)),
                bin_samples,
                earth_rotation=earth_rotation,
                memory_fraction=memory_fraction,
            )
        for chunk_indices in _count_chunks(bin_indices, effective_chunk):
            chunk_parameters = {key: np.asarray(values)[chunk_indices] for key, values in parameters.items()}
            batch = simulate_cbc_batch(
                approximant,
                detector_names,
                sampling_frequency=sampling_frequency,
                minimum_frequency=minimum_frequency,
                parameters=chunk_parameters,
                backend=bin_backend,
                earth_rotation=earth_rotation,
                output_grid=output_grid,
            )
            # Chain groups: inject each on top of the segments built from earlier ones.
            backgrounds = [stack.to_dict() for stack in segments] if segments is not None else None
            segments = assemble_segments(
                batch,
                segment_duration=segment_duration,
                segment_start_times=segment_start_times,
                backgrounds=backgrounds,
                interpolate_if_offset=interpolate_if_offset,
            )
    return segments if segments is not None else []