lisaanalysistools 1.0.0__cp312-cp312-macosx_10_9_x86_64.whl

lisatools/utils/pointeradjust.py

@@ -0,0 +1,106 @@
+# check to see if cupy is available for gpus
+import numpy as np
+
+try:
+    import cupy as xp
+
+    gpu = True
+
+except (ImportError, ModuleNotFoundError) as e:
+    import numpy as np
+
+    gpu = False
+
+
+def wrapper(*args, **kwargs):
+    """Function to convert array and C/C++ class arguments to ptrs
+
+    This function checks the object type. If it is a cupy or numpy array,
+    it will determine its pointer by calling the proper attributes. If you design
+    a Cython class to be passed through python, it must have a :code:`ptr`
+    attribute.
+
+    If you use this function, you must convert input arrays to size_t data type in Cython and
+    then properly cast the pointer as it enters the c++ function. See the
+    Cython codes
+    `here <https://github.com/BlackHolePerturbationToolkit/FastEMRIWaveforms/tree/master/src>`_
+    for examples.
+
+    args:
+        *args (list): list of the arguments for a function.
+        **kwargs (dict): dictionary of keyword arguments to be converted.
+
+    returns:
+        Tuple: (targs, tkwargs) where t indicates target (with pointer values
+            rather than python objects).
+
+    """
+    # declare target containers
+    targs = []
+    tkwargs = {}
+
+    # args first
+    for arg in args:
+        if gpu:
+            # cupy arrays
+            if isinstance(arg, xp.ndarray):
+                targs.append(arg.data.mem.ptr)
+                continue
+
+        # numpy arrays
+        if isinstance(arg, np.ndarray):
+            targs.append(arg.__array_interface__["data"][0])
+            continue
+
+        try:
+            # cython classes
+            targs.append(arg.ptr)
+            continue
+        except AttributeError:
+            # regular argument
+            targs.append(arg)
+
+    # kwargs next
+    for key, arg in kwargs.items():
+        if gpu:
+            # cupy arrays
+            if isinstance(arg, xp.ndarray):
+                tkwargs[key] = arg.data.mem.ptr
+                continue
+
+        if isinstance(arg, np.ndarray):
+            # numpy arrays
+            tkwargs[key] = arg.__array_interface__["data"][0]
+            continue
+
+        try:
+            # cython classes
+            tkwargs[key] = arg.ptr
+            continue
+        except AttributeError:
+            # other arguments
+            tkwargs[key] = arg
+
+    return (targs, tkwargs)
+
+
+def pointer_adjust(func):
+    """Decorator function for cupy/numpy agnostic cython
+
+    This decorator applies :func:`few.utils.utility.wrapper` to functions
+    via the decorator construction.
+
+    If you use this decorator, you must convert input arrays to size_t data type in Cython and
+    then properly cast the pointer as it enters the c++ function. See the
+    Cython codes
+    `here <https://github.com/BlackHolePerturbationToolkit/FastEMRIWaveforms/tree/master/src>`_
+    for examples.
+
+    """
+
+    def func_wrapper(*args, **kwargs):
+        # get pointers
+        targs, tkwargs = wrapper(*args, **kwargs)
+        return func(*targs, **tkwargs)
+
+    return func_wrapper

lisatools/utils/utility.py

@@ -0,0 +1,240 @@
+from multiprocessing.sharedctypes import Value
+from types import ModuleType, NoneType
+from typing import Tuple
+import numpy as np
+
+# from ..sensitivity import get_sensitivity
+
+try:
+    import cupy as cp
+
+except (ModuleNotFoundError, ImportError):
+    import numpy as cp
+
+    pass
+
+
+def get_array_module(arr: np.ndarray | cp.ndarray) -> ModuleType:
+    """Return array library of an array (np/cp).
+
+    Args:
+        arr: Numpy or Cupy array.
+
+    """
+    if isinstance(arr, np.ndarray):
+        return np
+    elif isinstance(arr, cp.ndarray):
+        return cp
+    else:
+        raise ValueError("arr must be a numpy or cupy array.")
+
+
+def generate_noise_fd(freqs, df, *sensitivity_args, func=None, **sensitivity_kwargs):
+    if func is None:
+        func = get_sensitivity
+
+    norm = 0.5 * (1.0 / df) ** 0.5
+    psd = func(freqs, *sensitivity_args, **sensitivity_kwargs)
+    noise_to_add = psd ** (1 / 2) * (
+        np.random.normal(0, norm, len(freqs))
+        + 1j * np.random.normal(0, norm, len(freqs))
+    )
+    return noise_to_add
+
+
+def AET(
+    X: float | np.ndarray, Y: float | np.ndarray, Z: float | np.ndarray
+) -> Tuple[float | np.ndarray, float | np.ndarray, float | np.ndarray]:
+    """Transform to AET from XYZ
+
+    .. math::
+
+        A = (Z - X) / \\sqrt(2)
+
+    .. math::
+
+        E = (X - 2Y + Z) / \\sqrt(6)
+
+    .. math::
+
+        T = (X + Y + Z) / \\sqrt(3)
+
+    Args:
+        X: X-channel information.
+        Y: Y-channel information.
+        Z: Z-channel information.
+
+    Returns:
+        A, E, T Channels.
+
+    """
+    return (
+        (Z - X) / np.sqrt(2.0),
+        (X - 2.0 * Y + Z) / np.sqrt(6.0),
+        (X + Y + Z) / np.sqrt(3.0),
+    )
+
+
+def searchsorted2d_vec(a, b, xp=None, gpu=None, **kwargs):
+    if xp is None:
+        xp = np
+    else:
+        try:
+            xp.cuda.runtime.setDevice(gpu)
+        except AttributeError:
+            pass
+
+    m, n = a.shape
+    max_num = xp.maximum(a.max() - a.min(), b.max() - b.min()) + 1
+    r = max_num * xp.arange(a.shape[0])[:, None]
+    p = xp.searchsorted((a + r).ravel(), (b + r).ravel(), **kwargs).reshape(m, -1)
+
+    out = p - n * (xp.arange(m)[:, None])
+    try:
+        xp.cuda.runtime.deviceSynchronize()
+    except AttributeError:
+        pass
+
+    return out
+
+
+def get_groups_from_band_structure(
+    f0, band_edges, f0_2=None, xp=None, num_groups_base=3, fix_f_test=None, inds=None
+):
+    if num_groups_base not in [2, 3, 4]:
+        raise ValueError("num_groups_base must be 2 or 3 or 4.")
+    if xp is None:
+        xp = np
+
+    else:
+        try:
+            xp.cuda.runtime.setDevice(xp.cuda.runtime.getDevice())
+
+        except AttributeError:
+            # it is numpy
+            pass
+
+    if not isinstance(f0, xp.ndarray) or not isinstance(band_edges, xp.ndarray):
+        raise TypeError(
+            "f0 and band_edges must be xp.ndarray with xp as numpy or cupy as given by the xp kwarg."
+        )
+
+    shape = f0.shape
+
+    # remove any above or below bands
+    bad = (f0 < band_edges.min()) | (f0 > band_edges.max())
+
+    band_indices = xp.searchsorted(band_edges, f0.flatten()).reshape(shape) - 1
+
+    # sort the bands in, but keep places with inds_band_indices
+    band_indices_sorted = xp.sort(band_indices, axis=-1)
+    inds_band_indices = xp.argsort(band_indices, axis=-1)
+
+    if f0_2 is not None:
+        assert f0_2.shape == f0.shape
+        band_indices_2 = xp.searchsorted(band_edges, f0_2.flatten()).reshape(shape) - 1
+        band_indices_2_sorted = xp.take_along_axis(
+            band_indices_2, inds_band_indices, axis=-1
+        )
+
+        # very important: ensures the proposed new point is not further than 1 band away.
+        diff = 1 if num_groups_base > 2 else 0
+        keep = (
+            np.abs(band_indices_2_sorted.flatten() - band_indices_sorted.flatten())
+            <= diff
+        )
+        if fix_f_test is not None:
+            keep[fix_f_test.flatten()] = False
+        remove = ~keep
+
+    else:
+        keep = np.ones(np.prod(band_indices_sorted.shape), dtype=bool)
+
+    # temperature index associated with each band
+    temp_inds = xp.repeat(
+        xp.arange(band_indices_sorted.shape[0]), np.prod(band_indices_sorted.shape[1:])
+    )[
+        keep
+    ]  # .reshape(shape)
+
+    # walker index associated with each band
+    walker_inds = (
+        xp.tile(
+            xp.arange(band_indices_sorted.shape[1]),
+            (band_indices_sorted.shape[0], band_indices_sorted.shape[2], 1),
+        )
+        .transpose((0, 2, 1))
+        .flatten()[keep]
+    )
+
+    if f0_2 is not None:
+        temp_inds_remove = xp.repeat(
+            xp.arange(band_indices_sorted.shape[0]),
+            np.prod(band_indices_sorted.shape[1:]),
+        )[
+            remove
+        ]  # .reshape(shape)
+
+        # walker index associated with each band
+        walker_inds_remove = (
+            xp.tile(
+                xp.arange(band_indices_sorted.shape[1]),
+                (band_indices_sorted.shape[0], band_indices_sorted.shape[2], 1),
+            )
+            .transpose((0, 2, 1))
+            .flatten()[remove]
+        )
+        inds_band_indices_remove = inds_band_indices.flatten()[remove]
+
+    # special indexing method
+    band_indices_sorted_special = (
+        band_indices_sorted.flatten()[keep]
+        + int(1e12) * temp_inds
+        + int(1e6) * walker_inds
+    )
+
+    # get the unique special indicators
+    (
+        unique_special,
+        unique_special_start_inds,
+        unique_special_reverse,
+        unique_special_counts,
+    ) = np.unique(
+        band_indices_sorted_special,
+        return_index=True,
+        return_inverse=True,
+        return_counts=True,
+    )
+
+    # this basically makes mini arange setups for each band
+    # the added_contribution for the first unique band index is removed
+    added_contribution = xp.arange(band_indices_sorted_special.shape[0])
+
+    # gets the groups
+    combined = band_indices_sorted_special + added_contribution
+    groups = (
+        combined - (combined[unique_special_start_inds])[unique_special_reverse]
+    )  # .reshape(shape)
+
+    groups_even_odd_tmp = xp.asarray(
+        [
+            (num_groups_base * groups + i)
+            * (band_indices_sorted.flatten()[keep] % num_groups_base == i)
+            for i in range(num_groups_base)
+        ]
+    )
+    groups_even_odd = xp.sum(groups_even_odd_tmp, axis=0)
+
+    groups_out = -2 * xp.ones_like(f0, dtype=int)
+    groups_out[
+        (temp_inds, walker_inds, inds_band_indices.flatten()[keep])
+    ] = groups_even_odd
+
+    groups_out[bad] = -1
+
+    """if f0_2 is not None and not np.all(keep):
+        fix = (temp_inds_remove, walker_inds_remove, inds_band_indices_remove)
+        fix_2 = band_indices_2[fix]
+        fix_1 = band_indices[fix]"""
+
+    return groups_out