Predicting Guide
===============

This guide explains how to use trained IANN models for making predictions.


Using the ML Calculators
--------------------

The ``MLCalculator`` provides a convenient ASE calculator interface:

.. code-block:: python

   from iann.calculators import MLCalculator
   from ase.io import read

   # Create calculator with trained model
   calc = MLCalculator("model.pt")
   
   # Read structures
   images = read("test_structures.traj", ":")

   # Make predictions
   for atoms in images:
      atoms.calc = calc
      energy = atoms.get_potential_energy()
      forces = atoms.get_forces()
      print(f"Energy: {energy} eV")
      print(f"Forces: {forces} eV/Å")

The calculator automatically:

* Determines the model type from the saved state dict
* Use the model architecture
* Do prediction

The ``EnsembleCalculator`` provides a convenient ASE calculator interface with uncertainty for each structure:

.. code-block:: python

   from iann.calculators import EnsembleCalculator
   from ase.io import read

   # Create calculator with trained models
   model_paths = ['model_1.pt', 'model_2.pt']
   calc = EnsembleCalculator(model_paths)

   # Read structures
   images = read("test_structures.traj", ":")

   # Make predictions
   for atoms in images:
      atoms.calc = calc
      energy = atoms.get_potential_energy()
      forces = atoms.get_forces()
      ensemble = atoms.calc.get_ensemble()
      print(f"Energy: {energy} eV")
      print(f"Forces: {forces} eV/Å")
      print(f"Energy variance: {ensemble['energy_var']} eV")
      print(f"Forces variance: {ensemble['forces_var']} eV/Å")

The ``AtomicEnsembleCalculator`` provides a convenient ASE calculator interface with uncertainty for each structure and each atom:

.. code-block:: python

   from iann.calculators import AtomicEnsembleCalculator
   from ase.io import read

   # Create calculator with trained models
   model_paths = ['model_1.pt', 'model_2.pt']
   calc = AtomicEnsembleCalculator(model_paths)

   # Read structures
   images = read("test_structures.traj", ":")

   # Make predictions
   for atoms in images:
      atoms.calc = calc
      energy = atoms.get_potential_energy()
      forces = atoms.get_forces()
      ensemble = atoms.calc.get_ensemble()
      print(f"Energy: {energy} eV")
      print(f"Forces: {forces} eV/Å")
      print(f"Energy variance: {ensemble['energy_var']} eV")
      print(f"Forces variance: {ensemble['forces_var']} eV/Å")
      print(f"Atomic energy variance: {ensemble['atomic_energy_var']} eV")
      print(f"Atomic forces variance: {ensemble['atomic_forces_var']} eV/Å")


      
Batch Prediction
--------------

For efficient batch prediction:

.. code-block:: python

   from iann.data import AseDataset
   from torch.utils.data import DataLoader

   # Create dataset
   dataset = AseDataset("test_structures.traj")
   dataloader = DataLoader(dataset, batch_size=32)

   model = torch.load("model.pt")

   # Make predictions
   model.eval()
   for batch in dataloader:
      energy, forces = model(batch)
      # Process predictions...

Performance Tips
--------------

1. **Memory Management**

   * Use appropriate batch sizes
   * Clear GPU cache if needed

2. **Speed Optimization**

   * Enable CUDA if available
   * Use batch processing when possible
   * Consider model quantization for deployment

3. **Accuracy Considerations**

   * Check cutoff radius matches training
   * Verify atomic numbers are correct


Applications: Geometric structure optimization
------------------------------------------

The ASE optimizers like ``BFGS`` can be used to optimize the geometry of a structure:

.. code-block:: python

   from iann.calculators import MLCalculator
   from ase.optimize import BFGS
   from ase.io import read

   # load structure
   atoms = read("atoms.traj") 

   # Create calculator
   calc = MLCalculator("model.pt")
   atoms.calc = calc

   # Create optimizer
   optimizer = BFGS(atoms, trajectory="opt.traj")

   # Run optimization
   optimizer.run(fmax=0.01)


Applications: Molecular dynamics
----------------------------

The ASE thermostats like ``Langevin`` can be used to perform molecular dynamics:

.. code-block:: python  

   from iann.calculators import MLCalculator
   from ase.io import read, Trajectory
   from ase.md.langevin import Langevin
   from ase import units
   from ase.md.velocitydistribution import MaxwellBoltzmannDistribution
   from ase.md import MDLogger

   # load structure
   atoms = read("atoms.traj") 

   # Create calculator
   calc = MLCalculator("model.pt")
   atoms.calc = calc

   # Set temperature and timestep
   temperature = 300
   timestep = 0.1

   # Initialize velocities from Maxwell-Boltzmann distribution
   MaxwellBoltzmannDistribution(atoms, temperature_K=temperature) 

   # Create Langevin thermostat
   dyn = Langevin(atoms, timestep=timestep * units.fs, temperature_K=temperature, friction=0.01 / units.fs)

   # Log and save trajectory
   dyn.attach(MDLogger(dyn, atoms, 'ase_md.log', header=True, stress=False, peratom=False, mode="a"), interval=1)
   traj = Trajectory('ase_md.traj', 'a', atoms)
   dyn.attach(traj.write, interval=10)

   # Run dynamics
   dyn.run(2000)



Integration with Other Tools
-------------------------

IANN models can be used with various tools:

* ASE for other applications
* LAMMPS for molecular dynamics (see :doc:`lammps`)
* Custom scripts for specific applications

For more details on the API and advanced usage, see the :doc:`api` reference.