API reference

Model

class net.e3gnn.E3GNN(mapper, cfg, *, legacy_separate_shifted_self=False)

Full network and training logic.

Forward returns a dict with three IrrepsBlockData entries (hamiltonian / overlap / density). Losses are computed inside training_step.

Parameters:

• mapper (BlockIrrepMapper)

• cfg (Config)

• legacy_separate_shifted_self (bool)

fit_training_data(dataset)

Fit declared train-only transforms once, using the dataset contract.

predict_matrices(x, *, physical=True)

Run inference and return predicted sparse block matrices.

This is a convenience wrapper around forward and predicted_irreps_to_block_matrices. It does not move x between devices; callers should prepare the input on the model’s device.

Parameters:

• x (Dict[str, Any])

• physical (bool)

Return type:

Dict[str, BlockMatrix]

class net.common.Config(pipeline=None, objective=None, optimization=None, schema_version=1, data=<factory>, training=<factory>, parameters=<factory>, species=<factory>, layout=<factory>, layout_rules=<factory>, codec=<factory>, reference=<factory>, cutoff_radius=11.0, l_max=6, hidden_base_dim=64, hidden_irreps='128x0e+16x0o+8x1e+64x1o+24x2e+8x2o+8x3e+24x3o+16x4e+4x4o+4x5e+12x5o+8x6e', edge_type_emb_dim=32, emb_use_odd_features=True, edge_encoder_style='rich', edge_encoder_use_sh_tensor_square=True, e3layernorm=False, num_layers_gnn=2, tp_type='separate_weight', use_self_connection=False, edge_update_node_combine='concat', edge_update_residual=False, node_update_message_agg='attention', node_update_attention_scalar_dim=64, node_update_attention_heads=4, node_update_residual=False, head_use_mlp_log_scale=False, neck_depth=2, internal_e3mlp_layers=1, head_e3mlp_layers=1, head_use_node_embeddings_for_self_edges=True, separate_shifted_self=False, head_use_tensor_square=False, head_pair_mode='split', head_diag_output_scale=1.0, head_offdiag_output_scale=1.0, head_log_scale_mlp_n_layers=1, e3mlp_variant='film', internal_e3mlp_variant='resnormact', head_e3mlp_variant='normact', e3mlp_output_scale=1.0, e3mlp_weight_init_scale=1.0, e3mlp_residual_scale=0.25, e3mlp_pre_norm=False, e3mlp_norm_eps=1e-08, e3mlp_film_hidden_dim=64, activation_odd_scalar='tanh', activation_odd_gate='tanh', nonlin_kind='normact', activation_scalar='leakyrelu', activation_gate='softplus', s2act_res=128, norm_kind='component', lr=0.010803383053097332, max_epochs=10000, max_wall_clock_seconds=43200.0, batch_size=1, use_lr_scheduler=True, lr_scheduler_factor=0.5, lr_scheduler_patience=60, lr_scheduler_min_lr=1e-08, lr_scheduler_target='val/hamiltonian_mae', revert_on_spike=True, revert_monitor='val/hamiltonian_mae', revert_decay_patience=4, revert_decay_rate=0.5, revert_spike_factor=2.0, checkpoint_monitor='val/hamiltonian_mae', dropout=0.0, l1_reg_coef=0.0, l2_reg_coef=0.0, init_weights_factor=1.0, grad_clip_algorithm='value', wall_clock_scope='session', grad_clip_val=2.0, accumulate_grad_batches=1, n_radial=128, radial_layers=<factory>, radial_weight_granularity='term', pair_conditioned_radial_mlp=False, pair_distance_normalization='off', enable_forces=False, enable_stress=False, enable_energy=True, enable_num_electrons=True, train_target='matrix', partial_train=None, train_on_forces=False, train_on_stress=False, train_on_energy=False, train_on_num_electrons=False, matrix_targets=<factory>, train_observables_on_gt=True, symmetrize_validation_output=True, symmetrize_overlap_targets=True, symmetrize_density_targets=True, symmetrize_output=False, symmetrize_hamiltonian_targets=True, rescale_density_to_num_electrons=False, loss_l1_fraction=1.0, matrix_loss_reduction='pair_sum', loss_coef_observables=0.0, allow_zero_observable_loss_control=False, observable_loss_kind='mae', loss_coef_forces=0.0, loss_coef_stress=0.0, hamiltonian_envelope_mode='off', hamiltonian_envelope_path=None, hamiltonian_envelope_eps=1e-12, loss_weighting_mode='off', loss_weight_min=0.0, loss_weight_max=100.0, spectral_loss_enabled=False, train_on_spectral=True, spectral_loss_coef=0.0, spectral_loss_kmesh='1x1x1', spectral_loss_window_ev=10.0, spectral_loss_taper_ev=2.0, spectral_loss_huber_delta_ev=0.1, spectral_loss_kind='huber', spectral_loss_overlap_psd_cleanup=False, spectral_loss_overlap_jitter=True, spectral_fermi_cache_path=None, freeze_backbone_train_heads_only=False, run_name='mandala-run', experiment_id=None, ablation_name=None, ablation_setting=None, ablation_setting_from=None, verbosity=1, bench_verbosity=1, log_partial_gt_observables=True, log_per_irrep_metrics=True, print_per_irrep_metrics=False, log_per_irrep_images=False, log_train_metrics=False, log_per_pair_loss_metrics=True, log_hamiltonian_irrep_contrib_metrics=True, log_hamiltonian_pair_contrib_metrics=True, log_activation_mag=False, log_grad_norm=False, wandb_project=None, log_every_n_steps=1, log_on_step=False, log_on_epoch=True, log_data=False, log_forward=False, benchmark=False, log_interval=10, adaptive_log_interval=True, video_max_atoms=6, safety_checks=False, dtype=torch.float32, device='cuda', gpus=1, num_workers=3, save_dir='checkpoints', log_model=False, snapshot_cache_dir=None, dataset_device=None, shuffle_snapshot_load_order=True, data_split_seed=42, seed=42, epoch_boundary_resume=False, precompute_edge_features=True, radial_embedding_scale='none', allow_incomplete_dataset=False, apply_cutoff_to_targets=True, graph_source='geometry_cutoff', require_exact_edge_match=True, allow_openmx_positions_box_from_out=False, tune=None, train_on_irrep_parts=False, compatibility=False, num_block=3, r_max=8.0, num_basis=128, use_sc=True, use_sbf=False)

