Skip to content

Geo Prediction

Prerequisites: The functions in this module read features directly from a hive-partitioned parquet store. The feature data for all requested targets, years, and spatial tiles must be present in the store before calling any prediction function — missing partitions will produce no rows for those combinations rather than raising an error.

The h2ml.geo module bridges trained FinalModels and hive-partitioned parquet feature stores (via h2mare.ParquetIndexer) to generate spatial predictions at scale.

Requires the [geo] optional dependencies:

uv pip install h2ml[geo]

predict_for_year

Generate predictions for a full calendar year across one or more targets.

from pathlib import Path
from h2ml.geo.geo_predict import predict_for_year

df = predict_for_year(
    target=["sparrow", "finch"],
    year=2023,
    root_dir=Path("project/"),
    input_parquet_dir=Path("data/features/"),
    schema="v1",
    geo_extent=(-10.0, 35.0, 30.0, 70.0),  # xmin, ymin, xmax, ymax
)

Models are loaded from root_dir/models/{target}_{schema}_final-model.pkl. Targets whose model file is missing or whose prediction raises an exception are skipped with a warning — the rest still complete.

Output columns: ['index', 'time', 'lon', 'lat', '{target}_{schema}', ...]

Rows with any null feature value appear as null in the prediction columns, preserving alignment with the original spatial grid.

Conformal interval columns

Pass alpha to add conformal bound columns for calibrated models. Uncalibrated models and multiclass classifiers always return point predictions only.

df = predict_for_year(
    target="sparrow",
    year=2023,
    root_dir=Path("project/"),
    input_parquet_dir=Path("data/features/"),
    schema="v1",
    geo_extent=(-10.0, 35.0, 30.0, 70.0),
    alpha=0.10,  # 90% coverage
)
# columns: sparrow_v1, sparrow_v1_pi_lower, sparrow_v1_pi_upper

Spatio-temporal varying intervals

When the model has a LocalConformalCalibration (built when store.coords and/or store.times were provided during training), passing local=True makes predict_for_year extract lon, lat, and time from the scan and produce interval widths that vary by location and season.

local defaults to False, so intervals use the global constant-width threshold unless you opt in:

# Global constant-width intervals (default)
df_global = predict_for_year(..., alpha=0.10)

# Spatio-temporally varying widths (requires a LocalConformalCalibration)
df_local = predict_for_year(..., alpha=0.10, local=True)

See Conformal Prediction — Spatio-temporal local calibration for details on how the local calibration is built and tuned.

The meaning of the bound columns differs by task type:

Regression — per-sample outcome bounds in the original (inverse-transformed) count scale. The bounds are computed in the transform space before inverting, which gives correct asymmetric intervals:

# without y_transform (symmetric)
pi_lower = max(0, ŷ − q)
pi_upper = ŷ + q

# with y_transform (asymmetric — e.g. log or sqrt)
pi_lower = max(0, inverse_fn(raw − q))
pi_upper = inverse_fn(raw + q)

where q is a per-sample threshold when local=True, or the global quantile when local=False.

Binary classification — per-sample calibration bands in probability space:

pi_lower = clip(p − q, 0, 1)
pi_upper = clip(p + q, 0, 1)

Note

Unlike regression, these bounds are not a formal coverage guarantee on the class label — they are a calibration band expressing how much the predicted probability could plausibly shift given past model errors.

predict_for_year_delta

Same as predict_for_year but loads DeltaFinalModels (two-component presence × abundance models). Model directories are expected at root_dir/models/{target}_{schema}_final-model/.

Unlike predict_for_year where alpha defaults to None (point predictions only), here alpha defaults to 0.10 — intervals are always attempted and omitted only when the model has no ConformalCalibration.

from h2ml.geo.geo_predict import predict_for_year_delta

df = predict_for_year_delta(
    target=["sparrow", "finch"],
    year=2023,
    root_dir=Path("project/"),
    input_parquet_dir=Path("data/features/"),
    schema="v1",
    geo_extent=(-10.0, 35.0, 30.0, 70.0),
    alpha=0.10,
    local=True,   # opt in to LocalConformalCalibration; default is False (global threshold)
)
# columns per target: {target}_{schema}, {target}_{schema}_pi_lower, {target}_{schema}_pi_upper

The delta prediction is P(present) × E(count | present). The conformal intervals are calibrated on the full combined delta output (not per component) and are always in the original count scale. Targets without a saved ConformalCalibration produce point predictions only.

predict_map

Quickly visualise model predictions aggregated over a spatial-temporal grid without saving results to disk. Useful for exploratory inspection of a single model.

from h2mare.storage import ParquetIndexer
from h2ml.geo.geo_predict import predict_map

indexer = ParquetIndexer(Path("data/features/"))

predict_map(
    model=final,                          # FinalModel
    indexer=indexer,
    dates=("2023-01-01", "2023-12-31"),
    bbox=(-10.0, 35.0, 30.0, 70.0),      # xmin, ymin, xmax, ymax
    target_col="sparrow_v1",
    agg_by="month",                       # "month" or "season"
)

Set save_path to write the figure to disk instead of calling plt.show():

predict_map(..., save_path="figures/sparrow_2023.png")

vminmax clips the colormap to a fixed range, useful when comparing maps across years:

predict_map(..., vminmax=(0.0, 50.0))

Differences from predict_for_year

predict_for_year predict_map
Output Polars DataFrame Plot only (None)
Targets Multiple Single model
Conformal intervals Yes (with alpha) No
Temporal aggregation None (raw time column) "month" or "season"
Typical use Production pipeline Exploratory / QA