DanielEmeka/FormationEnergyPredictor

1import pandas as pd2import gradio as gr3from sklearn.pipeline import Pipeline4from sklearn.ensemble import GradientBoostingRegressor5from sklearn.preprocessing import OneHotEncoder, StandardScaler6from sklearn.impute import SimpleImputer7from sklearn.compose import ColumnTransformer8from sklearn.model_selection import train_test_split9from sklearn.metrics import r2_score10import matplotlib.pyplot as plt11 12# Load dataset13df = pd.read_csv("data.csv")14 15# Keep only top features16categorical = ["A site #1", "B site #1", "X site"]17numerical = [18    "Number of elements",19    "Density_AB_avg",20    "Ionization Energy (kJ/mol)_AB_avg",21    "Atomic Volume (cm³/mol)_AB_avg"22]23target = "formation_energy (eV/atom)"24 25# Drop NaNs and prepare26df = df.dropna(subset=[target])27df = df[categorical + numerical + [target]]28X = df[categorical + numerical]29y = df[target]30 31# Preprocessing and model32preprocessor = ColumnTransformer([33    ("cat", Pipeline([34        ("imputer", SimpleImputer(strategy="most_frequent")),35        ("onehot", OneHotEncoder(handle_unknown="ignore"))36    ]), categorical),37    ("num", Pipeline([38        ("imputer", SimpleImputer(strategy="mean")),39        ("scaler", StandardScaler())40    ]), numerical)41])42 43model = Pipeline([44    ("prep", preprocessor),45    ("reg", GradientBoostingRegressor(n_estimators=300, learning_rate=0.05, max_depth=5, random_state=42))46])47 48# Train model49X_train, X_test, y_train, y_test = train_test_split(X, y, random_state=42, test_size=0.2)50model.fit(X_train, y_train)51r2 = r2_score(y_test, model.predict(X_test))52 53# Prediction + plot54def predict_and_plot(a1, b1, x, num_elem, density, ion_energy, atomic_vol):55    data = {56        "A site #1": a1,57        "B site #1": b1,58        "X site": x,59        "Number of elements": float(num_elem),60        "Density_AB_avg": float(density),61        "Ionization Energy (kJ/mol)_AB_avg": float(ion_energy),62        "Atomic Volume (cm³/mol)_AB_avg": float(atomic_vol)63    }64    df_input = pd.DataFrame([data])65    pred = model.predict(df_input)[0]66 67    # Stability logic68    if pred < -1.0:69        status = "🟢 Stable"70    elif -1.0 <= pred <= 0.5:71        status = "🟡 Metastable"72    else:73        status = "🔴 Unstable"74 75    # Plot76    fig, ax = plt.subplots()77    ax.barh(["Formation Energy"], [pred], color="green" if pred < -1 else "orange" if pred <= 0.5 else "red")78    ax.set_xlim(-3, 2)79    ax.set_xlabel("eV/atom")80    ax.set_title(f"Prediction: {round(pred, 4)} eV/atom — {status}")81    plt.tight_layout()82 83    return round(pred, 5), status, fig84 85# Inputs86inputs = [87    gr.Textbox(label="A site #1"),88    gr.Textbox(label="B site #1"),89    gr.Textbox(label="X site"),90    gr.Number(label="Number of elements", value=5),91    gr.Number(label="Density_AB_avg", value=5.5),92    gr.Number(label="Ionization Energy (kJ/mol)_AB_avg", value=700),93    gr.Number(label="Atomic Volume (cm³/mol)_AB_avg", value=10.0)94]95 96# Interface97demo = gr.Interface(98    fn=predict_and_plot,99    inputs=inputs,100    outputs=[101        gr.Number(label="Predicted Formation Energy (eV/atom)"),102        gr.Text(label="Stability Status"),103        gr.Plot(label="Stability Visualization")104    ],105    title="Formation Energy Predictor",106    description=(107        "🎯 This tool predicts the **formation energy** (eV/atom) of a compound "108        "based on elemental and physical properties.\n\n"109        "**Interpretation**:\n"110        "- 🟢 Low/Negative → Stable\n"111        "- 🟡 Close to Zero → Metastable\n"112        "- 🔴 Positive → Unstable\n\n"113        f"📈 Model trained with R² score: **{round(r2, 4)}**"114    )115)116 117if __name__ == "__main__":118    demo.launch()119