Exploring support for eager and lazy dataframe semantics

[aaraney]

from __future__ import annotations

import typing

import narwhals as nw
import numpy as np

if typing.TYPE_CHECKING:
    import numpy.typing as npt
    import narwhals.typing as nwt


def mean_error(
    y_true: npt.ArrayLike,
    y_pred: npt.ArrayLike,
    power: float = 1.0,
    root: bool = False,
) -> float:
    # Compute mean error
    _mean_error = np.sum(np.abs(np.subtract(y_true, y_pred)) ** power) / len(y_true)

    # Return mean error, optionally return root mean error
    if root:
        return np.sqrt(_mean_error)
    return _mean_error


def mean_error_narwals_eager(
    y_true: nwt.IntoSeriesT,
    y_pred: nwt.IntoSeriesT,
    power: float = 1.0,
    root: bool = False,
) -> float:
    true, pred = (
        nw.from_native(y_true, series_only=True),
        nw.from_native(y_pred, series_only=True),
    )
    _mean_error = ((true - pred).abs() ** power).sum() / true.len()

    # Return mean error, optionally return root mean error
    if root:
        return np.sqrt(_mean_error)
    return _mean_error


def mean_error_narwals_lazy(
    y_true: nw.Expr,
    y_pred: nw.Expr,
    power: float = 1.0,
    root: bool = False,
) -> nw.Expr:
    _mean_error = ((y_true - y_pred).abs() ** power).sum() / y_true.len()

    # Return mean error, optionally return root mean error
    if root:
        return _mean_error.sqrt()
    return _mean_error


if typing.TYPE_CHECKING:
    import pandas as pd
    import polars as pl

    IntoSeriesT = typing.TypeVar("IntoSeriesT", pd.Series, pl.Series, nw.Series)


@typing.overload
def mean_error_narwals_overload(
    y_true: IntoSeriesT,
    y_pred: IntoSeriesT,
    power: float = 1.0,
    root: bool = False,
) -> float: ...


@typing.overload
def mean_error_narwals_overload(
    y_true: nwt.Expr,
    y_pred: nwt.Expr,
    power: float = 1.0,
    root: bool = False,
) -> nw.Expr: ...


def mean_error_narwals_overload(
    y_true: IntoSeriesT | nwt.Expr,
    y_pred: IntoSeriesT | nwt.Expr,
    power: float = 1.0,
    root: bool = False,
) -> float | nw.Expr:
    if isinstance(y_true, nw.Expr):
        true, pred = y_true, y_pred
    else:
        true, pred = (
            nw.from_native(y_true, series_only=True),
            nw.from_native(y_pred, series_only=True),
        )

    _mean_error = ((true - pred).abs() ** power).sum() / true.len()

    # Return mean error, optionally return root mean error
    if root:
        if isinstance(_mean_error, nw.Expr):
            return _mean_error.sqrt()
        return np.sqrt(_mean_error)
    return _mean_error


def p(x: object) -> None:
    print(x)


import pandas as pd
import polars as pl

true = np.array([1, 2, 3], dtype="float64")
pred = np.array([1, 2, 3], dtype="float64")

pd_df = pd.DataFrame({"true": true, "pred": pred})
pl_df = pl.DataFrame({"true": true, "pred": pred})
pl_lf = pl.LazyFrame({"true": true, "pred": pred})
nw_df = nw.from_native(pd_df)

p(mean_error(pd_df["true"], pd_df["pred"]))

p(
    mean_error_narwals_eager(pd_df["true"], pd_df["pred"]),
)
p(
    mean_error_narwals_eager(pl_df["true"], pl_df["pred"]),
)
p(
    mean_error_narwals_eager(nw_df["true"], nw_df["pred"]),
)
try:
    p(
        mean_error_narwals_eager(
            pd_df["true"], pl_df["pred"]
        )  # we want this to type error
    )
except TypeError as e:
    p(e)

try:
    p(
        mean_error_narwals_eager(true, pred),  # should typing error
    )
except TypeError as e:
    p(e)

p(
    mean_error_narwals_lazy(nw.col("true"), nw.col("pred")),
)
# mean_error_narwals_lazy("true", "pred")  # likely want this too

