dycw-utilities 0.109.17__tar.gz → 0.109.19__tar.gz

{dycw_utilities-0.109.17 → dycw_utilities-0.109.19}/PKG-INFO

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

{dycw_utilities-0.109.17 → dycw_utilities-0.109.19}/pyproject.toml

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

{dycw_utilities-0.109.17 → dycw_utilities-0.109.19}/src/tests/test_numpy.py

@@ -19,6 +19,7 @@ from numpy import (
     ndarray,
     ones,
     pi,
+    where,
     zeros,
     zeros_like,
 )
@@ -40,12 +41,12 @@ from utilities.numpy import (
     _BoxCarLocationsError,
     _BoxCarLowerBoundSlopeError,
     _BoxCarUpperBoundSlopeError,
+    adjust_frequencies,
     array_indexer,
     as_int,
     boxcar,
     discretize,
     fillna,
-    filter_frequencies,
     flatn0,
     get_frequency_spectrum,
     has_dtype,
@@ -106,6 +107,32 @@ if TYPE_CHECKING:
     from collections.abc import Sequence
 
 
+class TestAdjustFrequencies:
+    def test_filter(self) -> None:
+        n = 1000
+        x = linspace(0, 2 * pi, n)
+        noise = DEFAULT_RNG.normal(scale=0.25, size=n)
+        y = x + noise
+        result = adjust_frequencies(y, filters=lambda f: np.abs(f) <= 0.02)
+        assert result.shape == (n,)
+        amplitudes = fft(result)
+        freqs = fftfreq(n)
+        assert np.allclose(amplitudes[np.abs(freqs) > 0.02], 0.0)
+
+    def test_weight(self) -> None:
+        n = 1000
+        x = linspace(0, 2 * pi, n)
+        noise = DEFAULT_RNG.normal(scale=0.25, size=n)
+        y = x + noise
+        result = adjust_frequencies(
+            y, weights=lambda f: where(np.abs(f) <= 0.02, 1.0, 0.0)
+        )
+        assert result.shape == (n,)
+        amplitudes = fft(result)
+        freqs = fftfreq(n)
+        assert np.allclose(amplitudes[np.abs(freqs) > 0.02], 0.0)
+
+
 class TestArrayIndexer:
     @mark.parametrize(
         ("i", "ndim", "expected"),
@@ -330,19 +357,6 @@ 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:
@@ -369,10 +383,10 @@ class TestGetFrequencySpectrum:
         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)
+        y2 = adjust_frequencies(y, filters=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)
+        assert np.allclose(result[np.abs(result[:, 0]) > 0.02, 1], 0.0)
 
 
 class TestHasDtype:

{dycw_utilities-0.109.17 → dycw_utilities-0.109.19}/src/tests/test_polars.py

@@ -12,6 +12,7 @@ from typing import TYPE_CHECKING, Any, ClassVar, Literal, cast
 from uuid import UUID, uuid4
 
 import hypothesis.strategies
+import numpy as np
 import polars as pl
 from hypothesis import given
 from hypothesis.strategies import (
@@ -26,6 +27,7 @@ from hypothesis.strategies import (
     sampled_from,
     timezones,
 )
+from numpy import allclose, linspace, pi
 from polars import (
     Boolean,
     DataFrame,
@@ -60,6 +62,7 @@ from utilities.hypothesis import (
     zoned_datetimes,
 )
 from utilities.math import number_of_decimals
+from utilities.numpy import DEFAULT_RNG
 from utilities.pathlib import PWD
 from utilities.polars import (
     AppendDataClassError,
@@ -104,6 +107,7 @@ from utilities.polars import (
     _InsertBetweenMissingColumnsError,
     _InsertBetweenNonConsecutiveError,
     _yield_struct_series_element_remove_nulls,
+    adjust_frequencies,
     append_dataclass,
     are_frames_equal,
     ceil_datetime,
@@ -113,6 +117,7 @@ from utilities.polars import (
     concat_series,
     convert_time_zone,
     cross,
+    cross_rolling_quantile,
     dataclass_to_dataframe,
     dataclass_to_schema,
     drop_null_struct_series,
@@ -121,6 +126,7 @@ from utilities.polars import (
     finite_ewm_mean,
     floor_datetime,
     get_data_type_or_series_time_zone,
+    get_frequency_spectrum,
     get_series_number_of_decimals,
     insert_after,
     insert_before,
@@ -162,6 +168,16 @@ if TYPE_CHECKING:
     from utilities.types import MaybeType, StrMapping, WeekDay
 
 
+class TestAdjustFrequencies:
+    def test_main(self) -> None:
+        n = 1000
+        x = linspace(0, 2 * pi, n)
+        noise = DEFAULT_RNG.normal(scale=0.25, size=n)
+        y = Series(values=x + noise)
+        result = adjust_frequencies(y, filters=lambda f: np.abs(f) <= 0.02)
+        assert isinstance(result, Series)
+
+
 class TestAppendDataClass:
     @given(
         data=fixed_dictionaries({
@@ -661,6 +677,45 @@ class TestCrossOrTouch:
         assert_series_equal(result, expected)
 
 
+class TestCrossRollingQuantile:
+    def test_main(self) -> None:
+        df = DataFrame(
+            data=[
+                (4, None, None),
+                (5, None, None),
+                (7, None, None),
+                (9, None, None),
+                (0, 5.0, False),
+                (1, 5.0, False),
+                (8, 7.0, True),
+                (9, 8.0, False),
+                (2, 2.0, False),
+                (3, 3.0, False),
+            ],
+            schema={"x": Int64, "median": Float64, "cross": Boolean},
+            orient="row",
+        )
+        assert_series_equal(
+            df["x"].rolling_quantile(0.5, window_size=5),
+            df["median"],
+            check_names=False,
+        )
+        assert_series_equal(
+            cross_rolling_quantile(df["x"], "up", 0.5, window_size=5),
+            df["cross"],
+            check_names=False,
+        )
+
+    def test_example(self) -> None:
+        close = Series(name="close", values=[8, 7, 8, 5, 0], dtype=Int64)
+        mid = Series(name="mid", values=[1, 2, 3, 4, 6], dtype=Int64)
+        result = cross(close, "down", mid).alias("result")
+        expected = Series(
+            name="result", values=[None, False, False, False, True], dtype=Boolean
+        )
+        assert_series_equal(result, expected)
+
+
 class TestDataClassToDataFrame:
     @given(data=data())
     def test_basic_type(self, *, data: DataObject) -> None:
@@ -1213,6 +1268,20 @@ class TestGetDataTypeOrSeriesTimeZone:
             _ = get_data_type_or_series_time_zone(Datetime)
 
 
+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 = Series(x + noise)
+        y2 = adjust_frequencies(y, filters=lambda f: np.abs(f) <= 0.02)
+        result = get_frequency_spectrum(y2)
+        check_polars_dataframe(
+            result, height=n, schema_list={"frequency": Float64, "amplitude": Float64}
+        )
+        assert allclose(result.filter(col("frequency").abs() > 0.02)["amplitude"], 0.0)
+
+
 class TestGetSeriesNumberOfDecimals:
     @given(data=data(), n=hypothesis.strategies.integers(1, 10), nullable=booleans())
     def test_main(self, *, data: DataObject, n: int, nullable: bool) -> None:

{dycw_utilities-0.109.17 → dycw_utilities-0.109.19}/src/utilities/__init__.py

@@ -1,3 +1,3 @@
 from __future__ import annotations
 
-__version__ = "0.109.17"
+__version__ = "0.109.19"

{dycw_utilities-0.109.17 → dycw_utilities-0.109.19}/src/utilities/numpy.py

@@ -16,11 +16,10 @@ from numpy import (
     errstate,
     exp,
     flatnonzero,
-    float64,
     floating,
     full_like,
     inf,
-    int64,
+    integer,
     isclose,
     isfinite,
     isinf,
@@ -39,11 +38,13 @@ from numpy.linalg import det, eig
 from numpy.random import default_rng
 from numpy.typing import NDArray
 
-from utilities.iterables import is_iterable_not_str
+from utilities.iterables import always_iterable, is_iterable_not_str
 
 if TYPE_CHECKING:
     from collections.abc import Callable, Iterable
 
+    from utilities.types import MaybeIterable
+
 
 ##
 
@@ -89,8 +90,9 @@ timedelta64as = dtype("timedelta64[as]")
 
 NDArrayA = NDArray[Any]
 NDArrayB = NDArray[bool_]
-NDArrayF = NDArray[float64]
-NDArrayI = NDArray[int64]
+NDArrayC128 = NDArray[complex128]
+NDArrayF = NDArray[floating[Any]]
+NDArrayI = NDArray[integer[Any]]
 NDArrayO = NDArray[object_]
 
 
@@ -135,7 +137,7 @@ class AsIntError(Exception): ...
 
 
 def boxcar(
-    array: NDArray[floating[Any]],
+    array: NDArrayF,
     /,
     *,
     loc_low: float = -1.0,
@@ -144,7 +146,7 @@ def boxcar(
     slope_high: float = 1.0,
     rtol: float | None = None,
     atol: float | None = None,
-) -> NDArray[floating[Any]]:
+) -> NDArrayF:
     """Construct a boxcar function."""
     if not is_at_most(loc_low, loc_high, rtol=rtol, atol=atol):
         raise _BoxCarLocationsError(low=loc_low, high=loc_high)
@@ -224,28 +226,43 @@ def fillna(array: NDArrayF, /, *, value: float = 0.0) -> NDArrayF:
 ##
 
 
-def filter_frequencies(
+def adjust_frequencies(
     array: NDArrayF,
     /,
-    *filters: Callable[[NDArray[floating[Any]]], NDArrayB],
+    *,
+    filters: MaybeIterable[Callable[[NDArrayF], NDArrayB]] | None = None,
+    weights: MaybeIterable[Callable[[NDArrayF], NDArrayF]] | None = None,
     d: int = 1,
 ) -> NDArrayF:
-    """Filter an array by the frequencies of its FFT."""
+    """Adjust an array via its FFT frequencies."""
     (n,) = array.shape
-    fft_vals = fft(array)
+    amplitudes = fft(array)
     freqs = fftfreq(n, d=d)
-    reduced = reduce(partial(_filter_frequencies_one, freqs=freqs), filters, fft_vals)
-    return ifft(reduced).real
+    if filters is not None:
+        amplitudes = reduce(
+            partial(_adjust_frequencies_filter_one, freqs=freqs),
+            always_iterable(filters),
+            amplitudes,
+        )
+    if weights is not None:
+        amplitudes = reduce(
+            partial(_adjust_frequencies_weight_one, freqs=freqs),
+            always_iterable(weights),
+            amplitudes,
+        )
+    return ifft(amplitudes).real
+
+
+def _adjust_frequencies_filter_one(
+    acc: NDArrayC128, el: Callable[[NDArrayF], NDArrayB], /, *, freqs: NDArrayF
+) -> NDArrayC128:
+    return where(el(freqs), acc, 0.0)
 
 
-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 _adjust_frequencies_weight_one(
+    acc: NDArrayC128, el: Callable[[NDArrayF], NDArrayF], /, *, freqs: NDArrayF
+) -> NDArrayC128:
+    return acc * el(freqs)
 
 
 ##
@@ -284,12 +301,12 @@ class FlatN0MultipleError(FlatN0Error):
 ##
 
 
-def get_frequency_spectrum(array: NDArrayF, /, *, d: int = 1) -> NDArray[floating[Any]]:
+def get_frequency_spectrum(array: NDArrayF, /, *, d: int = 1) -> NDArrayF:
     """Get the frequency spectrum."""
     (n,) = array.shape
-    fft_vals = fft(array)
+    amplitudes = fft(array)
     freqs = fftfreq(n, d=d)
-    amplitudes = np.abs(fft_vals)
+    amplitudes = np.abs(amplitudes)
     data = np.hstack([freqs.reshape(-1, 1), amplitudes.reshape(-1, 1)])
     return data[argsort(data[:, 0])]
 
@@ -908,14 +925,14 @@ def shift_bool(
 
 
 def sigmoid(
-    array: NDArray[floating[Any]],
+    array: NDArrayF,
     /,
     *,
     loc: float = 0.0,
     slope: float = 1.0,
     rtol: float | None = None,
     atol: float | None = None,
-) -> NDArray[floating[Any]]:
+) -> NDArrayF:
     """Construct a sigmoid function."""
     if is_zero(slope, rtol=rtol, atol=atol):
         raise SigmoidError
@@ -965,6 +982,7 @@ __all__ = [
     "NDArrayO",
     "ShiftError",
     "SigmoidError",
+    "adjust_frequencies",
     "array_indexer",
     "as_int",
     "boxcar",
@@ -983,7 +1001,6 @@ __all__ = [
     "datetime64us",
     "discretize",
     "fillna",
-    "filter_frequencies",
     "flatn0",
     "get_frequency_spectrum",
     "has_dtype",

{dycw_utilities-0.109.17 → dycw_utilities-0.109.19}/src/utilities/polars.py

@@ -115,10 +115,12 @@ if TYPE_CHECKING:
         JoinStrategy,  # pyright: ignore[reportPrivateImportUsage]
         JoinValidation,  # pyright: ignore[reportPrivateImportUsage]
         PolarsDataType,  # pyright: ignore[reportPrivateImportUsage]
+        RollingInterpolationMethod,  # pyright: ignore[reportPrivateImportUsage]
         SchemaDict,  # pyright: ignore[reportPrivateImportUsage]
         TimeUnit,  # pyright: ignore[reportPrivateImportUsage]
     )
 
+    from utilities.numpy import NDArrayB, NDArrayF
     from utilities.types import (
         Dataclass,
         MaybeIterable,
@@ -141,6 +143,27 @@ _FINITE_EWM_MIN_WEIGHT = 0.9999
 ##
 
 
+def adjust_frequencies(
+    series: Series,
+    /,
+    *,
+    filters: MaybeIterable[Callable[[NDArrayF], NDArrayB]] | None = None,
+    weights: MaybeIterable[Callable[[NDArrayF], NDArrayF]] | None = None,
+    d: int = 1,
+) -> Series:
+    """Adjust a Series via its FFT frequencies."""
+    import utilities.numpy
+
+    array = series.to_numpy()
+    adjusted = utilities.numpy.adjust_frequencies(
+        array, filters=filters, weights=weights, d=d
+    )
+    return Series(name=series.name, values=adjusted, dtype=Float64)
+
+
+##
+
+
 def append_dataclass(df: DataFrame, obj: Dataclass, /) -> DataFrame:
     """Append a dataclass object to a DataFrame."""
     non_null_fields = {k: v for k, v in asdict(obj).items() if v is not None}
@@ -658,6 +681,73 @@ def _cross_or_touch(
 ##
 
 
+@overload
+def cross_rolling_quantile(
+    expr: ExprLike,
+    up_or_down: Literal["up", "down"],
+    quantile: float,
+    /,
+    *,
+    interpolation: RollingInterpolationMethod = "nearest",
+    window_size: int = 2,
+    weights: list[float] | None = None,
+    min_samples: int | None = None,
+    center: bool = False,
+) -> Expr: ...
+@overload
+def cross_rolling_quantile(
+    expr: Series,
+    up_or_down: Literal["up", "down"],
+    quantile: float,
+    /,
+    *,
+    interpolation: RollingInterpolationMethod = "nearest",
+    window_size: int = 2,
+    weights: list[float] | None = None,
+    min_samples: int | None = None,
+    center: bool = False,
+) -> Series: ...
+@overload
+def cross_rolling_quantile(
+    expr: IntoExprColumn,
+    up_or_down: Literal["up", "down"],
+    quantile: float,
+    /,
+    *,
+    interpolation: RollingInterpolationMethod = "nearest",
+    window_size: int = 2,
+    weights: list[float] | None = None,
+    min_samples: int | None = None,
+    center: bool = False,
+) -> Expr | Series: ...
+def cross_rolling_quantile(
+    expr: IntoExprColumn,
+    up_or_down: Literal["up", "down"],
+    quantile: float,
+    /,
+    *,
+    interpolation: RollingInterpolationMethod = "nearest",
+    window_size: int = 2,
+    weights: list[float] | None = None,
+    min_samples: int | None = None,
+    center: bool = False,
+) -> Expr | Series:
+    """Compute when a column crosses its rolling quantile."""
+    expr = ensure_expr_or_series(expr)
+    rolling = expr.rolling_quantile(
+        quantile,
+        interpolation=interpolation,
+        window_size=window_size,
+        weights=weights,
+        min_samples=min_samples,
+        center=center,
+    )
+    return cross(expr, up_or_down, rolling)
+
+
+##
+
+
 def dataclass_to_dataframe(
     objs: MaybeIterable[Dataclass],
     /,
@@ -1027,6 +1117,20 @@ class _GetDataTypeOrSeriesTimeZoneNotZonedError(GetDataTypeOrSeriesTimeZoneError
 ##
 
 
+def get_frequency_spectrum(series: Series, /, *, d: int = 1) -> DataFrame:
+    """Get the frequency spectrum."""
+    import utilities.numpy
+
+    array = series.to_numpy()
+    spectrum = utilities.numpy.get_frequency_spectrum(array, d=d)
+    return DataFrame(
+        data=spectrum, schema={"frequency": Float64, "amplitude": Float64}, orient="row"
+    )
+
+
+##
+
+
 @overload
 def get_series_number_of_decimals(
     series: Series, /, *, nullable: Literal[True]
@@ -1736,6 +1840,7 @@ __all__ = [
     "SetFirstRowAsColumnsError",
     "StructFromDataClassError",
     "YieldStructSeriesElementsError",
+    "adjust_frequencies",
     "append_dataclass",
     "are_frames_equal",
     "ceil_datetime",
@@ -1753,6 +1858,7 @@ __all__ = [
     "finite_ewm_mean",
     "floor_datetime",
     "get_data_type_or_series_time_zone",
+    "get_frequency_spectrum",
     "get_series_number_of_decimals",
     "insert_after",
     "insert_before",

{dycw_utilities-0.109.17 → dycw_utilities-0.109.19}/src/utilities/polars_ols.py

@@ -10,7 +10,9 @@ from utilities.functions import is_sequence_of
 from utilities.polars import concat_series, ensure_expr_or_series
 
 if TYPE_CHECKING:
-    from polars._typing import IntoExprColumn
+    from polars._typing import (
+        IntoExprColumn,  # pyright: ignore[reportPrivateImportUsage]
+    )
     from polars_ols import NullPolicy
 
     from utilities.polars import ExprLike