Experiment with polars

.github/workflows/notebooks_test.yml

@@ -8,6 +8,8 @@ on:
 
 jobs:
   build:
+    # disable this job for now
+    if : false
 
     runs-on: ubuntu-latest
 

.gitignore

@@ -145,4 +145,6 @@ notebooks/sandbox.ipynb
 streamlit_app.py
 notes.md
 
-notebooks/Untitled.ipynb
+notebooks/Untitled.ipynb
+
+uv.lock

modelskill/comparison/_collection.py

@@ -1,4 +1,5 @@
 from __future__ import annotations
+from collections import defaultdict
 from copy import deepcopy
 import os
 from pathlib import Path
@@ -21,9 +22,9 @@
 import zipfile
 import numpy as np
 import pandas as pd
+import polars as pl
 
 
-from .. import metrics as mtr
 
 from ._collection_plotter import ComparerCollectionPlotter
 from ..skill import SkillTable
@@ -449,36 +450,50 @@ def skill(
         cc = self
 
         pmetrics = _parse_metric(metrics)
+        # TODO don't use hardcoded metrics
+        # pmetrics = ["n", "bias", "rmse", "urmse", "mae"]
 
         agg_cols = _parse_groupby(by, n_mod=cc.n_models, n_qnt=cc.n_quantities)
         agg_cols, attrs_keys = self._attrs_keys_in_by(agg_cols)
 
         df = cc._to_long_dataframe(attrs_keys=attrs_keys, observed=observed)
+        assert "model" in df.columns
 
         res = _groupby_df(df, by=agg_cols, metrics=pmetrics)
-        mtr_cols = [m.__name__ for m in pmetrics]  # type: ignore
-        res = res.dropna(subset=mtr_cols, how="all")  # TODO: ok to remove empty?
         res = self._append_xy_to_res(res, cc)
-        res = cc._add_as_col_if_not_in_index(df, skilldf=res)  # type: ignore
+
+        # TODO this should not be necessary
+        if "model" not in res.columns:
+            res = res.with_columns(pl.lit(cc.mod_names[0]).alias("model"))
         return SkillTable(res)
 
     def _to_long_dataframe(
         self, attrs_keys: Iterable[str] | None = None, observed: bool = False
-    ) -> pd.DataFrame:
+    ) -> pl.DataFrame:
         """Return a copy of the data as a long-format pandas DataFrame (for groupby operations)"""
         frames = []
         for cmp in self:
             frame = cmp._to_long_dataframe(attrs_keys=attrs_keys)
+
             if self.n_quantities > 1:
-                frame["quantity"] = cmp.quantity.name
+                frame = frame.with_columns(pl.lit(cmp.quantity.name).alias("quantity"))
             frames.append(frame)
-        res = pd.concat(frames)
 
-        cat_cols = res.select_dtypes(include=["object"]).columns
-        res[cat_cols] = res[cat_cols].astype("category")
+        # convert all ints and floats to f64
+        import polars.selectors as cs
+
+        # frames = [
+        #    df.with_columns([pl.col(pl.Float32, pl.Int32, pl.Int64).cast(pl.Float64)])
+        #    for df in frames
+        # ]
+        frames = [df.with_columns(cs.numeric().cast(pl.Float64)) for df in frames]
+
+        # TODO why doesn't all frames have the same columns?
+        res = pl.concat(frames, how="diagonal")
+        if observed and attrs_keys is not None:
+            # res = res.loc[~(res == False).any(axis=1)]  # noqa
+            res = res.filter(pl.col(attrs_keys).is_not_null())
 
-        if observed:
-            res = res.loc[~(res == False).any(axis=1)]  # noqa
         return res
 
     @staticmethod
@@ -495,21 +510,24 @@ def _attrs_keys_in_by(by: List[str | pd.Grouper]) -> Tuple[List[str], List[str]]
         return agg_cols, attrs_keys
 
     @staticmethod
-    def _append_xy_to_res(res: pd.DataFrame, cc: ComparerCollection) -> pd.DataFrame:
+    def _append_xy_to_res(res: pl.DataFrame, cc: ComparerCollection) -> pl.DataFrame:
         """skill() helper: Append x and y to res if possible"""
