[Version 2.0] Add pip install functionality and optimize code

.gitignore

@@ -1,2 +1,3 @@
 .DS_Store
-dev/*
+dev/*
+__pycache__/

atst_tools/__init__.py

@@ -0,0 +1,3 @@
+from atst_tools.abacus_autoneb import AbacusAutoNEB
+from atst_tools.abacus_neb import AbacusNEB
+from atst_tools.abacus_dimer import AbacusDimer

source/abacus_autoneb.py → atst_tools/abacus_autoneb.py

@@ -7,7 +7,7 @@
 from ase.io import read, write
 from ase.calculators.abacus import Abacus, AbacusProfile
 #from ase.mep.autoneb import AutoNEB # official
-from my_autoneb import AutoNEB
+from atst_tools.my_autoneb import AutoNEB
 from ase.parallel import world, parprint, paropen
 from ase.optimize import FIRE, BFGS
 from ase.constraints import FixAtoms

source/abacus_dimer.py → atst_tools/abacus_dimer.py

@@ -2,12 +2,14 @@
 # Dimer calculation workflow by using ABACUS calculator
 # part of ATST-Tools scripts
 
+import os
+import sys
+
+import numpy as np
 from ase.io import Trajectory, read, write
 from ase.mep import DimerControl, MinModeAtoms, MinModeTranslate
 #from my_dimer import DimerControl, MinModeAtoms, MinModeTranslate
 from ase.calculators.abacus import Abacus, AbacusProfile
-import os, sys
-import numpy as np
 
 class AbacusDimer:
     """Customize Dimer calculation workflow by using ABACUS"""

source/abacus_neb.py → atst_tools/abacus_neb.py

@@ -2,6 +2,7 @@
 # NEB calculation workflow by ASE-ABACUS
 # part of ATST-Tools scripts
 
+from typing import Optional
 import os
 from ase.calculators.abacus import Abacus, AbacusProfile
 from ase.optimize import FIRE, BFGS
@@ -114,7 +115,7 @@ def set_neb_chain(self, climb=True, fmax=0.05):
         return neb
 
 
-    def run(self, optimizer=FIRE, fmax=0.05, climb=True, outfile="neb.traj", properties=["energy", "forces", "stress"]):
+    def run(self, optimizer=FIRE, fmax=0.05, climb=True, outfile="neb.traj", properties=["energy", "forces", "stress"], max_steps: Optional[int] = None):
         """Run Abacus NEB
 
         optimizer (Optimizer object): defaults to FIRE. BFGS, LBFGS, GPMin, MDMin and QuasiNewton are supported, recommend FIRE method
@@ -126,5 +127,8 @@ def run(self, optimizer=FIRE, fmax=0.05, climb=True, outfile="neb.traj", propert
         # traj = Trajectory(outfile, 'w', neb, properties=properties) # cannot run for now
         traj = Trajectory(outfile, 'w', neb)
         opt = optimizer(neb, trajectory=traj)
-        opt.run(fmax)
+        if max_steps is not None and max_steps > 0:
+            opt.run(fmax=fmax, steps=max_steps)
+        else:
+            opt.run(fmax)
         print("----- NEB calculation finished -----")

source/my_autoneb.py → atst_tools/my_autoneb.py

@@ -1,16 +1,17 @@
 # Refined AutoNEB workflow by QuantumMisaka in 20231128
 # Parts of ATST-Tools
 
-import numpy as np
+
 import shutil
 import os
 import types
-from math import log
-from math import exp
+from math import log, exp
 from contextlib import ExitStack
 from pathlib import Path
 from warnings import warn
 
+import numpy as np
+
 from ase.io import Trajectory
 from ase.io import read
 from ase.mep import NEB

source/my_neb.py → atst_tools/my_neb.py

@@ -1160,7 +1160,7 @@ def plot_bands(self, constant_x=False, constant_y=False,
         with PdfPages(label + '.pdf') as pdf:
             for index in range(nebsteps):
                 sys.stdout.write('\rProcessing band {:10d} / {:10d}'
-                                 .format(index, nebsteps))
+                                 .format(index + 1, nebsteps))
                 sys.stdout.flush()
                 fig, ax = pyplot.subplots()
                 images = self.images[index * nimages:(index + 1) * nimages]

source/neb2vib.py → atst_tools/neb2vib.py

@@ -1,10 +1,12 @@
 # using neb chain to get the main vibrate atoms
 # JamesMisaka in 2024-02-11
 
+import sys
+
 from ase.io import read
 from ase.mep.neb import NEBTools
 import numpy as np
-import sys
+
 
 norm = np.linalg.norm
 

neb/neb_post.py

@@ -35,9 +35,9 @@ def get_TS_stru(self, name="TS_get"):
         """Get TS structure from NEB chain"""
         raw_barrier = NEBTools(self.neb_chain).get_barrier(fit=False, raw=True)
         for atoms in self.neb_chain:
-            if atoms.get_potential_energy() == raw_barrier:
+            if atoms.get_potential_energy() == raw_barrier[0]:
                 write(f"{name}.cif", atoms, format="cif")
-                write(f"{name}.stru", atoms, format="stru")
+                write(f"{name}.stru", atoms, format="abacus")
                 print(f"TS structure is saved as {name}.cif and {name}.stru")
                 return
 

setup.py

@@ -0,0 +1,9 @@
+from setuptools import setup, find_packages
+
+setup(
+    name="atst-tools",
+    version="0.1",
+    packages=find_packages(),
+    install_requires=[],
+    python_requires=">=3.9",
+)