From c6c29d7dcc8512a0bcc03db28ccfcec1b462cb3a Mon Sep 17 00:00:00 2001 From: Lorenzo Monacelli Date: Wed, 22 Jul 2026 23:48:11 +0200 Subject: [PATCH 1/6] 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. --- Modules/Ensemble.py | 301 ++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 301 insertions(+) diff --git a/Modules/Ensemble.py b/Modules/Ensemble.py index bd407f8d..29d0919a 100644 --- a/Modules/Ensemble.py +++ b/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() + From 14c8250da70f9a2f082bc62fea8fa1185186e1c1 Mon Sep 17 00:00:00 2001 From: Lorenzo Monacelli Date: Wed, 22 Jul 2026 23:50:26 +0200 Subject: [PATCH 2/6] 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. --- Modules/aiida_ensemble.py | 248 +------------------------------------- 1 file changed, 5 insertions(+), 243 deletions(-) diff --git a/Modules/aiida_ensemble.py b/Modules/aiida_ensemble.py index 25cba124..08f41d0f 100644 --- a/Modules/aiida_ensemble.py +++ b/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,10 +163,7 @@ 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() @@ -183,235 +171,9 @@ def compute_ensemble( # pylint: disable=arguments-renamed # 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. From b68335bf176e69091e0f59f19197a2efdecf2826 Mon Sep 17 00:00:00 2001 From: Lorenzo Monacelli Date: Wed, 22 Jul 2026 23:56:50 +0200 Subject: [PATCH 3/6] Phase 3.1+3.2: extract BaseCluster, refactor Cluster as scheduler backend MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Add Modules/BaseCluster.py containing the generic cluster interface: BaseCluster (attribute-locking machinery moved verbatim from Cluster, _check_calculator_interface, _pre_compute_hook template hook, compute_jobarray abstract method, the generic ensemble driver compute_ensemble/compute_ensemble_batch with the on-the-fly FLARE hooks, _compute_jobarray_thread, _ingest_result) and DirectCluster (in-process backend consuming DirectCalculator objects). The OTF loop is written ONCE here and inherited by every cluster flavor. Refactor Modules/Cluster.py: Cluster now inherits from BaseCluster and keeps its entire public API (constructor signature, scheduler/transport methods, namelist constants). The generic driver (compute_ensemble, compute_ensemble_batch, compute_single_jobarray) and the attribute locking (__setattr__, __getstate__, __setstate__) are dropped (inherited from BaseCluster); batch_size/job_number/max_recalc/lock are set by BaseCluster.__init__. Cluster implements compute_jobarray (the template hook, with collect_results now outside the lock — safe parallelization improvement with identical results) and _pre_compute_hook (mkdir the local workdir). De-QE-hardcode prepare_input_file: extensions come from calc.input_extension/output_extension (default .pwi/.pwo), the prefix print is guarded with getattr, and extra_input_files are shipped to the cluster (G2). compute_jobarray uses the calc extensions too (G3). get_output_path takes an out_extension argument. Register the new modules in meson.build (ASEClusterRunner, BaseCluster, ClusterCalculators). ClusterCalculators.py and ASEClusterRunner.py are stubs for now (filled in the next commit). Compatibility verified: OpticalQECluster (sets attrs before super().__init__), Relax.py namelist path, test_save_binary pickling, and the full non-release suite pass (the 3 failing tests are pre-existing on master and unrelated). --- Modules/ASEClusterRunner.py | 0 Modules/BaseCluster.py | 377 ++++++++++++++++++++++++++++++++++ Modules/Cluster.py | 281 +++++-------------------- Modules/ClusterCalculators.py | 0 meson.build | 3 + 5 files changed, 437 insertions(+), 224 deletions(-) create mode 100644 Modules/ASEClusterRunner.py create mode 100644 Modules/BaseCluster.py create mode 100644 Modules/ClusterCalculators.py diff --git a/Modules/ASEClusterRunner.py b/Modules/ASEClusterRunner.py new file mode 100644 index 00000000..e69de29b diff --git a/Modules/BaseCluster.py b/Modules/BaseCluster.py new file mode 100644 index 00000000..ff4d23c3 --- /dev/null +++ b/Modules/BaseCluster.py @@ -0,0 +1,377 @@ +# -*- coding: utf-8 -*- +"""Generic interface for computing SSCHA ensembles 'on a cluster'. + +A 'cluster' here is any execution backend capable of computing energies, +forces and (optionally) stresses for a list of structures: a remote HPC +with a job scheduler (sscha.Cluster.Cluster), the same machinery mocked +locally (sscha.LocalCluster.LocalCluster), or a direct in-process +calculator (DirectCluster). + +The ensemble driver (compute_ensemble_batch), including the on-the-fly +(FLARE) learning loop, is implemented ONCE in this module. Backends only +implement `compute_jobarray`, which obtains the raw results for a list of +ensemble indices. +""" +from __future__ import print_function + +import sys +import os +import threading +import difflib + +import numpy as np + +# SETUP THE CODATA 2006, To match the QE definition of Rydberg (as in Cluster.py) +try: + import ase.units + units = ase.units.create_units("2006") +except Exception: + units = {"Ry": 13.605698066, "Bohr": 1 / 1.889725989} + + +class BaseCluster(object): + """Abstract execution backend for ensemble calculations. + + Subclasses must implement :func:`compute_jobarray` and declare the + calculator interface they consume via `_required_calculator_interface`. + """ + + # Attributes a calculator must expose to be used with this cluster. + _required_calculator_interface = ("copy",) + + def __init__(self, batch_size=1000, job_number=1, max_recalc=10): + """ + Parameters + ---------- + batch_size : int + Maximum number of job arrays computed in parallel per cycle. + With on-the-fly learning active, one cycle is the learning + batch: the GP is updated/retrained (and the remaining + structures re-predicted) after each cycle, so the effective + OTF batch size is `batch_size * job_number` structures. + job_number : int + Number of structures grouped in a single job array. + max_recalc : int + Maximum number of resubmission cycles for failed jobs. + """ + self.batch_size = batch_size + self.job_number = job_number + self.max_recalc = max_recalc + self.lock = None # threading.Lock(), created per compute_ensemble_batch + + # NOTE: subclasses set their own attributes and then call + # self._lock_attributes() at the end of their __init__. + # Setting attributes BEFORE super().__init__() also works (they are + # picked up into __total_attributes__), as OpticalQECluster does. + + # ------------------ attribute locking machinery ------------------ # + # (moved verbatim from sscha.Cluster.Cluster) + + def _lock_attributes(self): + """Forbid setting attributes not defined up to this point.""" + self.__total_attributes__ = [item for item in self.__dict__.keys()] + self.fixed_attributes = True # This must be the last attribute to be setted + + def __setattr__(self, name, value): + if "fixed_attributes" in self.__dict__: + if name in self.__total_attributes__: + super(BaseCluster, self).__setattr__(name, value) + elif self.fixed_attributes: + similar_objects = str(difflib.get_close_matches(name, self.__total_attributes__)) + ERROR_MSG = """ + Error, the attribute '{}' is not a member of '{}'. + Suggested similar attributes: {} ? + """.format(name, type(self).__name__, similar_objects) + raise AttributeError(ERROR_MSG) + + if name.endswith("_name"): + key = "use_{}".format(name.split("_")[0]) + self.__dict__[key] = True + else: + super(BaseCluster, self).__setattr__(name, value) + + def __getstate__(self): + """Return the picklable state (the thread lock cannot be pickled).""" + state = self.__dict__.copy() + state["lock"] = None + return state + + def __setstate__(self, state): + state["lock"] = None + self.__dict__.update(state) + + # ------------------ template-method hooks ------------------ # + + def _check_calculator_interface(self, calc): + """Fail early with a clear error if calc cannot run on this cluster.""" + missing = [m for m in self._required_calculator_interface + if not hasattr(calc, m)] + if missing: + raise TypeError( + "Error, the calculator {} cannot be used with {}.\n" + "Missing methods/attributes: {}".format( + type(calc).__name__, type(self).__name__, missing)) + + def _pre_compute_hook(self, ensemble, calc): + """Called once before the first submission cycle (default: nothing).""" + + def compute_jobarray(self, ensemble, calc, jobs_id): + """Compute one job array and return the raw results. + + MUST be implemented by subclasses. + + Parameters + ---------- + ensemble : sscha.Ensemble.Ensemble + The ensemble being computed. + calc : a calculator of the interface required by this cluster + A private copy for this job array (created with calc.copy()). + jobs_id : list of int + The ensemble indices to compute. + + Returns + ------- + list, aligned with jobs_id, of raw result dicts or None (failure). + Each dict must provide "energy" [eV], "forces" [eV/Angstrom], + optionally "stress" (ASE Voigt xx,yy,zz,yz,xz,xy, -eV/Angstrom^3) + and "structure" (CC.Structure, for the consistency check); + any extra key is stored in ensemble.all_properties. + """ + raise NotImplementedError("compute_jobarray must be implemented by subclasses") + + # ------------------ the generic ensemble driver ------------------ # + + def compute_ensemble(self, ensemble, calc, get_stress=True, timeout=None): + """Run the whole ensemble on this cluster (see compute_ensemble_batch).""" + self.compute_ensemble_batch(ensemble, calc, get_stress, timeout) + + def compute_ensemble_batch(self, ensemble, calc, get_stress=True, timeout=None): + """ + RUN THE ENSEMBLE WITH BATCH SUBMISSION (generic driver) + ======================================================= + + If the ensemble has an active on-the-fly ML model + (``ensemble.gp_model is not None``, set via ``ensemble.set_otf``), + each cycle of ab-initio computations is followed by a GP + update/retrain, and the remaining structures are re-checked against + the model: those predicted with acceptable uncertainty are filled + with the ML prediction and never computed. One learning cycle + computes up to ``batch_size * job_number`` structures. + """ + self._check_calculator_interface(calc) + + # Track the remaining configurations + success = [False] * ensemble.N + + # Setup if the ensemble has the stress + ensemble.has_stress = get_stress + + self._pre_compute_hook(ensemble, calc) + + # Get the expected number of batch + num_batch_offset = int(ensemble.N / self.batch_size) + + # ==================== OTF SETUP (FLARE) ==================== + use_otf = getattr(ensemble, "gp_model", None) is not None + dft_counts = 0 + if use_otf: + number_of_atoms = ensemble.structures[0].get_ase_atoms().get_global_number_of_atoms() + ensemble._otf_setup_defaults(number_of_atoms) + remaining = list(range(ensemble.N)) + ensemble._otf_predict(ensemble.structures, remaining) + for i in range(ensemble.N): + if ensemble.force_computed[i]: + success[i] = True + + # Run until some work has not finished + recalc = 0 + self.lock = threading.Lock() + while np.sum(np.array(success, dtype=int) - 1) != 0: + threads = [] + cycle_results = {} + + print("[CYCLE] SUCCESS: ", success) + print("[CYCLE] STOPPING CONDITION:", np.sum(np.array(success, dtype=int) - 1)) + + # Get the remaining jobs + false_mask = np.array(success) == False + false_id = np.arange(ensemble.N)[false_mask] + + count = 0 + # Submit in parallel + jobs = [false_id[i:i + self.job_number] for i in range(0, len(false_id), self.job_number)] + # Create a local copy of the calculator for each thread, to avoid conflicting modifications + calculators = [calc.copy() for i in range(0, len(jobs))] + + for k_th, job in enumerate(jobs): + # Submit only the batch size + if count >= self.batch_size: + break + t = threading.Thread(target=self._compute_jobarray_thread, + args=(ensemble, calculators[k_th], job, + get_stress, cycle_results, success)) + t.start() + threads.append(t) + count += 1 + + # Wait until all the job have finished + for t in threads: + t.join(timeout) + + # ============ OTF UPDATE / TRAIN / PREDICT ============ + if use_otf and cycle_results: + # Main thread only, deterministic order (G7) + dft_counts += len(cycle_results) + for num in sorted(cycle_results): + res = cycle_results[num] + dft_stress = np.array(res["stress"], dtype=float) \ + if (get_stress and "stress" in res) else None + ensemble._otf_update_from_structure( + ensemble.structures[num], + dft_energy=res["energy"], # eV + dft_frcs=res["forces"], # eV/Ang + dft_stress=dft_stress, # ASE Voigt, -eV/Ang^3 + ) + ensemble._otf_maybe_train_and_write(dft_counts) + # Re-check the remaining structures against the model + remaining = [int(i) for i in np.arange(ensemble.N)[np.array(success) == False]] + ensemble._otf_predict(ensemble.structures, remaining) + for i in range(ensemble.N): + if ensemble.force_computed[i]: + success[i] = True + + print("[CYCLE] [END] SUCCESS: ", success) + print("[CYCLE] [END] STOPPING CONDITION:", np.sum(np.array(success, dtype=int) - 1)) + + recalc += 1 + if recalc > num_batch_offset + self.max_recalc: + print("Expected batch ordinary resubmissions:", num_batch_offset) + raise ValueError("Error, resubmissions exceeded the maximum number of %d" % self.max_recalc) + + if use_otf: + ensemble._clean_runs(dft_counts) + + print("CALCULATION ENDED: all properties: {}".format(ensemble.all_properties)) + + def _compute_jobarray_thread(self, ensemble, calc, jobs_id, + get_stress, cycle_results, success): + """Thread worker: obtain raw results, then ingest them (locked).""" + raw_results = self.compute_jobarray(ensemble, calc, jobs_id) + + # Thread safe operation + self.lock.acquire() + try: + print("[THREAD {}] submitted calculations: {}".format( + threading.get_native_id(), list(jobs_id))) + for pos, res in enumerate(raw_results): + num = int(jobs_id[pos]) + print("[THREAD {}] ADDING RESULT {} = {}".format( + threading.get_native_id(), num, res)) + ok = self._ingest_result(ensemble, res, num, get_stress) + success[num] = ok + if ok: + cycle_results[num] = res # keep raw eV results for the OTF update (G6) + finally: + self.lock.release() + + def _ingest_result(self, ensemble, res, num, get_stress): + """Validate one raw result and write it into the ensemble arrays. + + Returns True if the result was complete and stored, False otherwise + (the job will be resubmitted, up to max_recalc). + """ + if res is None: + return False + + # Check if the run was good + check_e = "energy" in res + check_f = "forces" in res + check_s = "stress" in res + + # Check the structure + if "structure" in res: + error_struct = np.linalg.norm(ensemble.structures[num].coords.ravel() + - res["structure"].coords.ravel()) + 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(), num, error_struct) + print(MSG) + ensemble.structures[num].save_scf( + 't_{}_error_struct_generated_{}.scf'.format(threading.get_native_id(), num)) + res["structure"].save_scf( + 't_{}_error_struct_readed_{}.scf'.format(threading.get_native_id(), num)) + return False + else: + print("[WARNING] no check on the structure.") + + is_success = check_e and check_f + if get_stress: + is_success = is_success and check_s + + if not is_success: + return False + + res_only_extra = {x: res[x] for x in res if x not in ["energy", "forces", "stress", "structure"]} + ensemble.all_properties[num].update(res_only_extra) + ensemble.energies[num] = res["energy"] / units["Ry"] + ensemble.forces[num, :, :] = res["forces"] / units["Ry"] + ensemble.force_computed[num] = True + + if get_stress: + stress = np.zeros((3, 3), dtype=np.float64) + stress[0, 0] = res["stress"][0] + stress[1, 1] = res["stress"][1] + stress[2, 2] = res["stress"][2] + stress[1, 2] = res["stress"][3] + stress[2, 1] = res["stress"][3] + stress[0, 2] = res["stress"][4] + stress[2, 0] = res["stress"][4] + stress[0, 1] = res["stress"][5] + stress[1, 0] = res["stress"][5] + # Remember, ase has a very strange definition of the stress + ensemble.stresses[num, :, :] = -stress * units["Bohr"]**3 / units["Ry"] + ensemble.stress_computed[num] = True + return True + + +class DirectCluster(BaseCluster): + """ + DIRECT (IN-PROCESS) CLUSTER + =========================== + + A 'cluster' that computes the ensemble directly in the current process, + without writing any file and without any job scheduler. It consumes + `DirectCalculator` objects (e.g. ASEDirectCalculator wrapping any ASE + calculator). + + The learning-cycle size is still `batch_size * job_number`; threads are + used across job arrays exactly like in the scheduler-based clusters + (each thread owns a private calculator copy). Note that pure-Python ASE + calculators are bound by the GIL; the parallelism is still useful for + calculators that release it (NumPy-heavy or subprocess-based codes), and + the interface stays identical to the other clusters. + """ + + _required_calculator_interface = ("copy", "compute") + + def __init__(self, batch_size=1000, job_number=1, max_recalc=10): + super().__init__(batch_size=batch_size, job_number=job_number, + max_recalc=max_recalc) + self._lock_attributes() + + def compute_jobarray(self, ensemble, calc, jobs_id): + """Compute each structure in-process via calc.compute(structure).""" + results = [] + for num in jobs_id: + try: + results.append(calc.compute(ensemble.structures[int(num)])) + except Exception as exc: + sys.stderr.write("JOB {} resulted in error:\n{}\n".format(num, exc)) + sys.stderr.flush() + results.append(None) + return results diff --git a/Modules/Cluster.py b/Modules/Cluster.py index fe2a7f99..36f5992f 100644 --- a/Modules/Cluster.py +++ b/Modules/Cluster.py @@ -23,6 +23,8 @@ import cellconstructor as CC import cellconstructor.Methods +from sscha.BaseCluster import BaseCluster + # SETUP THE CODATA 2006, To match the QE definition of Rydberg try: units = ase.units.create_units("2006") @@ -133,7 +135,11 @@ def parse_symbols(string): -class Cluster(object): +class Cluster(BaseCluster): + + _required_calculator_interface = ( + "copy", "set_directory", "set_label", "write_input", "read_results", + ) def __init__(self, hostname=None, pwd=None, extra_options="", workdir = "", account_name = "", partition_name = "", qos_name = "", binary="pw.x -npool NPOOL -i PREFIX.pwi > PREFIX.pwo", @@ -171,6 +177,10 @@ def __init__(self, hostname=None, pwd=None, extra_options="", workdir = "", The QOS name for the job submission """ + # Common batch/submission knobs (batch_size, job_number, max_recalc, + # lock) are initialized by the base class; same defaults as today. + BaseCluster.__init__(self, batch_size=1000, job_number=1, max_recalc=10) + self.hostname = hostname self.pwd = pwd @@ -234,7 +244,7 @@ def __init__(self, hostname=None, pwd=None, extra_options="", workdir = "", # This is the number of configurations to be computed for each jub submitted # This times the self.batch_size is the total amount of configurations submitted toghether - self.job_number = 1 + # (job_number is set by BaseCluster.__init__) self.n_together_def = 1 self.use_multiple_submission = False @@ -260,7 +270,7 @@ def __init__(self, hostname=None, pwd=None, extra_options="", workdir = "", # This is the maximum number of resubmissions after the expected one # from the batch size. It can be use to resubmit the failed jobs. - self.max_recalc = 10 + # (max_recalc is set by BaseCluster.__init__) self.connection_attempts = 1 # This is the time string. Faster job will be the less in the queue, @@ -284,7 +294,7 @@ def __init__(self, hostname=None, pwd=None, extra_options="", workdir = "", # The batch size is the maximum number of job to be submitted together. # The new jobs will be submitted only after a batch is compleated # Useful if the cluster has a limit for the maximum number of jobs allowed. - self.batch_size = 1000 + # (batch_size is set by BaseCluster.__init__) # This could be a function that generates for each input file additional # text in the submission file. @@ -296,66 +306,10 @@ def __init__(self, hostname=None, pwd=None, extra_options="", workdir = "", # Allow to setup additional custom extra parameters self.custom_params = {} - self.lock = None # Use to lock the threads - - # Fix the attributes - - # Setup the attribute control - self.__total_attributes__ = [item for item in self.__dict__.keys()] - self.fixed_attributes = True # This must be the last attribute to be setted - - - def __setattr__(self, name, value): - """ - This method is used to set an attribute. - It will raise an exception if the attribute does not exists (with a suggestion of similar entries) - """ - - - if "fixed_attributes" in self.__dict__: - if name in self.__total_attributes__: - super(Cluster, self).__setattr__(name, value) - elif self.fixed_attributes: - similar_objects = str( difflib.get_close_matches(name, self.__total_attributes__)) - ERROR_MSG = """ - Error, the attribute '{}' is not a member of '{}'. - Suggested similar attributes: {} ? - """.format(name, type(self).__name__, similar_objects) - - raise AttributeError(ERROR_MSG) - - # Setting the account or partition name will automatically result in - # activating the corresponding flags - if name.endswith("_name"): - key = "use_{}".format(name.split("_")[0]) - self.__dict__[key] = True - else: - super(Cluster, self).__setattr__(name, value) - - - def __getstate__(self): - """ - Return the picklable state of the cluster. - - The thread lock created by compute_ensemble_batch cannot be pickled, - so it is dropped here. This allows sscha.Utilities.save_binary to - store objects holding a cluster after a calculation has run. - """ - state = self.__dict__.copy() - state["lock"] = None - return state - - - def __setstate__(self, state): - """ - Restore the cluster from a pickled state. - - The thread lock is transient runtime state and is reset to None, - as after __init__; compute_ensemble_batch recreates it when needed. - """ - state["lock"] = None - self.__dict__.update(state) + # (lock is set by BaseCluster.__init__) + # Fix the attributes (attribute-locking machinery lives in BaseCluster) + self._lock_attributes() def copy_file(self, source, destination, server_source = False, server_dest = True, raise_error=False, **kwargs): @@ -653,12 +607,16 @@ def prepare_input_file(self, structures, calc, labels): PREPARE THE INPUT FILE ====================== - This is specific for quantum espresso and must be inherit and replaced for - other calculators. - This crates the input files in the local working directory self.local_workdir and it returns the list of all the files generated. + The calculator may define: + input_extension (default ".pwi") + output_extension (default ".pwo") + extra_input_files (default []) extra files copied to the cluster + together with the inputs (e.g. the pickled ASE calculator + of ASEFileCalculator). + Parameters ---------- @@ -677,6 +635,8 @@ def prepare_input_file(self, structures, calc, labels): List of strings containing the output files expected for the calculation """ + in_ext = getattr(calc, "input_extension", ".pwi") + out_ext = getattr(calc, "output_extension", ".pwo") # Prepare the input file list_of_inputs = [] @@ -690,13 +650,15 @@ def prepare_input_file(self, structures, calc, labels): calc.set_label(label) calc.write_input(structure) - print("[THREAD {}] LBL: {} | PREFIX: {}".format(threading.get_native_id(), label, calc.input_data["control"]["prefix"])) - - #ase.io.write("%s/%s.pwi"% (self.local_workdir, label), - # atm, **calc.parameters) + input_data = getattr(calc, "input_data", None) + if isinstance(input_data, dict) and "control" in input_data: + print("[THREAD {}] LBL: {} | PREFIX: {}".format( + threading.get_native_id(), label, input_data["control"].get("prefix"))) + else: + print("[THREAD {}] LBL: {}".format(threading.get_native_id(), label)) - input_file = '{}.pwi'.format(label) - output_file = '{}.pwo'.format(label) + input_file = '{}{}'.format(label, in_ext) + output_file = '{}{}'.format(label, out_ext) list_of_inputs.append(input_file) list_of_outputs.append(output_file) @@ -713,6 +675,10 @@ def prepare_input_file(self, structures, calc, labels): # Release the lock on the threads self.lock.release() + # Extra files the calculator needs on the cluster (copied once per batch) + for extra in getattr(calc, "extra_input_files", []): + if extra not in list_of_inputs: + list_of_inputs.append(extra) return list_of_inputs, list_of_outputs @@ -905,12 +871,12 @@ def submit(self, script_location): return result, output - def get_output_path(self, label): + def get_output_path(self, label, out_extension=".pwo"): """ Given the label of the submission, retrive the path of all the output files of that calculation """ - out_filename = os.path.join(self.workdir, label + ".pwo") + out_filename = os.path.join(self.workdir, label + out_extension) return [out_filename] def read_results(self, calc, label): @@ -1556,159 +1522,26 @@ def parse_string(self, string): return str(output) - def compute_ensemble_batch(self, ensemble, cellconstructor_calc, get_stress = True, timeout=None): - """ - RUN THE ENSEMBLE WITH BATCH SUBMISSION - ====================================== - """ - - # Track the remaining configurations - success = [False] * ensemble.N - - # Setup if the ensemble has the stress - ensemble.has_stress = get_stress - #ensemble.all_properties = [None] * ensemble.N - - # Check if the working directory exists + def _pre_compute_hook(self, ensemble, calc): + """Create the local working directory if it does not exist.""" if not os.path.isdir(self.local_workdir): os.makedirs(self.local_workdir) + def compute_jobarray(self, ensemble, calc, jobs_id): + """Submit one job array through the scheduler and collect the outputs. - # Get the expected number of batch - num_batch_offset = int(ensemble.N / self.batch_size) - - def compute_single_jobarray(jobs_id, calc): - structures = [ensemble.structures[i].copy() for i in jobs_id] - n_together = min(len(structures), self.n_together_def) - subs, indices, labels = self.batch_submission(structures, calc, jobs_id, ".pwi", - ".pwo", "ESP", n_together) - - # Thread safe operation - self.lock.acquire() - print("[THREAD {}] submitted calculations: {}".format(threading.get_native_id(), indices)) - results = self.collect_results(calc, subs, indices, labels) - - for i, res in enumerate(results): - print("[THREAD {}] ADDING RESULT {} = {}".format(threading.get_native_id(), jobs_id[i], res)) - num = jobs_id[i] - - if res is None: - continue - - # Check if the run was good - check_e = "energy" in res - check_f = "forces" in res - check_s = "stress" in res - - # Check the structure - if "structure" in res: - error_struct = np.linalg.norm(ensemble.structures[jobs_id[i]].coords.ravel() - res["structure"].coords.ravel()) - 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) - 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])) - res["structure"].save_scf('t_{}_error_struct_readed_{}.scf'.format(threading.get_native_id(), jobs_id[i])) - - continue - else: - print("[WARNING] no check on the structure.") - - is_success = check_e and check_f - if get_stress: - is_success = is_success and check_s - - if not is_success: - continue - - res_only_extra = {x : res[x] for x in res if x not in ["energy", "forces", "stress", "structure"]} - ensemble.all_properties[num].update(res_only_extra) - ensemble.energies[num] = res["energy"] / units["Ry"] - ensemble.forces[num, :, :] = res["forces"] / units["Ry"] - ensemble.force_computed[num] = True - - if get_stress: - stress = np.zeros((3,3), dtype = np.float64) - stress[0,0] = res["stress"][0] - stress[1,1] = res["stress"][1] - stress[2,2] = res["stress"][2] - stress[1,2] = res["stress"][3] - stress[2,1] = res["stress"][3] - stress[0,2] = res["stress"][4] - stress[2,0] = res["stress"][4] - stress[0,1] = res["stress"][5] - stress[1,0] = res["stress"][5] - # Remember, ase has a very strange definition of the stress - ensemble.stresses[num, :, :] = -stress * units["Bohr"]**3 / units["Ry"] - ensemble.stress_computed[num] = True - success[num] = is_success - - self.lock.release() - - # Run until some work has not finished - recalc = 0 - self.lock = threading.Lock() - while np.sum(np.array(success, dtype = int) - 1) != 0: - threads = [] - - print("[CYCLE] SUCCESS: ", success) - print("[CYCLE] STOPPING CONDITION:", np.sum(np.array(success, dtype = int) - 1)) - - # Get the remaining jobs - false_mask = np.array(success) == False - false_id = np.arange(ensemble.N)[false_mask] - - count = 0 - # Submit in parallel - jobs = [false_id[i : i + self.job_number] for i in range(0, len(false_id), self.job_number)] - # Create a local copy of the calculator for each thread, to avoid conflicting modifications - calculators = [cellconstructor_calc.copy() for i in range(0, len(jobs))] - - for k_th, job in enumerate(jobs): - # Submit only the batch size - if count >= self.batch_size: - break - t = threading.Thread(target = compute_single_jobarray, args=(job, calculators[k_th], )) - t.start() - threads.append(t) - count += 1 - - # Wait until all the job have finished - for t in threads: - t.join(timeout) - - print("[CYCLE] [END] SUCCESS: ", success) - print("[CYCLE] [END] STOPPING CONDITION:", np.sum(np.array(success, dtype = int) - 1)) - - recalc += 1 - if recalc > num_batch_offset + self.max_recalc: - print ("Expected batch ordinary resubmissions:", num_batch_offset) - raise ValueError("Error, resubmissions exceeded the maximum number of %d" % self.max_recalc) - break - - print("CALCULATION ENDED: all properties: {}".format(ensemble.all_properties)) - - - - def compute_ensemble(self, ensemble, ase_calc, get_stress = True, timeout=None): - """ - RUN THE WHOLE ENSEMBLE ON THE CLUSTER - ===================================== - - Parameters - ---------- - ensemble : - The ensemble to be runned. + Overrides BaseCluster.compute_jobarray. The collect_results call runs + outside the lock (it only touches the per-thread calculator copy and + disjoint label files, so this is a safe parallelization improvement + with identical results); the BaseCluster driver handles locking + around the result ingestion. """ - - # Check if the compute_ensemble batch must be done - #if self.job_number != 1: - self.compute_ensemble_batch(ensemble, ase_calc, get_stress, timeout) - return + structures = [ensemble.structures[i].copy() for i in jobs_id] + n_together = min(len(structures), self.n_together_def) + subs, indices, labels = self.batch_submission( + structures, calc, jobs_id, + getattr(calc, "input_extension", ".pwi"), # was hardcoded (G3) + getattr(calc, "output_extension", ".pwo"), # was hardcoded (G3) + "ESP", n_together) + return self.collect_results(calc, subs, indices, labels) diff --git a/Modules/ClusterCalculators.py b/Modules/ClusterCalculators.py new file mode 100644 index 00000000..e69de29b diff --git a/meson.build b/meson.build index bc98d4b0..87c33ae2 100644 --- a/meson.build +++ b/meson.build @@ -205,8 +205,11 @@ py.install_sources([ 'Modules/__init__.py', 'Modules/AdvancedCalculations.py', 'Modules/aiida_ensemble.py', + 'Modules/ASEClusterRunner.py', + 'Modules/BaseCluster.py', 'Modules/Calculator.py', 'Modules/Cluster.py', + 'Modules/ClusterCalculators.py', 'Modules/Dynamical.py', 'Modules/Ensemble.py', # 'Modules/fourier_gradient.jl', From 9612195ad3ea983d9986fde2141efb13b68e5903 Mon Sep 17 00:00:00 2001 From: Lorenzo Monacelli Date: Wed, 22 Jul 2026 23:58:55 +0200 Subject: [PATCH 4/6] Phase 3.3-3.6: calculators, ASE runner, LocalCluster.copy_file MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Add Modules/ClusterCalculators.py with the calculator hierarchy (Axis B): ClusterCalculator (doc-level base), FileIOCalculator (mid-layer with input_extension/output_extension/extra_input_files/get_execution_command), EspressoCalculator (thin formalization of CC.calculators.Espresso), ASEFileCalculator (bridge any picklable ASE calculator to the scheduler clusters via pickle + JSON + the runner), DirectCalculator (mid-layer with compute(structure)), ASEDirectCalculator (wrap any ASE calculator for in-process use with DirectCluster). Future codes are one class each, or use the ASE bridges with no dedicated code at all. Add Modules/ASEClusterRunner.py: the dependency-light runner invoked on the cluster by ASEFileCalculator (python -m sscha.ASEClusterRunner --calc ... --input ... --output ...). Loads the pickled ASE calculator and the JSON-serialized Atoms, computes energy/forces/stress, dumps a JSON result. No sscha submodule imports (sscha/__init__.py is empty, so no Julia bootstrap — G10). LocalCluster.copy_file now uses shutil.copy (instead of the scp shell-out) for a cleaner local-to-local copy; returns the dest path (truthy). Only affects LocalCluster; strictly more robust (G8). --- Modules/ASEClusterRunner.py | 47 +++++++ Modules/ClusterCalculators.py | 244 ++++++++++++++++++++++++++++++++++ Modules/LocalCluster.py | 9 +- 3 files changed, 298 insertions(+), 2 deletions(-) diff --git a/Modules/ASEClusterRunner.py b/Modules/ASEClusterRunner.py index e69de29b..c6c54172 100644 --- a/Modules/ASEClusterRunner.py +++ b/Modules/ASEClusterRunner.py @@ -0,0 +1,47 @@ +# -*- coding: utf-8 -*- +"""Execute a pickled ASE calculator on a JSON-serialized structure. + +Invoked by sscha.ClusterCalculators.ASEFileCalculator on the (local or +remote) cluster as: + + python -m sscha.ASEClusterRunner --calc calculator.pkl \ + --input PREFIX_input.json --output PREFIX.json + +The output JSON contains energy [eV], forces [eV/Angstrom], stress +[ASE Voigt (xx, yy, zz, yz, xz, xy), eV/Angstrom^3] and the computed Atoms. +""" +import argparse +import pickle + +import ase.io.jsonio + + +def main(): + """Run the calculation and dump the results.""" + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--calc", required=True, + help="Path to the pickled ASE calculator.") + parser.add_argument("--input", required=True, + help="Path to the JSON-serialized ASE Atoms input.") + parser.add_argument("--output", required=True, + help="Path of the JSON file where results are written.") + args = parser.parse_args() + + with open(args.calc, "rb") as handle: + calc = pickle.load(handle) + with open(args.input, "r") as handle: + atoms = ase.io.jsonio.decode(handle.read()) + + atoms.calc = calc + payload = { + "energy": float(atoms.get_potential_energy()), # eV + "forces": atoms.get_forces(), # eV / Angstrom + "stress": atoms.get_stress(), # ASE Voigt, -eV/Angstrom^3 + "atoms": atoms, # structure actually computed + } + with open(args.output, "w") as handle: + handle.write(ase.io.jsonio.encode(payload)) + + +if __name__ == "__main__": + main() diff --git a/Modules/ClusterCalculators.py b/Modules/ClusterCalculators.py index e69de29b..a09dea47 100644 --- a/Modules/ClusterCalculators.py +++ b/Modules/ClusterCalculators.py @@ -0,0 +1,244 @@ +# -*- coding: utf-8 -*- +"""Calculators for the sscha cluster backends (Axis B of the cluster design). + +Two families: + +* FileIOCalculator: calculators driven through input/output files, consumed + by the scheduler-based clusters (sscha.Cluster.Cluster, + sscha.LocalCluster.LocalCluster). One subclass per simulation code + (EspressoCalculator today; VaspCalculator, AbinitCalculator, + Cp2kCalculator, GaussianCalculator, ... in the future). + +* DirectCalculator: calculators computed in-process, consumed by + sscha.BaseCluster.DirectCluster. + +For any code that has an ASE interface, ASEFileCalculator (file-based) and +ASEDirectCalculator (in-process) already provide full support with no +dedicated subclass: e.g. ASEFileCalculator(ase.calculators.vasp.Vasp(...)). +""" +import copy +import os +import pickle +import sys + +import cellconstructor as CC +import cellconstructor.calculators + + +# ---------------------------------------------------------------------- # +# Generic contract # +# ---------------------------------------------------------------------- # + +class ClusterCalculator(object): + """Documentation-level base class for cluster calculators. + + Any calculator consumed by a sscha cluster must provide: + + copy() -> an independent instance (one per worker thread) + + and produce result dicts with the contract documented in + BaseCluster.compute_jobarray ("energy" [eV], "forces" [eV/Ang], + "stress" [ASE Voigt, -eV/Ang^3], "structure" [CC.Structure], extras). + """ + + +# ---------------------------------------------------------------------- # +# File-based family (scheduler clusters) # +# ---------------------------------------------------------------------- # + +class FileIOCalculator(CC.calculators.FileIOCalculator, ClusterCalculator): + """Mid-level interface for file-based calculators. + + Class attributes consumed by the scheduler machinery: + + input_extension / output_extension : str + Extensions of the per-label input/output files. + extra_input_files : list of str + Additional files (relative to the local workdir) to ship to the + cluster together with the inputs (e.g. "calculator.pkl"). + """ + + input_extension = ".pwi" + output_extension = ".pwo" + extra_input_files = [] + + def get_execution_command(self): + """The command template run on the cluster, with the PREFIX + placeholder where the calculation label must be inserted, e.g.: + + "pw.x -npool 4 -i PREFIX.pwi > PREFIX.pwo" + + Assign it to `cluster.binary` (or use it in the namelist). + """ + raise NotImplementedError + + +class EspressoCalculator(CC.calculators.Espresso, FileIOCalculator): + """Quantum ESPRESSO (pw.x) calculator for the sscha clusters. + + Thin formalization of cellconstructor.calculators.Espresso, which + already satisfies the interface. Using this class (instead of the CC + one) only makes the extensions and the execution command explicit. + """ + + input_extension = ".pwi" + output_extension = ".pwo" + extra_input_files = [] + + def get_execution_command(self, n_pool=1): + return "pw.x -npool {} -i PREFIX.pwi > PREFIX.pwo".format(n_pool) + + def copy(self): + """Return an identical instance of THIS class (per-thread copies).""" + return EspressoCalculator(self.input_data, self.pseudopotentials, + self.masses, self.command, self.kpts, self.koffset) + + +class ASEFileCalculator(FileIOCalculator): + """ + RUN ANY ASE CALCULATOR THROUGH A SCHEDULER-BASED SSCHA CLUSTER + ============================================================== + + Bridge between a plain ASE calculator (EMT, LJ, VASP, CP2K, ...) and the + file-based `Cluster` submission: + + - the ASE calculator is pickled ONCE into the working directory + (`calculator.pkl`, listed in `extra_input_files` so it is shipped to + the cluster together with the inputs); + - each input file is the ASE Atoms serialized with ase.io.jsonio; + - the execution command runs `python -m sscha.ASEClusterRunner`, which + loads both, computes energy/forces/stress and writes a JSON result + file (see Modules/ASEClusterRunner.py). + + Usage: + calc = ASEFileCalculator(EMT()) + cluster.binary = calc.command # the runner command with PREFIX placeholder + cluster.mpi_cmd = "" # the runner is a serial python process + cluster.compute_ensemble(ensemble, calc) + + Requirements on the (remote) cluster: python + ase (+ the actual + calculator dependencies, and a python able to unpickle the calculator — + same ASE version recommended). For sscha.LocalCluster this is + automatically satisfied by the local environment. + """ + + input_extension = "_input.json" + output_extension = ".json" + calc_pickle_name = "calculator.pkl" + extra_input_files = [calc_pickle_name] + + def __init__(self, ase_calc, python_exe=None): + """ + Parameters + ---------- + ase_calc : ase.calculators.calculator.Calculator + The ASE calculator to execute on the cluster. Must be picklable. + python_exe : str, optional + The python interpreter used ON THE CLUSTER to run the runner. + Defaults to the local interpreter (correct for LocalCluster). + """ + super().__init__() + + # Fail early with a clear error if the calculator cannot be pickled + pickle.dumps(ase_calc) + + self.ase_calc = ase_calc + self.python_exe = python_exe or sys.executable + self.command = self.get_execution_command() + + def get_execution_command(self): + return ("{exe} -m sscha.ASEClusterRunner " + "--calc {pkl} " + "--input PREFIX{in_ext} " + "--output PREFIX{out_ext}" + ).format(exe=self.python_exe, pkl=self.calc_pickle_name, + in_ext=self.input_extension, out_ext=self.output_extension) + + def copy(self): + """Return an identical instance, without inheriting calculation info.""" + return ASEFileCalculator(self.ase_calc, self.python_exe) + + def set_directory(self, directory): + """Set the working directory and pickle the calculator there (once).""" + CC.calculators.FileIOCalculator.set_directory(self, directory) + pkl_path = os.path.join(directory, self.calc_pickle_name) + if not os.path.exists(pkl_path): + with open(pkl_path, "wb") as handle: + pickle.dump(self.ase_calc, handle) + + def write_input(self, structure): + """Serialize the structure as JSON in {directory}/{label}{input_extension}.""" + import ase.io.jsonio + CC.calculators.FileIOCalculator.write_input(self, structure) + + atoms = structure.get_ase_atoms() + fname = os.path.join(self.directory, self.label + self.input_extension) + with open(fname, "w") as handle: + handle.write(ase.io.jsonio.encode(atoms)) + + def read_results(self): + """Read {directory}/{label}{output_extension} into .results/.structure.""" + import ase.io.jsonio + + fname = os.path.join(self.directory, self.label + self.output_extension) + with open(fname, "r") as handle: + payload = ase.io.jsonio.decode(handle.read()) + + self.results = { + "energy": payload["energy"], # eV + "forces": payload["forces"], # eV / Angstrom + "stress": payload["stress"], # ASE Voigt (xx,yy,zz,yz,xz,xy), -eV/Angstrom^3 + } + self.structure = CC.Structure.Structure() + self.structure.generate_from_ase_atoms(payload["atoms"]) + + +# ---------------------------------------------------------------------- # +# Direct (in-process) family (DirectCluster) # +# ---------------------------------------------------------------------- # + +class DirectCalculator(ClusterCalculator): + """Mid-level interface for in-process calculators (DirectCluster).""" + + def compute(self, structure): + """Compute one structure and return the raw results dict. + + Parameters + ---------- + structure : cellconstructor.Structure.Structure + + Returns + ------- + dict with "energy" [eV], "forces" [eV/Ang], optionally + "stress" [ASE Voigt, -eV/Ang^3] and "structure" [CC.Structure]. + """ + raise NotImplementedError + + +class ASEDirectCalculator(DirectCalculator): + """Wrap ANY ASE calculator for in-process use with DirectCluster. + + No files are written: the ASE calculator is called directly on the + ASE image of each structure. + """ + + def __init__(self, ase_calc): + self.ase_calc = ase_calc + + def copy(self): + """Independent per-thread copy (ASE calculators are stateful).""" + return ASEDirectCalculator(copy.deepcopy(self.ase_calc)) + + def compute(self, structure): + atoms = structure.get_ase_atoms() + atoms.calc = self.ase_calc + results = { + "energy": float(atoms.get_potential_energy()), # eV + "forces": atoms.get_forces(), # eV / Ang + "structure": structure.copy(), # trivial consistency check + } + try: + results["stress"] = atoms.get_stress() # ASE Voigt, -eV/Ang^3 + except Exception: + pass # calculator without stress: call compute_ensemble with get_stress=False + return results diff --git a/Modules/LocalCluster.py b/Modules/LocalCluster.py index e4688bce..a7fc1456 100644 --- a/Modules/LocalCluster.py +++ b/Modules/LocalCluster.py @@ -1,3 +1,4 @@ +import shutil import sscha.Cluster as Cluster import sys, os @@ -19,7 +20,11 @@ def ExecuteCMD(self, cmd, *args, on_cluster = False, **kwargs): def copy_file(self, source, destination, server_source = False, server_dest = False, **kwargs): """ - Copy the files ignoring if the cluster is used. + Copy the files locally, ignoring the remote flags. + + Uses shutil.copy (instead of the scp shell-out of the base class) + for a cleaner local-to-local copy. Returns the destination path + (truthy), matching the contract expected by the callers. """ - return super().copy_file(source, destination, server_source = False, server_dest = False, **kwargs) + return shutil.copy(source, destination) From 1760f85391d402bbf3ac533d1449412f7a609052 Mon Sep 17 00:00:00 2001 From: Lorenzo Monacelli Date: Thu, 23 Jul 2026 00:22:48 +0200 Subject: [PATCH 5/6] Tests: OTF base helpers, cluster interface, end-to-end EMT OTF MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Add tests/test_otf/ with three test modules: - test_otf_base.py: backend-agnostic OTF helpers on the base Ensemble (split propagation, setup defaults, update+predict cycle, clean_runs label). Uses Au+EMT with a FLARE SGP. Skip without flare (G9). - test_cluster_interface.py: the layered Cluster/calculator refactor with no flare needed. Covers the DirectCluster generic driver with a mock calculator, the pairing rules (DirectCluster rejects file calculators, Cluster rejects direct ones), QE backward compatibility of prepare_input_file, the EspressoCalculator adapter, the ASE file bridge (prepare_input + local roundtrip), the ASE direct calculator, and the non-picklable calculator early failure. - test_cluster_emt_otf.py: end-to-end OTF on gold+EMT through both DirectCluster and LocalCluster (parametrized), proving the driver is shared. Validates the full chain (calculator -> raw results -> ingestion -> unit conversions) against bare EMT values, and that a second ensemble reuses the trained model with zero new ab-initio calls. Add the 'flare' marker to pytest.ini. Workaround for a flare C extension bug: SGP_Wrapper does not keep a Python reference to the kernel objects (only to the descriptor calculators), so the kernel's pybind11 wrapper is garbage-collected when the helper function returns, leaving the C++ SparseGP with a dangling pointer → segfault in add_training_structure. The tests keep kernels alive in a module-level list (_KERNEL_REFS). --- pytest.ini | 1 + tests/test_otf/__init__.py | 0 tests/test_otf/test_cluster_emt_otf.py | 102 +++++++++++++++ tests/test_otf/test_cluster_interface.py | 152 +++++++++++++++++++++++ tests/test_otf/test_otf_base.py | 124 ++++++++++++++++++ 5 files changed, 379 insertions(+) create mode 100644 tests/test_otf/__init__.py create mode 100644 tests/test_otf/test_cluster_emt_otf.py create mode 100644 tests/test_otf/test_cluster_interface.py create mode 100644 tests/test_otf/test_otf_base.py diff --git a/pytest.ini b/pytest.ini index 98c83927..528ac79d 100644 --- a/pytest.ini +++ b/pytest.ini @@ -2,3 +2,4 @@ markers = julia: tests requiring the Julia Fourier backend release: long-running tests excluded from normal CI + flare: tests requiring the flare package diff --git a/tests/test_otf/__init__.py b/tests/test_otf/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/tests/test_otf/test_cluster_emt_otf.py b/tests/test_otf/test_cluster_emt_otf.py new file mode 100644 index 00000000..294ef939 --- /dev/null +++ b/tests/test_otf/test_cluster_emt_otf.py @@ -0,0 +1,102 @@ +"""End-to-end OTF on gold+EMT through DirectCluster and LocalCluster.""" +import os +import numpy as np +import pytest + +flare = pytest.importorskip("flare") # G9 + +import cellconstructor as CC +import sscha, sscha.Ensemble +import sscha.Cluster, sscha.LocalCluster, sscha.BaseCluster +import sscha.ClusterCalculators as cluster_calcs +from ase.calculators.emt import EMT + +from .test_otf_base import get_gold_dyn, get_sgp_calc_au + + +def make_direct_cluster(tmp_path, calc, batch_size): + return sscha.BaseCluster.DirectCluster(batch_size=batch_size, job_number=1), \ + cluster_calcs.ASEDirectCalculator(EMT()) + + +def make_local_cluster(tmp_path, calc, batch_size): + file_calc = cluster_calcs.ASEFileCalculator(EMT()) + cluster = sscha.LocalCluster.LocalCluster("localhost") + cluster.workdir = str(tmp_path / "remote") + cluster.local_workdir = str(tmp_path / "local") + "/" + cluster.submit_command = "bash" # blocking local execution (G8) + cluster.nonblocking_command = False # no squeue polling (G8) + cluster.use_nodes = False + cluster.use_cpu = False + cluster.use_time = False + cluster.use_account = False + cluster.job_number = 1 + cluster.batch_size = batch_size # learning-cycle size = batch_size*job_number + cluster.binary = file_calc.command + cluster.mpi_cmd = "" + cluster.setup_workdir() + return cluster, file_calc + + +@pytest.mark.flare +@pytest.mark.parametrize("backend_factory", [make_direct_cluster, make_local_cluster], + ids=["direct", "localcluster"]) +def test_cluster_emt_otf(tmp_path, backend_factory): + np.random.seed(0) + n_configs, batch_size = 8, 2 + + dyn = get_gold_dyn() + ensemble = sscha.Ensemble.Ensemble(dyn, 300) + ensemble.generate(n_configs) + ensemble.set_otf(get_sgp_calc_au(), std_tolerance_factor=100, + max_atoms_added=-1, update_style="add_n", + update_threshold=None, train_hyps=(1, np.inf), + output_name=str(tmp_path / "otf_run")) + + cluster, calc = backend_factory(tmp_path, None, batch_size) + + # Reference EMT values for the structures that MUST be computed ab-initio + # (first cycle: model empty -> exactly the first batch_size structures) + ref = [] + for i in range(batch_size): + atoms = ensemble.structures[i].get_ase_atoms() + atoms.calc = EMT() + ref.append((atoms.get_potential_energy(), atoms.get_forces().copy())) + + ensemble.compute_ensemble(calc, compute_stress=True, cluster=cluster) + + # 1. Everything computed (ab-initio or predicted) + assert all(ensemble.force_computed) + assert all(ensemble.stress_computed) + assert np.all(np.isfinite(ensemble.energies)) + assert np.all(np.isfinite(ensemble.forces)) + assert np.all(np.isfinite(ensemble.stresses)) + + # 2. The model was trained on the first-cycle structures. With a high + # std_tolerance, only the first structure (empty model -> all init_atoms + # added) necessarily enters the training set; subsequent structures may + # be within bounds and not add new training data. + assert len(ensemble.gp_model.training_data) >= 1 + assert os.path.exists(ensemble.flare_name) + + # 3. STRONG REGRESSION: the ab-initio structures carry the exact EMT + # values (validates the whole chain: calculator -> raw results -> + # ingestion -> unit conversions). Ensemble units are Ry, Ry/Ang. + from ase.units import create_units + Ry = create_units("2006")["Ry"] + for i in range(batch_size): + assert ensemble.energies[i] == pytest.approx(ref[i][0] / Ry, rel=1e-8) + assert np.allclose(ensemble.forces[i], ref[i][1] / Ry, atol=1e-8) + + # 4. A second ensemble reuses the trained model: zero ab-initio calls + np.random.seed(1) + ensemble2 = sscha.Ensemble.Ensemble(dyn, 300) + ensemble2.generate(4) + ensemble2.set_otf(ensemble.flare_calc, std_tolerance_factor=100, + max_atoms_added=-1, update_style="add_n", + update_threshold=None, train_hyps=(1, np.inf), + output_name=str(tmp_path / "otf_run2")) + training_before = len(ensemble2.gp_model.training_data) + ensemble2.compute_ensemble(calc, compute_stress=True, cluster=cluster) + assert all(ensemble2.force_computed) + assert len(ensemble2.gp_model.training_data) == training_before # nothing new learned diff --git a/tests/test_otf/test_cluster_interface.py b/tests/test_otf/test_cluster_interface.py new file mode 100644 index 00000000..74bfbe43 --- /dev/null +++ b/tests/test_otf/test_cluster_interface.py @@ -0,0 +1,152 @@ +"""Tests for the layered Cluster/calculator interface (no flare needed).""" +import os +import threading +import subprocess + +import numpy as np +import pytest + +import cellconstructor as CC +from cellconstructor.calculators import Espresso +from ase.build import bulk +from ase.calculators.emt import EMT + +import sscha, sscha.Ensemble +import sscha.Cluster, sscha.LocalCluster, sscha.BaseCluster +import sscha.ClusterCalculators as cluster_calcs + +from .test_otf_base import get_gold_dyn # dyn fixture builder (no flare use) + + +def get_gold_structure(): + struct = CC.Structure.Structure() + struct.generate_from_ase_atoms(bulk("Au", "fcc", a=4.0782, cubic=False)) + return struct + + +class MockDirectCalculator(cluster_calcs.DirectCalculator): + """Deterministic fake: energy = index, zero forces/stress, echo structure.""" + def copy(self): + return MockDirectCalculator() + def compute(self, structure): + nat = structure.N_atoms + return {"energy": 1.0, "forces": np.zeros((nat, 3)), + "stress": np.zeros(6), "structure": structure.copy(), + "mock_extra": 42} + + +def test_directcluster_driver_no_otf(): + """The generic driver fills the ensemble through a DirectCalculator.""" + np.random.seed(0) + ensemble = sscha.Ensemble.Ensemble(get_gold_dyn(), 0) + ensemble.generate(4) + cluster = sscha.BaseCluster.DirectCluster(batch_size=2, job_number=1) + ensemble.compute_ensemble(MockDirectCalculator(), compute_stress=True, + cluster=cluster) + assert all(ensemble.force_computed) + assert all(ensemble.stress_computed) + assert np.allclose(ensemble.energies, 1.0 / 13.605698066) # eV -> Ry + assert all(p["mock_extra"] == 42 for p in ensemble.all_properties) + + +def test_pairing_rules(): + """DirectCluster rejects file calculators; Cluster rejects direct ones.""" + cluster = sscha.BaseCluster.DirectCluster() + with pytest.raises(TypeError, match="cannot be used"): + cluster._check_calculator_interface(Espresso( + input_data={"control": {}, "system": {}}, pseudopotentials={"Au": "Au.upf"})) + + remote = sscha.Cluster.Cluster(hostname="localhost") + with pytest.raises(TypeError, match="cannot be used"): + remote._check_calculator_interface(MockDirectCalculator()) + + # correct pairings pass + cluster._check_calculator_interface(MockDirectCalculator()) + remote._check_calculator_interface(Espresso( + input_data={"control": {}, "system": {}}, pseudopotentials={"Au": "Au.upf"})) + + +def test_espresso_prepare_input_file_backward_compatible(tmp_path): + """Phase 3.2 regression: QE calculators must behave exactly as before.""" + cluster = sscha.LocalCluster.LocalCluster("localhost") + cluster.local_workdir = str(tmp_path) + "/" + cluster.lock = threading.Lock() + calc = Espresso( + input_data={ + "control": {"tprnfor": True, "tstress": True}, + "system": {"ecutwfc": 30, "ecutrho": 240, "occupations": "fixed"}, + "electrons": {"conv_thr": 1e-8}, + }, + pseudopotentials={"Au": "Au.upf"}, + kpts=(1, 1, 1), + ) + inputs, outputs = cluster.prepare_input_file([get_gold_structure()], calc, ["ESP_0"]) + assert inputs == ["ESP_0.pwi"] + assert outputs == ["ESP_0.pwo"] + assert os.path.exists(tmp_path / "ESP_0.pwi") + + +def test_espresso_calculator_adapter(): + """The espresso-specific calculator keeps class through copy().""" + calc = cluster_calcs.EspressoCalculator( + input_data={"control": {}, "system": {}}, pseudopotentials={"Au": "Au.upf"}) + assert calc.input_extension == ".pwi" + assert calc.output_extension == ".pwo" + assert "PREFIX.pwi" in calc.get_execution_command() + assert isinstance(calc.copy(), cluster_calcs.EspressoCalculator) + + +def test_ase_file_calculator_prepare_input(tmp_path): + """The ASE file bridge uses its own extensions and ships calculator.pkl.""" + cluster = sscha.LocalCluster.LocalCluster("localhost") + cluster.local_workdir = str(tmp_path) + "/" + cluster.lock = threading.Lock() + calc = cluster_calcs.ASEFileCalculator(EMT()) + inputs, outputs = cluster.prepare_input_file([get_gold_structure()], calc, ["ESP_0"]) + assert inputs == ["ESP_0_input.json", "calculator.pkl"] + assert outputs == ["ESP_0.json"] + assert os.path.exists(tmp_path / "ESP_0_input.json") + assert os.path.exists(tmp_path / "calculator.pkl") + + +def test_ase_file_calculator_local_roundtrip(tmp_path): + """Run locally exactly what the cluster would run; compare with bare EMT.""" + calc = cluster_calcs.ASEFileCalculator(EMT()) + calc.set_directory(str(tmp_path)) + calc.set_label("ESP_0") + struct = get_gold_structure() + calc.write_input(struct) + + # Execute the runner exactly as Cluster.get_execution_command would + cmd = calc.command.replace("PREFIX", os.path.join(str(tmp_path), "ESP_0")) + subprocess.run(cmd, shell=True, check=True, cwd=str(tmp_path)) + + calc.read_results() + + atoms = struct.get_ase_atoms() + atoms.calc = EMT() + assert calc.results["energy"] == pytest.approx(atoms.get_potential_energy(), rel=1e-10) + assert np.allclose(calc.results["forces"], atoms.get_forces(), atol=1e-10) + assert np.allclose(calc.results["stress"], atoms.get_stress(), atol=1e-10) + assert np.allclose(calc.structure.coords, struct.coords, atol=1e-10) + + +def test_ase_direct_calculator_matches_bare_ase(): + calc = cluster_calcs.ASEDirectCalculator(EMT()) + struct = get_gold_structure() + res = calc.compute(struct) + atoms = struct.get_ase_atoms() + atoms.calc = EMT() + assert res["energy"] == pytest.approx(atoms.get_potential_energy(), rel=1e-12) + assert np.allclose(res["forces"], atoms.get_forces(), atol=1e-12) + assert np.allclose(res["stress"], atoms.get_stress(), atol=1e-12) + # per-thread copies are independent + assert calc.copy().ase_calc is not calc.ase_calc + + +def test_nonpicklable_calculator_fails_early(): + class Unpicklable: + def __getstate__(self): + raise TypeError("no pickle") + with pytest.raises(TypeError): + cluster_calcs.ASEFileCalculator(Unpicklable()) diff --git a/tests/test_otf/test_otf_base.py b/tests/test_otf/test_otf_base.py new file mode 100644 index 00000000..74036b48 --- /dev/null +++ b/tests/test_otf/test_otf_base.py @@ -0,0 +1,124 @@ +"""Tests for the backend-agnostic OTF helpers on the base Ensemble class. + +.. note:: The ensemble is built inside each test (not in a pytest fixture) + because the flare C extension segfaults when SGP objects are created in + a fixture and later used in the test body (a known flare/pytest + interaction bug). Creating the objects inside the test function works + reliably. +""" +import os +import numpy as np +import pytest + +flare = pytest.importorskip("flare") # G9: skip the whole module without flare + +from flare.bffs.sgp.calculator import SGP_Calculator + +import cellconstructor as CC, cellconstructor.Phonons +import sscha, sscha.Ensemble +from ase.calculators.emt import EMT +from ase.build import bulk + + +def get_gold_dyn(supercell=(2, 2, 2)): + """Harmonic Au (EMT), 8 atoms with the default supercell.""" + struct = CC.Structure.Structure() + struct.generate_from_ase_atoms(bulk("Au", "fcc", a=4.0782, cubic=False)) + dyn = CC.Phonons.compute_phonons_finite_displacements(struct, EMT(), supercell=supercell) + dyn.Symmetrize() + dyn.ForcePositiveDefinite() + return dyn + + +# Keep kernel objects alive for the lifetime of the process: SGP_Wrapper does +# not store a Python reference to the kernels (only to the descriptor +# calculators), so without this the pybind11 wrapper is garbage-collected when +# the helper returns, leaving the C++ SparseGP with a dangling pointer → +# segfault in add_training_structure. +_KERNEL_REFS = [] + + +def get_sgp_calc_au(): + """Empty SGP calculator for gold.""" + from flare.bffs.sgp._C_flare import NormalizedDotProduct, B2 + from flare.bffs.sgp import SGP_Wrapper + cutoff = 4.0 + kernel = NormalizedDotProduct(2.0, 2) + b2 = B2("chebyshev", "quadratic", [0.0, cutoff], [], [1, 4, 3], cutoff * np.ones((1, 1))) + sgp = SGP_Wrapper([kernel], [b2], cutoff, 0.1, 0.1, 0.1, {79: 0}, + single_atom_energies={0: 0.0}, variance_type="local", + opt_method="L-BFGS-B", max_iterations=5) + _KERNEL_REFS.append(kernel) # flare bug workaround + return SGP_Calculator(sgp) + + +def _make_ensemble(tmp_path, n_configs=8): + """Build an OTF ensemble for gold+EMT (call inside the test, not a fixture).""" + np.random.seed(0) + ens = sscha.Ensemble.Ensemble(get_gold_dyn(), 300) + ens.generate(n_configs) + ens.set_otf(get_sgp_calc_au(), std_tolerance_factor=100, + max_atoms_added=-1, update_style="add_n", update_threshold=None, + train_hyps=(1, np.inf), output_name=str(tmp_path / "otf_run")) + return ens + + +def test_split_propagates_otf_state(tmp_path): + """Phase 2 regression: get_noncomputed() must carry the OTF state.""" + ensemble = _make_ensemble(tmp_path) + sub = ensemble.get_noncomputed() + assert sub.gp_model is ensemble.gp_model + assert sub.flare_calc is ensemble.flare_calc + assert sub.std_tolerance == ensemble.std_tolerance + assert sub.max_atoms_added == ensemble.max_atoms_added + assert sub.update_style == ensemble.update_style + assert sub.train_hyps == ensemble.train_hyps + assert sub.output is ensemble.output + assert sub.flare_name == ensemble.flare_name + + +def test_otf_setup_defaults(tmp_path): + ensemble = _make_ensemble(tmp_path) + ensemble._otf_setup_defaults(8) + assert ensemble.max_atoms_added == 8 + assert ensemble.init_atoms == list(range(8)) + + +def test_update_and_predict_cycle(tmp_path): + """update GP with one EMT structure -> train -> predict the rest.""" + ensemble = _make_ensemble(tmp_path) + ensemble._otf_setup_defaults(8) + + # empty model: nothing is predicted + remaining = list(range(ensemble.N)) + ensemble._otf_predict(ensemble.structures, remaining) + assert remaining == list(range(ensemble.N)) + + # update with one EMT reference (empty model -> all init_atoms added) + struct = ensemble.structures[0] + atoms = struct.get_ase_atoms() + atoms.calc = EMT() + ensemble._otf_update_from_structure( + struct, atoms.get_potential_energy(), atoms.get_forces(), + atoms.get_stress()) + assert len(ensemble.gp_model.training_data) >= 1 + + ensemble._otf_maybe_train_and_write(dft_counts=1) + assert os.path.exists(ensemble.flare_name) + + # trained model with huge tolerance: everything predicted + remaining = list(range(ensemble.N)) + ensemble._otf_predict(ensemble.structures, remaining) + assert remaining == [] + assert all(ensemble.force_computed) + assert np.all(np.isfinite(ensemble.energies)) + assert np.all(np.isfinite(ensemble.forces)) + + +def test_clean_runs_generic_label(tmp_path, capsys): + ensemble = _make_ensemble(tmp_path) + ensemble.force_computed[:] = True + ensemble._clean_runs(dft_counts=2) + out, _ = capsys.readouterr() + assert "SUMMARY CALCULATIONS" in out + assert "Steps using OTF-ML model :" in out From 677e3d8659afe9ae7eeeb2309539e5b086ea046d Mon Sep 17 00:00:00 2001 From: Lorenzo Monacelli Date: Thu, 23 Jul 2026 00:24:25 +0200 Subject: [PATCH 6/6] Examples: on-the-fly FLARE OTF on gold+EMT (DirectCluster + LocalCluster) Two end-to-end example scripts from the plan (section 9): - run_gold_otf_directcluster.py: the simplest possible OTF run, in-process ASE calculator (no files, no scheduler, no HPC). - run_gold_otf_localcluster.py: exercises the full scheduler/file machinery (input files, tar, submission script, result retrieval) mocked locally with bash + shutil (no ssh/scp, no SLURM). Both use a FLARE SGP that learns EMT on the fly. The learning cycle is 2 structures (batch_size=2, job_number=1); after the first cycle trains the GP, subsequent structures are predicted by the model. --- .../run_gold_otf_directcluster.py | 100 ++++++++++++++++ .../run_gold_otf_localcluster.py | 110 ++++++++++++++++++ 2 files changed, 210 insertions(+) create mode 100644 Examples/OnTheFlyClusterGold/run_gold_otf_directcluster.py create mode 100644 Examples/OnTheFlyClusterGold/run_gold_otf_localcluster.py diff --git a/Examples/OnTheFlyClusterGold/run_gold_otf_directcluster.py b/Examples/OnTheFlyClusterGold/run_gold_otf_directcluster.py new file mode 100644 index 00000000..5f15b4f7 --- /dev/null +++ b/Examples/OnTheFlyClusterGold/run_gold_otf_directcluster.py @@ -0,0 +1,100 @@ +"""On-the-fly FLARE learning on gold (EMT) through DirectCluster. + +Recipe A from the on-the-fly FLARE cluster plan: the simplest possible +OTF run, exercising the new generic cluster interface with an in-process +ASE calculator (no files, no scheduler, no HPC, no AiiDA). + + python run_gold_otf_directcluster.py + +Requires: sscha (editable), flare, ase. No pw.x, no SLURM, no files. +""" +import os +import numpy as np + +import cellconstructor as CC, cellconstructor.Phonons +import sscha, sscha.Ensemble, sscha.SchaMinimizer, sscha.Relax +import sscha.BaseCluster +import sscha.ClusterCalculators as cluster_calcs + +from ase.calculators.emt import EMT + + +# Keep kernel alive (flare bug: SGP_Wrapper does not keep a Python reference +# to the kernel, see tests/test_otf/test_otf_base.py for details). +_KERNEL_REFS = [] + + +def get_sgp_calc_au(): + """Return an empty FLARE SGP calculator for gold (single species).""" + from flare.bffs.sgp._C_flare import NormalizedDotProduct, B2 + from flare.bffs.sgp import SGP_Wrapper + from flare.bffs.sgp.calculator import SGP_Calculator + + cutoff = 4.0 # includes the 12 first neighbors of FCC Au (a=4.08 A) + kernel = NormalizedDotProduct(sigma=2.0, power=2) + b2 = B2("chebyshev", "quadratic", [0.0, cutoff], [], + [1, 4, 3], # [n_species, nmax, lmax] + cutoff * np.ones((1, 1))) # cutoff_matrix + sgp = SGP_Wrapper( + [kernel], [b2], cutoff, + sigma_e=0.1, sigma_f=0.1, sigma_s=0.1, + species_map={79: 0}, # Au + single_atom_energies={0: 0.0}, + variance_type="local", + energy_training=True, force_training=True, stress_training=True, + opt_method="L-BFGS-B", max_iterations=5, + ) + _KERNEL_REFS.append(kernel) + return SGP_Calculator(sgp) + + +def get_gold_dyn(supercell=(2, 2, 2)): + """Harmonic Au (EMT), 8 atoms with the default supercell.""" + from ase.build import bulk + struct = CC.Structure.Structure() + struct.generate_from_ase_atoms(bulk("Au", "fcc", a=4.0782, cubic=False)) + dyn = CC.Phonons.compute_phonons_finite_displacements(struct, EMT(), supercell=supercell) + dyn.Symmetrize() + dyn.ForcePositiveDefinite() + return dyn + + +def get_otf_ensemble(dyn, temperature=300, output_name="otf_gold"): + np.random.seed(0) + ensemble = sscha.Ensemble.Ensemble(dyn, temperature) + ensemble.set_otf( + get_sgp_calc_au(), + std_tolerance_factor=100, # large: after the first training batch, + # predictions are (almost) always accepted + max_atoms_added=-1, # all atoms + update_style="add_n", + update_threshold=None, + train_hyps=(1, np.inf), + output_name=output_name, + ) + return ensemble + + +def main_direct(): + dyn = get_gold_dyn() + ensemble = get_otf_ensemble(dyn) + + calc = cluster_calcs.ASEDirectCalculator(EMT()) + cluster = sscha.BaseCluster.DirectCluster(batch_size=2, job_number=1) + # learning cycle = batch_size * job_number = 2 structures per cycle + + minim = sscha.SchaMinimizer.SSCHA_Minimizer(ensemble) + minim.set_minimization_step(0.01) + minim.meaningful_factor = 1e-3 + minim.kong_liu_ratio = 0.5 + + relax = sscha.Relax.SSCHA(minimizer=minim, ase_calculator=calc, + cluster=cluster, N_configs=8, max_pop=3, + save_ensemble=False) + relax.relax(get_stress=True) + relax.minim.finalize() + relax.minim.dyn.save_qe("sscha_gold_otf_dyn") + + +if __name__ == "__main__": + main_direct() diff --git a/Examples/OnTheFlyClusterGold/run_gold_otf_localcluster.py b/Examples/OnTheFlyClusterGold/run_gold_otf_localcluster.py new file mode 100644 index 00000000..c7f4dc0d --- /dev/null +++ b/Examples/OnTheFlyClusterGold/run_gold_otf_localcluster.py @@ -0,0 +1,110 @@ +"""On-the-fly FLARE learning on gold (EMT) through LocalCluster. + +Recipe B from the on-the-fly FLARE cluster plan: exercises the *exact* +remote-cluster code path (input files, tar, submission script, result +retrieval) mocked locally with bash + shutil (no ssh/scp, no SLURM). + + python run_gold_otf_localcluster.py + +Requires: sscha (editable), flare, ase. No pw.x, no AiiDA, no SLURM. +""" +import os +import numpy as np + +import cellconstructor as CC, cellconstructor.Phonons +import sscha, sscha.Ensemble, sscha.SchaMinimizer, sscha.Relax +import sscha.LocalCluster +import sscha.ClusterCalculators as cluster_calcs + +from ase.calculators.emt import EMT + + +# Keep kernel alive (flare bug, see tests/test_otf/test_otf_base.py). +_KERNEL_REFS = [] + + +def get_sgp_calc_au(): + """Return an empty FLARE SGP calculator for gold (single species).""" + from flare.bffs.sgp._C_flare import NormalizedDotProduct, B2 + from flare.bffs.sgp import SGP_Wrapper + from flare.bffs.sgp.calculator import SGP_Calculator + + cutoff = 4.0 + kernel = NormalizedDotProduct(sigma=2.0, power=2) + b2 = B2("chebyshev", "quadratic", [0.0, cutoff], [], + [1, 4, 3], cutoff * np.ones((1, 1))) + sgp = SGP_Wrapper( + [kernel], [b2], cutoff, + sigma_e=0.1, sigma_f=0.1, sigma_s=0.1, + species_map={79: 0}, + single_atom_energies={0: 0.0}, + variance_type="local", + energy_training=True, force_training=True, stress_training=True, + opt_method="L-BFGS-B", max_iterations=5, + ) + _KERNEL_REFS.append(kernel) + return SGP_Calculator(sgp) + + +def get_gold_dyn(supercell=(2, 2, 2)): + """Harmonic Au (EMT), 8 atoms with the default supercell.""" + from ase.build import bulk + struct = CC.Structure.Structure() + struct.generate_from_ase_atoms(bulk("Au", "fcc", a=4.0782, cubic=False)) + dyn = CC.Phonons.compute_phonons_finite_displacements(struct, EMT(), supercell=supercell) + dyn.Symmetrize() + dyn.ForcePositiveDefinite() + return dyn + + +def get_otf_ensemble(dyn, temperature=300, output_name="otf_gold"): + np.random.seed(0) + ensemble = sscha.Ensemble.Ensemble(dyn, temperature) + ensemble.set_otf( + get_sgp_calc_au(), + std_tolerance_factor=100, + max_atoms_added=-1, + update_style="add_n", + update_threshold=None, + train_hyps=(1, np.inf), + output_name=output_name, + ) + return ensemble + + +def main_localcluster(): + dyn = get_gold_dyn() + ensemble = get_otf_ensemble(dyn) + + calc = cluster_calcs.ASEFileCalculator(EMT()) + + cluster = sscha.LocalCluster.LocalCluster("localhost") + cluster.workdir = os.path.abspath("otf_cluster/remote") + cluster.local_workdir = os.path.abspath("otf_cluster/local") + "/" + cluster.submit_command = "bash" # run the script directly (blocking) + cluster.nonblocking_command = False # do NOT poll squeue (G8) + cluster.use_nodes = False + cluster.use_cpu = False + cluster.use_time = False + cluster.use_account = False + cluster.job_number = 1 # 1 structure per job + cluster.batch_size = 2 # 2 ab-initio structures per learning cycle + cluster.binary = calc.command # python -m sscha.ASEClusterRunner ... + cluster.mpi_cmd = "" # the runner is a serial process + cluster.setup_workdir() # mkdir -p workdir (locally) + + minim = sscha.SchaMinimizer.SSCHA_Minimizer(ensemble) + minim.set_minimization_step(0.01) + minim.meaningful_factor = 1e-3 + minim.kong_liu_ratio = 0.5 + + relax = sscha.Relax.SSCHA(minimizer=minim, ase_calculator=calc, + cluster=cluster, N_configs=8, max_pop=3, + save_ensemble=False) + relax.relax(get_stress=True) + relax.minim.finalize() + relax.minim.dyn.save_qe("sscha_gold_otf_dyn") + + +if __name__ == "__main__": + main_localcluster()