Implement your own method

Use this page when you want to plug your own method into skcausal rather than treat the library as a fixed menu of estimators.

Noteskcausal Is Meant To Be Customized

The library is designed around small, composable base classes. You can implement your own synthetic benchmark datasets, treatment density estimators, average-response estimators, and benchmarking metrics without modifying the rest of the library.

What you can extend

Extension point Base class Typical use
Synthetic benchmark dataset BaseSyntheticDataset You know the data-generating process and want ground-truth response curves
Density estimator BaseDensityEstimator You want a new conditional density or stabilized density-ratio model
Average-response estimator BaseAverageCausalResponseEstimator You want a new ADRF or average-response algorithm
Benchmark metric AverageResponseMetric You want to compare estimators with a custom score on truth-known synthetic data

This separation is what makes skcausal flexible: a new estimator can reuse existing datasets and benchmarking utilities, while a new dataset can immediately benchmark existing estimators.

A practical workflow for researchers

  1. Start from the closest extension template rather than from a blank file.
  2. Implement the minimum private hooks required by that base class.
  3. Fit your method against an existing synthetic dataset first, so you can compare against known truth.
  4. Benchmark over multiple random seeds before tuning design choices.
  5. Only then expand to more realistic datasets, larger models, or new metrics.

Extension templates

Each template lives in extension_templates/ and the snippets below are read at render time from the installed skcausal checkout, so the page shows the real current template rather than a shortened copy.

Full template: average_causal_response_estimator.py

"""Extension template for average causal response estimators.

Purpose of this template
------------------------
Use this file as a starting point when adding a new estimator under
``skcausal.causal_estimators``.

How to use this template
------------------------
1. Copy the file to a module with a descriptive name.
2. Rename ``MyAverageCausalResponseEstimator``.
3. Update the module and class docstrings.
4. Set the estimator tags so they match the treatment types and backend you
   actually support.
5. Implement the mandatory private hooks ``_fit`` and ``_predict``.
6. Add ``get_test_params`` so the estimator can be instantiated by automated
   checks and local tests.

Repo-specific contract notes
----------------------------
* ``BaseAverageCausalResponseEstimator.fit`` converts ``X``, ``t``, and ``y``
  to the backend declared in ``_tags`` before calling ``_fit``.
* ``predict`` only accepts treatment values. Average-response estimators are
  expected to average over the covariate sample stored during ``fit``.
* Treatment schema checks happen in the base class, so ``_predict`` receives
  backend-native treatment rows that match the fit-time treatment metadata.
* If your estimator depends on scikit-learn regressors, categorical treatment
  encoding should usually live inside the supplied model or pipeline.
"""

from __future__ import annotations

import numpy as np
import pandas as pd

from skcausal.causal_estimators.base import BaseAverageCausalResponseEstimator

__all__ = ["MyAverageCausalResponseEstimator"]


class MyAverageCausalResponseEstimator(BaseAverageCausalResponseEstimator):
    """Template for average causal response estimators.

    Rename this class and replace the TODO markers with estimator-specific
    logic. The public ``fit`` and ``predict`` methods are inherited from
    ``BaseAverageCausalResponseEstimator``.

    Parameters
    ----------
    some_hyperparameter : float, default=1.0
            Example hyperparameter showing how constructor arguments are exposed to
            ``get_params``/``set_params``.
    """

    _tags = {
        "backend": "pandas",
        "capability:t_type": ["continuous"],
        "capability:multidimensional_treatment": False,
    }

    def __init__(self, some_hyperparameter: float = 1.0):
        self.some_hyperparameter = some_hyperparameter
        super().__init__()

    def _fit(self, X: pd.DataFrame, t: pd.DataFrame, y: pd.DataFrame):
        """Fit the estimator on backend-native inputs.

        Parameters
        ----------
        X : pd.DataFrame
                Covariate matrix already converted to the backend declared in
                ``_tags``.
        t : pd.DataFrame
                Treatment table already checked against the estimator's treatment
                capability tags.
        y : pd.DataFrame
                Outcome table already converted to the estimator backend.

        Returns
        -------
        self
                Fitted estimator.
        """

        # TODO: validate hyperparameters or training data here.
        # TODO: clone and fit any nuisance components here, and store fitted
        # objects in attributes ending with ``_``.
        # TODO: if your estimator needs transformed features, build them from
        # the backend-native ``X`` and ``t`` received here.
        raise NotImplementedError(
            "Replace the template _fit implementation with estimator-specific logic."
        )

    def _predict(self, t: pd.DataFrame) -> np.ndarray:
        """Return one average response per treatment row.

        Parameters
        ----------
        t : pd.DataFrame
                Treatment rows already converted to the estimator backend.

        Returns
        -------
        np.ndarray
                Array with one prediction per row in ``t``. Returning shape
                ``(n_t,)`` or ``(n_t, 1)`` is fine; the base class will coerce the
                output to a 2D column vector.
        """

        fit_X = self._get_fit_X()

        # TODO: use ``fit_X`` and ``t`` to compute one mean response per
        # requested treatment row. Most estimators repeat each treatment row
        # across the fitted covariate sample, evaluate an individual response,
        # and then average over rows.
        _ = fit_X
        raise NotImplementedError(
            "Replace the template _predict implementation with estimator-specific logic."
        )

    @classmethod
    def get_test_params(cls, parameter_set: str = "default"):
        """Return lightweight constructor arguments for local checks.

        Replace this with parameters that instantiate a fast, representative
        version of your estimator.
        """

        return [
            {"some_hyperparameter": 1.0},
            {"some_hyperparameter": 2.0},
        ]

