dycw-utilities 0.109.8__py3-none-any.whl → 0.109.10__py3-none-any.whl

{dycw_utilities-0.109.8.dist-info → dycw_utilities-0.109.10.dist-info}/METADATA

@@ -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.dist-info → dycw_utilities-0.109.10.dist-info}/RECORD

@@ -1,4 +1,4 @@
-utilities/__init__.py,sha256=HolaprMxFmj01s_l_QOeA1rUAo9639B4EnJ2_HF8s5o,60
+utilities/__init__.py,sha256=rSepLbrv6OLPYSa8JNcDoi9ph_2-kk82P4d1Ebs7rTE,61
 utilities/altair.py,sha256=NSyDsm8QlkAGmsGdxVwCkHnPxt_35yJBa9Lg7bz9Ays,9054
 utilities/astor.py,sha256=xuDUkjq0-b6fhtwjhbnebzbqQZAjMSHR1IIS5uOodVg,777
 utilities/asyncio.py,sha256=41oQUurWMvadFK5gFnaG21hMM0Vmfn2WS6OpC0R9mas,14757
@@ -45,7 +45,8 @@ utilities/pathlib.py,sha256=31WPMXdLIyXgYOMMl_HOI2wlo66MGSE-cgeelk-Lias,1410
 utilities/period.py,sha256=ikHXsWtDLr553cfH6p9mMaiCnIAP69B7q84ckWV3HaA,10884
 utilities/pickle.py,sha256=Bhvd7cZl-zQKQDFjUerqGuSKlHvnW1K2QXeU5UZibtg,657
 utilities/platform.py,sha256=NU7ycTvAXAG-fdYmDXaM1m4EOml2cGiaYwaUzfzSqyU,1767
-utilities/polars.py,sha256=USK_Rck8nmFYg2Rs-akqN9jV4w52lpz4rgkWUMQdLMk,49087
+utilities/polars.py,sha256=nB2pfK8N8HRpPE_tdbiTfFGLWC_TekAqgHlYDhnUzAM,52169
+utilities/polars_ols.py,sha256=AQe3RFOMv8CEI_ZCoscb_-PxB4JWjO0TAEmk8DKLeaI,2138
 utilities/pqdm.py,sha256=foRytQybmOQ05pjt5LF7ANyzrIa--4ScDE3T2wd31a4,3118
 utilities/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 utilities/pydantic.py,sha256=f6qtR5mO2YMuyvNmbaEj5YeD9eGA4YYfb7Bjzh9jUs0,1845
@@ -84,7 +85,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.8.dist-info/METADATA,sha256=MIDeCogv3Eg50hnzviSxygm6L6i4mPPDY7_VzeKIrR0,13004
-dycw_utilities-0.109.8.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-dycw_utilities-0.109.8.dist-info/licenses/LICENSE,sha256=gppZp16M6nSVpBbUBrNL6JuYfvKwZiKgV7XoKKsHzqo,1066
-dycw_utilities-0.109.8.dist-info/RECORD,,
+dycw_utilities-0.109.10.dist-info/METADATA,sha256=dn1-Yy6KJLMDkEyWztQtPMjCF7DKg5bjdo7wRbZA2oA,13038
+dycw_utilities-0.109.10.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+dycw_utilities-0.109.10.dist-info/licenses/LICENSE,sha256=gppZp16M6nSVpBbUBrNL6JuYfvKwZiKgV7XoKKsHzqo,1066
+dycw_utilities-0.109.10.dist-info/RECORD,,

utilities/__init__.py

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

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

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