Coverage for /builds/kinetik161/ase/ase/calculators/lammpsrun.py: 87.62%

1"""ASE calculator for the LAMMPS classical MD code"""

2# lammps.py (2011/03/29)

6# This library is free software; you can redistribute it and/or

7# modify it under the terms of the GNU Lesser General Public

8# License as published by the Free Software Foundation; either

9# version 2.1 of the License, or (at your option) any later version.

10#

11# This library is distributed in the hope that it will be useful,

12# but WITHOUT ANY WARRANTY; without even the implied warranty of

13# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU

14# Lesser General Public License for more details.

15#

16# You should have received a copy of the GNU Lesser General Public

17# License along with this file; if not, write to the Free Software

18# Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301

19# USA or see <http://www.gnu.org/licenses/>.

22import os

23import shlex

24import shutil

25import warnings

26from re import IGNORECASE

27from re import compile as re_compile

28import subprocess

29from tempfile import NamedTemporaryFile, mkdtemp

30from tempfile import mktemp as uns_mktemp

31from threading import Thread

32from typing import Any, Dict

34import numpy as np

36from ase.calculators.calculator import Calculator, all_changes

37from ase.calculators.lammps import (CALCULATION_END_MARK, Prism, convert,

38 write_lammps_in)

39from ase.data import atomic_masses, chemical_symbols

40from ase.io.lammpsdata import write_lammps_data

41from ase.io.lammpsrun import read_lammps_dump

43__all__ = ["LAMMPS"]

46class LAMMPS(Calculator):

47 """LAMMPS (https://lammps.sandia.gov/) calculator

49 The environment variable :envvar:`ASE_LAMMPSRUN_COMMAND` must be defined to

50 tell ASE how to call a LAMMPS binary. This should contain the path to the

51 lammps binary, or more generally, a command line possibly also including an

52 MPI-launcher command.

54 For example (in a Bourne-shell compatible environment):

56 .. code-block:: bash

58 export ASE_LAMMPSRUN_COMMAND=/path/to/lmp_binary

60 or possibly something similar to

62 .. code-block:: bash

64 export ASE_LAMMPSRUN_COMMAND="/path/to/mpirun --np 4 lmp_binary"

66 Parameters

67 ----------

68 files : list[str]

69 List of files needed by LAMMPS. Typically a list of potential files.

70 parameters : dict[str, Any]

71 Dictionary of settings to be passed into the input file for calculation.

72 specorder : list[str]

73 Within LAAMPS, atoms are identified by an integer value starting from 1.

74 This variable allows the user to define the order of the indices

75 assigned to the atoms in the calculation, with the default

76 if not given being alphabetical

77 keep_tmp_files : bool, default: False

78 Retain any temporary files created. Mostly useful for debugging.

79 tmp_dir : str, default: None

80 path/dirname (default None -> create automatically).

81 Explicitly control where the calculator object should create

82 its files. Using this option implies 'keep_tmp_files=True'.

83 no_data_file : bool, default: False

84 Controls whether an explicit data file will be used for feeding

85 atom coordinates into lammps. Enable it to lessen the pressure on

86 the (tmp) file system. THIS OPTION MIGHT BE UNRELIABLE FOR CERTAIN

87 CORNER CASES (however, if it fails, you will notice...).

88 keep_alive : bool, default: True

89 When using LAMMPS as a spawned subprocess, keep the subprocess

90 alive (but idling when unused) along with the calculator object.

91 always_triclinic : bool, default: False

92 Force LAMMPS to treat the cell as tilted, even if the cell is not

93 tilted, by printing ``xy``, ``xz``, ``yz`` in the data file.

94 reduce_cell : bool, default: False

95 If True, reduce cell to have shorter lattice vectors.

96 write_velocities : bool, default: False

97 If True, forward ASE velocities to LAMMPS.

98 verbose: bool, default: False

99 If True, print additional debugging output to STDOUT.

100

101 Examples

102 --------

103 Provided that the respective potential file is in the working directory,

104 one can simply run (note that LAMMPS needs to be compiled to work with EAM

105 potentials)

106

107 ::

108

109 from ase import Atom, Atoms

110 from ase.build import bulk

111 from ase.calculators.lammpsrun import LAMMPS

112

113 parameters = {'pair_style': 'eam/alloy',

114 'pair_coeff': ['* * NiAlH_jea.eam.alloy H Ni']}

115

116 files = ['NiAlH_jea.eam.alloy']

117

118 Ni = bulk('Ni', cubic=True)

119 H = Atom('H', position=Ni.cell.diagonal()/2)

120 NiH = Ni + H

121

122 lammps = LAMMPS(parameters=parameters, files=files)

123

124 NiH.calc = lammps

125 print("Energy ", NiH.get_potential_energy())

126 """

127

128 name = "lammpsrun"

129 implemented_properties = ["energy", "free_energy", "forces", "stress",

130 "energies"]

131

132 # parameters to choose options in LAMMPSRUN

133 ase_parameters: Dict[str, Any] = dict(

134 specorder=None,

135 atorder=True,

136 always_triclinic=False,

137 reduce_cell=False,

138 keep_alive=True,

139 keep_tmp_files=False,

140 no_data_file=False,

141 tmp_dir=None,

142 files=[], # usually contains potential parameters

143 verbose=False,

144 write_velocities=False,

145 binary_dump=True, # bool - use binary dump files (full

146 # precision but long long ids are casted to

147 # double)

148 lammps_options="-echo log -screen none -log /dev/stdout",

149 trajectory_out=None, # file object, if is not None the

150 # trajectory will be saved in it

151 )

152

153 # parameters forwarded to LAMMPS

154 lammps_parameters = dict(

155 boundary=None, # bounadry conditions styles

156 units="metal", # str - Which units used; some potentials

157 # require certain units

158 atom_style="atomic",

159 special_bonds=None,

160 # potential informations

161 pair_style="lj/cut 2.5",

162 pair_coeff=["* * 1 1"],

163 masses=None,

164 pair_modify=None,

165 # variables controlling the output

166 thermo_args=[

167 "step", "temp", "press", "cpu", "pxx", "pyy", "pzz",

168 "pxy", "pxz", "pyz", "ke", "pe", "etotal", "vol", "lx",

169 "ly", "lz", "atoms", ],

170 dump_properties=["id", "type", "x", "y", "z", "vx", "vy",

171 "vz", "fx", "fy", "fz", ],

172 dump_period=1, # period of system snapshot saving (in MD steps)

173 )

174

175 default_parameters = dict(ase_parameters, **lammps_parameters)

176

177 def __init__(self, label="lammps", **kwargs):

178 super().__init__(label=label, **kwargs)

179

180 self.prism = None

181 self.calls = 0

182 self.forces = None

183 # thermo_content contains data "written by" thermo_style.

184 # It is a list of dictionaries, each dict (one for each line

185 # printed by thermo_style) contains a mapping between each

186 # custom_thermo_args-argument and the corresponding

187 # value as printed by lammps. thermo_content will be

188 # re-populated by the read_log method.

189 self.thermo_content = []

190

191 if self.parameters['tmp_dir'] is not None:

192 # If tmp_dir is pointing somewhere, don't remove stuff!

193 self.parameters['keep_tmp_files'] = True

194 self._lmp_handle = None # To handle the lmp process

195

196 if self.parameters['tmp_dir'] is None:

197 self.parameters['tmp_dir'] = mkdtemp(prefix="LAMMPS-")

198 else:

199 self.parameters['tmp_dir'] = os.path.realpath(

200 self.parameters['tmp_dir'])

201 if not os.path.isdir(self.parameters['tmp_dir']):

202 os.mkdir(self.parameters['tmp_dir'], 0o755)

203

204 for f in self.parameters['files']:

205 shutil.copy(

206 f, os.path.join(self.parameters['tmp_dir'],

207 os.path.basename(f)))

208

209 def get_lammps_command(self):

210 cmd = self.parameters.get('command')

211 if cmd is None:

212 envvar = f'ASE_{self.name.upper()}_COMMAND'

213 cmd = os.environ.get(envvar)

214

215 if cmd is None:

216 cmd = 'lammps'

217

218 opts = self.parameters.get('lammps_options')

219

220 if opts is not None:

221 cmd = f'{cmd} {opts}'

222

223 return cmd

224

225 def clean(self, force=False):

