From e2311b8751e9771f8c51e680a6beeee451380122 Mon Sep 17 00:00:00 2001
From: Carlo Antonio Pignedoli <c.pignedoli@gmail.com>
Date: Thu, 12 Sep 2024 13:58:16 +0000
Subject: [PATCH] working

---
 .../workflows/gaussian/__init__.py            |   2 +
 .../workflows/gaussian/nics_workchain.py      | 219 +++++++++++-------
 .../workflows/gaussian/relax_workchain.py     |  38 ++-
 .../workflows/gaussian/scf_workchain.py       |  52 ++++-
 examples/workflows/example_gaussian_nics.py   |  59 +++--
 pyproject.toml                                |   1 +
 6 files changed, 263 insertions(+), 108 deletions(-)

diff --git a/aiida_nanotech_empa/workflows/gaussian/__init__.py b/aiida_nanotech_empa/workflows/gaussian/__init__.py
index d08b655b..0b5d156f 100644
--- a/aiida_nanotech_empa/workflows/gaussian/__init__.py
+++ b/aiida_nanotech_empa/workflows/gaussian/__init__.py
@@ -7,6 +7,7 @@
 from .relax_workchain import GaussianRelaxWorkChain
 from .scf_workchain import GaussianScfWorkChain
 from .spin_workchain import GaussianSpinWorkChain
+from .nics_workchain import GaussianNicsWorkChain
 
 __all__ = (
     "GaussianScfWorkChain",
@@ -18,4 +19,5 @@
     "GaussianConstrOptChainWorkChain",
     "GaussianCasscfWorkChain",
     "GaussianCasscfSeriesWorkChain",
+    "GaussianNicsWorkChain",
 )
diff --git a/aiida_nanotech_empa/workflows/gaussian/nics_workchain.py b/aiida_nanotech_empa/workflows/gaussian/nics_workchain.py
index 58d5f1ee..966a42fd 100644
--- a/aiida_nanotech_empa/workflows/gaussian/nics_workchain.py
+++ b/aiida_nanotech_empa/workflows/gaussian/nics_workchain.py
@@ -1,4 +1,30 @@
-class GaussianRelaxWorkChain(engine.WorkChain):
+import numpy as np
+import ase
+from aiida import engine, orm
+from .relax_workchain import GaussianRelaxWorkChain
+from .scf_workchain import GaussianScfWorkChain
+from ...utils import cycle_tools, nmr
+from ...utils import common_utils
+
+@engine.calcfunction
+def prepare_nmr_structure(structure,height):
+    ase_geo = structure.get_ase()
+    pos = ase_geo.positions
+    extents = np.array([
+        np.max(pos[:,0]) - np.min(pos[:,0]),
+        np.max(pos[:,1]) - np.min(pos[:,1]),
+        np.max(pos[:,2]) - np.min(pos[:,2]),
+    ])
+    inds = np.argsort(-extents)
+    new_pos = pos[:, inds]
+    ase_atoms = ase.Atoms(numbers=ase_geo.numbers, positions=new_pos)
+    ase_atoms_no_h = ase.Atoms([a for a in ase_atoms if a.symbol != 'H'])
+    cycles = cycle_tools.dumb_cycle_detection(ase_atoms_no_h, 8)
+    ref_p = nmr.find_ref_points(ase_atoms_no_h, cycles, height.value)
+    new_ase = ase_atoms + ref_p
+    return orm.StructureData(ase=new_ase)
+
+class GaussianNicsWorkChain(engine.WorkChain):
     @classmethod
     def define(cls, spec):
         super().define(spec)
@@ -11,6 +37,18 @@ def define(cls, spec):
             required=True,
             help="input geometry",
         )
+        spec.input(
+            "opt",
+            valid_type=orm.Bool,
+            required=False,
+            help="False do not optimize structure",
+        )
+        spec.input(
+            "height",
+            valid_type=orm.Float,
+            required=False,
+            help="Height of NICS centers",
+        )
         spec.input(
             "functional", valid_type=orm.Str, required=True, help="xc functional"
         )
