MeasureIABox

def __init__(
		self,
		data,
		output_file_name,
		simulation=None,
		snapshot=None,
		separation_limits=[0.1, 20.0],
		num_bins_r=8,
		num_bins_pi=20,
		pi_max=None,
		boxsize=None,
		periodicity=True,
		num_nodes=1,
		positions_density_sample_name="Position",
		positions_shape_sample_name="Position_shape_sample",
		axis_direction_name="Axis_Direction",
		axis_ratio_name="q",
		line_of_sight_index_name="LOS",
		weight_density_sample_name="weight",
		weight_shape_sample_name="weight_shape_sample",
		num_overlap=None,
):
	"""
	The __init__ method of the MeasureIABox class.

	Parameters
	----------
	num_nodes : int, optional
		Number of cores to be used in multiprocessing. Default is 1.
	positions_density_sample_name : str, optional
		Name of the key in the data dictionary that contains the positions of the density sample.
	positions_shape_sample_name : str, optional
		Name of the key in the data dictionary that contains the positions of the shape sample.
	axis_direction_name : str, optional
		Name of the key in the data dictionary that contains the axis direction vectors of the shape sample.
	axis_ratio_name : str, optional
		Name of the key in the data dictionary that contains the axis ratios of the shape sample.
	line_of_sight_index_name : str, optional
		Name of the key in the data dictionary that contains the column index of the line of sight in the
		position vectors.
	weight_density_sample_name : str, optional
		Name of the key in the data dictionary that contains the weights of the density sample.
	weight_shape_sample_name : str, optional
		Name of the key in the data dictionary that contains the weights of the shape sample.
	num_overlap : int or NoneType, optional
		Number of objects present in *both* the position and the shape sample. The analytic
		RR is normalised by ``Num_position * Num_shape - num_overlap``, because a shape
		galaxy cannot pair with itself and the pair loop already drops that self-pair (the
		separation window starts at ``r_min > 0``). Default None, which measures the overlap
		from the coordinates and is what you want for real data, where the shape sample is
		normally drawn from the position sample. Pass an integer to override it -- most
		usefully ``0``, which reproduces the convention external codes such as halotools
		and corr_pc use, where the two samples are treated as independent. An override is
		applied uniformly, so the per-jackknife-region adjustment is only made in the
		default (measured) mode.

	Notes
	-----
	Constructor parameters 'data', 'output_file_name', 'simulation', 'snapshot', 'separation_limits', 'num_bins_r',
	'num_bins_pi', 'pi_max', 'boxsize' and 'periodicity' are passed to MeasureIABase.
	The data dictionary (and any mask dictionaries passed to the measurement methods) may use any key names;
	they are given through the *_name parameters and translated to the internal default names on input.

	"""
	self._input_name_map = {
		positions_density_sample_name: "Position",
		positions_shape_sample_name: "Position_shape_sample",
		axis_direction_name: "Axis_Direction",
		axis_ratio_name: "q",
		line_of_sight_index_name: "LOS",
		weight_density_sample_name: "weight",
		weight_shape_sample_name: "weight_shape_sample",
	}
	if output_file_name is not None:
		self.check_paths([output_file_name])
	if data is not None:
		self.check_dict(data, [positions_density_sample_name, positions_shape_sample_name, axis_direction_name,
							   axis_ratio_name, line_of_sight_index_name])
		self.check_type_input_data(data,
								   (positions_density_sample_name, positions_shape_sample_name, axis_direction_name,
									axis_ratio_name, line_of_sight_index_name))
		data = self.rename_input_keys(data, self._input_name_map)
	super().__init__(data, output_file_name, simulation, snapshot, separation_limits, num_bins_r, num_bins_pi,
					 pi_max, boxsize, periodicity)
	if self.data is not None and self.boxsize is not None:
		self.check_units_coordinates(self.data["Position"], self.boxsize)
	if not (isinstance(num_nodes, (int, np.integer)) and not isinstance(num_nodes, bool) and num_nodes >= 1):
		raise ValueError(f"num_nodes must be an integer >= 1, got {num_nodes!r}.")
	self.num_nodes = num_nodes
	if num_overlap is not None and not (isinstance(num_overlap, (int, np.integer))
										and not isinstance(num_overlap, bool) and num_overlap >= 0):
		raise ValueError(f"num_overlap must be None or an integer >= 0, got {num_overlap!r}.")
	self._num_overlap_override = num_overlap
	self.randoms_data = None
	self.data_dir = None
	self.num_samples = None

	return

