Skip to content

Commit 677e3d8

Browse files
mesonepigrecomesonepigreco
committed
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.
1 parent 1760f85 commit 677e3d8

2 files changed

Lines changed: 210 additions & 0 deletions

File tree

‎Examples/OnTheFlyClusterGold/run_gold_otf_directcluster.py‎

Lines changed: 100 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,100 @@
1+
"""On-the-fly FLARE learning on gold (EMT) through DirectCluster.
2+
3+
Recipe A from the on-the-fly FLARE cluster plan: the simplest possible
4+
OTF run, exercising the new generic cluster interface with an in-process
5+
ASE calculator (no files, no scheduler, no HPC, no AiiDA).
6+
7+
python run_gold_otf_directcluster.py
8+
9+
Requires: sscha (editable), flare, ase. No pw.x, no SLURM, no files.
10+
"""
11+
import os
12+
import numpy as np
13+
14+
import cellconstructor as CC, cellconstructor.Phonons
15+
import sscha, sscha.Ensemble, sscha.SchaMinimizer, sscha.Relax
16+
import sscha.BaseCluster
17+
import sscha.ClusterCalculators as cluster_calcs
18+
19+
from ase.calculators.emt import EMT
20+
21+
22+
# Keep kernel alive (flare bug: SGP_Wrapper does not keep a Python reference
23+
# to the kernel, see tests/test_otf/test_otf_base.py for details).
24+
_KERNEL_REFS = []
25+
26+
27+
def get_sgp_calc_au():
28+
"""Return an empty FLARE SGP calculator for gold (single species)."""
29+
from flare.bffs.sgp._C_flare import NormalizedDotProduct, B2
30+
from flare.bffs.sgp import SGP_Wrapper
31+
from flare.bffs.sgp.calculator import SGP_Calculator
32+
33+
cutoff = 4.0 # includes the 12 first neighbors of FCC Au (a=4.08 A)
34+
kernel = NormalizedDotProduct(sigma=2.0, power=2)
35+
b2 = B2("chebyshev", "quadratic", [0.0, cutoff], [],
36+
[1, 4, 3], # [n_species, nmax, lmax]
37+
cutoff * np.ones((1, 1))) # cutoff_matrix
38+
sgp = SGP_Wrapper(
39+
[kernel], [b2], cutoff,
40+
sigma_e=0.1, sigma_f=0.1, sigma_s=0.1,
41+
species_map={79: 0}, # Au
42+
single_atom_energies={0: 0.0},
43+
variance_type="local",
44+
energy_training=True, force_training=True, stress_training=True,
45+
opt_method="L-BFGS-B", max_iterations=5,
46+
)
47+
_KERNEL_REFS.append(kernel)
48+
return SGP_Calculator(sgp)
49+
50+
51+
def get_gold_dyn(supercell=(2, 2, 2)):
52+
"""Harmonic Au (EMT), 8 atoms with the default supercell."""
53+
from ase.build import bulk
54+
struct = CC.Structure.Structure()
55+
struct.generate_from_ase_atoms(bulk("Au", "fcc", a=4.0782, cubic=False))
56+
dyn = CC.Phonons.compute_phonons_finite_displacements(struct, EMT(), supercell=supercell)
57+
dyn.Symmetrize()
58+
dyn.ForcePositiveDefinite()
59+
return dyn
60+
61+
62+
def get_otf_ensemble(dyn, temperature=300, output_name="otf_gold"):
63+
np.random.seed(0)
64+
ensemble = sscha.Ensemble.Ensemble(dyn, temperature)
65+
ensemble.set_otf(
66+
get_sgp_calc_au(),
67+
std_tolerance_factor=100, # large: after the first training batch,
68+
# predictions are (almost) always accepted
69+
max_atoms_added=-1, # all atoms
70+
update_style="add_n",
71+
update_threshold=None,
72+
train_hyps=(1, np.inf),
73+
output_name=output_name,
74+
)
75+
return ensemble
76+
77+
78+
def main_direct():
79+
dyn = get_gold_dyn()
80+
ensemble = get_otf_ensemble(dyn)
81+
82+
calc = cluster_calcs.ASEDirectCalculator(EMT())
83+
cluster = sscha.BaseCluster.DirectCluster(batch_size=2, job_number=1)
84+
# learning cycle = batch_size * job_number = 2 structures per cycle
85+
86+
minim = sscha.SchaMinimizer.SSCHA_Minimizer(ensemble)
87+
minim.set_minimization_step(0.01)
88+
minim.meaningful_factor = 1e-3
89+
minim.kong_liu_ratio = 0.5
90+
91+
relax = sscha.Relax.SSCHA(minimizer=minim, ase_calculator=calc,
92+
cluster=cluster, N_configs=8, max_pop=3,
93+
save_ensemble=False)
94+
relax.relax(get_stress=True)
95+
relax.minim.finalize()
96+
relax.minim.dyn.save_qe("sscha_gold_otf_dyn")
97+
98+
99+
if __name__ == "__main__":
100+
main_direct()

‎Examples/OnTheFlyClusterGold/run_gold_otf_localcluster.py‎