p(
    mean_error_narwals_overload(pd_df["true"], pd_df["pred"]),
)
p(
    mean_error_narwals_overload(pl_df["true"], pl_df["pred"]),
)
p(
    mean_error_narwals_overload(nw_df["true"], nw_df["pred"]),
)

p(
    nw.from_native(pl_lf)
    .select(
        mean_error_narwals_overload(nw.col("true"), nw.col("pred")),
    )
    .to_native()
    .collect()
)

[aaraney]

Fair warning, I had claude fable re-write the functionality of the metrics module using narwals.Expr. I've dropped the docstrings for the purposes of this issue.

code

from __future__ import annotations
from functools import partial
import typing

import narwhals as nw

if typing.TYPE_CHECKING:
    # A contingency table is a mapping from component label to either a
    # narwhals expression (aggregated count) or a plain number.
    ContingencyTable = typing.Mapping[str, nw.Expr | int | float]


def mean_error(
    y_true: nw.Expr,
    y_pred: nw.Expr,
    power: float = 1.0,
    root: bool = False,
) -> nw.Expr:
    # Compute mean error
    errors = (y_true - y_pred).abs() ** power
    _mean_error = errors.sum() / errors.count()

    # Return mean error, optionally return root mean error
    if root:
        return _mean_error.sqrt()
    return _mean_error


mean_squared_error = partial(mean_error, power=2.0, root=False)
root_mean_squared_error = partial(mean_error, power=2.0, root=True)


def mean_error_skill_score(
    y_true: nw.Expr,
    y_pred: nw.Expr,
    y_base: nw.Expr | float,
    power: float = 1.0,
    normalized: bool = False,
) -> nw.Expr:
    # Compute components
    numerator = mean_error(y_true, y_pred, power=power)
    if not isinstance(y_base, nw.Expr):
        y_base = nw.lit(y_base)
    denominator = mean_error(y_true, y_base, power=power)

    # Compute score, optionally normalize
    if normalized:
        return 1.0 / (1.0 + numerator / denominator)
    return 1.0 - numerator / denominator


volumetric_efficiency = partial(mean_error_skill_score, y_base=0.0, power=1.0)


def nash_sutcliffe_efficiency(
    y_true: nw.Expr,
    y_pred: nw.Expr,
    log: bool = False,
    power: float = 2.0,
    normalized: bool = False,
) -> nw.Expr:
    # Optionally transform components
    if log:
        y_true = y_true.log()
        y_pred = y_pred.log()

    # Compute score
    return mean_error_skill_score(
        y_true, y_pred, y_true.mean(), power=power, normalized=normalized
    )


def coefficient_of_persistence(
    y_true: nw.Expr,
    y_pred: nw.Expr,
    lag: int = 1,
    log: bool = False,
    power: float = 2.0,
    normalized: bool = False,
) -> nw.Expr:
    # Optionally transform components
    if log:
        y_true = y_true.log()
        y_pred = y_pred.log()

    # Compute baseline. shift(lag) yields null for the first `lag` rows,
    # which are excluded from the score (equivalent to slicing off the
    # first `lag` values in an eager implementation).
    y_base = y_true.shift(lag)

    # Mask y_true so the first `lag` rows are also excluded from the
    # numerator, keeping numerator and denominator over the same rows
    y_true_masked = nw.when(~y_base.is_null()).then(y_true)

    # Compute score
    return mean_error_skill_score(
        y_true_masked, y_pred, y_base, power=power, normalized=normalized
    )


def coefficient_of_extrapolation(
    y_true: nw.Expr,
    y_pred: nw.Expr,
    log: bool = False,
    power: float = 2.0,
    normalized: bool = False,
) -> nw.Expr:
    # Optionally transform components
    if log:
        y_true = y_true.log()
        y_pred = y_pred.log()

    # Compute baseline: y_base_i = y_true_{i-1} + (y_true_{i-1} - y_true_{i-2}).
    # The first two rows are null and thereby excluded from the score
    # (equivalent to slicing off the first two values in an eager
    # implementation).
    y_base = y_true.shift(1) + y_true.diff().shift(1)

    # Mask y_true so the first two rows are also excluded from the
    # numerator, keeping numerator and denominator over the same rows
    y_true_masked = nw.when(~y_base.is_null()).then(y_true)

    # Compute score
    return mean_error_skill_score(
        y_true_masked, y_pred, y_base, power=power, normalized=normalized
    )


