Merge pull request #253 from tovrstra/atom-line-template

Make atom lines configurable in the input writer

doc/index.rst

@@ -46,17 +46,17 @@ For the list of file formats that can be loaded or dumped by IOData, see
 :ref:`file_formats`. The two tables below summarize the file formats and
 features supported by IOData.
 
-======= ==========
-Code    Definition
-======= ==========
-**L**   loading is supported
-**D**   dumping is supported
-*(d)*   attribute may be derived from other attributes
-R       attribute is always read
-r       attribute is read if present
-W       attribute is always written
-w       attribute is is written if present
-======= ==========
+========= ==========
+Code      Definition
+========= ==========
+**L**     loading is supported
+**D**     dumping is supported
+*(d)*     attribute may be derived from other attributes
+R         attribute is always read
+r         attribute is read if present
+W         attribute is always written
+w         attribute is is written if present
+========= ==========
 
 .. include:: formats_tab.inc
 

iodata/api.py

@@ -24,7 +24,7 @@
 from importlib import import_module
 from pkgutil import iter_modules
 from types import ModuleType
-from typing import Optional
+from typing import Callable, Optional
 
 from .iodata import IOData
 from .utils import LineIterator
@@ -225,7 +225,14 @@ def dump_many(iodatas: Iterator[IOData], filename: str, fmt: Optional[str] = Non
         format_module.dump_many(f, iodatas, **kwargs)
 
 
-def write_input(iodata: IOData, filename: str, fmt: str, template: Optional[str] = None, **kwargs):
+def write_input(
+    iodata: IOData,
+    filename: str,
+    fmt: str,
+    template: Optional[str] = None,
+    atom_line: Optional[Callable] = None,
+    **kwargs,
+):
     """Write input file using an instance of IOData for the specified software format.
 
     Parameters
@@ -238,10 +245,16 @@ def write_input(iodata: IOData, filename: str, fmt: str, template: Optional[str]
         The name of the software for which input file is generated.
     template
         The template input string.
+        If not given, a default template for the selected software is used.
+    atom_line
+        A function taking two arguments: an IOData instance, and an index of
+        the atom. This function returns a formatted line for the corresponding
+        atom. When omitted, a default atom_line function for the selected
+        input format is used.
     **kwargs
         Keyword arguments are passed on to the input-specific write_input function.
 
     """
     input_module = _select_input_module(fmt)
-    with open(filename, "w") as f:
-        input_module.write_input(f, iodata, template=template, **kwargs)
+    with open(filename, "w") as fh:
+        input_module.write_input(fh, iodata, template, atom_line, **kwargs)

iodata/docstrings.py

@@ -335,12 +335,10 @@ def _document_write(
     fmt: str,
     required: list[str],
     optional: Optional[list[str]] = None,
-    kwdocs: Optional[dict[str, str]] = None,
     notes: Optional[str] = None,
 ):
-    if kwdocs is None:
-        kwdocs = {}
-    optional = optional or []
+    if optional is None:
+        optional = []
 
     def decorator(func):
         if optional:
@@ -355,16 +353,11 @@ def decorator(func):
             fmt=fmt,
             required=", ".join(f"``{word}``" for word in required),
             optional=optional_sentence,
-            kwdocs="\n".join(
-                "{}\n    {}".format(name, docu.replace("\n", " "))
-                for name, docu in sorted(kwdocs.items())
-            ),
             notes=(notes or ""),
         )
         func.fmt = fmt
         func.required = required
         func.optional = optional
-        func.kwdocs = kwdocs
         func.notes = notes
         return func
 
@@ -383,7 +376,13 @@ def decorator(func):
     {required}.{optional}
 template
     A template input string.
-{kwdocs}
+atom_line
+    A function taking two arguments: an IOData instance, and an index of
+    the atom. This function returns a formatted line for the corresponding
+    atom. When omitted, a default atom_line function for the selected
+    input format is used.
+**kwargs
+    Keyword arguments are passed on to the input-specific write_input function.
 Notes
 -----
 
@@ -396,7 +395,6 @@ def document_write_input(
     fmt: str,
     required: list[str],
     optional: Optional[list[str]] = None,
-    kwdocs: Optional[dict[str, str]] = None,
     notes: Optional[str] = None,
 ):
     """Decorate a write_input function to generate a docstring.
@@ -409,10 +407,6 @@ def document_write_input(
         A list of mandatory IOData attributes needed to write the file.
     optional
         A list of optional IOData attributes which can be include when writing the file.
-    kwdocs
-        A dictionary with documentation for keyword arguments. Each key is a
-        keyword argument name and the corresponding value is text explaining the
-        argument.
     notes
         Additional information to be added to the docstring.
 
@@ -422,6 +416,4 @@ def document_write_input(
         A decorator function.
 
     """
-    if kwdocs is None:
-        kwdocs = {}
-    return _document_write(WRITE_INPUT_DOC_TEMPLATE, fmt, required, optional, kwdocs, notes)
+    return _document_write(WRITE_INPUT_DOC_TEMPLATE, fmt, required, optional, notes)

