Source code for scm.pisa.input_def

#!/usr/bin/env amspython

from __future__ import annotations

import json
import os
import shutil
import site
from pathlib import Path

from scm.pisa import block as input_block
from scm.pisa import key as input_key
from scm.pisa.block_subclass_string import BlockSubclassAttribute, BlockSubclassString
from scm.pisa.utils import to_valid_identifier

INPUT_DEF_FOLDER = Path(os.environ["AMSBIN"]) / "input_def"
# these dictionaries control which files in INPUT_DEF_FOLDER get turned into python class and the values determine
# the name of the classes.
ENGINE_BLOCK_FILES = {
    INPUT_DEF_FOLDER / "adf.json": "ADF",
    INPUT_DEF_FOLDER / "ase.json": "ASE",
    INPUT_DEF_FOLDER / "band.json": "BAND",
    INPUT_DEF_FOLDER / "dftb.json": "DFTB",
    INPUT_DEF_FOLDER / "external.json": "External",
    INPUT_DEF_FOLDER / "forcefield.json": "ForceField",
    INPUT_DEF_FOLDER / "gfnff.json": "GFNFF",
    INPUT_DEF_FOLDER / "hybrid.json": "Hybrid",
    INPUT_DEF_FOLDER / "lennardjones.json": "LennardJones",
    INPUT_DEF_FOLDER / "mlpotential.json": "MLPotential",
    INPUT_DEF_FOLDER / "mopac.json": "Mopac",
    INPUT_DEF_FOLDER / "pipe.json": "Pipe",
    INPUT_DEF_FOLDER / "reaxff.json": "ReaxFF",
    INPUT_DEF_FOLDER / "quantumespresso.json": "QuantumESPRESSO",
    INPUT_DEF_FOLDER / "vasp.json": "VASP",
}
DRIVER_BLOCK_FILES = {
    INPUT_DEF_FOLDER / "acerxn.json": "ACERXN",
    INPUT_DEF_FOLDER / "adfnbo.json": "ADFNBO",
    INPUT_DEF_FOLDER / "amsbatch.json": "AMSBatch",
    INPUT_DEF_FOLDER / "ams_interactive.json": "AMSInteractive",
    INPUT_DEF_FOLDER / "ams.json": "AMS",
    INPUT_DEF_FOLDER / "analysis.json": "Analysis",
    INPUT_DEF_FOLDER / "atomtyping.json": "AtomTyping",
    INPUT_DEF_FOLDER / "chemtrayzer2.json": "Chemtrayzer2",
    INPUT_DEF_FOLDER / "conductance.json": "Conductance",
    INPUT_DEF_FOLDER / "conformers.json": "Conformers",
    INPUT_DEF_FOLDER / "cpl.json": "CPL",
    # INPUT_DEF_FOLDER / "crs.json",  TODO code breaks running on this file: ''?
    INPUT_DEF_FOLDER / "densf.json": "DENSF",
    INPUT_DEF_FOLDER / "dftbplus.json": "DFTBPlus",
    INPUT_DEF_FOLDER / "epr.json": "EPR",
    INPUT_DEF_FOLDER / "fcf.json": "FCF",
    INPUT_DEF_FOLDER / "glompo_genref.json": "ParAMSGenerateReference",
    INPUT_DEF_FOLDER / "glompo.json": "ParAMS",
    INPUT_DEF_FOLDER / "glompo_machinelearning.json": "ParAMSMachineLearning",
    INPUT_DEF_FOLDER / "glompo_optimize.json": "ParAMSOptimize",
    INPUT_DEF_FOLDER / "glompo_sensitivity.json": "ParAMSSensitivity",
    INPUT_DEF_FOLDER / "glompo_singlepoint.json": "ParAMSSinglePoint",
    INPUT_DEF_FOLDER / "green.json": "Green",
    INPUT_DEF_FOLDER / "lfdft.json": "LFDFT",
    INPUT_DEF_FOLDER / "lfdft_tdm.json": "LFDFT_TDM",
    INPUT_DEF_FOLDER / "nao.json": "NAO",
    INPUT_DEF_FOLDER / "nmr.json": "NMR",
    INPUT_DEF_FOLDER / "oldfcf.json": "OldFCF",
    INPUT_DEF_FOLDER / "oled-deposition.json": "OLEDDepostion",
    INPUT_DEF_FOLDER / "oled-properties.json": "OLEDProperties",
    INPUT_DEF_FOLDER / "pre_amsified_adf.json": "PreAmsifiedADF",
    INPUT_DEF_FOLDER / "reactions_discovery.json": "ReactionsDiscovery",
    INPUT_DEF_FOLDER / "redox_potential.json": "RedoxPotential",
    INPUT_DEF_FOLDER / "sgf.json": "SGF",
    INPUT_DEF_FOLDER / "symmetrize.json": "Symmetrize",
    INPUT_DEF_FOLDER / "UnifiedChemicalSystem.json": "UnifiedChemicalSystem",
    INPUT_DEF_FOLDER / "simple_active_learning.json": "SimpleActiveLearning",
}


[docs]def block_to_cls_str(block: input_block.Block, insert_underscore: bool = True) -> BlockSubclassString:
    """Generate a :class:`BlockSubclassString` object from a :class:`Block` object. Uses the __repr__ functionality to
    generate a line of python for every key/block attribute of the Block to convert.

    :param block: Block to convert
    :type block: input_block.Block
    :param insert_underscore: Whether to prefix the class name with '_'. Should be true except for the toplevel block. Used to prevent name clashes between the nested class definition and the instance attribute. Defaults to True
    :type insert_underscore: bool, optional
    :return: Class string object that can be used to generate a string containing python code defining the subclass that represents the block.
    :rtype: BlockSubclassString
    """
    attributes: list[BlockSubclassAttribute] = []
    for child_key in block.keys:
        cls_name, init_args = repr(child_key).split("(", 1)
        attributes.append(
            BlockSubclassAttribute(
                child_key.name,
                f"{to_valid_identifier(cls_name)}({init_args}",
                input_key.type_hint(child_key),
            )
        )
    for child_block in block.blocks:
        # the block __repr__ from dataclass will have the block class
        # but we need to overwrite it with the specific nested class name
        # e.g. from FixedBlock(name='_PerAtomType', ...)
        # to self.__PerAtomType(...)
        _, init_args = repr(child_block).split("(", 1)
        # nested class definitions need to be prefixed with '_'
        # to prevent name clashes with the instance attributes of the same name
        attributes.append(
            BlockSubclassAttribute(
                child_block.name,
                f"self._{to_valid_identifier(child_block.name)}({init_args}",
                input_block.type_hint(child_block),
            )
        )

    cls_str = BlockSubclassString(
        name=f"_{to_valid_identifier(block.name)}" if insert_underscore else to_valid_identifier(block.name),
        parent=block.__class__.__name__,
        attributes=attributes,
        doc_str=block.comment,
    )
    for child_block in block.blocks:
        block_cls_str = block_to_cls_str(child_block)
        cls_str.insert(block_cls_str)
    return cls_str


[docs]def write_block_class_to_file(block: input_block.Block, path: Path) -> None:
    """Generates python code that defines a specific block and write it to a file. Will insert the necessary imports as
    well.

    :param block: Block object to write to file
    :type block: input_block.Block
    :param path: Path to write the file to.
    :type path: Path
    """
    cls_str = block_to_cls_str(block, insert_underscore=False)
    # TODO maybe this could be made less fragile?
    # but as long as it is covered by tests, should be okay
    import_line_block = "from scm.pisa.block import "
    import_line_key = "from scm.pisa.key import "
    block_classes = []
    key_classes = []
    # perhaps just hardcoding the string is more readable?
    # there will not be any new block or key types often anyway
    for attr in dir(input_block):
        attr = getattr(input_block, attr)
        if isinstance(attr, type) and issubclass(attr, input_block.Block) and attr != input_block.Block:
            block_classes.append(attr.__name__)
    for attr in dir(input_key):
        attr = getattr(input_key, attr)
        if isinstance(attr, type) and issubclass(attr, input_key.Key) and attr != input_key.Key:
            key_classes.append(attr.__name__)
    key_classes += ["BoolType"]
    with open(path, "w", encoding="utf-8") as f:
        f.write("from __future__ import annotations\n")
        f.write("from pathlib import Path\n")
        f.write("from typing import Iterable, Literal, Sequence\n")
        f.write(import_line_block + ",".join(block_classes) + "\n")
        f.write(import_line_key + ",".join(key_classes) + "\n\n")
        f.write(str(cls_str))


