pair_kernel

def prepare_box_samples(data, masks, Num_position, Num_shape, *, shapes, ellipticity, base,
                        require_full_masks=False):
    """Apply masks and compute per-galaxy ellipticity size for a Box measurement.

    Reproduces, verbatim, the mask-application and ``e`` computation shared by
    every Box (rp, pi) / (r, mu_r) counting method: mask defaulting rules
    (``Position``/``Position_shape_sample`` masks default to "select all";
    ``weight``/``weight_shape_sample`` default to the coordinate mask and are
    written back into ``masks`` in place, matching the legacy fallback-default
    behaviour other code paths rely on), then ``e = f(q)`` for the requested
    ellipticity definition.

    ``require_full_masks`` selects the mask-indexing convention: the non-jk methods
    default a missing ``Position``/``Position_shape_sample``/``Axis_Direction``/``q``
    mask (``.get`` with "select all" / coordinate-mask fallbacks), whereas the box
    jackknife methods index ``masks["Position"]`` etc. **directly** and raise KeyError
    on a partial dict — pass ``require_full_masks=True`` to reproduce that (see
    REFACTOR_PLAN.md section 3.1). ``weight``/``weight_shape_sample`` still default to
    the coordinate mask in both modes.

    Parameters
    ----------
    data : dict
        The object's ``self.data``.
    masks : dict or None
        Per-call mask dict; mutated in place to inject default ``weight``/
        ``weight_shape_sample`` masks when absent (legacy behaviour).
    Num_position, Num_shape : int
        Full (unmasked) sample sizes, used as the default "select all" mask length.
    shapes : bool
        If False, skip ``Axis_Direction``/``q``/``e`` (DD-only / count_pairs paths).
    ellipticity : str
        'distortion' or 'ellipticity'; see ``MeasureIABase.get_ellipticity``.
    base : object
        The calling instance (unused today; accepted for interface symmetry
        with future steps that need e.g. responsivity_correction here).

    Returns
    -------
    SampleSet
    """
    if masks is None:
        positions = data["Position"]
        positions_shape_sample = data["Position_shape_sample"]
        weight = data["weight"]
        weight_shape = data["weight_shape_sample"]
        axis_direction_v = data["Axis_Direction"] if shapes else None
        q = data["q"] if shapes else None
    else:
        if require_full_masks:
            pos_mask = masks["Position"]
            shape_mask = masks["Position_shape_sample"]
        else:
            pos_mask = masks.get("Position", np.ones(Num_position, dtype=bool))
            shape_mask = masks.get("Position_shape_sample", np.ones(Num_shape, dtype=bool))
        positions = data["Position"][pos_mask]
        positions_shape_sample = data["Position_shape_sample"][shape_mask]
        if "weight" not in masks:
            masks["weight"] = pos_mask
        if "weight_shape_sample" not in masks:
            masks["weight_shape_sample"] = shape_mask
        weight = data["weight"][masks["weight"]]
        weight_shape = data["weight_shape_sample"][masks["weight_shape_sample"]]
        if shapes:
            if require_full_masks:
                dir_mask = masks["Axis_Direction"]
                q_mask = masks["q"]
            else:
                dir_mask = masks.get("Axis_Direction", shape_mask)
                q_mask = masks.get("q", shape_mask)
            axis_direction_v = data["Axis_Direction"][dir_mask]
            q = data["q"][q_mask]
        else:
            axis_direction_v = None
            q = None

    axis_direction = None
    e = None
    if shapes:
        axis_direction_len = np.sqrt(np.sum(axis_direction_v ** 2, axis=1))
        axis_direction = (axis_direction_v.transpose() / axis_direction_len).transpose()
        if ellipticity == 'distortion':
            e = (1 - q ** 2) / (1 + q ** 2)
        elif ellipticity == 'ellipticity':
            e = (1 - q) / (1 + q)
        else:
            raise ValueError("Invalid value for ellipticity. Choose 'distortion' or 'ellipticity'.")

    LOS_ind = data["LOS"]
    not_LOS = np.array([0, 1, 2])[np.isin([0, 1, 2], LOS_ind, invert=True)]

    # How many objects sit in both samples, after masking. The analytic RR needs this:
    # a shape galaxy cannot pair with itself, and the pair loop already drops that
    # self-pair because the separation window starts at r_min > 0. Recorded on the
    # caller so every RR call in the backends can use one consistent value.
    # Imported here rather than at module scope to keep the import graph acyclic.
    from .measure_IA_base import count_overlap
    override = getattr(base, "_num_overlap_override", None)
    base.num_overlap = (int(override) if override is not None
                        else count_overlap(positions, positions_shape_sample))

    return SampleSet(
        pos=positions, pos_shape=positions_shape_sample,
        weight=weight, weight_shape=weight_shape,
        axis_direction=axis_direction, e=e,
        LOS_ind=LOS_ind, not_LOS=not_LOS,
    )

def prepare_lightcone_samples(data, masks, *, shapes, cosmology, over_h,
                              responsivity_correction, base, print_num=True):
    """Apply masks and build the 3D comoving sky geometry for a Lightcone measurement.

    Reproduces, verbatim, the shared head of every lightcone (rp, pi) / (r, mu_r) counting
    method: RA/DEC/Redshift/e1/e2 mask-application (direct ``masks["RA"]``-style indexing —
    a partial dict raises KeyError, as in the legacy; ``weight``/``weight_shape_sample``
    default to the ``RA``/``RA_shape_sample`` coordinate masks and are written back into
    ``masks`` in place), redshift → comoving distance via ``pyccl`` (default cosmology when
    None; ``over_h`` scales by ``h``), the unit-sphere direction vectors, the position/shape
    comoving vectors ``s_pos``/``s_shape``, and — for ``shapes=True`` — the per-shape
    ellipticity ``e = (e1, e2)`` **pre-scaled by 1/(2R)** when ``responsivity_correction``
    (responsivity is baked into ``e`` here, not divided in the pair loop, unlike the box
    families — REFACTOR_PLAN.md section 3.2), plus the local ``east``/``north`` sky basis at
    each position.

    ``SampleSet`` layout for the lightcone geometry: ``pos = s_pos`` (N,3), ``pos_shape =
    s_shape`` (M,3), ``weight``/``weight_shape``; and when ``shapes``: ``e`` (M,2), ``east``
    (N,3), ``north`` (N,3), ``n_pos`` (N,3).

    Parameters
    ----------
    data : dict
        The object's ``self.data``.
    masks : dict or None
        Per-call mask dict; mutated in place to inject default ``weight``/
        ``weight_shape_sample`` masks when absent (legacy behaviour).
    shapes : bool
        If False, skip ``e1``/``e2``/responsivity and the ``east``/``north`` basis
        (DD-only / count_pairs paths).
    cosmology : pyccl.Cosmology or None
        Cosmology for redshift→comoving distance; a fixed default is built (and, when
        ``print_num``, announced) if None, matching the legacy.
    over_h : bool
        If True, multiply comoving distances by ``h`` (positions in cMpc/h).
    responsivity_correction : bool
        If True (and ``shapes``), pre-scale ``e1, e2`` by ``1/(2R)`` with ``R`` the
        weighted responsivity over the shape sample. Note the lightcone default is False
        (unlike the box families) — pass ``getattr(self, "responsivity_correction", False)``.
    base : object
        The calling instance (accepted for interface symmetry; unused here).
    print_num : bool
        Gate on the "No cosmology given" informational print, matching the legacy.

    Returns
    -------
    SampleSet
    """
    if masks is None:
        redshift = data["Redshift"]
        redshift_shape_sample = data["Redshift_shape_sample"]
        RA = data["RA"]
        RA_shape_sample = data["RA_shape_sample"]
        DEC = data["DEC"]
        DEC_shape_sample = data["DEC_shape_sample"]
        weight = data["weight"]
        weight_shape = data["weight_shape_sample"]
        if shapes:
            e1 = data["e1"]
            e2 = data["e2"]
    else:
        redshift = data["Redshift"][masks["Redshift"]]
        redshift_shape_sample = data["Redshift_shape_sample"][masks["Redshift_shape_sample"]]
        RA = data["RA"][masks["RA"]]
        RA_shape_sample = data["RA_shape_sample"][masks["RA_shape_sample"]]
        DEC = data["DEC"][masks["DEC"]]
        DEC_shape_sample = data["DEC_shape_sample"][masks["DEC_shape_sample"]]
        if shapes:
            e1 = data["e1"][masks["e1"]]
            e2 = data["e2"][masks["e2"]]
        if "weight" not in masks:
            masks["weight"] = masks["RA"]
        if "weight_shape_sample" not in masks:
            masks["weight_shape_sample"] = masks["RA_shape_sample"]
        weight = data["weight"][masks["weight"]]
        weight_shape = data["weight_shape_sample"][masks["weight_shape_sample"]]

    Num_position = len(RA)

    if cosmology is None:
        cosmology = ccl.Cosmology(Omega_c=0.225, Omega_b=0.045, sigma8=0.8, h=0.7, n_s=1.0)
        if print_num:
            print("No cosmology given, using Omega_m=0.27, Omega_b=0.045, sigma8=0.8, h=0.7, n_s=1.")
    h = cosmology["h"]

    LOS_all = ccl.comoving_radial_distance(cosmology, 1 / (1 + redshift))
    LOS_all_shape_sample = ccl.comoving_radial_distance(cosmology, 1 / (1 + redshift_shape_sample))
    if over_h:
        LOS_all *= h
        LOS_all_shape_sample *= h

    e = None
    east = None
    north = None
    if shapes:
        if responsivity_correction:
            R = sum(weight_shape * (1 - (e1 ** 2 + e2 ** 2) / 2.0)) / sum(weight_shape)
            e1, e2 = e1 / (2 * R), e2 / (2 * R)
        e = np.array([e1, e2]).transpose()

    RA_rad = RA / 180 * np.pi
    RA_shape_sample_rad = RA_shape_sample / 180 * np.pi
    DEC_rad = DEC / 180 * np.pi
    DEC_shape_sample_rad = DEC_shape_sample / 180 * np.pi
    n_shape = np.array([np.cos(DEC_shape_sample_rad) * np.cos(RA_shape_sample_rad),
                        np.cos(DEC_shape_sample_rad) * np.sin(RA_shape_sample_rad),
                        np.sin(DEC_shape_sample_rad)]).transpose()
    s_shape = n_shape * np.array([LOS_all_shape_sample]).transpose()
    n_pos = np.array([np.cos(DEC_rad) * np.cos(RA_rad),
                      np.cos(DEC_rad) * np.sin(RA_rad),
                      np.sin(DEC_rad)]).transpose()
    if shapes:
        east = np.array([-np.sin(RA_rad), np.cos(RA_rad), np.zeros(Num_position)]).transpose()
        north = np.array([
            -np.sin(DEC_rad) * np.cos(RA_rad),
            -np.sin(DEC_rad) * np.sin(RA_rad),
            np.cos(DEC_rad)
        ]).transpose()
    s_pos = np.array([LOS_all]).transpose() * n_pos

    return SampleSet(
        pos=s_pos, pos_shape=s_shape,
        weight=weight, weight_shape=weight_shape,
        e=e, east=east, north=north, n_pos=n_pos,
    )

def spatial_order(coords):
    """Visit order that makes consecutive chunks compact in space.

    ``accumulate`` processes the chunked sample 100 at a time, builds a KDTree of
    each chunk and queries it against the full tree of the other sample. A
    dual-tree traversal can only prune when the chunk occupies a small region, and
    nothing about a catalogue's storage order guarantees that: a sample stored
    grouped by halo, or by id, or shuffled, gives chunks whose bounding box spans
    essentially the whole volume, and the query degenerates towards brute force.

    Measured on the package's own mock at 100,000 galaxies, the chunk extent was
    619 Mpc inside a 711 Mpc box, the query took 6.20 s, and its cost scaled as
    N^1.63. Visiting the same galaxies in the order returned here made the chunks
    243 Mpc across, the query 0.34 s, and the scaling N^1.01 -- an 18x speed-up
    from the identical set of pairs (benchmarks/FINDINGS.md F5).

    The key is a Morton (Z-order) code on a 1024-per-axis grid spanning the
    sample, which interleaves the bits of the per-axis cell indices so that points
    close in space stay close in the ordering. Works for 2D coordinates (the
    (rp, pi) box binning projects out the line of sight) as well as 3D.

    Parameters
    ----------
    coords : (N, D) ndarray
        Coordinates in the same metric the KDTree is built on, i.e. whatever
        ``binning.tree_coords`` returns.

    Returns
    -------
    (N,) ndarray of int
        Indices into ``coords``, spatially ordered. A stable sort, so the order is
        deterministic for a given input.

    """
    coords = np.asarray(coords, dtype=float)
    n, dim = coords.shape
    bits = 10
    lo = coords.min(axis=0)
    span = coords.max(axis=0) - lo
    span[span <= 0] = 1.0
    cell = np.clip(((coords - lo) / span * (1 << bits)).astype(np.int64),
                   0, (1 << bits) - 1)
    key = np.zeros(n, dtype=np.int64)
    for b in range(bits):
        for d in range(dim):
            key |= ((cell[:, d] >> b) & 1) << (b * dim + d)
    return np.argsort(key, kind="stable")

def accumulate(sample_set, binning, *, base, R=None, shapes=True,
               chunk_axis="shape", chunk_size_outer=100, jk=False, num_box=None,
               pos_tree=None, shape_tree=None, backend="tree",
               per_galaxy=False, per_galaxy_proj=None, per_galaxy_jk=False,
               per_galaxy_jk_sparse=False):
    """Run the pair-accumulation loop and return the resulting grids.

    Implemented so far: box geometry, ``BoxRpPi`` / ``BoxRMuR`` binnings,
    ``chunk_axis="shape"``, backends ``"tree"`` and ``"brute"``, optional jackknife.
    Later migration steps (REFACTOR_PLAN.md section 6, steps 6-7) extend this same
    function to the lightcone geometry.

    ``jk=True`` (with ``num_box`` = number of jackknife realisations) additionally
    accumulates the union-deletion per-realisation grids ``DD_jk`` (and, when
    ``shapes``, the raw ``Splus_D_jk``): every pair contributes to the shape's patch,
    and to the position's patch only where the two patches differ. ``sample_set``
    must then carry ``jk_pos`` (position-aligned) and ``jk_shape`` (shape-aligned)
    patch indices. This is a per-batch quantity in the mp path — the parent sums the
    partial jk grids and computes ``R_jk`` separately via ``compute_R_jk``.

    Iteration order (outer loop over shape-sample chunks of ``chunk_size_outer``,
    inner loop over the chunk, vectorized ``np.add.at`` per shape galaxy) is fixed
    by the float-summation-order rule in REFACTOR_PLAN.md section 4 and must not
    change without re-deriving bit-identity against the legacy tree/mp paths.

    ``backend`` selects how each shape galaxy's candidate positions are chosen:
      - ``"tree"``: KDTree ball query at ``r_max`` against the position tree (the
        legacy tree/mp order — bit-identical). There is no inner ``r_min`` query:
        the binning mask already drops everything below ``r_bins[0]``, so the extra
        query and the set difference it fed were removed as redundant
        (benchmarks/FINDINGS.md F1). Candidates below ``r_min`` are therefore passed
        to ``bin_pairs`` and masked out there.
      - ``"brute"``: every position is a candidate (full cross-join per chunk);
        the ``[r_min, r_max)`` window is applied by the binning mask. This runs on
        the *same* shape-chunk order as the tree backend rather than the legacy
        brute's position-outer order, so it matches the legacy brute only to
        floating-point tolerance (``allclose``), not bit-identically — a
        deliberate consolidation choice (REFACTOR_PLAN.md section 4). It counts
        exactly the same pairs the legacy brute did (same window mask), so integer
        (unit-weight) DD grids still match exactly.

    ``per_galaxy=True`` (box path only) additionally resolves the same sums **per shape
    galaxy**, i.e. without summing over the shape sample. It changes nothing about the
    grids above: the per-galaxy arrays are accumulated in their own branch and satisfy
    ``per_galaxy_arrays.sum(axis=0) == the corresponding grid`` by construction. The
    galaxy axis is indexed as in ``sample_set.pos_shape`` (so, in the multiprocessing
    path, local to the batch — the caller concatenates batches in order).

      - ``per_galaxy_proj=None``: ``DD_gal``/``Splus_D_gal`` have shape
        ``(M, num_bins_r, num_bins_pi)``, the same axes as the grids.
      - ``per_galaxy_proj=W`` with ``W`` of shape ``(num_bins_r, num_bins_pi)``: the
        second bin axis is contracted with ``W`` as it is accumulated, giving
        ``Splus_D_gal`` of shape ``(M, num_bins_r)``. This is how the multipole kernel
        (Legendre weight / RR, both analytic in the box) is folded in without ever
        materialising the ``mu_r`` axis per galaxy. ``DD_gal`` is *not* contracted with
        ``W`` — it is the plain pair count per radial bin, which is what a per-galaxy
        regression design matrix needs — so it has shape ``(M, num_bins_r)`` too.

    ``per_galaxy_jk_sparse=True`` stores that decomposition only for the patches each
    galaxy actually has pairs in. A galaxy's neighbours span a ball of radius ``r_max``, so
    it reaches at most a handful of sub-boxes however many there are in total, and the rest
    of the patch axis is *structurally* zero. The dense ``(M, num_box, num_bins_r)`` form
    becomes ``gal_jk_patches`` ``(M, K)`` plus ``*_gal_jk_values`` ``(M, K, num_bins_r)``,
    with ``K`` the largest number of patches any one galaxy touches and unused slots padded
    with patch ``-1`` and zero values. ``DD`` and ``S+D`` share one patch array, since a
    pair contributes to both. This is a pure change of representation: the stored floats
    are the same values, so results are bit-identical (asserted by a test).

    ``per_galaxy_jk=True`` further decomposes the per-galaxy arrays by the jackknife
    patch of the *position-sample* partner, giving ``(M, num_box, num_bins_r)``. Combined
    with the patch of the shape galaxy itself this reproduces the union-deletion of
    ``DD_jk``/``Splus_D_jk`` exactly: realisation ``n`` drops every shape galaxy with
    ``jk_shape == n`` and subtracts column ``n`` from the rest. It requires
    ``per_galaxy_proj`` (and ``sample_set.jk_pos``/``jk_shape``), since the undecomposed
    ``mu_r`` axis would multiply the footprint by ``num_bins_pi``.

    All of this is inert when ``per_galaxy=False`` (the default): the pair loop, its
    iteration order and its float summation order are untouched, so the bit-identity
    guarantee of REFACTOR_PLAN.md section 4 still holds and normal measurements pay
    nothing beyond one branch test per shape galaxy.

    ``pos_tree`` may be a prebuilt ``KDTree`` over ``sample_set.pos[:, not_LOS]``
    (tree backend only). The multiprocessing path passes the tree it built once in
    the parent process (shared to every worker) rather than rebuilding it per
    batch; when None the tree is built here. ``sample_set.pos`` must be the same
    full position array the tree was built from either way.
    """
    if backend not in ("tree", "brute"):
        raise NotImplementedError(
            f"pair_kernel.accumulate: unknown backend {backend!r} (expected 'tree' or 'brute')."
        )
    if chunk_axis == "position":
        if not isinstance(binning, (SkyRpPi, SkyRMuR)):
            raise NotImplementedError(
                "pair_kernel.accumulate: chunk_axis='position' (lightcone) requires a "
                "SkyRpPi / SkyRMuR binning."
            )
        if jk and num_box is None:
            raise ValueError("pair_kernel.accumulate: jk=True requires num_box (num_jk).")
        if backend == "brute" and shape_tree is not None:
            raise ValueError(
                "pair_kernel.accumulate: shape_tree is meaningless with backend='brute'."
            )
        return _accumulate_lightcone(
            sample_set, binning, base=base, shapes=shapes,
            chunk_size_outer=chunk_size_outer, backend=backend,
            jk=jk, num_jk=num_box, shape_tree=shape_tree,
        )
    if chunk_axis != "shape":
        raise NotImplementedError(
            "pair_kernel.accumulate: only chunk_axis='shape' (box) and "
            "chunk_axis='position' (lightcone) are implemented."
        )
    if not isinstance(binning, (BoxRpPi, BoxRMuR)):
        raise NotImplementedError(
            "pair_kernel.accumulate: only BoxRpPi / BoxRMuR binnings are implemented "
            "for the box (chunk_axis='shape') path."
        )
    if backend == "brute" and pos_tree is not None:
        raise ValueError(
            "pair_kernel.accumulate: pos_tree is meaningless with backend='brute' "
            "(no KDTree is built); pass pos_tree only with backend='tree'."
        )
    if jk and num_box is None:
        raise ValueError("pair_kernel.accumulate: jk=True requires num_box.")
    if per_galaxy_jk and not per_galaxy:
        raise ValueError(
            "pair_kernel.accumulate: per_galaxy_jk=True requires per_galaxy=True."
        )
    if per_galaxy_jk and per_galaxy_proj is None:
        raise ValueError(
            "pair_kernel.accumulate: per_galaxy_jk=True requires per_galaxy_proj, so that "
            "the mu_r/pi axis is contracted before the per-patch decomposition (an "
            "undecomposed axis would multiply the per-galaxy footprint by num_bins_pi)."
        )
    if per_galaxy_jk and num_box is None:
        raise ValueError("pair_kernel.accumulate: per_galaxy_jk=True requires num_box.")
    if per_galaxy_jk_sparse and not per_galaxy_jk:
        raise ValueError(
            "pair_kernel.accumulate: per_galaxy_jk_sparse=True requires per_galaxy_jk=True."
        )
    if per_galaxy_jk and (sample_set.jk_pos is None or sample_set.jk_shape is None):
        raise ValueError(
            "pair_kernel.accumulate: per_galaxy_jk=True requires sample_set.jk_pos and "
            "sample_set.jk_shape."
        )
    if per_galaxy_proj is not None:
        per_galaxy_proj = np.asarray(per_galaxy_proj, dtype=float)
        expected = (binning.num_bins_r, binning.num_bins_pi)
        if per_galaxy_proj.shape != expected:
            raise ValueError(
                f"pair_kernel.accumulate: per_galaxy_proj has shape {per_galaxy_proj.shape}, "
                f"expected {expected} (num_bins_r, num_bins_pi)."
            )

    DD = np.array([[0.0] * binning.num_bins_pi] * binning.num_bins_r)
    Splus_D = np.array([[0.0] * binning.num_bins_pi] * binning.num_bins_r) if shapes else None
    Scross_D = np.array([[0.0] * binning.num_bins_pi] * binning.num_bins_r) if shapes else None
    DD_jk = np.zeros((num_box, binning.num_bins_r, binning.num_bins_pi)) if jk else None
    Splus_D_jk = np.zeros((num_box, binning.num_bins_r, binning.num_bins_pi)) if (jk and shapes) else None

    DD_gal = Splus_D_gal = DD_gal_jk = Splus_D_gal_jk = None
    if per_galaxy:
        num_gal = len(sample_set.pos_shape)
        if per_galaxy_proj is None:
            gal_shape = (num_gal, binning.num_bins_r, binning.num_bins_pi)
        else:
            gal_shape = (num_gal, binning.num_bins_r)
        DD_gal = np.zeros(gal_shape)
        Splus_D_gal = np.zeros(gal_shape) if shapes else None
        if per_galaxy_jk and not per_galaxy_jk_sparse:
            jk_gal_shape = (num_gal, num_box, binning.num_bins_r)
            DD_gal_jk = np.zeros(jk_gal_shape)
            Splus_D_gal_jk = np.zeros(jk_gal_shape) if shapes else None
    # Sparse jackknife decomposition: accumulate one outer chunk at a time into a small
    # dense buffer, then keep only the patches that chunk actually touched. The buffer is
    # (chunk_size_outer, num_box, num_bins_r), which stays small however large num_box is.
    sp_buffer_DD = sp_buffer_Splus = None
    sp_chunks_patches, sp_chunks_DD, sp_chunks_Splus = [], [], []
    if per_galaxy and per_galaxy_jk_sparse:
        sp_buffer_DD = np.zeros((chunk_size_outer, num_box, binning.num_bins_r))
        sp_buffer_Splus = np.zeros((chunk_size_outer, num_box, binning.num_bins_r)) if shapes else None

    positions = sample_set.pos
    positions_shape_sample = sample_set.pos_shape
    weight = sample_set.weight
    weight_shape = sample_set.weight_shape
    not_LOS = sample_set.not_LOS
    LOS_ind = sample_set.LOS_ind
    jk_pos = sample_set.jk_pos

    if backend == "brute":
        all_positions = np.arange(len(positions))
        # No tree, so nothing to prune and nothing to gain; keeping array order
        # here also keeps the brute path bit-identical to previous releases.
        order = np.arange(len(positions_shape_sample))
    else:
        if pos_tree is None:
            pos_tree = KDTree(binning.tree_coords(positions, not_LOS), boxsize=base.boxsize)
        # Visit shape galaxies in spatial order so each chunk's KDTree covers a
        # small region and the dual-tree query can prune (FINDINGS.md F5). The
        # order is computed on tree_coords, i.e. the same metric the trees use.
        # Results are unchanged up to float summation order; the per-galaxy
        # branches below index their outputs by the *original* galaxy id so that
        # what the caller gets back is still in the caller's order.
        order = spatial_order(binning.tree_coords(positions_shape_sample, not_LOS))
    for i in np.arange(0, len(positions_shape_sample), chunk_size_outer):
        i2 = min(len(positions_shape_sample), i + chunk_size_outer)
        sel = order[i:i2]
        if sp_buffer_DD is not None:
            sp_buffer_DD[:i2 - i] = 0.0
            if sp_buffer_Splus is not None:
                sp_buffer_Splus[:i2 - i] = 0.0
        positions_shape_sample_i = positions_shape_sample[sel]
        weight_shape_i = weight_shape[sel]
        if shapes:
            axis_direction_i = sample_set.axis_direction[sel]
            e_i = sample_set.e[sel]
        if jk:
            jk_shape_i = sample_set.jk_shape[sel]
        if backend == "brute":
            # every position is a candidate for every shape in the chunk
            ind_rbin_i = [all_positions] * len(positions_shape_sample_i)
        else:
            shape_tree = KDTree(binning.tree_coords(positions_shape_sample_i, not_LOS), boxsize=base.boxsize)
            # One query, at the outer radius only. The inner r_min query and the
            # set difference it fed were redundant: binning.bin_pairs already
            # drops every pair with separation < r_bins[0], and r_bins[0] == r_min
            # exactly, on the same metric this tree is built on. The extra
            # candidates are masked out in the same (sorted) order, so the
            # surviving pairs and their float summation order are unchanged.
            # Worth ~30% of a box measurement -- see benchmarks/FINDINGS.md F1.
            # asarray keeps the downstream fancy-indexing off Python lists.
            ind_rbin_i = [np.asarray(c, dtype=np.intp)
                          for c in shape_tree.query_ball_tree(pos_tree, binning.query_r_max)]

        for n in np.arange(0, len(positions_shape_sample_i)):
            if len(ind_rbin_i[n]) > 0:
                separation = positions_shape_sample_i[n] - positions[ind_rbin_i[n]]
                if base.periodicity:
                    separation[separation > base.L_0p5] -= base.boxsize
                    separation[separation < -base.L_0p5] += base.boxsize

                mask, ind_r, ind_pi, projected_sep, proj_len = binning.bin_pairs(
                    separation, not_LOS, LOS_ind
                )

                if shapes:
                    with np.errstate(invalid='ignore'):
                        separation_dir = (projected_sep.transpose() / proj_len).transpose()
                        phi = np.arccos(
                            separation_dir[:, 0] * axis_direction_i[n, 0]
                            + separation_dir[:, 1] * axis_direction_i[n, 1]
                        )
                    e_plus, e_cross = base.get_ellipticity(e_i[n], phi)
                    e_plus[np.isnan(e_plus)] = 0.0
                    e_cross[np.isnan(e_cross)] = 0.0
                    np.add.at(Splus_D, (ind_r, ind_pi),
                              (weight[ind_rbin_i[n]][mask] * weight_shape_i[n] * e_plus[mask]) / (2 * R))
                    np.add.at(Scross_D, (ind_r, ind_pi),
                              (weight[ind_rbin_i[n]][mask] * weight_shape_i[n] * e_cross[mask]) / (2 * R))
                np.add.at(DD, (ind_r, ind_pi), weight[ind_rbin_i[n]][mask] * weight_shape_i[n])

                if jk:
                    # union (two-sided) deletion: every pair contributes to the shape's
                    # patch, and to the position's patch only where the two patches differ.
                    # Splus_D_jk stores the raw (un-/2R) S+ contribution; responsivity is
                    # applied later in the wrapper reduction. Order matches the legacy jk
                    # loop (Splus_D_jk before DD_jk).
                    shape_patch = jk_shape_i[n]
                    pos_patches = jk_pos[ind_rbin_i[n]][mask]
                    pos_diff = np.where(pos_patches != shape_patch)[0]
                    w_pairs = weight[ind_rbin_i[n]][mask] * weight_shape_i[n]
                    if shapes:
                        np.add.at(Splus_D_jk, (shape_patch, ind_r, ind_pi), w_pairs * e_plus[mask])
                        np.add.at(Splus_D_jk,
                                  (pos_patches[pos_diff], ind_r[pos_diff], ind_pi[pos_diff]),
                                  (w_pairs * e_plus[mask])[pos_diff])
                    np.add.at(DD_jk, (shape_patch, ind_r, ind_pi), w_pairs)
                    np.add.at(DD_jk,
                              (pos_patches[pos_diff], ind_r[pos_diff], ind_pi[pos_diff]),
                              w_pairs[pos_diff])

                if per_galaxy:
                    # Same pairs, same weights, resolved on the shape-galaxy axis. Kept
                    # last and in its own branch so that nothing above changes when
                    # per_galaxy is off.
                    # the caller's galaxy id, not the visit position: the loop
                    # walks the sample in spatial order (see `order` above), so
                    # i + n would attribute each galaxy's row to a different one
                    gj = sel[n]
                    w_pairs_gal = weight[ind_rbin_i[n]][mask] * weight_shape_i[n]
                    if per_galaxy_proj is None:
                        np.add.at(DD_gal[gj], (ind_r, ind_pi), w_pairs_gal)
                        if shapes:
                            np.add.at(Splus_D_gal[gj], (ind_r, ind_pi),
                                      w_pairs_gal * e_plus[mask] / (2 * R))
                    else:
                        proj_gal = per_galaxy_proj[ind_r, ind_pi]
                        # DD_gal stays the plain pair count per radial bin: the projection
                        # belongs to the estimator (S+D), not to the design matrix.
                        np.add.at(DD_gal[gj], ind_r, w_pairs_gal)
                        if shapes:
                            np.add.at(Splus_D_gal[gj], ind_r,
                                      proj_gal * w_pairs_gal * e_plus[mask] / (2 * R))
                    if per_galaxy_jk:
                        # Decompose by the position partner's patch only: the shape
                        # galaxy's own patch is known from jk_shape, and delete-one drops
                        # such a galaxy wholesale. Reconstructing realisation n as
                        #   sum_{j: jk_shape[j] != n} (total_j - column_n_j)
                        # therefore reproduces the union deletion of DD_jk / Splus_D_jk.
                        # Splus_D_gal_jk is raw (not divided by 2R), matching Splus_D_jk.
                        pos_patches_gal = jk_pos[ind_rbin_i[n]][mask]
                        # dense target, or this chunk's buffer row when storing sparsely
                        target_DD = sp_buffer_DD[n] if per_galaxy_jk_sparse else DD_gal_jk[gj]
                        np.add.at(target_DD, (pos_patches_gal, ind_r), w_pairs_gal)
                        if shapes:
                            target_S = (sp_buffer_Splus[n] if per_galaxy_jk_sparse
                                        else Splus_D_gal_jk[gj])
                            np.add.at(target_S, (pos_patches_gal, ind_r),
                                      proj_gal * w_pairs_gal * e_plus[mask])

        if sp_buffer_DD is not None:
            patches_c, DD_c, Splus_c = _compress_jk_chunk(
                sp_buffer_DD, sp_buffer_Splus, i2 - i)
            sp_chunks_patches.append(patches_c)
            sp_chunks_DD.append(DD_c)
            if Splus_c is not None:
                sp_chunks_Splus.append(Splus_c)

    gal_jk_patches = DD_gal_jk_values = Splus_D_gal_jk_values = None
    if sp_buffer_DD is not None:
        gal_jk_patches = _pad_and_stack(sp_chunks_patches, fill=-1, dtype=np.int32)
        DD_gal_jk_values = _pad_and_stack(sp_chunks_DD, fill=0.0)
        if sp_chunks_Splus:
            Splus_D_gal_jk_values = _pad_and_stack(sp_chunks_Splus, fill=0.0)
        # The sparse arrays are built one chunk at a time and concatenated, so
        # their galaxy axis follows the *visit* order rather than the caller's.
        # The dense per-galaxy arrays do not need this because they are written
        # at gj = sel[n] as they go; these cannot be, so map them back here.
        # (order is the identity on the brute backend, where this is a no-op.)
        gal_jk_patches = _restore_galaxy_order(gal_jk_patches, order, fill=-1)
        DD_gal_jk_values = _restore_galaxy_order(DD_gal_jk_values, order, fill=0.0)
        if Splus_D_gal_jk_values is not None:
            Splus_D_gal_jk_values = _restore_galaxy_order(
                Splus_D_gal_jk_values, order, fill=0.0)

    return Grids(DD=DD, Splus_D=Splus_D, Scross_D=Scross_D, DD_jk=DD_jk, Splus_D_jk=Splus_D_jk,
                 DD_gal=DD_gal, Splus_D_gal=Splus_D_gal,
                 DD_gal_jk=DD_gal_jk, Splus_D_gal_jk=Splus_D_gal_jk,
                 gal_jk_patches=gal_jk_patches,
                 DD_gal_jk_values=DD_gal_jk_values,
                 Splus_D_gal_jk_values=Splus_D_gal_jk_values)