overhaul Aims extension

closes #4

Author: jhrmnn

mona/sci/aims/aims.py

@@ -1,23 +1,23 @@
 # This Source Code Form is subject to the terms of the Mozilla Public
 # License, v. 2.0. If a copy of the MPL was not distributed with this
 # file, You can obtain one at http://mozilla.org/MPL/2.0/.
-import shutil
+import os
 from collections import OrderedDict
 from copy import deepcopy
 from pathlib import Path
-from typing import Any, Callable, Dict, Tuple, cast
+from typing import Any, Callable, Dict, List, Tuple, Type, cast
 
 from ...dirtask import dir_task
-from ...errors import InvalidInput, MonaError
+from ...errors import InvalidInput
 from ...files import File
 from ...pluggable import Pluggable, Plugin
 from ...pyhash import hash_function
 from ...tasks import Task
 from ..geomlib import Atom, Molecule
 from .dsl import expand_dicts, parse_aims_input
 
-__version__ = '0.1.0'
-__all__ = ['Aims', 'SpeciesDefaults']
+__version__ = '0.2.0'
+__all__ = ()
 
 
 class AimsPlugin(Plugin['Aims']):
@@ -29,24 +29,62 @@ def _func_hash(self) -> str:
 
 
 class Aims(Pluggable):
-    """A task factory that creates FHI-aims directory tasks."""
+    """:term:`Task factory` that creates an FHI-aims :term:`directory task`.
+
+    The creation of a task is handeld by a series of plugins that process
+    keyword arguments passed to the instance of :class:`Aims` to generate
+    ``control.in`` and ``geometry.in``. The default plugins are created with no
+    extra arguments. Custom plugins can be registered by calling them with the
+    :class:`Aims` instance as an argument.
+
+    Aims tasks require two prerequisites to function:
+
+    - ``mona_aims`` executable in ``PATH``, which runs Aims and accepts these
+      environment variables:
+
+      - ``MONA_AIMS`` specifies the Aims executable to run.
+      - ``MONA_NCORES`` specifies the number of processor cores to run Aims
+        with.
+
+    - ``AIMS_SPECIES_DEFAULTS`` environment variable that points to the
+      ``species_defaults`` Aims directory.
+    """
+
+    plugin_factories: List[Type[AimsPlugin]] = []
+    """List of plugins run in this order."""
 
     def __init__(self) -> None:
         Pluggable.__init__(self)
-        for factory in default_plugins:
+        for factory in Aims.plugin_factories:
             factory()(self)
 
     def __call__(self, *, label: str = None, **kwargs: Any) -> Task[Dict[str, File]]:
-        """Create an FHI-aims.
+        """Create an Aims task.
 
-        :param kwargs: processed by individual plugins
+        :param kwargs: keyword arguments processed by plugins
+        :param label: passed to :func:`~mona.dirtask.dir_task`
+
+        Existing keywords (some are generated by plugins, see below):
+
+        - ``atoms``: List of two-tuples of a species and a coordinate.
+        - ``geom``: Instance of :class:`~mona.sci.geomlib.Molecule`.
+        - ``species_specs``: List of species specifications.
+        - ``tags``: Dictionary of tags for ``control.in``.
+        - ``aims``: Aims executable, must be present in ``PATH``.
+        - ``check``: Whether task should error out on abnormal Aims exit.
+        - ``ncores``: Number of cores to use, all available if not given.
         """
         self.run_plugins('process', kwargs)
         script = File.from_str('aims.sh', kwargs.pop('script'))
-        inputs = [File.from_str(name, cont) for name, cont in kwargs.pop('inputs')]
+        inputs = [
+            File.from_str(name, kwargs.pop(name))
+            for name in ['control.in', 'geometry.in']
+        ]
+        task = dir_task(script, inputs, label=label)
+        task.storage['ncores'] = kwargs.pop('ncores', -1)
         if kwargs:
             raise InvalidInput(f'Unknown Aims kwargs: {list(kwargs.keys())}')
-        return dir_task(script, inputs, label=label)
+        return task
 
     def _func_hash(self) -> str:
         return ','.join(
@@ -57,42 +95,35 @@ def _func_hash(self) -> str:
         )
 
 
-class SpeciesDir(AimsPlugin):
-    def __init__(self) -> None:
-        self._speciesdirs: Dict[Tuple[str, str], Path] = {}
-
-    def process(self, kwargs: Dict[str, Any]) -> None:
-        sp_def_key = aims, sp_def = kwargs['aims'], kwargs.pop('species_defaults')
-        speciesdir = self._speciesdirs.get(sp_def_key)
-        if not speciesdir:
-            pathname = shutil.which(aims)
-            if not pathname:
-                pathname = shutil.which('aims-master')
-            if not pathname:
-                raise MonaError(f'Aims "{aims}" not found')
-            path = Path(pathname)
-            speciesdir = path.parents[1] / 'aimsfiles/species_defaults' / sp_def
-            self._speciesdirs[sp_def_key] = speciesdir  # type: ignore
-        kwargs['speciesdir'] = speciesdir
+class Atoms(AimsPlugin):
+    """Aims plugin handling geometry.
 
+    Keywords processed: ``atoms``. Keywords added: ``geom``.
+    """
 
-class Atoms(AimsPlugin):
     def process(self, kwargs: Dict[str, Any]) -> None:
         if 'atoms' in kwargs:
             kwargs['geom'] = Molecule([Atom(*args) for args in kwargs.pop('atoms')])
 
 
 class SpeciesDefaults(AimsPlugin):
-    """Aims plugin that handles adding species defaults to control.in."""
+    """Aims plugin handling species specifications.
+
+    :param mod: Callable that is passed the species specifications for modification.
+
+    Keywords added: ``species_specs``. Keywords used: ``geom``, ``species_defaults``.
+    """
 
     def __init__(self, mod: Callable[..., Any] = None) -> None:
         self._species_defs: Dict[Tuple[Path, str], Dict[str, Any]] = {}
         self._mod = mod
 
     def process(self, kwargs: Dict[str, Any]) -> None:  # noqa: D102
-        speciesdir = kwargs.pop('speciesdir')
+        speciesdir = Path(os.environ['AIMS_SPECIES_DEFAULTS']) / kwargs.pop(
+            'species_defaults'
+        )
         all_species = {(a.number, a.species) for a in kwargs['geom'].centers}
-        species_defs = []
+        species_specs = []
         for Z, species in sorted(all_species):
             if (speciesdir, species) not in self._species_defs:
                 species_def = parse_aims_input(
@@ -101,11 +132,11 @@ def process(self, kwargs: Dict[str, Any]) -> None:  # noqa: D102
                 self._species_defs[speciesdir, species] = species_def
             else:
                 species_def = self._species_defs[speciesdir, species]
-            species_defs.append(species_def)
+            species_specs.append(species_def)
         if self._mod:
-            species_defs = deepcopy(species_defs)
-            self._mod(species_defs, kwargs)
-        kwargs['species_defs'] = species_defs
+            species_specs = deepcopy(species_specs)
+            self._mod(species_specs, kwargs)
+        kwargs['species_specs'] = species_specs
 
     def _func_hash(self) -> str:
         if not self._mod:
@@ -115,9 +146,14 @@ def _func_hash(self) -> str:
 
 
 class Control(AimsPlugin):
+    """Aims plugin generating ``control.in``.
+
+    Keywords processed: ``species_specs``, ``tags``. Keywords added: ``control.in``.
+    """
+
     def process(self, kwargs: Dict[str, Any]) -> None:
         species_tags = []
-        for spec in kwargs.pop('species_defs'):
+        for spec in kwargs.pop('species_specs'):
             spec = OrderedDict(spec)
             while spec:
                 tag, value = spec.popitem(last=False)
@@ -142,32 +178,34 @@ def process(self, kwargs: Dict[str, Any]) -> None:
                 lines.extend(f'{tag}  {p2f(v)}' for v in value)
             else:
                 lines.append(f'{tag}  {p2f(value)}')
-        kwargs['control'] = '\n'.join(lines)
+        kwargs['control.in'] = '\n'.join(lines)
 
 
 class Geom(AimsPlugin):
-    def process(self, kwargs: Dict[str, Any]) -> None:
-        kwargs['geometry'] = kwargs.pop('geom').dumps('aims')
+    """Aims plugin generating ``geometry.in``.
 
+    Keywords processed: ``geom``. Keywords added: ``geometry.in``.
+    """
 
-class Core(AimsPlugin):
     def process(self, kwargs: Dict[str, Any]) -> None:
-        kwargs['inputs'] = [
-            ('control.in', kwargs.pop('control')),
-            ('geometry.in', kwargs.pop('geometry')),
-        ]
+        kwargs['geometry.in'] = kwargs.pop('geom').dumps('aims')
 
 
 class Script(AimsPlugin):
+    """Aims plugin generating script for the Aims taks.
+
+    Keywords processed: ``aims``, ``check``. Keywords added: ``script``.
+    """
+
     def process(self, kwargs: Dict[str, Any]) -> None:
         aims, check = kwargs.pop('aims'), kwargs.pop('check', True)
-        lines = ['#!/bin/bash', 'set -e', f'AIMS={aims} run_aims']
+        lines = ['#!/bin/bash', 'set -e', f'MONA_AIMS={aims} mona_aims']
         if check:
             lines.append('egrep "Have a nice day|stop_if_parser" STDOUT >/dev/null')
         kwargs['script'] = '\n'.join(lines)
 
 
-default_plugins = [SpeciesDir, Atoms, SpeciesDefaults, Control, Geom, Core, Script]
+Aims.plugin_factories = [Atoms, SpeciesDefaults, Control, Geom, Script]
 
 
 def p2f(value: Any, nospace: bool = False) -> str: