CALCLOUD-474 Add script to calculate models for memory and runtime

Author: sashawood

scripts/manual_predict.py

@@ -0,0 +1,266 @@
+import boto3
+import joblib
+import matplotlib.pyplot as plt
+import numpy as np
+import pandas as pd
+from sklearn.ensemble import RandomForestRegressor, HistGradientBoostingRegressor
+from sklearn.metrics import mean_absolute_error, mean_squared_error, r2_score
+from sklearn.model_selection import train_test_split
+
+
+def get_dynamo_items(number_items):
+    """Get data from dynamo"""
+    session = boto3.Session(profile_name="aws-hst-repro-ops-Developer")
+    dynamodb = session.resource("dynamodb")
+    table = dynamodb.Table("calcloud-model-ops")
+    items = []
+    response = table.scan(Limit=number_items)
+    items.extend(response["Items"])
+
+    while "LastEvaluatedKey" in response and len(items) < number_items:
+        response = table.scan(ExclusiveStartKey=response["LastEvaluatedKey"], Limit=number_items)
+        items.extend(response["Items"])
+
+    items = items[:number_items]
+    return items
+
+
+input_keys = ("n_files", "total_mb", "drizcorr", "pctecorr", "crsplit", "subarray", "detector", "dtype", "instr")
+prediction_keys = ("mem_pred", "wall_pred", "bin_pred", "wc_mean", "wc_err", "wc_std", "x_mean", "x_files")
+actual_keys = ("memory", "wallclock", "mem_bin")
+
+
+def convert_elements_to_numeric_values(items):
+    """Clean up the data, converting strings to numeric values"""
+    for item in items:
+        item["inputs"] = {}
+        item["predictions"] = {}
+        item["actuals"] = {}
+        for key in input_keys:
+            value = item.pop(key)
+            try:
+                item["inputs"][key] = int(value)
+            except ValueError:
+                item["inputs"][key] = np.float64(value)
+        for key in prediction_keys:
+            value = item.pop(key, None)
+            if value is not None:
+                item["predictions"][key] = np.float64(value)
+        for key in actual_keys:
+            item["actuals"][key] = np.float64(item.pop(key))
+
+
+def calculate_memory_model(data):
+    """Calculate the model for memory, analyze result, present actual vs. predicted"""
+    df = pd.DataFrame(data)
+    original_df = pd.DataFrame(data)
+
+    df = pd.get_dummies(
+        original_df,
+        columns=["instr", "dtype", "detector", "drizcorr", "pctecorr", "crsplit", "subarray"],
+        drop_first=True,
+    )
+
+    X = df.drop(columns=["wallclock", "memory"])
+    y = df["memory"]
+
+    X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)
+
+    # Memory model
+    memory_model = RandomForestRegressor(
+        n_estimators=200,
+        random_state=42,
+    )
+    memory_model.fit(X_train, y_train)
+    actual_memory = y_test
+    predicted_memory = memory_model.predict(X_test)
+
+    memory_r2 = r2_score(y_test, predicted_memory)
+    memory_mae = mean_absolute_error(y_test, predicted_memory)
+    memory_rmse = np.sqrt(mean_squared_error(y_test, predicted_memory))
+    memory_mape = np.mean(np.abs((y_test - predicted_memory) / y_test)) * 100
+
+    print("\nMemory")
+    print(f"  R²   = {memory_r2:.3f}")
+    print(f"  MAE  = {memory_mae:.2f}")
+    print(f"  RMSE = {memory_rmse:.2f}")
+    print(f"  MAPE = {memory_mape:.1f}%")
+    print("")
+
+    # Save the model
+    joblib.dump(memory_model, "scripts/memory_rf_model.pkl")
+
+    # To use:
+    #     rf_model = joblib.load("scripts/memory_rf_model.pkl")
+    #     prediction = np.expm1(rf_model.predict(feature_df))
+
+    plot_actual_vs_predicted(
+        actual_memory,
+        predicted_memory,
+        xlabel="Actual Memory",
+        ylabel="Predicted Memory",
+        title="Memory Predicted vs Actual",
+    )
+
+
+def plot_actual_vs_predicted(actual, predicted, xlabel, ylabel, title, use_log_log_scale=False):
+    """Plot actual vs. predicted"""
+    plt.figure(figsize=(8, 8))
+
+    plt.scatter(actual, predicted, alpha=0.6, edgecolors="none")
+
+    if use_log_log_scale:
+        plt.xscale("log")
+        plt.yscale("log")
+
+    # Perfect prediction line
+    min_val = min(actual.min(), predicted.min())
+    max_val = max(actual.max(), predicted.max())
+
+    plt.plot([min_val, max_val], [min_val, max_val], "r--", linewidth=2, label="Perfect prediction")
+
+    plt.xlabel(xlabel)
+    plt.ylabel(ylabel)
+    plt.title(title)
+    plt.legend()
+    plt.grid(True, alpha=0.3)
+
+    plt.tight_layout()
+    plt.show()
+
+
+def calculate_wallclock_model(data):
+    """Calculate the model for wallclock, analyze result, present actual vs. predicted"""
+    df = pd.DataFrame(data)
+    original_df = pd.DataFrame(data)
+
+    df = pd.get_dummies(
+        original_df,
+        columns=["instr", "dtype", "detector", "drizcorr", "pctecorr", "crsplit", "subarray"],
+        drop_first=True,
+    )
+    df["log_total_mb"] = np.log1p(df["total_mb"])
+    df["log_n_files"] = np.log1p(df["n_files"])
+
+    X = df.drop(columns=["wallclock", "memory", "total_mb", "n_files"])
+    y = np.log1p(df["wallclock"])
+
+    X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)
+
+    # Wallclock model
+    wallclock_model = HistGradientBoostingRegressor(max_depth=5, learning_rate=0.05, max_iter=500, random_state=42)
+    wallclock_model.fit(X_train, y_train)
+
+    log_pred = wallclock_model.predict(X_test)
+
+    predicted_wallclock = np.expm1(log_pred)
+    actual_wallclock = np.expm1(y_test)
+
+    log_r2 = r2_score(y_test, log_pred)
+
+    mae = mean_absolute_error(y_test, log_pred)
+    rmse = np.sqrt(mean_squared_error(y_test, log_pred))
+
+    ape = (np.abs(predicted_wallclock - actual_wallclock) / actual_wallclock) * 100
+    mean_ape = np.mean(ape)
+    median_ape = np.median(ape)
+    ape_95 = np.percentile(ape, 95)
+    ape_99 = np.percentile(ape, 99)
+
+    print("\nWallclock")
+    print(f"  Log-space R²   = {log_r2:.3f}")
+    print(f"  Log-space MAE  = {mae:.2f}")
+    print(f"  Log-space RMSE = {rmse:.2f}")
+    print(f"  Mean APE       = {mean_ape:.1f}%")
+
+    print(f"  Median APE     = {median_ape:.1f}%")
+    print(f"  95th pct APE   = {ape_95:.1f}%")
+    print(f"  99th pct APE   = {ape_99:.1f}%")
+
+    joblib.dump({"model": wallclock_model, "columns": list(X.columns)}, "scripts/wallclock_model.pkl")
+
+    # To use:
+    #     saved = joblib.load("scripts/wallclock_model.pkl")
+    #     model = saved["model"]
+    #     prediction = np.expm1(model.predict(feature_df))
+
+    plot_actual_vs_predicted(
+        actual_wallclock,
+        predicted_wallclock,
+        xlabel="Actual Wallclock",
+        ylabel="Predicted Wallclock",
+        title="Wallclock: Predicted vs Actual (Log Scale)",
+        use_log_log_scale=True,
+    )
+
+
+def calculate_models(items):
+    """Create and evaluate the models"""
+    data = {}
+    for input_key in input_keys:
+        data[input_key] = [x["inputs"][input_key] for x in items]
+    for output_key in ("wallclock", "memory"):
+        data[output_key] = [x["actuals"][output_key] for x in items]
+    calculate_memory_model(data)
+    calculate_wallclock_model(data)
+
+
+def plot_bins(items):
+    """Plot memory and wallclock in bins to see distribution"""
+    fig, axs = plt.subplots(1, 2, figsize=(18, 6), layout="constrained")
+
+    deciles_array = np.arange(10, 100, 10)
+
+    # Plot wallclock
+    percent = 100
+    cutoff = 0
+    wallclock = [x["actuals"]["wallclock"] for x in items]
+    wallclock = [x for x in wallclock if x > cutoff]
+    deciles = np.percentile(wallclock, deciles_array)
+
+    label = f"{len(wallclock)} samples\n"
+    if cutoff != 0:
+        label += f"(wallclock > {cutoff} s)\n"
+    label += f"D10: {deciles[0]:.2f}\n" f"D50 (Med): {deciles[4]:.2f}\n" f"D90: {deciles[8]:.2f}"
+    axs[0].hist(wallclock, bins="auto", edgecolor="black", color="skyblue")
+    title = "Wallclock Distribution"
+    if percent != 100:
+        title += f" (top {percent}%)"
+    axs[0].set_title(title)
+    axs[0].set_xlabel("Wallclock (s)")
+    axs[0].set_ylabel("Frequency")
+    axs[0].text(0.60, 0.95, label, transform=axs[0].transAxes, horizontalalignment="left", verticalalignment="top")
+
+    # Plot memory
+    cutoff = 0
+    memory = [x["actuals"]["memory"] for x in items]
+    memory = [x for x in memory if x > cutoff]
+    deciles = np.percentile(memory, deciles_array)
+
+    label = f"{len(memory)} samples\n"
+    if cutoff != 0:
+        label += f"(memory > {cutoff} GB)\n"
+    label += f"D10: {deciles[0]:.2f}\n" f"D50 (Med): {deciles[4]:.2f}\n" f"D90: {deciles[8]:.2f}"
+    axs[1].hist(memory, bins="auto", edgecolor="black", color="salmon")
+    title = "Memory Distribution"
+    if percent != 100:
+        title += f" (top {percent}%)"
+    axs[1].set_title(title)
+    axs[1].set_xlabel("Memory (GB)")
+    axs[1].text(0.60, 0.95, label, transform=axs[1].transAxes, horizontalalignment="left", verticalalignment="top")
+
+    # 5. Clean up layout spacing and display
+    plt.tight_layout()
+    plt.show()
+
+
+def main():
+    items = get_dynamo_items(100000)
+    convert_elements_to_numeric_values(items)
+
+    plot_bins(items)
+    calculate_models(items)
+
+
+if __name__ == "__main__":
+    main()