feat: add permutation feature importance to JointFMClient

docs/api-reference.md

@@ -6,7 +6,7 @@ This reference covers the supported public Python surface exported by `jointfm_c
 
 | Name | Purpose |
 | --- | --- |
-| `JointFMClient` | Synchronous client for hosted or local JointFM endpoints. Use `from_env()` for `.env` and `config.yaml` backed hosted settings, `health()` for consensus typed service metadata, `health_instances()` for per-deployment probe results and pooled sample topology, `predict(payload)` for low-level JSON prediction, `forecast(...)` for validated tabular forecasts, and the `forecast_mean(...)`, `forecast_samples(...)`, and `forecast_quantiles(...)` convenience methods for typed forecast results. `health()` probes `GET /healthz` for local deployments and POSTs `{"request_type": "health"}` to `predict_url` for hosted DataRobot deployments because the DataRobot deployment gateway only proxies the unstructured prediction route. |
+| `JointFMClient` | Synchronous client for hosted or local JointFM endpoints. Use `from_env()` for `.env` and `config.yaml` backed hosted settings, `health()` for consensus typed service metadata, `health_instances()` for per-deployment probe results and pooled sample topology, `predict(payload)` for low-level JSON prediction, `forecast(...)` for validated tabular forecasts, and the `forecast_mean(...)`, `forecast_samples(...)`, and `forecast_quantiles(...)` convenience methods for typed forecast results. `health()` probes `GET /healthz` for local deployments and POSTs `{"request_type": "health"}` to `predict_url` for hosted DataRobot deployments because the DataRobot deployment gateway only proxies the unstructured prediction route. `feature_importance(...)` runs permutation feature importance: one baseline `forecast_samples` call plus one per shuffled feature column, returning a list of `{"feature", "mean", "distance"}` dicts, each holding that feature's absolute forecast-mean shift and centered squared 2-Wasserstein distance indexed by target and horizon. |
 
 `JointFMClient.from_env()` loads `config.yaml`, optional `.env` values, and process environment variables. `JointFMClient.health(cache=True)` caches health metadata only when requested. `JointFMClient.health_instances()` returns the same probe as a `HealthInstances` object: one `InstanceHealth` per configured deployment (including failures), `max_sample_count` as the sum of reachable caps (overall parallel capacity), and `topology` / `topology_label` grouping those caps (unavailable peers are listed but excluded from the sum and topology). Each endpoint's health payload describes only that endpoint; the client aggregates by calling each configured peer. `health()` still exposes the minimum reachable `max_sample_count`, which is the sample-batch cap used by forecast helpers. `JointFMClient.predict(payload)` requires `payload["model_version"]`; high-level forecast helpers resolve the configured model version when the caller does not pass one explicitly. When `forecast_samples(...)` requests an explicit `n_samples`, the client learns the deployment's `max_sample_count` from health metadata before the first prediction, splits oversized requests into capped prediction batches, and returns one merged `SampleForecastResult`. Clients configured without a reachable health route fall back to discovering the cap from the structured service error.
 

src/jointfm_client/client.py

@@ -24,7 +24,10 @@
 from typing import Any, Self, cast
 from urllib.parse import urlparse
 
-from jointfm_client.adapters import build_forecast_payload_from_dataframe
+from jointfm_client.adapters import (
+    build_forecast_payload_from_dataframe,
+    dataframe_to_history_rows,
+)
 from jointfm_client.configuration import (
     DATAROBOT_REQUEST_ID_HEADERS,
     DEFAULT_CONFIG_PATH,
@@ -52,6 +55,10 @@
     JointFMServiceError,
     UnsupportedModelVersionError,
 )
+from jointfm_client.feature_importance import (
+    feature_importance_entry,
+    permute_history_column,
+)
 from jointfm_client.contract import (
     ForecastDiagnostics,
     ForecastResponse,
@@ -474,6 +481,137 @@ def forecast_quantiles(
             ),
         )
 
+    def feature_importance(
+        self,
+        history: Any,
+        *,
+        query_times: Sequence[Any],
+        horizons: Sequence[int],
+        feature_columns: Sequence[str],
+        target_columns: Sequence[str],
+        schema: DataFrameSchema | None = None,
+        time_index_mode: TimeIndexMode = "ordinal",
+        columns: Sequence[ColumnSpec] | None = None,
+        time_column: str | None = None,
+        model_version: str | None = None,
+        n_samples: int = 1024,
+        seed: int = 7,
+    ) -> list[dict[str, Any]]:
+        """Permutation feature importance for each feature, target, and horizon.
+
+        For every column in ``feature_columns``, the column is shuffled across
+        history rows (preserving its marginal) and scored against one shared
+        baseline forecast sampled from the unpermuted history. ``horizons`` are
+        integer labels for the output, paired by position with ``query_times``
+        (``horizons[i]`` names the forecast step requested by
+        ``query_times[i]``).
+
+        Returns one dict per feature: ``{"feature": name, "mean": {target:
+        {horizon: value}}, "distance": {target: {horizon: value}}}``. ``mean``
+        is the absolute shift in forecast mean; ``distance`` is the centered
+        squared 2-Wasserstein distance between baseline and permuted forecast
+        samples, which also catches distributional changes a mean shift alone
+        would miss. Both are computed from the same permuted-forecast sample
+        set, at the same seed and sample count as the baseline.
+        """
+        if len(feature_columns) == 0:
+            raise ValueError("feature_columns must not be empty")
+        if len(target_columns) == 0:
+            raise ValueError("target_columns must not be empty")
+        if len(horizons) != len(query_times):
+            raise ValueError(
+                "horizons must have the same length as query_times: "
+                f"{len(horizons)} != {len(query_times)}"
+            )
+
+        baseline = self._sample_forecast_for_importance(
+            history,
+            query_times=query_times,
+            schema=schema,
+            time_index_mode=time_index_mode,
+            columns=columns,
+            time_column=time_column,
+            target_columns=target_columns,
+            feature_columns=feature_columns,
+            model_version=model_version,
+            n_samples=n_samples,
+            seed=seed,
+        )
+
+        entries: list[dict[str, Any]] = []
+        for feature_index, feature in enumerate(feature_columns):
+            permuted_history = permute_history_column(
+                history, feature, seed=seed + 1009 * (feature_index + 1)
+            )
+            permuted = self._sample_forecast_for_importance(
+                permuted_history,
+                query_times=query_times,
+                schema=schema,
+                time_index_mode=time_index_mode,
+                columns=columns,
+                time_column=time_column,
+                target_columns=target_columns,
+                feature_columns=feature_columns,
+                model_version=model_version,
+                n_samples=n_samples,
+                seed=seed,
+            )
+            entries.append(
+                feature_importance_entry(
+                    feature=feature,
+                    horizons=horizons,
+                    target_columns=target_columns,
+                    baseline_samples=baseline.samples,
+                    permuted_samples=permuted.samples,
+                    baseline_columns=baseline.requested_columns,
+                )
+            )
+        return entries
+
+    def _sample_forecast_for_importance(
+        self,
+        history: Any,
+        *,
+        query_times: Sequence[Any],
+        schema: DataFrameSchema | None,
+        time_index_mode: TimeIndexMode,
+        columns: Sequence[ColumnSpec] | None,
+        time_column: str | None,
+        target_columns: Sequence[str],
+        feature_columns: Sequence[str],
+        model_version: str | None,
+        n_samples: int,
+        seed: int,
+    ) -> SampleForecastResult:
+        """Sample-forecast one history for ``feature_importance``, role-aware."""
+        resolved_history = history
+        if schema is not None and not _is_history_row_sequence(history):
+            # forecast() takes the row-payload path whenever schema is set, so a
+            # DataFrame paired with an explicit schema must become rows first.
+            resolved_history = dataframe_to_history_rows(history, schema)
+
+        result = self.forecast(
+            resolved_history,
+            query_times=query_times,
+            schema=schema,
+            time_index_mode=time_index_mode,
+            columns=columns,
+            time_column=time_column,
+            requested_columns=target_columns,
+            return_mode="samples",
+            model_version=model_version,
+            n_samples=n_samples,
+            seed=seed,
+            target_columns=target_columns,
+            feature_columns=feature_columns,
+        )
+        if not isinstance(result, SampleForecastResult):
+            raise JointFMServiceError(
+                "JointFM forecast response violated the V1 contract: "
+                "feature_importance requires sample forecast responses"
+            )
+        return result
+
     def refresh_health(self, *, cache: bool = True) -> HealthMetadata:
         """Fetch fresh service metadata and update the explicit cache by default."""
         return self.health(cache=cache, refresh=True)

src/jointfm_client/feature_importance.py

@@ -0,0 +1,161 @@
+# Copyright 2026 DataRobot, Inc. and its affiliates.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Permutation feature importance helpers for ``JointFMClient.feature_importance``.
+
+Each feature column is shuffled across history rows while its marginal is
+preserved, and the permuted forecast is compared against one shared baseline
+forecast. Two scores come out of that comparison: the absolute shift in
+forecast mean, and the centered squared 2-Wasserstein distance between the
+baseline and permuted sample sets (which also catches shape changes a mean
+shift alone would miss).
+"""
+
+from __future__ import annotations
+
+from collections.abc import Sequence
+import random
+from typing import Any
+
+
+def sample_w2_distance(
+    baseline: Sequence[float],
+    permuted: Sequence[float],
+    *,
+    location_invariant: bool = True,
+) -> float:
+    """Centered squared 2-Wasserstein distance between two 1D sample sets.
+
+    Both sample vectors are one-dimensional draws at the same target and
+    horizon, so the distance has a closed form: sort both vectors and take the
+    mean squared gap (halved to match the usual W2 cost convention). When
+    ``location_invariant`` is true, both vectors are centered first so the
+    score reflects spread, skew, and tail movement rather than repeating the
+    mean-shift readout.
+    """
+    if len(baseline) == 0 or len(permuted) == 0:
+        return 0.0
+    if len(baseline) != len(permuted):
+        raise ValueError(
+            f"sample vectors must have equal length; got {len(baseline)} and {len(permuted)}"
+        )
+
+    left = list(baseline)
+    right = list(permuted)
+    if location_invariant:
+        left_mean = sum(left) / len(left)
+        right_mean = sum(right) / len(right)
+        left = [value - left_mean for value in left]
+        right = [value - right_mean for value in right]
+
+    left.sort()
+    right.sort()
+    squared_gaps = sum(
+        (left_value - right_value) ** 2
+        for left_value, right_value in zip(left, right, strict=True)
+    )
+    return squared_gaps / len(left) / 2.0
+
+
+def _is_history_row_sequence(history: Any) -> bool:
+    """Return whether ``history`` is a sequence of row mappings, not a DataFrame."""
+    return isinstance(history, Sequence) and not isinstance(
+        history, str | bytes | bytearray
+    )
+
+
+def permute_history_column(history: Any, feature: str, *, seed: int) -> Any:
+    """Return a copy of ``history`` with ``feature`` shuffled across rows."""
+    rng = random.Random(seed)
+    if _is_history_row_sequence(history):
+        rows = [dict(row) for row in history]
+        present_indices = [index for index, row in enumerate(rows) if feature in row]
+        if not present_indices:
+            raise ValueError(f"history rows are missing column {feature!r}")
+        values = [rows[index][feature] for index in present_indices]
+        rng.shuffle(values)
+        for index, value in zip(present_indices, values, strict=True):
+            rows[index][feature] = value
+        return rows
+
+    pandas_module = _require_pandas_module()
+    if not isinstance(history, pandas_module.DataFrame):
+        raise ValueError(
+            "history must be a pandas DataFrame or a sequence of row mappings"
+        )
+    if feature not in history.columns:
+        raise ValueError(f"history frame is missing column {feature!r}")
+    permuted = history.copy()
+    values = list(history[feature])
+    rng.shuffle(values)
+    permuted[feature] = values
+    return permuted
+
+
+def feature_importance_entry(
+    *,
+    feature: str,
+    horizons: Sequence[int],
+    target_columns: Sequence[str],
+    baseline_samples: Sequence[Sequence[Sequence[float]]],
+    permuted_samples: Sequence[Sequence[Sequence[float]]],
+    baseline_columns: Sequence[str],
+) -> dict[str, Any]:
+    """Score one permuted feature against the shared baseline forecast.
+
+    Returns ``{"feature": ..., "mean": {target: {horizon: value}}, "distance":
+    {target: {horizon: value}}}``, where ``mean`` is the absolute shift in
+    forecast mean and ``distance`` is ``sample_w2_distance``, both indexed by
+    every requested target and horizon.
+    """
+    mean_scores: dict[str, dict[int, float]] = {}
+    distance_scores: dict[str, dict[int, float]] = {}
+    for target in target_columns:
+        target_index = list(baseline_columns).index(target)
+        mean_by_horizon: dict[int, float] = {}
+        distance_by_horizon: dict[int, float] = {}
+        for horizon_index, horizon in enumerate(horizons):
+            baseline_values = [
+                sample[horizon_index][target_index] for sample in baseline_samples
+            ]
+            permuted_values = [
+                sample[horizon_index][target_index] for sample in permuted_samples
+            ]
+            baseline_mean = sum(baseline_values) / len(baseline_values)
+            permuted_mean = sum(permuted_values) / len(permuted_values)
+            mean_by_horizon[horizon] = abs(permuted_mean - baseline_mean)
+            distance_by_horizon[horizon] = sample_w2_distance(
+                baseline_values, permuted_values
+            )
+        mean_scores[target] = mean_by_horizon
+        distance_scores[target] = distance_by_horizon
+    return {"feature": feature, "mean": mean_scores, "distance": distance_scores}
+
+
+def _require_pandas_module() -> Any:
+    """Return pandas or raise with the SDK extra needed for DataFrame history."""
+    try:
+        import pandas as pandas_module
+    except ImportError as error:  # pragma: no cover - exercised only without extra
+        raise RuntimeError(
+            "pandas history support requires installing jointfm-client[notebooks]"
+        ) from error
+    return pandas_module
+
+
+__all__ = [
+    "feature_importance_entry",
+    "permute_history_column",
+    "sample_w2_distance",
+]