Adding a New Model

This guide walks through every file you need to create or modify to add a new ML forecasting model to the dashboard. The XGBoost implementation is used as the reference throughout.

Overview of What You Will Create

Step	Layer / Responsibility	File Path	Purpose
1	Model Pipeline	app/models/my_model.py	Implements feature engineering, training, prediction, and pipeline orchestration
2	Configuration	app/configs/MyModel_config_NiceGUI.yaml	Defines data schema, preprocessing, training, and model hyperparameters
3	Model Runner Integration	app/services/model_manager.py	Registers and executes the model within the system pipeline
4	Config Constant	app/configs/constants.py	Stores the config file path as a reusable constant
5	UI Integration	app/ui/control_tab_ui.py	Adds model button and configuration UI components
6	Event Wiring & Execution	app/ui/components/control_tab.py	Connects UI actions to model execution logic
7	Control Logic Integration	app/services/control_tab_logic.py	Adds model case to execution dispatcher

Step 1 — Create the Model Pipeline

Create app/models/my_model.py. Your file must implement three functions with these exact signatures:

add_engineered_features(df, target_col)

Takes a wide-format DataFrame and returns it with additional feature columns appended. At minimum add rolling means, differences, time-of-day encodings, and elapsed minutes — these are expected by the feature pipeline downstream.

def add_engineered_features(df: pd.DataFrame, target_col: str) -> pd.DataFrame:
    """Add engineered features to a wide-format DataFrame.

    Args:
        df (pd.DataFrame): Wide-format input with a DatetimeIndex or
            'time' column.
        target_col (str): Name of the target metric column.

    Returns:
        pd.DataFrame: DataFrame with engineered features appended and
            index reset.

    Raises:
        ValueError: If no datetime index or 'time' column is present.
    """
    if not isinstance(df.index, pd.DatetimeIndex):
        if "time" in df.columns:
            df = df.set_index("time")
        else:
            raise ValueError("df must contain a 'time' column.")

    df[f"{target_col}_roll3"] = df[target_col].rolling(3).mean()
    df["tod_sin"] = np.sin(2 * np.pi * (df.index.hour * 60 + df.index.minute) / 1440)
    df["tod_cos"] = np.cos(2 * np.pi * (df.index.hour * 60 + df.index.minute) / 1440)
    t0 = df.index[0]
    df["experiment_elapsed_min"] = ((df.index - t0).total_seconds() / 60.0)

    return (
        df.replace([np.inf, -np.inf], np.nan)
        .ffill().bfill().fillna(0.0)
        .reset_index()
    )

train_model(df, features, targets, cfg, ...)

Must be an async function. Trains the model and returns a model bundle dict. The bundle must contain at minimum:

{
    "models": { 0: trained_model_h0, 1: trained_model_h1, ... },  # one per horizon
    "scaler": fitted_scaler,
    "y_mean": float,
    "y_std": float,
    "use_log": bool,
    "gompertz_params": dict | None,
    "experiment_start": pd.Timestamp | None,
}

After training, log metrics to the database:

from app.backend.database import log_metrics

await log_metrics(
    "MyModel",
    mse, rmse, r2, spearman,
    mae=mae, mape=mape, direction_accuracy=direction_acc,
    smape=smape, max_error=max_err,
    n_train=len(x_train), n_test=len(x_test),
    metric_type="training",
    experiment_id=experiment_id,
    entity_id=entity_id,
)

Return None if training cannot proceed (not enough data, empty features).

predict_future(model, latest_row_df, features, n_forecast, ...)

Must be an async function. Generates n_forecast future predictions autoregressively and saves each one to the database:

from app.backend.database import save_prediction

await save_prediction(
    entity=entity,
    experiment_id=experiment_id,
    time=future_timestamp,
    predicted_value=predicted_value,
    prediction_time=pd.Timestamp.now().tz_localize(None),
    starting_time=last_time,
    starting_value=last_actual_val,
    model_name="MyModel",
    target_metric=target_col,
)

run_pipeline(cfg, data=None, simulation_mode=False, ...)

Must be an async function with this exact signature:

async def run_pipeline(
    cfg: dict,
    data: pd.DataFrame | None = None,
    simulation_mode: bool = False,
    simulation_df: pd.DataFrame | None = None,
    simulation_experiment_id: str | None = None,
    full_training_df: pd.DataFrame | None = None,
) -> None:

This is the entry point called by model_manager.py. It must:

1. Use data if provided, otherwise fetch via fetch_from_websocket(cfg)
2. Iterate over entities in the data
3. Call preprocess_to_full_wide() from app/backend/ml.py
4. Call add_engineered_features()
5. Call create_lagged_features_targets() from app/backend/ml.py
6. Call train_model() and then predict_future()
7. Call evaluate_predictions() from app/services/prediction_evaluator.py

The preprocessing utilities are already implemented and shared — do not rewrite them:

from app.backend.ml import (
    preprocess_to_full_wide,
    create_lagged_features_targets,
    calculate_microbial_metrics,
    safe_predict,
)
from app.services.prediction_evaluator import evaluate_predictions

Step 2 — Create the Config YAML

Create app/configs/MyModel_config_NiceGUI.yaml. Copy an existing config as a starting point and modify the model section for your architecture. The following top-level sections are required:

data:
  id_column: entity
  metric_column: metric
  target_column: optical_density
  time_column: time
  plot_history_minutes: 20

database:
  path: "metrics.db"

forecast:
  horizon_forecast: 15
  horizon_minutes: 1

lags:
  number: 60

preprocessing:
  fill_method: bfill_ffill
  resample_freq: 1min
  use_log: true
  value_column: value

feature_columns:
  - optical_density
  - temperature
  - agitation_rate
  - optical_density_roll3
  - optical_density_diff1
  - experiment_elapsed_min
  - tod_sin
  - tod_cos

training:
  window_train: 500
  window_test: 50

gompertz:
  enabled: false
  adaptive_window: 300
  refit_every: 10
  min_points: 20
  use_residual_learning: false

evaluation:
  max_horizon: 6
  rolling_window_size: 30
  horizon_unit: "min"
  min_points_for_correlation: 20
  min_variance_for_correlation: 1e-4

model:
  my_model_path: trained_models/my_model.pkl
  # Add your model-specific hyperparameters here

db_config: database_config.yaml

Step 3 — Register the Runner in model_manager.py

Open app/services/model_manager.py and add a runner function following the same pattern as run_xgboost_model:

async def run_my_model(
    config_path,
    use_uploaded_models=True,
    data=None,
    simulation_mode=False,
    simulation_df=None,
    simulation_experiment_id=None,
    full_training_df=None,
) -> None:
    from app.models.my_model import run_pipeline

    cfg = load_config(config_path)

    if use_uploaded_models:
        selected_models = get_selected_models()
        my_models = [m for m in selected_models if "my_model" in m.lower()]
        if my_models:
            cfg["uploaded_models"] = my_models

    if data is None:
        data = await fetch_selected_experiment_data(config_path)

    if data is None or data.empty:
        logger.warning("No data available for MyModel. Aborting.")
        return

    await run_pipeline(
        cfg,
        data=data,
        simulation_mode=simulation_mode,
        simulation_df=simulation_df,
        simulation_experiment_id=simulation_experiment_id,
        full_training_df=full_training_df,
    )

Then add it to run_all_models():

async def run_all_models(...):
    async with model_run_lock:
        # ... existing calls ...
        try:
            logger.info("Starting MyModel...")
            await run_my_model(
                CONFIG_FILE_MY_MODEL,
                use_uploaded_models,
                data=data,
                simulation_mode=simulation_mode,
                simulation_df=simulation_df,
                simulation_experiment_id=simulation_experiment_id,
                full_training_df=full_training_df,
            )
        except Exception as e:
            logger.error(f"MyModel failed: {e}", exc_info=True)

Step 4 — Add the Config Constant

Open app/configs/constants.py and add the config path:

CONFIG_FILE_MY_MODEL = os.path.join(
    PACKAGE_DIR, "configs", "MyModel_config_NiceGUI.yaml"
)

Step 5 — Add a UI Button in control_tab_ui.py

Open app/ui/control_tab_ui.py and find create_model_run(). Add a button for your model in the Execute Individual Models column:

my_model_btn = ui.button(
    "MyModel", icon="auto_awesome", color="secondary"
).props("unelevated").classes("flex-1")

Then attach the reference to the returned namespace object:

mr.my_model_btn = my_model_btn

Also add a configuration panel for your model in create_model_configurations() following the XGBoost or LSTM pattern.

Step 6 — Wire the Button in control_tab.py

Open app/ui/components/control_tab.py and add the click handler:

from app.configs.constants import CONFIG_FILE_MY_MODEL

model_run.my_model_btn.on_click(
    lambda: run_individual_model("my_model", CONFIG_FILE_MY_MODEL)
)

Then open app/services/control_tab_logic.py and add the "my_model" case to run_individual_model():

elif model_type == "my_model":
    asyncio.create_task(run_my_model(config_file, use_uploaded_models=True))

Wire the configuration save button:

model_configurations.my_cfg.save_btn.on_click(
    lambda: save_model_config(model_configurations.my_cfg, CONFIG_FILE_MY_MODEL)
)

Checklist

□ app/models/my_model.py
    □ add_engineered_features() implemented
    □ train_model() is async, returns model bundle, calls log_metrics()
    □ predict_future() is async, calls save_prediction() for each step
    □ run_pipeline() is async, matches required signature

□ app/configs/MyModel_config_NiceGUI.yaml
    □ All required sections present (data, forecast, lags, preprocessing,
      training, gompertz, evaluation, model, db_config)

□ app/configs/constants.py
    □ CONFIG_FILE_MY_MODEL path added

□ app/services/model_manager.py
    □ run_my_model() added following the existing pattern
    □ run_all_models() calls run_my_model()

□ app/ui/control_tab_ui.py
    □ Button added in create_model_run()
    □ Config panel added in create_model_configurations()

□ app/ui/components/control_tab.py
    □ Button click handler wired
    □ Config save button wired

□ app/services/control_tab_logic.py
    □ "my_model" case added to run_individual_model()

After completing all steps, restart the application and your model will appear as a button in Model Control → Run Models, its configuration panel will appear under Model Control → Model Configuration, and its predictions and metrics will populate the Experimental Insight and Performance Metrics pages automatically.