MatRiFFlow 2.0.1__py3-none-any.whl

IMDgroup/__init__.py

@@ -0,0 +1,82 @@
+"""Deprecated compatibility package for the MatRiFFlow rename.
+
+``IMDgroup`` was the import package name before the project was renamed to
+MatRiFFlow (import package ``matrifflow``).  Importing anything under
+``IMDgroup`` emits a :class:`DeprecationWarning` and resolves to the
+corresponding ``matrifflow`` module, so existing scripts keep working
+during the transition.
+"""
+
+import importlib
+import importlib.abc
+import importlib.util
+import sys
+import warnings
+
+_LEGACY = "IMDgroup"
+_TARGET = "matrifflow"
+
+# The 8 CLI subcommand modules were renamed ``imdg*`` -> ``mrf*`` during the
+# MatRiFFlow rename.  The prefix-only mapping above cannot recover those, so
+# map their basenames explicitly.  Everything else shares its basename.
+_CLI_REMAP = {
+    "imdg": "mrf",
+    "imdg_create": "mrf_create",
+    "imdg_derive": "mrf_derive",
+    "imdg_diff": "mrf_diff",
+    "imdg_analyze": "mrf_analyze",
+    "imdg_status": "mrf_status",
+    "imdg_visualize": "mrf_visualize",
+    "imdg_run": "mrf_run",
+}
+
+warnings.warn(
+    "IMDgroup is deprecated; use matrifflow instead.",
+    DeprecationWarning,
+    stacklevel=2,
+)
+
+
+class _AliasLoader(importlib.abc.Loader):
+    """Loader that returns the already-imported real module.
+
+    Reusing the real module object (rather than re-executing its source)
+    means ``IMDgroup.X`` *is* ``matrifflow.X``, so ``isinstance`` checks,
+    pickling, and pymatgen's MSONable registry keep working across the two
+    import paths.
+    """
+
+    def __init__(self, real_name: str):
+        self._real_name = real_name
+
+    def create_module(self, spec):
+        _ = spec # unused
+        return importlib.import_module(self._real_name)
+
+    def exec_module(self, module):
+        _ = module # unused
+        # Already executed by create_module()'s import_module() call.
+        pass
+
+
+class _AliasFinder(importlib.abc.MetaPathFinder):
+    def find_spec(self, fullname, path=None, target=None):
+        _, _ = path, target # unused
+        if fullname == _LEGACY or not fullname.startswith(_LEGACY + "."):
+            return None
+        real_name = _TARGET + fullname[len(_LEGACY):]
+        pkg, _, leaf = real_name.rpartition(".")
+        if leaf in _CLI_REMAP:
+            real_name = f"{pkg}.{_CLI_REMAP[leaf]}"
+        real_spec = importlib.util.find_spec(real_name)
+        if real_spec is None:
+            return None
+        spec = importlib.util.spec_from_loader(fullname, _AliasLoader(real_name))
+        assert spec is not None
+        if real_spec.submodule_search_locations is not None:
+            spec.submodule_search_locations = list(real_spec.submodule_search_locations)
+        return spec
+
+
+if not any(isinstance(f, _AliasFinder) for f in sys.meta_path):
+    sys.meta_path.insert(0, _AliasFinder())

matrifflow/__init__.py

@@ -0,0 +1,39 @@
+# MIT License
+#
+# Copyright (c) 2024-2026 Inverse Materials Design Group
+#
+# Author: Ihor Radchenko <yantar92@posteo.net>
+#
+# This file is a part of MatRiFFlow package
+#
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in all
+# copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
+
+
+"""Pymatgen add-ons and cluster job-submission tools.
+
+Top-level namespace for the Inverse Materials Design group's tooling.
+Subpackages:
+
+    - ``matrifflow.pymatgen``: pymatgen add-ons (structure, I/O, diffusion,
+      transformations, CLI).
+    - ``matrifflow.gorun``: Slurm job-submission helpers (VASP, MACE, ATAT).
+    - ``matrifflow.utils``: plotting helpers.
+
+See https://www.oimalyi.org/ and https://yantar92.github.io/MatRiFFlow/
+"""

matrifflow/common.py

