add md step bar charts

mlip_arena/tasks/stability/chgnet/chloride-salts.json

@@ -1 +1 @@
-[{"material_id":"mp-568540","formula":"Ag135Cl135Hg108P216","method":"CHGNet","natoms":594,"total_time_seconds":878.366427,"total_steps":340,"steps_per_second":0.3870821898,"seconds_per_step":2.5834306676,"seconds_per_step_per_atom":0.0043492099},{"material_id":"mp-28361","formula":"Al72Cd36Cl288","method":"CHGNet","natoms":396,"total_time_seconds":2024.840096,"total_steps":49990,"steps_per_second":24.6883692686,"seconds_per_step":0.0405049029,"seconds_per_step_per_atom":0.0001022851},{"material_id":"mp-566301","formula":"As48Cl16O192Pb80","method":"CHGNet","natoms":336,"total_time_seconds":2443.199102,"total_steps":49990,"steps_per_second":20.4608785093,"seconds_per_step":0.0488737568,"seconds_per_step_per_atom":0.0001454576},{"material_id":"mp-553919","formula":"Ag162Cl27F81Mo54O189","method":"CHGNet","natoms":513,"total_time_seconds":49.747983,"total_steps":210,"steps_per_second":4.2212766696,"seconds_per_step":0.2368951571,"seconds_per_step_per_atom":0.0004617839},{"material_id":"mp-567334","formula":"Ag64B384Cl384","method":"CHGNet","natoms":832,"total_time_seconds":6094.657858,"total_steps":4020,"steps_per_second":0.6595940402,"seconds_per_step":1.5160840443,"seconds_per_step_per_atom":0.0018222164},{"material_id":"mp-637206","formula":"Al64Bi64Cl256Se64","method":"CHGNet","natoms":448,"total_time_seconds":2102.670703,"total_steps":49990,"steps_per_second":23.7745263339,"seconds_per_step":0.0420618264,"seconds_per_step_per_atom":0.000093888},{"material_id":"mp-27346","formula":"Al48Cl192Hg72","method":"CHGNet","natoms":312,"total_time_seconds":620.864262,"total_steps":190,"steps_per_second":0.3060250229,"seconds_per_step":3.2677066421,"seconds_per_step_per_atom":0.0104734187},{"material_id":"mp-680722","formula":"As32Cl64O272Pb224","method":"CHGNet","natoms":592,"total_time_seconds":153.997307,"total_steps":540,"steps_per_second":3.5065548257,"seconds_per_step":0.2851801981,"seconds_per_step_per_atom":0.0004817233},{"material_id":"mp-567734","formula":"Al32As96Cl128Se128","method":"CHGNet","natoms":384,"total_time_seconds":1934.45237,"total_steps":49990,"steps_per_second":25.8419389256,"seconds_per_step":0.0386967868,"seconds_per_step_per_atom":0.0001007729},{"material_id":"mp-672352","formula":"Al48Cl288I144","method":"CHGNet","natoms":480,"total_time_seconds":697.476109,"total_steps":1940,"steps_per_second":2.7814572786,"seconds_per_step":0.3595237675,"seconds_per_step_per_atom":0.0007490078},{"material_id":"mp-27863","formula":"Al100Cl100O100","method":"CHGNet","natoms":300,"total_time_seconds":2977.600956,"total_steps":1181,"steps_per_second":0.3966280296,"seconds_per_step":2.5212539848,"seconds_per_step_per_atom":0.0084041799}]
+[{"material_id":"mp-560328","formula":"Ag60Cl12P16S64","method":"CHGNet","natoms":152,"total_time_seconds":2505.24287,"total_steps":49990,"steps_per_second":19.9541531876,"seconds_per_step":0.0501148804,"seconds_per_step_per_atom":0.0003297032}]

mlip_arena/tasks/stability/equiformer/chloride-salts.json

