Spaces:
Running
Running
File size: 4,370 Bytes
1fe249a e517f23 1fe249a e517f23 1fe249a 2e5c86d 1fe249a 2e5c86d e517f23 2e5c86d e517f23 2e5c86d e517f23 2e5c86d e517f23 2e5c86d e517f23 2e5c86d 1fe249a e517f23 2e5c86d e517f23 2e5c86d e517f23 2e5c86d e517f23 1fe249a 2e5c86d 1fe249a 2e5c86d 1fe249a 2e5c86d 1fe249a 2e5c86d |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 |
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)
|