def kling_gupta_efficiency(
    y_true: nw.Expr,
    y_pred: nw.Expr,
    r_scale: float = 1.0,
    a_scale: float = 1.0,
    b_scale: float = 1.0,
) -> nw.Expr:
    # Pearson correlation coefficient (population moments)
    covariance = ((y_pred - y_pred.mean()) * (y_true - y_true.mean())).mean()
    linear_correlation = covariance / (y_pred.std(ddof=0) * y_true.std(ddof=0))

    # Relative variability
    relative_variability = y_pred.std(ddof=0) / y_true.std(ddof=0)

    # Relative mean
    relative_mean = y_pred.mean() / y_true.mean()

    # Scaled Euclidean distance
    euclidean_distance = (
        (r_scale * (linear_correlation - 1.0)) ** 2.0
        + (a_scale * (relative_variability - 1.0)) ** 2.0
        + (b_scale * (relative_mean - 1.0)) ** 2.0
    ).sqrt()

    # Return KGE
    return 1.0 - euclidean_distance


def compute_contingency_table(
    observed: nw.Expr,
    simulated: nw.Expr,
    true_positive_key: str = "true_positive",
    false_positive_key: str = "false_positive",
    false_negative_key: str = "false_negative",
    true_negative_key: str = "true_negative",
) -> dict[str, nw.Expr]:
    # Validate boolean categorical
    observed = observed.cast(nw.Boolean)
    simulated = simulated.cast(nw.Boolean)

    # Cross tabulate
    return {
        true_positive_key: (observed & simulated).sum(),
        false_positive_key: (~observed & simulated).sum(),
        false_negative_key: (observed & ~simulated).sum(),
        true_negative_key: (~observed & ~simulated).sum(),
    }


def _as_float(value: nw.Expr | float) -> nw.Expr | float:
    """Cast a contingency table component to a float-valued expression or scalar."""
    if isinstance(value, nw.Expr):
        return value.cast(nw.Float64)
    return float(value)


def probability_of_detection(
    contingency_table: ContingencyTable,
    true_positive_key: str = "true_positive",
    false_negative_key: str = "false_negative",
) -> nw.Expr | float:
    a = _as_float(contingency_table[true_positive_key])
    c = _as_float(contingency_table[false_negative_key])
    return a / (a + c)


def probability_of_false_detection(
    contingency_table: ContingencyTable,
    false_positive_key: str = "false_positive",
    true_negative_key: str = "true_negative",
) -> nw.Expr | float:
    b = _as_float(contingency_table[false_positive_key])
    d = _as_float(contingency_table[true_negative_key])
    return b / (b + d)


def probability_of_false_alarm(
    contingency_table: ContingencyTable,
    true_positive_key: str = "true_positive",
    false_positive_key: str = "false_positive",
) -> nw.Expr | float:
    b = _as_float(contingency_table[false_positive_key])
    a = _as_float(contingency_table[true_positive_key])
    return b / (b + a)


def threat_score(
    contingency_table: ContingencyTable,
    true_positive_key: str = "true_positive",
    false_positive_key: str = "false_positive",
    false_negative_key: str = "false_negative",
) -> nw.Expr | float:
    a = _as_float(contingency_table[true_positive_key])
    b = _as_float(contingency_table[false_positive_key])
    c = _as_float(contingency_table[false_negative_key])
    return a / (a + b + c)


def frequency_bias(
    contingency_table: ContingencyTable,
    true_positive_key: str = "true_positive",
    false_positive_key: str = "false_positive",
    false_negative_key: str = "false_negative",
) -> nw.Expr | float:
    a = _as_float(contingency_table[true_positive_key])
    b = _as_float(contingency_table[false_positive_key])
    c = _as_float(contingency_table[false_negative_key])
    return (a + b) / (a + c)


def percent_correct(
    contingency_table: ContingencyTable,
    true_positive_key: str = "true_positive",
    false_positive_key: str = "false_positive",
    false_negative_key: str = "false_negative",
    true_negative_key: str = "true_negative",
) -> nw.Expr | float:
    a = _as_float(contingency_table[true_positive_key])
    b = _as_float(contingency_table[false_positive_key])
    c = _as_float(contingency_table[false_negative_key])
    d = _as_float(contingency_table[true_negative_key])
    return (a + d) / (a + b + c + d)


def base_chance(
    contingency_table: ContingencyTable,
    true_positive_key: str = "true_positive",
    false_positive_key: str = "false_positive",
    false_negative_key: str = "false_negative",
    true_negative_key: str = "true_negative",
) -> nw.Expr | float:
    a = _as_float(contingency_table[true_positive_key])
    b = _as_float(contingency_table[false_positive_key])
    c = _as_float(contingency_table[false_negative_key])
    d = _as_float(contingency_table[true_negative_key])
    return ((a + b) * (a + c)) / (a + b + c + d)


def equitable_threat_score(
    contingency_table: ContingencyTable,
    true_positive_key: str = "true_positive",
    false_positive_key: str = "false_positive",
    false_negative_key: str = "false_negative",
    true_negative_key: str = "true_negative",
) -> nw.Expr | float:
    a_r = base_chance(
        contingency_table,
        true_positive_key=true_positive_key,
        false_positive_key=false_positive_key,
        false_negative_key=false_negative_key,
        true_negative_key=true_negative_key,
    )
    a = _as_float(contingency_table[true_positive_key])
    b = _as_float(contingency_table[false_positive_key])
    c = _as_float(contingency_table[false_negative_key])
    return (a - a_r) / (a + b + c - a_r)

[jarq6c]

Mulling over interfaces, I think there may be some merit in borrowing patterns from visualization libraries like plotly.

Looping in @bolotinl and @christophertubbs

Let me know if you have any comments or a preference for one of the proposed interfaces.

Motivating example

import plotly.express as px
df = px.data.iris()
fig = px.scatter(df, x="sepal_width", y="sepal_length", color="species")
fig.show()

Metrics adaptation

import hydrotools.metrics as htm
df = pd.read_csv("mah_data.csv")
evaluation = htm.nash_sutcliffe_efficiency(
    df,
    observations="usgs_value",
    predictions="cfe_value",
    group_by="location_id"
)
evaluation.compute()

Assumptions

1. Something inside htm.nash_sutcliffe_efficiency is normalizing df to narwahls.
2. htm.nash_sutcliffe_efficiency is lazy by default.
3. evaluation.compute resolves in memory by default, but could have parameters that accommodate larger-than-memory computations that stream to disk (or wherever).
4. I think this is a good simple entry point. I can imagine a higher level method that computes multiple metrics.

Multiple metrics

Option 1: Super method

import hydrotools.metrics as htm
df = pd.read_csv("mah_data.csv")
evaluation = htm.evaluate(
    df,
    observations="usgs_value",
    predictions="cfe_value",
    group_by="location_id",
    metrics=[htm.Metric.NASH_SUTCLIFFE_EFFICIENCY, htm.Metric.PEARSON_CORRELATION_COEFFICIENT],
    metric_labels=["nse", "pearson_r"]
)
evaluation.compute()

Option 2: pandas-like split-apply-combine

import hydrotools.metrics as htm
df = pd.read_csv("mah_data.csv")
evaluation = htm.evaluate(
    df,
    observations="usgs_value",
    predictions="cfe_value"
).group_by(
    "location_id"
).apply(
    nse=htm.nash_sutcliffe_efficiency,
    pearson_r=htm.pearson_correlation_coefficient
)
evaluation.compute()

Option 3: Declarative objects

import hydrotools.metrics as htm
df = pd.read_csv("mah_data.csv")
evaluation = htm.Evaluation(
    df,
    observations="usgs_value",
    predictions="cfe_value"
)
evaluation.add_group("location_id")
evaluation.add_metric(htm.Metric.NASH_SUTCLIFFE_EFFICIENCY, label="nse")
evaluation.add_metric(htm.Metric.PEARSON_CORRELATION_COEFFICIENT, label="pearson_r")
evaluation.compute()