.gitignore
.pre-commit-config.yaml
LICENSE
README.md
environment.yml
mkdocs.yml
pyproject.toml
.github/workflows/actions.yml
.github/workflows/docs.yml
.github/workflows/publish.yml
.vscode/settings.json
docs/api-reference.md
docs/calculators.md
docs/cli.md
docs/data.md
docs/fine-tuning.md
docs/index.md
docs/installation.md
docs/jax-bundles.md
docs/model-wrappers.md
docs/phonon-finetuning-paper.md
docs/python-api.md
docs/quickstart.md
docs/reaction-relative-losses.md
docs/resources.md
docs/training-options.md
docs/stylesheets/extra.css
equitrain/__init__.py
equitrain/__main__.py
equitrain/_version.py
equitrain/argparser.py
equitrain/checkpoint.py
equitrain/evaluate.py
equitrain/logger.py
equitrain/loss_metrics.py
equitrain/model_wrappers.py
equitrain/predict.py
equitrain/preprocess.py
equitrain/train.py
equitrain.egg-info/PKG-INFO
equitrain.egg-info/SOURCES.txt
equitrain.egg-info/dependency_links.txt
equitrain.egg-info/entry_points.txt
equitrain.egg-info/requires.txt
equitrain.egg-info/scm_file_list.json
equitrain.egg-info/scm_version.json
equitrain.egg-info/top_level.txt
equitrain/backends/__init__.py
equitrain/backends/jax_backend.py
equitrain/backends/jax_checkpoint.py
equitrain/backends/jax_evaluate.py
equitrain/backends/jax_freeze.py
equitrain/backends/jax_loss.py
equitrain/backends/jax_loss_fn.py
equitrain/backends/jax_loss_metrics.py
equitrain/backends/jax_nnx_compat.py
equitrain/backends/jax_optimizer.py
equitrain/backends/jax_predict.py
equitrain/backends/jax_runtime.py
equitrain/backends/jax_scheduler.py
equitrain/backends/jax_utils.py
equitrain/backends/scheduler_common.py
equitrain/backends/torch_backend.py
equitrain/backends/torch_checkpoint.py
equitrain/backends/torch_evaluate.py
equitrain/backends/torch_freeze.py
equitrain/backends/torch_loss.py
equitrain/backends/torch_loss_fn.py
equitrain/backends/torch_loss_metrics.py
equitrain/backends/torch_model.py
equitrain/backends/torch_optimizer.py
equitrain/backends/torch_predict.py
equitrain/backends/torch_scheduler.py
equitrain/backends/torch_utils.py
equitrain/backends/jax_wrappers/__init__.py
equitrain/backends/jax_wrappers/ani.py
equitrain/backends/jax_wrappers/m3gnet.py
equitrain/backends/jax_wrappers/mace.py
equitrain/backends/torch_derivatives/force.py
equitrain/backends/torch_derivatives/stress.py
equitrain/backends/torch_wrappers/__init__.py
equitrain/backends/torch_wrappers/ani.py
equitrain/backends/torch_wrappers/base.py
equitrain/backends/torch_wrappers/m3gnet.py
equitrain/backends/torch_wrappers/mace.py
equitrain/backends/torch_wrappers/orb.py
equitrain/backends/torch_wrappers/sevennet.py
equitrain/calculators/__init__.py
equitrain/calculators/jax_wrapper.py
equitrain/calculators/torch_wrapper.py
equitrain/data/__init__.py
equitrain/data/atomic.py
equitrain/data/configuration.py
equitrain/data/neighborhood.py
equitrain/data/statistics_data.py
equitrain/data/backend_jax/__init__.py
equitrain/data/backend_jax/atoms_to_graphs.py
equitrain/data/backend_jax/atoms_to_graphs_impl.py
equitrain/data/backend_jax/loaders.py
equitrain/data/backend_jax/loaders_impl.py
equitrain/data/backend_jax/loaders_impl_cache.py
equitrain/data/backend_jax/statistics.py
equitrain/data/backend_torch/__init__.py
equitrain/data/backend_torch/atoms_to_graphs.py
equitrain/data/backend_torch/loaders.py
equitrain/data/backend_torch/loaders_impl.py
equitrain/data/backend_torch/loaders_reaction.py
equitrain/data/backend_torch/scatter.py
equitrain/data/backend_torch/statistics.py
equitrain/data/backend_torch/utility.py
equitrain/data/format_hdf5/__init__.py
equitrain/data/format_hdf5/dataset.py
equitrain/data/format_lmdb/__init__.py
equitrain/data/format_lmdb/lmdb.py
equitrain/data/format_xyz/__init__.py
equitrain/data/format_xyz/reader.py
equitrain/finetune/__init__.py
equitrain/finetune/_layer_selection.py
equitrain/finetune/_lora_common.py
equitrain/finetune/_torch_common.py
equitrain/finetune/delta_jax.py
equitrain/finetune/delta_torch.py
equitrain/finetune/freeze_torch.py
equitrain/finetune/lora_jax.py
equitrain/finetune/lora_torch.py
equitrain/scripts/__init__.py
equitrain/scripts/equitrain.py
equitrain/scripts/equitrain_evaluate.py
equitrain/scripts/equitrain_export.py
equitrain/scripts/equitrain_hdf5_benchmark.py
equitrain/scripts/equitrain_hdf5_info.py
equitrain/scripts/equitrain_inspect.py
equitrain/scripts/equitrain_predict.py
equitrain/scripts/equitrain_preprocess.py
equitrain/utility_test/__init__.py
equitrain/utility_test/mace_support.py
equitrain/utility_test/wrapper_ani.py
equitrain/utility_test/wrapper_m3gnet.py
equitrain/utility_test/wrapper_mace.py
equitrain/utility_test/wrapper_orb.py
equitrain/utility_test/wrapper_sevennet.py
resources/data/alexandria+mptraj-statistics.json
resources/data/alexandria/01_download.py
resources/data/alexandria/02_convert_to_xyz.py
resources/data/alexandria/03_join.py
resources/data/alexandria/04_xyz_to_hdf5.py
resources/data/alexandria/makefile
resources/data/alexandria+mptraj/01_alexandria_download.py
resources/data/alexandria+mptraj/02_alexandria_to_xyz.py
resources/data/alexandria+mptraj/02_mptraj_to_xyz.py
resources/data/alexandria+mptraj/03_join.py
resources/data/alexandria+mptraj/04_xyz_to_hdf5.py
resources/data/alexandria+mptraj/makefile
resources/data/mptraj/01_json_to_xyz.py
resources/data/mptraj/02_xyz_to_hdf5.py
resources/data/mptraj/makefile
resources/data/omat24/00_download_omat24.sh
resources/data/omat24/01_aselmdb_to_hdf5.py
resources/data/omat24/02_compute_statistics.py
resources/data/omat24/makefile
resources/models/ani/README.md
resources/models/ani/ani-initial-model.py
resources/models/m3gnet/README.md
resources/models/m3gnet/m3gnet-config.yaml
resources/models/m3gnet/m3gnet-initial-model.py
resources/models/mace/mace-initial-model.py
resources/models/mace-jax/README.md
resources/models/mace-jax/convert_foundation_to_jax.py
resources/models/orb/README.md
resources/models/orb/orb_config.yaml
resources/models/sevennet/sevennet-initial-model.py
resources/models/sevennet/sevennet-initial-model.yaml
resources/training/README.md
resources/training/mace-alex-mptraj-multigpu.sh
resources/training/mace-alex-mptraj-multigpu.yaml
resources/training/mace-alex-mptraj-multinode-slurm.sh
resources/training/mace-alex-mptraj-multinode.sh
resources/training/mace-alex-mptraj-multinode.yaml
resources/training/mace-alex-mptraj.py
tests/__init__.py
tests/conftest.py
tests/data.xyz
tests/test_backend_defaults.py
tests/test_backend_jax_statistics.py
tests/test_calculators_jax.py
tests/test_calculators_torch.py
tests/test_checkpoint_retention.py
tests/test_evaluate_mace.py
tests/test_export_torch_checkpoint.py
tests/test_finetune_delta_equitrain_0_1_0.py
tests/test_finetune_delta_torch.py
tests/test_finetune_freeze_torch.py
tests/test_finetune_lora_jax.py
tests/test_finetune_lora_torch.py
tests/test_finetune_mace.py
tests/test_finetune_mace_jax.py
tests/test_finetune_mace_readout.py
tests/test_jax_ani_wrapper.py
tests/test_jax_evaluate.py
tests/test_jax_freeze_mask.py
tests/test_jax_loader_deterministic.py
tests/test_jax_loss.py
tests/test_jax_m3gnet_wrapper.py
tests/test_jax_model_equivalence.py
tests/test_jax_nnx_compat.py
tests/test_jax_predict.py
tests/test_jax_shard_map.py
tests/test_jax_wrappers_registry.py
tests/test_lmdb_conversion.py
tests/test_mace_rmax_override.py
tests/test_mace_wrapper_compat.py
tests/test_niggli_predict.py
tests/test_polar_mace_data.py
tests/test_predict_ani_atoms.py
tests/test_predict_m3gnet.py
tests/test_predict_m3gnet_structures.py
tests/test_predict_mace.py
tests/test_predict_mace_atoms.py
tests/test_predict_mace_structures.py
tests/test_preprocess.py
tests/test_reaction_relative_loss.py
tests/test_scheduler_args.py
tests/test_train_ani_energy_only.py
tests/test_train_m3gnet.py
tests/test_train_mace.py
tests/test_train_mace_jax.py
tests/test_train_orb.py
tests/test_train_sevennet.py
tests/test_train_sevennet.yaml
tests/data/statistics.json
tests/data/train.h5
tests/data/valid.h5
tests/data/equitrain_0_1_0_delta/train.h5
tests/data/equitrain_0_1_0_delta/valid.h5
tests/data/lmdb/val_metadata.npz