Update funcionalities for ALMACE/MACE

pynta/calculator.py

@@ -1,4 +1,5 @@
 from xtb.ase.calculator import XTB
+from pynta.wrapper_almace import wrapperALMACE
 import numpy as np
 import ase
 from ase.atoms import Atoms
@@ -66,7 +67,33 @@ def get_energy_forces_site_bond(atoms,ind,site_pos,k,deq):
     energy = k*(d-deq)**2
     return energy,k*forces
 
-class HarmonicallyForcedXTB(XTB):
+class HarmonicallyForcedXTB(wrapperALMACE):
+    def get_energy_forces(self):
+        energy = 0.0
+        forces = np.zeros(self.atoms.positions.shape)
+        if hasattr(self.parameters,"atom_bond_potentials"):
+            for atom_bond_potential in self.parameters.atom_bond_potentials:
+                E,F = get_energy_forces_atom_bond(self.atoms,**atom_bond_potential)
+                energy += E
+                forces += F
+
+        if hasattr(self.parameters,"site_bond_potentials"):
+            for site_bond_potential in self.parameters.site_bond_potentials:
+                E,F = get_energy_forces_site_bond(self.atoms,**site_bond_potential)
+                energy += E
+                forces += F
+
+        return energy[0][0],forces
+
+    def calculate(self, atoms=None, properties=None, system_changes=calculator.all_changes):
+        wrapperALMACE.calculate(self,atoms=atoms,properties=properties,system_changes=system_changes)
+        energy,forces = self.get_energy_forces()
+        self.results["energy"] += energy
+        self.results["free_energy"] += energy
+        self.results["forces"] += forces
+
+
+class HarmonicallyForcedXTB2(XTB):
     def get_energy_forces(self):
         energy = 0.0
         forces = np.zeros(self.atoms.positions.shape)
