diff --git a/setup.cfg b/setup.cfg index ed9c6ee..407db6b 100644 --- a/setup.cfg +++ b/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/ diff --git a/workflows/aiidalab_atmospec_workchain/__init__.py b/workflows/aiidalab_atmospec_workchain/__init__.py index 05e85b6..fbb5efb 100644 --- a/workflows/aiidalab_atmospec_workchain/__init__.py +++ b/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) diff --git a/workflows/aiidalab_atmospec_workchain/optimization.py b/workflows/aiidalab_atmospec_workchain/optimization.py index ccb6534..5d72a18 100644 --- a/workflows/aiidalab_atmospec_workchain/optimization.py +++ b/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)