@@ -24,7 +62,6 @@ def define(cls, spec):
             default=lambda: orm.Int(0),
             help="spin multiplicity; 0 means RKS",
         )
-
         spec.input(
             "wfn_stable_opt",
             valid_type=orm.Bool,
@@ -32,83 +69,105 @@ def define(cls, spec):
             default=lambda: orm.Bool(False),
             help="if true, perform wfn stability optimization",
         )
+        spec.input(
+            "empirical_dispersion",
+            valid_type=orm.Str,
+            required=False,
+            default=lambda: orm.Str(""),
+            help=("Include empirical dispersion corrections" '(e.g. "GD3", "GD3BJ")'),
+        )
+        spec.input(
+            "options",
+            valid_type=orm.Dict,
+            required=False,
+            help="Use custom metadata.options instead of the automatic ones.",
+        )        
+        
+        spec.outline(
+            cls.setup,
+            engine.if_(cls.should_submit_opt)(cls.submit_opt,cls.inspect_opt),            
+            cls.submit_nics,
+            cls.inspect_nics,
+            cls.finalize,
+        )
 
-# Create a methane (CH4) molecule using the ASE molecular database
-ch4 = molecule('CH4')
-
-GaussianCalculation = CalculationFactory("gaussian")
-
-
-def example_dft(gaussian_code):
-    """Run a simple gaussian optimization"""
-
-    # structure
-    structure = StructureData(ase=ch4)
-
-    num_cores = 1
-    memory_mb = 300
-
-    # Main parameters: geometry optimization
-    parameters = Dict(
-        {
-            "link0_parameters": {
-                "%chk": "aiida.chk",
-                "%mem": "%dMB" % memory_mb,
-                "%nprocshared": num_cores,
-            },
-            "functional": "B3LYP",
-            "basis_set": "6-311g(d,p)",
-            "charge": 0,
-            "multiplicity": 1,
-            "dieze_tag": "#P",
-            "route_parameters": {
-                "scf": {
-                    "maxcycle":2048,
-                    "cdiis": None,
-                    "conver":8
-                },
-                "int":"superfine",
-                "guess":"mix",
-                "nosymm": None,
-                "opt": "tight",
-            },
-        }
-    )
-
-    # Construct process builder
-
-    builder = GaussianCalculation.get_builder()
-
-    builder.structure = structure
-    builder.parameters = parameters
-    builder.code = gaussian_code
-
-    builder.metadata.options.resources = {
-        "num_machines": 1,
-        "tot_num_mpiprocs": num_cores,
-    }
-
-    # Should ask for extra +25% extra memory
-    builder.metadata.options.max_memory_kb = int(1.25 * memory_mb) * 1024
-    builder.metadata.options.max_wallclock_seconds = 5 * 60
-
-    print("Running calculation...")
-    res, _node = run_get_node(builder)
-
-    print("Final scf energy: %.4f" % res["output_parameters"]["scfenergies"][-1])
-
-
-@click.command("cli")
-@click.argument("codelabel", default="gaussian@localhost")
-def cli(codelabel):
-    """Click interface"""
-    try:
-        code = Code.get_from_string(codelabel)
-    except NotExistent:
-        print(f"The code '{codelabel}' does not exist")
-        sys.exit(1)
-    example_dft(code)
-
+        spec.outputs.dynamic = True
 
