core-lens 0.1.dev186__tar.gz → 0.1.dev188__tar.gz

{core_lens-0.1.dev186 → core_lens-0.1.dev188}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.5
 Name: core-lens
-Version: 0.1.dev186
+Version: 0.1.dev188
 Summary: Query, analyse, and visualise CoreStack's microwatershed and Earth science data through a clean, composable Python API.
 Project-URL: Homepage, https://github.com/ApoorvaKashyap/core-lens
 Project-URL: Issues, https://github.com/ApoorvaKashyap/core-lens/issues

{core_lens-0.1.dev186 → core_lens-0.1.dev188}/src/core_lens/_version.py

@@ -18,7 +18,7 @@ version_tuple: tuple[int | str, ...]
 commit_id: str | None
 __commit_id__: str | None
 
-__version__ = version = '0.1.dev186'
-__version_tuple__ = version_tuple = (0, 1, 'dev186')
+__version__ = version = '0.1.dev188'
+__version_tuple__ = version_tuple = (0, 1, 'dev188')
 
 __commit_id__ = commit_id = None

{core_lens-0.1.dev186 → core_lens-0.1.dev188}/src/core_lens/base/namespaces/stats.py

@@ -11,6 +11,7 @@ import numpy as np
 import polars as pl
 
 from core_lens.utils.polars_utils import collect_lf, cached_read_schema
+from core_lens.base.result import Result
 
 if TYPE_CHECKING:
     from core_lens.base.result import Result
@@ -157,7 +158,7 @@ class StatsNamespace:
     Computed values always go in ``data``; method parameters go in ``metadata``.
     """
 
-    def __init__(self, result: "Result") -> None:
+    def __init__(self, result: Result) -> None:
         """Initialize StatsNamespace.
 
         Args:
@@ -178,7 +179,7 @@ class StatsNamespace:
         self,
         columns: list[str] | None = None,
         by: str = "column",
-    ) -> "Result":
+    ) -> Result:
         r"""Per-column or per-entity descriptive statistics.
 
         Uses polars' in-built methods for mean, std, min, max, quantiles etc.
@@ -238,26 +239,80 @@ class StatsNamespace:
         columns: list[str],
         method: CorrelateMethod = CorrelateMethod.PEARSON,
         across: str = "entity",
+        group_by: str | None = None,
+        min_obs: int = 3,
     ) -> "Result":
-        """Pairwise correlations between columns.
+        """Compute pairwise correlations between columns.
+
+        Correlations are computed with ``scipy.stats`` and returned as a
+        :class:`~core_lens.base.result.Result` whose ``data`` is a Polars
+        ``DataFrame``. Two modes are supported:
+
+        * **Pooled** (``group_by is None``): all rows are treated as one
+        observation set and a single correlation is computed per column pair.
+        * **Per-entity** (``group_by`` set): the result is grouped by the given
+        key column (e.g. ``"mws_id"``) and one correlation is computed per
+        entity across its remaining rows (typically its time series). This is
+        what enables questions such as "which MWS show the strongest
+        dependence of cropping intensity on annual rainfall?".
 
         Args:
-            columns (list[str]): At least 2 column names.
-            method (CorrelateMethod, optional): A :class:`~core_lens.base.namespaces.stats.CorrelateMethod` enum value.
-            across (str, optional): ``"entity"`` or ``"time"`` — recorded in metadata only.
+            columns: Column names to correlate. Must contain at least two
+                entries. In pooled mode every pair is computed; in per-entity
+                mode the first two columns are used as the pair.
+            method: The correlation coefficient to compute. One of
+                :class:`CorrelateMethod` (``PEARSON``, ``SPEARMAN``, ``KENDALL``).
+            across: ``"entity"`` or ``"time"``. Recorded in result metadata to
+                describe the intended axis of the relationship.
+            group_by: Optional key column to compute one correlation per group
+                (e.g. ``"mws_id"``). When ``None``, the pooled behaviour is
+                preserved for backward compatibility.
+            min_obs: Minimum number of non-null observations required per group
+                (per-entity mode) or overall (pooled mode) before a correlation
+                is computed. Groups below this threshold are skipped.
 
         Returns:
-            Result: Result whose data has columns
-            ``column_a | column_b | correlation | p_value``.
+            Result: A :class:`~core_lens.base.result.Result` whose ``data``
+                DataFrame has the following columns:
+
+                * Pooled mode: ``column_a | column_b | correlation | p_value``
+                * Per-entity mode: ``{group_by} | correlation | p_value | n_obs``
+
+                Geometry is dropped (``has_geometry=False``).
 
         Raises:
-            CorrelationError: If fewer than 2 columns supplied.
-            ValueError: If ``method`` is not recognised.
+            CorrelationError: If fewer than two columns are supplied.
+            ValueError: If ``method`` is not a valid :class:`CorrelateMethod`,
+                or if ``group_by`` is set but the column is not present in the
+                underlying data.
 
-        Under the hood:
-            Calls ``scipy.stats.pearsonr``, ``scipy.stats.spearmanr``, or ``scipy.stats.kendalltau``
-            to compute the corresponding correlation coefficients and p-values.
+        Examples:
+            Pooled correlation across all MWS in the AoI::
 
+                result = aoi.mws.annual.stats.correlate(
+                    columns=["dw_precipitation", "ci_cropping_intensity"],
+                    method=CorrelateMethod.PEARSON,
+                )
+                print(result.data.collect())
+
+            Per-MWS correlation, ranked by strength of dependence::
+
+                top5 = (
+                    aoi.mws.annual.stats
+                    .correlate(
+                        columns=["dw_precipitation", "ci_cropping_intensity"],
+                        method=CorrelateMethod.PEARSON,
+                        group_by="mws_id",
+                        min_obs=3,
+                    )
+                    .data
+                    .with_columns(pl.col("correlation").abs().alias("abs_corr"))
+                    .sort("abs_corr", descending=True)
+                    .select("mws_id", "correlation", "p_value", "n_obs")
+                    .limit(5)
+                    .collect()
+                )
+                print(top5)
         """
         if len(columns) < 2:
             raise CorrelationError(
@@ -265,27 +320,124 @@ class StatsNamespace:
             )
         if not isinstance(method, CorrelateMethod):
             raise ValueError(
-                f"StatsNamespace.correlate: method must be a CorrelateMethod. Valid options: {[e.name for e in CorrelateMethod]}."
+                f"StatsNamespace.correlate: method must be a CorrelateMethod. "
+                f"Valid options: {[e.name for e in CorrelateMethod]}."
             )
 
         import scipy.stats as sp
 
