Register a New Model

Add a custom estimator to the COREMat model registry so it participates in all training, HPO, and feature selection workflows automatically.

---

Available groups

Models belong to one of the following groups (used in the Model Zoo and models.level presets):

Group	Examples
ens	XGBoost, CatBoost, LightGBM, Random Forest, ExtraTrees, HGBT
fm	TabPFN variants
lin	Ridge, Lasso, ElasticNet
krr	Kernel Ridge Regression
svm	SVR
nbr	KNN
nn	MLP, Emb-MLP
tree	Decision Tree
cd	Partial Least Squares (cross decomposition)
gp	Gaussian Process Regression

Use one of these strings as the group field of your ModelSpec.

Model ID naming convention

Model IDs follow the schema <group>__<estimator>[__variant], where the third token appears only when one estimator has several variants and names what varies.

Part	Meaning	Example
category	Group prefix matching the group field	lin, ens, nn, fm, svm, krr, nbr, tree
family	Algorithm family within the group	ridge, cb, xgb, lgbm, rf, mlp
variant	Loss function or kernel suffix	squared, absolute, quantile, huber, rbf

Examples from the built-in zoo: lin__ridge, ens__cb__absolute, ens__xgb__squared, ens__lgbm__quantile.

---

Step 1: Create the model module

Create a new file under src/coremat/trainlab/models/:

# src/coremat/trainlab/models/my_gbr.py
from optuna.distributions import FloatDistribution, IntDistribution

from coremat import ModelSpec, PreprocessPolicy, model_description, register


def _make_gbr():
    from sklearn.ensemble import GradientBoostingRegressor
    return GradientBoostingRegressor()


def _param_space():
    # Keys are pipeline-prefixed ("regressor__"); values are Optuna distributions.
    return {
        "regressor__n_estimators": IntDistribution(100, 500),
        "regressor__max_depth": IntDistribution(3, 7),
        "regressor__learning_rate": FloatDistribution(0.01, 0.1, log=True),
    }


register(ModelSpec(
    name="ens__gbr__mse",
    make_estimator=_make_gbr,
    param_space=_param_space,
    supports_optuna=True,
    supports_feature_selection=True,
    preprocess=PreprocessPolicy(scale_numeric="standard", encode_low_card="onehot"),
    shap_explainer="tree",    # "tree" | "linear" | "black-box-permutation"
    group="ens",              # sklearn-taxonomy group; pick the most appropriate one
    description=model_description(
        summary="Scikit-learn Gradient Boosting Regressor (MSE)",
        use_when="Small datasets, interpretable ensemble, good baseline",
        avoid_when="Large datasets (slow training)",
    ),
))

ModelSpec fields

Field	Type	Required	Notes
name	str	Yes	Registry key; schema <category>__<family>__<variant>
make_estimator	callable	Yes	Factory returning a fresh estimator. It may take device and/or threads; it receives only the ones it declares (a factory with no parameters gets none)
param_space	callable	Yes	Returns the HPO search space: {"regressor__<param>": <Optuna distribution>}. Return {} for a model that does not search
supports_optuna	bool	No (default True)	False for a model with nothing to search (no inner loop)
supports_feature_selection	bool	No (default True)	False for models that do not support it
preprocess	PreprocessPolicy	No (default PreprocessPolicy())	Default preprocessing for this model
shap_explainer	str	Yes	"tree" | "linear" | "black-box-permutation" (keyword-only)
group	str	No (default None) — recommended	One of the groups listed above
description	str	No (default None) — recommended	Human-readable description (use model_description)

Only name, make_estimator, param_space and the keyword-only shap_explainer have no default and must be passed.

---

Step 2: Register via side-effect import

Edit src/coremat/trainlab/models/__init__.py to import your new module:

from . import my_gbr   # registers ens__gbr__mse via side-effect

The import alone triggers register(), which adds the model to the global registry.

---

Step 3: Verify

from coremat import get_model_spec, run_training, FrameworkConfig, PathsConfig
from coremat import DataPreprocessingConfig, ModelsConfig

spec = get_model_spec("ens__gbr__mse")
print(spec.group)          # "ens"
print(spec.supports_optuna)

fw = FrameworkConfig(
    paths=PathsConfig(data_file_path="data/alloys.csv", run_name="custom_model_test"),
    data=DataPreprocessingConfig(
        predict_column="yield_strength",
    ),
    models=ModelsConfig(selected=["ens__gbr__mse"]),
)

result = run_training(fw)