Skip to content

Use new ModelSelect widget #14

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
oerc0122 merged 1 commit into from
Oct 29, 2026
Merged

Use new ModelSelect widget #14

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
3 changes: 2 additions & 1 deletion src/aiidalab_mlip/models/code.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2,7 +2,7 @@

from aiida.orm import Code
from aiida_mlip.data.model import ModelData
from traitlets import Instance, Int, Unicode
from traitlets import Dict, Instance, Int, Unicode

from .base import Model

Expand All @@ -16,4 +16,5 @@ class CodeModel(Model):

model = Instance(klass=ModelData)
arch = Unicode()
calc_kwargs = Dict()
device = Unicode()
173 changes: 7 additions & 166 deletions src/aiidalab_mlip/steps/model.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,18 +1,14 @@
"""Results viewing wizard step."""

from pathlib import Path

import ipywidgets as ipw
from aiida import orm
from aiida.orm import Code, QueryBuilder
from aiida_mlip.data.model import ModelData
from alc_aiidalab_widgets.layouts import WizardStep
from alc_aiidalab_widgets.widgets import AiiDADatabaseQueryWidget
from alc_aiidalab_widgets.widgets.parameters import ParametersBlock
from alc_aiidalab_widgets.widgets.model_select import AiiDAMlipModelSelect
from traitlets import link

from aiidalab_mlip.models.code import CodeModel
from aiidalab_mlip.util import tab_from_dict


class ModelWizardStep(WizardStep):
Expand Down Expand Up @@ -56,127 +52,7 @@ def __init__(self, model: CodeModel, **kwargs):
)
link((self.ncpus_input, "value"), (self.model, "ncpus"))

self.model_from_uri = ipw.Text(
value="",
description="Model:",
placeholder="file:///Path/To/Model",
layout=ipw.Layout(width="80%"),
)
self.model_from_str = ipw.Text(
value="",
description="Str:",
placeholder="model",
layout=ipw.Layout(width="80%"),
)
self.model_from_node = AiiDADatabaseQueryWidget(
title="AiiDA Database",
query=[ModelData],
)
self.refresh_models_button = ipw.Button(
description="Refresh",
button_style="info",
tooltip="Refresh the list of available models",
icon="refresh",
layout={"width": "20%"},
)
self.refresh_models_button.on_click(self.update_models)

self.model_node_box = ipw.HBox(
[self.model_from_node, self.refresh_models_button], layout={"width": "100%"}
)

self.arch = ipw.Dropdown(
options=(
"mace",
"mace_mp",
"mace_off",
"chgnet",
"sevennet",
"nequip",
"dpa3",
"orb",
"mattersim",
"grace",
"upet",
"fairchem",
"mace_omol",
),
description="Arch:",
layout=ipw.Layout(width="80%"),
)
node_vbox = ipw.VBox([self.model_node_box, self.arch], layout={"width": "100%"})

uri_vbox = ipw.VBox([self.model_from_uri, self.arch], layout={"width": "100%"})
parameters_block = ParametersBlock(
"",
{"model": self.model_from_uri, "arch": self.arch},
{
"mace_mp": {
"model": "https://github.com/stfc/janus-core/raw/main/tests/models/mace_mp_small.model",
"arch": "mace_mp",
}
},
)
uri_vbox = ipw.HBox([uri_vbox, parameters_block])

### Str approach not supported by aiida-mlip
# str_vbox = ipw.VBox([self.model_from_str, self.arch], layout={"width": "100%"})
# parameters_block_str = ParametersBlock(
# "",
# {"model": self.model_from_str, "arch": self.arch},
# {
# "mace_mp": {
# "model": "small",
# "arch": "mace_mp",
# },
# "mace_off": {
# "model": "small",
# "arch": "mace_off",
# },
# "mace_omol": {
# "model": "extra_large",
# "arch": "mace_omol",
# },
# "mace_polar": {
# "model": "polar-1-m",
# "arch": "mace_polar",
# },
# "sevennet": {
# "model": "SevenNet-0_11July2024",
# "arch": "sevennet",
# },
# "orb": {
# "model": "orb_v3_conservative_20_omat",
# "arch": "orb",
# },
# "mattersim": {
# "model": "mattersim-v1.0.0-5M",
# "arch": "mattersim",
# },
# "grace": {
# "model": "GRACE-2L-OMAT",
# "arch": "grace",
# },
# "upet": {
# "model": "pet-mad-s",
# "arch": "upet",
# },
# "fairchem": {
# "model": "uma-m-1p1",
# "arch": "fairchem",
# },
# },
# )
# str_vbox = ipw.HBox([str_vbox, parameters_block_str])

self.model_type = tab_from_dict(
ipw.Tab,
{
# "From str": str_vbox,
"From URI": uri_vbox,
"AiiDA Database": node_vbox,
},
)
self.model_type = AiiDAMlipModelSelect()

self.device = ipw.Dropdown(
options=("cpu", "cuda", "mps", "xpu"),
Expand Down Expand Up @@ -214,44 +90,8 @@ def _try_load_code(self, code: str) -> Code | None:
Error: Code '{code}' not found. \n
Try: verdi code create core.code.installed --config janus.yml""")

def _load_from_str(self, model_str: str, arch: str) -> str | None: ...

def _load_from_uri(self, model_str: str, arch: str) -> ModelData | None:
if model_str.startswith(("http", "ftp", "sftp")):
model_uri = model_str
self.logspace.append_stdout("Loading model from web...")
else:
model_pth = Path(model_str).absolute()
self.logspace.append_stdout("Loading model from file...")

if not model_pth.is_file():
self.fail(f"File ({model_str}) not found.")
return None

model_uri = model_pth.as_uri()

try:
model_str = ModelData.from_uri(model_uri, architecture=arch)
model_str.label = f"{arch}:{model_uri}"
self.model.arch = arch
self.ok(f"Loaded model from {model_uri}.")
return model_str
except Exception as err:
self.fail(f"Unable to load model from {model_uri}. \nDue to: {err}")

def _load_from_node(self, node: ModelData, arch: str) -> ModelData | None:
self.model.arch = arch
return node

def _try_load_model(self, arch: str) -> ModelData | None:

typ = self.model_type.get_title(self.model_type.selected_index)

match typ:
case "From URI":
return self._load_from_uri(self.model_from_uri.value, arch)
case "AiiDA Database":
return self._load_from_node(self.model_from_node.data_object, arch)
def _try_load_model(self) -> tuple[str, ModelData] | None:
return self.model_type._try_load_model()

def submit(self, _: ipw.Button | None = None):
self.status.clear()
Expand All @@ -261,12 +101,13 @@ def submit(self, _: ipw.Button | None = None):
self.running()
if not (loaded_code := self._try_load_code(self.code.value)):
return
if not (loaded_model := self._try_load_model(self.arch.value)):
if not (loaded_model := self._try_load_model()):
return

self.model.code = loaded_code
self.model.model = loaded_model
self.model.arch, self.model.model = loaded_model
self.model.device = self.device.value
self.model.calc_kwargs = self.model_type.calc_kwargs
self.model.submitted = True
self.ok(f"Loaded code: {loaded_code}, model: {loaded_model}")

Expand Down
1 change: 1 addition & 0 deletions src/aiidalab_mlip/steps/run.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -137,6 +137,7 @@ def submit(self, _):
builder.model = model
builder.arch = orm.Str(self.model.mlip_model.arch)
builder.device = orm.Str(self.model.mlip_model.device)
builder.calc_kwargs = orm.Dict(self.model.mlip_model.calc_kwargs)

builder.metadata.options.resources = {
"num_mpiprocs_per_machine": self.model.mlip_model.ncpus,
Expand Down
Loading