pandas switch

app.py

@@ -22,44 +22,42 @@ subsets = [
     "compatible_scan",
 ]
 
-polars_dfs = {
-    subset: pl.read_parquet(
-        "hf://datasets/LeMaterial/LeMat1/{}/train-*.parquet".format(subset),
-        storage_options={
-            "token": HF_TOKEN,
-        },
-    )
-    for subset in subsets
-}
-
-# Load only the train split of the dataset
-
-# datasets = []
-# 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",
-#             "immutable_id",
-#             "elements",
-#             "functional",
-#         ],
+# polars_dfs = {
+#     subset: pl.read_parquet(
+#         "hf://datasets/LeMaterial/LeMat1/{}/train-*.parquet".format(subset),
+#         storage_options={
+#             "token": HF_TOKEN,
+#         },
 #     )
-#     datasets.append(dataset["train"])
+#     for subset in subsets
+# }
+
+# # Load only the train split of the dataset
+
+subsets_ds = {}
+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",
+            "immutable_id",
+            "elements",
+            "functional",
+        ],
+    )
+    subsets_ds[subset] = dataset["train"]
 
 # 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
 
-dataset = concatenate_datasets(datasets)
-
 # dataset_element_combination_dict = {}
 
 # isubset = lambda x: set(x).issubset(element_list)
@@ -90,33 +88,34 @@ def create_phase_diagram(
 
     # Filter entries based on functional
     if functional == "PBE":
-        df = polars_dfs["compatible_pbe"].clone()
+        entries_df = subsets_ds["compatible_pbe"].to_pandas()
         # entries_df = train_df[train_df["functional"] == "pbe"]
     elif functional == "PBESol":
-        df = polars_dfs["compatible_pbesol"].clone()
+        entries_df = subsets_ds["compatible_pbesol"].to_pandas()
         # entries_df = train_df[train_df["functional"] == "pbesol"]
     elif functional == "SCAN":
-        df = polars_dfs["compatible_scan"].clone()
+        entries_df = subsets_ds["compatible_scan"].to_pandas()
         # entries_df = train_df[train_df["functional"] == "scan"]
 
     # entries_df = df.to_pandas()
 
-    # isubset = lambda x: set(x).issubset(element_list)
-    # isintersection = lambda x: len(set(x).intersection(element_list)) > 0
-    # entries_df = entries_df[entries_df["elements"]](
-    #     lambda example: isintersection(example["elements"])
-    #     and isubset(example["elements"])
-    # )
+    entries_df = entries_df[~entries_df['immutable_id'].isna()]
 
-    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")
-    )
+    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()]
+    ]
+
+    # 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")
+    # )
 
-    entries_df = df.to_pandas()
+    # entries_df = df.to_pandas()
 
     # Fetch all entries from the Materials Project database
     entries = [