ase_md enhancements, write_atoms, minor cleanup

Author: mkphuthi

asimtools/asimmodules/ase_md/ase_md.py

@@ -18,6 +18,7 @@
 from ase.md.velocitydistribution import MaxwellBoltzmannDistribution
 from ase.md.langevin import Langevin
 from ase.md.npt import NPT
+from ase.md import MDLogger
 from asimtools.calculators import load_calc
 from asimtools.utils import (
     get_atoms,
@@ -31,6 +32,9 @@ def langevin_nvt(
     traj_file: str = None,
     friction: float = 1e-2,
     timestep: float = 1*fs,
+    properties: Optional[Sequence] = ('energy', 'forces', 'stress'),
+    log_interval: int = 1,
+    traj_interval: Optional[int] = 1,
 ):
     """Does Langevin dynamics
 
@@ -64,7 +68,19 @@ def langevin_nvt(
             mode='w',
             properties=['energy', 'forces']
         )
-        dyn.attach(traj.write)
+        dyn.attach(traj.write, interval=traj_interval)
+    stress = False
+    if 'stress' in properties:
+        stress = True
+    dyn.attach(MDLogger(
+        dyn,
+        atoms,
+        'md.log',
+        header=True,
+        stress=stress,
+        peratom=True,
+        mode="a"
+    ), interval=log_interval)
     dyn.run(nsteps)
     return atoms, traj
 
@@ -78,6 +94,8 @@ def npt(
     ttime: float = 25*fs,
     pfactor: Optional[float] = None, # (75*fs)**2 * 14*GPa, #Replace 14 with bulk modulus of material
     properties: Optional[Sequence] = ('energy', 'forces', 'stress'),
+    log_interval: int = 1,
+    traj_interval: Optional[int] = 1,
 ):
     """Does NPT dynamics
 
@@ -116,7 +134,20 @@ def npt(
             mode='w',
             properties=properties,
         )
-        dyn.attach(traj.write)
+        dyn.attach(traj.write, interval=traj_interval)
+
+    stress = False
+    if 'stress' in properties:
+        stress = True
+    dyn.attach(MDLogger(
+        dyn,
+        atoms,
+        'md.log',
+        header=True,
+        stress=stress,
+        peratom=True,
+        mode="a"
+    ), interval=log_interval)
     dyn.run(nsteps)
     traj = Trajectory(traj_file, 'r')
     return atoms, traj
@@ -176,6 +207,9 @@ def ase_md(
     plot: Optional[bool] = True,
     time_unit: Optional[str] = 'ase',
     plot_args: Optional[dict] = None,
+    properties: Optional[Sequence] = ('energy', 'forces', 'stress'),
+    log_interval: Optional[int] = 1,
+    traj_interval: Optional[int] = 1,
 ) -> 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
@@ -209,7 +243,7 @@ def ase_md(
 
     calc = load_calc(**calc_spec)
     atoms = get_atoms(**image)
-    atoms.set_calculator(calc)
+    atoms.calc = calc
 
     if time_unit == 'fs':
         timestep *= fs
@@ -224,6 +258,8 @@ def ase_md(
             traj_file='output.traj',
             timestep=timestep,
             friction=friction,
+            log_interval=log_interval,
+            traj_interval=traj_interval,
         )
     elif dynamics == 'npt':
         atoms, _ = npt(
@@ -235,6 +271,9 @@ def ase_md(
             pfactor=pfactor,
             externalstress=externalstress,
             ttime=ttime,
+            properties=properties,
+            log_interval=log_interval,
+            traj_interval=traj_interval,
         )
     elif dynamics == 'nvt':
         atoms, _ = npt(
@@ -245,7 +284,9 @@ def ase_md(
             timestep=timestep,
             pfactor=None,
             ttime=ttime,
-            properties=['energy', 'forces'],
+            properties=properties,
+            log_interval=log_interval,
+            traj_interval=traj_interval,
         )
 
 

asimtools/asimmodules/elastic_constants/cubic_energy_expansion.py

@@ -101,7 +101,7 @@ def cubic_energy_expansion(
     """
     calc = load_calc(calc_id)
     atoms = get_atoms(**image)
-    atoms.set_calculator(calc)
+    atoms.calc = calc
 
     # Start by getting the Bulk modulus and optimized cell from the EOS
     logging.info('Calculating EOS')

asimtools/asimmodules/geometry_optimization/atom_relax.py

@@ -40,7 +40,7 @@ def atom_relax(
     """
     calc = load_calc(calc_id)
     atoms = get_atoms(**image)
-    atoms.set_calculator(calc)
+    atoms.calc = calc
     logger = get_logger()
 
     if prefix is not None:

asimtools/asimmodules/geometry_optimization/cell_relax.py

@@ -46,7 +46,7 @@ def cell_relax(
     """
     calc = load_calc(calc_id)
     atoms = get_atoms(**image)
-    atoms.set_calculator(calc)
+    atoms.calc = calc
 
     if prefix is not None:
         prefix = prefix + '_'

asimtools/asimmodules/geometry_optimization/optimize.py

@@ -46,7 +46,7 @@ def optimize(
 
     calc = load_calc(calc_id)
     atoms = get_atoms(**image)
-    atoms.set_calculator(calc)
+    atoms.calc = calc
 
     ecf = ExpCellFilter(atoms, **expcellfilter_args)
 

asimtools/asimmodules/geometry_optimization/symmetric_cell_relax.py

@@ -47,7 +47,7 @@ def symmetric_cell_relax(
 
     calc = load_calc(calc_id)
     atoms = get_atoms(**image)
-    atoms.set_calculator(calc)
+    atoms.calc = calc
 
     atoms.set_constraint(FixSymmetry(atoms, **fixsymmetry_args))
     ecf = ExpCellFilter(atoms, **expcellfilter_args)

asimtools/asimmodules/vacancy_formation_energy/vacancy_formation_energy.py

@@ -49,7 +49,7 @@ def vacancy_formation_energy(
 
     calc = load_calc(calc_id)
     bulk = get_atoms(**image).repeat(repeat)
-    bulk.set_calculator(calc)
+    bulk.calc = calc
 
     vacant = bulk.copy()
     del vacant[vacancy_index]

asimtools/job.py

@@ -21,6 +21,7 @@
 from asimtools.utils import (
     read_yaml,
     write_yaml,
+    write_atoms,
     join_names,
     get_atoms,
     get_images,
@@ -396,9 +397,9 @@ def gen_input_files(
         if image and write_image:
             atoms = get_atoms(**image)
             input_image_file = 'image_input.xyz' # Relative to workdir
-            atoms.write(
+            write_atoms(
                 self.workdir / input_image_file,
-                format='extxyz'
+                atoms,
             )
             sim_input['args']['image'] = {
                 'image_file': str(input_image_file),

asimtools/utils.py

@@ -16,7 +16,7 @@
 import numpy as np
 import matplotlib.pyplot as plt
 import pandas as pd
-from ase.io import read
+from ase.io import read, write
 from ase.parallel import paropen
 import ase.db
 import ase.build
@@ -149,6 +149,52 @@ def join_names(substrs: Sequence[str]) -> str:
     name = '__'.join(final_substrs) + '__'
     return name
 
+def write_atoms(
+    image_file: str,
+    atoms: Atoms,
+    fmt: str = 'extxyz',
+    write_info: bool = True,
+    columns: Optional[Sequence] = None,
+    **kwargs
+):
+    """
+    
+    """
+    if kwargs.get('format', False):
+        fmt = kwargs.pop('format')
+
+    if fmt in ['extxyz']:
+        if kwargs.get('write_info', False):
+            write_info = kwargs.pop('write_info')
+        if kwargs.get('columns', False):
+            write_info = kwargs.pop('columns')
+        else:
+            reserved_ks = ['symbols', 'positions', 'numbers', 'species', 'pos']
+            columns = ['symbols', 'positions'] + kwargs.get(
+                'columns',
+                [k for k in atoms.arrays.keys() if k not in reserved_ks]
+            )
+
+            if len(atoms.constraints) > 0:
+                columns.append('move_mask')
+
+        write(
+            image_file,
+            atoms,
+            format=fmt,
+            write_info=write_info,
+            columns=columns,
+            **kwargs
+        )
+    else:
+        write(
+            image_file,
+            atoms,
+            format=fmt,
+            **kwargs
+        )
+
+
 def get_atoms(
     image_file: Optional[str] = None,
     interface: str = 'ase',
@@ -742,11 +788,21 @@ def get_str_btn(
     s = s[start_index:stop_index]
     while occurence - j >= 0:
         if s1 is not None:
-            i1 = s.index(s1) + len(s1)
+            try:
+                i1 = s.index(s1) + len(s1)
+            except:
+                raise ValueError(
+                    f'substring {s1} not found in {s}'
+                )
         else:
             i1 = 0
         if s2 is not None:
-            i2 = s[i1:].index(s2) + i1
+            try:
+                i2 = s[i1:].index(s2) + i1
+            except:
+                raise ValueError(
+                    f'substring {s2} not found in {s}'
+                )
         else:
             i2 = len(s)
 

tests/unit/test_utils.py

@@ -20,6 +20,7 @@
     get_nth_label,
     get_str_btn,
     expand_wildcards,
+    write_atoms,
 )
 import ase.build