Implement tasklets upon base ASE tasklets

The base tasklets provides a backend-agnostic interface to run ASE calculators asynchronously. To simplify the interface, we do not follow ASE’s conventions, such as Atoms object and return values’ units.

Before starting to implement a tasklet for a new ASE calculators, you might want to know what ASE-Grain tasklet is actually doing in addition to the ASE calculator. In short, a tasklet:

  1. Make running the calculation async, and
  2. Make the calculation resource-aware (i.e. processors and memory)

All the other things (e.g. generate input files from a somehow uniform parameter (atoms, method, basis, etc.) interface; run the calculation and parse the output to get energy, forces, etc.) are handled by ASE calculator. The first thing is nicely handled by the following base tasklets. However, the second thing is often not so straightforward, as backends tend to have diverse ways to define processors and memory constraints. The existing tasklets would be a good reference for implementing new ones.

await ase_grain.base.ase_fio_task(cid, calc, atcor, ian, cell=None, pbc=None)

Base task for calculators implementing the FileIOCalculator interface (e.g. Gaussian, Orca, QE).

Class FileIOCalculator is patched to enable async calculation through async subprocess. See monkey_patch.py for details.

Parameters:
  • cid (str) – Calculation ID, for time record only
  • calc – The calculator instance
  • atcor (ndarray[(N,3), float]) – Atom coordinate
  • ian (ndarray[(N,), int]) – Atomic numbers
  • pbc (bool or (bool, bool, bool)) – Periodic boundary conditions along each axis. Arg for ase.atoms.Atoms
  • cell (ndarray[(3,3), float] or see ASE doc) – Unit cell size. Arg for ase.atoms.Atoms
Returns:

Energy (in hartree) and forces (in hartree/angstrom)

Note

Any task calling this should make sure that the calculator itself’s calculate method get handled, as this only takes care of FileIOCalculator.calculate.

await ase_grain.base.ase_task(cid, calc, atcor, ian)

Base task for inproc calculators (e.g. Psi4).

For Parameters and Returns see above ase_fio_task().

ase_task() uses grain.subproc.subprocify to “asyncify” the calculator, as we assume that an inproc calculator’s calculate method is CPU intensive. subprocify requires an async context to maintain a subprocess pool. This could be simply achieved by:

from grain.subproc import subprocess_pool_scope
async with subprocess_pool_scope():
    # Safely call subprocified functions within

Note

Task calling this experiences an overhead for the first few runs due to subprocify’s subprocess startup cost, which will be amortized later as the subprocesses are reused.