Extending IAML

Implement a custom component when your study needs a method outside the existing catalogue. The examples below cover transformations, predictors and the interfaces around training. Save the components in your own project, then connect them through recipes or the corresponding IAML setting.

To combine or configure existing components, follow Build and adapt a pipeline. This page focuses on the Python interfaces a new implementation must satisfy.

The main building blocks

Block

Role

Transformations

Prepare or derive features while preserving observations.

Resamplers

Change the rows used for training, for example to balance classes.

Predictors

Learn from the prepared data and predict outcomes.

Pipeline organization

Define the order of steps and which alternatives to explore.

Metrics

Measure prediction performance and define the search objective.

Validation splitters

Separate training and validation observations according to the study.

Optimizers

Propose new pipeline configurations from evaluated candidates.

Descriptive statistics

Summarize the dataset independently of model performance.

Plots

Present data, performance or explanations as figures.

See Component availability for existing implementations and How it works for the relationships between the core classes.

How to use these examples

Save the definitions below in an importable module named study_components.py. Keep classes and functions at module level so worker processes can import them. Import the module before constructing IAML. Each example explains how it joins the workflow, and the final script connects the regression components.

For steps, @is_step(...) registers the class under tags used by the search. Define an explicit constructor with no required arguments. The decorator initializes the base classes, so the constructor only sets the component’s own configuration and state. Keep constructors lightweight and learn from data in fit, which must reset learned state on each call.

Use descriptive names and unique metric or statistic identifiers. Keep configurable choices in configuration, give defaults within the allowed ranges, and change values through configure. Add descriptions and method references through _description and refs when sharing a component.

Add a transformation

An Actionable can transform features. This example adds body mass index from height in centimetres and weight in kilograms. Missing or nonpositive measurements produce a missing BMI.

import pandas as pd
from iaml import Actionable
from iaml.decorators.is_step import is_step


@is_step("features_precleaning")
class StudyBMI(Actionable):
    name = "Study BMI"
    _description = "Add BMI from height in cm and weight in kg."

    def __init__(self):
        self.optimizable = False

    def suitable(self, dataset):
        return "bmi" not in dataset.X and all(
            column in dataset.X
            and pd.api.types.is_numeric_dtype(dataset.X[column])
            for column in ["height_cm", "weight_kg"]
        )

    def fit(self, dataset):
        return self

    def transform(self, X):
        result = X.copy()
        height = result["height_cm"]
        weight = result["weight_kg"]
        metres = height.where(height > 0) / 100
        result["bmi"] = weight.where(weight > 0) / metres**2
        return result

Connect it. Importing the module registers this transformation with the regular features_precleaning stage. It can also be placed explicitly in a recipe with use(StudyBMI). To insert it into the main default strategy:

from iaml import IAML
from iaml.flow import use

search = IAML()
search.pipeline.add(use(StudyBMI).named("bmi"), before="cleaning")

The anchor finds the sequence containing cleaning. The separate minimal branch remains available; use initial_preprocessor when a transformation must precede every branch.

Design constraints.

  • suitable(dataset) checks the required columns and types without changing the data.

  • fit(dataset) returns self. Learned transformations estimate their state from the supplied training data and reuse it in transform.

  • transform(X) returns a DataFrame with the same rows and index. Copy the input before changing it and keep the output columns consistent at prediction.

Add a resampler

A resampler is also an Actionable, but implements resample(X, y) to change training rows. This example duplicates minority-class observations using the existing imbalanced-learn dependency.

from imblearn.over_sampling import RandomOverSampler
from iaml import Actionable
from iaml.decorators.is_step import is_step


@is_step("imbalance")
class StudyOversampling(Actionable):
    name = "Study oversampling"

    def __init__(self):
        self.sampler = None

    def suitable(self, dataset):
        return dataset.type_of_target in ("binary", "multiclass")

    def fit(self, dataset):
        self.sampler = RandomOverSampler(
            sampling_strategy="minority", random_state=42
        )
        return self

    def resample(self, X, y):
        return self.sampler.fit_resample(X, y)

Connect it. The imbalance tag makes it available to the corresponding search stage. With the default genetic optimizer, resampling alternatives can be introduced by mutation. In an explicit classification workflow, place optional(use(StudyOversampling)) after preprocessing and before the predictor to explore both the resampled and original training data.

Design constraints.

  • Return aligned features and targets, keeping features as a DataFrame.

  • Preserve every supplied column. IAML temporarily includes group identifiers during resampling and extracts them afterwards. This duplication-based example carries those identifiers with their observations.

  • Resampling belongs inside each training fold. Validation and prediction observations retain their original rows.

Add a predictor

Inherit from Predictor to adapt a learning model. This small Ridge wrapper exposes one regularization parameter. IAML already provides Ridge regression, so the example demonstrates the adapter contract.

import numpy as np
import pandas as pd
from sklearn.linear_model import Ridge
from iaml.predictor import Predictor
from iaml.decorators.is_step import is_step


@is_step("predictor", "study_regression")
class StudyRidge(Predictor):
    name = "Study ridge"

    def __init__(self):
        self.configuration = {
            "alpha": {
                "description": "L2 regularization strength.",
                "default": 1.0,
                "range": [0.1, 10.0],
            }
        }

    def suitable(self, dataset):
        return (
            dataset.type_of_target == "continuous"
            and not dataset.X.empty
            and all(pd.api.types.is_numeric_dtype(dtype)
                    for dtype in dataset.X.dtypes)
            and not dataset.X.isna().any().any()
            and np.isfinite(dataset.X.to_numpy(dtype=float)).all()
        )

    def fit(self, dataset):
        self.model = Ridge(**self.passthrough_parameters())
        self.model.fit(dataset.X, dataset.y)
        return self

Connect it. Importing the class makes it available through the predictor tag. Use use(StudyRidge, alpha=0.1) in a recipe, or declare a domain with alpha=Float(0.1, 10.0, initial=1.0). The recipe’s local alternatives define which predictors the genetic optimizer can substitute. The additional study_regression tag allows discovery of just this component family.

Design constraints.

  • suitable declares the supported task and inputs. This wrapper requires finite numerical features and a continuous target.

  • Each fit creates a fresh estimator in self.model and returns self. Use only the supplied training data and preserve feature order.

  • Forward exposed parameters with passthrough_parameters(). Mark any wrapper-only parameter with passthrough=False.

  • The base class delegates prediction to the estimator. Probability or survival outputs are available only when that estimator supports them.

Compose a recipe for the extensions

A reusable recipe derives BMI, imputes missing values, scales features and explores two named Ridge configurations:

from iaml.flow import choice, use
from iaml.steps import SimpleImputer, StandardScaler


def build_study_pipeline():
    models = choice(
        use(StudyRidge, alpha=0.1).named("ridge_low"),
        use(StudyRidge, alpha=10.0).named("ridge_high"),
    ).named("predictor")
    return (
        use(StudyBMI).named("bmi")
        >> use(SimpleImputer).named("cleaning")
        >> use(StandardScaler).named("normalize")
        >> models
    )

Connect it. Pass pipeline=build_study_pipeline() to IAML. Only the declared branches are generated; IAML does not add a minimal strategy to an explicit recipe. The final script uses this factory.

Design constraints.

  • Put feature derivation before imputation and scaling when it needs the original measurements. Every completed branch ends with one predictor.

  • Tags identify eligible components, not execution order. Adding a new tag requires a stage that uses it. Registration alone does not make a step mandatory in every candidate.

See Composition and search scope for ordering and search-scope rules, and Editing methods to modify an attached recipe.

Ordinary transformations and predictors reuse the inherited run method. If a new orchestration strategy needs its own run(candidate), decorate it with @runner from iaml.decorators.runner and return a Candidate or list of candidates. This preserves IAML’s suitability checks and step cache. For a mandatory transformation across the default search branches, IAML(initial_preprocessor=...) accepts a clonable scikit-learn transformer that returns a numerical DataFrame with unchanged rows and index.

Add a metric

A Metric defines the score and its direction. This example implements mean absolute error to show the contract for a custom loss.

from sklearn.metrics import mean_absolute_error
from iaml import Metric


class StudyAbsoluteError(Metric):
    greater_is_better = False

    def __str__(self):
        return "study_absolute_error"

    def suitable(self, X, y, type_of_target):
        return type_of_target == "continuous"

    def compute(self, y, y_pred, **kwargs):
        return float(mean_absolute_error(y, y_pred))

Connect it. Pass the instance main_metric=StudyAbsoluteError() to IAML. Imported metric subclasses are also considered for secondary scores when suitable. To choose exact calculations and name their outputs, use metrics(use(StudyAbsoluteError).named("absolute_error")) and main_metric="absolute_error". See Configure a complete study for analytical collections.

Design constraints. Use a unique string identifier and a constructor with no required arguments. Return a finite scalar and declare the correct greater_is_better direction. Error values remain positive in reports. Metrics use predict by default. Override needed_prediction when another prediction method is required and ensure the selected estimators support it.

Add a validation splitter

A splitter is a callable receiving a Dataset. This example uses three folds with separate patient or site groups.

from sklearn.model_selection import GroupKFold


def study_splitter(dataset):
    if not dataset.has_groups:
        raise ValueError("This study requires patient or site groups.")
    folds = GroupKFold(n_splits=3)
    yield from dataset.split(
        folds.split, groups=dataset.groups.iloc[:, 0]
    )

Connect it. Pass splitter=study_splitter and provide groups to fit. The example requires at least three distinct groups.

Design constraints. Yield pairs of training and validation Dataset objects. Dataset.split adapts index-based splitters and preserves data and group metadata. Use groups explicitly when separation is required and check that training and validation groups are disjoint. Define reproducible seeds for randomized strategies. See Separate test data and cross-validation for study-level validation choices.

Add an optimizer

An Optimizer receives evaluated candidates and proposes new ones. This example keeps the best candidate and explores a small grid around the best StudyRidge pipeline in one optimization round.

from time import monotonic
from iaml import Optimizer
from iaml.search_policy import parameter_keys


class StudyOptimizer(Optimizer):
    def __init__(self, duration=None):
        super().__init__()
        self.deadline = float("inf") if duration is None else monotonic() + duration
        self.done = False

    @property
    def finished(self):
        return self.done or monotonic() >= self.deadline

    def run(self, candidates):
        self.done = True
        if not candidates:
            return []
        proposals = [max(candidates)]
        ridge_candidates = [
            candidate for candidate in candidates
            if candidate.pipeline.predictor is not None
            and isinstance(candidate.pipeline.predictor[1], StudyRidge)
        ]
        if not ridge_candidates:
            return proposals
        seed = max(ridge_candidates)
        predictor = seed.pipeline.predictor[1]
        if "alpha" not in parameter_keys(predictor):
            return proposals
        low, high = predictor.configuration["alpha"]["range"]
        for alpha in (0.1, 1.0, 10.0):
            if monotonic() >= self.deadline:
                break
            if not low <= alpha <= high or alpha == predictor.get_config("alpha"):
                continue
            proposal = seed.to_output()
            proposal.pipeline.predictor[1].configure({"alpha": alpha})
            proposals.append(proposal)
        return proposals

Connect it. Pass the class, optimizer=StudyOptimizer. IAML constructs it with the remaining search duration, or None for an unlimited search.

Design constraints.

  • Expose finished and implement run(candidates). Return proposals for IAML to evaluate, keeping model fitting outside the optimizer.

  • Compare candidates directly or use get_main_metric_score() to respect both maximization and minimization objectives. Retain a strong candidate while exploring alternatives.

  • Copy candidates with to_output() before editing them. Use configure so parameter changes update the cache fingerprints.

  • This example works on the internal candidate API. parameter_keys is the engine’s shared eligibility contract: honor Const and each compiled parameter domain. An optimizer that does not replace components cannot support partial start selections.

  • Bound proposal generation and check the remaining duration. Custom run code is not interrupted at its deadline. Its stopping condition does not limit initial candidate generation or final fitting.

Add a descriptive statistic

A Statistic summarizes the data. This example measures the width between the 10th and 90th percentiles of each numerical feature in a regression dataset.

import pandas as pd
from iaml import Statistic


class StudyCentralWidth(Statistic):
    _description = "Width between the 10th and 90th percentiles."

    def __str__(self):
        return "study_central_width"

    def suitable(self, dataset):
        return (
            dataset.type_of_target == "continuous"
            and not dataset.X.select_dtypes(include="number").empty
        )

    def compute(self, dataset, **kwargs):
        values = dataset.X.select_dtypes(include="number")
        width = values.quantile(0.9) - values.quantile(0.1)
        return pd.DataFrame([width], index=[str(self)])

Connect it. Import the class before calling search.get_descriptive_statistics(X, y). IAML discovers suitable subclasses in its default statistics family. To request this statistic explicitly, pass statistics=statistics(use(StudyCentralWidth).named("central_width")) to IAML. See Describe data before training for selecting statistics and using their results before or after training.

Design constraints. Provide a constructor with no required arguments and implement suitable(dataset). Return a DataFrame with unique statistic identifiers as rows and feature names as columns, without changing the dataset. For statistics grouped by classification label, follow the existing naming convention <feature>_all and <feature>_<class> so descriptive plots can group the columns.

Add a plot

Inherit directly from MetricPlot to add a model performance figure. This example compares observed and predicted values for single-output regression.

import io
import numpy as np
import matplotlib.pyplot as plt
from iaml import MetricPlot


class StudyObservedPredicted(MetricPlot):
    title = "Observed and predicted outcomes"
    description = "Compare predictions with the observed outcomes."

    @classmethod
    def suitable(cls, type_of_target):
        return type_of_target == "continuous"

    def compute(self, estimator, X, y, X_train=None, y_train=None, **kwargs):
        observed = np.asarray(y).ravel()
        predicted = np.asarray(estimator.predict(X)).ravel()
        limits = [min(observed.min(), predicted.min()),
                  max(observed.max(), predicted.max())]
        figure, axes = plt.subplots()
        axes.scatter(observed, predicted)
        axes.plot(limits, limits, "--", color="gray")
        axes.set(xlabel="Observed", ylabel="Predicted", title=self.title)
        self._binary_image = io.BytesIO()
        figure.savefig(self._binary_image, format="png", bbox_inches="tight")
        plt.close(figure)
        return self

Connect it. After import, the figure is included in chosen_model.explain_model_performance(X_test, y_test) for regression. For a configured on-demand collection, pass explanations=explanations(use(StudyObservedPredicted).named("observed")) to IAML and call chosen_model.explain(X_test, y_test). To generate only this figure directly, call its compute method as in the final example.

Design constraints. Inherit directly from MetricPlot for automatic discovery and provide a constructor with no required arguments. Declare the supported target type and accept the training-data keywords shown above. Use the fitted pipeline for prediction, without fitting on evaluation data. Return self after storing PNG bytes in self._binary_image as a BytesIO, and close the Matplotlib figure. The example expects nonempty, finite observed and predicted values.

For descriptive figures, inherit directly from StatisticPlot instead. Its compute(dataframe, **kwargs) receives the statistics table. Check the rows it needs and return an exported figure using the same image contract.

Connect and check the components

Save this separate script beside study_components.py. It uses 48 synthetic observations, three grouped validation folds and one worker. The explicit workflow limits the search to the small Ridge family. The classification resampler is not used in this regression example.

Before reusing an extension, check it on a few observations: expected outputs, row and group alignment, refitting on different data, and a clear rejection of unsupported inputs. Verify configured values reach the underlying estimator and that plotting produces a readable image without changing the fitted model. Continue with 02 / Evaluate for held-out assessment and Study reporting for reporting the methods and settings used in a study.