Network hyperparameters controlling representations, model depth, nonlinearity, regularization, radial basis, output heads, and loss weighting.

Parameters:

• pipeline (dict | None)

• objective (dict | None)

• optimization (dict | None)

• schema_version (int)

• data (dict)

• training (dict)

• parameters (dict)

• species (dict)

• layout (dict)

• layout_rules (dict)

• codec (dict)

• reference (dict)

• cutoff_radius (float)

• l_max (int)

• hidden_base_dim (int)

• hidden_irreps (str | None)

• edge_type_emb_dim (int)

• emb_use_odd_features (bool)

• edge_encoder_style (str)

• edge_encoder_use_sh_tensor_square (bool)

• e3layernorm (bool)

• num_layers_gnn (int)

• tp_type (str)

• use_self_connection (bool)

• edge_update_node_combine (str)

• edge_update_residual (bool)

• node_update_message_agg (str)

• node_update_attention_scalar_dim (int)

• node_update_attention_heads (int)

• node_update_residual (bool)

• head_use_mlp_log_scale (bool)

• neck_depth (int)

• internal_e3mlp_layers (int)

• head_e3mlp_layers (int)

• head_use_node_embeddings_for_self_edges (bool)

• separate_shifted_self (bool)

• head_use_tensor_square (bool)

• head_pair_mode (str)

• head_diag_output_scale (float)

• head_offdiag_output_scale (float)

• head_log_scale_mlp_n_layers (int)

• e3mlp_variant (str)

• internal_e3mlp_variant (str | None)

• head_e3mlp_variant (str | None)

• e3mlp_output_scale (float)

• e3mlp_weight_init_scale (float)

• e3mlp_residual_scale (float)

• e3mlp_pre_norm (bool)

• e3mlp_norm_eps (float)

• e3mlp_film_hidden_dim (int)

• activation_odd_scalar (str)

• activation_odd_gate (str)

• nonlin_kind (str)

• activation_scalar (str)

• activation_gate (str)

• s2act_res (int)

• norm_kind (str)

• lr (float)

• max_epochs (int)

• max_wall_clock_seconds (float | None)

• batch_size (int)

• use_lr_scheduler (bool)

• lr_scheduler_factor (float)

• lr_scheduler_patience (int)

• lr_scheduler_min_lr (float)

• lr_scheduler_target (str)

• revert_on_spike (bool)

• revert_monitor (str | None)

• revert_decay_patience (int)

• revert_decay_rate (float)

• revert_spike_factor (float)

• checkpoint_monitor (str)

• dropout (float)

• l1_reg_coef (float)

• l2_reg_coef (float)

• init_weights_factor (float)

• grad_clip_algorithm (str)

• wall_clock_scope (str)

• grad_clip_val (float | None)

• accumulate_grad_batches (int)

• n_radial (int)

• radial_layers (Sequence[int])

• radial_weight_granularity (str)

• pair_conditioned_radial_mlp (bool)

• pair_distance_normalization (str)

• enable_forces (bool)

• enable_stress (bool)

• enable_energy (bool)

• enable_num_electrons (bool)

• train_target (str)

• partial_train (str | None)

• train_on_forces (bool)

• train_on_stress (bool)

• train_on_energy (bool)

• train_on_num_electrons (bool)

• matrix_targets (list)

• train_observables_on_gt (bool)

• symmetrize_validation_output (bool)

• symmetrize_overlap_targets (bool)

• symmetrize_density_targets (bool)

• symmetrize_output (bool)

• symmetrize_hamiltonian_targets (bool)

• rescale_density_to_num_electrons (bool)

• loss_l1_fraction (float)

• matrix_loss_reduction (str)

• loss_coef_observables (float)

• allow_zero_observable_loss_control (bool)

• observable_loss_kind (str)

• loss_coef_forces (float)

• loss_coef_stress (float)

• hamiltonian_envelope_mode (str)

• hamiltonian_envelope_path (str | None)

• hamiltonian_envelope_eps (float)

• loss_weighting_mode (str)

• loss_weight_min (float)

• loss_weight_max (float)

• spectral_loss_enabled (bool)

• train_on_spectral (bool)

• spectral_loss_coef (float)

• spectral_loss_kmesh (str)

• spectral_loss_window_ev (float)

• spectral_loss_taper_ev (float)

• spectral_loss_huber_delta_ev (float)

• spectral_loss_kind (str)

• spectral_loss_overlap_psd_cleanup (bool)

• spectral_loss_overlap_jitter (bool)

• spectral_fermi_cache_path (str | None)

• freeze_backbone_train_heads_only (bool)

• run_name (str)

• experiment_id (str | None)

• ablation_name (str | None)

• ablation_setting (str | None)

• ablation_setting_from (str | None)

• verbosity (int)

• bench_verbosity (int)

• log_partial_gt_observables (bool)

• log_per_irrep_metrics (bool)

• print_per_irrep_metrics (bool)

• log_per_irrep_images (bool)

• log_train_metrics (bool)

• log_per_pair_loss_metrics (bool)

