From 585593a8cd01c4e42822cf445b8b9af17cbf71ff Mon Sep 17 00:00:00 2001
From: Shyue Ping Ong <shyuep@gmail.com>
Date: Tue, 19 Nov 2024 14:19:00 -0800
Subject: [PATCH] Update wioth natoms and nelements.

---
 app.py         | 84 +++++++++++++++++++++++++++++---------------------
 pyproject.toml | 12 ++++----
 2 files changed, 55 insertions(+), 41 deletions(-)

diff --git a/app.py b/app.py
index f425378..9b6c033 100644
--- a/app.py
+++ b/app.py
@@ -1,4 +1,5 @@
 """Main MatPES Explorer App."""
+
 from __future__ import annotations
 
 import collections
@@ -18,13 +19,16 @@
 # Set up MongoDB client and database
 CLIENT = MongoClient()
 DB = CLIENT["matpes"]
-#print(DB["PBE"].find_one())
+# print(DB["PBE"].find_one())
 RAW_DATA = {}
 for f in FUNCTIONALS:
     collection = DB[f]
     RAW_DATA[f] = pd.DataFrame(
-        collection.find({}, projection=["elements", "energy", "cohesive_energy_per_atom", "formation_energy", "natoms"])
+        collection.find(
+            {}, projection=["elements", "energy", "cohesive_energy_per_atom", "formation_energy", "natoms", "nelements"]
         )
+    )
+
 
 @functools.lru_cache
 def get_data(functional, el):
@@ -34,53 +38,63 @@ def get_data(functional, el):
         df = df[df["elements"].apply(lambda x: el in x)]
     return df
 
+
 # Initialize the Dash app with a Bootstrap theme
 external_stylesheets = [dbc.themes.CERULEAN]
 app = Dash(__name__, external_stylesheets=external_stylesheets)
 
 # Define the app layout
-app.layout = dbc.Container([
-    dbc.Row([
-        html.Div("MatPES Explorer", className="text-primary text-center fs-3")
-    ]),
-    dbc.Row([dcc.RadioItems(
-        options=[
-            {"label": f, "value": f} for f in FUNCTIONALS
-        ],
-        value="PBE",
-        id="functional"
-    ),
-    dcc.Dropdown(
-        options=[
-            {"label": el.symbol, "value": el.symbol} for el in Element
-        ],
-        id="el_filter"
-    )]),
-    dbc.Row([
-    dcc.Graph(
-        id="ptheatmap"
-    )]),
-    dbc.Row([
-    dcc.Graph(
-        id="coh_energy_hist"
-    ),
-    dcc.Graph(
-        id="form_energy_hist"
-    )
-    ])
-])
+app.layout = dbc.Container(
+    [
+        dbc.Row([html.Div("MatPES Explorer", className="text-primary text-center fs-3")]),
+        dbc.Row(
+            [
+                html.Label("Functional"),
+                dcc.RadioItems(options=[{"label": f, "value": f} for f in FUNCTIONALS], value="PBE", id="functional"),
+                html.Label("Element Filter"),
+                dcc.Dropdown(options=[{"label": el.symbol, "value": el.symbol} for el in Element], id="el_filter"),
+            ]
+        ),
+        dbc.Row([dcc.Graph(id="ptheatmap")]),
+        dbc.Row(
+            [
+                dbc.Col([dcc.Graph(id="coh_energy_hist")], width=6),
+                dbc.Col([dcc.Graph(id="form_energy_hist")], width=6),
+            ]
+        ),
+        dbc.Row(
+            [
+                dbc.Col([dcc.Graph(id="natoms_hist")], width=6),
+                dbc.Col([dcc.Graph(id="nelements_hist")], width=6),
+            ]
+        ),
+    ]
+)
+
 
 # Define callback to update the heatmap based on selected functional
 @callback(
-    [Output("ptheatmap", "figure"), Output("coh_energy_hist", "figure"), Output("form_energy_hist", "figure")],
-    [Input("functional", "value"), Input("el_filter", "value")]
+    [
+        Output("ptheatmap", "figure"),
+        Output("coh_energy_hist", "figure"),
+        Output("form_energy_hist", "figure"),
+        Output("natoms_hist", "figure"),
+        Output("nelements_hist", "figure"),
+    ],
+    [Input("functional", "value"), Input("el_filter", "value")],
 )
 def update_graph(functional, el_filter):
     """Update graph based on input."""
     df = get_data(functional, el_filter)
     el_count = collections.Counter(itertools.chain(*df["elements"]))
     heatmap_figure = pmv.ptable_heatmap_plotly(el_count, log=True)
-    return heatmap_figure, px.histogram(df, x="cohesive_energy_per_atom"), px.histogram(df, x="formation_energy")
+    return (
+        heatmap_figure,
+        px.histogram(df, x="cohesive_energy_per_atom"),
+        px.histogram(df, x="formation_energy"),
+        px.histogram(df, x="natoms"),
+        px.histogram(df, x="nelements"),
+    )
 
 
 # Run the app
diff --git a/pyproject.toml b/pyproject.toml
index 7cd02e4..b4c8465 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -59,9 +59,9 @@ default-tag = "0.0.1"
 line-length = 120
 
 [tool.ruff]
-target-version = "py38"
+target-version = "py310"
 line-length = 120
-select = [
+lint.select = [
   "B",    # flake8-bugbear
   "C4",   # flake8-comprehensions
   "D",    # pydocstyle
@@ -92,7 +92,7 @@ select = [
   "W",    # pycodestyle warning
   "YTT",  # flake8-2020
 ]
-ignore = [
+lint.ignore = [
   "B023",    # Function definition does not bind loop variable
   "B028",    # No explicit stacklevel keyword argument found
   "B904",    # Within an except clause, raise exceptions with ...
@@ -109,9 +109,9 @@ ignore = [
   "RUF012",  # Disable checks for mutable class args. This is a non-problem.
   "SIM105",  # Use contextlib.suppress(OSError) instead of try-except-pass
 ]
-pydocstyle.convention = "google"
-isort.required-imports = ["from __future__ import annotations"]
-isort.split-on-trailing-comma = false
+lint.pydocstyle.convention = "google"
+lint.isort.required-imports = ["from __future__ import annotations"]
+lint.isort.split-on-trailing-comma = false
 
 [tool.ruff.per-file-ignores]
 "__init__.py" = ["F401"]