iodata/formats/json.py

@@ -57,13 +57,13 @@
 The provenance field contains information about how the associated QCSchema object and its
 attributes were generated, provided, and manipulated. A provenance entry expects these fields:
 
-======= ===========
-Field   Description
-======= ===========
-creator **Required**. The program that generated, provided, or manipulated this file.
-version The version of the creator.
-routine The routine of the creator.
-======= ===========
+========= ===========
+Field     Description
+========= ===========
+creator   **Required**. The program that generated, provided, or manipulated this file.
+version   The version of the creator.
+routine   The routine of the creator.
+========= ===========
 
 In QCElemental, only a single provenance entry is permitted. When generating a QCSchema file for use
 with QCElemental, the easiest way to ensure compliance is to leave the provenance field blank, to

iodata/inputs/common.py

@@ -18,25 +18,30 @@
 # --
 """Utilities for writing input files."""
 
+from typing import Callable, TextIO
+
 import attrs
 import numpy as np
 
 from ..iodata import IOData
-from ..utils import angstrom
 
-__all__ = ["populate_fields"]
+__all__ = ["write_input_base"]
 
 
-def populate_fields(data: IOData) -> dict:
+def write_input_base(
+    fh: TextIO, data: IOData, template: str, atom_line: Callable, user_fields: dict
+):
     """Generate a dictionary with fields to replace in the template."""
-    # load IOData dict using attr.asdict because the IOData class uses __slots__
+    # Convert IOData instance to dict for formatting.
     fields = attrs.asdict(data, recurse=False)
-    # store atomic coordinates in angstrom
-    fields["atcoords"] = data.atcoords / angstrom
-    # set general defaults
+    # Set general defaults.
     fields["title"] = data.title if data.title is not None else "Input Generated by IOData"
-    fields["run_type"] = data.run_type if data.run_type is not None else "energy"
-    # convert spin polarization to multiplicity
+    # Convert spin polarization to multiplicity.
     fields["spinmult"] = int(abs(np.round(data.spinpol))) + 1 if data.spinpol is not None else 1
     fields["charge"] = int(data.charge) if data.charge is not None else 0
-    return fields
+    # User- or format-specific fields have priority.
+    fields.update(user_fields)
+    # Generate geometry.
+    geometry = [atom_line(data, iatom) for iatom in range(data.natom)]
+    fields["geometry"] = "\n".join(geometry)
+    print(template.format(**fields), file=fh)

iodata/inputs/gaussian.py

@@ -18,12 +18,13 @@
 # --
 """Gaussian Input Module."""
 
-from typing import Optional, TextIO
+from typing import Callable, Optional, TextIO
 
 from ..docstrings import document_write_input
 from ..iodata import IOData
 from ..periodic import num2sym
-from .common import populate_fields
+from ..utils import angstrom
+from .common import write_input_base
 
 __all__ = []
 
@@ -39,35 +40,43 @@
 """
 
 
+def default_atom_line(data: IOData, iatom: int):
+    """Format atom line for Gaussian input."""
+    symbol = num2sym[data.atnums[iatom]]
+    atcoord = data.atcoords[iatom] / angstrom
+    return f"{symbol:3s} {atcoord[0]:10.6f} {atcoord[1]:10.6f} {atcoord[2]:10.6f}"
+
+
 @document_write_input(
     "GAUSSIAN",
     ["atnums", "atcoords"],
     ["title", "run_type", "lot", "obasis_name", "spinmult", "charge"],
 )
-def write_input(f: TextIO, data: IOData, template: Optional[str] = None, **kwargs):
+def write_input(
+    fh: TextIO,
+    data: IOData,
+    template: Optional[str] = None,
+    atom_line: Optional[Callable] = None,
+    **kwargs,
+):
     """Do not edit this docstring. It will be overwritten."""
-    # initialize a dictionary with fields to replace in the template
-    fields = populate_fields(data)
-    # set format-specific defaults
-    fields["lot"] = data.lot if data.lot is not None else "hf"
-    fields["obasis_name"] = data.obasis_name if data.obasis_name is not None else "sto-3g"
-    # convert run type to Gaussian keywords
-    run_types = {
+    # Fill in some Gaussian-specific defaults and field names.
+    if template is None:
+        template = default_template
+    if atom_line is None:
+        atom_line = default_atom_line
+    gaussian_keywords = {
         "energy": "sp",
         "energy_force": "force",
         "opt": "opt",
         "scan": "scan",
         "freq": "freq",
     }
-    fields["run_type"] = run_types[fields["run_type"].lower()]
-    # generate geometry (in angstrom)
-    geometry = []
-    for num, coord in zip(fields["atnums"], fields["atcoords"]):
-        geometry.append(f"{num2sym[num]:3} {coord[0]:10.6f} {coord[1]:10.6f} {coord[2]:10.6f}")
-    fields["geometry"] = "\n".join(geometry)
-    # get template
-    if template is None:
-        template = default_template
-    # populate files & write input
+    fields = {
+        "lot": data.lot or "hf",
+        "obasis_name": data.obasis_name or "sto-3g",
+        "run_type": gaussian_keywords[(data.run_type or "energy").lower()],
+    }
+    # User-specifield fields have priority, may overwrite default ones.
     fields.update(kwargs)
-    print(template.format(**fields), file=f)
+    write_input_base(fh, data, template, atom_line, fields)

iodata/inputs/orca.py

@@ -18,12 +18,13 @@
 # --
 """Orca Input Module."""
 
-from typing import Optional, TextIO
+from typing import Callable, Optional, TextIO
 
 from ..docstrings import document_write_input
 from ..iodata import IOData
 from ..periodic import num2sym
-from .common import populate_fields
+from ..utils import angstrom
+from .common import write_input_base
 
 __all__ = []
 
@@ -36,33 +37,38 @@
 *"""
 
 
