Tests: OTF base helpers, cluster interface, end-to-end EMT OTF

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).

Author: mesonepigreco

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

tests/test_otf/__init__.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

tests/test_otf/test_cluster_emt_otf.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())

tests/test_otf/test_cluster_interface.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