@@ -0,0 +1,150 @@
+# MIT License
+#
+# Copyright (c) 2024-2026 Inverse Materials Design Group
+#
+# Author: Ihor Radchenko <yantar92@posteo.net>
+#
+# This file is a part of MatRiFFlow package
+#
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in all
+# copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
+
+
+"""Common routines used across the sub-libraries."""
+import logging
+from pathlib import Path
+
+logger = logging.getLogger(__name__)
+
+
+def short_path(path: str) -> str:
+    """Return the final path component for compact log messages.
+
+    gorun working directories live under long cluster mount points
+    (e.g. ``/net/storage/.../generations.Li/.../relax.2``).  Repeating
+    the full path on every line drowns out the useful information, so
+    per-step messages only show the final component.  The working
+    directory is announced once at the start of a run.
+    """
+    name = Path(path).name
+    return name if name else str(path)
+
+
+def configure_progress_output():
+    """Send ``alive_progress`` output to stderr.
+
+    ``alive_progress`` writes progress bars and their final receipt to
+    stdout by default.  For CLI tools, stdout is reserved for actual
+    results, so redirect progress to stderr.  This is a process-wide
+    setting; call it once at each CLI entry point.
+    """
+    import sys
+    from alive_progress import config_handler
+    # alive_progress 3.1.5 monkey-patches ``set_global`` onto the
+    # ``config_handler`` callable.  Pyright only sees the function type
+    # and cannot resolve the attached attribute.
+    config_handler.set_global(file=sys.stderr)  # pyright: ignore[reportFunctionMemberAccess]
+
+
+def suppress_spglib_warnings():
+    """Silence spglib's raw C-level warnings.
+
+    spglib prints internal retry diagnostics directly to stderr when
+    symmetry detection needs several attempts (e.g. ``"spglib:
+    ssm_get_exact_positions failed."``, ``"spglib: No point group was
+    found"``).  They are harmless and duplicate the structured warnings
+    surfaced by matrifflow's own diagnostics, but they interleave with
+    rich tables and corrupt the display.
+
+    spglib reads ``SPGLIB_WARNING`` on every call (see its
+    ``warning_enabled``), so setting it at runtime is enough.  We use
+    ``setdefault`` so a user who has already chosen a value keeps
+    control; only the absent-key case is forced to ``OFF``.
+    """
+    import os
+    os.environ.setdefault("SPGLIB_WARNING", "OFF")
+
+
+def groupby_cmp(lst, cmp_eq, title_function=None):
+    """Group elements of a list by pairwise equality.
+
+    Each item is compared only against the first member of every existing
+    group, so the number of ``cmp_eq`` calls is O(n * g), where ``g`` is
+    the number of groups produced, instead of O(n^2) over group members.
+
+    This is exact when ``cmp_eq`` is an equivalence relation (reflexive,
+    symmetric, and transitive).  A non-transitive ``cmp_eq`` still works,
+    but it is only compared against the first member of each group, so
+    items that would match a later member but not the representative may
+    be split into separate groups.
+
+    Args:
+        lst: List of elements to group.
+        cmp_eq: Function that takes two elements and returns True if they
+            should belong to the same group.  Ideally an equivalence
+            relation; see above for the behavior when it is not
+            transitive.
+        title_function: Optional function that takes an element and returns
+            a string used in log messages.  When None, "??" is logged.
+
+    Returns:
+        list[list]: List of grouped lists.  Each sublist contains elements
+        that are pairwise equivalent according to cmp_eq.
+    """
+    groups = []
+    for item in lst:
+        title = title_function(item) if title_function is not None else "??"
+        for group in groups:
+            if cmp_eq(item, group[0]):
+                logger.debug('Appending %s to existing group', title)
+                group.append(item)
+                break
+        else:
+            logger.debug('Creating a new group for %s', title)
+            groups.append([item])
+    return groups
+
+
+def deprecated_aliases(module_globals, **aliases):
+    """Return a module-level ``__getattr__`` resolving deprecated names.
+
+    ``aliases`` maps each legacy name to the current name present in
+    ``module_globals``.  Accessing a legacy name emits a
+    :class:`DeprecationWarning` and returns the current object, so old
+    ``IMD*`` / ``IMDG*`` class names keep working after the MatRiFFlow
+    rename.
+
+    Usage (at the bottom of a module)::
+
+        __getattr__ = deprecated_aliases(globals(), IMDGVaspDir="MRFVaspDir")
+    """
+    import warnings
+
+    def __getattr__(name):
+        if name in aliases:
+            warnings.warn(
+                f"{name} is deprecated; use {aliases[name]} instead.",
+                DeprecationWarning,
+                stacklevel=2,
+            )
+            return module_globals[aliases[name]]
+        raise AttributeError(
+            f"module {module_globals['__name__']!r} has no attribute {name!r}"
+        )
+
+    return __getattr__