[docs]def create_input_classes_module(destination_dir: Path, overwrite: bool = False) -> None:
    """Convert all driver and engine block json definitions to python object, and write python code that define these
    objects to files. Also inserts __init__.py files to turn `destination_dir` into a python module.

    Takes all json files defined in  module attributes :attr:`ENGINE_BLOCK_FILES` and :attr:`DRIVER_BLOCK_FILES`
    as input.

    :param destination_dir: Folder to write the python module to
    :type destination_dir: Path
    :param overwrite: Delete the directory if it already exists, defaults to False
    :type overwrite: bool, optional
    :raises Exception: When `overwrite` is set to `False`, but directory already exists.
    """

    if destination_dir.exists():
        if overwrite:
            if destination_dir.is_dir():
                shutil.rmtree(destination_dir)
            else:
                destination_dir.unlink()
        else:
            raise Exception(f"Target directory {destination_dir} exists and overwrite set to False")
    destination_dir.mkdir()
    drivers_dir = destination_dir / "drivers"
    drivers_dir.mkdir()
    engines_dir = destination_dir / "engines"
    engines_dir.mkdir()
    assert "AMSHOME" in os.environ
    blocks: list[input_block.Block] = []
    for fn in INPUT_DEF_FOLDER.iterdir():
        if fn in ENGINE_BLOCK_FILES:
            block_type = "engine"
            cls_name = ENGINE_BLOCK_FILES[fn]
        elif fn in DRIVER_BLOCK_FILES:
            block_type = "driver"
            cls_name = DRIVER_BLOCK_FILES[fn]
        else:
            continue
        with open(fn, "r", encoding="utf-8") as f:
            def_dict = json.load(f)
        def_dict["_category"] = "block"
        def_dict["_type"] = block_type
        def_dict["_comment"] = ""
        blocks.append(input_block.Block.from_def_dict(name=cls_name, def_dict=def_dict))
    blocks = sorted(blocks, key=lambda b: b.name)
    for b in blocks:
        module_dir = drivers_dir if isinstance(b, input_block.DriverBlock) else engines_dir
        write_block_class_to_file(b, module_dir / f"{to_valid_identifier(b.name.lower())}.py")
    # generate __init__.py
    driver_blocks = [b for b in blocks if isinstance(b, input_block.DriverBlock)]
    engine_blocks = [b for b in blocks if isinstance(b, input_block.EngineBlock)]
    with open(destination_dir / "__init__.py", "w", encoding="utf-8") as f:
        f.write("from . import drivers\n")
        f.write("from . import engines\n")
        for b in driver_blocks:
            cls_name = to_valid_identifier(b.name)
            f.write(f"from .drivers import {cls_name}\n")
        for b in engine_blocks:
            cls_name = to_valid_identifier(b.name)
            f.write(f"from .engines import {cls_name}\n")
        f.write("\n")
        f.write(f"__all__ = {[to_valid_identifier(block.name) for block in blocks]}")
    with open(drivers_dir / "__init__.py", "w", encoding="utf-8") as f:
        for b in driver_blocks:
            cls_name = to_valid_identifier(b.name)
            f.write(f"from .{cls_name.lower()} import {cls_name}\n")
        f.write("\n")
        f.write(f"__all__ = {[to_valid_identifier(block.name) for block in driver_blocks]}")
    with open(engines_dir / "__init__.py", "w", encoding="utf-8") as f:
        for b in engine_blocks:
            cls_name = to_valid_identifier(b.name)
            f.write(f"from .{cls_name.lower()} import {cls_name}\n")
        f.write("\n")
        f.write(f"__all__ = {[to_valid_identifier(block.name) for block in engine_blocks]}")


if __name__ == "__main__":
    # site packages directory differs between windows and linux platforms, better to get it programmaticallysvn st
    # we try to avoid the directory in the scm virtual environment
    # on linux this typically looks like: $AMSBIN/python3.8/lib/python3.8/site-packages
    # and on windows it looks like: $AMSBIN/python3.8/lib/site-packages
    site_packages_dir = Path([path for path in site.getsitepackages() if "venv" not in path][-1])
    create_input_classes_module(destination_dir=site_packages_dir / "scm" / "input_classes", overwrite=True)