From 8ce0ad8e24e764a8f168e66c0efbfea854cbcf1f Mon Sep 17 00:00:00 2001
From: Michael Taylor <mgt16@lanl.gov>
Date: Tue, 12 Sep 2023 09:28:15 -0600
Subject: [PATCH] Allow ASE kwargs for user methods.

---
 README.md                       |  2 ++
 architector/io_calc.py          | 16 +++++++++++-----
 architector/io_process_input.py |  9 +++++++++
 3 files changed, 22 insertions(+), 5 deletions(-)

diff --git a/README.md b/README.md
index 217666d..1c9c945 100644
--- a/README.md
+++ b/README.md
@@ -106,7 +106,9 @@ inputDict = {
         
     # Method parameters.
     "calculator":None, # ASE calculator class input for usage during construction or for optimization.
+    "calculator_kwargs":dict(), # ASE calculator kwargs.
     "ase_opt_method":None, # ASE optimizer class used for geometry optimizations. Default will use LBFGSLineSearch.
+    "ase_opt_kwargs":dict(), # ASE optimizer kwargs. Do not include "trajectory" nor "logfile" kwargs.
     "fmax":0.1, # eV/Angstrom maximum force to optimize to with ASE optimizer.
     "maxsteps":1000, # Maxmimum number of steps for ASE otpimizer to take.
     # Note that charge and multiplicity (uhf) are tracked by ase.Atoms.initial_charges and ase.Atoms.initial_magnetic_moments
diff --git a/architector/io_calc.py b/architector/io_calc.py
index 8643048..22ab415 100644
--- a/architector/io_calc.py
+++ b/architector/io_calc.py
@@ -52,6 +52,8 @@
 "metal_ox": None, # Oxidation State
 "metal_spin": None, # Spin State
 "xtb_solvent": 'none', # Add any named XTB solvent!
+"calculator_kwargs":dict(), # ASE calculator kwargs.
+"ase_opt_kwargs":dict(), # ASE optimizer kwargs.
 "xtb_accuracy":1.0, # Numerical Accuracy for XTB calculations
 "xtb_electronic_temperature":300, # In K -> fermi smearing - increase for convergence on harder systems
 "xtb_max_iterations":250, # Max iterations for xtb SCF.
@@ -224,7 +226,7 @@ def calculate(self):
                 self.mol.calc_suggested_spin(params=self.parameters)
             obabel_ff_requested = False
             if (self.calculator is not None) and ('custom' in self.method): # If ASE calculator passed use that by default
-                calc = self.calculator
+                calc = self.calculator(**self.calculator_kwargs)
                 # Here, if a calculator needs spin/charge information in another way we can assign.
                 # Or handle as a different use case.
                 uhf_vect = np.zeros(len(self.mol.ase_atoms))
@@ -280,16 +282,20 @@ def calculate(self):
                                 if self.logfile is not None:
                                     dyn = self.opt_method(self.mol.ase_atoms, 
                                                         trajectory='temp.traj',
-                                                        logfile=self.logfile)
+                                                        logfile=self.logfile,
+                                                        **self.ase_opt_kwargs)
                                 else:
                                     dyn = self.opt_method(self.mol.ase_atoms, 
-                                                        trajectory='temp.traj')
+                                                        trajectory='temp.traj',
+                                                        **self.ase_opt_kwargs)
                             else:
                                 if self.logfile is not None:
                                     dyn = self.opt_method(self.mol.ase_atoms,
-                                                        logfile=self.logfile)
+                                                        logfile=self.logfile,
+                                                        **self.ase_opt_kwargs)
                                 else:
-                                    dyn = self.opt_method(self.mol.ase_atoms)
+                                    dyn = self.opt_method(self.mol.ase_atoms,
+                                                          **self.ase_opt_kwargs)
                             dyn.run(fmax=self.fmax,steps=self.maxsteps)
                             if self.parameters['save_trajectories']:
                                 self.read_traj()
diff --git a/architector/io_process_input.py b/architector/io_process_input.py
index 7a78b48..92bcabc 100644
--- a/architector/io_process_input.py
+++ b/architector/io_process_input.py
@@ -772,6 +772,15 @@ def inparse(inputDict):
             "alternate_metal_spin": None, # Secondary spin state to check. 
                 
             # Method parameters.
+            "calculator":None, # ASE calculator class input for usage during construction or for optimization.
+            "calculator_kwargs":dict(), # ASE calculator kwargs.
+            "ase_opt_method":None, # ASE optimizer class used for geometry optimizations. Default will use LBFGSLineSearch.
+            "ase_opt_kwargs":dict(), # ASE optimizer kwargs.
+            "fmax":0.1, # eV/Angstrom maximum force to optimize to with ASE optimizer.
+            "maxsteps":1000, # Maxmimum number of steps for ASE otpimizer to take.
+            # Note that charge and multiplicity (uhf) are tracked by ase.Atoms.initial_charges and ase.Atoms.initial_magnetic_moments
+            # In the background by Architector. Make sure your ASE calculator either 1. uses these vectors for setting spin/charge or
+            # 2. has these quantities pre-specified for the spin/charge you'd like to run.
             "xtb_solvent": 'none', # Add any named XTB solvent!
             "xtb_accuracy":1.0, # Numerical Accuracy for XTB calculations
             "xtb_electronic_temperature":300, # In K -> fermi smearing - increase for convergence on harder systems