polars test

app.py

@@ -1,4 +1,5 @@
 import os
+import polars as pl
 
 import gradio as gr
 import numpy as np
@@ -6,7 +7,7 @@ import pandas as pd
 import plotly.graph_objs as go
 from datasets import concatenate_datasets, load_dataset
 from pymatgen.analysis.phase_diagram import PDPlotter, PhaseDiagram
-from pymatgen.core import Composition, Structure
+from pymatgen.core import Composition, Structure, Element
 from pymatgen.core.composition import Composition
 from pymatgen.entries.computed_entries import (
     ComputedStructureEntry,
@@ -21,26 +22,36 @@ 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",
-        ],
-    )
-    datasets.append(dataset["train"])
+# 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",
+#         ],
+#     )
+#     datasets.append(dataset["train"])
 
 # Convert the train split to a pandas DataFrame
 # df = pd.concat([x.to_pandas() for x in datasets])
@@ -49,6 +60,21 @@ for subset in subsets:
 
 dataset = concatenate_datasets(datasets)
 
+# 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"])
+#                 )
+
 
 def create_phase_diagram(
     elements,
@@ -64,23 +90,33 @@ def create_phase_diagram(
 
     # Filter entries based on functional
     if functional == "PBE":
-        ds_filter = dataset.filter(lambda example: example["functional"] == "pbe")
+        df = polars_dfs["compatible_pbe"].clone()
         # entries_df = train_df[train_df["functional"] == "pbe"]
     elif functional == "PBESol":
-        ds_filter = dataset.filter(lambda example: example["functional"] == "pbesol")
+        df = polars_dfs["compatible_pbesol"].clone()
         # entries_df = train_df[train_df["functional"] == "pbesol"]
     elif functional == "SCAN":
-        ds_filter = dataset.filter(lambda example: example["functional"] == "scan")
+        df = polars_dfs["compatible_scan"].clone()
         # entries_df = train_df[train_df["functional"] == "scan"]
 
-    isubset = lambda x: set(x).issubset(element_list)
-    isintersection = lambda x: len(set(x).intersection(element_list)) > 0
-    ds_filter = ds_filter.filter(
-        lambda example: isintersection(example["functional"])
-        and isubset(example["functional"])
+    # 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"])
+    # )
+
+    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 = ds_filter.to_pandas()
+    entries_df = df.to_pandas()
 
     # Fetch all entries from the Materials Project database
     entries = [