226

227 self._lmp_end()

228

229 if not self.parameters['keep_tmp_files'] or force:

230 shutil.rmtree(self.parameters['tmp_dir'])

231

232 def check_state(self, atoms, tol=1.0e-10):

233 # Transforming the unit cell to conform to LAMMPS' convention for

234 # orientation (c.f. https://lammps.sandia.gov/doc/Howto_triclinic.html)

235 # results in some precision loss, so we use bit larger tolerance than

236 # machine precision here. Note that there can also be precision loss

237 # related to how many significant digits are specified for things in

238 # the LAMMPS input file.

239 return Calculator.check_state(self, atoms, tol)

240

241 def calculate(self, atoms=None, properties=None, system_changes=None):

242 if properties is None:

243 properties = self.implemented_properties

244 if system_changes is None:

245 system_changes = all_changes

246 Calculator.calculate(self, atoms, properties, system_changes)

247 self.run()

248

249 def _lmp_alive(self):

250 # Return True if this calculator is currently handling a running

251 # lammps process

252 return self._lmp_handle and not isinstance(

253 self._lmp_handle.poll(), int

254 )

255

256 def _lmp_end(self):

257 # Close lammps input and wait for lammps to end. Return process

258 # return value

259 if self._lmp_alive():

260 # !TODO: handle lammps error codes

261 try:

262 self._lmp_handle.communicate(timeout=5)

263 except subprocess.TimeoutExpired:

264 self._lmp_handle.kill()

265 self._lmp_handle.communicate()

266 err = self._lmp_handle.poll()

267 assert err is not None

268 return err

269

270 def set_missing_parameters(self):

271 """Verify that all necessary variables are set.

272 """

273 symbols = self.atoms.get_chemical_symbols()

274 # If unspecified default to atom types in alphabetic order

275 if not self.parameters.get('specorder'):

276 self.parameters['specorder'] = sorted(set(symbols))

277

278 # !TODO: handle cases were setting masses actual lead to errors

279 if not self.parameters.get('masses'):

280 self.parameters['masses'] = []

281 for type_id, specie in enumerate(self.parameters['specorder']):

282 mass = atomic_masses[chemical_symbols.index(specie)]

283 self.parameters['masses'] += [

284 f"{type_id + 1:d} {mass:f}"

285 ]

286

287 # set boundary condtions

288 if not self.parameters.get('boundary'):

289 b_str = " ".join(["fp"[int(x)] for x in self.atoms.pbc])

290 self.parameters['boundary'] = b_str

291

292 def run(self, set_atoms=False):

293 # !TODO: split this function

294 """Method which explicitly runs LAMMPS."""

295 pbc = self.atoms.get_pbc()

296 if all(pbc):

297 cell = self.atoms.get_cell()

298 elif not any(pbc):

299 # large enough cell for non-periodic calculation -

300 # LAMMPS shrink-wraps automatically via input command

301 # "periodic s s s"

302 # below

303 cell = 2 * np.max(np.abs(self.atoms.get_positions())) * np.eye(3)

304 else:

305 warnings.warn(

306 "semi-periodic ASE cell detected - translation "

307 + "to proper LAMMPS input cell might fail"

308 )

309 cell = self.atoms.get_cell()

310 self.prism = Prism(cell)

311

312 self.set_missing_parameters()

313 self.calls += 1

314

315 # change into subdirectory for LAMMPS calculations

316 tempdir = self.parameters['tmp_dir']

317

318 # setup file names for LAMMPS calculation

319 label = f"{self.label}{self.calls:>06}"

320 lammps_in = uns_mktemp(

321 prefix="in_" + label, dir=tempdir

322 )

323 lammps_log = uns_mktemp(

324 prefix="log_" + label, dir=tempdir

325 )

326 lammps_trj_fd = NamedTemporaryFile(

327 prefix="trj_" + label,

328 suffix=(".bin" if self.parameters['binary_dump'] else ""),

329 dir=tempdir,

330 delete=(not self.parameters['keep_tmp_files']),

331 )

332 lammps_trj = lammps_trj_fd.name

333 if self.parameters['no_data_file']:

334 lammps_data = None

335 else:

