started mlip additions

asimtools/asimmodules/ase_md/ase_md.py

@@ -0,0 +1,236 @@
+#!/usr/bin/env python
+'''
+Describe the script briefly here. If this is a script that runs multiple steps,
+describe it here using reStructuredText to generate autodocs
+
+Cite the papers where the method/script was first introduced here as well
+
+Author: [email protected]
+'''
+
+from typing import Dict, Optional, Sequence
+from future import __annotations__
+import numpy as np
+import matplotlib.pyplot as plt
+import ase
+from ase.io.trajectory import Trajectory
+from ase.units import GPa, fs
+from ase.md.velocitydistribution import MaxwellBoltzmannDistribution
+from ase.md.langevin import Langevin
+from ase.md.npt import NPT
+from asimtools.calculators import load_calc
+from asimtools.utils import (
+    get_atoms,
+    get_images,
+)
+
+def langevin_nvt(
+    atoms: ase.Atoms,
+    temp: float,
+    nsteps: int,
+    traj_file: str = None,
+    friction: float = 1e-2,
+    timestep: float = 1*fs,
+):
+    """Does Langevin dynamics
+
+    :param atoms: atoms object
+    :type atoms: ase.Atoms
+    :param temp: Temperature in Kelvin
+    :type temp: float
+    :param nsteps: Number of steps to run
+    :type nsteps: int
+    :param traj_file: trajectory file name, defaults to None
+    :type traj_file: str, optional
+    :param friction: friction parameter, defaults to 1e-2
+    :type friction: float, optional
+    :param timestep: Timestep in ASE time units, defaults to 1*fs
+    :type timestep: float, optional
+    :return: final atoms object and trajectory object
+    :rtype: ase.Atoms, Trajectory
+    """
+    MaxwellBoltzmannDistribution(atoms, temperature_K=temp)
+    dyn = Langevin(
+        atoms,
+        timestep=timestep,
+        temperature_K=temp,
+        friction=friction,
+    )
+
+    if traj_file is not None:
+        traj = Trajectory(
+            traj_file,
+            atoms=atoms,
+            mode='w',
+            properties=['energy', 'forces', 'stress']
+        )
+        dyn.attach(traj.write)
+    dyn.run(nsteps)
+    return atoms, traj
+
+def npt(
+    atoms: ase.Atoms,
+    temp: float,
+    nsteps: int,
+    timestep: float,
+    externalstress: float = 0,
+    traj_file: str = None,
+    ttime: float = 25*fs,
+    pfactor: float = (75*fs)**2 * 14*GPa, #Replace 14 with bulk modulus of material
+):
+    """Does NPT dynamics
+
+    :param atoms: input atoms object
+    :type atoms: ase.Atoms
+    :param temp: Temperature in Kelvin
+    :type temp: float
+    :param nsteps: Number of steps to run
+    :type nsteps: int
+    :param timestep: Timestep in ASE time units
+    :type timestep: float
+    :param externalstress: External pressure to apply, defaults to 0
+    :type externalstress: float, optional
+    :param traj_file: trajectory file name, defaults to None
+    :type traj_file: str, optional
+    :param ttime: thermostat time constant, defaults to 25*fs
+    :type ttime: float, optional
+    :return: final atoms object and trajectory object
+    :rtype: ase.Atoms, Trajectory
+    """
+    MaxwellBoltzmannDistribution(atoms, temperature_K=temp)
+    dyn = NPT(
+        atoms,
+        timestep=timestep,
+        temperature_K=temp,
+        externalstress=externalstress,
+        ttime=ttime,
+        pfactor=pfactor,
+        # mask=np.diag([1, 1, 1]),
+    )
+    if traj_file is not None:
+        traj = Trajectory(
+            traj_file,
+            atoms=atoms,
+            mode='w',
+            properties=['energy', 'forces', 'stress'],
+        )
+        dyn.attach(traj.write)
+    dyn.run(nsteps)
+    traj = Trajectory(traj_file, 'r')
+    return atoms, traj
+
+def plot_thermo(
+    images: Dict,
+    props: Sequence = ['epot', 'temp', 'ekin', 'etot', 'press'],
+    prefix: str = 'ase_md',
+):
+    """Plots thermodynamic properties of a trajectory
+
+    :param images: Dictionary of images to read from
+    :type images: Dict
+    :param props: Properties to plot, defaults to ['epot', 'temp', 'ekin', 'etot', 'press']
+    :type props: Sequence, optional
+    :param prefix: Prefix for output file names, defaults to 'ase_md'
+    :type prefix: str, optional
+    """
+    atoms = get_images(**images)
+    prop_dict = {prop: [] for prop in props}
+
+    for a, atoms in enumerate(images):
+        epot = atoms.get_potential_energy()
+        ekin = atoms.get_kinetic_energy()
+        etot = epot+ekin
+        T = atoms.get_temperature()
+        prop_dict['epot'].append(epot)
+        prop_dict['ekin'].append(ekin)
+        prop_dict['etot'].append(etot)
+        prop_dict['temp'].append(T)
+        if 'press' in props:
+            pressure = np.mean(atoms.get_stress(include_ideal_gas=True)[:3])
+            prop_dict['press'].append(pressure)
+
+    stride = images.get('index', ':')
+
+    steps = np.arange(len(prop_dict['epots'])) * stride
+    for _, prop in enumerate(props):
+        _, ax = plt.subplots()
+        ax.plot(steps, prop_dict[prop])
+        ax.set_xlabel('Step')
+        ax.set_ylabel(prop)
+        plt.savefig(f'{prefix}_{prop}.png')
+
+    plt.close(fig='all')
+
+def ase_md(
+    calc_id: str,
+    image: Dict,
+    timestep: float,
+    temp: float,
+    dynamics: str = 'npt',
+    friction: float = 1e-2,
+    ttime: float = 25*fs,
+    pfactor: float = None, #(75*fs)**2 * 14*GPa, #14 is bulk modulus of material i.e. Li
+    nsteps: int = 100,
+    externalstress: float = 0,
+    plot: bool = True,
+) -> Dict:
+    """Runs ASE MD simulations. This is only recommended for small systems and
+    for testing. For larger systems, use LAMMPS or more purpose-built code
+
+    :param calc_id: calc_id specification
+    :type calc_id: str
+    :param image: Image specification, see :func:`asimtools.utils.get_atoms`
+    :type image: Dict
+    :param timestep: Timestep in ASE time units
+    :type timestep: float
+    :param temp: Temperature in Kelvin
+    :type temp: float
+    :param dynamics: Type of dynamics to run from 'nvt', 'langevin' and 'nvt',
+        defaults to 'npt'
+    :type dynamics: str, optional
+    :param friction: Friction parameter, defaults to 1e-2
+    :type friction: float, optional
+    :param ttime: Thermostat time constant, defaults to 25*fs
+    :type ttime: float, optional
+    :param pfactor: Pressure factor, defaults to None
+    :type pfactor: float, optional
+    :param nsteps: Number of steps to run, defaults to 100
+    :type nsteps: int, optional
+    :param externalstress: External stress to apply, defaults to 0
+    :type externalstress: float, optional
+    :param plot: Whether to plot thermodynamic properties, defaults to True
+    :type plot: bool, optional
+    :return: Results dictionary
+    :rtype: Dict
+    """
+
+    calc = load_calc(calc_id)
+    atoms = get_atoms(**image)
+    atoms.set_calculator(calc)
+
+    if dynamics == 'langevin':
+        atoms, _ = langevin_nvt(
+            atoms,
+            temp,
+            nsteps,
+            traj_file='output.traj',
+            timestep=timestep,
+            friction=friction,
+        )
+    elif dynamics == 'npt':
+        atoms, _ = npt(
+            atoms,
+            temp,
+            nsteps,
+            traj_file='output.traj',
+            timestep=timestep,
+            pfactor=pfactor,
+            externalstress=externalstress,
+            ttime=ttime,
+        )
+
+    if plot:
+        plot_thermo({'input_file': 'output.traj'})
+
+    results = {}
+    return results