Lines changed: 110 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,110 @@
1+
"""On-the-fly FLARE learning on gold (EMT) through LocalCluster.
2+
3+
Recipe B from the on-the-fly FLARE cluster plan: exercises the *exact*
4+
remote-cluster code path (input files, tar, submission script, result
5+
retrieval) mocked locally with bash + shutil (no ssh/scp, no SLURM).
6+
7+
python run_gold_otf_localcluster.py
8+
9+
Requires: sscha (editable), flare, ase. No pw.x, no AiiDA, no SLURM.
10+
"""
11+
import os
12+
import numpy as np
13+
14+
import cellconstructor as CC, cellconstructor.Phonons
15+
import sscha, sscha.Ensemble, sscha.SchaMinimizer, sscha.Relax
16+
import sscha.LocalCluster
17+
import sscha.ClusterCalculators as cluster_calcs
18+
19+
from ase.calculators.emt import EMT
20+
21+
22+
# Keep kernel alive (flare bug, see tests/test_otf/test_otf_base.py).
23+
_KERNEL_REFS = []
24+
25+
26+
def get_sgp_calc_au():
27+
"""Return an empty FLARE SGP calculator for gold (single species)."""
28+
from flare.bffs.sgp._C_flare import NormalizedDotProduct, B2
29+
from flare.bffs.sgp import SGP_Wrapper
30+
from flare.bffs.sgp.calculator import SGP_Calculator
31+
32+
cutoff = 4.0
33+
kernel = NormalizedDotProduct(sigma=2.0, power=2)
34+
b2 = B2("chebyshev", "quadratic", [0.0, cutoff], [],
35+
[1, 4, 3], cutoff * np.ones((1, 1)))
36+
sgp = SGP_Wrapper(
37+
[kernel], [b2], cutoff,
38+
sigma_e=0.1, sigma_f=0.1, sigma_s=0.1,
39+
species_map={79: 0},
40+
single_atom_energies={0: 0.0},
41+
variance_type="local",
42+
energy_training=True, force_training=True, stress_training=True,
43+
opt_method="L-BFGS-B", max_iterations=5,
44+
)
45+
_KERNEL_REFS.append(kernel)
46+
return SGP_Calculator(sgp)
47+
48+
49+
def get_gold_dyn(supercell=(2, 2, 2)):
50+
"""Harmonic Au (EMT), 8 atoms with the default supercell."""
51+
from ase.build import bulk
52+
struct = CC.Structure.Structure()
53+
struct.generate_from_ase_atoms(bulk("Au", "fcc", a=4.0782, cubic=False))
54+
dyn = CC.Phonons.compute_phonons_finite_displacements(struct, EMT(), supercell=supercell)
55+
dyn.Symmetrize()
56+
dyn.ForcePositiveDefinite()
57+
return dyn
58+
59+
60+
def get_otf_ensemble(dyn, temperature=300, output_name="otf_gold"):
61+
np.random.seed(0)
62+
ensemble = sscha.Ensemble.Ensemble(dyn, temperature)
63+
ensemble.set_otf(
64+
get_sgp_calc_au(),
65+
std_tolerance_factor=100,
66+
max_atoms_added=-1,
67+
update_style="add_n",
68+
update_threshold=None,
69+
train_hyps=(1, np.inf),
70+
output_name=output_name,
71+
)
72+
return ensemble
73+
74+
75+
def main_localcluster():
76+
dyn = get_gold_dyn()
77+
ensemble = get_otf_ensemble(dyn)
78+
79+
calc = cluster_calcs.ASEFileCalculator(EMT())
80+
81+
cluster = sscha.LocalCluster.LocalCluster("localhost")
82+
cluster.workdir = os.path.abspath("otf_cluster/remote")
83+
cluster.local_workdir = os.path.abspath("otf_cluster/local") + "/"
84+
cluster.submit_command = "bash" # run the script directly (blocking)
85+
cluster.nonblocking_command = False # do NOT poll squeue (G8)
86+
cluster.use_nodes = False
87+
cluster.use_cpu = False
88+
cluster.use_time = False
89+
cluster.use_account = False
90+
cluster.job_number = 1 # 1 structure per job
91+
cluster.batch_size = 2 # 2 ab-initio structures per learning cycle
92+
cluster.binary = calc.command # python -m sscha.ASEClusterRunner ...
93+
cluster.mpi_cmd = "" # the runner is a serial process
94+
cluster.setup_workdir() # mkdir -p workdir (locally)
95+
96+
minim = sscha.SchaMinimizer.SSCHA_Minimizer(ensemble)
97+
minim.set_minimization_step(0.01)
98+
minim.meaningful_factor = 1e-3
99+
minim.kong_liu_ratio = 0.5
100+
101+
relax = sscha.Relax.SSCHA(minimizer=minim, ase_calculator=calc,
102+
cluster=cluster, N_configs=8, max_pop=3,
103+
save_ensemble=False)
104+
relax.relax(get_stress=True)
105+
relax.minim.finalize()
106+
relax.minim.dyn.save_qe("sscha_gold_otf_dyn")
107+
108+
109+
if __name__ == "__main__":
110+
main_localcluster()

0 commit comments

Comments
 (0)