@worker_pool.pooled
def measure_xi_w(self, dataset_name, corr_type, num_jk=0, temp_file_path=None, masks=None,
				 ellipticity='distortion', chunk_size=1000, responsivity=True):
	r"""Measures $\xi_{gg}$, $\xi_{g+}$ and $w_{gg}$, $w_{g+}$ including jackknife covariance if desired.
	Manages the various _measure_xi_rp_pi_box method options in MeasureWBox and MeasureWBoxJackknife.

	Parameters
	----------
	dataset_name : str
		Name of the dataset in the output file.
	corr_type : str
		Type of correlation to be measured. Choose from [g+, gg, both].
	num_jk : int, optional
		Number of jackknife regions (needs to be x^3, with x an int) for the covariance measurement.
		Default is 0 (no covariance).
	temp_file_path : str or NoneType, optional
		Path to where the data is temporarily stored [file name generated automatically].
	masks : dict or NoneType, optional
		Directory of mask information in the same form as the data dictionary, where the masks are placed over
		the data to apply selections. Default is None.
	chunk_size: int, optional
		Size of the chunks of data sent to each multiprocessing node. If larger, more RAM is needed per node.
		Default is 1000.
	ellipticity : str, optional
		Definition of ellipticity. Choose from 'distortion', defined as (1-q^2)/(1+q^2), or 'ellipticity', defined
		 as (1-q)/(1+q). Default is 'distortion'.
	responsivity : bool, optional
		If True (default), the g+ shape signal is calibrated by the responsivity factor 2R, with
		R = <w (1 - e^2 / 2)> / <w> the weighted shear responsivity (R falls back to 0.5, i.e. no correction,
		when False or when the shape weights sum to zero). Use the default for raw shape distortions derived
		from axis ratios; set to False when the input shapes are already calibrated shears. Only the g+
		correlations are affected; the clustering (gg) signal is unchanged. Default is True.

	"""
	self._validate_measure_options(corr_type, ellipticity, num_jk)
	self.responsivity_correction = responsivity
	masks = self.rename_input_keys(masks, self._input_name_map)
	if num_jk > 0:
		try:
			assert sympy.integer_nthroot(num_jk, 3)[1]
			L = sympy.integer_nthroot(num_jk, 3)[0]
			self.check_jackknife_max_separation(num_jk, self.boxsize, self.r_max, self.num_bins_r)
		except AssertionError:
			raise ValueError(
				f"Use x^3 as input for num_jk, with x as an int. {float(int(num_jk ** (1. / 3)))},{num_jk ** (1. / 3)}")

	if temp_file_path == False:
		temp_storage = False
		temp_file_path = None
	else:
		temp_storage = True
	if temp_storage and temp_file_path == None:
		raise ValueError(
			"Input temp_file_path for faster computation. Do not want to save data temporarily? Input file_path_tree=False.")

	if self.data is not None and "RA" in self.data:
		raise TypeError("Given data is lightcone data (contains 'RA'). Use MeasureIALightcone instead.")

	if num_jk > 0:  # include covariance
		if corr_type == "gg":  # DD-only pair counting, skips shape/ellipticity computation
			if self.num_nodes > 1 and temp_storage:
				self._count_pairs_xi_rp_pi_box_jk_multiprocessing(masks=masks, L_subboxes=L,
																  dataset_name=dataset_name,
																  return_output=False,
																  num_nodes=self.num_nodes,
																  jk_group_name=f"{dataset_name}_jk{num_jk}",
																  chunk_size=chunk_size,
																  temp_file_path=temp_file_path)
			elif temp_storage:
				self._count_pairs_xi_rp_pi_box_jk_tree(masks=masks, L_subboxes=L, dataset_name=dataset_name,
													   return_output=False,
													   jk_group_name=f"{dataset_name}_jk{num_jk}")
			else:
				self._count_pairs_xi_rp_pi_box_jk_brute(masks=masks, L_subboxes=L, dataset_name=dataset_name,
														return_output=False,
														jk_group_name=f"{dataset_name}_jk{num_jk}")
		elif self.num_nodes > 1 and temp_storage:
			self._measure_xi_rp_pi_box_jk_multiprocessing(masks=masks, L_subboxes=L, dataset_name=dataset_name,
														  return_output=False,
														  num_nodes=self.num_nodes,
														  jk_group_name=f"{dataset_name}_jk{num_jk}",
														  chunk_size=chunk_size, ellipticity=ellipticity,
														  temp_file_path=temp_file_path)
		elif temp_storage:
			self._measure_xi_rp_pi_box_jk_tree(masks=masks, L_subboxes=L, dataset_name=dataset_name,
											   return_output=False, ellipticity=ellipticity,
											   jk_group_name=f"{dataset_name}_jk{num_jk}")
		else:
			self._measure_xi_rp_pi_box_jk_brute(masks=masks, L_subboxes=L, dataset_name=dataset_name,
												return_output=False, ellipticity=ellipticity,
												jk_group_name=f"{dataset_name}_jk{num_jk}")
		self._measure_w_g_i(corr_type=corr_type, dataset_name=dataset_name, return_output=False)
		for i in np.arange(num_jk):
			self._measure_w_g_i(corr_type=corr_type, dataset_name=f"{dataset_name}_{i}",
								jk_group_name=f"{dataset_name}_jk{num_jk}", return_output=False)
		if corr_type == "both":
			corr_group = ["w_g_plus", "w_gg"]
		elif corr_type == "g+":
			corr_group = ["w_g_plus"]
		elif corr_type == "gg":
			corr_group = ["w_gg"]
		else:
			raise KeyError("Unknown value for corr_type. Choose from [g+, gg, both]")
		self._combine_jackknife_information(dataset_name=dataset_name, jk_group_name=f"{dataset_name}_jk{num_jk}",
											corr_group=corr_group, num_box=num_jk)
	else:  # no covariance
		if corr_type == "gg":  # DD-only pair counting, skips shape/ellipticity computation
			if self.num_nodes > 1 and temp_storage:
				self._count_pairs_xi_rp_pi_box_multiprocessing(dataset_name=dataset_name,
															   temp_file_path=temp_file_path,
															   masks=masks, return_output=False,
															   num_nodes=self.num_nodes, chunk_size=chunk_size)
			elif temp_storage:
				self._count_pairs_xi_rp_pi_box_tree(masks=masks, dataset_name=dataset_name, return_output=False)
			else:
				self._count_pairs_xi_rp_pi_box_brute(masks=masks, dataset_name=dataset_name, return_output=False)
		elif self.num_nodes > 1 and temp_storage:
			self._measure_xi_rp_pi_box_multiprocessing(dataset_name=dataset_name, temp_file_path=temp_file_path,
													   masks=masks, return_output=False, num_nodes=self.num_nodes,
													   chunk_size=chunk_size, ellipticity=ellipticity)
		elif temp_storage:
			self._measure_xi_rp_pi_box_tree(masks=masks, dataset_name=dataset_name,
											return_output=False, ellipticity=ellipticity)
		else:
			self._measure_xi_rp_pi_box_brute(masks=masks, dataset_name=dataset_name,
											 return_output=False, ellipticity=ellipticity)
		self._measure_w_g_i(corr_type=corr_type, dataset_name=dataset_name, return_output=False)

	return