• log_hamiltonian_irrep_contrib_metrics (bool)

• log_hamiltonian_pair_contrib_metrics (bool)

• log_activation_mag (bool)

• log_grad_norm (bool)

• wandb_project (str | None)

• log_every_n_steps (int)

• log_on_step (bool)

• log_on_epoch (bool)

• log_data (bool)

• log_forward (bool)

• benchmark (bool)

• log_interval (int)

• adaptive_log_interval (bool)

• video_max_atoms (int | None)

• safety_checks (bool)

• dtype (dtype)

• device (str)

• gpus (int)

• num_workers (int | None)

• save_dir (str)

• log_model (bool)

• snapshot_cache_dir (str | None)

• dataset_device (str | None)

• shuffle_snapshot_load_order (bool)

• data_split_seed (int)

• seed (int)

• epoch_boundary_resume (bool)

• precompute_edge_features (bool)

• radial_embedding_scale (str)

• allow_incomplete_dataset (bool)

• apply_cutoff_to_targets (bool)

• graph_source (str)

• require_exact_edge_match (bool)

• allow_openmx_positions_box_from_out (bool)

• tune (str | None)

• train_on_irrep_parts (bool)

• compatibility (bool)

• num_block (int)

• r_max (float)

• num_basis (int)

• use_sc (bool)

• use_sbf (bool)

classmethod from_mapping(document)

Load component recipes and flat settings through the same Config.

to_mapping()

Plain serializable configuration for either graph representation.

classmethod resolve(source, *, species_shells, spinful=None, overrides=None)

net.configuration.resolve_configuration(source, *, species_shells, spinful=None, overrides=None)

Return (plain_resolved_configuration, native_shell_codec).

source is a YAML filename or mapping. overrides is a nested mapping merged before interpolation. Species order, shell channels, dimensions and optional declared irrep-layout rules are then resolved. The returned codec uses the caller’s current default torch dtype; set it before resolution.

layout_rules supports channel_harmonics and tensor_product_completion. They can be used by arbitrary custom recipes.

Parameters:

• source (str | Path | Mapping)

• species_shells (Mapping)

• spinful (bool | None)

• overrides (Mapping | None)

Pipeline

class net.pipeline.FeatureSpec(domain: 'str', representation: 'str' = 'scalar', irreps: 'str | None' = None, frame: 'str' = 'global', width: 'int | None' = None, support: 'str | None' = None, angular_basis: 'str | None' = None, batch: 'dict | None' = None)

Parameters:

• domain (str)

• representation (str)

• irreps (str | None)

• frame (str)

• width (int | None)

• support (str | None)

• angular_basis (str | None)

• batch (dict | None)

class net.pipeline.FeaturePipeline(specification, *, num_species, mapper=None)

A validated sequence; outputs never overwrite input or saved slots.

Parameters:

specification (Mapping)

net.pipeline.register_transform(name)

Register an operation factory; graph compilation never selects a model.

net.pipeline.build_transform(name, options, inputs, *, mapper=None)

Return (module, output specs). Inputs are positional, never inferred by name.

net.component_defaults.default_graph(cfg, num_species)

Layers

class net.encoders.NodeEncoder(node_one_hot_dim, cfg, info=None)

Node features = only element one-hot -> learned embedding (scalars).

The embedding is mapped (Linear) to irreps_out (scalars only in our current setup, but we keep it general).

Parameters:

• node_one_hot_dim (int)

• cfg (Config)

• info (dict)

class net.encoders.EdgeEncoder(n_edge_types, irreps_out, cfg, info=None)

Encode per-edge raw features into hidden irreps.

Inputs:

• edge_type_idx: LongTensor[E] of edge-type indices (n_edge_types).

• length_emb: Tensor[E, n_radial] radial distance embeddings.

• sh: Tensor[E, sh_irreps.dim] spherical harmonics coefficients.

Combines:

1. Learned edge-type embedding.

2. Radial distance embeddings.

3. Spherical harmonics coefficients.

Tensor[E, irreps_out.dim].

Parameters:

• n_edge_types (int)

• irreps_out (Irreps)

• cfg (Config)

• info (dict)

class net.layers.EquiConv(n_radial, irreps_in1, irreps_in2, irreps_out, cfg, nonlin=True, info=None, n_edge_types=None)

E(3)-equivariant convolution with radial MLP weighting.

Similar to DeepH-E3’s EquiConv. Combines: - Tensor product: irreps_in1 (x) irreps_in2 -> irreps_out - Radial MLP: generates weights from edge length embeddings - Optional Gate nonlinearity - Element-wise multiplication with learned weights

Args:

n_radial: Input dimension for radial MLP (edge length embedding size) irreps_in1: Input irreps for geometric features (node/edge concatenation) irreps_in2: Input irreps for spherical harmonics irreps_out: Output irreps cfg: Configuration object nonlin: Whether to use Gate nonlinearity after TP

Parameters:

• n_radial (int)

• irreps_in1 (Irreps)

• irreps_in2 (Irreps)

• irreps_out (Irreps)

• cfg (Config)

• nonlin (bool)

• info (dict | None)

• n_edge_types (int | None)

class net.layers.NodeUpdateBlock(node_irreps, edge_irreps, num_species, cfg, info=None, node_irreps_out=None, n_edge_types=None)

Node update block following DeepH-E3 architecture.

Aggregates edge messages to nodes:

• Concatenate node[i], node[j], and edge features

• Apply EquiConv to create edge messages

• Scatter-aggregate messages to destination nodes

• Apply transformations and optional self-connection

Parameters:

• node_irreps (Irreps)

• edge_irreps (Irreps)

• num_species (int)

• cfg (Config)

• info (dict)

• node_irreps_out (Irreps)

