Use variadic args to calcfunctions, new since AiiDA 2.3

setup.cfg

@@ -31,6 +31,7 @@ python_requires = >=3.8
 # at the end we will probably have to do the same thing as aiidalab-qe
 # and publish a separate package to PyPi.
 install_requires =
+    aiida-core>=2.3.0
     aiidalab-widgets-base[smiles]~=2.0.0b2
     aiida-orca~=0.6.1
     aiidalab_atmospec_workchain @ file:///home/jovyan/apps/aiidalab-ispg/workflows/

workflows/aiidalab_atmospec_workchain/__init__.py

@@ -55,7 +55,9 @@ def combine(self):
 
 
 @calcfunction
-def pick_wigner_structure(wigner_structures, index):
+def pick_wigner_structure(
+    wigner_structures: TrajectoryData, index: int
+) -> StructureData:
     return wigner_structures.get_step_structure(index.value)
 
 
@@ -69,7 +71,10 @@ def add_orca_wf_guess(orca_params: Dict) -> Dict:
 
 @calcfunction
 def generate_wigner_structures(
-    minimum_structure, orca_output_dict, nsample, low_freq_thr
+    minimum_structure: StructureData,
+    orca_output_dict: dict,
+    nsample: int,
+    low_freq_thr: float,
 ):
     seed = orca_output_dict.extras["_aiida_hash"]
     ase_molecule = minimum_structure.get_ase()
@@ -378,11 +383,8 @@ def collect(self):
         # TODO: Include energies in TrajectoryData for optimized structures
         # TODO: Calculate Boltzmann weights and append them to TrajectoryData
         if self.inputs.optimize:
-            relaxed_structures = {
-                f"struct_{i}": wc.outputs.relaxed_structure
-                for i, wc in enumerate(self.ctx.confs)
-            }
-            trajectory = structures_to_trajectory(**relaxed_structures)
+            relaxed_structures = [wc.outputs.relaxed_structure for wc in self.ctx.confs]
+            trajectory = structures_to_trajectory(*relaxed_structures)
             self.out("relaxed_structures", trajectory)
 
 

workflows/aiidalab_atmospec_workchain/optimization.py

@@ -25,14 +25,15 @@
 AUtoKJ = AUtoKCAL * KCALtoKJ
 EVtoKJ = AUtoKCAL * KCALtoKJ / AUtoEV
 
-# TODO: Switch to variadic arguments (supported since AiiDA 2.3)
+
 @calcfunction
-def structures_to_trajectory(arrays: Array = None, **structures) -> TrajectoryData:
+def structures_to_trajectory(*structures, arrays: Array = None) -> TrajectoryData:
     """Concatenate a list of StructureData to TrajectoryData
 
-    Optionally, set additional data as Arrays.
+    :param arrays: Optional Array node to be merged into TrajectoryData (including extras)
+    :returns: TrajectoryData node
     """
-    traj = TrajectoryData(list(structures.values()))
+    traj = TrajectoryData(structures)
     if arrays is None:
         return traj
 
@@ -47,12 +48,12 @@ def structures_to_trajectory(arrays: Array = None, **structures) -> TrajectoryDa
     return traj
 
 
-def calc_boltzmann_weights(energies: list, T: float):
+def calc_boltzmann_weights(energies: list, T: float) -> np.ndarray:
     """Compute Boltzmann weights for a list of energies.
 
-    param energies: list of energies / kJ per mole
-    param T: temperature / Kelvin
-    returns: Boltzmann weights as numpy array
+    :param energies: list of energies / kJ per mole
+    :param T: temperature / Kelvin
+    :returns: Boltzmann weights as numpy array
     """
     # Molar gas constant, Avogadro times Boltzmann
     R = 8.3144598
@@ -64,19 +65,18 @@ def calc_boltzmann_weights(energies: list, T: float):
 
 
 @calcfunction
-def extract_trajectory_arrays(**orca_output_parameters) -> Array:
-    """Extract Gibbs energies and other useful stuff from the list
-    of ORCA output parameter dictionaries.
+def extract_trajectory_arrays(*orca_output_parameters) -> Array:
+    """Extract Gibbs energies et al from ORCA output parameters.
 
-    Return Array node, which will be appended to TrajectoryData node.
+    :returns: Array node which will be later appended to TrajectoryData node.
     """
     gibbs_energies = np.array(
-        [params["freeenergy"] for params in orca_output_parameters.values()]
+        [params["freeenergy"] for params in orca_output_parameters]
     )
     en0 = min(gibbs_energies)
     relative_gibbs_energies_kj = AUtoKJ * (gibbs_energies - en0)
 
-    temperature = list(orca_output_parameters.values())[0]["temperature"]
+    temperature = orca_output_parameters[0]["temperature"]
 
     boltzmann_weights = calc_boltzmann_weights(relative_gibbs_energies_kj, temperature)
 
@@ -189,16 +189,16 @@ def inspect_conformer_optimization(self):
 
     def collect_optimized_conformers(self):
         """Combine all optimized geometries into single TrajectoryData"""
-        # TODO: Switch to lists in AiiDA 2.3
-        relaxed_structures = {}
-        orca_output_params = {}
+        # TODO: Use list comprehension once we ensure the order of self.ctx.confs
+        relaxed_structures = []
+        orca_output_params = []
         for wc in self.ctx.confs:
-            relaxed_structures[f"struct_{wc.pk}"] = wc.outputs.relaxed_structure
-            orca_output_params[f"params_{wc.pk}"] = wc.outputs.output_parameters
+            relaxed_structures.append(wc.outputs.relaxed_structure)
+            orca_output_params.append(wc.outputs.output_parameters)
 
         array_data = None
         if len(self.ctx.confs) > 1:
-            array_data = extract_trajectory_arrays(**orca_output_params)
+            array_data = extract_trajectory_arrays(*orca_output_params)
 
-        trajectory = structures_to_trajectory(arrays=array_data, **relaxed_structures)
+        trajectory = structures_to_trajectory(*relaxed_structures, arrays=array_data)
         self.out("relaxed_structures", trajectory)