flashcart.nn

Build equivariant message-passing blocks from the operations in flashcart.o3. InteractionLayer combines node features with edge directions, applies distance-dependent weights, and sums the messages at receiver atoms. ProductLayer builds higher-order correlations by repeatedly coupling the aggregated node features.

flashcart.nn.layers

Build interaction, product, gating, normalization, and readout layers.

In ProductLayer, correlation controls the polynomial degree of the recurrence, while the l_max parameters limit tensor rank. With path reduction, paths are summed within each product. The contributions from all degrees are added before a final linear map mixes channels. Otherwise, each recurrence step mixes the retained paths and channels with a linear map. Gates, normalization, and scalar readout layers complete the blocks used by FlashCartPotential.

The LinearLayer in this module is a dense feature map. For an equivariant map acting on tensor-rank blocks, use LinearLayer.

class flashcart.nn.layers.RescaledSiLULayer(scale=None)[source]

Bases: Module

SiLU rescaled to preserve the second moment of standard-normal inputs.

Parameters:

scale (float, optional) – Rescaling factor. Default: None, using 1 / sqrt(E[silu(x)**2]) for standard-normal x.

forward(x)[source]

Apply the rescaled SiLU activation elementwise.

Parameters:

x (torch.Tensor) – Input values of any shape.

Returns:

Activated values with the input shape.

Return type:

torch.Tensor

class flashcart.nn.layers.RescaledSigmoidLayer(scale=None)[source]

Bases: Module

Sigmoid rescaled to preserve the second moment of standard-normal inputs.

Parameters:

scale (float, optional) – Rescaling factor. Default: None, using 1 / sqrt(E[sigmoid(x)**2]) for standard-normal x.

forward(x)[source]

Apply the rescaled sigmoid activation elementwise.

Parameters:

x (torch.Tensor) – Input values of any shape.

Returns:

Activated values with the input shape.

Return type:

torch.Tensor

class flashcart.nn.layers.LinearLayer(in_features, out_features, bias=False, weight_scale=None)[source]

Bases: Module

Linear map with input-width normalization.

Weights are initialized from a standard normal distribution and scaled by 1 / sqrt(in_features) during evaluation. The optional bias is initialized to zero and excluded from weight decay.

Parameters:

  • in_features (int) – Input width.

  • out_features (int) – Output width.

  • bias (bool, optional) – Add a bias. Default: False.

  • weight_scale (float, optional) – Scale applied to the weights. Default: None, using 1 / sqrt(in_features).

forward(x)[source]

Apply the normalized linear map.

Parameters:

x (torch.Tensor) – Input features with final dimension in_features.

Returns:

Output features with final dimension out_features and

unchanged leading dimensions.

Return type:

torch.Tensor

class flashcart.nn.layers.InteractionLayer(in1_l_max, in2_l_max, out_l_max, in1_features, out_features, n_radial, hidden_radial=None, n_elements=None, element_features=None, avg_neighbors=1.0, use_sc=False, use_layer_norm=False, use_triton=False, recompute_radial=False)[source]

Bases: Module

Equivariant convolution with distance-dependent tensor-product weights.

A linear map transforms the node features before coupling them with edge attributes. A radial multilayer perceptron provides the coupling weights, optionally conditioned on the sender and receiver elements. Its inputs are multiplied by the cutoff envelope. Messages are summed at each receiver, divided by sqrt(avg_neighbors), and transformed by an equivariant linear map. An optional residual projection and normalization complete the layer.

Parameters:

  • in1_l_max (int) – Maximum tensor rank of the input node features.

  • in2_l_max (int) – Maximum tensor rank of the edge attributes.

  • out_l_max (int) – Maximum tensor rank of the output features.

  • in1_features (int) – Input feature channels per tensor rank.

  • out_features (int) – Output feature channels per tensor rank.

  • n_radial (int) – Number of radial basis functions.

  • hidden_radial (list[int], optional) – Radial-MLP hidden widths. Default: [64, 64].

  • n_elements (int, optional) – Number of elements for element-conditioned radial weights. None is element-agnostic. Default: None.

  • element_features (int, optional) – Width of the sender and receiver element embeddings. Default: None, allowed only when element conditioning is disabled.

  • avg_neighbors (float, optional) – Neighbor-count normalization. Default: 1.0.

  • use_sc (bool, optional) – Skip connection (projected residual). Default: False.

  • use_layer_norm (bool, optional) – RMS layer norm on the output. Default: False.

  • use_triton (bool, optional) – Triton tensor-product kernels. Default: False.

  • recompute_radial (bool, optional) – Recompute the final radial linear output during differentiation instead of retaining that tensor. Hidden activations remain stored. Default: False.