@worker_pool.pooled
def measure_xi_multipoles(self, dataset_name, corr_type, num_jk=0, temp_file_path=None, masks=None, rp_cut=None,
						  ellipticity='distortion', chunk_size=1000, responsivity=True):
	r"""Measures $\xi_{gg}$, $\xi_{g+}$ and $\tilde{\xi}_{gg,0}$, $\tilde{\xi}_{g+,2}$ including jackknife covariance
	if desired. Manages the various _measure_xi_r_mur_box method options in MeasureMultipolesBox and
	MeasureMultipolesBoxJackknife.

	Parameters
	----------
	dataset_name : str
		Name of the dataset in the output file.
	corr_type : str
		Type of correlation to be measured. Choose from [g+, gg, both].
	num_jk : int, optional
		Number of jackknife regions (needs to be x^3, with x an int) for the covariance measurement. Default is 0 (no covariance).
	temp_file_path : str or NoneType, optional
		Path to where the data is temporarily stored [file name generated automatically].
	masks : dict or NoneType, optional
		Directory of mask information in the same form as the data dictionary, where the masks are placed over
		the data to apply selections. Default is None.
	rp_cut : float or NoneType, optional
		Applies a minimum r_p value condition for pairs to be included. Default is None.
	chunk_size: int, optional
		Size of the chunks of data sent to each multiprocessing node. If larger, more RAM is needed per node.
	ellipticity : str, optional
		Definition of ellipticity. Choose from 'distortion', defined as (1-q^2)/(1+q^2), or 'ellipticity', defined
		 as (1-q)/(1+q). Default is 'distortion'.
	"""
	self._validate_measure_options(corr_type, ellipticity, num_jk)
	self.responsivity_correction = responsivity
	masks = self.rename_input_keys(masks, self._input_name_map)
	if num_jk > 0:
		try:
			assert sympy.integer_nthroot(num_jk, 3)[1]
			L = sympy.integer_nthroot(num_jk, 3)[0]
			self.check_jackknife_max_separation(num_jk, self.boxsize, self.r_max, self.num_bins_r)
		except AssertionError:
			raise ValueError(
				f"Use x^3 as input for num_jk, with x as an int. {float(int(num_jk ** (1. / 3)))},{num_jk ** (1. / 3)}")

	if temp_file_path == False:
		temp_storage = False
		temp_file_path = None
	else:
		temp_storage = True
	if temp_storage and temp_file_path == None:
		raise ValueError(
			"Input temp_file_path for faster computation. Do not want to save data temporarily? Input file_path_tree=False.")

	if self.data is not None and "RA" in self.data:
		raise TypeError("Given data is lightcone data (contains 'RA'). Use MeasureIALightcone instead.")

	if num_jk > 0:  # include covariance
		if corr_type == "gg":  # DD-only pair counting, skips shape/ellipticity computation
			if self.num_nodes > 1 and temp_storage:
				self._count_pairs_xi_r_mur_box_jk_multiprocessing(masks=masks, L_subboxes=L,
																  dataset_name=dataset_name,
																  return_output=False, rp_cut=rp_cut,
																  num_nodes=self.num_nodes,
																  jk_group_name=f"{dataset_name}_jk{num_jk}",
																  chunk_size=chunk_size,
																  temp_file_path=temp_file_path)
			elif temp_storage:
				self._count_pairs_xi_r_mur_box_jk_tree(masks=masks, L_subboxes=L, dataset_name=dataset_name,
													   return_output=False, rp_cut=rp_cut,
													   jk_group_name=f"{dataset_name}_jk{num_jk}")
			else:
				self._count_pairs_xi_r_mur_box_jk_brute(masks=masks, L_subboxes=L, dataset_name=dataset_name,
														return_output=False, rp_cut=rp_cut,
														jk_group_name=f"{dataset_name}_jk{num_jk}")
		elif self.num_nodes > 1 and temp_storage:
			self._measure_xi_r_mur_box_jk_multiprocessing(masks=masks, L_subboxes=L, dataset_name=dataset_name,
														  return_output=False, rp_cut=rp_cut,
														  num_nodes=self.num_nodes,
														  jk_group_name=f"{dataset_name}_jk{num_jk}",
														  chunk_size=chunk_size, ellipticity=ellipticity,
														  temp_file_path=temp_file_path)
		elif temp_storage:
			self._measure_xi_r_mur_box_jk_tree(masks=masks, L_subboxes=L, dataset_name=dataset_name,
											   return_output=False, rp_cut=rp_cut, ellipticity=ellipticity,
											   jk_group_name=f"{dataset_name}_jk{num_jk}")
		else:
			self._measure_xi_r_mur_box_jk_brute(masks=masks, L_subboxes=L, dataset_name=dataset_name,
												return_output=False, rp_cut=rp_cut, ellipticity=ellipticity,
												jk_group_name=f"{dataset_name}_jk{num_jk}")
		self._measure_multipoles(corr_type=corr_type, dataset_name=dataset_name, return_output=False)
		for i in np.arange(num_jk):
			self._measure_multipoles(corr_type=corr_type, dataset_name=f"{dataset_name}_{i}",
									 jk_group_name=f"{dataset_name}_jk{num_jk}", return_output=False)
		if corr_type == "both":
			corr_group = ["multipoles_g_plus", "multipoles_gg"]
		elif corr_type == "g+":
			corr_group = ["multipoles_g_plus"]
		elif corr_type == "gg":
			corr_group = ["multipoles_gg"]
		else:
			raise KeyError("Unknown value for corr_type. Choose from [g+, gg, both]")
		self._combine_jackknife_information(dataset_name=dataset_name, jk_group_name=f"{dataset_name}_jk{num_jk}",
											corr_group=corr_group, num_box=num_jk)
	else:  # no covariance
		if corr_type == "gg":  # DD-only pair counting, skips shape/ellipticity computation
			if self.num_nodes > 1 and temp_storage:
				self._count_pairs_xi_r_mur_box_multiprocessing(dataset_name=dataset_name,
															   temp_file_path=temp_file_path,
															   masks=masks, return_output=False, rp_cut=rp_cut,
															   num_nodes=self.num_nodes, chunk_size=chunk_size)
			elif temp_storage:
				self._count_pairs_xi_r_mur_box_tree(masks=masks, dataset_name=dataset_name,
													return_output=False, rp_cut=rp_cut)
			else:
				self._count_pairs_xi_r_mur_box_brute(masks=masks, dataset_name=dataset_name,
													 return_output=False, rp_cut=rp_cut)
		elif self.num_nodes > 1 and temp_storage:
			self._measure_xi_r_mur_box_multiprocessing(dataset_name=dataset_name, temp_file_path=temp_file_path,
													   masks=masks, return_output=False, rp_cut=rp_cut,
													   num_nodes=self.num_nodes,
													   chunk_size=chunk_size, ellipticity=ellipticity)
		elif temp_storage:
			self._measure_xi_r_mur_box_tree(masks=masks, dataset_name=dataset_name,
											return_output=False, rp_cut=rp_cut,
											ellipticity=ellipticity)
		else:
			self._measure_xi_r_mur_box_brute(masks=masks, dataset_name=dataset_name,
											 return_output=False, rp_cut=rp_cut, ellipticity=ellipticity)
		self._measure_multipoles(corr_type=corr_type, dataset_name=dataset_name, return_output=False)

	return

