causalig.models

Surrogate models, their training utilities, and a differentiable structural causal model. Surrogate models is the narrative guide, and Adding a new surrogate explains how to write your own.

Surrogates

class causalig.models.MaskedGlobalLinear(adjacency_matrix, predict_only_children=True, use_basis_expansion=True, X_train=None, n_knots=5, knot_power=2, polynomial_degree=3, fourier_features=0, learnable_knots=True)

Sparse linear (optionally basis-expanded) parent→child surrogate.

Reconstructs each node from its direct parents over a directed graph.

Mask convention: parents are rows, children are columns. For child j the prediction aggregates over parents i where mask[i, j] == 1:

pred[j] = bias[j] + sum_{i: mask[i,j]=1} <weight[i,j,:], phi(x_i)>

where phi is the per-feature BasisExpansion (or the identity, i.e. phi(x_i) = x_i, when use_basis_expansion=False). Only edges present in adjacency_matrix carry parameters, so the model is sparse in the number of edges E rather than dense in N^2.

Because the model is linear in phi, compute_input_jacobian() returns the exact input Jacobian in closed form (via derivative()), which the causal-attribution methods use through compute_surrogate_gradients(..., has_closed_form_jacobian=True).

Parameters:

• adjacency_matrix (ndarray) – (N, N) array or tensor; adjacency_matrix[i, j] non-zero means i is a parent of j. Must contain at least one edge.

• predict_only_children (bool) – If True (default), parentless features are output as 0; if False they are copied from the input.

• use_basis_expansion (bool) – Expand each parent with BasisExpansion (nonlinear per-edge functions). False gives a linear model.

• X_train (Tensor | None) – Optional training data (n, N); its per-feature mean and standard deviation standardize the basis inputs. Recommended when use_basis_expansion is True.

• n_knots – Truncated-power terms per feature (basis only).

• knot_power – Power of the truncated-power terms (basis only).

• polynomial_degree – Highest polynomial power (basis only).

• fourier_features – Number of sin/cos pairs (basis only).

• learnable_knots – Train the knot positions (basis only).

forward(x)

Predict every child from its parents: (B, N) -> (B, N).

Return type:

Tensor

compute_input_jacobian(x, dim_start, dim_end)

d pred[:, dim_start:dim_end] / d x → shape [B, chunk, N].

Return type:

Tensor

class causalig.models.surrogate.BasisExpansion(n_features, n_knots=4, knot_power=3, polynomial_degree=2, fourier_features=0, fourier_scale=2.0, learnable_knots=True, knot_init=None, x_loc=None, x_scale=None, activation_clip=8.0)

Per-feature, scale-adaptive basis expansion on standardized inputs.

Components, concatenated along the last dim:

polynomial:       z, z^2, ..., z^polynomial_degree         (D terms)
truncated power:  ReLU(z - c_k)^p   for k=1..n_knots       (K terms)
Fourier:          sin(f·z), cos(f·z)  for f in frequencies (2F terms)

Total n_basis = D + K + 2F.

Inputs are standardized as z = (x - x_loc) / x_scale and clipped to [-activation_clip, activation_clip] before expansion.

Parameters:

• n_features (int) – Number of input features N.

• n_knots (int) – Number of truncated-power terms K per feature.

• knot_power (int) – Power p of the truncated-power terms.

• polynomial_degree (int) – Highest polynomial power D.

• fourier_features (int) – Number of sin/cos frequency pairs F; frequency k is k * pi / fourier_scale.

• fourier_scale (float) – Period scale of the Fourier terms.

• learnable_knots (bool) – Train the knot positions instead of keeping them fixed.

• knot_init (Tensor | None) – Optional (N, K) initial knot positions. Defaults to K evenly spaced values on [-2, 2] for every feature.

• x_loc (Tensor | None) – Per-feature centre (N,) for standardization. Defaults to 0.

• x_scale (Tensor | None) – Per-feature scale (N,) for standardization. Defaults to 1.

