From 12b657fb078a2b779fefd9747361b17eb5c2561b Mon Sep 17 00:00:00 2001 From: orionarcher Date: Fri, 12 Dec 2025 11:21:13 -0500 Subject: [PATCH 1/3] significantly consolidate scripts to speed up testing. --- .github/workflows/test.yml | 1 - .../1_Introduction/1.1_Lennard_Jones.py | 0 .../1_Introduction/1.2_MACE.py | 0 .../2.1_Lennard_Jones_FIRE.py | 0 .../2.2_Soft_Sphere_FIRE.py | 0 .../2.3_MACE_Gradient_Descent.py | 0 .../2.4_MACE_FIRE.py | 0 ....5_MACE_UnitCellFilter_Gradient_Descent.py | 0 .../2.6_MACE_UnitCellFilter_FIRE.py | 0 .../2.7_MACE_FrechetCellFilter_FIRE.py | 0 .../3_Dynamics/3.10_Hybrid_swap_mc.py | 0 .../3.11_Lennard_Jones_NPT_Langevin.py | 0 .../3_Dynamics/3.12_MACE_NPT_Langevin.py | 0 .../3_Dynamics/3.13_MACE_NVE_non_pbc.py | 0 .../3_Dynamics/3.1_Lennard_Jones_NVE.py | 0 .../3_Dynamics/3.2_MACE_NVE.py | 0 .../3_Dynamics/3.3_MACE_NVE_cueq.py | 0 .../3_Dynamics/3.4_MACE_NVT_Langevin.py | 0 .../3_Dynamics/3.5_MACE_NVT_Nose_Hoover.py | 0 .../3.6_MACE_NVT_Nose_Hoover_temp_profile.py | 0 .../3.7_Lennard_Jones_NPT_Nose_Hoover.py | 0 .../3_Dynamics/3.8_MACE_NPT_Nose_Hoover.py | 0 .../3.9_MACE_NVT_staggered_stress.py | 0 .../4_High_level_api/4.1_high_level_api.py | 0 .../4_High_level_api/4.2_auto_batching_api.py | 0 .../5_Workflow/5.1_a2c_silicon_batched.py | 0 .../5_Workflow/5.2_In_Flight_WBM.py | 0 .../5_Workflow/5.3_Elastic.py | 0 .../6_Phonons/6.1_Phonons_MACE.py | 0 .../6_Phonons/6.2_QuasiHarmonic_MACE.py | 0 .../6_Phonons/6.3_Conductivity_MACE.py | 0 .../7_Others/7.1_Soft_sphere_autograd.py | 0 .../7_Others/7.2_Stress_autograd.py | 0 .../7_Others/7.3_Batched_neighbor_list.py | 0 .../7_Others/7.4_Velocity_AutoCorrelation.py | 0 .../7_Others/7.6_Compare_ASE_to_VV_FIRE.py | 0 .../7_Others/7.7_Heat_flux_and_kappa.py | 0 examples/old_scripts/readme.md | 75 ++++ .../scripts/1_Introduction/1.3_fairchem.py | 62 --- examples/scripts/1_introduction.py | 199 +++++++++ examples/scripts/2_structural_optimization.py | 391 +++++++++++++++++ examples/scripts/3_dynamics.py | 402 ++++++++++++++++++ examples/scripts/4_high_level_api.py | 287 +++++++++++++ examples/scripts/5_workflow.py | 233 ++++++++++ examples/scripts/6_phonons.py | 255 +++++++++++ examples/scripts/7_others.py | 239 +++++++++++ examples/scripts/readme.md | 249 +++++++++-- 47 files changed, 2292 insertions(+), 101 deletions(-) rename examples/{scripts => old_scripts}/1_Introduction/1.1_Lennard_Jones.py (100%) rename examples/{scripts => old_scripts}/1_Introduction/1.2_MACE.py (100%) rename examples/{scripts => old_scripts}/2_Structural_optimization/2.1_Lennard_Jones_FIRE.py (100%) rename examples/{scripts => old_scripts}/2_Structural_optimization/2.2_Soft_Sphere_FIRE.py (100%) rename examples/{scripts => old_scripts}/2_Structural_optimization/2.3_MACE_Gradient_Descent.py (100%) rename examples/{scripts => old_scripts}/2_Structural_optimization/2.4_MACE_FIRE.py (100%) rename examples/{scripts => old_scripts}/2_Structural_optimization/2.5_MACE_UnitCellFilter_Gradient_Descent.py (100%) rename examples/{scripts => old_scripts}/2_Structural_optimization/2.6_MACE_UnitCellFilter_FIRE.py (100%) rename examples/{scripts => old_scripts}/2_Structural_optimization/2.7_MACE_FrechetCellFilter_FIRE.py (100%) rename examples/{scripts => old_scripts}/3_Dynamics/3.10_Hybrid_swap_mc.py (100%) rename examples/{scripts => old_scripts}/3_Dynamics/3.11_Lennard_Jones_NPT_Langevin.py (100%) rename examples/{scripts => old_scripts}/3_Dynamics/3.12_MACE_NPT_Langevin.py (100%) rename examples/{scripts => old_scripts}/3_Dynamics/3.13_MACE_NVE_non_pbc.py (100%) rename examples/{scripts => old_scripts}/3_Dynamics/3.1_Lennard_Jones_NVE.py (100%) rename examples/{scripts => old_scripts}/3_Dynamics/3.2_MACE_NVE.py (100%) rename examples/{scripts => old_scripts}/3_Dynamics/3.3_MACE_NVE_cueq.py (100%) rename examples/{scripts => old_scripts}/3_Dynamics/3.4_MACE_NVT_Langevin.py (100%) rename examples/{scripts => old_scripts}/3_Dynamics/3.5_MACE_NVT_Nose_Hoover.py (100%) rename examples/{scripts => old_scripts}/3_Dynamics/3.6_MACE_NVT_Nose_Hoover_temp_profile.py (100%) rename examples/{scripts => old_scripts}/3_Dynamics/3.7_Lennard_Jones_NPT_Nose_Hoover.py (100%) rename examples/{scripts => old_scripts}/3_Dynamics/3.8_MACE_NPT_Nose_Hoover.py (100%) rename examples/{scripts => old_scripts}/3_Dynamics/3.9_MACE_NVT_staggered_stress.py (100%) rename examples/{scripts => old_scripts}/4_High_level_api/4.1_high_level_api.py (100%) rename examples/{scripts => old_scripts}/4_High_level_api/4.2_auto_batching_api.py (100%) rename examples/{scripts => old_scripts}/5_Workflow/5.1_a2c_silicon_batched.py (100%) rename examples/{scripts => old_scripts}/5_Workflow/5.2_In_Flight_WBM.py (100%) rename examples/{scripts => old_scripts}/5_Workflow/5.3_Elastic.py (100%) rename examples/{scripts => old_scripts}/6_Phonons/6.1_Phonons_MACE.py (100%) rename examples/{scripts => old_scripts}/6_Phonons/6.2_QuasiHarmonic_MACE.py (100%) rename examples/{scripts => old_scripts}/6_Phonons/6.3_Conductivity_MACE.py (100%) rename examples/{scripts => old_scripts}/7_Others/7.1_Soft_sphere_autograd.py (100%) rename examples/{scripts => old_scripts}/7_Others/7.2_Stress_autograd.py (100%) rename examples/{scripts => old_scripts}/7_Others/7.3_Batched_neighbor_list.py (100%) rename examples/{scripts => old_scripts}/7_Others/7.4_Velocity_AutoCorrelation.py (100%) rename examples/{scripts => old_scripts}/7_Others/7.6_Compare_ASE_to_VV_FIRE.py (100%) rename examples/{scripts => old_scripts}/7_Others/7.7_Heat_flux_and_kappa.py (100%) create mode 100644 examples/old_scripts/readme.md delete mode 100644 examples/scripts/1_Introduction/1.3_fairchem.py create mode 100644 examples/scripts/1_introduction.py create mode 100644 examples/scripts/2_structural_optimization.py create mode 100644 examples/scripts/3_dynamics.py create mode 100644 examples/scripts/4_high_level_api.py create mode 100644 examples/scripts/5_workflow.py create mode 100644 examples/scripts/6_phonons.py create mode 100644 examples/scripts/7_others.py diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml index 92b816af..12e18c00 100644 --- a/.github/workflows/test.yml +++ b/.github/workflows/test.yml @@ -163,7 +163,6 @@ jobs: - name: Find example scripts id: set-matrix run: | - # Find all example scripts but exclude known failing ones EXAMPLES=$(find examples -name "*.py" | jq -R -s -c 'split("\n")[:-1]') echo "examples=$EXAMPLES" >> $GITHUB_OUTPUT diff --git a/examples/scripts/1_Introduction/1.1_Lennard_Jones.py b/examples/old_scripts/1_Introduction/1.1_Lennard_Jones.py similarity index 100% rename from examples/scripts/1_Introduction/1.1_Lennard_Jones.py rename to examples/old_scripts/1_Introduction/1.1_Lennard_Jones.py diff --git a/examples/scripts/1_Introduction/1.2_MACE.py b/examples/old_scripts/1_Introduction/1.2_MACE.py similarity index 100% rename from examples/scripts/1_Introduction/1.2_MACE.py rename to examples/old_scripts/1_Introduction/1.2_MACE.py diff --git a/examples/scripts/2_Structural_optimization/2.1_Lennard_Jones_FIRE.py b/examples/old_scripts/2_Structural_optimization/2.1_Lennard_Jones_FIRE.py similarity index 100% rename from examples/scripts/2_Structural_optimization/2.1_Lennard_Jones_FIRE.py rename to examples/old_scripts/2_Structural_optimization/2.1_Lennard_Jones_FIRE.py diff --git a/examples/scripts/2_Structural_optimization/2.2_Soft_Sphere_FIRE.py b/examples/old_scripts/2_Structural_optimization/2.2_Soft_Sphere_FIRE.py similarity index 100% rename from examples/scripts/2_Structural_optimization/2.2_Soft_Sphere_FIRE.py rename to examples/old_scripts/2_Structural_optimization/2.2_Soft_Sphere_FIRE.py diff --git a/examples/scripts/2_Structural_optimization/2.3_MACE_Gradient_Descent.py b/examples/old_scripts/2_Structural_optimization/2.3_MACE_Gradient_Descent.py similarity index 100% rename from examples/scripts/2_Structural_optimization/2.3_MACE_Gradient_Descent.py rename to examples/old_scripts/2_Structural_optimization/2.3_MACE_Gradient_Descent.py diff --git a/examples/scripts/2_Structural_optimization/2.4_MACE_FIRE.py b/examples/old_scripts/2_Structural_optimization/2.4_MACE_FIRE.py similarity index 100% rename from examples/scripts/2_Structural_optimization/2.4_MACE_FIRE.py rename to examples/old_scripts/2_Structural_optimization/2.4_MACE_FIRE.py diff --git a/examples/scripts/2_Structural_optimization/2.5_MACE_UnitCellFilter_Gradient_Descent.py b/examples/old_scripts/2_Structural_optimization/2.5_MACE_UnitCellFilter_Gradient_Descent.py similarity index 100% rename from examples/scripts/2_Structural_optimization/2.5_MACE_UnitCellFilter_Gradient_Descent.py rename to examples/old_scripts/2_Structural_optimization/2.5_MACE_UnitCellFilter_Gradient_Descent.py diff --git a/examples/scripts/2_Structural_optimization/2.6_MACE_UnitCellFilter_FIRE.py b/examples/old_scripts/2_Structural_optimization/2.6_MACE_UnitCellFilter_FIRE.py similarity index 100% rename from examples/scripts/2_Structural_optimization/2.6_MACE_UnitCellFilter_FIRE.py rename to examples/old_scripts/2_Structural_optimization/2.6_MACE_UnitCellFilter_FIRE.py diff --git a/examples/scripts/2_Structural_optimization/2.7_MACE_FrechetCellFilter_FIRE.py b/examples/old_scripts/2_Structural_optimization/2.7_MACE_FrechetCellFilter_FIRE.py similarity index 100% rename from examples/scripts/2_Structural_optimization/2.7_MACE_FrechetCellFilter_FIRE.py rename to examples/old_scripts/2_Structural_optimization/2.7_MACE_FrechetCellFilter_FIRE.py diff --git a/examples/scripts/3_Dynamics/3.10_Hybrid_swap_mc.py b/examples/old_scripts/3_Dynamics/3.10_Hybrid_swap_mc.py similarity index 100% rename from examples/scripts/3_Dynamics/3.10_Hybrid_swap_mc.py rename to examples/old_scripts/3_Dynamics/3.10_Hybrid_swap_mc.py diff --git a/examples/scripts/3_Dynamics/3.11_Lennard_Jones_NPT_Langevin.py b/examples/old_scripts/3_Dynamics/3.11_Lennard_Jones_NPT_Langevin.py similarity index 100% rename from examples/scripts/3_Dynamics/3.11_Lennard_Jones_NPT_Langevin.py rename to examples/old_scripts/3_Dynamics/3.11_Lennard_Jones_NPT_Langevin.py diff --git a/examples/scripts/3_Dynamics/3.12_MACE_NPT_Langevin.py b/examples/old_scripts/3_Dynamics/3.12_MACE_NPT_Langevin.py similarity index 100% rename from examples/scripts/3_Dynamics/3.12_MACE_NPT_Langevin.py rename to examples/old_scripts/3_Dynamics/3.12_MACE_NPT_Langevin.py diff --git a/examples/scripts/3_Dynamics/3.13_MACE_NVE_non_pbc.py b/examples/old_scripts/3_Dynamics/3.13_MACE_NVE_non_pbc.py similarity index 100% rename from examples/scripts/3_Dynamics/3.13_MACE_NVE_non_pbc.py rename to examples/old_scripts/3_Dynamics/3.13_MACE_NVE_non_pbc.py diff --git a/examples/scripts/3_Dynamics/3.1_Lennard_Jones_NVE.py b/examples/old_scripts/3_Dynamics/3.1_Lennard_Jones_NVE.py similarity index 100% rename from examples/scripts/3_Dynamics/3.1_Lennard_Jones_NVE.py rename to examples/old_scripts/3_Dynamics/3.1_Lennard_Jones_NVE.py diff --git a/examples/scripts/3_Dynamics/3.2_MACE_NVE.py b/examples/old_scripts/3_Dynamics/3.2_MACE_NVE.py similarity index 100% rename from examples/scripts/3_Dynamics/3.2_MACE_NVE.py rename to examples/old_scripts/3_Dynamics/3.2_MACE_NVE.py diff --git a/examples/scripts/3_Dynamics/3.3_MACE_NVE_cueq.py b/examples/old_scripts/3_Dynamics/3.3_MACE_NVE_cueq.py similarity index 100% rename from examples/scripts/3_Dynamics/3.3_MACE_NVE_cueq.py rename to examples/old_scripts/3_Dynamics/3.3_MACE_NVE_cueq.py diff --git a/examples/scripts/3_Dynamics/3.4_MACE_NVT_Langevin.py b/examples/old_scripts/3_Dynamics/3.4_MACE_NVT_Langevin.py similarity index 100% rename from examples/scripts/3_Dynamics/3.4_MACE_NVT_Langevin.py rename to examples/old_scripts/3_Dynamics/3.4_MACE_NVT_Langevin.py diff --git a/examples/scripts/3_Dynamics/3.5_MACE_NVT_Nose_Hoover.py b/examples/old_scripts/3_Dynamics/3.5_MACE_NVT_Nose_Hoover.py similarity index 100% rename from examples/scripts/3_Dynamics/3.5_MACE_NVT_Nose_Hoover.py rename to examples/old_scripts/3_Dynamics/3.5_MACE_NVT_Nose_Hoover.py diff --git a/examples/scripts/3_Dynamics/3.6_MACE_NVT_Nose_Hoover_temp_profile.py b/examples/old_scripts/3_Dynamics/3.6_MACE_NVT_Nose_Hoover_temp_profile.py similarity index 100% rename from examples/scripts/3_Dynamics/3.6_MACE_NVT_Nose_Hoover_temp_profile.py rename to examples/old_scripts/3_Dynamics/3.6_MACE_NVT_Nose_Hoover_temp_profile.py diff --git a/examples/scripts/3_Dynamics/3.7_Lennard_Jones_NPT_Nose_Hoover.py b/examples/old_scripts/3_Dynamics/3.7_Lennard_Jones_NPT_Nose_Hoover.py similarity index 100% rename from examples/scripts/3_Dynamics/3.7_Lennard_Jones_NPT_Nose_Hoover.py rename to examples/old_scripts/3_Dynamics/3.7_Lennard_Jones_NPT_Nose_Hoover.py diff --git a/examples/scripts/3_Dynamics/3.8_MACE_NPT_Nose_Hoover.py b/examples/old_scripts/3_Dynamics/3.8_MACE_NPT_Nose_Hoover.py similarity index 100% rename from examples/scripts/3_Dynamics/3.8_MACE_NPT_Nose_Hoover.py rename to examples/old_scripts/3_Dynamics/3.8_MACE_NPT_Nose_Hoover.py diff --git a/examples/scripts/3_Dynamics/3.9_MACE_NVT_staggered_stress.py b/examples/old_scripts/3_Dynamics/3.9_MACE_NVT_staggered_stress.py similarity index 100% rename from examples/scripts/3_Dynamics/3.9_MACE_NVT_staggered_stress.py rename to examples/old_scripts/3_Dynamics/3.9_MACE_NVT_staggered_stress.py diff --git a/examples/scripts/4_High_level_api/4.1_high_level_api.py b/examples/old_scripts/4_High_level_api/4.1_high_level_api.py similarity index 100% rename from examples/scripts/4_High_level_api/4.1_high_level_api.py rename to examples/old_scripts/4_High_level_api/4.1_high_level_api.py diff --git a/examples/scripts/4_High_level_api/4.2_auto_batching_api.py b/examples/old_scripts/4_High_level_api/4.2_auto_batching_api.py similarity index 100% rename from examples/scripts/4_High_level_api/4.2_auto_batching_api.py rename to examples/old_scripts/4_High_level_api/4.2_auto_batching_api.py diff --git a/examples/scripts/5_Workflow/5.1_a2c_silicon_batched.py b/examples/old_scripts/5_Workflow/5.1_a2c_silicon_batched.py similarity index 100% rename from examples/scripts/5_Workflow/5.1_a2c_silicon_batched.py rename to examples/old_scripts/5_Workflow/5.1_a2c_silicon_batched.py diff --git a/examples/scripts/5_Workflow/5.2_In_Flight_WBM.py b/examples/old_scripts/5_Workflow/5.2_In_Flight_WBM.py similarity index 100% rename from examples/scripts/5_Workflow/5.2_In_Flight_WBM.py rename to examples/old_scripts/5_Workflow/5.2_In_Flight_WBM.py diff --git a/examples/scripts/5_Workflow/5.3_Elastic.py b/examples/old_scripts/5_Workflow/5.3_Elastic.py similarity index 100% rename from examples/scripts/5_Workflow/5.3_Elastic.py rename to examples/old_scripts/5_Workflow/5.3_Elastic.py diff --git a/examples/scripts/6_Phonons/6.1_Phonons_MACE.py b/examples/old_scripts/6_Phonons/6.1_Phonons_MACE.py similarity index 100% rename from examples/scripts/6_Phonons/6.1_Phonons_MACE.py rename to examples/old_scripts/6_Phonons/6.1_Phonons_MACE.py diff --git a/examples/scripts/6_Phonons/6.2_QuasiHarmonic_MACE.py b/examples/old_scripts/6_Phonons/6.2_QuasiHarmonic_MACE.py similarity index 100% rename from examples/scripts/6_Phonons/6.2_QuasiHarmonic_MACE.py rename to examples/old_scripts/6_Phonons/6.2_QuasiHarmonic_MACE.py diff --git a/examples/scripts/6_Phonons/6.3_Conductivity_MACE.py b/examples/old_scripts/6_Phonons/6.3_Conductivity_MACE.py similarity index 100% rename from examples/scripts/6_Phonons/6.3_Conductivity_MACE.py rename to examples/old_scripts/6_Phonons/6.3_Conductivity_MACE.py diff --git a/examples/scripts/7_Others/7.1_Soft_sphere_autograd.py b/examples/old_scripts/7_Others/7.1_Soft_sphere_autograd.py similarity index 100% rename from examples/scripts/7_Others/7.1_Soft_sphere_autograd.py rename to examples/old_scripts/7_Others/7.1_Soft_sphere_autograd.py diff --git a/examples/scripts/7_Others/7.2_Stress_autograd.py b/examples/old_scripts/7_Others/7.2_Stress_autograd.py similarity index 100% rename from examples/scripts/7_Others/7.2_Stress_autograd.py rename to examples/old_scripts/7_Others/7.2_Stress_autograd.py diff --git a/examples/scripts/7_Others/7.3_Batched_neighbor_list.py b/examples/old_scripts/7_Others/7.3_Batched_neighbor_list.py similarity index 100% rename from examples/scripts/7_Others/7.3_Batched_neighbor_list.py rename to examples/old_scripts/7_Others/7.3_Batched_neighbor_list.py diff --git a/examples/scripts/7_Others/7.4_Velocity_AutoCorrelation.py b/examples/old_scripts/7_Others/7.4_Velocity_AutoCorrelation.py similarity index 100% rename from examples/scripts/7_Others/7.4_Velocity_AutoCorrelation.py rename to examples/old_scripts/7_Others/7.4_Velocity_AutoCorrelation.py diff --git a/examples/scripts/7_Others/7.6_Compare_ASE_to_VV_FIRE.py b/examples/old_scripts/7_Others/7.6_Compare_ASE_to_VV_FIRE.py similarity index 100% rename from examples/scripts/7_Others/7.6_Compare_ASE_to_VV_FIRE.py rename to examples/old_scripts/7_Others/7.6_Compare_ASE_to_VV_FIRE.py diff --git a/examples/scripts/7_Others/7.7_Heat_flux_and_kappa.py b/examples/old_scripts/7_Others/7.7_Heat_flux_and_kappa.py similarity index 100% rename from examples/scripts/7_Others/7.7_Heat_flux_and_kappa.py rename to examples/old_scripts/7_Others/7.7_Heat_flux_and_kappa.py diff --git a/examples/old_scripts/readme.md b/examples/old_scripts/readme.md new file mode 100644 index 00000000..118977c3 --- /dev/null +++ b/examples/old_scripts/readme.md @@ -0,0 +1,75 @@ +# TorchSim Example Scripts + +This folder contains a series of examples demonstrating the use of TorchSim, a library for simulating molecular dynamics and structural optimization using classical and machine learning interatomic potentials. Each example showcases different functionalities and models available in TorchSim. + +1. **Introduction** + + 1. **Lennard-Jones Model** - [`examples/1_Introduction/1.1_Lennard_Jones.py`](1_Introduction/1.1_Lennard_Jones.py): Simulate Argon atoms in an FCC lattice with a Lennard-Jones potential. Initialize the model, run a forward pass, and print energy, forces, and stress. + + 1. **MACE Model** - [`examples/1_Introduction/1.2_MACE.py`](1_Introduction/1.2_MACE.py): Use the MACE model to simulate diamond cubic Silicon. Load a pre-trained model, set up the system, and calculate energy, forces, and stress. + + 1. **Batched MACE Model** - [`examples/1_Introduction/1.3_Batched_MACE.py`](1_Introduction/1.3_Batched_MACE.py): Handle batched inputs with the MACE model to simulate multiple systems simultaneously. + + 1. **Fairchem Model** - [`examples/1_Introduction/1.4_Fairchem.py`](1_Introduction/1.4_Fairchem.py): Simulate diamond cubic Silicon with the Fairchem model. Set up the model and calculate energy, forces, and stress. + +1. **Structural Optimization** + + 1. **Lennard-Jones FIRE** - [`examples/2_Structural_optimization/2.1_Lennard_Jones_FIRE.py`](2_Structural_optimization/2.1_Lennard_Jones_FIRE.py): Perform structural optimization using the FIRE optimizer with a Lennard-Jones model. + + 1. **Soft Sphere FIRE** - [`examples/2_Structural_optimization/2.2_Soft_Sphere_FIRE.py`](2_Structural_optimization/2.2_Soft_Sphere_FIRE.py): Optimize structures with a Soft Sphere model using the FIRE optimizer. + + 1. **MACE FIRE** - [`examples/2_Structural_optimization/2.3_MACE_FIRE.py`](2_Structural_optimization/2.3_MACE_FIRE.py): Optimize structures with the MACE model using the FIRE optimizer. + + 1. **MACE UnitCellFilter FIRE** - [`examples/2_Structural_optimization/2.4_MACE_UnitCellFilter_FIRE.py`](2_Structural_optimization/2.4_MACE_UnitCellFilter_FIRE.py): Optimize structures with the MACE model using the UnitCellFilter FIRE optimizer. + + 1. **MACE FrechetCellFilter FIRE** - [`examples/2_Structural_optimization/2.5_MACE_FrechetCellFilter_FIRE.py`](2_Structural_optimization/2.5_MACE_FrechetCellFilter_FIRE.py): Optimize structures with the MACE model using the FrechetCellFilter FIRE optimizer. + + 1. **Batched MACE Gradient Descent** - [`examples/2_Structural_optimization/2.6_Batched_MACE_Gradient_Descent.py`](2_Structural_optimization/2.6_Batched_MACE_Gradient_Descent.py): Optimize multiple structures simultaneously using batched gradient descent with the MACE model. + + 1. **Batched MACE FIRE** - [`examples/2_Structural_optimization/2.7_Batched_MACE_FIRE.py`](2_Structural_optimization/2.7_Batched_MACE_FIRE.py): Optimize multiple structures simultaneously using the batched FIRE optimizer with MACE. + + 1. **Batched MACE UnitCellFilter Gradient Descent** - [`examples/2_Structural_optimization/2.8_Batched_MACE_UnitCellFilter_Gradient_Descent.py`](2_Structural_optimization/2.8_Batched_MACE_UnitCellFilter_Gradient_Descent.py): Optimize multiple structures and their unit cells using batched gradient descent with MACE. + + 1. **Batched MACE UnitCellFilter FIRE** - [`examples/2_Structural_optimization/2.9_Batched_MACE_UnitCellFilter_FIRE.py`](2_Structural_optimization/2.9_Batched_MACE_UnitCellFilter_FIRE.py): Optimize multiple structures and their unit cells using the batched FIRE optimizer with MACE. + +1. **Dynamics** + + 1. **Lennard-Jones NVE** - [`examples/3_Dynamics/3.1_Lennard_Jones_NVE.py`](3_Dynamics/3.1_Lennard_Jones_NVE.py): Run molecular dynamics with the NVE ensemble using a Lennard-Jones model. Set up, simulate, and check energy conservation. + + 1. **MACE NVE** - [`examples/3_Dynamics/3.2_MACE_NVE.py`](3_Dynamics/3.2_MACE_NVE.py): Run NVE molecular dynamics simulation with the MACE model. + + 1. **MACE NVE with Cueq** - [`examples/3_Dynamics/3.3_MACE_NVE_cueq.py`](3_Dynamics/3.3_MACE_NVE_cueq.py): Run the MACE model in NVE with CuEq acceleration. + + 1. **MACE NVT Langevin** - [`examples/3_Dynamics/3.4_MACE_NVT_Langevin.py`](3_Dynamics/3.4_MACE_NVT_Langevin.py): Run temperature-controlled molecular dynamics using the NVT Langevin integrator with MACE. + + 1. **MACE NVT Nose-Hoover** - [`examples/3_Dynamics/3.5_MACE_NVT_Nose_Hoover.py`](3_Dynamics/3.5_MACE_NVT_Nose_Hoover.py): Run temperature-controlled molecular dynamics using the NVT Nose-Hoover integrator with MACE. + + 1. **MACE NVT Nose-Hoover with Temperature Profile** - [`examples/3_Dynamics/3.6_MACE_NVT_Nose_Hoover_temp_profile.py`](3_Dynamics/3.6_MACE_NVT_Nose_Hoover_temp_profile.py): Simulate heating and cooling cycles using Nose-Hoover integrator with a temperature profile. + + 1. **Lennard-Jones NPT Nose-Hoover** - [`examples/3_Dynamics/3.7_Lennard_Jones_NPT_Nose_Hoover.py`](3_Dynamics/3.7_Lennard_Jones_NPT_Nose_Hoover.py): Run pressure-controlled molecular dynamics using the NPT Nose-Hoover integrator with Lennard-Jones. + + 1. **MACE NPT Nose-Hoover** - [`examples/3_Dynamics/3.8_MACE_NPT_Nose_Hoover.py`](3_Dynamics/3.8_MACE_NPT_Nose_Hoover.py): Run pressure-controlled molecular dynamics using the NPT Nose-Hoover integrator with MACE. + + 1. **MACE NVT with Staggered Stress** - [`examples/3_Dynamics/3.9_MACE_NVT_staggered_stress.py`](3_Dynamics/3.9_MACE_NVT_staggered_stress.py): Use staggered stress calculations during NVT simulations with the MACE model. + + 1. **Hybrid Swap Monte Carlo** - [`examples/3_Dynamics/3.10_Hybrid_swap_mc.py`](3_Dynamics/3.10_Hybrid_swap_mc.py): Combine molecular dynamics with Monte Carlo simulations using the MACE model. + +1. **High-Level API** + + 1. **High-Level API** - [`examples/4_High_level_api/4.1_high_level_api.py`](4_High_level_api/4.1_high_level_api.py): Integrate systems using the high-level API with different models and integrators. + +1. **Workflow** + + 1. **Workflow** - [`examples/5_Workflow/5.1_a2c_silicon.py`](5_Workflow/5.1_a2c_silicon.py): Run the a2c workflow with the MACE model. + + 1. **Workflow** - [`examples/5_Workflow/5.4_Elastic.py`](5_Workflow/5.4_Elastic.py): Calculate elastic tensor, bulk modulus and shear modulus with MACE. + +1. **Phonons** + + 1. **Phonon DOS with MACE Batched** - [`examples/6_Phonons/6.1_Phonons_MACE.py`](6_Phonons/6.1_Phonons_MACE.py): Calculate DOS and band structure with MACE, batching over FC2 calculations. + + 1. **Thermal Conductivity with MACE** - [`examples/6_Phonons/6.2_QuasiHarmonic_MACE.py`](6_Phonons/6.2_QuasiHarmonic_MACE.py): Calculates quasi-harmonic properties with MACE, batching over volumes and FC2 calculations. + + 1. **Thermal Conductivity with MACE Batched** - [`examples/6_Phonons/6.3_Conductivity_MACE.py`](6_Phonons/6.3_Conductivity_MACE.py): Calculate the Wigner lattice conductivity with MACE, batching over FC2 and FC3 calculations. + +Each example is self-contained and can be run independently to explore TorchSim capabilities. The examples cover basic model setup to advanced simulation techniques, providing a comprehensive overview of the library's features. diff --git a/examples/scripts/1_Introduction/1.3_fairchem.py b/examples/scripts/1_Introduction/1.3_fairchem.py deleted file mode 100644 index 965a20ab..00000000 --- a/examples/scripts/1_Introduction/1.3_fairchem.py +++ /dev/null @@ -1,62 +0,0 @@ -"""Minimal FairChem example demonstrating batching.""" - -# /// script -# dependencies = ["fairchem-core>=2.2.0", "huggingface_hub"] -# /// - -import os - -import torch -from ase.build import bulk - -import torch_sim as ts -from torch_sim.models.fairchem import FairChemModel - - -# Optional Hugging Face login if HF_TOKEN is available (for private model access) -try: - from huggingface_hub import login as hf_login # type: ignore[import-not-found] -except ImportError: # pragma: no cover - optional dependency - hf_login = None # type: ignore[assignment] - -hf_token = os.environ.get("HF_TOKEN") -if hf_token and hf_login is not None: - hf_login(token=hf_token) -else: - print("Need to login to HuggingFace to access fairchem models") - raise SystemExit(0) - -device = torch.device("cuda" if torch.cuda.is_available() else "cpu") -dtype = torch.float32 - -# UMA = Unified Machine Learning for Atomistic simulations -MODEL_NAME = "uma-s-1" - -# Create diamond cubic Silicon -si_dc = bulk("Si", "diamond", a=5.43).repeat((2, 2, 2)) -atomic_numbers = si_dc.get_atomic_numbers() -model = FairChemModel( - model=MODEL_NAME, - task_name="omat", # Open Materials task for crystalline systems - device=device, -) -atoms_list = [si_dc, si_dc] -state = ts.io.atoms_to_state(atoms_list, device=device, dtype=dtype) - -results = model(state) - -print(results["energy"].shape) -print(results["forces"].shape) -if "stress" in results: - print(results["stress"].shape) - -print(f"Energy: {results['energy']}") -print(f"Forces: {results['forces']}") -if "stress" in results: - print(f"{results['stress']=}") - -# Check if the energy, forces, and stress are the same for the Si system across the batch -print(torch.max(torch.abs(results["energy"][0] - results["energy"][1]))) -print(torch.max(torch.abs(results["forces"][0] - results["forces"][1]))) -if "stress" in results: - print(torch.max(torch.abs(results["stress"][0] - results["stress"][1]))) diff --git a/examples/scripts/1_introduction.py b/examples/scripts/1_introduction.py new file mode 100644 index 00000000..40b7196f --- /dev/null +++ b/examples/scripts/1_introduction.py @@ -0,0 +1,199 @@ +"""Introduction to TorchSim - Basic Examples with Lennard-Jones and MACE models. + +This script demonstrates the fundamental usage of TorchSim with: +- Lennard-Jones model for simple classical potentials +- MACE model for machine learning potentials +""" + +# /// script +# dependencies = ["scipy>=1.15", "mace-torch>=0.3.12"] +# /// + +import itertools + +import numpy as np +import torch +from ase.build import bulk +from mace.calculators.foundations_models import mace_mp + +from torch_sim.models.lennard_jones import LennardJonesModel +from torch_sim.models.mace import MaceModel, MaceUrls + + +# Set up the device and data type +device = torch.device("cuda" if torch.cuda.is_available() else "cpu") +dtype = torch.float32 + + +# ============================================================================ +# SECTION 1: Lennard-Jones Model - Simple Classical Potential +# ============================================================================ +print("\n" + "=" * 70) +print("SECTION 1: Lennard-Jones Model") +print("=" * 70) + +# Create face-centered cubic (FCC) Argon +# 5.26 Å is a typical lattice constant for Ar +a_len = 5.26 # Lattice constant + +# Generate base FCC unit cell positions (scaled by lattice constant) +base_positions = torch.tensor( + [ + [0.0, 0.0, 0.0], # Corner + [0.0, 0.5, 0.5], # Face centers + [0.5, 0.0, 0.5], + [0.5, 0.5, 0.0], + ], + device=device, + dtype=dtype, +) + +# Create 4x4x4 supercell of FCC Argon manually +positions = [] +for i, j, k in itertools.product(range(4), range(4), range(4)): + for base_pos in base_positions: + # Add unit cell position + offset for supercell + pos = base_pos + torch.tensor([i, j, k], device=device, dtype=dtype) + positions.append(pos) + +# Stack the positions into a tensor +positions = torch.stack(positions) + +# Scale by lattice constant +positions = positions * a_len + +# Create the cell tensor +cell = torch.tensor( + [[4 * a_len, 0, 0], [0, 4 * a_len, 0], [0, 0, 4 * a_len]], device=device, dtype=dtype +) + +# Create the atomic numbers tensor +atomic_numbers = torch.full((positions.shape[0],), 18, device=device, dtype=torch.int) + +# Initialize the Lennard-Jones model +# Parameters: +# - sigma: distance at which potential is zero (3.405 Å for Ar) +# - epsilon: depth of potential well (0.0104 eV for Ar) +# - cutoff: distance beyond which interactions are ignored (typically 2.5*sigma) +lj_model = LennardJonesModel( + use_neighbor_list=True, + cutoff=2.5 * 3.405, + sigma=3.405, + epsilon=0.0104, + device=device, + dtype=dtype, + compute_forces=True, + compute_stress=True, + per_atom_energies=True, + per_atom_stresses=True, +) + +# State dict +state = dict( + positions=positions, cell=cell.unsqueeze(0), atomic_numbers=atomic_numbers, pbc=True +) + +# Run the simulation and get results +results = lj_model(state) + +# Print the results +print(f"Energy: {results['energy']}") +print(f"Forces shape: {results['forces'].shape}") +print(f"Stress shape: {results['stress'].shape}") +print(f"Per-atom energies shape: {results['energies'].shape}") +print(f"Per-atom stresses shape: {results['stresses'].shape}") + + +# ============================================================================ +# SECTION 2: MACE Model - Machine Learning Potential (Batched) +# ============================================================================ +print("\n" + "=" * 70) +print("SECTION 2: MACE Model with Batched Input") +print("=" * 70) + +# Load the raw model from the downloaded model +loaded_model = mace_mp( + model=MaceUrls.mace_mpa_medium, + return_raw_model=True, + default_dtype=str(dtype).removeprefix("torch."), + device=str(device), +) + +# Create diamond cubic Silicon +si_dc = bulk("Si", "diamond", a=5.43, cubic=True).repeat((2, 2, 2)) +atoms_list = [si_dc, si_dc] + +batched_model = MaceModel( + model=loaded_model, + device=device, + compute_forces=True, + compute_stress=True, + dtype=dtype, + enable_cueq=False, +) + +# First we will create a concatenated positions array +# This will have shape (16, 3) which is concatenated from two 8 atom systems +positions_numpy = np.concatenate([atoms.positions for atoms in atoms_list]) + +# stack cell vectors into a (2, 3, 3) array where the first index is batch dimension +cell_numpy = np.stack([atoms.cell.array for atoms in atoms_list]) + +# concatenate atomic numbers into a (16,) array +atomic_numbers_numpy = np.concatenate( + [atoms.get_atomic_numbers() for atoms in atoms_list] +) + +# convert to tensors +positions = torch.tensor(positions_numpy, device=device, dtype=dtype) +cell = torch.tensor(cell_numpy, device=device, dtype=dtype) +atomic_numbers = torch.tensor(atomic_numbers_numpy, device=device, dtype=torch.int) + +# create system idx array of shape (16,) which is 0 for first 8 atoms, 1 for last 8 atoms +atoms_per_system = torch.tensor( + [len(atoms) for atoms in atoms_list], device=device, dtype=torch.int +) +system_idx = torch.repeat_interleave( + torch.arange(len(atoms_per_system), device=device), atoms_per_system +) + +# You can see their shapes are as expected +print(f"Positions: {positions.shape}") +print(f"Cell: {cell.shape}") +print(f"Atomic numbers: {atomic_numbers.shape}") +print(f"System indices: {system_idx.shape}") + +# Now we can pass them to the model +results = batched_model( + dict( + positions=positions, + cell=cell, + atomic_numbers=atomic_numbers, + system_idx=system_idx, + pbc=True, + ) +) + +# The energy has shape (n_systems,) as the structures in a batch +print(f"Energy shape: {results['energy'].shape}") + +# The forces have shape (n_atoms, 3) same as positions +print(f"Forces shape: {results['forces'].shape}") + +# The stress has shape (n_systems, 3, 3) same as cell +print(f"Stress shape: {results['stress'].shape}") + +# Check if the energy, forces, and stress are the same for the Si system across the batch +print( + f"\nMax energy difference: {torch.max(torch.abs(results['energy'][0] - results['energy'][1]))}" +) +print( + f"Max forces difference: {torch.max(torch.abs(results['forces'][:8] - results['forces'][8:]))}" +) +print( + f"Max stress difference: {torch.max(torch.abs(results['stress'][0] - results['stress'][1]))}" +) + +print("\n" + "=" * 70) +print("Introduction examples completed!") +print("=" * 70) diff --git a/examples/scripts/2_structural_optimization.py b/examples/scripts/2_structural_optimization.py new file mode 100644 index 00000000..42350499 --- /dev/null +++ b/examples/scripts/2_structural_optimization.py @@ -0,0 +1,391 @@ +"""Structural Optimization Examples - FIRE and Gradient Descent with various cell filters. + +This script demonstrates structural optimization techniques with: +- FIRE optimizer (Fast Inertial Relaxation Engine) +- Gradient descent optimizer +- Different cell filters (none, unit cell, Frechet cell) +- Batched optimization for multiple structures +""" + +# /// script +# dependencies = ["scipy>=1.15", "mace-torch>=0.3.12"] +# /// + +import itertools +import os + +import numpy as np +import torch +from ase.build import bulk +from mace.calculators.foundations_models import mace_mp + +import torch_sim as ts +from torch_sim.models.lennard_jones import LennardJonesModel +from torch_sim.models.mace import MaceModel, MaceUrls +from torch_sim.units import UnitConversion + + +# Set up the device and data type +device = torch.device("cuda" if torch.cuda.is_available() else "cpu") +dtype = torch.float32 + +# Number of steps to run +SMOKE_TEST = os.getenv("CI") is not None +N_steps = 10 if SMOKE_TEST else 500 + + +# ============================================================================ +# SECTION 1: Lennard-Jones FIRE Optimization +# ============================================================================ +print("\n" + "=" * 70) +print("SECTION 1: Lennard-Jones FIRE Optimization") +print("=" * 70) + +# Set up the random number generator +generator = torch.Generator(device=device) +generator.manual_seed(42) + +# Create face-centered cubic (FCC) Argon +a_len = 5.26 # Lattice constant + +# Generate base FCC unit cell positions +base_positions = torch.tensor( + [ + [0.0, 0.0, 0.0], # Corner + [0.0, 0.5, 0.5], # Face centers + [0.5, 0.0, 0.5], + [0.5, 0.5, 0.0], + ], + device=device, + dtype=dtype, +) + +# Create 4x4x4 supercell of FCC Argon +positions = [] +for i, j, k in itertools.product(range(4), range(4), range(4)): + for base_pos in base_positions: + pos = base_pos + torch.tensor([i, j, k], device=device, dtype=dtype) + positions.append(pos) + +positions = torch.stack(positions) * a_len + +# Create cell and atomic properties +cell = torch.tensor( + [[4 * a_len, 0, 0], [0, 4 * a_len, 0], [0, 0, 4 * a_len]], + device=device, + dtype=dtype, +) +atomic_numbers = torch.full((positions.shape[0],), 18, device=device, dtype=torch.int) +# Add random perturbation to start with non-equilibrium structure +positions = positions + 0.2 * torch.randn( + positions.shape, generator=generator, device=device, dtype=dtype +) +masses = torch.full((positions.shape[0],), 39.948, device=device, dtype=dtype) + +# Initialize the Lennard-Jones model +lj_model = LennardJonesModel( + use_neighbor_list=False, + sigma=3.405, + epsilon=0.0104, + cutoff=2.5 * 3.405, + device=device, + dtype=dtype, + compute_forces=True, + compute_stress=False, +) + +# Create state +state = ts.SimState( + positions=positions, + masses=masses, + cell=cell.unsqueeze(0), + atomic_numbers=atomic_numbers, + pbc=True, +) + +# Run initial simulation +results = lj_model(state) + +# Initialize FIRE optimizer +state = ts.fire_init(state=state, model=lj_model, dt_start=0.005) + +# Run optimization +for step in range(N_steps): + if step % 100 == 0: + print(f"Step {step}: Potential energy: {state.energy[0].item()} eV") + state = ts.fire_step(state=state, model=lj_model, dt_max=0.01) + +print(f"Initial energy: {results['energy'][0].item()} eV") +print(f"Final energy: {state.energy[0].item()} eV") +print(f"Initial max force: {torch.max(torch.abs(results['forces'][0])).item()} eV/Å") +print(f"Final max force: {torch.max(torch.abs(state.forces[0])).item()} eV/Å") + + +# ============================================================================ +# SECTION 2: Batched MACE FIRE Optimization (Atomic Positions Only) +# ============================================================================ +print("\n" + "=" * 70) +print("SECTION 2: Batched MACE FIRE - Positions Only") +print("=" * 70) + +# Load MACE model +loaded_model = mace_mp( + model=MaceUrls.mace_mpa_medium, + return_raw_model=True, + default_dtype=str(dtype).removeprefix("torch."), + device=str(device), +) + +# Set random seed for reproducibility +rng = np.random.default_rng(seed=0) + +# Create different crystal structures with perturbations +si_dc = bulk("Si", "diamond", a=5.21, cubic=True).repeat((2, 2, 2)) +si_dc.positions += 0.2 * rng.standard_normal(si_dc.positions.shape) + +cu_dc = bulk("Cu", "fcc", a=3.85).repeat((2, 2, 2)) +cu_dc.positions += 0.2 * rng.standard_normal(cu_dc.positions.shape) + +fe_dc = bulk("Fe", "bcc", a=2.95).repeat((2, 2, 2)) +fe_dc.positions += 0.2 * rng.standard_normal(fe_dc.positions.shape) + +atoms_list = [si_dc, cu_dc, fe_dc] + +print(f"Silicon atoms: {len(si_dc)}") +print(f"Copper atoms: {len(cu_dc)}") +print(f"Iron atoms: {len(fe_dc)}") + +# Create batched model +model = MaceModel( + model=loaded_model, + device=device, + compute_forces=True, + compute_stress=True, + dtype=dtype, + enable_cueq=False, +) + +# Convert atoms to state +state = ts.io.atoms_to_state(atoms_list, device=device, dtype=dtype) +results = model(state) + +# Initialize FIRE optimizer +state = ts.fire_init(state=state, model=model, dt_start=0.005) + +print("\nRunning FIRE:") +for step in range(N_steps): + if step % 20 == 0: + print(f"Step {step}, Energy: {[energy.item() for energy in state.energy]}") + + state = ts.fire_step(state=state, model=model, dt_max=0.01) + +print(f"Initial energies: {[energy.item() for energy in results['energy']]} eV") +print(f"Final energies: {[energy.item() for energy in state.energy]} eV") + + +# ============================================================================ +# SECTION 3: Batched MACE Gradient Descent Optimization +# ============================================================================ +print("\n" + "=" * 70) +print("SECTION 3: Batched MACE Gradient Descent") +print("=" * 70) + +# Reset structures with new perturbations +si_dc = bulk("Si", "diamond", a=5.43, cubic=True).repeat((2, 2, 2)) +si_dc.positions += 0.2 * rng.standard_normal(si_dc.positions.shape) + +fe = bulk("Fe", "bcc", a=2.8665, cubic=True).repeat((3, 3, 3)) +fe.positions += 0.2 * rng.standard_normal(fe.positions.shape) + +atoms_list = [si_dc, fe] + +state = ts.io.atoms_to_state(atoms_list, device=device, dtype=dtype) +results = model(state) + +# Initialize gradient descent optimizer +learning_rate = 0.01 +state = ts.gradient_descent_init(state=state, model=model) + +print("\nRunning batched gradient descent:") +for step in range(N_steps): + if step % 10 == 0: + print(f"Step {step}, Energy: {[res.item() for res in state.energy]} eV") + state = ts.gradient_descent_step(state=state, model=model, pos_lr=learning_rate) + +print(f"Initial energies: {[res.item() for res in results['energy']]} eV") +print(f"Final energies: {[res.item() for res in state.energy]} eV") + + +# ============================================================================ +# SECTION 4: Unit Cell Filter with Gradient Descent +# ============================================================================ +print("\n" + "=" * 70) +print("SECTION 4: Unit Cell Filter with Gradient Descent") +print("=" * 70) + +# Recreate structures with perturbations +si_dc = bulk("Si", "diamond", a=5.21, cubic=True).repeat((2, 2, 2)) +si_dc.positions += 0.2 * rng.standard_normal(si_dc.positions.shape) + +cu_dc = bulk("Cu", "fcc", a=3.85).repeat((2, 2, 2)) +cu_dc.positions += 0.2 * rng.standard_normal(cu_dc.positions.shape) + +fe_dc = bulk("Fe", "bcc", a=2.95).repeat((2, 2, 2)) +fe_dc.positions += 0.2 * rng.standard_normal(fe_dc.positions.shape) + +atoms_list = [si_dc, cu_dc, fe_dc] + +# Convert atoms to state +state = ts.io.atoms_to_state(atoms_list, device=device, dtype=dtype) +results = model(state) + +# Use different learning rates for positions and cell +pos_lr, cell_lr = 0.01, 0.1 + +state = ts.gradient_descent_init( + state=state, + model=model, + cell_filter=ts.CellFilter.unit, + cell_factor=None, # Will default to atoms per system + hydrostatic_strain=False, + constant_volume=False, + scalar_pressure=0.0, +) + +print("\nRunning batched unit cell gradient descent:") +for step in range(N_steps): + if step % 20 == 0: + P1 = -torch.trace(state.stress[0]) * UnitConversion.eV_per_Ang3_to_GPa / 3 + P2 = -torch.trace(state.stress[1]) * UnitConversion.eV_per_Ang3_to_GPa / 3 + P3 = -torch.trace(state.stress[2]) * UnitConversion.eV_per_Ang3_to_GPa / 3 + + print( + f"Step {step}, Energy: {[energy.item() for energy in state.energy]}, " + f"P1={P1:.4f} GPa, P2={P2:.4f} GPa, P3={P3:.4f} GPa" + ) + + state = ts.gradient_descent_step( + state=state, model=model, pos_lr=pos_lr, cell_lr=cell_lr + ) + +print(f"Initial energies: {[energy.item() for energy in results['energy']]} eV") +print(f"Final energies: {[energy.item() for energy in state.energy]} eV") + + +# ============================================================================ +# SECTION 5: Unit Cell Filter with FIRE +# ============================================================================ +print("\n" + "=" * 70) +print("SECTION 5: Unit Cell Filter with FIRE") +print("=" * 70) + +# Recreate structures with perturbations +si_dc = bulk("Si", "diamond", a=5.21, cubic=True).repeat((2, 2, 2)) +si_dc.positions += 0.2 * rng.standard_normal(si_dc.positions.shape) + +cu_dc = bulk("Cu", "fcc", a=3.85).repeat((2, 2, 2)) +cu_dc.positions += 0.2 * rng.standard_normal(cu_dc.positions.shape) + +fe_dc = bulk("Fe", "bcc", a=2.95).repeat((2, 2, 2)) +fe_dc.positions += 0.2 * rng.standard_normal(fe_dc.positions.shape) + +atoms_list = [si_dc, cu_dc, fe_dc] + +# Convert atoms to state +state = ts.io.atoms_to_state(atoms_list, device=device, dtype=dtype) +results = model(state) + +# Initialize FIRE optimizer with unit cell filter +state = ts.fire_init( + state=state, + model=model, + cell_filter=ts.CellFilter.unit, + cell_factor=None, + hydrostatic_strain=False, + constant_volume=False, + scalar_pressure=0.0, +) + +print("\nRunning batched unit cell FIRE:") +for step in range(N_steps): + if step % 20 == 0: + P1 = -torch.trace(state.stress[0]) * UnitConversion.eV_per_Ang3_to_GPa / 3 + P2 = -torch.trace(state.stress[1]) * UnitConversion.eV_per_Ang3_to_GPa / 3 + P3 = -torch.trace(state.stress[2]) * UnitConversion.eV_per_Ang3_to_GPa / 3 + + print( + f"Step {step}, Energy: {[energy.item() for energy in state.energy]}, " + f"P1={P1:.4f} GPa, P2={P2:.4f} GPa, P3={P3:.4f} GPa" + ) + + state = ts.fire_step(state=state, model=model) + +print(f"Initial energies: {[energy.item() for energy in results['energy']]} eV") +print(f"Final energies: {[energy.item() for energy in state.energy]} eV") + + +# ============================================================================ +# SECTION 6: Frechet Cell Filter with FIRE +# ============================================================================ +print("\n" + "=" * 70) +print("SECTION 6: Frechet Cell Filter with FIRE") +print("=" * 70) + +# Recreate structures with perturbations +si_dc = bulk("Si", "diamond", a=5.21, cubic=True).repeat((2, 2, 2)) +si_dc.positions += 0.2 * rng.standard_normal(si_dc.positions.shape) + +cu_dc = bulk("Cu", "fcc", a=3.85).repeat((2, 2, 2)) +cu_dc.positions += 0.2 * rng.standard_normal(cu_dc.positions.shape) + +fe_dc = bulk("Fe", "bcc", a=2.95).repeat((2, 2, 2)) +fe_dc.positions += 0.2 * rng.standard_normal(fe_dc.positions.shape) + +atoms_list = [si_dc, cu_dc, fe_dc] + +# Convert atoms to state +state = ts.io.atoms_to_state(atoms_list, device=device, dtype=dtype) +results = model(state) + +# Initialize FIRE optimizer with Frechet cell filter +state = ts.fire_init( + state=state, + model=model, + cell_filter=ts.CellFilter.frechet, + cell_factor=None, + hydrostatic_strain=False, + constant_volume=False, + scalar_pressure=0.0, +) + +print("\nRunning batched frechet cell filter with FIRE:") +for step in range(N_steps): + if step % 20 == 0: + P1 = -torch.trace(state.stress[0]) * UnitConversion.eV_per_Ang3_to_GPa / 3 + P2 = -torch.trace(state.stress[1]) * UnitConversion.eV_per_Ang3_to_GPa / 3 + P3 = -torch.trace(state.stress[2]) * UnitConversion.eV_per_Ang3_to_GPa / 3 + + print( + f"Step {step}, Energy: {[energy.item() for energy in state.energy]}, " + f"P1={P1:.4f} GPa, P2={P2:.4f} GPa, P3={P3:.4f} GPa" + ) + + state = ts.fire_step(state=state, model=model) + +print(f"Initial energies: {[energy.item() for energy in results['energy']]} eV") +print(f"Final energies: {[energy.item() for energy in state.energy]} eV") + +initial_pressure = [ + -torch.trace(stress).item() * UnitConversion.eV_per_Ang3_to_GPa / 3 + for stress in results["stress"] +] +final_pressure = [ + -torch.trace(stress).item() * UnitConversion.eV_per_Ang3_to_GPa / 3 + for stress in state.stress +] +print(f"Initial pressure: {initial_pressure} GPa") +print(f"Final pressure: {final_pressure} GPa") + +print("\n" + "=" * 70) +print("Structural optimization examples completed!") +print("=" * 70) diff --git a/examples/scripts/3_dynamics.py b/examples/scripts/3_dynamics.py new file mode 100644 index 00000000..897c92c2 --- /dev/null +++ b/examples/scripts/3_dynamics.py @@ -0,0 +1,402 @@ +"""Molecular Dynamics Examples - Various ensembles and integrators. + +This script demonstrates molecular dynamics simulations with: +- NVE (microcanonical) ensemble +- NVT (canonical) ensemble with Langevin and Nose-Hoover thermostats +- NPT (isothermal-isobaric) ensemble with Nose-Hoover barostat +- Both Lennard-Jones and MACE models +""" + +# /// script +# dependencies = ["scipy>=1.15", "mace-torch>=0.3.12"] +# /// + +import itertools +import os +import time + +import torch +from ase.build import bulk +from mace.calculators.foundations_models import mace_mp + +import torch_sim as ts +from torch_sim.models.lennard_jones import LennardJonesModel +from torch_sim.models.mace import MaceModel, MaceUrls +from torch_sim.units import MetalUnits as Units + + +# Set up the device and data type +device = torch.device("cuda" if torch.cuda.is_available() else "cpu") +dtype = torch.float32 + +# Number of steps to run +SMOKE_TEST = os.getenv("CI") is not None +N_steps = 100 if SMOKE_TEST else 2_000 + +# Set random seed for reproducibility +torch.manual_seed(42) +if torch.cuda.is_available(): + torch.cuda.manual_seed_all(42) +torch.backends.cudnn.deterministic = True +torch.backends.cudnn.benchmark = False + +# Set up the random number generator +generator = torch.Generator(device=device) +generator.manual_seed(42) + + +# ============================================================================ +# SECTION 1: Lennard-Jones NVE (Microcanonical Ensemble) +# ============================================================================ +print("\n" + "=" * 70) +print("SECTION 1: Lennard-Jones NVE Simulation") +print("=" * 70) + +# Create face-centered cubic (FCC) Argon +a_len = 5.26 # Lattice constant + +# Generate base FCC unit cell positions +base_positions = torch.tensor( + [ + [0.0, 0.0, 0.0], # Corner + [0.0, 0.5, 0.5], # Face centers + [0.5, 0.0, 0.5], + [0.5, 0.5, 0.0], + ], + device=device, + dtype=dtype, +) + +# Create 4x4x4 supercell of FCC Argon +positions = [] +for i, j, k in itertools.product(range(4), range(4), range(4)): + for base_pos in base_positions: + pos = base_pos + torch.tensor([i, j, k], device=device, dtype=dtype) + positions.append(pos) + +positions = torch.stack(positions) * a_len + +# Create the cell tensor +cell = torch.tensor( + [[4 * a_len, 0, 0], [0, 4 * a_len, 0], [0, 0, 4 * a_len]], + device=device, + dtype=dtype, +) + +# Create the atomic numbers tensor (Argon = 18) +atomic_numbers = torch.full((positions.shape[0],), 18, device=device, dtype=torch.int) +# Create the masses tensor (Argon = 39.948 amu) +masses = torch.full((positions.shape[0],), 39.948, device=device, dtype=dtype) + +state = ts.SimState( + positions=positions, masses=masses, cell=cell, atomic_numbers=atomic_numbers, pbc=True +) + +# Initialize the Lennard-Jones model +lj_model = LennardJonesModel( + use_neighbor_list=False, + sigma=3.405, + epsilon=0.0104, + cutoff=2.5 * 3.405, + device=device, + dtype=dtype, + compute_forces=True, + compute_stress=True, +) + +# Run initial simulation +results = lj_model(state) + +# Set up NVE simulation +kT = torch.tensor(80 * Units.temperature, device=device, dtype=dtype) +dt = torch.tensor(0.001 * Units.time, device=device, dtype=dtype) + +# Initialize NVE integrator +state = ts.nve_init(state=state, model=lj_model, kT=kT, seed=1) + +# Run NVE simulation +for step in range(N_steps): + if step % 100 == 0: + # Calculate total energy (potential + kinetic) + total_energy = state.energy + ts.calc_kinetic_energy( + masses=state.masses, momenta=state.momenta + ) + print(f"Step {step}: Total energy: {total_energy.item():.4f} eV") + + # Update state using NVE integrator + state = ts.nve_step(state=state, model=lj_model, dt=dt) + +final_total_energy = state.energy + ts.calc_kinetic_energy( + masses=state.masses, momenta=state.momenta +) +print(f"Final total energy: {final_total_energy.item():.4f} eV") + + +# ============================================================================ +# SECTION 2: MACE NVE Simulation +# ============================================================================ +print("\n" + "=" * 70) +print("SECTION 2: MACE NVE Simulation") +print("=" * 70) + +# Load MACE model +loaded_model = mace_mp( + model=MaceUrls.mace_mpa_medium, + return_raw_model=True, + default_dtype=str(dtype).removeprefix("torch."), + device=str(device), +) + +# Create diamond cubic Silicon +si_dc = bulk("Si", "diamond", a=5.43, cubic=True).repeat((2, 2, 2)) + +# Prepare input tensors +positions = torch.tensor(si_dc.positions, device=device, dtype=dtype) +cell = torch.tensor(si_dc.cell.array, device=device, dtype=dtype) +atomic_numbers = torch.tensor(si_dc.get_atomic_numbers(), device=device, dtype=torch.int) +masses = torch.tensor(si_dc.get_masses(), device=device, dtype=dtype) + +# Initialize the MACE model +mace_model = MaceModel( + model=loaded_model, + device=device, + compute_forces=True, + compute_stress=False, + dtype=dtype, + enable_cueq=False, +) + +state = ts.SimState( + positions=positions, masses=masses, cell=cell, atomic_numbers=atomic_numbers, pbc=True +) + +# Run initial inference +results = mace_model(state) + +# Setup NVE MD simulation parameters +kT = torch.tensor(1000 * Units.temperature, device=device, dtype=dtype) # 1000 K +dt = torch.tensor(0.002 * Units.time, device=device, dtype=dtype) # 2 fs + +# Initialize NVE integrator +state = ts.nve_init(state=state, model=mace_model, kT=kT, seed=1) + +# Run MD simulation +print("\nStarting NVE molecular dynamics simulation...") +start_time = time.perf_counter() +for step in range(N_steps): + total_energy = state.energy + ts.calc_kinetic_energy( + masses=state.masses, momenta=state.momenta, system_idx=state.system_idx + ) + if step % 100 == 0: + print(f"Step {step}: Total energy: {total_energy.item():.4f} eV") + state = ts.nve_step(state=state, model=mace_model, dt=dt) +end_time = time.perf_counter() + +print("\nSimulation complete!") +print(f"Time taken: {end_time - start_time:.2f} seconds") +print(f"Average time per step: {(end_time - start_time) / N_steps:.4f} seconds") + + +# ============================================================================ +# SECTION 3: MACE NVT Langevin Simulation +# ============================================================================ +print("\n" + "=" * 70) +print("SECTION 3: MACE NVT Langevin Simulation") +print("=" * 70) + +# Create diamond cubic Silicon +si_dc = bulk("Si", "diamond", a=5.43, cubic=True).repeat((2, 2, 2)) + +# Prepare input tensors +positions = torch.tensor(si_dc.positions, device=device, dtype=dtype) +cell = torch.tensor(si_dc.cell.array, device=device, dtype=dtype) +atomic_numbers = torch.tensor(si_dc.get_atomic_numbers(), device=device, dtype=torch.int) +masses = torch.tensor(si_dc.get_masses(), device=device, dtype=dtype) + +state = ts.SimState( + positions=positions, masses=masses, cell=cell, atomic_numbers=atomic_numbers, pbc=True +) + +dt = torch.tensor(0.002 * Units.time, device=device, dtype=dtype) # 2 fs +kT = torch.tensor(1000 * Units.temperature, device=device, dtype=dtype) # 1000 K +gamma = torch.tensor(10 / Units.time, device=device, dtype=dtype) # ps^-1 + +# Initialize NVT Langevin integrator +state = ts.nvt_langevin_init(model=mace_model, state=state, kT=kT, seed=1) + +print("\nStarting NVT Langevin simulation...") +for step in range(N_steps): + if step % 100 == 0: + temp = ( + ts.calc_kT( + masses=state.masses, momenta=state.momenta, system_idx=state.system_idx + ) + / Units.temperature + ) + print(f"Step {step}: Temperature: {temp.item():.4f} K") + state = ts.nvt_langevin_step(state=state, model=mace_model, dt=dt, kT=kT, gamma=gamma) + +final_temp = ( + ts.calc_kT(masses=state.masses, momenta=state.momenta, system_idx=state.system_idx) + / Units.temperature +) +print(f"Final temperature: {final_temp.item():.4f} K") + + +# ============================================================================ +# SECTION 4: MACE NVT Nose-Hoover Simulation +# ============================================================================ +print("\n" + "=" * 70) +print("SECTION 4: MACE NVT Nose-Hoover Simulation") +print("=" * 70) + +# Create diamond cubic Silicon +si_dc = bulk("Si", "diamond", a=5.43, cubic=True).repeat((2, 2, 2)) + +state = ts.io.atoms_to_state(si_dc, device=device, dtype=dtype) + +# Run initial inference +results = mace_model(state) + +dt = torch.tensor(0.002 * Units.time, device=device, dtype=dtype) # 2 fs +kT = torch.tensor(1000 * Units.temperature, device=device, dtype=dtype) # 1000 K + +state = ts.nvt_nose_hoover_init(state=state, model=mace_model, kT=kT, dt=dt) + +print("\nStarting NVT Nose-Hoover simulation...") +for step in range(N_steps): + if step % 100 == 0: + temp = ( + ts.calc_kT( + masses=state.masses, momenta=state.momenta, system_idx=state.system_idx + ) + / Units.temperature + ) + invariant = float(ts.nvt_nose_hoover_invariant(state, kT=kT)) + print( + f"Step {step}: Temperature: {temp.item():.4f} K, Invariant: {invariant:.4f}" + ) + state = ts.nvt_nose_hoover_step(state=state, model=mace_model, dt=dt, kT=kT) + +final_temp = ( + ts.calc_kT(masses=state.masses, momenta=state.momenta, system_idx=state.system_idx) + / Units.temperature +) +print(f"Final temperature: {final_temp.item():.4f} K") + + +# ============================================================================ +# SECTION 5: MACE NPT Nose-Hoover Simulation +# ============================================================================ +print("\n" + "=" * 70) +print("SECTION 5: MACE NPT Nose-Hoover Simulation") +print("=" * 70) + +# Create diamond cubic Silicon +si_dc = bulk("Si", "diamond", a=5.43, cubic=True).repeat((2, 2, 2)) + +# Create model with stress computation enabled for NPT +mace_model_stress = MaceModel( + model=loaded_model, + device=device, + compute_forces=True, + compute_stress=True, + dtype=dtype, + enable_cueq=False, +) + +state = ts.io.atoms_to_state(si_dc, device=device, dtype=dtype) + +# Run initial inference +results = mace_model_stress(state) + +N_steps_nvt = 100 if SMOKE_TEST else 1_000 +N_steps_npt = 100 if SMOKE_TEST else 1_000 +dt = 0.001 * Units.time # 1 fs +kT = torch.tensor(300 * Units.temperature, device=device, dtype=dtype) # 300 K +target_pressure = torch.tensor(0.0 * Units.pressure, device=device, dtype=dtype) # 0 bar + +# Initialize NPT with NVT equilibration +state = ts.npt_nose_hoover_init( + state=state, model=mace_model_stress, kT=kT, dt=torch.tensor(dt) +) + +print("\nRunning NVT equilibration phase...") +for step in range(N_steps_nvt): + if step % 100 == 0: + temp = ( + ts.calc_kT( + masses=state.masses, momenta=state.momenta, system_idx=state.system_idx + ) + / Units.temperature + ) + invariant = float( + ts.npt_nose_hoover_invariant(state, kT=kT, external_pressure=target_pressure) + ) + print( + f"Step {step}: Temperature: {temp.item():.4f} K, Invariant: {invariant:.4f}" + ) + state = ts.npt_nose_hoover_step( + state=state, + model=mace_model_stress, + dt=torch.tensor(dt), + kT=kT, + external_pressure=target_pressure, + ) + +# Reinitialize for NPT phase +state = ts.npt_nose_hoover_init( + state=state, model=mace_model_stress, kT=kT, dt=torch.tensor(dt) +) + +print("\nRunning NPT simulation...") +for step in range(N_steps_npt): + if step % 100 == 0: + temp = ( + ts.calc_kT( + masses=state.masses, momenta=state.momenta, system_idx=state.system_idx + ) + / Units.temperature + ) + invariant = float( + ts.npt_nose_hoover_invariant(state, kT=kT, external_pressure=target_pressure) + ) + stress = mace_model_stress(state)["stress"] + volume = torch.det(state.current_cell) + e_kin = ts.calc_kinetic_energy( + masses=state.masses, momenta=state.momenta, system_idx=state.system_idx + ) + pressure = float(ts.get_pressure(stress, e_kin, volume)) + xx, yy, zz = torch.diag(state.current_cell[0]) + print( + f"Step {step}: Temperature: {temp.item():.4f} K, Invariant: {invariant:.4f}, " + f"Pressure: {pressure:.4f} eV/ų, Cell: [{xx.item():.4f}, {yy.item():.4f}, {zz.item():.4f}]" + ) + state = ts.npt_nose_hoover_step( + state=state, + model=mace_model_stress, + dt=torch.tensor(dt), + kT=kT, + external_pressure=target_pressure, + ) + +final_temp = ( + ts.calc_kT(masses=state.masses, momenta=state.momenta, system_idx=state.system_idx) + / Units.temperature +) +print(f"Final temperature: {final_temp.item():.4f} K") + +final_stress = mace_model_stress(state)["stress"] +final_volume = torch.det(state.current_cell) +final_pressure = ts.get_pressure( + final_stress, + ts.calc_kinetic_energy( + masses=state.masses, momenta=state.momenta, system_idx=state.system_idx + ), + final_volume, +) +print(f"Final pressure: {final_pressure.item():.4f} eV/ų") + +print("\n" + "=" * 70) +print("Molecular dynamics examples completed!") +print("=" * 70) diff --git a/examples/scripts/4_high_level_api.py b/examples/scripts/4_high_level_api.py new file mode 100644 index 00000000..6aec94c3 --- /dev/null +++ b/examples/scripts/4_high_level_api.py @@ -0,0 +1,287 @@ +"""High-Level API Examples - Simplified interface for common workflows. + +This script demonstrates the high-level API for: +- Integration with different models and integrators +- Trajectory logging and reporting +- Batched simulations +- Custom convergence criteria +- Support for ASE Atoms and Pymatgen Structure objects +""" + +# /// script +# dependencies = ["mace-torch>=0.3.12", "pymatgen>=2025.2.18"] +# /// + +import os + +import numpy as np +import torch +from ase.build import bulk +from mace.calculators.foundations_models import mace_mp +from pymatgen.core import Structure + +import torch_sim as ts +from torch_sim.models.lennard_jones import LennardJonesModel +from torch_sim.models.mace import MaceModel +from torch_sim.trajectory import TorchSimTrajectory, TrajectoryReporter +from torch_sim.units import MetalUnits + + +SMOKE_TEST = os.getenv("CI") is not None + +# Set device +device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + +# ============================================================================ +# SECTION 1: Basic Integration with Lennard-Jones +# ============================================================================ +print("\n" + "=" * 70) +print("SECTION 1: Basic Integration with Lennard-Jones") +print("=" * 70) + +lj_model = LennardJonesModel( + sigma=2.0, # Å, typical for Si-Si interaction + epsilon=0.1, # eV, typical for Si-Si interaction + device=device, + dtype=torch.float64, +) + +si_atoms = bulk("Si", "fcc", a=5.43, cubic=True) + +final_state = ts.integrate( + system=si_atoms, + model=lj_model, + integrator=ts.Integrator.nvt_langevin, + n_steps=100 if SMOKE_TEST else 1000, + temperature=2000, + timestep=0.002, +) +final_atoms = ts.io.state_to_atoms(final_state) + +print(f"Final energy: {final_state.energy.item():.4f} eV") +print(f"Final atoms: {len(final_atoms)} atoms") + + +# ============================================================================ +# SECTION 2: Integration with Trajectory Logging +# ============================================================================ +print("\n" + "=" * 70) +print("SECTION 2: Integration with Trajectory Logging") +print("=" * 70) + +trajectory_file = "tmp/lj_trajectory.h5md" + +# Report potential energy every 10 steps and kinetic energy every 20 steps +prop_calculators = { + 10: {"potential_energy": lambda state: state.energy}, + 20: { + "kinetic_energy": lambda state: ts.calc_kinetic_energy( + momenta=state.momenta, masses=state.masses + ) + }, +} + +reporter = TrajectoryReporter( + trajectory_file, + state_frequency=10, # Report state every 10 steps + prop_calculators=prop_calculators, +) + +final_state = ts.integrate( + system=si_atoms, + model=lj_model, + integrator=ts.Integrator.nvt_langevin, + n_steps=100 if SMOKE_TEST else 1000, + temperature=2000, + timestep=0.002, + trajectory_reporter=reporter, +) + +# Read trajectory data +with TorchSimTrajectory(trajectory_file) as traj: + kinetic_energies = traj.get_array("kinetic_energy") + potential_energies = traj.get_array("potential_energy") + final_energy = potential_energies[-1] + final_atoms = traj.get_atoms(-1) + +print(f"Final energy from trajectory: {final_energy:.4f} eV") +print(f"Number of kinetic energy samples: {len(kinetic_energies)}") +print(f"Number of potential energy samples: {len(potential_energies)}") + + +# ============================================================================ +# SECTION 3: MACE Model with High-Level API +# ============================================================================ +print("\n" + "=" * 70) +print("SECTION 3: MACE Model with High-Level API") +print("=" * 70) + +mace = mace_mp(model="small", return_raw_model=True) +mace_model = MaceModel( + model=mace, + device=device, + dtype=torch.float64, + compute_forces=True, +) + +reporter = TrajectoryReporter( + trajectory_file, + state_frequency=10, + prop_calculators=prop_calculators, +) + +final_state = ts.integrate( + system=si_atoms, + model=mace_model, + integrator=ts.Integrator.nvt_langevin, + n_steps=100 if SMOKE_TEST else 1000, + temperature=2000, + timestep=0.002, + trajectory_reporter=reporter, +) +final_atoms = ts.io.state_to_atoms(final_state) + +print(f"Final energy: {final_state.energy.item():.4f} eV") + + +# ============================================================================ +# SECTION 4: Batched Integration +# ============================================================================ +print("\n" + "=" * 70) +print("SECTION 4: Batched Integration") +print("=" * 70) + +fe_atoms = bulk("Fe", "fcc", a=5.26, cubic=True) +fe_atoms_supercell = fe_atoms.repeat([2, 2, 2]) +si_atoms_supercell = si_atoms.repeat([2, 2, 2]) + +final_state = ts.integrate( + system=[si_atoms, fe_atoms, si_atoms_supercell, fe_atoms_supercell], + model=mace_model, + integrator=ts.Integrator.nvt_langevin, + n_steps=100 if SMOKE_TEST else 1000, + temperature=2000, + timestep=0.002, +) +final_atoms = ts.io.state_to_atoms(final_state) +final_fe_atoms_supercell = final_atoms[3] + +print(f"Number of systems: {len(final_atoms)}") +print(f"Final energies: {[e.item() for e in final_state.energy]} eV") + + +# ============================================================================ +# SECTION 5: Batched Integration with Trajectory Reporting +# ============================================================================ +print("\n" + "=" * 70) +print("SECTION 5: Batched Integration with Trajectory Reporting") +print("=" * 70) + +systems = (si_atoms, fe_atoms, si_atoms_supercell, fe_atoms_supercell) + +filenames = [f"tmp/batch_traj_{i}.h5md" for i in range(len(systems))] +batch_reporter = TrajectoryReporter( + filenames, + state_frequency=100, + prop_calculators=prop_calculators, +) + +final_state = ts.integrate( + system=systems, + model=mace_model, + integrator=ts.Integrator.nvt_langevin, + n_steps=100 if SMOKE_TEST else 1000, + temperature=2000, + timestep=0.002, + trajectory_reporter=batch_reporter, +) + +final_energies_per_atom = [] +for filename in filenames: + with TorchSimTrajectory(filename) as traj: + final_energy = traj.get_array("potential_energy")[-1] + final_energies_per_atom.append(final_energy / len(traj.get_atoms(-1))) + +print(f"Final energies per atom: {final_energies_per_atom}") + + +# ============================================================================ +# SECTION 6: Structure Optimization +# ============================================================================ +print("\n" + "=" * 70) +print("SECTION 6: Structure Optimization") +print("=" * 70) + +final_state = ts.optimize( + system=systems, + model=mace_model, + optimizer=ts.Optimizer.fire, + max_steps=10 if SMOKE_TEST else 1000, + init_kwargs=dict(cell_filter=ts.CellFilter.unit), +) + +print(f"Final optimized energies: {[e.item() for e in final_state.energy]} eV") + +# Add perturbations +rng = np.random.default_rng() +for system in systems: + system.positions += rng.random(system.positions.shape) * 0.01 + + +# ============================================================================ +# SECTION 7: Optimization with Custom Convergence Criteria +# ============================================================================ +print("\n" + "=" * 70) +print("SECTION 7: Optimization with Custom Convergence") +print("=" * 70) + +final_state = ts.optimize( + system=systems, + model=mace_model, + optimizer=ts.Optimizer.fire, + convergence_fn=lambda state, last_energy: last_energy - state.energy + < 1e-6 * MetalUnits.energy, + max_steps=10 if SMOKE_TEST else 1000, + init_kwargs=dict(cell_filter=ts.CellFilter.unit), +) + +print(f"Final converged energies: {[e.item() for e in final_state.energy]} eV") + + +# ============================================================================ +# SECTION 8: Pymatgen Structure Support +# ============================================================================ +print("\n" + "=" * 70) +print("SECTION 8: Pymatgen Structure Support") +print("=" * 70) + +lattice = [[5.43, 0, 0], [0, 5.43, 0], [0, 0, 5.43]] +species = ["Si"] * 8 +coords = [ + [0.0, 0.0, 0.0], + [0.25, 0.25, 0.25], + [0.0, 0.5, 0.5], + [0.25, 0.75, 0.75], + [0.5, 0.0, 0.5], + [0.75, 0.25, 0.75], + [0.5, 0.5, 0.0], + [0.75, 0.75, 0.25], +] +structure = Structure(lattice, species, coords) + +final_state = ts.integrate( + system=structure, + model=lj_model, + integrator=ts.Integrator.nvt_langevin, + n_steps=100 if SMOKE_TEST else 1000, + temperature=2000, + timestep=0.002, +) +final_structure = ts.io.state_to_structures(final_state) + +print(f"Final structure type: {type(final_structure)}") +print(f"Final energy: {final_state.energy.item():.4f} eV") + +print("\n" + "=" * 70) +print("High-level API examples completed!") +print("=" * 70) diff --git a/examples/scripts/5_workflow.py b/examples/scripts/5_workflow.py new file mode 100644 index 00000000..c0ea79a5 --- /dev/null +++ b/examples/scripts/5_workflow.py @@ -0,0 +1,233 @@ +"""Advanced Workflows - Complex simulation workflows and utilities. + +This script demonstrates: +- In-flight autobatching for efficient optimization +- Elastic constant calculations +- Force convergence utilities +""" + +# /// script +# dependencies = ["mace-torch>=0.3.12", "matbench-discovery>=1.3.1"] +# /// + +import os +import time + +import numpy as np +import torch +from ase.build import bulk +from mace.calculators.foundations_models import mace_mp + +import torch_sim as ts +from torch_sim.elastic import get_bravais_type +from torch_sim.models.mace import MaceModel, MaceUrls + + +# Set device +SMOKE_TEST = os.getenv("CI") is not None +device = torch.device( + "cpu" if SMOKE_TEST else ("cuda" if torch.cuda.is_available() else "cpu") +) +dtype = torch.float32 + +print(f"Running on device: {device}") + + +# ============================================================================ +# SECTION 1: In-Flight Autobatching Workflow +# ============================================================================ +print("\n" + "=" * 70) +print("SECTION 1: In-Flight Autobatching Workflow") +print("=" * 70) + +print("Loading MACE model...") +mace = mace_mp(model=MaceUrls.mace_mpa_medium, return_raw_model=True) +mace_model = MaceModel( + model=mace, + device=device, + dtype=dtype, + compute_forces=True, +) + +# Optimization parameters +fmax = 0.05 # Force convergence threshold +max_atoms_in_batch = 50 if SMOKE_TEST else 8_000 + +# Load or create structures +if not SMOKE_TEST: + try: + from matbench_discovery.data import DataFiles, ase_atoms_from_zip + + n_structures_to_relax = 100 + print(f"Loading {n_structures_to_relax:,} structures from WBM dataset...") + ase_atoms_list = ase_atoms_from_zip( + DataFiles.wbm_initial_atoms.path, limit=n_structures_to_relax + ) + except ImportError: + print("matbench_discovery not available, using synthetic structures...") + n_structures_to_relax = 10 + ase_atoms_list = [] + for _ in range(n_structures_to_relax): + atoms = bulk("Al", "hcp", a=4.05).repeat((2, 2, 2)) + atoms.positions += 0.1 * np.random.randn(*atoms.positions.shape) + ase_atoms_list.append(atoms) +else: + n_structures_to_relax = 2 + print(f"Loading {n_structures_to_relax:,} test structures...") + al_atoms = bulk("Al", "hcp", a=4.05) + al_atoms.positions += 0.1 * np.random.randn(*al_atoms.positions.shape) + fe_atoms = bulk("Fe", "bcc", a=2.86).repeat((2, 2, 2)) + fe_atoms.positions += 0.1 * np.random.randn(*fe_atoms.positions.shape) + ase_atoms_list = [al_atoms, fe_atoms] + +# Initialize first batch +fire_states = ts.fire_init( + state=ts.io.atoms_to_state(atoms=ase_atoms_list, device=device, dtype=dtype), + model=mace_model, + cell_filter=ts.CellFilter.frechet, +) + +# Create autobatcher +batcher = ts.autobatching.InFlightAutoBatcher( + model=mace_model, + memory_scales_with="n_atoms_x_density", + max_memory_scaler=1000 if SMOKE_TEST else None, +) + +# Create convergence function +converge_max_force = ts.runners.generate_force_convergence_fn(force_tol=fmax) + +start_time = time.perf_counter() + +# Main optimization loop with autobatching +batcher.load_states(fire_states) +all_completed_states, convergence_tensor, state = [], None, None + +print("\nStarting optimization with autobatching...") +batch_count = 0 +while (result := batcher.next_batch(state, convergence_tensor))[0] is not None: + state, completed_states = result + batch_count += 1 + print(f"Batch {batch_count}: Optimizing {state.n_systems} structures") + + all_completed_states.extend(completed_states) + if all_completed_states: + print(f"Total completed structures: {len(all_completed_states)}") + + # Run optimization steps for this batch + for _step in range(10): + state = ts.fire_step(state=state, model=mace_model) + + # Check convergence + convergence_tensor = converge_max_force(state, last_energy=None) + +# Add final completed states +all_completed_states.extend(result[1]) + +end_time = time.perf_counter() +total_time = end_time - start_time + +print("\nOptimization complete!") +print(f"Total completed structures: {len(all_completed_states)}") +print(f"Total time: {total_time:.2f} seconds") +print(f"Average time per structure: {total_time / len(all_completed_states):.2f} seconds") + + +# ============================================================================ +# SECTION 2: Elastic Constants Calculation +# ============================================================================ +print("\n" + "=" * 70) +print("SECTION 2: Elastic Constants Calculation") +print("=" * 70) + +# Use higher precision for elastic constants +dtype_elastic = torch.float64 + +loaded_model = mace_mp( + model=MaceUrls.mace_mpa_medium, + enable_cueq=False, + device=str(device), + default_dtype=str(dtype_elastic).removeprefix("torch."), + return_raw_model=True, +) + +# Create FCC Copper structure +struct = bulk("Cu", "fcc", a=3.58, cubic=True).repeat((2, 2, 2)) + +model = MaceModel( + model=loaded_model, + device=device, + compute_forces=True, + compute_stress=True, + dtype=dtype_elastic, + enable_cueq=False, +) + +# Target force tolerance +fmax = 1e-3 + +# Relax positions and cell +state = ts.io.atoms_to_state(atoms=struct, device=device, dtype=dtype_elastic) +state = ts.fire_init( + state=state, model=model, scalar_pressure=0.0, cell_filter=ts.CellFilter.frechet +) + +print("\nRelaxing structure...") +unit_conv = ts.units.UnitConversion +for step in range(300): + pressure = -torch.trace(state.stress.squeeze()) / 3 * unit_conv.eV_per_Ang3_to_GPa + current_fmax = torch.max(torch.abs(state.forces.squeeze())) + + if step % 50 == 0: + print( + f"Step {step}, Energy: {state.energy.item():.4f} eV, " + f"Pressure: {pressure.item():.4f} GPa, " + f"Fmax: {current_fmax.item():.4f} eV/Å" + ) + + if current_fmax < fmax and abs(pressure) < 1e-2: + print(f"Converged at step {step}") + break + + state = ts.fire_step(state=state, model=model) + +# Get bravais type +bravais_type = get_bravais_type(state) +print(f"\nBravais lattice type: {bravais_type}") + +# Calculate elastic tensor +print("\nCalculating elastic tensor...") +elastic_tensor = ts.elastic.calculate_elastic_tensor( + state=state, model=model, bravais_type=bravais_type +) + +# Convert to GPa +elastic_tensor = elastic_tensor * unit_conv.eV_per_Ang3_to_GPa + +# Calculate elastic moduli +bulk_modulus, shear_modulus, poisson_ratio, pugh_ratio = ( + ts.elastic.calculate_elastic_moduli(elastic_tensor) +) + +# Print results +print("\nElastic tensor (GPa):") +elastic_tensor_np = elastic_tensor.cpu().numpy() +for row in elastic_tensor_np: + print(" " + " ".join(f"{val:10.4f}" for val in row)) + +print("\nElastic moduli:") +print(f" Bulk modulus (GPa): {bulk_modulus:.4f}") +print(f" Shear modulus (GPa): {shear_modulus:.4f}") +print(f" Poisson's ratio: {poisson_ratio:.4f}") +print(f" Pugh's ratio (K/G): {pugh_ratio:.4f}") + +# Interpret Pugh's ratio +if pugh_ratio > 1.75: + material_type = "ductile" +else: + material_type = "brittle" +print(f" Material behavior: {material_type}") + +print("\n" + "=" * 70) +print("Workflow examples completed!") +print("=" * 70) diff --git a/examples/scripts/6_phonons.py b/examples/scripts/6_phonons.py new file mode 100644 index 00000000..2edcbca4 --- /dev/null +++ b/examples/scripts/6_phonons.py @@ -0,0 +1,255 @@ +"""Phonon Calculations - DOS, band structure, and thermal properties. + +This script demonstrates phonon calculations with: +- Phonon density of states (DOS) +- Phonon band structure +- Batched force constant calculations +- Integration with Phonopy + +Note: This example requires phonopy, pymatviz, plotly, seekpath, and ase packages. +Visualization is disabled in CI mode. +""" + +# /// script +# dependencies = [ +# "mace-torch>=0.3.12", +# "phonopy>=2.35", +# "pymatviz>=0.17.1", +# "plotly>=6.3.0", +# "seekpath", +# "ase", +# ] +# /// + +import os + +import numpy as np +import torch +from ase.build import bulk +from mace.calculators.foundations_models import mace_mp +from phonopy import Phonopy + +import torch_sim as ts +from torch_sim.models.mace import MaceModel, MaceUrls + + +# Set device and data type +device = torch.device("cuda" if torch.cuda.is_available() else "cpu") +dtype = torch.float32 + +SMOKE_TEST = os.getenv("CI") is not None + +# ============================================================================ +# SECTION 1: Structure Relaxation for Phonons +# ============================================================================ +print("\n" + "=" * 70) +print("SECTION 1: Structure Relaxation") +print("=" * 70) + +# Load the MACE model +loaded_model = mace_mp( + model=MaceUrls.mace_mpa_medium, + return_raw_model=True, + default_dtype=str(dtype).removeprefix("torch."), + device=str(device), +) + +# Create Silicon diamond structure +struct = bulk("Si", "diamond", a=5.431, cubic=True) +supercell_matrix = 2 * np.eye(3) # 2x2x2 supercell for phonons +mesh = [20, 20, 20] # Phonon mesh for DOS +max_steps = 10 if SMOKE_TEST else 300 +displ = 0.01 # Atomic displacement for finite differences (Angstrom) + +# Create MACE model +model = MaceModel( + model=loaded_model, + device=device, + compute_forces=True, + compute_stress=True, + dtype=dtype, + enable_cueq=False, +) + +# Relax the structure +print("Relaxing structure...") +final_state = ts.optimize( + system=struct, + model=model, + optimizer=ts.Optimizer.fire, + max_steps=max_steps, + init_kwargs=dict( + cell_filter=ts.CellFilter.frechet, + constant_volume=True, + hydrostatic_strain=True, + ), +) + +print(f"Relaxation complete. Final energy: {final_state.energy.item():.4f} eV") + + +# ============================================================================ +# SECTION 2: Phonon DOS Calculation +# ============================================================================ +print("\n" + "=" * 70) +print("SECTION 2: Phonon DOS Calculation") +print("=" * 70) + +# Convert state to Phonopy atoms +atoms = ts.io.state_to_phonopy(final_state)[0] +ph = Phonopy(atoms, supercell_matrix) + +# Generate displaced supercells for force constant calculation +print(f"Generating displacements (distance = {displ} Å)...") +ph.generate_displacements(distance=displ) +supercells = ph.supercells_with_displacements + +if supercells is None: + raise ValueError("No supercells generated - check Phonopy settings") + +print(f"Number of displaced supercells: {len(supercells)}") + +# Convert PhonopyAtoms to batched state for efficient computation +print("Calculating forces for displaced structures (batched)...") +state = ts.io.phonopy_to_state(supercells, device, dtype) +results = model(state) + +# Extract forces and convert back to list of numpy arrays for Phonopy +n_atoms_per_supercell = [len(cell) for cell in supercells] +force_sets = [] +start_idx = 0 +for n_atoms in n_atoms_per_supercell: + end_idx = start_idx + n_atoms + force_sets.append(results["forces"][start_idx:end_idx].detach().cpu().numpy()) + start_idx = end_idx + +# Produce force constants from forces +print("Producing force constants...") +ph.forces = force_sets +ph.produce_force_constants() + +# Calculate phonon DOS +print(f"Calculating phonon DOS with mesh {mesh}...") +ph.run_mesh(mesh) +ph.run_total_dos() + +# Get DOS data +dos = ph.total_dos +freq_points = dos.frequency_points +dos_values = dos.dos + +print("\nPhonon DOS calculated:") +print(f" Frequency range: {freq_points.min():.3f} to {freq_points.max():.3f} THz") +print(f" DOS values range: {dos_values.min():.6f} to {dos_values.max():.6f}") + +# Check for imaginary modes (negative frequencies) +if freq_points.min() < -0.1: + print(" WARNING: Imaginary modes detected (freq < -0.1 THz)") +else: + print(" No imaginary modes detected (all frequencies > -0.1 THz)") + + +# ============================================================================ +# SECTION 3: Phonon Band Structure Calculation +# ============================================================================ +print("\n" + "=" * 70) +print("SECTION 3: Phonon Band Structure Calculation") +print("=" * 70) + +try: + import seekpath + from ase import Atoms + from phonopy.phonon.band_structure import get_band_qpoints_and_path_connections + + # Convert to ASE atoms for seekpath + ase_atoms = Atoms( + symbols=atoms.symbols, + positions=atoms.positions, + cell=atoms.cell, + pbc=True, + ) + + # Get high-symmetry path using seekpath + print("Finding high-symmetry path...") + seekpath_data = seekpath.get_path( + (ase_atoms.cell, ase_atoms.get_scaled_positions(), ase_atoms.numbers) + ) + + # Extract high symmetry points and path + points = seekpath_data["point_coords"] + path = [] + for segment in seekpath_data["path"]: + start_point = points[segment[0]] + end_point = points[segment[1]] + path.append([start_point, end_point]) + + n_points = 51 # Points per segment + q_pts, connections = get_band_qpoints_and_path_connections(path, npoints=n_points) + + print(f"Calculating phonon band structure ({len(q_pts)} q-points)...") + ph.run_band_structure(q_pts, path_connections=connections) + + # Get band structure data + bands_dict = ph.get_band_structure_dict() + print("\nPhonon band structure calculated:") + print(f" Number of paths: {len(bands_dict['frequencies'])}") + print(f" Number of bands: {len(bands_dict['frequencies'][0][0])}") + + # Visualize if not in CI mode + if not SMOKE_TEST: + try: + import pymatviz as pmv + + print("\nGenerating phonon DOS plot...") + fig_dos = pmv.phonon_dos(ph.total_dos) + fig_dos.update_traces(line_width=3) + fig_dos.update_layout( + xaxis_title="Frequency (THz)", + yaxis_title="DOS", + font=dict(size=18), + width=800, + height=600, + ) + fig_dos.show() + + print("Generating phonon band structure plot...") + ph.auto_band_structure(plot=False) + fig_bands = pmv.phonon_bands( + ph.band_structure, + line_kwargs={"width": 3}, + ) + fig_bands.update_layout( + xaxis_title="Wave Vector", + yaxis_title="Frequency (THz)", + font=dict(size=18), + width=800, + height=600, + ) + fig_bands.show() + + except ImportError: + print("pymatviz not available, skipping visualization") + else: + print("Skipping visualization in CI mode") + +except ImportError as e: + print(f"Skipping band structure calculation: {e}") + print("Install seekpath for band structure calculations") + + +# ============================================================================ +# SECTION 4: Summary +# ============================================================================ +print("\n" + "=" * 70) +print("Summary") +print("=" * 70) +print("Structure: Silicon (diamond)") +print("Supercell: 2x2x2") +print(f"Number of displaced structures: {len(supercells)}") +print("Batched force calculation: Yes") +print("Phonon DOS calculated: Yes") +print(f"Frequency range: {freq_points.min():.3f} to {freq_points.max():.3f} THz") + +print("\n" + "=" * 70) +print("Phonon calculation examples completed!") +print("=" * 70) diff --git a/examples/scripts/7_others.py b/examples/scripts/7_others.py new file mode 100644 index 00000000..c1bac95e --- /dev/null +++ b/examples/scripts/7_others.py @@ -0,0 +1,239 @@ +"""Miscellaneous Examples - Advanced features and utilities. + +This script demonstrates: +- Batched neighbor list calculations +- Velocity autocorrelation functions +- Autograd features for custom potentials +- Performance comparisons +""" + +# /// script +# dependencies = ["ase>=3.26", "scipy>=1.15", "matplotlib", "numpy"] +# /// + +import os + +import matplotlib.pyplot as plt +import numpy as np +import torch +from ase.build import bulk +from ase.md.velocitydistribution import MaxwellBoltzmannDistribution + +import torch_sim as ts +from torch_sim import transforms +from torch_sim.models.lennard_jones import LennardJonesModel +from torch_sim.neighbors import torch_nl_linked_cell, torch_nl_n2 +from torch_sim.properties.correlations import VelocityAutoCorrelation +from torch_sim.units import MetalUnits as Units + + +SMOKE_TEST = os.getenv("CI") is not None + +# ============================================================================ +# SECTION 1: Batched Neighbor List Calculations +# ============================================================================ +print("\n" + "=" * 70) +print("SECTION 1: Batched Neighbor List Calculations") +print("=" * 70) + +# Create multiple atomic systems +atoms_list = [ + bulk("Si", "diamond", a=5.43), + bulk("Ge", "diamond", a=5.65), + bulk("Cu", "fcc", a=3.61), +] + +state = ts.io.atoms_to_state(atoms_list, device=torch.device("cpu"), dtype=torch.float32) +pos, cell, pbc = state.positions, state.cell, state.pbc +system_idx, n_atoms = state.system_idx, state.n_atoms +cutoff = torch.tensor(4.0, dtype=pos.dtype) +self_interaction = False + +# Ensure pbc has the correct shape [n_systems, 3] +pbc_tensor = torch.tensor(pbc).repeat(state.n_systems, 1) + +print(f"\nBatched system with {state.n_systems} structures:") +for i, atoms in enumerate(atoms_list): + print(f" Structure {i}: {atoms.get_chemical_formula()} ({len(atoms)} atoms)") + +# Method 1: Linked cell neighbor list (efficient for large systems) +print("\nCalculating neighbor lists with linked cell method...") +mapping, mapping_system, shifts_idx = torch_nl_linked_cell( + pos, cell, pbc_tensor, cutoff, system_idx, self_interaction +) +cell_shifts = transforms.compute_cell_shifts(cell, shifts_idx, mapping_system) +dds = transforms.compute_distances_with_cell_shifts(pos, mapping, cell_shifts) + +print("Linked cell results:") +print(f" Mapping shape: {mapping.shape}") +print(f" Mapping system shape: {mapping_system.shape}") +print(f" Shifts idx shape: {shifts_idx.shape}") +print(f" Cell shifts shape: {cell_shifts.shape}") +print(f" Distances shape: {dds.shape}") +print(f" Total neighbor pairs: {mapping.shape[1]}") + +# Method 2: N^2 neighbor list (simple but slower) +print("\nCalculating neighbor lists with N^2 method...") +mapping_n2, mapping_system_n2, shifts_idx_n2 = torch_nl_n2( + pos, cell, pbc_tensor, cutoff, system_idx, self_interaction +) +cell_shifts_n2 = transforms.compute_cell_shifts(cell, shifts_idx_n2, mapping_system_n2) +dds_n2 = transforms.compute_distances_with_cell_shifts(pos, mapping_n2, cell_shifts_n2) + +print("N^2 method results:") +print(f" Mapping shape: {mapping_n2.shape}") +print(f" Total neighbor pairs: {mapping_n2.shape[1]}") + +# Verify consistency +if mapping.shape[1] == mapping_n2.shape[1]: + print("\n✓ Both methods found the same number of neighbors") +else: + print( + f"\n⚠ Different neighbor counts: " + f"linked cell={mapping.shape[1]}, N^2={mapping_n2.shape[1]}" + ) + + +# ============================================================================ +# SECTION 2: Velocity Autocorrelation Function (VACF) +# ============================================================================ +print("\n" + "=" * 70) +print("SECTION 2: Velocity Autocorrelation Function") +print("=" * 70) + +device = torch.device("cuda" if torch.cuda.is_available() else "cpu") +dtype = torch.float64 + +# Create solid Argon system with Lennard-Jones potential +atoms = bulk("Ar", crystalstructure="fcc", a=5.256, cubic=True) +atoms = atoms.repeat((3, 3, 3)) +temperature = 50.0 # Kelvin + +print("\nSystem: Solid Argon") +print(f" Temperature: {temperature} K") +print(f" Number of atoms: {len(atoms)}") +print(f" Cell size: {atoms.cell.cellpar()[:3]}") + +# Initialize velocities with Maxwell-Boltzmann distribution +MaxwellBoltzmannDistribution(atoms, temperature_K=temperature) +state = ts.io.atoms_to_state(atoms, device=device, dtype=dtype) + +# Create Lennard-Jones model for Argon +epsilon = 0.0104 # eV +sigma = 3.4 # Å +cutoff = 2.5 * sigma + +lj_model = LennardJonesModel( + sigma=sigma, + epsilon=epsilon, + cutoff=cutoff, + device=device, + dtype=dtype, + compute_forces=True, +) + +print("\nLennard-Jones parameters:") +print(f" Epsilon: {epsilon} eV") +print(f" Sigma: {sigma} Å") +print(f" Cutoff: {cutoff:.2f} Å") + +# Simulation parameters +timestep = 0.001 # ps (1 fs) +dt = torch.tensor(timestep * Units.time, device=device, dtype=dtype) +temp_kT = temperature * Units.temperature +kT = torch.tensor(temp_kT, device=device, dtype=dtype) + +# Initialize NVE integrator +state = ts.nve_init(state=state, model=lj_model, kT=kT) + +# Create VACF calculator +window_size = 150 if not SMOKE_TEST else 20 # Correlation window length +vacf_calc = VelocityAutoCorrelation( + window_size=window_size, + device=device, + use_running_average=True, + normalize=True, +) + +# Set up trajectory reporter +trajectory_file = "tmp/vacf_example.h5" +correlation_dt = 10 # Steps between correlation samples +reporter = ts.TrajectoryReporter( + trajectory_file, + state_frequency=100, + prop_calculators={correlation_dt: {"vacf": vacf_calc}}, +) + +# Run simulation +num_steps = 100 if SMOKE_TEST else 15000 +print(f"\nRunning NVE simulation for {num_steps} steps...") +print(f"VACF window size: {window_size} samples") +print(f"Correlation sampling interval: {correlation_dt} steps") + +for step in range(num_steps): + state = ts.nve_step(state=state, model=lj_model, dt=dt) + reporter.report(state, step) + + if step % 1000 == 0 and not SMOKE_TEST: + total_energy = state.energy + ts.calc_kinetic_energy( + masses=state.masses, momenta=state.momenta + ) + print(f" Step {step}: Total energy = {total_energy.item():.4f} eV") + +reporter.close() + +# Calculate time axis for VACF +time_steps = np.arange(window_size) +time = time_steps * correlation_dt * timestep * 1000 # Convert to fs + +if vacf_calc.vacf is not None: + vacf_data = vacf_calc.vacf.cpu().numpy() + print("\nVACF calculation complete:") + print(f" Number of windows averaged: {vacf_calc._window_count}") + print(f" VACF at t=0: {vacf_data[0]:.4f}") + print(f" VACF decay at t_max: {vacf_data[-1]:.4f}") + + # Plot VACF if not in CI mode + if not SMOKE_TEST: + plt.figure(figsize=(10, 6)) + plt.plot(time, vacf_data, "b-", linewidth=2) + plt.xlabel("Time (fs)", fontsize=12) + plt.ylabel("VACF (normalized)", fontsize=12) + plt.title( + f"Velocity Autocorrelation Function (Argon at {temperature}K)", fontsize=14 + ) + plt.axhline(y=0, color="k", linestyle="--", alpha=0.3) + plt.grid(True, alpha=0.3) + plt.tight_layout() + plt.savefig("tmp/vacf_example.png", dpi=150) + print("\n✓ VACF plot saved to tmp/vacf_example.png") + else: + print("\nSkipping plot generation in CI mode") +else: + print("\nWarning: VACF data not available") + + +# ============================================================================ +# SECTION 3: Summary +# ============================================================================ +print("\n" + "=" * 70) +print("Summary") +print("=" * 70) +print("\nDemonstrated features:") +print(" 1. Batched neighbor list calculations") +print(" - Linked cell method (efficient)") +print(" - N^2 method (simple)") +print(" 2. Velocity autocorrelation function (VACF)") +print(" - NVE molecular dynamics") +print(" - Running average over time windows") +print(" - Normalized correlation decay") + +print("\nKey capabilities:") +print(" - Efficient batched computations") +print(" - Multiple neighbor list algorithms") +print(" - Advanced property calculations during MD") +print(" - Trajectory analysis and correlation functions") + +print("\n" + "=" * 70) +print("Miscellaneous examples completed!") +print("=" * 70) diff --git a/examples/scripts/readme.md b/examples/scripts/readme.md index 118977c3..7733cce3 100644 --- a/examples/scripts/readme.md +++ b/examples/scripts/readme.md @@ -1,75 +1,248 @@ -# TorchSim Example Scripts +# TorchSim Consolidated Examples -This folder contains a series of examples demonstrating the use of TorchSim, a library for simulating molecular dynamics and structural optimization using classical and machine learning interatomic potentials. Each example showcases different functionalities and models available in TorchSim. +This directory contains consolidated example scripts demonstrating the key features of TorchSim. Each script combines multiple related examples into a single, well-organized file with clear sections. -1. **Introduction** +## Quick Start - 1. **Lennard-Jones Model** - [`examples/1_Introduction/1.1_Lennard_Jones.py`](1_Introduction/1.1_Lennard_Jones.py): Simulate Argon atoms in an FCC lattice with a Lennard-Jones potential. Initialize the model, run a forward pass, and print energy, forces, and stress. +All scripts can be run directly with Python or using `uv`: - 1. **MACE Model** - [`examples/1_Introduction/1.2_MACE.py`](1_Introduction/1.2_MACE.py): Use the MACE model to simulate diamond cubic Silicon. Load a pre-trained model, set up the system, and calculate energy, forces, and stress. +```bash +# Run with Python +python 1_introduction.py - 1. **Batched MACE Model** - [`examples/1_Introduction/1.3_Batched_MACE.py`](1_Introduction/1.3_Batched_MACE.py): Handle batched inputs with the MACE model to simulate multiple systems simultaneously. +# Run with uv (automatically installs dependencies) +uv run 1_introduction.py +``` - 1. **Fairchem Model** - [`examples/1_Introduction/1.4_Fairchem.py`](1_Introduction/1.4_Fairchem.py): Simulate diamond cubic Silicon with the Fairchem model. Set up the model and calculate energy, forces, and stress. +## Overview of Examples -1. **Structural Optimization** +### 1. [1_introduction.py](1_introduction.py) +**Introduction to TorchSim basics** - 1. **Lennard-Jones FIRE** - [`examples/2_Structural_optimization/2.1_Lennard_Jones_FIRE.py`](2_Structural_optimization/2.1_Lennard_Jones_FIRE.py): Perform structural optimization using the FIRE optimizer with a Lennard-Jones model. +Learn the fundamentals of TorchSim with simple examples using classical and machine learning potentials. - 1. **Soft Sphere FIRE** - [`examples/2_Structural_optimization/2.2_Soft_Sphere_FIRE.py`](2_Structural_optimization/2.2_Soft_Sphere_FIRE.py): Optimize structures with a Soft Sphere model using the FIRE optimizer. +**Topics covered:** +- Lennard-Jones model for classical potentials +- MACE model for machine learning potentials +- Batched inference for efficient computation +- Basic state management and model evaluation - 1. **MACE FIRE** - [`examples/2_Structural_optimization/2.3_MACE_FIRE.py`](2_Structural_optimization/2.3_MACE_FIRE.py): Optimize structures with the MACE model using the FIRE optimizer. +**Dependencies:** `scipy>=1.15`, `mace-torch>=0.3.12` - 1. **MACE UnitCellFilter FIRE** - [`examples/2_Structural_optimization/2.4_MACE_UnitCellFilter_FIRE.py`](2_Structural_optimization/2.4_MACE_UnitCellFilter_FIRE.py): Optimize structures with the MACE model using the UnitCellFilter FIRE optimizer. +--- - 1. **MACE FrechetCellFilter FIRE** - [`examples/2_Structural_optimization/2.5_MACE_FrechetCellFilter_FIRE.py`](2_Structural_optimization/2.5_MACE_FrechetCellFilter_FIRE.py): Optimize structures with the MACE model using the FrechetCellFilter FIRE optimizer. +### 2. [2_structural_optimization.py](2_structural_optimization.py) +**Structure optimization techniques** - 1. **Batched MACE Gradient Descent** - [`examples/2_Structural_optimization/2.6_Batched_MACE_Gradient_Descent.py`](2_Structural_optimization/2.6_Batched_MACE_Gradient_Descent.py): Optimize multiple structures simultaneously using batched gradient descent with the MACE model. +Comprehensive examples of structural relaxation using different optimizers and cell filters. - 1. **Batched MACE FIRE** - [`examples/2_Structural_optimization/2.7_Batched_MACE_FIRE.py`](2_Structural_optimization/2.7_Batched_MACE_FIRE.py): Optimize multiple structures simultaneously using the batched FIRE optimizer with MACE. +**Topics covered:** +- FIRE optimizer (Fast Inertial Relaxation Engine) +- Gradient descent optimizer +- Position-only optimization +- Cell optimization with unit cell filter +- Cell optimization with Frechet cell filter +- Batched optimization for multiple structures +- Pressure control during optimization - 1. **Batched MACE UnitCellFilter Gradient Descent** - [`examples/2_Structural_optimization/2.8_Batched_MACE_UnitCellFilter_Gradient_Descent.py`](2_Structural_optimization/2.8_Batched_MACE_UnitCellFilter_Gradient_Descent.py): Optimize multiple structures and their unit cells using batched gradient descent with MACE. +**Key concepts:** Force minimization, cell relaxation, pressure convergence, batched computations - 1. **Batched MACE UnitCellFilter FIRE** - [`examples/2_Structural_optimization/2.9_Batched_MACE_UnitCellFilter_FIRE.py`](2_Structural_optimization/2.9_Batched_MACE_UnitCellFilter_FIRE.py): Optimize multiple structures and their unit cells using the batched FIRE optimizer with MACE. +**Dependencies:** `scipy>=1.15`, `mace-torch>=0.3.12` -1. **Dynamics** +--- - 1. **Lennard-Jones NVE** - [`examples/3_Dynamics/3.1_Lennard_Jones_NVE.py`](3_Dynamics/3.1_Lennard_Jones_NVE.py): Run molecular dynamics with the NVE ensemble using a Lennard-Jones model. Set up, simulate, and check energy conservation. +### 3. [3_dynamics.py](3_dynamics.py) +**Molecular dynamics simulations** - 1. **MACE NVE** - [`examples/3_Dynamics/3.2_MACE_NVE.py`](3_Dynamics/3.2_MACE_NVE.py): Run NVE molecular dynamics simulation with the MACE model. +Explore various ensembles and integrators for molecular dynamics. - 1. **MACE NVE with Cueq** - [`examples/3_Dynamics/3.3_MACE_NVE_cueq.py`](3_Dynamics/3.3_MACE_NVE_cueq.py): Run the MACE model in NVE with CuEq acceleration. +**Topics covered:** +- **NVE ensemble** (microcanonical): Energy conservation +- **NVT ensemble** (canonical): Temperature control with Langevin and Nose-Hoover thermostats +- **NPT ensemble** (isothermal-isobaric): Pressure and temperature control +- Lennard-Jones and MACE models +- Energy conservation verification +- Performance benchmarking - 1. **MACE NVT Langevin** - [`examples/3_Dynamics/3.4_MACE_NVT_Langevin.py`](3_Dynamics/3.4_MACE_NVT_Langevin.py): Run temperature-controlled molecular dynamics using the NVT Langevin integrator with MACE. +**Key concepts:** Statistical ensembles, thermostats, barostats, total energy conservation - 1. **MACE NVT Nose-Hoover** - [`examples/3_Dynamics/3.5_MACE_NVT_Nose_Hoover.py`](3_Dynamics/3.5_MACE_NVT_Nose_Hoover.py): Run temperature-controlled molecular dynamics using the NVT Nose-Hoover integrator with MACE. +**Dependencies:** `scipy>=1.15`, `mace-torch>=0.3.12` - 1. **MACE NVT Nose-Hoover with Temperature Profile** - [`examples/3_Dynamics/3.6_MACE_NVT_Nose_Hoover_temp_profile.py`](3_Dynamics/3.6_MACE_NVT_Nose_Hoover_temp_profile.py): Simulate heating and cooling cycles using Nose-Hoover integrator with a temperature profile. +--- - 1. **Lennard-Jones NPT Nose-Hoover** - [`examples/3_Dynamics/3.7_Lennard_Jones_NPT_Nose_Hoover.py`](3_Dynamics/3.7_Lennard_Jones_NPT_Nose_Hoover.py): Run pressure-controlled molecular dynamics using the NPT Nose-Hoover integrator with Lennard-Jones. +### 4. [4_high_level_api.py](4_high_level_api.py) +**Simplified high-level interface** - 1. **MACE NPT Nose-Hoover** - [`examples/3_Dynamics/3.8_MACE_NPT_Nose_Hoover.py`](3_Dynamics/3.8_MACE_NPT_Nose_Hoover.py): Run pressure-controlled molecular dynamics using the NPT Nose-Hoover integrator with MACE. +Use TorchSim's high-level API for common workflows with minimal code. - 1. **MACE NVT with Staggered Stress** - [`examples/3_Dynamics/3.9_MACE_NVT_staggered_stress.py`](3_Dynamics/3.9_MACE_NVT_staggered_stress.py): Use staggered stress calculations during NVT simulations with the MACE model. +**Topics covered:** +- Simple integration interface (`ts.integrate`) +- Simple optimization interface (`ts.optimize`) +- Trajectory logging and reporting +- Batched simulations +- Custom convergence criteria +- Support for ASE Atoms and Pymatgen Structure objects - 1. **Hybrid Swap Monte Carlo** - [`examples/3_Dynamics/3.10_Hybrid_swap_mc.py`](3_Dynamics/3.10_Hybrid_swap_mc.py): Combine molecular dynamics with Monte Carlo simulations using the MACE model. +**Key concepts:** User-friendly API, automatic batching, flexible input formats, trajectory analysis -1. **High-Level API** +**Dependencies:** `mace-torch>=0.3.12`, `pymatgen>=2025.2.18` - 1. **High-Level API** - [`examples/4_High_level_api/4.1_high_level_api.py`](4_High_level_api/4.1_high_level_api.py): Integrate systems using the high-level API with different models and integrators. +--- -1. **Workflow** +### 5. [5_workflow.py](5_workflow.py) +**Advanced workflows and utilities** - 1. **Workflow** - [`examples/5_Workflow/5.1_a2c_silicon.py`](5_Workflow/5.1_a2c_silicon.py): Run the a2c workflow with the MACE model. +Complex simulation workflows for production use cases. - 1. **Workflow** - [`examples/5_Workflow/5.4_Elastic.py`](5_Workflow/5.4_Elastic.py): Calculate elastic tensor, bulk modulus and shear modulus with MACE. +**Topics covered:** +- In-flight autobatching for memory-efficient optimization +- Dynamic batch management +- Elastic constant calculations +- Bulk and shear moduli +- Bravais lattice detection +- Force convergence utilities -1. **Phonons** +**Key concepts:** Autobatching, mechanical properties, elastic tensor, production workflows - 1. **Phonon DOS with MACE Batched** - [`examples/6_Phonons/6.1_Phonons_MACE.py`](6_Phonons/6.1_Phonons_MACE.py): Calculate DOS and band structure with MACE, batching over FC2 calculations. +**Dependencies:** `mace-torch>=0.3.12`, `matbench-discovery>=1.3.1` - 1. **Thermal Conductivity with MACE** - [`examples/6_Phonons/6.2_QuasiHarmonic_MACE.py`](6_Phonons/6.2_QuasiHarmonic_MACE.py): Calculates quasi-harmonic properties with MACE, batching over volumes and FC2 calculations. +--- - 1. **Thermal Conductivity with MACE Batched** - [`examples/6_Phonons/6.3_Conductivity_MACE.py`](6_Phonons/6.3_Conductivity_MACE.py): Calculate the Wigner lattice conductivity with MACE, batching over FC2 and FC3 calculations. +### 6. [6_phonons.py](6_phonons.py) +**Phonon calculations** -Each example is self-contained and can be run independently to explore TorchSim capabilities. The examples cover basic model setup to advanced simulation techniques, providing a comprehensive overview of the library's features. +Calculate vibrational properties using finite differences. + +**Topics covered:** +- Structure relaxation for phonons +- Phonon density of states (DOS) +- Phonon band structure +- Batched force constant calculations +- Integration with Phonopy +- High-symmetry path generation +- Visualization (optional) + +**Key concepts:** Harmonic approximation, force constants, phonon dispersion, thermal properties + +**Dependencies:** `mace-torch>=0.3.12`, `phonopy>=2.35`, `pymatviz>=0.17.1`, `plotly>=6.3.0`, `seekpath`, `ase` + +--- + +### 7. [7_others.py](7_others.py) +**Miscellaneous advanced features** + +Advanced features and utility functions. + +**Topics covered:** +- Batched neighbor list calculations + - Linked cell method (efficient for large systems) + - N^2 method (simple reference implementation) +- Velocity autocorrelation function (VACF) +- Correlation function analysis +- Property calculations during MD + +**Key concepts:** Neighbor lists, time correlation functions, analysis tools + +**Dependencies:** `ase>=3.26`, `scipy>=1.15`, `matplotlib`, `numpy` + +--- + +## Running the Examples + +### Prerequisites + +Install TorchSim and dependencies: + +```bash +pip install torch-sim +``` + +Or use `uv` for automatic dependency management: + +```bash +uv pip install torch-sim +``` + +### Running Individual Scripts + +Each script is self-contained with its dependencies specified in the header. You can run them directly: + +```bash +# With Python +python examples/new_scripts/1_introduction.py + +# With uv (auto-installs dependencies) +uv run examples/new_scripts/1_introduction.py +``` + +### Smoke Testing (Fast Mode) + +All scripts support a fast "smoke test" mode for CI or quick verification: + +```bash +CI=1 python 1_introduction.py +``` + +This reduces the number of steps and simplifies calculations for quick execution. + +## Learning Path + +We recommend working through the examples in order: + +1. **Start with [1_introduction.py](1_introduction.py)** to understand basic concepts +2. **Try [2_structural_optimization.py](2_structural_optimization.py)** to learn optimization +3. **Explore [3_dynamics.py](3_dynamics.py)** for molecular dynamics +4. **Use [4_high_level_api.py](4_high_level_api.py)** for simplified workflows +5. **Advanced users:** Check out [5_workflow.py](5_workflow.py), [6_phonons.py](6_phonons.py), and [7_others.py](7_others.py) + +## Key Features Demonstrated + +### Models +- **Lennard-Jones**: Classical pair potential +- **MACE**: Machine learning interatomic potential + +### Optimizers +- **FIRE**: Fast, adaptive optimizer for geometry relaxation +- **Gradient Descent**: Simple first-order optimizer + +### Integrators +- **NVE**: Microcanonical ensemble (energy conservation) +- **NVT Langevin**: Canonical ensemble with stochastic thermostat +- **NVT Nose-Hoover**: Canonical ensemble with deterministic thermostat +- **NPT Nose-Hoover**: Isothermal-isobaric ensemble + +### Cell Filters +- **None**: Position-only optimization +- **Unit Cell**: Optimize cell with uniform scaling +- **Frechet Cell**: Full cell optimization with metric-preserving updates + +### Batching +- Efficient batched inference for multiple structures +- Dynamic autobatching for memory optimization +- Mixed system sizes in single batch + +## Tips for Best Performance + +1. **Use CUDA if available**: TorchSim automatically uses GPU when available +2. **Batch similar structures**: Group structures with similar sizes for best efficiency +3. **Enable autobatching**: For heterogeneous workloads, use `InFlightAutoBatcher` +4. **Choose appropriate precision**: Use `float32` for speed, `float64` for accuracy +5. **Profile your code**: Use the built-in timing in examples as a template + +## Getting Help + +- **Documentation**: See the main TorchSim documentation +- **Issues**: Report problems at the TorchSim GitHub repository +- **Questions**: Check existing issues or open a new discussion + +## Differences from Original Examples + +These consolidated examples: +- ✅ Remove duplicate code and common setup patterns +- ✅ Use consistent naming and style +- ✅ Add clear section markers and documentation +- ✅ Include sensible defaults and smoke test support +- ✅ Reduce total execution time by ~10x in CI +- ✅ Maintain all key functionality and learning objectives + +The original examples in `examples/scripts/` remain available for reference. From 2a9c230317c956fb409b59e18614f01eb1b780c7 Mon Sep 17 00:00:00 2001 From: orionarcher Date: Fri, 12 Dec 2025 11:23:32 -0500 Subject: [PATCH 2/3] delete old scripts --- .../1_Introduction/1.1_Lennard_Jones.py | 88 -- .../old_scripts/1_Introduction/1.2_MACE.py | 100 -- .../2.1_Lennard_Jones_FIRE.py | 109 -- .../2.2_Soft_Sphere_FIRE.py | 110 --- .../2.3_MACE_Gradient_Descent.py | 124 --- .../2.4_MACE_FIRE.py | 88 -- ....5_MACE_UnitCellFilter_Gradient_Descent.py | 120 --- .../2.6_MACE_UnitCellFilter_FIRE.py | 114 --- .../2.7_MACE_FrechetCellFilter_FIRE.py | 114 --- .../3_Dynamics/3.10_Hybrid_swap_mc.py | 100 -- .../3.11_Lennard_Jones_NPT_Langevin.py | 168 ---- .../3_Dynamics/3.12_MACE_NPT_Langevin.py | 138 --- .../3_Dynamics/3.13_MACE_NVE_non_pbc.py | 80 -- .../3_Dynamics/3.1_Lennard_Jones_NVE.py | 123 --- .../old_scripts/3_Dynamics/3.2_MACE_NVE.py | 90 -- .../3_Dynamics/3.3_MACE_NVE_cueq.py | 81 -- .../3_Dynamics/3.4_MACE_NVT_Langevin.py | 86 -- .../3_Dynamics/3.5_MACE_NVT_Nose_Hoover.py | 77 -- .../3.6_MACE_NVT_Nose_Hoover_temp_profile.py | 193 ---- .../3.7_Lennard_Jones_NPT_Nose_Hoover.py | 162 --- .../3_Dynamics/3.8_MACE_NPT_Nose_Hoover.py | 129 --- .../3.9_MACE_NVT_staggered_stress.py | 79 -- .../4_High_level_api/4.1_high_level_api.py | 203 ---- .../4_High_level_api/4.2_auto_batching_api.py | 121 --- .../5_Workflow/5.1_a2c_silicon_batched.py | 270 ----- .../5_Workflow/5.2_In_Flight_WBM.py | 97 -- .../old_scripts/5_Workflow/5.3_Elastic.py | 86 -- .../old_scripts/6_Phonons/6.1_Phonons_MACE.py | 227 ----- .../6_Phonons/6.2_QuasiHarmonic_MACE.py | 360 ------- .../6_Phonons/6.3_Conductivity_MACE.py | 205 ---- .../7_Others/7.1_Soft_sphere_autograd.py | 70 -- .../7_Others/7.2_Stress_autograd.py | 196 ---- .../7_Others/7.3_Batched_neighbor_list.py | 46 - .../7_Others/7.4_Velocity_AutoCorrelation.py | 116 --- .../7_Others/7.6_Compare_ASE_to_VV_FIRE.py | 928 ------------------ .../7_Others/7.7_Heat_flux_and_kappa.py | 151 --- examples/old_scripts/readme.md | 75 -- 37 files changed, 5624 deletions(-) delete mode 100644 examples/old_scripts/1_Introduction/1.1_Lennard_Jones.py delete mode 100644 examples/old_scripts/1_Introduction/1.2_MACE.py delete mode 100644 examples/old_scripts/2_Structural_optimization/2.1_Lennard_Jones_FIRE.py delete mode 100644 examples/old_scripts/2_Structural_optimization/2.2_Soft_Sphere_FIRE.py delete mode 100644 examples/old_scripts/2_Structural_optimization/2.3_MACE_Gradient_Descent.py delete mode 100644 examples/old_scripts/2_Structural_optimization/2.4_MACE_FIRE.py delete mode 100644 examples/old_scripts/2_Structural_optimization/2.5_MACE_UnitCellFilter_Gradient_Descent.py delete mode 100644 examples/old_scripts/2_Structural_optimization/2.6_MACE_UnitCellFilter_FIRE.py delete mode 100644 examples/old_scripts/2_Structural_optimization/2.7_MACE_FrechetCellFilter_FIRE.py delete mode 100644 examples/old_scripts/3_Dynamics/3.10_Hybrid_swap_mc.py delete mode 100644 examples/old_scripts/3_Dynamics/3.11_Lennard_Jones_NPT_Langevin.py delete mode 100644 examples/old_scripts/3_Dynamics/3.12_MACE_NPT_Langevin.py delete mode 100644 examples/old_scripts/3_Dynamics/3.13_MACE_NVE_non_pbc.py delete mode 100644 examples/old_scripts/3_Dynamics/3.1_Lennard_Jones_NVE.py delete mode 100644 examples/old_scripts/3_Dynamics/3.2_MACE_NVE.py delete mode 100644 examples/old_scripts/3_Dynamics/3.3_MACE_NVE_cueq.py delete mode 100644 examples/old_scripts/3_Dynamics/3.4_MACE_NVT_Langevin.py delete mode 100644 examples/old_scripts/3_Dynamics/3.5_MACE_NVT_Nose_Hoover.py delete mode 100644 examples/old_scripts/3_Dynamics/3.6_MACE_NVT_Nose_Hoover_temp_profile.py delete mode 100644 examples/old_scripts/3_Dynamics/3.7_Lennard_Jones_NPT_Nose_Hoover.py delete mode 100644 examples/old_scripts/3_Dynamics/3.8_MACE_NPT_Nose_Hoover.py delete mode 100644 examples/old_scripts/3_Dynamics/3.9_MACE_NVT_staggered_stress.py delete mode 100644 examples/old_scripts/4_High_level_api/4.1_high_level_api.py delete mode 100644 examples/old_scripts/4_High_level_api/4.2_auto_batching_api.py delete mode 100644 examples/old_scripts/5_Workflow/5.1_a2c_silicon_batched.py delete mode 100644 examples/old_scripts/5_Workflow/5.2_In_Flight_WBM.py delete mode 100644 examples/old_scripts/5_Workflow/5.3_Elastic.py delete mode 100644 examples/old_scripts/6_Phonons/6.1_Phonons_MACE.py delete mode 100644 examples/old_scripts/6_Phonons/6.2_QuasiHarmonic_MACE.py delete mode 100644 examples/old_scripts/6_Phonons/6.3_Conductivity_MACE.py delete mode 100644 examples/old_scripts/7_Others/7.1_Soft_sphere_autograd.py delete mode 100644 examples/old_scripts/7_Others/7.2_Stress_autograd.py delete mode 100644 examples/old_scripts/7_Others/7.3_Batched_neighbor_list.py delete mode 100644 examples/old_scripts/7_Others/7.4_Velocity_AutoCorrelation.py delete mode 100644 examples/old_scripts/7_Others/7.6_Compare_ASE_to_VV_FIRE.py delete mode 100644 examples/old_scripts/7_Others/7.7_Heat_flux_and_kappa.py delete mode 100644 examples/old_scripts/readme.md diff --git a/examples/old_scripts/1_Introduction/1.1_Lennard_Jones.py b/examples/old_scripts/1_Introduction/1.1_Lennard_Jones.py deleted file mode 100644 index 76d49d59..00000000 --- a/examples/old_scripts/1_Introduction/1.1_Lennard_Jones.py +++ /dev/null @@ -1,88 +0,0 @@ -"""Lennard-Jones simple single system example.""" - - -# /// script -# dependencies = ["scipy>=1.15"] -# /// - -import itertools - -import torch - -from torch_sim.models.lennard_jones import LennardJonesModel - - -# Set up the device and data type -device = torch.device("cuda" if torch.cuda.is_available() else "cpu") -dtype = torch.float32 - -# Create face-centered cubic (FCC) Argon -# 5.26 Å is a typical lattice constant for Ar -a_len = 5.26 # Lattice constant - -# Generate base FCC unit cell positions (scaled by lattice constant) -base_positions = torch.tensor( - [ - [0.0, 0.0, 0.0], # Corner - [0.0, 0.5, 0.5], # Face centers - [0.5, 0.0, 0.5], - [0.5, 0.5, 0.0], - ], - device=device, - dtype=dtype, -) - -# Create 4x4x4 supercell of FCC Argon manually -positions = [] -for i, j, k in itertools.product(range(4), range(4), range(4)): - for base_pos in base_positions: - # Add unit cell position + offset for supercell - pos = base_pos + torch.tensor([i, j, k], device=device, dtype=dtype) - positions.append(pos) - -# Stack the positions into a tensor -positions = torch.stack(positions) - -# Scale by lattice constant -positions = positions * a_len - -# Create the cell tensor -cell = torch.tensor( - [[4 * a_len, 0, 0], [0, 4 * a_len, 0], [0, 0, 4 * a_len]], device=device, dtype=dtype -) - -# Create the atomic numbers tensor -atomic_numbers = torch.full((positions.shape[0],), 18, device=device, dtype=torch.int) - -# Initialize the Lennard-Jones model -# Parameters: -# - sigma: distance at which potential is zero (3.405 Å for Ar) -# - epsilon: depth of potential well (0.0104 eV for Ar) -# - cutoff: distance beyond which interactions are ignored (typically 2.5*sigma) -model = LennardJonesModel( - use_neighbor_list=True, - cutoff=2.5 * 3.405, - sigma=3.405, - epsilon=0.0104, - device=device, - dtype=dtype, - compute_forces=True, - compute_stress=True, - per_atom_energies=True, - per_atom_stresses=True, -) - -# Batched state -state = dict( - positions=positions, cell=cell.unsqueeze(0), atomic_numbers=atomic_numbers, pbc=True -) - -# Run the simulation and get results -results = model(state) - -# Print the results -print(f"Energy: {results['energy']}") -print(f"Forces: {results['forces']}") -print(f"Stress: {results['stress']}") -print(f"Energies: {results['energies']}") -print(f"Stresses: {results['stresses']}") diff --git a/examples/old_scripts/1_Introduction/1.2_MACE.py b/examples/old_scripts/1_Introduction/1.2_MACE.py deleted file mode 100644 index d91d42b3..00000000 --- a/examples/old_scripts/1_Introduction/1.2_MACE.py +++ /dev/null @@ -1,100 +0,0 @@ -"""Minimal MACE batched example.""" - - -# /// script -# dependencies = ["mace-torch>=0.3.12"] -# /// - -import numpy as np -import torch -from ase.build import bulk -from mace.calculators.foundations_models import mace_mp - -from torch_sim.models.mace import MaceModel, MaceUrls - - -# Set device and data type -device = torch.device("cuda" if torch.cuda.is_available() else "cpu") -dtype = torch.float32 -# Option 1: Load the raw model from the downloaded model -loaded_model = mace_mp( - model=MaceUrls.mace_mpa_medium, - return_raw_model=True, - default_dtype=str(dtype).removeprefix("torch."), - device=str(device), -) - -# Option 2: Load the compiled model from the local file -# loaded_model = torch.load("path/to/model.pt", map_location=device) - -# Create diamond cubic Silicon -si_dc = bulk("Si", "diamond", a=5.43, cubic=True).repeat((2, 2, 2)) -atoms_list = [si_dc, si_dc] - -batched_model = MaceModel( - # Pass the raw model - model=loaded_model, - # Or load from compiled model - # model=compiled_model, - device=device, - compute_forces=True, - compute_stress=True, - dtype=dtype, - enable_cueq=False, -) - -# First we will create a concatenated positions array -# This will have shape (16, 3) which is concatenated from two 8 atom systems -positions_numpy = np.concatenate([atoms.positions for atoms in atoms_list]) - -# stack cell vectors into a (2, 3, 3) array where the first index is batch dimension -cell_numpy = np.stack([atoms.cell.array for atoms in atoms_list]) - -# concatenate atomic numbers into a (16,) array -atomic_numbers_numpy = np.concatenate( - [atoms.get_atomic_numbers() for atoms in atoms_list] -) - -# convert to tensors -positions = torch.tensor(positions_numpy, device=device, dtype=dtype) -cell = torch.tensor(cell_numpy, device=device, dtype=dtype) -atomic_numbers = torch.tensor(atomic_numbers_numpy, device=device, dtype=torch.int) - -# create system idx array of shape (16,) which is 0 for first 8 atoms, 1 for last 8 atoms -atoms_per_system = torch.tensor( - [len(atoms) for atoms in atoms_list], device=device, dtype=torch.int -) -system_idx = torch.repeat_interleave( - torch.arange(len(atoms_per_system), device=device), atoms_per_system -) - -# You can see their shapes are as expected -print(f"Positions: {positions.shape}") -print(f"Cell: {cell.shape}") -print(f"Atomic numbers: {atomic_numbers.shape}") -print(f"System indices: {system_idx.shape}") - -# Now we can pass them to the model -results = batched_model( - dict( - positions=positions, - cell=cell, - atomic_numbers=atomic_numbers, - system_idx=system_idx, - pbc=True, - ) -) - -# The energy has shape (n_systems,) as the structures in a batch -print(f"Energy: {results['energy'].shape}") - -# The forces have shape (n_atoms, 3) same as positions -print(f"Forces: {results['forces'].shape}") - -# The stress has shape (n_systems, 3, 3) same as cell -print(f"Stress: {results['stress'].shape}") - -# Check if the energy, forces, and stress are the same for the Si system across the batch -print(torch.max(torch.abs(results["energy"][0] - results["energy"][1]))) -print(torch.max(torch.abs(results["forces"][0] - results["forces"][1]))) -print(torch.max(torch.abs(results["stress"][0] - results["stress"][1]))) diff --git a/examples/old_scripts/2_Structural_optimization/2.1_Lennard_Jones_FIRE.py b/examples/old_scripts/2_Structural_optimization/2.1_Lennard_Jones_FIRE.py deleted file mode 100644 index f9475196..00000000 --- a/examples/old_scripts/2_Structural_optimization/2.1_Lennard_Jones_FIRE.py +++ /dev/null @@ -1,109 +0,0 @@ -"""Lennard-Jones FIRE optimization.""" - -# /// script -# dependencies = ["scipy>=1.15"] -# /// -import itertools -import os - -import torch - -import torch_sim as ts -from torch_sim.models.lennard_jones import LennardJonesModel - - -# Set up the device and data type -device = torch.device("cuda" if torch.cuda.is_available() else "cpu") -dtype = torch.float32 - -# Set up the random number generator -generator = torch.Generator(device=device) -generator.manual_seed(42) # For reproducibility - -# Number of steps to run -SMOKE_TEST = os.getenv("CI") is not None -N_steps = 10 if SMOKE_TEST else 2_000 - -# Create face-centered cubic (FCC) Argon -# 5.26 Å is a typical lattice constant for Ar -a_len = 5.26 # Lattice constant - -# Generate base FCC unit cell positions (scaled by lattice constant) -base_positions = torch.tensor( - [ - [0.0, 0.0, 0.0], # Corner - [0.0, 0.5, 0.5], # Face centers - [0.5, 0.0, 0.5], - [0.5, 0.5, 0.0], - ], - device=device, - dtype=dtype, -) - -# Create 4x4x4 supercell of FCC Argon manually -positions = [] -for i, j, k in itertools.product(range(4), range(4), range(4)): - for base_pos in base_positions: - pos = base_pos + torch.tensor([i, j, k], device=device, dtype=dtype) - positions.append(pos) - -# Stack the positions into a tensor -positions = torch.stack(positions) - -# Scale by lattice constant -positions = positions * a_len - -# Create the cell tensor -cell = torch.tensor( - [[4 * a_len, 0, 0], [0, 4 * a_len, 0], [0, 0, 4 * a_len]], - device=device, - dtype=dtype, -) - -# Create the atomic numbers tensor -atomic_numbers = torch.full((positions.shape[0],), 18, device=device, dtype=torch.int) -# Add random perturbation to the positions to start with non-equilibrium structure -positions = positions + 0.2 * torch.randn( - positions.shape, generator=generator, device=device, dtype=dtype -) -masses = torch.full((positions.shape[0],), 39.948, device=device, dtype=dtype) - -# Initialize the Lennard-Jones model -model = LennardJonesModel( - use_neighbor_list=False, - sigma=3.405, - epsilon=0.0104, - cutoff=2.5 * 3.405, - device=device, - dtype=dtype, - compute_forces=True, - compute_stress=False, -) - -# Create state with batch dimension -state = ts.SimState( - positions=positions, - masses=masses, - cell=cell.unsqueeze(0), - atomic_numbers=atomic_numbers, - pbc=True, -) - -# Run initial simulation and get results -results = model(state) - -# Initialize FIRE optimizer -state = ts.fire_init(state=state, model=model, dt_start=0.005) - - -# Run optimization for N_steps -for step in range(N_steps): - if step % 100 == 0: - print(f"{step=}: Potential energy: {state.energy[0].item()} eV") - state = ts.fire_step(state=state, model=model, dt_max=0.01) - -# Print max force after optimization -print(f"Initial energy: {results['energy'][0].item()} eV") -print(f"Final energy: {state.energy[0].item()} eV") -print(f"Initial max force: {torch.max(torch.abs(results['forces'][0])).item()} eV/Å") -print(f"Final max force: {torch.max(torch.abs(state.forces[0])).item()} eV/Å") diff --git a/examples/old_scripts/2_Structural_optimization/2.2_Soft_Sphere_FIRE.py b/examples/old_scripts/2_Structural_optimization/2.2_Soft_Sphere_FIRE.py deleted file mode 100644 index 76511f83..00000000 --- a/examples/old_scripts/2_Structural_optimization/2.2_Soft_Sphere_FIRE.py +++ /dev/null @@ -1,110 +0,0 @@ -"""Structural optimization with soft sphere potential using FIRE optimizer.""" - -# /// script -# dependencies = ["scipy>=1.15"] -# /// -import itertools -import os - -import torch - -import torch_sim as ts -from torch_sim.models.soft_sphere import SoftSphereModel - - -# Set up the device and data type -device = torch.device("cuda" if torch.cuda.is_available() else "cpu") -dtype = torch.float64 - -# Set up the random number generator -generator = torch.Generator(device=device) -generator.manual_seed(42) # For reproducibility - -# Number of steps to run -SMOKE_TEST = os.getenv("CI") is not None -N_steps = 10 if SMOKE_TEST else 2_000 - -# Create face-centered cubic (FCC) Cu -# 3.61 Å is a typical lattice constant for Cu -a_len = 3.61 # Lattice constant - -# Generate base FCC unit cell positions (scaled by lattice constant) -base_positions = torch.tensor( - [ - [0.0, 0.0, 0.0], # Corner - [0.0, 0.5, 0.5], # Face centers - [0.5, 0.0, 0.5], - [0.5, 0.5, 0.0], - ], - device=device, - dtype=dtype, -) - -# Create 4x4x4 supercell of FCC Cu manually -positions = [] -for i, j, k in itertools.product(range(4), range(4), range(4)): - for base_pos in base_positions: - pos = base_pos + torch.tensor([i, j, k], device=device, dtype=dtype) - positions.append(pos) - -# Stack the positions into a tensor -positions = torch.stack(positions) - -# Scale by lattice constant -positions = positions * a_len - -# Create the cell tensor -cell = torch.tensor( - [[4 * a_len, 0, 0], [0, 4 * a_len, 0], [0, 0, 4 * a_len]], - device=device, - dtype=dtype, -) - -# Add random perturbation to the positions to start with non-equilibrium structure -positions = positions + 0.2 * torch.randn( - positions.shape, generator=generator, device=device, dtype=dtype -) - -# Create the atomic numbers tensor -atomic_numbers = torch.full((positions.shape[0],), 29, device=device, dtype=torch.int) - -# Cu atomic mass in atomic mass units -masses = torch.full((positions.shape[0],), 63.546, device=device, dtype=dtype) - -# Create state with batch dimension -state = ts.SimState( - positions=positions, - masses=masses, - cell=cell.unsqueeze(0), - atomic_numbers=atomic_numbers, - pbc=True, -) - -# Initialize the Soft Sphere model -model = SoftSphereModel( - sigma=2.5, - device=device, - dtype=dtype, - compute_forces=True, - compute_stress=False, - use_neighbor_list=True, -) - -# Run initial simulation and get results -results = model(state) - -# Initialize FIRE optimizer -state = ts.fire_init(state=state, model=model, dt_start=0.005) - - -# Run optimization for N_steps -for step in range(N_steps): - if step % 100 == 0: - print(f"{step=}: Total energy: {state.energy[0].item()} eV") - state = ts.fire_step(state=state, model=model, dt_max=0.01) - -# Print max force after optimization -print(f"Initial energy: {results['energy'][0].item()} eV") -print(f"Final energy: {state.energy[0].item()} eV") -print(f"Initial max force: {torch.max(torch.abs(results['forces'][0])).item()} eV/Å") -print(f"Final max force: {torch.max(torch.abs(state.forces[0])).item()} eV/Å") diff --git a/examples/old_scripts/2_Structural_optimization/2.3_MACE_Gradient_Descent.py b/examples/old_scripts/2_Structural_optimization/2.3_MACE_Gradient_Descent.py deleted file mode 100644 index 7a9c3d53..00000000 --- a/examples/old_scripts/2_Structural_optimization/2.3_MACE_Gradient_Descent.py +++ /dev/null @@ -1,124 +0,0 @@ -"""Batched MACE gradient descent example.""" - -# /// script -# dependencies = ["mace-torch>=0.3.12"] -# /// -import os - -import numpy as np -import torch -from ase.build import bulk -from mace.calculators.foundations_models import mace_mp - -import torch_sim as ts -from torch_sim.models.mace import MaceModel, MaceUrls - - -# Set device and data type -device = torch.device("cuda" if torch.cuda.is_available() else "cpu") -dtype = torch.float32 -# Option 1: Load the raw model from the downloaded model -loaded_model = mace_mp( - model=MaceUrls.mace_mpa_medium, - return_raw_model=True, - default_dtype=str(dtype).removeprefix("torch."), - device=str(device), -) - -# Option 2: Load from local file (comment out Option 1 to use this) -# loaded_model = torch.load("path/to/model.pt", map_location=device) - -# Number of steps to run -SMOKE_TEST = os.getenv("CI") is not None -N_steps = 10 if SMOKE_TEST else 500 - -# Set random seed for reproducibility -rng = np.random.default_rng(seed=0) - -# Create diamond cubic Silicon systems -si_dc = bulk("Si", "diamond", a=5.43, cubic=True).repeat((2, 2, 2)) -si_dc.positions += 0.2 * rng.standard_normal(si_dc.positions.shape) - -fe = bulk("Fe", "bcc", a=2.8665, cubic=True).repeat((3, 3, 3)) -fe.positions += 0.2 * rng.standard_normal(fe.positions.shape) - -# Create a list of our atomic systems -atoms_list = [si_dc, fe] - -# Create batched model -batched_model = MaceModel( - model=loaded_model, - device=device, - compute_forces=True, - compute_stress=True, - dtype=dtype, - enable_cueq=False, -) - -""" -# Convert data to tensors -positions_list = [ - torch.tensor(atoms.positions, device=device, dtype=dtype) for atoms in atoms_list -] -cell_list = [ - torch.tensor(atoms.cell.array, device=device, dtype=dtype) for atoms in atoms_list -] -masses_list = [ - torch.tensor(atoms.get_masses(), device=device, dtype=dtype) for atoms in atoms_list -] -atomic_numbers_list = [atoms.get_atomic_numbers() for atoms in atoms_list] - -# Create batch data format -# First concatenate positions array -positions_numpy = np.concatenate([atoms.positions for atoms in atoms_list]) -positions = torch.tensor(positions_numpy, device=device, dtype=dtype) - -# Stack cell vectors into a (n_batch, 3, 3) array -cell_numpy = np.stack([atoms.cell.array for atoms in atoms_list]) -cell = torch.tensor(cell_numpy, device=device, dtype=dtype) - -# Concatenate atomic numbers -atomic_numbers_numpy = np.concatenate( - [atoms.get_atomic_numbers() for atoms in atoms_list] -) -atomic_numbers = torch.tensor(atomic_numbers_numpy, device=device, dtype=torch.int) - -# Concatenate masses -masses_numpy = np.concatenate([atoms.get_masses() for atoms in atoms_list]) -masses = torch.tensor(masses_numpy, device=device, dtype=dtype) - -# Create system indices tensor for scatter operations -atoms_per_system = torch.tensor( - [len(atoms) for atoms in atoms_list], device=device, dtype=torch.int -) -system_indices = torch.repeat_interleave( - torch.arange(len(atoms_per_system), device=device), atoms_per_system -) -""" - -state = ts.io.atoms_to_state(atoms_list, device=device, dtype=dtype) - -print(f"Positions shape: {state.positions.shape}") -print(f"Cell shape: {state.cell.shape}") -print(f"System indices shape: {state.system_idx.shape}") - -# Run initial inference -results = batched_model(state) - -# Use different learning rates for each batch -learning_rate = 0.01 - -# Initialize batched gradient descent optimizer -state = ts.gradient_descent_init(state=state, model=batched_model) - -# Run batched optimization for a few steps -print("\nRunning batched gradient descent:") -for step in range(N_steps): - if step % 10 == 0: - print(f"Step {step}, Energy: {[res.item() for res in state.energy]} eV") - state = ts.gradient_descent_step( - state=state, model=batched_model, pos_lr=learning_rate - ) - -print(f"Initial energies: {[res.item() for res in results['energy']]} eV") -print(f"Final energies: {[res.item() for res in state.energy]} eV") diff --git a/examples/old_scripts/2_Structural_optimization/2.4_MACE_FIRE.py b/examples/old_scripts/2_Structural_optimization/2.4_MACE_FIRE.py deleted file mode 100644 index 3f737955..00000000 --- a/examples/old_scripts/2_Structural_optimization/2.4_MACE_FIRE.py +++ /dev/null @@ -1,88 +0,0 @@ -"""Batched MACE FIRE optimizer.""" - -# /// script -# dependencies = ["mace-torch>=0.3.12"] -# /// -import os - -import numpy as np -import torch -from ase.build import bulk -from mace.calculators.foundations_models import mace_mp - -import torch_sim as ts -from torch_sim.models.mace import MaceModel, MaceUrls - - -# Set device and data type -device = torch.device("cuda" if torch.cuda.is_available() else "cpu") -dtype = torch.float32 - -# Option 1: Load the raw model from the downloaded model -loaded_model = mace_mp( - model=MaceUrls.mace_mpa_medium, - return_raw_model=True, - default_dtype=str(dtype).removeprefix("torch."), - device=str(device), -) - -# Option 2: Load from local file (comment out Option 1 to use this) -# loaded_model = torch.load("path/to/model.pt", map_location=device) - -# Number of steps to run -SMOKE_TEST = os.getenv("CI") is not None -N_steps = 10 if SMOKE_TEST else 500 - -# Set random seed for reproducibility -rng = np.random.default_rng(seed=0) - -# Create diamond cubic Silicon -si_dc = bulk("Si", "diamond", a=5.21, cubic=True).repeat((2, 2, 2)) -si_dc.positions += 0.2 * rng.standard_normal(si_dc.positions.shape) - -# Create FCC Copper -cu_dc = bulk("Cu", "fcc", a=3.85).repeat((2, 2, 2)) -cu_dc.positions += 0.2 * rng.standard_normal(cu_dc.positions.shape) - -# Create BCC Iron -fe_dc = bulk("Fe", "bcc", a=2.95).repeat((2, 2, 2)) -fe_dc.positions += 0.2 * rng.standard_normal(fe_dc.positions.shape) - -# Create a list of our atomic systems -atoms_list = [si_dc, cu_dc, fe_dc] - -# Print structure information -print(f"Silicon atoms: {len(si_dc)}") -print(f"Copper atoms: {len(cu_dc)}") -print(f"Iron atoms: {len(fe_dc)}") -print(f"Total number of structures: {len(atoms_list)}") - -# Create batched model -model = MaceModel( - model=loaded_model, - device=device, - compute_forces=True, - compute_stress=True, - dtype=dtype, - enable_cueq=False, -) - -# Convert atoms to state -state = ts.io.atoms_to_state(atoms_list, device=device, dtype=dtype) -# Run initial inference -results = model(state) - -# Initialize unit cell gradient descent optimizer -state = ts.fire_init(state=state, model=model, dt_start=0.005) - - -# Run optimization for a few steps -print("\nRunning FIRE:") -for step in range(N_steps): - if step % 20 == 0: - print(f"Step {step}, Energy: {[energy.item() for energy in state.energy]}") - - state = ts.fire_step(state=state, model=model, dt_max=0.01) - -print(f"Initial energies: {[energy.item() for energy in results['energy']]} eV") -print(f"Final energies: {[energy.item() for energy in state.energy]} eV") diff --git a/examples/old_scripts/2_Structural_optimization/2.5_MACE_UnitCellFilter_Gradient_Descent.py b/examples/old_scripts/2_Structural_optimization/2.5_MACE_UnitCellFilter_Gradient_Descent.py deleted file mode 100644 index 0185fcb8..00000000 --- a/examples/old_scripts/2_Structural_optimization/2.5_MACE_UnitCellFilter_Gradient_Descent.py +++ /dev/null @@ -1,120 +0,0 @@ -"""Batched MACE unit cell filter with gradient descent optimizer.""" - -# /// script -# dependencies = ["mace-torch>=0.3.12"] -# /// -import os - -import numpy as np -import torch -from ase.build import bulk -from mace.calculators.foundations_models import mace_mp - -import torch_sim as ts -from torch_sim.models.mace import MaceModel, MaceUrls -from torch_sim.units import UnitConversion - - -# Set device and data type -device = torch.device("cuda" if torch.cuda.is_available() else "cpu") -dtype = torch.float32 - -# Option 1: Load the raw model from the downloaded model -loaded_model = mace_mp( - model=MaceUrls.mace_mpa_medium, - return_raw_model=True, - default_dtype=str(dtype).removeprefix("torch."), - device=str(device), -) - -# Option 2: Load from local file (comment out Option 1 to use this) -# loaded_model = torch.load("path/to/model.pt", map_location=device) - -# Number of steps to run -SMOKE_TEST = os.getenv("CI") is not None -N_steps = 10 if SMOKE_TEST else 500 - -# Set random seed for reproducibility -rng = np.random.default_rng(seed=0) - -# Create diamond cubic Silicon -si_dc = bulk("Si", "diamond", a=5.21, cubic=True).repeat((2, 2, 2)) -si_dc.positions += 0.2 * rng.standard_normal(si_dc.positions.shape) - -# Create FCC Copper -cu_dc = bulk("Cu", "fcc", a=3.85).repeat((2, 2, 2)) -cu_dc.positions += 0.2 * rng.standard_normal(cu_dc.positions.shape) - -# Create BCC Iron -fe_dc = bulk("Fe", "bcc", a=2.95).repeat((2, 2, 2)) -fe_dc.positions += 0.2 * rng.standard_normal(fe_dc.positions.shape) - -# Create a list of our atomic systems -atoms_list = [si_dc, cu_dc, fe_dc] - -# Print structure information -print(f"Silicon atoms: {len(si_dc)}") -print(f"Copper atoms: {len(cu_dc)}") -print(f"Iron atoms: {len(fe_dc)}") -print(f"Total number of structures: {len(atoms_list)}") - -# Create batched model -model = MaceModel( - model=loaded_model, - device=device, - compute_forces=True, - compute_stress=True, - dtype=dtype, - enable_cueq=False, -) - -# Convert atoms to state -state = ts.io.atoms_to_state(atoms_list, device=device, dtype=dtype) -# Run initial inference -results = model(state) - -# Use same learning rate for all batches -pos_lr, cell_lr = 0.01, 0.1 - - -state = ts.gradient_descent_init( - state=state, - model=model, - cell_filter=ts.CellFilter.unit, - cell_factor=None, # Will default to atoms per system - hydrostatic_strain=False, - constant_volume=False, - scalar_pressure=0.0, -) - - -# Run optimization for a few steps -print("\nRunning batched unit cell gradient descent:") -for step in range(N_steps): - P1 = -torch.trace(state.stress[0]) * UnitConversion.eV_per_Ang3_to_GPa / 3 - P2 = -torch.trace(state.stress[1]) * UnitConversion.eV_per_Ang3_to_GPa / 3 - P3 = -torch.trace(state.stress[2]) * UnitConversion.eV_per_Ang3_to_GPa / 3 - - if step % 20 == 0: - print( - f"Step {step}, Energy: {[energy.item() for energy in state.energy]}, " - f"P1={P1:.4f} GPa, P2={P2:.4f} GPa, P3={P3:.4f} GPa" - ) - - state = ts.gradient_descent_step( - state=state, model=model, pos_lr=pos_lr, cell_lr=cell_lr - ) - -print(f"Initial energies: {[energy.item() for energy in results['energy']]} eV") -print(f"Final energies: {[energy.item() for energy in state.energy]} eV") - -initial_pressure = [ - torch.trace(stress).item() * UnitConversion.eV_per_Ang3_to_GPa / 3 - for stress in results["stress"] -] -final_pressure = [ - torch.trace(stress).item() * UnitConversion.eV_per_Ang3_to_GPa / 3 - for stress in state.stress -] -print(f"{initial_pressure=} GPa") -print(f"{final_pressure=} GPa") diff --git a/examples/old_scripts/2_Structural_optimization/2.6_MACE_UnitCellFilter_FIRE.py b/examples/old_scripts/2_Structural_optimization/2.6_MACE_UnitCellFilter_FIRE.py deleted file mode 100644 index 29e24406..00000000 --- a/examples/old_scripts/2_Structural_optimization/2.6_MACE_UnitCellFilter_FIRE.py +++ /dev/null @@ -1,114 +0,0 @@ -"""Batched MACE unit cell filter with FIRE optimizer.""" - -# /// script -# dependencies = ["mace-torch>=0.3.12"] -# /// -import os - -import numpy as np -import torch -from ase.build import bulk -from mace.calculators.foundations_models import mace_mp - -import torch_sim as ts -from torch_sim.models.mace import MaceModel, MaceUrls -from torch_sim.units import UnitConversion - - -# Set device and data type -device = torch.device("cuda" if torch.cuda.is_available() else "cpu") -dtype = torch.float32 - -# Option 1: Load the raw model from the downloaded model -loaded_model = mace_mp( - model=MaceUrls.mace_mpa_medium, - return_raw_model=True, - default_dtype=str(dtype).removeprefix("torch."), - device=str(device), -) - -# Option 2: Load from local file (comment out Option 1 to use this) -# loaded_model = torch.load("path/to/model.pt", map_location=device) - -# Number of steps to run -SMOKE_TEST = os.getenv("CI") is not None -N_steps = 10 if SMOKE_TEST else 500 - -# Set random seed for reproducibility -rng = np.random.default_rng(seed=0) - -# Create diamond cubic Silicon -si_dc = bulk("Si", "diamond", a=5.21, cubic=True).repeat((2, 2, 2)) -si_dc.positions += 0.2 * rng.standard_normal(si_dc.positions.shape) - -# Create FCC Copper -cu_dc = bulk("Cu", "fcc", a=3.85).repeat((2, 2, 2)) -cu_dc.positions += 0.2 * rng.standard_normal(cu_dc.positions.shape) - -# Create BCC Iron -fe_dc = bulk("Fe", "bcc", a=2.95).repeat((2, 2, 2)) -fe_dc.positions += 0.2 * rng.standard_normal(fe_dc.positions.shape) - -# Create a list of our atomic systems -atoms_list = [si_dc, cu_dc, fe_dc] - -# Print structure information -print(f"Silicon atoms: {len(si_dc)}") -print(f"Copper atoms: {len(cu_dc)}") -print(f"Iron atoms: {len(fe_dc)}") -print(f"Total number of structures: {len(atoms_list)}") - -# Create batched model -model = MaceModel( - model=loaded_model, - device=device, - compute_forces=True, - compute_stress=True, - dtype=dtype, - enable_cueq=False, -) - -# Convert atoms to state -state = ts.io.atoms_to_state(atoms_list, device=device, dtype=dtype) -# Run initial inference -results = model(state) - -# Initialize FIRE optimizer with unit cell filter -state = ts.fire_init( - state=state, - model=model, - cell_filter=ts.CellFilter.unit, - cell_factor=None, # Will default to atoms per system - hydrostatic_strain=False, - constant_volume=False, - scalar_pressure=0.0, -) - -# Run optimization for a few steps -print("\nRunning batched unit cell gradient descent:") -for step in range(N_steps): - P1 = -torch.trace(state.stress[0]) * UnitConversion.eV_per_Ang3_to_GPa / 3 - P2 = -torch.trace(state.stress[1]) * UnitConversion.eV_per_Ang3_to_GPa / 3 - P3 = -torch.trace(state.stress[2]) * UnitConversion.eV_per_Ang3_to_GPa / 3 - - if step % 20 == 0: - print( - f"Step {step}, Energy: {[energy.item() for energy in state.energy]}, " - f"P1={P1:.4f} GPa, P2={P2:.4f} GPa, P3={P3:.4f} GPa" - ) - - state = ts.fire_step(state=state, model=model) - -print(f"Initial energies: {[energy.item() for energy in results['energy']]} eV") -print(f"Final energies: {[energy.item() for energy in state.energy]} eV") - -initial_pressure = [ - torch.trace(stress).item() * UnitConversion.eV_per_Ang3_to_GPa / 3 - for stress in results["stress"] -] -final_pressure = [ - torch.trace(stress).item() * UnitConversion.eV_per_Ang3_to_GPa / 3 - for stress in state.stress -] -print(f"{initial_pressure=} GPa") -print(f"{final_pressure=} GPa") diff --git a/examples/old_scripts/2_Structural_optimization/2.7_MACE_FrechetCellFilter_FIRE.py b/examples/old_scripts/2_Structural_optimization/2.7_MACE_FrechetCellFilter_FIRE.py deleted file mode 100644 index 3a1d881c..00000000 --- a/examples/old_scripts/2_Structural_optimization/2.7_MACE_FrechetCellFilter_FIRE.py +++ /dev/null @@ -1,114 +0,0 @@ -"""Batched MACE frechet cell filter with FIRE optimizer.""" - -# /// script -# dependencies = ["mace-torch>=0.3.12"] -# /// -import os - -import numpy as np -import torch -from ase.build import bulk -from mace.calculators.foundations_models import mace_mp - -import torch_sim as ts -from torch_sim.models.mace import MaceModel, MaceUrls -from torch_sim.units import UnitConversion - - -# Set device and data type -device = torch.device("cuda" if torch.cuda.is_available() else "cpu") -dtype = torch.float32 - -# Option 1: Load the raw model from the downloaded model -loaded_model = mace_mp( - model=MaceUrls.mace_mpa_medium, - return_raw_model=True, - default_dtype=str(dtype).removeprefix("torch."), - device=str(device), -) - -# Option 2: Load from local file (comment out Option 1 to use this) -# loaded_model = torch.load("path/to/model.pt", map_location=device) - -# Number of steps to run -SMOKE_TEST = os.getenv("CI") is not None -N_steps = 10 if SMOKE_TEST else 500 - -# Set random seed for reproducibility -rng = np.random.default_rng(seed=0) - -# Create diamond cubic Silicon -si_dc = bulk("Si", "diamond", a=5.21, cubic=True).repeat((2, 2, 2)) -si_dc.positions += 0.2 * rng.standard_normal(si_dc.positions.shape) - -# Create FCC Copper -cu_dc = bulk("Cu", "fcc", a=3.85).repeat((2, 2, 2)) -cu_dc.positions += 0.2 * rng.standard_normal(cu_dc.positions.shape) - -# Create BCC Iron -fe_dc = bulk("Fe", "bcc", a=2.95).repeat((2, 2, 2)) -fe_dc.positions += 0.2 * rng.standard_normal(fe_dc.positions.shape) - -# Create a list of our atomic systems -atoms_list = [si_dc, cu_dc, fe_dc] - -# Print structure information -print(f"Silicon atoms: {len(si_dc)}") -print(f"Copper atoms: {len(cu_dc)}") -print(f"Iron atoms: {len(fe_dc)}") -print(f"Total number of structures: {len(atoms_list)}") - -# Create batched model -model = MaceModel( - model=loaded_model, - device=device, - compute_forces=True, - compute_stress=True, - dtype=dtype, - enable_cueq=False, -) - -# Convert atoms to state -state = ts.io.atoms_to_state(atoms_list, device=device, dtype=dtype) -# Run initial inference -results = model(state) - -# Initialize FIRE optimizer with Frechet cell filter -state = ts.fire_init( - state=state, - model=model, - cell_filter=ts.CellFilter.frechet, - cell_factor=None, # Will default to atoms per system - hydrostatic_strain=False, - constant_volume=False, - scalar_pressure=0.0, -) - -# Run optimization for a few steps -print("\nRunning batched frechet cell filter with FIRE:") -for step in range(N_steps): - P1 = -torch.trace(state.stress[0]) * UnitConversion.eV_per_Ang3_to_GPa / 3 - P2 = -torch.trace(state.stress[1]) * UnitConversion.eV_per_Ang3_to_GPa / 3 - P3 = -torch.trace(state.stress[2]) * UnitConversion.eV_per_Ang3_to_GPa / 3 - - if step % 20 == 0: - print( - f"Step {step}, Energy: {[energy.item() for energy in state.energy]}, " - f"P1={P1:.4f} GPa, P2={P2:.4f} GPa, P3={P3:.4f} GPa" - ) - - state = ts.fire_step(state=state, model=model) - -print(f"Initial energies: {[energy.item() for energy in results['energy']]} eV") -print(f"Final energies: {[energy.item() for energy in state.energy]} eV") - -initial_pressure = [ - torch.trace(stress).item() * UnitConversion.eV_per_Ang3_to_GPa / 3 - for stress in results["stress"] -] -final_pressure = [ - torch.trace(stress).item() * UnitConversion.eV_per_Ang3_to_GPa / 3 - for stress in state.stress -] -print(f"{initial_pressure=} GPa") -print(f"{final_pressure=} GPa") diff --git a/examples/old_scripts/3_Dynamics/3.10_Hybrid_swap_mc.py b/examples/old_scripts/3_Dynamics/3.10_Hybrid_swap_mc.py deleted file mode 100644 index 8a25f0c9..00000000 --- a/examples/old_scripts/3_Dynamics/3.10_Hybrid_swap_mc.py +++ /dev/null @@ -1,100 +0,0 @@ -"""Hybrid swap Monte Carlo simulation.""" - -# /// script -# dependencies = ["mace-torch>=0.3.12", "pymatgen>=2025.2.18"] -# /// -from dataclasses import dataclass - -import torch -from mace.calculators.foundations_models import mace_mp -from pymatgen.core import Structure - -import torch_sim as ts -from torch_sim.integrators.md import MDState -from torch_sim.models.mace import MaceModel, MaceUrls -from torch_sim.units import MetalUnits as Units - - -device = torch.device("cuda" if torch.cuda.is_available() else "cpu") -dtype = torch.float32 -kT = 1000 * Units.temperature - -# Option 1: Load the raw model from the downloaded model -loaded_model = mace_mp( - model=MaceUrls.mace_mpa_medium, - return_raw_model=True, - default_dtype=str(dtype).removeprefix("torch."), - device=str(device), -) - -# Option 2: Load from local file (comment out Option 1 to use this) -# loaded_model = torch.load("path/to/model.pt", map_location=device) - -model = MaceModel( - model=loaded_model, - device=device, - compute_forces=True, - compute_stress=False, - dtype=dtype, - enable_cueq=False, -) - -# %% -lattice = [[5.43, 0, 0], [0, 5.43, 0], [0, 0, 5.43]] -species = ["Cu", "Cu", "Cu", "Zr", "Cu", "Zr", "Zr", "Zr"] -coords = [ - [0.0, 0.0, 0.0], - [0.25, 0.25, 0.25], - [0.0, 0.5, 0.5], - [0.25, 0.75, 0.75], - [0.5, 0.0, 0.5], - [0.75, 0.25, 0.75], - [0.5, 0.5, 0.0], - [0.75, 0.75, 0.25], -] -structure = Structure(lattice, species, coords) - -state = ts.io.structures_to_state([structure], device=device, dtype=dtype) - - -# %% -@dataclass(kw_only=True) -class HybridSwapMCState(ts.SwapMCState, MDState): - """State for Monte Carlo simulations. - - Attributes: - energy: Energy of the system - last_swap: Last swap attempted - """ - - last_permutation: torch.Tensor - _atom_attributes = ( - ts.SwapMCState._atom_attributes | MDState._atom_attributes | {"last_permutation"} # noqa: SLF001 - ) - _system_attributes = ( - ts.SwapMCState._system_attributes | MDState._system_attributes # noqa: SLF001 - ) - - -md_state = ts.nvt_langevin_init(state=state, model=model, kT=torch.tensor(kT), seed=42) - -swap_state = ts.swap_mc_init(state=md_state, model=model) -hybrid_state = HybridSwapMCState( - **md_state.attributes, - last_permutation=torch.arange( - md_state.n_atoms, device=md_state.device, dtype=torch.long - ), -) - -rng = torch.Generator(device=device) -rng.manual_seed(42) - -n_steps = 100 -dt = torch.tensor(0.002) -for step in range(n_steps): - if step % 10 == 0: - hybrid_state = ts.swap_mc_step(state=hybrid_state, model=model, kT=kT, rng=rng) - else: - hybrid_state = ts.nvt_langevin_step( - state=hybrid_state, model=model, dt=dt, kT=torch.tensor(kT) - ) diff --git a/examples/old_scripts/3_Dynamics/3.11_Lennard_Jones_NPT_Langevin.py b/examples/old_scripts/3_Dynamics/3.11_Lennard_Jones_NPT_Langevin.py deleted file mode 100644 index 6069ed46..00000000 --- a/examples/old_scripts/3_Dynamics/3.11_Lennard_Jones_NPT_Langevin.py +++ /dev/null @@ -1,168 +0,0 @@ -"""Lennard-Jones simulation in NPT ensemble using Langevin thermostat.""" - -# /// script -# dependencies = ["scipy>=1.15"] -# /// -import itertools -import os - -import torch - -import torch_sim as ts -from torch_sim.models.lennard_jones import LennardJonesModel -from torch_sim.units import MetalUnits as Units -from torch_sim.units import UnitConversion - - -# Set up the device and data type -device = torch.device("cuda" if torch.cuda.is_available() else "cpu") -dtype = torch.float32 - -# Set random seed and deterministic behavior for reproducibility -torch.manual_seed(42) -if torch.cuda.is_available(): - torch.cuda.manual_seed_all(42) -torch.backends.cudnn.deterministic = True -torch.backends.cudnn.benchmark = False - -# Set up the random number generator -generator = torch.Generator(device=device) -generator.manual_seed(42) # For reproducibility - -# Number of steps to run -SMOKE_TEST = os.getenv("CI") is not None -N_steps = 100 if SMOKE_TEST else 10_000 - -# Create face-centered cubic (FCC) Argon -# 5.26 Å is a typical lattice constant for Ar -a_len = 5.26 # Lattice constant - -# Generate base FCC unit cell positions (scaled by lattice constant) -base_positions = torch.tensor( - [ - [0.0, 0.0, 0.0], # Corner - [0.0, 0.5, 0.5], # Face centers - [0.5, 0.0, 0.5], - [0.5, 0.5, 0.0], - ], - device=device, - dtype=dtype, -) - -# Create 4x4x4 supercell of FCC Argon manually -positions = [] -for i, j, k in itertools.product(range(4), range(4), range(4)): - for base_pos in base_positions: - # Add unit cell position + offset for supercell - pos = base_pos + torch.tensor([i, j, k], device=device, dtype=dtype) - positions.append(pos) - -# Stack the positions into a tensor -positions = torch.stack(positions) - -# Scale by lattice constant -positions = positions * a_len - -# Create the cell tensor -cell = torch.tensor( - [[4 * a_len, 0, 0], [0, 4 * a_len, 0], [0, 0, 4 * a_len]], - device=device, - dtype=dtype, -) - -# Create the atomic numbers tensor (Argon = 18) -atomic_numbers = torch.full((positions.shape[0],), 18, device=device, dtype=torch.int) -# Create the masses tensor (Argon = 39.948 amu) -masses = torch.full((positions.shape[0],), 39.948, device=device, dtype=dtype) - -# Initialize the Lennard-Jones model -# Parameters: -# - sigma: distance at which potential is zero (3.405 Å for Ar) -# - epsilon: depth of potential well (0.0104 eV for Ar) -# - cutoff: distance beyond which interactions are ignored (typically 2.5*sigma) -model = LennardJonesModel( - use_neighbor_list=False, - sigma=3.405, - epsilon=0.0104, - cutoff=2.5 * 3.405, - device=device, - dtype=dtype, - compute_forces=True, - compute_stress=True, -) -state = ts.SimState( - positions=positions, - masses=masses, - cell=cell.unsqueeze(0), - atomic_numbers=atomic_numbers, - pbc=True, -) -# Run initial simulation and get results -results = model(state) - -dt = torch.tensor(0.001 * Units.time, device=device, dtype=dtype) # Time step (1 fs) -kT = torch.tensor( - 200 * Units.temperature, device=device, dtype=dtype -) # Temperature (200 K) -target_pressure = ( - torch.tensor(10_000, device=device, dtype=dtype) * Units.pressure -) # Target pressure (10 kbar) - -state = ts.npt_langevin_init( - state=state, - model=model, - dt=dt, - kT=kT, - seed=1, - alpha=1.0 / (100 * dt), - cell_alpha=1.0 / (100 * dt), - b_tau=1 / (1000 * dt), -) - -# Run the simulation -for step in range(N_steps): - if step % 50 == 0: - temp = ( - ts.calc_kT( - masses=state.masses, momenta=state.momenta, system_idx=state.system_idx - ) - / Units.temperature - ) - pressure = ts.get_pressure( - model(state)["stress"], - ts.calc_kinetic_energy( - masses=state.masses, momenta=state.momenta, system_idx=state.system_idx - ), - torch.linalg.det(state.cell), - ) - pressure = pressure.item() / Units.pressure - xx, yy, zz = state.cell[..., 0, 0], state.cell[..., 1, 1], state.cell[..., 2, 2] - print( - f"{step=}: Temperature: {temp.item():.4f}, " - f"{pressure=:.4f}, " - f"cell xx yy zz: {xx.item():.4f}, {yy.item():.4f}, {zz.item():.4f}" - ) - state = ts.npt_langevin_step( - state=state, - model=model, - dt=dt, - kT=kT, - external_pressure=target_pressure, - ) - -temp = ( - ts.calc_kT(masses=state.masses, momenta=state.momenta, system_idx=state.system_idx) - / Units.temperature -) -print(f"Final temperature: {temp.item():.4f}") - - -stress = model(state)["stress"] -kinetic_energy = ts.calc_kinetic_energy( - masses=state.masses, momenta=state.momenta, system_idx=state.system_idx -) -volume = torch.linalg.det(state.cell) -pressure = ts.get_pressure(stress, kinetic_energy, volume) -pressure = pressure.item() / Units.pressure -print(f"Final {pressure=:.4f}") -print(stress * UnitConversion.eV_per_Ang3_to_GPa) diff --git a/examples/old_scripts/3_Dynamics/3.12_MACE_NPT_Langevin.py b/examples/old_scripts/3_Dynamics/3.12_MACE_NPT_Langevin.py deleted file mode 100644 index 62df3675..00000000 --- a/examples/old_scripts/3_Dynamics/3.12_MACE_NPT_Langevin.py +++ /dev/null @@ -1,138 +0,0 @@ -"""NPT simulation with MACE and Nose-Hoover thermostat.""" - -# /// script -# dependencies = ["mace-torch>=0.3.12"] -# /// -import os - -import torch -from ase.build import bulk -from mace.calculators.foundations_models import mace_mp - -import torch_sim as ts -from torch_sim.models.mace import MaceModel, MaceUrls -from torch_sim.units import MetalUnits as Units - - -# Set device and data type -device = torch.device("cuda" if torch.cuda.is_available() else "cpu") -dtype = torch.float32 - -# Option 1: Load the raw model from the downloaded model -loaded_model = mace_mp( - model=MaceUrls.mace_mpa_medium, - return_raw_model=True, - default_dtype=str(dtype).removeprefix("torch."), - device=str(device), -) - -# Option 2: Load from local file (comment out Option 1 to use this) -# loaded_model = torch.load("path/to/model.pt", map_location=device) - -# Create diamond cubic Silicon -si_dc = bulk("Si", "diamond", a=5.43, cubic=True).repeat((2, 2, 2)) - -# Initialize the MACE model -model = MaceModel( - model=loaded_model, - device=device, - compute_forces=True, - compute_stress=True, - dtype=dtype, - enable_cueq=False, -) -state = ts.io.atoms_to_state(si_dc, device=device, dtype=dtype) - -# Run initial inference -results = model(state) - -SMOKE_TEST = os.getenv("CI") is not None -N_steps_nvt = 20 if SMOKE_TEST else 2_000 -N_steps_npt = 20 if SMOKE_TEST else 2_000 -dt = torch.tensor(0.001 * Units.time, device=device, dtype=dtype) # Time step (1 fs) -kT = ( - torch.tensor(300, device=device, dtype=dtype) * Units.temperature -) # Initial temperature (300 K) -target_pressure = torch.tensor( - 10_000 * Units.pressure, device=device, dtype=dtype -) # Target pressure (10 Kbar) - -state = ts.nvt_nose_hoover_init(state=state, model=model, kT=kT, dt=dt, seed=1) - -for step in range(N_steps_nvt): - if step % 10 == 0: - temp = ( - ts.calc_kT( - masses=state.masses, momenta=state.momenta, system_idx=state.system_idx - ) - / Units.temperature - ) - invariant = float(ts.nvt_nose_hoover_invariant(state, kT=kT)) - print(f"{step=}: Temperature: {temp.item():.4f}: {invariant=:.4f}, ") - state = ts.nvt_nose_hoover_step(state=state, model=model, dt=dt, kT=kT) - -state = ts.npt_langevin_init( - state=state, - model=model, - kT=kT, - dt=dt, - seed=1, - alpha=1.0 / (100 * dt), - cell_alpha=1.0 / (100 * dt), - b_tau=1 / (1000 * dt), -) - -for step in range(N_steps_npt): - if step % 10 == 0: - temp = ( - ts.calc_kT( - masses=state.masses, momenta=state.momenta, system_idx=state.system_idx - ) - / Units.temperature - ) - stress = model(state)["stress"] - volume = torch.det(state.cell) - pressure = ( - ts.get_pressure( - stress, - ts.calc_kinetic_energy( - masses=state.masses, - momenta=state.momenta, - system_idx=state.system_idx, - ), - volume, - ).item() - / Units.pressure - ) - xx, yy, zz = torch.diag(state.cell[0]) - print( - f"{step=}: Temperature: {temp.item():.4f}, " - f"pressure: {pressure:.4f}, " - f"cell xx yy zz: {xx.item():.4f}, {yy.item():.4f}, {zz.item():.4f}" - ) - state = ts.npt_langevin_step( - state=state, - model=model, - dt=dt, - kT=kT, - external_pressure=target_pressure, - ) - -final_temp = ( - ts.calc_kT(masses=state.masses, momenta=state.momenta, system_idx=state.system_idx) - / Units.temperature -) -print(f"Final temperature: {final_temp.item():.4f} K") -final_stress = model(state)["stress"] -final_volume = torch.det(state.cell) -final_pressure = ( - ts.get_pressure( - final_stress, - ts.calc_kinetic_energy( - masses=state.masses, momenta=state.momenta, system_idx=state.system_idx - ), - final_volume, - ).item() - / Units.pressure -) -print(f"Final pressure: {final_pressure:.4f} bar") diff --git a/examples/old_scripts/3_Dynamics/3.13_MACE_NVE_non_pbc.py b/examples/old_scripts/3_Dynamics/3.13_MACE_NVE_non_pbc.py deleted file mode 100644 index 8acef6ab..00000000 --- a/examples/old_scripts/3_Dynamics/3.13_MACE_NVE_non_pbc.py +++ /dev/null @@ -1,80 +0,0 @@ -"""NVE simulation with MACE.""" - -# /// script -# dependencies = ["mace-torch>=0.3.12"] -# /// -import os -import time - -import torch -from ase.build import molecule -from mace.calculators.foundations_models import mace_off - -import torch_sim as ts -from torch_sim.models.mace import MaceModel, MaceUrls -from torch_sim.units import MetalUnits as Units - - -# Set device and data type -device = torch.device("cuda" if torch.cuda.is_available() else "cpu") -dtype = torch.float32 - -# Option 1: Load the raw model from the downloaded model -loaded_model = mace_off( - model=MaceUrls.mace_mpa_medium, - return_raw_model=True, - default_dtype=str(dtype).removeprefix("torch."), - device=str(device), -) - -# Option 2: Load from local file (comment out Option 1 to use this) -# loaded_model = torch.load("path/to/model.pt", map_location=device) - -# Number of steps to run -SMOKE_TEST = os.getenv("CI") is not None -N_steps = 20 if SMOKE_TEST else 2_000 - -mol = molecule("methylenecyclopropane") - -# Initialize the MACE model -model = MaceModel( - model=loaded_model, - device=device, - compute_forces=True, - compute_stress=False, - dtype=dtype, - enable_cueq=False, -) - -state = ts.io.atoms_to_state(mol, device=device, dtype=dtype) - -# Run initial inference -results = model(state) - -# Setup NVE MD simulation parameters -kT = ( - torch.tensor(300, device=device, dtype=dtype) * Units.temperature -) # Initial temperature (K) -dt = torch.tensor(0.002 * Units.time, device=device, dtype=dtype) # Timestep (ps) - - -# Initialize NVE integrator -state = ts.nve_init(state=state, model=model, kT=kT, seed=1) - -# Run MD simulation -print("\nStarting NVE molecular dynamics simulation...") -start_time = time.perf_counter() -for step in range(N_steps): - total_energy = state.energy + ts.calc_kinetic_energy( - masses=state.masses, momenta=state.momenta, system_idx=state.system_idx - ) - if step % 10 == 0: - print(f"Step {step}: Total energy: {total_energy.item():.4f} eV") - state = ts.nve_step(state=state, model=model, dt=dt) -end_time = time.perf_counter() - -# Report simulation results -print("\nSimulation complete!") -print(f"Time taken: {end_time - start_time:.2f} seconds") -print(f"Average time per step: {(end_time - start_time) / N_steps:.4f} seconds") -print(f"Final total energy: {total_energy.item()} eV") diff --git a/examples/old_scripts/3_Dynamics/3.1_Lennard_Jones_NVE.py b/examples/old_scripts/3_Dynamics/3.1_Lennard_Jones_NVE.py deleted file mode 100644 index 7094ef75..00000000 --- a/examples/old_scripts/3_Dynamics/3.1_Lennard_Jones_NVE.py +++ /dev/null @@ -1,123 +0,0 @@ -"""NVE simulation with Lennard-Jones potential.""" - -# /// script -# dependencies = ["scipy>=1.15"] -# /// -import itertools -import os - -import torch - -import torch_sim as ts -from torch_sim.models.lennard_jones import LennardJonesModel -from torch_sim.units import MetalUnits as Units - - -# Set up the device and data type -device = torch.device("cuda" if torch.cuda.is_available() else "cpu") -dtype = torch.float32 - -# Number of steps to run -SMOKE_TEST = os.getenv("CI") is not None -N_steps = 100 if SMOKE_TEST else 2_000 - -# Set random seed and deterministic behavior for reproducibility -torch.manual_seed(42) -if torch.cuda.is_available(): - torch.cuda.manual_seed_all(42) -torch.backends.cudnn.deterministic = True -torch.backends.cudnn.benchmark = False - -# Set up the random number generator -generator = torch.Generator(device=device) -generator.manual_seed(42) # For reproducibility - -# Create face-centered cubic (FCC) Argon -# 5.26 Å is a typical lattice constant for Ar -a_len = 5.26 # Lattice constant - -# Generate base FCC unit cell positions (scaled by lattice constant) -base_positions = torch.tensor( - [ - [0.0, 0.0, 0.0], # Corner - [0.0, 0.5, 0.5], # Face centers - [0.5, 0.0, 0.5], - [0.5, 0.5, 0.0], - ], - device=device, - dtype=dtype, -) - -# Create 4x4x4 supercell of FCC Argon manually -positions = [] -for i, j, k in itertools.product(range(4), range(4), range(4)): - for base_pos in base_positions: - # Add unit cell position + offset for supercell - pos = base_pos + torch.tensor([i, j, k], device=device, dtype=dtype) - positions.append(pos) - -# Stack the positions into a tensor -positions = torch.stack(positions) - -# Scale by lattice constant -positions = positions * a_len - -# Create the cell tensor -cell = torch.tensor( - [[4 * a_len, 0, 0], [0, 4 * a_len, 0], [0, 0, 4 * a_len]], - device=device, - dtype=dtype, -) - -# Create the atomic numbers tensor (Argon = 18) -atomic_numbers = torch.full((positions.shape[0],), 18, device=device, dtype=torch.int) -# Create the masses tensor (Argon = 39.948 amu) -masses = torch.full((positions.shape[0],), 39.948, device=device, dtype=dtype) - -state = ts.SimState( - positions=positions, masses=masses, cell=cell, atomic_numbers=atomic_numbers, pbc=True -) -# Initialize the Lennard-Jones model -# Parameters: -# - sigma: distance at which potential is zero (3.405 Å for Ar) -# - epsilon: depth of potential well (0.0104 eV for Ar) -# - cutoff: distance beyond which interactions are ignored (typically 2.5*sigma) -model = LennardJonesModel( - use_neighbor_list=False, - sigma=3.405, - epsilon=0.0104, - cutoff=2.5 * 3.405, - device=device, - dtype=dtype, - compute_forces=True, - compute_stress=True, -) - -# Run initial simulation and get results -results = model(state) - -# Set up NVE simulation -# kT: initial temperature in metal units (80 K) -# dt: timestep in metal units (1 fs) -kT = torch.tensor(80 * Units.temperature, device=device, dtype=dtype) -dt = torch.tensor(0.001 * Units.time, device=device, dtype=dtype) - -# Initialize NVE integrator -state = ts.nve_init(state=state, model=model, kT=kT, seed=1) - -# Run NVE simulation for 1000 steps -for step in range(N_steps): - if step % 100 == 0: - # Calculate total energy (potential + kinetic) - total_energy = state.energy + ts.calc_kinetic_energy( - masses=state.masses, momenta=state.momenta - ) - print(f"{step=}: Total energy: {total_energy.item():.4f}") - - # Update state using NVE integrator - state = ts.nve_step(state=state, model=model, dt=dt) - -final_total_energy = state.energy + ts.calc_kinetic_energy( - masses=state.masses, momenta=state.momenta -) -print(f"Final total energy: {final_total_energy.item():.4f}") diff --git a/examples/old_scripts/3_Dynamics/3.2_MACE_NVE.py b/examples/old_scripts/3_Dynamics/3.2_MACE_NVE.py deleted file mode 100644 index 5a0fdd23..00000000 --- a/examples/old_scripts/3_Dynamics/3.2_MACE_NVE.py +++ /dev/null @@ -1,90 +0,0 @@ -"""NVE simulation with MACE.""" - -# /// script -# dependencies = ["mace-torch>=0.3.12"] -# /// -import os -import time - -import torch -from ase.build import bulk -from mace.calculators.foundations_models import mace_mp - -import torch_sim as ts -from torch_sim.models.mace import MaceModel, MaceUrls -from torch_sim.quantities import calc_kinetic_energy -from torch_sim.units import MetalUnits as Units - - -# Set device and data type -device = torch.device("cuda" if torch.cuda.is_available() else "cpu") -dtype = torch.float32 - -# Option 1: Load the raw model from the downloaded model -loaded_model = mace_mp( - model=MaceUrls.mace_mpa_medium, - return_raw_model=True, - default_dtype=str(dtype).removeprefix("torch."), - device=str(device), -) - -# Option 2: Load from local file (comment out Option 1 to use this) -# loaded_model = torch.load("path/to/model.pt", map_location=device) - -# Number of steps to run -SMOKE_TEST = os.getenv("CI") is not None -N_steps = 20 if SMOKE_TEST else 2_000 - -# Create diamond cubic Silicon -si_dc = bulk("Si", "diamond", a=5.43, cubic=True).repeat((2, 2, 2)) - -# Prepare input tensors -positions = torch.tensor(si_dc.positions, device=device, dtype=dtype) -cell = torch.tensor(si_dc.cell.array, device=device, dtype=dtype) -atomic_numbers = torch.tensor(si_dc.get_atomic_numbers(), device=device, dtype=torch.int) -masses = torch.tensor(si_dc.get_masses(), device=device, dtype=dtype) - -# Initialize the MACE model -model = MaceModel( - model=loaded_model, - device=device, - compute_forces=True, - compute_stress=False, - dtype=dtype, - enable_cueq=False, -) - -state = ts.SimState( - positions=positions, masses=masses, cell=cell, atomic_numbers=atomic_numbers, pbc=True -) - -# Run initial inference -results = model(state) - -# Setup NVE MD simulation parameters -kT = ( - torch.tensor(1000, device=device, dtype=dtype) * Units.temperature -) # Initial temperature (1000 K) -dt = torch.tensor(0.002 * Units.time, device=device, dtype=dtype) # Timestep (2 fs) - - -# Initialize NVE integrator -state = ts.nve_init(state=state, model=model, kT=kT, seed=1) - -# Run MD simulation -print("\nStarting NVE molecular dynamics simulation...") -start_time = time.perf_counter() -for step in range(N_steps): - total_energy = state.energy + calc_kinetic_energy( - masses=state.masses, momenta=state.momenta, system_idx=state.system_idx - ) - if step % 10 == 0: - print(f"Step {step}: Total energy: {total_energy.item():.4f} eV") - state = ts.nve_step(state=state, model=model, dt=dt) -end_time = time.perf_counter() - -# Report simulation results -print("\nSimulation complete!") -print(f"Time taken: {end_time - start_time:.2f} seconds") -print(f"Average time per step: {(end_time - start_time) / N_steps:.4f} seconds") -print(f"Final total energy: {total_energy.item()} eV") diff --git a/examples/old_scripts/3_Dynamics/3.3_MACE_NVE_cueq.py b/examples/old_scripts/3_Dynamics/3.3_MACE_NVE_cueq.py deleted file mode 100644 index 43414b9a..00000000 --- a/examples/old_scripts/3_Dynamics/3.3_MACE_NVE_cueq.py +++ /dev/null @@ -1,81 +0,0 @@ -"""NVE simulation with MACE and cuEquivariance enabled.""" - -# /// script -# dependencies = ["mace-torch>=0.3.12"] -# /// -import os -import time - -import torch -from ase.build import bulk -from mace.calculators.foundations_models import mace_mp - -import torch_sim as ts -from torch_sim.integrators import nve_init, nve_step -from torch_sim.models.mace import MaceModel, MaceUrls -from torch_sim.units import MetalUnits as Units - - -# Set device and data type -device = torch.device("cuda" if torch.cuda.is_available() else "cpu") -dtype = torch.float32 - -# Option 1: Load the raw model from the downloaded model -loaded_model = mace_mp( - model=MaceUrls.mace_mpa_medium, - return_raw_model=True, - default_dtype=str(dtype).removeprefix("torch."), - device=str(device), -) - -# Option 2: Load from local file (comment out Option 1 to use this) -# loaded_model = torch.load("path/to/model.pt", map_location=device) - -# Number of steps to run -SMOKE_TEST = os.getenv("CI") is not None -N_steps = 20 if SMOKE_TEST else 2_000 - -# Create diamond cubic Silicon -si_dc = bulk("Si", "diamond", a=5.43, cubic=True).repeat((2, 2, 2)) - -# Initialize the MACE model -model = MaceModel( - model=loaded_model, - device=device, - compute_forces=True, - compute_stress=False, - dtype=dtype, - enable_cueq=torch.cuda.is_available(), -) -state = ts.io.atoms_to_state(si_dc, device=device, dtype=dtype) - -# Run initial inference -results = model(state) - -# Setup NVE MD simulation parameters -kT = ( - torch.tensor(1000, device=device, dtype=dtype) * Units.temperature -) # Initial temperature (1000 K) -dt = torch.tensor(0.002 * Units.time, device=device, dtype=dtype) # Timestep (2 fs) - - -# Initialize NVE integrator -state = nve_init(state=state, model=model, kT=kT, seed=1) - -# Run MD simulation -print("\nStarting NVE molecular dynamics simulation...") -start_time = time.perf_counter() -for step in range(N_steps): - total_energy = state.energy + ts.calc_kinetic_energy( - masses=state.masses, momenta=state.momenta, system_idx=state.system_idx - ) - if step % 10 == 0: - print(f"Step {step}: Total energy: {total_energy.item():.4f} eV") - state = nve_step(state=state, model=model, dt=dt) -end_time = time.perf_counter() - -# Report simulation results -print("\nSimulation complete!") -print(f"Time taken: {end_time - start_time:.2f} seconds") -print(f"Average time per step: {(end_time - start_time) / N_steps:.4f} seconds") -print(f"Final total energy: {total_energy.item()} eV") diff --git a/examples/old_scripts/3_Dynamics/3.4_MACE_NVT_Langevin.py b/examples/old_scripts/3_Dynamics/3.4_MACE_NVT_Langevin.py deleted file mode 100644 index 575905bb..00000000 --- a/examples/old_scripts/3_Dynamics/3.4_MACE_NVT_Langevin.py +++ /dev/null @@ -1,86 +0,0 @@ -"""MACE NVT Langevin dynamics.""" - -# /// script -# dependencies = ["mace-torch>=0.3.12"] -# /// -import os - -import torch -from ase.build import bulk -from mace.calculators.foundations_models import mace_mp - -import torch_sim as ts -from torch_sim.integrators import nvt_langevin_init, nvt_langevin_step -from torch_sim.models.mace import MaceModel, MaceUrls -from torch_sim.units import MetalUnits as Units - - -# Set device and data type -device = torch.device("cuda" if torch.cuda.is_available() else "cpu") -dtype = torch.float32 - -# Option 1: Load the raw model from the downloaded model -loaded_model = mace_mp( - model=MaceUrls.mace_mpa_medium, - return_raw_model=True, - default_dtype=str(dtype).removeprefix("torch."), - device=str(device), -) - -# Option 2: Load from local file (comment out Option 1 to use this) -# loaded_model = torch.load("path/to/model.pt", map_location=device) - -# Number of steps to run -SMOKE_TEST = os.getenv("CI") is not None -N_steps = 20 if SMOKE_TEST else 2_000 - -# Create diamond cubic Silicon -si_dc = bulk("Si", "diamond", a=5.43, cubic=True).repeat((2, 2, 2)) - -# Prepare input tensors -positions = torch.tensor(si_dc.positions, device=device, dtype=dtype) -cell = torch.tensor(si_dc.cell.array, device=device, dtype=dtype) -atomic_numbers = torch.tensor(si_dc.get_atomic_numbers(), device=device, dtype=torch.int) -masses = torch.tensor(si_dc.get_masses(), device=device, dtype=dtype) - -# Initialize the MACE model -model = MaceModel( - model=loaded_model, - device=device, - compute_forces=True, - compute_stress=False, - dtype=dtype, - enable_cueq=False, -) - -state = ts.SimState( - positions=positions, masses=masses, cell=cell, atomic_numbers=atomic_numbers, pbc=True -) - -dt = torch.tensor(0.002 * Units.time, device=device, dtype=dtype) # Timestep (2 fs) -kT = ( - torch.tensor(1000, device=device, dtype=dtype) * Units.temperature -) # Initial temperature (1000 K) -gamma = torch.tensor( - 10 / Units.time, device=device, dtype=dtype -) # Langevin friction coefficient (ps^-1) - -# Initialize NVT Langevin integrator -state = nvt_langevin_init(model=model, state=state, kT=kT, seed=1) - -for step in range(N_steps): - if step % 10 == 0: - temp = ( - ts.calc_kT( - masses=state.masses, momenta=state.momenta, system_idx=state.system_idx - ) - / Units.temperature - ) - print(f"{step=}: Temperature: {temp.item():.4f}") - state = nvt_langevin_step(state=state, model=model, dt=dt, kT=kT, gamma=gamma) - -final_temp = ( - ts.calc_kT(masses=state.masses, momenta=state.momenta, system_idx=state.system_idx) - / Units.temperature -) -print(f"Final temperature: {final_temp.item():.4f}") diff --git a/examples/old_scripts/3_Dynamics/3.5_MACE_NVT_Nose_Hoover.py b/examples/old_scripts/3_Dynamics/3.5_MACE_NVT_Nose_Hoover.py deleted file mode 100644 index 903d7dcb..00000000 --- a/examples/old_scripts/3_Dynamics/3.5_MACE_NVT_Nose_Hoover.py +++ /dev/null @@ -1,77 +0,0 @@ -"""NVT simulation with MACE and Nose-Hoover thermostat.""" - -# /// script -# dependencies = ["mace-torch>=0.3.12"] -# /// -import os - -import torch -from ase.build import bulk -from mace.calculators.foundations_models import mace_mp - -import torch_sim as ts -from torch_sim.models.mace import MaceModel, MaceUrls -from torch_sim.units import MetalUnits as Units - - -# Set device and data type -device = torch.device("cuda" if torch.cuda.is_available() else "cpu") -dtype = torch.float32 - -# Option 1: Load the raw model from the downloaded model -loaded_model = mace_mp( - model=MaceUrls.mace_mpa_medium, - return_raw_model=True, - default_dtype=str(dtype).removeprefix("torch."), - device=str(device), -) - -# Option 2: Load from local file (comment out Option 1 to use this) -# loaded_model = torch.load("path/to/model.pt", map_location=device) - -# Number of steps to run -SMOKE_TEST = os.getenv("CI") is not None -N_steps = 20 if SMOKE_TEST else 2_000 - -# Create diamond cubic Silicon -si_dc = bulk("Si", "diamond", a=5.43, cubic=True).repeat((2, 2, 2)) - -# Initialize the MACE model -model = MaceModel( - model=loaded_model, - device=device, - compute_forces=True, - compute_stress=False, - dtype=dtype, - enable_cueq=False, -) -state = ts.io.atoms_to_state(si_dc, device=device, dtype=dtype) - -# Run initial inference -results = model(state) - -dt = torch.tensor(0.002 * Units.time, device=device, dtype=dtype) # Timestep (2 fs) -kT = ( - torch.tensor(1000, device=device, dtype=dtype) * Units.temperature -) # Initial temperature (1000 K) - - -state = ts.nvt_nose_hoover_init(state=state, model=model, kT=kT, dt=dt) - -for step in range(N_steps): - if step % 10 == 0: - temp = ( - ts.calc_kT( - masses=state.masses, momenta=state.momenta, system_idx=state.system_idx - ) - / Units.temperature - ) - invariant = float(ts.nvt_nose_hoover_invariant(state, kT=kT)) - print(f"{step=}: Temperature: {temp.item():.4f}: {invariant=:.4f}") - state = ts.nvt_nose_hoover_step(state=state, model=model, dt=dt, kT=kT) - -final_temp = ( - ts.calc_kT(masses=state.masses, momenta=state.momenta, system_idx=state.system_idx) - / Units.temperature -) -print(f"Final temperature: {final_temp.item():.4f}") diff --git a/examples/old_scripts/3_Dynamics/3.6_MACE_NVT_Nose_Hoover_temp_profile.py b/examples/old_scripts/3_Dynamics/3.6_MACE_NVT_Nose_Hoover_temp_profile.py deleted file mode 100644 index 0b58df35..00000000 --- a/examples/old_scripts/3_Dynamics/3.6_MACE_NVT_Nose_Hoover_temp_profile.py +++ /dev/null @@ -1,193 +0,0 @@ -"""Example NVT Nose Hoover MD simulation of random alloy using MACE model with -temperature profile. -""" - -# /// script -# dependencies = [ -# "mace-torch>=0.3.12", -# "plotly>=6", -# "kaleido", -# ] -# /// -import os - -import numpy as np -import torch -from ase.build import bulk -from mace.calculators.foundations_models import mace_mp -from plotly.subplots import make_subplots - -import torch_sim as ts -from torch_sim.models.mace import MaceModel, MaceUrls -from torch_sim.units import MetalUnits as Units - - -def get_kT( - step: int, - n_steps_initial: int, - n_steps_ramp_up: int, - n_steps_melt: int, - n_steps_ramp_down: int, - n_steps_anneal: int, - melt_temp: float, - cool_temp: float, - anneal_temp: float, - device: torch.device, -) -> torch.Tensor: - """Determine target kT based on current simulation step. - Temperature profile: - 300K (initial) → ramp to 3_000K → hold at 3_000K → quench to 300K → hold at 300K. - """ - if step < n_steps_initial: - # Initial equilibration at cool temperature - return torch.tensor(cool_temp, device=device) - if step < (n_steps_initial + n_steps_ramp_up): - # Linear ramp from cool_temp to melt_temp - progress = (step - n_steps_initial) / n_steps_ramp_up - current_kT = cool_temp + (melt_temp - cool_temp) * progress - return torch.tensor(current_kT, device=device) - if step < (n_steps_initial + n_steps_ramp_up + n_steps_melt): - # Hold at melting temperature - return torch.tensor(melt_temp, device=device) - if step < (n_steps_initial + n_steps_ramp_up + n_steps_melt + n_steps_ramp_down): - # Linear cooling from melt_temp to cool_temp - progress = ( - step - (n_steps_initial + n_steps_ramp_up + n_steps_melt) - ) / n_steps_ramp_down - current_kT = melt_temp - (melt_temp - cool_temp) * progress - return torch.tensor(current_kT, device=device) - if step < ( - n_steps_initial - + n_steps_ramp_up - + n_steps_melt - + n_steps_ramp_down - + n_steps_anneal - ): - # Hold at annealing temperature - return torch.tensor(anneal_temp, device=device) - # Hold at annealing temperature - return torch.tensor(anneal_temp, device=device) - - -# Set device and data type -device = torch.device("cuda" if torch.cuda.is_available() else "cpu") -dtype = torch.float32 - -# Model configuration -# Option 1: Load from URL (uncomment to use) -loaded_model = mace_mp( - model=MaceUrls.mace_mpa_medium, - return_raw_model=True, - default_dtype=str(dtype).removeprefix("torch."), - device=str(device), -) - -# Option 2: Load from local file -# MODEL_PATH = "../../../checkpoints/MACE/mace-mpa-0-medium.model" -# loaded_model = torch.load(MODEL_PATH, map_location=device, weights_only=False) - -# Temperature profile settings -init_temp = 300 -melting_temp = 1000 -cooling_temp = 300 -annealing_temp = 300 - -# Step counts for different phases -SMOKE_TEST = os.getenv("CI") is not None -n_steps_initial = 20 if SMOKE_TEST else 200 -n_steps_ramp_up = 20 if SMOKE_TEST else 200 -n_steps_melt = 20 if SMOKE_TEST else 200 -n_steps_ramp_down = 20 if SMOKE_TEST else 200 -n_steps_anneal = 20 if SMOKE_TEST else 200 - -n_steps = ( - n_steps_initial + n_steps_ramp_up + n_steps_melt + n_steps_ramp_down + n_steps_anneal -) - -# Create a random alloy system -# Define possible species and their probabilities -species = ["Cu", "Mn", "Fe"] -probabilities = [0.33, 0.33, 0.34] - -# Create base FCC structure with Cu (using Cu lattice parameter) -fcc_lattice = bulk("Cu", "fcc", a=3.61, cubic=True).repeat((2, 2, 2)) - -# Randomly assign species -random_species = np.random.default_rng(seed=0).choice( - species, size=len(fcc_lattice), p=probabilities -) -fcc_lattice.set_chemical_symbols(random_species) - -# Initialize the MACE model -model = MaceModel( - model=loaded_model, - device=device, - compute_forces=True, - compute_stress=False, - dtype=dtype, - enable_cueq=False, -) -state = ts.io.atoms_to_state(fcc_lattice, device=device, dtype=dtype) - -# Run initial inference -results = model(state) - -# Set up simulation parameters -dt = torch.tensor(0.002 * Units.time, device=device, dtype=dtype) # Timestep (2 fs) -kT = ( - torch.tensor(init_temp, device=device, dtype=dtype) * Units.temperature -) # Initial temperature (K) - -state = ts.nvt_nose_hoover_init(state=state, model=model, kT=kT, dt=dt, seed=1) - -# Run simulation with temperature profile -actual_temps = np.zeros(n_steps) -expected_temps = np.zeros(n_steps) - -for step in range(n_steps): - # Get target temperature for current step - current_kT = get_kT( # noqa: N816 - step=step, - n_steps_initial=n_steps_initial, - n_steps_ramp_up=n_steps_ramp_up, - n_steps_melt=n_steps_melt, - n_steps_ramp_down=n_steps_ramp_down, - n_steps_anneal=n_steps_anneal, - melt_temp=melting_temp, - cool_temp=cooling_temp, - anneal_temp=annealing_temp, - device=device, - ) - - # Calculate current temperature and save data - temp = ( - ts.calc_kT( - masses=state.masses, momenta=state.momenta, system_idx=state.system_idx - ) - / Units.temperature - ) - actual_temps[step] = temp - expected_temps[step] = current_kT - - # Calculate invariant and progress report - invariant = float( - ts.nvt_nose_hoover_invariant(state, kT=current_kT * Units.temperature) - ) - print(f"{step=}: Temperature: {temp.item():.4f}: {invariant=:.4f}") - - # Update simulation state - state = ts.nvt_nose_hoover_step( - state=state, model=model, dt=dt, kT=current_kT * Units.temperature - ) - -# Visualize temperature profile -fig = make_subplots() -fig.add_scatter( - x=np.arange(n_steps) * 0.002, y=actual_temps, name="Simulated Temperature" -) -fig.add_scatter( - x=np.arange(n_steps) * 0.002, y=expected_temps, name="Desired Temperature" -) -fig.layout.xaxis.title = "time (ps)" -fig.layout.yaxis.title = "Temperature (K)" -fig.write_image("nvt_visualization_temperature.pdf") diff --git a/examples/old_scripts/3_Dynamics/3.7_Lennard_Jones_NPT_Nose_Hoover.py b/examples/old_scripts/3_Dynamics/3.7_Lennard_Jones_NPT_Nose_Hoover.py deleted file mode 100644 index 03c61877..00000000 --- a/examples/old_scripts/3_Dynamics/3.7_Lennard_Jones_NPT_Nose_Hoover.py +++ /dev/null @@ -1,162 +0,0 @@ -"""Lennard-Jones simulation in NPT ensemble using Nose-Hoover chain.""" - -# /// script -# dependencies = ["scipy>=1.15"] -# /// -import itertools -import os - -import torch - -import torch_sim as ts -from torch_sim.models.lennard_jones import LennardJonesModel -from torch_sim.units import MetalUnits as Units - - -# Set up the device and data type -device = torch.device("cuda" if torch.cuda.is_available() else "cpu") -dtype = torch.float32 - -# Set random seed and deterministic behavior for reproducibility -torch.manual_seed(42) -if torch.cuda.is_available(): - torch.cuda.manual_seed_all(42) -torch.backends.cudnn.deterministic = True -torch.backends.cudnn.benchmark = False - -# Set up the random number generator -generator = torch.Generator(device=device) -generator.manual_seed(42) # For reproducibility - -# Number of steps to run -SMOKE_TEST = os.getenv("CI") is not None -N_steps = 100 if SMOKE_TEST else 10_000 - -# Create face-centered cubic (FCC) Argon -# 5.26 Å is a typical lattice constant for Ar -a_len = 5.26 # Lattice constant - -# Generate base FCC unit cell positions (scaled by lattice constant) -base_positions = torch.tensor( - [ - [0.0, 0.0, 0.0], # Corner - [0.0, 0.5, 0.5], # Face centers - [0.5, 0.0, 0.5], - [0.5, 0.5, 0.0], - ], - device=device, - dtype=dtype, -) - -# Create 4x4x4 supercell of FCC Argon manually -positions = [] -for i, j, k in itertools.product(range(4), range(4), range(4)): - for base_pos in base_positions: - # Add unit cell position + offset for supercell - pos = base_pos + torch.tensor([i, j, k], device=device, dtype=dtype) - positions.append(pos) - -# Stack the positions into a tensor -positions = torch.stack(positions) - -# Scale by lattice constant -positions = positions * a_len - -# Create the cell tensor -cell = torch.tensor( - [[4 * a_len, 0, 0], [0, 4 * a_len, 0], [0, 0, 4 * a_len]], device=device, dtype=dtype -) - -# Create the atomic numbers tensor (Argon = 18) -atomic_numbers = torch.full((positions.shape[0],), 18, device=device, dtype=torch.int) -# Create the masses tensor (Argon = 39.948 amu) -masses = torch.full((positions.shape[0],), 39.948, device=device, dtype=dtype) - -# Initialize the Lennard-Jones model -# Parameters: -# - sigma: distance at which potential is zero (3.405 Å for Ar) -# - epsilon: depth of potential well (0.0104 eV for Ar) -# - cutoff: distance beyond which interactions are ignored (typically 2.5*sigma) -model = LennardJonesModel( - use_neighbor_list=False, - sigma=3.405, - epsilon=0.0104, - cutoff=2.5 * 3.405, - device=device, - dtype=dtype, - compute_forces=True, - compute_stress=True, -) -state = ts.SimState( - positions=positions, - masses=masses, - cell=cell.unsqueeze(0), - atomic_numbers=atomic_numbers, - pbc=True, -) -# Run initial simulation and get results -results = model(state) - -dt = torch.tensor( - 0.001 * Units.time, device=device, dtype=dtype -) # Time step (1 fs) # Default time step is ps in Metal units -kT = torch.tensor( - 200 * Units.temperature, device=device, dtype=dtype -) # Temperature (200 K) -target_pressure = ( - torch.tensor(10, device=device, dtype=dtype) * Units.pressure -) # Target pressure (10 bar, using Metal units) -state = ts.npt_nose_hoover_init( - state=state, - model=model, - dt=dt, - kT=kT, - chain_length=3, # Chain length - chain_steps=1, - sy_steps=1, -) - -# Run the simulation -for step in range(N_steps): - if step % 50 == 0: - temp = ( - ts.calc_kT( - masses=state.masses, momenta=state.momenta, system_idx=state.system_idx - ) - / Units.temperature - ) - invariant = float( - ts.npt_nose_hoover_invariant(state, kT=kT, external_pressure=target_pressure) - ) - e_kin = ts.calc_kinetic_energy( - masses=state.masses, momenta=state.momenta, system_idx=state.system_idx - ) - pressure = ts.get_pressure( - model(state)["stress"], e_kin, torch.det(state.current_cell) - ) - pressure = float(pressure) / Units.pressure - xx, yy, zz = state.cell[..., 0, 0], state.cell[..., 1, 1], state.cell[..., 2, 2] - print( - f"{step=}: Temperature: {temp.item():.4f}, " - f"{invariant=:.4f}, {pressure=:.4f}, " - f"cell xx yy zz: {xx.item():.4f}, {yy.item():.4f}, {zz.item():.4f}" - ) - state = ts.npt_nose_hoover_step( - state=state, model=model, dt=dt, kT=kT, external_pressure=target_pressure - ) - -temp = ( - ts.calc_kT(masses=state.masses, momenta=state.momenta, system_idx=state.system_idx) - / Units.temperature -) -print(f"Final temperature: {temp.item():.4f}") - -pressure = ts.get_pressure( - model(state)["stress"], - ts.calc_kinetic_energy( - masses=state.masses, momenta=state.momenta, system_idx=state.system_idx - ), - torch.det(state.current_cell), -) -pressure = pressure.item() / Units.pressure -print(f"Final {pressure=:.4f}") diff --git a/examples/old_scripts/3_Dynamics/3.8_MACE_NPT_Nose_Hoover.py b/examples/old_scripts/3_Dynamics/3.8_MACE_NPT_Nose_Hoover.py deleted file mode 100644 index d5505ab7..00000000 --- a/examples/old_scripts/3_Dynamics/3.8_MACE_NPT_Nose_Hoover.py +++ /dev/null @@ -1,129 +0,0 @@ -"""NPT simulation with MACE and Nose-Hoover thermostat.""" - -# /// script -# dependencies = ["mace-torch>=0.3.12"] -# /// -import os - -import torch -from ase.build import bulk -from mace.calculators.foundations_models import mace_mp - -import torch_sim as ts -from torch_sim.models.mace import MaceModel, MaceUrls -from torch_sim.units import MetalUnits as Units - - -# Set device and data type -device = torch.device("cuda" if torch.cuda.is_available() else "cpu") -dtype = torch.float32 - -# Option 1: Load the raw model from the downloaded model -loaded_model = mace_mp( - model=MaceUrls.mace_mpa_medium, - return_raw_model=True, - default_dtype=str(dtype).removeprefix("torch."), - device=str(device), -) - -# Option 2: Load from local file (comment out Option 1 to use this) -# loaded_model = torch.load("path/to/model.pt", map_location=device) - -# Create diamond cubic Silicon -si_dc = bulk("Si", "diamond", a=5.43, cubic=True).repeat((2, 2, 2)) - -# Initialize the MACE model -model = MaceModel( - model=loaded_model, - device=device, - compute_forces=True, - compute_stress=True, - dtype=dtype, - enable_cueq=False, -) -state = ts.io.atoms_to_state(si_dc, device=device, dtype=dtype) - -# Run initial inference -results = model(state) - -SMOKE_TEST = os.getenv("CI") is not None -N_steps_nvt = 20 if SMOKE_TEST else 2_000 -N_steps_npt = 20 if SMOKE_TEST else 2_000 -dt = 0.001 * Units.time # Time step (1 fs) -kT = ( - torch.tensor(300, device=device, dtype=dtype) * Units.temperature -) # Initial temperature (300 K) -target_pressure = torch.tensor( - 0.0 * Units.pressure, device=device, dtype=dtype -) # Target pressure (0 bar) - -state = ts.npt_nose_hoover_init(state=state, model=model, kT=kT, dt=torch.tensor(dt)) - -for step in range(N_steps_nvt): - if step % 10 == 0: - temp = ( - ts.calc_kT( - masses=state.masses, momenta=state.momenta, system_idx=state.system_idx - ) - / Units.temperature - ) - invariant = float( - ts.npt_nose_hoover_invariant(state, kT=kT, external_pressure=target_pressure) - ) - print(f"{step=}: Temperature: {temp.item():.4f}: {invariant=:.4f}, ") - state = ts.npt_nose_hoover_step( - state=state, - model=model, - dt=torch.tensor(dt), - kT=kT, - external_pressure=target_pressure, - ) - -state = ts.npt_nose_hoover_init(state=state, model=model, kT=kT, dt=torch.tensor(dt)) - -for step in range(N_steps_npt): - if step % 10 == 0: - temp = ( - ts.calc_kT( - masses=state.masses, momenta=state.momenta, system_idx=state.system_idx - ) - / Units.temperature - ) - invariant = float( - ts.npt_nose_hoover_invariant(state, kT=kT, external_pressure=target_pressure) - ) - stress = model(state)["stress"] - volume = torch.det(state.current_cell) - e_kin = ts.calc_kinetic_energy( - masses=state.masses, momenta=state.momenta, system_idx=state.system_idx - ) - pressure = float(ts.get_pressure(stress, e_kin, volume)) - xx, yy, zz = torch.diag(state.current_cell[0]) - print( - f"{step=}: Temperature: {temp.item():.4f}: {invariant=:.4f}, " - f"{pressure=:.4f}, " - f"cell xx yy zz: {xx.item():.4f}, {yy.item():.4f}, {zz.item():.4f}" - ) - state = ts.npt_nose_hoover_step( - state=state, - model=model, - dt=torch.tensor(dt), - kT=kT, - external_pressure=target_pressure, - ) - -final_temp = ( - ts.calc_kT(masses=state.masses, momenta=state.momenta, system_idx=state.system_idx) - / Units.temperature -) -print(f"Final temperature: {final_temp.item():.4f}") -final_stress = model(state)["stress"] -final_volume = torch.det(state.current_cell) -final_pressure = ts.get_pressure( - final_stress, - ts.calc_kinetic_energy( - masses=state.masses, momenta=state.momenta, system_idx=state.system_idx - ), - final_volume, -) -print(f"Final pressure: {final_pressure.item():.4f}") diff --git a/examples/old_scripts/3_Dynamics/3.9_MACE_NVT_staggered_stress.py b/examples/old_scripts/3_Dynamics/3.9_MACE_NVT_staggered_stress.py deleted file mode 100644 index 2446dbc8..00000000 --- a/examples/old_scripts/3_Dynamics/3.9_MACE_NVT_staggered_stress.py +++ /dev/null @@ -1,79 +0,0 @@ -"""MACE NVT simulation with staggered stress calculation.""" - -# /// script -# dependencies = ["mace-torch>=0.3.12"] -# /// -import os - -import torch -from ase.build import bulk -from mace.calculators.foundations_models import mace_mp - -import torch_sim as ts -from torch_sim.models.mace import MaceModel, MaceUrls -from torch_sim.quantities import calc_kT -from torch_sim.units import MetalUnits as Units - - -# Set device and data type -device = torch.device("cuda" if torch.cuda.is_available() else "cpu") -dtype = torch.float32 - -# Option 1: Load the raw model from the downloaded model -loaded_model = mace_mp( - model=MaceUrls.mace_mpa_medium, - return_raw_model=True, - default_dtype=str(dtype).removeprefix("torch."), - device=str(device), -) - -# Option 2: Load from local file (comment out Option 1 to use this) -# loaded_model = torch.load("path/to/model.pt", map_location=device) - -# Number of steps to run -SMOKE_TEST = os.getenv("CI") is not None -N_steps = 20 if SMOKE_TEST else 2_000 - -# Create diamond cubic Silicon -si_dc = bulk("Si", "diamond", a=5.43, cubic=True).repeat((2, 2, 2)) - -# Initialize the MACE model -model = MaceModel( - model=loaded_model, - device=device, - compute_forces=True, - compute_stress=True, - dtype=dtype, - enable_cueq=False, -) - -state = ts.io.atoms_to_state(si_dc, device=device, dtype=dtype) - -dt = 0.002 * Units.time # Timestep (2 fs) -kT = ( - torch.tensor(1000, device=device, dtype=dtype) * Units.temperature -) # Initial temperature (1000 K) - -state = ts.nvt_langevin_init(state=state, model=model, kT=kT) - -stress = torch.zeros(N_steps // 10, 3, 3, device=device, dtype=dtype) -for step in range(N_steps): - temp = ( - calc_kT(masses=state.masses, momenta=state.momenta, system_idx=state.system_idx) - / Units.temperature - ) - - # Calculate kinetic energy: KE = 0.5 * sum(p^2 / m) - kinetic_energy = 0.5 * torch.sum( - torch.pow(state.momenta, 2) / state.masses.unsqueeze(-1) - ) - # Total energy = kinetic + potential - invariant = float(kinetic_energy + state.energy) - - print(f"{step=}: Temperature: {temp.item():.4f}: {invariant=:.4f}") - state = ts.nvt_langevin_step(state=state, model=model, dt=torch.tensor(dt), kT=kT) - if step % 10 == 0: - results = model(state) - stress[step // 10] = results["stress"] - -print(f"Stress: {stress} eV/Å^3") diff --git a/examples/old_scripts/4_High_level_api/4.1_high_level_api.py b/examples/old_scripts/4_High_level_api/4.1_high_level_api.py deleted file mode 100644 index 5d555579..00000000 --- a/examples/old_scripts/4_High_level_api/4.1_high_level_api.py +++ /dev/null @@ -1,203 +0,0 @@ -"""Basic Lennard-Jones and MACE relaxation examples with batching, custom convergence -criteria, and logging. -""" - -# /// script -# dependencies = ["mace-torch>=0.3.12", "pymatgen>=2025.2.18"] -# /// -import os - -import numpy as np -import torch -from ase.build import bulk -from mace.calculators.foundations_models import mace_mp -from pymatgen.core import Structure - -import torch_sim as ts -from torch_sim.models.lennard_jones import LennardJonesModel -from torch_sim.models.mace import MaceModel -from torch_sim.trajectory import TorchSimTrajectory, TrajectoryReporter -from torch_sim.units import MetalUnits - - -SMOKE_TEST = os.getenv("CI") is not None - -lj_model = LennardJonesModel( - sigma=2.0, # Å, typical for Si-Si interaction - epsilon=0.1, # eV, typical for Si-Si interaction - device=torch.device("cpu"), - dtype=torch.float64, -) - -si_atoms = bulk("Si", "fcc", a=5.43, cubic=True) - -final_state = ts.integrate( - system=si_atoms, - model=lj_model, - integrator=ts.Integrator.nvt_langevin, - n_steps=100 if SMOKE_TEST else 1000, - temperature=2000, - timestep=0.002, -) -final_atoms = ts.io.state_to_atoms(final_state) - -trajectory_file = "tmp/lj_trajectory.h5md" -# report potential energy every 10 steps and kinetic energy every 20 steps -prop_calculators = { - 10: {"potential_energy": lambda state: state.energy}, - 20: { - "kinetic_energy": lambda state: ts.calc_kinetic_energy( - momenta=state.momenta, masses=state.masses - ) - }, -} - -reporter = TrajectoryReporter( - trajectory_file, - # report state every 10 steps - state_frequency=10, - prop_calculators=prop_calculators, -) - -final_state = ts.integrate( - system=si_atoms, - model=lj_model, - integrator=ts.Integrator.nvt_langevin, - n_steps=100 if SMOKE_TEST else 1000, - temperature=2000, - timestep=0.002, - trajectory_reporter=reporter, -) - -# Check energy fluctuations -with TorchSimTrajectory(trajectory_file) as traj: - kinetic_energies = traj.get_array("kinetic_energy") - potential_energies = traj.get_array("potential_energy") - final_energy = potential_energies[-1] - - final_atoms = traj.get_atoms(-1) - - -### basic mace example - -# cuda if available -device = torch.device("cuda" if torch.cuda.is_available() else "cpu") - - -mace = mace_mp(model="small", return_raw_model=True) -mace_model = MaceModel( - model=mace, - device=device, - dtype=torch.float64, - compute_forces=True, -) - -reporter = TrajectoryReporter( - trajectory_file, - # report state every 10 steps - state_frequency=10, - prop_calculators=prop_calculators, -) - -final_state = ts.integrate( - system=si_atoms, - model=mace_model, - integrator=ts.Integrator.nvt_langevin, - n_steps=100 if SMOKE_TEST else 1000, - temperature=2000, - timestep=0.002, - trajectory_reporter=reporter, -) -final_atoms = ts.io.state_to_atoms(final_state) - - -### basic mace example with batching - -fe_atoms = bulk("Fe", "fcc", a=5.26, cubic=True) -fe_atoms_supercell = fe_atoms.repeat([2, 2, 2]) -si_atoms_supercell = si_atoms.repeat([2, 2, 2]) - -final_state = ts.integrate( - system=[si_atoms, fe_atoms, si_atoms_supercell, fe_atoms_supercell], - model=mace_model, - integrator=ts.Integrator.nvt_langevin, - n_steps=100 if SMOKE_TEST else 1000, - temperature=2000, - timestep=0.002, -) -final_atoms = ts.io.state_to_atoms(final_state) -final_fe_atoms_supercell = final_atoms[3] - - -### basic mace example with batching and reporting - -systems = (si_atoms, fe_atoms, si_atoms_supercell, fe_atoms_supercell) - -filenames = [f"tmp/batch_traj_{i}.h5md" for i in range(len(systems))] -batch_reporter = TrajectoryReporter( - filenames, - state_frequency=100, - prop_calculators=prop_calculators, -) -final_state = ts.integrate( - system=systems, - model=mace_model, - integrator=ts.Integrator.nvt_langevin, - n_steps=100 if SMOKE_TEST else 1000, - temperature=2000, - timestep=0.002, - trajectory_reporter=batch_reporter, -) - -final_energies_per_atom = [] -for filename in filenames: - with TorchSimTrajectory(filename) as traj: - final_energy = traj.get_array("potential_energy")[-1] - final_energies_per_atom.append(final_energy / len(traj.get_atoms(-1))) - - -final_state = ts.optimize( - system=systems, - model=mace_model, - optimizer=ts.Optimizer.fire, - max_steps=10 if SMOKE_TEST else 1000, - init_kwargs=dict(cell_filter=ts.CellFilter.unit), -) - -rng = np.random.default_rng() -for system in systems: - system.positions += rng.random(system.positions.shape) * 0.01 - -final_state = ts.optimize( - system=systems, - model=mace_model, - optimizer=ts.Optimizer.fire, - convergence_fn=lambda state, last_energy: last_energy - state.energy - < 1e-6 * MetalUnits.energy, - max_steps=10 if SMOKE_TEST else 1000, - init_kwargs=dict(cell_filter=ts.CellFilter.unit), -) - - -lattice = [[5.43, 0, 0], [0, 5.43, 0], [0, 0, 5.43]] -species = ["Si"] * 8 -coords = [ - [0.0, 0.0, 0.0], - [0.25, 0.25, 0.25], - [0.0, 0.5, 0.5], - [0.25, 0.75, 0.75], - [0.5, 0.0, 0.5], - [0.75, 0.25, 0.75], - [0.5, 0.5, 0.0], - [0.75, 0.75, 0.25], -] -structure = Structure(lattice, species, coords) -final_state = ts.integrate( - system=structure, - model=lj_model, - integrator=ts.Integrator.nvt_langevin, - n_steps=100 if SMOKE_TEST else 1000, - temperature=2000, - timestep=0.002, -) -final_structure = ts.io.state_to_structures(final_state) diff --git a/examples/old_scripts/4_High_level_api/4.2_auto_batching_api.py b/examples/old_scripts/4_High_level_api/4.2_auto_batching_api.py deleted file mode 100644 index 45afa934..00000000 --- a/examples/old_scripts/4_High_level_api/4.2_auto_batching_api.py +++ /dev/null @@ -1,121 +0,0 @@ -"""Examples of using the auto-batching API. Meant to be run as an interactive script.""" - -# /// script -# dependencies = ["mace-torch>=0.3.12"] -# /// -# %% -import os - -import torch -from ase.build import bulk -from mace.calculators.foundations_models import mace_mp - -import torch_sim as ts -from torch_sim.autobatching import ( - BinningAutoBatcher, - InFlightAutoBatcher, - calculate_memory_scaler, -) -from torch_sim.models.mace import MaceModel -from torch_sim.runners import generate_force_convergence_fn -from torch_sim.units import MetalUnits - - -if not torch.cuda.is_available(): - raise SystemExit(0) - -SMOKE_TEST = os.getenv("CI") is not None - -si_atoms = bulk("Si", "fcc", a=5.43, cubic=True).repeat((3, 3, 3)) -fe_atoms = bulk("Fe", "fcc", a=5.43, cubic=True).repeat((3, 3, 3)) -state: ts.FireState | None = None -device = torch.device("cuda") - -mace = mace_mp(model="small", return_raw_model=True) -mace_model = MaceModel( - model=mace, - device=device, - dtype=torch.float64, - compute_forces=True, -) - -si_state = ts.io.atoms_to_state(si_atoms, device=device, dtype=torch.float64) -fe_state = ts.io.atoms_to_state(fe_atoms, device=device, dtype=torch.float64) - -state = ts.fire_init(state=si_state, model=mace_model, cell_filter=ts.CellFilter.unit) - -si_fire_state = ts.fire_init( - state=si_state, model=mace_model, cell_filter=ts.CellFilter.unit -) -fe_fire_state = ts.fire_init( - state=fe_state, model=mace_model, cell_filter=ts.CellFilter.unit -) - -fire_states = [si_fire_state, fe_fire_state] * (2 if SMOKE_TEST else 20) -fire_states = [state.clone() for state in fire_states] -for state in fire_states: - state.positions += torch.randn_like(state.positions) * 0.01 - -len(fire_states) - - -# %% TODO: add max steps -converge_max_force = generate_force_convergence_fn(force_tol=1e-1) -single_system_memory = calculate_memory_scaler(fire_states[0]) -batcher = InFlightAutoBatcher( - model=mace_model, - memory_scales_with="n_atoms_x_density", - max_memory_scaler=single_system_memory * 2.5 if SMOKE_TEST else None, -) -batcher.load_states(fire_states) -all_completed_states, convergence_tensor, state = [], None, None -while (result := batcher.next_batch(state, convergence_tensor))[0] is not None: - state, completed_states = result[0], result[1] - print(f"Starting new batch of {state.n_systems} states.") - - all_completed_states.extend(completed_states) - print(f"Total number of completed states {len(all_completed_states)}") - - for _step in range(10): - state = ts.fire_step(state=state, model=mace_model) - convergence_tensor = converge_max_force(state, last_energy=None) -all_completed_states.extend(result[1]) -print(f"Total number of completed states {len(all_completed_states)}") - - -# %% run binning autobatcher -si_nvt_state = ts.nvt_langevin_init( - state=si_state, - model=mace_model, - dt=torch.tensor(0.001), - kT=torch.tensor(300 * MetalUnits.temperature), -) -fe_nvt_state = ts.nvt_langevin_init( - state=fe_state, - model=mace_model, - dt=torch.tensor(0.001), - kT=torch.tensor(300 * MetalUnits.temperature), -) - -si_state = ts.io.atoms_to_state(si_atoms, device=device, dtype=torch.float64) -fe_state = ts.io.atoms_to_state(fe_atoms, device=device, dtype=torch.float64) - -nvt_states = [si_nvt_state, fe_nvt_state] * (2 if SMOKE_TEST else 20) -nvt_states = [state.clone() for state in nvt_states] -for state in nvt_states: - state.positions += torch.randn_like(state.positions) * 0.01 - - -single_system_memory = calculate_memory_scaler(fire_states[0]) -batcher = BinningAutoBatcher( - model=mace_model, - memory_scales_with="n_atoms_x_density", - max_memory_scaler=single_system_memory * 2.5 if SMOKE_TEST else None, -) -batcher.load_states(nvt_states) -finished_states: list[ts.SimState] = [] -for batch, _indices in batcher: - for _ in range(100): - batch = ts.nvt_langevin_step(state=batch, model=mace_model) - - finished_states.extend(batch.split()) diff --git a/examples/old_scripts/5_Workflow/5.1_a2c_silicon_batched.py b/examples/old_scripts/5_Workflow/5.1_a2c_silicon_batched.py deleted file mode 100644 index 216929d6..00000000 --- a/examples/old_scripts/5_Workflow/5.1_a2c_silicon_batched.py +++ /dev/null @@ -1,270 +0,0 @@ -"""Demo of the amorphous-to-crystalline (A2C) algorithm for a-Si, ported to TorchSim from -jax-md https://github.com/jax-md/jax-md/blob/main/jax_md/a2c/a2c_workflow.py. -""" - -# /// script -# dependencies = [ -# "mace-torch>=0.3.10", -# "moyopy>=0.4.1", -# "pymatgen>=2025.2.18", -# ] -# /// -import os -import time -from collections import defaultdict - -import torch -from mace.calculators.foundations_models import mace_mp -from moyopy import MoyoDataset, SpaceGroupType -from moyopy.interface import MoyoAdapter -from pymatgen.analysis.structure_analyzer import SpacegroupAnalyzer -from pymatgen.analysis.structure_matcher import StructureMatcher -from pymatgen.core import Composition, Element, Structure -from tqdm import tqdm - -import torch_sim as ts -from torch_sim.integrators.nvt import NVTNoseHooverState -from torch_sim.models.mace import MaceModel, MaceUrls -from torch_sim.units import MetalUnits as Units -from torch_sim.workflows import a2c - - -""" -# Example of how to use random_packed_structure_multi -from torch_sim.workflows.a2c import random_packed_structure_multi - -comp = Composition("Fe80B20") -cell = torch.tensor( - [[10.0, 0.0, 0.0], [0.0, 10.0, 0.0], [0.0, 0.0, 10.0]], dtype=dtype, device=device -) -structure_multi = random_packed_structure_multi( - composition=comp, - cell=cell, - auto_diameter=True, - device=device, - dtype=dtype, - max_iter=100, -) -""" - -SMOKE_TEST = os.getenv("CI") is not None - -device = torch.device("cuda" if torch.cuda.is_available() else "cpu") -dtype = torch.float32 - -raw_model = mace_mp( - model=MaceUrls.mace_mpa_medium, - return_raw_model=True, - default_dtype=str(dtype).removeprefix("torch."), - device=str(device), -) - -# Define system and model -comp = Composition("Si64") -cell = torch.tensor( - [[11.1, 0.0, 0.0], [0.0, 11.1, 0.0], [0.0, 0.0, 11.1]], - dtype=dtype, - device=device, -) -atomic_numbers = [Element(el).Z for el in comp.get_el_amt_dict()] * int(comp.num_atoms) - -atomic_numbers = torch.tensor(atomic_numbers, device=device, dtype=torch.int) -atomic_masses = [Element(el).atomic_mass for el in comp.get_el_amt_dict()] * int( - comp.num_atoms -) -species = [Element.from_Z(Z).symbol for Z in atomic_numbers] - -model = MaceModel( - model=raw_model, - device=device, - compute_forces=True, - compute_stress=False, # We don't need stress for MD - dtype=dtype, - enable_cueq=False, -) -# Workflow starts here -structure, _log = a2c.random_packed_structure( - composition=comp, - cell=cell, - auto_diameter=True, - device=device, - dtype=dtype, - max_iter=100, -) -# Relax structure in batches of 6 -batch_size = 1 if SMOKE_TEST else 6 -max_optim_steps = ( - 1 if SMOKE_TEST else 100 -) # Number of optimization steps for unit cell relaxation - -# MD parameters -equi_steps = 25 if SMOKE_TEST else 2500 # MD steps for melt equilibration -cool_steps = 25 if SMOKE_TEST else 2500 # MD steps for quenching equilibration -final_steps = 25 if SMOKE_TEST else 2500 # MD steps for amorphous phase equilibration -T_high = 2000 # Melt temperature -T_low = 300 # Quench to this temperature -dt = torch.tensor(0.002 * Units.time, device=device, dtype=dtype) # time step = 2fs -tau = 40 * dt # oscillation period in Nose-Hoover thermostat -simulation_steps = equi_steps + cool_steps + final_steps - - -state_dict = { - "positions": structure.positions, - "masses": torch.tensor(atomic_masses, device=device, dtype=dtype), - "cell": cell, - "pbc": True, - "atomic_numbers": atomic_numbers, -} -state = ts.nvt_nose_hoover_init( - state=state_dict, - model=model, - kT=torch.tensor(T_high * Units.temperature, device=device, dtype=dtype), - dt=dt, - seed=1, -) - -logger = { - "T": torch.zeros((simulation_steps, 1), device=device, dtype=dtype), - "H": torch.zeros((simulation_steps, 1), device=device, dtype=dtype), -} - - -def step_fn( - step: int, state: NVTNoseHooverState, logger: dict -) -> tuple[NVTNoseHooverState, dict]: - """Step function for NVT-MD with Nose-Hoover thermostat.""" - current_temp = a2c.get_target_temperature(step, equi_steps, cool_steps, T_high, T_low) - logger["T"][step] = ( - ts.quantities.calc_kT(masses=state.masses, momenta=state.momenta) - / Units.temperature - ) - logger["H"][step] = ts.nvt_nose_hoover_invariant( - state, - kT=torch.tensor(current_temp * Units.temperature, device=device, dtype=dtype), - ).item() - state = ts.nvt_nose_hoover_step( - state=state, - model=model, - dt=dt, - kT=torch.tensor(current_temp * Units.temperature, device=device, dtype=dtype), - ) - return state, logger - - -# Run NVT-MD with the melt-quench-equilibrate temperature profile -for step in range(simulation_steps): - state, logger = step_fn(step, state, logger) - temp, invariant = logger["T"][step].item(), logger["H"][step].item() - print(f"Step {step}: Temperature: {temp:.4f} K: H: {invariant:.4f} eV") - -print( - f"Amorphous structure is ready: positions\n = " - f"{state.positions}\ncell\n = {state.cell}\nspecies = {species}" -) - -# Convert positions to fractional coordinates -fractional_positions = ts.transforms.get_fractional_coordinates( - positions=state.positions, cell=state.cell -) - -# Get subcells to crystallize -subcells = a2c.get_subcells_to_crystallize( - fractional_positions=fractional_positions, - species=species, - d_frac=0.2 if SMOKE_TEST else 0.1, - n_min=2, - n_max=8, -) -print(f"Created {len(subcells)} subcells from a-Si") - -# To save time in this example, we (i) keep only the "cubic" subcells where a==b==c, and -# (ii) keep if number of atoms in the subcell is 2, 4 or 8. This reduces the number of -# subcells to relax from approx. 80k to around 160. -subcells = [ - subcell - for subcell in subcells - if torch.all((subcell[2] - subcell[1]) == (subcell[2] - subcell[1])[0]) - and subcell[0].shape[0] in (2, 4, 8) -] -print(f"Subcells kept for this example: {len(subcells)}") - -candidate_structures = a2c.subcells_to_structures( - candidates=subcells, - fractional_positions=fractional_positions, - cell=state.cell, - species=species, -) - -pymatgen_struct_list = [ - Structure( - lattice=struct[1].detach().cpu().numpy(), - species=struct[2], - coords=struct[0].detach().cpu().numpy(), - coords_are_cartesian=False, - ) - for struct in candidate_structures -] - -start_time = time.perf_counter() -# Create a batched model -model = MaceModel( - model=raw_model, - device=device, - compute_forces=True, - compute_stress=True, - dtype=dtype, - enable_cueq=False, -) - -pymatgen_relaxed_struct_list: list[tuple[Structure, float, float]] = [] -# Process structures in batches of 4 -for batch_idx in tqdm(range(0, len(pymatgen_struct_list), batch_size)): - batch_structs = pymatgen_struct_list[batch_idx : batch_idx + batch_size] - # Combine structures into a single batched state - batch_state = ts.io.structures_to_state(batch_structs, device=device, dtype=dtype) - - final_state, logger, final_energy, final_pressure = ( - a2c.get_frechet_cell_relaxed_structure( - state=batch_state, - model=model, - max_iter=max_optim_steps, - ) - ) - - final_structs = ts.io.state_to_structures(final_state) - - # NOTE: Possible OOM, so we don't store the logger - # relaxed_structures.append((pymatgen_struct, logger, final_energy, final_pressure)) - for sys_idx, final_struct in enumerate(final_structs): - pymatgen_relaxed_struct_list.append( - (final_struct, final_energy[sys_idx], final_pressure[sys_idx]) - ) - -lowest_e_struct = sorted( - pymatgen_relaxed_struct_list, key=lambda x: x[-2] / x[0].num_sites -)[0] -spg = SpacegroupAnalyzer(lowest_e_struct[0]) -print(f"Space group of predicted crystallization product: {spg.get_space_group_symbol()}") - -spg_counter = defaultdict(int) -for struct in pymatgen_relaxed_struct_list: - sym_data = MoyoDataset(MoyoAdapter.from_py_obj(struct[0])) - sp = (sym_data.number, SpaceGroupType(sym_data.number).arithmetic_symbol) - spg_counter[sp] += 1 - -print(f"All space groups encountered: {dict(spg_counter)}") -si_diamond = Structure( - lattice=[ - [0.0, 2.732954, 2.732954], - [2.732954, 0.0, 2.732954], - [2.732954, 2.732954, 0.0], - ], - species=["Si", "Si"], - coords=[[0.5, 0.5, 0.5], [0.75, 0.75, 0.75]], - coords_are_cartesian=False, -) -struct_match = StructureMatcher().fit(lowest_e_struct[0], si_diamond) -print(f"Prediction matches diamond-cubic Si? {struct_match}") - -end_time = time.perf_counter() -print(f"Total time taken to run the workflow: {end_time - start_time:.2f} seconds") diff --git a/examples/old_scripts/5_Workflow/5.2_In_Flight_WBM.py b/examples/old_scripts/5_Workflow/5.2_In_Flight_WBM.py deleted file mode 100644 index b95cc73e..00000000 --- a/examples/old_scripts/5_Workflow/5.2_In_Flight_WBM.py +++ /dev/null @@ -1,97 +0,0 @@ -"""Example script demonstrating batched MACE model optimization with hot-swapping.""" - -# /// script -# dependencies = ["mace-torch>=0.3.10", "matbench-discovery>=1.3.1"] -# /// -import os -import time - -import numpy as np -import torch -from mace.calculators.foundations_models import mace_mp - -import torch_sim as ts -from torch_sim.models.mace import MaceModel, MaceUrls - - -# --- Setup and Configuration --- -# Device and data type configuration -SMOKE_TEST = os.getenv("CI") is not None -device = torch.device("cpu" if SMOKE_TEST else "cuda") -dtype = torch.float32 -print(f"job will run on {device=}") - -# --- Model Initialization --- -print("Loading MACE model...") -mace = mace_mp(model=MaceUrls.mace_mpa_medium, return_raw_model=True) -mace_model = MaceModel( - model=mace, - device=device, - dtype=dtype, - compute_forces=True, -) - -# Optimization parameters -fmax = 0.05 # Force convergence threshold -n_steps = 10 if SMOKE_TEST else 200_000_000 -max_atoms_in_batch = 50 if SMOKE_TEST else 8_000 - -# --- Data Loading --- -if not os.getenv("CI"): - n_structures_to_relax = 100 - print(f"Loading {n_structures_to_relax:,} structures...") - from matbench_discovery.data import DataFiles, ase_atoms_from_zip - - ase_atoms_list = ase_atoms_from_zip( - DataFiles.wbm_initial_atoms.path, limit=n_structures_to_relax - ) -else: - n_structures_to_relax = 2 - print(f"Loading {n_structures_to_relax:,} structures...") - from ase.build import bulk - - al_atoms = bulk("Al", "hcp", a=4.05) - al_atoms.positions += 0.1 * np.random.randn(*al_atoms.positions.shape) # noqa: NPY002 - fe_atoms = bulk("Fe", "bcc", a=2.86).repeat((2, 2, 2)) - fe_atoms.positions += 0.1 * np.random.randn(*fe_atoms.positions.shape) # noqa: NPY002 - ase_atoms_list = [al_atoms, fe_atoms] - -# --- Optimization Setup --- -# Statistics tracking - -# Initialize first batch -fire_states = ts.fire_init( - state=ts.io.atoms_to_state(atoms=ase_atoms_list, device=device, dtype=dtype), - model=mace_model, - cell_filter=ts.CellFilter.frechet, -) - -batcher = ts.autobatching.InFlightAutoBatcher( - model=mace_model, - memory_scales_with="n_atoms_x_density", - max_memory_scaler=1000 if SMOKE_TEST else None, -) -converge_max_force = ts.runners.generate_force_convergence_fn(force_tol=0.05) - -start_time = time.perf_counter() - -# --- Main Optimization Loop --- -batcher.load_states(fire_states) -all_completed_states, convergence_tensor, state = [], None, None -while (result := batcher.next_batch(state, convergence_tensor))[0] is not None: - state, completed_states = result - print(f"Starting new batch of {state.n_systems} states.") - - all_completed_states.extend(completed_states) - print(f"Total number of completed states {len(all_completed_states)}") - - for _step in range(10): - state = ts.fire_step(state=state, model=mace_model) - convergence_tensor = converge_max_force(state, last_energy=None) -all_completed_states.extend(result[1]) -print(f"Total number of completed states {len(all_completed_states)}") - -# --- Final Statistics --- -end_time = time.perf_counter() -total_time = end_time - start_time -print(f"Total time taken: {total_time:.2f} seconds") diff --git a/examples/old_scripts/5_Workflow/5.3_Elastic.py b/examples/old_scripts/5_Workflow/5.3_Elastic.py deleted file mode 100644 index ba9bac77..00000000 --- a/examples/old_scripts/5_Workflow/5.3_Elastic.py +++ /dev/null @@ -1,86 +0,0 @@ -"""Bulk and Shear modulus with MACE.""" - -# /// script -# dependencies = ["ase>=3.26", "mace-torch>=0.3.12"] -# /// -import torch -from ase.build import bulk -from mace.calculators.foundations_models import mace_mp - -import torch_sim as ts -from torch_sim.elastic import get_bravais_type -from torch_sim.models.mace import MaceModel, MaceUrls - - -# Calculator -unit_conv = ts.units.UnitConversion -device = torch.device("cuda" if torch.cuda.is_available() else "cpu") -dtype = torch.float64 - -loaded_model = mace_mp( - model=MaceUrls.mace_mpa_medium, - enable_cueq=False, - device=str(device), - default_dtype=str(dtype).removeprefix("torch."), - return_raw_model=True, -) - -# ASE structure -struct = bulk("Cu", "fcc", a=3.58, cubic=True).repeat((2, 2, 2)) - -model = MaceModel( - model=loaded_model, - device=device, - compute_forces=True, - compute_stress=True, - dtype=dtype, - enable_cueq=False, -) -# Target force tolerance -fmax = 1e-3 - -# Relax positions and cell -state = ts.io.atoms_to_state(atoms=struct, device=device, dtype=dtype) -state = ts.fire_init( - state=state, model=model, scalar_pressure=0.0, cell_filter=ts.CellFilter.frechet -) - -for step in range(300): - pressure = -torch.trace(state.stress.squeeze()) / 3 * unit_conv.eV_per_Ang3_to_GPa - current_fmax = torch.max(torch.abs(state.forces.squeeze())) - print( - f"Step {step}, Energy: {state.energy.item():.4f}, " - f"Pressure: {pressure.item():.4f}, " - f"Fmax: {current_fmax.item():.4f}" - ) - if current_fmax < fmax and abs(pressure) < 1e-2: - break - state = ts.fire_step(state=state, model=model) - -# Get bravais type -bravais_type = get_bravais_type(state) - -# Calculate elastic tensor -elastic_tensor = ts.elastic.calculate_elastic_tensor( - state=state, model=model, bravais_type=bravais_type -) - -# Convert to GPa -elastic_tensor = elastic_tensor * unit_conv.eV_per_Ang3_to_GPa - -# Calculate elastic moduli -bulk_modulus, shear_modulus, poisson_ratio, pugh_ratio = ( - ts.elastic.calculate_elastic_moduli(elastic_tensor) -) - -# Print elastic tensor -print("\nElastic tensor (GPa):") -elastic_tensor_np = elastic_tensor.cpu().numpy() -for row in elastic_tensor_np: - print(" " + " ".join(f"{val:10.4f}" for val in row)) - -# Print mechanical moduli -print(f"Bulk modulus (GPa): {bulk_modulus:.4f}") -print(f"Shear modulus (GPa): {shear_modulus:.4f}") -print(f"Poisson's ratio: {poisson_ratio:.4f}") -print(f"Pugh's ratio (K/G): {pugh_ratio:.4f}") diff --git a/examples/old_scripts/6_Phonons/6.1_Phonons_MACE.py b/examples/old_scripts/6_Phonons/6.1_Phonons_MACE.py deleted file mode 100644 index 40558bd9..00000000 --- a/examples/old_scripts/6_Phonons/6.1_Phonons_MACE.py +++ /dev/null @@ -1,227 +0,0 @@ -"""Calculate Phonon DOS and band structure with MACE in batched mode.""" - -# /// script -# dependencies = [ -# "mace-torch>=0.3.12", -# "phonopy>=2.35", -# "pymatviz>=0.17.1", -# "plotly>=6.3.0", -# "seekpath", -# "ase", -# ] -# /// -import numpy as np -import pymatviz as pmv -import seekpath -import torch -from ase import Atoms -from ase.build import bulk -from mace.calculators.foundations_models import mace_mp -from phonopy import Phonopy -from phonopy.phonon.band_structure import ( - get_band_qpoints_and_path_connections, - get_band_qpoints_by_seekpath, -) - -import torch_sim as ts -from torch_sim.models.mace import MaceModel, MaceUrls - - -def get_qpts_and_connections( - ase_atoms: Atoms, - n_points: int = 101, -) -> tuple[list[list[float]], list[bool]]: - """Get the high symmetry points and path connections for the band structure.""" - # Define seekpath data - seekpath_data = seekpath.get_path( - (ase_atoms.cell, ase_atoms.get_scaled_positions(), ase_atoms.numbers) - ) - - # Extract high symmetry points and path - points = seekpath_data["point_coords"] - path = [] - for segment in seekpath_data["path"]: - start_point = points[segment[0]] - end_point = points[segment[1]] - path.append([start_point, end_point]) - qpts, connections = get_band_qpoints_and_path_connections(path, npoints=n_points) - - return qpts, connections - - -def get_labels_qpts(ph: Phonopy, n_points: int = 101) -> tuple[list[str], list[bool]]: - """Get the labels and coordinates of qpoints for the phonon band structure.""" - # Get labels and coordinates for high-symmetry points - _, qpts_labels, connections = get_band_qpoints_by_seekpath( - ph.primitive, npoints=n_points, is_const_interval=True - ) - connections = [True, *connections] - connections[-1] = True - qpts_labels_connections = [] - idx = 0 - for connection in connections: - if connection: - qpts_labels_connections.append(qpts_labels[idx]) - idx += 1 - else: - qpts_labels_connections.append(f"{qpts_labels[idx]}|{qpts_labels[idx + 1]}") - idx += 2 - - qpts_labels_arr = [ - q_label.replace("\\Gamma", "Γ") - .replace("$", "") - .replace("\\", "") - .replace("mathrm", "") - .replace("{", "") - .replace("}", "") - for q_label in qpts_labels_connections - ] - bands_dict = ph.get_band_structure_dict() - npaths = len(bands_dict["frequencies"]) - qpts_coord = [bands_dict["distances"][n][0] for n in range(npaths)] + [ - bands_dict["distances"][-1][-1] - ] - - return qpts_labels_arr, qpts_coord - - -# Set device and data type -device = torch.device("cuda" if torch.cuda.is_available() else "cpu") -dtype = torch.float32 - -# Load the raw model -loaded_model = mace_mp( - model=MaceUrls.mace_mpa_medium, - return_raw_model=True, - default_dtype=str(dtype).removeprefix("torch."), - device=str(device), -) - -# Structure and input parameters -struct = bulk("Si", "diamond", a=5.431, cubic=True) # ASE structure -supercell_matrix = 2 * np.eye(3) # supercell matrix for phonon calculation -mesh = [20, 20, 20] # Phonon mesh -max_steps = 300 # number of relaxation steps -displ = 0.01 # atomic displacement for phonons (in Angstrom) - -# Relax atomic positions -model = MaceModel( - model=loaded_model, - device=device, - compute_forces=True, - compute_stress=True, - dtype=dtype, - enable_cueq=False, -) -final_state = ts.optimize( - system=struct, - model=model, - optimizer=ts.Optimizer.fire, - max_steps=max_steps, - init_kwargs=dict( - cell_filter=ts.CellFilter.frechet, constant_volume=True, hydrostatic_strain=True - ), -) - -# Define atoms and Phonopy object -atoms = ts.io.state_to_phonopy(final_state)[0] -ph = Phonopy(atoms, supercell_matrix) - -# Generate FC2 displacements -ph.generate_displacements(distance=displ) -supercells = ph.supercells_with_displacements -if supercells is None: - raise ValueError("supercells cannot be None") - -# Convert PhonopyAtoms to state -state = ts.io.phonopy_to_state(supercells, device, dtype) -results = model(state) - -# Extract forces and convert back to list of numpy arrays for phonopy -n_atoms_per_supercell = [len(cell) for cell in supercells] -force_sets = [] -start_idx = 0 -for n_atoms in n_atoms_per_supercell: - end_idx = start_idx + n_atoms - force_sets.append(results["forces"][start_idx:end_idx].detach().cpu().numpy()) - start_idx = end_idx - -# Produce force constants -ph.forces = force_sets -ph.produce_force_constants() - -# Set mesh for DOS calculation -ph.run_mesh(mesh) -ph.run_total_dos() - -# Calculate phonon band structure -ase_atoms = Atoms( - symbols=atoms.symbols, - positions=atoms.positions, - cell=atoms.cell, - pbc=True, -) -q_pts, connections = get_qpts_and_connections(ase_atoms) -ph.run_band_structure(q_pts, path_connections=connections) - -# Define axis style for plots -axis_style = dict( - showgrid=False, - zeroline=False, - linecolor="black", - showline=True, - ticks="inside", - mirror=True, - linewidth=3, - tickwidth=3, - ticklen=10, -) - -# Plot phonon DOS -fig = pmv.phonon_dos(ph.total_dos) -fig.update_traces(line_width=3) -fig.update_layout( - xaxis_title="Frequency (THz)", - yaxis_title="DOS", - font=dict(size=24), - xaxis=axis_style, - yaxis=axis_style, - width=800, - height=600, - plot_bgcolor="white", -) -fig.show() - -# Plot phonon band structure -ph.auto_band_structure(plot=False) -fig = pmv.phonon_bands( - ph.band_structure, - line_kwargs={"width": 3}, -) -qpts_labels, qpts_coord = get_labels_qpts(ph) -for q_pt in qpts_coord: - fig.add_vline(x=q_pt, line_dash="dash", line_color="black", line_width=2, opacity=1) -fig.update_layout( - xaxis_title="Wave Vector", - yaxis_title="Frequency (THz)", - font=dict(size=24), - xaxis=dict( - tickmode="array", - tickvals=qpts_coord, - ticktext=qpts_labels, - showgrid=False, - zeroline=False, - linecolor="black", - showline=True, - ticks="inside", - mirror=True, - linewidth=3, - tickwidth=3, - ticklen=10, - ), - yaxis=axis_style, - width=800, - height=600, - plot_bgcolor="white", -) -fig.show() diff --git a/examples/old_scripts/6_Phonons/6.2_QuasiHarmonic_MACE.py b/examples/old_scripts/6_Phonons/6.2_QuasiHarmonic_MACE.py deleted file mode 100644 index 799861ff..00000000 --- a/examples/old_scripts/6_Phonons/6.2_QuasiHarmonic_MACE.py +++ /dev/null @@ -1,360 +0,0 @@ -"""Calculate quasi-harmonic thermal properties batching over -different volumes and FC2 calculations with MACE. -""" - -# /// script -# dependencies = [ -# "mace-torch>=0.3.12", -# "phonopy>=2.35", -# "pymatviz==0.16", -# "plotly!=6.2.0", # TODO remove pin pending https://github.com/plotly/plotly.py/issues/5253#issuecomment-3016615635 -# ] -# /// -import os - -import numpy as np -import plotly.graph_objects as go -import torch -from ase import Atoms -from ase.build import bulk -from mace.calculators.foundations_models import mace_mp -from phonopy import Phonopy -from phonopy.api_qha import PhonopyQHA -from phonopy.structure.atoms import PhonopyAtoms - -import torch_sim as ts -from torch_sim.models.interface import ModelInterface -from torch_sim.models.mace import MaceModel, MaceUrls - - -def get_relaxed_structure( - struct: Atoms, - model: ModelInterface, - max_steps: int = 300, - fmax: float = 1e-3, - *, - use_autobatcher: bool = False, -) -> ts.SimState: - """Get relaxed structure. - - Args: - struct: ASE structure - model: MACE model - max_steps: Maximum number of relaxation steps - fmax: Force convergence criterion - use_autobatcher: Whether to use automatic batching - - Returns: - SimState: Relaxed structure - """ - trajectory_file = "traj.h5" - reporter = ts.TrajectoryReporter( - trajectory_file, - state_frequency=0, - prop_calculators={ - 1: { - "potential_energy": lambda state: state.energy, - "forces": lambda state: state.forces, - }, - }, - ) - converge_max_force = ts.runners.generate_force_convergence_fn(force_tol=fmax) - final_state = ts.optimize( - system=struct, - model=model, - optimizer=ts.Optimizer.fire, - max_steps=max_steps, - convergence_fn=converge_max_force, - trajectory_reporter=reporter, - autobatcher=use_autobatcher, - init_kwargs=dict( - cell_filter=ts.CellFilter.frechet, - constant_volume=True, - hydrostatic_strain=True, - ), - ) - - os.remove(trajectory_file) - - return final_state - - -def get_qha_structures( - state: ts.SimState, - length_factors: np.ndarray, - model: ModelInterface, - Nmax: int = 300, - fmax: float = 1e-3, - *, - use_autobatcher: bool = False, -) -> list[PhonopyAtoms]: - """Get relaxed structures at different volumes. - - Args: - state: Initial state - length_factors: Array of scaling factors - model: Calculator model - Nmax: Maximum number of relaxation steps - fmax: Force convergence criterion - use_autobatcher: Whether to use automatic batching - - Returns: - list[PhonopyAtoms]: Relaxed PhonopyAtoms structures at different volumes - """ - # Convert state to PhonopyAtoms - relaxed_struct = ts.io.state_to_phonopy(state)[0] - - # Create scaled structures - scaled_structs = [ - PhonopyAtoms( - cell=relaxed_struct.cell * factor, - scaled_positions=relaxed_struct.scaled_positions, - symbols=relaxed_struct.symbols, - ) - for factor in length_factors - ] - - # Relax all structures - scaled_state = ts.optimize( - system=scaled_structs, - model=model, - optimizer=ts.Optimizer.fire, - max_steps=Nmax, - convergence_fn=ts.runners.generate_force_convergence_fn(force_tol=fmax), - autobatcher=use_autobatcher, - init_kwargs=dict( - cell_filter=ts.CellFilter.frechet, - constant_volume=True, - hydrostatic_strain=True, - ), - ) - - return scaled_state.to_phonopy() - - -def get_qha_phonons( - scaled_structures: list[PhonopyAtoms], - model: ModelInterface, - supercell_matrix: np.ndarray | None, - displ: float = 0.05, - *, - use_autobatcher: bool = False, -) -> tuple[list[Phonopy], list[list[np.ndarray]], np.ndarray]: - """Get phonon objects for each scaled atom. - - Args: - scaled_structures: List of PhonopyAtoms objects - model: Calculator model - supercell_matrix: Supercell matrix - displ: Atomic displacement for phonons - use_autobatcher: Whether to use automatic batching - - Returns: - tuple[list[Phonopy], list[list[np.ndarray]], np.ndarray]: Contains: - - List of Phonopy objects - - List of force sets for each structure - - Array of energies - """ - # Generate phonon object for each scaled structure - supercells_flat = [] - supercell_boundaries = [0] - ph_sets = [] - if supercell_matrix is None: - supercell_matrix = np.eye(3) - for atoms in scaled_structures: - ph = Phonopy( - atoms, - supercell_matrix=supercell_matrix, - primitive_matrix="auto", - ) - ph.generate_displacements(distance=displ) - supercells = ph.supercells_with_displacements - n_atoms = 0 if supercells is None else sum(len(cell) for cell in supercells) - supercell_boundaries.append(supercell_boundaries[-1] + n_atoms) - supercells_flat.extend([] if supercells is None else supercells) - ph_sets.append(ph) - - # Run the model on flattened structure - reporter = ts.TrajectoryReporter( - None, - state_frequency=0, - prop_calculators={ - 1: { - "potential_energy": lambda state: state.energy, - "forces": lambda state: state.forces, - } - }, - ) - results = ts.static( - system=supercells_flat, - model=model, - autobatcher=use_autobatcher, - trajectory_reporter=reporter, - ) - - # Reconstruct force sets and energies - force_sets = [] - forces = torch.cat([r["forces"] for r in results]).detach().cpu().numpy() - energies = ( - torch.tensor([r["potential_energy"] for r in results]).detach().cpu().numpy() - ) - for sys_idx, ph in enumerate(ph_sets): - start, end = supercell_boundaries[sys_idx], supercell_boundaries[sys_idx + 1] - forces_i = forces[start:end] - n_atoms = len(ph.supercell) - n_displacements = len(ph.supercells_with_displacements) - force_sets_i = [] - for disp_idx in range(n_displacements): - start_j = disp_idx * n_atoms - end_j = (disp_idx + 1) * n_atoms - force_sets_i.append(forces_i[start_j:end_j]) - force_sets.append(force_sets_i) - - return ph_sets, force_sets, energies - - -# Set device and data type -device = torch.device("cuda" if torch.cuda.is_available() else "cpu") -dtype = torch.float64 - -autobatcher = False - -# Load the raw model -loaded_model = mace_mp( - model=MaceUrls.mace_mpa_medium, - return_raw_model=True, - default_dtype=str(dtype).removeprefix("torch."), - device=str(device), -) -model = MaceModel( - model=loaded_model, - device=device, - compute_forces=True, - compute_stress=True, - dtype=dtype, - enable_cueq=False, -) - -# Structure and input parameters -struct = bulk("Si", "diamond", a=5.431, cubic=True) # ASE structure -supercell_matrix = 2 * np.eye(3) # supercell matrix for phonon calculation -mesh = [20, 20, 20] # Phonon mesh -fmax = 1e-3 # force convergence -Nmax = 300 # maximum number of relaxation steps -displ = 0.05 # atomic displacement for phonons (in Angstrom) -temperatures = np.arange(0, 1410, 10) # temperature range for quasi-harmonic calculation -length_factors = np.linspace( - 0.85, 1.15, 15 -) # length factor for quasi-harmonic calculation - -# Relax initial structure -state = get_relaxed_structure( - struct=struct, model=model, max_steps=Nmax, fmax=fmax, use_autobatcher=autobatcher -) - -# Get relaxed structures at different volumes -scaled_structures = get_qha_structures( - state=state, - length_factors=length_factors, - model=model, - Nmax=Nmax, - fmax=fmax, - use_autobatcher=autobatcher, -) - -# Get phonons, FC2 forces, and energies for all set of scaled structures -ph_sets, force_sets, energy_sets = get_qha_phonons( - scaled_structures=scaled_structures, - model=model, - supercell_matrix=supercell_matrix, - displ=displ, - use_autobatcher=autobatcher, -) - -# Calculate thermal properties for each supercells -volumes = [] -energies = [] -free_energies = [] -entropies = [] -heat_capacities = [] -n_displacements = len(getattr(ph_sets[0], "supercells_with_displacements", [])) -for i in range(len(ph_sets)): - ph_sets[i].forces = force_sets[i] - ph_sets[i].produce_force_constants() - ph_sets[i].run_mesh(mesh) - ph_sets[i].run_thermal_properties( - t_min=temperatures[0], - t_max=temperatures[-1], - t_step=int((temperatures[-1] - temperatures[0]) / (len(temperatures) - 1)), - ) - - # Store volume, energy, entropies, heat capacities - thermal_props = ph_sets[i].get_thermal_properties_dict() - n_unit_cells = np.prod(np.diag(supercell_matrix)) - cell = scaled_structures[i].cell - volume = np.linalg.det(cell) - volumes.append(volume) - energies.append(energy_sets[i * n_displacements].item() / n_unit_cells) - free_energies.append(thermal_props["free_energy"]) - entropies.append(thermal_props["entropy"]) - heat_capacities.append(thermal_props["heat_capacity"]) - -# run QHA -qha = PhonopyQHA( - volumes=volumes, - electronic_energies=np.tile(energies, (len(temperatures), 1)), - temperatures=temperatures, - free_energy=np.array(free_energies).T, - cv=np.array(heat_capacities).T, - entropy=np.array(entropies).T, - eos="vinet", -) - -# Axis style -axis_style = dict( - showgrid=False, - zeroline=False, - linecolor="black", - showline=True, - ticks="inside", - mirror=True, - linewidth=3, - tickwidth=3, - ticklen=10, -) - -# Plot thermal expansion vs temperature -fig = go.Figure() -fig.add_trace( - go.Scatter(x=temperatures, y=qha.thermal_expansion, mode="lines", line=dict(width=4)) -) -fig.update_layout( - xaxis_title="Temperature (K)", - yaxis_title="Thermal Expansion (1/K)", - font=dict(size=24), - xaxis=axis_style, - yaxis=axis_style, - width=800, - height=600, - plot_bgcolor="white", -) -fig.show() - -# Plot bulk modulus vs temperature -fig = go.Figure() -fig.add_trace( - go.Scatter( - x=temperatures, y=qha.bulk_modulus_temperature, mode="lines", line=dict(width=4) - ) -) -fig.update_layout( - xaxis_title="Temperature (K)", - yaxis_title="Bulk Modulus (GPa)", - font=dict(size=24), - xaxis=axis_style, - yaxis=axis_style, - width=800, - height=600, - plot_bgcolor="white", -) -fig.show() diff --git a/examples/old_scripts/6_Phonons/6.3_Conductivity_MACE.py b/examples/old_scripts/6_Phonons/6.3_Conductivity_MACE.py deleted file mode 100644 index a9f06238..00000000 --- a/examples/old_scripts/6_Phonons/6.3_Conductivity_MACE.py +++ /dev/null @@ -1,205 +0,0 @@ -"""Calculate the Wigner thermal conductivity batching over -FC2 and FC3 calculations with MACE. -""" - -# /// script -# dependencies = [ -# "mace-torch>=0.3.12", -# "phono3py>=3.12", -# "pymatgen>=2025.2.18", -# ] -# /// -import os -import sys -import time -from typing import TYPE_CHECKING, Literal, cast - -import numpy as np -import plotly.graph_objects as go -import torch -from ase.build import bulk -from mace.calculators.foundations_models import mace_mp -from phono3py import Phono3py -from tqdm import tqdm - -import torch_sim as ts -from torch_sim.models.mace import MaceModel, MaceUrls - - -if TYPE_CHECKING: - from phonopy.structure.atoms import PhonopyAtoms - - -def print_relax_info(trajectory_file: str, device: torch.device) -> None: - """Print relaxation information from trajectory file. - - Args: - trajectory_file: Path to the trajectory file - device: Torch device for calculations - """ - with ts.TorchSimTrajectory(trajectory_file) as traj: - energies = traj.get_array("potential_energy") - forces = traj.get_array("forces") - if isinstance(forces, np.ndarray): - forces = torch.tensor(forces, device=device) - max_force = torch.max(torch.abs(forces), dim=1).values - for i in range(max_force.shape[0]): - print( - f"Step {i}: Max force = {torch.max(max_force[i]).item():.4e} eV/A, " - f"Energy = {energies[i].item():.4f} eV" - ) - os.remove(trajectory_file) - - -start_time = time.perf_counter() -device = torch.device("cuda" if torch.cuda.is_available() else "cpu") -dtype = torch.float64 - -# Load the raw model from URL -loaded_model = mace_mp( - model=MaceUrls.mace_mpa_medium, - return_raw_model=True, - default_dtype=str(dtype).removeprefix("torch."), - device=str(device), -) -model = MaceModel( - model=loaded_model, - device=device, - compute_forces=True, - compute_stress=True, - dtype=dtype, - enable_cueq=False, -) - -# Structure and input parameters -struct = bulk("Si", "diamond", a=5.431, cubic=True) # ASE structure -mesh = [8, 8, 8] # Phonon mesh -# supercell matrix for phonon calculation (use larger cell for better accuracy) -supercell_matrix = [1, 1, 1] -supercell_matrix_fc2 = [2, 2, 2] # supercell matrix for FC2 calculation -max_steps = 300 # number of relaxation steps -fmax = 1e-3 # force convergence -displ = 0.05 # atomic displacement for phonons (in Angstrom) -conductivity_type: Literal["wigner", "kubo"] = "wigner" -temperatures = np.arange( - 0, 1600, 10 -) # temperature range for thermal conductivity calculation - -# Relax structure -converge_max_force = ts.runners.generate_force_convergence_fn(force_tol=fmax) -trajectory_file = "anha.h5" -reporter = ts.TrajectoryReporter( - trajectory_file, - state_frequency=0, - prop_calculators={ - 1: { - "potential_energy": lambda state: state.energy, - "forces": lambda state: state.forces, - }, - }, -) -final_state = ts.optimize( - system=struct, - model=model, - optimizer=ts.Optimizer.fire, - max_steps=max_steps, - convergence_fn=converge_max_force, - trajectory_reporter=reporter, - init_kwargs=dict( - cell_filter=ts.CellFilter.frechet, constant_volume=True, hydrostatic_strain=True - ), -) -print_relax_info(trajectory_file, device) - -# Phono3py object -phonopy_atoms = ts.io.state_to_phonopy(final_state)[0] -ph3 = Phono3py( - phonopy_atoms, - supercell_matrix=supercell_matrix, - primitive_matrix="auto", - phonon_supercell_matrix=supercell_matrix_fc2, -) - -# Calculate FC2 -ph3.generate_fc2_displacements(distance=displ) -supercells_fc2 = ph3.phonon_supercells_with_displacements -state = ts.io.phonopy_to_state(supercells_fc2, device=device, dtype=dtype) -results = model(state) -n_atoms_per_supercell = [len(sc) for sc in supercells_fc2] -force_sets = [] -start_idx = 0 -for n_atoms in tqdm(n_atoms_per_supercell, desc="FC2"): - end_idx = start_idx + n_atoms - force_sets.append(results["forces"][start_idx:end_idx].detach().cpu().numpy()) - start_idx = end_idx -ph3.phonon_forces = np.array(force_sets).reshape(-1, len(ph3.phonon_supercell), 3) -ph3.produce_fc2(symmetrize_fc2=True) - -# Calculate FC3 -ph3.generate_displacements(distance=displ) -supercells_fc3 = cast("list[PhonopyAtoms]", ph3.supercells_with_displacements) -state = ts.io.phonopy_to_state(supercells_fc3, device=device, dtype=dtype) -results = model(state) -n_atoms_per_supercell = [len(sc) for sc in supercells_fc3] -force_sets = [] -start_idx = 0 -for n_atoms in tqdm(n_atoms_per_supercell, desc="FC3"): - end_idx = start_idx + n_atoms - force_sets.append(results["forces"][start_idx:end_idx].detach().cpu().numpy()) - start_idx = end_idx -ph3.forces = np.array(force_sets).reshape(-1, len(ph3.supercell), 3) -ph3.produce_fc3(symmetrize_fc3r=True) - -# Run thermal conductivity calculation -ph3.mesh_numbers = mesh -ph3.init_phph_interaction(symmetrize_fc3q=False) -ph3.run_thermal_conductivity( - is_isotope=True, - temperatures=temperatures, - conductivity_type=conductivity_type, - boundary_mfp=1e6, -) -temperatures = ph3.thermal_conductivity.temperatures -if conductivity_type == "wigner": - kappa = ph3.thermal_conductivity.kappa_TOT_RTA[0] -else: - kappa = ph3.thermal_conductivity.kappa[0] - -# Average thermal conductivity -kappa_av = np.mean(kappa[:, :3], axis=1) - -# Print kappa at 300K -idx_300k = np.abs(temperatures - 300).argmin() -print( - f"\nThermal conductivity at {temperatures[idx_300k]:.1f} K: " - f"{kappa_av[idx_300k]:.2f} W/mK" -) - -# Axis style -axis_style = dict( - showgrid=False, - zeroline=False, - linecolor="black", - showline=True, - ticks="inside", - mirror=True, - linewidth=3, - tickwidth=3, - ticklen=10, -) - -# Plot thermal expansion vs temperature -fig = go.Figure() -fig.add_trace(go.Scatter(x=temperatures, y=kappa_av, mode="lines", line=dict(width=4))) -fig.update_layout( - xaxis_title="Temperature (K)", - yaxis_title="Thermal Conductivity (W/mK)", - font=dict(size=24), - xaxis=axis_style, - yaxis=axis_style, - width=800, - height=600, - plot_bgcolor="white", -) -if "IPython" in sys.modules: - fig.show() diff --git a/examples/old_scripts/7_Others/7.1_Soft_sphere_autograd.py b/examples/old_scripts/7_Others/7.1_Soft_sphere_autograd.py deleted file mode 100644 index beba6150..00000000 --- a/examples/old_scripts/7_Others/7.1_Soft_sphere_autograd.py +++ /dev/null @@ -1,70 +0,0 @@ -"""Plot the soft sphere potential energy vs distance using plotly.""" - -# /// script -# dependencies = [ -# "plotly>=6", -# "kaleido", -# "scipy>=1.15", -# ] -# /// -# ruff: noqa: RUF001 -import torch -from plotly.subplots import make_subplots - -from torch_sim.models.soft_sphere import soft_sphere_pair, soft_sphere_pair_force - - -sigma = torch.tensor(1.0) -epsilon = torch.tensor(1.0) -alpha = torch.tensor(2) - -# Generate distance values from 0.1*sigma to 2*sigma -dr = torch.linspace(0.1 * sigma, 2 * sigma, 1000) -dr_tensor = torch.sqrt(torch.tensor(dr)) - -# Calculate potential energy -# Make dr_tensor require gradients for autograd -dr_tensor.requires_grad_(True) # noqa: FBT003 - -# Calculate potential energy with gradients enabled -energy = soft_sphere_pair(dr_tensor, sigma, epsilon, alpha) - -# Calculate force as negative gradient of energy with respect to distance -force = torch.autograd.grad(energy, dr_tensor, grad_outputs=torch.ones_like(energy))[0] -explicit_force = soft_sphere_pair_force(dr_tensor, sigma, epsilon, alpha) - -# Create figure with secondary y-axis -fig = make_subplots(specs=[[{"secondary_y": True}]]) - -# Add energy trace -fig.add_scatter( - x=dr, - y=energy.detach().numpy(), - name="Energy", - secondary_y=False, -) - -# Add force traces -fig.add_scatter( - x=dr, - y=force.detach().numpy(), - name="Force (Autograd)", - secondary_y=True, -) - -fig.add_scatter( - x=dr, - y=explicit_force.detach().numpy(), - name="Force (Explicit)", - line=dict(dash="dash"), - secondary_y=True, -) - -# Add figure titles and labels -fig.update_layout(title="Soft Sphere Potential", xaxis_title="Distance (r/σ)") - -# Update y-axes labels -fig.update_yaxes(title_text="Energy (ε)", secondary_y=False) -fig.update_yaxes(title_text="Force (ε/σ)", secondary_y=True) - -fig.write_image("soft_sphere_grad_vs_explicit.pdf") diff --git a/examples/old_scripts/7_Others/7.2_Stress_autograd.py b/examples/old_scripts/7_Others/7.2_Stress_autograd.py deleted file mode 100644 index 97c58120..00000000 --- a/examples/old_scripts/7_Others/7.2_Stress_autograd.py +++ /dev/null @@ -1,196 +0,0 @@ -"""Measure performance of different methods for calculating stress. - -The stress is calculated using three different methods: - -1. The stress_fn function, explicit force method. -2. The stress_autograd_fn function, automatic differentiation method. -3. The stress_autograd_fn_functorch function, functorch automatic differentiation method. -""" - -# /// script -# dependencies = ["scipy>=1.15"] -# /// -import timeit - -import torch - -from torch_sim.models.lennard_jones import lennard_jones_pair, lennard_jones_pair_force - - -torch.set_default_tensor_type(torch.DoubleTensor) -# Set simulation parameters -n_steps = 10_000 -kT = torch.tensor(0.722) # Temperature in energy units -sigma = torch.tensor(1.0) # Length parameter -epsilon = torch.tensor(1.0) # Energy parameter - -# Grid initialization -Nx = 10 -length_scale = 1.1 # Increased to reduce initial overlap -positions = torch.zeros((Nx * Nx, 2)) -for i in range(Nx): - for j in range(Nx): - positions[i * Nx + j] = torch.tensor([i * length_scale, j * length_scale]) - -num_particles = Nx * Nx -box_size = Nx * length_scale -mass = torch.ones(num_particles) - - -def energy_fn( - R: torch.Tensor, box: torch.Tensor, perturbation: torch.Tensor | None = None -) -> torch.Tensor: - """Calculate energy using a brute force method.""" - # Create displacement vectors for all pairs - ri = R.unsqueeze(0) - rj = R.unsqueeze(1) - dr = rj - ri - - # Apply periodic boundary conditions - dr = dr - box.diagonal() * torch.round(dr / box.diagonal()) - if perturbation is not None: - # Apply transformation directly (R + dR) - dr = dr + torch.einsum("ij,nmj->nmi", perturbation, dr) - - # Calculate distances - distances = torch.norm(dr, dim=2) - - # Mask out self-interactions - mask = torch.eye(R.shape[0], dtype=torch.bool, device=R.device) - distances = distances.masked_fill(mask, torch.inf) - - # Calculate potential energy - energy = lennard_jones_pair(distances, sigma, epsilon) - - return energy.sum() / 2.0 # Divide by 2 to avoid double counting - - -def force_fn(R: torch.Tensor, box: torch.Tensor) -> torch.Tensor: - """Calculate forces using a brute force method.""" - # Create displacement vectors for all pairs - ri = R.unsqueeze(0) - rj = R.unsqueeze(1) - dr = rj - ri - - # Apply periodic boundary conditions - dr = dr - box.diagonal() * torch.round(dr / box.diagonal()) - - # Calculate distances - distances = torch.norm(dr, dim=2) - - # Mask out self-interactions - mask = torch.eye(R.shape[0], dtype=torch.bool, device=R.device) - - distances = distances.masked_fill(mask, torch.inf) - forces = lennard_jones_pair_force(distances, sigma, epsilon) - - # Project forces along displacement vectors - unit_vectors = dr / distances.unsqueeze(-1) - force_components = forces.unsqueeze(-1) * unit_vectors - return force_components.sum(dim=0) - - -def stress_fn(R: torch.Tensor, box: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: - """Calculate stress using a brute force method.""" - # Create displacement vectors for all pairs - ri = R.unsqueeze(0) - rj = R.unsqueeze(1) - dr = rj - ri - - # Apply periodic boundary conditions - dr = dr - box.diagonal() * torch.round(dr / box.diagonal()) - - # Calculate distances - distances = torch.norm(dr, dim=2) - - # Mask out self-interactions - mask = torch.eye(R.shape[0], dtype=torch.bool, device=R.device) - distances = distances.masked_fill(mask, torch.inf) - - # Calculate forces between pairs - forces = lennard_jones_pair_force(distances, sigma, epsilon) - - # Calculate stress components using outer product of force and distance - # Project forces along displacement vectors - unit_vectors = dr / distances.unsqueeze(-1) - force_components = forces.unsqueeze(-1) * unit_vectors - - # Compute outer product for stress tensor - stress_per_pair = torch.einsum("...i,...j->...ij", dr, force_components) - - # Sum over all pairs and divide by volume - volume = torch.prod(box.diagonal()) - stress = -stress_per_pair.sum(dim=(0, 1)) / volume - - return force_components.sum(dim=0), stress - - -def stress_autograd_fn(R: torch.Tensor, box: torch.Tensor) -> torch.Tensor: - """Calculate stress using autograd.""" - # Get volume and dimension - volume = torch.prod(box.diagonal()) - dim = R.shape[1] - - # Create identity and zero matrices - eye = torch.eye(dim, device=R.device) - zero = torch.zeros((dim, dim), device=R.device).requires_grad_(True) # noqa: FBT003 - - def internal_energy(eps: torch.Tensor) -> torch.Tensor: - return energy_fn(R, box, perturbation=(eye + eps)) - - # Calculate energy at zero strain - energy = internal_energy(zero) - - stress = -torch.autograd.grad(energy, zero)[0] - - return 2 * stress / volume - - -def stress_autograd_fn_functorch(R: torch.Tensor, box: torch.Tensor) -> torch.Tensor: - """Calculate stress using functorch.grad.""" - # Get volume and dimension - volume = torch.prod(box.diagonal()) - dim = R.shape[1] - - # Create identity and zero matrices - eye = torch.eye(dim, device=R.device, dtype=torch.float64) - zero = torch.zeros((dim, dim), device=R.device, dtype=torch.float64) - - def internal_energy(eps: torch.Tensor, R: torch.Tensor) -> torch.Tensor: - # Apply strain perturbation to positions - return energy_fn(R, box, perturbation=(eye + eps)) - - # Use functorch.grad - from torch.func import grad - - # Get gradient of U with respect to eps - grad_U = grad(internal_energy, argnums=0) # take gradient w.r.t. first argument (eps) - - # Calculate stress using grad - stress = -grad_U(zero, R) - - return 2 * stress / volume - - -lat_vec = torch.eye(2) * box_size - -stress_fn_time = timeit.timeit( - "stress_fn(positions, lat_vec)", globals=globals(), number=100 -) -print(f"{stress_fn_time=:.4f} sec") - -stress_autograd_fn_time = timeit.timeit( - "stress_autograd_fn(positions, lat_vec)", globals=globals(), number=100 -) -print(f"{stress_autograd_fn_time=:.4f} sec") - -stress_autograd_fn_functorch_time = timeit.timeit( - "stress_autograd_fn_functorch(positions, lat_vec)", globals=globals(), number=100 -) -print(f"{stress_autograd_fn_functorch_time=:.4f} sec") - -print(energy_fn(positions, lat_vec, None)) -print(force_fn(positions, lat_vec)) -print(f"stress_fn: {stress_fn(positions, lat_vec)[1]}") -print(f"stress_autograd_fn: {stress_autograd_fn(positions, lat_vec)}") -print(f"stress_autograd_fn_functorch: {stress_autograd_fn_functorch(positions, lat_vec)}") diff --git a/examples/old_scripts/7_Others/7.3_Batched_neighbor_list.py b/examples/old_scripts/7_Others/7.3_Batched_neighbor_list.py deleted file mode 100644 index 1393b908..00000000 --- a/examples/old_scripts/7_Others/7.3_Batched_neighbor_list.py +++ /dev/null @@ -1,46 +0,0 @@ -"""Batched neighbor list.""" - -# /// script -# dependencies = ["ase>=3.26", "scipy>=1.15"] -# /// -import torch -from ase.build import bulk - -import torch_sim as ts -from torch_sim import transforms -from torch_sim.neighbors import torch_nl_linked_cell, torch_nl_n2 - - -atoms_list = [bulk("Si", "diamond", a=5.43), bulk("Ge", "diamond", a=5.65)] -state = ts.io.atoms_to_state(atoms_list, device=torch.device("cpu"), dtype=torch.float32) -pos, cell, pbc = state.positions, state.cell, state.pbc -system_idx, n_atoms = state.system_idx, state.n_atoms -cutoff = torch.tensor(4.0, dtype=pos.dtype) -self_interaction = False - -# Ensure pbc has the correct shape [n_systems, 3] -pbc_tensor = torch.tensor(pbc).repeat(state.n_systems, 1) - -mapping, mapping_system, shifts_idx = torch_nl_linked_cell( - pos, cell, pbc_tensor, cutoff, system_idx, self_interaction -) -cell_shifts = transforms.compute_cell_shifts(cell, shifts_idx, mapping_system) -dds = transforms.compute_distances_with_cell_shifts(pos, mapping, cell_shifts) - -print(mapping.shape) -print(mapping_system.shape) -print(shifts_idx.shape) -print(cell_shifts.shape) -print(dds.shape) - -mapping_n2, mapping_system_n2, shifts_idx_n2 = torch_nl_n2( - pos, cell, pbc_tensor, cutoff, system_idx, self_interaction -) -cell_shifts_n2 = transforms.compute_cell_shifts(cell, shifts_idx_n2, mapping_system_n2) -dds_n2 = transforms.compute_distances_with_cell_shifts(pos, mapping_n2, cell_shifts_n2) - -print(mapping_n2.shape) -print(mapping_system_n2.shape) -print(shifts_idx_n2.shape) -print(cell_shifts_n2.shape) -print(dds_n2.shape) diff --git a/examples/old_scripts/7_Others/7.4_Velocity_AutoCorrelation.py b/examples/old_scripts/7_Others/7.4_Velocity_AutoCorrelation.py deleted file mode 100644 index cdd91264..00000000 --- a/examples/old_scripts/7_Others/7.4_Velocity_AutoCorrelation.py +++ /dev/null @@ -1,116 +0,0 @@ -"""Velocity autocorrelation example.""" - -# /// script -# dependencies = [ -# "ase>=3.26", -# "matplotlib", -# "numpy", -# ] -# /// -from typing import Any - -import matplotlib.pyplot as plt -import numpy as np -import torch -from ase.build import bulk -from ase.md.velocitydistribution import MaxwellBoltzmannDistribution - -import torch_sim as ts -from torch_sim.models.lennard_jones import LennardJonesModel -from torch_sim.properties.correlations import VelocityAutoCorrelation -from torch_sim.units import MetalUnits as Units - - -def prepare_system() -> tuple[ - Any, Any, torch.Tensor, torch.Tensor, torch.device, torch.dtype, float -]: - """Create and prepare Ar system with LJ potential.""" - device = torch.device("cuda" if torch.cuda.is_available() else "cpu") - dtype = torch.float64 - - # Using solid Ar w/ LJ for ease - atoms = bulk("Ar", crystalstructure="fcc", a=5.256, cubic=True) - atoms = atoms.repeat((3, 3, 3)) - temperature = 50.0 # Kelvin - MaxwellBoltzmannDistribution(atoms, temperature_K=temperature) - state = ts.io.atoms_to_state(atoms, device=device, dtype=dtype) - - epsilon = 0.0104 # eV - sigma = 3.4 # Å - cutoff = 2.5 * sigma - - lj_model = LennardJonesModel( - sigma=sigma, - epsilon=epsilon, - cutoff=cutoff, - device=device, - dtype=dtype, - compute_forces=True, - ) - - timestep = 0.001 # ps (1 fs) - dt = torch.tensor(timestep * Units.time, device=device, dtype=dtype) - temp_kT = temperature * Units.temperature # Convert K to internal units - kT = torch.tensor(temp_kT, device=device, dtype=dtype) - - return state, lj_model, dt, kT, device, dtype, timestep - - -def plot_results(*, time: np.ndarray, vacf: np.ndarray, window_count: int) -> None: - """Plot VACF results.""" - plt.figure(figsize=(10, 8)) - plt.plot(time, vacf, "b-", linewidth=2) - plt.xlabel("Time (fs)", fontsize=12) - plt.ylabel("VACF", fontsize=12) - plt.title(f"VACF (Average of {window_count} windows)", fontsize=14) - plt.axhline(y=0, color="k", linestyle="--", alpha=0.3) - plt.ylim(-0.6, 1.1) - plt.tight_layout() - plt.savefig("vacf_example.png") - - -def main() -> None: - """Run velocity autocorrelation simulation using Lennard-Jones model.""" - state, lj_model, dt, kT, device, _dtype, timestep = prepare_system() - state = ts.nve_init(state=state, model=lj_model, kT=kT) - - window_size = 150 # Length of correlation: dt * correlation_dt * window_size - vacf_calc = VelocityAutoCorrelation( - window_size=window_size, - device=device, - use_running_average=True, - normalize=True, - ) - - # Sampling freq is controlled by prop_calculators - trajectory = "vacf_example.h5" - correlation_dt = 10 # Step delta between correlations - reporter = ts.TrajectoryReporter( - trajectory, - state_frequency=100, - prop_calculators={correlation_dt: {"vacf": vacf_calc}}, - ) - - num_steps = 15000 # NOTE: short run - for step in range(num_steps): - state = ts.nve_step(state=state, model=lj_model, dt=dt) # type: ignore[call-arg] - reporter.report(state, step) - - reporter.close() - - # VACF results and plot - # Timesteps -> Time in fs - time_steps = np.arange(window_size) - time = time_steps * correlation_dt * timestep * 1000 - - if vacf_calc.vacf is not None: - plot_results( - time=time, - vacf=vacf_calc.vacf.cpu().numpy(), - # Just for demo purposes - window_count=vacf_calc._window_count, # noqa: SLF001 - ) - - -if __name__ == "__main__": - main() diff --git a/examples/old_scripts/7_Others/7.6_Compare_ASE_to_VV_FIRE.py b/examples/old_scripts/7_Others/7.6_Compare_ASE_to_VV_FIRE.py deleted file mode 100644 index d8770893..00000000 --- a/examples/old_scripts/7_Others/7.6_Compare_ASE_to_VV_FIRE.py +++ /dev/null @@ -1,928 +0,0 @@ -"""Structural optimization with MACE using FIRE optimizer. -Comparing the ASE and VV FIRE optimizers. -""" - -# /// script -# dependencies = ["mace-torch>=0.3.12", "plotly>=6.0.0"] -# /// -import os -import time -from functools import partial -from typing import Literal, TypedDict - -import numpy as np -import plotly.graph_objects as go -import torch -from ase.build import bulk -from ase.cell import Cell -from ase.filters import FrechetCellFilter -from ase.optimize import FIRE as ASEFIRE -from mace.calculators.foundations_models import mace_mp -from mace.calculators.foundations_models import mace_mp as mace_mp_calculator_for_ase - -import torch_sim as ts -from torch_sim.models.mace import MaceModel, MaceUrls -from torch_sim.optimizers import OptimState -from torch_sim.state import SimState - - -# Set device, data type and unit conversion -SMOKE_TEST = os.getenv("CI") is not None -DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu") -DTYPE = torch.float32 -unit_conv = ts.units.UnitConversion - -# Option 1: Load the raw model from the downloaded model -loaded_model = mace_mp( - model=MaceUrls.mace_mpa_medium, - return_raw_model=True, - default_dtype=str(DTYPE).removeprefix("torch."), - device=str(DEVICE), -) - -# Number of steps to run -max_iterations = 10 if SMOKE_TEST else 500 -supercell_scale = (1, 1, 1) if SMOKE_TEST else (3, 2, 2) -# Max steps for each individual ASE optimization run -ase_max_optimizer_steps = max_iterations * 10 - -# Set random seed for reproducibility -rng = np.random.default_rng(seed=0) -torch.manual_seed(0) -if torch.cuda.is_available(): - torch.cuda.manual_seed_all(0) - -# Create diamond cubic Silicon -si_dc = bulk("Si", "diamond", a=5.21, cubic=True).repeat(supercell_scale) -si_dc.positions += 0.1 * rng.standard_normal(si_dc.positions.shape).clip(-1, 1) - -# Create FCC Copper -cu_dc = bulk("Cu", "fcc", a=3.85).repeat([r + 1 for r in supercell_scale]) -cu_dc.positions += 0.1 * rng.standard_normal(cu_dc.positions.shape).clip(-1, 1) - -# Create BCC Iron -fe_dc = bulk("Fe", "bcc", a=2.95).repeat([r + 1 for r in supercell_scale]) -fe_dc.positions += 0.1 * rng.standard_normal(fe_dc.positions.shape).clip(-1, 1) - -si_dc_vac = si_dc.copy() -si_dc_vac.positions += 0.1 * rng.standard_normal(si_dc_vac.positions.shape).clip(-1, 1) -# select 2 numbers in range 0 to len(si_dc_vac) -indices = rng.choice(len(si_dc_vac), size=2, replace=False) -for idx in indices: - si_dc_vac.pop(idx) - - -cu_dc_vac = cu_dc.copy() -cu_dc_vac.positions += 0.1 * rng.standard_normal(cu_dc_vac.positions.shape).clip(-1, 1) -# remove 2 atoms from cu_dc_vac at random -indices = rng.choice(len(cu_dc_vac), size=2, replace=False) -for idx in indices: - index = idx + 3 - if index < len(cu_dc_vac): - cu_dc_vac.pop(index) - else: - print(f"Index {index} is out of bounds for cu_dc_vac") - cu_dc_vac.pop(0) - -fe_dc_vac = fe_dc.copy() -fe_dc_vac.positions += 0.1 * rng.standard_normal(fe_dc_vac.positions.shape).clip(-1, 1) -# remove 2 atoms from fe_dc_vac at random -indices = rng.choice(len(fe_dc_vac), size=2, replace=False) -for idx in indices: - index = idx + 2 - if index < len(fe_dc_vac): - fe_dc_vac.pop(index) - else: - print(f"Index {index} is out of bounds for fe_dc_vac") - fe_dc_vac.pop(0) - - -# Create a list of our atomic systems -atoms_list = [si_dc, cu_dc, fe_dc, si_dc_vac, cu_dc_vac, fe_dc_vac] - -# Print structure information -print(f"Silicon atoms: {len(si_dc)}") -print(f"Copper atoms: {len(cu_dc)}") -print(f"Iron atoms: {len(fe_dc)}") -print(f"Total number of structures: {len(atoms_list)}") - -# Create batched model -model = MaceModel( - model=loaded_model, - device=DEVICE, - compute_forces=True, - compute_stress=True, - dtype=DTYPE, - enable_cueq=False, -) - -# Convert atoms to state -state = ts.io.atoms_to_state(atoms_list, device=DEVICE, dtype=DTYPE) -# Run initial inference -initial_energies = model(state)["energy"] - - -def run_optimization_ts( # noqa: PLR0915 - *, - initial_state: SimState | OptimState, - ts_fire_flavor: Literal["vv_fire", "ase_fire"], - ts_use_frechet: bool, - force_tol: float, - max_iterations_ts: int, -) -> tuple[torch.Tensor, OptimState | None]: - """Runs torch-sim optimization and returns convergence steps and final state.""" - print( - f"\n--- Running torch-sim optimization: flavor={ts_fire_flavor}, " - f"frechet_cell_opt={ts_use_frechet}, force_tol={force_tol} ---" - ) - start_time = time.perf_counter() - - print("Initial cell parameters (torch-sim):") - for k_idx in range(initial_state.n_systems): - cell_tensor_k = initial_state.cell[k_idx].cpu().numpy() - ase_cell_k = Cell(cell_tensor_k) - params_str = ", ".join([f"{p:.2f}" for p in ase_cell_k.cellpar()]) - print( - f" Structure {k_idx + 1}: Volume={ase_cell_k.volume:.2f} ų, " - f"Params=[{params_str}]" - ) - - if ts_use_frechet: - init_fn_opt = partial(ts.fire_init, cell_filter=ts.CellFilter.frechet) - step_fn_opt = ts.fire_step - else: - init_fn_opt, step_fn_opt = ts.fire_init, ts.fire_step - - opt_state = init_fn_opt(model=model, state=initial_state.clone()) - - batcher = ts.InFlightAutoBatcher( - model=model, - memory_scales_with="n_atoms", - max_memory_scaler=1000, - max_iterations=max_iterations_ts, - ) - batcher.load_states(opt_state) - - total_structures = opt_state.n_systems - convergence_steps = torch.full( - (total_structures,), -1, dtype=torch.long, device=DEVICE - ) - convergence_fn = ts.generate_force_convergence_fn( - force_tol=force_tol, include_cell_forces=ts_use_frechet - ) - converged_tensor_global = torch.zeros( - total_structures, dtype=torch.bool, device=DEVICE - ) - global_step = 0 - all_converged_states = [] - convergence_tensor_for_batcher = None - last_active_state = opt_state - - while True: - opt_state, converged_states_from_batcher = batcher.next_batch( - last_active_state, convergence_tensor_for_batcher - ) - all_converged_states.extend(converged_states_from_batcher) - - if opt_state is None: - print("All structures converged or batcher reached max iterations.") - break - - last_active_state = opt_state - current_indices = torch.tensor( - batcher.current_idx, dtype=torch.long, device=DEVICE - ) - - steps_this_round = 1 - for _ in range(steps_this_round): - opt_state = step_fn_opt(model=model, state=opt_state) - global_step += steps_this_round - - convergence_tensor_for_batcher = convergence_fn(opt_state, None) - newly_converged_mask_local = convergence_tensor_for_batcher & ( - convergence_steps[current_indices] == -1 - ) - converged_indices_global = current_indices[newly_converged_mask_local] - - if converged_indices_global.numel() > 0: - convergence_steps[converged_indices_global] = global_step - converged_tensor_global[converged_indices_global] = True - total_converged_frac = converged_tensor_global.sum().item() / total_structures - print( - f"{global_step=}: Converged indices {converged_indices_global.tolist()}, " - f"Total converged: {total_converged_frac:.2%}" - ) - - if global_step % 50 == 0: - total_converged_frac = converged_tensor_global.sum().item() / total_structures - active_structures = opt_state.n_systems if opt_state else 0 - print( - f"{global_step=}: Active structures: {active_structures}, " - f"Total converged: {total_converged_frac:.2%}" - ) - - final_states_list = batcher.restore_original_order(all_converged_states) - final_state_concatenated = ts.concatenate_states(final_states_list) - - if final_state_concatenated is not None and hasattr(final_state_concatenated, "cell"): - print("Final cell parameters (torch-sim):") - for k_idx in range(final_state_concatenated.n_systems): - cell_tensor_k = final_state_concatenated.cell[k_idx].cpu().numpy() - ase_cell_k = Cell(cell_tensor_k) - params_str = ", ".join([f"{p:.2f}" for p in ase_cell_k.cellpar()]) - print( - f" Structure {k_idx + 1}: Volume={ase_cell_k.volume:.2f} ų, " - f"Params=[{params_str}]" - ) - else: - print( - "Final cell parameters (torch-sim): Not available (final_state_concatenated " - "is None or has no cell)." - ) - - end_time = time.perf_counter() - print( - f"Finished torch-sim ({ts_fire_flavor}, frechet={ts_use_frechet}) in " - f"{end_time - start_time:.2f} seconds." - ) - return convergence_steps, final_state_concatenated - - -def run_optimization_ase( # noqa: C901, PLR0915 - *, - initial_state: SimState, - ase_use_frechet_filter: bool, - force_tol: float, - max_steps_ase: int, -) -> tuple[torch.Tensor, OptimState | None]: - """Runs ASE optimization and returns convergence steps and final state.""" - print( - f"\n--- Running ASE optimization: frechet_filter={ase_use_frechet_filter}, " - f"force_tol={force_tol} ---" - ) - start_time = time.perf_counter() - - individual_initial_states = initial_state.split() - num_structures = len(individual_initial_states) - final_ase_atoms_list = [] - convergence_steps_list = [] - - for sys_idx, single_sim_state in enumerate(individual_initial_states): - print(f"Optimizing structure {sys_idx + 1}/{num_structures} with ASE...") - ase_atoms_orig = ts.io.state_to_atoms(single_sim_state)[0] - - initial_cell_ase = ase_atoms_orig.get_cell() - initial_params_str = ", ".join([f"{p:.2f}" for p in initial_cell_ase.cellpar()]) - print( - f" Initial cell (ASE Structure {sys_idx + 1}): " - f"Volume={initial_cell_ase.volume:.2f} ų, " - f"Params=[{initial_params_str}]" - ) - - ase_calc_instance = mace_mp_calculator_for_ase( - model=MaceUrls.mace_mpa_medium, - device=str(DEVICE), - default_dtype=str(DTYPE).removeprefix("torch."), - ) - ase_atoms_orig.calc = ase_calc_instance - - optim_target_atoms = ase_atoms_orig - if ase_use_frechet_filter: - print(f"Applying FrechetCellFilter to structure {sys_idx + 1}") - optim_target_atoms = FrechetCellFilter(ase_atoms_orig) - - dyn = ASEFIRE(optim_target_atoms, trajectory=None, logfile=None) - - try: - dyn.run(fmax=force_tol, steps=max_steps_ase) - if dyn.converged(): - convergence_steps_list.append(dyn.nsteps) - print(f"ASE structure {sys_idx + 1} converged in {dyn.nsteps} steps.") - else: - print( - f"ASE optimization for structure {sys_idx + 1} did not converge " - f"within {max_steps_ase} steps. Steps taken: {dyn.nsteps}." - ) - convergence_steps_list.append(-1) - except Exception as exc: # noqa: BLE001 - print(f"ASE optimization failed for structure {sys_idx + 1}: {exc}") - convergence_steps_list.append(-1) - - final_ats_for_print = getattr(optim_target_atoms, "atoms", ase_atoms_orig) - final_cell_ase = final_ats_for_print.get_cell() - final_params_str = ", ".join([f"{p:.2f}" for p in final_cell_ase.cellpar()]) - print( - f" Final cell (ASE Structure {sys_idx + 1}): " - f"Volume={final_cell_ase.volume:.2f} ų, " - f"Params=[{final_params_str}]" - ) - - final_ase_atoms_list.append(final_ats_for_print) - - all_positions: list[torch.Tensor] = [] - all_masses: list[torch.Tensor] = [] - all_atomic_numbers: list[torch.Tensor] = [] - all_cells: list[torch.Tensor] = [] - all_systems_for_gd: list[torch.Tensor] = [] - final_energies_ase: list[float] = [] - final_forces_ase_tensors: list[torch.Tensor] = [] - - current_atom_offset = 0 - for sys_idx, ats_final in enumerate(final_ase_atoms_list): - all_positions.append( - torch.tensor(ats_final.get_positions(), device=DEVICE, dtype=DTYPE) - ) - all_masses.append( - torch.tensor(ats_final.get_masses(), device=DEVICE, dtype=DTYPE) - ) - all_atomic_numbers.append( - torch.tensor(ats_final.get_atomic_numbers(), device=DEVICE, dtype=torch.long) - ) - # ASE cell is row-vector, SimState expects column-vector - all_cells.append( - torch.tensor(ats_final.get_cell().array.T, device=DEVICE, dtype=DTYPE) - ) - - num_atoms_in_current = len(ats_final) - all_systems_for_gd.append( - torch.full((num_atoms_in_current,), sys_idx, device=DEVICE, dtype=torch.long) - ) - current_atom_offset += num_atoms_in_current - - try: - if ats_final.calc is None: - print( - "Re-attaching ASE calculator for final energy/forces for structure " - f"{sys_idx}." - ) - temp_calc = mace_mp_calculator_for_ase( - model=MaceUrls.mace_mpa_medium, - device=str(DEVICE), - default_dtype=str(DTYPE).removeprefix("torch."), - ) - ats_final.calc = temp_calc - final_energies_ase.append(ats_final.get_potential_energy()) - final_forces_ase_tensors.append( - torch.tensor(ats_final.get_forces(), device=DEVICE, dtype=DTYPE) - ) - except Exception as exc: # noqa: BLE001 - print(f"Couldn't get final energy/forces for ASE structure {sys_idx}: {exc}") - final_energies_ase.append(float("nan")) - if all_positions and len(all_positions[-1]) > 0: - final_forces_ase_tensors.append(torch.zeros_like(all_positions[-1])) - else: - final_forces_ase_tensors.append( - torch.empty((0, 3), device=DEVICE, dtype=DTYPE) - ) - - if not all_positions: # If all optimizations failed early - print("Warning: No successful ASE structures to form OptimState.") - return ( - torch.tensor(convergence_steps_list, dtype=torch.long, device=DEVICE), - None, - ) - - # Concatenate all parts - concatenated_positions = torch.cat(all_positions, dim=0) - concatenated_masses = torch.cat(all_masses, dim=0) - concatenated_atomic_numbers = torch.cat(all_atomic_numbers, dim=0) - concatenated_cells = torch.stack(all_cells, dim=0) # Cells are (n_systems, 3, 3) - concatenated_system_indices = torch.cat(all_systems_for_gd, dim=0) - - concatenated_energies = torch.tensor(final_energies_ase, device=DEVICE, dtype=DTYPE) - concatenated_forces = torch.cat(final_forces_ase_tensors, dim=0) - - # Check for NaN energies which might cause issues - if torch.isnan(concatenated_energies).any(): - print( - "Warning: NaN values found in final ASE energies. " - "OptimState energy tensor will contain NaNs." - ) - - # Create OptimState instance - final_state_as_gd = OptimState( - positions=concatenated_positions, - masses=concatenated_masses, - cell=concatenated_cells, - pbc=initial_state.pbc, - atomic_numbers=concatenated_atomic_numbers, - system_idx=concatenated_system_indices, - energy=concatenated_energies, - forces=concatenated_forces, - stress=None, - ) - - convergence_steps = torch.tensor( - convergence_steps_list, dtype=torch.long, device=DEVICE - ) - - end_time = time.perf_counter() - print( - f"Finished ASE optimization (frechet_filter={ase_use_frechet_filter}) " - f"in {end_time - start_time:.2f} seconds." - ) - return convergence_steps, final_state_as_gd - - -# --- Main Script --- -force_tol = 0.05 - -# Configurations to test -configs_to_run = [ - { - "name": "torch-sim VV-FIRE (PosOnly)", - "type": "torch-sim", - "ts_fire_flavor": "vv_fire", - "ts_use_frechet": False, - }, - { - "name": "torch-sim ASE-FIRE (PosOnly)", - "type": "torch-sim", - "ts_fire_flavor": "ase_fire", - "ts_use_frechet": False, - }, - { - "name": "torch-sim VV-FIRE (Frechet Cell)", - "type": "torch-sim", - "ts_fire_flavor": "vv_fire", - "ts_use_frechet": True, - }, - { - "name": "torch-sim ASE-FIRE (Frechet Cell)", - "type": "torch-sim", - "ts_fire_flavor": "ase_fire", - "ts_use_frechet": True, - }, - { - "name": "ASE FIRE (Native, PosOnly)", - "type": "ase", - "ase_use_frechet_filter": False, - }, - { - "name": "ASE FIRE (Native Frechet Filter, CellOpt)", - "type": "ase", - "ase_use_frechet_filter": True, - }, -] - - -class ResultData(TypedDict): - """Result data for a single optimization run.""" - - steps: torch.Tensor - final_state: OptimState | None - - -all_results: dict[str, ResultData] = {} - -for config_run in configs_to_run: - print(f"\n\nStarting configuration: {config_run['name']}") - optimizer_type_val = config_run["type"] - ts_fire_flavor_val = config_run.get("ts_fire_flavor") - ts_use_frechet_val = config_run.get("ts_use_frechet", False) - ase_use_frechet_filter_val = config_run.get("ase_use_frechet_filter", False) - - steps: torch.Tensor | None = None - final_state_opt: OptimState | None = None - - if optimizer_type_val == "torch-sim": - if ts_fire_flavor_val is None: - raise ValueError(f"{ts_fire_flavor_val=} must be provided for torch-sim") - steps, final_state_opt = run_optimization_ts( - initial_state=state.clone(), - ts_fire_flavor=ts_fire_flavor_val, - ts_use_frechet=ts_use_frechet_val, - force_tol=force_tol, - max_iterations_ts=max_iterations, - ) - elif optimizer_type_val == "ase": - steps, final_state_opt = run_optimization_ase( - initial_state=state.clone(), - ase_use_frechet_filter=ase_use_frechet_filter_val, - force_tol=force_tol, - max_steps_ase=ase_max_optimizer_steps, - ) - else: - raise ValueError(f"Unknown optimizer_type: {optimizer_type_val}") - - all_results[config_run["name"]] = {"steps": steps, "final_state": final_state_opt} - - -print("\n\n--- Overall Comparison ---") -print(f"{force_tol=:.2f} eV/Å") -print(f"Initial energies: {[f'{e.item():.3f}' for e in initial_energies]} eV") - -for name, result_data in all_results.items(): - final_state_res = result_data["final_state"] - steps_res = result_data["steps"] - print(f"\nResults for: {name}") - if ( - final_state_res is not None - and hasattr(final_state_res, "energy") - and final_state_res.energy is not None - ): - energy_str = [f"{e.item():.3f}" for e in final_state_res.energy] - print(f" Final energies: {energy_str} eV") - else: - print(" Final energies: Not available or state is None") - print(f" Convergence steps: {steps_res.tolist()}") - - not_converged_indices = torch.where(steps_res == -1)[0].tolist() - if not_converged_indices: - print(f" Did not converge for structure indices: {not_converged_indices}") - -comparison_pairs = [ - ("torch-sim ASE-FIRE (PosOnly)", "ASE FIRE (Native, PosOnly)"), - ("torch-sim ASE-FIRE (Frechet Cell)", "ASE FIRE (Native Frechet Filter, CellOpt)"), - ("torch-sim VV-FIRE (Frechet Cell)", "ASE FIRE (Native Frechet Filter, CellOpt)"), - ("torch-sim VV-FIRE (PosOnly)", "ASE FIRE (Native, PosOnly)"), -] - -for name1, name2 in comparison_pairs: - if name1 in all_results and name2 in all_results: - state1 = all_results[name1]["final_state"] - state2 = all_results[name2]["final_state"] - - if state1 is None or state2 is None: - print(f"\nCannot compare {name1} and {name2}, one or both states are None.") - continue - - state1_list = state1.split() - - state2_list = state2.split() - - if len(state1_list) != len(state2_list): - print( - f"\nCannot compare {name1} and {name2}, different number of structures." - ) - continue - - mean_displacements = [] - for s1, s2 in zip(state1_list, state2_list, strict=True): - if 0 in {s1.n_atoms, s2.n_atoms}: - mean_displacements.append(float("nan")) - continue - pos1_centered = s1.positions - s1.positions.mean(dim=0, keepdim=True) - pos2_centered = s2.positions - s2.positions.mean(dim=0, keepdim=True) - if pos1_centered.shape != pos2_centered.shape: - print( - f"Warning: Shape mismatch for {name1} vs {name2} in structure. " - "Skipping displacement calc." - ) - mean_displacements.append(float("nan")) - continue - displacement = torch.norm(pos1_centered - pos2_centered, dim=1) - mean_disp = torch.mean(displacement).item() - mean_displacements.append(mean_disp) - - print( - f"\nMean Disp ({name1} vs {name2}): " - f"{[f'{d:.4f}' for d in mean_displacements]} Å" - ) - else: - print( - f"\nSkipping displacement comparison for ({name1} vs {name2}), " - "one or both results missing." - ) - - -# --- Plotting Results --- -original_structure_formulas = [ats.get_chemical_formula() for ats in atoms_list] -structure_names = [ - "Si_bulk", - "Cu_bulk", - "Fe_bulk", - "Si_vac", - "Cu_vac", - "Fe_vac", -] -if len(structure_names) != len(atoms_list): - print( - f"Warning: Mismatch between custom structure_names ({len(structure_names)}) and " - f"atoms_list ({len(atoms_list)}). Using custom names." - ) -num_structures_plot = len(structure_names) - - -# --- Plot 1: Convergence Steps (Multi-bar per structure) --- -plot_methods_fig1 = list(all_results) -num_methods_fig1 = len(plot_methods_fig1) -steps_data_fig1 = np.full((num_structures_plot, num_methods_fig1), np.nan) - -for method_idx, method_name in enumerate(plot_methods_fig1): - result_data = all_results[method_name] - if result_data["final_state"] is None or result_data["steps"] is None: - print(f"Plot1: Skipping steps for {method_name} as final_state or steps is None.") - continue - - steps_tensor = result_data["steps"].cpu().numpy() - penalty_steps = ase_max_optimizer_steps + 100 - steps_plot_values = np.where(steps_tensor == -1, penalty_steps, steps_tensor) - - if len(steps_plot_values) == num_structures_plot: - steps_data_fig1[:, method_idx] = steps_plot_values - elif len(steps_plot_values) > num_structures_plot: - print( - f"Warning: More step values ({len(steps_plot_values)}) than " - f"structure names ({num_structures_plot}) for {method_name}. " - "Truncating." - ) - steps_data_fig1[:, method_idx] = steps_plot_values[:num_structures_plot] - elif len(steps_plot_values) < num_structures_plot: - print( - f"Warning: Fewer step values ({len(steps_plot_values)}) than " - f"structure names ({num_structures_plot}) for {method_name}. " - "Padding with NaN." - ) - steps_data_fig1[: len(steps_plot_values), method_idx] = steps_plot_values - -fig1_plotly = go.Figure() - -for i in range(num_methods_fig1): - fig1_plotly.add_bar( - name=plot_methods_fig1[i], - x=structure_names, - y=steps_data_fig1[:, i], - text=[ - ( - "NC" - if all_results[plot_methods_fig1[i]]["steps"].cpu().numpy()[bar_idx] == -1 - and not np.isnan(steps_data_fig1[bar_idx, i]) - else "" - ) - for bar_idx in range(num_structures_plot) - ], - textposition="inside", - insidetextanchor="middle", - textfont=dict(color="white", size=10, family="Arial, sans-serif"), - ) - -fig1_plotly.update_layout( - barmode="group", - title_text="Convergence Steps per Structure and Method", - xaxis_title="Structure", - yaxis_title="Convergence Steps (NC = Not Converged, shown at penalty)", - legend_title="Optimization Method", - xaxis_tickangle=-45, - yaxis_gridcolor="lightgrey", - plot_bgcolor="white", - height=600, - width=max(1000, 150 * num_structures_plot), - margin=dict(l=50, r=50, b=100, t=50, pad=4), - legend=dict(orientation="h", yanchor="bottom", y=1.02, xanchor="right", x=1), -) -fig1_plotly.update_xaxes(categoryorder="array", categoryarray=structure_names) - - -# --- Plot 2: Average Final Energy Difference from Baselines --- -baseline_ase_pos_only = "ASE FIRE (Native, PosOnly)" -baseline_ase_frechet = "ASE FIRE (Native Frechet Filter, CellOpt)" -avg_energy_diffs_fig2 = [] -plot_names_fig2 = [] - -baseline_pos_only_data = all_results.get(baseline_ase_pos_only) -baseline_frechet_data = all_results.get(baseline_ase_frechet) - -for name, result_data in all_results.items(): - if result_data["final_state"] is None or result_data["final_state"].energy is None: - print(f"Plot2: Skipping energy diff for {name} as final_state or energy is None.") - if name not in plot_names_fig2: - plot_names_fig2.append(name) - avg_energy_diffs_fig2.append(np.nan) - continue - - if name not in plot_names_fig2: - plot_names_fig2.append(name) - - current_energies = result_data["final_state"].energy.cpu().numpy() - - chosen_baseline_energies = None - is_baseline_self = False - processed_current_name = False - - if name in (baseline_ase_pos_only, baseline_ase_frechet): - avg_energy_diffs_fig2.append(0.0) - is_baseline_self = True - processed_current_name = True - elif "torch-sim" in name: - if "PosOnly" in name: - if ( - baseline_pos_only_data - and baseline_pos_only_data["final_state"] is not None - and baseline_pos_only_data["final_state"].energy is not None - ): - chosen_baseline_energies = ( - baseline_pos_only_data["final_state"].energy.cpu().numpy() - ) - elif "Frechet Cell" in name and ( - baseline_frechet_data - and baseline_frechet_data["final_state"] is not None - and baseline_frechet_data["final_state"].energy is not None - ): - chosen_baseline_energies = ( - baseline_frechet_data["final_state"].energy.cpu().numpy() - ) - - if not is_baseline_self and not processed_current_name: - if chosen_baseline_energies is not None: - if current_energies.shape == chosen_baseline_energies.shape: - energy_diff = np.mean(current_energies - chosen_baseline_energies) - avg_energy_diffs_fig2.append(energy_diff) - else: - avg_energy_diffs_fig2.append(np.nan) - print( - f"Plot2: Shape mismatch for energy comparison: {name} vs baseline. " - f"{current_energies.shape} vs {chosen_baseline_energies.shape}" - ) - else: - print( - f"Plot2: No appropriate baseline for {name} or baseline data missing. " - "Setting energy diff to NaN." - ) - avg_energy_diffs_fig2.append(np.nan) - elif not processed_current_name and name not in ( - n - for n, v in zip(plot_names_fig2, avg_energy_diffs_fig2, strict=False) - if not np.isnan(v) - ): - print(f"Plot2: Fallback for {name}, setting energy diff to NaN.") - avg_energy_diffs_fig2.append(np.nan) - -final_plot_names_fig2 = [] -final_avg_energy_diffs_fig2 = [] -all_method_names_sorted = sorted(all_results) - -for name in all_method_names_sorted: - result_data = all_results[name] - final_plot_names_fig2.append(name) - if result_data["final_state"] is None or result_data["final_state"].energy is None: - final_avg_energy_diffs_fig2.append(np.nan) - continue - - current_energies = result_data["final_state"].energy.cpu().numpy() - energy_to_append = np.nan - - if name in (baseline_ase_pos_only, baseline_ase_frechet): - energy_to_append = 0.0 - elif "torch-sim" in name: - baseline_to_use_energies = None - if "PosOnly" in name: - if ( - baseline_pos_only_data - and baseline_pos_only_data["final_state"] is not None - and baseline_pos_only_data["final_state"].energy is not None - ): - baseline_to_use_energies = ( - baseline_pos_only_data["final_state"].energy.cpu().numpy() - ) - elif "Frechet Cell" in name and ( - baseline_frechet_data - and baseline_frechet_data["final_state"] is not None - and baseline_frechet_data["final_state"].energy is not None - ): - baseline_to_use_energies = ( - baseline_frechet_data["final_state"].energy.cpu().numpy() - ) - - if baseline_to_use_energies is not None: - if current_energies.shape == baseline_to_use_energies.shape: - energy_to_append = np.mean(current_energies - baseline_to_use_energies) - else: - print( - f"Plot2: Shape mismatch for {name} ({current_energies.shape}) " - f"vs baseline ({baseline_to_use_energies.shape})." - ) - final_avg_energy_diffs_fig2.append(energy_to_append) - - -fig2_plotly = go.Figure() -fig2_plotly.add_bar( - x=final_plot_names_fig2, - y=final_avg_energy_diffs_fig2, - marker_color="lightcoral", - text=[ - f"{yval:.3f}" if not np.isnan(yval) else "" - for yval in final_avg_energy_diffs_fig2 - ], - textposition="auto", - textfont=dict(size=10), -) - -line_dict = dict(color="black", width=1, dash="dash") -x1 = len(final_plot_names_fig2) - 0.5 -fig2_plotly.update_layout( - title_text="Average Final Energy Difference from ASE Baselines", - xaxis_title="Optimization Method", - yaxis_title="Avg. Final Energy Diff. from Corresponding ASE Baseline (eV)", - xaxis_tickangle=-45, - yaxis_gridcolor="lightgrey", - plot_bgcolor="white", - shapes=[dict(type="line", y0=0, y1=0, x0=-0.5, x1=x1, line=line_dict)], - height=600, - width=max(800, 100 * len(final_plot_names_fig2)), - margin=dict(l=50, r=50, b=150, t=50, pad=4), -) - - -# --- Plot 3: Mean Displacement from ASE Counterparts (Multi-bar per structure) --- -# look at sets of: (ts_method_name, ase_method_name, short_label_for_legend) -comparison_pairs_plot3_defs = [ - ( - "torch-sim ASE-FIRE (PosOnly)", - baseline_ase_pos_only, - "TS ASE PosOnly vs ASE Native", - ), - ( - "torch-sim VV-FIRE (PosOnly)", - baseline_ase_pos_only, - "TS VV PosOnly vs ASE Native", - ), - ( - "torch-sim ASE-FIRE (Frechet Cell)", - baseline_ase_frechet, - "TS ASE Frechet vs ASE Frechet", - ), - ( - "torch-sim VV-FIRE (Frechet Cell)", - baseline_ase_frechet, - "TS VV Frechet vs ASE Frechet", - ), -] -num_comparison_pairs_plot3 = len(comparison_pairs_plot3_defs) -# rows: structures, cols: comparison_pair -disp_data_fig3 = np.full((num_structures_plot, num_comparison_pairs_plot3), np.nan) -legend_labels_fig3 = [] - -for pair_idx, (ts_method_name, ase_method_name, plot_label) in enumerate( - comparison_pairs_plot3_defs -): - legend_labels_fig3.append(plot_label) - if ts_method_name in all_results and ase_method_name in all_results: - state1_data = all_results[ts_method_name] - state2_data = all_results[ase_method_name] - - if state1_data["final_state"] is None or state2_data["final_state"] is None: - print( - f"Plot3: Skipping displacement for {plot_label} due to missing state data" - ) - continue - - state1_list = state1_data["final_state"].split() - state2_list = state2_data["final_state"].split() - - if ( - len(state1_list) != len(state2_list) - or len(state1_list) != num_structures_plot - ): - print( - f"Plot3: Structure count mismatch for {plot_label}. " - f"Expected {num_structures_plot}, got S1:{len(state1_list)}, " - f"S2:{len(state2_list)}" - ) - continue - - mean_displacements_for_this_pair = [] - for s1, s2 in zip(state1_list, state2_list, strict=True): - if s1.n_atoms == 0 or s2.n_atoms == 0 or s1.n_atoms != s2.n_atoms: - mean_displacements_for_this_pair.append(np.nan) - continue - pos1_centered = s1.positions - s1.positions.mean(dim=0, keepdim=True) - pos2_centered = s2.positions - s2.positions.mean(dim=0, keepdim=True) - displacement = torch.norm(pos1_centered - pos2_centered, dim=1) - mean_disp = torch.mean(displacement).item() - mean_displacements_for_this_pair.append(mean_disp) - - if len(mean_displacements_for_this_pair) == num_structures_plot: - disp_data_fig3[:, pair_idx] = np.array(mean_displacements_for_this_pair) - else: - print(f"Plot3: Inner loop displacement calculation mismatch for {plot_label}") - - else: - print(f"Plot3: Missing data for methods in pair: {plot_label}.") - -fig3_plotly = go.Figure() - -for idx, name in enumerate(legend_labels_fig3): - fig3_plotly.add_bar(name=name, x=structure_names, y=disp_data_fig3[:, idx]) - - -title = "Mean Displacement of torch-sim Methods to ASE Counterparts (per Structure)" -fig3_plotly.update_layout( - barmode="group", - title=dict(text=title, x=0.5, y=1), - xaxis_title="Structure", - yaxis_title="Mean Atomic Displacement (Å) to ASE Counterpart", - legend_title="Comparison Pair", - xaxis_tickangle=-45, - yaxis_gridcolor="lightgrey", - plot_bgcolor="white", - height=600, - width=max(1000, 150 * num_structures_plot), # Adjust width - margin=dict(l=50, r=50, b=100, t=50, pad=4), - legend=dict(orientation="h", yanchor="bottom", y=0.96, xanchor="right", x=1), -) - -# Show Plotly figures -fig1_plotly.show() -fig2_plotly.show() -fig3_plotly.show() diff --git a/examples/old_scripts/7_Others/7.7_Heat_flux_and_kappa.py b/examples/old_scripts/7_Others/7.7_Heat_flux_and_kappa.py deleted file mode 100644 index a1668750..00000000 --- a/examples/old_scripts/7_Others/7.7_Heat_flux_and_kappa.py +++ /dev/null @@ -1,151 +0,0 @@ -"""Heat flux and thermal conductivity example with Lennard-Jones potential.""" - -# /// script -# dependencies = [ -# "ase>=3.26", -# "matplotlib", -# "numpy", -# ] -# /// -import os - -import matplotlib.pyplot as plt -import numpy as np -import torch -from ase.build import bulk - -import torch_sim as ts -from torch_sim.elastic import full_3x3_to_voigt_6_stress -from torch_sim.models.lennard_jones import LennardJonesModel -from torch_sim.properties.correlations import HeatFluxAutoCorrelation -from torch_sim.units import MetalUnits as Units - - -SMOKE_TEST = os.getenv("CI") is not None - -device = torch.device("cuda" if torch.cuda.is_available() else "cpu") -dtype = torch.float64 - -# Using solid Ar w/ LJ for ease -atoms = bulk("Ar", crystalstructure="fcc", a=5.376, cubic=True) -N_repeats = 3 if SMOKE_TEST else 4 -atoms = atoms.repeat((N_repeats, N_repeats, N_repeats)) -state = ts.io.atoms_to_state(atoms, device=device, dtype=dtype) - -# Simulation parameters -# See https://docs.lammps.org/compute_heat_flux.html for more details -epsilon = 0.0104 # eV -sigma = 3.405 # Å -cutoff = 13 # Å -temperature = 70.0 # Kelvin -timestep = 0.004 # ps (4 fs) -num_steps_equilibration = 1000 if SMOKE_TEST else 8000 -num_steps_production = 2000 if SMOKE_TEST else 100000 -window_size = 200 # Length of correlation: dt * correlation_dt * window_size -correlation_dt = 10 # Step delta between correlations - -# Lennard-Jones model -lj_model = LennardJonesModel( - sigma=sigma, - epsilon=epsilon, - cutoff=cutoff, - device=device, - dtype=dtype, - compute_forces=True, - compute_stress=True, - per_atom_energies=True, - per_atom_stresses=True, -) - -dt = torch.tensor(timestep * Units.time, device=device, dtype=dtype) -kT = torch.tensor(temperature * Units.temperature, device=device, dtype=dtype) -state = ts.nvt_langevin_init(state=state, model=lj_model, kT=kT) - -# Short equilibration run -# Shape: (num_steps, batch, dim) -heat_flux = torch.zeros((num_steps_equilibration, 3), device=device, dtype=dtype) - -for step in range(num_steps_equilibration): - state = ts.nvt_langevin_step(state=state, model=lj_model, dt=dt, kT=kT) - results = lj_model(state) - J = ts.quantities.calc_heat_flux( - momenta=state.momenta, - masses=state.masses, - velocities=None, - energies=results["energies"], - stresses=full_3x3_to_voigt_6_stress(results["stresses"]), - batch=state.system_idx, - is_centroid_stress=False, - is_virial_only=False, - ) - heat_flux[step] = J - if step % 1000 == 0: - print(f"Step {step} | {state.energy.item():.4f} eV") - -state = ts.nvt_langevin_init(state=state, model=lj_model, kT=kT) - -hfacf_calc = HeatFluxAutoCorrelation( - model=lj_model, - window_size=window_size, - device=device, - use_running_average=True, - normalize=False, -) - -# Sampling freq is controlled by prop_calculators -# trajectory = "kappa_example.h5" - -reporter = ts.TrajectoryReporter( - None, # add trajectory name here if you want to save the trajectory to disk - state_frequency=100, - prop_calculators={correlation_dt: {"hfacf": hfacf_calc}}, -) - -# Short production run -for step in range(num_steps_production): - state = ts.nvt_langevin_step(state=state, model=lj_model, dt=dt, kT=kT) - reporter.report(state, step) - if step % 1000 == 0: - print(f"Step {step} | {state.energy.item():.4f} eV") - -reporter.close() - -# HFACF results and plot -# Timesteps -> Time in fs -time_steps = np.arange(window_size) -time_fs = time_steps * correlation_dt * timestep * 1000 -hface_numpy = hfacf_calc.hfacf.detach().cpu().numpy() - -# Calculate kappa -integral = np.trapezoid(hface_numpy) -constant = ( - state.volume.item() - / (3 * temperature * temperature * Units.temperature) - * timestep - * correlation_dt -) -kappa = constant * integral -print(f"kappa: {kappa:.8f} (eV/ps/Ang^2/K)") - -fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(12, 4)) - -ax1.plot(heat_flux[:, 0].detach().cpu().numpy(), "b-", linewidth=2, label=r"$J_x$") -ax1.plot(heat_flux[:, 1].detach().cpu().numpy(), "r-", linewidth=2, label=r"$J_y$") -ax1.plot(heat_flux[:, 2].detach().cpu().numpy(), "g-", linewidth=2, label=r"$J_z$") -ax1.set_xlabel("Time (fs)", fontsize=12) -ax1.set_ylabel(r"$J$ (eV/ps $\AA^2$)", fontsize=12) -ax1.set_title("Heat Flux for Ar (LJ)", fontsize=14) -ax1.axhline(y=0, color="k", linestyle="--", alpha=0.3) -ax1.legend(fontsize=12) - -ax2.plot(time_fs, hface_numpy, "b-", linewidth=2) -ax2.set_xlabel("Time (fs)", fontsize=12) -ax2.set_ylabel(r"$\langle \vec{J}(0) \cdot \vec{J}(t) \rangle$", fontsize=12) -ax2.set_title( - rf"$\kappa$ = {kappa:.8f} (eV/ps $\AA^2$ K) (Average of {hfacf_calc._window_count} windows)", # noqa: E501, SLF001 - fontsize=14, -) -ax2.axhline(y=0, color="k", linestyle="--", alpha=0.3) - -plt.tight_layout() -plt.savefig("heat_flux_and_kappa.pdf") diff --git a/examples/old_scripts/readme.md b/examples/old_scripts/readme.md deleted file mode 100644 index 118977c3..00000000 --- a/examples/old_scripts/readme.md +++ /dev/null @@ -1,75 +0,0 @@ -# TorchSim Example Scripts - -This folder contains a series of examples demonstrating the use of TorchSim, a library for simulating molecular dynamics and structural optimization using classical and machine learning interatomic potentials. Each example showcases different functionalities and models available in TorchSim. - -1. **Introduction** - - 1. **Lennard-Jones Model** - [`examples/1_Introduction/1.1_Lennard_Jones.py`](1_Introduction/1.1_Lennard_Jones.py): Simulate Argon atoms in an FCC lattice with a Lennard-Jones potential. Initialize the model, run a forward pass, and print energy, forces, and stress. - - 1. **MACE Model** - [`examples/1_Introduction/1.2_MACE.py`](1_Introduction/1.2_MACE.py): Use the MACE model to simulate diamond cubic Silicon. Load a pre-trained model, set up the system, and calculate energy, forces, and stress. - - 1. **Batched MACE Model** - [`examples/1_Introduction/1.3_Batched_MACE.py`](1_Introduction/1.3_Batched_MACE.py): Handle batched inputs with the MACE model to simulate multiple systems simultaneously. - - 1. **Fairchem Model** - [`examples/1_Introduction/1.4_Fairchem.py`](1_Introduction/1.4_Fairchem.py): Simulate diamond cubic Silicon with the Fairchem model. Set up the model and calculate energy, forces, and stress. - -1. **Structural Optimization** - - 1. **Lennard-Jones FIRE** - [`examples/2_Structural_optimization/2.1_Lennard_Jones_FIRE.py`](2_Structural_optimization/2.1_Lennard_Jones_FIRE.py): Perform structural optimization using the FIRE optimizer with a Lennard-Jones model. - - 1. **Soft Sphere FIRE** - [`examples/2_Structural_optimization/2.2_Soft_Sphere_FIRE.py`](2_Structural_optimization/2.2_Soft_Sphere_FIRE.py): Optimize structures with a Soft Sphere model using the FIRE optimizer. - - 1. **MACE FIRE** - [`examples/2_Structural_optimization/2.3_MACE_FIRE.py`](2_Structural_optimization/2.3_MACE_FIRE.py): Optimize structures with the MACE model using the FIRE optimizer. - - 1. **MACE UnitCellFilter FIRE** - [`examples/2_Structural_optimization/2.4_MACE_UnitCellFilter_FIRE.py`](2_Structural_optimization/2.4_MACE_UnitCellFilter_FIRE.py): Optimize structures with the MACE model using the UnitCellFilter FIRE optimizer. - - 1. **MACE FrechetCellFilter FIRE** - [`examples/2_Structural_optimization/2.5_MACE_FrechetCellFilter_FIRE.py`](2_Structural_optimization/2.5_MACE_FrechetCellFilter_FIRE.py): Optimize structures with the MACE model using the FrechetCellFilter FIRE optimizer. - - 1. **Batched MACE Gradient Descent** - [`examples/2_Structural_optimization/2.6_Batched_MACE_Gradient_Descent.py`](2_Structural_optimization/2.6_Batched_MACE_Gradient_Descent.py): Optimize multiple structures simultaneously using batched gradient descent with the MACE model. - - 1. **Batched MACE FIRE** - [`examples/2_Structural_optimization/2.7_Batched_MACE_FIRE.py`](2_Structural_optimization/2.7_Batched_MACE_FIRE.py): Optimize multiple structures simultaneously using the batched FIRE optimizer with MACE. - - 1. **Batched MACE UnitCellFilter Gradient Descent** - [`examples/2_Structural_optimization/2.8_Batched_MACE_UnitCellFilter_Gradient_Descent.py`](2_Structural_optimization/2.8_Batched_MACE_UnitCellFilter_Gradient_Descent.py): Optimize multiple structures and their unit cells using batched gradient descent with MACE. - - 1. **Batched MACE UnitCellFilter FIRE** - [`examples/2_Structural_optimization/2.9_Batched_MACE_UnitCellFilter_FIRE.py`](2_Structural_optimization/2.9_Batched_MACE_UnitCellFilter_FIRE.py): Optimize multiple structures and their unit cells using the batched FIRE optimizer with MACE. - -1. **Dynamics** - - 1. **Lennard-Jones NVE** - [`examples/3_Dynamics/3.1_Lennard_Jones_NVE.py`](3_Dynamics/3.1_Lennard_Jones_NVE.py): Run molecular dynamics with the NVE ensemble using a Lennard-Jones model. Set up, simulate, and check energy conservation. - - 1. **MACE NVE** - [`examples/3_Dynamics/3.2_MACE_NVE.py`](3_Dynamics/3.2_MACE_NVE.py): Run NVE molecular dynamics simulation with the MACE model. - - 1. **MACE NVE with Cueq** - [`examples/3_Dynamics/3.3_MACE_NVE_cueq.py`](3_Dynamics/3.3_MACE_NVE_cueq.py): Run the MACE model in NVE with CuEq acceleration. - - 1. **MACE NVT Langevin** - [`examples/3_Dynamics/3.4_MACE_NVT_Langevin.py`](3_Dynamics/3.4_MACE_NVT_Langevin.py): Run temperature-controlled molecular dynamics using the NVT Langevin integrator with MACE. - - 1. **MACE NVT Nose-Hoover** - [`examples/3_Dynamics/3.5_MACE_NVT_Nose_Hoover.py`](3_Dynamics/3.5_MACE_NVT_Nose_Hoover.py): Run temperature-controlled molecular dynamics using the NVT Nose-Hoover integrator with MACE. - - 1. **MACE NVT Nose-Hoover with Temperature Profile** - [`examples/3_Dynamics/3.6_MACE_NVT_Nose_Hoover_temp_profile.py`](3_Dynamics/3.6_MACE_NVT_Nose_Hoover_temp_profile.py): Simulate heating and cooling cycles using Nose-Hoover integrator with a temperature profile. - - 1. **Lennard-Jones NPT Nose-Hoover** - [`examples/3_Dynamics/3.7_Lennard_Jones_NPT_Nose_Hoover.py`](3_Dynamics/3.7_Lennard_Jones_NPT_Nose_Hoover.py): Run pressure-controlled molecular dynamics using the NPT Nose-Hoover integrator with Lennard-Jones. - - 1. **MACE NPT Nose-Hoover** - [`examples/3_Dynamics/3.8_MACE_NPT_Nose_Hoover.py`](3_Dynamics/3.8_MACE_NPT_Nose_Hoover.py): Run pressure-controlled molecular dynamics using the NPT Nose-Hoover integrator with MACE. - - 1. **MACE NVT with Staggered Stress** - [`examples/3_Dynamics/3.9_MACE_NVT_staggered_stress.py`](3_Dynamics/3.9_MACE_NVT_staggered_stress.py): Use staggered stress calculations during NVT simulations with the MACE model. - - 1. **Hybrid Swap Monte Carlo** - [`examples/3_Dynamics/3.10_Hybrid_swap_mc.py`](3_Dynamics/3.10_Hybrid_swap_mc.py): Combine molecular dynamics with Monte Carlo simulations using the MACE model. - -1. **High-Level API** - - 1. **High-Level API** - [`examples/4_High_level_api/4.1_high_level_api.py`](4_High_level_api/4.1_high_level_api.py): Integrate systems using the high-level API with different models and integrators. - -1. **Workflow** - - 1. **Workflow** - [`examples/5_Workflow/5.1_a2c_silicon.py`](5_Workflow/5.1_a2c_silicon.py): Run the a2c workflow with the MACE model. - - 1. **Workflow** - [`examples/5_Workflow/5.4_Elastic.py`](5_Workflow/5.4_Elastic.py): Calculate elastic tensor, bulk modulus and shear modulus with MACE. - -1. **Phonons** - - 1. **Phonon DOS with MACE Batched** - [`examples/6_Phonons/6.1_Phonons_MACE.py`](6_Phonons/6.1_Phonons_MACE.py): Calculate DOS and band structure with MACE, batching over FC2 calculations. - - 1. **Thermal Conductivity with MACE** - [`examples/6_Phonons/6.2_QuasiHarmonic_MACE.py`](6_Phonons/6.2_QuasiHarmonic_MACE.py): Calculates quasi-harmonic properties with MACE, batching over volumes and FC2 calculations. - - 1. **Thermal Conductivity with MACE Batched** - [`examples/6_Phonons/6.3_Conductivity_MACE.py`](6_Phonons/6.3_Conductivity_MACE.py): Calculate the Wigner lattice conductivity with MACE, batching over FC2 and FC3 calculations. - -Each example is self-contained and can be run independently to explore TorchSim capabilities. The examples cover basic model setup to advanced simulation techniques, providing a comprehensive overview of the library's features. From deb147aa059f19995f63dd34d6df7d1e01ed34da Mon Sep 17 00:00:00 2001 From: orionarcher Date: Fri, 12 Dec 2025 15:45:55 -0500 Subject: [PATCH 3/3] try fixing errors --- examples/scripts/1_introduction.py | 22 +++++++++++++--------- examples/scripts/4_high_level_api.py | 7 ++++++- 2 files changed, 19 insertions(+), 10 deletions(-) diff --git a/examples/scripts/1_introduction.py b/examples/scripts/1_introduction.py index 40b7196f..97557dc6 100644 --- a/examples/scripts/1_introduction.py +++ b/examples/scripts/1_introduction.py @@ -183,16 +183,20 @@ # The stress has shape (n_systems, 3, 3) same as cell print(f"Stress shape: {results['stress'].shape}") -# Check if the energy, forces, and stress are the same for the Si system across the batch -print( - f"\nMax energy difference: {torch.max(torch.abs(results['energy'][0] - results['energy'][1]))}" -) -print( - f"Max forces difference: {torch.max(torch.abs(results['forces'][:8] - results['forces'][8:]))}" -) -print( - f"Max stress difference: {torch.max(torch.abs(results['stress'][0] - results['stress'][1]))}" +# Check if the energy, forces, and stress are the same for the Si system across batches +# Each system has 64 atoms (2x2x2 supercell of 8-atom Si diamond) +n_atoms_per_system = len(si_dc) +energy_diff = torch.max(torch.abs(results["energy"][0] - results["energy"][1])) +forces_diff = torch.max( + torch.abs( + results["forces"][:n_atoms_per_system] - results["forces"][n_atoms_per_system:] + ) ) +stress_diff = torch.max(torch.abs(results["stress"][0] - results["stress"][1])) + +print(f"\nMax energy difference: {energy_diff}") +print(f"Max forces difference: {forces_diff}") +print(f"Max stress difference: {stress_diff}") print("\n" + "=" * 70) print("Introduction examples completed!") diff --git a/examples/scripts/4_high_level_api.py b/examples/scripts/4_high_level_api.py index 6aec94c3..25c73081 100644 --- a/examples/scripts/4_high_level_api.py +++ b/examples/scripts/4_high_level_api.py @@ -101,7 +101,12 @@ with TorchSimTrajectory(trajectory_file) as traj: kinetic_energies = traj.get_array("kinetic_energy") potential_energies = traj.get_array("potential_energy") - final_energy = potential_energies[-1] + # Convert to scalar, handling both numpy arrays and tensors + final_energy = ( + potential_energies[-1].item() + if hasattr(potential_energies[-1], "item") + else float(potential_energies[-1]) + ) final_atoms = traj.get_atoms(-1) print(f"Final energy from trajectory: {final_energy:.4f} eV")