• n_edge_types (int | None)

class net.layers.EdgeUpdateBlock(node_irreps, edge_irreps, num_species, cfg, info=None, edge_irreps_out=None, n_edge_types=None)

Edge update block following DeepH-E3 architecture.

For each edge (i -> j):

• Concatenate node[i], node[j], and edge features

• Apply EquiConv with spherical harmonics and radial weighting

• Apply linear transformation and nonlinearity

• Optional self-connection and residual

Parameters:

• node_irreps (Irreps)

• edge_irreps (Irreps)

• num_species (int)

• cfg (Config)

• info (dict)

• edge_irreps_out (Irreps | None)

• n_edge_types (int | None)

class net.heads.DeepHead(irreps_diag_in, irreps_edge_in, irreps_neck, pair_keys, mapper, cfg, info=None)

Onsite and pair readouts; the separate shifted-self path is legacy-only.

New trainings route every non-onsite block through the pair readout. Saved configurations can reconstruct the historical modules unchanged.

Parameters:

• irreps_diag_in (Irreps)

• irreps_edge_in (Irreps)

• irreps_neck (Irreps)

• pair_keys (List[str])

• mapper (BlockIrrepMapper)

• cfg (Config)

• info (dict)

class net.common.E3MLP(input_irreps, hidden_irreps, output_irreps, num_layers, cfg, activate_last=False, *, variant=None, output_scale=None, weight_init_scale=None, residual_scale=None, pre_norm=None, film_hidden_dim=None, post_scale=1.0)

A generic N-layer E(3)-equivariant feed-forward network.

This module is a sequence of e3nn.o3.Linear layers alternating with equivariant non-linearities.

Parameters:

• input_irreps (e3nn.o3.Irreps) – Input irreducible representations.

• hidden_irreps (e3nn.o3.Irreps) – Irreducible representations for the hidden layers. This is ignored if num_layers is 1.

• output_irreps (e3nn.o3.Irreps) – Output irreducible representations.

• num_layers (int) – Total number of linear layers in the MLP. Must be >= 1.

• cfg (Config) – Configuration object for activation functions.

• activate_last (bool, optional) – Whether to apply a non-linearity after the final layer. Defaults to False.

• variant (str | None)

• output_scale (float | None)

• weight_init_scale (float | None)

• residual_scale (float | None)

• pre_norm (bool | None)

• film_hidden_dim (int | None)

• post_scale (float)

class net.common.RadialMLP(in_dim, out_dim, *, layers=(128,), act='silu', dtype=torch.float32)

Generic N-layer feed-forward network applied to the radial distance embedding (soft one-hot, Bessel basis, …).

Parameters:

• in_dim (int) – Input size (n_radial from dataset).

• layers (Sequence[int]) – Sequence of hidden layer widths. The final width out_dim is appended automatically if not already present.

• out_dim (int) – Output size produced per edge (number of scalar coefficients).

• act (str) – Name of scalar activation (“silu”, “relu”, …).

• dtype (torch.dtype)

Example

>>> mlp = RadialMLP(64, out_dim=32, layers=(128, 64))

class net.common.SeparateWeightTensorProduct(irreps_in1, irreps_in2, irreps_out, **kwargs)

Tensor product with separate learnable weights for each input.

From DeepH-E3: z_i = W’_{ij} x_j (x) W’’_{ik} y_k

This differs from FullyConnectedTensorProduct by having separate weight matrices for each input irrep, which can be more expressive.

net.activations.make_nonlinearity(irreps, cfg)

Build an equivariant normalisation + activation module.

Parameters:

• irreps (Irreps) – Input/output representation (unchanged by the non-linearity).

• cfg (Config) – Any object that exposes the attributes used below (nonlin_kind, activation_scalar, norm_kind).

Returns:

Callable that maps (…, irreps.dim) -> (…, irreps.dim)

Return type:

torch.nn.Module

class net.activations.GateMagnitudes(irreps, nonlin_scalars, nonlin_gate)

A module that 1. Applies a non-linearity to the scalars in the input 2. Computes the magnitudes of the non-scalars 3. Applies an activation function to the magnitudes 4. Multiplies the non-scalars by the activation output

Parameters:

• irreps (Irreps)

• nonlin_scalars (nn.Module)

• nonlin_gate (nn.Module)

class net.layer_norm.E3LayerNorm(irreps_in, eps=1e-05, affine=True, normalization='component')

DeepH-E3 style equivariant layer normalization.

This is intentionally close to the study implementation so the silicon sweeps can reproduce the same normalization path.

Parameters:

• irreps_in (Irreps)

• eps (float)

• affine (bool)

• normalization (str)

net.e3mlp_variants.make_variant_block(variant, irreps_in, irreps_hidden, cfg)

Parameters:

• variant (str)

• irreps_in (Irreps)

• irreps_hidden (Irreps)

• cfg (VariantConfig)

Return type:

Module

class net.transforms.equivariant.GaussianBasis(start=0.0, stop=5.0, n_gaussians=50, centered=False, trainable=False)

Unnormalized Gaussian features with one-grid-spacing standard deviation.

distances has arbitrary leading shape; the feature axis is appended. Parameters use the current default dtype, as in the reference. There is no cutoff envelope or implicit sqrt(n_gaussians) normalization.

Parameters:

• start (float)

• stop (float)

• n_gaussians (int)

• centered (bool)

• trainable (bool)

class net.transforms.equivariant.GatedRadialTensorProduct(fc_len_in, irreps_in1, irreps_in2, irreps_out, norm='', nonlin=True, act=None, act_gates=None, radial_hidden=(64, 64))

Factorized TP, optional gate, then invariant radial channel scaling.