Use this when your method estimates an average response curve such as an ADRF. The base class handles input conversion and treatment-schema checks before calling your private hooks.

Full template: density_estimator.py

"""Extension template for density estimators.

Purpose of this template
------------------------
Use this file as a starting point when adding a new estimator under
``skcausal.density``.

How to use this template
------------------------
1. Copy the file to a module with a descriptive name.
2. Rename ``MyDensityEstimator``.
3. Update the module and class docstrings.
4. Set the estimator tags so they match the treatment types, backend, and
   density quantity you actually support.
5. Implement the mandatory private hooks ``_fit`` and ``_predict_density``.
6. Add ``get_test_params`` so the estimator can be instantiated by automated
   checks and local tests.

Repo-specific contract notes
----------------------------
* ``BaseDensityEstimator.fit`` converts ``X`` and ``t`` to the backend declared
  in ``_tags`` before calling ``_fit``.
* ``predict_density`` must return an out-of-sample score for every input row.
* The estimator may return either the conditional density ``p(t | x)`` or a
  stabilized ratio ``p(t | x) / p(t)``; declare which one you return through
  the ``density_kind`` tag.
* ``_predict_density`` may return shape ``(n_samples,)`` or ``(n_samples, 1)``;
  the base class coerces both to a 2D float array.
"""

from __future__ import annotations

import numpy as np
import pandas as pd

from skcausal.density.base import BaseDensityEstimator

__all__ = ["MyDensityEstimator"]


class MyDensityEstimator(BaseDensityEstimator):
    """Template for conditional or stabilized treatment density estimators.

    Rename this class and replace the TODO markers with estimator-specific
    logic.

    Parameters
    ----------
    some_hyperparameter : float, default=1.0
            Example hyperparameter showing how constructor arguments are exposed to
            ``get_params``/``set_params``.
    """

    _tags = {
        "backend": "pandas",
        "capability:t_type": ["continuous"],
        "capability:multidimensional_treatment": False,
        "density_kind": "conditional",
        "soft_dependencies": [],
    }

    def __init__(self, some_hyperparameter: float = 1.0):
        self.some_hyperparameter = some_hyperparameter
        super().__init__()

    def _fit(self, X: pd.DataFrame, t: pd.DataFrame):
        """Fit the estimator on backend-native covariates and treatments.

        Parameters
        ----------
        X : pd.DataFrame
                Covariate matrix already converted to the backend declared in
                ``_tags``.
        t : pd.DataFrame
                Treatment table already checked against the estimator's treatment
                capability tags.

        Returns
        -------
        self
                Fitted estimator.
        """

        # TODO: validate hyperparameters or preprocess X/t here.
        # TODO: clone and fit any wrapped models here, and store fitted objects
        # in attributes ending with ``_``.
        raise NotImplementedError(
            "Replace the template _fit implementation with estimator-specific logic."
        )

    def _predict_density(self, X: pd.DataFrame, t: pd.DataFrame) -> np.ndarray:
        """Return one density-like score per row.

        Parameters
        ----------
        X : pd.DataFrame
                Covariate matrix already converted to the estimator backend.
        t : pd.DataFrame
                Treatment rows already converted to the estimator backend.

        Returns
        -------
        np.ndarray
                One density or stabilized-ratio value per input row.
        """

        # TODO: compute one out-of-sample score for each row in ``X`` and ``t``.
        raise NotImplementedError(
            "Replace the template _predict_density implementation with estimator-specific logic."
        )

    @classmethod
    def get_test_params(cls, parameter_set: str = "default"):
        """Return lightweight constructor arguments for local checks."""

        return [
            {"some_hyperparameter": 1.0},
            {"some_hyperparameter": 2.0},
        ]

