cyrusyc commited on
Commit
b5af995
1 Parent(s): 4cf6e46

add combution rank file

Browse files
Files changed (1) hide show
  1. serve/ranks/combustion.py +49 -0
serve/ranks/combustion.py ADDED
@@ -0,0 +1,49 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from pathlib import Path
2
+
3
+ import numpy as np
4
+ import pandas as pd
5
+ import streamlit as st
6
+
7
+ from mlip_arena.models import REGISTRY as MODELS
8
+
9
+ valid_models = [
10
+ model
11
+ for model, metadata in MODELS.items()
12
+ if Path(__file__).stem in metadata.get("gpu-tasks", [])
13
+ ]
14
+
15
+ DATA_DIR = Path("mlip_arena/tasks/combustion")
16
+
17
+ @st.cache_data
18
+ def get_data(models):
19
+ families = [MODELS[str(model)]["family"] for model in models]
20
+ dfs = [
21
+ pd.read_json(DATA_DIR / family.lower() / "hydrogen.json") for family in families
22
+ ]
23
+ df = pd.concat(dfs, ignore_index=True)
24
+ df.drop_duplicates(inplace=True, subset=["formula", "method"])
25
+ return df
26
+
27
+ df = get_data(valid_models)
28
+
29
+ @st.cache_data
30
+ def get_com_drifts(df):
31
+ df_exploded = df.explode(["timestep", "com_drifts"]).reset_index(drop=True)
32
+
33
+ # Convert the 'com_drifts' column (which are arrays) into separate columns for x, y, and z components
34
+ df_exploded[["com_drift_x", "com_drift_y", "com_drift_z"]] = pd.DataFrame(
35
+ df_exploded["com_drifts"].tolist(), index=df_exploded.index
36
+ )
37
+
38
+ # Drop the original 'com_drifts' column
39
+ df_flat = df_exploded.drop(columns=["com_drifts"])
40
+
41
+ df_flat["total_com_drift"] = np.sqrt(
42
+ df_flat["com_drift_x"] ** 2 + df_flat["com_drift_y"] ** 2 + df_flat["com_drift_z"] ** 2
43
+ )
44
+
45
+ return df_flat
46
+
47
+ df_exploded = get_com_drifts(df)
48
+
49
+ table = pd.DataFrame()