Unlike a tensor product whose path weights come from the radial MLP, this transformation applies the MLP’s channel scales after its gate. Optional normalization acts on the resulting edge rows per graph.

Parameters:

radial_hidden (Sequence[int])

class net.transforms.equivariant.GraphIrrepNorm(irreps_in, eps=1e-05, affine=True, normalization='component', subtract_mean=True, divide_norm=False)

Normalize each irrep using rows and channels within each graph.

x is [rows, irreps_in.dim] and batch assigns every row to a graph. With the reference defaults, means are subtracted for all irreps, but variance normalization is applied only to l=0. The denominator is sqrt(variance) + eps. Only even scalars receive learned biases. This is graph normalization, not independent normalization of each atom.

Parameters:

• eps (float)

• affine (bool)

• normalization (str)

• subtract_mean (bool)

• divide_norm (bool)

class net.transforms.cluster.NeighborDensityInteraction(node_attrs_irreps, node_feats_irreps, edge_attrs_irreps, edge_feats_irreps, target_irreps, hidden_irreps, avg_num_neighbors, radial_layers=(64, 64, 64), *, emit_expansion=False, normalization='none')

Radially weighted neighbor tensor products, summed before correlation.

skip is a species-conditioned tensor product of the input node features. density is channel-first. An optional independently learned expansion is flat and is saved for a later angular-degree stream.

With normalization='none' both post-sum projections are divided by avg_num_neighbors. With 'graph_component' the summed intermediate is normalized before either projection, and no degree divisor is used. This distinction reproduces the two actual upstream transformations.

Parameters retain upstream names and creation order, including the radial MLP’s e3nn initialization and external tensor-product path weights.

Parameters:

• avg_num_neighbors (float)

• radial_layers (Sequence[int])

• emit_expansion (bool)

• normalization (str)

class net.transforms.cluster.SymmetricProductContraction(irreps_in, irreps_out, correlation, num_elements, *, use_skip=True)

Learned species-conditioned symmetric products of summed densities.

Uses the source real Clebsch–Gordan coupling basis and Horner contraction, including all orders from one through correlation. The source tracing example RNG draws are intentionally retained: removing them would change seeded from-scratch initialization of all later stages.

Parameters:

• correlation (int)

• num_elements (int)

• use_skip (bool)

class net.transforms.cluster.AngularDegreeExpansion(hidden_irreps, hidden_post_irreps, up_irreps, up_post_irreps, irreps_out, num_species, *, use_self_connection=True, nonlinearity=True, normalization='graph_component')

Couple one hidden node stream to a separately carried angular stream.

Order: independent projections -> factorized tensor product -> optional gate -> biased linear -> optional species self-connection -> graph norm -> biased projection of the original hidden input. Its output can become the next layer’s angular stream without changing the main hidden stream.

Parameters:

• num_species (int)

• use_self_connection (bool)

• nonlinearity (bool)

• normalization (str)

class net.transforms.composition.ProjectedAttention(node_irreps, message_irreps, radial_width, *, scalar_dim=64, heads=4, receiver=1, projection=None)

The original projected attention reducer, independent of its message map.

class net.training_policies.SpeciesIrrepAffine(num_species, irreps, *, gradient_mode='physical', eps=1e-12, target='hamiltonian', feature_basis=None, statistics_basis='yzx')

Species-pair affine output transformation with selectable backward rule.

gradient_mode='identity' uses the MACE-H paper’s intentional surrogate gradient. physical differentiates the actual affine function and must be used for forces, stress, and other derivatives of predicted observables.

Operators and data

class data.snapshot.Snapshot(hamiltonian, overlap, density, *, positions=None, forces=None, box=None, stress=None, matrix_path=None, info_path=None, cutoff_radius=None, cfg=None, info=None)

Bundle H, S, D for one configuration and provide physics helpers.

Parameters:

• hamiltonian (BlockMatrix)

• overlap (BlockMatrix)

• density (BlockMatrix)

• positions (torch.Tensor | None)

• forces (torch.Tensor | None)

• box (torch.Tensor | None)

• stress (torch.Tensor | None)

• cutoff_radius (float | None)

• cfg (Config | None)

• info (Any)

get_number_of_electrons()

Return scalar Tr(D·S).

Return type:

Tensor

get_energy()

Return scalar Tr(D·H).

Return type:

Tensor

symmetrize_matrices(*, hamiltonian=True, overlap=True, density=True)

Return a new snapshot with selected matrices symmetrized by 0.5 * (M + M^T).

Parameters:

• hamiltonian (bool)

• overlap (bool)

• density (bool)

Return type:

Snapshot

to_e3nn()

Return a (possibly new) Snapshot in the E3NN convention.

Return type:

Snapshot

to_openmx()

Return a (possibly new) Snapshot in the OpenMX convention.

Return type:

Snapshot

rotate(R)

Return a new rotated snapshot where all three matrices have been rotated by the same 3x3 matrix R.

Parameters:

R (Tensor)

Return type:

Snapshot

filter_by_distance(cutoff, which='density')

Return a new snapshot where edges whose minimal-image distance exceeds cutoff (Å) are removed in **all* three matrices*.

Parameters:

• cutoff (float)

• which (str)

Return type:

Snapshot

reduce_orbitals(selection=None, *, strict=True, keep_unspecified=True)

Return a new Snapshot with reduced orbital basis per element.

Selection examples: - “1s1p”: apply to all elements - {“O”: “2s3p”, “H”: “1s”}: per-element selection

Selection is prefix/in-order per orbital type: if an element has multiple sets of a given l, the first n sets are kept.

Parameters:

• selection (str | Mapping[str, str] | None)

• strict (bool)

• keep_unspecified (bool)

Return type:

Snapshot

