UniPath/scripts_analysis/novelty.py

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

from __future__ import annotations

import io
from argparse import ArgumentParser, Namespace
from contextlib import redirect_stdout
from pathlib import Path

import numpy as np
import pandas as pd
from pymatgen.analysis.structure_matcher import StructureMatcher
from pymatgen.core import Structure
from toolz import compose

from flowmm.joblib_ import joblib_map
from flowmm.old_eval.core import save_metrics_only_overwrite_newly_computed
from flowmm.pandas_ import (
    filter_prerelaxed,
    get_intersection,
    maybe_get_missing_columns,
)
from flowmm.pymatgen_ import COLUMNS_COMPUTATIONS, get_chemsys, to_structure
from flowmm.tabular import VALID_STAGES, VALID_TABULAR_DATASETS, get_tabular_dataset

trap = io.StringIO()


def get_matches(
    structure: Structure, alternatives: pd.Series, matcher: StructureMatcher
) -> tuple[list[int], list[float]]:
    with redirect_stdout(trap):
        structure = to_structure(structure)
    matches, rms_dists = [], []
    for ind, alt in alternatives.items():
        with redirect_stdout(trap):
            alt_structure = to_structure(alt)
        rms_dist = matcher.get_rms_dist(structure, alt_structure)
        if rms_dist is not None:
            rms_dist, *_ = rms_dist
            rms_dists.append(rms_dist)
            matches.append(ind)
    return matches, rms_dists


