diff --git a/README.md b/README.md index 71db181..4c3c3bb 100644 --- a/README.md +++ b/README.md @@ -25,7 +25,31 @@ uv pip install .[aiida] ### Model checkpoints for retrained versions of MatterGen -Model checkpoints for the retrained versions of MatterGen used in our work can be downloaded from [Hugging Face](https://huggingface.co/t-reents/XtalPaint). Currently, the repository contains the `pos-only` and `TD-pos-only` models discussed in the paper. +Model checkpoints for the retrained versions of MatterGen used in our work are hosted on [Hugging Face](https://huggingface.co/t-reents/XtalPaint). Currently, the repository contains the `pos-only` and `TD-pos-only` models discussed in the paper. + +> [!IMPORTANT] +> **The example below uses the `TD-pos-only` model — the core model of XtalPaint — which we recommend for accurate inpainting**, together with the time-dependent (`TD`) predictor-corrector. See the [documentation](https://psi-lms.github.io/XtalPaint/configuration/) for the available model and predictor-corrector combinations. + +Just like MatterGen's own checkpoints, the XtalPaint models are **downloaded automatically** the first time you select them by name — simply set `pretrained_name` to `"TD-pos-only"` (or `"pos-only"`): + +```python +config = InpaintingConfig( + pretrained_name="TD-pos-only", # auto-downloaded & cached from Hugging Face + predictor_corrector="TD", # time-dependent predictor-corrector + N_steps=50, + coordinates_snr=0.2, + n_corrector_steps=1, + batch_size=100, +) +``` + +You can also download a checkpoint explicitly and point `model_path` at it: + +```python +from xtalpaint.models import download_pretrained_model + +model_path = download_pretrained_model("TD-pos-only") +``` ## Example @@ -65,10 +89,13 @@ from xtalpaint.inpainting.inpainting_process import ( run_inpainting_pipeline, ) -# Parameters for the inpainting, please adjust to reasonable values +# Parameters for the inpainting, please adjust to reasonable values. +# This uses our recommended `TD-pos-only` model, which is downloaded +# automatically from Hugging Face the first time it is selected. See the +# documentation for the available model/predictor-corrector combinations. config = InpaintingConfig( - pretrained_name="mattergen_base", - predictor_corrector="baseline", + pretrained_name="TD-pos-only", + predictor_corrector="TD", N_steps=50, coordinates_snr=0.2, n_corrector_steps=1, diff --git a/docs/configuration.md b/docs/configuration.md index abc7b9c..9e20cda 100644 --- a/docs/configuration.md +++ b/docs/configuration.md @@ -111,11 +111,11 @@ The core diffusion stage. All sampling parameters live in one flat block. ```python inpainting=InpaintingConfig( # Model — provide exactly one of these: - pretrained_name="mattergen_base", # use a bundled pretrained checkpoint - # model_path="/path/to/checkpoint", # or point to a local file + pretrained_name="TD-pos-only", # XtalPaint model (auto-downloaded) or a MatterGen checkpoint + # model_path="/path/to/model_dir", # or point to a local checkpoint directory # Sampling - predictor_corrector="baseline", # see supported keys below + predictor_corrector="TD", # see supported combinations below N_steps=5, coordinates_snr=0.2, n_corrector_steps=1, @@ -128,7 +128,27 @@ inpainting=InpaintingConfig( ) ``` -**Supported `predictor_corrector` values:** +!!! tip "Recommended model: `TD-pos-only`" + `pretrained_name` accepts the XtalPaint models hosted on + [Hugging Face](https://huggingface.co/t-reents/XtalPaint) (`TD-pos-only`, + `pos-only`) as well as MatterGen's own checkpoints. The XtalPaint models are + downloaded automatically and cached the first time they are selected. For + accurate inpainting we **recommend the `TD-pos-only` model** — the core + model of XtalPaint — with the time-dependent (`TD`) predictor-corrector. + + Alternatively, download a checkpoint explicitly and pass it as `model_path`: + + ```python + from xtalpaint.models import download_pretrained_model + + model_path = download_pretrained_model("TD-pos-only") + ``` + + Selecting `predictor_corrector="TD"` without the `TD-pos-only` model (or + pointing `model_path` at a checkpoint that has not been downloaded) raises + an error with these instructions. + +**Supported `predictor_corrector` values** — the `TD` variant requires the `TD-pos-only` model; the others are used with the `pos-only` model or a MatterGen checkpoint: | Key | Description | |---|---| @@ -137,7 +157,7 @@ inpainting=InpaintingConfig( | `baseline-store-scores` | Records score function outputs | | `repaint-v1` | RePaint resampling (legacy) | | `repaint-v2` | RePaint resampling (v2) | -| `TD` | Time-dependent (TD-Paint) variant | +| `TD` | Time-dependent (TD-Paint) variant — requires the `TD-pos-only` model | !!! note "Repaint variants" When using `repaint-v1` or `repaint-v2`, you must also set `n_resample_steps` and `jump_length`: @@ -245,8 +265,8 @@ from xtalpaint.utils.relaxation_utils import relax_structures config = XtalPaintConfig( candidate_generation=CandidateGenerationConfig(n_inp=2, element="H"), inpainting=InpaintingConfig( - model_path="/path/to/checkpoint.ckpt", - predictor_corrector="baseline", + pretrained_name="TD-pos-only", # auto-downloaded from Hugging Face + predictor_corrector="TD", N_steps=5, coordinates_snr=0.2, n_corrector_steps=1, @@ -376,8 +396,8 @@ wg.submit() element="H", ), inpainting=InpaintingConfig( - pretrained_name="mattergen_base", - predictor_corrector="baseline", + pretrained_name="TD-pos-only", + predictor_corrector="TD", N_steps=5, coordinates_snr=0.2, n_corrector_steps=1, @@ -437,8 +457,8 @@ wg.submit() element="H", ), inpainting=InpaintingConfig( - pretrained_name="mattergen_base", - predictor_corrector="baseline", + pretrained_name="TD-pos-only", + predictor_corrector="TD", N_steps=5, coordinates_snr=0.2, n_corrector_steps=1, diff --git a/docs/index.md b/docs/index.md index 93a8a5c..274e1b0 100644 --- a/docs/index.md +++ b/docs/index.md @@ -39,7 +39,24 @@ uv pip install .[aiida] ### Model checkpoints for retrained versions of MatterGen -Model checkpoints for the retrained versions of MatterGen used in our work can be downloaded from [Hugging Face](https://huggingface.co/t-reents/XtalPaint). Currently, the repository contains the `pos-only` and `TD-pos-only` models discussed in the paper. +Model checkpoints for the retrained versions of MatterGen used in our work are hosted on [Hugging Face](https://huggingface.co/t-reents/XtalPaint). Currently, the repository contains the `pos-only` and `TD-pos-only` models discussed in the paper. + +!!! tip "Recommended model: `TD-pos-only`" + The examples use the **retrained `TD-pos-only` model** — the core model of XtalPaint — together with the time-dependent (`TD`) predictor-corrector, which we recommend for accurate inpainting. See the [Configuration Guide](configuration.md) for the available model and predictor-corrector combinations. + + Like MatterGen's own checkpoints, the XtalPaint models are downloaded automatically the first time you select them by name: + + ```python + InpaintingConfig( + pretrained_name="TD-pos-only", # auto-downloaded & cached from Hugging Face + predictor_corrector="TD", + ..., + ) + ``` + + You can also download a checkpoint explicitly with + `xtalpaint.models.download_pretrained_model("TD-pos-only")` and pass the + returned path as `model_path`. ## Acknowledgements diff --git a/examples/README.md b/examples/README.md index cc491af..22e285c 100644 --- a/examples/README.md +++ b/examples/README.md @@ -3,8 +3,9 @@ ## Without using AiiDA -To be added. +The `running-wo-AiiDA.ipynb` notebook shows how to use the XtalPaint package without AiiDA. It is a good starting point for users who want to get familiar with the package and its functionality. All the methods that are otherwise wrapped in the AiiDA WorkGraphs are just executed serially in this notebook. +The latest symmetry refinement and uniqueness analysis steps are not included. However, they are not specific to this package and can easily be integrated by the user if needed. ## Using AiiDA -To be added. +The `running-with-AiiDA.ipynb` notebook shows how to use the XtalPaint package with AiiDA. The advantage of using AiiDA is that it keeps track of the data and also enables remote submission to different HPC clusters from your local machine. Moreover, WorkGraphs are already predefined with several optional steps. diff --git a/examples/running-with-AiiDA.ipynb b/examples/running-with-AiiDA.ipynb index 525ec7d..7fada8c 100644 --- a/examples/running-with-AiiDA.ipynb +++ b/examples/running-with-AiiDA.ipynb @@ -2,7 +2,7 @@ "cells": [ { "cell_type": "code", - "execution_count": 2, + "execution_count": 1, "id": "3ae5d5ce", "metadata": {}, "outputs": [ @@ -21,7 +21,7 @@ "Loaded AiiDA DB environment - profile name: presto.\n" ] }, - "execution_count": 2, + "execution_count": 1, "metadata": {}, "output_type": "execute_result" } @@ -33,7 +33,7 @@ }, { "cell_type": "code", - "execution_count": 3, + "execution_count": 2, "id": "e6d0ac82", "metadata": {}, "outputs": [ @@ -41,15 +41,7 @@ "name": "stdout", "output_type": "stream", "text": [ - "MODELS_PROJECT_ROOT: /Users/treents/project/paper-missing-hydrogen-archive/git/XtalPaint-dev/.venv/lib/python3.10/site-packages/mattergen\n" - ] - }, - { - "name": "stderr", - "output_type": "stream", - "text": [ - "/Users/treents/project/paper-missing-hydrogen-archive/git/XtalPaint-dev/.venv/lib/python3.10/site-packages/lightning_fabric/__init__.py:36: UserWarning: pkg_resources is deprecated as an API. See https://setuptools.pypa.io/en/latest/pkg_resources.html. The pkg_resources package is slated for removal as early as 2025-11-30. Refrain from using this package or pin to Setuptools<81.\n", - " __import__(\"pkg_resources\").declare_namespace(__name__)\n" + "MODELS_PROJECT_ROOT: /home/reents_t/project/test-xtalpaint/mattergen-clean/mattergen\n" ] } ], @@ -64,7 +56,7 @@ }, { "cell_type": "code", - "execution_count": 4, + "execution_count": 3, "id": "058f90ca", "metadata": {}, "outputs": [], @@ -77,7 +69,7 @@ }, { "cell_type": "code", - "execution_count": 5, + "execution_count": 4, "id": "39c1c807", "metadata": {}, "outputs": [ @@ -87,7 +79,7 @@ "5" ] }, - "execution_count": 5, + "execution_count": 4, "metadata": {}, "output_type": "execute_result" } @@ -98,7 +90,7 @@ }, { "cell_type": "code", - "execution_count": 6, + "execution_count": 5, "id": "671d0715", "metadata": {}, "outputs": [], @@ -111,9 +103,17 @@ "}" ] }, + { + "cell_type": "markdown", + "id": "343ee890", + "metadata": {}, + "source": [ + "> **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](https://psi-lms.github.io/XtalPaint/configuration/) for the available model and predictor-corrector combinations." + ] + }, { "cell_type": "code", - "execution_count": 7, + "execution_count": 6, "id": "619f3cb4", "metadata": {}, "outputs": [], @@ -133,9 +133,9 @@ " \"num_samples\": 1,\n", " },\n", " inpainting={\n", - " \"predictor_corrector\": \"baseline\",\n", + " \"predictor_corrector\": \"TD\",\n", " **param_grid,\n", - " \"pretrained_name\": \"mattergen_base\",\n", + " \"pretrained_name\": \"TD-pos-only\",\n", " \"sampling_config_path\": f\"{Path().resolve().parent}/mattergen-configs/sampling_conf\",\n", " },\n", " relaxation={\n", @@ -200,7 +200,7 @@ { "data": { "application/vnd.jupyter.widget-view+json": { - "model_id": "f5bbfaa9d6364c4fb06d28caa34dd4ee", + "model_id": "8db5625daaea4de087e2f9862a8b1a02", "version_major": 2, "version_minor": 1 }, @@ -268,7 +268,7 @@ " const { RenderUtils } = ReteRenderUtils;\n", " const styled = window.styled;\n", "\n", - " const nodegraphData = {"name": "InpaintingWorkGraph", "uuid": "04aa8b18-66ef-11f1-b91a-2a4049579a32", "state": "CREATED", "nodes": {"graph_inputs": {"identifier": "graph_inputs", "label": "graph_inputs", "node_type": "NORMAL", "inputs": [], "properties": {"structures": {"identifier": "workgraph.any", "value": null}, "inputs": {"identifier": "workgraph.any", "value": null}}, "outputs": [{"name": "structures"}], "position": [30, 30], "children": []}, "graph_outputs": {"identifier": "graph_outputs", "label": "graph_outputs", "node_type": "NORMAL", "inputs": [{"name": "inpainted_structures", "identifier": "workgraph.any"}, {"name": "inpainted_constrained_relaxation", "identifier": "workgraph.namespace"}, {"name": "unrelaxed_inpainted_full_relaxation", "identifier": "workgraph.namespace"}, {"name": "pre_relaxed_inpainted_full_relaxation", "identifier": "workgraph.namespace"}, {"name": "inpainting_candidates"}, {"name": "inpainted_structures"}, {"name": "inpainted_constrained_relaxation.structures"}, {"name": "inpainted_constrained_relaxation.final_energies"}, {"name": "pre_relaxed_inpainted_full_relaxation.structures"}, {"name": "pre_relaxed_inpainted_full_relaxation.final_energies"}], "properties": {"inpainted_structures": {"identifier": "workgraph.any", "value": null}, "inpainting_candidates": {"identifier": "workgraph.any", "value": null}}, "outputs": [], "position": [30, 30], "children": []}, "generate_inpainting_candidates": {"identifier": "generate_inpainting_candidates_task", "label": "generate_inpainting_candidates", "node_type": "PYTHONJOB", "inputs": [{"name": "structures", "identifier": "workgraph.annotated"}, {"name": "n_inp", "identifier": "workgraph.annotated"}, {"name": "element", "identifier": "workgraph.annotated"}, {"name": "function_data", "identifier": "workgraph.namespace"}, {"name": "structures"}], "properties": {"structures": {"identifier": "workgraph.any", "value": null}, "n_inp": {"identifier": "workgraph.any", "value": null}, "element": {"identifier": "workgraph.any", "value": null}, "num_samples": {"identifier": "workgraph.int", "value": 1}, "code": {"identifier": "workgraph.any", "value": {"input_plugin": "pythonjob.pythonjob", "append_text": "", "prepend_text": "", "use_double_quotes": false, "with_mpi": null, "wrap_cmdline_params": false, "filepath_executable": "/Users/treents/project/paper-missing-hydrogen-archive/git/XtalPaint-dev/.venv/bin/python", "node_type": "data.core.code.installed.InstalledCode."}}, "remote_folder": {"identifier": "workgraph.any", "value": null}, "process_label": {"identifier": "workgraph.any", "value": null}, "parent_folder": {"identifier": "workgraph.any", "value": null}, "parent_folder_name": {"identifier": "workgraph.any", "value": null}, "parent_output_folder": {"identifier": "workgraph.any", "value": null}, "additional_retrieve_list": {"identifier": "workgraph.any", "value": null}, "computer": {"identifier": "workgraph.string", "value": null}, "command_info": {"identifier": "workgraph.dict", "value": null}, "register_pickle_by_value": {"identifier": "workgraph.bool", "value": null}, "_wait": {"identifier": "workgraph.any", "value": null}}, "outputs": [{"name": "candidates"}, {"name": "candidates"}], "position": [30, 30], "children": []}, "inpainting": {"identifier": "inpainting_pipeline_task", "label": "inpainting", "node_type": "PYTHONJOB", "inputs": [{"name": "structures", "identifier": "workgraph.annotated"}, {"name": "config", "identifier": "workgraph.dict"}, {"name": "function_data", "identifier": "workgraph.namespace"}, {"name": "structures"}], "properties": {"structures": {"identifier": "workgraph.any", "value": null}, "config": {"identifier": "workgraph.dict", "value": null}, "usempi": {"identifier": "workgraph.bool", "value": false}, "code": {"identifier": "workgraph.any", "value": {"input_plugin": "pythonjob.pythonjob", "append_text": "", "prepend_text": "", "use_double_quotes": false, "with_mpi": null, "wrap_cmdline_params": false, "filepath_executable": "/Users/treents/project/paper-missing-hydrogen-archive/git/XtalPaint-dev/.venv/bin/python", "node_type": "data.core.code.installed.InstalledCode."}}, "remote_folder": {"identifier": "workgraph.any", "value": null}, "process_label": {"identifier": "workgraph.any", "value": null}, "parent_folder": {"identifier": "workgraph.any", "value": null}, "parent_folder_name": {"identifier": "workgraph.any", "value": null}, "parent_output_folder": {"identifier": "workgraph.any", "value": null}, "additional_retrieve_list": {"identifier": "workgraph.any", "value": null}, "computer": {"identifier": "workgraph.string", "value": null}, "command_info": {"identifier": "workgraph.dict", "value": null}, "register_pickle_by_value": {"identifier": "workgraph.bool", "value": null}, "_wait": {"identifier": "workgraph.any", "value": null}}, "outputs": [{"name": "structures"}, {"name": "structures"}], "position": [30, 30], "children": []}, "inpainted_constrained_relaxation": {"identifier": "relaxation_graph", "label": "inpainted_constrained_relaxation", "node_type": "GRAPH", "inputs": [{"name": "structures", "identifier": "workgraph.any"}, {"name": "relax_config", "identifier": "workgraph.namespace"}, {"name": "structures"}], "properties": {"structures": {"identifier": "workgraph.any", "value": null}, "command_info": {"identifier": "workgraph.any", "value": null}, "_wait": {"identifier": "workgraph.any", "value": null}}, "outputs": [{"name": "structures"}, {"name": "structures"}, {"name": "final_energies"}], "position": [30, 30], "children": []}, "pre_relaxed_inpainted_full_relaxation": {"identifier": "relaxation_graph", "label": "pre_relaxed_inpainted_full_relaxation", "node_type": "GRAPH", "inputs": [{"name": "structures", "identifier": "workgraph.any"}, {"name": "relax_config", "identifier": "workgraph.namespace"}, {"name": "structures"}], "properties": {"structures": {"identifier": "workgraph.any", "value": null}, "command_info": {"identifier": "workgraph.any", "value": null}, "_wait": {"identifier": "workgraph.any", "value": null}}, "outputs": [{"name": "structures"}, {"name": "final_energies"}], "position": [30, 30], "children": []}}, "links": [{"from_node": "graph_inputs", "to_node": "generate_inpainting_candidates", "from_socket": "structures", "to_socket": "structures"}, {"from_node": "generate_inpainting_candidates", "to_node": "inpainting", "from_socket": "candidates", "to_socket": "structures"}, {"from_node": "inpainting", "to_node": "inpainted_constrained_relaxation", "from_socket": "structures", "to_socket": "structures"}, {"from_node": "inpainted_constrained_relaxation", "to_node": "pre_relaxed_inpainted_full_relaxation", "from_socket": "structures", "to_socket": "structures"}, {"from_node": "generate_inpainting_candidates", "to_node": "graph_outputs", "from_socket": "candidates", "to_socket": "inpainting_candidates"}, {"from_node": "inpainting", "to_node": "graph_outputs", "from_socket": "structures", "to_socket": "inpainted_structures"}, {"from_node": "inpainted_constrained_relaxation", "to_node": "graph_outputs", "from_socket": "structures", "to_socket": "inpainted_constrained_relaxation.structures"}, {"from_node": "inpainted_constrained_relaxation", "to_node": "graph_outputs", "from_socket": "final_energies", "to_socket": "inpainted_constrained_relaxation.final_energies"}, {"from_node": "pre_relaxed_inpainted_full_relaxation", "to_node": "graph_outputs", "from_socket": "structures", "to_socket": "pre_relaxed_inpainted_full_relaxation.structures"}, {"from_node": "pre_relaxed_inpainted_full_relaxation", "to_node": "graph_outputs", "from_socket": "final_energies", "to_socket": "pre_relaxed_inpainted_full_relaxation.final_energies"}]}\n", + " const nodegraphData = {"name": "InpaintingWorkGraph", "uuid": "ca040eda-6a2c-11f1-9806-c91ad9856c5f", "state": "CREATED", "nodes": {"graph_inputs": {"identifier": "graph_inputs", "label": "graph_inputs", "node_type": "NORMAL", "inputs": [], "properties": {"structures": {"identifier": "workgraph.any", "value": null}, "inputs": {"identifier": "workgraph.any", "value": null}}, "outputs": [{"name": "structures"}], "position": [30, 30], "children": []}, "graph_outputs": {"identifier": "graph_outputs", "label": "graph_outputs", "node_type": "NORMAL", "inputs": [{"name": "inpainted_structures", "identifier": "workgraph.any"}, {"name": "inpainted_constrained_relaxation", "identifier": "workgraph.namespace"}, {"name": "unrelaxed_inpainted_full_relaxation", "identifier": "workgraph.namespace"}, {"name": "pre_relaxed_inpainted_full_relaxation", "identifier": "workgraph.namespace"}, {"name": "inpainting_candidates"}, {"name": "inpainted_structures"}, {"name": "inpainted_constrained_relaxation.structures"}, {"name": "inpainted_constrained_relaxation.final_energies"}, {"name": "pre_relaxed_inpainted_full_relaxation.structures"}, {"name": "pre_relaxed_inpainted_full_relaxation.final_energies"}], "properties": {"inpainted_structures": {"identifier": "workgraph.any", "value": null}, "inpainting_candidates": {"identifier": "workgraph.any", "value": null}}, "outputs": [], "position": [30, 30], "children": []}, "generate_inpainting_candidates": {"identifier": "generate_inpainting_candidates_task", "label": "generate_inpainting_candidates", "node_type": "PYTHONJOB", "inputs": [{"name": "structures", "identifier": "workgraph.annotated"}, {"name": "n_inp", "identifier": "workgraph.annotated"}, {"name": "element", "identifier": "workgraph.annotated"}, {"name": "function_data", "identifier": "workgraph.namespace"}, {"name": "structures"}], "properties": {"structures": {"identifier": "workgraph.any", "value": null}, "n_inp": {"identifier": "workgraph.any", "value": null}, "element": {"identifier": "workgraph.any", "value": null}, "num_samples": {"identifier": "workgraph.int", "value": 1}, "code": {"identifier": "workgraph.any", "value": {"input_plugin": "pythonjob.pythonjob", "append_text": "", "prepend_text": "", "use_double_quotes": false, "with_mpi": null, "wrap_cmdline_params": false, "filepath_executable": "/home/reents_t/.aiida_venvs/test-xtalpaint/bin/python", "node_type": "data.core.code.installed.InstalledCode."}}, "remote_folder": {"identifier": "workgraph.any", "value": null}, "process_label": {"identifier": "workgraph.any", "value": null}, "parent_folder": {"identifier": "workgraph.any", "value": null}, "parent_folder_name": {"identifier": "workgraph.any", "value": null}, "parent_output_folder": {"identifier": "workgraph.any", "value": null}, "additional_retrieve_list": {"identifier": "workgraph.any", "value": null}, "computer": {"identifier": "workgraph.string", "value": null}, "command_info": {"identifier": "workgraph.dict", "value": null}, "register_pickle_by_value": {"identifier": "workgraph.bool", "value": null}, "_wait": {"identifier": "workgraph.any", "value": null}}, "outputs": [{"name": "candidates"}, {"name": "candidates"}], "position": [30, 30], "children": []}, "inpainting": {"identifier": "inpainting_pipeline_task", "label": "inpainting", "node_type": "PYTHONJOB", "inputs": [{"name": "structures", "identifier": "workgraph.annotated"}, {"name": "config", "identifier": "workgraph.dict"}, {"name": "function_data", "identifier": "workgraph.namespace"}, {"name": "structures"}], "properties": {"structures": {"identifier": "workgraph.any", "value": null}, "config": {"identifier": "workgraph.dict", "value": null}, "usempi": {"identifier": "workgraph.bool", "value": false}, "code": {"identifier": "workgraph.any", "value": {"input_plugin": "pythonjob.pythonjob", "append_text": "", "prepend_text": "", "use_double_quotes": false, "with_mpi": null, "wrap_cmdline_params": false, "filepath_executable": "/home/reents_t/.aiida_venvs/test-xtalpaint/bin/python", "node_type": "data.core.code.installed.InstalledCode."}}, "remote_folder": {"identifier": "workgraph.any", "value": null}, "process_label": {"identifier": "workgraph.any", "value": null}, "parent_folder": {"identifier": "workgraph.any", "value": null}, "parent_folder_name": {"identifier": "workgraph.any", "value": null}, "parent_output_folder": {"identifier": "workgraph.any", "value": null}, "additional_retrieve_list": {"identifier": "workgraph.any", "value": null}, "computer": {"identifier": "workgraph.string", "value": null}, "command_info": {"identifier": "workgraph.dict", "value": null}, "register_pickle_by_value": {"identifier": "workgraph.bool", "value": null}, "_wait": {"identifier": "workgraph.any", "value": null}}, "outputs": [{"name": "structures"}, {"name": "structures"}], "position": [30, 30], "children": []}, "inpainted_constrained_relaxation": {"identifier": "relaxation_graph", "label": "inpainted_constrained_relaxation", "node_type": "GRAPH", "inputs": [{"name": "structures", "identifier": "workgraph.any"}, {"name": "relax_config", "identifier": "workgraph.namespace"}, {"name": "structures"}], "properties": {"structures": {"identifier": "workgraph.any", "value": null}, "command_info": {"identifier": "workgraph.any", "value": null}, "_wait": {"identifier": "workgraph.any", "value": null}}, "outputs": [{"name": "structures"}, {"name": "structures"}, {"name": "final_energies"}], "position": [30, 30], "children": []}, "pre_relaxed_inpainted_full_relaxation": {"identifier": "relaxation_graph", "label": "pre_relaxed_inpainted_full_relaxation", "node_type": "GRAPH", "inputs": [{"name": "structures", "identifier": "workgraph.any"}, {"name": "relax_config", "identifier": "workgraph.namespace"}, {"name": "structures"}], "properties": {"structures": {"identifier": "workgraph.any", "value": null}, "command_info": {"identifier": "workgraph.any", "value": null}, "_wait": {"identifier": "workgraph.any", "value": null}}, "outputs": [{"name": "structures"}, {"name": "final_energies"}], "position": [30, 30], "children": []}}, "links": [{"from_node": "graph_inputs", "to_node": "generate_inpainting_candidates", "from_socket": "structures", "to_socket": "structures"}, {"from_node": "generate_inpainting_candidates", "to_node": "inpainting", "from_socket": "candidates", "to_socket": "structures"}, {"from_node": "inpainting", "to_node": "inpainted_constrained_relaxation", "from_socket": "structures", "to_socket": "structures"}, {"from_node": "inpainted_constrained_relaxation", "to_node": "pre_relaxed_inpainted_full_relaxation", "from_socket": "structures", "to_socket": "structures"}, {"from_node": "generate_inpainting_candidates", "to_node": "graph_outputs", "from_socket": "candidates", "to_socket": "inpainting_candidates"}, {"from_node": "inpainting", "to_node": "graph_outputs", "from_socket": "structures", "to_socket": "inpainted_structures"}, {"from_node": "inpainted_constrained_relaxation", "to_node": "graph_outputs", "from_socket": "structures", "to_socket": "inpainted_constrained_relaxation.structures"}, {"from_node": "inpainted_constrained_relaxation", "to_node": "graph_outputs", "from_socket": "final_energies", "to_socket": "inpainted_constrained_relaxation.final_energies"}, {"from_node": "pre_relaxed_inpainted_full_relaxation", "to_node": "graph_outputs", "from_socket": "structures", "to_socket": "pre_relaxed_inpainted_full_relaxation.structures"}, {"from_node": "pre_relaxed_inpainted_full_relaxation", "to_node": "graph_outputs", "from_socket": "final_energies", "to_socket": "pre_relaxed_inpainted_full_relaxation.final_energies"}]}\n", "\n", " // Define Schemes to use in vanilla JS\n", " const Schemes = {\n", @@ -498,7 +498,7 @@ "\" width=\"100%\" height=\"600px\" frameborder=\"0\" allowfullscreen>" ], "text/plain": [ - "" + "" ] }, "execution_count": 9, @@ -516,75 +516,15 @@ "id": "3254cf67", "metadata": {}, "outputs": [ - { - "name": "stderr", - "output_type": "stream", - "text": [ - "06/13/2026 08:13:39 AM <36552> aiida.broker.rabbitmq: [WARNING] RabbitMQ v4.1.0 is not supported and will cause unexpected problems!\n", - "WARNING:aiida.broker.rabbitmq:RabbitMQ v4.1.0 is not supported and will cause unexpected problems!\n", - "06/13/2026 08:13:39 AM <36552> aiida.broker.rabbitmq: [WARNING] It can cause long-running workflows to crash and jobs to be submitted multiple times.\n", - "WARNING:aiida.broker.rabbitmq:It can cause long-running workflows to crash and jobs to be submitted multiple times.\n", - "06/13/2026 08:13:39 AM <36552> aiida.broker.rabbitmq: [WARNING] See https://github.com/aiidateam/aiida-core/wiki/RabbitMQ-version-to-use for details.\n", - "WARNING:aiida.broker.rabbitmq:See https://github.com/aiidateam/aiida-core/wiki/RabbitMQ-version-to-use for details.\n", - "06/13/2026 08:13:40 AM <36552> aiida.orm.nodes.process.workflow.workchain.WorkChainNode: [REPORT] [44|WorkGraphEngine|continue_workgraph]: tasks ready to run: generate_inpainting_candidates\n", - "REPORT:aiida.orm.nodes.process.workflow.workchain.WorkChainNode:[44|WorkGraphEngine|continue_workgraph]: tasks ready to run: generate_inpainting_candidates\n", - "06/13/2026 08:13:40 AM <36552> aiida.orm.nodes.process.workflow.workchain.WorkChainNode: [REPORT] [44|WorkGraphEngine|on_wait]: Process status: Waiting for child processes: 46\n", - "REPORT:aiida.orm.nodes.process.workflow.workchain.WorkChainNode:[44|WorkGraphEngine|on_wait]: Process status: Waiting for child processes: 46\n", - "06/13/2026 08:13:44 AM <36552> aiida.orm.nodes.process.workflow.workchain.WorkChainNode: [REPORT] [44|WorkGraphEngine|update_task_state]: Task: generate_inpainting_candidates, type: PYTHONJOB, finished.\n", - "REPORT:aiida.orm.nodes.process.workflow.workchain.WorkChainNode:[44|WorkGraphEngine|update_task_state]: Task: generate_inpainting_candidates, type: PYTHONJOB, finished.\n", - "06/13/2026 08:13:45 AM <36552> aiida.orm.nodes.process.workflow.workchain.WorkChainNode: [REPORT] [44|WorkGraphEngine|continue_workgraph]: tasks ready to run: inpainting\n", - "REPORT:aiida.orm.nodes.process.workflow.workchain.WorkChainNode:[44|WorkGraphEngine|continue_workgraph]: tasks ready to run: inpainting\n", - "06/13/2026 08:13:45 AM <36552> aiida.orm.nodes.process.workflow.workchain.WorkChainNode: [REPORT] [44|WorkGraphEngine|on_wait]: Process status: Waiting for child processes: 52\n", - "REPORT:aiida.orm.nodes.process.workflow.workchain.WorkChainNode:[44|WorkGraphEngine|on_wait]: Process status: Waiting for child processes: 52\n", - "06/13/2026 08:14:05 AM <36552> aiida.orm.nodes.process.workflow.workchain.WorkChainNode: [REPORT] [44|WorkGraphEngine|update_task_state]: Task: inpainting, type: PYTHONJOB, finished.\n", - "REPORT:aiida.orm.nodes.process.workflow.workchain.WorkChainNode:[44|WorkGraphEngine|update_task_state]: Task: inpainting, type: PYTHONJOB, finished.\n", - "06/13/2026 08:14:05 AM <36552> aiida.orm.nodes.process.workflow.workchain.WorkChainNode: [REPORT] [44|WorkGraphEngine|continue_workgraph]: tasks ready to run: inpainted_constrained_relaxation\n", - "REPORT:aiida.orm.nodes.process.workflow.workchain.WorkChainNode:[44|WorkGraphEngine|continue_workgraph]: tasks ready to run: inpainted_constrained_relaxation\n", - "06/13/2026 08:14:05 AM <36552> aiida.orm.nodes.process.workflow.workchain.WorkChainNode: [REPORT] [44|WorkGraphEngine|on_wait]: Process status: Waiting for child processes: 59\n", - "REPORT:aiida.orm.nodes.process.workflow.workchain.WorkChainNode:[44|WorkGraphEngine|on_wait]: Process status: Waiting for child processes: 59\n", - "06/13/2026 08:14:06 AM <36552> aiida.orm.nodes.process.workflow.workchain.WorkChainNode: [REPORT] [59|WorkGraphEngine|continue_workgraph]: tasks ready to run: relaxation_task\n", - "REPORT:aiida.orm.nodes.process.workflow.workchain.WorkChainNode:[59|WorkGraphEngine|continue_workgraph]: tasks ready to run: relaxation_task\n", - "06/13/2026 08:14:06 AM <36552> aiida.orm.nodes.process.workflow.workchain.WorkChainNode: [REPORT] [59|WorkGraphEngine|on_wait]: Process status: Waiting for child processes: 61\n", - "REPORT:aiida.orm.nodes.process.workflow.workchain.WorkChainNode:[59|WorkGraphEngine|on_wait]: Process status: Waiting for child processes: 61\n", - "06/13/2026 08:15:10 AM <36552> aiida.orm.nodes.process.workflow.workchain.WorkChainNode: [REPORT] [59|WorkGraphEngine|update_task_state]: Task: relaxation_task, type: PYTHONJOB, finished.\n", - "REPORT:aiida.orm.nodes.process.workflow.workchain.WorkChainNode:[59|WorkGraphEngine|update_task_state]: Task: relaxation_task, type: PYTHONJOB, finished.\n", - "06/13/2026 08:15:10 AM <36552> aiida.orm.nodes.process.workflow.workchain.WorkChainNode: [REPORT] [59|WorkGraphEngine|continue_workgraph]: tasks ready to run: \n", - "REPORT:aiida.orm.nodes.process.workflow.workchain.WorkChainNode:[59|WorkGraphEngine|continue_workgraph]: tasks ready to run: \n", - "06/13/2026 08:15:10 AM <36552> aiida.orm.nodes.process.workflow.workchain.WorkChainNode: [REPORT] [59|WorkGraphEngine|finalize]: Finalize workgraph.\n", - "REPORT:aiida.orm.nodes.process.workflow.workchain.WorkChainNode:[59|WorkGraphEngine|finalize]: Finalize workgraph.\n", - "06/13/2026 08:15:11 AM <36552> aiida.orm.nodes.process.workflow.workchain.WorkChainNode: [REPORT] [44|WorkGraphEngine|update_task_state]: Task: inpainted_constrained_relaxation, type: GRAPH, finished.\n", - "REPORT:aiida.orm.nodes.process.workflow.workchain.WorkChainNode:[44|WorkGraphEngine|update_task_state]: Task: inpainted_constrained_relaxation, type: GRAPH, finished.\n", - "06/13/2026 08:15:11 AM <36552> aiida.orm.nodes.process.workflow.workchain.WorkChainNode: [REPORT] [44|WorkGraphEngine|continue_workgraph]: tasks ready to run: pre_relaxed_inpainted_full_relaxation\n", - "REPORT:aiida.orm.nodes.process.workflow.workchain.WorkChainNode:[44|WorkGraphEngine|continue_workgraph]: tasks ready to run: pre_relaxed_inpainted_full_relaxation\n", - "06/13/2026 08:15:11 AM <36552> aiida.orm.nodes.process.workflow.workchain.WorkChainNode: [REPORT] [44|WorkGraphEngine|on_wait]: Process status: Waiting for child processes: 69\n", - "REPORT:aiida.orm.nodes.process.workflow.workchain.WorkChainNode:[44|WorkGraphEngine|on_wait]: Process status: Waiting for child processes: 69\n", - "06/13/2026 08:15:11 AM <36552> aiida.orm.nodes.process.workflow.workchain.WorkChainNode: [REPORT] [69|WorkGraphEngine|continue_workgraph]: tasks ready to run: relaxation_task\n", - "REPORT:aiida.orm.nodes.process.workflow.workchain.WorkChainNode:[69|WorkGraphEngine|continue_workgraph]: tasks ready to run: relaxation_task\n", - "06/13/2026 08:15:12 AM <36552> aiida.orm.nodes.process.workflow.workchain.WorkChainNode: [REPORT] [69|WorkGraphEngine|on_wait]: Process status: Waiting for child processes: 71\n", - "REPORT:aiida.orm.nodes.process.workflow.workchain.WorkChainNode:[69|WorkGraphEngine|on_wait]: Process status: Waiting for child processes: 71\n", - "06/13/2026 08:16:07 AM <36552> aiida.orm.nodes.process.workflow.workchain.WorkChainNode: [REPORT] [69|WorkGraphEngine|update_task_state]: Task: relaxation_task, type: PYTHONJOB, finished.\n", - "REPORT:aiida.orm.nodes.process.workflow.workchain.WorkChainNode:[69|WorkGraphEngine|update_task_state]: Task: relaxation_task, type: PYTHONJOB, finished.\n", - "06/13/2026 08:16:07 AM <36552> aiida.orm.nodes.process.workflow.workchain.WorkChainNode: [REPORT] [69|WorkGraphEngine|continue_workgraph]: tasks ready to run: \n", - "REPORT:aiida.orm.nodes.process.workflow.workchain.WorkChainNode:[69|WorkGraphEngine|continue_workgraph]: tasks ready to run: \n", - "06/13/2026 08:16:07 AM <36552> aiida.orm.nodes.process.workflow.workchain.WorkChainNode: [REPORT] [69|WorkGraphEngine|finalize]: Finalize workgraph.\n", - "REPORT:aiida.orm.nodes.process.workflow.workchain.WorkChainNode:[69|WorkGraphEngine|finalize]: Finalize workgraph.\n", - "06/13/2026 08:16:07 AM <36552> aiida.orm.nodes.process.workflow.workchain.WorkChainNode: [REPORT] [44|WorkGraphEngine|update_task_state]: Task: pre_relaxed_inpainted_full_relaxation, type: GRAPH, finished.\n", - "REPORT:aiida.orm.nodes.process.workflow.workchain.WorkChainNode:[44|WorkGraphEngine|update_task_state]: Task: pre_relaxed_inpainted_full_relaxation, type: GRAPH, finished.\n", - "06/13/2026 08:16:07 AM <36552> aiida.orm.nodes.process.workflow.workchain.WorkChainNode: [REPORT] [44|WorkGraphEngine|continue_workgraph]: tasks ready to run: \n", - "REPORT:aiida.orm.nodes.process.workflow.workchain.WorkChainNode:[44|WorkGraphEngine|continue_workgraph]: tasks ready to run: \n", - "06/13/2026 08:16:07 AM <36552> aiida.orm.nodes.process.workflow.workchain.WorkChainNode: [REPORT] [44|WorkGraphEngine|finalize]: Finalize workgraph.\n", - "REPORT:aiida.orm.nodes.process.workflow.workchain.WorkChainNode:[44|WorkGraphEngine|finalize]: Finalize workgraph.\n" - ] - }, { "data": { "text/plain": [ - "{'inpainted_structures': ,\n", - " 'inpainting_candidates': ,\n", - " 'inpainted_constrained_relaxation': {'structures': ,\n", - " 'final_energies': },\n", - " 'pre_relaxed_inpainted_full_relaxation': {'structures': ,\n", - " 'final_energies': }}" + "{'inpainted_structures': ,\n", + " 'inpainting_candidates': ,\n", + " 'inpainted_constrained_relaxation': {'structures': ,\n", + " 'final_energies': },\n", + " 'pre_relaxed_inpainted_full_relaxation': {'structures': ,\n", + " 'final_energies': }}" ] }, "execution_count": 10, @@ -607,7 +547,7 @@ ], "metadata": { "kernelspec": { - "display_name": "XtalPaint-dev (3.10.17.final.0)", + "display_name": "test-xtalpaint (3.10.12.final.0)", "language": "python", "name": "python3" }, @@ -621,7 +561,7 @@ "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3", - "version": "3.10.17" + "version": "3.10.12" } }, "nbformat": 4, diff --git a/examples/running-wo-AiiDA.ipynb b/examples/running-wo-AiiDA.ipynb index 3abbd3d..9391292 100644 --- a/examples/running-wo-AiiDA.ipynb +++ b/examples/running-wo-AiiDA.ipynb @@ -2,7 +2,7 @@ "cells": [ { "cell_type": "code", - "execution_count": 2, + "execution_count": 1, "id": "e6d0ac82", "metadata": {}, "outputs": [], @@ -17,7 +17,7 @@ }, { "cell_type": "code", - "execution_count": 3, + "execution_count": 2, "id": "ac593022", "metadata": {}, "outputs": [], @@ -30,7 +30,7 @@ }, { "cell_type": "code", - "execution_count": 4, + "execution_count": 3, "id": "ef123e24", "metadata": {}, "outputs": [ @@ -40,7 +40,7 @@ "5" ] }, - "execution_count": 4, + "execution_count": 3, "metadata": {}, "output_type": "execute_result" } @@ -51,7 +51,7 @@ }, { "cell_type": "code", - "execution_count": 5, + "execution_count": 4, "id": "671d0715", "metadata": {}, "outputs": [], @@ -64,33 +64,25 @@ "}" ] }, + { + "cell_type": "markdown", + "id": "36020a00", + "metadata": {}, + "source": [ + "> **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](https://psi-lms.github.io/XtalPaint/configuration/) for the available model and predictor-corrector combinations." + ] + }, { "cell_type": "code", - "execution_count": 6, + "execution_count": 5, "id": "619f3cb4", "metadata": {}, "outputs": [ - { - "name": "stderr", - "output_type": "stream", - "text": [ - "/Users/treents/project/paper-missing-hydrogen-archive/git/XtalPaint-dev/.venv/lib/python3.10/site-packages/tqdm/auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n", - " from .autonotebook import tqdm as notebook_tqdm\n" - ] - }, { "name": "stdout", "output_type": "stream", "text": [ - "MODELS_PROJECT_ROOT: /Users/treents/project/paper-missing-hydrogen-archive/git/XtalPaint-dev/.venv/lib/python3.10/site-packages/mattergen\n" - ] - }, - { - "name": "stderr", - "output_type": "stream", - "text": [ - "/Users/treents/project/paper-missing-hydrogen-archive/git/XtalPaint-dev/.venv/lib/python3.10/site-packages/lightning_fabric/__init__.py:36: UserWarning: pkg_resources is deprecated as an API. See https://setuptools.pypa.io/en/latest/pkg_resources.html. The pkg_resources package is slated for removal as early as 2025-11-30. Refrain from using this package or pin to Setuptools<81.\n", - " __import__(\"pkg_resources\").declare_namespace(__name__)\n" + "MODELS_PROJECT_ROOT: /home/reents_t/project/test-xtalpaint/mattergen-clean/mattergen\n" ] } ], @@ -106,9 +98,9 @@ " },\n", " inpainting={\n", " \"record_trajectories\": False,\n", - " \"predictor_corrector\": \"baseline\",\n", + " \"predictor_corrector\": \"TD\",\n", " **param_grid,\n", - " \"pretrained_name\": \"mattergen_base\",\n", + " \"pretrained_name\": \"TD-pos-only\",\n", " \"sampling_config_path\": \"../mattergen-configs/sampling_conf\",\n", " },\n", " relaxation={\n", @@ -144,7 +136,7 @@ }, { "cell_type": "code", - "execution_count": 7, + "execution_count": 6, "id": "18ca9e1b", "metadata": {}, "outputs": [ @@ -170,7 +162,7 @@ }, { "cell_type": "code", - "execution_count": 8, + "execution_count": 7, "id": "21c7c7a1", "metadata": {}, "outputs": [], @@ -182,7 +174,7 @@ }, { "cell_type": "code", - "execution_count": 9, + "execution_count": 8, "id": "32d0c0ce", "metadata": {}, "outputs": [ @@ -215,7 +207,7 @@ }, { "cell_type": "code", - "execution_count": 10, + "execution_count": 9, "id": "0680591c", "metadata": {}, "outputs": [ @@ -271,15 +263,15 @@ " '662c7351_ee76_48ea_bab7_b733e1fdf607': Structure Summary\n", " Lattice\n", " abc : 8.414296245395509 8.414296245395509 8.414296245395509\n", - " angles : 93.51670444494472 93.51670444494472 151.32085705391836\n", + " angles : 93.51670444494472 93.51670444494472 151.3208570539183\n", " volume : 276.98877799736266\n", " A : -5.7644391936362 5.7644391936362 2.0839536633569\n", " B : 5.7644391936362 -5.7644391936362 2.0839536633569\n", " C : 5.7644391936362 5.7644391936362 -2.0839536633569\n", " pbc : True True True\n", " PeriodicSite: O (4.413e-16, 8.82, -0.1928) [0.7188, -0.04627, 0.7651]\n", - " PeriodicSite: O (7.374e-16, 3.056, 1.235) [0.5613, 0.2963, 0.2651]\n", - " PeriodicSite: O (2.987e-16, 2.709, -0.1928) [0.1887, -0.04627, 0.2349]\n", + " PeriodicSite: O (7.129e-17, 3.056, 1.235) [0.5613, 0.2963, 0.2651]\n", + " PeriodicSite: O (7.664e-17, 2.709, -0.1928) [0.1887, -0.04627, 0.2349]\n", " PeriodicSite: O (7.102e-16, 8.473, 1.235) [1.031, 0.2963, 0.7349]\n", " PeriodicSite: O (3.056, 5.764, 2.277) [1.046, 0.8113, 0.7651]\n", " PeriodicSite: O (8.82, 5.764, 0.8491) [0.7037, 0.9688, 1.265]\n", @@ -351,7 +343,7 @@ " PeriodicSite: H (nan, nan, nan) [nan, nan, nan]}" ] }, - "execution_count": 10, + "execution_count": 9, "metadata": {}, "output_type": "execute_result" } @@ -370,7 +362,7 @@ }, { "cell_type": "code", - "execution_count": 11, + "execution_count": 10, "id": "e1dc0e8c", "metadata": {}, "outputs": [], @@ -383,7 +375,7 @@ }, { "cell_type": "code", - "execution_count": 12, + "execution_count": 11, "id": "fb6ee614", "metadata": {}, "outputs": [], @@ -399,7 +391,7 @@ }, { "cell_type": "code", - "execution_count": 13, + "execution_count": 12, "id": "2c6cb21b", "metadata": {}, "outputs": [], @@ -415,7 +407,7 @@ }, { "cell_type": "code", - "execution_count": 14, + "execution_count": 13, "id": "8d60cef4", "metadata": {}, "outputs": [ @@ -435,47 +427,47 @@ " PeriodicSite: N (4.894, 2.86, 2.309) [0.9467, 0.5533, 0.4467]\n", " PeriodicSite: N (2.86, 2.309, 4.894) [0.5533, 0.4467, 0.9467]\n", " PeriodicSite: N (0.2756, 0.2756, 0.2756) [0.0533, 0.0533, 0.0533]\n", - " PeriodicSite: H (1.14, 4.434, 0.5381) [0.2206, 0.8578, 0.1041]\n", - " PeriodicSite: H (1.937, 1.15, 4.583) [0.3747, 0.2224, 0.8865]\n", - " PeriodicSite: H (1.436, 4.945, 2.987) [0.2778, 0.9566, 0.5778]\n", - " PeriodicSite: H (3.769, 2.298, 0.758) [0.729, 0.4444, 0.1466]\n", - " PeriodicSite: H (4.507, 2.123, 4.563) [0.8718, 0.4107, 0.8826]\n", - " PeriodicSite: H (3.157, 0.6855, 2.533) [0.6106, 0.1326, 0.49]\n", - " PeriodicSite: H (3.407, 1.12, 0.2044) [0.659, 0.2167, 0.03953]\n", - " PeriodicSite: H (1.048, 2.116, 3.06) [0.2027, 0.4093, 0.5919]\n", - " PeriodicSite: H (3.138, 3.168, 1.91) [0.607, 0.6127, 0.3694]\n", - " PeriodicSite: H (4.154, 3.005, 3.073) [0.8034, 0.5812, 0.5943]\n", - " PeriodicSite: H (0.1894, 3.669, 2.09) [0.03663, 0.7097, 0.4043]\n", - " PeriodicSite: H (0.2451, 3.923, 4.987) [0.04742, 0.7589, 0.9647],\n", + " PeriodicSite: H (4.658, 5.029, 1.039) [0.9011, 0.9727, 0.2009]\n", + " PeriodicSite: H (4.013, 2.973, 1.974) [0.7762, 0.5751, 0.3819]\n", + " PeriodicSite: H (1.729, 4.65, 2.002) [0.3345, 0.8995, 0.3873]\n", + " PeriodicSite: H (2.914, 3.283, 0.1321) [0.5636, 0.635, 0.02556]\n", + " PeriodicSite: H (0.4254, 1.122, 0.5066) [0.08228, 0.217, 0.098]\n", + " PeriodicSite: H (0.5707, 2.943, 2.357) [0.1104, 0.5693, 0.4559]\n", + " PeriodicSite: H (3.371, 1.632, 0.07947) [0.6522, 0.3157, 0.01537]\n", + " PeriodicSite: H (2.711, 4.801, 3.502) [0.5245, 0.9286, 0.6774]\n", + " PeriodicSite: H (1.961, 2.167, 4.502) [0.3794, 0.4192, 0.8708]\n", + " PeriodicSite: H (4.554, 2.437, 3.395) [0.881, 0.4715, 0.6567]\n", + " PeriodicSite: H (0.263, 4.842, 4.726) [0.05087, 0.9365, 0.9142]\n", + " PeriodicSite: H (2.051, 0.9003, 2.706) [0.3967, 0.1741, 0.5234],\n", " '47b9a869_9b1e_438b_8c93_f5ac654bfdd8': Structure Summary\n", " Lattice\n", - " abc : 4.38880300521851 4.401449203491215 7.175978183746338\n", - " angles : 90.00000250447809 90.00000250447802 119.90500920011145\n", - " volume : 120.16235473216483\n", + " abc : 4.38880300521851 4.401448726654057 7.175978183746338\n", + " angles : 90.000002504478 90.00000250447802 119.90499541135148\n", + " volume : 120.16235834626013\n", " A : 4.388803005218506 0.0 -1.9184066957222967e-07\n", - " B : -2.1944020434734735 3.815410169629585 -1.9239345760979631e-07\n", + " B : -2.1944008875261014 3.8154102843847926 -1.9239342918808688e-07\n", " C : 0.0 0.0 7.175978183746338\n", " pbc : True True True\n", " PeriodicSite: O (2.194, 2.549, 3.886) [0.8341, 0.6681, 0.5415]\n", - " PeriodicSite: O (-4.87e-08, 1.266, 0.2978) [0.1659, 0.3319, 0.0415]\n", + " PeriodicSite: O (2.041e-07, 1.266, 0.2978) [0.1659, 0.3319, 0.0415]\n", " PeriodicSite: O (2.194, 2.544, 6.579) [0.8334, 0.6667, 0.9168]\n", - " PeriodicSite: O (-3.111e-07, 1.272, 2.991) [0.1666, 0.3333, 0.4168]\n", - " PeriodicSite: H (-1.41, 2.653, 3.074) [0.02646, 0.6955, 0.4283]\n", - " PeriodicSite: H (-0.8308, 2.039, 1.796) [0.07786, 0.5343, 0.2503]\n", - " PeriodicSite: H (2.268, 2.906, 4.844) [0.8975, 0.7616, 0.675]\n", - " PeriodicSite: H (0.8597, 3.518, 3.477) [0.6569, 0.922, 0.4846]\n", - " PeriodicSite: H (-0.516, 1.239, 6.565) [0.04476, 0.3247, 0.9149]\n", - " PeriodicSite: H (1.068, 3.239, 6.799) [0.6679, 0.849, 0.9475]\n", - " PeriodicSite: H (0.3006, 1.175, 4.255) [0.2225, 0.3079, 0.593]\n", - " PeriodicSite: H (0.9108, 0.9392, 0.4745) [0.3306, 0.2462, 0.06613],\n", + " PeriodicSite: O (2.05e-07, 1.272, 2.991) [0.1666, 0.3333, 0.4168]\n", + " PeriodicSite: H (-0.00404, 1.27, 1.633) [0.1655, 0.3328, 0.2275]\n", + " PeriodicSite: H (2.18, 2.575, 5.204) [0.8343, 0.6749, 0.7252]\n", + " PeriodicSite: H (1.241, 2.012, 3.439) [0.5466, 0.5274, 0.4792]\n", + " PeriodicSite: H (2.154, 3.561, 3.503) [0.9575, 0.9332, 0.4882]\n", + " PeriodicSite: H (2.998, 2.102, 6.862) [0.9585, 0.5509, 0.9562]\n", + " PeriodicSite: H (0.8265, 1.755, 7.096) [0.4183, 0.46, 0.9888]\n", + " PeriodicSite: H (0.01103, 0.1311, 7.014) [0.01969, 0.03436, 0.9775]\n", + " PeriodicSite: H (-0.7893, 1.721, 3.241) [0.04569, 0.4511, 0.4516],\n", " '662c7351_ee76_48ea_bab7_b733e1fdf607': Structure Summary\n", " Lattice\n", - " abc : 4.167909019574006 8.414296283494188 8.414295196533203\n", - " angles : 86.4832955066596 75.66042236811894 75.66042279114727\n", - " volume : 276.9888425278154\n", - " A : 4.038057327270508 0.0 1.0322589874267578\n", - " B : 2.019028748729142 8.152147915404894 0.5161290764808655\n", - " C : 0.0 0.0 8.414295196533203\n", + " abc : 4.167909392982183 8.414297301182946 8.41429615020752\n", + " angles : 86.48328861312336 75.66042875469425 75.66042198607414\n", + " volume : 276.98893913632514\n", + " A : 4.038057804107666 0.0 1.0322586297988892\n", + " B : 2.0190288528503957 8.152148872105393 0.5161301493644714\n", + " C : 0.0 0.0 8.41429615020752\n", " pbc : True True True\n", " PeriodicSite: O (0.6612, 1.915, 6.606) [0.04627, 0.2349, 0.7651]\n", " PeriodicSite: O (4.326, 5.991, 3.336) [0.7037, 0.7349, 0.2651]\n", @@ -485,47 +477,47 @@ " PeriodicSite: O (3.751, 5.991, 7.143) [0.5613, 0.7349, 0.7349]\n", " PeriodicSite: O (2.307, 6.237, 7.027) [0.1887, 0.7651, 0.7651]\n", " PeriodicSite: O (0.6612, 2.161, 2.399) [0.03121, 0.2651, 0.2651]\n", - " PeriodicSite: H (2.003, 6.885, 7.471) [0.07385, 0.8445, 0.827]\n", - " PeriodicSite: H (4.384, 7.511, 8.59) [0.6251, 0.9214, 0.8877]\n", - " PeriodicSite: H (3.426, 5.829, 2.848) [0.491, 0.715, 0.2343]\n", - " PeriodicSite: H (3.292, 2.8, 3.087) [0.6434, 0.3435, 0.2669]\n", - " PeriodicSite: H (1.211, 2.332, 7.627) [0.1569, 0.2861, 0.8697]\n", - " PeriodicSite: H (1.483, 1.609, 3.175) [0.2686, 0.1974, 0.3322]\n", - " PeriodicSite: H (1.13, 3.42, 5.603) [0.07015, 0.4195, 0.6316]\n", - " PeriodicSite: H (4.362, 6.361, 6.174) [0.6902, 0.7803, 0.6012],\n", + " PeriodicSite: H (3.678, 2.716, 7.429) [0.7442, 0.3332, 0.7712]\n", + " PeriodicSite: H (0.1682, 0.5427, 6.398) [0.008369, 0.06657, 0.7552]\n", + " PeriodicSite: H (3.914, 6.039, 6.108) [0.5988, 0.7408, 0.607]\n", + " PeriodicSite: H (5.223, 6.384, 3.671) [0.9019, 0.7831, 0.2776]\n", + " PeriodicSite: H (4.78, 3.176, 3.444) [0.9889, 0.3895, 0.264]\n", + " PeriodicSite: H (2.005, 2.087, 6.873) [0.3685, 0.256, 0.756]\n", + " PeriodicSite: H (2.442, 6.189, 7.609) [0.2252, 0.7592, 0.8301]\n", + " PeriodicSite: H (2.547, 1.818, 2.858) [0.5193, 0.223, 0.2622],\n", " '7fa282c5_4971_46f4_8b3b_776595a0fa06': Structure Summary\n", " Lattice\n", - " abc : 6.332112789154059 6.332114219665534 6.603632926940918\n", - " angles : 90.00000250447802 90.00000250447808 120.00000394341146\n", - " volume : 229.3036252628656\n", + " abc : 6.332112789154059 6.332114696502691 6.603632926940918\n", + " angles : 90.00000250447809 90.00000250447808 119.99999014147143\n", + " volume : 229.303674421453\n", " A : 6.332112789154053 0.0 -2.767854425655969e-07\n", - " B : -3.1660574872561824 5.483771555989494 -2.7678549940901576e-07\n", + " B : -3.1660564046927258 5.483772731611829 -2.767855278307252e-07\n", " C : 0.0 0.0 6.603632926940918\n", " pbc : True True True\n", " PeriodicSite: Y (2.134, 0.0, 4.953) [0.3369, 0.0, 0.75]\n", " PeriodicSite: Y (1.067, 1.848, 1.651) [0.3369, 0.3369, 0.25]\n", - " PeriodicSite: Y (-2.099, 3.636, 1.651) [0.0, 0.6631, 0.25]\n", + " PeriodicSite: Y (-2.099, 3.636, 1.651) [6.806e-10, 0.6631, 0.25]\n", " PeriodicSite: Y (-1.067, 1.848, 4.953) [0.0, 0.3369, 0.75]\n", - " PeriodicSite: Y (4.199, 0.0, 1.651) [0.6631, 0.0, 0.25]\n", + " PeriodicSite: Y (4.199, 4.059e-09, 1.651) [0.6631, 7.402e-10, 0.25]\n", " PeriodicSite: Y (2.099, 3.636, 4.953) [0.6631, 0.6631, 0.75]\n", - " PeriodicSite: H (0.794, 2.63, 3.703) [0.3652, 0.4797, 0.5608]\n", - " PeriodicSite: H (1.941, 1.615, 4.938) [0.4539, 0.2946, 0.7478]\n", - " PeriodicSite: H (5.916, 0.1827, 5.518) [0.951, 0.03332, 0.8356]\n", - " PeriodicSite: H (-0.06808, 4.299, 4.622) [0.3812, 0.7839, 0.6999]\n", - " PeriodicSite: H (1.295, 4.365, 0.2861) [0.6025, 0.7961, 0.04332]\n", - " PeriodicSite: H (2.724, 4.428, 2.062) [0.834, 0.8074, 0.3122]\n", - " PeriodicSite: H (3.671, 2.61, 4.224) [0.8177, 0.4759, 0.6396]\n", - " PeriodicSite: H (3.874, 3.342, 6.056) [0.9165, 0.6095, 0.9171]\n", - " PeriodicSite: H (-0.8797, 4.082, 0.5196) [0.2333, 0.7444, 0.07868]\n", - " PeriodicSite: H (4.303, 1.536, 5.382) [0.8195, 0.28, 0.815]\n", - " PeriodicSite: H (3.395, 0.7929, 0.246) [0.6084, 0.1446, 0.03725]\n", - " PeriodicSite: H (-0.2419, 2.32, 0.3047) [0.1733, 0.4231, 0.04614]\n", - " PeriodicSite: H (2.826, 2.619, 1.502) [0.6851, 0.4775, 0.2275]\n", - " PeriodicSite: H (5.496, 0.5895, 3.986) [0.9216, 0.1075, 0.6036]\n", - " PeriodicSite: H (1.116, 0.3201, 0.266) [0.2054, 0.05837, 0.04029]\n", - " PeriodicSite: H (2.371, 1.049, 2.397) [0.4701, 0.1913, 0.3629]\n", - " PeriodicSite: H (0.1526, 4.11, 2.184) [0.3988, 0.7494, 0.3307]\n", - " PeriodicSite: H (2.168, 4.776, 0.008604) [0.7778, 0.8709, 0.001303],\n", + " PeriodicSite: H (2.617, 4.687, 0.3188) [0.8406, 0.8547, 0.04828]\n", + " PeriodicSite: H (3.336, 5.011, 4.241) [0.9836, 0.9138, 0.6422]\n", + " PeriodicSite: H (1.738, 3.59, 0.5792) [0.6018, 0.6547, 0.0877]\n", + " PeriodicSite: H (4.316, 0.07872, 5.541) [0.6888, 0.01435, 0.8391]\n", + " PeriodicSite: H (-0.1905, 4.877, 0.4051) [0.4146, 0.8894, 0.06134]\n", + " PeriodicSite: H (0.8336, 3.08, 3.427) [0.4125, 0.5617, 0.519]\n", + " PeriodicSite: H (4.981, 1.859, 2.67) [0.9562, 0.339, 0.4044]\n", + " PeriodicSite: H (2.164, 0.2364, 2.583) [0.3634, 0.04311, 0.3912]\n", + " PeriodicSite: H (-1.189, 3.809, 3.433) [0.1596, 0.6946, 0.5199]\n", + " PeriodicSite: H (3.06, 1.971, 1.342) [0.663, 0.3594, 0.2032]\n", + " PeriodicSite: H (6.034, 0.05884, 3.022) [0.9583, 0.01073, 0.4576]\n", + " PeriodicSite: H (3.116, 1.799, 4.557) [0.6561, 0.328, 0.69]\n", + " PeriodicSite: H (1.331, 0.2379, 0.3926) [0.2319, 0.04337, 0.05945]\n", + " PeriodicSite: H (0.2431, 3.149, 5.876) [0.3255, 0.5743, 0.8898]\n", + " PeriodicSite: H (-1.052, 1.826, 0.5455) [0.0003599, 0.333, 0.08261]\n", + " PeriodicSite: H (2.353, 4.268, 1.745) [0.7608, 0.7782, 0.2643]\n", + " PeriodicSite: H (4.193, 3.567, 5.83) [0.9874, 0.6505, 0.8829]\n", + " PeriodicSite: H (0.9715, 1.36, 4.971) [0.2774, 0.248, 0.7527],\n", " 'c436bbf4_9aef_44a8_8960_00227f79a32f': Structure Summary\n", " Lattice\n", " abc : 4.403401374816899 4.962053775787357 5.689428329467773\n", @@ -538,20 +530,20 @@ " PeriodicSite: V (2.02, 2.481, 1.451) [0.4587, 0.5, 0.2551]\n", " PeriodicSite: V (2.02, 4.962, 4.238) [0.4587, 1.0, 0.7449]\n", " PeriodicSite: V (2.02, 2.481, 4.238) [0.4587, 0.5, 0.7449]\n", - " PeriodicSite: V (2.02, 2.07e-06, 1.451) [0.4587, 4.172e-07, 0.2551]\n", + " PeriodicSite: V (2.02, 2.13e-06, 1.451) [0.4587, 4.293e-07, 0.2551]\n", " PeriodicSite: V (4.251, 2.481, 2.845) [0.9654, 0.5, 0.5]\n", " PeriodicSite: V (4.251, 0.0, 2.845) [0.9654, 0.0, 0.5]\n", " PeriodicSite: V (4.243, 2.481, -2.939e-07) [0.9636, 0.5, 0.0]\n", " PeriodicSite: V (4.243, 0.0, -1.855e-07) [0.9636, 0.0, 0.0]\n", - " PeriodicSite: H (2.522, 1.9, 2.592) [0.5727, 0.3828, 0.4555]\n", - " PeriodicSite: H (0.7042, 3.341, 4.811) [0.1599, 0.6734, 0.8457]\n", - " PeriodicSite: H (1.551, 0.6233, 3.259) [0.3523, 0.1256, 0.5728]\n", - " PeriodicSite: H (0.4838, 4.28, 1.176) [0.1099, 0.8626, 0.2067]\n", - " PeriodicSite: H (3.972, 0.5121, 4.549) [0.902, 0.1032, 0.7995]\n", - " PeriodicSite: H (3.513, 1.376, 0.883) [0.7977, 0.2773, 0.1552]}" + " PeriodicSite: H (2.721, 1.204, 5.614) [0.6179, 0.2427, 0.9868]\n", + " PeriodicSite: H (2.758, 3.671, 2.776) [0.6263, 0.7399, 0.4878]\n", + " PeriodicSite: H (0.8412, 3.666, 0.8015) [0.191, 0.7389, 0.1409]\n", + " PeriodicSite: H (0.5002, 0.9219, 1.349) [0.1136, 0.1858, 0.2371]\n", + " PeriodicSite: H (0.5522, 1.211, 4.23) [0.1254, 0.244, 0.7434]\n", + " PeriodicSite: H (2.739, 1.2, 2.812) [0.6219, 0.2419, 0.4942]}" ] }, - "execution_count": 14, + "execution_count": 13, "metadata": {}, "output_type": "execute_result" } @@ -570,7 +562,7 @@ }, { "cell_type": "code", - "execution_count": 15, + "execution_count": 14, "id": "c00b5e12", "metadata": {}, "outputs": [], @@ -580,7 +572,7 @@ }, { "cell_type": "code", - "execution_count": 16, + "execution_count": 15, "id": "76d337c2", "metadata": {}, "outputs": [ @@ -604,15 +596,13 @@ " \"return_initial_forces\": false,\n", " \"return_final_forces\": false\n", "}\n", - "\u001b[32m2026-06-13 07:54:57.830\u001b[0m | \u001b[1mINFO \u001b[0m | \u001b[36mmattersim.forcefield.potential\u001b[0m:\u001b[36mfrom_checkpoint\u001b[0m:\u001b[36m891\u001b[0m - \u001b[1mLoading the pre-trained mattersim-v1.0.0-5M.pth model\u001b[0m\n" + "\u001b[32m2026-06-17 11:16:50.461\u001b[0m | \u001b[1mINFO \u001b[0m | \u001b[36mmattersim.forcefield.potential\u001b[0m:\u001b[36mfrom_checkpoint\u001b[0m:\u001b[36m891\u001b[0m - \u001b[1mLoading the pre-trained mattersim-v1.0.0-5M.pth model\u001b[0m\n" ] }, { "name": "stderr", "output_type": "stream", "text": [ - "/Users/treents/project/paper-missing-hydrogen-archive/git/XtalPaint-dev/.venv/lib/python3.10/site-packages/mattersim/forcefield/potential.py:896: FutureWarning: You are using `torch.load` with `weights_only=False` (the current default value), which uses the default pickle module implicitly. It is possible to construct malicious pickle data which will execute arbitrary code during unpickling (See https://github.com/pytorch/pytorch/blob/main/SECURITY.md#untrusted-models for more details). In a future release, the default value for `weights_only` will be flipped to `True`. This limits the functions that could be executed during unpickling. Arbitrary objects will no longer be allowed to be loaded via this mode unless they are explicitly allowlisted by the user via `torch.serialization.add_safe_globals`. We recommend you start setting `weights_only=True` for any use case where you don't have full control of the loaded file. Please open an issue on GitHub for any issues related to this experimental feature.\n", - " checkpoint = torch.load(load_path, map_location=device)\n", " 0%| | 0/5 [00:00\n", " \n", " 20acc66e_8e38_4e5e_9e7a_c2400262cdc8\n", - " 1.146412\n", + " 0.823770\n", " False\n", " \n", " \n", " 47b9a869_9b1e_438b_8c93_f5ac654bfdd8\n", - " 1.030685\n", - " False\n", + " 0.182364\n", + " True\n", " \n", " \n", " 662c7351_ee76_48ea_bab7_b733e1fdf607\n", - " 1.624403\n", + " 1.808948\n", " False\n", " \n", " \n", " 7fa282c5_4971_46f4_8b3b_776595a0fa06\n", - " 1.225103\n", + " 0.922381\n", " False\n", " \n", " \n", " c436bbf4_9aef_44a8_8960_00227f79a32f\n", - " 1.199534\n", + " 1.243808\n", " False\n", " \n", " \n", @@ -1540,14 +1365,14 @@ "text/plain": [ " rmsd match\n", "keys \n", - "20acc66e_8e38_4e5e_9e7a_c2400262cdc8 1.146412 False\n", - "47b9a869_9b1e_438b_8c93_f5ac654bfdd8 1.030685 False\n", - "662c7351_ee76_48ea_bab7_b733e1fdf607 1.624403 False\n", - "7fa282c5_4971_46f4_8b3b_776595a0fa06 1.225103 False\n", - "c436bbf4_9aef_44a8_8960_00227f79a32f 1.199534 False" + "20acc66e_8e38_4e5e_9e7a_c2400262cdc8 0.823770 False\n", + "47b9a869_9b1e_438b_8c93_f5ac654bfdd8 0.182364 True\n", + "662c7351_ee76_48ea_bab7_b733e1fdf607 1.808948 False\n", + "7fa282c5_4971_46f4_8b3b_776595a0fa06 0.922381 False\n", + "c436bbf4_9aef_44a8_8960_00227f79a32f 1.243808 False" ] }, - "execution_count": 19, + "execution_count": 18, "metadata": {}, "output_type": "execute_result" } @@ -1558,7 +1383,7 @@ }, { "cell_type": "code", - "execution_count": 20, + "execution_count": 19, "id": "776bbbcd", "metadata": {}, "outputs": [ @@ -1595,27 +1420,27 @@ " \n", " \n", " 20acc66e_8e38_4e5e_9e7a_c2400262cdc8\n", - " 0.870491\n", + " 0.778736\n", " False\n", " \n", " \n", " 47b9a869_9b1e_438b_8c93_f5ac654bfdd8\n", - " 0.026524\n", + " 0.009466\n", " True\n", " \n", " \n", " 662c7351_ee76_48ea_bab7_b733e1fdf607\n", - " 1.741756\n", + " 1.777973\n", " False\n", " \n", " \n", " 7fa282c5_4971_46f4_8b3b_776595a0fa06\n", - " 0.619338\n", + " 0.575861\n", " False\n", " \n", " \n", " c436bbf4_9aef_44a8_8960_00227f79a32f\n", - " 1.550880\n", + " 1.408238\n", " False\n", " \n", " \n", @@ -1625,14 +1450,14 @@ "text/plain": [ " rmsd match\n", "keys \n", - "20acc66e_8e38_4e5e_9e7a_c2400262cdc8 0.870491 False\n", - "47b9a869_9b1e_438b_8c93_f5ac654bfdd8 0.026524 True\n", - "662c7351_ee76_48ea_bab7_b733e1fdf607 1.741756 False\n", - "7fa282c5_4971_46f4_8b3b_776595a0fa06 0.619338 False\n", - "c436bbf4_9aef_44a8_8960_00227f79a32f 1.550880 False" + "20acc66e_8e38_4e5e_9e7a_c2400262cdc8 0.778736 False\n", + "47b9a869_9b1e_438b_8c93_f5ac654bfdd8 0.009466 True\n", + "662c7351_ee76_48ea_bab7_b733e1fdf607 1.777973 False\n", + "7fa282c5_4971_46f4_8b3b_776595a0fa06 0.575861 False\n", + "c436bbf4_9aef_44a8_8960_00227f79a32f 1.408238 False" ] }, - "execution_count": 20, + "execution_count": 19, "metadata": {}, "output_type": "execute_result" } @@ -1640,19 +1465,11 @@ "source": [ "constrained_relaxation_evaluation" ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "25ee1e07", - "metadata": {}, - "outputs": [], - "source": [] } ], "metadata": { "kernelspec": { - "display_name": "XtalPaint-dev (3.10.17.final.0)", + "display_name": "test-xtalpaint (3.10.12.final.0)", "language": "python", "name": "python3" }, @@ -1666,7 +1483,7 @@ "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3", - "version": "3.10.17" + "version": "3.10.12" } }, "nbformat": 4, diff --git a/src/xtalpaint/inpainting/config_schema.py b/src/xtalpaint/inpainting/config_schema.py index 1be6a45..8f28010 100644 --- a/src/xtalpaint/inpainting/config_schema.py +++ b/src/xtalpaint/inpainting/config_schema.py @@ -295,9 +295,9 @@ class XtalPaintConfig(BaseModel): XtalPaintConfig( inpainting=InpaintingConfig( - pretrained_name="mattergen_base", - predictor_corrector="baseline", - N_steps=5, coordinates_snr=0.2, + pretrained_name="TD-pos-only", + predictor_corrector="TD", + N_steps=50, coordinates_snr=0.2, n_corrector_steps=1, batch_size=1000, ), ) diff --git a/src/xtalpaint/inpainting/inpainting_process.py b/src/xtalpaint/inpainting/inpainting_process.py index 8fed493..f045fb7 100644 --- a/src/xtalpaint/inpainting/inpainting_process.py +++ b/src/xtalpaint/inpainting/inpainting_process.py @@ -19,6 +19,7 @@ generate_reconstructed_structures, ) from xtalpaint.inpainting.config_schema import InpaintingConfig +from xtalpaint.models import resolve_inpainting_model from xtalpaint.utils.data_utils import create_dataloader XTALPAINT_BASE = "xtalpaint.predictor_corrector" @@ -207,6 +208,14 @@ def _run_inpainting( Tuple of (inpainted_structures, trajectories, mean_trajectories). mean_trajectories is None if not recorded. """ + # Resolve the model selection, auto-downloading XtalPaint checkpoints + # (e.g. TD-pos-only) from Hugging Face when selected by name. + pretrained_name, model_path = resolve_inpainting_model( + predictor_corrector=predictor_corrector, + pretrained_name=pretrained_name, + model_path=model_path, + ) + inpainting_model_params: dict[str, Any] = { "N_steps": N_steps, "coordinates_snr": coordinates_snr, diff --git a/src/xtalpaint/models.py b/src/xtalpaint/models.py new file mode 100644 index 0000000..d3fac7a --- /dev/null +++ b/src/xtalpaint/models.py @@ -0,0 +1,134 @@ +"""Download and resolve the retrained XtalPaint model checkpoints. + +The retrained MatterGen checkpoints used in our work are hosted on Hugging +Face (https://huggingface.co/t-reents/XtalPaint) rather than bundled with the +package. The ``TD-pos-only`` model is the recommended core model of XtalPaint; +``pos-only`` is also provided for comparison. Both are downloaded automatically +(and cached) when selected by name via ``pretrained_name``, mirroring how +MatterGen resolves its own pretrained checkpoints. +""" + +from pathlib import Path +from typing import Optional + +XTALPAINT_HF_REPO = "t-reents/XtalPaint" +XTALPAINT_PRETRAINED_MODELS = ("TD-pos-only", "pos-only") +RECOMMENDED_MODEL = "TD-pos-only" + + +def download_pretrained_model( + model_name: str = RECOMMENDED_MODEL, + local_dir: Optional[str] = None, + repository_name: str = XTALPAINT_HF_REPO, +) -> Path: + """Download a retrained XtalPaint checkpoint from Hugging Face. + + Args: + model_name: Model to download, one of ``XTALPAINT_PRETRAINED_MODELS``. + local_dir: Directory to download into. If ``None``, the shared Hugging + Face cache is used. + repository_name: Hugging Face repository to download from. + + Returns: + Path to the local model directory containing ``config.yaml`` and + ``checkpoints/last.ckpt``, usable as ``model_path``. + """ + from huggingface_hub import hf_hub_download + + if model_name not in XTALPAINT_PRETRAINED_MODELS: + raise ValueError( + f"Unknown XtalPaint model '{model_name}'. Available models: " + f"{list(XTALPAINT_PRETRAINED_MODELS)}." + ) + + download_kwargs = {"repo_id": repository_name} + if local_dir is not None: + download_kwargs["local_dir"] = local_dir + + hf_hub_download( + filename=f"{model_name}/checkpoints/last.ckpt", **download_kwargs + ) + config_path = hf_hub_download( + filename=f"{model_name}/config.yaml", **download_kwargs + ) + + return Path(config_path).parent + + +def _has_checkpoint(model_path: str) -> bool: + """Return whether ``model_path`` is a directory containing a checkpoint. + + Mirrors how ``MatterGenCheckpointInfo`` consumes ``model_path``: it expects + a directory holding ``config.yaml`` and one or more ``*.ckpt`` files. + """ + path = Path(model_path) + return path.is_dir() and any(path.rglob("*.ckpt")) + + +def _auto_download_hint(model_name: str = RECOMMENDED_MODEL) -> str: + """Return a hint describing how to obtain ``model_name``.""" + return ( + f"The '{model_name}' model is hosted on Hugging Face " + f"(https://huggingface.co/{XTALPAINT_HF_REPO}). Select it by name to " + "download it automatically:\n\n" + f' InpaintingConfig(pretrained_name="{model_name}", ...)\n\n' + "or download it explicitly and pass the returned path as " + "`model_path`:\n\n" + " from xtalpaint.models import download_pretrained_model\n" + f' model_path = download_pretrained_model("{model_name}")' + ) + + +def resolve_inpainting_model( + predictor_corrector: str, + pretrained_name: Optional[str], + model_path: Optional[str], + local_dir: Optional[str] = None, +) -> tuple[Optional[str], Optional[str]]: + """Resolve the inpainting model, auto-downloading XtalPaint checkpoints. + + An XtalPaint model selected via ``pretrained_name`` is downloaded from + Hugging Face and returned as a local ``model_path``. Otherwise the model + selection is validated and passed through unchanged. + + Args: + predictor_corrector: The selected predictor-corrector key. + pretrained_name: A bundled MatterGen checkpoint or XtalPaint model + name. + model_path: Path to a local checkpoint directory. + local_dir: Directory to download XtalPaint models into. + + Returns: + A ``(pretrained_name, model_path)`` tuple ready for the generator. + For an XtalPaint model, ``pretrained_name`` is ``None`` and + ``model_path`` points to the downloaded checkpoint. + + Raises: + FileNotFoundError: If ``model_path`` is set but the checkpoint is + missing. + ValueError: If the ``TD`` predictor-corrector is used without the + ``TD-pos-only`` model. + """ + if pretrained_name in XTALPAINT_PRETRAINED_MODELS: + downloaded = download_pretrained_model( + pretrained_name, local_dir=local_dir + ) + return None, str(downloaded) + + if model_path is not None: + if _has_checkpoint(model_path): + return pretrained_name, model_path + raise FileNotFoundError( + f"No model checkpoint was found at '{model_path}'.\n\n" + "This is expected if you have not downloaded the retrained " + "XtalPaint model yet.\n\n" + _auto_download_hint(RECOMMENDED_MODEL) + ) + + if predictor_corrector == "TD": + raise ValueError( + "The 'TD' predictor-corrector requires the retrained " + "'TD-pos-only' model, which is not bundled with MatterGen.\n\n" + + _auto_download_hint("TD-pos-only") + ) + + return pretrained_name, model_path