static from_deeph_e3(snapshot_dir, *, convention='e3nn', symmetrize_density=True, cutoff_radius=None, dtype=torch.float32, cfg=None, validate_hermitian=True, validate_support=True, canonicalize=True, matrix_files=None, source_basis=None, source_cartesian_axes=None)

Load one snapshot in DeepH-E3’s processed OpenMX HDF5 format.

Parameters:

• snapshot_dir (str | PathLike)

• convention (str)

• symmetrize_density (bool)

• cutoff_radius (float | None)

• dtype (dtype)

• cfg (Config | None)

• validate_hermitian (bool)

• validate_support (bool)

• canonicalize (bool)

• matrix_files (dict | None)

• source_basis (str | None)

• source_cartesian_axes (str | None)

Return type:

Snapshot

dos(sigma=0.005, bin_width=0.001, E_min=-1.0, E_max=1.0)

Compute the density of states (DOS) for this snapshot.

Parameters:

• sigma (float) – The broadening parameter for the DOS.

• bin_width (float) – The width of each energy bin.

• E_min (float) – The minimum energy to consider.

• E_max (float) – The maximum energy to consider.

• keys (Returns a dict with)

• "energies" (torch.Tensor) – The energy bins.

• "dos" (torch.Tensor) – The computed density of states for each energy bin.

export_to_deephe3(path)

Export this snapshot using DeepH-E3’s eV/Angstrom convention.

Parameters:

path (str | PathLike)

class data.block_matrix.BlockMatrix(atoms: 'Tuple[str, ...]', atom_counts: 'Dict[str, int]', pair_blocks: 'Dict[PairKey, torch.Tensor]', pair_edges: 'Dict[PairKey, torch.Tensor]', lookup: 'Dict[Tuple[int, int], Tuple[PairKey, int]]', orbital_cfg: 'OrbitalIrrepConfig', basis: 'str', _coverage_masks: 'dict | None' = None)

Parameters:

• atoms (Tuple[str, ...])

• atom_counts (Dict[str, int])

• pair_blocks (Dict[str, Tensor])

• pair_edges (Dict[str, Tensor])

• lookup (Dict[Tuple[int, int], Tuple[str, int]])

• orbital_cfg (OrbitalIrrepConfig)

• basis (str)

• _coverage_masks (dict | None)

diag()

Return a dict key -> tensor that contains only the blocks whose global source/target atom are identical (self-edges).

Off-diagonal blocks are omitted; keys that would become empty disappear.

Return type:

Dict[str, Tensor]

offdiag()

Same as diag() but returns the off-diagonal blocks ( i ≠ j ).

Return type:

Dict[str, Tensor]

to_vectors(mapper, *, layout=None)

Convert blocks to irreducible representation vectors using a mapper.

Parameters:

mapper (BlockIrrepMapper)

Return type:

IrrepsBlockData

transpose()

Return a new snapshot representing the transposed matrix.

Return type:

BlockMatrix

to_dense()

Assemble a full dense matrix of shape (Σ d_i, Σ d_i) where d_i is orbital dimension of atom i. Summing all periodic images into the central cell.

Return type:

Tensor

classmethod from_dense(matrix, orbital_cfg, atoms, *, diagonal=False, sparsity_threshold=0.0, basis)

Build a block snapshot from a fully dense matrix in global atom order. Primarily for tests / debugging.

Parameters:

• matrix (Tensor)

• atoms (Tuple[str, ...] | list[str])

• diagonal (bool)

• sparsity_threshold (float)

• basis (str)

Return type:

BlockMatrix

sparsify(threshold)

Return a new snapshot in which only blocks whose root-mean-square (RMS) magnitude is strictly greater than threshold are kept.

A block’s RMS is computed as rms = sqrt( mean( block**2 ) ) across its last two dimensions (orbital rows & columns).

Parameters:

threshold (float) – Cut-off applied per block. A Python float is accepted; it is internally cast to the block’s dtype and device.

Return type:

BlockMatrix

Notes

• Pruning is performed per block (first tensor dimension). Keys for which all blocks are removed disappear entirely.

rotate(R)

Return a new BlockMatrix whose orbital reference frame has been rotated by the (3 x 3) matrix R (active rotation).

Notes

• The operation is block-wise:

  M’_(A,B) = U_A · M_(A,B) · U_Bᵀ

  where U_A = irreps_A.D_from_matrix(R) is block-diagonal with one Wigner-D copy per orbital.

Parameters:

R (Tensor)

Return type:

BlockMatrix

class data.block_matrix.IrrepsBlockData(atoms: 'Tuple[str, ...]', atom_counts: 'Dict[str, int]', pair_vectors: 'Dict[PairKey, torch.Tensor]', pair_edges: 'Dict[PairKey, torch.Tensor]', lookup: 'Dict[Tuple[int, int], Tuple[PairKey, int]]', orbital_cfg: 'OrbitalIrrepConfig', basis: 'str' = 'e3nn', layout: 'object | None' = None, packed_values: 'torch.Tensor | None' = None, packed_edges: 'torch.Tensor | None' = None, frames: 'torch.Tensor | None' = None, operator: 'str | None' = None)

Parameters:

• atoms (Tuple[str, ...])

• atom_counts (Dict[str, int])

• pair_vectors (Dict[str, Tensor])

• pair_edges (Dict[str, Tensor])

• lookup (Dict[Tuple[int, int], Tuple[str, int]])

• orbital_cfg (OrbitalIrrepConfig)

• basis (str)

• layout (object | None)

• packed_values (Tensor | None)

• packed_edges (Tensor | None)

• frames (Tensor | None)

• operator (str | None)

static pack_matrix(matrix, edges, layout, *, frames=None, operator=None)

Express physical blocks in a declared supervision frame and unit.

classmethod from_packed(values, x, mapper, layout, *, frame='global', operator=None)

Represent packed orbital channels through MANDALA’s public container.

rotate(R, mapper)

Return a new IrrepsBlockData whose orbital reference frame has been rotated by the (3 x 3) matrix R (active rotation).

Parameters:

• R (Tensor) – The (3 x 3) rotation matrix.

• mapper (BlockIrrepMapper) – The BlockIrrepMapper that provides the pair irreps for computing Wigner-D matrices.

Return type:

IrrepsBlockData

Notes

• For irrep vectors representing the tensor product of orbitals from atoms A and B, we apply the Wigner-D matrix for the output irreps:

  v’_(A,B) = D_out(R) · v_(A,B)

  where D_out = pair_irreps.D_from_matrix(R) acts on the tensor product irreps.

class data.factory.DatasetFactory(cfg, convention=None, preprocessing=None)

Collects snapshot paths -> builds datasets + shared mapper.

Parameters:

• cfg (Config)

• convention (str | None)

add_snapshot(matrix_path, info_path=None, purpose='train')

Register a snapshot directory or an explicit (matrix, info) pair.

Parameters:

• matrix_path (str | PathLike) – Path to the *.scfout file with H/S/D blocks.

• info_path (str | PathLike | None) – Path to the corresponding *.info / *.out file.

• purpose (Literal['train', 'val', 'test']) – One of "train" (default), "val", or "test".

Return type:

None

create(*, include_test=False)

Build datasets and a shared mapper.

Returns:

• train_ds – E3GNNDataset for training (batch size is handled elsewhere).

• val_ds – Validation dataset, or None if no validation snapshots were added.

• mapper – The single BlockIrrepMapper used by both datasets.

Parameters:

include_test (bool)

class data.gnn_dataset.E3GNNDataset(snapshot_paths, mapper, cfg, convention='e3nn', *, preprocessing=None, split='train')

Fully in-memory dataset that yields (x, y) tuples.

x - no self-edges y - targets

Parameters:

• snapshot_paths (Sequence[Tuple[Path, Path]])

• mapper (BlockIrrepMapper)

• cfg (Config)

• convention (str)

data.structure_inference.build_model_input_from_structure(*, atoms, positions, box, cfg, mapper)

Parameters:

• atoms (tuple[str, ...])

• positions (Tensor)

• box (Tensor | None)

• cfg (Config)

• mapper (BlockIrrepMapper)

Return type:

dict[str, Any]

Orbitals and irreps

class core.orbital_irrep_config.OrbitalIrrepConfig(element_to_irreps=<factory>)

Holds per-element orbital irreps and exposes convenience accessors.

Parameters:

element_to_irreps (Dict[str, Irreps])

classmethod from_yaml(yaml_str)

Parse YAML string or bytes into a validated instance.

Parameters:

yaml_str (str | bytes)

Return type:

OrbitalIrrepConfig

to_dict()

Return a plain-python mapping of the original irrep strings so the config can be safely written to YAML / JSON.

Return type:

Dict[str, List[str]]

elements()

Return the list of element symbols in deterministic order.

Return type:

List[str]

block_dims(pair)

Return (d_i, d_j) - the orbital dimensions of the two atoms that form its matrix block.

Parameters:

pair (Tuple[str, str] | str) – Either a hyphen-joined string "Si-H" or a two-tuple ("Si", "H").

Return type:

Tuple[int, int]

class core.block_irrep_mapper.BlockIrrepMapper(orbital_cfg, *, diagonal=False, device='cpu', dtype=torch.float32)

Manages all per-pair _IrrepToMatrix instances.

Parameters:

• orbital_cfg (OrbitalIrrepConfig) – Instance built from user YAML / dict.

• diagonal (bool) – If True and only for same element pairs, use "ij=ji" (symmetric) reduced tensor product.

• device (torch.device | str) – Where the Q matrices live (CPU by default).

• dtype (torch.dtype)

blocks_to_vectors(pair, blocks, *, layout=None)

Convert matrix blocks for a given atom pair to irrep vectors.

pair can be ("Si","H") or “Si-H” (string with hyphen).

Parameters:

• pair (Tuple[str, str] | str)

• blocks (Tensor)

Return type:

Tensor

vectors_to_blocks(pair, vectors, *, layout=None)

Inverse mapping vector -> block.

Parameters:

• pair (Tuple[str, str] | str)

• vectors (Tensor)

Return type:

Tensor

vector_dim(pair)

Number of irrep coefficients for one block of pair.

Parameters:

pair (Tuple[str, str] | str)

Return type:

int

block_dims(pair)

Underlying dimensions (d_i, d_j) of the orbital block.

Parameters:

pair (Tuple[str, str] | str)

Return type:

Tuple[int, int]

get_pair_irreps(pair)

Return the output Irreps for a given pair.

Parameters:

pair (Tuple[str, str] | str)

Return type:

Irreps

class core.operator_layout.ShellTensorCodec(shell_pairs, *, spinful=False, ignore_parity=False, sort_irreps=False, dtype=torch.float32, device='cpu')

Invertible source-compatible shell coupling, optionally shared by species.

shell_pairs fixes the ordered output channels, with one (l1,l2) per channel. from_species_shells additionally provides selection, packing, masks and reconstruction of full native blocks. Scalar shifts, standard deviations and Hermitian projection are separate operations.

decode(channels)

Real coupled coefficients → packed native shell entries.

encode(packed)

Packed native shell entries → real coupled coefficients.

layout_signature()

Serializable identity of native shell slots and coupled conventions.

class core.operator_layout.OrbitalFrameRotation(dtype=torch.float32, device='cpu', spinful=False)

Native OpenMX ↔ cyclic real-harmonic basis and active frame changes.

rotate_blocks follows the reference row-vector rotation convention: its spatial R acts on positions on the right. For column-vector active rotations, pass R.T. SOC full matrices are spin-major (all up orbitals, then all down orbitals). Proper/improper O(3) orbital transformations are supported; spin transforms use the proper part, as in the reference.

Sparse algebra and reciprocal space

core.sparse_math.trace_matmul_sparse_block_matrix(A, B)

Scalar trace via per-edge loop using lookup table. Works for any BlockMatrix / IrrepsBlockData combination.

Parameters:

• A (BlockMatrix)

• B (BlockMatrix)

Return type:

Tensor

core.sparse_math.trace_matmul_sparse_block_matrix_aligned(A, B, alignment)

Vectorized sparse trace using a precomputed reverse-edge alignment.

Parameters:

• A (BlockMatrix)

• B (BlockMatrix)

• alignment (Dict[str, Tuple[str, Tensor]])

Return type:

Tensor

core.sparse_math.build_trace_alignment(A, B)

Convenience wrapper building alignment from BlockMatrix edge dictionaries.

Parameters:

• A (BlockMatrix)

• B (BlockMatrix)

Return type:

Dict[str, Tuple[str, Tensor]]

core.periodic_fourier.shiftspace_to_kspace_dense(matrices_shift, *, kpoints_abs, shifts, box, phase=None)

Parameters:

• matrices_shift (Tensor)

• kpoints_abs (Tensor)

• shifts (Tensor)

• box (Tensor)

• phase (Tensor | None)

Return type:

Tensor

core.periodic_fourier.kspace_to_shiftspace_dense(matrices_k, *, kpoints_abs, shifts, box)

Parameters:

• matrices_k (Tensor)

• kpoints_abs (Tensor)

• shifts (Tensor)

• box (Tensor)

Return type:

Tensor

Objectives

net.objectives.register_objective_term(name)

Register a function of keyword arguments predictions, targets, inputs.

Predictions and targets map operator names to physical BlockMatrix objects in the same orbital basis (Hamiltonians in Hartree). inputs contains one structure’s graph metadata. Additional keyword arguments come from the term’s YAML options. Return one real scalar tensor, retaining the prediction’s autograd graph. Terms have no learnable parameters.

net.spectral_loss.compute_spectral_eigenvalue_loss(*, pred_hamiltonian, spectral_payload, box, loss_kind, huber_delta_ev, overlap_psd_cleanup, overlap_jitter)

Parameters:

• pred_hamiltonian (BlockMatrix)

• spectral_payload (dict[str, Tensor])

• box (Tensor)

• loss_kind (str)

• huber_delta_ev (float)

• overlap_psd_cleanup (bool)

• overlap_jitter (bool)

Return type:

tuple[Tensor, dict[str, Tensor]]

net.observable_metrics.observable_loss(*, cfg, observable_values, energy_target, num_electrons_target, mse_fn, device)

Parameters:

• observable_values (dict[str, Tensor])

• energy_target (Tensor | None)

• num_electrons_target (Tensor | None)

• device (device)

Return type:

tuple[Tensor, Tensor]

data.envelope.load_pair_envelope_table(path, *, pair_order=None, dtype=torch.float32, device='cpu')

Parameters:

• path (str | Path)

• pair_order (tuple[str, ...] | list[str] | None)

• dtype (dtype)

• device (device | str)

Return type:

PairEnvelopeTable

data.envelope.evaluate_slater_soft_cutoff(edge_lengths, pair_params)

Parameters:

• edge_lengths (Tensor)

• pair_params (Tensor)

Return type:

Tensor

data.envelope.evaluate_slater_exp_quad_soft_wall(edge_lengths, pair_params, reference_x_max)

Parameters:

• edge_lengths (Tensor)

• pair_params (Tensor)

• reference_x_max (Tensor)

Return type:

Tensor

Evaluation and export

Small, stateless helpers for checkpoint and snapshot evaluation.

These functions deliberately compose the existing Config, Snapshot, BlockIrrepMapper, and E3GNN classes. They do not introduce a new model wrapper or change the Lightning checkpoint format.

net.evaluation.evaluate_sample(model, x, y)

Evaluate one already-aligned dataset sample and its observables.

Parameters:

• model (E3GNN)

• x (dict[str, Any])

• y (dict[str, Any])

Return type:

dict[str, Any]

net.evaluation.load_checkpoint(path)

Load a Lightning checkpoint on CPU without changing its payload.

Parameters:

path (str | Path)

Return type:

dict[str, Any]

net.evaluation.predict_snapshot(model, snapshot, *, physical=True)

Run structure-only inference and return predicted block matrices.

Parameters:

• model (E3GNN)

• snapshot (Snapshot)

• physical (bool)

Return type:

dict[str, Any]

net.evaluation.predictions_to_snapshot(predictions, reference, *, reference_matrices=('overlap', 'density'))

Combine predicted matrices with explicitly selected reference matrices.

Parameters:

• predictions (dict[str, Any])

• reference (Snapshot)

• reference_matrices (tuple[str, ...])

Return type:

Snapshot

net.evaluation.restore_config(checkpoint)

Restore the saved Config object from a Mandala checkpoint.

Parameters:

checkpoint (dict[str, Any])

Return type:

Config

net.evaluation.restore_model_for_snapshot(checkpoint_path, snapshot, *, device='cpu')

Restore a model using the orbital mapper inferred from snapshot.

Parameters:

• checkpoint_path (str | Path)

• snapshot (Snapshot)

• device (str | device)

Return type:

E3GNN

data.openmx_writer.export_model_openmx(model, *, atoms, positions, lattice, output_path)

Infer/export using ONLY a loaded model and native Cartesian geometry.