Refactor/hpo from model to protocol

chap_core/assessment/evaluation.py

@@ -7,9 +7,10 @@
 
 import datetime
 import json
+import logging
 from abc import ABC, abstractmethod
 from collections.abc import Iterable
-from dataclasses import dataclass
+from dataclasses import asdict, dataclass
 from pathlib import Path
 from typing import cast
 
@@ -27,6 +28,7 @@
     convert_backtest_to_flat_forecasts,
     max_horizon_distance,
 )
+from chap_core.assessment.prediction_evaluator import Estimator
 from chap_core.assessment.weather_providers import (
     DEFAULT_WEATHER_PROVIDER_ID,
     LEGACY_WEATHER_PROVIDER_ID,
@@ -36,10 +38,17 @@
 from chap_core.database.model_templates_and_config_tables import ConfiguredModelDB
 from chap_core.database.tables import Backtest, BacktestForecast, BacktestSpecification
 from chap_core.datatypes import SamplesWithTruth
+from chap_core.external.ExtendedPredictor import ExtendedPredictor
 from chap_core.external.model_configuration import ModelTemplateConfigV2
+from chap_core.hpo.hyperparameter_optimizer import HyperparameterOptimizer
+from chap_core.hpo.meta_learner import MetaLearner
+from chap_core.hpo.types import FlatHyperparameterOptimization
+from chap_core.models.configured_model import ConfiguredModel
 from chap_core.rest_api.data_models import BacktestCreate
 from chap_core.time_period import Month, TimePeriod
 
+logger = logging.getLogger(__name__)
+
 try:
     from chap_core import __version__ as CHAP_VERSION
 except ImportError:
@@ -114,6 +123,9 @@ def _flat_data_to_xarray(flat_data: "FlatEvaluationData", model_metadata: dict)
     if "model_info" in model_metadata:
         attrs["model_info"] = model_metadata["model_info"]
 
+    if flat_data.hpo is not None:
+        attrs["hpo"] = json.dumps(asdict(flat_data.hpo))
+
     ds.attrs.update(attrs)
 
     return ds
@@ -163,10 +175,16 @@ def _xarray_to_flat_data(ds: xr.Dataset) -> "FlatEvaluationData":
         if not historical_df.empty:
             historical_observations = FlatObserved.validate(historical_df)
 
+    # Load HPO data if present (backwards compatible)
+    hpo = None
+    if "hpo" in ds.attrs:
+        hpo = FlatHyperparameterOptimization(**json.loads(ds.attrs["hpo"]))
+
     return FlatEvaluationData(
         forecasts=FlatForecasts.validate(forecasts_df),
         observations=FlatObserved.validate(observations_df),
         historical_observations=historical_observations,
+        hpo=hpo,
     )
 
 
@@ -188,6 +206,7 @@ class FlatEvaluationData:
     forecasts: pa.typing.DataFrame[FlatForecasts]
     observations: pa.typing.DataFrame[FlatObserved]
     historical_observations: pa.typing.DataFrame[FlatObserved] | None = None
+    hpo: FlatHyperparameterOptimization | None = None
 
 
 class EvaluationBase(ABC):
@@ -271,6 +290,7 @@ def __init__(
         backtest: "Backtest",
         historical_observations: list[Observation] | None = None,
         historical_context_periods: int = 0,
+        hpo: FlatHyperparameterOptimization | None = None,
     ):
         """
         Initialize Evaluation with a Backtest object.
@@ -284,6 +304,7 @@ def __init__(
         self._backtest = backtest
         self._historical_observations = historical_observations or []
         self._historical_context_periods = historical_context_periods
+        self._hpo = hpo
         self._flat_data_cache: FlatEvaluationData | None = None
 
     @property
@@ -318,6 +339,7 @@ def from_samples_with_truth(
         historical_observations: list[Observation] | None = None,
         historical_context_periods: int = 0,
         specification: BacktestSpecification | None = None,
+        hpo: FlatHyperparameterOptimization | None = None,
     ) -> "Evaluation":
         info.created = datetime.datetime.now()
         # The parameters live on the specification now and are computed fields on
@@ -405,13 +427,14 @@ def from_samples_with_truth(
             backtest,
             historical_observations=historical_observations,
             historical_context_periods=historical_context_periods,
+            hpo=hpo,
         )
 
     @classmethod
     def create(
         cls,
         configured_model: ConfiguredModelDB,
-        estimator,
+        estimator: Estimator | MetaLearner,
         dataset: _DataSet,
         backtest_params: BacktestParams,
         backtest_name: str = "evaluation",
@@ -449,9 +472,30 @@ def create(
             future_weather_provider=backtest_params.future_weather_provider,
         )
 
+        hpo_data = None
+        if isinstance(estimator, HyperparameterOptimizer):
+            hpo_data = estimator.meta_learn(train_set)
+            model = hpo_data.objective.model_template.get_model(hpo_data.model_configuration)  # type: ignore[arg-type]
+            tuned_estimator = model()  # type: ignore[assignment]
+        elif isinstance(estimator, MetaLearner):
+            raise TypeError(f"Unsupported MetaLearner: {type(estimator).__name__}")
+        else:
+            tuned_estimator = estimator
+
+        # also used by hpo objective call
+        if (
+            isinstance(tuned_estimator, ConfiguredModel) and tuned_estimator.model_information is not None
+        ):  # ensembleModel returns None, NaiveModel has no model_information
+            max_periods = tuned_estimator.model_information.max_prediction_periods
+            if max_periods is not None and max_periods < backtest_params.n_periods:
+                logger.warning(
+                    f"Wrapping model to extend prediction length from {max_periods} to {backtest_params.n_periods}. This is done iteratively, and may worsen model performance"
+                )
+                tuned_estimator = ExtendedPredictor(tuned_estimator, backtest_params.n_periods)
+
         # Run backtest
         evaluation_results = backtest(
-            estimator=estimator,
+            estimator=tuned_estimator,
             train_set=train_set,
             test_generator=test_generator,
             n_test_sets=backtest_params.n_splits,
@@ -493,6 +537,7 @@ def create(
             info=backtest_info,
             historical_observations=historical_observations,
             historical_context_periods=historical_context_periods,
+            hpo=hpo_data.to_flat() if hpo_data is not None else None,
         )
 
     @classmethod
@@ -623,6 +668,7 @@ def to_flat(self) -> FlatEvaluationData:
                 forecasts=FlatForecasts.validate(forecasts_df),
                 observations=FlatObserved.validate(observations_df),
                 historical_observations=historical_observations,
+                hpo=self._hpo,
             )
         return self._flat_data_cache
 
@@ -791,6 +837,7 @@ def from_file(cls, filepath: str | Path) -> "Evaluation":
             backtest,
             historical_observations=historical_observations,
             historical_context_periods=historical_context_periods,
+            hpo=flat_data.hpo,
         )
 
     @staticmethod

chap_core/cli_endpoints/_common.py

@@ -13,7 +13,7 @@
     from chap_core.api_types import BacktestParams, EstimatorOptions
     from chap_core.database.model_templates_and_config_tables import ModelConfiguration
     from chap_core.external.model_configuration import ModelTemplateConfigV2
-    from chap_core.hpo.hpoModel import HpoModel
+    from chap_core.hpo.hyperparameter_optimizer import HyperparameterOptimizer
     from chap_core.models.external_model import ExternalModel
     from chap_core.models.model_template import ModelTemplate
     from chap_core.spatio_temporal_data.temporal_dataclass import DataSet
@@ -323,7 +323,7 @@ def get_hpo_estimator(
     configuration: ModelConfiguration | None,
     backtest_params: BacktestParams,
     options: EstimatorOptions,
-) -> HpoModel:
+) -> HyperparameterOptimizer:
     """
     Build an HPO-backend estimator from either:
     - an explicit YAML search space, or
@@ -333,7 +333,7 @@ def get_hpo_estimator(
 
     from chap_core.api_types import SearcherType
     from chap_core.hpo.base import load_search_space_from_config
-    from chap_core.hpo.hpoModel import HpoModel
+    from chap_core.hpo.hyperparameter_optimizer import HyperparameterOptimizer
     from chap_core.hpo.objective import Objective
     from chap_core.hpo.searcher import GridSearcher, RandomSearcher, Searcher, TPESearcher
     from chap_core.hpo.types import DEFAULT_HPO_TRIALS
@@ -372,7 +372,7 @@ def get_hpo_estimator(
     else:
         raise ValueError(f"Unknown searcher: {searcher_type!r}")
 
-    return HpoModel(
+    return HyperparameterOptimizer(
         objective=objective,
         searcher=searcher,
         configuration=configuration,

chap_core/cli_endpoints/evaluate.py

@@ -26,10 +26,10 @@
     resolve_csv_path,
     warn_unused_covariates,
 )
+from chap_core.hpo.meta_learner import MetaLearner
 
 if TYPE_CHECKING:
     from chap_core.external.ExtendedPredictor import ExtendedPredictor
-    from chap_core.hpo.hpoModel import HpoModel
     from chap_core.models.external_model import ExternalModel
 
 logger = logging.getLogger(__name__)
@@ -130,7 +130,7 @@ def _run_eval(
             "Use --estimator-options.mode=normal for a normal evaluation run. "
             "Use --estimator-options.mode=hpo for hyperparameter optimization. "
             "Use --estimator-options.mode=ensemble for ensemble learning. "
-            "Optionally --estimator-options.search_space_yaml=<path> for hpo. "
+            "Optionally --estimator-options.search-space=<path> for hpo. "
             "Optionally --estimator-options.metric=<metric> for hpo. "
             "Optionally --estimator-options.searcher=<searcher> for hpo. "
             "Optionally --estimator-options.max-trials=<max_trials> for hpo. "
@@ -169,7 +169,7 @@ def _run_eval(
         # Evaluate with hyperparameter optimization
         chap eval --model-name https://github.com/chap-models/minimal_template_example \\
             --dataset-csv ./example_data/vietnam_monthly.csv --output-file ./chap_core/hpo/eval.nc \\
-            --estimator-options.mode hpo --estimator-options.search-space-yaml ./chap_core/hpo/config3.yaml \\
+            --estimator-options.mode hpo --estimator-options.search-space ./chap_core/hpo/config3.yaml \\
             --estimator-options.metric rmse --estimator-options.searcher tpe
     """
     from chap_core.assessment.evaluation import Evaluation
@@ -213,7 +213,7 @@ def _run_eval(
         logger.warning(
             "Dry run does not support estimator_options.mode=%s; forcing mode='normal'.", estimator_options.mode.value
         )
-        estimator_options = EstimatorOptions(mode=EstimatorMode.NORMAL, metric=estimator_options.metric)
+        estimator_options = EstimatorOptions(mode=EstimatorMode.NORMAL)
 
     logger.info(f"Loading model template from {model_name}")
     template = ModelTemplate.from_directory_or_github_url(
@@ -227,7 +227,7 @@ def _run_eval(
 
     with template:
         configuration = get_configuration(model_configuration_yaml)
-        estimator: ExternalModel | HpoModel | ExtendedPredictor
+        estimator: ExternalModel | MetaLearner | ExtendedPredictor
         if estimator_options.mode == EstimatorMode.NORMAL:
             estimator = get_estimator(template=template, configuration=configuration)
         elif estimator_options.mode == EstimatorMode.HPO:
@@ -254,14 +254,6 @@ def _run_eval(
             raise ValueError(
                 f"The desired prediction length of {backtest_params.n_periods} is less than the model's minimum prediction length of {model_info.min_prediction_periods}"
             )
-        if (
-            model_info.max_prediction_periods is not None
-            and model_info.max_prediction_periods < backtest_params.n_periods
-        ):
-            logger.warning(
-                f"Wrapping model to extend prediction length from {model_info.max_prediction_periods} to {backtest_params.n_periods}. This is done iteratively, and may worsen model performance"
-            )
-            estimator = ExtendedPredictor(estimator, backtest_params.n_periods)
 
         model_template_db = ModelTemplateDB(
             id=template.model_template_config.name,
@@ -273,6 +265,7 @@ def _run_eval(
             id="cli_eval",
             model_template_id=model_template_db.id,
             model_template=model_template_db,
+            # dumps input/user configuration even in hpo mode, dumps input config.yaml in normal mode even if config doesn't fit and isn't used.
             **configuration.model_dump() if configuration else {},
         )
 
@@ -285,6 +278,15 @@ def _run_eval(
             from chap_core.assessment.dataset_splitting import train_test_generator
             from chap_core.assessment.prediction_evaluator import backtest
 
+            assert not isinstance(estimator, MetaLearner)
+
+            max_periods = estimator.model_information.max_prediction_periods
+            if max_periods is not None and max_periods < backtest_params.n_periods:
+                logger.warning(
+                    f"Wrapping model to extend prediction length from {max_periods} to {backtest_params.n_periods}. This is done iteratively, and may worsen model performance"
+                )
+                estimator = ExtendedPredictor(estimator, backtest_params.n_periods)
+
             train_set, test_generator = train_test_generator(
                 dataset=dataset,
                 prediction_length=backtest_params.n_periods,

chap_core/database/model_templates_and_config_tables.py

@@ -88,7 +88,7 @@ class ModelTemplateInformation(SQLModel):
     hpo_search_space: dict | None = Field(
         default=None,
         sa_column=Column(JSON),
-        description="Search space used by HPO when training this template in `hpo` mode.",
+        description="Search space used by HPO when tuning this template in `hpo` mode.",
     )
     required_covariates: list[str] = Field(
         default_factory=list,