+Faster search -Higher memory

app.py

@@ -3,6 +3,7 @@ import os
 import gradio as gr
 import numpy as np
 import pandas as pd
+import periodictable
 import plotly.graph_objs as go
 import polars as pl
 from datasets import concatenate_datasets, load_dataset
@@ -51,27 +52,29 @@ for subset in subsets:
             "functional",
         ],
     )
-    subsets_ds[subset] = dataset["train"]
+    subsets_ds[subset] = dataset["train"].to_pandas()
 
-# Convert the train split to a pandas DataFrame
-# df = pd.concat([x.to_pandas() for x in datasets])
-# train_df = dataset.to_pandas()
-# del dataset
+elements_df = {k: subset["elements"] for k, subset in subsets_ds.items()}
 
-# dataset_element_combination_dict = {}
 
-# isubset = lambda x: set(x).issubset(element_list)
-# isintersection = lambda x: len(set(x).intersection(element_list)) > 0
-# for element_1 in Element:
-#     for element_2 in Element:
-#         for element_3 in Element:
-#             if element_1 != element_2 and element_2 != element_3 and element_3 != element_1:
-#                 print("processing {},{},{}".format(*element_list))
-#                 element_list = [element_1.name, element_2.name, element_3.name]
-#                 dataset_element_combination_dict(sorted(tuple(element_list))) = dataset.filter(
-#                     lambda example: isintersection(example["elements"])
-#                     and isubset(example["elements"])
-#                 )
+all_elements = {str(el): i for i, el in enumerate(periodictable.elements)}
+elements_indices = {}
+for subset, df in elements_df.items():
+    print("Processing subset: ", subset)
+    elements_indices[subset] = np.zeros((len(df), len(all_elements)))
+
+    def map_elements(row):
+        index, xs = row["index"], row["elements"]
+        for x in xs:
+            elements_indices[subset][index, all_elements[x]] = 1
+
+    df = df.reset_index().apply(map_elements, axis=1)
+
+map_functional = {
+    "PBE": "compatible_pbe",
+    "PBESol": "compatible_pbesol",
+    "SCAN": "compatible_scan",
+}
 
 
 def create_phase_diagram(
@@ -85,36 +88,31 @@ def create_phase_diagram(
     # Split elements and remove any whitespace
     element_list = [el.strip() for el in elements.split("-")]
 
-    # Filter entries based on functional
-    if functional == "PBE":
-        entries_df = subsets_ds["compatible_pbe"].to_pandas()
-        # entries_df = train_df[train_df["functional"] == "pbe"]
-    elif functional == "PBESol":
-        entries_df = subsets_ds["compatible_pbesol"].to_pandas()
-        # entries_df = train_df[train_df["functional"] == "pbesol"]
-    elif functional == "SCAN":
-        entries_df = subsets_ds["compatible_scan"].to_pandas()
-        # entries_df = train_df[train_df["functional"] == "scan"]
+    subset_name = map_functional[functional]
 
-    # entries_df = df.to_pandas()
+    element_list_vector = np.zeros(len(all_elements))
+    for el in element_list:
+        element_list_vector[all_elements[el]] = 1
 
-    entries_df = entries_df[~entries_df["immutable_id"].isna()]
-
-    isubset = lambda x: set(x).issubset(element_list)
-    isintersection = lambda x: len(set(x).intersection(element_list)) > 0
-    entries_df = entries_df[
-        [isintersection(l) and isubset(l) for l in entries_df.elements.values.tolist()]
+    n_elements = elements_indices[subset_name].sum(axis=1)
+    n_elements_query = elements_indices[subset_name][
+        :, element_list_vector.astype(bool)
     ]
 
-    # df = df.filter((df.col("elements").list.contains(x) for x in element_list))
-    # df = df.filter(
-    #     pl.col("elements")
-    #     .list.eval(pl.element().is_in(element_list))
-    #     .list.any()
-    #     .alias("check")
-    # )
+    if n_elements_query.shape[1] == 0:
+        indices_with_only_elements = []
+    else:
+        indices_with_only_elements = np.where(
+            n_elements_query.sum(axis=1) == n_elements
+        )[0]
+
+    print(indices_with_only_elements)
+
+    entries_df = subsets_ds[subset_name].loc[indices_with_only_elements]
+
+    entries_df = entries_df[~entries_df["immutable_id"].isna()]
 
-    # entries_df = df.to_pandas()
+    print(entries_df)
 
     # Fetch all entries from the Materials Project database
     def get_energy_correction(energy_correction, row):
@@ -153,6 +151,7 @@ def create_phase_diagram(
     try:
         phase_diagram = PhaseDiagram(entries)
     except ValueError as e:
+        print(e)
         return go.Figure().add_annotation(text=str(e))
 
     # Generate plotly figure
@@ -188,7 +187,10 @@ elements_input = gr.Textbox(
 #     minimum=0, maximum=1, value=0.1, label="Maximum Energy Above Hull (eV)"
 # )
 energy_correction_dropdown = gr.Dropdown(
-    choices=["The 110 PBE Method", "Database specific, or MP2020",],
+    choices=[
+        "The 110 PBE Method",
+        "Database specific, or MP2020",
+    ],
     label="Energy correction",
 )
 plot_style_dropdown = gr.Dropdown(choices=["2D", "3D"], label="Plot Style")
@@ -210,9 +212,7 @@ warning_message += " from <a href='https://chemrxiv.org/engage/api-gateway/chemr
 message = '<div class="alert"><span class="closebtn" onclick="this.parentElement.style.display="none";">&times;</span>{}</div>Generate a phase diagram for a set of elements using LeMat-Bulk data.'.format(
     warning_message
 )
-message += (
-    "<br>Built with <a href='https://pymatgen.org/' target='_blank'>Pymatgen</a> and  <a href='https://docs.crystaltoolkit.org/' target='_blank'>Crystal Toolkit</a>.<br>"
-)
+message += "<br>Built with <a href='https://pymatgen.org/' target='_blank'>Pymatgen</a> and  <a href='https://docs.crystaltoolkit.org/' target='_blank'>Crystal Toolkit</a>.<br>"
 
 # Create Gradio interface
 iface = gr.Interface(