-        res["x"] = np.nan
-        res["y"] = np.nan
-
-        # for MultiIndex in res find "observation" level and
-        # insert x, y if gtype=point for that observation
-        if "observation" in res.index.names:
-            idx_names = res.index.names
-            res = res.reset_index()
-            for cmp in cc:
-                if cmp.gtype == "point":
-                    res.loc[res.observation == cmp.name, "x"] = cmp.x
-                    res.loc[res.observation == cmp.name, "y"] = cmp.y
-            res = res.set_index(idx_names)
+
+        if "observation" not in res.columns:
+            return res
+        xs = defaultdict(lambda: np.nan)
+        ys = defaultdict(lambda: np.nan)
+        for cmp in cc:
+            if cmp.gtype == "point":
+                xs[cmp.name] = float(cmp.x)
+                ys[cmp.name] = float(cmp.y)
+
+        # add x and y to res based on observation name based on the xs and ys dicts
+        res = res.with_columns(
+            pl.col("observation").map_elements(lambda name: xs.get(name)).alias("x"),
+            pl.col("observation").map_elements(lambda name: ys.get(name)).alias("y"),
+        )
+
         return res
 
     def _add_as_col_if_not_in_index(
@@ -602,6 +620,9 @@ def gridded_skill(
 
         metrics = _parse_metric(metrics)
 
+        # TODO avoid hardcoded metrics
+        # metrics = ["n", "bias", "rmse", "mae"]
+
         df = cmp._to_long_dataframe()
         df = _add_spatial_grid_to_df(df=df, bins=bins, binsize=binsize)
 
@@ -611,13 +632,17 @@ def gridded_skill(
         if "y" not in agg_cols:
             agg_cols.insert(0, "y")
 
-        df = df.drop(columns=["x", "y"]).rename(columns=dict(xBin="x", yBin="y"))
-        res = _groupby_df(df, by=agg_cols, metrics=metrics, n_min=n_min)
-        ds = res.to_xarray().squeeze()
+        df = df.drop(["x", "y"]).rename(dict(xBin="x", yBin="y"))
+        res = _groupby_df(df, by=agg_cols, metrics=metrics, n_min=n_min)  # type: ignore
+
+        potential_cols = ["x", "y", "model", "observation"]
+        cols = [c for c in potential_cols if c in res.columns]
+
+        ds = res.to_pandas().set_index(cols).to_xarray().squeeze()
 
         # change categorial index to coordinates
-        for dim in ("x", "y"):
-            ds[dim] = ds[dim].astype(float)
+        # for dim in ("x", "y"):
+        #    ds[dim] = ds[dim].astype(float)
         return SkillGrid(ds)
 
     def mean_skill(
@@ -675,39 +700,66 @@ def mean_skill(
         cc = self
 
         df = cc._to_long_dataframe()  # TODO: remove
+        assert "model" in df.columns
         mod_names = cc.mod_names
         # obs_names = cmp.obs_names  # df.observation.unique()
         qnt_names = cc.quantity_names
 
         # skill assessment
         pmetrics = _parse_metric(metrics)
+
+        # TODO avoid hardcoded metrics
+        # pmetrics = ["n", "bias", "rmse", "mae"]
         sk = cc.skill(metrics=pmetrics)
         if sk is None:
             return None
         skilldf = sk.to_dataframe()
+        assert "model" in skilldf.columns
 
         # weights
         weights = cc._parse_weights(weights, sk.obs_names)
-        skilldf["weights"] = (
-            skilldf.n if weights is None else np.tile(weights, len(mod_names))  # type: ignore
-        )
+        # skilldf["weights"] = (
+        #    skilldf.n if weights is None else np.tile(weights, len(mod_names))  # type: ignore
+        # )
+        if weights is None:
+            skilldf = skilldf.with_columns(pl.col("n").alias("weights"))
+        else:
+            wdict = {o: w for o, w in zip(sk.obs_names, weights)}
+            skilldf = skilldf.with_columns(
+                pl.col("observation")
+                .map_elements(lambda name: wdict.get(name))
+                .alias("weights")
+            )
 
-        def weighted_mean(x: Any) -> Any:
-            return np.average(x, weights=skilldf.loc[x.index, "weights"])
+        # def weighted_mean(x: Any) -> Any:
+        #    return np.average(x, weights=skilldf.loc[x.index, "weights"])
+        # weighted_mean = pl.
 
         # group by
         by = cc._mean_skill_by(skilldf, mod_names, qnt_names)  # type: ignore
-        agg = {"n": "sum"}
-        for metric in pmetrics:  # type: ignore
-            agg[metric.__name__] = weighted_mean  # type: ignore
-        res = skilldf.groupby(by, observed=False).agg(agg)
+        # agg = {"n": "sum"}
+        # for metric in pmetrics:  # type: ignore
+        #    agg[metric.__name__] = weighted_mean  # type: ignore
+        # (pl.col("values") * pl.col("weights")).sum() / pl.col("weights").sum()).alias("weighted_mean")
+        exprs = [
+            (pl.col(metric) * pl.col("weights")).sum() / pl.col("weights").sum()
+            for metric in pmetrics
+        ]
+        # res = skilldf.groupby(by, observed=False).agg(agg)
+
+        # Numpy ufuncs are supported, e.g.
+        # df.select(np.log(pl.all()).name.suffix("_log"))
+        res = skilldf.group_by(by).agg(exprs)
+
+        # res = skilldf.group_by(by).agg(np.average(pl.all(), weights=pl.col("weights")))
 
         # TODO is this correct?
-        res.index.name = "model"
+        # res.index.name = "model"
 
         # output
-        res = cc._add_as_col_if_not_in_index(df, res, fields=["model", "quantity"])  # type: ignore
-        return SkillTable(res.astype({"n": int}))
+        # res = cc._add_as_col_if_not_in_index(df, res, fields=["model", "quantity"])  # type: ignore
+        return SkillTable(res)
+        # return SkillTable(res.astype({"n": int}))
 
     # def mean_skill_points(
     #     self,
@@ -760,7 +812,8 @@ def weighted_mean(x: Any) -> Any:
     #     return cmp.skill(metrics=metrics)  # NOT CORRECT - SEE ABOVE
 
     def _mean_skill_by(self, skilldf, mod_names, qnt_names):  # type: ignore
-        by = []
+        # TODO clean up this mess
+        by = ["model"]
         if len(mod_names) > 1:
             by.append("model")
         if len(qnt_names) > 1:
@@ -771,9 +824,11 @@ def _mean_skill_by(self, skilldf, mod_names, qnt_names):  # type: ignore
             elif "model" in skilldf:
                 by.append("model")
             else:
-                by = [mod_names[0]] * len(skilldf)
-        return by
+                # by = [mod_names[0]] * len(skilldf)
+                by = None
+        return list(set(by))
 
+    # TODO add useful type hints
     def _parse_weights(self, weights: Any, observations: Any) -> Any:
         if observations is None:
             observations = self.obs_names
@@ -819,7 +874,7 @@ def _parse_weights(self, weights: Any, observations: Any) -> Any:
 
     def score(
         self,
-        metric: str | Callable = mtr.rmse,
+        metric: str | Callable = "rmse",
         **kwargs: Any,
     ) -> Dict[str, float]:
         """Weighted mean score of model(s) over all observations
@@ -878,21 +933,15 @@ def score(
         if not (callable(metric) or isinstance(metric, str)):
             raise ValueError("metric must be a string or a function")
 
-        assert kwargs == {}, f"Unknown keyword arguments: {kwargs}"
-
         cmp = self
 
         if cmp.n_points == 0:
             raise ValueError("Dataset is empty, no data to compare.")
 
         sk = cmp.mean_skill(weights=weights, metrics=[metric])
-        df = sk.to_dataframe()
+        assert isinstance(metric, str)
 
-        metric_name = metric if isinstance(metric, str) else metric.__name__
-        ser = df[metric_name]
-        score = {str(col): float(value) for col, value in ser.items()}
-
-        return score
+        return {k: v for k, v in sk.sort("model").select(["model", metric]).rows()}
 
     def save(self, filename: Union[str, Path]) -> None:
         """Save the ComparerCollection to a zip file.
@@ -963,3 +1012,5 @@ def _load_comparer(folder: str, f: str) -> Comparer:
         cmp = Comparer.load(f)
         os.remove(f)
         return cmp
+
+