UniPath/src/flowmm/old_eval/core.py

"""Copyright (c) Meta Platforms, Inc. and affiliates."""

from __future__ import annotations

import json
import os
import sys
from collections import Counter
from pathlib import Path

import numpy as np
import pandas as pd
import torch
from matminer.featurizers.composition.composite import ElementProperty
from matminer.featurizers.site.fingerprint import CrystalNNFingerprint
from pymatgen.core.composition import Composition
from pymatgen.core.lattice import Lattice
from pymatgen.core.structure import Structure
from torch_geometric.data import Data

from diffcsp.eval_utils import (
    get_crystals_list,
    load_data,
    smact_validity,
    structure_validity,
)
from flowmm.data import NUM_ATOMIC_BITS, NUM_ATOMIC_TYPES
from flowmm.joblib_ import joblib_map
from flowmm.rfm.manifold_getter import ManifoldGetter
from flowmm.rfm.manifolds.analog_bits import analog_bits_to_int

CrysArrayListType = list[dict[str, np.ndarray]]

CrystalNNFP = CrystalNNFingerprint.from_preset("ops")
CompFP = ElementProperty.from_preset("magpie")


class Crystal(object):
    def __init__(
        self,
        crys_array_dict,
        make_atom_types_discrete: bool = True,
    ):
        self.frac_coords = crys_array_dict["frac_coords"]
        self.atom_types = crys_array_dict["atom_types"]

        if make_atom_types_discrete and (self.frac_coords.ndim == self.atom_types.ndim):
            if self.atom_types.shape[-1] == NUM_ATOMIC_TYPES:
                self.atom_types = ManifoldGetter._inverse_atomic_one_hot(
                    self.atom_types
                )
            elif self.atom_types.shape[-1] == NUM_ATOMIC_BITS:
                self.atom_types = ManifoldGetter._inverse_atomic_bits(self.atom_types)
            else:
                raise ValueError()

        self.lengths = crys_array_dict["lengths"].squeeze()
        assert self.lengths.ndim == 1
        self.angles = crys_array_dict["angles"].squeeze()
        assert self.lengths.ndim == 1
        self.dict = crys_array_dict

        self.get_structure()
        if self.constructed:
            self.get_composition()
            self.get_validity()
            self.get_fingerprints()
        else:
            self.valid = False

    def get_structure(self):
        if min(self.lengths) < 0:
            self.constructed = False
            self.invalid_reason = "non_positive_lattice"
        if (
            np.isnan(self.lengths).any()
            or np.isnan(self.angles).any()
            or np.isnan(self.frac_coords).any()
        ):
            self.constructed = False
            self.invalid_reason = "nan_value"
        # this catches validity failures down the line
        elif (1 > self.atom_types).any() or (self.atom_types > 104).any():
            self.constructed = False
            self.invalid_reason = f"{self.atom_types=} are not with range"
        else:
            try:
                self.structure = Structure(
                    lattice=Lattice.from_parameters(
                        *(self.lengths.tolist() + self.angles.tolist())
                    ),
                    species=self.atom_types,
                    coords=self.frac_coords,
                    coords_are_cartesian=False,
                )
                self.constructed = True
                if self.structure.volume < 0.1:
                    self.constructed = False
                    self.invalid_reason = "unrealistically_small_lattice"
            except TypeError:
                self.constructed = False
                self.invalid_reason = f"{self.atom_types=} are not possible"
            except Exception:
                self.constructed = False
                self.invalid_reason = "construction_raises_exception"

    def get_composition(self):
        elem_counter = Counter(self.atom_types)
        composition = [
            (elem, elem_counter[elem]) for elem in sorted(elem_counter.keys())
        ]
        elems, counts = list(zip(*composition))
        counts = np.array(counts)
        counts = counts / np.gcd.reduce(counts)
        self.elems = elems
        self.comps = tuple(counts.astype("int").tolist())

    def get_validity(self):
        self.comp_valid = smact_validity(self.elems, self.comps)
        if self.constructed:
            self.struct_valid = structure_validity(self.structure)
        else:
            self.struct_valid = False
        self.valid = self.comp_valid and self.struct_valid

    def get_fingerprints(self):
        elem_counter = Counter(self.atom_types)
        comp = Composition(elem_counter)
        self.comp_fp = CompFP.featurize(comp)
        try:
            site_fps = [
                CrystalNNFP.featurize(self.structure, i)
                for i in range(len(self.structure))
            ]
        except Exception:
            # counts crystal as invalid if fingerprint cannot be constructed.
            self.valid = False
            self.comp_fp = None
            self.struct_fp = None
            return
        self.struct_fp = np.array(site_fps).mean(axis=0)


def get_file_paths(root_path, task, label="", suffix="pt"):
    if label == "":
        out_name = f"eval_{task}.{suffix}"
    else:
        out_name = f"eval_{task}_{label}.{suffix}"
    out_name = os.path.join(root_path, out_name)
    return out_name


def get_crystal_array_list(
    file_path: Path,
    batch_idx: int = 0,
) -> CrysArrayListType:
    """batch_idx == -1, diffcsp format
    batch_idx == -2, cdvae format
    """
    data = load_data(str(file_path.resolve()))
    if batch_idx == -1:
        # batch_size = data["frac_coords"].shape[0]
        batch_size = len(data["frac_coords"])
        crys_array_list = []
        for i in range(batch_size):
            tmp_crys_array_list = get_crystals_list(
                data["frac_coords"][i],
                data["atom_types"][i],
                data["lengths"][i],
                data["angles"][i],
                data["num_atoms"][i],
            )
            crys_array_list.append(tmp_crys_array_list)
    elif batch_idx == -2:
        crys_array_list = get_crystals_list(
            data["frac_coords"],
            data["atom_types"],
            data["lengths"],
            data["angles"],
            data["num_atoms"],
        )
    # elif batch_idx == -3:
    #     batch_size = data['frac_coords'].shape[0]
    #     crys_array_list = []
    #     for i in range(batch_size):
    #         tmp_crys_array_list = {k: v[i].detach().cpu().numpy() for k, v in data.items() if k != "input_data_batch"}
    #         crys_array_list.append(tmp_crys_array_list)
    else:
        crys_array_list = get_crystals_list(
            data["frac_coords"][batch_idx],
            data["atom_types"][batch_idx],
            data["lengths"][batch_idx],
            data["angles"][batch_idx],
            data["num_atoms"][batch_idx],
        )

    if "input_data_batch" in data:
        batch_by_eval = data["input_data_batch"]
        if isinstance(batch_by_eval, dict):
            true_crystal_array_list = get_crystals_list(
                batch_by_eval["frac_coords"],
                batch_by_eval["atom_types"],
                batch_by_eval["lengths"],
                batch_by_eval["angles"],
                batch_by_eval["num_atoms"],
            )
        elif hasattr(batch_by_eval, "frac_coords"):
            true_crystal_array_list = get_crystals_list(
                batch_by_eval.frac_coords,
                batch_by_eval.atom_types,
                batch_by_eval.lengths,
                batch_by_eval.angles,
                batch_by_eval.num_atoms,
            )
        elif isinstance(batch_by_eval, list) and not isinstance(batch_by_eval[0], Data):
            if isinstance(batch_by_eval[0][0], Data) and hasattr(
                batch_by_eval[0][0], "frac_coords"
            ):
                true_crystal_array_list = []
                features = ["frac_coords", "atom_types", "lengths", "angles"]

                # we collect crystals from the first eval
                crystals_in_first_eval = batch_by_eval[0]
                for crystal in crystals_in_first_eval:
                    true_crystal_array_list.append(
                        {k: crystal[k].detach().cpu().numpy() for k in features}
                    )

                # ... and make sure that every eval is in the same order
                for crystals_in_batch in batch_by_eval[1:]:
                    for i, crystal in enumerate(crystals_in_batch):
                        for k in features:
                            t = true_crystal_array_list[i][k]
                            other = crystal[k].detach().cpu().numpy()
                            assert np.allclose(
                                t, other
                            ), f"{t=} != {other=}, that means the first eval is not in the same order as the subsequent"
            else:
                true_crystal_array_list = None
                print("batch_by_eval[0][0]=", batch_by_eval[0][0])
                print("Testing order in generation mode is not supported.")
        else:
            true_crystal_array_list = None
    else:
        true_crystal_array_list = None

    return crys_array_list, true_crystal_array_list


class HiddenPrints:
    def __enter__(self):
        self._original_stdout = sys.stdout
        sys.stdout = open(os.devnull, "w")

    def __exit__(self, exc_type, exc_val, exc_tb):
        sys.stdout.close()
        sys.stdout = self._original_stdout


def cif_to_crys_array_dict(cif: str) -> dict[str, np.ndarray]:
    with HiddenPrints():
        structure = Structure.from_str(cif, fmt="cif", primitive=True)
    lattice = structure.lattice
    return {
        "frac_coords": structure.frac_coords,
        "atom_types": np.array([_.Z for _ in structure.species]),
        "lengths": np.array(lattice.abc),
        "angles": np.array(lattice.angles),
    }


def cif_to_crystal(cif: str) -> Crystal:
    return Crystal(cif_to_crys_array_dict(cif))