Use this when your estimator produces either a conditional treatment density p(t | x) or a stabilized ratio such as p(t | x) / p(t).

Full template: synthetic_dataset.py

"""Extension template for synthetic datasets.

Purpose of this template
------------------------
Use this file as a starting point when adding a new dataset under
``skcausal.datasets``.

How to use this template
------------------------
1. Copy the file to a module with a descriptive name.
2. Rename ``MySyntheticDataset``.
3. Update the module and class docstrings.
4. Define ``column_types`` for the treatment table using ``"continuous"`` or
    ``"categorical"`` values.
5. Implement ``_get_covariates``, ``_get_treatments``, and ``_predict_y``.
6. Call ``self._prepare(self.n)`` in ``__init__`` if ``load()`` should work
   immediately after construction.
7. Add ``get_test_params`` so the dataset can be instantiated in local tests.

Repo-specific contract notes
----------------------------
* ``BaseSyntheticDataset.load`` returns covariates, treatments, and outcomes as
  backend-native frames.
* ``predict_y`` should expose the noiseless structural response and accept the
  treatment tables returned by ``load``.
* ``predict_curve`` is already implemented in the base class and averages
  ``predict_y`` over a treatment grid.
* Continuous datasets typically expose ``get_grid``; categorical datasets often
  replace it with a helper such as ``get_levels``.
"""

from __future__ import annotations

import numpy as np
import polars as pl

from skcausal.datasets.base import BaseSyntheticDataset

__all__ = ["MySyntheticDataset"]


class MySyntheticDataset(BaseSyntheticDataset):
    """Template for synthetic datasets with a known response surface.

    Rename this class and replace the TODO markers with dataset-specific logic.

    Parameters
    ----------
    n : int, default=1000
            Number of rows generated by the dataset.
    outcome_noise : float, default=1.0
            Scale parameter for the default Gaussian noise model.
    random_state : int, default=0
            Seed forwarded to the dataset RNG.
    """

    column_types = {"t_0": "continuous"}

    def __init__(
        self,
        n: int = 1000,
        outcome_noise: float = 1.0,
        random_state: int = 0,
    ):
        self.outcome_noise = outcome_noise

        super().__init__(n=n, random_state=random_state)

        # TODO: create any dataset-specific random coefficients here before
        # calling ``_prepare``.
        self._prepare(self.n)

    def _get_covariates(self) -> np.ndarray:
        """Return the covariate matrix used to generate the dataset.

        Returns
        -------
        np.ndarray
                Array with ``self.n`` rows. Named-column datasets may instead
                return a ``polars.DataFrame``.
        """

        # TODO: generate the observed covariates.
        raise NotImplementedError(
            "Replace the template _get_covariates implementation with dataset-specific logic."
        )

    def _get_treatments(self, covariates) -> np.ndarray | pl.DataFrame:
        """Return one treatment row per covariate row.

        Parameters
        ----------
        covariates : np.ndarray or pl.DataFrame
                Covariates returned by ``_get_covariates``.

        Returns
        -------
        np.ndarray or pl.DataFrame
                Treatment assignments. If you return a numpy array, make sure it is
            compatible with ``column_types``.
        """

        # TODO: sample or deterministically generate treatments from the
        # covariates.
        raise NotImplementedError(
            "Replace the template _get_treatments implementation with dataset-specific logic."
        )

    def _predict_y(self, covariates, treatments) -> np.ndarray:
        """Return the noiseless structural response ``E[Y | X, T]``.

        Parameters
        ----------
        covariates : np.ndarray or pl.DataFrame
                Covariates supplied by the public ``predict_y`` wrapper.
        treatments : np.ndarray or pl.DataFrame
                Treatment table supplied by the public ``predict_y`` wrapper.

        Returns
        -------
        np.ndarray
                One mean outcome per row in ``treatments``.
        """

        # TODO: implement the structural response surface.
        raise NotImplementedError(
            "Replace the template _predict_y implementation with dataset-specific logic."
        )

    def _get_outcomes(self, covariates, treatments) -> np.ndarray:
        """Optionally override the observation noise model.

        The default implementation here adds homoskedastic Gaussian noise to the
        structural response. Replace or remove this method if your dataset uses
        a different sampling model.
        """

        expected_outcomes = np.asarray(
            self.predict_y(covariates=covariates, treatments=treatments),
            dtype=float,
        )
        return expected_outcomes + self._rng.normal(
            scale=self.outcome_noise,
            size=expected_outcomes.shape,
        )

    def get_grid(self, n: int = 100) -> pl.DataFrame:
        """Return a default treatment grid for response-curve evaluation.

        Replace this helper for continuous treatments, or swap it for a method
        such as ``get_levels`` when the treatment is categorical.
        """

        return self._coerce_treatment_frame(
            pl.DataFrame({"t_0": np.linspace(0.0, 1.0, n)})
        )

    @classmethod
    def get_test_params(cls, parameter_set: str = "default"):
        """Return lightweight constructor arguments for local checks."""

        return [
            {"n": 32, "random_state": 0},
            {"n": 16, "outcome_noise": 0.5, "random_state": 1},
        ]

Use this when you want a controlled benchmark with known truth. Once your dataset subclasses BaseSyntheticDataset, existing estimator benchmarks can score against its true response curve.

Use AI with the templates

Tip

AI is a good accelerator here, especially for the boring first pass. Give the model the right template, the mathematical definition of your method, and the constraints you need it to preserve.

A practical prompt looks like this:

Here is the skcausal extension template for BaseAverageCausalResponseEstimator.
Fill in the TODOs for my algorithm, but do not change the public method
signatures, the base class, or the tag names. My estimator should:

- fit an outcome model m(x, t)
- optionally use a density model for weighting
- return one average response per treatment row in _predict
- include lightweight get_test_params values

Method description:
[paste your identifying assumptions, nuisance models, and prediction formula here]

Use AI to scaffold, not to skip validation. After it fills the template, run the estimator on a synthetic benchmark and check whether the returned curve has the right shape, scale, and seed-to-seed stability.

Benchmark a new average-response estimator

If your algorithm subclasses BaseAverageCausalResponseEstimator, you can benchmark it directly on any synthetic dataset that subclasses BaseSyntheticDataset.

from skcausal.causal_estimators.benchmarking import evaluate_multiple_dataset_seeds
from skcausal.causal_estimators.benchmarking.metrics import MAE, RMSE

dataset = MySyntheticDataset(n=2000, random_state=0)
estimator = MyAverageCausalResponseEstimator(...)

results = evaluate_multiple_dataset_seeds(
    dataset=dataset,
    estimator=estimator,
    metrics=[MAE(n_treatments=128), RMSE(n_treatments=128)],
    random_states=[0, 1, 2, 3, 4],
)

summary = results.drop(columns=["dataset_seed"]).mean().to_frame("mean_score")
summary

This helper clones the dataset across random seeds, fits a fresh clone of your estimator on each seed, and returns one row per seed. The built-in average-response metrics compare your estimator’s predicted curve against the dataset’s known truth.

Important

Average-response benchmarking requires a BaseSyntheticDataset because the metric needs access to the true response function, not just observed outcomes.

Benchmark a new density estimator

If your new component subclasses BaseDensityEstimator, you can evaluate it with cross-validation using the density benchmarking utility.

from sklearn.model_selection import KFold

from skcausal.density.performance_evaluation.evaluate import evaluate
from skcausal.density.performance_evaluation.metrics.likelihood import (
    LogLikelihoodMetric,
)

X, t, _ = dataset.load()
density_estimator = MyDensityEstimator(...)

density_results = evaluate(
    estimator=density_estimator,
    cv=KFold(n_splits=5, shuffle=True, random_state=0),
    X=X,
    t=t,
    scoring=[LogLikelihoodMetric()],
)

density_results

This returns one row per fold with the requested metric plus fit and score times. That makes it easy to compare density-model quality and computational cost side by side.

Where to go next