Source code for structure_factor.utils

"""Collection of secondary functions used in the principal modules."""

import numpy as np
from scipy import stats
from scipy.special import j0, j1, jn_zeros, jv, y0, y1, yv


[docs]def get_random_number_generator(seed=None): """Turn seed into a np.random.Generator instance.""" return np.random.default_rng(seed)
[docs]def set_nan_inf_to_zero(array, nan=0, posinf=0, neginf=0): """Set nan, posinf, and neginf values of ``array`` to the corresponding input arguments. Defaults to zero.""" return np.nan_to_num(array, nan=nan, posinf=posinf, neginf=neginf)
[docs]def bessel1(order, x): """Evaluate `Bessel function of the first kind <https://en.wikipedia.org/wiki/Bessel_function>`_.""" if order == 0: return j0(x) if order == 1: return j1(x) return jv(order, x)
[docs]def bessel1_zeros(order, nb_zeros): """Evaluate zeros of the `Bessel function of the first kind <https://en.wikipedia.org/wiki/Bessel_function>`_.""" return jn_zeros(order, nb_zeros)
[docs]def bessel2(order, x): """Evaluate `Bessel function of the second kind <https://en.wikipedia.org/wiki/Bessel_function>`_.""" if order == 0: return y0(x) if order == 1: return y1(x) return yv(order, x)
def norm(k): return np.linalg.norm(k, axis=-1)
[docs]def meshgrid_to_column_matrix(X): r"""Transform output ``X`` of numpy.meshgrid to a 2d numpy array with columns formed by flattened versions of the elements of ``X``. .. code-block:: python np.column_stack([x.ravel() for x in X]) Args: X (list): output of numpy.meshgrid. Returns: np.ndarray: 2d array. """ return np.column_stack([x.ravel() for x in X])
[docs]def sort_by_keys(keys, *arrays, **argsort_params): """Return a sorted version of ``arrays`` according to the indices that would sort ``keys`` by calling ``numpy.argsort(keys, **argsort_params)``. Args: keys (array_like): Array to extract the sorting indices from. arrays (array_like): Sequence of arrays to be sorted. These arrays must have a length larger or equal to the length of keys. Returns: list: sorted version of ``arrays``. """ idx = np.argsort(keys, **argsort_params) return [arr[idx] for arr in arrays]
[docs]def _sort_vectors(k, x_k, y_k=None): """Sort ``k`` by increasing order and rearranging the associated vectors to ``k``, ``x_k``and ``y_k``. Args: k (numpy.ndarray): Vector to be sorted by increasing order. x_k (numpy.ndarray): Vector of evaluations associated with ``k``. y_k (numpy.ndarray, optional): Vector of evaluations associated with ``k``. Defaults to None. Returns: (numpy.ndarray, numpy.ndarray, numpy.ndarray): ``k`` sorted by increasing order and the associated vectors ``x_k``and ``y_k``. """ if y_k is None: return (*sort_by_keys(k, k, x_k), y_k) return sort_by_keys(k, k, x_k, y_k)
[docs]def _bin_statistics(x, y, **params): """Divide ``x`` into bins and evaluate the mean and the standard deviation of the corresponding elements of ``y`` over each bin. Args: x (numpy.ndarray): Vector of data. y (numpy.ndarray): Vector of data associated with the vector ``x``. Keyword Args: params (dict): Keyword arguments (except ``"x"``, ``"values"`` and ``"statistic"``) of `scipy.stats.binned_statistic <https://docs.scipy.org/doc/scipy/reference/generated/scipy.stats.binned_statistic.html>`_. Returns: tuple(numpy.ndarray, numpy.ndarray, numpy.ndarray): - ``bin_centers``: Vector of centers of the bins associated to ``x``. - ``bin_mean``: Vector of means of ``y`` over the bins. - ``std_mean``: Vector of standard deviations of ``y`` over the bins. """ bin_mean, bin_edges, _ = stats.binned_statistic(x, y, statistic="mean", **params) bin_centers = np.convolve(bin_edges, np.ones(2), "valid") bin_centers /= 2 count, _, _ = stats.binned_statistic(x, y, statistic="count", **params) bin_std, _, _ = stats.binned_statistic(x, y, statistic="std", **params) bin_std /= np.sqrt(count) return bin_centers, bin_mean, bin_std