-        df = self._r.df()
+        lf = (
+            self._r.lazy()
+        )  # LazyFrame — stays lazy until an explicit collect boundary below
+
+        def _t_pvals(corr_col: "pl.Expr", n_col: "pl.Expr") -> "pl.Expr":
+            r_safe = corr_col.clip(-0.9999999, 0.9999999)
+            return (corr_col * ((n_col - 2) / (1 - r_safe**2)).sqrt()).alias("t_stat")
+
+        # Per-entity mode
+        if group_by is not None:
+            # Schema check only — does NOT execute the query plan.
+            if group_by not in lf.collect_schema().names():
+                raise ValueError(
+                    f"StatsNamespace.correlate: group_by column '{group_by}' "
+                    f"not found in data. Available: {lf.collect_schema().names()}"
+                )
+
+            col_a, col_b = columns[0], columns[1]
+
+            if method in (CorrelateMethod.PEARSON, CorrelateMethod.SPEARMAN):
+                # Fully lazy path: rank (if Spearman) -> group_by -> agg -> filter.
+                # No collect until the very end, where scipy needs numpy anyway.
+                if method is CorrelateMethod.SPEARMAN:
+                    src = lf.with_columns(
+                        pl.col(col_a)
+                        .rank(method="average")
+                        .over(group_by)
+                        .alias("__a"),
+                        pl.col(col_b)
+                        .rank(method="average")
+                        .over(group_by)
+                        .alias("__b"),
+                    )
+                    a_col, b_col = "__a", "__b"
+                else:
+                    src, a_col, b_col = lf, col_a, col_b
+
+                lazy_result = (
+                    src.group_by(group_by)
+                    .agg(
+                        pl.corr(pl.col(a_col), pl.col(b_col), method="pearson").alias(
+                            "correlation"
+                        ),
+                        pl.len().alias("n_obs"),
+                    )
+                    .filter(pl.col("correlation").is_finite())
+                    .filter(pl.col("n_obs") >= min_obs)
+                    .with_columns(_t_pvals(pl.col("correlation"), pl.col("n_obs")))
+                    .filter(pl.col("t_stat").is_finite())
+                )
+
+                # --- single collect boundary: scipy.t.sf needs eager numpy ---
+                result = lazy_result.collect()
+                t_vals = result["t_stat"].to_numpy()
+                df_vals = result["n_obs"].to_numpy() - 2
+                p_vals = 2.0 * sp.t.sf(np.abs(t_vals), df_vals)
+                result = result.with_columns(pl.Series("p_value", p_vals))
+
+                data = result.select([group_by, "correlation", "p_value", "n_obs"])
+
+            else:
+                # Kendall: scipy has no batched/lazy equivalent, so this branch
+                # is inherently eager. Collect ONCE here — not scattered filters.
+                df = lf.collect()
+                rows: list[dict[str, Any]] = []
+                for sub_full in df.partition_by(group_by):
+                    key = sub_full[group_by][0]
+                    sub = sub_full.select([col_a, col_b]).drop_nulls()
+                    a = sub[col_a].to_numpy().astype(float)
+                    b = sub[col_b].to_numpy().astype(float)
+                    n = len(a)
+                    if n < min_obs or a.std() == 0 or b.std() == 0:
+                        continue
+                    corr, pval = sp.kendalltau(a, b)
+                    rows.append(
+                        {
+                            group_by: key,
+                            "correlation": float(cast(float, corr)),
+                            "p_value": float(cast(float, pval)),
+                            "n_obs": n,
+                        }
+                    )
+                data = pl.DataFrame(rows)
+
+            metadata: dict[str, Any] = {
+                "method": method.value,
+                "columns": [col_a, col_b],
+                "across": across,
+                "group_by": group_by,
+                "min_obs": min_obs,
+                "n_entities_computed": len(data),
+            }
+            return self._r._replace(data=data, has_geometry=False, metadata=metadata)
+
+        # Pooled mode — eager needed regardless (scipy.pearsonr/spearmanr/
+        # kendalltau all take raw numpy arrays). Collect ONCE here.
+        df = lf.collect()
         n_obs = len(df)
