AlexWelcing/glim-mlip-bench

1"""glim-mlip-bench - HF ZeroGPU Space exposing MLIP elastic predictions."""2from __future__ import annotations3import os4import gradio as gr5from elastic import ElasticResult, elastic_constants6 7try:8    import spaces9    HAS_SPACES = hasattr(spaces, "GPU")10except ImportError:11    HAS_SPACES = False12 13DEFAULT_MLIPS = ("chgnet", "mace_mp0", "m3gnet")14DEFAULT_ELEMENTS = ("Al", "Cu", "Ni", "Ag", "Au", "Pt", "Pd", "Pb",15                    "Fe", "Cr", "Mo", "W", "V", "Nb", "Ta")16_CALC_CACHE: dict[str, object] = {}17 18 19def _gpu_decorator(fn):20    if HAS_SPACES:21        return spaces.GPU(duration=120)(fn)22    return fn23 24 25def _get_calculator(mlip_id: str):26    if mlip_id in _CALC_CACHE:27        return _CALC_CACHE[mlip_id]28    from calculators import make_calculator29    calc = make_calculator(mlip_id)30    _CALC_CACHE[mlip_id] = calc31    return calc32 33 34def _result_to_record(result: ElasticResult, mlip_id: str, mlip_label: str,35                      references: dict[str, float] | None) -> list[dict]:36    from datetime import datetime, timezone37    ts = datetime.now(timezone.utc).isoformat()38    refs = references or {}39    pred_map = [("C11", result.c11, "GPa"), ("C12", result.c12, "GPa"),40                ("C44", result.c44, "GPa"), ("a0", result.a0, "A")]41    records = []42    for prop, pred, unit in pred_map:43        ref = refs.get(prop)44        if ref is None:45            continue46        record = {47            "record_id": f"{result.element}_{mlip_id}_{prop}_{ts.replace(':','-').replace('.','-')}",48            "element": result.element,49            "potential_id": mlip_id,50            "potential_label": mlip_label,51            "pair_style": "mlip",52            "property": prop,53            "reference": ref,54            "predicted": pred,55            "unit": unit,56            "provenance": "hf-space",57            "agent_id": "glim-mlip-bench",58            "timestamp": ts,59        }60        if ref is not None:61            record["error"] = (pred - ref) / ref if ref != 0 else None62        records.append(record)63    return records64 65 66def predict(element: str, mlip: str = "chgnet"):67    try:68        calc = _get_calculator(mlip)69        result = elastic_constants(element, calc)70        return {71            "element": result.element,72            "structure": result.structure,73            "a0": result.a0,74            "c11": result.c11,75            "c12": result.c12,76            "c44": result.c44,77        }78    except Exception as e:79        return {"error": str(e), "element": element, "mlip": mlip}80 81 82@_gpu_decorator83def predict_batch(elements_csv: str, mlips_csv: str = "chgnet",84                  references_json: str = "{}") -> list[dict]:85    import json86    elements = [e.strip() for e in elements_csv.split(",") if e.strip()]87    mlips = [m.strip() for m in mlips_csv.split(",") if m.strip()]88    refs = json.loads(references_json) if references_json.strip() else {}89    out: list[dict] = []90    for mlip in mlips:91        try:92            calc = _get_calculator(mlip)93        except Exception as e:94            for el in elements:95                out.append({"error": f"calculator init: {e}", "element": el, "mlip": mlip})96            continue97        for el in elements:98            try:99                result = elastic_constants(el, calc)100                if refs.get(el):101                    label = f"{mlip} (HF Space)"102                    out.extend(_result_to_record(result, mlip, label, refs[el]))103                else:104                    out.append({"element": result.element, "mlip": mlip,105                                "structure": result.structure, "a0": result.a0,106                                "c11": result.c11, "c12": result.c12, "c44": result.c44})107            except Exception as e:108                out.append({"error": str(e), "element": el, "mlip": mlip})109    return out110 111 112with gr.Blocks(title="glim-mlip-bench") as demo:113    gr.Markdown("# glim-mlip-bench\nZeroGPU MLIP elastic-constant predictions (C11, C12, C44, a0)")114    with gr.Tabs():115        with gr.TabItem("Single prediction"):116            single_element = gr.Dropdown(choices=list(DEFAULT_ELEMENTS), value="Al", label="element")117            single_mlip = gr.Dropdown(choices=list(DEFAULT_MLIPS) + ["emt"], value="chgnet", label="mlip")118            single_out = gr.JSON(label="result")119            single_btn = gr.Button("Predict", variant="primary")120            single_btn.click(fn=predict, inputs=[single_element, single_mlip],121                             outputs=single_out, api_name="predict")122        with gr.TabItem("Batch prediction"):123            batch_elements = gr.Textbox(value=",".join(DEFAULT_ELEMENTS), label="elements (comma-separated)")124            batch_mlips = gr.Textbox(value="chgnet", label="mlips (comma-separated)")125            batch_refs = gr.Textbox(value="{}", label="references JSON")126            batch_out = gr.JSON(label="records")127            batch_btn = gr.Button("Predict batch", variant="primary")128            batch_btn.click(fn=predict_batch, inputs=[batch_elements, batch_mlips, batch_refs],129                            outputs=batch_out, api_name="predict_batch")130 131if __name__ == "__main__":132    demo.launch(133        server_name="0.0.0.0",134        server_port=int(os.environ.get("PORT", "7860")),135    )136