MeasureIABase

class MeasureIABase(SimInfo):
	"""Base class for MeasureIA package that includes some general methods used throughout the package.

	Attributes
	----------
	Num_position : int
		Number of objects in the position sample. This value is updated in jackknife realisations.
	Num_shape : int
		Number of objects in the shape sample. This value is updated in jackknife realisations.
	r_min : float
		Minimum bound of (projected) separation length; bin edge. Default is 0.1.
	r_max : float
		Maximum bound of (projected) separation length; bin edge. Default is 20.
	binning : str
		Spacing of the separation bins, 'log' (default) or 'linear'.
	r_bins : ndarray
		Bin edges of the (projected) separation length (r_p or r).
	pi_bins : ndarray
		Bin edges of the line of sight (pi).
	mu_r_bins : ndarray
		Bin edges of the mu_r.

	Methods
	-------
	calculate_dot_product_arrays()
		Calculates dot product of elements of two arrays
	get_ellipticity()
		Given e and phi, e_+ and e_x components of ellipticity are returned.
	get_ellipticity_from_direction()
		As get_ellipticity, but from cos(phi)/sin(phi) rather than phi itself.
	get_random_pairs()
		Analytical RR for a (rp,pi) bin.
	get_volume_spherical_cap()
		Volume of an (r,mu_r) bin.
	get_random_pairs_r_mur()
		Analytical RR for a (r,mu_r) bin.
	_measure_w_g_i()
		Measure wgg or wg+ from xi grid provided by MeasureWBox or MeasureWLightcone class methods.
	_measure_multipoles()
		Measure multipoles (gg or g+) from xi grid provided by MeasureMultipolesBox or MeasureMultipolesLightcone
		class methods.
	_obs_estimator()
		Combines elements (DD, RR, etc) of xi estimators into xi_gg or xi_g+ for MeasureIALightcone.
	assign_jackknife_patches()
		Given positions of multiple samples, defines jackknife patches and returns index of every object in the sample.

	Notes
	-----
	Inherits attributes from 'SimInfo', where 'boxsize', 'L_0p5' and 'snap_group' are used in this class.

	"""

	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,
			binning="log",
	):
		"""
		The __init__ method of the MeasureIABase class.

		Parameters
		----------
		data : dict or NoneType
			Dictionary with data needed for calculations.
			For cartesian coordinates, the keywords are:
			'Position' and 'Position_shape_sample': (N_p,3), (N_s,3) ndarrays with the x, y, z coordinates
			of the N_p, N_s objects in the position and shape samples, respectively.
			'Axis_Direction': (N_s,2) ndarray with the two elements of the unit vectors describing the
			axis direction of the projected axis of the object shape.
			'LOS': index referring back to the column number in the 'Position' samples that contains the
			line-of-sight coordinate. (e.g. if the shapes are projected over the z-axis, LOS=2)
			'q': (N_s) array containing the axis ratio q=b/a for each object in the shape sample.
			For lightcone coordinates, the keywords are:
			'Redshift' and 'Redshift_shape_sample': (N_p) and (N_s) ndarray with redshifts of position and shape samples.
			'RA' and 'RA_shape_sample': (N_p) and (N_s) ndarray with RA coordinate of position and shape samples.
			'DEC' and 'DEC_shape_sample': (N_p) and (N_s) ndarray with DEC coordinate of position and shape samples.
			'e1' and 'e2': (N_s) arrays with the two ellipticity components e1 and e2 of the shape sample objects.
		output_file_name : str
			Name and filepath of the file where the output should be stored. Needs to be hdf5-type.
		simulation : str or NoneType, optional
			Indicator of simulation, obtaining correct boxsize in cMpc/h automatically. 
			Choose from [TNG100, TNG100_2, TNG300, EAGLE, HorizonAGN, FLAMINGO_L1, FLAMINGO_L2p8, COLIBRE_L400,
			COLIBRE_L200].
			Default is None, in which case boxsize needs to be added manually; or in the case of observational data, 
			the pi_max.
		snapshot : int or str or NoneType, optional
			Number of the snapshot, which, if given, will ensure that the output file to contains a group
			'Snapshot_[snapshot]'. If None, the group is omitted from the output file structure. Default is None.
		separation_limits : iterable of 2 entries, optional
			Bounds of the (projected) separation vector length bins in cMpc/h (so, r or r_p). Default is [0.1,20].
		num_bins_r : int, optional
			Number of bins for (projected) separation vector. Default is 8.
		num_bins_pi : int, optional
			Number of bins for line of sight (LOS) vector, pi or mu_r when multipoles are measured. Default is 20.
		pi_max : int or float, optional
			Bound for line of sight bins. Bounds will be [-pi_max, pi_max]. Default is None, in which case half the
			boxsize will be used.
		boxsize : int or float or NoneType, optional
			If simulation is not included in SimInfo, a manual boxsize can be added here. Make sure simulation=None
			and the boxsize units are equal to those in the data dictionary. Default is None.
		periodicity : bool, optional
			If True, the periodic boundary conditions of the simulation box are taken into account. If False, they are
			ignored. Note that because this code used analytical randoms for the simulations, the correlations will not
			be correct in this case and only DD and S+D terms should be studied. Non-periodic randoms can be measured by
			providing random data to the code and considering the DD term that is measured. Correlations and covariance
			matrix will need to be reconstructed from parts. [Please add a request for teh integration of this method of
			this if you would like to use this option often.] Default is True.

		"""
		SimInfo.__init__(self, simulation, snapshot, boxsize)
		# --- validate numeric configuration parameters ---
		if not (isinstance(num_bins_r, (int, np.integer)) and not isinstance(num_bins_r, bool) and num_bins_r >= 1):
			raise ValueError(f"num_bins_r must be an integer >= 1, got {num_bins_r!r}.")
		if not (isinstance(num_bins_pi, (int, np.integer)) and not isinstance(num_bins_pi, bool) and num_bins_pi >= 1):
			raise ValueError(f"num_bins_pi must be an integer >= 1, got {num_bins_pi!r}.")
		try:
			r_min_check, r_max_check = separation_limits
		except (TypeError, ValueError):
			raise ValueError(
				f"separation_limits must be a length-2 sequence [r_min, r_max], got {separation_limits!r}.")
		if not (0 < r_min_check < r_max_check):
			raise ValueError(
				f"separation_limits must satisfy 0 < r_min < r_max, got {separation_limits}.")
		if binning not in ("log", "linear"):
			raise ValueError(f"binning must be 'log' or 'linear', got {binning!r}.")
		if self.boxsize is not None and self.boxsize is not False and not (self.boxsize > 0):
			raise ValueError(f"boxsize must be > 0, got {self.boxsize!r}.")
		if pi_max is not None and not (pi_max > 0):
			raise ValueError(f"pi_max must be > 0, got {pi_max!r}.")
		self.data = data
		self.output_file_name = output_file_name
		self.periodicity = periodicity
		if periodicity:
			periodic = "periodic "
		else:
			periodic = ""
		try:
			self.Num_position = len(data["Position"])  # number of halos in position sample
			self.Num_shape = len(data["Position_shape_sample"])  # number of halos in shape sample
		except (KeyError, TypeError):  # TypeError: data is None
			try:
				self.Num_position = len(data["RA"])
				self.Num_shape = len(data["RA_shape_sample"])
			except (KeyError, TypeError):
				self.Num_position = 0
				self.Num_shape = 0
				print("Warning: no Position or Position_shape_sample given.")
		if self.Num_position > 0:
			if "weight" not in self.data:
				self.data["weight"] = np.ones(self.Num_position)
			if "weight_shape_sample" not in self.data:
				self.data["weight_shape_sample"] = np.ones(self.Num_shape)
		self.r_min = separation_limits[0]  # cMpc/h
		self.r_max = separation_limits[1]  # cMpc/h
		self.num_bins_r = num_bins_r
		self.num_bins_pi = num_bins_pi
		self.binning = binning
		if binning == "log":
			self.r_bins = np.logspace(np.log10(self.r_min), np.log10(self.r_max), self.num_bins_r + 1)
		else:
			self.r_bins = np.linspace(self.r_min, self.r_max, self.num_bins_r + 1)
		if pi_max == None:
			if self.L_0p5 is None:
				raise ValueError(
					"Both pi_max and boxsize are None. Provide input on one of them to determine the integration limit pi_max.")
			else:
				pi_max = self.L_0p5
		self.pi_bins = np.linspace(-pi_max, pi_max, self.num_bins_pi + 1)
		self.mu_r_bins = np.linspace(-1, 1, self.num_bins_pi + 1)
		if simulation == False:
			print(f"MeasureIA object initialised with:\n \
					observational data.\n \
					There are {self.Num_shape} galaxies in the shape sample and {self.Num_position} galaxies in the position sample.\n\
					The separation bin edges are given by {self.r_bins} Mpc.\n \
					There are {num_bins_r} {binning} r or r_p bins and {num_bins_pi} pi bins.\n \
					The maximum pi used for binning is {pi_max}.\n \
					The data will be written to {self.output_file_name}")
		else:
			print(f"MeasureIA object initialised with:\n \
			simulation {simulation} that has a {periodic}boxsize of {self.boxsize} cMpc/h.\n \
			There are {self.Num_shape} galaxies in the shape sample and {self.Num_position} galaxies in the position sample.\n\
			The separation bin edges are given by {self.r_bins} cMpc/h.\n \
			There are {num_bins_r} {binning} r or r_p bins and {num_bins_pi} pi bins.\n \
			The maximum pi used for binning is {pi_max}.\n \
			The data will be written to {self.output_file_name}")
		return

	@staticmethod
	def calculate_dot_product_arrays(a1, a2):
		"""Calculates the dot product over 2 2D arrays across axis 1 so that dot_product[i] = np.dot(a1[i],a2[i])

		Parameters
		----------
		a1 : ndarray
			First array
		a2 : ndarray
			Second array

		Returns
		-------
		ndarray
			Dot product of columns of arrays

		"""
		dot_product = np.zeros(np.shape(a1)[0])
		for i in np.arange(0, np.shape(a1)[1]):
			dot_product += a1[:, i] * a2[:, i]
		return dot_product

	@staticmethod
	def get_ellipticity(e, phi):
		"""Calculates the radial (+) and cross (x) components of the ellipticity with respect to the
		separation vector, following the intrinsic-alignment sign convention: e_+ > 0 means the
		major axis points along the separation vector (radial alignment), so radial alignment
		gives w_g+ > 0.

		For 2D input, e1 and e2 must follow the standard survey shear-catalogue convention
		(components defined on the local (RA, DEC) axes as delivered by e.g. lensfit/metacal-style
		catalogues and expected by TreeCorr). Note that the IA e_+ has the opposite sign to the
		lensing tangential shear: e_+ = -gamma_t.

		Parameters
		----------
		e : ndarray
			if 1D: size of the ellipticity vector; if 2D: e1,e2 components of the ellipticity vector
			in the survey shear-catalogue convention
		phi : ndarray
			if e 1D: angle between semimajor/semiminor axis and separation vector;
			if 2D: angle of the projected separation vector in the internal (east, north) frame,
			arctan2(north, east)

		Returns
		-------
		ndarray
			e_+ and e_x

		"""
		if len(np.shape(e)) > 1:
			e_plus = e[:, 0] * np.cos(2 * phi) - e[:, 1] * np.sin(2 * phi)
			e_cross = e[:, 0] * np.sin(2 * phi) + e[:, 1] * np.cos(2 * phi)
		else:
			e_plus, e_cross = e * np.cos(2 * phi), e * np.sin(2 * phi)
		return e_plus, e_cross

	@staticmethod
	def get_ellipticity_from_direction(e, cos_phi, sin_phi):
		"""Radial (+) and cross (x) ellipticity components from the *signed* angle between the
		semi-major axis and the separation vector, supplied as its cosine and sine.

		This is the scalar-``e`` counterpart of :meth:`get_ellipticity`, and exists because the
		semi-major axis direction has a physically meaningless sign: ``a`` and ``-a`` describe the
		same shape. Recovering the angle with ``arccos`` folds it into ``[0, pi]``, which maps
		``phi -> pi - phi`` under that flip. ``cos 2phi`` survives it but ``sin 2phi`` changes sign,
		so ``e_x`` — and any product of two ``e_x`` — would depend on an arbitrary input
		convention. Passing the signed pair ``(cos phi, sin phi)`` built from the dot and 2D cross
		products of the two unit vectors instead maps the flip to ``phi -> phi + pi``, under which
		both components are invariant.

		The double-angle identities are applied directly, so no trigonometric call is made.

		Parameters
		----------
		e : float or ndarray
			Size of the ellipticity vector. A scalar broadcasts against ``cos_phi``/``sin_phi``.
		cos_phi : ndarray
			Dot product of the unit semi-major axis with the unit projected separation direction.
		sin_phi : ndarray
			2D cross product of the same two unit vectors, ``a_0 s_1 - a_1 s_0``; the rotation
			carrying the axis onto the separation direction.

		Returns
		-------
		ndarray
			e_+ and e_x

		"""
		cos_2phi = 2.0 * cos_phi * cos_phi - 1.0
		sin_2phi = 2.0 * cos_phi * sin_phi
		return e * cos_2phi, e * sin_2phi

	@staticmethod
	def sample_responsivity(e, weight, responsivity_correction=True):
		"""Weighted responsivity ``R = <w(1 - e^2/2)>/<w>`` for one sample, or 0.5 when the
		correction is off (so that ``2R = 1`` and no calibration is applied).

		Factored out because a shape-shape correlation needs one of these per sample: the
		density sample carries its own shapes and generally its own shape-noise properties.
		"""
		if not responsivity_correction or e is None or sum(weight) <= 0:
			return 0.5
		return sum(weight * (1 - e ** 2 / 2.0)) / sum(weight)

	def write_shape_shape_grids(self, output_file, statistic, coords, grids, RR, RR_denom,
								separation_bins, second_bins, dataset_name, jk_group_name=""):
		r"""Writes the three shape-shape products and their estimators to the output file.

		Additive: it touches none of the existing groups, so the g+/gg datasets and their
		float summation order are unaffected by a shape-shape run.

		``Splus_Splus`` and ``Scross_Scross`` are the two II signals ($\xi_{++}$,
		$\xi_{\times\times}$); ``Splus_Scross`` is the symmetrised parity-odd null, the
		shape-shape analogue of $\xi_{g\times}$. All three share the g+ random-pair grid,
		because they count exactly the same pairs.

		Parameters
		----------
		output_file : h5py.File
			Open output file.
		statistic : str
			'w' or 'multipoles' -- the top-level group.
		coords : tuple of 2 str
			Suffixes of the two bin-centre datasets, ('rp', 'pi') or ('r', 'mu_r').
		grids : pair_kernel.Grids
			Must carry the shape-shape products (i.e. the kernel ran with ``shapes="both"``).
		RR : ndarray
			Random-pair grid, written to file as-is.
		RR_denom : ndarray
			The same grid with zeros replaced by ones, used as the divisor.
		separation_bins, second_bins : ndarray
			Bin centres of the two axes.
		dataset_name : str
			Name of the dataset in the output file.
		jk_group_name : str, optional
			Sub-group for jackknife realisations; empty for the full sample.

		"""
		coord_1, coord_2 = coords
		for group_name, raw, suffix in (
				("xi_plus_plus", grids.Splus_Splus, "_SplusSplus"),
				("xi_cross_cross", grids.Scross_Scross, "_ScrossScross"),
				("xi_plus_cross", grids.Splus_Scross, "_SplusScross")):
			group = create_group_hdf5(
				output_file, f"{self.snap_group}/{statistic}/{group_name}/{jk_group_name}")
			write_dataset_hdf5(group, dataset_name, data=raw / RR_denom)
			write_dataset_hdf5(group, dataset_name + suffix, data=raw)
			write_dataset_hdf5(group, dataset_name + "_RR", data=RR)
			write_dataset_hdf5(group, dataset_name + f"_{coord_1}", data=separation_bins)
			write_dataset_hdf5(group, dataset_name + f"_{coord_2}", data=second_bins)
		return

	def write_shape_shape_counts(self, output_file, statistic, coords, grids,
								 separation_bins, second_bins, dataset_name, jk_group_name=""):
		"""Writes the raw shape-shape pair sums for the lightcone.

		The lightcone backends write counts and let ``_obs_estimator`` assemble the
		estimator from them once every pass has run, so unlike the box writer this one
		does no division. Only the S+D pass produces these sums -- the S+R pass swaps the
		density sample for randoms, which carry no shapes.

		Parameters
		----------
		output_file : h5py.File
			Open output file.
		statistic : str
			'w' or 'multipoles'.
		coords : tuple of 2 str
			Suffixes of the two bin-centre datasets, ('rp', 'pi') or ('r', 'mu_r').
		grids : pair_kernel.Grids
			Must carry the shape-shape products.
		separation_bins, second_bins : ndarray
			Bin centres of the two axes.
		dataset_name : str
			Name of the dataset in the output file.
		jk_group_name : str, optional
			Sub-group for jackknife realisations; empty for the full sample.

		"""
		coord_1, coord_2 = coords
		for group_name, raw, suffix in (
				("xi_plus_plus", grids.Splus_Splus, "_SplusSplus"),
				("xi_cross_cross", grids.Scross_Scross, "_ScrossScross"),
				("xi_plus_cross", grids.Splus_Scross, "_SplusScross")):
			group = create_group_hdf5(
				output_file, f"{self.snap_group}/{statistic}/{group_name}/{jk_group_name}")
			write_dataset_hdf5(group, dataset_name + suffix, data=raw)
			write_dataset_hdf5(group, dataset_name + f"_{coord_1}", data=separation_bins)
			write_dataset_hdf5(group, dataset_name + f"_{coord_2}", data=second_bins)
		return

	def write_shape_shape_counts_jk(self, output_file, statistic, coords, grids,
									separation_bins, second_bins, dataset_name,
									jk_group_name, num_jk):
		"""Lightcone jackknife twin of :meth:`write_shape_shape_counts`.

		Mirrors how the lightcone writes ``S_+D`` under jackknife: the full-sample sums go
		to the top-level group and the delete-one sums, ``full - jk[i]``, into the
		realisation group. There is no responsivity juggling here -- the lightcone bakes
		1/(2R) into ``e`` and ``e_pos`` up front, so the subtraction is plain, exactly as
		it is for ``S_+D``.

		Parameters
		----------
		output_file : h5py.File
			Open output file.
		statistic : str
			'w' or 'multipoles'.
		coords : tuple of 2 str
			Suffixes of the two bin-centre datasets, ('rp', 'pi') or ('r', 'mu_r').
		grids : pair_kernel.Grids
			Must carry the shape-shape products and their ``_jk`` twins.
		separation_bins, second_bins : ndarray
			Bin centres of the two axes.
		dataset_name : str
			Name of the dataset in the output file.
		jk_group_name : str
			Sub-group holding the realisations.
		num_jk : int
			Number of jackknife realisations.

		"""
		coord_1, coord_2 = coords
		for group_name, full, raw_jk, suffix in (
				("xi_plus_plus", grids.Splus_Splus, grids.Splus_Splus_jk, "_SplusSplus"),
				("xi_cross_cross", grids.Scross_Scross, grids.Scross_Scross_jk, "_ScrossScross"),
				("xi_plus_cross", grids.Splus_Scross, grids.Splus_Scross_jk, "_SplusScross")):
			group = create_group_hdf5(
				output_file, f"{self.snap_group}/{statistic}/{group_name}/")
			write_dataset_hdf5(group, dataset_name + suffix, data=full)
			write_dataset_hdf5(group, dataset_name + f"_{coord_1}", data=separation_bins)
			write_dataset_hdf5(group, dataset_name + f"_{coord_2}", data=second_bins)
			group = create_group_hdf5(
				output_file, f"{self.snap_group}/{statistic}/{group_name}/{jk_group_name}")
			for i in np.arange(0, num_jk):
				write_dataset_hdf5(group, dataset_name + f"_{i}" + suffix,
								   data=full - raw_jk[i])
				write_dataset_hdf5(group, dataset_name + f"_{i}_{coord_1}", data=separation_bins)
				write_dataset_hdf5(group, dataset_name + f"_{i}_{coord_2}", data=second_bins)
		return

	def write_shape_shape_jk_realisations(self, output_file, statistic, coords, grids, RR_jk,
										  resp_full, R_jk, R_pos_jk, separation_bins,
										  second_bins, dataset_name, jk_group_name, num_box):
		r"""Writes the delete-one shape-shape realisations, mirroring the g+ jackknife block.

		The ``_jk`` grids from the kernel are raw sums, while the full-sample grids already
		carry the ``(2R)(2R_pos)`` division, so the full grid is multiplied back up before the
		subtraction. Both samples' responsivities change from realisation to realisation, so
		the retained-sample pair ``R_jk[i]``/``R_pos_jk[i]`` is applied per realisation --
		the shape-shape analogue of the single ``2 R_jk[i]`` the g+ block applies.

		Parameters
		----------
		output_file : h5py.File
			Open output file.
		statistic : str
			'w' or 'multipoles'.
		coords : tuple of 2 str
			Suffixes of the two bin-centre datasets, ('rp', 'pi') or ('r', 'mu_r').
		grids : pair_kernel.Grids
			Must carry the shape-shape products and their ``_jk`` twins.
		RR_jk : ndarray
			(num_box, nr, n2) random-pair grids of the delete-one samples.
		resp_full : float
			``(2R)(2R_pos)`` of the full sample, undoing the division in the full-sample grids.
		R_jk, R_pos_jk : ndarray
			Per-realisation responsivities of the shape and density samples.
		separation_bins, second_bins : ndarray
			Bin centres of the two axes.
		dataset_name : str
			Name of the dataset in the output file.
		jk_group_name : str
			Sub-group holding the realisations.
		num_box : int
			Number of jackknife realisations.

		"""
		coord_1, coord_2 = coords
		for group_name, full, raw_jk, suffix in (
				("xi_plus_plus", grids.Splus_Splus, grids.Splus_Splus_jk, "_SplusSplus"),
				("xi_cross_cross", grids.Scross_Scross, grids.Scross_Scross_jk, "_ScrossScross"),
				("xi_plus_cross", grids.Splus_Scross, grids.Splus_Scross_jk, "_SplusScross")):
			group = create_group_hdf5(
				output_file, f"{self.snap_group}/{statistic}/{group_name}/{jk_group_name}")
			for i in np.arange(0, num_box):
				RR_jk_denom = RR_jk[i].copy()  # guard against empty realisations/bins
				RR_jk_denom[RR_jk_denom == 0] = 1
				retained = full * resp_full - raw_jk[i]
				resp_i = (2 * R_jk[i]) * (2 * R_pos_jk[i])
				write_dataset_hdf5(group, dataset_name + f"_{i}",
								   data=retained / (RR_jk_denom * resp_i))
				write_dataset_hdf5(group, dataset_name + f"_{i}" + suffix, data=retained / resp_i)
				write_dataset_hdf5(group, dataset_name + f"_{i}_RR", data=RR_jk[i])
				write_dataset_hdf5(group, dataset_name + f"_{i}_{coord_1}", data=separation_bins)
				write_dataset_hdf5(group, dataset_name + f"_{i}_{coord_2}", data=second_bins)
		return

	@staticmethod
	def get_random_pairs(rp_max, rp_min, pi_max, pi_min, L3, corrtype, Num_position, Num_shape,
						 num_overlap=0):
		"""Returns analytical value of the number of pairs expected in an r_p, pi bin for a random uniform distribution.
		(Singh et al. 2023)

		Parameters
		----------
		rp_max : float
			Upper bound of projected separation vector bin
		rp_min : float
			Lower bound of projected separation vector bin
		pi_max : float
			Upper bound of line of sight vector bin
		pi_min : float
			Lower bound of line of sight vector bin
		L3 : float or int
			Volume of the simulation box
		corrtype : str
			Correlation type, auto or cross. RR for auto is RR_cross/2.
		Num_position : int
			Number of objects in the position sample.
		Num_shape : int
			Number of objects in the shape sample.
		num_overlap : int, optional
			Number of objects present in *both* samples. A shape galaxy cannot pair with
			itself, and the pair loop drops that self-pair automatically because the
			separation window starts at ``r_min > 0``, so the available pair count is
			``Num_position * Num_shape - num_overlap`` rather than the plain product. Pass 0
			(the default) for genuinely disjoint samples; pass ``Num_shape`` when the shape
			sample is drawn from the position sample, which is the usual IA configuration.
			See :func:`available_pairs`.


		Returns
		-------
		float
			number of pairs in r_p, pi bin

		"""
		num_pairs = available_pairs(Num_position, Num_shape, num_overlap, corrtype)
		# volume is cylindrical pi*dr^2 * height
		return num_pairs * np.pi * (rp_max ** 2 - rp_min ** 2) * abs(pi_max - pi_min) / L3

	def get_random_pairs_r_mur(self, r_max, r_min, mur_max, mur_min, L3, corrtype, Num_position,
							   Num_shape, num_overlap=0):
		"""Returns analytical value of the number of pairs expected in an r, mu_r bin for a random uniform distribution.

		Parameters
		----------
		r_max : float
			Upper bound of separation vector bin
		r_min : float
			Lower bound of separation vector bin
		mur_max : float
			Upper bound of mu_r bin
		mur_min : float
			Lower bound of mu_r bin
		L3 : float
			Volume of the simulation box
		corrtype : str
			Correlation type, auto or cross. RR for auto is RR_cross/2.
		Num_position : int
			Number of objects in the position sample.
		Num_shape : int
			Number of objects in the shape sample.
		num_overlap : int, optional
			Number of objects present in *both* samples. A shape galaxy cannot pair with
			itself, and the pair loop drops that self-pair automatically because the
			separation window starts at ``r_min > 0``, so the available pair count is
			``Num_position * Num_shape - num_overlap`` rather than the plain product. Pass 0
			(the default) for genuinely disjoint samples; pass ``Num_shape`` when the shape
			sample is drawn from the position sample, which is the usual IA configuration.
			See :func:`available_pairs`.


		Returns
		-------
		float
			number of pairs in r, mu_r bin

		"""

		num_pairs = available_pairs(Num_position, Num_shape, num_overlap, corrtype)
		# volume is big cap - small cap for large - small radius
		return abs(num_pairs * 2. * np.pi / 3. * (r_max ** 3 - r_min ** 3)
				   * (mur_max - mur_min) / L3)

	def _measure_w_g_i(self, dataset_name, corr_type="both", return_output=False, jk_group_name=""):
		"""Measures w_gg or w_g+ for a given xi_gi dataset that has been calculated with the _measure_xi_rp_pi_sims
		methods. Integrates over pi bins via sum * dpi. Stores rp, and w_gg or w_g+.

		Parameters
		----------
		dataset_name : str
			Name of xi_gg or xi_g+ dataset and name given to w_gg or w_g+ dataset when stored.
		return_output : bool, optional
			Output is returned if True, saved to file if False. Default value = False
		corr_type : str, optional
			Type of correlation function. Choose from [g+,gg,both,++,all]. Default value = "both".
			'++' covers the three shape-shape products (plus-plus, cross-cross and the parity null).
		jk_group_name : str, optional
			Name of subgroup in hdf5 file where jackknife realisations are stored. Default value = ""

		Returns
		-------
		ndarray
			[rp, wgg] or [rp, wg+] if return_output is True

		"""
		try:
			xi_data, wg_data = zip(*W_PRODUCTS[corr_type])
		except KeyError:
			raise KeyError(f"Unknown value for corr_type. Choose from {sorted(W_PRODUCTS)}")
		xi_data, wg_data = list(xi_data), list(wg_data)
		for i in np.arange(0, len(xi_data)):
			correlation_data_file = h5py.File(self.output_file_name, "a")
			group = correlation_data_file[f"{self.snap_group}w/{xi_data[i]}/{jk_group_name}"]
			correlation_data = group[dataset_name][:]
			pi = group[dataset_name + "_pi"][:]
			rp = group[dataset_name + "_rp"][:]
			dpi = (self.pi_bins[1:] - self.pi_bins[:-1])
			pi_bins = self.pi_bins[:-1] + abs(dpi) / 2.0  # middle of bins
			if sum(np.isin(pi, pi_bins)) == len(pi):
				dpi = np.array([dpi] * len(correlation_data[:, 0]))
				correlation_data = correlation_data * abs(dpi)
			else:
				raise ValueError("Update pi bins in initialisation of object to match xi_g_plus dataset.")
			# xi carries NaN/inf in bins with zero empirical RR (undefined estimator, already
			# warned); the pi-integral over such a bin is legitimately NaN, so silence the
			# redundant reduce/invalid-value warning.
			with np.errstate(invalid='ignore'):
				w_g_i = np.sum(correlation_data, axis=1)  # sum over pi values
			if return_output:
				output_data = np.array([rp, w_g_i]).transpose()
				correlation_data_file.close()
				return output_data
			else:
				group_out = create_group_hdf5(correlation_data_file,
											  f"{self.snap_group}{wg_data[i]}/{jk_group_name}")
				write_dataset_hdf5(group_out, dataset_name + "_rp", data=rp)
				write_dataset_hdf5(group_out, dataset_name, data=w_g_i)
				correlation_data_file.close()
		return

	def _measure_multipoles(self, dataset_name, corr_type="both", return_output=False, jk_group_name=""):
		"""Measures multipoles for a given xi_g+ or xi_gg measured by _measure_xi_r_pi_sims methods.
		The data assumes xi_g+ and xi_gg to be measured in bins of r and mu_r.

		Parameters
		----------
		dataset_name : str
			Name of xi_gg or xi_g+ dataset and name given to multipoles dataset when stored.
		corr_type : str, optional
			Type of correlation function. Choose from [g+,gg,both,++,all]. Default value = "both".
			'++' covers the three shape-shape products (plus-plus, cross-cross and the parity null).
		return_output : bool, optional
			Output is returned if True, saved to file if False. Default value = False.
		jk_group_name : str, optional
			Name of subgroup in hdf5 file where jackknife realisations are stored. Default value = ""

		Returns
		-------
		ndarray
			[r, multipoles_gg] or [r, multipoles_g+] if return_output is True
		"""
		correlation_data_file = h5py.File(self.output_file_name, "a")
		# The (l, s_ab) pairs live in M_PRODUCTS: (0,0) for gg, (2,2) for g+ and (4,4) for
		# ++ (Singh et al. 2023 -- two shapes in the correlation, so spin 4). l == s_ab for
		# every product currently offered, which is why the two share one list; a general
		# l > s_ab moment would break that coupling here.
		try:
			products = M_PRODUCTS[corr_type]
		except KeyError:
			raise KeyError(f"Unknown value for corr_type. Choose from {sorted(M_PRODUCTS)}")
		correlation_data_list, r_list, mu_r_list, sab_list, corr_type_list = [], [], [], [], []
		for xi_name, multipole_name, sab in products:
			group = correlation_data_file[f"{self.snap_group}multipoles/{xi_name}/{jk_group_name}"]
			correlation_data_list.append(group[dataset_name][:])   # xi in the (r, mu_r) grid
			r_list.append(group[dataset_name + "_r"][:])
			mu_r_list.append(group[dataset_name + "_mu_r"][:])
			sab_list.append(sab)
			corr_type_list.append(multipole_name.replace("multipoles_", ""))
		l_list = sab_list
		for i in np.arange(0, len(sab_list)):
			corr_type_i = corr_type_list[i]
			correlation_data = correlation_data_list[i]
			r = r_list[i]
			mu_r = mu_r_list[i]
			sab = sab_list[i]
			l = l_list[i]
			mu_r = np.array(list(mu_r) * len(r)).reshape((len(r), len(mu_r)))  # make pi into grid for mu

			r = np.array(list(r) * len(mu_r)).reshape((len(r), len(mu_r)))
			r = r.transpose()
			# Associated Legendre polynomial grid. scipy's lpmn (removed in 1.17) returned a
			# (l+1, sab+1) grid whose [-1, -1] element was order l, degree sab; that value is
			# assoc_legendre_p(degree=sab, order=l, z). The call vectorises over the whole
			# mu_r grid (leading axis is the order, hence [0]); output is bit-identical.
			L = assoc_legendre_p(sab, l, mu_r)[0]
			dmur = (self.mu_r_bins[1:] - self.mu_r_bins[:-1])
			dmu_r_array = np.array(list(dmur) * len(r)).reshape((len(r), len(dmur)))
			# correlation_data carries NaN/inf in bins with zero empirical RR (the estimator
			# is undefined there and the user was already warned); the integral over such a
			# bin is legitimately NaN, so silence the redundant reduce/invalid-value warnings.
			with np.errstate(invalid='ignore'):
				multipoles = (
						(2 * l + 1)
						/ 2.0
						* math.factorial(l - sab)
						/ math.factorial(l + sab)
						* L
						* correlation_data
						* dmu_r_array
				)
				multipoles = np.sum(multipoles, axis=1)
			dsep = (self.r_bins[1:] - self.r_bins[:-1]) / 2.0
			separation = self.r_bins[:-1] + abs(dsep)  # middle of bins
			if return_output:
				correlation_data_file.close()
				return np.array([separation, multipoles]).transpose()
			else:
				group_out = create_group_hdf5(
					correlation_data_file, f"{self.snap_group}multipoles_{corr_type_i}/{jk_group_name}"
				)
				write_dataset_hdf5(group_out, dataset_name + "_r", data=separation)
				write_dataset_hdf5(group_out, dataset_name, data=multipoles)
		correlation_data_file.close()
		return

	def _obs_estimator(self, corr_type, IA_estimator, dataset_name, num_samples,
					   jk_group_name=""):
		"""Reads various components of xi and combines into correct estimator for cluster or galaxy
		lightcone alignment correlations. It then writes the xi_gg or xi_g+ in the correct place in the output file.

		Parameters
		----------
		corr_type : list of 2 str elements
			First element: ['gg', 'g+', 'both'], second: 'w' or 'multipoles'
		IA_estimator : str
			Chooser from 'clusters' or 'galaxies' for different estimator definition.
		dataset_name : str
			Name of the dataset
		num_samples : dict
			Dictionary of samples sizes for position, shape and random samples. Keywords: D, S, R_D, R_S
		jk_group_name : str
			Name of subgroup in hdf5 file where jackknife realisations are stored. Default value = ""

		Returns
		-------

		"""
		# which families this call has to assemble; 'both' keeps its original meaning
		# (g+ and gg) and 'all' adds the shape-shape products on top
		want_gp = corr_type[0] in ("g+", "both", "all")
		want_gg = corr_type[0] in ("gg", "both", "all")
		want_pp = corr_type[0] in ("++", "all")
		output_file = h5py.File(self.output_file_name, "a")
		group_gg = output_file[f"{self.snap_group}{corr_type[1]}/xi_gg/{jk_group_name}"]
		if want_pp:
			group_pp = output_file[f"{self.snap_group}{corr_type[1]}/xi_plus_plus/{jk_group_name}"]
			group_xx = output_file[f"{self.snap_group}{corr_type[1]}/xi_cross_cross/{jk_group_name}"]
			group_px = output_file[f"{self.snap_group}{corr_type[1]}/xi_plus_cross/{jk_group_name}"]
			# same normalisation as S+D: the shape-shape sums run over the same D x S pairs
			pp_norm = max(num_samples["S"] * num_samples["D"] - num_samples["D_S"], 1)
			SpSp = group_pp[f"{dataset_name}_SplusSplus"][:] / pp_norm
			SxSx = group_xx[f"{dataset_name}_ScrossScross"][:] / pp_norm
			SpSx = group_px[f"{dataset_name}_SplusScross"][:] / pp_norm
		if want_gp:
			group_gp = output_file[
				f"{self.snap_group}{corr_type[1]}/xi_g_plus/{jk_group_name}"]
			SpD = group_gp[f"{dataset_name}_SplusD"][:]
			SpD /= max(num_samples["S"] * num_samples["D"] - num_samples["D_S"], 1)
			SpR = group_gp[f"{dataset_name}_SplusR"][:]
			SpR /= max(num_samples["S"] * num_samples["R_D"], 1)
			if jk_group_name == "":  # cross (parity null test) only for the full sample
				group_gc = output_file[f"{self.snap_group}{corr_type[1]}/xi_g_cross/"]
				ScD = group_gc[f"{dataset_name}_ScrossD"][:]
				ScD /= max(num_samples["S"] * num_samples["D"] - num_samples["D_S"], 1)
				ScR = group_gc[f"{dataset_name}_ScrossR"][:]
				ScR /= max(num_samples["S"] * num_samples["R_D"], 1)
		if want_gg or IA_estimator == "clusters":
			SR = group_gg[f"{dataset_name}_SR"][:]
			SR /= max(num_samples["S"] * num_samples["R_D"], 1)
		if want_gg:
			RD = group_gg[f"{dataset_name}_RD"][:]
			RD /= max(num_samples["D"] * num_samples["R_S"], 1)
		if IA_estimator == 'clusters' or want_gg:
			DD = group_gg[f"{dataset_name}_DD"][:]
			DD /= max(num_samples["D"] * num_samples["S"] - num_samples["D_S"], 1)
			DD_denom = DD.copy()  # guard for the clusters division; raw DD keeps 0 in the gg numerator
			DD_denom[DD_denom == 0] = 1.
		if IA_estimator == "galaxies" or want_gg or want_pp:
			RR = group_gg[f"{dataset_name}_RR"][:]
			if jk_group_name == "" and np.any(RR == 0):
				warnings.warn(
					f"{int(np.sum(RR == 0))} bin(s) of '{dataset_name}' contain zero random-random pairs; "
					"the estimator is undefined (NaN) there and any w/multipole integral over these bins "
					"will be NaN. Increase the number of randoms to fill all bins.", RuntimeWarning)
			RR /= max(num_samples["R_D"] * num_samples["R_S"], 1)

		# Bins with zero empirical RR (or SR) divide to NaN/inf on purpose: the estimator
		# is undefined there. The zero-RR RuntimeWarning above already tells the user, so
		# silence the redundant numpy divide/invalid-value warnings from these divisions.
		with np.errstate(invalid='ignore', divide='ignore'):
			if IA_estimator == "clusters":
				if want_gp:
					correlation_gp = SpD / DD_denom - SpR / SR
					write_dataset_hdf5(group_gp, dataset_name, correlation_gp)
					if jk_group_name == "":
						correlation_gc = ScD / DD_denom - ScR / SR
						write_dataset_hdf5(group_gc, dataset_name, correlation_gc)
				if want_gg:
					correlation_gg = (DD - RD - SR) / RR + 1
					write_dataset_hdf5(group_gg, dataset_name, correlation_gg)
			elif IA_estimator == "galaxies":
				if want_gp:
					correlation_gp = (SpD - SpR) / RR
					write_dataset_hdf5(group_gp, dataset_name, correlation_gp)
					if jk_group_name == "":
						correlation_gc = (ScD - ScR) / RR
						write_dataset_hdf5(group_gc, dataset_name, correlation_gc)
				if want_gg:
					correlation_gg = (DD - RD - SR) / RR + 1
					write_dataset_hdf5(group_gg, dataset_name, correlation_gg)
				if want_pp:
					# xi_++ = S+S+/RR, and likewise for the cross-cross signal and the
					# parity-odd null. There is no S+R analogue to subtract: the randoms
					# carry no shapes, so no shape-shape term can be built from them.
					for group, raw, suffix in (
							(group_pp, SpSp, "_SplusSplus"),
							(group_xx, SxSx, "_ScrossScross"),
							(group_px, SpSx, "_SplusScross")):
						write_dataset_hdf5(group, dataset_name, raw / RR)
			else:
				raise ValueError("Unknown input for IA_estimator, choose from [clusters, galaxies].")
		output_file.close()
		return