More responsive

app.py

@@ -1,112 +1,73 @@
-import os
-import re
-
 import crystal_toolkit.components as ctc
 import dash
 import dash_mp_components as dmp
 import numpy as np
-import pandas as pd
 import periodictable
 from crystal_toolkit.settings import SETTINGS
 from dash import dcc, html
 from dash.dependencies import Input, Output, State
 from dash_breakpoints import WindowBreakpoints
-from datasets import concatenate_datasets, load_dataset
 from pymatgen.analysis.structure_analyzer import SpacegroupAnalyzer
 from pymatgen.core import Structure
 
-HF_TOKEN = os.environ.get("HF_TOKEN")
-top_k = 500
-
-subsets = ["compatible_pbe", "compatible_pbesol", "compatible_scan", "non_compatible"]
+from components import (
+    get_display_table,
+    get_dropdown,
+    get_materials_display,
+    get_periodic_table,
+    get_upload_div,
+)
+from data_utils import (
+    build_embeddings_index,
+    build_formula_index,
+    get_crystal_plot,
+    get_dataset,
+    get_properties_table,
+    search_materials,
+)
 
-# Load only the train split of the dataset
+EMPTY_DATA = False
+CACHE_PATH = None
 
-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", # not yet available in LeMat-Bulk
-            "immutable_id",
-            "elements",
-            "functional",
-            "stress_tensor",
-            "magnetic_moments",
-            "forces",
-            # "band_gap_direct", #future release
-            # "band_gap_indirect", #future release
-            "dos_ef",
-            # "charges", #future release
-            "functional",
-            "chemical_formula_reduced",
-            "chemical_formula_descriptive",
-            "total_magnetization",
-            "entalpic_fingerprint"
-        ],
-    )
-    datasets.append(dataset["train"])
+dataset = get_dataset()
 
-display_columns = [
+display_columns_query = [
     "chemical_formula_descriptive",
     "functional",
     "immutable_id",
     "energy",
 ]
-display_names = {
+display_names_query = {
     "chemical_formula_descriptive": "Formula",
     "functional": "Functional",
     "immutable_id": "Material ID",
     "energy": "Energy (eV)",
 }
 
+
 mapping_table_idx_dataset_idx = {}
+available_similar_materials = []
 
 map_periodic_table = {v.symbol: k for k, v in enumerate(periodictable.elements)}
-n_elements = len(map_periodic_table)
 
-# Preprocessing step to create an index for the dataset
-# df = pd.concat([x.to_pandas() for x in datasets])
-dataset = concatenate_datasets(datasets)
-train_df = dataset.select_columns(["chemical_formula_descriptive"]).to_pandas()
-
-pattern = re.compile(r"(?P<element>[A-Z][a-z]?)(?P<count>\d*)")
-extracted = train_df["chemical_formula_descriptive"].str.extractall(pattern)
-extracted["count"] = extracted["count"].replace("", "1").astype(int)
-
-wide_df = extracted.reset_index().pivot_table(  # Move index to columns for pivoting
-    index="level_0",  # original row index
-    columns="element",
-    values="count",
-    aggfunc="sum",
-    fill_value=0,
+# dataset_index, immutable_id_to_idx = build_formula_index(dataset, cache_path=None)
+dataset_index, immutable_id_to_idx = build_formula_index(
+    dataset, cache_path=CACHE_PATH, empty_data=EMPTY_DATA
 )
 
-all_elements = [el.symbol for el in periodictable.elements]  # full element list
-wide_df = wide_df.reindex(columns=all_elements, fill_value=0)
-
-
-dataset_index = wide_df.values
-
-dataset_index = dataset_index / np.sum(dataset_index, axis=1)[:, None]
-dataset_index = (
-    dataset_index / np.linalg.norm(dataset_index, axis=1)[:, None]
-)  # Normalize vectors
-
-del train_df, extracted, wide_df
-
 # Initialize the Dash app
-app = dash.Dash(__name__, assets_folder=SETTINGS.ASSETS_PATH)
+external_stylesheets = [
+    "/assets/styles.css",
+]
+
+app = dash.Dash(
+    __name__,
+    external_stylesheets=external_stylesheets,
+)
 server = app.server  # Expose the server for deployment
 
 # Define the app layout
-layout = html.Div(
+app.layout = html.Div(
     [
         WindowBreakpoints(
             id="breakpoints",
@@ -119,178 +80,26 @@ layout = html.Div(
         ),
         html.Div(
             [
-                html.Div(
-                    [
-                        html.Div(
-                            "Search a material to display its structure and properties",
-                            style={"textAlign": "center"},
-                        ),
-                    ],
-                    id="structure-container",
-                    style={
-                        "width": "44%",
-                        "verticalAlign": "top",
-                        "boxShadow": "0px 4px 8px rgba(0, 0, 0, 0.1)",
-                        "borderRadius": "10px",
-                        "backgroundColor": "#f9f9f9",
-                        "padding": "20px",
-                        "textAlign": "center",
-                        "display": "flex",
-                        "justifyContent": "center",
-                        "alignItems": "center",
-                    },
-                ),
-                html.Div(
-                    id="properties-container",
-                    style={
-                        "width": "55%",
-                        "paddingLeft": "4%",
-                        "verticalAlign": "top",
-                        "boxShadow": "0px 4px 8px rgba(0, 0, 0, 0.1)",
-                        "borderRadius": "10px",
-                        "backgroundColor": "#f9f9f9",
-                        "padding": "20px",
-                        "overflow": "auto",
-                        "maxHeight": "600px",
-                        "display": "flex",
-                        "justifyContent": "center",
-                        "wordWrap": "break-word",
-                    },
-                    children=[
-                        html.Div(
-                            "Properties will be displayed here",
-                            style={"textAlign": "center"},
-                        ),
-                    ],
-                ),
+                get_materials_display(
+                    "",
+                    "Structure will be displayed here",
+                    "Properties will be displayed here",
+                )
             ],
-            style={
-                "marginTop": "20px",
-                "display": "flex",
-                "justifyContent": "space-between",  # Ensure the two sections are responsive
-                "flexWrap": "wrap",
-            },
+            className="container-row",
         ),
         html.Div(
             [
-                html.Div(
-                    [
-                        html.H3("Search Materials (eg. 'Ac,Cd,Ge' or 'Ac2CdGe3')"),
-                        html.Div(
-                            [
-                                html.Div(
-                                    [
-                                        dmp.MaterialsInput(
-                                            allowedInputTypes=["elements", "formula"],
-                                            hidePeriodicTable=False,
-                                            periodicTableMode="toggle",
-                                            hideWildcardButton=True,
-                                            showSubmitButton=True,
-                                            submitButtonText="Search",
-                                            type="elements",
-                                            id="materials-input",
-                                        ),
-                                    ],
-                                    id="materials-input-container",
-                                    style={
-                                        "width": "100%",
-                                    },
-                                ),
-                            ],
-                            style={
-                                "display": "flex",
-                                "justifyContent": "center",
-                                "width": "100%",
-                            },
-                        ),
-                    ],
-                    style={
-                        "width": "48%",
-                        "verticalAlign": "top",
-                    },
-                ),
-                html.Div(
-                    [
-                        html.Label(
-                            "Select a row to display the material's structure and properties",
-                            style={"margin-bottom": "20px"},
-                        ),
-                        # dcc.Dropdown(
-                        #     id="material-dropdown",
-                        #     options=[],  # Empty options initially
-                        #     value=None,
-                        # ),
-                        dash.dash_table.DataTable(
-                            id="table",
-                            columns=[
-                                (
-                                    {"name": display_names[col], "id": col}
-                                    if col != "energy"
-                                    else {
-                                        "name": display_names[col],
-                                        "id": col,
-                                        "type": "numeric",
-                                        "format": {"specifier": ".2f"},
-                                    }
-                                )
-                                for col in display_columns
-                            ],
-                            data=[{}],
-                            style_cell={
-                                "fontFamily": "Arial",
-                                "padding": "10px",
-                                "border": "1px solid #ddd",  # Subtle border for elegance
-                                "textAlign": "left",
-                                "fontSize": "14px",
-                            },
-                            style_header={
-                                "backgroundColor": "#f5f5f5",  # Light grey header
-                                "fontWeight": "bold",
-                                "textAlign": "left",
-                                "borderBottom": "2px solid #ddd",
-                            },
-                            style_data={
-                                "backgroundColor": "#ffffff",
-                                "color": "#333333",
-                                "borderBottom": "1px solid #ddd",
-                            },
-                            style_data_conditional=[
-                                {
-                                    "if": {"state": "active"},
-                                    "backgroundColor": "#e6f7ff",
-                                    "border": "1px solid #1890ff",
-                                },
-                            ],
-                            style_table={
-                                "maxHeight": "400px",
-                                "overflowX": "auto",
-                                "overflowY": "auto",
-                            },
-                            style_as_list_view=True,
-                            row_selectable="single",
-                            selected_rows=[],
-                        ),
-                    ],
-                    style={
-                        "width": "48%",
-                        # "maxWidth": "800px",
-                        "margin": "0 auto",
-                        "padding": "20px",
-                        "backgroundColor": "#ffffff",
-                        "borderRadius": "10px",
-                        "boxShadow": "0px 4px 8px rgba(0, 0, 0, 0.1)",
-                    },
+                get_periodic_table("materials-input", {}),
+                get_display_table(
+                    "table",
+                    display_names_query,
+                    display_columns_query,
+                    "Select a row to display the material's structure and properties",
                 ),
             ],
-            style={
-                "margin-top": "20px",
-                "margin-bottom": "20px",
-                "display": "flex",
-                "flexDirection": "row",
-                "alignItems": "center",
-            },
+            className="container-row-periodic",
         ),
-        # acknowledgements to mp dash components and crystal toolkit
         html.Footer(
             [
                 html.P(
@@ -308,16 +117,6 @@ layout = html.Div(
                     style={"textAlign": "center"},
                 )
             ],
-            style={
-                "display": "flex",
-                "justifyContent": "center",
-                "alignItems": "center",
-                "flexWrap": "wrap",
-                "padding": "1rem 0",
-                "backgroundColor": "#f1f1f1",  # Optional: light gray footer background
-                "borderTop": "1px solid #ddd",  # Optional: subtle border at the top
-                "width": "100%",
-            },
         ),
     ],
     style={
@@ -327,34 +126,6 @@ layout = html.Div(
 )
 
 
-def search_materials(query):
-    query_vector = np.zeros(n_elements)
-
-    if "," in query:
-        element_list = [el.strip() for el in query.split(",")]
-        for el in element_list:
-            query_vector[map_periodic_table[el]] = 1
-    else:
-        # Formula
-        import re
-
-        matches = re.findall(r"([A-Z][a-z]{0,2})(\d*)", query)
-        for el, numb in matches:
-            numb = int(numb) if numb else 1
-            query_vector[map_periodic_table[el]] = numb
-
-    similarity = np.dot(dataset_index, query_vector) / (np.linalg.norm(query_vector))
-    indices = np.argsort(similarity)[::-1][:top_k]
-
-    options = [dataset[int(i)] for i in indices]
-
-    mapping_table_idx_dataset_idx.clear()
-    for i, idx in enumerate(indices):
-        mapping_table_idx_dataset_idx[int(i)] = int(idx)
-
-    return options
-
-
 # Callback to update the table based on search
 @app.callback(
     Output("table", "data"),
@@ -365,9 +136,11 @@ def on_submit_materials_input(n_clicks, query):
     if n_clicks is None or not query:
         return []
 
-    entries = search_materials(query)
+    entries = search_materials(
+        query, dataset, dataset_index, mapping_table_idx_dataset_idx, map_periodic_table
+    )
 
-    return [{col: entry[col] for col in display_columns} for entry in entries]
+    return [{col: entry[col] for col in display_columns_query} for entry in entries]
 
 
 # Callback to display the selected material
@@ -376,7 +149,6 @@ def on_submit_materials_input(n_clicks, query):
         Output("structure-container", "children"),
         Output("properties-container", "children"),
     ],
-    # Input("display-button", "n_clicks"),
     Input("table", "active_cell"),
     Input("table", "derived_virtual_selected_rows"),
 )
@@ -408,69 +180,17 @@ def display_material(active_cell, selected_rows):
     if row["magnetic_moments"]:
         structure.add_site_property("magmom", row["magnetic_moments"])
 
-    sga = SpacegroupAnalyzer(structure)
-
-    # Create the StructureMoleculeComponent
-    structure_component = ctc.StructureMoleculeComponent(structure)
+    structure_layout, sga = get_crystal_plot(structure)
 
     # Extract key properties
-    properties = {
-        "Material ID": row["immutable_id"],
-        "Formula": row["chemical_formula_descriptive"],
-        "Energy per atom (eV/atom)": round(
-            row["energy"] / len(row["species_at_sites"]), 3
-        ),
-        # "Band Gap (eV)": row["band_gap_direct"] or row["band_gap_indirect"], #future release
-        "Total Magnetization (μB)": round(row["total_magnetization"], 3) if row['total_magnetization'] is not None else None,
-        "Density (g/cm^3)": round(structure.density, 3),
-        "Fermi energy level (eV)": round(row["dos_ef"],3) if row['dos_ef'] is not None else None,
-        "Crystal system": sga.get_crystal_system(),
-        "International Spacegroup": sga.get_symmetry_dataset().international,
-        "Magnetic moments (μB)": np.round(row["magnetic_moments"], 3),
-        "Stress tensor (kB)": np.round(row["stress_tensor"], 3),
-        "Forces on atoms (eV/A)": np.round(row["forces"], 3),
-        # "Bader charges (e-)": np.round(row["charges"], 3), # future release
-        "DFT Functional": row["functional"],
-        "Entalpic fingerprint": row['entalpic_fingerprint'],
-    }
-
-    # Format properties as an HTML table
-    properties_html = html.Table(
-        [
-            html.Tbody(
-                [
-                    html.Tr(
-                        [
-                            html.Th(
-                                key,
-                                style={
-                                    "padding": "10px",
-                                    "verticalAlign": "middle",
-                                },
-                            ),
-                            html.Td(
-                                str(value),
-                                style={
-                                    "padding": "10px",
-                                    "borderBottom": "1px solid #ddd",
-                                },
-                            ),
-                        ],
-                    )
-                    for key, value in properties.items()
-                ],
-            )
-        ],
-        style={
-            "width": "100%",
-            "borderCollapse": "collapse",
-            "fontFamily": "'Arial', sans-serif",
-            "fontSize": "14px",
-            "color": "#333333",
-        },
+    properties_html = get_properties_table(
+        row, structure, sga, [None, None], container_type="results"
     )
 
-    return structure_component.layout(), properties_html
+    return (
+        structure_layout,
+        properties_html,
+    )
 
 
 @app.callback(
@@ -505,7 +225,7 @@ def update_materials_input_layout(breakpoint_name, width):
 
 
 # Register crystal toolkit with the app
-ctc.register_crystal_toolkit(app, layout)
+ctc.register_crystal_toolkit(app, app.layout)
 
 if __name__ == "__main__":
     app.run_server(debug=True, port=7860, host="0.0.0.0")

assets/styles.css

@@ -0,0 +1,124 @@
+h1 {
+  font-size: 24px;
+  font-weight: 700;
+  color: #333;
+}
+
+.body {
+  background-color: #4a4a4a;
+}
+
+.header-container {
+  display: flex;
+  flex-direction: row;
+  justify-content: space-between;
+  align-items: center;
+  padding: 20px;
+  margin-bottom: 20px;
+}
+
+.container {
+  box-shadow: 0px 4px 8px rgba(0, 0, 0, 0.1);
+  border-radius: 10px;
+  background-color: rgb(249, 249, 249);
+  padding: 20px;
+  margin-left: 10px;
+  margin-right: 10px;
+  max-height: 600px;
+  justify-content: center;
+}
+
+.container-visu {
+  width: 45%;
+  align-items: center;
+}
+
+.container-table {
+  width: 50%;
+  align-items: center;
+  overflow: auto;
+}
+
+/* remove background in periodical table */
+.periodic-table {
+  background-color: transparent;
+  box-shadow: none;
+}
+
+.container-row {
+  width: 100%;
+  display: flex;
+  flex-direction: row;
+  justify-content: space-between;
+  padding: 10px;
+  margin-bottom: 10px;
+  margin-top: 10px;
+}
+
+.container-row-periodic {
+  width: 100%;
+  display: flex;
+  flex-direction: row;
+  align-items: center;
+}
+
+.container-col {
+  width: 100%;
+  display: flex;
+  flex-direction: column;
+  justify-content: space-between;
+  padding: 10px;
+  margin-bottom: 10px;
+  margin-top: 10px;
+}
+
+body {
+  font-family: "Arial", sans-serif;
+  font-size: 16px;
+}
+
+@media (max-width: 800px) {
+  .container {
+    width: 100%;
+    margin: 5px;
+    margin-top: 10px;
+    margin-bottom: 10px;
+  }
+
+  .container-row {
+    flex-direction: column;
+  }
+
+  .container-row-periodic {
+    flex-direction: column;
+  }
+
+  .container-visu {
+    width: 100%;
+  }
+
+  .container-table {
+    width: 100%;
+  }
+}
+
+@media (max-width: 1000px) and (min-width: 800px) {
+  .container-visu {
+    width: 60%;
+  }
+
+  .container-table {
+    width: 39%;
+  }
+}
+
+footer {
+  display: flex;
+  justify-content: center;
+  align-items: center;
+  flex-wrap: wrap;
+  margin-top: 40px;
+  background-color: #ffffff;
+  border-top: "1px solid #ddd";
+  width: 100%;
+}

components.py

@@ -0,0 +1,184 @@
+import dash
+import dash_mp_components as dmp
+from dash import dcc, html
+
+display_columns = [
+    "chemical_formula_descriptive",
+    "functional",
+    "immutable_id",
+    "energy",
+]
+
+display_names = {
+    "chemical_formula_descriptive": "Formula",
+    "functional": "Functional",
+    "immutable_id": "Material ID",
+    "energy": "Energy (eV)",
+}
+
+
+def get_periodic_table(id, table_kwargs, **style_kwargs):
+
+    return html.Div(
+        [
+            html.H3("Search Materials (eg. 'Ac,Cd,Ge' or 'Ac2CdGe3')"),
+            html.Div(
+                [
+                    dmp.MaterialsInput(
+                        allowedInputTypes=[
+                            "elements",
+                            "formula",
+                        ],
+                        hidePeriodicTable=False,
+                        periodicTableMode="toggle",
+                        hideWildcardButton=True,
+                        showSubmitButton=True,
+                        submitButtonText="Search",
+                        type="elements",
+                        **table_kwargs,
+                        id=id,
+                    ),
+                ],
+                id="materials-input-container",
+                style={
+                    "width": "100%",
+                },
+            ),
+        ],
+        className="container periodic-table",
+    )
+
+
+def get_dropdown(id, options, **style_kwargs):
+    return dcc.Dropdown(
+        id=id,
+        options=options,
+        placeholder="Embedder",
+        value=None,
+        clearable=False,
+        style=style_kwargs,
+    )
+
+
+def get_upload_div(id, **style_kwargs):
+    return html.Div(
+        [
+            html.H3("Upload a CIF file"),
+            dcc.Upload(
+                id=id,
+                children=html.Div(
+                    [
+                        "Drag and Drop or ",
+                        html.A("Select a CIF file"),
+                    ]
+                ),
+                style={
+                    "width": "100%",
+                    "height": "60px",
+                    "lineHeight": "60px",
+                    "borderWidth": "1px",
+                    "borderStyle": "dashed",
+                    "borderRadius": "5px",
+                    "textAlign": "center",
+                    "margin": "10px",
+                },
+                multiple=False,
+            ),
+        ],
+        className="container",
+    )
+
+
+def get_display_table(id, display_names, display_columns, text, **style_kwargs):
+
+    return html.Div(
+        [
+            html.Label(
+                text,
+                style={"margin-bottom": "20px"},
+            ),
+            dash.dash_table.DataTable(
+                id=id,
+                columns=[
+                    (
+                        {
+                            "name": display_names[col],
+                            "id": col,
+                        }
+                        if col != "energy"
+                        else {
+                            "name": display_names[col],
+                            "id": col,
+                            "type": "numeric",
+                            "format": {"specifier": ".2f"},
+                        }
+                    )
+                    for col in display_columns
+                ],
+                data=[{}],
+                style_cell={
+                    "fontFamily": "Arial",
+                    "padding": "10px",
+                    "border": "1px solid #ddd",  # Subtle border for elegance
+                    "textAlign": "left",
+                    "fontSize": "14px",
+                },
+                style_header={
+                    "backgroundColor": "#f5f5f5",  # Light grey header
+                    "fontWeight": "bold",
+                    "textAlign": "left",
+                    "borderBottom": "2px solid #ddd",
+                },
+                style_data={
+                    "backgroundColor": "#ffffff",
+                    "color": "#333333",
+                    "borderBottom": "1px solid #ddd",
+                },
+                style_data_conditional=[
+                    {
+                        "if": {"state": "active"},
+                        "backgroundColor": "#e6f7ff",
+                        "border": "1px solid #1890ff",
+                    },
+                ],
+                style_table={
+                    "maxHeight": "400px",
+                    "overflowX": "auto",
+                    "overflowY": "auto",
+                },
+                style_as_list_view=True,
+                row_selectable="single",
+                selected_rows=[],
+            ),
+        ],
+        className="container",
+    )
+
+
+def get_materials_display(id, text_materials_div, text_table_div, **style_kwargs):
+    return html.Div(
+        [
+            html.Div(
+                [
+                    html.Div(
+                        text_materials_div,
+                        style={"textAlign": "center"},
+                    ),
+                ],
+                id=f"structure-container{id}",
+                className="container container-visu",
+            ),
+            html.Div(
+                id=f"properties-container{id}",
+                className="container container-table",
+                style={"width": "100%"},
+                children=[
+                    html.Div(
+                        text_table_div,
+                        style={"textAlign": "center"},
+                    ),
+                ],
+            ),
+        ],
+        className="container-row",
+    )

data_utils.py

@@ -0,0 +1,261 @@
+import os
+import re
+
+import crystal_toolkit.components as ctc
+import numpy as np
+import periodictable
+from dash import dcc, html
+from datasets import concatenate_datasets, load_dataset
+from pymatgen.analysis.structure_analyzer import SpacegroupAnalyzer
+
+HF_TOKEN = os.environ.get("HF_TOKEN")
+top_k = 500
+
+
+def get_dataset():
+    # Load only the train split of the dataset
+    datasets = []
+    subsets = [
+        "compatible_pbe",
+        "compatible_pbesol",
+        "compatible_scan",
+        "non_compatible",
+    ]
+
+    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", # not yet available in LeMat-Bulk
+                "immutable_id",
+                "elements",
+                "functional",
+                "stress_tensor",
+                "magnetic_moments",
+                "forces",
+                # "band_gap_direct", #future release
+                # "band_gap_indirect", #future release
+                "dos_ef",
+                # "charges", #future release
+                "functional",
+                "chemical_formula_reduced",
+                "chemical_formula_descriptive",
+                "total_magnetization",
+                "entalpic_fingerprint",
+            ],
+        )
+        datasets.append(dataset["train"])
+
+    return concatenate_datasets(datasets)
+
+
+display_columns = [
+    "chemical_formula_descriptive",
+    "functional",
+    "immutable_id",
+    "energy",
+]
+display_names = {
+    "chemical_formula_descriptive": "Formula",
+    "functional": "Functional",
+    "immutable_id": "Material ID",
+    "energy": "Energy (eV)",
+}
+
+# Global shared variables
+mapping_table_idx_dataset_idx = {}
+
+
+def build_formula_index(dataset, index_range=None, cache_path=None, empty_data=False):
+    if empty_data:
+        return np.zeros((1, 1)), {}
+
+    use_dataset = dataset
+    if index_range is not None:
+        use_dataset = dataset.select(index_range)
+
+    # Preprocessing step to create an index for the dataset
+    if cache_path is not None:
+        train_df = pickle.load(open(f"{cache_path}/train_df.pkl", "rb"))
+
+        dataset_index = pickle.load(open(f"{cache_path}/dataset_index.pkl", "rb"))
+    else:
+        train_df = use_dataset.select_columns(
+            ["chemical_formula_descriptive", "immutable_id"]
+        ).to_pandas()
+
+        pattern = re.compile(r"(?P<element>[A-Z][a-z]?)(?P<count>\d*)")
+        extracted = train_df["chemical_formula_descriptive"].str.extractall(pattern)
+        extracted["count"] = extracted["count"].replace("", "1").astype(int)
+
+        wide_df = (
+            extracted.reset_index().pivot_table(  # Move index to columns for pivoting
+                index="level_0",  # original row index
+                columns="element",
+                values="count",
+                aggfunc="sum",
+                fill_value=0,
+            )
+        )
+
+        all_elements = [el.symbol for el in periodictable.elements]  # full element list
+        wide_df = wide_df.reindex(columns=all_elements, fill_value=0)
+
+        dataset_index = wide_df.values
+
+        dataset_index = dataset_index / np.sum(dataset_index, axis=1)[:, None]
+        dataset_index = (
+            dataset_index / np.linalg.norm(dataset_index, axis=1)[:, None]
+        )  # Normalize vectors
+
+    immutable_id_to_idx = train_df["immutable_id"].to_dict()
+    immutable_id_to_idx = {v: k for k, v in immutable_id_to_idx.items()}
+
+    return dataset_index, immutable_id_to_idx
+
+
+import pickle
+from pathlib import Path
+
+
+# TODO: Just load the index from a file
+def build_embeddings_index(empty_data=False):
+    if empty_data:
+        return None, {}, {}
+
+    features_dict = pickle.load(open("features_dict.pkl", "rb"))
+
+    from indexer import FAISSIndex
+
+    index = FAISSIndex()
+    for key in features_dict:
+        index.index.add(features_dict[key].reshape(1, -1))
+
+    idx_to_immutable_id = {i: key for i, key in enumerate(features_dict)}
+
+    # index = FAISSIndex.from_store("index.faiss")
+
+    return index, features_dict, idx_to_immutable_id
+
+
+def search_materials(
+    query, dataset, dataset_index, mapping_table_idx_dataset_idx, map_periodic_table
+):
+    n_elements = len(map_periodic_table)
+    query_vector = np.zeros(n_elements)
+
+    if "," in query:
+        element_list = [el.strip() for el in query.split(",")]
+        for el in element_list:
+            query_vector[map_periodic_table[el]] = 1
+    else:
+        # Formula
+        import re
+
+        matches = re.findall(r"([A-Z][a-z]{0,2})(\d*)", query)
+        for el, numb in matches:
+            numb = int(numb) if numb else 1
+            query_vector[map_periodic_table[el]] = numb
+
+    similarity = np.dot(dataset_index, query_vector) / (np.linalg.norm(query_vector))
+    indices = np.argsort(similarity)[::-1][:top_k]
+
+    options = [dataset[int(i)] for i in indices]
+
+    mapping_table_idx_dataset_idx.clear()
+    for i, idx in enumerate(indices):
+        mapping_table_idx_dataset_idx[int(i)] = int(idx)
+
+    return options
+
+
+def get_properties_table(
+    row, structure, sga, properties_container_update, container_type="query"
+):
+    properties = {
+        "Material ID": row["immutable_id"],
+        "Formula": row["chemical_formula_descriptive"],
+        "Energy per atom (eV/atom)": round(
+            row["energy"] / len(row["species_at_sites"]), 3
+        ),
+        # "Band Gap (eV)": row["band_gap_direct"] or row["band_gap_indirect"], #future release
+        "Total Magnetization (μB)": (
+            round(row["total_magnetization"], 3)
+            if row["total_magnetization"] is not None
+            else None
+        ),
+        "Density (g/cm^3)": round(structure.density, 3),
+        "Fermi energy level (eV)": (
+            round(row["dos_ef"], 3) if row["dos_ef"] is not None else None
+        ),
+        "Crystal system": sga.get_crystal_system(),
+        "International Spacegroup": sga.get_symmetry_dataset().international,
+        "Magnetic moments (μB)": np.round(row["magnetic_moments"], 3),
+        "Stress tensor (kB)": np.round(row["stress_tensor"], 3),
+        "Forces on atoms (eV/A)": np.round(row["forces"], 3),
+        # "Bader charges (e-)": np.round(row["charges"], 3), # future release
+        "DFT Functional": row["functional"],
+        "Entalpic fingerprint": row["entalpic_fingerprint"],
+    }
+
+    style = {
+        "padding": "10px",
+        "borderBottom": "1px solid #ddd",
+    }
+
+    if container_type == "query":
+        properties_container_update[0] = properties
+    else:
+        properties_container_update[1] = properties
+        # if (type(value) in [str, float]) and (
+        #     properties_container_update[0][key] == properties_container_update[1][key]
+        # ):
+        #     style["backgroundColor"] = "#e6f7ff"
+
+    # Format properties as an HTML table
+    properties_html = html.Table(
+        [
+            html.Tbody(
+                [
+                    html.Tr(
+                        [
+                            html.Th(
+                                key,
+                                style={
+                                    "padding": "10px",
+                                    "verticalAlign": "middle",
+                                },
+                            ),
+                            html.Td(
+                                str(value),
+                                style=style,
+                            ),
+                        ],
+                    )
+                    for key, value in properties.items()
+                ],
+            )
+        ],
+        style={
+            "width": "100%",
+            "borderCollapse": "collapse",
+            "fontFamily": "'Arial', sans-serif",
+            "fontSize": "14px",
+            "color": "#333333",
+        },
+    )
+
+    return properties_html
+
+
+def get_crystal_plot(structure):
+    sga = SpacegroupAnalyzer(structure)
+    # Create the StructureMoleculeComponent
+    structure_component = ctc.StructureMoleculeComponent(structure)
+    return structure_component.layout(), sga