336 lammps_data_fd = NamedTemporaryFile(

337 prefix="data_" + label,

338 dir=tempdir,

339 delete=(not self.parameters['keep_tmp_files']),

340 mode='w',

341 encoding='ascii'

342 )

343 write_lammps_data(

344 lammps_data_fd,

345 self.atoms,

346 specorder=self.parameters['specorder'],

347 force_skew=self.parameters['always_triclinic'],

348 reduce_cell=self.parameters['reduce_cell'],

349 velocities=self.parameters['write_velocities'],

350 prismobj=self.prism,

351 units=self.parameters['units'],

352 atom_style=self.parameters['atom_style'],

353 )

354 lammps_data = lammps_data_fd.name

355 lammps_data_fd.flush()

356

357 # see to it that LAMMPS is started

358 if not self._lmp_alive():

359 command = self.get_lammps_command()

360 # Attempt to (re)start lammps

361 self._lmp_handle = subprocess.Popen(

362 shlex.split(command, posix=(os.name == "posix")),

363 stdin=subprocess.PIPE,

364 stdout=subprocess.PIPE,

365 encoding='ascii',

366 )

367 lmp_handle = self._lmp_handle

368

369 # Create thread reading lammps stdout (for reference, if requested,

370 # also create lammps_log, although it is never used)

371 if self.parameters['keep_tmp_files']:

372 lammps_log_fd = open(lammps_log, "w")

373 fd = SpecialTee(lmp_handle.stdout, lammps_log_fd)

374 else:

375 fd = lmp_handle.stdout

376 thr_read_log = Thread(target=self.read_lammps_log, args=(fd,))

377 thr_read_log.start()

378

379 # write LAMMPS input (for reference, also create the file lammps_in,

380 # although it is never used)

381 if self.parameters['keep_tmp_files']:

382 lammps_in_fd = open(lammps_in, "w")

383 fd = SpecialTee(lmp_handle.stdin, lammps_in_fd)

384 else:

385 fd = lmp_handle.stdin

386 write_lammps_in(

387 lammps_in=fd,

388 parameters=self.parameters,

389 atoms=self.atoms,

390 prismobj=self.prism,

391 lammps_trj=lammps_trj,

392 lammps_data=lammps_data,

393 )

394

395 if self.parameters['keep_tmp_files']:

396 lammps_in_fd.close()

397

398 # Wait for log output to be read (i.e., for LAMMPS to finish)

399 # and close the log file if there is one

400 thr_read_log.join()

401 if self.parameters['keep_tmp_files']:

402 lammps_log_fd.close()

403

404 if not self.parameters['keep_alive']:

405 self._lmp_end()

406

407 exitcode = lmp_handle.poll()

408 if exitcode and exitcode != 0:

409 raise RuntimeError(

410 "LAMMPS exited in {} with exit code: {}."

411 "".format(tempdir, exitcode)

412 )

413

414 # A few sanity checks

415 if len(self.thermo_content) == 0:

416 raise RuntimeError("Failed to retrieve any thermo_style-output")

417 if int(self.thermo_content[-1]["atoms"]) != len(self.atoms):

418 # This obviously shouldn't happen, but if prism.fold_...() fails,

419 # it could

420 raise RuntimeError("Atoms have gone missing")

421

422 trj_atoms = read_lammps_dump(

423 infileobj=lammps_trj,

424 order=self.parameters['atorder'],

425 index=-1,

426 prismobj=self.prism,

427 specorder=self.parameters['specorder'],

428 )

429

430 if set_atoms:

431 self.atoms = trj_atoms.copy()

432

433 self.forces = trj_atoms.get_forces()

434 # !TODO: trj_atoms is only the last snapshot of the system; Is it

435 # desirable to save also the inbetween steps?

436 if self.parameters['trajectory_out'] is not None:

437 # !TODO: is it advisable to create here temporary atoms-objects

438 self.trajectory_out.write(trj_atoms)

439

440 tc = self.thermo_content[-1]

441 self.results["energy"] = convert(

442 tc["pe"], "energy", self.parameters["units"], "ASE"

443 )

444 self.results["free_energy"] = self.results["energy"]

445 self.results['forces'] = convert(self.forces.copy(),

446 'force',

447 self.parameters['units'],

448 'ASE')

449 stress = np.array(

450 [-tc[i] for i in ("pxx", "pyy", "pzz", "pyz", "pxz", "pxy")]

451 )

452

453 # We need to apply the Lammps rotation stuff to the stress:

454 xx, yy, zz, yz, xz, xy = stress

455 stress_tensor = np.array([[xx, xy, xz],

456 [xy, yy, yz],

457 [xz, yz, zz]])

458 stress_atoms = self.prism.tensor2_to_ase(stress_tensor)

459 stress_atoms = stress_atoms[[0, 1, 2, 1, 0, 0],

460 [0, 1, 2, 2, 2, 1]]

461 stress = stress_atoms

462

463 self.results["stress"] = convert(

464 stress, "pressure", self.parameters["units"], "ASE"

465 )

466

467 lammps_trj_fd.close()

468 if not self.parameters['no_data_file']:

469 lammps_data_fd.close()

470

471 def __enter__(self):

472 return self

473

474 def __exit__(self, *args):

475 self._lmp_end()

476

477 def read_lammps_log(self, fileobj):

478 # !TODO: somehow communicate 'thermo_content' explicitly

479 """Method which reads a LAMMPS output log file."""

480

481 # read_log depends on that the first (three) thermo_style custom args

482 # can be capitalized and matched against the log output. I.e.

483 # don't use e.g. 'ke' or 'cpu' which are labeled KinEng and CPU.

484 mark_re = r"^\s*" + r"\s+".join(

485 [x.capitalize() for x in self.parameters['thermo_args'][0:3]]

486 )

487 _custom_thermo_mark = re_compile(mark_re)

488

489 # !TODO: regex-magic necessary?

490 # Match something which can be converted to a float

491 f_re = r"([+-]?(?:(?:\d+(?:\.\d*)?|\.\d+)(?:e[+-]?\d+)?|nan|inf))"

492 n_args = len(self.parameters["thermo_args"])

493 # Create a re matching exactly N white space separated floatish things

494 _custom_thermo_re = re_compile(

495 r"^\s*" + r"\s+".join([f_re] * n_args) + r"\s*$", flags=IGNORECASE

496 )

497

498 thermo_content = []

499 line = fileobj.readline()

500 while line and line.strip() != CALCULATION_END_MARK:

501 # check error

502 if 'ERROR:' in line:

503 raise RuntimeError(f'LAMMPS exits with error message: {line}')

504

505 # get thermo output

506 if _custom_thermo_mark.match(line):

507 while True:

508 line = fileobj.readline()

509 if 'WARNING:' in line:

510 continue

511

512 bool_match = _custom_thermo_re.match(line)

513 if not bool_match:

514 break

515

516 # create a dictionary between each of the

517 # thermo_style args and it's corresponding value

518 thermo_content.append(

519 dict(

520 zip(

521 self.parameters['thermo_args'],

522 map(float, bool_match.groups()),

523 )

524 )

525 )

526 else:

527 line = fileobj.readline()

528

529 self.thermo_content = thermo_content

530

531

532class SpecialTee:

533 """A special purpose, with limited applicability, tee-like thing.

534

535 A subset of stuff read from, or written to, orig_fd,

536 is also written to out_fd.

537 It is used by the lammps calculator for creating file-logs of stuff

538 read from, or written to, stdin and stdout, respectively.

539 """

540

541 def __init__(self, orig_fd, out_fd):

542 self._orig_fd = orig_fd

543 self._out_fd = out_fd

544 self.name = orig_fd.name

545

546 def write(self, data):

547 self._orig_fd.write(data)

548 self._out_fd.write(data)

549 self.flush()

550

551 def read(self, *args, **kwargs):

552 data = self._orig_fd.read(*args, **kwargs)

553 self._out_fd.write(data)

554 return data

555

556 def readline(self, *args, **kwargs):

557 data = self._orig_fd.readline(*args, **kwargs)

558 self._out_fd.write(data)

559 return data

560

561 def readlines(self, *args, **kwargs):

562 data = self._orig_fd.readlines(*args, **kwargs)

563 self._out_fd.write("".join(data))

564 return data

565

566 def flush(self):

567 self._orig_fd.flush()

568 self._out_fd.flush()