Enhance tests for DGL and PyG backends, including new mock implementations and improved coverage for MolecularDynamics. Update cache handling tests and ensure proper error handling for model loading. Refactor model prediction logic in ModelLightningModule to support both single and batched inputs. (#764)

Author: shyuep

src/matgl/utils/_training_pyg.py

@@ -206,20 +206,37 @@ def forward(
         l_g: Data | None = None,
         state_attr: torch.Tensor | None = None,
     ):
-        """Args:
-            g: dgl Graph
-            lat: lattice
-            l_g: Line graph
-            state_attr: State attribute.
+        """Compute model predictions.
+
+        Attaches per-node ``pos`` and per-edge ``pbc_offshift`` tensors derived from
+        ``frac_coords``/``pbc_offset`` and the supplied lattice(s), then delegates to
+        the wrapped model. Works for both single-graph (``lat`` shape ``(3, 3)``) and
+        batched (``lat`` shape ``(B, 3, 3)``) inputs.
+
+        Args:
+            g: PyG graph (single ``Data`` or batched ``Batch``).
+            lat: Lattice tensor.
+            l_g: Optional line graph.
+            state_attr: Optional state attribute.
 
         Returns:
             Model prediction.
         """
-        g.edata["lattice"] = torch.repeat_interleave(lat, g.batch_num_edges(), dim=0)  # type:ignore[arg-type]
-        g.edata["pbc_offshift"] = (g.edata["pbc_offset"].unsqueeze(dim=-1) * g.edata["lattice"]).sum(dim=1)
-        g.ndata["pos"] = (
-            g.ndata["frac_coords"].unsqueeze(dim=-1) * torch.repeat_interleave(lat, g.batch_num_nodes(), dim=0)  # type:ignore[arg-type]
-        ).sum(dim=1)
+        if lat is not None:
+            # Normalize lat to (B, 3, 3) regardless of whether the caller supplied a
+            # single-graph (3, 3) tensor or a (B, 3, 3) batched tensor.
+            if lat.dim() == 2:
+                lat = lat.unsqueeze(0)
+            # PyG batches assign each node a graph index via ``g.batch``; for a
+            # non-batched ``Data`` this attribute is missing/None, in which case
+            # every node belongs to graph 0.
+            batch = getattr(g, "batch", None)
+            if batch is None:
+                batch = torch.zeros(g.num_nodes, dtype=torch.long, device=g.frac_coords.device)
+            node_lat = lat[batch]
+            g.pos = (g.frac_coords.unsqueeze(dim=-1) * node_lat).sum(dim=1)
+            edge_lat = lat[batch[g.edge_index[0]]]
+            g.pbc_offshift = (g.pbc_offset.unsqueeze(dim=-1) * edge_lat).sum(dim=1)
         if self.include_line_graph:
             return self.model(g=g, l_g=l_g, state_attr=state_attr)
         return self.model(g, state_attr=state_attr)

tests/ext/test_ase_dgl.py

@@ -1,18 +1,21 @@
 from __future__ import annotations
 
 import os.path
+from unittest.mock import MagicMock, patch
 
 import numpy as np
 import pytest
 import torch
 from ase.build import molecule
+from ase.calculators.calculator import Calculator
 from pymatgen.io.ase import AseAtomsAdaptor
 
 import matgl
 from matgl import load_model
 
 if matgl.config.BACKEND != "DGL":
     pytest.skip("Skipping DGL tests", allow_module_level=True)
+from matgl.ext import _ase_dgl as ase_mod
 from matgl.ext.ase import Atoms2Graph, M3GNetCalculator, MolecularDynamics, PESCalculator, Relaxer
 
 
@@ -220,7 +223,7 @@ def test_get_graph_from_atoms_mol():
     assert np.allclose(state, [0.0, 0.0])
 
 
-@pytest.mark.skipif(True, reason="Too slow.")
+@pytest.mark.skipif(True, reason="Slow; superseded by test_molecular_dynamics_branches which mocks the calculator.")
 def test_molecular_dynamics(MoS2):
     pot = load_model("TensorNetDGL-PES-MatPES-PBE-2025.2")
     for ensemble in [
@@ -243,3 +246,72 @@ def test_molecular_dynamics(MoS2):
     md.run(1)
     with pytest.raises(ValueError, match="Ensemble not supported"):
         MolecularDynamics(MoS2, potential=pot, ensemble="notanensemble")
+
+
+def _fake_pes_calculator():
+    """Build a stand-in for ``PESCalculator`` that satisfies ASE's ``set_calculator``.
+
+    The real ``PESCalculator`` requires a trained Potential. We avoid downloading a
+    pretrained model in tests by substituting a minimal ASE ``Calculator`` whose
+    ``calculate`` is never invoked because the tests never run MD steps.
+    """
+    calc = MagicMock(spec=Calculator)
+    calc.results = {}
+    calc.parameters = {}
+    return calc
+
+
+def test_molecular_dynamics_branches(MoS2):
+    """Exercise every supported MD ensemble branch (and the invalid one) without
+    requiring a pretrained potential.
+
+    The MD constructors don't run any forces at init time, so a mock PESCalculator
+    is sufficient to hit all branches in ``MolecularDynamics.__init__`` and the
+    ``set_atoms`` helper.
+    """
+    fake_potential = MagicMock(spec=torch.nn.Module)
+
+    ensembles = [
+        "nvt",
+        "nve",
+        "nvt_langevin",
+        "nvt_andersen",
+        "nvt_bussi",
+        "nvt_nose_hoover_chain",
+        "npt",
+        "npt_berendsen",
+        "npt_nose_hoover",
+        "npt_nose_hoover_chain",
+    ]
+
+    with patch.object(ase_mod, "PESCalculator", side_effect=lambda **_kw: _fake_pes_calculator()):
+        for ensemble in ensembles:
+            md = MolecularDynamics(
+                MoS2,
+                potential=fake_potential,
+                ensemble=ensemble,
+                taut=0.1,
+                taup=0.1,
+                compressibility_au=10,
+            )
+            assert md.dyn is not None
+            # Re-wire calculator/atoms via the public helper.
+            md.set_atoms(MoS2)
+
+        # ``taut`` / ``taup`` defaults branch.
+        md_default = MolecularDynamics(
+            MoS2,
+            potential=fake_potential,
+            ensemble="npt_nose_hoover_chain",
+            taut=None,
+            taup=None,
+            compressibility_au=10,
+        )
+        assert md_default.dyn is not None
+
+        with pytest.raises(ValueError, match="Ensemble not supported"):
+            MolecularDynamics(MoS2, potential=fake_potential, ensemble="notanensemble")
+
+        # ``stress_weight`` is rejected with a warning.
+        with pytest.warns(UserWarning, match="Relaxer does not support user-defined stress_weight"):
+            MolecularDynamics(MoS2, potential=fake_potential, ensemble="nve", stress_weight=0.1)

tests/graph/test_converters.py

@@ -0,0 +1,21 @@
+"""Tests for the backend-dispatching ``matgl.graph.converters`` shim."""
+
+from __future__ import annotations
+
+import matgl
+from matgl.graph.converters import GraphConverter
+
+
+def test_converters_shim_dispatches_to_active_backend():
+    """``matgl.graph.converters.GraphConverter`` must alias the active backend's
+    converter base class.
+
+    The shim is a thin re-export and is imported by other modules (notably for
+    type-checking), so we just need to make sure the alias works on either backend.
+    """
+    if matgl.config.BACKEND == "DGL":
+        from matgl.graph._converters_dgl import GraphConverter as Expected
+    else:
+        from matgl.graph._converters_pyg import GraphConverter as Expected
+
+    assert GraphConverter is Expected

tests/test_config.py

@@ -1,18 +1,55 @@
 from __future__ import annotations
 
+import importlib.util
 import os.path
+from unittest.mock import patch
 
 import pytest
 
 import matgl
-from matgl.config import MATGL_CACHE, clear_cache
+from matgl.config import MATGL_CACHE, clear_cache, ensure_backend
 
 
 def test_clear_cache():
     clear_cache(False)
     assert not os.path.exists(MATGL_CACHE)
 
 
+def test_clear_cache_missing_dir(capsys):
+    """A second ``clear_cache`` call after the cache was already deleted must not raise."""
+    clear_cache(False)
+    clear_cache(False)
+    captured = capsys.readouterr()
+    assert "not found" in captured.out
+
+
+def test_clear_cache_no_when_user_says_no(monkeypatch):
+    """If the user answers 'n', the cache directory must remain untouched."""
+    os.makedirs(MATGL_CACHE, exist_ok=True)
+    answers = iter(["n"])
+    monkeypatch.setattr("builtins.input", lambda _prompt: next(answers))
+    clear_cache(confirm=True)
+    assert os.path.exists(MATGL_CACHE)
+
+
+def test_ensure_backend_dgl_missing_raises_runtime_error():
+    """Mocked-missing DGL must raise a ``RuntimeError`` from ``ensure_backend('DGL')``."""
+    with (
+        patch.object(importlib.util, "find_spec", side_effect=ImportError("nope")),
+        pytest.raises(RuntimeError, match="Please install DGL"),
+    ):
+        ensure_backend("DGL")
+
+
+def test_ensure_backend_pyg_missing_raises_runtime_error():
+    """Mocked-missing PyG must raise a ``RuntimeError`` from ``ensure_backend('PYG')``."""
+    with (
+        patch.object(importlib.util, "find_spec", side_effect=ImportError("nope")),
+        pytest.raises(RuntimeError, match="Please install torch_geometric"),
+    ):
+        ensure_backend("PYG")
+
+
 def test_set_backend():
     with pytest.raises(ValueError, match="Invalid backend"):
         matgl.set_backend("nonsense")

tests/utils/test_io.py

@@ -217,3 +217,149 @@ def test_load_bad_model():
             load_model("bad_serialized_model")
     finally:
         shutil.rmtree("bad_serialized_model")
+
+
+def test_load_model_wraps_unknown_error_as_runtime_error(tmp_path):
+    """Non-(ImportError, ValueError) exceptions should be re-raised as ``RuntimeError``."""
+    serialized_dir = tmp_path / "weird_model"
+    serialized_dir.mkdir()
+    for fn in ("model.pt", "state.pt", "model.json"):
+        (serialized_dir / fn).write_text("doesn't matter")
+
+    def boom(*_args, **_kwargs):
+        raise KeyboardInterrupt("boom")
+
+    with (
+        patch.object(matgl_io, "_get_file_paths", side_effect=boom),
+        pytest.raises(RuntimeError, match="Unknown error occurred while loading model"),
+    ):
+        load_model(serialized_dir)
+
+
+def test_get_file_paths_malformed_identifier_raises_value_error(tmp_path):
+    """An identifier that contains ``/`` but doesn't match a valid HF repo id must
+    fail with a clear ``ValueError`` rather than attempting a network call."""
+    from matgl.utils.io import _get_file_paths
+
+    # Doesn't exist locally and doesn't match ``owner/name`` (leading slash, double slash).
+    bogus = "//bad//repo//id"
+    with pytest.raises(ValueError, match=r"No valid model found locally or at Hugging Face Hub"):
+        _get_file_paths(tmp_path / bogus, str_path=bogus)
+
+
+def test_get_file_paths_bare_name_hub_failure_raises_value_error(tmp_path):
+    """When a bare name fails to download from the materialyze HF org, raise ``ValueError``."""
+    from matgl.utils.io import _get_file_paths
+
+    def boom(*_args, **_kwargs):
+        raise RuntimeError("simulated hub failure")
+
+    with (
+        patch.object(matgl_io, "_download_from_hf_hub", side_effect=boom),
+        pytest.raises(ValueError, match=r"No valid model found locally or at Hugging Face repo"),
+    ):
+        _get_file_paths(tmp_path / "BareName", str_path="BareName")
+
+
+def test_iomixin_load_dgl_class_under_pyg_warns():
+    """Loading a model whose nested kwargs reference a DGL-only class under PYG must warn.
+
+    Triggers the branch that auto-flips the backend to DGL when a serialized model has a
+    nested component class name containing ``m3gnet`` / ``megnet`` / ``chgnet`` / ``qet``.
+    """
+    import matgl as _matgl
+
+    if _matgl.config.BACKEND != "PYG":
+        pytest.skip("Only meaningful on the PyG backend.")
+
+    # IOMixIn.load expects a dict-of-paths or a Path. Pre-build the artifacts on disk.
+    import json as _json
+    import tempfile
+
+    with tempfile.TemporaryDirectory() as tmpdir:
+        tmp_path = Path(tmpdir)
+        nested = {
+            "@class": "M3GNet",
+            "@module": "definitely.not.a.real.module.at.all",
+            "@model_version": 1,
+            "init_args": {},
+        }
+        init_args = {"n": 1, "submodel": nested}
+
+        torch.save(init_args, tmp_path / "model.pt")
+        torch.save({}, tmp_path / "state.pt")
+        (tmp_path / "model.json").write_text(
+            _json.dumps(
+                {
+                    "@class": "OldModel",
+                    "@module": "tests.utils.test_io",
+                    "@model_version": 1,
+                    "metadata": None,
+                    "kwargs": init_args,
+                }
+            )
+        )
+
+        # ``matgl.set_backend("DGL")`` would mutate global state and may fail if DGL
+        # isn't installed; patch it to a no-op so the test is self-contained.
+        with (
+            patch.object(_matgl, "set_backend") as mock_set_backend,
+            pytest.warns(UserWarning, match=r"Setting the backend to DGL"),
+            pytest.raises((ImportError, ValueError, ModuleNotFoundError)),
+        ):
+            OldModel.load(tmp_path)
+
+        mock_set_backend.assert_called_with("DGL")
+
+
+def test_generate_hf_model_card_with_unserializable_metadata():
+    """``_generate_hf_model_card`` must swallow ``TypeError`` from non-serializable metadata.
+
+    Forces the ``json.dumps`` fallback to fail even with ``default=str`` by using an
+    object whose ``__repr__`` raises (and therefore so does ``str(obj)``).
+    """
+    from matgl.utils.io import _generate_hf_model_card
+
+    class Unserializable:
+        def __repr__(self):
+            raise TypeError("repr exploded")
+
+    model = OldModel(1)
+    card = _generate_hf_model_card(model, metadata={"oops": Unserializable()})
+
+    assert "## Metadata" not in card
+    assert "OldModel" in card
+
+
+def test_get_available_pretrained_models_handles_hub_errors():
+    """If the HF hub call fails, ``get_available_pretrained_models`` returns an empty list."""
+
+    class _BoomApi:
+        def list_models(self, **_kwargs):
+            raise RuntimeError("network is down")
+
+    with patch.object(matgl_io, "HfApi", return_value=_BoomApi()):
+        names = get_available_pretrained_models()
+
+    assert names == []
+
+
+def test_get_available_pretrained_models_strips_owner_prefix():
+    """Returned names should be bare (no ``"owner/"`` prefix) and sorted."""
+
+    class _FakeModelInfo:
+        def __init__(self, repo_id: str):
+            self.id = repo_id
+
+    class _FakeApi:
+        def list_models(self, **_kwargs):
+            return [
+                _FakeModelInfo("materialyze/Zeta"),
+                _FakeModelInfo("materialyze/Alpha"),
+                _FakeModelInfo("no-slash-id"),  # malformed entries are silently skipped
+            ]
+
+    with patch.object(matgl_io, "HfApi", return_value=_FakeApi()):
+        names = get_available_pretrained_models()
+
+    assert names == ["Alpha", "Zeta"]