forward(node_feats, edge_attrs, edge_feats, edge_index, envelope, atom_types=None, out_nodes=None)[source]

Aggregate tensor-product messages onto the receiver nodes.

Parameters:

  • node_feats (torch.Tensor) – Node features with shape (n_nodes, (in1_l_max + 1)**2 * in1_features).

  • edge_attrs (torch.Tensor) – Edge attributes with shape (n_edges, (in2_l_max + 1)**2).

  • edge_feats (torch.Tensor) – Radial features with shape (n_edges, n_radial).

  • edge_index (torch.Tensor) – Shape (2, n_edges). Row 0 contains receivers and row 1 contains senders.

  • envelope (torch.Tensor) – Cutoff values with shape (n_edges,).

  • atom_types (torch.Tensor, optional) – Element indices with shape (n_nodes,), required for element-conditioned interactions. Default: None.

  • out_nodes (int, optional) – Number of receiver nodes when the senders include ghost atoms. Default: None, retaining all nodes.

Returns:

Receiver features with shape

(n_receivers, (out_l_max + 1)**2 * out_features), where n_receivers is out_nodes when supplied and otherwise n_nodes.

Return type:

torch.Tensor

set_avg_neighbors(avg_neighbors)[source]

Overwrite the neighbor normalization (called by pre_fit).

Parameters:

avg_neighbors (float) – Positive average neighbor count.

class flashcart.nn.layers.RMSLayerNorm(l_max, n_features, eps=1e-12, affine=True, centering=False, use_triton=False)[source]

Bases: Module

Normalize irreducible Cartesian features using an invariant mean square.

For each tensor rank, a scalar tensor product computes the invariant squared norm. The result is divided by the number of components and averaged over ranks and feature channels. This accounts for the non-orthonormal component layout at ranks of two and above.

Optional centering subtracts the channel mean from the scalar block. Learnable gains act separately on each rank and channel. A scalar bias is included when both affine transformation and centering are enabled.

Parameters:

  • l_max (int) – Maximum tensor rank of the features.

  • n_features (int) – Feature channels per tensor rank.

  • eps (float, optional) – Stabilizer added to the mean square. Default: 1e-12.

  • affine (bool, optional) – Per-rank learnable gains (and an l=0 bias when centering). Default: True.

  • centering (bool, optional) – Mean-center the l=0 block. Default: False.

  • use_triton (bool, optional) – Triton tensor-product kernels. Default: False.

forward(node_feats)[source]

Normalize the features of each node independently.

Parameters:

node_feats (torch.Tensor) – Features with shape (n_nodes, (l_max + 1)**2 * n_features).

Returns:

Normalized features with the input shape.

Return type:

torch.Tensor

class flashcart.nn.layers.EquivariantGatedLayer(l_max, in_features, hidden_features=None, gate_activation='silu', gate_bias=True, use_triton=False)[source]

Bases: Module

Equivariant gated residual layer initialized to the identity.

The input scalar channels produce one gate for each tensor rank and hidden channel. The gated values pass through an output linear map whose weights are initialized to zero. The layer therefore initially returns its input. Sigmoid gates are bounded by their rescaling factor.

Parameters:

  • l_max (int) – Maximum tensor rank of the features.

  • in_features (int) – Feature channels per tensor rank.

  • hidden_features (int, optional) – Width of the gated branch. Default: in_features.

  • gate_activation (str, optional) – “silu” or “sigmoid” (bounded). Default: “silu”.

  • gate_bias (bool, optional) – Include a bias in the linear map that produces the gates. Default: True.

  • use_triton (bool, optional) – Triton kernels. Default: False.

forward(node_feats)[source]

Add the gated equivariant update to the input features.

Parameters:

node_feats (torch.Tensor) – Features with shape (n_nodes, (l_max + 1)**2 * in_features).

Returns:

Updated features with the input shape.

Return type:

torch.Tensor

class flashcart.nn.layers.IrrepsSelector(in_l_max, out_l_max, in_features, in_paths=None)[source]

Bases: Module

Select tensor-rank blocks up to out_l_max.

Parameters:

  • in_l_max (int) – Maximum tensor rank of the input layout.

  • out_l_max (int) – Maximum tensor rank kept (<= in_l_max).

  • in_features (int) – Feature channels per tensor rank.

  • in_paths (list[int], optional) – Paths per tensor rank of the input layout. Default: one path per tensor rank.

