"""Module for utilities.
.. currentmodule:: typed_lisa_toolkit.utils
Functions
---------
.. autofunction:: get_subset_slice
.. autofunction:: get_subset_mask
.. autofunction:: get_support_slice
.. autofunction:: promote_slice
.. autofunction:: extend_to
Decorators
----------
.. autofunction:: warn_external
.. autofunction:: deprecated
"""
from __future__ import annotations
import contextlib
import functools
import inspect
import logging
import os
import warnings
from collections.abc import Callable
from typing import Any
import array_api_compat as xpc
from .types.misc import AnyArray, AnyAxis
log = logging.getLogger(__name__)
[docs]
def warn_external(
message: str,
category: type[Warning],
*,
package_prefix: str = "typed_lisa_toolkit",
) -> None:
"""Emit a warning pointing to the first caller outside this package."""
frame = inspect.currentframe()
if frame is None:
warnings.warn(message, category, stacklevel=2)
return
internal_depth = 0
frame = frame.f_back
while frame is not None:
module_name = str(frame.f_globals.get("__name__", ""))
if not module_name.startswith(package_prefix):
break
internal_depth += 1
frame = frame.f_back
warnings.warn(message, category, stacklevel=internal_depth + 2)
def _deprecated(
message: str,
category: type[Warning] = DeprecationWarning,
*,
package_prefix: str = "typed_lisa_toolkit",
) -> Callable[[Callable[..., Any]], Callable[..., Any]]:
"""Return a decorator that emits a caller-facing deprecation warning."""
def _decorator(func: Callable[..., Any]) -> Callable[..., Any]:
@functools.wraps(func)
def _wrapped(*args: Any, **kwargs: Any):
warn_external(
message,
category=category,
package_prefix=package_prefix,
)
return func(*args, **kwargs)
return _wrapped
return _decorator
[docs]
def deprecated(
name: str,
role: str,
eol: str,
alternative: str | None = None,
*,
package_prefix: str = "typed_lisa_toolkit",
):
"""Return a decorator that emits a caller-facing deprecation warning with a standard message format.""" # noqa: E501
alt_msg = f" Use '{package_prefix}.{alternative}' instead." if alternative else ""
message = (
f"The {role} '{name}' is deprecated and will be removed in {eol}. {alt_msg}"
)
return _deprecated(message, package_prefix=package_prefix)
[docs]
def get_subset_slice(increasing_array: AnyArray, /, min: float, max: float):
"""Return the index slice for the subset [min, max] of the increasing array.
Examples
--------
>>> get_subset_slice(0.5 * np.arange(10), 1.0, 3.0)
slice(2, 7, None)
"""
xp = xpc.get_namespace(increasing_array)
start_idx = xp.searchsorted(increasing_array, min, side="left")
end_idx = xp.searchsorted(increasing_array, max, side="right")
return slice(start_idx, end_idx)
[docs]
def get_subset_mask[ArrayT: AnyArray](
array: ArrayT, /, min: float, max: float
) -> ArrayT:
"""Return a boolean mask for the subset [min, max] of the array.
Examples
--------
>>> get_subset_mask(0.5 * np.arange(10), 1.0, 3.0)
array([False, False, True, True, True, True, True, False, False, False])
"""
xp = xpc.get_namespace(array)
return xp.logical_and(array >= min, array <= max)
[docs]
def get_support_slice(array: AnyArray):
"""Return the index slice for the support of the array.
If the array is all zeros, the slice returned is empty.
Examples
--------
>>> get_support_slice(np.array([0, 0, 1, 2, 0, 0]))
slice(2, 4, None)
"""
xp = xpc.get_namespace(array)
non_zero_indices = xp.nonzero(xp.ravel(array))[0]
if len(non_zero_indices) == 0:
return slice(0, 0)
return slice(non_zero_indices[0], non_zero_indices[-1] + 1)
[docs]
def extend_to[AT: AnyAxis](
target_grid: tuple[AT, ...] | AT,
*,
known_slices: tuple[slice, ...] | None = None,
):
"""Return a function that extends the entries to the target grid.
The returned function has the signature:
.. code-block:: python
def get_extension(grid: tuple[AnyAxis, ...], entries: AnyArray) -> AnyArray:
...
The function extends the entries to the target grid by setting the entries
outside the input grid to zero. Both the input grid and the target grid are
assumed to be increasing, and the input grid is assumed to be a connected
subset of the target grid. Both grids should be tuples of arrays.
The entries are assumed to have canonical shape:
``(n_batches, n_channels, n_harmonics, n_features, *grid_dims)``
Arguments
---------
target_grid: The grid to which the entries should be extended. Can be a tuple
of arrays or a single array (for 1D grid).
known_slices: Optional tuple of slices for
the support of the input grid in the target grid.
If provided, it will be used directly instead of computing the
support slices from the input grid.
Examples
--------
>>> grid = (2 + np.arange(5),)
>>> entries = np.array([1, 2, 3, 4, 5]).reshape(1, 1, 1, 1, 5)
>>> target_grid = (np.arange(10),)
>>> result = extend_to(target_grid)(grid, entries)
>>> result.shape
(1, 1, 1, 1, 10)
"""
# One might consider pre-allocate memory for extended_entries, i.e. designing
# a class that stores extended_entries and reinitializes it with zeros before
# each call to get_extension. This would accelerate the function, but I fail
# to guarantee the correctness of the implementation, probably due to some
# issue with np.memmap.fill.
_target_grid = target_grid if isinstance(target_grid, tuple) else (target_grid,)
def get_extension(
grid: tuple[AnyAxis, ...] | AnyAxis,
entries: AnyArray,
) -> AnyArray:
_grid = grid if isinstance(grid, tuple) else (grid,)
xp = xpc.get_namespace(entries)
# Handle canonical shape:
# (n_batches, n_channels, n_harmonics, n_features, *grid_dims)
# Create extended array with target_grid lengths in the grid dimensions
extended_shape = entries.shape[:4] + tuple(len(g) for g in _target_grid)
extended_entries = xp.zeros(extended_shape, dtype=entries.dtype)
if known_slices is not None:
support_slices = known_slices
else:
support_slices = tuple(
get_subset_slice(xp.asarray(target_g), g[0], g[-1])
for target_g, g in zip(_target_grid, _grid, strict=True)
)
# Build full indexing tuple for canonical shape
index = promote_slice(support_slices)
try:
extended_entries[index] = entries
except TypeError:
# Fall back to functional update (JAX immutable arrays)
extended_entries = extended_entries.at[index].set(entries)
return extended_entries
return get_extension
@contextlib.contextmanager
def set_env(**environ: Any):
"""Temporarily set the process environment variables.
Examples
--------
>>> with set_env(TEST_VAR="123"):
... print(os.environ.get("TEST_VAR"))
123
>>> print(os.environ.get("TEST_VAR")) # Should be None or not set again
None
"""
old_environ = dict(os.environ)
os.environ.update(environ)
try:
yield
finally:
os.environ.clear()
os.environ.update(old_environ)
# def trim_interp(interpolator: Interpolator):
# """Return decorated interpolator.
# The decorated interpolate function will first trim the input
# array to its support before calling the original interpolator.
# Outside the support, the interpolated values are set to zero.
# """
# @functools.wraps(interpolator)
# def _interpolator(
# grid: AnyArray,
# entries: AnyArray,
# ) -> ArrayFunc:
# support_slice = get_support_slice(entries)
# xp = xpc.get_namespace(entries)
# trimmed_grid = grid[support_slice]
# if trimmed_grid.size == 0:
# log.warning("Empty array. Returning a zero interpolator.")
# return xp.zeros_like
# min = float(trimmed_grid[0])
# max = float(trimmed_grid[-1])
# def _interpolated(target_grid: AnyArray) -> AnyArray:
# """Return the interpolated entries at the target grid, zero outside support.""" # noqa: E501
# target_support_slice = get_subset_slice(target_grid, min, max)
# interp_grid = target_grid[target_support_slice]
# interpolated = interpolator(grid[support_slice], entries[support_slice])(
# interp_grid,
# )
# return extend_to(target_grid)(interp_grid, interpolated)
# return _interpolated
# return _interpolator