+def default_atom_line(data: IOData, iatom: int):
+    """Format atom line for ORCA input."""
+    symbol = num2sym[data.atnums[iatom]]
+    atcoord = data.atcoords[iatom] / angstrom
+    return f"{symbol:3s} {atcoord[0]:10.6f} {atcoord[1]:10.6f} {atcoord[2]:10.6f}"
+
+
 @document_write_input(
     "ORCA",
     ["atnums", "atcoords"],
     ["title", "run_type", "lot", "obasis_name", "spinmult", "charge"],
 )
-def write_input(f: TextIO, data: IOData, template: Optional[str] = None, **kwargs):
+def write_input(
+    fh: TextIO,
+    data: IOData,
+    template: Optional[str] = None,
+    atom_line: Optional[Callable] = None,
+    **kwargs,
+):
     """Do not edit this docstring. It will be overwritten."""
-    # initialize a dictionary with fields to replace in the template
-    fields = populate_fields(data)
-    # set format-specific defaults
-    fields["lot"] = data.lot if data.lot is not None else "HF"
-    fields["obasis_name"] = data.obasis_name if data.obasis_name is not None else "STO-3G"
-    # convert run type to Orca keywords
-    run_types = {"energy": "Energy", "freq": "Freq", "opt": "Opt"}
-    fields["run_type"] = run_types[fields["run_type"].lower()]
-    # generate geometry (in angstrom)
-    geometry = []
-    for num, coord in zip(fields["atnums"], fields["atcoords"]):
-        sym = f"{num2sym[num]:3}"
-        # check if template has a %coords block
-        if template is not None and "%coords" in template:
-            sym = f"{sym:>11}"  # adding an appropiate indentation
-        geometry.append(f"{sym} {coord[0]:10.6f} {coord[1]:10.6f} {coord[2]:10.6f}")
-    fields["geometry"] = "\n".join(geometry)
-    # get template
+    # Fill in some ORCA-specific defaults and field names.
     if template is None:
         template = default_template
-    # populate files & write input
+    if atom_line is None:
+        atom_line = default_atom_line
+    orca_keywords = {"energy": "Energy", "freq": "Freq", "opt": "Opt"}
+    # Set format-specific defaults.
+    fields = {
+        "lot": data.lot or "HF",
+        "obasis_name": data.obasis_name or "STO-3G",
+        "run_type": orca_keywords[(data.run_type or "energy").lower()],
+    }
+    # User-specifield fields have priority, may overwrite default ones.
     fields.update(kwargs)
-    print(template.format(**fields), file=f)
+    write_input_base(fh, data, template, atom_line, fields)

iodata/test/data/input_gaussian_bsse.com

@@ -0,0 +1,16 @@
+# hf/sto-3g Counterpoise=2
+
+Counterpoise calculation on 4144_02WaterMeOH
+
+0,1 0,1 0,1
+O(Fragment=1)  -0.525330  -0.050971  -0.314517
+H(Fragment=1)  -0.942007   0.747902   0.011253
+H(Fragment=1)   0.403697   0.059786  -0.073568
+O(Fragment=2)   2.316633   0.045501   0.071858
+H(Fragment=2)   2.684616  -0.526577   0.749387
+C(Fragment=2)   2.781638  -0.426129  -1.190301
+H(Fragment=2)   2.350821   0.224965  -1.943415
+H(Fragment=2)   3.867602  -0.375336  -1.264613
+H(Fragment=2)   2.453296  -1.445999  -1.389381
+
+

iodata/test/data/s66_4114_02WaterMeOH.xyz

@@ -0,0 +1,11 @@
+       9
+name="4144_02WaterMeOH" Properties=species:S:1:pos:R:3:fragment_ids:I:1
+ O -0.525329794 -0.050971084 -0.314516861 1
+ H -0.942006633  0.747901631  0.011252816 1
+ H  0.403696525  0.059785981 -0.073568368 1
+ O  2.316633291  0.045500849  0.071858389 2
+ H  2.684616115 -0.526576554  0.749386716 2
+ C  2.781638362 -0.426129067 -1.190300721 2
+ H  2.350821267  0.224964624 -1.943414753 2
+ H  3.867602049 -0.375336206 -1.264612649 2
+ H  2.453295744 -1.445998564 -1.389381355 2