matrifflow/gorun/__init__.py

@@ -0,0 +1,36 @@
+# MIT License
+#
+# Copyright (c) 2024-2026 Inverse Materials Design Group
+#
+# Author: Ihor Radchenko <yantar92@posteo.net>
+#
+# This file is a part of MatRiFFlow package
+#
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in all
+# copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
+
+
+"""Slurm job submission helpers tailored for Inverse Materials Design group.
+
+Supports:
+    - VASP: ``gorun vasp`` (or plain ``gorun`` for backward compat)
+    - MACE fine-tuning: ``gorun mace`` (or ``gorun-mace``)
+    - ATAT cluster expansion: ``gorun-maps``, ``gorun-atat-local``
+
+See https://www.oimalyi.org/ and https://yantar92.github.io/MatRiFFlow/
+"""

matrifflow/gorun/adapters/__init__.py

@@ -0,0 +1,32 @@
+# MIT License
+#
+# Copyright (c) 2024-2026 Inverse Materials Design Group
+#
+# Author: Ihor Radchenko <yantar92@posteo.net>
+#
+# This file is a part of MatRiFFlow package
+#
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in all
+# copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
+
+
+"""Software adapters for gorun.
+
+Each module in this package implements the `SoftwareAdapter` protocol
+defined in `adapters.base`.
+"""

matrifflow/gorun/adapters/_mace_train_masked.py

@@ -0,0 +1,218 @@
+# MIT License
+#
+# Copyright (c) 2024-2026 Inverse Materials Design Group
+#
+# Author: Ihor Radchenko <yantar92@posteo.net>
+#
+# This file is a part of MatRiFFlow package
+#
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in all
+# copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
+
+"""Wrapper around ``mace.cli.run_train`` that replaces the
+energy/forces/stress loss with a mask-aware variant.
+
+When ``ref.atom_mask`` is present on a batch (per-atom float tensor,
+1 = supervised, 0 = masked), masked atoms' site energies and atomic
+stresses are detached from autograd, preventing direct energy/stress
+training on those atoms while preserving indirect training through
+neighbour message passing.  When ``ref.atom_mask`` is absent (e.g.
+pt_head replay data), the standard loss is used unchanged.
+
+Usage:  python -u -m matrifflow.gorun.adapters._mace_train_masked [ARGS...]
+
+Does *not* require a custom ``--loss`` value -- keep using
+``--loss stress``.  The ``configure_model`` stress check and SWA
+loss selection both work because the original class is replaced
+in-place.
+"""
+
+from __future__ import annotations
+
+import logging
+
+import torch
+
+import mace.modules.loss as mace_loss
+from mace.tools.scatter import scatter_sum
+
+logger = logging.getLogger(__name__)
+
+#: Threshold above which a mask entry is considered "supervised"
+#: (contributes to the loss).  Mask values are 0 (masked) or 1
+#: (supervised), but a float comparison avoids strict equality.
+MASK_THRESHOLD: float = 0.5
+
+
+## Masked loss class (replaces WeightedEnergyForcesStressLoss)
+
+class MaskedEnergyForcesStressLoss(
+    mace_loss.WeightedEnergyForcesStressLoss,
+):
+    """Standard E+F+S loss with per-atom masking when ``ref.atom_mask``
+    is present on the batch."""
+
+    # Pyright infers torch.Tensor | Module which should never happen
+    # in runtime.
+    energy_weight: torch.Tensor
+    forces_weight: torch.Tensor
+    stress_weight: torch.Tensor
+
+    def forward(
+        self,
+        ref,
+        pred,
+        ddp: bool | None = None,
+    ) -> torch.Tensor:
+        mask = getattr(ref, 'atom_mask', None)
+        if mask is None:
+            return super().forward(ref, pred, ddp)
+
+        return (
+            self.energy_weight * _masked_energy(ref, pred, mask, ddp)
+            + self.forces_weight * _masked_forces(ref, pred, mask, ddp)
+            + self.stress_weight * _masked_stress(ref, pred, mask, ddp)
+        )
+
+
+## Replace the class in-place so that get_loss_fn and get_swa
+## pick up the masked variant automatically.
+
+mace_loss.WeightedEnergyForcesStressLoss = MaskedEnergyForcesStressLoss
+
+
+## Masked helpers
+
+def _masked_energy(
+    ref,
+    pred,
+    mask: torch.Tensor,
+    ddp: bool | None,
+) -> torch.Tensor:
+    """Compute masked energy loss.
+
+    Node energies of masked atoms are detached from autograd
+    so they contribute to the per-structure total (and thus
+    indirectly to training via message passing) but their
+    site-energy gradient is zero.
+    """
+    supervised_float = (
+        (mask > MASK_THRESHOLD)
+        .to(dtype=pred['node_energy'].dtype)
+    )
+    node_energy = pred['node_energy']
+    hybrid_node_energy = (
+        node_energy * supervised_float
+        + node_energy.detach() * (1.0 - supervised_float)
+    )
+
+    predicted_energy = scatter_sum(
+        hybrid_node_energy,
+        ref['batch'],
+        dim=0,
+        dim_size=ref.num_graphs,
+    )
+    natoms_per_config = ref.ptr[1:] - ref.ptr[:-1]
+    raw_loss = (
+        ref.weight
+        * ref.energy_weight
+        * torch.square(
+            (ref['energy'] - predicted_energy) / natoms_per_config,
+        )
+    )
+    return mace_loss.reduce_loss(raw_loss, ddp)
+
+
+def _masked_forces(
+    ref,
+    pred,
+    mask: torch.Tensor,
+    ddp: bool | None,
+) -> torch.Tensor:
+    """Compute masked forces loss.
+
+    Masked atoms get zero gradient for the force residual --
+    only supervised atoms drive the force loss.
+    """
+    mask_2d = mask.view(-1, 1)
+    natoms_per_config = ref.ptr[1:] - ref.ptr[:-1]
+    config_weight = torch.repeat_interleave(
+        ref.weight, natoms_per_config,
+    ).unsqueeze(-1)
+    config_forces_weight = torch.repeat_interleave(
+        ref.forces_weight, natoms_per_config,
+    ).unsqueeze(-1)
+    raw_loss = (
+        config_weight
+        * config_forces_weight
+        * mask_2d
+        * torch.square(ref['forces'] - pred['forces'])
+    )
+    return mace_loss.reduce_loss(raw_loss, ddp)
+
+
+def _masked_stress(
+    ref,
+    pred,
+    mask: torch.Tensor,
+    ddp: bool | None,
+) -> torch.Tensor:
+    """Compute masked stress loss.
+
+    When atomic stresses are available, masked atoms' stress
+    contributions are detached.  Falls back to the standard
+    per-structure stress loss when atomic stresses are absent.
+    """
+    if (
+        'atomic_stresses' not in pred
+        or pred['atomic_stresses'] is None
+    ):
+        return mace_loss.weighted_mean_squared_stress(ref, pred, ddp)
+
+    supervised_float = (
+        (mask > MASK_THRESHOLD)
+        .to(dtype=pred['atomic_stresses'].dtype)
+        .view(-1, 1, 1)
+    )
+    atomic_stresses = pred['atomic_stresses']
+    hybrid_atomic_stresses = (
+        atomic_stresses * supervised_float
+        + atomic_stresses.detach() * (1.0 - supervised_float)
+    )
+
+    predicted_stress = scatter_sum(
+        hybrid_atomic_stresses,
+        ref['batch'],
+        dim=0,
+        dim_size=ref.num_graphs,
+    )
+    config_weight = ref.weight.view(-1, 1, 1)
+    config_stress_weight = ref.stress_weight.view(-1, 1, 1)
+    raw_loss = (
+        config_weight
+        * config_stress_weight
+        * torch.square(ref['stress'] - predicted_stress)
+    )
+    return mace_loss.reduce_loss(raw_loss, ddp)
+
+
+## Delegate to MACE
+
+if __name__ == '__main__':
+    from mace.cli.run_train import main
+    main()