dycw-utilities 0.109.8__tar.gz → 0.109.10__tar.gz

{dycw_utilities-0.109.8 → dycw_utilities-0.109.10}/PKG-INFO

@@ -1,9 +1,10 @@
 Metadata-Version: 2.4
 Name: dycw-utilities
-Version: 0.109.8
+Version: 0.109.10
 Author-email: Derek Wan <d.wan@icloud.com>
 License-File: LICENSE
 Requires-Python: >=3.12
+Requires-Dist: polars-ols>=0.3.5
 Requires-Dist: typing-extensions<4.14,>=4.13.1
 Provides-Extra: test
 Requires-Dist: hypothesis<6.132,>=6.131.6; extra == 'test'

{dycw_utilities-0.109.8 → dycw_utilities-0.109.10}/pyproject.toml

@@ -40,6 +40,7 @@ dev = [
   "orjson >= 3.10.16, < 3.11",
   "pathvalidate >= 3.2.3, < 3.3",
   "polars-lts-cpu >= 1.27.1, < 1.28",
+  "polars-ols >= 0.3.5, < 0.4",
   "pqdm >= 0.2.0, < 0.3",
   "psycopg2-binary >= 2.9.10, < 2.10",  # for sqlalchemy
   "pydantic >= 2.11.2, < 2.12",
@@ -50,6 +51,7 @@ dev = [
   "python-dotenv >= 1.1.0, < 1.2",
   "redis >= 5.2.1, < 5.3",
   "rich >= 14.0.0, < 14.1",
+  "scikit-learn >= 1.6.1, < 1.7",
   "scipy >= 1.15.2, < 1.16",
   "slack-sdk >= 3.35.0, < 3.36",
   "sqlalchemy >= 2.0.40, < 2.1",
@@ -84,12 +86,13 @@ dev = [
 [project]
 authors = [{name = "Derek Wan", email = "d.wan@icloud.com"}]
 dependencies = [
+  "polars-ols>=0.3.5",
   "typing-extensions >= 4.13.1, < 4.14",
 ]
 name = "dycw-utilities"
 readme = "README.md"
 requires-python = ">= 3.12"
-version = "0.109.8"
+version = "0.109.10"
 
 [project.optional-dependencies]
 test = [
@@ -332,7 +335,7 @@ zzz-test-zoneinfo = [
 # bump-my-version
 [tool.bumpversion]
 allow_dirty = true
-current_version = "0.109.8"
+current_version = "0.109.10"
 
 [[tool.bumpversion.files]]
 filename = "src/utilities/__init__.py"

{dycw_utilities-0.109.8 → dycw_utilities-0.109.10}/src/tests/test_polars.py

@@ -11,6 +11,7 @@ from pathlib import Path
 from typing import TYPE_CHECKING, Any, ClassVar, Literal, cast
 from uuid import UUID, uuid4
 
+import hypothesis.strategies
 import polars as pl
 from hypothesis import given
 from hypothesis.strategies import (
@@ -20,7 +21,6 @@ from hypothesis.strategies import (
     data,
     fixed_dictionaries,
     floats,
-    integers,
     lists,
     none,
     sampled_from,
@@ -50,10 +50,12 @@ from polars import (
 from polars.testing import assert_frame_equal, assert_series_equal
 from pytest import raises
 
+import utilities.polars
 from utilities.datetime import get_now, get_today
 from utilities.hypothesis import (
     assume_does_not_raise,
     int64s,
+    pairs,
     text_ascii,
     zoned_datetimes,
 )
@@ -129,11 +131,13 @@ from utilities.polars import (
     map_over_columns,
     nan_sum_agg,
     nan_sum_cols,
+    normal,
     replace_time_zone,
     set_first_row_as_columns,
     struct_dtype,
     struct_from_dataclass,
     touch,
+    uniform,
     unique_element,
     week_num,
     yield_struct_series_dataclasses,
@@ -1210,7 +1214,7 @@ class TestGetDataTypeOrSeriesTimeZone:
 
 
 class TestGetSeriesNumberOfDecimals:
-    @given(data=data(), n=integers(1, 10), nullable=booleans())
+    @given(data=data(), n=hypothesis.strategies.integers(1, 10), nullable=booleans())
     def test_main(self, *, data: DataObject, n: int, nullable: bool) -> None:
         strategy = int64s() | none() if nullable else int64s()
         ints_or_none = data.draw(lists(strategy, min_size=1, max_size=10))
@@ -1312,6 +1316,37 @@ class TestInsertBetween:
             _ = insert_between(self.df, "a", "c", lit(None).alias("new"))
 
 
+class TestIntegers:
+    @given(
+        length=hypothesis.strategies.integers(0, 10),
+        high=hypothesis.strategies.integers(1, 10),
+    )
+    def test_int(self, *, length: int, high: int) -> None:
+        series = utilities.polars.integers(length, high)
+        assert series.len() == length
+        assert series.is_between(0, high, closed="left").all()
+
+    @given(
+        length=hypothesis.strategies.integers(0, 10),
+        high=hypothesis.strategies.integers(1, 10),
+    )
+    def test_series(self, *, length: int, high: int) -> None:
+        orig = int_range(end=length, eager=True)
+        series = utilities.polars.integers(orig, high)
+        assert series.len() == length
+        assert series.is_between(0, high, closed="left").all()
+
+    @given(
+        length=hypothesis.strategies.integers(0, 10),
+        high=hypothesis.strategies.integers(1, 10),
+    )
+    def test_dataframe(self, *, length: int, high: int) -> None:
+        df = int_range(end=length, eager=True).to_frame()
+        series = utilities.polars.integers(df, high)
+        assert series.len() == length
+        assert series.is_between(0, high, closed="left").all()
+
+
 class TestIsNullAndIsNotNullStructSeries:
     @given(
         case=sampled_from([
@@ -1526,6 +1561,28 @@ class TestNanSumCols:
         assert df["z"].item() == expected
 
 
+class TestNormal:
+    @given(length=hypothesis.strategies.integers(0, 10))
+    def test_int(self, *, length: int) -> None:
+        series = normal(length)
+        assert series.len() == length
+        assert series.is_finite().all()
+
+    @given(length=hypothesis.strategies.integers(0, 10))
+    def test_series(self, *, length: int) -> None:
+        orig = int_range(end=length, eager=True)
+        series = normal(orig)
+        assert series.len() == length
+        assert series.is_finite().all()
+
+    @given(length=hypothesis.strategies.integers(0, 10))
+    def test_dataframe(self, *, length: int) -> None:
+        df = int_range(end=length, eager=True).to_frame()
+        series = normal(df)
+        assert series.len() == length
+        assert series.is_finite().all()
+
+
 class TestReplaceTimeZone:
     def test_datetime(self) -> None:
         now_utc = get_now()
@@ -1711,6 +1768,40 @@ class TestStructFromDataClass:
             _ = struct_from_dataclass(Example)
 
 
+class TestUniform:
+    @given(
+        length=hypothesis.strategies.integers(0, 10),
+        bounds=pairs(floats(0.0, 1.0), sorted=True),
+    )
+    def test_int(self, *, length: int, bounds: tuple[float, float]) -> None:
+        low, high = bounds
+        series = uniform(length, low=low, high=high)
+        assert series.len() == length
+        assert series.is_between(low, high).all()
+
+    @given(
+        length=hypothesis.strategies.integers(0, 10),
+        bounds=pairs(floats(0.0, 1.0), sorted=True),
+    )
+    def test_series(self, *, length: int, bounds: tuple[float, float]) -> None:
+        low, high = bounds
+        orig = int_range(end=length, eager=True)
+        series = uniform(orig, low=low, high=high)
+        assert series.len() == length
+        assert series.is_between(low, high).all()
+
+    @given(
+        length=hypothesis.strategies.integers(0, 10),
+        bounds=pairs(floats(0.0, 1.0), sorted=True),
+    )
+    def test_dataframe(self, *, length: int, bounds: tuple[float, float]) -> None:
+        low, high = bounds
+        df = int_range(end=length, eager=True).to_frame()
+        series = uniform(df, low=low, high=high)
+        assert series.len() == length
+        assert series.is_between(low, high).all()
+
+
 class TestUniqueElement:
     def test_main(self) -> None:
         series = Series(

dycw_utilities-0.109.10/src/tests/test_polars_ols.py

@@ -0,0 +1,103 @@
+from __future__ import annotations
+
+from hypothesis import given
+from hypothesis.strategies import sampled_from
+from numpy import isclose
+from polars import DataFrame, Float64, Struct, UInt32, col
+from polars.testing import assert_frame_equal
+from sklearn.linear_model import LinearRegression
+
+from utilities.polars import concat_series, integers, normal
+from utilities.polars_ols import compute_rolling_ols
+
+
+class TestComputeRollingOLS:
+    def test_ols_results(self) -> None:
+        result = self._df["ols"].struct.unnest().with_row_index()
+        # fmt: off
+        exp_values = [
+            (10, {"x1": 0.333396198442681, "x2": 0.6845517746145712, "const": 0.2808021232120448}, 59.921571424913495, 1.67032007883995, 0.9955364659986504),
+            (11, {"x1": 0.322785525889542, "x2": 0.6896341527044252, "const": 0.5401793852579858}, 15.974446064929626, -0.3429678871268038, 0.9961762567103958),
+            (12, {"x1": 0.31042868991153927, "x2": 0.7055685710743383, "const": 1.145326562525439}, -31.310827706894123, -0.45191863996575066, 0.998022262986332),
+            (13, {"x1": 0.33311466967931097, "x2": 0.684137842579758, "const": -0.7961518480794516}, 50.66821598287034, -2.975371834066671, 0.9974533939791341),
+            (14, {"x1": 0.35299385150914864, "x2": 0.6758890569593843, "const": -0.9377907849336107}, -0.8749325340834626, 1.0581261048863142, 0.9973453833170313),
+            (15, {"x1": 0.351300641938209, "x2": 0.6456834722890913, "const": -1.859577387752822}, 1.6809655259738476, 0.3217076349681922, 0.9951571413022856),
+            (16, {"x1": 0.3583378199895871, "x2": 0.6588347796692774, "const": -1.109675446287481}, 26.65448170418155, 2.496480675700724, 0.9933751737130443),
+            (17, {"x1": 0.36195632084801765, "x2": 0.658322951215466, "const": -1.3860237983291754}, 7.43453332939416, -0.791818995629618, 0.9905085882663488),
+            (18, {"x1": 0.35564162435283225, "x2": 0.6656931556738634, "const": -0.562680573000551}, -51.90315462605006, 0.6592474497562932, 0.9973576833556038),
+            (19, {"x1": 0.3100421300754357, "x2": 0.675357816881863, "const": 0.48493124625501927}, -36.70039604095908, -0.6071841038068868, 0.9978541580643828),
+        ]
+        # fmt: on
+        expected = DataFrame(
+            data=exp_values,
+            schema={
+                "index": UInt32,
+                "coefficients": Struct({
+                    "x1": Float64,
+                    "x2": Float64,
+                    "const": Float64,
+                }),
+                "predictions": Float64,
+                "residuals": Float64,
+                "R2": Float64,
+            },
+            orient="row",
+        )
+        assert_frame_equal(result[-10:], expected)
+
+    @given(
+        case=sampled_from([
+            (
+                slice(-7, -2),
+                [0.3619563208480195, 0.6583229512154678],
+                -1.386023798329262,
+                7.434533329394103,
+                0.994253238813284,
+            ),
+            (
+                slice(-6, -1),
+                [0.35564162435283264, 0.6656931556738643],
+                -0.5626805730005437,
+                -51.903154626050124,
+                0.9979752966843768,
+            ),
+            (
+                slice(-5, None),
+                [0.3100421300754358, 0.6753578168818635],
+                0.48493124625502837,
+                -36.70039604095908,
+                0.9977272526713715,
+            ),
+        ])
+    )
+    def test_tail(
+        self, *, case: tuple[slice, list[float], float, float, float]
+    ) -> None:
+        slice_, coeffs, intercept, prediction, r2 = case
+        df = self._df[slice_]
+        X = df.select("x1", "x2").to_numpy()  # noqa: N806
+        y = df.select("y").to_numpy()
+        model = LinearRegression()
+        model = model.fit(X, y)
+        assert isclose(model.coef_, coeffs).all()
+        assert isclose(model.intercept_, intercept)
+        assert isclose(model.predict(X)[-1], prediction).all()
+        assert isclose(model.score(X, y), r2)
+
+    @property
+    def _df(self) -> DataFrame:
+        n = 20
+        return (
+            concat_series(
+                integers(n, -100, high=100, seed=0).alias("x1"),
+                integers(n, -100, high=100, seed=1).alias("x2"),
+            )
+            .with_columns(
+                y=(col("x1") + 2 * col("x2") + normal(n, scale=10.0, seed=2)) / 3
+            )
+            .with_columns(
+                compute_rolling_ols(
+                    "y", "x1", "x2", window_size=5, min_periods=5, add_intercept=True
+                )
+            )
+        )

{dycw_utilities-0.109.8 → dycw_utilities-0.109.10}/src/utilities/__init__.py

@@ -1,3 +1,3 @@
 from __future__ import annotations
 
-__version__ = "0.109.8"
+__version__ = "0.109.10"

{dycw_utilities-0.109.8 → dycw_utilities-0.109.10}/src/utilities/polars.py

@@ -1165,6 +1165,52 @@ class _InsertBetweenNonConsecutiveError(InsertBetweenError):
 ##
 
 
+def integers(
+    obj: int | Series | DataFrame,
+    low: int,
+    /,
+    *,
+    high: int | None = None,
+    seed: int | None = None,
+    endpoint: bool = False,
+    name: str | None = None,
+    dtype: PolarsDataType = Int64,
+) -> Series:
+    """Construct a series of normally-distributed numbers."""
+    match obj:
+        case int() as height:
+            from numpy.random import default_rng
+
+            rng = default_rng(seed=seed)
+            values = rng.integers(low, high=high, size=height, endpoint=endpoint)
+            return Series(name=name, values=values, dtype=dtype)
+        case Series() as series:
+            return integers(
+                series.len(),
+                low,
+                high=high,
+                seed=seed,
+                endpoint=endpoint,
+                name=name,
+                dtype=dtype,
+            )
+        case DataFrame() as df:
+            return integers(
+                df.height,
+                low,
+                high=high,
+                seed=seed,
+                endpoint=endpoint,
+                name=name,
+                dtype=dtype,
+            )
+        case _ as never:
+            assert_never(never)
+
+
+##
+
+
 def is_not_null_struct_series(series: Series, /) -> Series:
     """Check if a struct-dtype Series is not null as per the <= 1.1 definition."""
     try:
@@ -1324,6 +1370,39 @@ def nan_sum_cols(
 ##
 
 
+def normal(
+    obj: int | Series | DataFrame,
+    /,
+    *,
+    loc: float = 0.0,
+    scale: float = 1.0,
+    seed: int | None = None,
+    name: str | None = None,
+    dtype: PolarsDataType = Float64,
+) -> Series:
+    """Construct a series of normally-distributed numbers."""
+    match obj:
+        case int() as height:
+            from numpy.random import default_rng
+
+            rng = default_rng(seed=seed)
+            values = rng.normal(loc=loc, scale=scale, size=height)
+            return Series(name=name, values=values, dtype=dtype)
+        case Series() as series:
+            return normal(
+                series.len(), loc=loc, scale=scale, seed=seed, name=name, dtype=dtype
+            )
+        case DataFrame() as df:
+            return normal(
+                df.height, loc=loc, scale=scale, seed=seed, name=name, dtype=dtype
+            )
+        case _ as never:
+            assert_never(never)
+
+
+##
+
+
 @overload
 def replace_time_zone(
     obj: Series, /, *, time_zone: TimeZoneLike | None = UTC
@@ -1461,6 +1540,39 @@ class _StructFromDataClassTypeError(StructFromDataClassError):
 ##
 
 
+def uniform(
+    obj: int | Series | DataFrame,
+    /,
+    *,
+    low: float = 0.0,
+    high: float = 1.0,
+    seed: int | None = None,
+    name: str | None = None,
+    dtype: PolarsDataType = Float64,
+) -> Series:
+    """Construct a series of uniformly-distributed numbers."""
+    match obj:
+        case int() as height:
+            from numpy.random import default_rng
+
+            rng = default_rng(seed=seed)
+            values = rng.uniform(low=low, high=high, size=height)
+            return Series(name=name, values=values, dtype=dtype)
+        case Series() as series:
+            return uniform(
+                series.len(), low=low, high=high, seed=seed, name=name, dtype=dtype
+            )
+        case DataFrame() as df:
+            return uniform(
+                df.height, low=low, high=high, seed=seed, name=name, dtype=dtype
+            )
+        case _ as never:
+            assert_never(never)
+
+
+##
+
+
 def unique_element(column: ExprLike, /) -> Expr:
     """Get the unique element in a list."""
     column = ensure_expr_or_series(column)
@@ -1645,17 +1757,20 @@ __all__ = [
     "insert_after",
     "insert_before",
     "insert_between",
+    "integers",
     "is_not_null_struct_series",
     "is_null_struct_series",
     "join",
     "map_over_columns",
     "nan_sum_agg",
     "nan_sum_cols",
+    "normal",
     "replace_time_zone",
     "set_first_row_as_columns",
     "struct_dtype",
     "struct_from_dataclass",
     "touch",
+    "uniform",
     "unique_element",
     "yield_struct_series_dataclasses",
     "yield_struct_series_elements",

dycw_utilities-0.109.10/src/utilities/polars_ols.py

@@ -0,0 +1,71 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+from polars import struct
+from polars_ols import RollingKwargs, compute_rolling_least_squares
+
+from utilities.polars import ensure_expr_or_series
+
+if TYPE_CHECKING:
+    from polars import Expr
+    from polars_ols import NullPolicy
+
+    from utilities.polars import ExprLike
+
+
+def compute_rolling_ols(
+    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)
+    rolling_kwargs = RollingKwargs(
+        null_policy=null_policy,
+        window_size=window_size,
+        min_periods=min_periods,
+        use_woodbury=use_woodbury,
+        alpha=alpha,
+    )
+    coefficients = compute_rolling_least_squares(
+        target,
+        *features,
+        sample_weights=sample_weights,
+        add_intercept=add_intercept,
+        mode="coefficients",
+        rolling_kwargs=rolling_kwargs,
+    ).alias("coefficients")
+    predictions = compute_rolling_least_squares(
+        target,
+        *features,
+        sample_weights=sample_weights,
+        add_intercept=add_intercept,
+        mode="predictions",
+        rolling_kwargs=rolling_kwargs,
+    ).alias("predictions")
+    residuals = compute_rolling_least_squares(
+        target,
+        *features,
+        sample_weights=sample_weights,
+        add_intercept=add_intercept,
+        mode="residuals",
+        rolling_kwargs=rolling_kwargs,
+    ).alias("residuals")
+    ssr = (residuals**2).rolling_sum(window_size, min_samples=min_periods).alias("SSR")
+    sst = (
+        ((target - target.rolling_mean(window_size, min_samples=min_periods)) ** 2)
+        .rolling_sum(window_size, min_samples=min_periods)
+        .alias("SST")
+    )
+    r2 = (1 - ssr / sst).alias("R2")
+    return struct(coefficients, predictions, residuals, r2).alias("ols")
+
+
+__all__ = ["compute_rolling_ols"]