plot_optknock.py

"""
AI was used to generate plotting and summary code
"""

import matplotlib.pyplot as plt
import numpy as np
import pandas as pd


def _sel(df, **kwargs):
    mask = pd.Series(True, index=df.index)
    for col, val in kwargs.items():
        if val is not None:
            mask &= df[col] == val
    return df[mask]


def _heatmap(pivot, title, cbar_label, cmap="coolwarm", fmt=".3g"):
    fig = plt.figure()
    k_vals = pivot.index.tolist()
    t_vals = pivot.columns.tolist()
    data = pivot.values.astype(float)
    log_data = np.log(np.clip(data, 1e-12, None))

    im = plt.imshow(log_data, aspect="auto", cmap=cmap, origin="lower")
    plt.colorbar(im, label=f"log {cbar_label}")
    plt.xticks(range(len(t_vals)), [str(t) for t in t_vals])
    plt.yticks(range(len(k_vals)), [str(k) for k in k_vals])
    plt.xlabel("Fermentation time  t  [h]")
    plt.ylabel("Knockout budget  k")
    plt.title(title)

    vmin, vmax = np.nanmin(log_data), np.nanmax(log_data)
    mid = (vmin + vmax) / 2
    for i in range(len(k_vals)):
        for j in range(len(t_vals)):
            val = data[i, j]
            if np.isfinite(val):
                color = "white" if log_data[i, j] < mid else "black"
                plt.text(
                    j,
                    i,
                    f"{val:{fmt}}",
                    ha="center",
                    va="center",
                    fontsize=8,
                    color=color,
                )
    return fig


def plot_production_heatmap(df, target, method, objective=None):
    if method == "greedy":
        sub = _sel(df, target=target, method="greedy")
        prod_col = "production"
        label = "greedy"
    else:
        sub = _sel(df, target=target, method="optknock", objective=objective)
        prod_col = "milp_production"
        label = f"optknock ({objective})"

    pivot = sub.pivot_table(
        index="k", columns="fermentation_time", values=prod_col, aggfunc="first"
    )
    return _heatmap(
        pivot,
        title=f"{target.capitalize()}  —  production  [{label}]",
        cbar_label="P(t)  [mmol/L]",
    )


def plot_solve_time_heatmap(df, target, objective):
    sub = _sel(df, target=target, method="optknock", objective=objective)
    pivot = sub.pivot_table(
        index="k", columns="fermentation_time", values="solve_time", aggfunc="first"
    )
    return _heatmap(
        pivot,
        title=f"{target.capitalize()}  —  solve time  [optknock ({objective})]",
        cbar_label="solve time  [s]",
        cmap="coolwarm",
    )


def plot_error_heatmap(df, target, objective):
    sub = _sel(df, target=target, method="optknock", objective=objective).copy()
    if "error" not in sub.columns:
        sub["error"] = (sub["milp_production"] - sub["production"]).abs()

    pivot = sub.pivot_table(
        index="k", columns="fermentation_time", values="error", aggfunc="first"
    )
    return _heatmap(
        pivot,
        title=f"{target.capitalize()}  —  MILP vs actual error  [optknock ({objective})]",
        cbar_label="|MILP − actual|  [mmol/L]",
        cmap="coolwarm",
    )


def plot_improvement_heatmap(df, target, objective):

    greedy_prod = _sel(df, target=target, method="greedy").set_index(
        ["k", "fermentation_time"]
    )["production"]
    ok_prod = _sel(df, target=target, method="optknock", objective=objective).set_index(
        ["k", "fermentation_time"]
    )["milp_production"]

    ratio = (ok_prod / greedy_prod).reset_index()
    ratio.columns = ["k", "fermentation_time", "ratio"]

    pivot = ratio.pivot(index="k", columns="fermentation_time", values="ratio")
    return _heatmap(
        pivot,
        title=f"{target.capitalize()}  —  OptKnock / greedy production  [({objective})]",
        cbar_label="production ratio",
        cmap="coolwarm",
    )


def summarize(df):
    stats = {}

    def _agg(series):
        return pd.DataFrame({"mean": series.mean(), "median": series.median()})

    def _print_block(title, grouped):
        print(title)
        print(grouped.to_string())

    # Use milp_production for optknock, post-FBA production for greedy
    df2 = df.copy()
    df2["display_production"] = np.where(
        df2["method"] == "optknock", df2["milp_production"], df2["production"]
    )
    prod = (
        df2.dropna(subset=["display_production"])
        .groupby(["method", "objective"])["display_production"]
        .agg(["mean", "median"])
    )
    stats["production"] = prod

    greedy = df2[df2["method"] == "greedy"].set_index(
        ["target", "k", "fermentation_time"]
    )["production"]
    ok = df2[df2["method"] == "optknock"].copy()
    ok = ok.join(
        greedy.rename("greedy_production"), on=["target", "k", "fermentation_time"]
    )
    ok["pct_improvement"] = (
        (ok["milp_production"] - ok["greedy_production"])
        / (ok["greedy_production"].replace(0, np.nan))
        * 100
    )
    improvement = ok.groupby("objective")["pct_improvement"].agg(["mean", "median"])
    stats["improvement_over_greedy"] = improvement

    ok_err = df[df["method"] == "optknock"].dropna(subset=["error"])
    error = ok_err.groupby("objective")["error"].agg(["mean", "median"])
    stats["error"] = error

    ok_pct = ok_err[ok_err["production"] > 1.0].copy()
    ok_pct["pct_error"] = (ok_pct["error"] / ok_pct["production"]) * 100
    pct_error = ok_pct.groupby("objective")["pct_error"].agg(["mean", "median"])
    stats["pct_error"] = pct_error

    solve_time = (
        df[df["method"] == "optknock"]
        .groupby("objective")["solve_time"]
        .agg(["mean", "median"])
    )
    stats["solve_time"] = solve_time

    sep = "-" * 60
    print(sep)
    _print_block("PRODUCTION  (by method / objective)", prod)
    print(sep)
    _print_block("% IMPROVEMENT OVER GREEDY", improvement)
    print(sep)
    _print_block("MILP ERROR  |milp_production - production|", error)
    print(sep)
    _print_block(f"MILP % ERROR  (production > 1 mmol/gDW, n={len(ok_pct)})", pct_error)
    print(sep)
    _print_block("SOLVE TIME  [s]", solve_time)
    print(sep)

    return stats