Spaces:
Running
Running
| from pathlib import Path | |
| import numpy as np | |
| import pandas as pd | |
| import plotly.colors as pcolors | |
| import plotly.express as px | |
| import plotly.graph_objects as go | |
| import streamlit as st | |
| from mlip_arena.models import REGISTRY | |
| DATA_DIR = Path("mlip_arena/tasks/stability") | |
| st.markdown( | |
| """ | |
| # Stability | |
| """ | |
| ) | |
| st.markdown("### Methods") | |
| container = st.container(border=True) | |
| 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.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) | |
| color_sequence = color_palettes[palette_name] | |
| if not models: | |
| st.stop() | |
| families = [REGISTRY[str(model)]["family"] for model in models] | |
| 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()) | |
| } | |
| fig = px.scatter( | |
| 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, | |
| 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) | |
| # 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 | |
| # )) | |
| # # Update layout for clarity | |
| # fig.update_layout( | |
| # title="Histogram of Total Steps", | |
| # xaxis_title="Count", | |
| # yaxis_title="Total Steps Range" | |
| # ) | |
| st.plotly_chart(fig) | |