forward(x)[source]

Select the leading tensor-rank blocks.

Parameters:

x (torch.Tensor) – Features with shape (n_nodes, in_dim) in increasing rank order.

Returns:

A view containing ranks through out_l_max.

Return type:

torch.Tensor

class flashcart.nn.layers.ProductPathWeights(n_weights, n_elements=None, learnable=False)[source]

Bases: Module

Multiplicative per-path weights 1 + delta for product tensor products.

delta is zero-initialized, so the weights start at exactly one. It is a fixed buffer when not learnable, a shared parameter, or (when n_elements is set) a per-element embedding queried by atom type (element-dependent weights are necessarily learnable).

Parameters:

  • n_weights (int) – Number of weights (paths times features).

  • n_elements (int, optional) – Number of elements for element-dependent weights. Requires learnable. Default: None.

  • learnable (bool, optional) – Train the deltas. Default: False.

forward(atom_types=None)[source]

Return shared or element-dependent coupling weights.

Parameters:

atom_types (torch.Tensor, optional) – Element indices with shape (n_nodes,), required for element-dependent weights. Default: None.

Returns:

Weights with shape (n_weights,) when shared or

(n_nodes, n_weights) when element-dependent.

Return type:

torch.Tensor

class flashcart.nn.layers.ProductLayer(in_l_max, out_l_max, in_features, n_elements=None, correlation=2, path_reduced=False, use_sc=False, use_layer_norm=False, use_triton=False)[source]

Bases: Module

Construct repeated equivariant products of the aggregated node features.

The recurrence starts with a linear contribution and forms terms of polynomial degree up to correlation by multiplying the current features with the original input. Intermediate states retain ranks up to in_l_max. Each contribution to the output retains ranks up to out_l_max.

With path_reduced=True, each tensor product sums its coupling paths and normalizes each output rank by the square root of its path count. The contributions from all degrees are added before a final linear map. Otherwise, each recurrence step uses a linear map to mix its retained paths and channels. Path weights are learnable when paths are reduced or depend on the element. An optional residual projection and normalization complete the layer.

Parameters:

  • in_l_max (int) – Maximum tensor rank of the input features.

  • out_l_max (int) – Maximum tensor rank of the output features.

  • in_features (int) – Feature channels per tensor rank.

  • n_elements (int, optional) – Number of elements for element-dependent path weights. None is element-agnostic. Default: None.

  • correlation (int, optional) – Maximum polynomial degree of the recurrence in the input features. Must be at least 2. Default: 2.

  • path_reduced (bool, optional) – Sum coupling paths inside the kernel. Default: False.

  • use_sc (bool, optional) – Skip connection. Default: False.

  • use_layer_norm (bool, optional) – RMS layer norm on the output. Default: False.

  • use_triton (bool, optional) – Triton kernels. Default: False.

forward(node_feats, atom_types=None)[source]

Sum the contributions from the repeated equivariant products.

Parameters:

  • node_feats (torch.Tensor) – Node features with shape (n_nodes, (in_l_max + 1)**2 * in_features).

  • atom_types (torch.Tensor, optional) – Element indices with shape (n_nodes,), required for element-dependent path weights. Default: None.

Returns:

Output features with shape

(n_nodes, (out_l_max + 1)**2 * in_features).

Return type:

torch.Tensor

class flashcart.nn.layers.EnergyReadoutLayer(in_features, n_layers, hidden_readout=None)[source]

Bases: Module

MLP readout of per-atom energies from every layer’s scalar channels.

Parameters:

  • in_features (int) – Scalar channels contributed by each layer.

  • n_layers (int) – Number of contributing layers.

  • hidden_readout (list[int], optional) – MLP hidden widths. Default: [64, 64].

forward(node_feats)[source]

Predict site energies from the scalar features of all layers.

Parameters:

node_feats (list[torch.Tensor]) – One feature tensor per layer, each with n_nodes rows and at least in_features leading scalar channels.

Returns:

Site energies with shape (n_nodes,).

Return type:

torch.Tensor

class flashcart.nn.layers.ScaleShiftEnergyLayer(n_elements)[source]

Bases: Module

Scale site energies and store the atomic energy shifts.

forward multiplies each site energy by its element-specific scale. The scales and shifts are stored as non-trainable buffers. Training loaders subtract the shifts from reference energies. Calculators may restore them when reporting energies.

Parameters:

n_elements (int) – Number of elements.

