Phase 5: slim AiiDAEnsemble to use the shared OTF helpers

Delete _predict_with_model, _compute_properties, _write_model, _update_gp,
_train_gp and _clean_runs from AiiDAEnsemble (now inherited unchanged
from Ensemble). Replace the inline default setup + empty-model guard and
the TRAIN SECTION in compute_ensemble with calls to _otf_setup_defaults,
_otf_predict and _otf_maybe_train_and_write. Keep the 'AIIDA CALCULATIONS'
summary header via _clean_runs(..., label='AIIDA CALCULATIONS'). Drop the
now-unused flare imports (compute_mae, get_env_indices, is_std_in_bound,
deepcopy); keep FLARE_Atoms for the workchain monitoring loop.

Behavior-preserving: the existing tests/aiida_ensemble/ suite (which
calls the inherited helpers on an AiiDAEnsemble) passes unchanged.

Author: mesonepigreco

Modules/aiida_ensemble.py

@@ -3,7 +3,7 @@
 from __future__ import annotations
 
 from typing import Literal
-from copy import copy, deepcopy
+from copy import copy
 import time
 import sys
 
@@ -26,8 +26,6 @@
 
 try:
     from flare.atoms import FLARE_Atoms
-    from flare.io.output import compute_mae
-    from flare.learners.utils import get_env_indices, is_std_in_bound
 except ImportError:
     pass
 
