On-the-fly FLARE learning for all cluster backends

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

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

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