Phase 1+2: hoist FLARE OTF helpers to Ensemble, propagate OTF state in split()

Move the backend-agnostic FLARE on-the-fly helpers (_predict_with_model,
_compute_properties, _write_model, _update_gp, _train_gp, _clean_runs)
from AiiDAEnsemble up to the base Ensemble class, so any cluster backend
can use them. Add _otf_setup_defaults, _otf_predict,
_otf_update_from_structure and _otf_maybe_train_and_write as the public
entry points the generic cluster driver will call, plus a
_import_flare_learners() lazy importer that keeps flare optional (G9).
Use Rydberg/Bohr from the Ensemble module units (G4) and add the
'deepcopy' import (G5).

Propagate the 16 OTF state attributes through Ensemble.split() via the
new OTF_STATE_ATTRIBUTES class constant. Without this, the sub-ensemble
built by compute_ensemble via get_noncomputed() would have gp_model=None
and OTF would be silently disabled on any backend going through the
compute_ensemble dispatcher (G1). The AiiDA path never splits with OTF
active, so this is behavior-preserving.

Author: mesonepigreco

Modules/Ensemble.py

@@ -2,6 +2,7 @@
 from __future__ import print_function, annotations
 import sys, os
 import warnings
+from copy import deepcopy
 import numpy as np
 import time
 #from scipy.special import tanh, sinh, cosh
@@ -122,6 +123,29 @@ def Main(self):
 
 __DEBUG_RHO__ = False
 
+
+def _import_flare_learners():
+    """Lazy import of the FLARE objects used by the on-the-fly helpers.
+
+    Returns
+    -------
+        (FLARE_Atoms, compute_mae, get_env_indices, is_std_in_bound)
+
+    Raises
+    ------
+        ImportError: with an actionable message if flare is not installed.
+    """
+    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 as exc:
+        raise ImportError(
+            "On-the-fly learning requires the 'flare' package "
+            "(https://github.com/mir-group/flare). Original error: {}".format(exc)
+        ) from exc
+    return FLARE_Atoms, compute_mae, get_env_indices, is_std_in_bound
+
 """
 This source contains the Ensemble class
 It is used to Load and Save info about the ensemble.
@@ -3962,6 +3986,14 @@ def split(self, split_mask):
 
         ens.all_properties = [self.all_properties[x] for x in np.arange(len(split_mask))[split_mask]]
 
+        # Propagate the on-the-fly (FLARE) state. Without this, the
+        # `computing_ensemble` built by compute_ensemble via get_noncomputed()
+        # would have gp_model=None and OTF would be silently disabled
+        # on any backend that goes through the compute_ensemble dispatcher
+        # (i.e. every Cluster flavor and plain ASE calculators).
+        for attr in self.OTF_STATE_ATTRIBUTES:
+            setattr(ens, attr, getattr(self, attr))
+
         return ens
 
 
@@ -4295,6 +4327,275 @@ def set_otf(
             train_hyps[1] = np.inf
         self.train_hyps = train_hyps
 
+    # ------------------------------------------------------------------ #
+    #  On-the-fly (FLARE) helpers — backend-agnostic, shared by all       #
+    #  clusters (BaseCluster driver) and the AiiDAEnsemble path.          #
+    # ------------------------------------------------------------------ #
+
+    OTF_STATE_ATTRIBUTES = (
+        "gp_model", "flare_calc", "std_tolerance", "max_atoms_added",
+        "update_style", "update_threshold", "build_mode", "output",
+        "output_name", "checkpt_name", "flare_name", "atoms_name",
+        "checkpt_files", "write_model", "init_atoms", "train_hyps",
+    )
+
+    def _otf_setup_defaults(self, number_of_atoms):
+        """Resolve the OTF settings that depend on the system size.
+
+        Extracted verbatim from AiiDAEnsemble.compute_ensemble.
+        """
+        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))
+
+    def _otf_predict(self, structures, candidate_indices):
+        """Predict with the ML model the structures within uncertainty.
+
+        No-op while the model is empty (first ab-initio batch), mirroring
+        the guard in AiiDAEnsemble.compute_ensemble.
+        """
+        if len(self.gp_model.training_data) > 0:
+            self._predict_with_model(structures, candidate_indices)
+
+    def _otf_update_from_structure(self, structure, dft_energy, dft_frcs,
+                                   dft_stress=None):
+        """Update the GP from a cellconstructor Structure and plain-eV data.
+
+        Backend-agnostic wrapper around _update_gp: callers (e.g. Cluster)
+        do not need to import flare at all.
+
+        Args:
+        ----
+            structure: cellconstructor.Structure.Structure that was computed.
+            dft_energy: total energy in eV.
+            dft_frcs: forces, shape (nat, 3), in eV/Angstrom.
+            dft_stress: stress in ASE convention and Voigt order
+                (xx, yy, zz, yz, xz, xy), in eV/Angstrom^3, or None.
+        """
+        FLARE_Atoms, _, _, _ = _import_flare_learners()
+        self._update_gp(
+            FLARE_Atoms.from_ase_atoms(structure.get_ase_atoms()),
+            dft_frcs=dft_frcs,
+            dft_energy=dft_energy,
+            dft_stress=dft_stress,
+        )
+
+    def _otf_maybe_train_and_write(self, dft_counts):
+        """Train the hyperparameters (inside the train_hyps window) and dump.
+
+        Exactly the 'TRAIN SECTION' of AiiDAEnsemble.compute_ensemble.
+        """
+        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()
+
+    def _predict_with_model(self, structures, candidate_indices):
+        """Predict on all the structures and estimate errors.
+
+        Structures whose predicted uncertainty is acceptable are removed from
+        `candidate_indices` (in place) and filled with the ML prediction.
+
+        Args:
+        ----
+            structures: list of cellconstructor.Structure.Structure
+            candidate_indices: indices of the structures still to be computed
+                ab-initio; modified in place.
+        """
+        FLARE_Atoms, _, get_env_indices, is_std_in_bound = _import_flare_learners()
+
+        sub_indices = deepcopy(candidate_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"[AB-INITIO CALLED] For structure with id={index}")
+            else:
+                print(f"[BFFS  USED] For structure with id={index}")
+                candidate_indices.remove(index)  # remove index computed via ML-FF
+
+                self.energies[index] = deepcopy(atoms.get_potential_energy()) / Rydberg # eV -> Ry
+                self.forces[index] = deepcopy(atoms.get_forces()) / Rydberg # eV/Ang -> Ry/Ang
+                if self.has_stress:
+                    self.stresses[index, :, :] = -1 * deepcopy(atoms.get_stress(voigt=False)) * (Bohr**3 / Rydberg) # -eV/(Ang^3) -> Ry/(Bohr^3)
+                    self.stress_computed[index] = True
+
+                self.force_computed[index] = True
+
+            sys.stdout.flush()
+
+    def _compute_properties(self, atoms):
+        """Compute energies, forces, stresses, and their uncertainties."""
+        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):
+        """Write the current model in a JSON file."""
+        self.flare_calc.write_model(self.flare_name)
+
+    def _update_gp(self, atoms, dft_frcs, dft_energy=None, dft_stress=None):
+        """Update the current GP model.
+
+        Args:
+        ----
+            atoms: :class:`flare.atoms.FLARE_Atoms` instance whose local
+                environments will be added to the training set.
+            dft_frcs (np.ndarray): ab-initio forces in eV/Angstrom.
+            dft_energy (float): total energy of the structure, in eV.
+            dft_stress (np.ndarray): ab-initio stress, ASE sign convention,
+                eV/Angstrom^3, Voigt order (xx, yy, zz, yz, xz, xy), or None.
+        """
+        from ase.calculators.singlepoint import SinglePointCalculator
+        _, compute_mae, get_env_indices, is_std_in_bound = _import_flare_learners()
+
+        tic = time.time()
+        is_empty_model = len(self.gp_model.training_data) == 0
+
+        # Skip adding environments if the stds are within the user-defined
+        # boundaries, even if the ab-initio calculation was performed.
+        if is_empty_model:
+            std_in_bound = False
+            train_atoms = self.init_atoms
+        else:
+            self._compute_properties(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):
+        """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, label="CALCULATIONS"):
+        """Clean the failed runs and print summary.
+
+        Args:
+        ----
+            dft_counts (int): number of performed ab-initio calculations.
+            label (str): identifier printed in the summary header.
+        """
+        n_calcs = np.sum(self.force_computed.astype(int))
+        print('=============== SUMMARY {} =============== \n'.format(label))
+        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()
+