asimtools/asimmodules/lammps/lammps.py

@@ -1,6 +1,6 @@
 #!/usr/bin/env python
 '''
-Runs a user defined lammps script or template 
+Runs a user defined lammps script or template. LAMMPS must be installed 
 
 Author: [email protected]
 '''

asimtools/asimmodules/mlips/active_learning.py

@@ -0,0 +1,8 @@
+from pathlib import Path
+from typing import Dict, Optional, Sequence
+from asimtools.job import UnitJob
+
+def active_learning(
+    image: Dict,
+) -> Dict:
+

asimtools/asimmodules/mlips/compute_deviation.py

@@ -0,0 +1,136 @@
+from typing import Dict, List, Sequence, Optional
+import os
+from glob import glob
+from natsort import natsorted
+import numpy as np
+import pandas as pd
+import matplotlib.pyplot as plt
+from asimtools.calculators import load_calc
+from asimtools.utils import (
+    change_dict_value,
+    get_images,
+)
+
+def compute_deviation(
+    images: Dict,
+    template_calc_input: Optional[Dict] = None,
+    model_weights_key_sequence: Optional[Sequence] = None,
+    model_weights_pattern: Optional[os.PathLike] = None,
+    calc_ids: Optional[Sequence] = None,
+) -> Dict:
+    """Computes variance of properties from a trajectory file
+
+    :param images: Images specification, see :func:`asimtools.utils.get_images`
+    :type images: Dict
+    :param template_calc_input: Template calc_input, defaults to None
+    :type template_calc_input: Optional[Dict]
+    :param model_weights_key_sequence: Sequence of keys to change in the
+        template calc_input
+    :type model_weights_key_sequence: Optional[Sequence]
+    :param model_weights_pattern: Pattern of model weights files, defaults to
+        None
+    :type model_weights_pattern: Optional[os.PathLike]
+    :param calc_ids: List of calc_ids to use, if provided, all other arguments
+        are ignored, defaults to None
+    :type calc_ids: Optional[Sequence]
+
+    """
+    properties = ['energy', 'forces', 'stress']
+    if calc_ids is None:
+        model_weights_files = natsorted(glob(model_weights_pattern))
+
+        calc_dict = {}
+        for i, model_weights_file in enumerate(model_weights_files):
+            new_calc_input = change_dict_value(
+                template_calc_input,
+                model_weights_file,
+                key_sequence=model_weights_key_sequence,
+                return_copy=True
+            )
+
+            calc_dict[f'calc-{i}'] = new_calc_input
+    else:
+        calc_dict = {calc_id: calc_id for calc_id in calc_ids}
+
+    variances = {prop: {} for prop in properties}
+
+    images = get_images(**images)
+    assert len(images) > 0, 'No images found'
+
+    prop_dict = {
+        prop: {calc_id: [] for calc_id in calc_dict} for prop in properties
+    }
+    for prop in properties:
+        if prop == 'forces':
+            prop_dict[prop]['mean_std'] = []
+            prop_dict[prop]['max_std'] = []
+        else:
+            prop_dict[prop]['std'] = []
+
+    for i, atoms in enumerate(images):
+        atom_results = {prop: [] for prop in properties}
+        for calc_id in calc_dict:
+            # Some calculators behave badly if not reloaded unfortunately
+            if isinstance(calc_dict[calc_id], str):
+                calc = load_calc(calc_id=calc_dict[calc_id])
+            else:
+                calc = load_calc(calc_input=calc_dict[calc_id])
+
+            atoms.set_calculator(calc)
+            energy = atoms.get_potential_energy(atoms)
+            forces = np.linalg.norm(atoms.get_forces(), axis=1)
+            stress = -np.sum(
+                atoms.get_stress(voigt=True, include_ideal_gas=False)[:3]
+            ) / 3
+
+            prop_dict['energy'][calc_id].append(energy)
+            prop_dict['forces'][calc_id].append(forces)
+            prop_dict['stress'][calc_id].append(stress)
+            atom_results['energy'].append(energy)
+            atom_results['forces'].append(forces)
+            atom_results['stress'].append(stress)
+
+        prop_dict['energy']['std'].append(np.std(atom_results['energy']))
+        prop_dict['forces']['mean_std'].append(
+            np.mean(np.std(atom_results['forces'], axis=1))
+        )
+        prop_dict['forces']['max_std'].append(
+            np.max(np.std(atom_results['forces'], axis=1))
+        )
+        prop_dict['stress']['std'].append(np.std(atom_results['stress']))
+
+    df = pd.DataFrame({
+        'energy_std': prop_dict['energy']['std'],
+        'force_mean_std': prop_dict['forces']['mean_std'],
+        'force_max_std': prop_dict['forces']['max_std'],
+        'stress_std': prop_dict['stress']['std']
+
+    })
+    df.to_csv('deviations.csv')
+
+    for prop in properties:
+        if prop not in ['forces']:
+            # df = pd.DataFrame(prop_dict[prop])
+            fig, ax = plt.subplots()
+            for calc_id in calc_dict:
+                ax.plot(prop_dict[prop][calc_id], label=calc_id)
+            ax.set_xlabel('Image index')
+            ax.set_ylabel(f'{prop} [ASE units]')
+            ax.legend()
+            plt.savefig(f'{prop}.png')
+            plt.close()
+
+        fig, ax = plt.subplots()
+        if prop == 'forces':
+            ax.plot(prop_dict[prop]['mean_std'], label='mean_std')
+            ax.plot(prop_dict[prop]['max_std'], label='max_std')
+            ax.legend()
+        else:
+            ax.plot(prop_dict[prop]['std'])
+        ax.set_xlabel('Image index')
+        ax.set_ylabel(f'{prop} std [ASE units]')
+        plt.savefig(f'{prop}_std.png')
+        plt.close()
+
+    return {}
+