• activation_clip (float | None) – Clip standardized inputs to this magnitude; None disables clipping. The derivative is zero where clipping is active.

forward(x)

Expand x of shape (B, N) into basis values (B, N, n_basis).

Return type:

Tensor

derivative(x)

Exact derivative of each basis value with respect to x, (B, N, n_basis).

Return type:

Tensor

Training

causalig.models.train_surrogate_model(surrogate, X, epochs=100, batch_size=256, lr=0.001, weight_decay=1e-05, train_frac=0.8, val_frac=0.1, seed=42, device='cuda', grad_clip_norm=1.0, num_workers=4, logger=None, save_dir=None)

Train a masked graph-structured surrogate (e.g. MaskedGlobalLinear).

Minimizes the MSE between surrogate(x) and x on the columns in surrogate.predictable_children (the children reconstructed from their parents), on unmasked inputs. This is only meaningful for surrogates whose output j reads nothing but the parents of j by construction.

Uses AdamW with a per-step cosine schedule (down to lr / 100), gradient clipping, and keeps the weights with the best validation loss. Batches with non-finite losses or gradients are skipped. weight_decay is not applied to parameters whose name contains bias or knots.

Parameters:

• surrogate – Model to train; must expose predictable_children.

• X – Training data (n_samples, d), a tensor or array (cast to float32).

• epochs – Number of passes over the training split.

• batch_size – Mini-batch size.

• lr – Peak learning rate.

• weight_decay – AdamW weight decay.

• train_frac – Fraction of the rows of X used for training.

• val_frac – Fraction used for validation (model selection); the rest is the test split.

• seed – Seed for the random split.

• device – Training device. Defaults to "cuda".

• grad_clip_norm – Maximum gradient norm.

• num_workers – DataLoader worker processes.

• logger – Logger for per-epoch losses; defaults to this module’s logger.

• save_dir – If given, the best surrogate.weight is saved there as best_masked_weight.npy (requires a weight attribute).

Returns:

The trained surrogate, on device, in eval mode.

Structural causal models

class causalig.models.DifferentiableSCM(adjacency_matrix, assignments, topological_order)

Differentiable Structural Causal Model.

This model implements a DAG-based structural causal model where each node is computed from its parents using assignment functions.

Initialize the SCM.

Parameters:

• adjacency_matrix – adjacency matrix (p x p) as numpy or torch, [i, j] nonzero if feature i causes feature j

• assignments – list of p functions, each taking a tensor of parent values

• topological_order – list of ints representing a valid ordering

forward(X_input=None, structural=False)

Forward pass of the SCM.

Parameters:

• X_input – (batch, p) tensor of exogenous inputs or proposed node values

• structural (bool) –

  • False: standard SCM sampling (nodes computed sequentially from parents)

  • True: structural mode, all nodes treated as independent inputs;

    outputs = direct mechanisms applied to inputs

Returns:

X – (batch, p) tensor of node values

forward_do(X_input, noise, do_nodes=None, do_values=None)

Forward pass with multiple do-interventions.

Parameters:

• X_input – (batch, p) tensor of exogenous inputs

• noise – (batch, p) tensor of noise values

• do_nodes – list or tensor of intervened node indices

• do_values – (batch, len(do_nodes)) tensor with intervention values

Returns:

X – (batch, p) tensor of node values under intervention

causalig.models.create_autoencoder_from_scm(scm, structural=False)

Create an autoencoder wrapper from a Structural Causal Model (SCM).

Parameters:

• scm (Callable) – SCM callable that takes input tensor and returns output tensor

• structural (bool) – Whether to use structural forward pass in SCM. Use True when the result is passed to CIG/CEG as a surrogate: each node’s mechanism is then applied to the parent values it is given. With False, non-root parents are re-simulated from their own (masked) ancestors, so indirect effects are lost on graphs deeper than one edge.

Returns:

Callable autoencoder function

Return type:

Callable