@@ -119,9 +146,9 @@ def run_harmonically_forced_xtb(atoms,atom_bond_potentials,site_bond_potentials,
             out_constraints.append(FixAtoms(
                 indices=list(range(n))
                 ))
-            
+
     atoms.set_constraint(out_constraints)
-    
+
     hfxtb = HarmonicallyForcedXTB(method="GFN1-xTB",
                               atom_bond_potentials=atom_bond_potentials,
                              site_bond_potentials=site_bond_potentials)
@@ -258,7 +285,7 @@ def run_harmonically_forced_xtb_no_pbc(atoms,atom_bond_potentials,site_bond_pote
             if "a2" in c:
                 c["a2"] += new_nslab-nslab
             new_constraints.append(c)
-    
+
     out_constraints = []
     for c in new_constraints:
         if isinstance(c,dict):
@@ -282,7 +309,7 @@ def run_harmonically_forced_xtb_no_pbc(atoms,atom_bond_potentials,site_bond_pote
             out_constraints.append(FixAtoms(
                 indices=list(range(n))
                 ))
-    
+
     hfxtb = HarmonicallyForcedXTB(method="GFN1-xTB",
                                   atom_bond_potentials=new_atom_bond_potentials,
                                  site_bond_potentials=new_site_potentials)

pynta/main.py

@@ -549,9 +549,10 @@ def launch(self,single_job=False):
         """
         if self.queue:
             rapidfirequeue(self.launchpad,self.fworker,self.qadapter,njobs_queue=self.njobs_queue,nlaunches="infinite")
-        elif not self.queue and (self.num_jobs == 1 or single_job):
-            rapidfire(self.launchpad,self.fworker,nlaunches="infinite")
+        #elif not self.queue and (self.num_jobs == 1 or single_job):
+        #    rapidfire(self.launchpad,self.fworker,nlaunches="infinite")
         else:
+            print(" my own multilauncher")
             listfworkers = createFWorkers(self.num_jobs)
             launch_multiprocess2(self.launchpad,listfworkers,"INFO",0,self.num_jobs,5)
 
@@ -613,6 +614,7 @@ def execute_from_initial_ad_guesses(self):
 
 #restart option: RHE
     def reset(self):
+        print(" >> Restarting ")
         # Get the information of the workflow
         wf1 = self.launchpad.get_wf_summary_dict(1, mode='more')
 
@@ -638,6 +640,8 @@ def reset(self):
                             '{{pynta.tasks.MolecularVibrationsTask}}',
                             '{{pynta.tasks.MolecularIRC}}']
 
+                print('TaskID:{2:5d} Name task {0:^20s} {1:^12s}'.format(task_name, task_state, task_id))
+
                 #if nameTask in nameTasks:
                 #    opt_method = newd['_tasks'][0]['opt_method'] if 'opt_method' in newd['_tasks'][0] else None
 
@@ -651,35 +655,55 @@ def reset(self):
                 # We load the trajectory file and save this structure in the
                 # tree of each uncompleted task
 
-                if 'opt' in task_name:
-                    dirs = wf_launchers[task_name]
-                    if dirs != []:
-                        src = dirs[0]
-                        print(' Name task {0:^20s} {1:^12s} {2}'.format(task_name, task_state, src))
-                        file_traj = [name for name in os.listdir(src) if name.endswith(".traj")]
-                        if len(file_traj) > 1:
-                            file_traj = file_traj[0]
-                            base, ext = os.path.splitext(file_traj)
+                #if 'opt' in task_name:
+                #    dirs = wf_launchers[task_name]
+                #    if dirs != []:
+                #        src = dirs[0]
+                #        print(' Name task {0:^20s} {1:^12s} {2}'.format(task_name, task_state, src))
+                #        file_traj = [name for name in os.listdir(src) if name.endswith(".traj")]
+                #        if len(file_traj) > 1:
+                #            file_traj = file_traj[0]
+                #            base, ext = os.path.splitext(file_traj)
 
-                            with open (os.path.join(src, "FW.json")) as file:
-                                filejson = json.load(file)
+                #            with open (os.path.join(src, "FW.json")) as file:
+                #                filejson = json.load(file)
 
-                            if opt_method == 'QuasiNewton':
-                                namexyz = f'weakopt_{base}.xyz'
-                            else:
-                                namexyz = f'{base}_init.xyz'
+                #            if opt_method == 'QuasiNewton':
+                #                namexyz = f'weakopt_{base}.xyz'
+                #            else:
+                #                namexyz = f'{base}_init.xyz'
 
-                            atoms = read(os.path.join(src, file_traj), index=-1)
-                            dst = os.path.dirname(filejson['spec']['_tasks'][0]['xyz'])
+                #            atoms = read(os.path.join(src, file_traj), index=-1)
+                #            dst = os.path.dirname(filejson['spec']['_tasks'][0]['xyz'])
 
-                            write(f'{src}/{namexyz}', atoms, format='xyz')
+                #            write(f'{src}/{namexyz}', atoms, format='xyz')
 
-                            copyDataAndSave(src, dst, namexyz)
-                            copyDataAndSave(src, dst, f'{base}.traj')
+                #            copyDataAndSave(src, dst, namexyz)
+                #            copyDataAndSave(src, dst, f'{base}.traj')
 
                 # Keep on with the task_state != 'COMPLETED'
                 self.launchpad.rerun_fw(task_id)
                 self.launchpad.update_spec([task_id], newd)
+            else:
+                task_id = int(task_name.split('--')[-1])
+                if task_id in [3, 5, 7, 9, 12, 14, 16, 18]:
+                    d = self.launchpad.get_fw_dict_by_id(task_id)
+                    newd = deepcopy(d['spec'])
+                    from pprint import pprint
+                    nameTask = newd['_tasks'][0]['_fw_name']
+                    print("  TASKID ",task_id)
+                    pprint(newd)
+
+                    nameTasks = ['{{pynta.tasks.MolecularOptimizationTask}}',
+                                '{{pynta.tasks.MolecularVibrationsTask}}',
+                                '{{pynta.tasks.MolecularIRC}}']
+
+                    print('TaskID:{2:5d} Name task {0:^20s} {1:^12s}'.format(task_name, task_state, task_id))
+
+                    self.launchpad.rerun_fw(task_id)
+                    self.launchpad.update_spec([task_id], newd)
+
+
 
 
         self.launch()

pynta/polaris.py

@@ -8,7 +8,7 @@
 def createCommand(node, software):
     binary = os.environ.get("EXE")
     if binary is None:
-        binary = '/lus/eagle/projects/catalysis_aesp/raymundo/soft/qe-7.1/build_cuda_nompiaware/bin/pw.x'
+        binary = '/soft/applications/quantum_espresso/7.3.1-nvhpc23.1-libxc610/bin/pw.x'
 
     if software == 'Espresso':
         command = 'mpiexec --hosts {} -n 4 --ppn 4 --depth=8 --cpu-bind depth --env OMP_NUM_THREADS=8 --env CUDA_VISIBLE_DEVICES=0,1,2,3 {} -nk 4 -in PREFIX.pwi > PREFIX.pwo'.format(node, binary)
@@ -32,6 +32,8 @@ def createFWorkers(num_jobs):
             for line in f:
                 nodes_list.append(line.strip('\n'))
 
+    #nodes_list.pop(0) # delete the first node for dispatch
+
     fworkers = []
 
     for i in range(num_jobs):

pynta/tasks.py

@@ -65,6 +65,7 @@ def optimize_firework(xyz,software,label,opt_method=None,sella=None,socket=False
     if opt_method: d["opt_method"] = opt_method
     if software_kwargs: d["software_kwargs"] = software_kwargs
     if opt_kwargs: d["opt_kwargs"] = opt_kwargs
+    print("opt kwargs", opt_kwargs)
     if run_kwargs: d["run_kwargs"] = run_kwargs
     if constraints: d["constraints"] = constraints
     sella = True if sella or (sella is None and order != 0) else False
@@ -80,9 +81,16 @@ def optimize_firework(xyz,software,label,opt_method=None,sella=None,socket=False
     d["facet"] = facet
     t1 = MolecularOptimizationTask(d)
     directory = os.path.dirname(xyz)
+    subname = directory.split('/')
     if out_path is None: out_path = os.path.join(directory,label+".xyz")
     t2 = FileTransferTask({'files': [{'src': label+'.xyz', 'dest': out_path}, {'src': label+'.traj', 'dest': os.path.join(directory,label+".traj")}],
             'mode': 'copy', 'ignore_errors' : ignore_errors})
+    if software == 'ALMACE' and opt_method == 'MDMin':
+        storage = d['software_kwargs']['storage']
+        t2 = FileTransferTask({'files': [{'src': label+'.xyz', 'dest': out_path}, {'src': label+'.traj', 'dest': os.path.join(directory,label+".traj")},
+                                        {'src':label+'.traj', 'dest':os.path.join(storage,subname[-2]+'_'+label+'.traj')},
+                                        ],
+                                'mode': 'copy', 'ignore_errors' : ignore_errors})
     return Firework([t1,t2],parents=parents,name=label+"opt",spec={"_allow_fizzled_parents": allow_fizzled_parents,"_priority": priority})
 
 @explicit_serialize
@@ -97,6 +105,14 @@ def run_task(self, fw_spec):
         if self["software"] == "Espresso" or self["software"] == "PWDFT":
             software_kwargs["command"] = createCommand(fw_spec["_fw_env"]["host"], self["software"])
 
+        if self["software"] == "MACE" :
+            software_kwargs["host"] = fw_spec["_fw_env"]["host"] # RHE
+
+        if self["software"] == "ALMACE" :
+            software_kwargs["host"] = fw_spec["_fw_env"]["host"] # RHE
+            software_kwargs['sub_software_kwargs']['command'] = createCommand(fw_spec['_fw_env']['host'], software_kwargs['sub_software'])
+            software_kwargs['opt_method'] = name_to_ase_opt(self["opt_method"]) if "opt_method" in self.keys() else BFGS
+
         socket = self["socket"] if "socket" in self.keys() else False
         if socket:
             unixsocket = "ase_"+self["software"].lower()+"_"+self["label"]+"_"+self["xyz"].replace("/","_").replace(".","_")
@@ -121,6 +137,7 @@ def run_task(self, fw_spec):
         label = self["label"]
         xyz = self['xyz']
         suffix = os.path.split(xyz)[-1].split(".")[-1]
+        print("\033[31m xyz: {0}\033[0m".format(xyz))
 
         try:
             if suffix == "xyz":
@@ -312,6 +329,10 @@ def run_task(self, fw_spec):
         # Mapping nodes to fworkers
         if self["software"] == "Espresso" or self["software"] == "PWDFT":
             software_kwargs["command"] = createCommand(fw_spec["_fw_env"]["host"], self["software"])
+
+        if self["software"] == "MACE" :
+            software_kwargs["host"] = fw_spec["_fw_env"]["host"]
+
         software = name_to_ase_software(self["software"])(**software_kwargs)
 
         opt_kwargs = deepcopy(self["opt_kwargs"]) if "opt_kwargs" in self.keys() else dict()
@@ -360,6 +381,10 @@ def run_task(self, fw_spec):
         if self["software"] == "Espresso" or self["software"] == "PWDFT":
             software_kwargs["command"] = createCommand(fw_spec["_fw_env"]["host"], self["software"])
 
+        if self["software"] == "MACE" :
+            software_kwargs["host"] = fw_spec["_fw_env"]["host"]
+
+
         try:
             sp = read(xyz)
             sp.calc = software(**software_kwargs)
@@ -405,6 +430,15 @@ def run_task(self, fw_spec):
         # Mapping nodes to fworkers
         if self["software"] == "Espresso" or self["software"] == "PWDFT":
             software_kwargs["command"] = createCommand(fw_spec["_fw_env"]["host"], self["software"])
+
+        if self["software"] == "MACE" :
+            software_kwargs["host"] = fw_spec["_fw_env"]["host"]
+
+        if self["software"] == "ALMACE" :
+            software_kwargs["host"] = fw_spec["_fw_env"]["host"] # RHE
+            software_kwargs['opt_method'] = name_to_ase_opt(self["opt_method"]) if "opt_method" in self.keys() else BFGS #RHE
+            software_kwargs['sub_software_kwargs']['command'] = createCommand(fw_spec['_fw_env']['host'], software_kwargs['sub_software'])
+
         software = name_to_ase_software(self["software"])(**software_kwargs)
         ignore_errors = deepcopy(self["ignore_errors"]) if "ignore_errors" in self.keys() else False
         socket = self["socket"] if "socket" in self.keys() else False
@@ -516,7 +550,7 @@ def run_task(self, fw_spec):
 
         cas = SlabAdsorptionSites(slab,facet,allow_6fold=False,composition_effect=False,
                             label_sites=True,tol=0.5,
-                        surrogate_metal='Pt')
+                        surrogate_metal=metal)
 
         adsorbates_path = self["adsorbates_path"]
 
@@ -781,6 +815,14 @@ def run_task(self, fw_spec):
         # Mapping nodes to fworkers
         if self["software"] == "Espresso" or self["software"] == "PWDFT":
             software_kwargs["command"] = createCommand(fw_spec["_fw_env"]["host"], self["software"])
+
+        if self["software"] == "MACE" :
+            software_kwargs["host"] = fw_spec["_fw_env"]["host"]
+
+        if self["software"] == "ALMACE" :
+            software_kwargs["host"] = fw_spec["_fw_env"]["host"] # RHE
+            software_kwargs['sub_software_kwargs']['command'] = createCommand(fw_spec['_fw_env']['host'], software_kwargs['sub_software'])
+
         socket = self["socket"] if "socket" in self.keys() else False
         if socket:
             unixsocket = "ase_"+self["software"].lower()+"_"+self["label"]+"_"+self["xyz"].replace("/","_").replace(".","_")

pynta/utils.py

@@ -213,6 +213,12 @@ def name_to_ase_software(software_name):
     elif software_name == "PWDFT":
         module = import_module("pynta.ase_pwdft.pwdft")
         return getattr(module, software_name)
+    elif software_name == "ALMACE":
+        module = import_module("pynta.wrapper_almace")
+        return getattr(module,"wrapperALMACE")
+    elif software_name == "MACE":
+        module = import_module("pynta.wrapper_mace")
+        return getattr(module,"wrapperMACE")
     else:
         module = import_module("ase.calculators."+software_name.lower())
         return getattr(module, software_name)