faster filtering

app.py

@@ -1,6 +1,7 @@
 import os
 
 import gradio as gr
+import numpy as np
 import plotly.graph_objs as go
 from datasets import load_dataset
 from pymatgen.analysis.phase_diagram import PDPlotter, PhaseDiagram
@@ -39,41 +40,41 @@ 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 = entries_df[train_df["functional"] == "pbe"]
+    elif functional == "PBESol":
+        entries_df = entries_df[train_df["functional"] == "pbe"]
+    elif functional == "SCAN":
+        entries_df = entries_df[train_df["functional"] == "pbe"]
+
+    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()]
+    ]
+
     # Fetch all entries from the Materials Project database
     entries = [
         ComputedStructureEntry(
             Structure(
-                [x.tolist() for x in train_df.iloc[0]["lattice_vectors"].tolist()],
+                [x.tolist() for x in row["lattice_vectors"].tolist()],
                 row["species_at_sites"],
                 row["cartesian_site_positions"],
                 coords_are_cartesian=True,
             ),
             energy=row["energy"],
-            correction=row["energy_corrected"] - row["energy"],
+            correction=row["energy_corrected"] - row["energy"]
+            if not np.isnan(row["energy_corrected"])
+            else 0,
             entry_id=row["immutable_id"],
             parameters={"run_type": row["functional"]},
         )
-        for n, row in train_df.iterrows()
-        if len(set(row["elements"]).intersection(element_list)) > 0
-        and set(row["elements"]).issubset(element_list)
+        for n, row in entries_df.iterrows()
     ]
     # Fetch elemental entries (they are usually GGA calculations)
     elemental_entries = [e for e in entries if e.composition.is_element]
 
-    # Filter entries based on functional
-    if functional == "PBE":
-        entries = [e for e in entries if e.parameters.get("run_type", "") == "pbe"]
-        entries.extend([e for e in elemental_entries if e not in entries])
-    elif functional == "PBESol":
-        entries = [e for e in entries if e.parameters.get("run_type", "") == "pbesol"]
-        # Add elemental entries to ensure they are included
-        entries.extend([e for e in elemental_entries if e not in entries])
-
-    elif functional == "SCAN":
-        entries = [e for e in entries if e.parameters.get("run_type", "") == "scan"]
-        # Add elemental entries to ensure they are included
-        entries.extend([e for e in elemental_entries if e not in entries])
-
     if finite_temp:
         entries = GibbsComputedStructureEntry.from_entries(entries)