scan, pbesol, pbe filter

app.py

@@ -6,10 +6,8 @@ from datasets import load_dataset
 from pymatgen.analysis.phase_diagram import PDPlotter, PhaseDiagram
 from pymatgen.core import Composition, Structure
 from pymatgen.core.composition import Composition
-from pymatgen.entries.computed_entries import (
-    ComputedStructureEntry,
-    GibbsComputedStructureEntry,
-)
+from pymatgen.entries.computed_entries import (ComputedStructureEntry,
+                                               GibbsComputedStructureEntry)
 
 # Load only the train split of the dataset
 dataset = load_dataset("LeMaterial/leDataset", split="train")
@@ -36,6 +34,7 @@ def create_phase_diagram(
             energy=row["energy"],
             correction=row["energy_corrected"] - row["energy"],
             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
@@ -45,16 +44,16 @@ def create_phase_diagram(
     elemental_entries = [e for e in entries if e.composition.is_element]
 
     # Filter entries based on functional
-    if functional == "GGA":
-        entries = [
-            e
-            for e in entries
-            if not e.parameters.get("run_type", "").startswith("GGA+U")
-        ]
-    elif functional == "GGA+U":
-        entries = [
-            e for e in entries if e.parameters.get("run_type", "").startswith("GGA+U")
-        ]
+    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])
 
@@ -103,7 +102,7 @@ color_scheme_dropdown = gr.Dropdown(
     choices=["Energy Above Hull", "Formation Energy"], label="Color Scheme"
 )
 plot_style_dropdown = gr.Dropdown(choices=["2D", "3D"], label="Plot Style")
-functional_dropdown = gr.Dropdown(choices=["GGA", "GGA+U", "Both"], label="Functional")
+functional_dropdown = gr.Dropdown(choices=["PBE", "PBESol", "SCAN"], label="Functional")
 finite_temp_toggle = gr.Checkbox(label="Enable Finite Temperature Estimation")
 
 # Create Gradio interface