@@ -1 +0,0 @@
-[{"material_id":"mp-25469","formula":"Al54Cl162","method":"EquiformerV2(OC22)","natoms":216,"total_time_seconds":1289.686627,"total_steps":49990,"steps_per_second":38.7613540789,"seconds_per_step":0.0257988923,"seconds_per_step_per_atom":0.0001194393},{"material_id":"mp-568540","formula":"Ag135Cl135Hg108P216","method":"EquiformerV2(OC22)","natoms":594,"total_time_seconds":855.203801,"total_steps":340,"steps_per_second":0.3975660534,"seconds_per_step":2.5153052971,"seconds_per_step_per_atom":0.0042345207},{"material_id":"mp-568923","formula":"B96Ba192Cl96N192","method":"EquiformerV2(OC22)","natoms":576,"total_time_seconds":799.132371,"total_steps":11900,"steps_per_second":14.891149992,"seconds_per_step":0.0671539808,"seconds_per_step_per_atom":0.0001165868},{"material_id":"mp-568443","formula":"Al32Cl128V16","method":"EquiformerV2(OC22)","natoms":176,"total_time_seconds":1179.514714,"total_steps":49990,"steps_per_second":42.3818367051,"seconds_per_step":0.0235950133,"seconds_per_step_per_atom":0.0001340626},{"material_id":"mp-28361","formula":"Al72Cd36Cl288","method":"EquiformerV2(OC22)","natoms":396,"total_time_seconds":516.631174,"total_steps":1301,"steps_per_second":2.5182375077,"seconds_per_step":0.3971031314,"seconds_per_step_per_atom":0.0010027857},{"material_id":"mp-1214746","formula":"Ba48Cl168Dy24","method":"EquiformerV2(OC22)","natoms":240,"total_time_seconds":1189.535266,"total_steps":40720,"steps_per_second":34.2318560566,"seconds_per_step":0.0292125556,"seconds_per_step_per_atom":0.000121719},{"material_id":"mp-22993","formula":"Ag64Cl64O256","method":"EquiformerV2(OC22)","natoms":384,"total_time_seconds":2998.510846,"total_steps":49990,"steps_per_second":16.6716088643,"seconds_per_step":0.0599822134,"seconds_per_step_per_atom":0.0001562037},{"material_id":"mp-566301","formula":"As48Cl16O192Pb80","method":"EquiformerV2(OC22)","natoms":336,"total_time_seconds":7681.971428,"total_steps":3123,"steps_per_second":0.4065362686,"seconds_per_step":2.4598051322,"seconds_per_step_per_atom":0.0073208486},{"material_id":"mp-555805","formula":"Al36Cl252S36","method":"EquiformerV2(OC22)","natoms":324,"total_time_seconds":2034.579734,"total_steps":49990,"steps_per_second":24.5701847731,"seconds_per_step":0.0406997346,"seconds_per_step_per_atom":0.0001256165},{"material_id":"mp-1214988","formula":"Ag16Cl208Mo96","method":"EquiformerV2(OC22)","natoms":320,"total_time_seconds":1770.09089,"total_steps":49990,"steps_per_second":28.2414876447,"seconds_per_step":0.0354088996,"seconds_per_step_per_atom":0.0001106528},{"material_id":"mp-567334","formula":"Ag64B384Cl384","method":"EquiformerV2(OC22)","natoms":832,"total_time_seconds":3567.962024,"total_steps":4214,"steps_per_second":1.1810663823,"seconds_per_step":0.8466924594,"seconds_per_step_per_atom":0.0010176592},{"material_id":"mp-638061","formula":"Al48Cl192In48","method":"EquiformerV2(OC22)","natoms":288,"total_time_seconds":1576.40944,"total_steps":49990,"steps_per_second":31.7113046468,"seconds_per_step":0.0315344957,"seconds_per_step_per_atom":0.0001094948},{"material_id":"mp-560833","formula":"B72Cl24O240Pb144","method":"EquiformerV2(OC22)","natoms":480,"total_time_seconds":3217.219664,"total_steps":49990,"steps_per_second":15.5382613626,"seconds_per_step":0.0643572647,"seconds_per_step_per_atom":0.0001340776},{"material_id":"mp-560549","formula":"Au64Cl64O192Se64","method":"EquiformerV2(OC22)","natoms":384,"total_time_seconds":2518.062072,"total_steps":49990,"steps_per_second":19.8525685907,"seconds_per_step":0.0503713157,"seconds_per_step_per_atom":0.0001311753},{"material_id":"mp-637206","formula":"Al64Bi64Cl256Se64","method":"EquiformerV2(OC22)","natoms":448,"total_time_seconds":1762.787761,"total_steps":1895,"steps_per_second":1.0750017909,"seconds_per_step":0.9302310084,"seconds_per_step_per_atom":0.0020764085},{"material_id":"mp-510724","formula":"As72Cl24O216Pb120","method":"EquiformerV2(OC22)","natoms":432,"total_time_seconds":2746.961559,"total_steps":49990,"steps_per_second":18.1982888826,"seconds_per_step":0.0549502212,"seconds_per_step_per_atom":0.0001271996},{"material_id":"mp-551456","formula":"Ba80Cl40Cu40O80","method":"EquiformerV2(OC22)","natoms":240,"total_time_seconds":48569.923768,"total_steps":1294,"steps_per_second":0.0266420019,"seconds_per_step":37.5347169768,"seconds_per_step_per_atom":0.1563946541},{"material_id":"mp-23575","formula":"As32Cl96F192Se96","method":"EquiformerV2(OC22)","natoms":416,"total_time_seconds":7242.212461,"total_steps":3086,"steps_per_second":0.4261128787,"seconds_per_step":2.3467960016,"seconds_per_step_per_atom":0.0056413365},{"material_id":"mp-27346","formula":"Al48Cl192Hg72","method":"EquiformerV2(OC22)","natoms":312,"total_time_seconds":56.74321,"total_steps":320,"steps_per_second":5.6394412653,"seconds_per_step":0.1773225312,"seconds_per_step_per_atom":0.0005683414},{"material_id":"mp-30023","formula":"Al48Ba24Cl192","method":"EquiformerV2(OC22)","natoms":264,"total_time_seconds":1444.744189,"total_steps":49990,"steps_per_second":34.6012812376,"seconds_per_step":0.0289006639,"seconds_per_step_per_atom":0.0001094722},{"material_id":"mp-628665","formula":"Al64Bi64Cl256Te64","method":"EquiformerV2(OC22)","natoms":448,"total_time_seconds":1954.564983,"total_steps":49990,"steps_per_second":25.576023532,"seconds_per_step":0.0390991195,"seconds_per_step_per_atom":0.0000872748},{"material_id":"mp-28192","formula":"Au32Cl160S32","method":"EquiformerV2(OC22)","natoms":224,"total_time_seconds":9731.339255,"total_steps":59,"steps_per_second":0.0060628859,"seconds_per_step":164.9379534746,"seconds_per_step_per_atom":0.7363301494},{"material_id":"mp-28372","formula":"Au32Cl224Se32","method":"EquiformerV2(OC22)","natoms":288,"total_time_seconds":1591.957321,"total_steps":49990,"steps_per_second":31.4015955959,"seconds_per_step":0.0318455155,"seconds_per_step_per_atom":0.0001105747},{"material_id":"mp-1204489","formula":"Al32Br32Cl128Sb64Te64","method":"EquiformerV2(OC22)","natoms":320,"total_time_seconds":1641.572844,"total_steps":49990,"steps_per_second":30.4525018081,"seconds_per_step":0.0328380245,"seconds_per_step_per_atom":0.0001026188},{"material_id":"mp-556418","formula":"Al36Cl108O72S36","method":"EquiformerV2(OC22)","natoms":252,"total_time_seconds":1476.543336,"total_steps":49990,"steps_per_second":33.856100787,"seconds_per_step":0.0295367741,"seconds_per_step_per_atom":0.0001172094},{"material_id":"mp-556587","formula":"Au32Cl224S32","method":"EquiformerV2(OC22)","natoms":288,"total_time_seconds":1547.017976,"total_steps":49990,"steps_per_second":32.3137809486,"seconds_per_step":0.0309465488,"seconds_per_step_per_atom":0.0001074533},{"material_id":"mp-561268","formula":"As72Cl24O144Pb48","method":"EquiformerV2(OC22)","natoms":288,"total_time_seconds":2060.350209,"total_steps":49990,"steps_per_second":24.2628654981,"seconds_per_step":0.0412152472,"seconds_per_step_per_atom":0.0001431085},{"material_id":"mp-32780","formula":"Au108Cl108","method":"EquiformerV2(OC22)","natoms":216,"total_time_seconds":1337.166632,"total_steps":49990,"steps_per_second":37.3850190423,"seconds_per_step":0.0267486824,"seconds_per_step_per_atom":0.0001238365},{"material_id":"mp-556705","formula":"As32Cl96F192S32","method":"EquiformerV2(OC22)","natoms":352,"total_time_seconds":2386.766314,"total_steps":49990,"steps_per_second":20.9446562518,"seconds_per_step":0.0477448753,"seconds_per_step_per_atom":0.0001356389},{"material_id":"mp-27647","formula":"Au72Cl216","method":"EquiformerV2(OC22)","natoms":288,"total_time_seconds":1436.317195,"total_steps":49990,"steps_per_second":34.80428987,"seconds_per_step":0.0287320903,"seconds_per_step_per_atom":0.0000997642},{"material_id":"mp-651358","formula":"Ba48Cl168Er24","method":"EquiformerV2(OC22)","natoms":240,"total_time_seconds":1420.480535,"total_steps":49990,"steps_per_second":35.1923160989,"seconds_per_step":0.0284152938,"seconds_per_step_per_atom":0.0001183971},{"material_id":"mp-540628","formula":"Al64Cl256Te128","method":"EquiformerV2(OC22)","natoms":448,"total_time_seconds":2025.528538,"total_steps":49990,"steps_per_second":24.6799781204,"seconds_per_step":0.0405186745,"seconds_per_step_per_atom":0.0000904435},{"material_id":"mp-557103","formula":"B160Cl32O288Pb64","method":"EquiformerV2(OC22)","natoms":544,"total_time_seconds":3527.414988,"total_steps":40590,"steps_per_second":11.5070101301,"seconds_per_step":0.0869035474,"seconds_per_step_per_atom":0.0001597492},{"material_id":"mp-573751","formula":"Al48Cl336Te48","method":"EquiformerV2(OC22)","natoms":432,"total_time_seconds":1963.111202,"total_steps":49990,"steps_per_second":25.4646807318,"seconds_per_step":0.0392700781,"seconds_per_step_per_atom":0.000090903},{"material_id":"mp-680722","formula":"As32Cl64O272Pb224","method":"EquiformerV2(OC22)","natoms":592,"total_time_seconds":152.850709,"total_steps":540,"steps_per_second":3.5328589807,"seconds_per_step":0.2830568685,"seconds_per_step_per_atom":0.0004781366},{"material_id":"mp-567734","formula":"Al32As96Cl128Se128","method":"EquiformerV2(OC22)","natoms":384,"total_time_seconds":1609.555767,"total_steps":41500,"steps_per_second":25.7835117309,"seconds_per_step":0.0387844763,"seconds_per_step_per_atom":0.0001010012},{"material_id":"mp-570340","formula":"Ag48Au48Cl192","method":"EquiformerV2(OC22)","natoms":288,"total_time_seconds":1607.842826,"total_steps":49990,"steps_per_second":31.0913474822,"seconds_per_step":0.0321632892,"seconds_per_step_per_atom":0.0001116781},{"material_id":"mp-1190390","formula":"Ba54Cl108O432","method":"EquiformerV2(OC22)","natoms":594,"total_time_seconds":126.959062,"total_steps":370,"steps_per_second":2.9143252492,"seconds_per_step":0.3431326,"seconds_per_step_per_atom":0.0005776643},{"material_id":"mp-672352","formula":"Al48Cl288I144","method":"EquiformerV2(OC22)","natoms":480,"total_time_seconds":10.902331,"total_steps":390,"steps_per_second":35.7721665211,"seconds_per_step":0.0279546949,"seconds_per_step_per_atom":0.0000582389},{"material_id":"mp-27265","formula":"Au72Cl72O72","method":"EquiformerV2(OC22)","natoms":216,"total_time_seconds":1590.120361,"total_steps":49990,"steps_per_second":31.437871765,"seconds_per_step":0.031808769,"seconds_per_step_per_atom":0.0001472628},{"material_id":"mp-27863","formula":"Al100Cl100O100","method":"EquiformerV2(OC22)","natoms":300,"total_time_seconds":4939.446747,"total_steps":2315,"steps_per_second":0.4686759709,"seconds_per_step":2.1336703011,"seconds_per_step_per_atom":0.0071122343}]

mlip_arena/tasks/stability/mace-mp/chloride-salts.json

@@ -1 +1 @@
-[{"material_id":"mp-567334","formula":"Ag64B384Cl384","method":"MACE-MP(M)","natoms":832,"total_time_seconds":3578.503151,"total_steps":40470,"steps_per_second":11.3091978105,"seconds_per_step":0.0884236015,"seconds_per_step_per_atom":0.0001062784},{"material_id":"mp-30273","formula":"Ac96Cl96O96","method":"MACE-MP(M)","natoms":288,"total_time_seconds":1974.313713,"total_steps":49990,"steps_per_second":25.3201908445,"seconds_per_step":0.0394941731,"seconds_per_step_per_atom":0.0001371325},{"material_id":"mp-573465","formula":"Ag320Cl64Te128","method":"MACE-MP(M)","natoms":512,"total_time_seconds":2296.972442,"total_steps":39000,"steps_per_second":16.9788715297,"seconds_per_step":0.0588967293,"seconds_per_step_per_atom":0.0001150327},{"material_id":"mp-27971","formula":"Ac48Cl144","method":"MACE-MP(M)","natoms":192,"total_time_seconds":1228.139194,"total_steps":49990,"steps_per_second":40.7038552668,"seconds_per_step":0.0245676974,"seconds_per_step_per_atom":0.0001279568},{"material_id":"mp-560328","formula":"Ag60Cl12P16S64","method":"MACE-MP(M)","natoms":152,"total_time_seconds":350.457564,"total_steps":14360,"steps_per_second":40.9750037525,"seconds_per_step":0.0244051228,"seconds_per_step_per_atom":0.00016056}]
+[{"material_id":"mp-553919","formula":"Ag162Cl27F81Mo54O189","method":"MACE-MP(M)","natoms":513,"total_time_seconds":50.711128,"total_steps":210,"steps_per_second":4.1411029153,"seconds_per_step":0.2414815619,"seconds_per_step_per_atom":0.0004707243},{"material_id":"mp-567334","formula":"Ag64B384Cl384","method":"MACE-MP(M)","natoms":832,"total_time_seconds":3578.503151,"total_steps":40470,"steps_per_second":11.3091978105,"seconds_per_step":0.0884236015,"seconds_per_step_per_atom":0.0001062784},{"material_id":"mp-30273","formula":"Ac96Cl96O96","method":"MACE-MP(M)","natoms":288,"total_time_seconds":1974.313713,"total_steps":49990,"steps_per_second":25.3201908445,"seconds_per_step":0.0394941731,"seconds_per_step_per_atom":0.0001371325},{"material_id":"mp-573465","formula":"Ag320Cl64Te128","method":"MACE-MP(M)","natoms":512,"total_time_seconds":2951.833216,"total_steps":49990,"steps_per_second":16.9352386609,"seconds_per_step":0.059048474,"seconds_per_step_per_atom":0.0001153291},{"material_id":"mp-27971","formula":"Ac48Cl144","method":"MACE-MP(M)","natoms":192,"total_time_seconds":1228.139194,"total_steps":49990,"steps_per_second":40.7038552668,"seconds_per_step":0.0245676974,"seconds_per_step_per_atom":0.0001279568},{"material_id":"mp-864635","formula":"Ag108Cl162","method":"MACE-MP(M)","natoms":270,"total_time_seconds":1579.32538,"total_steps":49990,"steps_per_second":31.6527554316,"seconds_per_step":0.0315928262,"seconds_per_step_per_atom":0.0001170105}]

