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 setup_inpainting_wg
from xtalpaint.inpainting.config_schema import InpaintingWorkGraphConfig
from xtalpaint.aiida.data import BatchedStructures
from ase.io import read
from pymatgen.io.ase import AseAtomsAdaptor

from xtalpaint.aiida.workgraphs.inpainting import setup_inpainting_wg from xtalpaint.inpainting.config_schema import InpaintingWorkGraphConfig from xtalpaint.aiida.data import BatchedStructures from ase.io import read from pymatgen.io.ase import AseAtomsAdaptor

MODELS_PROJECT_ROOT: /home/reents_t/project/dev-mattergen-inpainting/git/mattergen/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, }

In [6]:

ENV_ACTIVATION_CMD = "source ~/.aiida_venvs/dev-mattergen-inpainting/bin/activate"

inputs = InpaintingWorkGraphConfig(
    inpainting_pipeline_params={
        "predictor_corrector": "baseline",
        "inpainting_model_params": param_grid,
        "pretrained_name": "mattergen_base",
        "sampling_config_path": "/home/reents_t/project/dev-mattergen-inpainting/git/mattergen/sampling_conf",
    },
    structures=BatchedStructures(
        {k.replace("-", "_"): s for k, s in structures.items()}
    ),
    gen_inpainting_candidates_params={
        "n_inp": {
            k.replace("-", "_"): int(s.composition["H"])
            for k, s in structures.items()
        },
        "element": "H",
        "num_samples": 1,
    },
    relax=True,
    full_relax=True,
    full_relax_wo_pre_relax=False,
    run_inpainting=True,
    relax_kwargs={
        "elements_to_relax": ["H"],
        "fmax": 0.01,
        "max_natoms_per_batch": 5000,
        "load_path": "MatterSim-v1.0.0-5M.pth",
        "max_n_steps": 50,
        "device": "cuda",
        "mlip": "mattersim",
        "optimizer": "BFGS",
        "return_initial_energies": False,
        "return_initial_forces": False,
        "return_final_forces": False,
    },
    gen_inpainting_candidates_options={
        "custom_scheduler_commands": f"{ENV_ACTIVATION_CMD}\nexport PYTHONBREAKPOINT=0",
    },
    options={
        "prepend_text": f"{ENV_ACTIVATION_CMD}\nexport PYTHONBREAKPOINT=0",
    },
    evaluate_params={"max_workers": 5, "metrics": ["match", "rmsd"]},
    evaluate=True,
)

ENV_ACTIVATION_CMD = "source ~/.aiida_venvs/dev-mattergen-inpainting/bin/activate" inputs = InpaintingWorkGraphConfig( inpainting_pipeline_params={ "predictor_corrector": "baseline", "inpainting_model_params": param_grid, "pretrained_name": "mattergen_base", "sampling_config_path": "/home/reents_t/project/dev-mattergen-inpainting/git/mattergen/sampling_conf", }, structures=BatchedStructures( {k.replace("-", "_"): s for k, s in structures.items()} ), gen_inpainting_candidates_params={ "n_inp": { k.replace("-", "_"): int(s.composition["H"]) for k, s in structures.items() }, "element": "H", "num_samples": 1, }, relax=True, full_relax=True, full_relax_wo_pre_relax=False, run_inpainting=True, relax_kwargs={ "elements_to_relax": ["H"], "fmax": 0.01, "max_natoms_per_batch": 5000, "load_path": "MatterSim-v1.0.0-5M.pth", "max_n_steps": 50, "device": "cuda", "mlip": "mattersim", "optimizer": "BFGS", "return_initial_energies": False, "return_initial_forces": False, "return_final_forces": False, }, gen_inpainting_candidates_options={ "custom_scheduler_commands": f"{ENV_ACTIVATION_CMD}\nexport PYTHONBREAKPOINT=0", }, options={ "prepend_text": f"{ENV_ACTIVATION_CMD}\nexport PYTHONBREAKPOINT=0", }, evaluate_params={"max_workers": 5, "metrics": ["match", "rmsd"]}, evaluate=True, )

In [7]:

wg = setup_inpainting_wg(
    inputs
)