@@ -103,15 +101,8 @@ def compute_ensemble( # pylint: disable=arguments-renamed
             # Predict only the ones that are within uncertainty, the rest do via DFT/AiiDA.
             if self.gp_model is not None:
                 number_of_atoms = structures[0].get_ase_atoms().get_global_number_of_atoms()
-                
-                if self.max_atoms_added < 0:
-                    self.max_atoms_added = number_of_atoms
-
-                if self.init_atoms is None:
-                    self.init_atoms = list(range(number_of_atoms))
-                
-                if len(self.gp_model.training_data) > 0:
-                    self._predict_with_model(structures, dft_indices)
+                self._otf_setup_defaults(number_of_atoms)
+                self._otf_predict(structures, dft_indices)
 
             dft_counts += len(dft_indices)
 
@@ -172,246 +163,17 @@ def compute_ensemble( # pylint: disable=arguments-renamed
 
                 # ================ TRAIN SECTION ================ #
                 if self.gp_model is not None:
-                    if dft_counts > 0:
-                        if self.train_hyps[0] <= len(self.gp_model.training_data) <= self.train_hyps[1]:
-                            self._train_gp()
-                        self._write_model()
+                    self._otf_maybe_train_and_write(dft_counts)
 
                 sys.stdout.flush()
 
         # ================ FINALIZE ================ #
         # if self.has_stress:
         #     self.stress_computed = copy(self.force_computed)
 
-        self._clean_runs(dft_counts)
+        self._clean_runs(dft_counts, label="AIIDA CALCULATIONS")
         self.init()
 
-    def _predict_with_model(
-        self,
-        structures: list[Structure],
-        dft_indices: list[int],
-    ) -> None:
-        """Predict on all the structures and estimate errors.
-
-        This is used to remove the structures indecis to not compute via AiiDA/DFT.
-
-        Args:
-        ----
-            structures: list of :class:`~cellconstructor.Structure.Structure` to simulate
-            dft_indices: list of integers related to the structures
-
-        """
-        sub_indices = deepcopy(dft_indices)
-
-        for index in sub_indices:
-            structure = structures[index]
-            atoms = FLARE_Atoms.from_ase_atoms(structure.get_ase_atoms())
-            self._compute_properties(atoms)
-
-            # get max uncertainty atoms
-            if self.build_mode == 'bayesian':
-                env_selection = is_std_in_bound
-            elif self.build_mode == 'direct':
-                env_selection = get_env_indices
-
-            tic = time.time()
-
-            std_in_bound, _ = env_selection(
-                self.std_tolerance,
-                self.gp_model.force_noise,
-                atoms,
-                max_atoms_added=self.max_atoms_added,
-                update_style=self.update_style,
-                update_threshold=self.update_threshold,
-            )
-
-            self.output.write_wall_time(tic, task='Env Selection')
-
-            if not std_in_bound:
-                print(f"[DFT CALLED] For structure with id={index}")
-            else:
-                print(f"[BFFS  USED] For structure with id={index}")
-                dft_indices.remove(index)  # remove index computed via ML-FF
-    
-                self.energies[index] = deepcopy(atoms.get_potential_energy()) / units.Ry # eV -> Ry
-                self.forces[index] = deepcopy(atoms.get_forces()) / units.Ry # eV/Ang -> Ry/Ang
-                if self.has_stress:
-                    self.stresses[index, :, :] = -1 * deepcopy(atoms.get_stress(voigt=False)) * (units.Bohr**3 / units.Ry) # -eV/(Ang^3) -> Ry/(Bohr^3)
-                    self.stress_computed[index] = True
-
-                self.force_computed[index] = True
-            
-            sys.stdout.flush()
-    
-
-    def _compute_properties(self, atoms: FLARE_Atoms) -> None:
-        """Compute energies, forces, stresses, and their uncertainties.
-
-        The FLARE ASE calculator is used, and write the results.
-
-        Args:
-        ----
-            atoms: a :class:`flare.atoms.FLARE_Atoms` instance for which to compute properties
-
-        """
-        tic = time.time()
-
-        atoms.calc = self.flare_calc
-        atoms.calc.calculate(atoms)
-
-        self.output.write_wall_time(tic, task='Compute Properties')
-
-    def _write_model(self) -> None:
-        """Write the current model in a JSON file."""
-        self.flare_calc.write_model(self.flare_name)
-
-    def _update_gp(
-        self,
-        atoms: FLARE_Atoms,
-        dft_frcs: ndarray,
-        dft_energy: float | None = None,
-        dft_stress: ndarray | None = None,
-    ) -> None:
-        """Update the current GP model.
-
-        Args:
-        ----
-            atoms (FLARE_Atoms): :class:`flare.atoms.FLARE_Atoms`` instance whose
-                local environments will be added to the training set.
-            dft_frcs (np.ndarray): DFT forces on all atoms in the structure, in eV/Angstrom.
-            dft_energy (float): total energy of the entire structure, in eV.
-            dft_stress (np.ndarray): DFT stress on structure.
-                Sign as in ASE (-1 in respect with QE), units in eV/Angstrom^3,
-                and in Voigt notation, i.e. (xx, yy, zz, yz, xz, xy).
-
-        """
-        from ase.calculators.singlepoint import SinglePointCalculator
-
-        tic = time.time()
-        is_empty_model = len(self.gp_model.training_data) == 0
-
-        # Here we make the decision to skip adding environments, if the stds 
-        # are within the user-defined boundaries, even if the ab-initio calculation
-        # was performed. This avoids slowing down the model, while the SSCHA
-        # is feeded with the DFT results.
-        if is_empty_model:
-            std_in_bound = False
-            train_atoms = self.init_atoms
-        else:
-            self._compute_properties(atoms)
-
-            # get max uncertainty atoms
-            if self.build_mode == 'bayesian':
-                env_selection = is_std_in_bound
-            elif self.build_mode == 'direct':
-                env_selection = get_env_indices
-
-            tic = time.time()
-
-            std_in_bound, train_atoms = env_selection(
-                self.std_tolerance,
-                self.gp_model.force_noise,
-                atoms,
-                max_atoms_added=self.max_atoms_added,
-                update_style=self.update_style,
-                update_threshold=self.update_threshold,
-            )
-
-            self.output.write_wall_time(tic, task='Env Selection')
-
-            # compute mae and write to output
-            e_mae, e_mav, f_mae, f_mav, s_mae, s_mav = compute_mae(
-                atoms,
-                self.output.basename,
-                atoms.potential_energy,
-                atoms.forces,
-                atoms.stress,
-                dft_energy,
-                dft_frcs,
-                dft_stress,
-                False,
-            )
-
-        if not std_in_bound:
-            if not is_empty_model:
-                stds = self.flare_calc.results.get('stds', np.zeros_like(dft_frcs))
-                self.output.add_atom_info(train_atoms, stds)
-
-            # Convert ASE stress (xx, yy, zz, yz, xz, xy) to FLARE stress
-            # (xx, xy, xz, yy, yz, zz).
-            flare_stress = None
-            if dft_stress is not None:
-                flare_stress = -np.array([
-                    dft_stress[0],
-                    dft_stress[5],
-                    dft_stress[4],
-                    dft_stress[1],
-                    dft_stress[3],
-                    dft_stress[2],
-                ])
-
-            results = {
-                'forces': dft_frcs,
-                'energy': dft_energy,
-                'free_energy': dft_energy,
-                'stress': dft_stress,
-            }
-
-            atoms.calc = SinglePointCalculator(atoms, **results)
-
-            # update gp model
-            self.gp_model.update_db(
-                atoms,
-                dft_frcs,
-                custom_range=train_atoms,
-                energy=dft_energy,
-                stress=flare_stress,
-            )
-
-            self.gp_model.set_L_alpha()
-            self.output.write_wall_time(tic, task='Update GP')
-        
-
-    def _train_gp(self) -> None:
-        """Optimize the hyperparameters of the current GP model."""
-        tic = time.time()
-
-        self.gp_model.train(logger_name=self.output_name + 'hyps.dat')
-
-        self.output.write_wall_time(tic, task='Train Hyps')
-
-        hyps, labels = self.gp_model.hyps_and_labels
-        if labels is None:
-            labels = self.gp_model.hyp_labels
-
-        self.output.write_hyps(
-            labels,
-            hyps,
-            tic, # actually here there should be the actual start time of the entire simulation
-            self.gp_model.likelihood,
-            self.gp_model.likelihood_gradient,
-            hyps_mask=self.gp_model.hyps_mask,
-        )
-        
-    def _clean_runs(self, dft_counts: int) -> None:
-        """Clean the failed runs and print summary.
-        
-        Args:
-        ----
-            dft_counts (int): number of performed DFT calculations.
- 
-        """
-        n_calcs = np.sum(self.force_computed.astype(int))
-        print('=============== SUMMARY AIIDA CALCULATIONS =============== \n')
-        print('Total structures included: ', n_calcs)
-        print('Structures not included  : ', self.N-n_calcs)
-        if self.gp_model is not None:
-            print('Steps using OTF-ML model : ', self.N-dft_counts)
-        print()
-        print('===================== END OF SUMMARY ===================== \n')
-        if n_calcs != self.N:
-            self.remove_noncomputed()
-
 
 def get_running_workchains(workchains: list[WorkChainNode], success: list[bool]) -> list:
     """Get the running workchains popping the finished ones.