data_utils.py

import os
import re

import crystal_toolkit.components as ctc
import numpy as np
import periodictable
from dash import dcc, html
from datasets import concatenate_datasets, load_dataset
from pymatgen.analysis.structure_analyzer import SpacegroupAnalyzer

HF_TOKEN = os.environ.get("HF_TOKEN")
top_k = 500


def get_dataset():
    # Load only the train split of the dataset
    datasets = []
    subsets = [
        "compatible_pbe",
        "compatible_pbesol",
        "compatible_scan",
        "non_compatible",
    ]

    for subset in subsets:
        dataset = load_dataset(
            "LeMaterial/LeMat-Bulk",
            subset,
#            token=HF_TOKEN,
            columns=[
                "lattice_vectors",
                "species_at_sites",
                "cartesian_site_positions",
                "energy",
                # "energy_corrected", # not yet available in LeMat-Bulk
                "immutable_id",
                "elements",
                "stress_tensor",
                "magnetic_moments",
                "forces",
                # "band_gap_direct", #future release
                # "band_gap_indirect", #future release
                "dos_ef",
                # "charges", #future release
                "functional",
                "chemical_formula_reduced",
                "chemical_formula_descriptive",
                "total_magnetization",
                "entalpic_fingerprint",
            ],
        )
        datasets.append(dataset["train"])

    return concatenate_datasets(datasets)


display_columns = [
    "chemical_formula_descriptive",
    "functional",
    "immutable_id",
    "energy",
]
display_names = {
    "chemical_formula_descriptive": "Formula",
    "functional": "Functional",
    "immutable_id": "Material ID",
    "energy": "Energy (eV)",
}

# Global shared variables
mapping_table_idx_dataset_idx = {}


def build_formula_index(dataset, index_range=None, cache_path=None, empty_data=False):
    print("Building formula index")
    if empty_data:
        return np.zeros((1, 1)), {}

    use_dataset = dataset
    if index_range is not None:
        use_dataset = dataset.select(index_range)

    # Preprocessing step to create an index for the dataset
    from scipy.sparse import load_npz

    if cache_path is not None and os.path.exists(f"{cache_path}/train_df.pkl"):
        train_df = pickle.load(open(f"{cache_path}/train_df.pkl", "rb"))
        dataset_index = load_npz(f"{cache_path}/dataset_index.npz")
    else:
        train_df = use_dataset.select_columns(
            ["species_at_sites", "immutable_id", "functional"]
        ).to_pandas()

        import tqdm

        all_elements = {
            str(el.symbol): i for i, el in enumerate(periodictable.elements)
        }  # full element list
        dataset_index = np.zeros((len(train_df), len(all_elements)))

        for idx, species in tqdm.tqdm(enumerate(train_df["species_at_sites"].values)):
            for el in species:
                dataset_index[idx, all_elements[el]] += 1

        dataset_index = dataset_index / np.sum(dataset_index, axis=1)[:, None]
        dataset_index = (
            dataset_index / np.linalg.norm(dataset_index, axis=1)[:, None]
        )  # Normalize vectors

        from scipy.sparse import csr_matrix, save_npz

        dataset_index = csr_matrix(dataset_index)

        if cache_path is not None:
            pickle.dump(train_df, open(f"{cache_path}/train_df.pkl", "wb"))
            save_npz(f"{cache_path}/dataset_index.npz", dataset_index)

    immutable_id_to_idx = train_df["immutable_id"].to_dict()
    del train_df
    immutable_id_to_idx = {v: k for k, v in immutable_id_to_idx.items()}

    return dataset_index, immutable_id_to_idx


import pickle
from pathlib import Path


# TODO: Just load the index from a file
def build_embeddings_index(empty_data=False):
    if empty_data:
        return None, {}, {}

    features_dict = pickle.load(open("features_dict.pkl", "rb"))

    from indexer import FAISSIndex

    index = FAISSIndex()
    for key in features_dict:
        index.index.add(features_dict[key].reshape(1, -1))

    idx_to_immutable_id = {i: key for i, key in enumerate(features_dict)}

    # index = FAISSIndex.from_store("index.faiss")

    return index, features_dict, idx_to_immutable_id


def search_materials(
    query, dataset, dataset_index, mapping_table_idx_dataset_idx, map_periodic_table
):
    n_elements = len(map_periodic_table)
    query_vector = np.zeros(n_elements)

    if "," in query:
        element_list = [el.strip() for el in query.split(",")]
        for el in element_list:
            query_vector[map_periodic_table[el]] = 1
    else:
        # Formula
        import re

        matches = re.findall(r"([A-Z][a-z]{0,2})(\d*)", query)
        for el, numb in matches:
            numb = int(numb) if numb else 1
            query_vector[map_periodic_table[el]] = numb

    similarity = dataset_index.dot(query_vector) / (np.linalg.norm(query_vector))
    indices = np.argsort(similarity)[::-1][:top_k]

    options = [dataset[int(i)] for i in indices]

    mapping_table_idx_dataset_idx.clear()
    for i, idx in enumerate(indices):
        mapping_table_idx_dataset_idx[int(i)] = int(idx)

    return options


def get_properties_table(
    row, structure, sga, properties_container_update, container_type="query"
):
    properties = {
        "Material ID": row["immutable_id"],
        "Formula": row["chemical_formula_descriptive"],
        "Energy per atom (eV/atom)": round(
            row["energy"] / len(row["species_at_sites"]), 3
        ),
        # "Band Gap (eV)": row["band_gap_direct"] or row["band_gap_indirect"], #future release
        "Total Magnetization (μB)": (
            round(row["total_magnetization"], 3)
            if row["total_magnetization"] is not None
            else None
        ),
        "Density (g/cm^3)": round(structure.density, 3),
        "Fermi energy level (eV)": (
            round(row["dos_ef"], 3) if row["dos_ef"] is not None else None
        ),
        "Crystal system": sga.get_crystal_system(),
        "International Spacegroup": sga.get_symmetry_dataset().international,
        "Magnetic moments (μB)": np.round(row["magnetic_moments"], 3),
        "Stress tensor (kB)": np.round(row["stress_tensor"], 3),
        "Forces on atoms (eV/A)": np.round(row["forces"], 3),
        # "Bader charges (e-)": np.round(row["charges"], 3), # future release
        "DFT Functional": row["functional"],
        "Entalpic fingerprint": row["entalpic_fingerprint"],
    }

    style = {
        "padding": "10px",
        "borderBottom": "1px solid #ddd",
    }

    if container_type == "query":
        properties_container_update[0] = properties
    else:
        properties_container_update[1] = properties
        # if (type(value) in [str, float]) and (
        #     properties_container_update[0][key] == properties_container_update[1][key]
        # ):
        #     style["backgroundColor"] = "#e6f7ff"

    # Format properties as an HTML table
    properties_html = html.Table(
        [
            html.Tbody(
                [
                    html.Tr(
                        [
                            html.Th(
                                key,
                                style={
                                    "padding": "10px",
                                    "verticalAlign": "middle",
                                },
                            ),
                            html.Td(
                                str(value),
                                style=style,
                            ),
                        ],
                    )
                    for key, value in properties.items()
                ],
            )
        ],
        style={
            "width": "100%",
            "borderCollapse": "collapse",
            "fontFamily": "'Arial', sans-serif",
            "fontSize": "14px",
            "color": "#333333",
        },
    )

    return properties_html


def get_crystal_plot(structure):
    sga = SpacegroupAnalyzer(structure)
    # Create the StructureMoleculeComponent
    structure_component = ctc.StructureMoleculeComponent(structure)
    return structure_component.layout(), sga