Working version

Author: bastonero

Modules/Cluster.py

@@ -1511,10 +1511,11 @@ def compute_single_jobarray(jobs_id, calc):
                     if error_struct > 1e-2:
                         print("ERROR IDENTIFYING STRUCTURE!")
                         MSG = """
-Error in thread {}.
-     Displacement between the expected structure {} and the one readed from the calculator
-     is of {} A.
-""".format(threading.get_native_id(), jobs_id[i], error_struct)
+                            Error in thread {}.
+                            Displacement between the expected structure {} 
+                            and the one readed from the calculator
+                            is of {} A.
+                        """.format(threading.get_native_id(), jobs_id[i], error_struct)
                         print(MSG)
                         ensemble.structures[jobs_id[i]].save_scf('t_{}_error_struct_generated_{}.scf'.format(threading.get_native_id(), jobs_id[i]))
                         structures[i].save_scf('t_{}_error_struct_cmp_local_{}.scf'.format(threading.get_native_id(), jobs_id[i]))

Modules/aiida_ensemble.py

@@ -3,11 +3,16 @@
 
 import time
 import numpy as np
+from copy import copy
+import cellconstructor
 
 from .Ensemble import Ensemble
 try:
     import aiida
     import aiida_quantumespresso
+    from qe_tools import CONSTANTS
+    
+    gpa_to_rybohr3 = 1. / (CONSTANTS.ry_si/ CONSTANTS.bohr_si**3 / 1.0e9) # GPa -> Ry/Bohr^3
 except ImportError:
     import warnings
     warnings.warn('aiida or aiida-quantumespresso are not installed')
@@ -19,10 +24,10 @@ class AiiDAEnsemble(Ensemble):
     def compute_ensemble(
         self, 
         pw_code: str,
-        profile: str | None = None, 
         protocol: str = 'moderate',
         options: dict = None,
         overrides: dict = None,
+        group_label: str = None,
         **kwargs
     ):
         """Get ensemble properties.
@@ -35,39 +40,35 @@ def compute_ensemble(
         ---------
             pw_code:
                 The string associated with the AiiDA code for `pw.x`
-            profile: 
-                AiiDA profile to use during the calculations. If None, the default is used
             protocol:
                 The protocol to be used; available protocols are 'fast', 'moderate' and 'precise'
             options:
                 The options for the calculations, such as the resources, wall-time, etc.
             overrides:
                 The overrides for the get_builder_from_protocol
+            group_label:
+                The group label where to add the submitted nodes for eventual future inspection
             kwargs:
                 The kwargs for the get_builder_from_protocol
         """
-        from aiida import load_profile
-        from aiida.engine import submit
-        from aiida.plugins import WorkflowFactory
-        from aiida.orm import StructureData
+        from aiida.orm import load_group
 
-        PwBaseWorkChain = WorkflowFactory('quantumespresso.pw')
-    
-        if profile:
-            load_profile(profile)
-        else:
-            load_profile()
+        group = None if group_label is None else load_group(group_label)        
 
         # Check if not all the calculation needs to be done
-        n_calcs = np.sum( self.force_computed.astype(int)) 
+        if self.force_computed is None:
+            self.force_computed = np.array([False] * self.N, dtype=bool)
+
+        n_calcs = np.sum(self.force_computed.astype(int)) 
         computing_ensemble = self
 
+        self.has_stress = True # by default we calculate stresses with the `get_builder_from_protocol`
         if overrides:
             try:
                 tstress = overrides['pw']['parameters']['CONTROL']['tstress']
                 self.has_stress = tstress
             except KeyError:
-                self.has_stress = True # by default we calculate stresses with the `get_builder_from_protocol`
+                pass
 
         # Check wheter compute the whole ensemble, or just a small part
         should_i_merge = False
@@ -76,43 +77,37 @@ def compute_ensemble(
             computing_ensemble = self.get_noncomputed()
             self.remove_noncomputed()
 
-        # ============= HERE AIIDA PART ============= #
-        structures_data = [StructureData(ase=cc.get_ase_atoms()) for cc in self.structures]
-        workchains = []
-        
-        for i, structure in enumerate(structures_data):
-            builder = PwBaseWorkChain.get_builder_from_protocol(
-                code=pw_code,
-                structure=structure,
-                options=options,
-                overrides=overrides,
-                **kwargs
-            )
-            builder.label = str(i)
-            workchains.append(submit(builder))
+        # ============= AIIDA SECTION ============= #
+        workchains = submit_and_get_workchains(
+            structures=computing_ensemble.structures,
+            pw_code=pw_code,
+            temperature=self.current_T,
+            protocol=protocol,
+            options=options,
+            overrides=overrides,
+            **kwargs
+        )
 
-        workchains_success = []
-        workchains_failed = []
-
-        while(workchains):
-            for workchain in workchains:
-                if workchain.is_finished:
-                    if workchain.is_failed:
-                        workchains_failed.append(workchain)
-                    else:
-                        workchains_success.append(workchain)
-                    workchains.remove(workchain) # here it may be critical
-
-            time.sleep(60) # wait before checking again
+        if group:
+            group.add_nodes(workchains)
+
+        workchains_copy = copy(workchains)
+        while(workchains_copy):
+            workchains_copy = get_running_workchains(workchains_copy, computing_ensemble.force_computed)
+            if workchains_copy:
+                time.sleep(60) # wait before checking again
 
-        # ---- TO DO
-        #   * Take energies, forces, stresses
-        #   * Print which one was successfull
-        #   * Add correct labels to workchains to gather afterwards
-        #   * For sure something else
-        #   * Add the possibility of adding wc to group
-        #   * Keep track is some better way (printing?) of the workchains submitted
-        # ============= HERE AIIDA PART ============= #
+        for i, is_computed in enumerate(computing_ensemble.force_computed):
+            if is_computed:
+                out = workchains[i].outputs
+                computing_ensemble.energies[i] = out.output_parameters.dict.energy / CONSTANTS.ry_to_ev
+                computing_ensemble.forces[i] = out.output_trajectory.get_array('forces')[-1] / CONSTANTS.ry_to_ev
+                if self.has_stress:
+                    computing_ensemble.stresses[i] = out.output_trajectory.get_array('stress')[-1] * gpa_to_rybohr3
+        # ============= AIIDA SECTION ============= #
+
+        if self.has_stress:
+            computing_ensemble.stress_computed = copy(computing_ensemble.force_computed)
 
         print("CE BEFORE MERGE:", len(self.force_computed))
 
@@ -121,5 +116,76 @@ def compute_ensemble(
             self.merge(computing_ensemble)
         print("CE AFTER MERGE:", len(self.force_computed))
 
-        print('ENSEMBLE ALL PROPERTIES:', self.all_properties)
 
+def get_running_workchains(workchains: list, success: list[bool]) -> list:
+    """Get the running workchains popping the finished ones.
+    
+    Two extra array should be given to populate the successfully finished runs.
+    """
+    wcs_left = copy(workchains)
+
+    for workchain in workchains:
+        if workchain.is_finished:
+            if workchain.is_failed:
+                print(f'[FAILURE] for <PwBaseWorkChain> with PK={workchain.pk}')
+            else:
+                index = int(workchain.label.split('_')[-1])
+                success[index] = True
+                print(f'[SUCCESS] for <PwBaseWorkChain> with PK={workchain.pk}')
+                
+            wcs_left.remove(workchain) # here it may be critical
+    
+    return wcs_left
+
+
+def submit_and_get_workchains(
+    structures: list[cellconstructor.Structure.Structure],
+    pw_code: str,
+    temperature: float | int,
+    protocol: str = 'moderate',
+    options: dict = None,
+    overrides: dict = None,
+    **kwargs
+):
+    """Submit and return the workchains for a list of :class:`~cellconstructor.Structure.Structure`.
+
+    Parameters
+    ---------
+        structures:
+            a list of :class:`~cellconstructor.Structure.Structure` to run via PwBaseWorkChain.
+        pw_code:
+            The string associated with the AiiDA code for `pw.x`
+        temperature:
+            The temperature corresponding to the structures ensemble
+        protocol:
+            The protocol to be used; available protocols are 'fast', 'moderate' and 'precise'
+        options:
+            The options for the calculations, such as the resources, wall-time, etc.
+        overrides:
+            The overrides for the get_builder_from_protocol
+        kwargs:
+            The kwargs for the get_builder_from_protocol
+    """
+    from aiida.engine import submit
+    from aiida.plugins import WorkflowFactory
+    from aiida.orm import StructureData
+
+    PwBaseWorkChain = WorkflowFactory('quantumespresso.pw.base')
+
+    structures_data = [StructureData(ase=cc.get_ase_atoms()) for cc in structures]
+    workchains = []
+
+    for i, structure in enumerate(structures_data):
+        builder = PwBaseWorkChain.get_builder_from_protocol(
+            code=pw_code,
+            structure=structure,
+            protocol=protocol,
+            options=options,
+            overrides=overrides,
+            **kwargs
+        )
+        builder.metadata.label = 'T_{}_id_{}'.format(temperature, i)
+        workchains.append(submit(builder))
+        print(f'Launched <PwBaseWorkChain> with PK={workchains[-1].pk}')
+    
+    return workchains

tests/aiida_ensemble/__init__.py

@@ -0,0 +1,55 @@
+"""Tests for :mod:`sscha.aiida_ensemble`."""
+import pytest
+
+
+@pytest.mark.usefixtures('aiida_profile')
+def test_get_running_workchains(generate_workchain_pw_node):
+    """Test the :func:`sscha.aiida_ensemble.get_running_workchains` method."""
+    from plumpy import ProcessState
+    from sscha.aiida_ensemble import get_running_workchains
+
+    finished = ProcessState.FINISHED
+    running  = ProcessState.RUNNING
+
+    workchains = [
+        generate_workchain_pw_node(process_state=finished, exit_status=0, label='T_300_id_0'), 
+        generate_workchain_pw_node(process_state=finished, exit_status=300, label='T_300_id_1'),
+    ]
+    success = [False, False]
+    
+    wcs_left = get_running_workchains(workchains=workchains, success=success)
+    
+    assert not wcs_left
+    assert success == [True, False]
+
+    workchains = [
+        generate_workchain_pw_node(process_state=running, label='T_300_id_0'), 
+        generate_workchain_pw_node(process_state=finished, label='T_300_id_1'),
+        generate_workchain_pw_node(process_state=running, label='T_300_id_2'),
+    ]
+    success = [False, False, False]
+    
+    wcs_left = get_running_workchains(workchains=workchains, success=success)
+    
+    assert wcs_left == [workchains[0], workchains[2]]
+    assert success == [False, True, False]
+
+
+@pytest.mark.usefixtures('aiida_profile')
+def test_submit_and_get_workchains(fixture_code):
+    """Test the :func:`sscha.aiida_ensemble.submit_and_get_workchains` method."""
+    from cellconstructor.Structure import Structure
+    from sscha.aiida_ensemble import submit_and_get_workchains
+    
+    pw_code = fixture_code('quantumespresso.pw')
+    structures = [Structure(nat=1) for _ in range(5)]
+    
+    workchains = submit_and_get_workchains(
+        structures=structures, 
+        pw_code=pw_code,
+        temperature=300,
+    )
+    
+    assert len(workchains) == 5
+    
+