-if __name__ == "__main__":
-    cli()  # pylint: disable=no-value-for-parameter
+        spec.exit_code(
+            390,
+            "ERROR_TERMINATION",
+            message="One or more steps of the work chain failed.",
+        )
+    def setup(self):
+        self.report("Setting up...")
+        self.ctx.nmr_structure = self.inputs.structure
+        self.ctx_should_opt = getattr(self.inputs, "opt", True)
+        
+    def should_submit_opt(self):
+        return self.ctx_should_opt
+    
+    def submit_opt(self):
+        self.report("Submitting optimization...")
+        label = "geo_opt"
+        builder = GaussianRelaxWorkChain.get_builder()
+        builder.gaussian_code = self.inputs.gaussian_code
+        builder.structure = self.inputs.structure
+        builder.functional = self.inputs.functional
+        builder.empirical_dispersion = self.inputs.empirical_dispersion
+        builder.basis_set = self.inputs.basis_set
+        builder.multiplicity = self.inputs.multiplicity
+        builder.tight = orm.Bool(True)
+        builder.int = orm.Str("superfine")
+        builder.cdiis = orm.Bool(True)
+        builder.maxcycle = orm.Int(2048)
+        builder.conver = orm.Int(8)
+        if "options" in self.inputs:
+            builder.options = self.inputs.options
+
+        submitted_node = self.submit(builder)
+        submitted_node.description = label
+        self.to_context(**{label: submitted_node})
+        
+    def inspect_opt(self):
+        self.report("Inspecting optimization...")
+        label = "geo_opt"
+        # check if everything finished nicely
+        if not common_utils.check_if_calc_ok(self, self.ctx[label]):
+            return self.exit_codes.ERROR_TERMINATION        
+        self.ctx.nmr_structure = self.ctx[label].outputs.output_structure
+        return engine.ExitCode(0)
+    
+    def submit_nics(self):
+        self.report("Submitting NICS calculation...")
+        label = "nics"
+        builder = GaussianScfWorkChain.get_builder()
+        height = getattr(self.inputs, "height", 1.0)
+        builder.gaussian_code = self.inputs.gaussian_code
+        builder.structure = prepare_nmr_structure(self.ctx.nmr_structure,orm.Float(height))
+        builder.functional = self.inputs.functional
+        builder.basis_set = self.inputs.basis_set
+        builder.multiplicity = self.inputs.multiplicity
+        builder.int = orm.Str("superfine")
+        builder.cdiis = orm.Bool(True)
+        builder.nmr = orm.Bool(True)
+        builder.maxcycle = orm.Int(2048)
+        builder.conver = orm.Int(8)        
+        if "options" in self.inputs:
+            builder.options = self.inputs.options
+
+        submitted_node = self.submit(builder)
+        submitted_node.description = label
+        self.to_context(**{label: submitted_node})
+        
+    def inspect_nics(self):
+        self.report("Inspecting nics...")
+        label = "nics"
+        # check if everything finished nicely
+        if not common_utils.check_if_calc_ok(self, self.ctx[label]):
+            return self.exit_codes.ERROR_TERMINATION
+        
+        self.out("output_parameters", self.ctx[label].outputs.output_parameters)
+        return engine.ExitCode(0)   
+         
+    def finalize(self):
+        self.report("Finalizing...")
\ No newline at end of file
diff --git a/aiida_nanotech_empa/workflows/gaussian/relax_workchain.py b/aiida_nanotech_empa/workflows/gaussian/relax_workchain.py
index 9c035731..1ea2d764 100644
--- a/aiida_nanotech_empa/workflows/gaussian/relax_workchain.py
+++ b/aiida_nanotech_empa/workflows/gaussian/relax_workchain.py
@@ -49,7 +49,30 @@ def define(cls, spec):
             default=lambda: orm.Bool(False),
             help="Use tight optimization criteria.",
         )
-
+        spec.input(
+            "maxcycle",
+            valid_type=orm.Int,
+            required=False,
+            help="the maximum number of scf cycles",
+        )
+        spec.input(
+            "cdiis",
+            valid_type=orm.Bool,
+            required=False,
+            help="Conjugate Direct Inversion in the Iterative Subspace",
+        )        
+        spec.input(
+            "conver",
+            valid_type=orm.Int,
+            required=False,
+            help="the scf convergece threshold",
+        )        
+        spec.input(
+            "int",
+            valid_type=orm.Str,
+            required=False,
+            help="the integral grid",
+        )
         spec.input(
             "freq",
             valid_type=orm.Bool,
@@ -290,7 +313,18 @@ def optimization(self):
 
         if len(opt_dict) != 0:
             parameters["route_parameters"]["opt"] = opt_dict
-
+        conver = getattr(self.inputs,"conver",None)
+        maxcycle = getattr(self.inputs,"maxcycle",None)
+        grid = getattr(self.inputs,"int",None)
+        cdiis = getattr(self.inputs,"cdiis",None)
+        if grid:
+            parameters["route_parameters"]["int"] = grid
+        if maxcycle:
+            parameters["route_parameters"]["scf"]["maxcycle"] = maxcycle
+        if cdiis:
+            parameters["route_parameters"]["scf"]["cdiis"] = None
+        if conver:
+            parameters["route_parameters"]["scf"]["conver"] = conver
         builder.gaussian.parameters = parameters
         builder.gaussian.structure = self.inputs.structure
         builder.gaussian.code = self.inputs.gaussian_code
diff --git a/aiida_nanotech_empa/workflows/gaussian/scf_workchain.py b/aiida_nanotech_empa/workflows/gaussian/scf_workchain.py
index f38857ae..db8644e7 100644
--- a/aiida_nanotech_empa/workflows/gaussian/scf_workchain.py
+++ b/aiida_nanotech_empa/workflows/gaussian/scf_workchain.py
@@ -60,6 +60,36 @@ def define(cls, spec):
             required=False,
             help="the folder of a completed gaussian calculation",
         )
+        spec.input(
+            "maxcycle",
+            valid_type=orm.Int,
+            required=False,
+            help="the maximum number of scf cycles",
+        )
+        spec.input(
+            "conver",
+            valid_type=orm.Int,
+            required=False,
+            help="the scf convergece threshold",
+        )  
+        spec.input(
+            "cdiis",
+            valid_type=orm.Bool,
+            required=False,
+            help="Conjugate Direct Inversion in the Iterative Subspace",
+        )               
+        spec.input(
+            "int",
+            valid_type=orm.Str,
+            required=False,
+            help="the integral grid",
+        )
+        spec.input(
+            "nmr",
+            valid_type=orm.Bool,
+            required=False,
+            help="nmr calculation",
+        )
 
         # Inputs for cubes generation.
         spec.input("formchk_code", valid_type=orm.Code, required=False)
@@ -167,7 +197,7 @@ def setup(self):
 
         # Use default convergence criterion at the start
         # but switch to conver=7 in case of failure.
-        self.ctx.conver = None
+        self.ctx.conver = getattr(self.inputs, "conver", None)
         self.ctx.scf_label = "scf"
 
         return engine.ExitCode(0)
@@ -224,10 +254,28 @@ def scf(self):
                 "charge": 0,
                 "multiplicity": self.ctx.mult,
                 "route_parameters": {
-                    "scf": {"maxcycle": 140},
+                    "scf": {"maxcycle": 140 },
                 },
             }
         )
+        nmr = getattr(self.inputs,"nmr",None)
+        conver = getattr(self.inputs,"conver",None)
+        maxcycle = getattr(self.inputs,"maxcycle",None)
+        grid = getattr(self.inputs,"int",None)
+        cdiis = getattr(self.inputs,"cdiis",None)
+        if grid:
+            parameters["route_parameters"]["int"] = grid
+        if maxcycle:
+            parameters["route_parameters"]["scf"]["maxcycle"] = maxcycle
+        if cdiis:
+            parameters["route_parameters"]["scf"]["cdiis"] = None
+        if conver:
+            parameters["route_parameters"]["scf"]["conver"] = conver
+        if nmr:
+            if conver is None:
+                conver = 8
+            parameters["route_parameters"]["nmr"] = None
+            parameters["route_parameters"]["cphf"] = {"conver": conver}
         if self.inputs.wfn_stable_opt:
             parameters["route_parameters"]["stable"] = "opt"
         else:
diff --git a/examples/workflows/example_gaussian_nics.py b/examples/workflows/example_gaussian_nics.py
index 541f9ffa..9c43f4b8 100644
--- a/examples/workflows/example_gaussian_nics.py
+++ b/examples/workflows/example_gaussian_nics.py
@@ -1,39 +1,50 @@
-mport os
-
-import ase.io
+import os
+import click
+from ase.build import molecule
 import numpy as np
