dycw-utilities 0.109.14__py3-none-any.whl → 0.109.16__py3-none-any.whl

{dycw_utilities-0.109.14.dist-info → dycw_utilities-0.109.16.dist-info}/METADATA

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

{dycw_utilities-0.109.14.dist-info → dycw_utilities-0.109.16.dist-info}/RECORD

@@ -1,4 +1,4 @@
-utilities/__init__.py,sha256=9sl1XiZAUQGAGfCyJEHhTC7Kfmx8TtcYJHtVOZdUugw,61
+utilities/__init__.py,sha256=4AFXsQvJMy-nibpq-XnsnHXbkRajpf73GAF_okhFQgY,61
 utilities/altair.py,sha256=Gpja-flOo-Db0PIPJLJsgzAlXWoKUjPU1qY-DQ829ek,9156
 utilities/astor.py,sha256=xuDUkjq0-b6fhtwjhbnebzbqQZAjMSHR1IIS5uOodVg,777
 utilities/asyncio.py,sha256=41oQUurWMvadFK5gFnaG21hMM0Vmfn2WS6OpC0R9mas,14757
@@ -36,7 +36,7 @@ utilities/math.py,sha256=TexfvLCI12d9Sw5_W4pKVBZ3nRr3zk2iPkcEU7xdEWU,26771
 utilities/memory_profiler.py,sha256=tf2C51P2lCujPGvRt2Rfc7VEw5LDXmVPCG3z_AvBmbU,962
 utilities/modules.py,sha256=SnhsRHRUS1po_acejrINauihGQpPvVsp8RDNCei1OLQ,3173
 utilities/more_itertools.py,sha256=CPUxrMAcTwRxbzbhiqPKi3Xx9hxqI0t6gkWjutaibGk,5534
-utilities/numpy.py,sha256=rA1b0_GkBUSMjnv77tinRM70KRnkcmZxI9xbrsXFDRg,21819
+utilities/numpy.py,sha256=-VCZZBUs9GaLxcPOHGVc_iLkVP_SEKKRL7YjV05jrS4,22961
 utilities/operator.py,sha256=0M2yZJ0PODH47ogFEnkGMBe_cfxwZR02T_92LZVZvHo,3715
 utilities/optuna.py,sha256=loyJGWTzljgdJaoLhP09PT8Jz6o_pwBOwehY33lHkhw,1923
 utilities/orjson.py,sha256=Wj5pzG_VdgoAy14a7Luhem-BgYrRtRFvvl_POiszRd0,36930
@@ -47,7 +47,7 @@ utilities/period.py,sha256=ikHXsWtDLr553cfH6p9mMaiCnIAP69B7q84ckWV3HaA,10884
 utilities/pickle.py,sha256=Bhvd7cZl-zQKQDFjUerqGuSKlHvnW1K2QXeU5UZibtg,657
 utilities/platform.py,sha256=NU7ycTvAXAG-fdYmDXaM1m4EOml2cGiaYwaUzfzSqyU,1767
 utilities/polars.py,sha256=aOQNVyV04qYZjg7Exi6zYERhSQoCMzBP74oufxqANFY,52167
-utilities/polars_ols.py,sha256=AQe3RFOMv8CEI_ZCoscb_-PxB4JWjO0TAEmk8DKLeaI,2138
+utilities/polars_ols.py,sha256=Uc9V5kvlWZ5cU93lKZ-cfAKdVFFw81tqwLW9PxtUvMs,5618
 utilities/pqdm.py,sha256=foRytQybmOQ05pjt5LF7ANyzrIa--4ScDE3T2wd31a4,3118
 utilities/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 utilities/pydantic.py,sha256=f6qtR5mO2YMuyvNmbaEj5YeD9eGA4YYfb7Bjzh9jUs0,1845
@@ -86,7 +86,7 @@ utilities/warnings.py,sha256=yUgjnmkCRf6QhdyAXzl7u0qQFejhQG3PrjoSwxpbHrs,1819
 utilities/whenever.py,sha256=TjoTAJ1R27-rKXiXzdE4GzPidmYqm0W58XydDXp-QZM,17786
 utilities/zipfile.py,sha256=24lQc9ATcJxHXBPc_tBDiJk48pWyRrlxO2fIsFxU0A8,699
 utilities/zoneinfo.py,sha256=-DQz5a0Ikw9jfSZtL0BEQkXOMC9yGn_xiJYNCLMiqEc,1989
-dycw_utilities-0.109.14.dist-info/METADATA,sha256=sjPmXifZcBMY4js2lgCXzweOgdeXLZ1qihJdZ8lFsTA,13005
-dycw_utilities-0.109.14.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-dycw_utilities-0.109.14.dist-info/licenses/LICENSE,sha256=gppZp16M6nSVpBbUBrNL6JuYfvKwZiKgV7XoKKsHzqo,1066
-dycw_utilities-0.109.14.dist-info/RECORD,,
+dycw_utilities-0.109.16.dist-info/METADATA,sha256=u1Fl275S55bRfV6fTrMT8qHsS8_bu4aispHkIR6QNz8,13005
+dycw_utilities-0.109.16.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+dycw_utilities-0.109.16.dist-info/licenses/LICENSE,sha256=gppZp16M6nSVpBbUBrNL6JuYfvKwZiKgV7XoKKsHzqo,1066
+dycw_utilities-0.109.16.dist-info/RECORD,,

utilities/__init__.py

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

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",

