diff --git a/ml_peg/app/data/element_coverage.json b/ml_peg/app/data/element_coverage.json index 2ee2f9e4c..2a73f0838 100644 --- a/ml_peg/app/data/element_coverage.json +++ b/ml_peg/app/data/element_coverage.json @@ -1,6 +1,192 @@ { "datasets": { - "MPtrj/Alexandria/OMAT": { + "MPtrj": { + "supported": [ + "Ac", + "Ag", + "Al", + "Ar", + "As", + "Au", + "B", + "Ba", + "Be", + "Bi", + "Br", + "C", + "Ca", + "Cd", + "Ce", + "Cl", + "Co", + "Cr", + "Cs", + "Cu", + "Dy", + "Er", + "Eu", + "F", + "Fe", + "Ga", + "Gd", + "Ge", + "H", + "He", + "Hf", + "Hg", + "Ho", + "I", + "In", + "Ir", + "K", + "Kr", + "La", + "Li", + "Lu", + "Mg", + "Mn", + "Mo", + "N", + "Na", + "Nb", + "Nd", + "Ne", + "Ni", + "Np", + "O", + "Os", + "P", + "Pa", + "Pb", + "Pd", + "Pm", + "Pr", + "Pt", + "Pu", + "Rb", + "Re", + "Rh", + "Ru", + "S", + "Sb", + "Sc", + "Se", + "Si", + "Sm", + "Sn", + "Sr", + "Ta", + "Tb", + "Tc", + "Te", + "Th", + "Ti", + "Tl", + "Tm", + "U", + "V", + "W", + "Xe", + "Y", + "Yb", + "Zn", + "Zr" + ] + }, + "Alexandria": { + "supported": [ + "Ac", + "Ag", + "Al", + "Ar", + "As", + "Au", + "B", + "Ba", + "Be", + "Bi", + "Br", + "C", + "Ca", + "Cd", + "Ce", + "Cl", + "Co", + "Cr", + "Cs", + "Cu", + "Dy", + "Er", + "Eu", + "F", + "Fe", + "Ga", + "Gd", + "Ge", + "H", + "He", + "Hf", + "Hg", + "Ho", + "I", + "In", + "Ir", + "K", + "Kr", + "La", + "Li", + "Lu", + "Mg", + "Mn", + "Mo", + "N", + "Na", + "Nb", + "Nd", + "Ne", + "Ni", + "Np", + "O", + "Os", + "P", + "Pa", + "Pb", + "Pd", + "Pm", + "Pr", + "Pt", + "Pu", + "Rb", + "Re", + "Rh", + "Ru", + "S", + "Sb", + "Sc", + "Se", + "Si", + "Sm", + "Sn", + "Sr", + "Ta", + "Tb", + "Tc", + "Te", + "Th", + "Ti", + "Tl", + "Tm", + "U", + "V", + "W", + "Xe", + "Y", + "Yb", + "Zn", + "Zr" + ] + }, + "OMAT": { "supported": [ "Ac", "Ag", diff --git a/ml_peg/app/filters.py b/ml_peg/app/filters.py index 04afb2267..571ac43ad 100644 --- a/ml_peg/app/filters.py +++ b/ml_peg/app/filters.py @@ -5,6 +5,7 @@ from functools import lru_cache import json from pathlib import Path +from typing import TypedDict from dash import ALL, Input, Output, State, callback, ctx, no_update from dash.dcc import Dropdown, Store @@ -18,6 +19,8 @@ PERIODIC_TABLE_SYMBOLS, ) from ml_peg.app import APP_ROOT +from ml_peg.models import current_models +from ml_peg.models.get_models import get_model_names, load_model_configs def get_model_filter(models) -> Details: @@ -134,6 +137,22 @@ def get_model_filter(models) -> Details: "cursor": "pointer", } +_APPLY_BUTTON_STYLE = { + "padding": "4px 12px", + "fontSize": "12px", + "backgroundColor": "#228be6", + "color": "#fff", + "border": "none", + "borderRadius": "4px", + "cursor": "pointer", +} + +# Apply style when the staged selection differs from the committed filter. +_APPLY_BUTTON_STYLE_PENDING = { + **_APPLY_BUTTON_STYLE, + "boxShadow": "0 0 0 3px rgba(34, 139, 230, 0.35)", +} + _ELEMENT_COVERAGE_PATH = APP_ROOT / "data" / "element_coverage.json" @@ -159,28 +178,123 @@ def _load_element_coverage(path: Path = _ELEMENT_COVERAGE_PATH) -> dict: return {} -def _dataset_presets() -> tuple[dict, ...]: +class ElementPreset(TypedDict): + """A named element-filter preset and the elements it keeps.""" + + id: str + label: str + include: tuple[str, ...] + title: str + + +def _dataset_presets() -> tuple[ElementPreset, ...]: """ Build preset entries from element_coverage.json dataset keys. + Datasets that share an identical supported-element set are merged into a + single preset (label = their names joined with "/"), so families with the + same coverage (e.g. MPtrj/Alexandria/OMAT) render as one button and future + identical-coverage datasets fold in automatically. + Returns ------- - tuple[dict, ...] - One preset dict per dataset entry in the coverage file. + tuple[ElementPreset, ...] + One preset per group of identical-coverage datasets. """ datasets = _load_element_coverage().get("datasets", {}) - return tuple( - { - "id": name.lower().replace("/", "-").replace(" ", "-"), - "label": name, - "include": tuple( - s for s in PERIODIC_TABLE_SYMBOLS if s in set(data["supported"]) - ), - "title": f"Keep elements covered by {name} training data", - } - for name, data in datasets.items() - if data.get("supported") + groups: dict[frozenset[str], list[str]] = {} + order: list[frozenset[str]] = [] + for name, data in datasets.items(): + supported = data.get("supported") + if not supported: + continue + key = frozenset(supported) + if key not in groups: + groups[key] = [] + order.append(key) + groups[key].append(name) + + presets = [] + for key in order: + label = "/".join(groups[key]) + presets.append( + { + "id": label.lower().replace("/", "-").replace(" ", "-"), + "label": label, + "include": tuple(s for s in PERIODIC_TABLE_SYMBOLS if s in key), + "title": f"Keep elements covered by {label} training data", + } + ) + return tuple(presets) + + +@lru_cache(maxsize=1) +def _model_supported_elements() -> dict[str, frozenset[str]]: + """ + Map each model to the elements covered by its training datasets. + + A model's elements are the union of the ``supported`` sets of every dataset + listed under its ``datasets`` key in ``models.yml``. Models with no + ``datasets`` tag, or whose datasets are absent from the coverage file, are + omitted (treated as untagged). + + Returns + ------- + dict[str, frozenset[str]] + Supported elements keyed by model name, for tagged models only. + """ + coverage = _load_element_coverage().get("datasets", {}) + models = get_model_names(current_models) + configs, _ = load_model_configs(tuple(models)) + + result: dict[str, frozenset[str]] = {} + for model in models: + elements: set[str] = set() + for dataset in configs.get(model, {}).get("datasets") or []: + supported = coverage.get(dataset, {}).get("supported") + if supported: + elements.update(supported) + if elements: + result[model] = frozenset(elements) + return result + + +def _selected_models_excluded_symbols( + selected_models: list[str] | None, + mode: str, +) -> list[str] | None: + """ + Return symbols to exclude to keep the union/intersection of model coverage. + + Parameters + ---------- + selected_models + Currently selected model names. + mode + Either ``"union"`` (elements supported by any selected model) or + ``"intersection"`` (elements supported by every selected model). + + Returns + ------- + list[str] | None + Symbols outside the combined coverage, or ``None`` if no selected model + is tagged with usable dataset coverage. + """ + model_elements = _model_supported_elements() + sets = [ + model_elements[model] + for model in (selected_models or []) + if model in model_elements + ] + if not sets: + return None + + keep = ( + set(sets[0]).union(*sets) + if mode == "union" + else set(sets[0]).intersection(*sets) ) + return [symbol for symbol in PERIODIC_TABLE_SYMBOLS if symbol not in keep] _ELEMENT_FILTER_PRESET_GROUPS = ( @@ -300,19 +414,68 @@ def _build_preset_section() -> Div: ) ) - return Div( + # Row of buttons deriving the keep-set from the currently selected models' + # dataset coverage, rather than a fixed preset. + selected_models_row = Div( [ Span( - "Presets", + "Selected models", style={ "fontSize": "11px", "fontWeight": "600", - "textTransform": "uppercase", - "letterSpacing": "0.07em", "color": "#6c757d", + "minWidth": "92px", }, ), + Div( + [ + Button( + "Union", + id="element-filter-union", + n_clicks=0, + title="Keep elements covered by any selected model", + style=_PRESET_BUTTON_STYLE, + ), + Button( + "Intersection", + id="element-filter-intersection", + n_clicks=0, + title="Keep elements covered by every selected model", + style=_PRESET_BUTTON_STYLE, + ), + ], + style={"display": "flex", "gap": "6px", "flexWrap": "wrap"}, + ), + ], + style={ + "display": "flex", + "alignItems": "center", + "gap": "8px", + "flexWrap": "wrap", + }, + ) + + return Div( + [ + Div( + [ + Span( + "Presets", + style={ + "fontWeight": "600", + "textTransform": "uppercase", + "letterSpacing": "0.07em", + }, + ), + Span( + " — keep only the chosen set", + style={"fontStyle": "italic"}, + ), + ], + style={"fontSize": "11px", "color": "#6c757d"}, + ), *groups, + selected_models_row, ], style={ "display": "flex", @@ -409,15 +572,7 @@ def get_element_filter() -> Div: "Apply", id="element-filter-apply", n_clicks=0, - style={ - "padding": "4px 12px", - "fontSize": "12px", - "backgroundColor": "#228be6", - "color": "#fff", - "border": "none", - "borderRadius": "4px", - "cursor": "pointer", - }, + style=_APPLY_BUTTON_STYLE, ), Button( "Exclude all", @@ -455,12 +610,35 @@ def get_element_filter() -> Div: "alignItems": "center", }, ) + grid_caption = Span( + "Click on elements to exclude them, then Apply to update tables.", + style={ + "display": "block", + "fontSize": "12px", + "color": "#6c757d", + "marginBottom": "6px", + }, + ) + + total = len(PERIODIC_TABLE_SYMBOLS) + count_display = Span( + f"Keeping {total} / {total} elements", + id="element-filter-count", + style={ + "display": "block", + "fontSize": "12px", + "color": "#6c757d", + "marginTop": "8px", + }, + ) + filter_body = Div( [ Div( - [grid_with_legend, _build_preset_section()], + [grid_caption, grid_with_legend, _build_preset_section()], style={"width": "fit-content"}, ), + count_display, actions, ] ) @@ -469,7 +647,7 @@ def get_element_filter() -> Div: [ Details( [ - Summary("Exclude elements", style=_SUMMARY_STYLE), + Summary("Element filtering", style=_SUMMARY_STYLE), Div(filter_body, style={"padding": "8px 12px"}), ], id="element-filter-details", @@ -521,12 +699,13 @@ def toggle_element(_n_clicks, pending): @callback( Output({"type": "element-btn", "index": ALL}, "style"), + Output("element-filter-count", "children"), Input("element-filter-pending", "data"), prevent_initial_call=True, ) def sync_button_styles(pending): """ - Synchronise element button styles with the pending selection. + Synchronise element button styles and kept-count with the selection. Parameters ---------- @@ -535,11 +714,40 @@ def sync_button_styles(pending): Returns ------- - list - Style dictionaries for all periodic-table element buttons. + tuple + Style dictionaries for all element buttons, and the kept-count text. """ excluded = set(pending or []) - return [_btn_style(sym, sym in excluded) for sym in all_symbols] + styles = [_btn_style(sym, sym in excluded) for sym in all_symbols] + kept = len(all_symbols) - len(excluded & set(all_symbols)) + return styles, f"Keeping {kept} / {len(all_symbols)} elements" + + @callback( + Output("element-filter-apply", "children"), + Output("element-filter-apply", "style"), + Input("element-filter-pending", "data"), + Input("element-filter", "data"), + prevent_initial_call=False, + ) + def flag_unapplied(pending, committed): + """ + Mark the Apply button when the staged selection is not yet committed. + + Parameters + ---------- + pending + Currently pending element exclusion selection. + committed + Currently committed element filter. + + Returns + ------- + tuple + Apply button label and style reflecting whether changes are pending. + """ + if set(pending or []) != set(committed or []): + return "Apply •", _APPLY_BUTTON_STYLE_PENDING + return "Apply", _APPLY_BUTTON_STYLE @callback( Output("element-filter", "data"), @@ -548,11 +756,24 @@ def sync_button_styles(pending): Input("element-filter-exclude-all", "n_clicks"), Input("element-filter-clear", "n_clicks"), Input({"type": "element-filter-preset", "index": ALL}, "n_clicks"), + Input("element-filter-union", "n_clicks"), + Input("element-filter-intersection", "n_clicks"), State("element-filter", "data"), State("element-filter-pending", "data"), + State("selected-models-store", "data"), prevent_initial_call=True, ) - def apply_or_clear(_apply, _exclude_all, _clear, _presets, current, pending): + def apply_or_clear( + _apply, + _exclude_all, + _clear, + _presets, + _union, + _intersection, + current, + pending, + selected_models, + ): """ Handle "Apply", clear, or bulk-update element exclusion actions. @@ -566,10 +787,16 @@ def apply_or_clear(_apply, _exclude_all, _clear, _presets, current, pending): Click count for the Clear button. _presets Click counts for preset element-set buttons. + _union + Click count for the selected-models Union button. + _intersection + Click count for the selected-models Intersection button. current Currently committed element exclusion selection. pending Currently pending element exclusion selection. + selected_models + Models currently selected in the global model filter. Returns ------- @@ -596,4 +823,10 @@ def apply_or_clear(_apply, _exclude_all, _clear, _presets, current, pending): if sorted(preset_excluded) == pending: raise PreventUpdate return no_update, preset_excluded + if trigger in ("element-filter-union", "element-filter-intersection"): + mode = "union" if trigger == "element-filter-union" else "intersection" + excluded = _selected_models_excluded_symbols(selected_models, mode) + if excluded is None or sorted(excluded) == pending: + raise PreventUpdate + return no_update, excluded raise PreventUpdate diff --git a/ml_peg/models/models.yml b/ml_peg/models/models.yml index 4976e8a52..5599ee8b5 100644 --- a/ml_peg/models/models.yml +++ b/ml_peg/models/models.yml @@ -1,6 +1,7 @@ mace-mp-0a: module: mace.calculators class_name: mace_mp + datasets: [MPtrj] device: "auto" trained_on_dispersion: false level_of_theory: PBE @@ -10,6 +11,7 @@ mace-mp-0a: mace-mp-0b3: module: mace.calculators class_name: mace_mp + datasets: [MPtrj] device: "auto" trained_on_dispersion: false level_of_theory: PBE @@ -19,6 +21,7 @@ mace-mp-0b3: mace-mpa-0: module: mace.calculators class_name: mace_mp + datasets: [MPtrj, Alexandria] device: "auto" trained_on_dispersion: false level_of_theory: PBE @@ -28,6 +31,7 @@ mace-mpa-0: mace-omat-0: module: mace.calculators class_name: mace_mp + datasets: [OMAT] device: "auto" trained_on_dispersion: false level_of_theory: PBE @@ -60,6 +64,7 @@ mace-matpes-r2scan: orb-v3-consv-inf-omat: module: orb_models.inference.calculator class_name: OrbCalc + datasets: [OMAT] device: "cpu" trained_on_dispersion: false level_of_theory: PBE @@ -69,6 +74,7 @@ orb-v3-consv-inf-omat: orb-v3-consv-omol: module: orb_models.forcefield.inference.calculator class_name: OrbCalc + datasets: [OMOL] device: "cpu" trained_on_dispersion: true level_of_theory: PBE