Index

_

__call__() (lagrangebench.evaluate.metrics.MetricsComputer method) (lagrangebench.models.base.BaseModel method) (lagrangebench.models.egnn.EGNN method) (lagrangebench.models.egnn.EGNNLayer method) (lagrangebench.models.gns.GNS method) (lagrangebench.models.linear.Linear method) (lagrangebench.models.painn.GatedEquivariantBlock method) (lagrangebench.models.painn.PaiNN method) (lagrangebench.models.painn.PaiNNLayer method) (lagrangebench.models.segnn.O3TensorProduct method) (lagrangebench.models.segnn.SEGNN method) (lagrangebench.models.segnn.SEGNNLayer method) __delattr__() (lagrangebench.case_setup.case.CaseSetupFn method) __eq__() (lagrangebench.case_setup.case.CaseSetupFn method) __getitem__() (lagrangebench.data.data.H5Dataset method) __getnewargs__() (lagrangebench.models.painn.NodeFeatures method) __hash__() (lagrangebench.case_setup.case.CaseSetupFn method) __init__() (lagrangebench.case_setup.case.CaseSetupFn method) (lagrangebench.data.data.DAM2D method) (lagrangebench.data.data.H5Dataset method) (lagrangebench.data.data.LDC2D method) (lagrangebench.data.data.LDC3D method) (lagrangebench.data.data.RPF2D method) (lagrangebench.data.data.RPF3D method) (lagrangebench.data.data.TGV2D method) (lagrangebench.data.data.TGV3D method) (lagrangebench.evaluate.metrics.MetricsComputer method) (lagrangebench.models.egnn.EGNN method) (lagrangebench.models.egnn.EGNNLayer method) (lagrangebench.models.gns.GNS method) (lagrangebench.models.linear.Linear method) (lagrangebench.models.painn.GatedEquivariantBlock method) (lagrangebench.models.painn.PaiNN method) (lagrangebench.models.painn.PaiNNLayer method) (lagrangebench.models.segnn.O3TensorProduct method) (lagrangebench.models.segnn.SEGNN method) (lagrangebench.models.segnn.SEGNNLayer method) (lagrangebench.models.utils.LinearXav method) (lagrangebench.models.utils.MLPXav method)

__new__() (lagrangebench.models.painn.NodeFeatures static method) __setattr__() (lagrangebench.case_setup.case.CaseSetupFn method)

A

add_gns_noise() (in module lagrangebench.train.strats) additional_message_features (lagrangebench.models.utils.SteerableGraphsTuple attribute)

allocate (lagrangebench.case_setup.case.CaseSetupFn attribute) allocate_eval (lagrangebench.case_setup.case.CaseSetupFn attribute) averaged_metrics() (in module lagrangebench.evaluate.metrics)

B

BaseModel (class in lagrangebench.models.base) broadcast_from_batch() (in module lagrangebench.utils)

broadcast_to_batch() (in module lagrangebench.utils) build_mlp() (in module lagrangebench.models.utils)

C

case_builder() (in module lagrangebench.case_setup.case) CaseSetupFn (class in lagrangebench.case_setup.case)

check_cfg() (in module lagrangebench.defaults) cosine_cutoff() (in module lagrangebench.models.painn)

D

DAM2D (class in lagrangebench.data.data)

displacement (lagrangebench.case_setup.case.CaseSetupFn attribute) download() (lagrangebench.data.data.H5Dataset method)

E

e_kin() (lagrangebench.evaluate.metrics.MetricsComputer method) edge_attributes (lagrangebench.models.utils.SteerableGraphsTuple attribute)

EGNN (class in lagrangebench.models.egnn) EGNNLayer (class in lagrangebench.models.egnn) eval_rollout() (in module lagrangebench.evaluate.rollout)

F

features_2d_to_3d() (in module lagrangebench.models.utils)

G

GatedEquivariantBlock (class in lagrangebench.models.painn) gaussian_rbf() (in module lagrangebench.models.painn) get_dataset_name_from_path() (in module lagrangebench.data.data) get_dataset_stats() (in module lagrangebench.data.utils) get_kinematic_mask() (in module lagrangebench.utils)