-from aiida.engine import run_get_node
-from aiida.orm import Bool,Int, Str, StructureData, load_code
+from aiida import engine, orm
 from aiida.plugins import WorkflowFactory
 
 import aiida_nanotech_empa.utils.gaussian_wcs_postprocess as pp
 
-GaussianRelaxWorkChain = WorkflowFactory("nanotech_empa.gaussian.relax")
-
-DATA_DIR = os.path.dirname(os.path.abspath(__file__))
-OUTPUT_DIR = os.path.dirname(os.path.abspath(__file__))
-
+GaussianNicsWorkChain = WorkflowFactory("nanotech_empa.gaussian.nics")
 
-def _example_gaussian_spin(gaussian_code):#, formchk_code, cubegen_code):
-    ase_geom = ase.io.read(os.path.join(DATA_DIR, "benzene.xyz"))
+def _example_gaussian_nics(gaussian_code,opt):
+    ase_geom = molecule('C6H6')
     ase_geom.cell = np.diag([10.0, 10.0, 10.0])
 
-    builder = GaussianRelaxWorkChain.get_builder()
+    builder = GaussianNicsWorkChain.get_builder()
+    builder.gaussian_code = gaussian_code
+    builder.structure = orm.StructureData(ase=ase_geom)
+    builder.functional = orm.Str("B3LYP")
+    builder.basis_set = orm.Str("6-311G(d,p)")
+    builder.opt = orm.Bool(opt)
+    builder.options = orm.Dict(
+        {
+            "resources": {
+                "tot_num_mpiprocs": 4,
+                "num_machines": 1,
+            },
+            "max_wallclock_seconds": 1 * 60 * 60,
+            "max_memory_kb": 8 * 1024 * 1024, #GB
+        }
+    )    
     
+    _, wc_node = engine.run_get_node(builder)
     
-    
-    
-    _, wc_node = run_get_node(builder)
-
     assert wc_node.is_finished_ok
+    return wc_node.pk
 
-    #pp.make_report(wc_node, nb=False, save_image_loc=OUTPUT_DIR)
+@click.command("cli")
+@click.argument("gaussian_code", default="gaussian@localhost")
+def run_nics(gaussian_code):
+    #print("#### Running Gaussian NICS WorkChain with geo_opt ####")
+    #uuid = _example_gaussian_nics(orm.load_code(gaussian_code),True)
+    #print(f"WorkChain completed uuid: {uuid}")
+    print("#### Running Gaussian NICS WorkChain without geo_opt ####")
+    uuid = _example_gaussian_nics(orm.load_code(gaussian_code),False)
+    print(f"WorkChain completed uuid: {uuid}")
 
 
 if __name__ == "__main__":
-    _example_gaussian_spin(
-        load_code("gaussian@tigu"),
-        #load_code("formchk@localhost"),
-        #load_code("cubegen@localhost"),
-    )
-    
\ No newline at end of file
+    run_nics()
\ No newline at end of file
diff --git a/pyproject.toml b/pyproject.toml
index 17a63dd0..19116f1d 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -70,6 +70,7 @@ dev = [
 "nanotech_empa.gaussian.constr_opt_chain" = "aiida_nanotech_empa.workflows.gaussian:GaussianConstrOptChainWorkChain"
 "nanotech_empa.gaussian.casscf" = "aiida_nanotech_empa.workflows.gaussian:GaussianCasscfWorkChain"
 "nanotech_empa.gaussian.casscf_series" = "aiida_nanotech_empa.workflows.gaussian:GaussianCasscfSeriesWorkChain"
+"nanotech_empa.gaussian.nics" = "aiida_nanotech_empa.workflows.gaussian:GaussianNicsWorkChain"
 "nanotech_empa.cp2k.geo_opt" = "aiida_nanotech_empa.workflows.cp2k:Cp2kGeoOptWorkChain"
 "nanotech_empa.cp2k.fragment_separation" = "aiida_nanotech_empa.workflows.cp2k:Cp2kFragmentSeparationWorkChain"
 "nanotech_empa.cp2k.ads_gw_ic" = "aiida_nanotech_empa.workflows.cp2k:Cp2kAdsorbedGwIcWorkChain"