utilities/polars_ols.py

@@ -1,19 +1,22 @@
 from __future__ import annotations
 
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, overload
 
-from polars import struct
+from polars import Expr, Series, struct
 from polars_ols import RollingKwargs, compute_rolling_least_squares
 
-from utilities.polars import ensure_expr_or_series
+from utilities.errors import ImpossibleCaseError
+from utilities.functions import is_sequence_of
+from utilities.polars import concat_series, ensure_expr_or_series
 
 if TYPE_CHECKING:
-    from polars import Expr
+    from polars._typing import IntoExprColumn
     from polars_ols import NullPolicy
 
     from utilities.polars import ExprLike
 
 
+@overload
 def compute_rolling_ols(
     target: ExprLike,
     *features: ExprLike,
@@ -24,9 +27,106 @@ def compute_rolling_ols(
     min_periods: int | None = None,
     use_woodbury: bool | None = None,
     alpha: float | None = None,
+) -> Expr: ...
+@overload
+def compute_rolling_ols(
+    target: Series,
+    *features: Series,
+    sample_weights: Series | None = None,
+    add_intercept: bool = False,
+    null_policy: NullPolicy = "drop_window",
+    window_size: int = 1000000,
+    min_periods: int | None = None,
+    use_woodbury: bool | None = None,
+    alpha: float | None = None,
+) -> Series: ...
+@overload
+def compute_rolling_ols(
+    target: IntoExprColumn,
+    *features: IntoExprColumn,
+    sample_weights: IntoExprColumn | None = None,
+    add_intercept: bool = False,
+    null_policy: NullPolicy = "drop_window",
+    window_size: int = 1000000,
+    min_periods: int | None = None,
+    use_woodbury: bool | None = None,
+    alpha: float | None = None,
+) -> Expr | Series: ...
+def compute_rolling_ols(
+    target: IntoExprColumn,
+    *features: IntoExprColumn,
+    sample_weights: IntoExprColumn | None = None,
+    add_intercept: bool = False,
+    null_policy: NullPolicy = "drop_window",
+    window_size: int = 1000000,
+    min_periods: int | None = None,
+    use_woodbury: bool | None = None,
+    alpha: float | None = None,
+) -> Expr | Series:
+    """Compute a rolling OLS."""
+    target = ensure_expr_or_series(target)
+    features2 = tuple(map(ensure_expr_or_series, features))
+    sample_weights = (
+        None if sample_weights is None else ensure_expr_or_series(sample_weights)
+    )
+    if (
+        isinstance(target, Expr)
+        and is_sequence_of(features2, Expr)
+        and ((sample_weights is None) or isinstance(sample_weights, Expr))
+    ):
+        return _compute_rolling_ols_expr(
+            target,
+            *features2,
+            sample_weights=sample_weights,
+            add_intercept=add_intercept,
+            null_policy=null_policy,
+            window_size=window_size,
+            min_periods=min_periods,
+            use_woodbury=use_woodbury,
+            alpha=alpha,
+        )
+    if (
+        isinstance(target, Series)
+        and is_sequence_of(features2, Series)
+        and ((sample_weights is None) or isinstance(sample_weights, Series))
+    ):
+        return concat_series(
+            target, *features2, *([] if sample_weights is None else [sample_weights])
+        ).with_columns(
+            _compute_rolling_ols_expr(
+                target.name,
+                *(f.name for f in features2),
+                sample_weights=None if sample_weights is None else sample_weights.name,
+                add_intercept=add_intercept,
+                null_policy=null_policy,
+                window_size=window_size,
+                min_periods=min_periods,
+                use_woodbury=use_woodbury,
+                alpha=alpha,
+            )
+        )["ols"]
+    raise ImpossibleCaseError(  # pragma: no cover
+        case=[f"{target=}", f"{features2=}", f"{sample_weights=}"]
+    )
+
+
+def _compute_rolling_ols_expr(
+    target: ExprLike,
+    *features: ExprLike,
+    sample_weights: ExprLike | None = None,
+    add_intercept: bool = False,
+    null_policy: NullPolicy = "drop_window",
+    window_size: int = 1000000,
+    min_periods: int | None = None,
+    use_woodbury: bool | None = None,
+    alpha: float | None = None,
 ) -> Expr:
     """Compute a rolling OLS."""
     target = ensure_expr_or_series(target)
+    features2 = tuple(map(ensure_expr_or_series, features))
+    sample_weights = (
+        None if sample_weights is None else ensure_expr_or_series(sample_weights)
+    )
     rolling_kwargs = RollingKwargs(
         null_policy=null_policy,
         window_size=window_size,
@@ -36,7 +136,7 @@ def compute_rolling_ols(
     )
     coefficients = compute_rolling_least_squares(
         target,
-        *features,
+        *features2,
         sample_weights=sample_weights,
         add_intercept=add_intercept,
         mode="coefficients",
@@ -44,7 +144,7 @@ def compute_rolling_ols(
     ).alias("coefficients")
     predictions = compute_rolling_least_squares(
         target,
-        *features,
+        *features2,
         sample_weights=sample_weights,
         add_intercept=add_intercept,
         mode="predictions",
@@ -52,7 +152,7 @@ def compute_rolling_ols(
     ).alias("predictions")
     residuals = compute_rolling_least_squares(
         target,
-        *features,
+        *features2,
         sample_weights=sample_weights,
         add_intercept=add_intercept,
         mode="residuals",