mlip_arena/tasks/stability/run.py

@@ -198,6 +198,7 @@ def md(
         time_step = 0.5 if has_h_isotope else 2.0
 
     n_steps = int(total_time / time_step)
+    target_steps = n_steps
 
     t_schedule, p_schedule = _get_ensemble_schedule(
         ensemble=ensemble,
@@ -280,6 +281,7 @@ def md(
             dyn.atoms.info["restart"] = last_step
             dyn.atoms.info["datetime"] = datetime.now()
             dyn.atoms.info["step"] = step
+            dyn.atoms.info["target_steps"] = target_steps
             if ensemble == "nve":
                 return
             dyn.set_temperature(temperature_K=t_schedule[step])

requirements.txt

@@ -6,4 +6,4 @@ ase==3.23.0
 torch==2.2.1
 pymatgen==2024.4.13
 bokeh==2.4.3
-statsmodels
+statsmodels==0.14.2

serve/tasks/homonuclear-diatomics.py

@@ -11,7 +11,14 @@ from scipy.interpolate import CubicSpline
 
 # from mlip_arena.models.utils import MLIPMap
 
-st.markdown("# Homonuclear diatomics")
+st.markdown(
+"""
+# Homonuclear diatomics
+
+Homonuclear diatomics are molecules composed of two atoms of the same element.
+The potential energy curves of homonuclear diatomics are the most fundamental interactions between atoms in quantum chemistry.
+"""
+)
 
 st.markdown("### Methods")
 container = st.container(border=True)
@@ -22,7 +29,7 @@ st.markdown("### Settings")
 vis = st.container(border=True)
 energy_plot = vis.checkbox("Show energy curves", value=True)
 force_plot = vis.checkbox("Show force curves", value=False)
-ncols = vis.select_slider("Number of columns", options=[1, 2, 3, 4], value=3)
+ncols = vis.select_slider("Number of columns", options=[1, 2, 3, 4], value=2)
 
 # Get all attributes from pcolors.qualitative
 all_attributes = dir(pcolors.qualitative)
@@ -140,7 +147,7 @@ for i, symbol in enumerate(chemical_symbols[1:]):
     )
 
     # Set x-axis title
-    fig.update_xaxes(title_text="Bond length [Å]")
+    fig.update_xaxes(title_text="Distance [Å]")
 
     # Set y-axes titles
     if energy_plot:
@@ -165,4 +172,4 @@ for i, symbol in enumerate(chemical_symbols[1:]):
             ),
         )
 
-    cols[i % ncols].plotly_chart(fig, use_container_width=True, height=250)
+    cols[i % ncols].plotly_chart(fig, use_container_width=True)

serve/tasks/stability.py

@@ -12,57 +12,147 @@ from mlip_arena.models import REGISTRY
 DATA_DIR = Path("mlip_arena/tasks/stability")
 
 
-st.markdown("# Stability")
+st.markdown(
+    """
+# Stability
+
+"""
+)
 
 st.markdown("### Methods")
 container = st.container(border=True)
-models = container.multiselect("MLIPs", REGISTRY.keys(), ['MACE-MP(M)', "CHGNet", "EquiformerV2(OC22)"])
+models = container.multiselect("MLIPs", REGISTRY.keys(), ["MACE-MP(M)", "CHGNet"])
 
 st.markdown("### Settings")
 vis = st.container(border=True)
 # Get all attributes from pcolors.qualitative
 all_attributes = dir(pcolors.qualitative)
-color_palettes = {attr: getattr(pcolors.qualitative, attr) for attr in all_attributes if isinstance(getattr(pcolors.qualitative, attr), list)}
+color_palettes = {
+    attr: getattr(pcolors.qualitative, attr)
+    for attr in all_attributes
+    if isinstance(getattr(pcolors.qualitative, attr), list)
+}
 color_palettes.pop("__all__", None)
 
 palette_names = list(color_palettes.keys())
 palette_colors = list(color_palettes.values())
 