@worker_pool.pooled
def measure_galaxy_contributions(self, dataset_name, num_jk=0, masks=None, rp_cut=None,
                                 ellipticity='distortion', responsivity=True, ell=2,
                                 statistic="multipoles", temp_file_path=None,
                                 chunk_size=1000, return_output=False):
    r"""Measure the per-shape-galaxy contributions to the box IA estimator.

    Runs one pair traversal and returns, per shape galaxy and radial bin, the
    projected alignment contribution ``Y`` and the pair count ``P`` described in the
    module docstring. With ``num_jk > 0`` it also returns the decomposition of both by
    the jackknife sub-box of the *position-sample* partner, which is enough to rebuild
    every delete-one realisation without re-counting any pairs.

    Parameters
    ----------
    dataset_name : str
        Name of the dataset in the output file.
    num_jk : int, optional
        Number of jackknife regions; must be x^3 with x an int. Default 0 (none).
    masks : dict or NoneType, optional
        Mask dictionary in the same form as the data dictionary. Default None.
    rp_cut : float or NoneType, optional
        Minimum projected separation for a pair to be included. Multipoles only, as in
        ``measure_xi_multipoles``. Default None.
    ellipticity : str, optional
        'distortion' ((1-q^2)/(1+q^2)) or 'ellipticity' ((1-q)/(1+q)). Default
        'distortion'.
    responsivity : bool, optional
        If True (default), divide the shape signal by the responsivity 2R, as in the
        correlation-function measurement.
    ell : int, optional
        Multipole order; the spin is taken equal to it. Ignored when
        ``statistic="w"``. Default 2.
    statistic : str, optional
        Which estimator to project onto: ``"multipoles"`` for
        $\tilde{\xi}_{g+,\ell}(r)$ (r, mu_r binning) or ``"w"`` for $w_{g+}(r_p)$
        (rp, pi binning). Default "multipoles".
    temp_file_path : str or NoneType, optional
        Path where the data is temporarily offloaded during multiprocessing, so the
        parent does not hold a second copy in RAM. Required when ``num_nodes > 1``.
    chunk_size : int, optional
        Number of shape galaxies per multiprocessing task. Default 1000.
    return_output : bool, optional
        If True, return the results dict instead of writing it to the output file.
        Default False.

    Returns
    -------
    dict
        Only if ``return_output``. Keys: ``Y`` (M, num_bins_r), ``P`` (M, num_bins_r),
        ``r`` (num_bins_r,), and — when ``num_jk > 0`` — ``Y_jk``/``P_jk``
        (M, num_jk, num_bins_r) decomposed by the position partner's patch,
        ``jk_shape`` (M,) the patch of each shape galaxy, ``R_jk`` (num_jk,) the
        delete-one responsivities, ``rr_ratio`` (num_jk,) the scalar ``RR_jk / RR``
        amplitude ratios, and ``R`` the full-sample responsivity.

    Notes
    -----
    ``Y`` has the responsivity divided out and the full-sample ``RR`` folded in, so
    ``Y.sum(axis=0)`` is the ordinary estimator. ``Y_jk`` is stored **raw** — neither
    ``2R`` nor the per-realisation ``RR`` amplitude applied — matching the convention
    of the package's own ``Splus_D_jk`` grids. Delete-one realisation ``n`` is then

    .. code-block:: python

        keep = jk_shape != n
        est_n = (Y_jk[keep].sum(axis=1) - Y_jk[keep, n]).sum(axis=0)
        est_n /= rr_ratio[n] * 2 * R_jk[n]

    Runs on ``num_nodes`` cores (set at initialisation); ``num_nodes > 1`` requires
    ``temp_file_path`` and the usual ``if __name__ == "__main__":`` guard.
    """
    if self.data is not None and "RA" in self.data:
        raise TypeError("Given data is lightcone data (contains 'RA'). Use MeasureIALightcone instead.")
    if statistic not in _STATISTICS:
        raise ValueError(f"Unknown statistic {statistic!r}. Choose from {list(_STATISTICS)}.")
    if ellipticity not in ('distortion', 'ellipticity'):
        raise ValueError("Unknown value for ellipticity. Choose from ['distortion', 'ellipticity'].")
    if num_jk < 0:
        raise ValueError("num_jk must be >= 0.")
    num_nodes = getattr(self, "num_nodes", 1)
    if num_nodes > 1 and temp_file_path is None:
        raise ValueError(
            "measure_galaxy_contributions: num_nodes > 1 requires temp_file_path, where the "
            "catalogue is offloaded while the workers run."
        )
    L_subboxes = None
    if num_jk > 0:
        root, exact = sympy.integer_nthroot(num_jk, 3)
        if not exact:
            raise ValueError(f"Use x^3 as input for num_jk, with x an int. Got {num_jk}.")
        L_subboxes = root
        self.check_jackknife_max_separation(num_jk, self.boxsize, self.r_max, self.num_bins_r)

    self.responsivity_correction = responsivity
    masks = self.rename_input_keys(masks, self._input_name_map)

    sample_set = pair_kernel.prepare_box_samples(
        self.data, masks, self.Num_position, self.Num_shape,
        shapes=True, ellipticity=ellipticity, base=self,
        require_full_masks=num_nodes > 1,
    )
    Num_position = len(sample_set.pos)
    Num_shape = len(sample_set.pos_shape)
    e = sample_set.e
    weight_shape = sample_set.weight_shape
    R = sum(weight_shape * (1 - e ** 2 / 2.0)) / sum(weight_shape) \
        if responsivity and sum(weight_shape) > 0 else 0.5
    self.rp_cut = 0.0 if rp_cut is None else rp_cut
    print(f"There are {Num_shape} galaxies in the shape sample and {Num_position} galaxies in the position sample.")

    L3 = self.boxsize ** 3
    RR_g_plus = self._galaxy_analytic_RR(statistic, L3, Num_position, Num_shape)
    K = self._galaxy_projection_kernel(statistic, RR_g_plus, ell)

    jk_pos = jk_shape = None
    if num_jk > 0:
        jk_pos, jk_shape = self._get_jackknife_region_indices(masks, L_subboxes)
        sample_set.jk_pos = jk_pos
        sample_set.jk_shape = jk_shape

    binning = self._galaxy_binning(statistic, self.rp_cut)
    pos_tree = KDTree(binning.tree_coords(sample_set.pos, sample_set.not_LOS), boxsize=self.boxsize)

    if num_nodes > 1:
        Y, P, Y_jk, P_jk, jk_patches = self._galaxy_contributions_multiprocessing(
            sample_set, statistic, K, R, num_jk, pos_tree, dataset_name,
            temp_file_path, chunk_size, num_nodes, masks,
        )
    else:
        grids = pair_kernel.accumulate(
            sample_set, binning, base=self, R=R, shapes=True,
            chunk_axis="shape", chunk_size_outer=100, backend="tree", pos_tree=pos_tree,
            per_galaxy=True, per_galaxy_proj=K, per_galaxy_jk=num_jk > 0,
            per_galaxy_jk_sparse=num_jk > 0,
            num_box=num_jk if num_jk > 0 else None,
        )
        Y, P = grids.Splus_D_gal, grids.DD_gal
        Y_jk, P_jk = grids.Splus_D_gal_jk_values, grids.DD_gal_jk_values
        jk_patches = grids.gal_jk_patches

    out = {"Y": Y, "P": P, "r": self._galaxy_separation_bins()}
    if num_jk > 0:
        R_jk = pair_kernel.compute_R_jk(e, weight_shape, jk_shape, num_jk, responsivity)
        # RR is analytic and separable: every (r, mu_r) or (rp, pi) bin carries the
        # same geometric factor, so RR_jk[n] / RR is a scalar and the projection
        # kernel K simply rescales by 1 / rr_ratio[n].
        volume_jk = L3 * (num_jk - 1) / num_jk
        rr_ratio = np.zeros(num_jk)
        for n in np.arange(num_jk):
            n_pos_jk = int(np.count_nonzero(jk_pos != n))
            n_shape_jk = int(np.count_nonzero(jk_shape != n))
            overlap_jk = int(self.num_overlap - self.overlap_jk_counts[n])
            rr_ratio[n] = self._galaxy_rr_ratio(statistic, volume_jk, n_pos_jk, n_shape_jk,
                                                L3, Num_position, Num_shape, overlap_jk)
        out.update({"Y_jk_values": Y_jk, "P_jk_values": P_jk, "jk_patches": jk_patches,
                    "jk_shape": jk_shape, "R_jk": R_jk, "rr_ratio": rr_ratio, "R": R})

    if return_output:
        return out

    output_file = h5py.File(self.output_file_name, "a")
    group = create_group_hdf5(output_file, f"{self.snap_group}galaxy_contributions/{statistic}")
    write_dataset_hdf5(group, dataset_name + "_Y", data=out["Y"])
    write_dataset_hdf5(group, dataset_name + "_P", data=out["P"])
    write_dataset_hdf5(group, dataset_name + "_r", data=out["r"])
    if num_jk > 0:
        jk_group = create_group_hdf5(group, f"{dataset_name}_jk{num_jk}")
        for key in ("Y_jk_values", "P_jk_values", "jk_patches", "jk_shape",
                    "R_jk", "rr_ratio"):
            write_dataset_hdf5(jk_group, key, data=out[key])
        jk_group.attrs["R"] = out["R"]
    output_file.close()
    return None

