EquFlash 0.0.1__tar.gz

equflash-0.0.1/GGNN/_cli.py

@@ -0,0 +1,113 @@
+"""
+This file is derived from fairchem-core.
+Original source: https://github.com/facebookresearch/fairchem
+
+Copyright (c) Meta Platforms, Inc. and affiliates.
+Licensed under the MIT License. See LICENSE for details.
+
+Modifications copyright (c) 2026 Samsung Electronics.
+Licensed under CC BY-NC-SA 4.0. See LICENSE for details.
+"""
+
+from __future__ import annotations
+
+import copy
+import logging
+from typing import TYPE_CHECKING
+
+from submitit import AutoExecutor
+from submitit.helpers import Checkpointable, DelayedSubmission
+from torch.distributed.launcher.api import LaunchConfig, elastic_launch
+
+from fairchem.core.common import distutils
+from fairchem.core.common.flags import flags
+from fairchem.core.common.utils import (
+    build_config,
+    create_grid,
+    save_experiment_log,
+    setup_logging,
+)
+from GGNN.common.utils import new_trainer_context
+
+if TYPE_CHECKING:
+    import argparse
+
+
+class Runner(Checkpointable):
+    def __init__(self) -> None:
+        self.config = None
+
+    def __call__(self, config: dict) -> None:
+        with new_trainer_context(config=config) as ctx:
+            self.config = ctx.config
+            self.task = ctx.task
+            self.trainer = ctx.trainer
+            self.task.setup(self.trainer)
+            self.task.run()
+
+    def checkpoint(self, *args, **kwargs):
+        new_runner = Runner()
+        self.trainer.save(checkpoint_file="checkpoint.pt", training_state=True)
+        self.config["checkpoint"] = self.task.chkpt_path
+        self.config["timestamp_id"] = self.trainer.timestamp_id
+        if self.trainer.logger is not None:
+            self.trainer.logger.mark_preempting()
+        logging.info(
+            f'Checkpointing callback is triggered, checkpoint saved to: {self.config["checkpoint"]}, timestamp_id: {self.config["timestamp_id"]}'
+        )
+        return DelayedSubmission(new_runner, self.config)
+
+
+def runner_wrapper(config: dict):
+    Runner()(config)
+
+
+def main(
+    args: argparse.Namespace | None = None, override_args: list[str] | None = None
+):
+    """Run the main fairchem program."""
+    setup_logging()
+
+    if args is None:
+        parser: argparse.ArgumentParser = flags.get_parser()
+        args, override_args = parser.parse_known_args()
+
+    assert (
+        args.num_gpus > 0
+    ), "num_gpus is used to determine number ranks, so it must be at least 1"
+    config = build_config(args, override_args)
+
+    if args.submit:  # Run on cluster
+        slurm_add_params = config.get("slurm", None)  # additional slurm arguments
+        configs = create_grid(config, args.sweep_yml) if args.sweep_yml else [config]
+
+        logging.info(f"Submitting {len(configs)} jobs")
+        executor = AutoExecutor(folder=args.logdir / "%j", slurm_max_num_timeout=3)
+        executor.update_parameters(
+            name=args.identifier,
+            mem_gb=args.slurm_mem,
+            timeout_min=args.slurm_timeout * 60,
+            slurm_partition=args.slurm_partition,
+            gpus_per_node=args.num_gpus,
+            cpus_per_task=(config["optim"]["num_workers"] + 1),
+            tasks_per_node=args.num_gpus,
+            nodes=args.num_nodes,
+            slurm_additional_parameters=slurm_add_params,
+            slurm_qos=args.slurm_qos,
+            slurm_account=args.slurm_account,
+        )
+        for config in configs:
+            config["slurm"] = copy.deepcopy(executor.parameters)
+            config["slurm"]["folder"] = str(executor.folder)
+        jobs = executor.map_array(Runner(), configs)
+        logging.info(f"Submitted jobs: {', '.join([job.job_id for job in jobs])}")
+        log_file = save_experiment_log(args, jobs, configs)
+        logging.info(f"Experiment log saved to: {log_file}")
+
+    else:  # Run locally on a single node, n-processes
+        runner_wrapper(config)
+
+
+
+if __name__ == "__main__":
+    main()

equflash-0.0.1/GGNN/common/calculator.py

@@ -0,0 +1,206 @@
+"""
+This file is derived from fairchem-core.
+Original source: https://github.com/facebookresearch/fairchem
+
+Copyright (c) Meta Platforms, Inc. and affiliates.
+Licensed under the MIT License. See LICENSE for details.
+
+Modifications copyright (c) 2026 Samsung Electronics.
+Licensed under CC BY-NC-SA 4.0. See LICENSE for details.
+"""
+
+from fairchem.core.common.relaxation.ase_utils import OCPCalculator
+from ..trainer.trainer import Trainer
+
+import copy
+import logging
+from types import MappingProxyType
+from typing import TYPE_CHECKING
+
+import torch
+from ase import Atoms
+from ase.calculators.calculator import Calculator
+from ase.calculators.singlepoint import SinglePointCalculator
+from ase.constraints import FixAtoms
+from ase.geometry import wrap_positions
+
+from fairchem.core.common.registry import registry
+from fairchem.core.common.utils import (
+    load_config,
+    setup_imports,
+    setup_logging,
+    update_config,
+)
+from fairchem.core.datasets import data_list_collater
+from fairchem.core.models.model_registry import model_name_to_local_file
+from fairchem.core.preprocessing import AtomsToGraphs
+from collections import OrderedDict
+
+
+def convert_compiled_ckpt(ckpt):
+    if ckpt["config"]["model"]["name"].endswith("_comp"):
+        ckpt["config"]["model"]["name"] = ckpt["config"]["model"]["name"].replace(
+            "_comp", ""
+        )
+        new_dict = OrderedDict()
+        for k, v in ckpt["state_dict"].items():
+            if k.startswith("module.model") and "z_to_onehot_tensor" not in k:
+                new_dict[k.replace(".model.", ".")] = v
+        ckpt["state_dict"] = new_dict
+    if "task" in ckpt["config"]:
+        ckpt["config"].pop("task")
+    return ckpt
+
+
+if TYPE_CHECKING:
+    from pathlib import Path
+
+    from torch_geometric.data import Batch
+
+
+class UCalculator(OCPCalculator):
+    """ASE based calculator using an OCP model"""
+
+    def __init__(
+        self,
+        config_yml: str | None = None,
+        checkpoint_path: str | None = None,
+        model_name: str | None = None,
+        local_cache: str | None = None,
+        trainer: str | None = None,
+        cpu: bool = True,
+        seed: int | None = None,
+        only_output: list[str] | None = None,
+    ) -> None:
+
+        setup_imports()
+        setup_logging()
+        Calculator.__init__(self)
+
+        if model_name is not None:
+            if checkpoint_path is not None:
+                raise RuntimeError(
+                    "model_name and checkpoint_path were both specified, please use only one at a time"
+                )
+            if local_cache is None:
+                raise NotImplementedError(
+                    "Local cache must be set when specifying a model name"
+                )
+            checkpoint_path = model_name_to_local_file(
+                model_name=model_name, local_cache=local_cache
+            )
+
+        # Either the config path or the checkpoint path needs to be provided
+        assert config_yml or checkpoint_path is not None
+
+        checkpoint = None
+        if config_yml is not None:
+            if isinstance(config_yml, str):
+                config, duplicates_warning, duplicates_error = load_config(config_yml)
+                if len(duplicates_warning) > 0:
+                    logging.warning(
+                        f"Overwritten config parameters from included configs "
+                        f"(non-included parameters take precedence): {duplicates_warning}"
+                    )
+                if len(duplicates_error) > 0:
+                    raise ValueError(
+                        f"Conflicting (duplicate) parameters in simultaneously "
+                        f"included configs: {duplicates_error}"
+                    )
+            else:
+                config = config_yml
+            if config["model"]["name"].startswith("compiled"):
+                config["model"]["name"] = config["model"]["name"].replace(
+                    "compiled_", ""
+                )
+
+            # Only keeps the train data that might have normalizer values
+            if isinstance(config["dataset"], list):
+                config["dataset"] = config["dataset"][0]
+            elif isinstance(config["dataset"], dict):
+                config["dataset"] = config["dataset"].get("train", None)
+        else:
+            # Loads the config from the checkpoint directly (always on CPU).
+            checkpoint = torch.load(
+                checkpoint_path, map_location=torch.device("cpu"), weights_only=False
+            )
+            checkpoint = convert_compiled_ckpt(checkpoint)
+            config = checkpoint["config"]
+
+        if trainer is not None:
+            config["trainer"] = trainer
+        else:
+            config["trainer"] = config.get("trainer", "default")
+
+        if "model_attributes" in config:
+            config["model_attributes"]["name"] = config.pop("model")
+            config["model"] = config["model_attributes"]
+
+        # Calculate the edge indices on the fly
+        config["model"]["otf_graph"] = True
+
+        ### backwards compatability with OCP v<2.0
+        config = update_config(config)
+
+        self.config = copy.deepcopy(config)
+        self.config["checkpoint"] = str(checkpoint_path)
+        del config["dataset"]["src"]
+
+        # some models that are published have configs that include tasks
+        # which are not output by the model
+        if only_output is not None:
+            assert isinstance(
+                only_output, list
+            ), "only output must be a list of targets to output"
+            for key in only_output:
+                assert (
+                    key in config["outputs"]
+                ), f"{key} listed in only_outputs is not present in current model outputs {config['outputs'].keys()}"
+            remove_outputs = set(config["outputs"].keys()) - set(only_output)
+            for key in remove_outputs:
+                config["outputs"].pop(key)
+
+        trainer_cls = Trainer
+
+        self.trainer = trainer_cls(
+            task=config.get("task", {}),
+            model=config["model"],
+            dataset=[config["dataset"]],
+            outputs=config["outputs"],
+            loss_functions=config["loss_functions"],
+            evaluation_metrics=config["evaluation_metrics"],
+            optimizer=config["optim"],
+            identifier="",
+            slurm=config.get("slurm", {}),
+            local_rank=config.get("local_rank", 0),
+            is_debug=config.get("is_debug", True),
+            cpu=cpu,
+            amp=config.get("amp", False),
+            inference_only=True,
+        )
+
+        if checkpoint_path is not None:
+            self.load_checkpoint(checkpoint_path=checkpoint_path, checkpoint=checkpoint)
+
+        seed = seed if seed is not None else self.trainer.config["cmd"]["seed"]
+        if seed is None:
+            logging.warning(
+                "No seed has been set in modelcheckpoint or OCPCalculator! Results may not be reproducible on re-run"
+            )
+        else:
+            self.trainer.set_seed(seed)
+
+        if "otf_graph" in config["model"]:  # temporalily impemented.
+            r_edges = not config["model"]["otf_graph"]
+        else:
+            r_edges = not self.trainer.model.otf_graph
+
+        self.a2g = AtomsToGraphs(
+            r_energy=False,
+            r_forces=False,
+            r_distances=False,
+            r_pbc=True,
+            # r_edges=not self.trainer.model.otf_graph,  # otf graph should not be a property of the model
+            r_edges=r_edges,
+        )
+        self.implemented_properties = list(self.config["outputs"].keys())

equflash-0.0.1/GGNN/common/lammps_mliap/lmp_mliap_wrapper.py

@@ -0,0 +1,211 @@
+# This file is a part of the `nequip` package. Please see LICENSE and README at the root for information on using it.
+
+import torch
+import os
+
+
+try:
+    from lammps.mliap.mliap_unified_abc import MLIAPUnified
+except ModuleNotFoundError:
+    raise ImportError(
+        "LAMMPS ML-IAP has to be installed in the Python environment for NequIP's ML-IAP integration. "
+        "See https://nequip.readthedocs.io/en/latest/integrations/lammps/mliap.html for installation instructions."
+    )
+
+from typing import List
+from GGNN.common.calculator import UCalculator,convert_compiled_ckpt
+from GGNN.model.SevenNet.sevennet import SevenNet, convert_compiled_to_original_sevennet
+
+from GGNN.model.EquFlash import EquFlash
+from fairchem.core.common.utils import (
+    match_state_dict,
+)
+from fairchem.core.common.registry import registry
+from fairchem.core.modules.normalization.element_references import (
+    create_element_references,
+)
+from fairchem.core.modules.normalization.normalizer import (
+    create_normalizer,
+)
+
+from torch_geometric.data import Batch
+
+
+class GGNNLAMMPSMLIAPWrapper(MLIAPUnified):
+    """LAMMPS-MLIAP interface for NequIP framework models."""
+
+    def __init__(
+        self,
+        ckpt_path: str,
+        **kwargs,
+    ):
+        # this is a white lie, unsure if strictly necessary, but just in case
+        super().__init__()
+        self.ckpt_path = ckpt_path
+        ckpt = torch.load(ckpt_path, map_location="cpu", weights_only=False)
+        self.model = None
+        self.device = None
+        # model_name = ckpt["config"]["model"].pop("name")
+        # self.model = registry.get_model_class(model_name)(self.ckpt["config"]["model"])
+        # self.model = registry.get_model_class(model_name)(ckpt["config"]["model"])
+        # import pdb;pdb.set_trace()
+        # to placate the interface
+        self.nparams = 1
+        self.ndescriptors = 1
+
+        # === set model-depnedent params ===
+        # what is rcutfac
+        # self.rcutfac = 0.5 * float(model.metadata[graph_model.R_MAX_KEY])
+        self.rcutfac = 3.0
+        import ase
+
+        self.element_types = (
+            ase.data.chemical_symbols
+        )  # ckpt["config"]["model"]["chemical_species"]
+        self.ckpt = ckpt
+
+    def compute_forces(self, lmp_data):
+        # === lazily load model ===
+        if self.model == None:
+            self.device = (
+                "cuda" if "kokkos" in lmp_data.__class__.__module__.lower() else "cpu"
+            )
+            self.ckpt = convert_compiled_ckpt(self.ckpt)
+            model_name = self.ckpt["config"]["model"].pop("name")
+            # self.model = registry.get_model_class(model_name)(self.ckpt["config"]["model"])
+            if "conv_kwargs" not in self.ckpt["config"]["model"].keys():
+                self.ckpt["config"]["model"]["conv_kwargs"] = {}
+            self.ckpt["config"]["model"]["conv_kwargs"]["use_lammps_mliap"] = True
+            
+
+            self.model = registry.get_model_class(model_name)(
+                self.ckpt["config"]["model"], lammps_mliap=True
+            ).to(self.device)
+
+            new_dict = match_state_dict(
+                self.model.state_dict(), self.ckpt["state_dict"]
+            )
+            self.model.load_state_dict(new_dict, strict=False)
+            # print(os.environ.get("OMPI_COMM_WORLD_RANK", "0"))
+            rank = eval(os.environ.get("OMPI_COMM_WORLD_RANK", "0"))
+            if rank == 0:
+                print(self.model)
+            self.normalizers = {}
+            for key, state_dict in self.ckpt.get("normalizers", {}).items():
+                ### Convert old normalizer keys to new target keys
+                if key == "target":
+                    target_key = "energy"
+                elif key == "grad_target":
+                    target_key = "forces"
+                else:
+                    target_key = key
+
+                if target_key not in self.normalizers:
+                    self.normalizers[target_key] = create_normalizer(
+                        state_dict=state_dict
+                    )
+                self.normalizers[target_key].to(self.device)
+            self.elementrefs = {}
+            for key, state_dict in self.ckpt.get("elementrefs", {}).items():
+                if key not in self.elementrefs:
+                    self.elementrefs[key] = create_element_references(
+                        state_dict=state_dict
+                    )
+                else:
+                    mkeys = self.elementrefs[key].load_state_dict(state_dict)
+                    assert len(mkeys.missing_keys) == 0
+                    assert len(mkeys.unexpected_keys) == 0
+
+                self.elementrefs[key].to(self.device)
+            self.model.regress_forces = False
+            self.model.regress_stress = False
+
+        if lmp_data.nlocal == 0 or lmp_data.npairs <= 1:
+            return
+
+        # === create input data ===
+
+        # NOTE
+        # This LAMMPS ML-IAP integration introduces a new dimension of having `num_local` vs `num_local + num_ghost` number of nodes.
+        # There are three crucial dimensions to be aware of `num_edges`, `num_local`, `num_local + num_ghost`.
+        # The following input tensors have the following shapes.
+        # - `edge_vectors`: (num_edges, 3)
+        # - `edge_idxs`: (2, num_edges)
+        # - `atom_types`: (num_local + num_ghost)
+
+        # This LAMMPS ML-IAP wrapper can handle output `atomic_energy` having either shape `num_local` or `num_local + num_ghost` based on the (uncompiled) size check.
+
+        # Models can perform optimizations based on an understanding of when ghost atoms matter or not and are responsible for carefully handling internal shape logic.
+        # Examples include:
+        # - edge -> node scatter operations / nodewise operations (e.g. in `nequip/nn/interaction_block.py`)
+        # - nodewise operations that involve `atom_types` (since `atom_types` is `num_local + num_ghost`), e.g. in `PerTypeScaleShift` and `ZBL`.
+
+        # TODO: we have yet to exploit per-edge-type cutoffs by pruning the edge vectors and neighborlist
+        # make sure edge vectors `requires_grad`
+        edge_vectors = (
+            torch.as_tensor(lmp_data.rij, dtype=torch.float64).to(self.device).float()
+        )
+        edge_vectors.requires_grad_(True)
+        edge_index = torch.vstack(
+            [
+                torch.as_tensor(lmp_data.pair_i, dtype=torch.int64).to(self.device),
+                torch.as_tensor(lmp_data.pair_j, dtype=torch.int64).to(self.device),
+            ],
+        )
+
+        atomic_type = torch.as_tensor(lmp_data.elems, dtype=torch.int64).to(self.device)
+
+        num_atoms = torch.tensor(
+            [lmp_data.nlocal, lmp_data.ntotal - lmp_data.nlocal], dtype=torch.int64
+        ).to(self.device)
+        natoms = torch.tensor(lmp_data.nlocal, dtype=torch.int64).to(self.device)
+        # import pdb;pdb.set_trace()
+        # === run model ===
+        # run model and backwards for edge forces
+        batch = Batch(
+            edge_vec=edge_vectors,
+            atomic_numbers=(atomic_type[: lmp_data.nlocal].long()),
+            edge_index=edge_index,
+            num_atoms=num_atoms,
+            natoms=natoms,
+            lmp_data=lmp_data,
+        )
+        out = self.model(batch)
+        # correct sign convention for consistency with LAMMPS
+        edge_forces = torch.autograd.grad(
+            [out["energy"].sum()],
+            [edge_vectors],
+        )[0]
+        # === pass outputs to LAMMPS ===
+        # handle ghosts
+        atomic_energies = out["atomic_energy"].squeeze()
+
+        # = nequip_data_out[
+        #     AtomicDataDict.PER_ATOM_ENERGY_KEY
+        # ].view(-1)
+
+        # shape-dependent control flow, but should be outside of compiled model
+        if atomic_energies.size(0) != lmp_data.nlocal:
+            atomic_energies = torch.narrow(atomic_energies, 0, 0, lmp_data.nlocal)
+            total_energy = torch.sum(atomic_energies)
+        else:
+            total_energy = torch.sum(atomic_energies)
+        if "forces" in self.normalizers:
+            edge_forces = self.normalizers["forces"](edge_forces)
+        if "energy" in self.normalizers:
+            total_energy = self.normalizers["energy"](total_energy)
+        if "energy" in self.elementrefs:
+            batch.batch = torch.zeros(natoms, device=self.device, dtype=torch.int)
+            total_energy = self.elementrefs["energy"](
+                total_energy.reshape(1, -1), batch
+            )
+
+        # update LAMMPS variables
+        lmp_data.energy = total_energy
+        lmp_data.update_pair_forces_gpu(edge_forces.double())
+
+    def compute_descriptors(self, lmp_data):
+        pass
+
+    def compute_gradients(self, lmp_data):
+        pass

equflash-0.0.1/GGNN/common/logger.py

@@ -0,0 +1,102 @@
+import os
+import logging
+import torch
+import yaml
+import math
+from typing import Any
+
+from fairchem.core.common.logger import Logger
+from fairchem.core.common.registry import registry
+
+import datetime
+
+def parse_logs(update_dict):
+    ts = datetime.datetime.now().timestamp()  
+    dt = datetime.datetime.fromtimestamp(ts)  
+    formatted = dt.strftime("%Y-%m-%d %H:%M:%S")
+    ss = f"{formatted} "
+    if "epoch" in update_dict:
+        ep = update_dict["epoch"]
+        ss += f"epoch {ep:.1f}"
+    if "step" in update_dict:
+        step = update_dict["step"]
+        ss += f" (step {int(step)})"
+    if ss != "":
+        ss += ":"
+    for key, val in update_dict.items():
+        if key in ["epoch", "step"]:
+            continue
+
+        mse_metric = "mse" in key
+        if mse_metric:
+            key = key.replace("mse", "rmse")
+
+        if torch.is_tensor(val):
+            if mse_metric:
+                val = torch.sqrt(val)
+            ss += f" {key} {val.item():.5f}"
+        elif isinstance(val, float):
+            if key == "lr":
+                ss += f" {key} {val:.2e}"
+            else:
+                if mse_metric:
+                    val = math.sqrt(val)
+                ss += f" {key} {val:.5f}"
+        else:
+            ss += f" {key} {val}"
+    return ss
+
+
+@registry.register_logger("files")
+class FilesLogger(Logger):
+    def __init__(self, config):
+        super().__init__(config)
+
+        logdir = self.config["cmd"]["logs_dir"] 
+        self.log_path = {"train" : os.path.join(logdir, "train.log")}
+        self.log_path["val"] = os.path.join(logdir, "val.log")
+        self.log_path["test"] = os.path.join(logdir, "test.log")
+
+    def watch(self, model):
+        logging.warning(
+            "Model gradient logging to files is not supported."
+        )
+        return False
+
+    def log(self, update_dict, step=None, split=""):
+        assert split in ["train", "val", "test"], f"Split {split} is not supported"
+        outfile = open(self.log_path[split], 'a')
+
+        ss = parse_logs(update_dict)
+        outfile.write(ss + "\n")
+        outfile.close()
+
+    def log_plots(self, plots):
+        pass
+
+    def mark_preempting(self):
+        pass
+
+    def log_model_training_info(self, model=None):
+        model_log_path = os.path.join(self.config["cmd"]["logs_dir"], "model_training_info.yml")
+        outfile = open(model_log_path, 'w')
+        outfile.write(yaml.dump(self.config, default_flow_style=False))
+        outfile.write("\n")
+        
+        if model:
+            outfile.write(str(model)+ "\n")
+            outfile.write(f"model num of parameters: {model.num_params}\n")
+        outfile.close()
+
+    def log_final_metrics(self, table, time=None):
+        log_path = os.path.join(self.config["cmd"]["logs_dir"], "final_metrics.log")
+        outfile = open(log_path, 'w')
+        outfile.write(str(table)+"\n")
+        if time:
+            outfile.write(f"train() elapsed time: {time:.1f} sec ({time/3600.0:.1f} h)\n")
+        outfile.close()
+    def log_summary(self, summary_dict: dict[str, Any]) -> None:
+        logging.warning("log_summary for Files not supported")
+
+    def log_artifact(self, name: str, type: str, file_location: str) -> None:
+        logging.warning("log_artifact for Files not supported")