def load_gt_crystal_array_list(ground_truth_path: Path) -> CrysArrayListType:
    ground_truth_path = Path(ground_truth_path)
    if ground_truth_path.suffix == ".csv":
        csv = pd.read_csv(ground_truth_path)
        gt_crys = joblib_map(
            cif_to_crys_array_dict,
            csv["cif"],
            n_jobs=-4,
            inner_max_num_threads=1,
            desc="gt_cryst",
            total=len(csv["cif"]),
        )
    elif ground_truth_path.suffix == ".pt":
        pkl = torch.load(ground_truth_path)
        gt_crys = joblib_map(
            cif_to_crys_array_dict,
            [p["cif"] for p in pkl],
            n_jobs=-4,
            inner_max_num_threads=1,
            desc="gt_cryst",
            total=len(pkl),
        )
    else:
        raise ValueError(f"{ground_truth_path=} had an unrecognized suffix.")
    return gt_crys


def load_gt_crystals(ground_truth_path: Path) -> CrysArrayListType:
    ground_truth_path = Path(ground_truth_path)
    if ground_truth_path.suffix == ".csv":
        csv = pd.read_csv(ground_truth_path)
        gt_crys = joblib_map(
            cif_to_crystal,
            csv["cif"],
            n_jobs=-4,
            inner_max_num_threads=1,
            desc="gt_cryst",
            total=len(csv["cif"]),
        )
    elif ground_truth_path.suffix == ".pt":
        pkl = torch.load(ground_truth_path)
        gt_crys = joblib_map(
            cif_to_crystal,
            [p["cif"] for p in pkl],
            n_jobs=-4,
            inner_max_num_threads=1,
            desc="gt_cryst",
            total=len(pkl),
        )
    else:
        raise ValueError(f"{ground_truth_path=} had an unrecognized suffix.")
    return gt_crys


def safe_crystal(x: dict[str, np.ndarray]) -> Crystal:
    if np.all(50 < x["angles"]) and np.all(x["angles"] < 130):
        return Crystal(x)
    else:
        # returns an absurd crystal
        atom_types = np.zeros((1, 100))
        atom_types[0, -1] = 1
        return Crystal(
            {
                "frac_coords": np.zeros((1, 3)),
                "atom_types": atom_types,
                "lengths": 100 * np.ones((3,)),
                "angles": np.ones((3,)) * 90,
            }
        )


def get_Crystal_obj_lists(
    path: Path,
    multi_eval: bool,
    ground_truth_path: Path | None = None,
) -> tuple[list[Crystal], list[Crystal]]:
    batch_idx = -1 if multi_eval else 0
    crys_array_list, true_crystal_array_list = get_crystal_array_list(
        path,
        batch_idx=batch_idx,
    )

    if ground_truth_path is not None:
        gt_crys = load_gt_crystals(ground_truth_path)
    else:
        gt_crys = joblib_map(
            lambda x: Crystal(x),
            true_crystal_array_list,
            n_jobs=-4,
            inner_max_num_threads=1,
            desc="gt_cryst",
            total=len(true_crystal_array_list),
        )

    # TIMEOUT: this MAYBE could be speed up (with timeout) using https://docs.python.org/3/library/multiprocessing.html#using-a-pool-of-workers
    # I could not find a way to have timeouts on single processes using joblib without crashing the whole array.
    # the issue is that if there's a crystal that does not process, the whole thing doesn't return...
    #
    # I tried this but it doesn't work
    # https://github.com/joblib/joblib/pull/366#issuecomment-267603530
    #
    # instead of using timeout, I just replace all crystals with unrealistic angles with a null crystal
    if multi_eval:
        pred_crys = []
        num_evals = len(crys_array_list)
        for i in range(num_evals):
            print(f"Processing batch {i}")
            pred_crys.append(
                joblib_map(
                    safe_crystal,  # instead I remove unrealistic crystals
                    crys_array_list[i],
                    n_jobs=-4,
                    inner_max_num_threads=1,
                    desc="pred_cryst",
                    total=len(crys_array_list),
                )
            )
    else:
        num_evals = 1
        pred_crys = joblib_map(
            safe_crystal,  # instead I remove unrealistic crystals
            crys_array_list,
            n_jobs=-4,
            inner_max_num_threads=1,
            desc="pred_cryst",
            total=len(crys_array_list),
        )

    return pred_crys, gt_crys, num_evals


def save_metrics_only_overwrite_newly_computed(
    path: Path, metrics: dict[str, float]
) -> None:
    # only overwrite metrics computed in the new run.
    if Path(path).exists():
        with open(path, "r") as f:
            written_metrics = json.load(f)
            if isinstance(written_metrics, dict):
                written_metrics.update(metrics)
            else:
                with open(path, "w") as f:
                    json.dump(metrics, f)
        if isinstance(written_metrics, dict):
            with open(path, "w") as f:
                json.dump(written_metrics, f)
    else:
        with open(path, "w") as f:
            json.dump(metrics, f)