diff --git a/docs/api-reference.md b/docs/api-reference.md index 688ff76..4665ef4 100644 --- a/docs/api-reference.md +++ b/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. diff --git a/src/jointfm_client/client.py b/src/jointfm_client/client.py index ef04f8a..248f054 100644 --- a/src/jointfm_client/client.py +++ b/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) diff --git a/src/jointfm_client/feature_importance.py b/src/jointfm_client/feature_importance.py new file mode 100644 index 0000000..5d06e92 --- /dev/null +++ b/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", +] diff --git a/tests/test_feature_importance.py b/tests/test_feature_importance.py new file mode 100644 index 0000000..8ad8d7c --- /dev/null +++ b/tests/test_feature_importance.py @@ -0,0 +1,297 @@ +# 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. + +"""Tests for permutation feature importance.""" + +from __future__ import annotations + +from collections.abc import Mapping +from typing import Any + +import pytest + +from jointfm_client import ColumnSpec, DataFrameSchema, JointFMClient +from jointfm_client.feature_importance import ( + permute_history_column, + sample_w2_distance, +) + +_MODEL_VERSION = "jointfm-inference:0.2.0+ckpt.sdk-test" + + +class _ImportanceTransport: + """Returns one queued sample response per `post_json` call, in order.""" + + def __init__(self, sample_batches: list[list[list[list[float]]]]) -> None: + """Init.""" + self.sample_batches = sample_batches + self.payloads: list[Mapping[str, Any]] = [] + + def get_json(self, url: str) -> Mapping[str, Any]: + """Get json.""" + raise AssertionError(f"unexpected health request: {url}") + + def post_json(self, url: str, payload: Mapping[str, Any]) -> Mapping[str, Any]: + """Post json.""" + self.payloads.append(dict(payload)) + samples = self.sample_batches[len(self.payloads) - 1] + return { + "schema_version": "v1", + "image_version": "0.2.0", + "model_version": _MODEL_VERSION, + "checkpoint_version": "sdk-test", + "head": "studentt", + "query_mode": "forecast", + "return_mode": "samples", + "outputs": { + "query_times": [2, 3], + "requested_columns": ["target"], + "mean": None, + "samples": samples, + "quantiles": None, + }, + "diagnostics": { + "history_rows": 2, + "horizon_count": 2, + "seed": payload.get("seed"), + }, + "errors": [], + } + + +def test_feature_importance_scores_mean_shift_and_distribution_distance() -> None: + """One feature's mean shift and distance are scored per target and horizon.""" + transport = _ImportanceTransport( + sample_batches=[ + [[[10.0], [20.0]], [[12.0], [22.0]]], + [[[11.0], [15.0]], [[13.0], [27.0]]], + ] + ) + client = JointFMClient( + predict_url="http://localhost:8080/predict", + transport=transport, + ) + columns = ( + ColumnSpec(name="target", modality="numeric", role="target"), + ColumnSpec(name="feat", modality="numeric", role="feature"), + ) + history = [{"target": 10.0, "feat": 1.0}, {"target": 11.0, "feat": 2.0}] + + result = client.feature_importance( + history, + query_times=[2, 3], + horizons=[1, 2], + feature_columns=["feat"], + target_columns=["target"], + columns=columns, + model_version=_MODEL_VERSION, + n_samples=2, + seed=7, + ) + + # horizon 1: means 11.0 -> 12.0 (shift 1.0), identical shape once centered (distance 0). + # horizon 2: means 21.0 -> 21.0 (no shift), spread widens from +-1 to +-6 (distance 12.5). + assert result == [ + { + "feature": "feat", + "mean": {"target": {1: 1.0, 2: 0.0}}, + "distance": {"target": {1: 0.0, 2: 12.5}}, + } + ] + assert len(transport.payloads) == 2 + assert [payload["seed"] for payload in transport.payloads] == [7, 7] + + baseline_feature_values = [ + row["feat"] for row in transport.payloads[0]["history_rows"] + ] + permuted_feature_values = [ + row["feat"] for row in transport.payloads[1]["history_rows"] + ] + assert baseline_feature_values == [1.0, 2.0] + assert sorted(permuted_feature_values) == sorted(baseline_feature_values) + + +def test_feature_importance_accepts_dataframe_history_with_explicit_schema() -> None: + """A DataFrame history paired with an explicit schema reaches the service. + + forecast() takes the row-payload path whenever schema is set, even when + history is a DataFrame rather than a row sequence, so feature_importance + must convert the frame to rows itself before delegating. + """ + pandas_module = pytest.importorskip("pandas") + transport = _ImportanceTransport( + sample_batches=[ + [[[10.0], [20.0]], [[12.0], [22.0]]], + [[[11.0], [15.0]], [[13.0], [27.0]]], + ] + ) + client = JointFMClient( + predict_url="http://localhost:8080/predict", + transport=transport, + ) + schema = DataFrameSchema( + columns=( + ColumnSpec(name="target", modality="numeric", role="target"), + ColumnSpec(name="feat", modality="numeric", role="feature"), + ), + time_index_mode="ordinal", + time_column="t", + ) + frame = pandas_module.DataFrame( + {"t": [0, 1], "target": [10.0, 11.0], "feat": [1.0, 2.0]} + ) + + result = client.feature_importance( + frame, + query_times=[2, 3], + horizons=[1, 2], + feature_columns=["feat"], + target_columns=["target"], + schema=schema, + model_version=_MODEL_VERSION, + n_samples=2, + seed=7, + ) + + assert result == [ + { + "feature": "feat", + "mean": {"target": {1: 1.0, 2: 0.0}}, + "distance": {"target": {1: 0.0, 2: 12.5}}, + } + ] + assert len(transport.payloads) == 2 + # the original frame is untouched by the conversion/permutation + assert frame["feat"].tolist() == [1.0, 2.0] + + +def test_feature_importance_requires_non_empty_feature_columns() -> None: + """feature_importance rejects an empty feature_columns list.""" + client = JointFMClient(predict_url="http://localhost:8080/predict") + + with pytest.raises(ValueError, match="feature_columns must not be empty"): + client.feature_importance( + [{"target": 10.0}], + query_times=[2], + horizons=[1], + feature_columns=[], + target_columns=["target"], + model_version=_MODEL_VERSION, + ) + + +def test_feature_importance_requires_non_empty_target_columns() -> None: + """feature_importance rejects an empty target_columns list.""" + client = JointFMClient(predict_url="http://localhost:8080/predict") + + with pytest.raises(ValueError, match="target_columns must not be empty"): + client.feature_importance( + [{"target": 10.0, "feat": 1.0}], + query_times=[2], + horizons=[1], + feature_columns=["feat"], + target_columns=[], + model_version=_MODEL_VERSION, + ) + + +def test_feature_importance_requires_horizons_matching_query_times_length() -> None: + """feature_importance rejects a horizons/query_times length mismatch.""" + client = JointFMClient(predict_url="http://localhost:8080/predict") + + with pytest.raises(ValueError, match="horizons must have the same length"): + client.feature_importance( + [{"target": 10.0, "feat": 1.0}], + query_times=[2, 3], + horizons=[1], + feature_columns=["feat"], + target_columns=["target"], + model_version=_MODEL_VERSION, + ) + + +def test_permute_history_column_shuffles_row_sequence_values() -> None: + """permute_history_column reshuffles one column across row mappings.""" + history = [{"feat": 1.0, "target": 10.0}, {"feat": 2.0, "target": 11.0}] + + permuted = permute_history_column(history, "feat", seed=1) + + assert sorted(row["feat"] for row in permuted) == [1.0, 2.0] + assert [row["target"] for row in permuted] == [10.0, 11.0] + # original history is left untouched + assert history == [{"feat": 1.0, "target": 10.0}, {"feat": 2.0, "target": 11.0}] + + +def test_permute_history_column_tolerates_rows_missing_the_column() -> None: + """A row that omits the feature key (sparse/nullable history) is left as-is.""" + history = [ + {"feat": 1.0, "target": 10.0}, + {"target": 11.0}, # no "feat" key, same as forecast() already accepts + {"feat": 3.0, "target": 12.0}, + ] + + permuted = permute_history_column(history, "feat", seed=1) + + assert "feat" not in permuted[1] + assert sorted(row["feat"] for row in permuted if "feat" in row) == [1.0, 3.0] + assert [row["target"] for row in permuted] == [10.0, 11.0, 12.0] + + +def test_permute_history_column_shuffles_dataframe_values() -> None: + """permute_history_column reshuffles one column across DataFrame rows.""" + pandas_module = pytest.importorskip("pandas") + frame = pandas_module.DataFrame({"feat": [1.0, 2.0], "target": [10.0, 11.0]}) + + permuted = permute_history_column(frame, "feat", seed=1) + + assert sorted(permuted["feat"].tolist()) == [1.0, 2.0] + assert permuted["target"].tolist() == [10.0, 11.0] + assert frame["feat"].tolist() == [1.0, 2.0] + + +def test_permute_history_column_rejects_missing_column() -> None: + """permute_history_column raises when the feature column is absent.""" + with pytest.raises(ValueError, match="missing column"): + permute_history_column([{"target": 10.0}], "feat", seed=1) + + +def test_permute_history_column_rejects_unsupported_history_type() -> None: + """permute_history_column raises for a history that is neither rows nor a frame.""" + with pytest.raises(ValueError, match="pandas DataFrame or a sequence"): + permute_history_column({"target": 10.0}, "feat", seed=1) + + +def test_permute_history_column_rejects_dataframe_missing_column() -> None: + """permute_history_column raises when the DataFrame lacks the feature column.""" + pandas_module = pytest.importorskip("pandas") + frame = pandas_module.DataFrame({"target": [10.0, 11.0]}) + + with pytest.raises(ValueError, match="missing column"): + permute_history_column(frame, "feat", seed=1) + + +def test_sample_w2_distance_is_zero_for_identically_shaped_shift() -> None: + """A pure mean shift scores zero distance once centered.""" + assert sample_w2_distance([10.0, 12.0], [11.0, 13.0]) == 0.0 + + +def test_sample_w2_distance_rejects_mismatched_lengths() -> None: + """sample_w2_distance requires equal-length sample vectors.""" + with pytest.raises(ValueError, match="equal length"): + sample_w2_distance([1.0, 2.0], [1.0]) + + +def test_sample_w2_distance_returns_zero_for_empty_inputs() -> None: + """sample_w2_distance returns zero rather than dividing by zero.""" + assert sample_w2_distance([], []) == 0.0