-palette_name = vis.selectbox(
-    "Color sequence",
-    options=palette_names, index=22
-)
+palette_name = vis.selectbox("Color sequence", options=palette_names, index=22)
 
 color_sequence = color_palettes[palette_name]
 
 if not models:
     st.stop()
 
-families = [REGISTRY[str(model)]['family'] for model in models]
+families = [REGISTRY[str(model)]["family"] for model in models]
 
-dfs = [pd.read_json(DATA_DIR / family.lower() /  "chloride-salts.json") for family in families]
+dfs = [
+    pd.read_json(DATA_DIR / family.lower() / "chloride-salts.json")
+    for family in families
+]
 df = pd.concat(dfs, ignore_index=True)
 df.drop_duplicates(inplace=True, subset=["material_id", "formula", "method"])
 
-method_color_mapping = {method: color_sequence[i % len(color_sequence)] for i, method in enumerate(df["method"].unique())}
+method_color_mapping = {
+    method: color_sequence[i % len(color_sequence)]
+    for i, method in enumerate(df["method"].unique())
+}
 
 
-
-# fig = px.scatter(df, x="natoms", y="seconds_per_step", trendline="ols", trendline_options=dict(log_y=True), log_y=True)
 fig = px.scatter(
-    df, x="natoms", y="steps_per_second", 
+    df,
+    x="natoms",
+    y="steps_per_second",
     color="method",
+    size="total_steps",
+    hover_data=["material_id", "formula"],
     color_discrete_map=method_color_mapping,
-    trendline="ols", trendline_options=dict(log_x=True), log_x=True
+    trendline="ols",
+    trendline_options=dict(log_x=True),
+    log_x=True,
+    title="Inference Speed",
+    labels={"steps_per_second": "Steps per second", "natoms": "Number of atoms"},
+)
+st.plotly_chart(fig)
+
+###
+
+fig = go.Figure()
+
+# Determine the bin edges for the entire dataset to keep them consistent across groups
+# bins = np.histogram_bin_edges(df['total_steps'], bins=10)
+
+max_steps = df["total_steps"].max()
+
+bins = np.append(np.arange(0, max_steps - 1, max_steps // 10), max_steps)
+bin_labels = [f"{bins[i]}-{bins[i+1]}" for i in range(len(bins)-1)]
+
+num_bins = len(bin_labels)
+colormap = px.colors.sequential.Redor
+indices = np.linspace(0, len(colormap) - 1, num_bins, dtype=int)
+bin_colors = [colormap[i] for i in indices]
+
+# Initialize a dictionary to hold the counts for each method and bin range
+counts_per_method = {method: [0] * len(bin_labels) for method in df['method'].unique()}
+
+# Populate the dictionary with counts
+for method, group in df.groupby('method'):
+    counts, _ = np.histogram(group['total_steps'], bins=bins)
+    counts_per_method[method] = counts
+
+# Create a figure
+fig = go.Figure()
+
+# Add a bar for each bin range across all methods
+for i, bin_label in enumerate(bin_labels):
+    for method, counts in counts_per_method.items():
+        fig.add_trace(go.Bar(
+            # name=method,  # This will be the legend entry
+            x=[counts[i]],  # Count for this bin
+            y=[method],  # Method as the y-axis category
+            # name=bin_label,
+            orientation='h',  # Horizontal bars
+            marker=dict(
+                color=bin_colors[i],
+                line=dict(color='rgb(248, 248, 249)', width=1)
+            ),
+            text=bin_label,
+            width=0.5
+        ))
+
+# Update the layout to stack the bars
+fig.update_layout(
+    barmode='stack',  # Stack the bars
+    title="Total MD Steps",
+    xaxis_title="Count",
+    yaxis_title="Method",
+    showlegend=False
 )
 
+# bins = np.linspace(0, 0.9, 10)
 
-event = st.plotly_chart(
-    fig, 
-    key="stability", 
-    on_select="rerun"
-)
+# for method, data in df.groupby("method"):
+
+#     # print(method, data)
+#     counts, bins = np.histogram(data['total_steps'])
+
+#     bin_labels = [f"{int(bins[i])}-{int(bins[i+1])}" for i in range(len(bins)-1)]
+
+#     # Create a horizontal bar chart
+#     fig = go.Figure(go.Bar(
+#         x=[counts[i]],  # Count for this bin
+#         y=[method],  # Method as the y-axis category
+#         # x=counts,  # Bar lengths
+#         # y=bin_labels,  # Bin labels as y-tick labels
+#         orientation='h'  # Horizontal bars
+#     ))
 
-event
 
+# # Update layout for clarity
+# fig.update_layout(
+#     title="Histogram of Total Steps",
+#     xaxis_title="Count",
+#     yaxis_title="Total Steps Range"
+# )
 
+st.plotly_chart(fig)