|
| 1 | +# -*- coding: utf-8 -*- |
| 2 | +"""Calculators for the sscha cluster backends (Axis B of the cluster design). |
| 3 | +
|
| 4 | +Two families: |
| 5 | +
|
| 6 | +* FileIOCalculator: calculators driven through input/output files, consumed |
| 7 | + by the scheduler-based clusters (sscha.Cluster.Cluster, |
| 8 | + sscha.LocalCluster.LocalCluster). One subclass per simulation code |
| 9 | + (EspressoCalculator today; VaspCalculator, AbinitCalculator, |
| 10 | + Cp2kCalculator, GaussianCalculator, ... in the future). |
| 11 | +
|
| 12 | +* DirectCalculator: calculators computed in-process, consumed by |
| 13 | + sscha.BaseCluster.DirectCluster. |
| 14 | +
|
| 15 | +For any code that has an ASE interface, ASEFileCalculator (file-based) and |
| 16 | +ASEDirectCalculator (in-process) already provide full support with no |
| 17 | +dedicated subclass: e.g. ASEFileCalculator(ase.calculators.vasp.Vasp(...)). |
| 18 | +""" |
| 19 | +import copy |
| 20 | +import os |
| 21 | +import pickle |
| 22 | +import sys |
| 23 | + |
| 24 | +import cellconstructor as CC |
| 25 | +import cellconstructor.calculators |
| 26 | + |
| 27 | + |
| 28 | +# ---------------------------------------------------------------------- # |
| 29 | +# Generic contract # |
| 30 | +# ---------------------------------------------------------------------- # |
| 31 | + |
| 32 | +class ClusterCalculator(object): |
| 33 | + """Documentation-level base class for cluster calculators. |
| 34 | +
|
| 35 | + Any calculator consumed by a sscha cluster must provide: |
| 36 | +
|
| 37 | + copy() -> an independent instance (one per worker thread) |
| 38 | +
|
| 39 | + and produce result dicts with the contract documented in |
| 40 | + BaseCluster.compute_jobarray ("energy" [eV], "forces" [eV/Ang], |
| 41 | + "stress" [ASE Voigt, -eV/Ang^3], "structure" [CC.Structure], extras). |
| 42 | + """ |
| 43 | + |
| 44 | + |
| 45 | +# ---------------------------------------------------------------------- # |
| 46 | +# File-based family (scheduler clusters) # |
| 47 | +# ---------------------------------------------------------------------- # |
| 48 | + |
| 49 | +class FileIOCalculator(CC.calculators.FileIOCalculator, ClusterCalculator): |
| 50 | + """Mid-level interface for file-based calculators. |
| 51 | +
|
| 52 | + Class attributes consumed by the scheduler machinery: |
| 53 | +
|
| 54 | + input_extension / output_extension : str |
| 55 | + Extensions of the per-label input/output files. |
| 56 | + extra_input_files : list of str |
| 57 | + Additional files (relative to the local workdir) to ship to the |
| 58 | + cluster together with the inputs (e.g. "calculator.pkl"). |
| 59 | + """ |
| 60 | + |
| 61 | + input_extension = ".pwi" |
| 62 | + output_extension = ".pwo" |
| 63 | + extra_input_files = [] |
| 64 | + |
| 65 | + def get_execution_command(self): |
| 66 | + """The command template run on the cluster, with the PREFIX |
| 67 | + placeholder where the calculation label must be inserted, e.g.: |
| 68 | +
|
| 69 | + "pw.x -npool 4 -i PREFIX.pwi > PREFIX.pwo" |
| 70 | +
|
| 71 | + Assign it to `cluster.binary` (or use it in the namelist). |
| 72 | + """ |
| 73 | + raise NotImplementedError |
| 74 | + |
| 75 | + |
| 76 | +class EspressoCalculator(CC.calculators.Espresso, FileIOCalculator): |
| 77 | + """Quantum ESPRESSO (pw.x) calculator for the sscha clusters. |
| 78 | +
|
| 79 | + Thin formalization of cellconstructor.calculators.Espresso, which |
| 80 | + already satisfies the interface. Using this class (instead of the CC |
| 81 | + one) only makes the extensions and the execution command explicit. |
| 82 | + """ |
| 83 | + |
| 84 | + input_extension = ".pwi" |
| 85 | + output_extension = ".pwo" |
| 86 | + extra_input_files = [] |
| 87 | + |
| 88 | + def get_execution_command(self, n_pool=1): |
| 89 | + return "pw.x -npool {} -i PREFIX.pwi > PREFIX.pwo".format(n_pool) |
| 90 | + |
| 91 | + def copy(self): |
| 92 | + """Return an identical instance of THIS class (per-thread copies).""" |
| 93 | + return EspressoCalculator(self.input_data, self.pseudopotentials, |
| 94 | + self.masses, self.command, self.kpts, self.koffset) |
| 95 | + |
| 96 | + |
| 97 | +class ASEFileCalculator(FileIOCalculator): |
| 98 | + """ |
| 99 | + RUN ANY ASE CALCULATOR THROUGH A SCHEDULER-BASED SSCHA CLUSTER |
| 100 | + ============================================================== |
| 101 | +
|
| 102 | + Bridge between a plain ASE calculator (EMT, LJ, VASP, CP2K, ...) and the |
| 103 | + file-based `Cluster` submission: |
| 104 | +
|
| 105 | + - the ASE calculator is pickled ONCE into the working directory |
| 106 | + (`calculator.pkl`, listed in `extra_input_files` so it is shipped to |
| 107 | + the cluster together with the inputs); |
| 108 | + - each input file is the ASE Atoms serialized with ase.io.jsonio; |
| 109 | + - the execution command runs `python -m sscha.ASEClusterRunner`, which |
| 110 | + loads both, computes energy/forces/stress and writes a JSON result |
| 111 | + file (see Modules/ASEClusterRunner.py). |
| 112 | +
|
| 113 | + Usage: |
| 114 | + calc = ASEFileCalculator(EMT()) |
| 115 | + cluster.binary = calc.command # the runner command with PREFIX placeholder |
| 116 | + cluster.mpi_cmd = "" # the runner is a serial python process |
| 117 | + cluster.compute_ensemble(ensemble, calc) |
| 118 | +
|
| 119 | + Requirements on the (remote) cluster: python + ase (+ the actual |
| 120 | + calculator dependencies, and a python able to unpickle the calculator — |
| 121 | + same ASE version recommended). For sscha.LocalCluster this is |
| 122 | + automatically satisfied by the local environment. |
| 123 | + """ |
| 124 | + |
| 125 | + input_extension = "_input.json" |
| 126 | + output_extension = ".json" |
| 127 | + calc_pickle_name = "calculator.pkl" |
| 128 | + extra_input_files = [calc_pickle_name] |
| 129 | + |
| 130 | + def __init__(self, ase_calc, python_exe=None): |
| 131 | + """ |
| 132 | + Parameters |
| 133 | + ---------- |
| 134 | + ase_calc : ase.calculators.calculator.Calculator |
| 135 | + The ASE calculator to execute on the cluster. Must be picklable. |
| 136 | + python_exe : str, optional |
| 137 | + The python interpreter used ON THE CLUSTER to run the runner. |
| 138 | + Defaults to the local interpreter (correct for LocalCluster). |
| 139 | + """ |
| 140 | + super().__init__() |
| 141 | + |
| 142 | + # Fail early with a clear error if the calculator cannot be pickled |
| 143 | + pickle.dumps(ase_calc) |
| 144 | + |
| 145 | + self.ase_calc = ase_calc |
| 146 | + self.python_exe = python_exe or sys.executable |
| 147 | + self.command = self.get_execution_command() |
| 148 | + |
| 149 | + def get_execution_command(self): |
| 150 | + return ("{exe} -m sscha.ASEClusterRunner " |
| 151 | + "--calc {pkl} " |
| 152 | + "--input PREFIX{in_ext} " |
| 153 | + "--output PREFIX{out_ext}" |
| 154 | + ).format(exe=self.python_exe, pkl=self.calc_pickle_name, |
| 155 | + in_ext=self.input_extension, out_ext=self.output_extension) |
| 156 | + |
| 157 | + def copy(self): |
| 158 | + """Return an identical instance, without inheriting calculation info.""" |
| 159 | + return ASEFileCalculator(self.ase_calc, self.python_exe) |
| 160 | + |
| 161 | + def set_directory(self, directory): |
| 162 | + """Set the working directory and pickle the calculator there (once).""" |
| 163 | + CC.calculators.FileIOCalculator.set_directory(self, directory) |
| 164 | + pkl_path = os.path.join(directory, self.calc_pickle_name) |
| 165 | + if not os.path.exists(pkl_path): |
| 166 | + with open(pkl_path, "wb") as handle: |
| 167 | + pickle.dump(self.ase_calc, handle) |
| 168 | + |
| 169 | + def write_input(self, structure): |
| 170 | + """Serialize the structure as JSON in {directory}/{label}{input_extension}.""" |
| 171 | + import ase.io.jsonio |
| 172 | + CC.calculators.FileIOCalculator.write_input(self, structure) |
| 173 | + |
| 174 | + atoms = structure.get_ase_atoms() |
| 175 | + fname = os.path.join(self.directory, self.label + self.input_extension) |
| 176 | + with open(fname, "w") as handle: |
| 177 | + handle.write(ase.io.jsonio.encode(atoms)) |
| 178 | + |
| 179 | + def read_results(self): |
| 180 | + """Read {directory}/{label}{output_extension} into .results/.structure.""" |
| 181 | + import ase.io.jsonio |
| 182 | + |
| 183 | + fname = os.path.join(self.directory, self.label + self.output_extension) |
| 184 | + with open(fname, "r") as handle: |
| 185 | + payload = ase.io.jsonio.decode(handle.read()) |
| 186 | + |
| 187 | + self.results = { |
| 188 | + "energy": payload["energy"], # eV |
| 189 | + "forces": payload["forces"], # eV / Angstrom |
| 190 | + "stress": payload["stress"], # ASE Voigt (xx,yy,zz,yz,xz,xy), -eV/Angstrom^3 |
| 191 | + } |
| 192 | + self.structure = CC.Structure.Structure() |
| 193 | + self.structure.generate_from_ase_atoms(payload["atoms"]) |
| 194 | + |
| 195 | + |
| 196 | +# ---------------------------------------------------------------------- # |
| 197 | +# Direct (in-process) family (DirectCluster) # |
| 198 | +# ---------------------------------------------------------------------- # |
| 199 | + |
| 200 | +class DirectCalculator(ClusterCalculator): |
| 201 | + """Mid-level interface for in-process calculators (DirectCluster).""" |
| 202 | + |
| 203 | + def compute(self, structure): |
| 204 | + """Compute one structure and return the raw results dict. |
| 205 | +
|
| 206 | + Parameters |
| 207 | + ---------- |
| 208 | + structure : cellconstructor.Structure.Structure |
| 209 | +
|
| 210 | + Returns |
| 211 | + ------- |
| 212 | + dict with "energy" [eV], "forces" [eV/Ang], optionally |
| 213 | + "stress" [ASE Voigt, -eV/Ang^3] and "structure" [CC.Structure]. |
| 214 | + """ |
| 215 | + raise NotImplementedError |
| 216 | + |
| 217 | + |
| 218 | +class ASEDirectCalculator(DirectCalculator): |
| 219 | + """Wrap ANY ASE calculator for in-process use with DirectCluster. |
| 220 | +
|
| 221 | + No files are written: the ASE calculator is called directly on the |
| 222 | + ASE image of each structure. |
| 223 | + """ |
| 224 | + |
| 225 | + def __init__(self, ase_calc): |
| 226 | + self.ase_calc = ase_calc |
| 227 | + |
| 228 | + def copy(self): |
| 229 | + """Independent per-thread copy (ASE calculators are stateful).""" |
| 230 | + return ASEDirectCalculator(copy.deepcopy(self.ase_calc)) |
| 231 | + |
| 232 | + def compute(self, structure): |
| 233 | + atoms = structure.get_ase_atoms() |
| 234 | + atoms.calc = self.ase_calc |
| 235 | + results = { |
| 236 | + "energy": float(atoms.get_potential_energy()), # eV |
| 237 | + "forces": atoms.get_forces(), # eV / Ang |
| 238 | + "structure": structure.copy(), # trivial consistency check |
| 239 | + } |
| 240 | + try: |
| 241 | + results["stress"] = atoms.get_stress() # ASE Voigt, -eV/Ang^3 |
| 242 | + except Exception: |
| 243 | + pass # calculator without stress: call compute_ensemble with get_stress=False |
| 244 | + return results |
0 commit comments