dycw-utilities 0.109.15__tar.gz → 0.109.16__tar.gz

{dycw_utilities-0.109.15 → dycw_utilities-0.109.16}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: dycw-utilities
-Version: 0.109.15
+Version: 0.109.16
 Author-email: Derek Wan <d.wan@icloud.com>
 License-File: LICENSE
 Requires-Python: >=3.12

{dycw_utilities-0.109.15 → dycw_utilities-0.109.16}/pyproject.toml

@@ -92,7 +92,7 @@ dependencies = [
 name = "dycw-utilities"
 readme = "README.md"
 requires-python = ">= 3.12"
-version = "0.109.15"
+version = "0.109.16"
 
 [project.optional-dependencies]
 test = [
@@ -335,7 +335,7 @@ zzz-test-zoneinfo = [
 # bump-my-version
 [tool.bumpversion]
 allow_dirty = true
-current_version = "0.109.15"
+current_version = "0.109.16"
 
 [[tool.bumpversion.files]]
 filename = "src/utilities/__init__.py"

{dycw_utilities-0.109.15 → dycw_utilities-0.109.16}/src/tests/test_numpy.py

@@ -3,6 +3,7 @@ from __future__ import annotations
 from re import escape
 from typing import TYPE_CHECKING, Any, Literal
 
+import numpy as np
 from hypothesis import assume, given
 from hypothesis.strategies import DataObject, data, floats, integers, none
 from numpy import (
@@ -12,13 +13,16 @@ from numpy import (
     full,
     inf,
     isclose,
+    linspace,
     median,
     nan,
     ndarray,
     ones,
+    pi,
     zeros,
     zeros_like,
 )
+from numpy.fft import fft, fftfreq
 from numpy.random import Generator
 from numpy.testing import assert_equal
 from pytest import mark, param, raises
@@ -36,7 +40,9 @@ from utilities.numpy import (
     as_int,
     discretize,
     fillna,
+    filter_frequencies,
     flatn0,
+    get_frequency_spectrum,
     has_dtype,
     is_at_least,
     is_at_least_or_nan,
@@ -268,6 +274,19 @@ class TestFillNa:
         assert_equal(result, expected)
 
 
+class TestFilterFrequencies:
+    def test_main(self) -> None:
+        n = 1000
+        x = linspace(0, 2 * pi, n)
+        noise = DEFAULT_RNG.normal(scale=0.25, size=n)
+        y = x + noise
+        result = filter_frequencies(y, lambda f: np.abs(f) <= 0.02)
+        assert result.shape == (n,)
+        fft_vals = fft(result)
+        freqs = fftfreq(n)
+        assert np.allclose(fft_vals[np.abs(freqs) > 0.02], 0.0)
+
+
 class TestFlatN0:
     @given(data=data(), n=integers(1, 10))
     def test_main(self, *, data: DataObject, n: int) -> None:
@@ -288,6 +307,18 @@ class TestFlatN0:
             _ = flatn0(ones(2, dtype=bool))
 
 
+class TestGetFrequencySpectrum:
+    def test_main(self) -> None:
+        n = 1000
+        x = linspace(0, 2 * pi, n)
+        noise = DEFAULT_RNG.normal(scale=0.25, size=n)
+        y = x + noise
+        y2 = filter_frequencies(y, lambda f: np.abs(f) <= 0.02)
+        result = get_frequency_spectrum(y2)
+        assert result.shape == (n, 2)
+        assert np.allclose(result[result[:, 0] > 0.02, 1], 0.0)
+
+
 class TestHasDtype:
     @mark.parametrize(("dtype", "expected"), [param(float, True), param(int, False)])
     @mark.parametrize("is_tuple", [param(True), param(False)])

{dycw_utilities-0.109.15 → dycw_utilities-0.109.16}/src/utilities/__init__.py

@@ -1,3 +1,3 @@
 from __future__ import annotations
 
-__version__ = "0.109.15"
+__version__ = "0.109.16"

{dycw_utilities-0.109.15 → dycw_utilities-0.109.16}/src/utilities/numpy.py

@@ -1,19 +1,22 @@
 from __future__ import annotations
 
 from dataclasses import dataclass
-from functools import reduce
+from functools import partial, reduce
 from itertools import repeat
 from typing import TYPE_CHECKING, Any, overload, override
 
 import numpy as np
 from numpy import (
+    argsort,
     array,
     bool_,
+    complex128,
     digitize,
     dtype,
     errstate,
     flatnonzero,
     float64,
+    floating,
     full_like,
     inf,
     int64,
@@ -30,6 +33,7 @@ from numpy import (
     roll,
     where,
 )
+from numpy.fft import fft, fftfreq, ifft
 from numpy.linalg import det, eig
 from numpy.random import default_rng
 from numpy.typing import NDArray
@@ -37,7 +41,7 @@ from numpy.typing import NDArray
 from utilities.iterables import is_iterable_not_str
 
 if TYPE_CHECKING:
-    from collections.abc import Iterable
+    from collections.abc import Callable, Iterable
 
 
 ##
@@ -160,6 +164,33 @@ def fillna(array: NDArrayF, /, *, value: float = 0.0) -> NDArrayF:
 ##
 
 
+def filter_frequencies(
+    array: NDArrayF,
+    /,
+    *filters: Callable[[NDArray[floating[Any]]], NDArrayB],
+    d: int = 1,
+) -> NDArrayF:
+    """Filter an array by the frequencies of its FFT."""
+    (n,) = array.shape
+    fft_vals = fft(array)
+    freqs = fftfreq(n, d=d)
+    reduced = reduce(partial(_filter_frequencies_one, freqs=freqs), filters, fft_vals)
+    return ifft(reduced).real
+
+
+def _filter_frequencies_one(
+    acc: NDArray[complex128],
+    el: Callable[[NDArray[floating[Any]]], NDArrayB],
+    /,
+    *,
+    freqs: NDArray[floating[Any]],
+) -> NDArray[complex128]:
+    return where(el(freqs), acc, 0.0)
+
+
+##
+
+
 def flatn0(array: NDArrayB, /) -> int:
     """Return the index of the unique True element."""
     if not array.any():
@@ -193,6 +224,19 @@ class FlatN0MultipleError(FlatN0Error):
 ##
 
 
+def get_frequency_spectrum(array: NDArrayF, /, *, d: int = 1) -> NDArray[floating[Any]]:
+    """Get the frequency spectrum."""
+    (n,) = array.shape
+    fft_vals = fft(array)
+    freqs = fftfreq(n, d=d)
+    amplitudes = np.abs(fft_vals)
+    data = np.hstack([freqs.reshape(-1, 1), amplitudes.reshape(-1, 1)])
+    return data[argsort(data[:, 0])]
+
+
+##
+
+
 def has_dtype(x: Any, dtype: Any, /) -> bool:
     """Check if an object has the required dtype."""
     if is_iterable_not_str(dtype):
@@ -851,7 +895,9 @@ __all__ = [
     "datetime64us",
     "discretize",
     "fillna",
+    "filter_frequencies",
     "flatn0",
+    "get_frequency_spectrum",
     "has_dtype",
     "is_at_least",
     "is_at_least_or_nan",