get_num_params() (in module lagrangebench.utils) get_trajectory() (lagrangebench.data.data.H5Dataset method) get_window() (lagrangebench.data.data.H5Dataset method) GNS (class in lagrangebench.models.gns) graph (lagrangebench.models.utils.SteerableGraphsTuple attribute)

H

H5Dataset (class in lagrangebench.data.data)

I

infer() (in module lagrangebench.evaluate.rollout)

integrate (lagrangebench.case_setup.case.CaseSetupFn attribute)

L

lagrangebench.case_setup.case module lagrangebench.case_setup.features module lagrangebench.data.data module lagrangebench.data.utils module lagrangebench.defaults module lagrangebench.evaluate.metrics module lagrangebench.evaluate.rollout module lagrangebench.evaluate.utils module lagrangebench.models.base module lagrangebench.models.egnn module lagrangebench.models.gns module

lagrangebench.models.linear module lagrangebench.models.painn module lagrangebench.models.segnn module lagrangebench.models.utils module lagrangebench.train.strats module lagrangebench.train.trainer module lagrangebench.utils module LDC2D (class in lagrangebench.data.data) LDC3D (class in lagrangebench.data.data) Linear (class in lagrangebench.models.linear) LinearXav (class in lagrangebench.models.utils) load_haiku() (in module lagrangebench.utils) load_pytree() (in module lagrangebench.utils)

M

mae() (lagrangebench.evaluate.metrics.MetricsComputer method) MetricsComputer (class in lagrangebench.evaluate.metrics) MLPXav (class in lagrangebench.models.utils) module lagrangebench.case_setup.case lagrangebench.case_setup.features lagrangebench.data.data lagrangebench.data.utils lagrangebench.defaults lagrangebench.evaluate.metrics lagrangebench.evaluate.rollout lagrangebench.evaluate.utils lagrangebench.models.base lagrangebench.models.egnn lagrangebench.models.gns lagrangebench.models.linear lagrangebench.models.painn lagrangebench.models.segnn lagrangebench.models.utils lagrangebench.train.strats lagrangebench.train.trainer lagrangebench.utils

mse() (lagrangebench.evaluate.metrics.MetricsComputer method)

N

node_attributes (lagrangebench.models.utils.SteerableGraphsTuple attribute) node_irreps() (in module lagrangebench.models.utils) NodeFeatures (class in lagrangebench.models.painn)

NodeType (class in lagrangebench.utils) normalization_stats (lagrangebench.case_setup.case.CaseSetupFn attribute) numpy_collate() (in module lagrangebench.data.utils)

O

O3Decoder() (in module lagrangebench.models.segnn) O3Embedding() (in module lagrangebench.models.segnn)

O3TensorProduct (class in lagrangebench.models.segnn) O3TensorProductGate() (in module lagrangebench.models.segnn)

P

PaiNN (class in lagrangebench.models.painn) PaiNNLayer (class in lagrangebench.models.painn) PaiNNReadout() (in module lagrangebench.models.painn) physical_feature_builder() (in module lagrangebench.case_setup.features)

pkl2vtk() (in module lagrangebench.evaluate.utils) preprocess (lagrangebench.case_setup.case.CaseSetupFn attribute) preprocess_eval (lagrangebench.case_setup.case.CaseSetupFn attribute) push_forward_build() (in module lagrangebench.train.strats) push_forward_sample_steps() (in module lagrangebench.train.strats)

R

RPF2D (class in lagrangebench.data.data)

RPF3D (class in lagrangebench.data.data)

S

save_haiku() (in module lagrangebench.utils) save_pytree() (in module lagrangebench.utils) SEGNN (class in lagrangebench.models.segnn) SEGNNLayer (class in lagrangebench.models.segnn)

set_defaults() (in module lagrangebench.defaults) set_seed() (in module lagrangebench.utils) sinkhorn() (lagrangebench.evaluate.metrics.MetricsComputer method) SteerableGraphsTuple (class in lagrangebench.models.utils)

T

TGV2D (class in lagrangebench.data.data) TGV3D (class in lagrangebench.data.data)

train() (lagrangebench.train.trainer.Trainer method) Trainer (class in lagrangebench.train.trainer)

W

weight_balanced_irreps() (in module lagrangebench.models.segnn)

write_vtk() (in module lagrangebench.evaluate.utils)