forward(node_output, atom_types)[source]

Apply the stored scale for each atom’s element.

Parameters:

  • node_output (torch.Tensor) – Site energies with shape (n_nodes,).

  • atom_types (torch.Tensor) – Element indices with shape (n_nodes,).

Returns:

Scaled site energies with shape (n_nodes,). Atomic shifts

are not added.

Return type:

torch.Tensor

initialize(train_loader, atomic_shifts=None, fit_atomic_shifts=True)[source]

Initialize atomic energy shifts and energy scales from the training set.

The reference force-component RMS supplies a common scale for all elements.

Parameters:

  • train_loader (DataLoader) – Loader over the training set.

  • atomic_shifts (Sequence[float], optional) – Initial per-element shifts. Default: None, using a zero prior when fitting. Required when fit_atomic_shifts=False.

  • fit_atomic_shifts (bool, optional) – Fit residual shifts on top of the prior by ridge-regularized least squares on the training energies. Default: True.

flashcart.nn.radial

Compute radial basis expansions and cutoff functions.

class flashcart.nn.radial.BesselBasis(n_radial, r_max)[source]

Bases: Module

Evaluate a radial basis proportional to spherical Bessel functions.

For n = 1, ..., n_radial, the basis is sqrt(2 / r_max) * sin(n * pi * r / r_max) / r. Evaluation uses the continuous limit at zero. The basis is not masked beyond r_max. A separate cutoff provides that restriction.

Parameters:

  • n_radial (int) – Number of basis functions.

  • r_max (float) – Cutoff radius.

forward(distances)[source]

Evaluate the basis.

Parameters:

distances (torch.Tensor) – Edge lengths, shape (n_edges,).

Returns:

Basis values, shape (n_edges, n_radial).

Return type:

torch.Tensor

class flashcart.nn.radial.ChebyshevBasis(n_radial, r_max, r_min=0.0, mapping='affine')[source]

Bases: Module

Evaluate Chebyshev polynomials T_0..T_{n-1} on mapped distances.

Distances in [r_min, r_max] are mapped to t in [0, 1]. The Chebyshev input is 2*t - 1 for the affine mapping or cos(pi*t) for the cosine mapping. The polynomials are evaluated with the three-term recurrence.

Distances outside the mapping interval are evaluated without clipping. The cutoff envelope is applied separately.

Parameters:

  • n_radial (int) – Number of basis functions.

  • r_max (float) – Upper bound of the mapping interval (cutoff radius).

  • r_min (float, optional) – Lower bound of the mapping interval. Default: 0.0.

  • mapping (str, optional) – “affine” or “cosine” distance mapping. Default: “affine”.

forward(distances)[source]

Evaluate the basis.

Parameters:

distances (torch.Tensor) – Edge lengths, shape (n_edges,).

Returns:

Basis values, shape (n_edges, n_radial).

Return type:

torch.Tensor

class flashcart.nn.radial.PolynomialCutoff(r_max, poly_order=6)[source]

Bases: Module

Polynomial envelope with continuous first and second derivatives at the cutoff.

For t = r / r_max, the envelope is 1 - c1*t**p + c2*t**(p+1) - c3*t**(p+2). The coefficients make the value and its first two derivatives vanish at r_max. Negative distances are treated as zero, and values at or beyond the cutoff are set to zero.

Parameters:

  • r_max (float) – Cutoff radius.

  • poly_order (int, optional) – Polynomial order p (>= 2). Default: 6.

forward(distances)[source]

Evaluate the cutoff envelope (zero at and beyond r_max).

Parameters:

distances (torch.Tensor) – Edge lengths, shape (n_edges,).

Returns:

Envelope values with shape (n_edges,).

Return type:

torch.Tensor

static calculate_envelope(distances, r_max, poly_order, coeff1=None, coeff2=None, coeff3=None)[source]

Evaluate the envelope without a module instance.

Parameters:

  • distances (torch.Tensor) – Distances, clamped to >= 0.

  • r_max (torch.Tensor) – Cutoff radius.

  • poly_order (int) – Polynomial order p.

  • coeff1 (torch.Tensor, optional) – Precomputed (p+1)(p+2)/2. Default: None, computed from poly_order.

  • coeff2 (torch.Tensor, optional) – Precomputed p(p+2). Default: None, computed from poly_order.

  • coeff3 (torch.Tensor, optional) – Precomputed p(p+1)/2. Default: None, computed from poly_order.

Returns:

Envelope values with the same shape as distances.

Return type:

torch.Tensor