diff --git a/README.md b/README.md index d917689..8fcdd1a 100644 --- a/README.md +++ b/README.md @@ -13,14 +13,12 @@ git clone https://github.com/corentinravoux/flip.git cd flip pip install . ``` -For now, the package requires you to install manually cosmoprimo: pip install git+https://github.com/cosmodesi/cosmoprimo - ## Required packages -Mandatory: numpy, scipy, matplotlib, [cosmoprimo](https://github.com/adematti/cosmoprimo), iminuit, emcee, sympy +Mandatory: numpy, scipy, matplotlib, iminuit, emcee, sympy -Optional: classy, pyccl, pypower, GPy, tensorflow +Optional: classy, pyccl, pmesh, GPy, tensorflow, [cosmoprimo](https://github.com/adematti/cosmoprimo) ## Examples diff --git a/flip/comparison/__init__.py b/flip/comparison/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/flip/comparison/likelihood.py b/flip/comparison/likelihood.py new file mode 100644 index 0000000..653965a --- /dev/null +++ b/flip/comparison/likelihood.py @@ -0,0 +1,242 @@ +import abc +from functools import partial + +from flip.data_vector.basic import DensMesh +from flip.utils import create_log, prior_sum, return_prior + +from .._config import __use_jax__ + +if __use_jax__: + try: + import jax.numpy as jnp + from jax import grad, jit + + jax_installed = True + + except ImportError: + import numpy as jnp + + jax_installed = False +else: + + import numpy as jnp + + jax_installed = False + + +# CR - Very first idea of the likelihood class + +log = create_log() + + +_available_field_expressions = ["interpolation", "remesh"] + + +def express_field_on_grid( + scaling_catalog, + parameter_values_dict, + fixed_grid, + method, + **kwargs, +): + if method not in _available_field_expressions: + raise ValueError( + f"Method {method} not recognized. Available methods: {_available_field_expressions}" + ) + + if method == "interpolation": + if isinstance(scaling_catalog, DensMesh): + log.warning( + "Scaling catalog is a dict. Interpolation method may not be appropriate." + ) + + # CR - not implemented yet - need to add interpolation method to DensMesh class + scaled_mesh = scaling_catalog.interpolate( + parameter_values_dict, + fixed_grid, + ) + else: + raise ValueError( + "Interpolation method only implemented for DensMesh catalogs." + ) + elif method == "remesh": + + scaling_catalog["rcom_zobs"] = ( + parameter_values_dict["H_0"] / 100 + ) * scaling_catalog["rcom_zobs"] + + rcom_max = kwargs.get("rcom_max", None) + grid_size = kwargs.get("grid_size", None) + grid_type = kwargs.get("grid_type", "sphere") + kind = kwargs.get("kind", "ngc") + + scaled_mesh = DensMesh.init_from_catalog( + scaling_catalog, + rcom_max, + grid_size, + grid_type, + kind, + **kwargs, + ) + + return scaled_mesh + + +class BaseLikelihood(abc.ABC): + + _default_likelihood_properties = { + "use_jit": False, + "use_gradient": False, + } + + def __init__( + self, + parameter_names=None, + likelihood_properties={}, + ): + self.parameter_names = parameter_names + + self.likelihood_properties = { + **self._default_likelihood_properties, + **likelihood_properties, + } + self.prior = self.initialize_prior() + + self.likelihood_call, self.likelihood_grad = self._init_likelihood() + + def __call__(self, parameter_values): + """Evaluate likelihood at parameter values. + + Args: + parameter_values (array-like): Parameter vector aligned with `parameter_names`. + + Returns: + float: Likelihood value, sign controlled by `negative_log_likelihood`. + """ + return self.likelihood_call(parameter_values) + + @abc.abstractmethod + def _init_likelihood(self, *args): + """Initialize likelihood and optional gradient. + + Returns: + tuple[Callable, Callable|None]: `(likelihood_call, likelihood_grad)`. + """ + likelihood_fun = None + likelihood_grad = None + return likelihood_fun, likelihood_grad + + def initialize_prior( + self, + ): + """Build prior function from likelihood properties. + + Returns: + Callable: Prior function mapping parameter dict to log-prior. + + Raises: + ValueError: If an unsupported prior type is requested. + """ + if "prior" not in self.likelihood_properties.keys(): + return lambda x: 0 + else: + prior_dict = self.likelihood_properties["prior"] + priors = [] + for parameter_name, prior_properties in prior_dict.items(): + prior = return_prior( + parameter_name, + prior_properties, + ) + priors.append(prior) + + prior_function = partial(prior_sum, priors) + return prior_function + + +class GaussianFieldComparisonLikelihood(BaseLikelihood): + + def __init__( + self, + basefield=None, + scaling_catalog=None, + parameter_names=None, + covariance_basefield=None, + likelihood_properties={}, + ): + + super(GaussianFieldComparisonLikelihood, self).__init__( + parameter_names=parameter_names, + likelihood_properties=likelihood_properties, + ) + self.basefield = basefield + self.scaling_catalog = scaling_catalog + self.covariance_basefield = covariance_basefield + + def _init_likelihood(self): + """Build callable likelihood and optional gradient for Gaussian model. + + Returns: + tuple[Callable, Callable|None]: `(likelihood_call, likelihood_grad)`. + """ + + use_jit = self.likelihood_properties["use_jit"] + + if jax_installed & use_jit: + prior = jit(self.prior) + else: + prior = self.prior + + def likelihood_evaluation( + parameter_values, + ): + """Evaluate likelihood for given parameters. + + Args: + parameter_values (array-like): Parameter vector aligned to names. + + Returns: + float: Likelihood value (sign depends on `negative_log_likelihood`). + """ + parameter_values_dict = dict(zip(self.parameter_names, parameter_values)) + + common_grid_x = self.basefield.data["x"] + common_grid_y = self.basefield.data["y"] + common_grid_z = self.basefield.data["z"] + + fixed_grid = jnp.stack( + jnp.meshgrid( + common_grid_x, common_grid_y, common_grid_z, indexing="ij" + ), + axis=-1, + ) + + scaled_field = express_field_on_grid( + self.scaling_catalog, + parameter_values_dict, + fixed_grid, + method=self.likelihood_properties["scaling_method"], + **self.likelihood_properties["scaling_kwargs"], + ) + + likelihood = jnp.sum( + (scaled_field["density"] - self.basefield["density"]) ** 2 + / ( + self.scaled_field["density_error"] ** 2 + + self.covariance_basefield["density_error"] ** 2 + ) + ) + + likelihood_value = likelihood + prior(parameter_values_dict) + + if self.likelihood_properties["negative_log_likelihood"]: + likelihood_value *= -1 + return likelihood_value + + if jax_installed: + likelihood_grad = grad(likelihood_evaluation) + if use_jit: + likelihood_evaluation = jit(likelihood_evaluation) + likelihood_grad = jit(likelihood_grad) + else: + likelihood_grad = None + return likelihood_evaluation, likelihood_grad diff --git a/flip/covariance/analytical/lai22/generator.py b/flip/covariance/analytical/lai22/generator.py index 17c40ef..0d3977f 100644 --- a/flip/covariance/analytical/lai22/generator.py +++ b/flip/covariance/analytical/lai22/generator.py @@ -2,13 +2,13 @@ import multiprocessing as mp from functools import partial -import cosmoprimo import numpy as np from scipy import integrate from scipy.interpolate import interp1d from scipy.special import factorial, spherical_jn from flip.covariance import cov_utils +from flip.covariance.hankel import PowerToCorrelation def compute_correlation_coefficient_simple_integration(p, q, ell, r, k, pk): @@ -41,7 +41,6 @@ def compute_correlation_coefficient_hankel( ): """Compute correlation coefficient using FFTLog Hankel transform. - Uses cosmoprimo's ``PowerToCorrelation`` to accelerate the computation. Falls back to direct integration for small ``r`` values where the Hankel result may be unreliable, controlled by ``hankel_overhead_coefficient``. @@ -58,8 +57,7 @@ def compute_correlation_coefficient_hankel( Array of correlation coefficient values at separations ``r``. """ integrand = k ** (2 * (p + q)) * pk - Hankel = cosmoprimo.fftlog.PowerToCorrelation(k, ell=ell, q=0, complex=False) - Hankel.set_fft_engine("numpy") + Hankel = PowerToCorrelation(k, ell=ell, q=0, complex=False) r_hankel, xi_hankel = Hankel(integrand) mask = r < np.min(r_hankel) * hankel_overhead_coefficient output = np.empty_like(r) diff --git a/flip/covariance/contraction.py b/flip/covariance/contraction.py index a7d1eee..7df2896 100644 --- a/flip/covariance/contraction.py +++ b/flip/covariance/contraction.py @@ -94,7 +94,6 @@ def init_from_flip( coordinate_type=coordinate_type, additional_parameters_values=additional_parameters_values, basis_definition=basis_definition, - covariance_prefactor_dict=covariance_prefactor_dict, **kwargs, ) diff --git a/flip/covariance/generator.py b/flip/covariance/generator.py index 1ca9a61..ad53084 100644 --- a/flip/covariance/generator.py +++ b/flip/covariance/generator.py @@ -2,7 +2,6 @@ import multiprocessing as mp from functools import partial -import cosmoprimo import mpmath import numpy as np from scipy import integrate @@ -10,6 +9,7 @@ from scipy.special import spherical_jn from flip.covariance import cov_utils +from flip.covariance.hankel import PowerToCorrelation from flip.utils import create_log log = create_log() @@ -65,11 +65,10 @@ def correlation_integration(ell, r, k, integrand): def correlation_hankel(ell, r, k, integrand, hankel_overhead_coefficient=2, kmin=None): """ - The correlation_hankel function is a wrapper for the cosmoprimo.fftlog.PowerToCorrelation function, + The correlation_hankel function is a wrapper for the PowerToCorrelation function, which computes the correlation function from power spectrum using FFTLog (Hamilton 2000). The PowerToCorrelation class takes in an array of k values and an array of P(k) values, and returns - an array of r values and an array of xi(r) values. The PowerToCorrelation class has two methods: set_fft_engine() - and __call__(). The set_fft_engine() method sets which fft engine to use; it can be either "n + an array of r values and an array of xi(r) values. Args: l: Determine the parity of the integrand, and therefore whether to add a 1j term @@ -84,8 +83,7 @@ def correlation_hankel(ell, r, k, integrand, hankel_overhead_coefficient=2, kmin Note: If l is odd, count a 1j term in the integrand, without the need for adding it """ - Hankel = cosmoprimo.fftlog.PowerToCorrelation(k, ell=ell, q=0, complex=False) - Hankel.set_fft_engine("numpy") + Hankel = PowerToCorrelation(k, ell=ell, q=0, complex=False) r_hankel, xi_hankel = Hankel(integrand) mask = r < np.min(r_hankel) * hankel_overhead_coefficient if np.any(r > np.max(r_hankel)): diff --git a/flip/covariance/hankel.py b/flip/covariance/hankel.py new file mode 100644 index 0000000..070694e --- /dev/null +++ b/flip/covariance/hankel.py @@ -0,0 +1,532 @@ +### This code is adapted from cosmoprimo (https://github.com/cosmodesi/cosmoprimo) +### Special thanks to the GOAT Arnaud de Mattia. + +import os + +import numpy as np +from scipy.special import loggamma + + +class FFTlog(object): + r""" + Implementation of the FFTlog algorithm presented in https://jila.colorado.edu/~ajsh/FFTLog/, which computes the generic integral: + + .. math:: + + G(y) = \int_{0}^{\infty} x dx F(x) K(xy) + + with :math:`F(x)` input function, :math:`K(xy)` a kernel. + + This transform is (mathematically) invariant under a power law transformation: + + .. math:: + + G_{q}(y) = \int_{0}^{\infty} x dx F_{q}(x) K_{q}(xy) + + where :math:`F_{q}(x) = G(x)x^{-q}`, :math:`K_{q}(t) = K(t)t^{q}` and :math:`G_{q}(y) = G(y)y^{q}`. + """ + + def __init__( + self, + x, + kernel, + q=0, + minfolds=2, + lowring=True, + xy=1, + check_level=0, + engine="numpy", + **engine_kwargs, + ): + r""" + Initialize :class:`FFTlog`, which can perform several transforms at once. + + Parameters + ---------- + x : array_like + Input log-spaced coordinates. Must be strictly increasing. + If 1D, is broadcast to the number of provided kernels. + + kernel : callable, list of callables + Mellin transform of the kernel: + .. math:: U_{K}(z) = \int_{0}^{\infty} t^{z-1} K(t) dt + If a list of kernels is provided, will perform all transforms at once. + + q : float, list of floats, default=0 + Power-law tilt(s) to regularise integration. + + minfolds : int, default=2 + The c is chosen with minimum :math:`n` chosen such that ``2**n > minfolds * x.size``. + + lowring : bool, default=True + If ``True`` set output coordinates according to the low-ringing condition, otherwise set it with ``xy``. + + xy : float, list of floats, default=1 + Enforce the reciprocal product (i.e. ``x[0] * y[-1]``) of the input ``x`` and output ``y`` coordinates. + + check_level : int, default=0 + If non-zero run sanity checks on input. + + engine : string, default='numpy' + FFT engine. See :meth:`set_fft_engine`. + + engine_kwargs : dict + Arguments for FFT engine. + + Note + ---- + Kernel definition is different from that of https://jila.colorado.edu/~ajsh/FFTLog/, which uses (eq. 10): + + .. math:: U_{K}(z) = \int_{0}^{\infty} t^{z} K(t) dt + + Therefore, one should use :math:`q = 1` for Bessel functions to match :math:`q = 0` in https://jila.colorado.edu/~ajsh/FFTLog/. + """ + self.inparallel = isinstance(kernel, (tuple, list)) + if not self.inparallel: + kernel = [kernel] + self.kernel = list(kernel) + if np.ndim(q) == 0: + q = [q] * self.nparallel + self.q = list(q) + self.x = np.asarray(x) + if not self.inparallel: + self.x = self.x[None, :] + elif self.x.ndim == 1: + self.x = np.tile(self.x[None, :], (self.nparallel, 1)) + if np.ndim(xy) == 0: + xy = [xy] * self.nparallel + self.xy = list(xy) + self.check_level = check_level + if self.check_level: + if len(self.x) != self.nparallel: + raise ValueError("x and kernel must of same length") + if len(self.q) != self.nparallel: + raise ValueError("q and kernel must be lists of same length") + if len(self.xy) != self.nparallel: + raise ValueError("xy and kernel must be lists of same length") + self.minfolds = minfolds + self.lowring = lowring + self.setup() + self.set_fft_engine(engine, **engine_kwargs) + + def set_fft_engine(self, engine="numpy", **engine_kwargs): + """ + Set up FFT engine. + See :func:`get_fft_engine` + + Parameters + ---------- + engine : BaseEngine, string, default='numpy' + FFT engine, or one of ['numpy', 'fftw']. + + engine_kwargs : dict + Arguments for FFT engine. + """ + self._engine = get_fft_engine( + engine, size=self.padded_size, nparallel=self.nparallel, **engine_kwargs + ) + + @property + def nparallel(self): + """Number of transforms performed in parallel.""" + return len(self.kernel) + + def setup(self): + """Set up u funtions.""" + self.size = self.x.shape[-1] + self.delta = np.log(self.x[:, -1] / self.x[:, 0]) / (self.size - 1) + + nfolds = (self.size * self.minfolds - 1).bit_length() + self.padded_size = 2**nfolds + npad = self.padded_size - self.size + self.padded_size_in_left, self.padded_size_in_right = ( + npad // 2, + npad - npad // 2, + ) + self.padded_size_out_left, self.padded_size_out_right = ( + npad - npad // 2, + npad // 2, + ) + + if self.check_level: + if not np.allclose( + np.log(self.x[:, 1:] / self.x[:, :-1]), self.delta, rtol=1e-3 + ): + raise ValueError("Input x must be log-spaced") + if self.padded_size < self.size: + raise ValueError("Convolution size must be larger than input x size") + + if self.lowring: + self.lnxy = np.array( + [ + delta / np.pi * np.angle(kernel(q + 1j * np.pi / delta)) + for kernel, delta, q in zip(self.kernel, self.delta, self.q) + ] + ) + else: + self.lnxy = np.log(self.xy) + self.delta + + self.y = np.exp(self.lnxy - self.delta)[:, None] / self.x[:, ::-1] + + m = np.arange(0, self.padded_size // 2 + 1) + self.padded_u, self.padded_prefactor, self.padded_postfactor = [], [], [] + self.padded_x = pad( + self.x, + (self.padded_size_in_left, self.padded_size_in_right), + axis=-1, + extrap="log", + ) + self.padded_y = pad( + self.y, + (self.padded_size_out_left, self.padded_size_out_right), + axis=-1, + extrap="log", + ) + prev_kernel, prev_q, prev_delta, prev_u = None, None, None, None + for kernel, padded_x, padded_y, lnxy, delta, q in zip( + self.kernel, self.padded_x, self.padded_y, self.lnxy, self.delta, self.q + ): + self.padded_prefactor.append(padded_x ** (-q)) + self.padded_postfactor.append(padded_y ** (-q)) + if kernel is prev_kernel and q == prev_q and delta == prev_delta: + u = prev_u + else: + u = prev_u = kernel(q + 2j * np.pi / self.padded_size / delta * m) + self.padded_u.append( + u * np.exp(-2j * np.pi * lnxy / self.padded_size / delta * m) + ) + prev_kernel, prev_q, prev_delta = kernel, q, delta + self.padded_u = np.array(self.padded_u) + self.padded_prefactor = np.array(self.padded_prefactor) + self.padded_postfactor = np.array(self.padded_postfactor) + + def __call__(self, fun, extrap=0, keep_padding=False): + """ + Perform the transforms. + + Parameters + ---------- + fun : array_like + Function to be transformed. + Last dimensions should match (:attr:`nparallel`,len(x)) where ``len(x)`` is the size of the input x-coordinates. + (if :attr:`nparallel` is 1, the only requirement is the last dimension to be (len(x))). + + extrap : float, string, default=0 + How to extrapolate function outside of ``x`` range to fit the integration range. + If 'log', performs a log-log extrapolation. + If 'edge', pad ``fun`` with its edge values. + Else, pad ``fun`` with the provided value. + Pass a tuple to differentiate between left and right sides. + + keep_padding : bool, default=False + Whether to return function padded to the number of points in the integral. + By default, crop it to its original size. + + Returns + ------- + y : array + Array of new coordinates. + + fftloged : array + Transformed function. + """ + fun = np.asarray(fun) + padded_fun = pad( + fun, + (self.padded_size_in_left, self.padded_size_in_right), + axis=-1, + extrap=extrap, + ) + fftloged = ( + self._engine.backward( + self._engine.forward(padded_fun * self.padded_prefactor) * self.padded_u + ) + * self.padded_postfactor + ) + + if not keep_padding: + y = self.y + fftloged = fftloged[ + ..., self.padded_size_out_left : self.padded_size_out_left + self.size + ] + else: + y = self.padded_y + if not self.inparallel: + y = y[0] + fftloged = np.reshape( + fftloged, + fun.shape if not keep_padding else fun.shape[:-1] + (self.padded_size,), + ) + return y, fftloged + + def inv(self): + """Inverse the transform.""" + self.x, self.y = self.y, self.x + self.padded_x, self.padded_y = self.y, self.x + self.padded_prefactor, self.padded_postfactor = ( + 1 / self.padded_postfactor, + 1 / self.padded_prefactor, + ) + self.padded_u = 1 / self.padded_u.conj() + + +class PowerToCorrelation(FFTlog): + r""" + Power spectrum to correlation function transform, defined as: + + .. math:: + \xi_{\ell}(s) = \frac{(-i)^{\ell}}{2 \pi^{2}} \int dk k^{2} P_{\ell}(k) j_{\ell}(ks) + + """ + + def __init__(self, k, ell=0, q=0, complex=False, **kwargs): + """ + Initialize power to correlation transform. + + Parameters + ---------- + k : array_like + Input log-spaced wavenumbers. + If 1D, is broadcast to the number of provided ``ell``. + + ell : int, list of int, default=0 + Poles. If a list is provided, will perform all transforms at once. + + q : float, list of floats, default=0 + Power-law tilt(s) to regularise integration. + + complex : bool, default=False + ``False`` assumes the imaginary part of odd power spectrum poles is provided. + + kwargs : dict + Arguments for :class:`FFTlog`. + """ + if np.ndim(ell) == 0: + kernel = SphericalBesselJKernel(ell) + else: + kernel = [SphericalBesselJKernel(ell_) for ell_ in ell] + FFTlog.__init__(self, k, kernel, q=1.5 + q, **kwargs) + self.padded_prefactor *= self.padded_x**3 / (2 * np.pi) ** 1.5 + # Convention is (-i)^ell/(2 pi^2) + ell = np.atleast_1d(ell) + if complex: + phase = (-1j) ** ell + else: + # Prefactor is (-i)^ell, but we take in the imaginary part of odd power spectra, hence: + # (-i)^ell = (-1)^(ell/2) if ell is even + # (-i)^ell i = (-1)^(ell//2) if ell is odd + phase = (-1) ** (ell // 2) + # Not in-place as phase (and hence padded_postfactor) may be complex instead of float + self.padded_postfactor = self.padded_postfactor * phase[:, None] + + +class CorrelationToPower(FFTlog): + r""" + Correlation function to power spectrum transform, defined as: + + .. math:: + P_{\ell}(k) = 4 \pi i^{\ell} \int ds s^{2} \xi_{\ell}(s) j_{\ell}(ks) + + """ + + def __init__(self, s, ell=0, q=0, complex=False, **kwargs): + """ + Initialize power to correlation transform. + + Parameters + ---------- + s : array_like + Input log-spaced separations. + If 1D, is broadcast to the number of provided ``ell``. + + ell : int, list of int, default=0 + Poles. If a list is provided, will perform all transforms at once. + + q : float, list of floats, default=0 + Power-law tilt(s) to regularise integration. + + complex : bool, default=False + ``False`` returns the real part of even poles, and the imaginary part of odd poles. + + kwargs : dict + Arguments for :class:`FFTlog`. + """ + if np.ndim(ell) == 0: + kernel = SphericalBesselJKernel(ell) + else: + kernel = [SphericalBesselJKernel(ell_) for ell_ in ell] + FFTlog.__init__(self, s, kernel, q=1.5 + q, **kwargs) + self.padded_prefactor *= self.padded_x**3 * (2 * np.pi) ** 1.5 + # Convention is 4 \pi i^ell, and we return imaginary part of odd poles + ell = np.atleast_1d(ell) + if complex: + phase = (-1j) ** ell + else: + # We return imaginary part of odd poles + phase = (-1) ** (ell // 2) + self.padded_postfactor = self.padded_postfactor * phase[:, None] + + +def pad(array, pad_width, axis=-1, extrap=0): + """ + Pad array along ``axis``. + + Parameters + ---------- + array : array_like + Input array to be padded. + + pad_width : int, tuple of ints + Number of points to be added on both sides of the array. + Pass a tuple to differentiate between left and right sides. + + axis : int, default=-1 + Axis along which padding is to be applied. + + extrap : string, float, default=0 + If 'log', performs a log-log extrapolation. + If 'edge', pad ``array`` with its edge values. + Else, pad ``array`` with the provided value. + Pass a tuple to differentiate between left and right sides. + + Returns + ------- + array : array + Padded array. + """ + array = np.asarray(array) + + try: + pad_width_left, pad_width_right = pad_width + except (TypeError, ValueError): + pad_width_left = pad_width_right = pad_width + + try: + extrap_left, extrap_right = extrap + except (TypeError, ValueError): + extrap_left = extrap_right = extrap + + axis = axis % array.ndim + to_axis = [1] * array.ndim + to_axis[axis] = -1 + + if extrap_left == "edge": + end = np.take(array, [0], axis=axis) + pad_left = np.repeat(end, pad_width_left, axis=axis) + elif extrap_left == "log": + end = np.take(array, [0], axis=axis) + ratio = np.take(array, [1], axis=axis) / end + exp = np.arange(-pad_width_left, 0).reshape(to_axis) + pad_left = end * ratio**exp + else: + pad_left = np.full( + array.shape[:axis] + (pad_width_left,) + array.shape[axis + 1 :], + extrap_left, + ) + + if extrap_right == "edge": + end = np.take(array, [-1], axis=axis) + pad_right = np.repeat(end, pad_width_right, axis=axis) + elif extrap_right == "log": + end = np.take(array, [-1], axis=axis) + ratio = np.take(array, [-2], axis=axis) / end + exp = np.arange(1, pad_width_right + 1).reshape(to_axis) + pad_right = end / ratio**exp + else: + pad_right = np.full( + array.shape[:axis] + (pad_width_right,) + array.shape[axis + 1 :], + extrap_right, + ) + + return np.concatenate([pad_left, array, pad_right], axis=axis) + + +class BaseFFTEngine(object): + """Base FFT engine.""" + + def __init__(self, size, nparallel=1, nthreads=None): + """ + Initialize FFT engine. + + Parameters + ---------- + size : int + Array size. + + nparallel : int, default=1 + Number of FFTs to be performed in parallel. + + nthreads : int, default=None + Number of threads. + """ + self.size = size + self.nparallel = nparallel + if nthreads is not None: + os.environ["OMP_NUM_THREADS"] = str(nthreads) + self.nthreads = int(os.environ.get("OMP_NUM_THREADS", 1)) + + +class NumpyFFTEngine(BaseFFTEngine): + """FFT engine based on :mod:`numpy.fft`.""" + + def forward(self, fun): + """Forward transform of ``fun``.""" + return np.fft.rfft(fun, axis=-1) + + def backward(self, fun): + """Backward transform of ``fun``.""" + return np.fft.irfft(fun.conj(), n=self.size, axis=-1) + + +def get_fft_engine(engine, *args, **kwargs): + """ + Return FFT engine. + + Parameters + ---------- + engine : BaseFFTEngine, string + FFT engine, or one of ['numpy', 'fftw']. + + args, kwargs : tuple, dict + Arguments for FFT engine. + + Returns + ------- + engine : BaseFFTEngine + """ + if isinstance(engine, str): + if engine.lower() == "numpy": + return NumpyFFTEngine(*args, **kwargs) + raise ValueError("FFT engine {} is unknown".format(engine)) + return engine + + +class BaseKernel(object): + """Base kernel.""" + + def __call__(self, z): + return self.eval(z) + + def __eq__(self, other): + return other.__class__ == self.__class__ + + +class BaseBesselKernel(BaseKernel): + """Base Bessel kernel.""" + + def __init__(self, nu): + self.nu = nu + + def __eq__(self, other): + return other.__class__ == self.__class__ and other.nu == self.nu + + +class SphericalBesselJKernel(BaseBesselKernel): + """(Mellin transform of) spherical Bessel kernel.""" + + def eval(self, z): + return np.exp( + np.log(2) * (z - 1.5) + + loggamma(0.5 * (self.nu + z)) + - loggamma(0.5 * (3 + self.nu - z)) + ) diff --git a/flip/covariance/likelihood.py b/flip/covariance/likelihood.py index ac6e2ad..e27469d 100644 --- a/flip/covariance/likelihood.py +++ b/flip/covariance/likelihood.py @@ -4,7 +4,7 @@ import numpy as np import scipy as sc -from flip.utils import create_log +from flip.utils import create_log, prior_sum, return_prior from .._config import __use_jax__ @@ -31,9 +31,6 @@ log = create_log() - -_available_priors = ["gaussian", "positive", "uniform"] - _available_inversion_methods = [ "inverse", "solve", @@ -151,31 +148,6 @@ def log_likelihood_gaussian_cholesky_regularized(vector, covariance_sum): log_likelihood_gaussian_solve_jit = jit(log_likelihood_gaussian_solve) -def no_prior(x): - """Return zero prior contribution. - - Args: - x (Any): Ignored parameter values input. - - Returns: - int: Zero. - """ - return 0 - - -def prior_sum(priors, x): - """Sum multiple prior contributions. - - Args: - priors (list[Callable]): List of prior callables accepting a parameter dict. - x (dict): Parameter values dictionary. - - Returns: - float: Sum of all prior log-probabilities. - """ - return sum(prior(x) for prior in priors) - - class BaseLikelihood(abc.ABC): """Abstract base class for likelihood evaluation. @@ -297,31 +269,15 @@ def initialize_prior( ValueError: If an unsupported prior type is requested. """ if "prior" not in self.likelihood_properties.keys(): - return no_prior + return lambda x: 0 else: prior_dict = self.likelihood_properties["prior"] priors = [] for parameter_name, prior_properties in prior_dict.items(): - if prior_properties["type"].lower() not in _available_priors: - raise ValueError( - f"""The prior type {prior_properties["type"]} is not available""" - f"""Please choose between {_available_priors}""" - ) - elif prior_properties["type"].lower() == "gaussian": - prior = GaussianPrior( - parameter_name=parameter_name, - prior_mean=prior_properties["mean"], - prior_standard_deviation=prior_properties["standard_deviation"], - ) - elif prior_properties["type"].lower() == "positive": - prior = PositivePrior( - parameter_name=parameter_name, - ) - elif prior_properties["type"].lower() == "uniform": - prior = UniformPrior( - parameter_name=parameter_name, - range=prior_properties["range"], - ) + prior = return_prior( + parameter_name, + prior_properties, + ) priors.append(prior) prior_function = partial(prior_sum, priors) @@ -717,107 +673,3 @@ def __call__( if self.likelihood_properties["negative_log_likelihood"]: return -likelihood_function(vector, covariance_sum) - prior_value return likelihood_function(vector, covariance_sum) + prior_value - - -class Prior: - """Base prior class encapsulating parameter-specific priors. - - Attributes: - parameter_name (str): Name of the parameter this prior applies to. - """ - - def __init__( - self, - parameter_name=None, - ): - self.parameter_name = parameter_name - - -class GaussianPrior(Prior): - """Univariate Gaussian prior on a parameter. - - Models $p(\theta) \propto \exp\{-\tfrac{1}{2}[(\theta-\mu)^2/\sigma^2]\}$. - """ - - def __init__( - self, - parameter_name=None, - prior_mean=None, - prior_standard_deviation=None, - ): - super().__init__(parameter_name=parameter_name) - self.prior_mean = prior_mean - self.prior_standard_deviation = prior_standard_deviation - - def __call__( - self, - parameter_values_dict, - ): - """Return Gaussian log-prior for the parameter value. - - Args: - parameter_values_dict (dict): Map of parameter names to values. - - Returns: - float: Log-prior value. - """ - return -0.5 * ( - np.log(2 * jnp.pi * self.prior_standard_deviation**2) - + (parameter_values_dict[self.parameter_name] - self.prior_mean) ** 2 - / self.prior_standard_deviation**2 - ) - - -class PositivePrior(Prior): - """Log-prior enforcing parameter positivity via Heaviside function. - - Returns `log(Heaviside(value))`, which is `0` for positive values and `-inf` otherwise. - """ - - def __init__( - self, - parameter_name=None, - ): - super().__init__(parameter_name=parameter_name) - - def __call__( - self, - parameter_values_dict, - ): - """Return log-prior that is zero for positive values, -inf otherwise. - - Args: - parameter_values_dict (dict): Map of parameter names to values. - - Returns: - float: Log-prior (0 or -inf). - """ - return jnp.log(jnp.heaviside(parameter_values_dict[self.parameter_name], 0)) - - -class UniformPrior(Prior): - """Uniform prior over a finite interval for a parameter. - - Uses `scipy.stats.uniform.logpdf` for the specified range. - """ - - def __init__(self, parameter_name=None, range=None): - super().__init__(parameter_name=parameter_name) - self.range = range - - def __call__( - self, - parameter_values_dict, - ): - """Return uniform log-prior over `[range[0], range[1]]`. - - Args: - parameter_values_dict (dict): Map of parameter names to values. - - Returns: - float: Log-prior value (constant inside range, -inf outside). - """ - value = parameter_values_dict[self.parameter_name] - return jsc.stats.uniform.logpdf( - value, loc=self.range[0], scale=self.range[1] - self.range[0] - ) diff --git a/flip/covariance/plot_utils.py b/flip/covariance/plot_utils.py index 47e01e6..03e69e0 100644 --- a/flip/covariance/plot_utils.py +++ b/flip/covariance/plot_utils.py @@ -4,6 +4,7 @@ import matplotlib.pyplot as plt import numpy as np +from scipy.stats import median_abs_deviation from flip import utils from flip.covariance import cov_utils @@ -294,6 +295,7 @@ def plot_all_mean_fits( remove_higher=None, plot=True, use_minos=False, + use_mad=False, **kwargs, ): """Plot mean parameter values and errors across grouped fits. @@ -448,7 +450,10 @@ def plot_all_mean_fits( else: error_mean_param = np.mean(errors) / np.sqrt(len(mask[mask])) mean_error_param = np.mean(errors) - std_param = np.std(params) + if use_mad: + std_param = median_abs_deviation(params) + else: + std_param = np.std(params) count = len(params) mean_param_dict[param_name].append(np.array(mean_param)) diff --git a/flip/data_vector/__init__.py b/flip/data_vector/__init__.py index 24d2de6..bdfefd4 100644 --- a/flip/data_vector/__init__.py +++ b/flip/data_vector/__init__.py @@ -1,4 +1,4 @@ """Init file of the flip.data_vector package.""" -from . import cosmo_utils, galaxypv_vectors, gridding, snia_vectors, vector_utils +from . import cosmo_utils, galaxypv_vectors, snia_vectors, vector_utils, gw_vectors, mesh from .basic import * diff --git a/flip/data_vector/basic.py b/flip/data_vector/basic.py index d90dd98..8c282d9 100644 --- a/flip/data_vector/basic.py +++ b/flip/data_vector/basic.py @@ -5,6 +5,7 @@ import numpy as np from flip.covariance import CovMatrix +from flip.data_vector import mesh from flip.utils import create_log from .._config import __use_jax__ @@ -285,6 +286,95 @@ def __init__(self, data, covariance_observation=None): self._covariance_observation = velocity_variance +class DensMesh(Dens): + _kind = "density" + _needed_keys = ["density", "density_error"] + _free_par = [] + _number_dimension_observation_covariance = 1 + _parameters_observation_covariance = ["density"] + + def give_data_and_variance(self, *args): + """Return density data and diagonal variance from `density_error`. + + Returns: + tuple: (density, density_error^2). + """ + + if self._covariance_observation is not None: + return self._data["density"], self._covariance_observation + return self._data["density"], self._data["density_error"] ** 2 + + def __init__(self, data, covariance_observation=None): + super().__init__(data, covariance_observation=covariance_observation) + + @classmethod + def init_from_catalog( + cls, + data_position_sky, + rcom_max, + grid_size, + grid_type, + kind, + **kwargs, + ): + grid = mesh.grid_data_density( + data_position_sky, + rcom_max, + grid_size, + grid_type, + kind, + **kwargs, + ) + + return cls(grid) + + +class DirectVelMesh(DirectVel): + _kind = "velocity" + _needed_keys = ["velocity", "velocity_error"] + _free_par = [] + _number_dimension_observation_covariance = 1 + _parameters_observation_covariance = ["velocity"] + + def give_data_and_variance(self, *args): + """Return velocity data and diagonal variance from `velocity_error`. + + Returns: + tuple: (velocity, velocity_error^2). + """ + + if self._covariance_observation is not None: + return self._data["velocity"], self._covariance_observation + return self._data["velocity"], self._data["velocity_error"] ** 2 + + def __init__(self, data, covariance_observation=None): + super().__init__(data, covariance_observation=covariance_observation) + + @classmethod + def init_from_catalog( + cls, + data_position_sky, + data, + rcom_max, + grid_size, + grid_type, + kind, + **kwargs, + ): + grid_velocity = mesh.grid_data_velocity( + data_position_sky, + rcom_max, + grid_size, + grid_type, + kind, + data["velocity_error"] ** 2, + velocity=data["velocity"], + **kwargs, + ) + + return cls(grid_velocity) + + class VelFromHDres(DataVector): _kind = "velocity" _needed_keys = ["dmu", "zobs"] diff --git a/flip/data_vector/gridding.py b/flip/data_vector/gridding.py deleted file mode 100644 index 78d81dc..0000000 --- a/flip/data_vector/gridding.py +++ /dev/null @@ -1,994 +0,0 @@ -import numpy as np - -from flip import utils - -log = utils.create_log() -try: - from pypower import CatalogMesh -except ImportError: - log.add("No pypower module detected, gridding with this method is unavailable") - -# CR - No cut in healpix implemented with randoms - -_GRID_KIND = ["ngp", "ngp_errw", "cic", "tsc", "pcs"] - - -def _compute_grid_window(grid_size, k, order, n): - """Numerically compute isotropic grid assignment window. - - Uses spherical averaging over directions to produce a 1D window for - resampler order (NGP/CIC/TSC/PCS) given grid size. - - Args: - grid_size (float): Grid cell size. - k (array-like): Wavenumbers at which to evaluate the window. - order (int): Assignment order (1:NGP, 2:CIC, 3:TSC, 4:PCS). - n (int): Number of angular samples for spherical averaging. - - Returns: - numpy.ndarray: Window values for each `k`. - """ - window = np.zeros_like(k) - theta = np.linspace(0, np.pi, n) - phi = np.linspace(0, 2 * np.pi, n) - kx = np.outer(np.sin(theta), np.cos(phi)) - ky = np.outer(np.sin(theta), np.sin(phi)) - kz = np.outer(np.cos(theta), np.ones(n)) - - # Forgotten in Howlett et al. formula - # we add spherical coordinate solid angle element - dthetaphi = np.outer(np.sin(theta), np.ones(phi.size)) - for i in range(k.size): - # the factor here has an extra np.pi because of the definition of np.sinc - fact = (k[i] * grid_size) / (2 * np.pi) - func = ( - np.sinc(fact * kx) * np.sinc(fact * ky) * np.sinc(fact * kz) - ) ** order * dthetaphi - win_theta = np.trapz(func, x=phi) - window[i] = np.trapz(win_theta, x=theta) - window *= 1 / (4 * np.pi) - return window - - -def compute_grid_window(grid_size, kh, kind="ngp", n=1000): - """Compute grid assignment window for a given resampler kind. - - Args: - grid_size (float): Grid cell size. - kh (array-like): Wavenumbers at which to evaluate the window. - kind (str): One of `ngp`, `ngp_errw`, `cic`, `tsc`, `pcs`. - n (int): Angular samples for spherical averaging. - - Returns: - numpy.ndarray|None: Window values, or None if `grid_size==0`. - """ - _order_dic = { - "ngp": 1, - "ngp_errw": 1, - "cic": 2, - "tsc": 3, - "pcs": 4, - } - - if kind not in _GRID_KIND: - allowed = ", ".join(_GRID_KIND) - raise ValueError(f"INVALID GRID TYPE! Allowed kinds: {allowed}") - if grid_size == 0: - return None - return _compute_grid_window(grid_size, kh, _order_dic[kind], n) - - -def construct_grid_regular_sphere(grid_size, rcom_max): - """Construct a regular spherical grid of voxel centers. - - Args: - grid_size (float): Cell size. - rcom_max (float): Maximum comoving radius cutoff. - - Returns: - dict: Grid dictionary with keys `ra`, `dec`, `rcom`, `x`, `y`, `z` after cut. - """ - - # Number of grid voxels per axis - n_grid = np.floor(rcom_max / grid_size).astype(np.int64) - - # Determine the coordinates of the voxel centers - cp_x, cp_y, cp_z = np.reshape( - np.meshgrid( - *(np.linspace(-n_grid * grid_size, n_grid * grid_size, 2 * n_grid + 1),) - * 3, - indexing="ij", - ), - (3, (2 * n_grid + 1) ** 3), - ) - - # Convert to ra, dec, r_comov - center_r_comov, center_ra, center_dec = utils.cart2radec(cp_x, cp_y, cp_z) - - # Cut the grid with rcom_max - grid = { - "ra": center_ra, - "dec": center_dec, - "rcom_zobs": center_r_comov, - "x": cp_x, - "y": cp_y, - "z": cp_z, - } - cut_grid(grid, n_cut=None, weight_min=None, rcom_max=rcom_max) - - return grid - - -def construct_grid_regular_rectangular(grid_size, rcom_max): - """Construct a regular rectangular grid of voxel centers. - - Args: - grid_size (float): Cell size. - rcom_max (float): Half-side length of the cube. - - Returns: - dict: Grid dictionary with `ra`, `dec`, `rcom`, `x`, `y`, `z`. - """ - - # Number of grid voxels per axis - n_grid = np.floor(rcom_max / grid_size).astype(np.int64) - - # Determine the coordinates of the voxel centers - cp_x, cp_y, cp_z = np.reshape( - np.meshgrid( - *(np.linspace(-n_grid * grid_size, n_grid * grid_size, 2 * n_grid + 1),) - * 3, - indexing="ij", - ), - (3, (2 * n_grid + 1) ** 3), - ) - - # Convert to ra, dec, r_comov - center_r_comov, center_ra, center_dec = utils.cart2radec(cp_x, cp_y, cp_z) - - # Cut the grid with rcom_max - grid = { - "ra": center_ra, - "dec": center_dec, - "rcom_zobs": center_r_comov, - "x": cp_x, - "y": cp_y, - "z": cp_z, - } - return grid - - -def ngp_weight(ds): - """Nearest Grid Point.""" - abs_ds = np.abs(ds) - w = 1.0 * (abs_ds < 1 / 2) - return w - - -def ngp_errw_weight(ds): - """Nearest Grid Point with Weighted error.""" - return ngp_weight(ds) - - -def cic_weight(ds): - """Cloud In Cell.""" - abs_ds = np.abs(ds) - w = (1.0 - abs_ds) * (abs_ds < 1) - return w - - -def tsc_weight(ds): - """Triangular Shaped Cloud.""" - abs_ds = np.abs(ds) - w = (3 / 4 - ds**2) * (abs_ds < 1 / 2) - w += 1 / 2 * (3 / 2 - abs_ds) ** 2 * ((1 / 2 <= abs_ds) & (abs_ds < 3 / 2)) - return w - - -def pcs_weight(ds): - """Triangular Shaped Cloud.""" - abs_ds = np.abs(ds) - w = 1 / 6 * (4 - 6 * ds**2 + 3 * abs_ds**3) * (abs_ds < 1) - w += 1 / 6 * (2 - abs_ds) ** 3 * ((1 <= abs_ds) & (abs_ds < 2)) - return w - - -def attribute_weight_density( - grid_size, - xobj, - yobj, - zobj, - xgrid, - ygrid, - zgrid, - weight_fun, -): - """Accumulate assignment weights and counts for objects onto a grid. - - Args: - grid_size (float): Cell size. - xobj (array-like): Object x positions. - yobj (array-like): Object y positions. - zobj (array-like): Object z positions. - xgrid (array-like): Grid x centers. - ygrid (array-like): Grid y centers. - zgrid (array-like): Grid z centers. - weight_fun (Callable): Resampler weight function of distance. - - Returns: - tuple[numpy.ndarray, numpy.ndarray, numpy.ndarray]: - sum of weights, sum of squared weights, and objects-per-cell counts. - """ - - Nobj = len(xobj) - Ngrid = len(xgrid) - - # Init array - grid_val = np.zeros(Ngrid, dtype="float") - grid_var = np.zeros(Ngrid, dtype="float") - nobj_in_cell = np.zeros(Ngrid, dtype="int") - - for i in range(Nobj): - dX = (xgrid - xobj[i]) / grid_size - dY = (ygrid - yobj[i]) / grid_size - dZ = (zgrid - zobj[i]) / grid_size - - w = weight_fun(dX) * weight_fun(dY) * weight_fun(dZ) - - grid_val += w - grid_var += w**2 - - nobj_in_cell += (w != 0).astype("int") - - return grid_val, grid_var, nobj_in_cell - - -def define_randoms( - random_method, - xobj, - yobj, - zobj, - raobj, - decobj, - rcomobj, - Nrandom=None, - coord_randoms=None, - max_coordinates=None, -): - """Generate random positions for density estimation. - - Supports cartesian uniform, choice-based on observed distributions, or from file. - - Args: - random_method (str): `cartesian`, `choice`, `choice_redshift`, or `file`. - xobj (array-like): Data x positions. - yobj (array-like): Data y positions. - zobj (array-like): Data z positions. - raobj (array-like): Data right ascensions. - decobj (array-like): Data declinations. - rcomobj (array-like): Data comoving distances. - Nrandom (int): Number of randoms per object. - coord_randoms (tuple, optional): `(ra, dec, rcom)` for `file` method. - max_coordinates (float, optional): Coordinate cutoff for `file` method. - - Returns: - tuple[numpy.ndarray, numpy.ndarray, numpy.ndarray]: Random x, y, z positions. - """ - N = xobj.size - - # Uniform in X,Y,Z - if random_method == "cartesian": - xobj_random = (np.max(xobj) - np.min(xobj)) * np.random.random( - size=Nrandom * N - ) + np.min(xobj) - yobj_random = (np.max(yobj) - np.min(yobj)) * np.random.random( - size=Nrandom * N - ) + np.min(yobj) - zobj_random = (np.max(zobj) - np.min(zobj)) * np.random.random( - size=Nrandom * N - ) + np.min(zobj) - - # Choice in the ra, dec, and redshift data coordinates to create the random. - elif random_method == "choice": - counts_ra, bins_ra = np.histogram(raobj, bins=1000) - ra_random = np.random.choice( - (bins_ra[1:] + bins_ra[:-1]) / 2, - p=counts_ra / float(counts_ra.sum()), - size=Nrandom * N, - ) - counts_dec, bins_dec = np.histogram(decobj, bins=1000) - dec_random = np.random.choice( - (bins_dec[1:] + bins_dec[:-1]) / 2, - p=counts_dec / float(counts_dec.sum()), - size=Nrandom * N, - ) - counts_rcom, bins_rcom = np.histogram(rcomobj, bins=1000) - rcom_random = np.random.choice( - (bins_rcom[1:] + bins_rcom[:-1]) / 2, - p=counts_rcom / float(counts_rcom.sum()), - size=Nrandom * N, - ) - - xobj_random, yobj_random, zobj_random = utils.radec2cart( - rcom_random, ra_random, dec_random - ) - - # Uniform in ra and dec. Choice in the redshift data coordinates. - elif random_method == "choice_redshift": - ra_random = (np.max(raobj) - np.min(raobj)) * np.random.random( - size=Nrandom * N - ) + np.min(raobj) - dec_random = (np.max(decobj) - np.min(decobj)) * np.random.random( - size=Nrandom * N - ) + np.min(decobj) - counts_rcom, bins_rcom = np.histogram(rcomobj, bins=1000) - rcom_random = np.random.choice( - (bins_rcom[1:] + bins_rcom[:-1]) / 2, - p=counts_rcom / float(counts_rcom.sum()), - size=Nrandom * N, - ) - - xobj_random, yobj_random, zobj_random = utils.radec2cart( - rcom_random, ra_random, dec_random - ) - - # From random file - elif random_method == "file": - ra_random, dec_random, rcom_random = ( - coord_randoms[0], - coord_randoms[1], - coord_randoms[2], - ) - - xobj_random, yobj_random, zobj_random = utils.radec2cart( - rcom_random, ra_random, dec_random - ) - if max_coordinates is not None: - mask_random = np.abs(xobj_random) < max_coordinates - mask_random &= np.abs(yobj_random) < max_coordinates - mask_random &= np.abs(zobj_random) < max_coordinates - xobj_random = xobj_random[mask_random] - yobj_random = yobj_random[mask_random] - zobj_random = zobj_random[mask_random] - - return xobj_random, yobj_random, zobj_random - - -def grid_data_density( - grid, - grid_size, - ra, - dec, - rcom, - kind="ngp", - n_cut=None, - weight_min=None, - verbose=False, - compute_density=True, - Nrandom=10, - random_method="cartesian", - coord_randoms=None, -): - """Grid objects onto a mesh and compute density contrast and errors. - - Args: - grid (dict): Grid coordinates dictionary from constructors. - grid_size (float): Cell size. - ra (array-like): Object right ascensions. - dec (array-like): Object declinations. - rcom (array-like): Object comoving distances. - kind (str): Resampler kind (`ngp`, `cic`, `tsc`, `pcs`, `ngp_errw`). - n_cut (int, optional): Minimum objects per cell. - weight_min (float, optional): Minimum weight per cell. - verbose (bool): Print number of cells. - compute_density (bool): Whether to compute density contrast. - Nrandom (int): Randoms per object if computing density. - random_method (str): Random generation method. - coord_randoms (tuple, optional): Randoms coordinates for `file` method. - - Returns: - dict: Updated grid with weights, counts, density, and errors. - """ - # Check valid input grid kind - kind = kind.lower() - if kind not in _GRID_KIND: - raise ValueError( - "INVALID GRID TYPE ! \n Grid allowed : " + f"{k}" for k in _GRID_KIND - ) - - xobj, yobj, zobj = utils.radec2cart(rcom, ra, dec) - xgrid, ygrid, zgrid = utils.radec2cart(grid["rcom_zobs"], grid["ra"], grid["dec"]) - - # Compute weight - weight_fun = globals()[kind + "_weight"] - sum_weights, var_weights, n_in_cell = attribute_weight_density( - grid_size, xobj, yobj, zobj, xgrid, ygrid, zgrid, weight_fun - ) - - grid["sum_weights"] = sum_weights - grid["var_weights"] = var_weights - grid["nincell"] = n_in_cell - - if compute_density: - ( - xobj_random, - yobj_random, - zobj_random, - ) = define_randoms( - random_method, - xobj, - yobj, - zobj, - ra, - dec, - rcom, - Nrandom=Nrandom, - coord_randoms=coord_randoms, - ) - - ( - sum_weights_random, - var_weights_random, - n_in_cell_random, - ) = attribute_weight_density( - grid_size, - xobj_random, - yobj_random, - zobj_random, - xgrid, - ygrid, - zgrid, - weight_fun, - ) - grid["sum_weights_random"] = sum_weights_random - grid["var_weights_random"] = var_weights_random - grid["nincell_random"] = n_in_cell_random - - grid["density"] = ( - (sum_weights / np.sum(sum_weights)) - / (sum_weights_random / np.sum(sum_weights_random)) - ) - 1 - - grid["density_error"] = np.sqrt( - 1 / (n_in_cell_random * (np.sum(n_in_cell) / np.sum(n_in_cell_random))) - ) - - grid["density_nincell"] = ( - (n_in_cell / np.sum(n_in_cell)) - / (n_in_cell_random / np.sum(n_in_cell_random)) - ) - 1 - - cut_grid(grid, n_cut=n_cut, weight_min=weight_min) - - if verbose: - print(f"N cells in grid = {len(grid['ra'])}") - return grid - - -def cut_grid( - grid, - remove_nan_density=True, - remove_empty_cells=False, - n_cut=None, - weight_min=None, - rcom_max=None, - xmax=None, - ymax=None, - zmax=None, - remove_origin=False, -): - """Apply selection cuts to a grid in-place. - - Args: - grid (dict): Grid dictionary to modify. - remove_nan_density (bool): Drop NaN density and errors. - remove_empty_cells (bool): Drop cells with zero counts. - n_cut (int, optional): Minimum objects per cell. - weight_min (float, optional): Minimum weight per cell. - rcom_max (float, optional): Radial cutoff. - xmax (float, optional): X cutoff. - ymax (float, optional): Y cutoff. - zmax (float, optional): Z cutoff. - remove_origin (bool): Remove the origin cell. - - Returns: - None: Modifies `grid` in-place. - """ - mask = np.full(grid["ra"].shape, True) - if n_cut is not None: - mask &= grid["nincell"] > n_cut - if weight_min is not None: - mask &= grid["sum_weights"] > weight_min - if rcom_max is not None: - mask &= grid["rcom_zobs"] < rcom_max - if xmax is not None: - mask &= np.abs(grid["x"]) < xmax - if ymax is not None: - mask &= np.abs(grid["y"]) < ymax - if zmax is not None: - mask &= np.abs(grid["z"]) < zmax - if remove_nan_density: - if "density" in grid: - mask &= ~(np.isnan(grid["density"])) - if "density_error" in grid: - mask &= ~(np.isnan(grid["density_error"])) - if remove_empty_cells: - if "N_in_cell" in grid: - mask &= ~(grid["N_in_cell"].astype(int) == 0) - if remove_origin: - mask = mask & ((grid["x"] != 0.0) | (grid["y"] != 0.0) | (grid["z"] != 0.0)) - for field in grid: - grid[field] = grid[field][mask] - - -def grid_data_density_pypower( - raobj, - decobj, - rcomobj, - rcom_max, - grid_size, - grid_type, - kind, - Nrandom=10, - random_method="cartesian", - interlacing=2, - compensate=False, - coord_randoms=None, - min_count_random=0, - overhead=20, -): - """Grid data with pypower and compute density contrast on a mesh. - - Args: - raobj (array-like): Right ascensions. - decobj (array-like): Declinations. - rcomobj (array-like): Comoving distances. - rcom_max (float): Outer cutoff for grid. - grid_size (float): Cell size. - grid_type (str): `rect` or `sphere` cut behavior. - kind (str): Resampler passed to CatalogMesh. - Nrandom (int): Randoms per data object. - random_method (str): Random generation method. - interlacing (int): Interlacing factor. - compensate (bool): Apply resampler compensation. - coord_randoms (tuple, optional): Randoms coordinates for `file` method. - min_count_random (int): Minimum random count for valid error. - overhead (float): Extra margin around cutoff. - - Returns: - dict: Grid with positions, density contrast, errors, and counts. - """ - xobj, yobj, zobj = utils.radec2cart(rcomobj, raobj, decobj) - mask = np.abs(xobj) < rcom_max + overhead - mask &= np.abs(yobj) < rcom_max + overhead - mask &= np.abs(zobj) < rcom_max + overhead - xobj, yobj, zobj = xobj[mask], yobj[mask], zobj[mask] - raobj, decobj, rcomobj = raobj[mask], decobj[mask], rcomobj[mask] - - ( - xobj_random, - yobj_random, - zobj_random, - ) = define_randoms( - random_method, - xobj, - yobj, - zobj, - raobj, - decobj, - rcomobj, - Nrandom=Nrandom, - coord_randoms=coord_randoms, - max_coordinates=rcom_max + overhead, - ) - - data_positions = np.array([xobj, yobj, zobj]).T - randoms_positions = np.array([xobj_random, yobj_random, zobj_random]).T - - data_weights = np.ones((data_positions.shape[0],)) - randoms_weights = np.ones((randoms_positions.shape[0],)) - - catalog_mesh = CatalogMesh( - data_positions=data_positions, - data_weights=data_weights, - randoms_positions=randoms_positions, - randoms_weights=randoms_weights, - interlacing=interlacing, - boxsize=2 * (rcom_max + overhead), - boxcenter=0.0, - cellsize=grid_size, - resampler=kind, - position_type="pos", - ) - - catalog_mesh_count = CatalogMesh( - data_positions=data_positions, - data_weights=data_weights, - randoms_positions=randoms_positions, - randoms_weights=randoms_weights, - interlacing=0, - boxsize=2 * (rcom_max + overhead), - boxcenter=0.0, - cellsize=grid_size, - resampler="ngp", - position_type="pos", - ) - - mesh_data = catalog_mesh.to_mesh(field="data", compensate=compensate) - mesh_randoms = catalog_mesh.to_mesh( - field="data-normalized_randoms", compensate=compensate - ) - - mesh_count_randoms = catalog_mesh_count.to_mesh(field="data-normalized_randoms") - - density_contrast = np.ravel(mesh_data.value / mesh_randoms.value - 1) - - count_randoms = np.ravel(mesh_count_randoms.value).astype(int) - density_contrast_err = np.full(count_randoms.shape, np.nan) - mask = count_randoms > min_count_random - density_contrast_err[mask] = np.sqrt(1 / (count_randoms[mask])) - - coord_mesh = np.array( - np.meshgrid( - np.sort(mesh_data.x[0][:, 0, 0]), - np.sort(mesh_data.x[1][0, :, 0]), - np.sort(mesh_data.x[2][0, 0, :]), - indexing="ij", - ) - ) - xgrid = np.ravel(coord_mesh[0, :, :, :]) + grid_size / 2 - ygrid = np.ravel(coord_mesh[1, :, :, :]) + grid_size / 2 - zgrid = np.ravel(coord_mesh[2, :, :, :]) + grid_size / 2 - - rcomgrid, ragrid, decgrid = utils.cart2radec(xgrid, ygrid, zgrid) - - grid = { - "x": xgrid, - "y": ygrid, - "z": zgrid, - "ra": ragrid, - "dec": decgrid, - "rcom_zobs": rcomgrid, - "density": density_contrast, - "density_error": density_contrast_err, - "count_random": count_randoms, - } - - if grid_type == "rect": - cut_grid( - grid, - remove_nan_density=True, - xmax=rcom_max, - ymax=rcom_max, - zmax=rcom_max, - remove_origin=True, - ) - - if grid_type == "sphere": - cut_grid( - grid, - remove_nan_density=True, - rcom_max=rcom_max, - remove_origin=True, - ) - - return grid - - -def grid_data_velocity_pypower( - raobj, - decobj, - rcomobj, - rcom_max, - variance, - velocity, - grid_size, - grid_type, - kind, - interlacing=0, - compensate=False, - overhead=20, -): - """Grid velocity catalog with pypower and compute weighted means and variance. - - Args: - raobj (array-like): Right ascensions. - decobj (array-like): Declinations. - rcomobj (array-like): Comoving distances. - rcom_max (float): Outer cutoff. - variance (array-like): Per-object variances. - velocity (array-like|None): Per-object velocities (optional for variance-only). - grid_size (float): Cell size. - grid_type (str): `rect` or `sphere` cut behavior. - kind (str): Resampler passed to CatalogMesh. - interlacing (int): Interlacing factor. - compensate (bool): Apply resampler compensation. - overhead (float): Extra margin around cutoff. - - Returns: - dict: Grid with positions, velocity (optional), variance, and counts. - """ - xobj, yobj, zobj = utils.radec2cart(rcomobj, raobj, decobj) - mask = np.abs(xobj) < rcom_max + overhead - mask &= np.abs(yobj) < rcom_max + overhead - mask &= np.abs(zobj) < rcom_max + overhead - xobj, yobj, zobj = xobj[mask], yobj[mask], zobj[mask] - raobj, decobj, rcomobj = raobj[mask], decobj[mask], rcomobj[mask] - - data_positions = np.array([xobj, yobj, zobj]).T - count_weights = np.ones((data_positions.shape[0],)) - - catalog_mesh_var = CatalogMesh( - data_positions=data_positions, - data_weights=variance, - interlacing=interlacing, - boxsize=2 * (rcom_max + overhead), - boxcenter=0.0, - cellsize=grid_size, - resampler=kind, - position_type="pos", - ) - catalog_mesh_count = CatalogMesh( - data_positions=data_positions, - data_weights=count_weights, - interlacing=interlacing, - boxsize=2 * (rcom_max + overhead), - boxcenter=0.0, - cellsize=grid_size, - resampler=kind, - position_type="pos", - ) - if velocity is not None: - weights_weighted_mean = velocity / variance - catalog_mesh_weighted_vel = CatalogMesh( - data_positions=data_positions, - data_weights=weights_weighted_mean, - interlacing=interlacing, - boxsize=2 * (rcom_max + overhead), - boxcenter=0.0, - cellsize=grid_size, - resampler=kind, - position_type="pos", - ) - catalog_mesh_invvar = CatalogMesh( - data_positions=data_positions, - data_weights=1 / variance, - interlacing=interlacing, - boxsize=2 * (rcom_max + overhead), - boxcenter=0.0, - cellsize=grid_size, - resampler=kind, - position_type="pos", - ) - mesh_var = catalog_mesh_var.to_mesh(field="data", compensate=compensate) - mesh_count = catalog_mesh_count.to_mesh(field="data") - N_in_cell = np.ravel(mesh_count.value) - variance_grid = np.ravel(mesh_var.value) / ( - N_in_cell**2 - ) # *N_in_cell/np.abs(N_in_cell) - if velocity is not None: - mesh_weighted_vel = catalog_mesh_weighted_vel.to_mesh( - field="data", compensate=compensate - ) - mesh_invvar = catalog_mesh_invvar.to_mesh(field="data", compensate=compensate) - velocity_grid = np.ravel(mesh_weighted_vel.value) / np.ravel(mesh_invvar.value) - else: - velocity_grid = None - - coord_mesh = np.array( - np.meshgrid( - np.sort(mesh_var.x[0][:, 0, 0]), - np.sort(mesh_var.x[1][0, :, 0]), - np.sort(mesh_var.x[2][0, 0, :]), - indexing="ij", - ) - ) - xgrid = np.ravel(coord_mesh[0, :, :, :]) + grid_size / 2 - ygrid = np.ravel(coord_mesh[1, :, :, :]) + grid_size / 2 - zgrid = np.ravel(coord_mesh[2, :, :, :]) + grid_size / 2 - rcomgrid, ragrid, decgrid = utils.cart2radec(xgrid, ygrid, zgrid) - if velocity is not None: - grid = { - "x": xgrid, - "y": ygrid, - "z": zgrid, - "ra": ragrid, - "dec": decgrid, - "rcom_zobs": rcomgrid, - "velocity": velocity_grid, - "velocity_variance": variance_grid, - "N_in_cell": N_in_cell, - } - else: - grid = { - "x": xgrid, - "y": ygrid, - "z": zgrid, - "ra": ragrid, - "dec": decgrid, - "rcom_zobs": rcomgrid, - "velocity_variance": variance_grid, - "N_in_cell": N_in_cell, - } - if grid_type == "rect": - cut_grid( - grid, - remove_nan_density=True, - remove_empty_cells=True, - xmax=rcom_max, - ymax=rcom_max, - zmax=rcom_max, - remove_origin=True, - ) - - if grid_type == "sphere": - cut_grid( - grid, - remove_nan_density=True, - remove_empty_cells=True, - rcom_max=rcom_max, - remove_origin=True, - ) - - return grid - - -# CR - First try unconnecting pypower - - -# def _get_compensation_window(resampler="cic", shotnoise=False): -# r""" -# Return the compensation function, which corrects for the particle-mesh assignment (resampler) kernel. - -# Taken from https://github.com/bccp/nbodykit/blob/master/nbodykit/source/mesh/catalog.py, -# following https://arxiv.org/abs/astro-ph/0409240. -# ("shotnoise" formula for pcs has been checked with WolframAlpha). - -# Parameters -# ---------- -# resampler : string, default='cic' -# Resampler used to assign particles to the mesh. -# Choices are ['ngp', 'cic', 'tcs', 'pcs']. - -# shotnoise : bool, default=False -# If ``False``, return expression for eq. 18 in https://arxiv.org/abs/astro-ph/0409240. -# This the correct choice when applying interlacing, as aliased images (:math:`\mathbf{n} \neq (0,0,0)`) are suppressed in eq. 17. -# If ``True``, return expression for eq. 19. - -# Returns -# ------- -# window : callable -# Window function, taking as input :math:`\pi k_{i} / k_{N} = k / c` -# where :math:`k_{N}` is the Nyquist wavenumber and :math:`c` is the cell size, -# for each :math:`x`, :math:`y`, :math:`z`, axis. -# """ -# resampler = resampler.lower() - -# if shotnoise: -# if resampler == "ngp": - -# def window(*x): -# return 1.0 - -# elif resampler == "cic": - -# def window(*x): -# toret = 1.0 -# for xi in x: -# toret = toret * (1 - 2.0 / 3 * np.sin(0.5 * xi) ** 2) ** 0.5 -# return toret - -# elif resampler == "tsc": - -# def window(*x): -# toret = 1.0 -# for xi in x: -# s = np.sin(0.5 * xi) ** 2 -# toret = toret * (1 - s + 2.0 / 15 * s**2) ** 0.5 -# return toret - -# elif resampler == "pcs": - -# def window(*x): -# toret = 1.0 -# for xi in x: -# s = np.sin(0.5 * xi) ** 2 -# toret = ( -# toret -# * ( -# 1 -# - 4.0 / 3.0 * s -# + 2.0 / 5.0 * s**2 -# - 4.0 / 315.0 * s**3 -# ) -# ** 0.5 -# ) -# return toret - -# else: -# p = {"ngp": 1, "cic": 2, "tsc": 3, "pcs": 4}[resampler] - -# def window(*x): -# toret = 1.0 -# for xi in x: -# toret = toret * np.sinc(0.5 / np.pi * xi) ** p -# return toret - -# return window - - -# def _get_resampler(resampler): -# conversions = {"ngp": "nnb", "cic": "cic", "tsc": "tsc", "pcs": "pcs"} -# resampler = conversions[resampler] -# from pmesh.window import FindResampler - -# return FindResampler(resampler) - - -# def paint( -# positions, weights, out, resampler, scaling=None, transform=None, offset=None -# ): -# from pmesh.pm import ParticleMesh - -# pm = ParticleMesh(BoxSize=2 * rcom_max, Nmesh=(10, 10, 10)) - -# positions = positions - offset -# factor = bool(interlacing) + 0.5 -# scalar_weights = weights is None - -# if scaling is not None: -# if scalar_weights: -# weights = scaling -# else: -# weights = weights * scaling - -# def paint_slab(sl): -# # Decompose positions such that they live in the same region as the mesh in the current process -# p = positions[sl] -# size = len(p) -# layout = pm.decompose(p, smoothing=factor * self.resampler.support) -# p = layout.exchange(p) -# w = weights if scalar_weights else layout.exchange(weights[sl]) -# # hold = True means no zeroing of out -# pm.paint( -# p, mass=w, resampler=resampler, transform=transform, hold=True, out=out -# ) -# return size - -# islab = 0 -# slab_npoints_max = int(1024 * 1024 * 4) -# slab_npoints = slab_npoints_max -# sizes = pm.comm.allgather(len(positions)) -# local_size_max = max(sizes) -# painted_size = 0 - -# while islab < local_size_max: -# sl = slice(islab, islab + slab_npoints) -# painted_size_slab = paint_slab(sl) -# painted_size += painted_size_slab -# islab += slab_npoints -# slab_npoints = min(slab_npoints_max, int(slab_npoints * 1.2)) - - -# def compensate(self, cfield): -# if self.mpicomm.rank == 0: -# self.log_info("Applying compensation {}.".format(self.compensation)) -# # Apply compensation window for particle-assignment scheme -# window = _get_compensation_window(**self.compensation) - -# cellsize = self.boxsize / self.nmesh -# for k, slab in zip(cfield.slabs.x, cfield.slabs): -# kc = tuple(ki * ci for ki, ci in zip(k, cellsize)) -# slab[...] /= window(*kc) - - -# def compute_density(positions, weights): -# out = pm.create(type="real", value=0.0) -# for p, w in zip(positions, weights): -# paint(p, *w, out) - -# if compensate: -# out = out.r2c() -# compensate(out) -# out = out.c2r() -# return out diff --git a/flip/data_vector/gw_vectors.py b/flip/data_vector/gw_vectors.py new file mode 100644 index 0000000..261a0fd --- /dev/null +++ b/flip/data_vector/gw_vectors.py @@ -0,0 +1,46 @@ +from . import mesh +from .basic import Dens + + +class GWDensMesh(Dens): + _kind = "density" + _needed_keys = ["density", "density_error"] + _free_par = [] + _number_dimension_observation_covariance = 1 + _parameters_observation_covariance = ["density"] + + def give_data_and_variance(self, *args): + """Return density data and diagonal variance from `density_error`. + + Returns: + tuple: (density, density_error^2). + """ + + if self._covariance_observation is not None: + return self._data["density"], self._covariance_observation + return self._data["density"], self._data["density_error"] ** 2 + + def __init__(self, data, covariance_observation=None): + super().__init__(data, covariance_observation=covariance_observation) + + @classmethod + def init_from_catalog( + cls, + data_position_sky, + data_position_sky_bandwidth, + rcom_max, + grid_size, + grid_type, + kind, + **kwargs, + ): + grid = mesh.grid_data_density_multivariate_kernel( + data_position_sky, + data_position_sky_bandwidth, + rcom_max, + grid_size, + grid_type, + kind, + **kwargs, + ) + return cls(grid) diff --git a/flip/data_vector/mesh.py b/flip/data_vector/mesh.py new file mode 100644 index 0000000..a58360f --- /dev/null +++ b/flip/data_vector/mesh.py @@ -0,0 +1,1160 @@ +try: + from pmesh.pm import ParticleMesh +except ImportError: + ParticleMesh = None +import numpy as np + +from flip import utils + +_grid_kind_avail = ["ngp", "ngp_errw", "cic", "tsc", "pcs"] + +_grid_order_dict = { + "ngp": 1, + "ngp_errw": 1, + "cic": 2, + "tsc": 3, + "pcs": 4, +} + + +def _compute_grid_window(grid_size, k, order, n): + """Numerically compute isotropic grid assignment window. + + Uses spherical averaging over directions to produce a 1D window for + resampler order (NGP/CIC/TSC/PCS) given grid size. + + Args: + grid_size (float): Grid cell size. + k (array-like): Wavenumbers at which to evaluate the window. + order (int): Assignment order (1:NGP, 2:CIC, 3:TSC, 4:PCS). + n (int): Number of angular samples for spherical averaging. + + Returns: + numpy.ndarray: Window values for each `k`. + """ + window = np.zeros_like(k) + theta = np.linspace(0, np.pi, n) + phi = np.linspace(0, 2 * np.pi, n) + kx = np.outer(np.sin(theta), np.cos(phi)) + ky = np.outer(np.sin(theta), np.sin(phi)) + kz = np.outer(np.cos(theta), np.ones(n)) + + # Forgotten in Howlett et al. formula + # we add spherical coordinate solid angle element + dthetaphi = np.outer(np.sin(theta), np.ones(phi.size)) + for i in range(k.size): + # the factor here has an extra np.pi because of the definition of np.sinc + fact = (k[i] * grid_size) / (2 * np.pi) + func = ( + np.sinc(fact * kx) * np.sinc(fact * ky) * np.sinc(fact * kz) + ) ** order * dthetaphi + win_theta = np.trapz(func, x=phi) + window[i] = np.trapz(win_theta, x=theta) + window *= 1 / (4 * np.pi) + return window + + +def compute_grid_window(grid_size, kh, kind="ngp", n=1000): + """Compute grid assignment window for a given resampler kind. + + Args: + grid_size (float): Grid cell size. + kh (array-like): Wavenumbers at which to evaluate the window. + kind (str): One of `ngp`, `ngp_errw`, `cic`, `tsc`, `pcs`. + n (int): Angular samples for spherical averaging. + + Returns: + numpy.ndarray|None: Window values, or None if `grid_size==0`. + """ + + if kind not in _grid_kind_avail: + allowed = ", ".join(_grid_kind_avail) + raise ValueError(f"INVALID GRID TYPE! Allowed kinds: {allowed}") + if grid_size == 0: + return None + return _compute_grid_window(grid_size, kh, _grid_order_dict[kind], n) + + +def ngp_weight(ds): + """Return nearest-grid-point assignment weights. + + Args: + ds (array-like): Distance to the cell center in cell-size units. + + Returns: + numpy.ndarray: NGP weights evaluated at `ds`. + """ + abs_ds = np.abs(ds) + w = 1.0 * (abs_ds < 1 / 2) + return w + + +def ngp_errw_weight(ds): + """Return NGP weights for the error-weighted assignment scheme. + + Args: + ds (array-like): Distance to the cell center in cell-size units. + + Returns: + numpy.ndarray: Assignment weights evaluated at `ds`. + """ + return ngp_weight(ds) + + +def cic_weight(ds): + """Return cloud-in-cell assignment weights. + + Args: + ds (array-like): Distance to the cell center in cell-size units. + + Returns: + numpy.ndarray: CIC weights evaluated at `ds`. + """ + abs_ds = np.abs(ds) + w = (1.0 - abs_ds) * (abs_ds < 1) + return w + + +def tsc_weight(ds): + """Return triangular-shaped-cloud assignment weights. + + Args: + ds (array-like): Distance to the cell center in cell-size units. + + Returns: + numpy.ndarray: TSC weights evaluated at `ds`. + """ + abs_ds = np.abs(ds) + w = (3 / 4 - ds**2) * (abs_ds < 1 / 2) + w += 1 / 2 * (3 / 2 - abs_ds) ** 2 * ((1 / 2 <= abs_ds) & (abs_ds < 3 / 2)) + return w + + +def pcs_weight(ds): + """Return piecewise-cubic-spline assignment weights. + + Args: + ds (array-like): Distance to the cell center in cell-size units. + + Returns: + numpy.ndarray: PCS weights evaluated at `ds`. + """ + abs_ds = np.abs(ds) + w = 1 / 6 * (4 - 6 * ds**2 + 3 * abs_ds**3) * (abs_ds < 1) + w += 1 / 6 * (2 - abs_ds) ** 3 * ((1 <= abs_ds) & (abs_ds < 2)) + return w + + +def _get_mesh_attrs( + boxsize, + cellsize, +): + """Derive mesh dimensions and box metadata from target scales. + + Args: + boxsize (float): Target box size along each Cartesian axis. + cellsize (float): Desired cell size along each Cartesian axis. + + Returns: + tuple[numpy.ndarray, numpy.ndarray, numpy.ndarray]: Mesh shape, + adjusted box size, and box center. + """ + boxcenter = 0.0 + boxsize = np.full((3,), boxsize, dtype="f8") + + cellsize = np.full((3,), cellsize, dtype="f8") + nmesh = boxsize / cellsize + nmesh = np.ceil(nmesh).astype("i8") + nmesh += nmesh % 2 # to make it even + boxsize = nmesh * cellsize # enforce exact cellsize + + nmesh = np.full((3,), nmesh, dtype="i4") + boxcenter = np.full((3,), boxcenter, dtype="f8") + return nmesh, boxsize, boxcenter + + +def define_randoms( + random_method, + xobj, + yobj, + zobj, + raobj, + decobj, + rcomobj, + Nrandom=None, + coord_randoms=None, + max_coordinates=None, + seed=None, +): + """Generate random Cartesian positions for density gridding. + + Supports uniform Cartesian sampling, random catalogs drawn from the + observed angular and radial distributions, or externally supplied randoms. + + Args: + random_method (str): `cartesian`, `choice`, `choice_redshift`, or `file`. + xobj (array-like): Data x positions. + yobj (array-like): Data y positions. + zobj (array-like): Data z positions. + raobj (array-like): Data right ascensions. + decobj (array-like): Data declinations. + rcomobj (array-like): Data comoving distances. + Nrandom (int): Number of randoms per object. + coord_randoms (tuple, optional): `(ra, dec, rcom)` for `file` method. + max_coordinates (float, optional): Coordinate cutoff for `file` method. + seed (int, optional): Random seed used for reproducible draws. + + Returns: + tuple[numpy.ndarray, numpy.ndarray, numpy.ndarray]: Random x, y, z positions. + """ + N = xobj.size + if seed is not None: + np.random.seed(seed) + + # Uniform in X,Y,Z + if random_method == "cartesian": + xobj_random = (np.max(xobj) - np.min(xobj)) * np.random.random( + size=Nrandom * N + ) + np.min(xobj) + yobj_random = (np.max(yobj) - np.min(yobj)) * np.random.random( + size=Nrandom * N + ) + np.min(yobj) + zobj_random = (np.max(zobj) - np.min(zobj)) * np.random.random( + size=Nrandom * N + ) + np.min(zobj) + + # Choice in the ra, dec, and redshift data coordinates to create the random. + elif random_method == "choice": + counts_ra, bins_ra = np.histogram(raobj, bins=1000) + ra_random = np.random.choice( + (bins_ra[1:] + bins_ra[:-1]) / 2, + p=counts_ra / float(counts_ra.sum()), + size=Nrandom * N, + ) + counts_dec, bins_dec = np.histogram(decobj, bins=1000) + dec_random = np.random.choice( + (bins_dec[1:] + bins_dec[:-1]) / 2, + p=counts_dec / float(counts_dec.sum()), + size=Nrandom * N, + ) + counts_rcom, bins_rcom = np.histogram(rcomobj, bins=1000) + rcom_random = np.random.choice( + (bins_rcom[1:] + bins_rcom[:-1]) / 2, + p=counts_rcom / float(counts_rcom.sum()), + size=Nrandom * N, + ) + + xobj_random, yobj_random, zobj_random = utils.radec2cart( + rcom_random, ra_random, dec_random + ) + + # Uniform in ra and dec. Choice in the redshift data coordinates. + elif random_method == "choice_redshift": + ra_random = (np.max(raobj) - np.min(raobj)) * np.random.random( + size=Nrandom * N + ) + np.min(raobj) + dec_random = (np.max(decobj) - np.min(decobj)) * np.random.random( + size=Nrandom * N + ) + np.min(decobj) + counts_rcom, bins_rcom = np.histogram(rcomobj, bins=1000) + rcom_random = np.random.choice( + (bins_rcom[1:] + bins_rcom[:-1]) / 2, + p=counts_rcom / float(counts_rcom.sum()), + size=Nrandom * N, + ) + + xobj_random, yobj_random, zobj_random = utils.radec2cart( + rcom_random, ra_random, dec_random + ) + + # From random file + elif random_method == "file": + ra_random, dec_random, rcom_random = ( + coord_randoms[0], + coord_randoms[1], + coord_randoms[2], + ) + + xobj_random, yobj_random, zobj_random = utils.radec2cart( + rcom_random, ra_random, dec_random + ) + if max_coordinates is not None: + mask_random = np.abs(xobj_random) < max_coordinates + mask_random &= np.abs(yobj_random) < max_coordinates + mask_random &= np.abs(zobj_random) < max_coordinates + xobj_random = xobj_random[mask_random] + yobj_random = yobj_random[mask_random] + zobj_random = zobj_random[mask_random] + + return xobj_random, yobj_random, zobj_random + + +def create_mesh( + positions, + boxsize, + cellsize, + assignement="ngp", + weights=None, + scaling=None, +): + """Paint weighted positions onto a `pmesh` mesh. + + Args: + positions (array-like): Cartesian positions with shape `(N, 3)`. + boxsize (float): Box size used to create the mesh. + cellsize (float): Target cell size. + assignement (str): Assignment scheme among `ngp`, `ngp_errw`, `cic`, + `tsc`, and `pcs`. + weights (array-like, optional): Per-object weights to paint on the mesh. + scaling (float, optional): Extra multiplicative factor applied to weights. + + Returns: + pmesh.pm.RealField: Painted real-space mesh. + + Raises: + ValueError: If positions or weights contain non-finite values. + """ + + conversions = { + "ngp": "nnb", + "ngp_errw": "nnb", + "cic": "cic", + "tsc": "tsc", + "pcs": "pcs", + } + resampler = conversions[assignement] + + nmesh, boxsize, boxcenter = _get_mesh_attrs( + boxsize, + cellsize, + ) + + pm = ParticleMesh( + BoxSize=boxsize, + Nmesh=nmesh, + dtype="f8", + ) + offset = boxcenter - boxsize / 2.0 + _slab_npoints_max = int(1024 * 1024 * 4) + + def paint(positions, weights, scaling, out, transform=None): + positions = positions - offset + factor = 0.5 + if not np.isfinite(positions).all(): + raise ValueError("Some positions are NaN/inf") + if not np.isfinite(weights).all(): + raise ValueError("Some weights are NaN/inf") + if scaling is not None: + weights = weights * scaling + + def paint_slab(sl): + p = positions[sl] + size = len(p) + layout = pm.decompose(p, smoothing=factor * _grid_order_dict[assignement]) + p = layout.exchange(p) + w = layout.exchange(weights[sl]) + pm.paint( + p, + mass=w, + resampler=resampler, + transform=transform, + hold=True, + out=out, + ) + return size + + islab = 0 + slab_npoints = _slab_npoints_max + while islab < len(positions): + sl = slice(islab, islab + slab_npoints) + _ = paint_slab(sl) + islab += slab_npoints + slab_npoints = min(_slab_npoints_max, int(slab_npoints * 1.2)) + + out = pm.create(type="real", value=0.0) + paint(positions, weights, scaling, out) + + return out + + +def cut_grid( + grid, + remove_nan_density=True, + remove_empty_cells=False, + n_cut=None, + weight_min=None, + rcom_max=None, + xmax=None, + ymax=None, + zmax=None, + remove_origin=False, +): + """Apply selection cuts to a grid in-place. + + Args: + grid (dict): Grid dictionary to modify. + remove_nan_density (bool): Drop NaN density and errors. + remove_empty_cells (bool): Drop cells with zero counts. + n_cut (int, optional): Minimum objects per cell. + weight_min (float, optional): Minimum weight per cell. + rcom_max (float, optional): Radial cutoff. + xmax (float, optional): X cutoff. + ymax (float, optional): Y cutoff. + zmax (float, optional): Z cutoff. + remove_origin (bool): Remove the origin cell. + + Returns: + None: Modifies `grid` in-place. + """ + mask = np.full(grid["ra"].shape, True) + if n_cut is not None: + mask &= grid["nincell"] > n_cut + if weight_min is not None: + mask &= grid["sum_weights"] > weight_min + if rcom_max is not None: + mask &= grid["rcom_zobs"] < rcom_max + if xmax is not None: + mask &= np.abs(grid["x"]) < xmax + if ymax is not None: + mask &= np.abs(grid["y"]) < ymax + if zmax is not None: + mask &= np.abs(grid["z"]) < zmax + if remove_nan_density: + if "density" in grid: + mask &= ~(np.isnan(grid["density"])) + if "density_error" in grid: + mask &= ~(np.isnan(grid["density_error"])) + if remove_empty_cells: + if "N_in_cell" in grid: + mask &= ~(grid["N_in_cell"].astype(int) == 0) + if remove_origin: + mask = mask & ((grid["x"] != 0.0) | (grid["y"] != 0.0) | (grid["z"] != 0.0)) + for field in grid: + grid[field] = grid[field][mask] + + +def cut_grid_type( + grid, + grid_type, + rcom_max, +): + if grid_type == "rect": + cut_grid( + grid, + remove_nan_density=True, + xmax=rcom_max, + ymax=rcom_max, + zmax=rcom_max, + remove_origin=True, + ) + + if grid_type == "sphere": + cut_grid( + grid, + remove_nan_density=True, + rcom_max=rcom_max, + remove_origin=True, + ) + + +def prepare_data_position( + data_position_sky, + rcom_max, + overhead, + random_method=None, + Nrandom=None, + coord_randoms=None, + seed=None, + data_position_sky_bandwidth=None, +): + """Convert sky coordinates to Cartesian positions and prepare randoms. + + Args: + data_position_sky (array-like): Sky coordinates with columns `(ra, dec, rcom)`. + rcom_max (float): Maximum comoving radius of the retained region. + overhead (float): Padding added to the mesh extent. + random_method (str, optional): Random generation method passed to + `define_randoms`. + Nrandom (int, optional): Number of random points per object. + coord_randoms (tuple, optional): Input random sky coordinates for the + `file` random method. + seed (int, optional): Random seed used when generating randoms. + data_position_sky_bandwidth (array-like, optional): Per-object bandwidth + matrices expressed in sky coordinates. + + Returns: + tuple: Cartesian data positions, transformed bandwidth matrices or `None`, + and random positions or `None`. + """ + + raobj = data_position_sky[:, 0] + decobj = data_position_sky[:, 1] + rcomobj = data_position_sky[:, 2] + xobj, yobj, zobj = utils.radec2cart(rcomobj, raobj, decobj) + mask = np.abs(xobj) < rcom_max + overhead + mask &= np.abs(yobj) < rcom_max + overhead + mask &= np.abs(zobj) < rcom_max + overhead + xobj, yobj, zobj = xobj[mask], yobj[mask], zobj[mask] + + if data_position_sky_bandwidth is not None: + + jacobian = utils.radec2cart_jacobian(rcomobj[mask], raobj[mask], decobj[mask]) + + data_position_sky_bandwidth = data_position_sky_bandwidth[mask, :, :] + + data_position_bandwith = ( + jacobian.T + @ data_position_sky_bandwidth + @ np.transpose(jacobian.T, (0, 2, 1)) + ) + else: + data_position_bandwith = None + + raobj, decobj, rcomobj = raobj[mask], decobj[mask], rcomobj[mask] + if random_method is not None: + ( + xobj_random, + yobj_random, + zobj_random, + ) = define_randoms( + random_method, + xobj, + yobj, + zobj, + raobj, + decobj, + rcomobj, + Nrandom=Nrandom, + coord_randoms=coord_randoms, + max_coordinates=rcom_max + overhead, + seed=seed, + ) + randoms_positions = np.array([xobj_random, yobj_random, zobj_random]).T + else: + randoms_positions = None + + data_positions = np.array([xobj, yobj, zobj]).T + + return data_positions, data_position_bandwith, randoms_positions + + +def define_grid_from_mesh(mesh_data, grid_size): + """Extract voxel-center coordinates and sky coordinates from a mesh. + + Args: + mesh_data: `pmesh` field used to define the grid geometry. + grid_size (float): Cell size used to shift coordinates to voxel centers. + + Returns: + tuple[numpy.ndarray, ...]: Cartesian and sky coordinates of all grid cells. + """ + coord_mesh = np.array( + np.meshgrid( + np.sort(mesh_data.x[0][:, 0, 0]), + np.sort(mesh_data.x[1][0, :, 0]), + np.sort(mesh_data.x[2][0, 0, :]), + indexing="ij", + ) + ) + xgrid = np.ravel(coord_mesh[0, :, :, :]) + grid_size / 2 + ygrid = np.ravel(coord_mesh[1, :, :, :]) + grid_size / 2 + zgrid = np.ravel(coord_mesh[2, :, :, :]) + grid_size / 2 + + rcomgrid, ragrid, decgrid = utils.cart2radec(xgrid, ygrid, zgrid) + + return xgrid, ygrid, zgrid, ragrid, decgrid, rcomgrid + + +def grid_data_density( + data_position_sky, + rcom_max, + grid_size, + grid_type, + kind, + Nrandom=10, + random_method="cartesian", + coord_randoms=None, + min_count_random=0, + overhead=20, + seed=None, +): + """Grid data with pypower and compute density contrast on a mesh. + + Args: + raobj (array-like): Right ascensions. + decobj (array-like): Declinations. + rcomobj (array-like): Comoving distances. + rcom_max (float): Outer cutoff for grid. + grid_size (float): Cell size. + grid_type (str): `rect` or `sphere` cut behavior. + kind (str): Resampler passed to CatalogMesh. + Nrandom (int): Randoms per data object. + random_method (str): Random generation method. + interlacing (int): Interlacing factor. + compensate (bool): Apply resampler compensation. + coord_randoms (tuple, optional): Randoms coordinates for `file` method. + min_count_random (int): Minimum random count for valid error. + overhead (float): Extra margin around cutoff. + + Returns: + dict: Grid with positions, density contrast, errors, and counts. + """ + + kind = kind.lower() + if kind not in _grid_kind_avail: + allowed = ", ".join(_grid_kind_avail) + raise ValueError(f"INVALID GRID TYPE! Allowed kinds: {allowed}") + + data_positions, _, randoms_positions = prepare_data_position( + data_position_sky, + rcom_max, + overhead, + random_method=random_method, + Nrandom=Nrandom, + coord_randoms=coord_randoms, + seed=seed, + ) + + data_weights = np.ones((data_positions.shape[0],)) + randoms_weights = np.ones((randoms_positions.shape[0],)) + + mesh_data = create_mesh( + data_positions, + 2 * (rcom_max + overhead), + grid_size, + assignement=kind, + weights=data_weights, + ) + + mesh_randoms = create_mesh( + randoms_positions, + 2 * (rcom_max + overhead), + grid_size, + assignement=kind, + weights=randoms_weights, + scaling=np.sum(data_weights) / np.sum(randoms_weights), + ) + if kind == "ngp": + mesh_count_randoms = mesh_randoms.copy() + else: + mesh_count_randoms = create_mesh( + randoms_positions, + 2 * (rcom_max + overhead), + grid_size, + assignement="ngp", + weights=randoms_weights, + scaling=np.sum(data_weights) / np.sum(randoms_weights), + ) + + density_contrast = np.zeros_like(mesh_data.value) + mask_randoms_nonzero = mesh_randoms.value != 0.0 + density_contrast[mask_randoms_nonzero] = np.ravel( + mesh_data.value[mask_randoms_nonzero] / mesh_randoms.value[mask_randoms_nonzero] + - 1 + ) + density_contrast[~mask_randoms_nonzero] = np.nan + density_contrast = np.ravel(density_contrast) + + count_randoms = np.ravel(mesh_count_randoms.value).astype(int) + density_contrast_err = np.full(count_randoms.shape, np.nan) + mask = count_randoms > min_count_random + density_contrast_err[mask] = np.sqrt(1 / (count_randoms[mask])) + + xgrid, ygrid, zgrid, ragrid, decgrid, rcomgrid = define_grid_from_mesh( + mesh_data, + grid_size, + ) + + grid = { + "x": xgrid, + "y": ygrid, + "z": zgrid, + "ra": ragrid, + "dec": decgrid, + "rcom_zobs": rcomgrid, + "density": density_contrast, + "density_error": density_contrast_err, + "count_random": count_randoms, + } + + cut_grid_type( + grid, + grid_type, + rcom_max, + ) + + return grid + + +def grid_data_density_multivariate_kernel( + data_position_sky, + data_position_sky_bandwidth, + rcom_max, + grid_size, + grid_type, + kind, + Nrandom=10, + random_method="cartesian", + coord_randoms=None, + min_count_random=0, + overhead=20, + seed=None, + kernel="gaussian", + cutoff_type=None, + threshold=1e-5, +): + """Grid density data after smoothing each object with a local kernel. + + Args: + data_position_sky (array-like): Sky coordinates with columns `(ra, dec, rcom)`. + data_position_sky_bandwidth (array-like): Per-object bandwidth matrices in + sky coordinates. + rcom_max (float): Outer cutoff for the final grid. + grid_size (float): Cell size of the mesh. + grid_type (str): Either `rect` or `sphere` for the final cut. + kind (str): Assignment kernel among the supported mesh schemes. + Nrandom (int): Number of random points per object. + random_method (str): Strategy used to generate random catalogs. + coord_randoms (tuple, optional): Input random sky coordinates for the + `file` random method. + min_count_random (int): Minimum random count required to define errors. + overhead (float): Padding added to the mesh extent. + seed (int, optional): Random seed used when generating randoms. + kernel (str): Name of the multivariate kernel. + cutoff_type (str, optional): Optional truncation criterion applied to the + kernel weights. + threshold (float): Threshold associated with `cutoff_type`. + + Returns: + dict: Gridded density field, density errors, and grid coordinates. + """ + + kind = kind.lower() + if kind not in _grid_kind_avail: + allowed = ", ".join(_grid_kind_avail) + raise ValueError(f"INVALID GRID TYPE! Allowed kinds: {allowed}") + + data_positions, data_position_bandwith, randoms_positions = prepare_data_position( + data_position_sky, + rcom_max, + overhead, + random_method=random_method, + Nrandom=Nrandom, + coord_randoms=coord_randoms, + seed=seed, + data_position_sky_bandwidth=data_position_sky_bandwidth, + ) + + data_weights = np.ones((data_positions.shape[0],)) + randoms_weights = np.ones((randoms_positions.shape[0],)) + + # First mesh for grid definition + mesh_data = create_mesh( + data_positions, + 2 * (rcom_max + overhead), + grid_size, + assignement=kind, + weights=data_weights, + ) + + xgrid, ygrid, zgrid, ragrid, decgrid, rcomgrid = define_grid_from_mesh( + mesh_data, + grid_size, + ) + + grid_positions = np.array([xgrid, ygrid, zgrid]).T + + data_kernel_positions = [] + data_kernel_weights = [] + for i in range(data_positions.shape[0]): + kernel_weights = multivariate_kernel_density_estimation( + data_positions[i], + data_position_bandwith[i], + grid_positions, + grid_size, + kernel=kernel, + cutoff_type=cutoff_type, + threshold=threshold, + ) + mask = kernel_weights != 0.0 + data_kernel_positions.append(grid_positions[mask]) + data_kernel_weights.append(kernel_weights[mask] / np.sum(kernel_weights[mask])) + + data_kernel_positions = np.concatenate(data_kernel_positions, axis=0) + data_kernel_weights = np.concatenate(data_kernel_weights, axis=0) + + data_kernel_weights = ( + data_positions.shape[0] * data_kernel_weights / np.sum(data_kernel_weights) + ) + + mesh_data_kernel = create_mesh( + data_kernel_positions, + 2 * (rcom_max + overhead), + grid_size, + assignement=kind, + weights=data_kernel_weights, + ) + + mesh_randoms = create_mesh( + randoms_positions, + 2 * (rcom_max + overhead), + grid_size, + assignement=kind, + weights=randoms_weights, + scaling=np.sum(data_weights) / np.sum(randoms_weights), + ) + if kind == "ngp": + mesh_count_randoms = mesh_randoms.copy() + else: + mesh_count_randoms = create_mesh( + randoms_positions, + 2 * (rcom_max + overhead), + grid_size, + assignement="ngp", + weights=randoms_weights, + scaling=np.sum(data_weights) / np.sum(randoms_weights), + ) + + density_contrast = np.zeros_like(mesh_data_kernel.value) + mask_randoms_nonzero = mesh_randoms.value != 0.0 + density_contrast[mask_randoms_nonzero] = np.ravel( + mesh_data_kernel.value[mask_randoms_nonzero] + / mesh_randoms.value[mask_randoms_nonzero] + - 1 + ) + density_contrast[~mask_randoms_nonzero] = np.nan + density_contrast = np.ravel(density_contrast) + + count_randoms = np.ravel(mesh_count_randoms.value).astype(int) + density_contrast_err = np.full(count_randoms.shape, np.nan) + mask = count_randoms > min_count_random + density_contrast_err[mask] = np.sqrt(1 / (count_randoms[mask])) + + grid = { + "x": xgrid, + "y": ygrid, + "z": zgrid, + "ra": ragrid, + "dec": decgrid, + "rcom_zobs": rcomgrid, + "density": density_contrast, + "density_error": density_contrast_err, + "count_random": count_randoms, + } + + cut_grid_type( + grid, + grid_type, + rcom_max, + ) + + return grid + + +def multivariate_kernel_density_estimation( + data_position, + bandwidth, + grid_positions, + grid_size, + kernel="gaussian", + cutoff_type=None, + threshold=1e-5, +): + """Evaluate multivariate kernel weights on a Cartesian grid. + + Args: + data_position (array-like): Cartesian position of the object center. + bandwidth (array-like): $3 \times 3$ covariance matrix of the kernel. + grid_positions (array-like): Cartesian positions of voxel centers. + grid_size (float): Cell size used to clip distances to voxel boundaries. + kernel (str): Kernel family to evaluate. + cutoff_type (str, optional): Optional truncation mode among + `kernel_unnormalized`, `kernel`, or `distance`. + threshold (float): Threshold used by `cutoff_type`. + + Returns: + numpy.ndarray: Kernel weights evaluated at each grid position. + + Raises: + ValueError: If the kernel or cutoff type is unsupported. + """ + + if kernel == "gaussian": + + distances_to_voxel = np.maximum( + np.abs(data_position - grid_positions) - grid_size / 2, 0 + ) + + if bandwidth[0, 0] == 0 and bandwidth[1, 1] == 0 and bandwidth[2, 2] == 0: + # null error case, assign all weight to the closest voxel + normalized_kernel = np.zeros((grid_positions.shape[0],)) + mask = np.sqrt(np.sum(distances_to_voxel**2, axis=1)) == 0 + normalized_kernel[mask] = 1.0 + else: + + kernel = np.exp( + -0.5 + * np.sum( + (distances_to_voxel @ np.linalg.inv(bandwidth)) + * distances_to_voxel, + axis=1, + ) + ) + norm = 1 / ((2 * np.pi) ** (1.5) * np.sqrt(np.linalg.det(bandwidth))) + normalized_kernel = kernel * norm + + else: + raise ValueError(f"Unsupported kernel: {kernel}") + + if cutoff_type is not None: + if cutoff_type == "kernel_unnormalized": + normalized_kernel[kernel < threshold] = 0.0 + elif cutoff_type == "kernel": + normalized_kernel[normalized_kernel < threshold] = 0.0 + elif cutoff_type == "distance": + distances = np.sqrt(np.sum(distances_to_voxel**2, axis=1)) + normalized_kernel[distances > threshold] = 0.0 + else: + raise ValueError(f"Unsupported cutoff type: {cutoff_type}") + + return normalized_kernel + + +def grid_data_density_kernel_sampling( + data_position_sky, + data_kernel_list, + rcom_max, + grid_size, + grid_type, + kind, + Nrandom=10, + random_method="cartesian", + coord_randoms=None, + min_count_random=0, + overhead=20, + seed=None, +): + """Grid data with pypower and compute density contrast on a mesh. + + Args: + raobj (array-like): Right ascensions. + decobj (array-like): Declinations. + rcomobj (array-like): Comoving distances. + rcom_max (float): Outer cutoff for grid. + grid_size (float): Cell size. + grid_type (str): `rect` or `sphere` cut behavior. + kind (str): Resampler passed to CatalogMesh. + Nrandom (int): Randoms per data object. + random_method (str): Random generation method. + interlacing (int): Interlacing factor. + compensate (bool): Apply resampler compensation. + coord_randoms (tuple, optional): Randoms coordinates for `file` method. + min_count_random (int): Minimum random count for valid error. + overhead (float): Extra margin around cutoff. + + Returns: + dict: Grid with positions, density contrast, errors, and counts. + """ + + kind = kind.lower() + if kind not in _grid_kind_avail: + allowed = ", ".join(_grid_kind_avail) + raise ValueError(f"INVALID GRID TYPE! Allowed kinds: {allowed}") + + data_positions, _, randoms_positions = prepare_data_position( + data_position_sky, + rcom_max, + overhead, + random_method=random_method, + Nrandom=Nrandom, + coord_randoms=coord_randoms, + seed=seed, + ) + + data_weights = np.ones((data_positions.shape[0],)) + randoms_weights = np.ones((randoms_positions.shape[0],)) + + mesh_data = create_mesh( + data_positions, + 2 * (rcom_max + overhead), + grid_size, + assignement=kind, + weights=data_weights, + ) + + mesh_randoms = create_mesh( + randoms_positions, + 2 * (rcom_max + overhead), + grid_size, + assignement=kind, + weights=randoms_weights, + scaling=np.sum(data_weights) / np.sum(randoms_weights), + ) + if kind == "ngp": + mesh_count_randoms = mesh_randoms.copy() + else: + mesh_count_randoms = create_mesh( + randoms_positions, + 2 * (rcom_max + overhead), + grid_size, + assignement="ngp", + weights=randoms_weights, + scaling=np.sum(data_weights) / np.sum(randoms_weights), + ) + + density_contrast = np.zeros_like(mesh_data.value) + mask_randoms_nonzero = mesh_randoms.value != 0.0 + density_contrast[mask_randoms_nonzero] = np.ravel( + mesh_data.value[mask_randoms_nonzero] / mesh_randoms.value[mask_randoms_nonzero] + - 1 + ) + density_contrast[~mask_randoms_nonzero] = np.nan + density_contrast = np.ravel(density_contrast) + + count_randoms = np.ravel(mesh_count_randoms.value).astype(int) + density_contrast_err = np.full(count_randoms.shape, np.nan) + mask = count_randoms > min_count_random + density_contrast_err[mask] = np.sqrt(1 / (count_randoms[mask])) + + xgrid, ygrid, zgrid, ragrid, decgrid, rcomgrid = define_grid_from_mesh( + mesh_data, + grid_size, + ) + + grid = { + "x": xgrid, + "y": ygrid, + "z": zgrid, + "ra": ragrid, + "dec": decgrid, + "rcom_zobs": rcomgrid, + "density": density_contrast, + "density_error": density_contrast_err, + "count_random": count_randoms, + } + + cut_grid_type( + grid, + grid_type, + rcom_max, + ) + + return grid + + +def grid_data_velocity( + data_position_sky, + rcom_max, + grid_size, + grid_type, + kind, + variance, + velocity=None, + overhead=20, +): + """Grid velocity data and compute per-cell means and errors. + + Args: + data_position_sky (array-like): Sky coordinates with columns `(ra, dec, rcom)`. + rcom_max (float): Outer cutoff. + grid_size (float): Cell size. + grid_type (str): `rect` or `sphere` cut behavior. + kind (str): Assignment kernel among the supported mesh schemes. + variance (array-like): Per-object variances. + velocity (array-like|None): Per-object velocities when computing the + weighted mean velocity. + overhead (float): Extra margin around cutoff. + + Returns: + dict: Grid with positions, velocity (optional), variance, and counts. + """ + + data_positions, _, _ = prepare_data_position( + data_position_sky, + rcom_max, + overhead, + ) + + count_weights = np.ones((data_positions.shape[0],)) + + mesh_variance = create_mesh( + data_positions, + 2 * (rcom_max + overhead), + grid_size, + assignement=kind, + weights=variance, + ) + + mesh_count = create_mesh( + data_positions, + 2 * (rcom_max + overhead), + grid_size, + assignement=kind, + weights=count_weights, + ) + + N_in_cell = np.ravel(mesh_count.value) + variance_grid = np.ravel(mesh_variance.value) / (N_in_cell**2) + + if velocity is not None: + weights_weighted_mean = velocity / variance + + mesh_weighted_velocity = create_mesh( + data_positions, + 2 * (rcom_max + overhead), + grid_size, + assignement=kind, + weights=weights_weighted_mean, + ) + + mesh_inverse_variance = create_mesh( + data_positions, + 2 * (rcom_max + overhead), + grid_size, + assignement=kind, + weights=1 / variance, + ) + velocity_grid = np.ravel(mesh_weighted_velocity.value) / np.ravel( + mesh_inverse_variance.value + ) + else: + velocity_grid = None + + xgrid, ygrid, zgrid, ragrid, decgrid, rcomgrid = define_grid_from_mesh( + mesh_variance, + grid_size, + ) + grid = { + "x": xgrid, + "y": ygrid, + "z": zgrid, + "ra": ragrid, + "dec": decgrid, + "rcom_zobs": rcomgrid, + "velocity_error": np.sqrt(variance_grid), + "N_in_cell": N_in_cell, + } + if velocity is not None: + grid["velocity"] = velocity_grid + + if grid_type == "rect": + cut_grid( + grid, + remove_nan_density=True, + remove_empty_cells=True, + xmax=rcom_max, + ymax=rcom_max, + zmax=rcom_max, + remove_origin=True, + ) + + if grid_type == "sphere": + cut_grid( + grid, + remove_nan_density=True, + remove_empty_cells=True, + rcom_max=rcom_max, + remove_origin=True, + ) + + return grid diff --git a/flip/data_vector/snia_vectors.py b/flip/data_vector/snia_vectors.py index a2bb448..3e4efe0 100644 --- a/flip/data_vector/snia_vectors.py +++ b/flip/data_vector/snia_vectors.py @@ -24,7 +24,7 @@ log = create_log() -class VelFromSALTfit(DataVector): +class VelTrippRelation(DataVector): _kind = "velocity" _needed_keys = ["zobs", "mb", "x1", "c", "rcom_zobs"] _free_par = ["alpha", "beta", "M_0", "sigma_M"] @@ -116,12 +116,30 @@ def compute_observed_distance_modulus(self, parameter_values_dict): Returns: ndarray: Distance modulus per object. """ + + if ( + "alpha_low" in parameter_values_dict + and "alpha_high" in parameter_values_dict + and "x1_treshold" in parameter_values_dict + ): + alpha = jnp.where( + self._data["x1"] < parameter_values_dict["x1_treshold"], + parameter_values_dict["alpha_low"], + parameter_values_dict["alpha_high"], + ) + else: + alpha = parameter_values_dict["alpha"] + observed_distance_modulus = ( self._data["mb"] - + parameter_values_dict["alpha"] * self._data["x1"] + + alpha * self._data["x1"] - parameter_values_dict["beta"] * self._data["c"] - parameter_values_dict["M_0"] ) + if "p" in self._data and "gamma" in parameter_values_dict: + observed_distance_modulus += ( + parameter_values_dict["gamma"] * self._data["p"] + ) if "host_logmass" in self.data: mask = self._data["host_logmass"] > 10 @@ -262,3 +280,7 @@ def give_data_and_variance(self, parameter_values_dict): ) return velocities, velocity_variance + + +# Placeholder for backward compatibility, to be removed in future versions +VelFromSALTfit = VelTrippRelation diff --git a/flip/utils.py b/flip/utils.py index 385d997..47aab35 100644 --- a/flip/utils.py +++ b/flip/utils.py @@ -5,8 +5,30 @@ import matplotlib.pyplot as plt import numpy as np +from ._config import __use_jax__ + +if __use_jax__: + try: + import jax.numpy as jnp + import jax.scipy as jsc + + jax_installed = True + + except ImportError: + import numpy as jnp + import scipy as jsc + + jax_installed = False +else: + import numpy as jnp + + jax_installed = False + + _C_LIGHT_KMS_ = acst.c.to("km/s").value +_available_priors = ["gaussian", "positive", "uniform"] + def Du(k, sigmau): """Damping function Du(k, sigma_u) = sin(k sigma_u) / (k sigma_u). @@ -38,6 +60,30 @@ def radec2cart(rcom, ra, dec): return x, y, z +def radec2cart_jacobian(rcom, ra, dec): + + if isinstance(rcom, np.ndarray): + null_component = np.zeros_like(rcom) + else: + null_component = 0 + + return np.array( + [ + [ + np.cos(ra) * np.cos(dec), + -rcom * np.sin(ra) * np.cos(dec), + -rcom * np.cos(ra) * np.sin(dec), + ], + [ + np.sin(ra) * np.cos(dec), + rcom * np.cos(ra) * np.cos(dec), + -rcom * np.sin(ra) * np.sin(dec), + ], + [np.sin(dec), null_component, rcom * np.cos(dec)], + ] + ) + + def cart2radec(x, y, z): """Convert Cartesian (x, y, z) to spherical (r, ra, dec). @@ -197,3 +243,148 @@ def add_array_statistics(arr, char): def close(): """Shut down the logging module cleanly.""" logging.shutdown() + + +def return_prior( + parameter_name, + prior_properties, +): + if prior_properties["type"].lower() not in _available_priors: + raise ValueError( + f"""The prior type {prior_properties["type"]} is not available""" + f"""Please choose between {_available_priors}""" + ) + elif prior_properties["type"].lower() == "gaussian": + prior = GaussianPrior( + parameter_name=parameter_name, + prior_mean=prior_properties["mean"], + prior_standard_deviation=prior_properties["standard_deviation"], + ) + elif prior_properties["type"].lower() == "positive": + prior = PositivePrior( + parameter_name=parameter_name, + ) + elif prior_properties["type"].lower() == "uniform": + prior = UniformPrior( + parameter_name=parameter_name, + range=prior_properties["range"], + ) + + return prior + + +class Prior: + """Base prior class encapsulating parameter-specific priors. + + Attributes: + parameter_name (str): Name of the parameter this prior applies to. + """ + + def __init__( + self, + parameter_name=None, + ): + self.parameter_name = parameter_name + + +class GaussianPrior(Prior): + """Univariate Gaussian prior on a parameter. + + Models $p(\theta) \propto \exp\{-\tfrac{1}{2}[(\theta-\mu)^2/\sigma^2]\}$. + """ + + def __init__( + self, + parameter_name=None, + prior_mean=None, + prior_standard_deviation=None, + ): + super().__init__(parameter_name=parameter_name) + self.prior_mean = prior_mean + self.prior_standard_deviation = prior_standard_deviation + + def __call__( + self, + parameter_values_dict, + ): + """Return Gaussian log-prior for the parameter value. + + Args: + parameter_values_dict (dict): Map of parameter names to values. + + Returns: + float: Log-prior value. + """ + return -0.5 * ( + np.log(2 * jnp.pi * self.prior_standard_deviation**2) + + (parameter_values_dict[self.parameter_name] - self.prior_mean) ** 2 + / self.prior_standard_deviation**2 + ) + + +class PositivePrior(Prior): + """Log-prior enforcing parameter positivity via Heaviside function. + + Returns `log(Heaviside(value))`, which is `0` for positive values and `-inf` otherwise. + """ + + def __init__( + self, + parameter_name=None, + ): + super().__init__(parameter_name=parameter_name) + + def __call__( + self, + parameter_values_dict, + ): + """Return log-prior that is zero for positive values, -inf otherwise. + + Args: + parameter_values_dict (dict): Map of parameter names to values. + + Returns: + float: Log-prior (0 or -inf). + """ + return jnp.log(jnp.heaviside(parameter_values_dict[self.parameter_name], 0)) + + +class UniformPrior(Prior): + """Uniform prior over a finite interval for a parameter. + + Uses `scipy.stats.uniform.logpdf` for the specified range. + """ + + def __init__(self, parameter_name=None, range=None): + super().__init__(parameter_name=parameter_name) + self.range = range + + def __call__( + self, + parameter_values_dict, + ): + """Return uniform log-prior over `[range[0], range[1]]`. + + Args: + parameter_values_dict (dict): Map of parameter names to values. + + Returns: + float: Log-prior value (constant inside range, -inf outside). + """ + value = parameter_values_dict[self.parameter_name] + return jsc.stats.uniform.logpdf( + value, loc=self.range[0], scale=self.range[1] - self.range[0] + ) + + +def prior_sum(priors, x): + """Sum multiple prior contributions. + + Args: + priors (list[Callable]): List of prior callables accepting a parameter dict. + x (dict): Parameter values dictionary. + + Returns: + float: Sum of all prior log-probabilities. + """ + return sum(prior(x) for prior in priors) diff --git a/notebook/flip_mesh.ipynb b/notebook/flip_mesh.ipynb new file mode 100644 index 0000000..29947c0 --- /dev/null +++ b/notebook/flip_mesh.ipynb @@ -0,0 +1,133 @@ +{ + "cells": [ + { + "cell_type": "code", + "execution_count": 1, + "id": "4d7edb85-f8c5-4ec8-8335-f5ff3dd73b75", + "metadata": {}, + "outputs": [ + { + "name": "stderr", + "output_type": "stream", + "text": [ + "[ 000000.00 ]: 03-19 07:04 root WARNING Install GPy to use the gpmatrix emulator\n" + ] + } + ], + "source": [ + "import numpy as np\n", + "import pandas as pd\n", + "import matplotlib.pyplot as plt\n", + "\n", + "from pathlib import Path\n", + "from flip import data_vector, __flip_dir_path__\n", + "\n", + "flip_base = Path(__flip_dir_path__)\n", + "data_path = flip_base / \"data\"\n", + "plt.style.use(data_path / \"style.mplstyle\")" + ] + }, + { + "cell_type": "code", + "execution_count": 2, + "id": "617decad-6faf-460d-afec-dd887aeee887", + "metadata": {}, + "outputs": [], + "source": [ + "# Density data - Galaxies already gridded\n", + "\n", + "grid = pd.read_parquet(data_path / \"data_density.parquet\")\n", + "DataDensity = data_vector.DensMesh(grid.to_dict(orient='list'))" + ] + }, + { + "cell_type": "code", + "execution_count": 3, + "id": "8f8ad654-1e37-4037-baf0-bb1b6adf5f4e", + "metadata": {}, + "outputs": [], + "source": [ + "# Density data - Gridding velocity positions\n", + "\n", + "rcom_max = 200\n", + "grid_size = 30\n", + "grid_type = \"rect\"\n", + "kind = \"ngp\"\n", + "\n", + "\n", + "data_velocity = pd.read_parquet(data_path / \"data_velocity.parquet\")\n", + "coordinates_velocity = np.array([data_velocity[\"ra\"], data_velocity[\"dec\"], data_velocity[\"rcom_zobs\"]])\n", + "\n", + "DataDensityVel = data_vector.basic.DensMesh.init_from_catalog(\n", + " coordinates_velocity.T,\n", + " rcom_max,\n", + " grid_size,\n", + " grid_type,\n", + " kind,\n", + " )" + ] + }, + { + "cell_type": "code", + "execution_count": 4, + "id": "28b8fb1d-f151-4ae3-b6f5-b951f018827a", + "metadata": {}, + "outputs": [], + "source": [ + "rcom_max = 200\n", + "grid_size = 30\n", + "grid_type = \"rect\"\n", + "kind = \"ngp\"\n", + "\n", + "\n", + "# Density data - Gridding velocity positions\n", + "data_velocity = pd.read_parquet(data_path / \"data_velocity.parquet\")\n", + "coordinates_velocity = np.array([data_velocity[\"ra\"], data_velocity[\"dec\"], data_velocity[\"rcom_zobs\"]])\n", + "error_ra = np.radians(0.0)\n", + "error_dec = np.radians(0.0)\n", + "error_rcom = 0.0\n", + "bandwidth = np.array([np.diag([error_ra, error_dec, error_rcom]) for i in range(len(coordinates_velocity[0]))])\n", + "\n", + "\n", + "\n", + "DataDensityVel = data_vector.gw_vectors.GWDensMesh.init_from_catalog(\n", + " coordinates_velocity.T,\n", + " bandwidth,\n", + " rcom_max,\n", + " grid_size,\n", + " grid_type,\n", + " kind,\n", + " )" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "89766841-c44c-4f29-bf9a-bd27b10e4a12", + "metadata": {}, + "outputs": [], + "source": [] + } + ], + "metadata": { + "kernelspec": { + "display_name": "corentin2", + "language": "python", + "name": "corentin2" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.12.9" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +}