Skip to content

Commit

Permalink
add combution rank file
Browse files Browse the repository at this point in the history
  • Loading branch information
chiang-yuan committed Sep 23, 2024
1 parent 4cf6e46 commit b5af995
Showing 1 changed file with 49 additions and 0 deletions.
49 changes: 49 additions & 0 deletions serve/ranks/combustion.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,49 @@
from pathlib import Path

import numpy as np
import pandas as pd
import streamlit as st

from mlip_arena.models import REGISTRY as MODELS

valid_models = [
model
for model, metadata in MODELS.items()
if Path(__file__).stem in metadata.get("gpu-tasks", [])
]

DATA_DIR = Path("mlip_arena/tasks/combustion")

@st.cache_data
def get_data(models):
families = [MODELS[str(model)]["family"] for model in models]
dfs = [
pd.read_json(DATA_DIR / family.lower() / "hydrogen.json") for family in families
]
df = pd.concat(dfs, ignore_index=True)
df.drop_duplicates(inplace=True, subset=["formula", "method"])
return df

df = get_data(valid_models)

@st.cache_data
def get_com_drifts(df):
df_exploded = df.explode(["timestep", "com_drifts"]).reset_index(drop=True)

# Convert the 'com_drifts' column (which are arrays) into separate columns for x, y, and z components
df_exploded[["com_drift_x", "com_drift_y", "com_drift_z"]] = pd.DataFrame(
df_exploded["com_drifts"].tolist(), index=df_exploded.index
)

# Drop the original 'com_drifts' column
df_flat = df_exploded.drop(columns=["com_drifts"])

df_flat["total_com_drift"] = np.sqrt(
df_flat["com_drift_x"] ** 2 + df_flat["com_drift_y"] ** 2 + df_flat["com_drift_z"] ** 2
)

return df_flat

df_exploded = get_com_drifts(df)

table = pd.DataFrame()

0 comments on commit b5af995

Please sign in to comment.