def assign_jackknife_patches(self, data, randoms_data, num_jk, seed=None):
	"""Assigns jackknife patches to data and randoms given a number of patches.

	The patch centres are fitted to the position randoms with k-means on the sphere
	(see `measureia.kmeans_sphere`), so the patches are compact sky regions; every
	other sample is then assigned to its nearest centre.

	Parameters
	----------
	data : dict
		Dictionary containing position and shape sample data. Keywords: "RA", "DEC", "RA_shape_sample",
		"DEC_shape_sample"
	randoms_data : dict
		Dictionary containing position and shape sample data of randoms. Keywords: "RA", "DEC", "RA_shape_sample",
		"DEC_shape_sample"
	num_jk : int
		Number of jackknife patches. Cannot exceed the number of position randoms, since the
		patch centres are fitted to them.
	seed : int or NoneType, optional
		Seed for the k-means initialisation, making the patch assignment reproducible. If None (default),
		the patches differ between runs. The global random state is never touched.

	Returns
	-------
	dict
		Dictionary with patch numbers for each sample. Keywords: 'position', 'shape', 'randoms_position',
		'randoms_shape'

	Warns
	-----
	UserWarning
		If any sample ends up with patches holding fewer than 10 objects, which makes the
		jackknife covariance unreliable, or if the k-means fit does not converge.

	"""

	jk_patches = {}

	# Read the randoms file from which the jackknife regions will be created
	RA = randoms_data['RA']
	DEC = randoms_data['DEC']

	# Define a number of jackknife regions and find their centres using k-means
	X = np.column_stack((RA, DEC))
	if not (isinstance(num_jk, (int, np.integer)) and not isinstance(num_jk, bool) and num_jk >= 1):
		raise ValueError(f"num_jk must be an integer >= 1, got {num_jk!r}.")
	if num_jk > len(X):
		raise ValueError(
			f"num_jk ({num_jk}) cannot exceed the number of position randoms ({len(X)}), since the "
			f"patch centres are fitted to them. Lower num_jk or provide more randoms.")
	km = kmeans_sample(X, num_jk, maxiter=100, tol=1.0e-5, seed=seed)
	jk_labels = km.labels

	jk_patches['randoms_position'] = jk_labels

	RA = randoms_data['RA_shape_sample']
	DEC = randoms_data['DEC_shape_sample']
	X2 = np.column_stack((RA, DEC))
	jk_labels = km.find_nearest(X2)

	jk_patches['randoms_shape'] = jk_labels

	RA_data = data['RA']
	DEC_data = data['DEC']
	X3 = np.column_stack((RA_data, DEC_data))
	jk_labels = km.find_nearest(X3)

	jk_patches['position'] = jk_labels

	RA_data = data['RA_shape_sample']
	DEC_data = data['DEC_shape_sample']
	X4 = np.column_stack((RA_data, DEC_data))
	jk_labels = km.find_nearest(X4)

	jk_patches['shape'] = jk_labels

	self._warn_on_sparse_patches(jk_patches, num_jk)

	return jk_patches