aiida-mlip/examples/calculations/submit_descriptors.py at main · stfc/aiida-mlip · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
"""Example code for submitting descriptors calculation."""

from __future__ import annotations

import ast

from aiida.common import NotExistent
from aiida.engine import run_get_node
from aiida.orm import Bool, Dict, Str, load_code
from aiida.plugins import CalculationFactory
import click

from aiida_mlip.helpers.help_load import load_model, load_structure


def descriptors(params: dict) -> None:
    """
    Prepare inputs and run a descriptors calculation.

    Parameters
    ----------
    params : dict
        A dictionary containing the input parameters for the calculations

    Returns
    -------
        None
    """
    structure = load_structure(params["struct"])

    # Select model to use
    model = load_model(params["model"], params["arch"])

    # Select calculation to use
    DescriptorsCalc = CalculationFactory("mlip.descriptors")

    # Define inputs
    inputs = {
        "metadata": {"options": {"resources": {"num_machines": 1}}},
        "code": params["code"],
        "arch": Str(params["arch"]),
        "struct": structure,
        "model": model,
        "device": Str(params["device"]),
        "invariants_only": Bool(params["invariants_only"]),
        "calc_per_element": Bool(params["calc_per_element"]),
        "calc_per_atom": Bool(params["calc_per_atom"]),
    }

    # Only calc_kwargs add if set
    if params["calc_kwargs"]:
        inputs["calc_kwargs"] = Dict(params["calc_kwargs"])

    # Run calculation
    result, node = run_get_node(DescriptorsCalc, **inputs)
    print(f"Printing results from calculation: {result}")
    print(f"Printing node of calculation: {node}")


# Arguments and options to give to the cli when running the script
@click.command("cli")
@click.argument("codelabel", type=str)
@click.option(
    "--struct",
    default=None,
    type=str,
    help="Specify the structure (aiida node or path to a structure file)",
)
@click.option(
    "--model",
    default=None,
    type=str,
    help="Specify path or URI of the model to use",
)
@click.option(
    "--arch",
    default="mace_mp",
    type=str,
    help="MLIP architecture to use for calculations.",
)
@click.option(
    "--device", default="cpu", type=str, help="Device to run calculations on."
)
@click.option(
    "--calc-kwargs",
    default="{}",
    type=str,
    help="Keyword arguments to pass to calculator.",
)
@click.option(
    "--invariants-only",
    default=False,
    type=bool,
    help="Only calculate invariant descriptors.",
)
@click.option(
    "--calc-per-element",
    default=False,
    type=bool,
    help="Calculate mean descriptors for each element.",
)
@click.option(
    "--calc-per-atom",
    default=False,
    type=bool,
    help="Calculate descriptors for each atom.",
)
def cli(
    codelabel,
    struct,
    model,
    arch,
    device,
    calc_kwargs,
    invariants_only,
    calc_per_element,
    calc_per_atom,
) -> None:
    """Click interface."""
    calc_kwargs = ast.literal_eval(calc_kwargs)

    try:
        code = load_code(codelabel)
    except NotExistent as exc:
        print(f"The code '{codelabel}' does not exist.")
        raise SystemExit from exc

    params = {
        "code": code,
        "struct": struct,
        "model": model,
        "arch": arch,
        "device": device,
        "calc_kwargs": calc_kwargs,
        "invariants_only": invariants_only,
        "calc_per_element": calc_per_element,
        "calc_per_atom": calc_per_atom,
    }

    # Submit descriptors
    descriptors(params)


if __name__ == "__main__":
    cli()