MeasureIA

class MeasureIABox(MeasureJackknife):
	r"""Manages the IA correlation function measurement methods used in the MeasureIA package based on speed and input.
	This class is used to call the methods that measure w_gg, w_g+ and multipoles for simulations in cartesian coordinates.
	Depending on the input parameters, various correlations incl covariance estimates are measured for given data.

	Methods
	-------
	measure_xi_w()
		Compute projected correlations $w_{gg}$ and/or $w_{g+}$.
	measure_xi_multipoles()
		Compute multipoles of the correlation functions, $\tilde{\xi}_{gg,0}$ and/or $\tilde{\xi}_{g+,2}$.

	Notes
	-----
	Inherits attributes from 'SimInfo', where 'boxsize', 'L_0p5' and 'snap_group' are used in this class.
	Inherits attributes from 'MeasureIABase', where 'data', 'output_file_name', 'periodicity', 'Num_position',
	'Num_shape', 'r_min', 'r_max', 'num_bins_r', 'num_bins_pi', 'r_bins', 'pi_bins', 'mu_r_bins' are used.

	"""

	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,
	):
		"""
		The __init__ method of the MeasureIABox class.

		Parameters
		----------
		num_nodes : int, optional
			Number of cores to be used in multiprocessing. Default is 1.

		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.

		"""
		super().__init__(data, output_file_name, simulation, snapshot, separation_limits, num_bins_r, num_bins_pi,
						 pi_max, boxsize, periodicity)
		self.num_nodes = num_nodes
		self.randoms_data = None
		self.data_dir = None
		self.num_samples = None

		return

	def measure_xi_w(self, dataset_name, corr_type, num_jk=0, measure_cov=True, file_tree_path=None, masks=None,
					 remove_tree_file=True, save_jk_terms=False):
		"""Measures xi_gg, xi_g+ and w_gg, w_g+ including jackknife covariance if desired.
		Manages the various _measure_xi_rp_pi_sims and _measure_jackknife_covariance_sims options in MeasureWSimulations
		and MeasureJackknife.

		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.
		measure_cov : bool, optional
			If True, jackknife covariance is measured. Default is True
		file_tree_path : str or NoneType, optional
			Path to where the tree information is temporarily stored [file name generated automatically].
			If None (default), no trees are used in the calculation.
			Note that the use of trees speeds up the calculations significantly.
		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.
		remove_tree_file : bool, optional
			If True (default), the file that stores the tree information is removed after the measurements.
		save_jk_terms : bool, optional
			If True, DD and S+D terms of the jackknife realisations are also saved in the output file.
			These terms are automatically saved when only 1 core is used in the measurements. Default is False.

		"""
		if measure_cov:
			try:
				assert sympy.integer_nthroot(num_jk, 3)[1]
				L = sympy.integer_nthroot(num_jk, 3)[0]
			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 self.num_nodes == 1:
				multiproc_bool = False
				save_tree = True
			elif num_jk > 0.5 * self.num_nodes:
				multiproc_bool = True
				save_tree = True
			else:
				multiproc_bool = True
				save_tree = False
		else:
			if self.num_nodes == 1:
				multiproc_bool = False
				save_tree = False
			elif self.num_nodes > 2:
				multiproc_bool = True
				save_tree = False
			else:
				multiproc_bool = True
				save_tree = True
		if save_tree and file_tree_path == False:
			save_tree = False
			file_tree_path = None
		elif save_tree and file_tree_path == None:
			raise ValueError(
				"Input file_tree_path for faster computation. Do not want to use trees? Input file_path_tree=False.")
		try:
			RA = self.data["RA"]
			sim_bool = False
		except:
			sim_bool = True
		if not sim_bool:
			print("Given data is observational, use measure_xi_w_obs method instead.")
		else:
			if multiproc_bool and save_tree:
				self._measure_xi_rp_pi_sims_tree(tree_input=None, masks=masks, dataset_name=dataset_name,
												 return_output=False, print_num=True, dataset_name_tree=None,
												 save_tree=save_tree, file_tree_path=file_tree_path)
				self._measure_w_g_i(corr_type=corr_type, dataset_name=dataset_name, return_output=False)
				if measure_cov:
					self._measure_jackknife_covariance_sims_multiprocessing(masks=masks, corr_type=[corr_type, "w"],
																			dataset_name=dataset_name, L_subboxes=L,
																			rp_cut=None,
																			num_nodes=self.num_nodes, twoD=False,
																			tree=True,
																			tree_saved=True,
																			file_tree_path=file_tree_path,
																			remove_tree_file=remove_tree_file,
																			save_jk_terms=save_jk_terms)
			elif not multiproc_bool and save_tree:
				self._measure_xi_rp_pi_sims_tree(tree_input=None, masks=masks, dataset_name=dataset_name,
												 return_output=False, print_num=True, dataset_name_tree=None,
												 save_tree=save_tree, file_tree_path=file_tree_path)
				self._measure_w_g_i(corr_type=corr_type, dataset_name=dataset_name, return_output=False)
				if measure_cov:
					self._measure_jackknife_covariance_sims(masks=masks, corr_type=[corr_type, "w"],
															dataset_name=dataset_name, L_subboxes=L, rp_cut=None,
															tree_saved=True, file_tree_path=file_tree_path,
															remove_tree_file=remove_tree_file)
			elif multiproc_bool and not save_tree:
				print("yes")
				self._measure_xi_rp_pi_sims_multiprocessing(num_nodes=self.num_nodes, masks=masks,
															dataset_name=dataset_name, return_output=False,
															print_num=True)
				self._measure_w_g_i(corr_type=corr_type, dataset_name=dataset_name, return_output=False)
				if measure_cov:
					self._measure_jackknife_covariance_sims(masks=masks, corr_type=[corr_type, "w"],
															dataset_name=dataset_name, L_subboxes=L, rp_cut=None,
															num_nodes=self.num_nodes, tree_saved=False)
			else:
				self._measure_xi_rp_pi_sims_brute(masks=masks, dataset_name=dataset_name,
												  return_output=False, print_num=True)
				self._measure_w_g_i(corr_type=corr_type, dataset_name=dataset_name, return_output=False)
				if measure_cov:
					self._measure_jackknife_covariance_sims(masks=masks, corr_type=[corr_type, "w"],
															dataset_name=dataset_name, L_subboxes=L, rp_cut=None,
															num_nodes=self.num_nodes, tree_saved=False)

		return

	def measure_xi_multipoles(self, dataset_name, corr_type, num_jk, measure_cov=True, file_tree_path=None, masks=None,
							  remove_tree_file=True, rp_cut=None):
		"""Measures multipoles including jackknife covariance if desired.
		Manages the various _measure_xi_r_mu_r_sims and _measure_jackknife_covariance_sims options in
		MeasureMultipolesSimulations and MeasureJackknife.

		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.
		measure_cov : bool, optional
			If True, jackknife covariance is measured. Default is True
		file_tree_path : str or NoneType, optional
			Path to where the tree information is temporarily stored [file name generated automatically].
			If None (default), no trees are used in the calculation.
			Note that the use of trees speeds up the calculations significantly.
		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.
		remove_tree_file : bool, optional
			If True (default), the file that stores the tree information is removed after the measurements.
		rp_cut : float or NoneType, optional
			Applies a minimum r_p value condition for pairs to be included. Default is None.

		"""
		if measure_cov:
			try:
				assert sympy.integer_nthroot(num_jk, 3)[1]
				L = sympy.integer_nthroot(num_jk, 3)[0]
			except AssertionError:
				raise ValueError("Use x^3 as input for num_jk, with x as an int.")
			if self.num_nodes == 1:
				multiproc_bool = False
				save_tree = True
			elif num_jk > 0.5 * self.num_nodes:
				multiproc_bool = True
				save_tree = True
			else:
				multiproc_bool = True
				save_tree = False
		else:
			if self.num_nodes == 1:
				multiproc_bool = False
				save_tree = False
			elif self.num_nodes > 2:
				multiproc_bool = True
				save_tree = False
			else:
				multiproc_bool = True
				save_tree = True
		if save_tree and file_tree_path == None:
			raise ValueError(
				"Input file_tree_path for faster computation. Do not want to use trees? Input file_path_tree=False.")
		elif save_tree and file_tree_path == False:
			save_tree = False
			file_tree_path = None
		try:
			RA = self.data["RA"]
			raise KeyError(
				"Lightcone input provided, use measure_xi_w_obs and measure_xi_multipoles_obs for measurements or "
				"provide carthesian coordinate data.")
		except KeyError:
			pass
		if multiproc_bool and save_tree:
			self._measure_xi_r_mur_sims_tree(tree_input=None, masks=masks,
											 dataset_name=dataset_name,
											 return_output=False, print_num=True,
											 dataset_name_tree=None, rp_cut=rp_cut,
											 save_tree=save_tree, file_tree_path=file_tree_path)
			self._measure_multipoles(corr_type=corr_type, dataset_name=dataset_name, return_output=False)
			if measure_cov:
				self._measure_jackknife_covariance_sims_multiprocessing(masks=masks,
																		corr_type=[corr_type, "multipoles"],
																		dataset_name=dataset_name, L_subboxes=L,
																		rp_cut=rp_cut,
																		num_nodes=self.num_nodes, twoD=False,
																		tree=True,
																		tree_saved=True,
																		file_tree_path=file_tree_path,
																		remove_tree_file=remove_tree_file)
		elif not multiproc_bool and save_tree:
			self._measure_xi_r_mur_sims_tree(tree_input=None, masks=masks,
											 dataset_name=dataset_name,
											 return_output=False, print_num=True,
											 dataset_name_tree=None, rp_cut=rp_cut,
											 save_tree=save_tree, file_tree_path=file_tree_path)
			self._measure_multipoles(corr_type=corr_type, dataset_name=dataset_name, return_output=False)
			if measure_cov:
				self._measure_jackknife_covariance_sims(masks=masks, corr_type=[corr_type, "multipoles"],
														dataset_name=dataset_name, L_subboxes=L, rp_cut=rp_cut,
														tree_saved=True, file_tree_path=file_tree_path,
														remove_tree_file=remove_tree_file)
		elif multiproc_bool and not save_tree:
			self._measure_xi_r_mur_sims_multiprocessing(num_nodes=self.num_nodes,
														masks=masks,
														dataset_name=dataset_name,
														return_output=False, rp_cut=rp_cut,
														print_num=True)
			self._measure_multipoles(corr_type=corr_type, dataset_name=dataset_name, return_output=False)
			if measure_cov:
				self._measure_jackknife_covariance_sims(masks=masks, corr_type=[corr_type, "multipoles"],
														dataset_name=dataset_name, L_subboxes=L,
														rp_cut=rp_cut, num_nodes=self.num_nodes,
														tree_saved=False)
		else:
			self._measure_xi_r_mur_sims_brute(masks=masks,
											  dataset_name=dataset_name,
											  return_output=False, print_num=True,
											  rp_cut=rp_cut)
			self._measure_multipoles(corr_type=corr_type, dataset_name=dataset_name, return_output=False)
			if measure_cov:
				self._measure_jackknife_covariance_sims(masks=masks, corr_type=[corr_type, "multipoles"],
														dataset_name=dataset_name, L_subboxes=L,
														rp_cut=rp_cut, num_nodes=self.num_nodes,
														tree_saved=False)

		return