Running with AiiDA

In [1]:

%load_ext aiida
%aiida

%load_ext aiida %aiida

Out[1]:

Loaded AiiDA DB environment - profile name: presto.

In [2]:

from xtalpaint.aiida.workgraphs.inpainting import InpaintingWorkGraph
from xtalpaint.inpainting.config_schema import XtalPaintConfig
from xtalpaint.data import BatchedStructures
from ase.io import read
from pymatgen.io.ase import AseAtomsAdaptor
from pathlib import Path

from xtalpaint.aiida.workgraphs.inpainting import InpaintingWorkGraph from xtalpaint.inpainting.config_schema import XtalPaintConfig from xtalpaint.data import BatchedStructures from ase.io import read from pymatgen.io.ase import AseAtomsAdaptor from pathlib import Path

MODELS_PROJECT_ROOT: /home/reents_t/project/test-xtalpaint/mattergen-clean/mattergen

In [3]:

structures = read("test-structures.extxyz", index=':')
structures = {
    a.info['uuid']: AseAtomsAdaptor.get_structure(a) for a in structures
}

structures = read("test-structures.extxyz", index=':') structures = { a.info['uuid']: AseAtomsAdaptor.get_structure(a) for a in structures }

In [4]:

len(structures)

len(structures)

Out[4]:

5

In [5]:

param_grid = {
    "N_steps": 5,
    "coordinates_snr": 0.2,
    "n_corrector_steps": 1,
    "batch_size": 1000,
}

param_grid = { "N_steps": 5, "coordinates_snr": 0.2, "n_corrector_steps": 1, "batch_size": 1000, }

Note on the model. This example uses our recommended TD-pos-only model — the core model of XtalPaint — with the time-dependent (TD) predictor-corrector. The XtalPaint models are downloaded automatically from Hugging Face the first time they are selected (you can also fetch one explicitly via xtalpaint.models.download_pretrained_model). See the configuration guide for the available model and predictor-corrector combinations.

In [6]:

structures_batched = BatchedStructures(
    {k.replace("-", "_"): s for k, s in structures.items()}
)

inputs = XtalPaintConfig(
    run_inpainting=True,
    candidate_generation={
        "n_inp": {
            k.replace("-", "_"): int(s.composition["H"])
            for k, s in structures.items()
        },
        "element": "H",
        "num_samples": 1,
    },
    inpainting={
        "predictor_corrector": "TD",
        **param_grid,
        "pretrained_name": "TD-pos-only",
        "sampling_config_path": f"{Path().resolve().parent}/mattergen-configs/sampling_conf",
    },
    relaxation={
        "constrained": True,
        "full": True,
        "relax_config": {
            "params": {
                "mlip": "mattersim",
                "optimizer": "BFGS",
                "load_path": "MatterSim-v1.0.0-5M.pth",
                "fmax": 0.01,
                "max_n_steps": 50,
                "max_natoms_per_batch": 5000,
                "device": "cpu",
                "elements_to_relax": ["H"],
                "return_initial_energies": False,
                "return_initial_forces": False,
                "return_final_forces": False,
            },
            "aiida": {
                "relax_code_label": "python@localhost",
                "relax_options": {
                    "resources": {},
                    "withmpi": False,
                },
            },
        },
    },
    aiida={
        "default_code_label": "python@localhost",
        "inpainting_options": {
            "resources": {},
            "withmpi": False,
        },
        "candidate_generation_options": {
            "resources": {},
            "withmpi": False,
        },
    },
)

structures_batched = BatchedStructures( {k.replace("-", "_"): s for k, s in structures.items()} ) inputs = XtalPaintConfig( run_inpainting=True, candidate_generation={ "n_inp": { k.replace("-", "_"): int(s.composition["H"]) for k, s in structures.items() }, "element": "H", "num_samples": 1, }, inpainting={ "predictor_corrector": "TD", **param_grid, "pretrained_name": "TD-pos-only", "sampling_config_path": f"{Path().resolve().parent}/mattergen-configs/sampling_conf", }, relaxation={ "constrained": True, "full": True, "relax_config": { "params": { "mlip": "mattersim", "optimizer": "BFGS", "load_path": "MatterSim-v1.0.0-5M.pth", "fmax": 0.01, "max_n_steps": 50, "max_natoms_per_batch": 5000, "device": "cpu", "elements_to_relax": ["H"], "return_initial_energies": False, "return_initial_forces": False, "return_final_forces": False, }, "aiida": { "relax_code_label": "python@localhost", "relax_options": { "resources": {}, "withmpi": False, }, }, }, }, aiida={ "default_code_label": "python@localhost", "inpainting_options": { "resources": {}, "withmpi": False, }, "candidate_generation_options": { "resources": {}, "withmpi": False, }, }, )

In [8]:

wg = InpaintingWorkGraph.build(
    structures=structures_batched,
    inputs=inputs,
)

wg = InpaintingWorkGraph.build( structures=structures_batched, inputs=inputs, )

In [9]:

wg

wg

Out[9]:

In [10]:

wg.run()

wg.run()

Out[10]:

{'inpainted_structures': <BatchedStructuresData: uuid: 9dd834f5-f796-4f9b-b22c-89ba6a1da145 (pk: 2996)>,
 'inpainting_candidates': <BatchedStructuresData: uuid: b227065b-5399-4734-8e03-50a67e05a8ea (pk: 2990)>,
 'inpainted_constrained_relaxation': {'structures': <BatchedStructuresData: uuid: 49068112-ebe6-46cb-9605-eddcee3d161e (pk: 3005)>,
  'final_energies': <PandasDataFrameData: uuid: e0d99156-ec6a-413f-b6ed-f128f8ec2d4f (pk: 3006)>},
 'pre_relaxed_inpainted_full_relaxation': {'structures': <BatchedStructuresData: uuid: 3f1ad5f4-92a0-4781-8687-99cf177e47a5 (pk: 3015)>,
  'final_energies': <PandasDataFrameData: uuid: d0c4590f-adff-4a82-81fb-af4bbe6c9404 (pk: 3016)>}}

In [ ]: