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

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

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

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

@@ -56,6 +56,7 @@ dev = [
   "scipy >= 1.15.2, < 1.16",
   "slack-sdk >= 3.35.0, < 3.36",
   "sqlalchemy >= 2.0.40, < 2.1",
+  "statsmodels >= 0.14.4, < 0.15",
   "streamlit >= 1.44.1, < 1.45",
   "tenacity >= 8.5.0, < 9.0",  # limited by luigi
   "tomlkit >= 0.13.2, < 0.14",
@@ -92,7 +93,7 @@ dependencies = [
 name = "dycw-utilities"
 readme = "README.md"
 requires-python = ">= 3.12"
-version = "0.109.19"
+version = "0.109.21"
 
 [project.optional-dependencies]
 test = [
@@ -335,7 +336,7 @@ zzz-test-zoneinfo = [
 # bump-my-version
 [tool.bumpversion]
 allow_dirty = true
-current_version = "0.109.19"
+current_version = "0.109.21"
 
 [[tool.bumpversion.files]]
 filename = "src/utilities/__init__.py"

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

@@ -43,6 +43,7 @@ from polars import (
     Series,
     String,
     Struct,
+    UInt32,
     col,
     date_range,
     datetime_range,
@@ -107,6 +108,7 @@ from utilities.polars import (
     _InsertBetweenMissingColumnsError,
     _InsertBetweenNonConsecutiveError,
     _yield_struct_series_element_remove_nulls,
+    acf,
     adjust_frequencies,
     append_dataclass,
     are_frames_equal,
@@ -168,6 +170,59 @@ if TYPE_CHECKING:
     from utilities.types import MaybeType, StrMapping, WeekDay
 
 
+class TestACF:
+    def test_main(self) -> None:
+        series = Series(linspace(0, 2 * pi, 1000))
+        df = acf(series)
+        check_polars_dataframe(
+            df, height=31, schema_list={"lag": UInt32, "autocorrelation": Float64}
+        )
+
+    def test_alpha(self) -> None:
+        series = Series(linspace(0, 2 * pi, 1000))
+        df = acf(series, alpha=0.5)
+        check_polars_dataframe(
+            df,
+            height=31,
+            schema_list={
+                "lag": UInt32,
+                "autocorrelation": Float64,
+                "lower": Float64,
+                "upper": Float64,
+            },
+        )
+
+    def test_qstat(self) -> None:
+        series = Series(linspace(0, 2 * pi, 1000))
+        df = acf(series, qstat=True)
+        check_polars_dataframe(
+            df,
+            height=31,
+            schema_list={
+                "lag": UInt32,
+                "autocorrelation": Float64,
+                "qstat": Float64,
+                "pvalue": Float64,
+            },
+        )
+
+    def test_alpha_and_qstat(self) -> None:
+        series = Series(linspace(0, 2 * pi, 1000))
+        df = acf(series, alpha=0.5, qstat=True)
+        check_polars_dataframe(
+            df,
+            height=31,
+            schema_list={
+                "lag": UInt32,
+                "autocorrelation": Float64,
+                "lower": Float64,
+                "upper": Float64,
+                "qstat": Float64,
+                "pvalue": Float64,
+            },
+        )
+
+
 class TestAdjustFrequencies:
     def test_main(self) -> None:
         n = 1000

dycw_utilities-0.109.21/src/tests/test_statsmodel.py

@@ -0,0 +1,44 @@
+from __future__ import annotations
+
+from numpy import allclose, isclose, linspace, pi
+
+from utilities.statsmodels import acf
+
+
+class TestACF:
+    def test_main(self) -> None:
+        x = linspace(0, 2 * pi, 1000)
+        acfs = acf(x)
+        assert acfs.shape == (31,)
+        assert isclose(acfs[-1], 0.9100539400539398)
+
+    def test_alpha(self) -> None:
+        x = linspace(0, 2 * pi, 1000)
+        result = acf(x, alpha=0.5)
+        acfs, confint = result
+        assert acfs.shape == (31,)
+        assert confint.shape == (31, 2)
+        assert allclose(confint[-1, :], [0.7534104755915192, 1.0666974045163604])
+
+    def test_qstat(self) -> None:
+        x = linspace(0, 2 * pi, 1000)
+        result = acf(x, qstat=True)
+        acfs, qstat, pvalues = result
+        assert acfs.shape == (31,)
+        assert qstat.shape == (30,)
+        assert isclose(qstat[-1], 27769.955439016478)
+        assert pvalues.shape == (30,)
+        assert isclose(pvalues[-1], 0.0)
+
+    def test_alpha_and_qstat(self) -> None:
+        x = linspace(0, 2 * pi, 1000)
+        result = acf(x, alpha=0.5, qstat=True)
+        acfs, confint, qstat, pvalues = result
+        assert acfs.shape == (31,)
+        assert isclose(acfs[-1], 0.9100539400539398)
+        assert confint.shape == (31, 2)
+        assert allclose(confint[-1, :], [0.7534104755915192, 1.0666974045163604])
+        assert qstat.shape == (30,)
+        assert isclose(qstat[-1], 27769.955439016478)
+        assert pvalues.shape == (30,)
+        assert isclose(pvalues[-1], 0.0)

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

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

{dycw_utilities-0.109.19 → dycw_utilities-0.109.21}/src/utilities/lightweight_charts.py

@@ -13,7 +13,7 @@ if TYPE_CHECKING:
     from lightweight_charts import AbstractChart, Chart
     from lightweight_charts.abstract import SeriesCommon
     from polars import DataFrame
-    from polars._typing import SchemaDict
+    from polars._typing import SchemaDict  # pyright: ignore[reportPrivateImportUsage]
 
     from utilities.types import PathLike
 

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

@@ -121,6 +121,7 @@ if TYPE_CHECKING:
     )
 
     from utilities.numpy import NDArrayB, NDArrayF
+    from utilities.statsmodels import ACFMissing
     from utilities.types import (
         Dataclass,
         MaybeIterable,
@@ -143,6 +144,85 @@ _FINITE_EWM_MIN_WEIGHT = 0.9999
 ##
 
 
+def acf(
+    series: Series,
+    /,
+    *,
+    adjusted: bool = False,
+    nlags: int | None = None,
+    qstat: bool = False,
+    fft: bool = True,
+    alpha: float | None = None,
+    bartlett_confint: bool = True,
+    missing: ACFMissing = "none",
+) -> DataFrame:
+    """Compute the autocorrelations of a series."""
+    from numpy import ndarray
+
+    import utilities.statsmodels
+
+    array = series.to_numpy()
+    result = utilities.statsmodels.acf(
+        array,
+        adjusted=adjusted,
+        nlags=nlags,
+        qstat=qstat,
+        fft=fft,
+        alpha=alpha,
+        bartlett_confint=bartlett_confint,
+        missing=missing,
+    )
+    match result:
+        case ndarray() as acfs:
+            return _acf_process_acfs(acfs)
+        case ndarray() as acfs, ndarray() as confints:
+            df_acfs = _acf_process_acfs(acfs)
+            df_confints = _acf_process_confints(confints)
+            return df_acfs.join(df_confints, on=["lag"])
+        case ndarray() as acfs, ndarray() as qstats, ndarray() as pvalues:
+            df_acfs = _acf_process_acfs(acfs)
+            df_qstats_pvalues = _acf_process_qstats_pvalues(qstats, pvalues)
+            return df_acfs.join(df_qstats_pvalues, on=["lag"], how="left")
+        case (
+            ndarray() as acfs,
+            ndarray() as confints,
+            ndarray() as qstats,
+            ndarray() as pvalues,
+        ):
+            df_acfs = _acf_process_acfs(acfs)
+            df_confints = _acf_process_confints(confints)
+            df_qstats_pvalues = _acf_process_qstats_pvalues(qstats, pvalues)
+            return join(df_acfs, df_confints, df_qstats_pvalues, on=["lag"], how="left")
+        case _ as never:
+            assert_never(never)
+
+
+def _acf_process_acfs(acfs: NDArrayF, /) -> DataFrame:
+    return (
+        Series(name="autocorrelation", values=acfs, dtype=Float64)
+        .to_frame()
+        .with_row_index(name="lag")
+    )
+
+
+def _acf_process_confints(confints: NDArrayF, /) -> DataFrame:
+    return DataFrame(
+        data=confints, schema={"lower": Float64, "upper": Float64}
+    ).with_row_index(name="lag")
+
+
+def _acf_process_qstats_pvalues(qstats: NDArrayF, pvalues: NDArrayF, /) -> DataFrame:
+    from numpy import hstack
+
+    data = hstack([qstats.reshape(-1, 1), pvalues.reshape(-1, 1)])
+    return DataFrame(
+        data=data, schema={"qstat": Float64, "pvalue": Float64}
+    ).with_row_index(name="lag", offset=1)
+
+
+##
+
+
 def adjust_frequencies(
     series: Series,
     /,
@@ -1840,6 +1920,7 @@ __all__ = [
     "SetFirstRowAsColumnsError",
     "StructFromDataClassError",
     "YieldStructSeriesElementsError",
+    "acf",
     "adjust_frequencies",
     "append_dataclass",
     "are_frames_equal",

dycw_utilities-0.109.21/src/utilities/statsmodels.py

@@ -0,0 +1,117 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any, Literal, cast, overload
+
+import statsmodels.tsa.stattools
+
+if TYPE_CHECKING:
+    from utilities.numpy import NDArrayF
+
+
+type ACFMissing = Literal["none", "raise", "conservative", "drop"]
+
+
+@overload
+def acf(
+    array: NDArrayF,
+    /,
+    *,
+    adjusted: bool = False,
+    nlags: int | None = None,
+    qstat: Literal[False] = False,
+    fft: bool = True,
+    alpha: None = None,
+    bartlett_confint: bool = True,
+    missing: ACFMissing = "none",
+) -> NDArrayF: ...
+@overload
+def acf(
+    array: NDArrayF,
+    /,
+    *,
+    adjusted: bool = False,
+    nlags: int | None = None,
+    qstat: Literal[False] = False,
+    fft: bool = True,
+    alpha: float,
+    bartlett_confint: bool = True,
+    missing: ACFMissing = "none",
+) -> tuple[NDArrayF, NDArrayF]: ...
+@overload
+def acf(
+    array: NDArrayF,
+    /,
+    *,
+    adjusted: bool = False,
+    nlags: int | None = None,
+    qstat: Literal[True],
+    fft: bool = True,
+    alpha: float,
+    bartlett_confint: bool = True,
+    missing: ACFMissing = "none",
+) -> tuple[NDArrayF, NDArrayF, NDArrayF, NDArrayF]: ...
+@overload
+def acf(
+    array: NDArrayF,
+    /,
+    *,
+    adjusted: bool = False,
+    nlags: int | None = None,
+    qstat: Literal[True],
+    fft: bool = True,
+    alpha: None = None,
+    bartlett_confint: bool = True,
+    missing: ACFMissing = "none",
+) -> tuple[NDArrayF, NDArrayF, NDArrayF]: ...
+@overload
+def acf(
+    array: NDArrayF,
+    /,
+    *,
+    adjusted: bool = False,
+    nlags: int | None = None,
+    qstat: bool = False,
+    fft: bool = True,
+    alpha: float | None = None,
+    bartlett_confint: bool = True,
+    missing: ACFMissing = "none",
+) -> (
+    NDArrayF
+    | tuple[NDArrayF, NDArrayF]
+    | tuple[NDArrayF, NDArrayF, NDArrayF]
+    | tuple[NDArrayF, NDArrayF, NDArrayF, NDArrayF]
+): ...
+def acf(
+    array: NDArrayF,
+    /,
+    *,
+    adjusted: bool = False,
+    nlags: int | None = None,
+    qstat: bool = False,
+    fft: bool = True,
+    alpha: float | None = None,
+    bartlett_confint: bool = True,
+    missing: ACFMissing = "none",
+) -> (
+    NDArrayF
+    | tuple[NDArrayF, NDArrayF]
+    | tuple[NDArrayF, NDArrayF, NDArrayF]
+    | tuple[NDArrayF, NDArrayF, NDArrayF, NDArrayF]
+):
+    """Typed version of `acf`."""
+    return cast(
+        "Any",
+        statsmodels.tsa.stattools.acf(
+            array,
+            adjusted=adjusted,
+            nlags=nlags,
+            qstat=qstat,
+            fft=fft,
+            alpha=alpha,
+            bartlett_confint=bartlett_confint,
+            missing=missing,
+        ),
+    )
+
+
+__all__ = ["ACFMissing", "acf"]