Continuous Treatment Densities

Use this guide when the treatment is numeric and you need either a conditional density \(p(t \mid x)\) or a stabilized density ratio \(p(t \mid x) / p(t)\).

Choose an estimator

from skcausal.utils.lookup import all_density_estimators

all_density_estimators(
    as_dataframe=True,
    return_tags=[
        "density_kind",
        "capability:t_type",
        "supports_multidimensional_treatment",
    ],
)
name object density_kind capability:t_type supports_multidimensional_treatment
0 CompositeFactorizedDensityEstimator <class 'skcausal.density.compose.CompositeFact... conditional [continuous, categorical] None
1 IntegratedMarginalAndConditional <class 'skcausal.density.stabilized_from_condi... stabilized [continuous, categorical] None
2 KernelMarginalAndConditional <class 'skcausal.density.stabilized_from_condi... stabilized [continuous, categorical] None
3 NaiveDensityEstimator <class 'skcausal.density.naive.NaiveDensityEst... conditional [continuous, categorical] None
4 OptunaSearchDensityEstimator <class 'skcausal.density.optuna.OptunaSearchDe... conditional [continuous, categorical] None
5 PermutationWeighting <class 'skcausal.density.permutation_weighting... stabilized [continuous, categorical] None
6 Pipeline <class 'skcausal.density.pipeline.Pipeline'> conditional [continuous, categorical] None
7 SklearnCategoricalDensity <class 'skcausal.density.sklearn.SklearnCatego... conditional [categorical] None
8 SkproDensityEstimator <class 'skcausal.density.skpro.SkproDensityEst... conditional [continuous] None

Example data

The examples below use a single continuous treatment column stored in a polars.DataFrame.

import numpy as np
import polars as pl

rng = np.random.default_rng(0)
n = 400

x0 = rng.normal(size=n)
x1 = rng.normal(size=n)
noise_scale = 0.2 + 0.5 * np.abs(x1)
t_values = 1.5 * x0 - 0.8 * x1 + rng.normal(scale=noise_scale, size=n)

X = pl.DataFrame({"x0": x0, "x1": x1})
t = pl.DataFrame({"t": t_values})

SkproDensityEstimator

SkproDensityEstimator is an adapter around a skpro probabilistic regressor. It fits the wrapped regressor on X and t, calls predict_proba(X), and evaluates the resulting predictive distribution at the supplied treatment values.

The current implementation supports floating-point treatment columns (pl.Float32 and pl.Float64).

from sklearn.ensemble import RandomForestRegressor
from sklearn.linear_model import LinearRegression
from skpro.regression.residual import ResidualDouble

from skcausal.density.skpro import SkproDensityEstimator

regressor = ResidualDouble(
    estimator=LinearRegression(),
    estimator_resid=RandomForestRegressor(n_estimators=25, random_state=0),
)

density_estimator = SkproDensityEstimator(estimator=regressor)
density_estimator.fit(X, t)

conditional_density = density_estimator.predict_density(X, t)
conditional_density[:5]
array([[0.92599426],
       [0.67714325],
       [3.80554233],
       [0.42835012],
       [1.07535137]])

In the test suite, the adapter matches the density obtained by calling the wrapped skpro regressor directly and then evaluating its pdf on the same rows.

PermutationWeighting

PermutationWeighting solves a different problem. It creates a binary classification dataset in which observed (X, t) pairs are labeled 1 and pairs built from observed X with permuted treatment values are labeled 0. The estimator then returns posterior-odds scores from that classifier, which are used as a stabilized density-ratio estimate.

from sklearn.linear_model import LogisticRegression

from skcausal.density.permutation_weighting import PermutationWeighting

ratio_estimator = PermutationWeighting(
    classifier=LogisticRegression(max_iter=1000),
    max_trials=3,
    random_state=0,
)

ratio_estimator.fit(X, t)

stabilized_ratio = ratio_estimator.predict_density(X, t)
stabilized_ratio[:5]
array([[1.],
       [1.],
       [1.],
       [1.],
       [1.]])

The key parameters control how the internal classifier trains and aggregates:

Parameter Default What it does
max_trials 5 Controls how many trial classifiers are fit in "ensemble" or "convergence" mode

Evaluation

For conditional density estimators, skcausal.density.performance_evaluation.evaluate.evaluate runs cross-validation and uses LogLikelihoodMetric() by default.

from sklearn.model_selection import KFold

from skcausal.density.performance_evaluation.evaluate import evaluate
from skcausal.density.performance_evaluation.metrics import LogLikelihoodMetric

results = evaluate(
    estimator=density_estimator,
    cv=KFold(n_splits=3),
    X=X,
    t=t,
    scoring=LogLikelihoodMetric(),
)

results[["test_LogLikelihoodMetric", "fit_time", "score_time"]]
test_LogLikelihoodMetric fit_time score_time
0 -0.672842 0.030195 0.005787
1 -1.983167 0.028946 0.005366
2 -1.325541 0.028996 0.005298

LogLikelihoodMetric() computes the mean of log(predict_density(X, t)) after clipping values below epsilon. That interpretation is natural for SkproDensityEstimator. If you evaluate PermutationWeighting, the metric is still well defined numerically, but it is scoring a stabilized ratio rather than a normalized treatment density.

Multi-column treatments

The continuous estimators in this guide work with a single treatment column. If t contains multiple treatment columns, use CompositeFactorizedDensityEstimator to combine per-column models instead.