def main(args: Namespace) -> None:
    df = pd.read_json(args.json_in)
    df = maybe_get_missing_columns(df, COLUMNS_COMPUTATIONS)

    if args.ehulls is not None:
        df_hull = pd.read_json(args.ehulls)
        df = df.join(df_hull, how="inner")
        # filter out high energy structures
        df = df[df[args.e_above_hull_column] <= args.e_above_hull_maximum]

    df = filter_prerelaxed(
        df,
        args.num_structures,
        maximum_nary=args.maximum_nary,
        minimum_nary=args.minimum_nary - 1,
    )

    path_json_sun_count = args.json_out.parent / args.json_sun_count
    save_metrics_only_overwrite_newly_computed(
        path_json_sun_count, {"num_stable": len(df)}
    )

    matcher = StructureMatcher()  # MatterGen Novelty settings
    # matcher = StructureMatcher(stol=0.5, angle_tol=10, ltol=0.3)  # CDVAE settings

    # uniqueness
    matches_rms_dists_s = joblib_map(
        lambda structure: get_matches(structure, df["structure"], matcher),
        df["structure"].array,
        n_jobs=-4,
        inner_max_num_threads=1,
        desc="Matching for uniqueness",
        total=len(df),
    )
    # place those lists into a dataframe
    records = []
    for j, (matches, rms_dists) in enumerate(matches_rms_dists_s):
        assert len(matches) == len(rms_dists)
        ind_self = df.index[j]

        if len(matches) == 0:
            record = {
                f"uniq_match_ind-0": pd.NA,
                f"uniq_rms_dist_to-0": float("nan"),
            }
        elif len(matches) == 1:
            if ind_self == matches[0]:
                record = {
                    f"uniq_match_ind-0": pd.NA,
                    f"uniq_rms_dist_to-0": float("nan"),
                }
            else:
                print(
                    f"did not match self! Matched {ind_self} to {matches[0]} with RMSD {rms_dists[0]}"
                )
                record = {
                    f"uniq_match_ind-0": matches[0],
                    f"uniq_rms_dist_to-0": rms_dists[0],
                }
        else:
            record = {}
            for i, (match, rms_dist) in enumerate(zip(matches, rms_dists)):
                if ind_self == match:
                    record[f"uniq_match_ind-{i}"] = pd.NA
                    record[f"uniq_rms_dist_to-{i}"] = float("nan")
                else:
                    record[f"uniq_match_ind-{i}"] = pd.NA if np.isnan(match) else match
                    record[f"uniq_rms_dist_to-{i}"] = rms_dist
        records.append(record)
    uniq_out = pd.DataFrame.from_records(records, index=df.index)
    uniq_match_cols = [col for col in uniq_out.columns if col.startswith("uniq_match")]
    uniq_out[uniq_match_cols] = uniq_out[uniq_match_cols].astype("Int64")

    # load tabular data to compare to
    tds = get_tabular_dataset(args.tabular_dataset)

    # novelty
    outs = []
    for stage in VALID_STAGES:
        if args.reprocess:
            tds.process(stage)
        tabular: pd.DataFrame = getattr(tds, stage + "_df")
        # compositions must match to compare the resulting structure
        gen_chemsys = df["composition"].map(compose(tuple, sorted, get_chemsys))
        tab_chemsys = tabular["composition"].map(compose(tuple, sorted, get_chemsys))
        intersection = get_intersection(gen_chemsys, tab_chemsys)
        gen_to_compare = df["structure"][gen_chemsys.isin(intersection)]
        tab_to_compare = tab_chemsys.isin(intersection)

        # now do pairwise comparisons between these filtered groups
        matches_rms_dists_s = joblib_map(
            lambda structure: get_matches(
                structure, tabular["cif"][tab_to_compare], matcher
            ),
            gen_to_compare.array,
            n_jobs=-4,
            inner_max_num_threads=1,
            desc="Matching for novelty",
            total=len(gen_to_compare),
        )

        # place those lists into a dataframe
        records = []
        for matches, rms_dists in matches_rms_dists_s:
            assert len(matches) == len(rms_dists)
            if len(matches) == 0:
                record = {
                    f"match_ind_{stage}-0": pd.NA,
                    f"rms_dist_to_{stage}-0": float("nan"),
                }
            else:
                record = {}
                for i, (match, rms_dist) in enumerate(zip(matches, rms_dists)):
                    record[f"match_ind_{stage}-{i}"] = match
                    record[f"rms_dist_to_{stage}-{i}"] = rms_dist
            records.append(record)
        out = pd.DataFrame.from_records(records, index=gen_to_compare.index)
        outs.append(out)
    out = pd.concat(outs, axis=1)
    out = pd.concat([uniq_out, out], axis=1)

    print(f"{len(df)=}")
    print(f"{len(out)=}")
    not_in_train = out[out["match_ind_train-0"].isna()]
    print(f"{len(not_in_train)=}")

    # remove duplicates that are not in the training set
    has_a_generated_dupe = pd.concat(
        [
            ~not_in_train[col].isna()
            for col in not_in_train.columns
            if col.startswith("uniq_match")
        ],
        axis=1,
    ).any(axis=1)
    not_in_train_is_dupe = not_in_train[has_a_generated_dupe]
    # mark the duplicates, avoiding the first one that appears
    dupes = []
    cols = [col for col in not_in_train_is_dupe.columns if col.startswith("uniq_match")]
    for i, row in not_in_train_is_dupe[cols].iterrows():
        if i not in dupes:
            dupes.extend(row.array.dropna().tolist())

    sun_materials = not_in_train.drop(dupes)
    print(f"{len(sun_materials)=}")

    save_metrics_only_overwrite_newly_computed(
        path_json_sun_count, {"num_sun_materials": len(sun_materials)}
    )

    out["sun"] = False
    out.loc[sun_materials.index, "sun"] = True
    out.to_json(args.json_out)


if __name__ == "__main__":
    parser = ArgumentParser()
    parser.add_argument("json_in", type=Path, help="prerelaxed dataframe")
    parser.add_argument("json_out", type=Path, help="new dataframe")
    parser.add_argument(
        "--tabular_dataset",
        type=str,
        choices=VALID_TABULAR_DATASETS,
        default="diffcsp_mp20",
    )
    parser.add_argument("-n", "--num_structures", type=int, default=None)
    parser.add_argument("--slurm_partition", type=str, default="ocp")
    parser.add_argument(
        "--maximum_nary",
        type=int,
        default=None,  # we know there aren't structures in the dataset with more than this
        help="Any queries to structures with higher nary are avoided.",
    )
    parser.add_argument(
        "--minimum_nary",
        type=int,
        default=2,
        help="Any queries to structures with lower nary are avoided.",
    )
    parser.add_argument("--ehulls", type=str, default=None)
    parser.add_argument(
        "--e_above_hull_column", type=str, default="e_above_hull_per_atom_dft_corrected"
    )
    parser.add_argument("--e_above_hull_maximum", type=float, default=0.0)
    parser.add_argument("--reprocess", action="store_true")
    parser.add_argument("--json_sun_count", type=str, default="sun_count.json")
    args = parser.parse_args()

    main(args)