Use sparse matrix and don't store df

app.py

@@ -38,7 +38,7 @@ subsets = [
 subsets_ds = {}
 for subset in subsets:
     dataset = load_dataset(
-        "LeMaterial/leMat1",
+        "LeMaterial/LeMat-Bulk",
         subset,
         token=HF_TOKEN,
         columns=[
@@ -52,10 +52,13 @@ for subset in subsets:
             "functional",
         ],
     )
-    subsets_ds[subset] = dataset["train"].to_pandas()
+    subsets_ds[subset] = dataset["train"]
 
-elements_df = {k: subset["elements"] for k, subset in subsets_ds.items()}
+elements_df = {
+    k: subset.select_columns("elements").to_pandas() for k, subset in subsets_ds.items()
+}
 
+from scipy.sparse import csr_matrix
 
 all_elements = {str(el): i for i, el in enumerate(periodictable.elements)}
 elements_indices = {}
@@ -70,6 +73,12 @@ for subset, df in elements_df.items():
 
     df = df.reset_index().apply(map_elements, axis=1)
 
+    import ipdb
+
+    ipdb.set_trace()
+
+    elements_indices[subset] = csr_matrix(elements_indices[subset])
+
 map_functional = {
     "PBE": "compatible_pbe",
     "PBESol (No correction scheme)": "compatible_pbesol",
@@ -108,7 +117,7 @@ def create_phase_diagram(
 
     print(indices_with_only_elements)
 
-    entries_df = subsets_ds[subset_name].loc[indices_with_only_elements]
+    entries_df = subsets_ds[subset_name].select(indices_with_only_elements).to_pandas()
 
     entries_df = entries_df[~entries_df["immutable_id"].isna()]
 
@@ -116,18 +125,18 @@ def create_phase_diagram(
 
     # Fetch all entries from the Materials Project database
     def get_energy_correction(energy_correction, row):
-        if energy_correction == "Database specific, or MP2020" and functional == 'PBE':
-            print('applying MP corrections')
+        if energy_correction == "Database specific, or MP2020" and functional == "PBE":
+            print("applying MP corrections")
             return (
                 row["energy_corrected"] - row["energy"]
                 if not np.isnan(row["energy_corrected"])
                 else 0
             )
-        elif energy_correction == "The 110 PBE Method" and functional == 'PBE':
-            print('applying PBE110 corrections')
-            return row["energy"] * 1.1-row['energy']
-        elif map_functional[functional] != 'pbe':
-            print('not applying any corrections')
+        elif energy_correction == "The 110 PBE Method" and functional == "PBE":
+            print("applying PBE110 corrections")
+            return row["energy"] * 1.1 - row["energy"]
+        elif map_functional[functional] != "pbe":
+            print("not applying any corrections")
             return 0
 
     entries = [
@@ -199,7 +208,10 @@ energy_correction_dropdown = gr.Dropdown(
     label="Energy correction",
 )
 plot_style_dropdown = gr.Dropdown(choices=["2D", "3D"], label="Plot Style")
-functional_dropdown = gr.Dropdown(choices=["PBE", "PBESol (No correction scheme)", "SCAN (No correction scheme)"], label="Functional")
+functional_dropdown = gr.Dropdown(
+    choices=["PBE", "PBESol (No correction scheme)", "SCAN (No correction scheme)"],
+    label="Functional",
+)
 finite_temp_toggle = gr.Checkbox(label="Enable Finite Temperature Estimation")
 
 warning_message = "⚠️ This application uses energy correction schemes directly"