wg = setup_inpainting_wg( inputs )

In [8]:

wg

wg

Out[8]:

NodeGraphWidget(settings={'minimap': True}, style={'width': '90%', 'height': '600px'}, value={'name': 'WorkGra…

In [9]:

wg.run()

wg.run()

01/07/2026 03:51:49 PM <183876> aiida.broker.rabbitmq: [WARNING] RabbitMQ v3.9.27 is not supported and will cause unexpected problems!
WARNING:aiida.broker.rabbitmq:RabbitMQ v3.9.27 is not supported and will cause unexpected problems!
01/07/2026 03:51:49 PM <183876> aiida.broker.rabbitmq: [WARNING] It can cause long-running workflows to crash and jobs to be submitted multiple times.
WARNING:aiida.broker.rabbitmq:It can cause long-running workflows to crash and jobs to be submitted multiple times.
01/07/2026 03:51:49 PM <183876> aiida.broker.rabbitmq: [WARNING] See https://github.com/aiidateam/aiida-core/wiki/RabbitMQ-version-to-use for details.
WARNING:aiida.broker.rabbitmq:See https://github.com/aiidateam/aiida-core/wiki/RabbitMQ-version-to-use for details.
01/07/2026 03:51:59 PM <183876> aiida.orm.nodes.process.workflow.workchain.WorkChainNode: [REPORT] [250993|WorkGraphEngine|continue_workgraph]: tasks ready to run: generate_inpainting_candidates
REPORT:aiida.orm.nodes.process.workflow.workchain.WorkChainNode:[250993|WorkGraphEngine|continue_workgraph]: tasks ready to run: generate_inpainting_candidates
01/07/2026 03:52:02 PM <183876> aiida.orm.nodes.process.workflow.workchain.WorkChainNode: [REPORT] [250993|WorkGraphEngine|on_wait]: Process status: Waiting for child processes: 250997
REPORT:aiida.orm.nodes.process.workflow.workchain.WorkChainNode:[250993|WorkGraphEngine|on_wait]: Process status: Waiting for child processes: 250997
01/07/2026 03:52:09 PM <183876> aiida.orm.nodes.process.workflow.workchain.WorkChainNode: [REPORT] [250993|WorkGraphEngine|update_task_state]: Task: generate_inpainting_candidates, type: PythonJob, finished.
REPORT:aiida.orm.nodes.process.workflow.workchain.WorkChainNode:[250993|WorkGraphEngine|update_task_state]: Task: generate_inpainting_candidates, type: PythonJob, finished.
01/07/2026 03:52:12 PM <183876> aiida.orm.nodes.process.workflow.workchain.WorkChainNode: [REPORT] [250993|WorkGraphEngine|continue_workgraph]: tasks ready to run: inpainting
REPORT:aiida.orm.nodes.process.workflow.workchain.WorkChainNode:[250993|WorkGraphEngine|continue_workgraph]: tasks ready to run: inpainting
01/07/2026 03:52:14 PM <183876> aiida.orm.nodes.process.workflow.workchain.WorkChainNode: [REPORT] [250993|WorkGraphEngine|on_wait]: Process status: Waiting for child processes: 251006
REPORT:aiida.orm.nodes.process.workflow.workchain.WorkChainNode:[250993|WorkGraphEngine|on_wait]: Process status: Waiting for child processes: 251006
01/07/2026 03:52:38 PM <183876> aiida.orm.nodes.process.workflow.workchain.WorkChainNode: [REPORT] [250993|WorkGraphEngine|update_task_state]: Task: inpainting, type: PythonJob, finished.
REPORT:aiida.orm.nodes.process.workflow.workchain.WorkChainNode:[250993|WorkGraphEngine|update_task_state]: Task: inpainting, type: PythonJob, finished.
01/07/2026 03:52:41 PM <183876> aiida.orm.nodes.process.workflow.workchain.WorkChainNode: [REPORT] [250993|WorkGraphEngine|continue_workgraph]: tasks ready to run: inpainted_constrained_relaxation,evaluate_inpainting_rmsd_inpainting,evaluate_inpainting_match_inpainting
REPORT:aiida.orm.nodes.process.workflow.workchain.WorkChainNode:[250993|WorkGraphEngine|continue_workgraph]: tasks ready to run: inpainted_constrained_relaxation,evaluate_inpainting_rmsd_inpainting,evaluate_inpainting_match_inpainting
01/07/2026 03:52:46 PM <183876> aiida.orm.nodes.process.workflow.workchain.WorkChainNode: [REPORT] [250993|WorkGraphEngine|on_wait]: Process status: Waiting for child processes: 251015, 251019, 251023
REPORT:aiida.orm.nodes.process.workflow.workchain.WorkChainNode:[250993|WorkGraphEngine|on_wait]: Process status: Waiting for child processes: 251015, 251019, 251023
/home/reents_t/project/dev-mattergen-inpainting/git/dbcsi-cleanup/src/xtalpaint/aiida/data.py:315: AiidaDeprecationWarning: `PandasDataFrameData.attributes` is deprecated, use `PandasDataFrameData.base.attributes.all` instead. (this will be removed in v3)
  handle, self.attributes.get("filename")
/home/reents_t/project/dev-mattergen-inpainting/git/dbcsi-cleanup/src/xtalpaint/aiida/data.py:315: AiidaDeprecationWarning: `PandasDataFrameData.attributes` is deprecated, use `PandasDataFrameData.base.attributes.all` instead. (this will be removed in v3)
  handle, self.attributes.get("filename")
01/07/2026 03:52:57 PM <183876> aiida.orm.nodes.process.workflow.workchain.WorkChainNode: [REPORT] [250993|WorkGraphEngine|update_task_state]: Task: evaluate_inpainting_rmsd_inpainting, type: PythonJob, finished.
REPORT:aiida.orm.nodes.process.workflow.workchain.WorkChainNode:[250993|WorkGraphEngine|update_task_state]: Task: evaluate_inpainting_rmsd_inpainting, type: PythonJob, finished.
01/07/2026 03:52:58 PM <183876> aiida.orm.nodes.process.workflow.workchain.WorkChainNode: [REPORT] [250993|WorkGraphEngine|update_task_state]: Task: evaluate_inpainting_match_inpainting, type: PythonJob, finished.
REPORT:aiida.orm.nodes.process.workflow.workchain.WorkChainNode:[250993|WorkGraphEngine|update_task_state]: Task: evaluate_inpainting_match_inpainting, type: PythonJob, finished.
01/07/2026 03:53:01 PM <183876> aiida.orm.nodes.process.workflow.workchain.WorkChainNode: [REPORT] [250993|WorkGraphEngine|continue_workgraph]: tasks ready to run: 
REPORT:aiida.orm.nodes.process.workflow.workchain.WorkChainNode:[250993|WorkGraphEngine|continue_workgraph]: tasks ready to run: 
01/07/2026 03:53:01 PM <183876> aiida.orm.nodes.process.workflow.workchain.WorkChainNode: [REPORT] [250993|WorkGraphEngine|on_wait]: Process status: Waiting for child processes: 251015
REPORT:aiida.orm.nodes.process.workflow.workchain.WorkChainNode:[250993|WorkGraphEngine|on_wait]: Process status: Waiting for child processes: 251015
/home/reents_t/project/dev-mattergen-inpainting/git/dbcsi-cleanup/src/xtalpaint/aiida/data.py:315: AiidaDeprecationWarning: `PandasDataFrameData.attributes` is deprecated, use `PandasDataFrameData.base.attributes.all` instead. (this will be removed in v3)
  handle, self.attributes.get("filename")
01/07/2026 03:53:26 PM <183876> aiida.orm.nodes.process.workflow.workchain.WorkChainNode: [REPORT] [250993|WorkGraphEngine|update_task_state]: Task: inpainted_constrained_relaxation, type: PythonJob, finished.
REPORT:aiida.orm.nodes.process.workflow.workchain.WorkChainNode:[250993|WorkGraphEngine|update_task_state]: Task: inpainted_constrained_relaxation, type: PythonJob, finished.
01/07/2026 03:53:29 PM <183876> aiida.orm.nodes.process.workflow.workchain.WorkChainNode: [REPORT] [250993|WorkGraphEngine|continue_workgraph]: tasks ready to run: pre_relaxed_inpainted_full_relaxation,evaluate_inpainting_rmsd_inpainted_constrained_relaxation,evaluate_inpainting_match_inpainted_constrained_relaxation
REPORT:aiida.orm.nodes.process.workflow.workchain.WorkChainNode:[250993|WorkGraphEngine|continue_workgraph]: tasks ready to run: pre_relaxed_inpainted_full_relaxation,evaluate_inpainting_rmsd_inpainted_constrained_relaxation,evaluate_inpainting_match_inpainted_constrained_relaxation
01/07/2026 03:53:34 PM <183876> aiida.orm.nodes.process.workflow.workchain.WorkChainNode: [REPORT] [250993|WorkGraphEngine|on_wait]: Process status: Waiting for child processes: 251041, 251045, 251049
REPORT:aiida.orm.nodes.process.workflow.workchain.WorkChainNode:[250993|WorkGraphEngine|on_wait]: Process status: Waiting for child processes: 251041, 251045, 251049
/home/reents_t/project/dev-mattergen-inpainting/git/dbcsi-cleanup/src/xtalpaint/aiida/data.py:315: AiidaDeprecationWarning: `PandasDataFrameData.attributes` is deprecated, use `PandasDataFrameData.base.attributes.all` instead. (this will be removed in v3)
  handle, self.attributes.get("filename")
/home/reents_t/project/dev-mattergen-inpainting/git/dbcsi-cleanup/src/xtalpaint/aiida/data.py:315: AiidaDeprecationWarning: `PandasDataFrameData.attributes` is deprecated, use `PandasDataFrameData.base.attributes.all` instead. (this will be removed in v3)
  handle, self.attributes.get("filename")
01/07/2026 03:53:45 PM <183876> aiida.orm.nodes.process.workflow.workchain.WorkChainNode: [REPORT] [250993|WorkGraphEngine|update_task_state]: Task: evaluate_inpainting_rmsd_inpainted_constrained_relaxation, type: PythonJob, finished.
REPORT:aiida.orm.nodes.process.workflow.workchain.WorkChainNode:[250993|WorkGraphEngine|update_task_state]: Task: evaluate_inpainting_rmsd_inpainted_constrained_relaxation, type: PythonJob, finished.
01/07/2026 03:53:46 PM <183876> aiida.orm.nodes.process.workflow.workchain.WorkChainNode: [REPORT] [250993|WorkGraphEngine|update_task_state]: Task: evaluate_inpainting_match_inpainted_constrained_relaxation, type: PythonJob, finished.
REPORT:aiida.orm.nodes.process.workflow.workchain.WorkChainNode:[250993|WorkGraphEngine|update_task_state]: Task: evaluate_inpainting_match_inpainted_constrained_relaxation, type: PythonJob, finished.
01/07/2026 03:53:49 PM <183876> aiida.orm.nodes.process.workflow.workchain.WorkChainNode: [REPORT] [250993|WorkGraphEngine|continue_workgraph]: tasks ready to run: 
REPORT:aiida.orm.nodes.process.workflow.workchain.WorkChainNode:[250993|WorkGraphEngine|continue_workgraph]: tasks ready to run: 
01/07/2026 03:53:50 PM <183876> aiida.orm.nodes.process.workflow.workchain.WorkChainNode: [REPORT] [250993|WorkGraphEngine|on_wait]: Process status: Waiting for child processes: 251041
REPORT:aiida.orm.nodes.process.workflow.workchain.WorkChainNode:[250993|WorkGraphEngine|on_wait]: Process status: Waiting for child processes: 251041
/home/reents_t/project/dev-mattergen-inpainting/git/dbcsi-cleanup/src/xtalpaint/aiida/data.py:315: AiidaDeprecationWarning: `PandasDataFrameData.attributes` is deprecated, use `PandasDataFrameData.base.attributes.all` instead. (this will be removed in v3)
  handle, self.attributes.get("filename")
01/07/2026 03:54:04 PM <183876> aiida.orm.nodes.process.workflow.workchain.WorkChainNode: [REPORT] [250993|WorkGraphEngine|update_task_state]: Task: pre_relaxed_inpainted_full_relaxation, type: PythonJob, finished.
REPORT:aiida.orm.nodes.process.workflow.workchain.WorkChainNode:[250993|WorkGraphEngine|update_task_state]: Task: pre_relaxed_inpainted_full_relaxation, type: PythonJob, finished.
01/07/2026 03:54:07 PM <183876> aiida.orm.nodes.process.workflow.workchain.WorkChainNode: [REPORT] [250993|WorkGraphEngine|continue_workgraph]: tasks ready to run: evaluate_inpainting_rmsd_pre_relaxed_inpainted_full_relaxation,evaluate_inpainting_match_pre_relaxed_inpainted_full_relaxation
REPORT:aiida.orm.nodes.process.workflow.workchain.WorkChainNode:[250993|WorkGraphEngine|continue_workgraph]: tasks ready to run: evaluate_inpainting_rmsd_pre_relaxed_inpainted_full_relaxation,evaluate_inpainting_match_pre_relaxed_inpainted_full_relaxation
01/07/2026 03:54:12 PM <183876> aiida.orm.nodes.process.workflow.workchain.WorkChainNode: [REPORT] [250993|WorkGraphEngine|on_wait]: Process status: Waiting for child processes: 251066, 251070
REPORT:aiida.orm.nodes.process.workflow.workchain.WorkChainNode:[250993|WorkGraphEngine|on_wait]: Process status: Waiting for child processes: 251066, 251070
/home/reents_t/project/dev-mattergen-inpainting/git/dbcsi-cleanup/src/xtalpaint/aiida/data.py:315: AiidaDeprecationWarning: `PandasDataFrameData.attributes` is deprecated, use `PandasDataFrameData.base.attributes.all` instead. (this will be removed in v3)
  handle, self.attributes.get("filename")
/home/reents_t/project/dev-mattergen-inpainting/git/dbcsi-cleanup/src/xtalpaint/aiida/data.py:315: AiidaDeprecationWarning: `PandasDataFrameData.attributes` is deprecated, use `PandasDataFrameData.base.attributes.all` instead. (this will be removed in v3)
  handle, self.attributes.get("filename")
01/07/2026 03:54:21 PM <183876> aiida.orm.nodes.process.workflow.workchain.WorkChainNode: [REPORT] [250993|WorkGraphEngine|update_task_state]: Task: evaluate_inpainting_rmsd_pre_relaxed_inpainted_full_relaxation, type: PythonJob, finished.
REPORT:aiida.orm.nodes.process.workflow.workchain.WorkChainNode:[250993|WorkGraphEngine|update_task_state]: Task: evaluate_inpainting_rmsd_pre_relaxed_inpainted_full_relaxation, type: PythonJob, finished.
01/07/2026 03:54:22 PM <183876> aiida.orm.nodes.process.workflow.workchain.WorkChainNode: [REPORT] [250993|WorkGraphEngine|update_task_state]: Task: evaluate_inpainting_match_pre_relaxed_inpainted_full_relaxation, type: PythonJob, finished.
REPORT:aiida.orm.nodes.process.workflow.workchain.WorkChainNode:[250993|WorkGraphEngine|update_task_state]: Task: evaluate_inpainting_match_pre_relaxed_inpainted_full_relaxation, type: PythonJob, finished.
01/07/2026 03:54:25 PM <183876> aiida.orm.nodes.process.workflow.workchain.WorkChainNode: [REPORT] [250993|WorkGraphEngine|continue_workgraph]: tasks ready to run: 
REPORT:aiida.orm.nodes.process.workflow.workchain.WorkChainNode:[250993|WorkGraphEngine|continue_workgraph]: tasks ready to run: 
01/07/2026 03:54:26 PM <183876> aiida.orm.nodes.process.workflow.workchain.WorkChainNode: [REPORT] [250993|WorkGraphEngine|finalize]: Finalize workgraph.
REPORT:aiida.orm.nodes.process.workflow.workchain.WorkChainNode:[250993|WorkGraphEngine|finalize]: Finalize workgraph.

In [ ]: