Fix/molecule gw

aiida_nanotech_empa/workflows/cp2k/adsorbed_gw_ic_workchain.py

@@ -1,6 +1,7 @@
 import numpy as np
 from aiida import engine, orm
 
+from ...utils import common_utils
 from .geo_opt_workchain import Cp2kGeoOptWorkChain
 from .molecule_gw_workchain import Cp2kMoleculeGwWorkChain
 
@@ -9,14 +10,24 @@
 }
 
 
-def geometrical_analysis(ase_geo, substr_elem):
+def geometrical_analysis(ase_geo, substr_elem, mol_atoms=None):
     """Simple geometry analysis that returns in the case of
     1) an isolated molecule -> geometry, None
     2) adsorbed system -> molecular geometry, top substr. layer z
     """
-    chem_symbols_arr = np.array(ase_geo.get_chemical_symbols())
-    s_atoms = ase_geo[chem_symbols_arr == substr_elem]
-    non_s_atoms = ase_geo[chem_symbols_arr != substr_elem]
+
+    if mol_atoms is not None:
+        s_atoms = ase_geo[
+            [
+                atom.index
+                for atom in ase_geo
+                if atom.index not in mol_atoms and atom.symbol == substr_elem
+            ]
+        ]
+    else:
+        chem_symbols_arr = np.array(ase_geo.get_chemical_symbols())
+        s_atoms = ase_geo[chem_symbols_arr == substr_elem]
+        non_s_atoms = ase_geo[chem_symbols_arr != substr_elem]
     if len(s_atoms) == 0:
         return ase_geo, None
     layer_dz = 1.0  # ang
@@ -25,17 +36,24 @@ def geometrical_analysis(ase_geo, substr_elem):
     ]
     surf_z = np.mean(substr_top_layer.positions[:, 2])
 
-    mol_atoms = non_s_atoms[non_s_atoms.positions[:, 2] > surf_z]
+    if mol_atoms is None:
+        mol_atoms = non_s_atoms[non_s_atoms.positions[:, 2] > surf_z]
+    else:
+        mol_atoms = ase_geo[mol_atoms]
 
     return mol_atoms, surf_z
 
 
 @engine.calcfunction
-def analyze_structure(structure, substrate, mag_per_site, ads_h=None):
+def analyze_structure(
+    structure, substrate, mag_per_site, ads_h=None, molecule_atoms=None
+):
     ase_geo = structure.get_ase()
     substr_elem = substrate.value.split("(")[0]
 
-    mol_atoms, surf_z = geometrical_analysis(ase_geo, substr_elem)
+    mol_atoms, surf_z = geometrical_analysis(
+        ase_geo, substr_elem, mol_atoms=molecule_atoms
+    )
 
     if surf_z is None:
         if ads_h is None:
@@ -117,6 +135,7 @@ class Cp2kAdsorbedGwIcWorkChain(engine.WorkChain):
     Two different ways to run:
     1) geometry of a molecule adsorbed on a substrate
     2) isolated molecule & adsorption height
+    Magnetization specified on molecule atoms is relabelled for the gas phase calculation
     """
 
     @classmethod
@@ -151,18 +170,14 @@ def define(cls, spec):
             required=False,
             help="Protocol supported by the Cp2kMoleculeGwWorkChain.",
         )
-        spec.input(
-            "multiplicity",
-            valid_type=orm.Int,
-            default=lambda: orm.Int(0),
-            required=False,
-        )
-        spec.input(
-            "magnetization_per_site",
-            valid_type=orm.List,
-            default=lambda: orm.List(list=[]),
-            required=False,
-        )
+        spec.input("dft_params", valid_type=orm.Dict)
+        spec.input("sys_params", valid_type=orm.Dict)
+        # spec.input(
+        #    "molecule_atoms",
+        #    valid_type=orm.List,
+        #    default=lambda: orm.List(list=[]),
+        #    required=False,
+        # )
         spec.input_namespace(
             "options",
             valid_type=dict,
@@ -198,6 +213,13 @@ def define(cls, spec):
             required=False,
             help="Possibilities: ads_geo, gas_opt",
         )
+        spec.input(
+            "debug",
+            valid_type=orm.Bool,
+            default=lambda: orm.Bool(False),
+            required=False,
+            help="Run with fast parameters for debugging.",
+        )
 
         spec.outline(
             cls.setup,
@@ -227,20 +249,30 @@ def define(cls, spec):
     def setup(self):
         self.report("Inspecting input and setting up things.")
 
+        self.ctx.sys_params = self.inputs.sys_params.get_dict()
+        self.ctx.dft_params = self.inputs.dft_params.get_dict()
+
         n_atoms = len(self.inputs.structure.get_ase())
-        n_mags = len(list(self.inputs.magnetization_per_site))
+        mags = self.ctx.dft_params.get("magnetization_per_site", [])
+        n_mags = len(mags)
         if n_mags not in (0, n_atoms):
             self.report("If set, magnetization_per_site needs a value for every atom.")
             return self.exit_codes.ERROR_TERMINATION
 
         if self.inputs.substrate.value not in IC_PLANE_HEIGHTS:
             return self.exit_codes.ERROR_SUBSTR_NOT_SUPPORTED
 
+        molecule_atoms = self.ctx.sys_params.get(
+            "molecule_atoms", []
+        )  # self.inputs.molecule_atoms
+        if len(molecule_atoms) == 0:
+            molecule_atoms = None
         an_out = analyze_structure(
             self.inputs.structure,
             self.inputs.substrate,
-            self.inputs.magnetization_per_site,
+            mags,
             None if "ads_height" not in self.inputs else self.inputs.ads_height,
+            molecule_atoms=molecule_atoms,
         )
 
         if "mol_struct" not in an_out:
@@ -250,6 +282,10 @@ def setup(self):
         self.ctx.mol_struct = an_out["mol_struct"]
         self.ctx.image_plane_z = an_out["image_plane_z"]
         self.ctx.mol_mag_per_site = an_out["mol_mag_per_site"]
+        if "magnetization_per_site" in self.ctx.dft_params:
+            self.ctx.dft_params["magnetization_per_site"] = an_out["mol_mag_per_site"]
+
+        ###
 
         return engine.ExitCode(0)
 
@@ -260,12 +296,14 @@ def gas_opt(self):
         builder = Cp2kGeoOptWorkChain.get_builder()
         builder.code = self.inputs.code
         builder.structure = self.ctx.mol_struct
-        builder.multiplicity = self.inputs.multiplicity
-        builder.magnetization_per_site = self.ctx.mol_mag_per_site
-        builder.vdw = orm.Bool(True)
+        builder.sys_params = orm.Dict(dict=self.ctx.sys_params)
+        builder.dft_params = orm.Dict(dict=self.ctx.dft_params)
         builder.protocol = orm.Str("standard")
+        builder.metadata.label = "CP2K_GW_GAS_GeoOpt"
+        if self.inputs.debug:
+            builder.protocol = orm.Str("debug")
         builder.options = self.inputs.options.scf
-        builder.metadata.description = "gas_opt"
+        builder.metadata.description = "Gas phase geomery optimization"
         submitted_node = self.submit(builder)
         return engine.ToContext(gas_opt=submitted_node)
 
@@ -291,12 +329,13 @@ def ic(self):
         builder.protocol = self.inputs.protocol
         builder.structure = self.ctx.mol_struct
         builder.magnetization_per_site = self.ctx.mol_mag_per_site
-        builder.multiplicity = self.inputs.multiplicity
+        builder.multiplicity = orm.Int(self.ctx.dft_params.get("multiplicity", 0))
         builder.run_image_charge = orm.Bool(True)
         builder.z_ic_plane = self.ctx.image_plane_z
         builder.options.scf = self.inputs.options.scf
         builder.options.gw = self.inputs.options.ic
         builder.metadata.description = "ic"
+        builder.debug = self.inputs.debug
         submitted_node = self.submit(builder)
         return engine.ToContext(ic=submitted_node)
 
@@ -311,7 +350,7 @@ def gw(self):
         builder.protocol = self.inputs.protocol
         builder.structure = self.ctx.mol_struct
         builder.magnetization_per_site = self.ctx.mol_mag_per_site
-        builder.multiplicity = self.inputs.multiplicity
+        builder.multiplicity = orm.Int(self.ctx.dft_params.get("multiplicity", 0))
         builder.options.scf = self.inputs.options.scf
         builder.options.gw = self.inputs.options.gw
         builder.metadata.description = "gw"
@@ -323,6 +362,10 @@ def finalize(self):
         if not self.ctx.gw.is_finished_ok:
             return self.exit_codes.ERROR_TERMINATION
 
+        if hasattr(self.ctx, "gas_opt") and self.ctx.gas_opt is not None:
+            gas_geo_opt_params = self.ctx.gas_opt.outputs.output_parameters
+            self.out("gas_geo_opt_parameters", gas_geo_opt_params)
+
         gw_out_params = self.ctx.gw.outputs.gw_output_parameters
         ic_out_params = self.ctx.ic.outputs.gw_output_parameters
 
@@ -334,13 +377,8 @@ def finalize(self):
             "gw_ic_parameters", calc_gw_ic_parameters(gw_out_params, ic_out_params)
         )
 
-        # Add the workchain pk to the input structure extras
-        extras_label = "Cp2kAdsorbedGwIcWorkChain_pks"
-        if extras_label not in self.inputs.structure.base.extras.all:
-            extras_list = []
-        else:
-            extras_list = self.inputs.structure.base.extras.all[extras_label]
-        extras_list.append(self.node.pk)
-        self.inputs.structure.base.extras.set(extras_label, extras_list)
+        # Add extras.
+        struc = self.inputs.structure
+        common_utils.add_extras(struc, "surfaces", self.node.uuid)
 
         return engine.ExitCode(0)

aiida_nanotech_empa/workflows/cp2k/molecule_gw_workchain.py

@@ -119,7 +119,7 @@ def setup(self):
             self.report("Error: protocol not supported.")
             return self.exit_codes.ERROR_TERMINATION
 
-        self.ctx.protocols = cp2k_utils.load_cp2k_protocol("gw_protocols.yml")
+        self.ctx.protocols = cp2k_utils.load_protocol("gw_protocols.yml")
         structure = self.inputs.structure
         self.ctx.cutoff = cp2k_utils.get_cutoff(structure=structure)
         magnetization_per_site = copy.deepcopy(self.inputs.magnetization_per_site)

aiida_nanotech_empa/workflows/cp2k/protocols/gw_protocols.yml

@@ -93,6 +93,7 @@ gpw_std_scf_step:
       CELL:
         PERIODIC: NONE
   GLOBAL:
+    ELPA_KERNEL: AVX2_BLOCK2
     RUN_TYPE: ENERGY
     PRINT_LEVEL: MEDIUM
     EXTENDED_FFT_LENGTHS: ''
@@ -143,6 +144,7 @@ gpw_std_gw_step:
       CELL:
         PERIODIC: NONE
   GLOBAL:
+    ELPA_KERNEL: AVX2_BLOCK2
     RUN_TYPE: ENERGY
     PRINT_LEVEL: MEDIUM
     EXTENDED_FFT_LENGTHS: ''
@@ -197,6 +199,7 @@ gpw_std_ic_step:
       CELL:
         PERIODIC: NONE
   GLOBAL:
+    ELPA_KERNEL: AVX2_BLOCK2
     RUN_TYPE: ENERGY
     PRINT_LEVEL: MEDIUM
     EXTENDED_FFT_LENGTHS: ''
@@ -236,6 +239,7 @@ gapw_std_scf_step:
       CELL:
         PERIODIC: NONE
   GLOBAL:
+    ELPA_KERNEL: AVX2_BLOCK2
     RUN_TYPE: ENERGY
     PRINT_LEVEL: MEDIUM
     EXTENDED_FFT_LENGTHS: ''
@@ -284,6 +288,7 @@ gapw_std_gw_step:
       CELL:
         PERIODIC: NONE
   GLOBAL:
+    ELPA_KERNEL: AVX2_BLOCK2
     RUN_TYPE: ENERGY
     PRINT_LEVEL: MEDIUM
     EXTENDED_FFT_LENGTHS: ''
@@ -337,6 +342,7 @@ gapw_std_ic_step:
       CELL:
         PERIODIC: NONE
   GLOBAL:
+    ELPA_KERNEL: AVX2_BLOCK2
     RUN_TYPE: ENERGY
     PRINT_LEVEL: MEDIUM
     EXTENDED_FFT_LENGTHS: ''
@@ -376,6 +382,7 @@ gapw_hq_scf_step:
       CELL:
         PERIODIC: NONE
   GLOBAL:
+    ELPA_KERNEL: AVX2_BLOCK2
     RUN_TYPE: ENERGY
     PRINT_LEVEL: MEDIUM
     EXTENDED_FFT_LENGTHS: ''
@@ -424,6 +431,7 @@ gapw_hq_gw_step:
       CELL:
         PERIODIC: NONE
   GLOBAL:
+    ELPA_KERNEL: AVX2_BLOCK2
     RUN_TYPE: ENERGY
     PRINT_LEVEL: MEDIUM
     EXTENDED_FFT_LENGTHS:  ''
@@ -477,6 +485,7 @@ gapw_hq_ic_step:
       CELL:
         PERIODIC: NONE
   GLOBAL:
+    ELPA_KERNEL: AVX2_BLOCK2
     RUN_TYPE: ENERGY
     PRINT_LEVEL: MEDIUM
     EXTENDED_FFT_LENGTHS: ''