-        rows: list[dict[str, Any]] = []
 
+        rows = []
         for col_a, col_b in combinations(columns, 2):
             sub = df.select([col_a, col_b]).drop_nulls()
             a = sub[col_a].to_numpy().astype(float)
             b = sub[col_b].to_numpy().astype(float)
-
+            if len(a) < min_obs or a.std() == 0 or b.std() == 0:
+                continue
             if method is CorrelateMethod.PEARSON:
                 corr, pval = sp.pearsonr(a, b)
             elif method is CorrelateMethod.SPEARMAN:
                 corr, pval = sp.spearmanr(a, b)
             else:
                 corr, pval = sp.kendalltau(a, b)
-
             rows.append(
                 {
                     "column_a": col_a,
@@ -296,8 +448,8 @@ class StatsNamespace:
             )
 
         data = pl.DataFrame(rows)
-        metadata: dict[str, Any] = {
-            "method": method.value if method is not None else None,
+        metadata = {
+            "method": method.value,
             "columns": columns,
             "across": across,
             "n_observations": n_obs,
@@ -312,7 +464,7 @@ class StatsNamespace:
         against: float | None = None,
         method: TestMethod | None = None,
         significance_level: float = 0.05,
-    ) -> "Result":
+    ) -> Result:
         """Hypothesis test in three modes: group-based, period-based, single-sample.
 
         Args:
@@ -470,7 +622,7 @@ class StatsNamespace:
         from_period: int,
         to_period: int,
         method: ChangeMethod = ChangeMethod.ABSOLUTE,
-    ) -> "Result":
+    ) -> Result:
         """Change between two time periods per entity.
 
         Args:
@@ -596,7 +748,7 @@ class StatsNamespace:
         method: AnomalyCrossMethod | AnomalyTsMethod,
         baseline: tuple[int, int] | None = None,
         threshold: float = 2.0,
-    ) -> "Result":
+    ) -> Result:
         """Anomaly detection in cross-sectional or timeseries mode.
 
         Args:
@@ -863,7 +1015,7 @@ class StatsNamespace:
         columns: dict[str, Any],
         method: SimilarityMethod = SimilarityMethod.EUCLIDEAN,
         top_n: int = 10,
-    ) -> "Result":
+    ) -> Result:
         """Find entities most similar to ``target`` across ``columns``.
 
         Args:

{core_lens-0.1.dev186 → core_lens-0.1.dev188}/uv.lock

@@ -1409,14 +1409,14 @@ wheels = [
 
 [[package]]
 name = "gitpython"
-version = "3.1.58"
+version = "3.1.62"
 source = { registry = "https://pypi.org/simple" }
 dependencies = [
     { name = "gitdb" },
 ]
-sdist = { url = "https://files.pythonhosted.org/packages/26/d6/5f358ff283325580c2003a6d953aea18cfe10ae87b46f5ebc80fa3a386dc/gitpython-3.1.58.tar.gz", hash = "sha256:621416df10ef3fd0e19fabf9172ddeed0fa704d353d04f194eec56a625a95b22", size = 228498, upload-time = "2026-08-04T15:05:49.47Z" }
+sdist = { url = "https://files.pythonhosted.org/packages/e0/db/3ca813cbacb23ab6fe46ff38a9b5ef8e73e970c8051f2ce903aacafe0446/gitpython-3.1.62.tar.gz", hash = "sha256:1791de66309bc0c7cfca40bf8d2e3de7ca091cbf94e6051be1ad0722c61062af", size = 231728, upload-time = "2026-09-07T02:57:21.155Z" }
 wheels = [
-    { url = "https://files.pythonhosted.org/packages/ec/0c/9d8752098bc442f0726e64aa6135940b3a96809915d1aa4206c1bb97881d/gitpython-3.1.58-py3-none-any.whl", hash = "sha256:d331e722577f0fd7fc1f857419b3ecc07af66282b933d2a4d95f84a042fdd50f", size = 220183, upload-time = "2026-08-04T15:05:48.025Z" },
+    { url = "https://files.pythonhosted.org/packages/d6/0b/29d7965215f8ef830a7ca1f42997fe13e5693d85e9edb18f938d063ef5f2/gitpython-3.1.62-py3-none-any.whl", hash = "sha256:7002251225e10e29d2e1f49e6532613fe5d5d9f0b6f1f02997a52b38fe56899e", size = 222753, upload-time = "2026-09-07T02:57:19.762Z" },
 ]
 
 [[package]]