gridfm-graphkit 0.0.4__tar.gz → 0.0.6__tar.gz

{gridfm_graphkit-0.0.4 → gridfm_graphkit-0.0.6}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: gridfm-graphkit
-Version: 0.0.4
+Version: 0.0.6
 Summary: Grid Foundation Model
 Author-email: Matteo Mazzonelli <matteo.mazzonelli1@ibm.com>, Alban Puech <apuech@seas.harvard.edu>, Tamara Govindasamy <tamara.govindasamy@ibm.com>, Mangaliso Mngomezulu <mngomezulum@ibm.com>, Etienne Vos <etienne.vos@ibm.com>, Celia Cintas <celia.cintas@ibm.com>, Jonas Weiss <jwe@zurich.ibm.com>
 Maintainer-email: Matteo Mazzonelli <matteo.mazzonelli1@ibm.com>
@@ -14,30 +14,39 @@ Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
 Requires-Python: <3.13,>=3.10
 Description-Content-Type: text/markdown
 License-File: LICENSE
-Requires-Dist: mlflow>=3.1.0
-Requires-Dist: nbformat>=5.10.4
-Requires-Dist: networkx>=3.4.2
-Requires-Dist: numpy>=2.2.6
-Requires-Dist: pandas>=2.3.0
-Requires-Dist: plotly>=6.1.2
-Requires-Dist: pyyaml>=6.0.2
-Requires-Dist: torch>=2.7.1
-Requires-Dist: torch-geometric>=2.6.1
-Requires-Dist: torchaudio>=2.7.1
-Requires-Dist: torchvision>=0.22.1
+Requires-Dist: torch>2.0
+Requires-Dist: torch-geometric
+Requires-Dist: mlflow
+Requires-Dist: nbformat
+Requires-Dist: networkx
+Requires-Dist: numpy
+Requires-Dist: pandas
+Requires-Dist: plotly
+Requires-Dist: pyyaml
 Requires-Dist: lightning
+Requires-Dist: seaborn
 Provides-Extra: dev
 Requires-Dist: mkdocs-material; extra == "dev"
 Requires-Dist: mkdocstrings[python]; extra == "dev"
-Requires-Dist: pre-commit>=4.2.0; extra == "dev"
-Requires-Dist: bandit>=1.8.5; extra == "dev"
+Requires-Dist: pre-commit; extra == "dev"
+Requires-Dist: bandit; extra == "dev"
 Requires-Dist: build; extra == "dev"
 Provides-Extra: test
 Requires-Dist: pytest; extra == "test"
 Requires-Dist: pytest-cov; extra == "test"
 Dynamic: license-file
 
-# gridfm-graphkit
+<p align="center">
+  <img src="https://raw.githubusercontent.com/gridfm/gridfm-graphkit/refs/heads/main/docs/figs/KIT.png" alt="GridFM logo" style="width: 40%; height: auto;"/>
+  <br/>
+</p>
+
+<p align="center" style="font-size: 25px;">
+    <b>gridfm-graphkit</b>
+</p>
+
+
+[![DOI](https://zenodo.org/badge/1007159095.svg)](https://doi.org/10.5281/zenodo.17016737)
 [![Docs](https://img.shields.io/badge/docs-available-brightgreen)](https://gridfm.github.io/gridfm-graphkit/)
 ![Coverage](https://img.shields.io/badge/coverage-83%25-yellowgreen)
 ![Python](https://img.shields.io/badge/python-3.10%20%E2%80%93%203.12-blue)
@@ -47,11 +56,6 @@ This library is brought to you by the GridFM team to train, finetune and interac
 
 ---
 
-<p align="center">
-  <img src="https://raw.githubusercontent.com/gridfm/gridfm-graphkit/refs/heads/main/docs/figs/pre_training.png" alt="GridFM logo"/>
-  <br/>
-</p>
-
 # Installation
 
 You can install `gridfm-graphkit` directly from PyPI:

{gridfm_graphkit-0.0.4 → gridfm_graphkit-0.0.6}/README.md

@@ -1,4 +1,14 @@
-# gridfm-graphkit
+<p align="center">
+  <img src="https://raw.githubusercontent.com/gridfm/gridfm-graphkit/refs/heads/main/docs/figs/KIT.png" alt="GridFM logo" style="width: 40%; height: auto;"/>
+  <br/>
+</p>
+
+<p align="center" style="font-size: 25px;">
+    <b>gridfm-graphkit</b>
+</p>
+
+
+[![DOI](https://zenodo.org/badge/1007159095.svg)](https://doi.org/10.5281/zenodo.17016737)
 [![Docs](https://img.shields.io/badge/docs-available-brightgreen)](https://gridfm.github.io/gridfm-graphkit/)
 ![Coverage](https://img.shields.io/badge/coverage-83%25-yellowgreen)
 ![Python](https://img.shields.io/badge/python-3.10%20%E2%80%93%203.12-blue)
@@ -8,11 +18,6 @@ This library is brought to you by the GridFM team to train, finetune and interac
 
 ---
 
-<p align="center">
-  <img src="https://raw.githubusercontent.com/gridfm/gridfm-graphkit/refs/heads/main/docs/figs/pre_training.png" alt="GridFM logo"/>
-  <br/>
-</p>
-
 # Installation
 
 You can install `gridfm-graphkit` directly from PyPI:

gridfm_graphkit-0.0.6/gridfm_graphkit/datasets/postprocessing.py

@@ -0,0 +1,83 @@
+import numpy as np
+from scipy.sparse import csr_matrix
+
+
+def compute_branch_currents_kA(Yf, Yt, V, Vf_base_kV, Vt_base_kV, sn_mva):
+    """
+    TODO docstrings
+    """
+
+    If_pu = Yf @ V  # From-end currents in per-unit (I_f = Y_f V)
+    If_kA = np.abs(If_pu) * sn_mva / (np.sqrt(3) * Vf_base_kV)  # Conversion to kA
+
+    # Construct to-end admittance matrix Yt:
+    # Yt[b, :] = y_tf_b * e_f + y_tt_b * e_t
+    It_pu = Yt @ V  # To-end currents in per-unit (I_t = Y_t V)
+    It_kA = np.abs(It_pu) * sn_mva / (np.sqrt(3) * Vt_base_kV)  # Conversion to kA
+
+    return If_kA, It_kA
+
+
+def compute_loading(If_kA, It_kA, Vf_base_kV, Vt_base_kV, rate_a):
+    """
+    Compute per-branch loading using current magnitudes and branch ratings.
+
+    Parameters:
+    - edge_index: np.ndarray of shape (n_edges, 2), each row is [from_bus, to_bus]
+    - If_kA: np.ndarray of from-side current magnitudes in kA
+    - It_kA: np.ndarray of to-side current magnitudes in kA
+    - base_kv: np.ndarray of shape (n_buses,), base voltage in kV per bus
+    - edge_attr: np.ndarray of shape (n_edges, >=5), edge features, column 4 = RATE_A
+
+    Returns:
+    - loading: np.ndarray of shape (n_edges,), max of from and to side loading
+    """
+
+    limitf = rate_a / (Vf_base_kV * np.sqrt(3))
+    limitt = rate_a / (Vt_base_kV * np.sqrt(3))
+
+    loadingf = If_kA / limitf
+    loadingt = It_kA / limitt
+
+    return np.maximum(loadingf, loadingt)
+
+
+def create_admittance_matrix(bus_params, edge_params, sn_mva=100):
+    """
+    TODO Docstrings
+
+    Parameters:
+    - bus_params: pandas df
+    - edge_params: pandas df
+
+    """
+
+    base_kv = bus_params["baseKV"].values
+
+    # Extract from-bus and to-bus indices for each branch
+
+    f = edge_params["from_bus"].values.astype(np.int32)
+    t = edge_params["to_bus"].values.astype(np.int32)
+
+    # Extract branch admittance coefficients
+    Yff = edge_params["Yff_r"].values + 1j * edge_params["Yff_i"].values
+    Yft = edge_params["Yft_r"].values + 1j * edge_params["Yft_i"].values
+    Ytf = edge_params["Ytf_r"].values + 1j * edge_params["Ytf_i"].values
+    Ytt = edge_params["Ytt_r"].values + 1j * edge_params["Ytt_i"].values
+
+    # Get base voltages for the from and to buses (for kA conversion)
+    Vf_base_kV = base_kv[f]
+    Vt_base_kV = base_kv[t]
+
+    nl = edge_params.shape[0]
+    nb = bus_params.shape[0]
+
+    # i = [0, 1, ..., nl-1, 0, 1, ..., nl-1], used for constructing Yf and Yt
+    i = np.hstack([np.arange(nl), np.arange(nl)])
+
+    # Construct from-end admittance matrix Yf using the linear combination:
+    # Yf[b, :] = y_ff_b * e_f + y_ft_b * e_t
+    Yf = csr_matrix((np.hstack([Yff, Yft]), (i, np.hstack([f, t]))), shape=(nl, nb))
+    Yt = csr_matrix((np.hstack([Ytf, Ytt]), (i, np.hstack([f, t]))), shape=(nl, nb))
+
+    return Yf, Yt, Vf_base_kV, Vt_base_kV

gridfm_graphkit-0.0.6/gridfm_graphkit/utils/utils.py

@@ -0,0 +1,42 @@
+def compute_cm_metrics(y_test, y_pred, model_name, label_plot):
+    """
+    Compute confusion matrix (TP,FP,TN,FN) for predicted overleads along with their respective rates and accuracy metric.
+
+    Parameters:
+    - y_pred: predicted overlads
+    - y_test: ground truth overloads
+    - prediction_dir:
+    - label_plot:
+    """
+
+    TP = (y_test & y_pred).sum()
+    FP = ((~y_test) & y_pred).sum()
+    TN = ((~y_test) & (~y_pred)).sum()
+    FN = (y_test & (~y_pred)).sum()
+
+    # accuracy
+    accuracy = (TP + TN) / (TP + FP + TN + FN)
+    print(f"Accuracy: {accuracy:.3f}")
+
+    TPR = TP / (TP + FN)
+    FPR = FP / (FP + TN)
+    TNR = TN / (TN + FP)
+    FNR = FN / (FN + TP)
+    # TODO change text to fit both overloadings and voltage violations
+    print("Confusion Matrix:")
+    print(f"TP: {TP}, FP: {FP}, TN: {TN}, FN: {FN}")
+    print(
+        f"GridFM\nTPR: {TPR:.3f} (percentage of overloadings correctly predicted)\nFPR: {FPR:.3f} (percentage of non-overloadings predicted as overloadings)\nTNR: {TNR:.2f}\nFNR: {FNR:.2f}",
+    )
+    with open(f"metrics_overloading_{model_name}.txt", "w") as f:
+        f.write(f"Accuracy: {accuracy:.3f}\n")
+        f.write("Confusion Matrix:\n")
+        f.write(f"TP: {TP}, FP: {FP}, TN: {TN}, FN: {FN}\n")
+        f.write(f"{label_plot} Metrics:\n")
+        f.write(f"TPR: {TPR:.5f} (percentage of overloadings correctly predicted)\n")
+        f.write(
+            f"FPR: {FPR:.5f} (percentage of non-overloadings predicted as overloadings)\n",
+        )
+        f.write(f"TNR: {TNR:.5f}\n")
+        f.write(f"FNR: {FNR:.5f}\n")
+    return TP, FP, TN, FN

gridfm_graphkit-0.0.6/gridfm_graphkit/utils/visualization.py

@@ -0,0 +1,513 @@
+from gridfm_graphkit.training.loss import PBELoss
+from gridfm_graphkit.datasets.globals import PQ, PV, REF
+
+import networkx as nx
+import matplotlib.pyplot as plt
+from matplotlib.colors import LogNorm
+from scipy.stats import pearsonr
+import seaborn as sns
+import numpy as np
+import copy
+
+
+def visualize_error(data_point, output, node_normalizer):
+    loss = PBELoss(visualization=True)
+
+    loss_dict = loss(
+        output,
+        data_point.y,
+        data_point.edge_index,
+        data_point.edge_attr,
+        data_point.mask,
+    )
+    active_loss = loss_dict["Nodal Active Power Loss in p.u."]
+    active_loss = active_loss.cpu() * node_normalizer.baseMVA
+
+    # Create a graph
+    G = nx.Graph()
+    edges = [
+        (u, v)
+        for u, v in zip(
+            data_point.edge_index[0].tolist(),
+            data_point.edge_index[1].tolist(),
+        )
+        if u != v
+    ]
+    G.add_edges_from(edges)
+
+    # Assign labels based on node type
+    node_shapes = {"REF": "s", "PV": "H", "PQ": "o"}
+    num_nodes = data_point.x.shape[0]
+    mask_PQ = data_point.x[:, PQ] == 1
+    mask_PV = data_point.x[:, PV] == 1
+    mask_REF = data_point.x[:, REF] == 1
+    node_labels = {}
+    for i in range(num_nodes):
+        if mask_REF[i]:
+            node_labels[i] = "REF"
+        elif mask_PV[i]:
+            node_labels[i] = "PV"
+        elif mask_PQ[i]:
+            node_labels[i] = "PQ"
+
+    # Set node positions
+    pos = nx.spring_layout(G, seed=42)
+
+    # Define colormap
+    cmap = plt.cm.viridis
+    vmin = min(active_loss)
+    vmax = max(active_loss)
+    norm = plt.Normalize(vmin=vmin, vmax=vmax)
+
+    # Create a figure and axis
+    fig, ax = plt.subplots(figsize=(13, 7))
+
+    # Draw nodes with heatmap coloring
+    for node_type, shape in node_shapes.items():
+        nodes = [i for i in node_labels if node_labels[i] == node_type]
+        nx.draw_networkx_nodes(
+            G,
+            pos,
+            nodelist=nodes,
+            node_color=[active_loss[i] for i in nodes],
+            cmap=cmap,
+            node_size=800,
+            ax=ax,
+            vmin=vmin,
+            vmax=vmax,
+            node_shape=shape,
+        )
+
+    # Draw edges
+    nx.draw_networkx_edges(G, pos, edge_color="gray", alpha=0.5, ax=ax)
+
+    # Draw labels (node types)
+    nx.draw_networkx_labels(
+        G,
+        pos,
+        labels=node_labels,
+        font_size=10,
+        font_color="white",
+        font_weight="bold",
+        ax=ax,
+    )
+
+    # Add colorbar
+    cbar = plt.colorbar(plt.cm.ScalarMappable(cmap=cmap, norm=norm), ax=ax)
+    cbar.set_label("Active Power Residuals (MW)", fontsize=12)
+    cbar.ax.tick_params(labelsize=12)
+
+    for spine in ax.spines.values():
+        spine.set_linewidth(2)  # Adjust thickness here (e.g., 2 or any value)
+
+    # Show plot
+    plt.title("Nodal Active Power Residuals", fontsize=14, fontweight="bold")
+    plt.show()
+
+
+def visualize_quantity_heatmap(
+    data_point,
+    output,
+    quantity,
+    quantity_name,
+    unit,
+    node_normalizer,
+):
+    """
+    Visualizes a heatmap of a specified quantity (VM, PD, QD, PG, QG, VA) for a given dataset and model.
+
+    Parameters:
+        data_point: Power grid data.
+        model: The trained model used for inference.
+        quantity: The quantity to visualize (e.g., VM, PD, QD, PG, QG, VA).
+    """
+    data_point = copy.deepcopy(data_point)
+    output = copy.deepcopy(output)
+    mask_PQ = data_point.x[:, PQ] == 1
+    mask_PV = data_point.x[:, PV] == 1
+    mask_REF = data_point.x[:, REF] == 1
+
+    output = node_normalizer.inverse_transform(output)
+    denormalized_gt = node_normalizer.inverse_transform(data_point.y)
+
+    gt_values = denormalized_gt[:, quantity]
+    predicted_values = output[:, quantity]
+    predicted_values[~data_point.mask[:, quantity]] = denormalized_gt[
+        ~data_point.mask[:, quantity],
+        quantity,
+    ]
+
+    num_nodes = data_point.x.shape[0]
+
+    node_shapes = {"REF": "s", "PV": "H", "PQ": "o"}
+
+    # Create graph
+    G = nx.Graph()
+    edges = [
+        (u, v)
+        for u, v in zip(
+            data_point.edge_index[0].tolist(),
+            data_point.edge_index[1].tolist(),
+        )
+        if u != v
+    ]
+    G.add_edges_from(edges)
+
+    node_labels = {}
+    for i in range(num_nodes):
+        if mask_REF[i]:
+            node_labels[i] = "REF"
+        elif mask_PV[i]:
+            node_labels[i] = "PV"
+        elif mask_PQ[i]:
+            node_labels[i] = "PQ"
+
+    pos = nx.spring_layout(G, seed=42)
+    cmap = plt.cm.viridis
+    vmin = min(predicted_values)
+    vmax = max(predicted_values)
+    norm = plt.Normalize(vmin=vmin, vmax=vmax)
+
+    masked_node_indices = np.where(data_point.mask[:, quantity].cpu())[0]
+
+    # Create subplots for side-by-side layout (3 plots)
+    fig, axes = plt.subplots(1, 3, figsize=(22, 8))
+
+    # First plot (ground truth values)
+    ax = axes[0]
+    for node_type, shape in node_shapes.items():
+        nodes = [i for i in node_labels if node_labels[i] == node_type]
+        node_size = 390 if node_type == "REF" else 600
+        nx.draw_networkx_nodes(
+            G,
+            pos,
+            nodelist=nodes,
+            node_color=[gt_values[i] for i in nodes],
+            cmap=cmap,
+            node_size=node_size,
+            ax=ax,
+            vmin=vmin,
+            vmax=vmax,
+            node_shape=shape,
+        )
+
+    nx.draw_networkx_edges(G, pos, edge_color="gray", alpha=0.5, ax=ax, width=2)
+    nx.draw_networkx_labels(
+        G,
+        pos,
+        labels=node_labels,
+        font_size=10,
+        font_color="white",
+        font_weight="bold",
+        ax=ax,
+    )
+    ax.set_title(f"Ground truth {quantity_name}", fontsize=14, fontweight="bold")
+
+    for spine in ax.spines.values():
+        spine.set_linewidth(2)  # Adjust thickness
+
+    # Second plot (with masked nodes in gray)
+    ax = axes[1]
+    for node_type, shape in node_shapes.items():
+        nodes = [i for i in node_labels if node_labels[i] == node_type]
+        node_size = 390 if node_type == "REF" else 600
+        nx.draw_networkx_nodes(
+            G,
+            pos,
+            nodelist=nodes,
+            node_color=[gt_values[i] for i in nodes],
+            cmap=cmap,
+            node_size=node_size,
+            ax=ax,
+            vmin=vmin,
+            vmax=vmax,
+            node_shape=shape,
+        )
+
+    nx.draw_networkx_nodes(
+        G,
+        pos,
+        nodelist=masked_node_indices,
+        node_color="#D3D3D3",
+        node_size=750,
+        ax=ax,
+    )
+    nx.draw_networkx_edges(G, pos, edge_color="gray", alpha=0.5, ax=ax, width=2)
+    nx.draw_networkx_labels(
+        G,
+        pos,
+        labels=node_labels,
+        font_size=10,
+        font_color="white",
+        font_weight="bold",
+        ax=ax,
+    )
+    ax.set_title(f"Masked {quantity_name}", fontsize=14, fontweight="bold")
+
+    for spine in ax.spines.values():
+        spine.set_linewidth(2)  # Adjust thickness
+
+    # Third plot (predicted values without masking)
+    ax = axes[2]
+    for node_type, shape in node_shapes.items():
+        nodes = [i for i in node_labels if node_labels[i] == node_type]
+        node_size = 390 if node_type == "REF" else 600
+        nx.draw_networkx_nodes(
+            G,
+            pos,
+            nodelist=nodes,
+            node_color=[predicted_values[i] for i in nodes],
+            cmap=cmap,
+            node_size=node_size,
+            ax=ax,
+            vmin=vmin,
+            vmax=vmax,
+            node_shape=shape,
+        )
+
+    nx.draw_networkx_edges(G, pos, edge_color="gray", alpha=0.5, ax=ax, width=2)
+    nx.draw_networkx_labels(
+        G,
+        pos,
+        labels=node_labels,
+        font_size=10,
+        font_color="white",
+        font_weight="bold",
+        ax=ax,
+    )
+    ax.set_title(f"Reconstructed {quantity_name}", fontsize=14, fontweight="bold")
+
+    for spine in ax.spines.values():
+        spine.set_linewidth(2)  # Adjust thickness
+
+    # Colorbar placement
+    cbar_ax = fig.add_axes([0.93, 0.1, 0.02, 0.8])
+    cbar = plt.colorbar(plt.cm.ScalarMappable(cmap=cmap, norm=norm), cax=cbar_ax)
+    cbar.set_label(f"{quantity_name} ({unit})", fontsize=12)
+    cbar.ax.tick_params(labelsize=12)
+
+    plt.subplots_adjust(right=0.9)
+    plt.show()
+
+
+def plot_mass_correlation_density(
+    true_vals,
+    gfm_vals,
+    model_name,
+    label_plot,
+    x_max=2,
+    y_max=3,
+):
+    """
+    TODO docstring
+
+    """
+    # TODO check if these parameters need to be passed by func or default behavior
+    vmin = 1
+    x_min = 0
+    y_min = 0
+    bin_width = 0.01  # consistent bin width for both plots
+
+    # Generate consistent bins
+    x_bins = np.arange(x_min, x_max + bin_width, bin_width)
+    y_bins = np.arange(y_min, y_max + bin_width, bin_width)
+
+    # estimate vmax on mean count of elements across bins
+    counts, _, _ = np.histogram2d(true_vals, gfm_vals, bins=[x_bins, y_bins])
+
+    counts[counts == 0] = np.nan
+    means = np.nanmean(counts)
+    std = np.nanstd(counts)
+    vmax = means + 3 * std
+
+    # Pearson correlations
+    corr_gfm, _ = pearsonr(true_vals, gfm_vals)
+
+    # Create figure with shared x-axis
+    fig, ax1 = plt.subplots(figsize=(8, 6), dpi=400)
+
+    # --- GridFM Mass Correlation ---
+    h1 = ax1.hist2d(
+        true_vals,
+        gfm_vals,
+        bins=[x_bins, y_bins],
+        norm=LogNorm(vmin=vmin, vmax=vmax),
+        cmap="inferno",
+    )
+    ax1.axvline(1, color="black", linestyle="--", linewidth=2.0)
+    ax1.axhline(1, color="black", linestyle="--", linewidth=2.0)
+    ax1.plot([0, 5], [0, 5], "k--", linewidth=0.5)
+    ax1.set_xlabel("True Loadings", fontsize=12)
+    ax1.set_ylabel("Predicted Loadings", fontsize=12)
+    ax1.set_title(label_plot, fontsize=14)
+    ax1.text(
+        x_max - 1.5,
+        0.93,
+        f"r = {corr_gfm:.5f}",
+        transform=ax1.transAxes,
+        fontsize=13,
+        weight="bold",
+    )
+
+    # Colorbar
+    cbar = fig.colorbar(h1[3], ax=ax1, pad=0.02)
+    cbar.set_label("Number of samples", fontsize=10)
+
+    # Style adjustments
+    ax1.set_xlim(x_min, x_max)
+    ax1.set_ylim(y_min, y_max)
+    ax1.grid(True, linewidth=0.3)
+    ax1.tick_params(axis="both", labelsize=10)
+
+    plt.tight_layout()
+    plt.savefig(f"mass_correlation_density_{model_name}.png", bbox_inches="tight")
+    plt.show()
+
+
+def plot_cm(TN, FP, FN, TP, model_name, label_plot):
+    """
+    TODO docstring
+    """
+    cm = np.array([[TN, FP], [FN, TP]])
+
+    cm_labels = ["Non-overload", "Overload"]
+
+    fig_cm, ax_cm = plt.subplots(figsize=(6, 6))
+
+    sns.heatmap(
+        cm,
+        annot=True,
+        fmt="d",
+        cbar=False,
+        square=True,
+        linewidths=0.5,
+        cmap="Blues",
+        xticklabels=cm_labels,
+        yticklabels=cm_labels,
+        ax=ax_cm,
+        annot_kws={"size": 14},
+    )
+
+    ax_cm.set_xlabel("Predicted", fontsize=12)
+    ax_cm.set_ylabel("True", fontsize=12)
+    ax_cm.set_title(f"Confusion Matrix {label_plot}", fontsize=14)
+    ax_cm.tick_params(axis="both", labelsize=12)
+
+    plt.tight_layout()
+    plt.savefig(f"confusion_matrix_overload_{model_name}.png", bbox_inches="tight")
+    plt.show()
+
+
+def plot_loading_predictions(
+    loadings_pred,
+    loadings_dc,
+    loadings_gt,
+    prediction_dir,
+    label_plot,
+):
+    """
+    TODO docstrings
+    """
+    plt.hist(
+        loadings_pred,
+        alpha=0.5,
+        label=label_plot,
+        density=True,
+        bins=100,
+    )
+    plt.hist(loadings_dc, alpha=0.5, label="DC Solver", density=True, bins=100)
+    plt.hist(loadings_gt, alpha=0.5, label="Ground truth", density=True, bins=100)
+
+    plt.xlabel("Loading Values")
+    plt.ylabel("Density")
+    plt.yscale("log")
+    plt.legend()
+
+    plt.savefig(f"distribution_loading_predictions_{prediction_dir}.png")
+    plt.show()
+
+
+def plot_mass_correlation_density_voltage(
+    pf_node,
+    prediction_dir,
+    label_plot,
+    x_min=0.85,
+    y_min=0.85,
+    x_max=1.15,
+    y_max=1.15,
+):
+    """
+    TODO docstrings
+    TODO refactor if we pass by parameters a few more plot deets we can use plot_mass_correlation_density for both
+
+    """
+    # Get the global min and max for color scaling (avoid log(0) by setting min to at least 1)
+    vmin = 1
+    bin_width = 0.001  # consistent bin width for both plots
+
+    # Generate consistent bins
+    x_bins = np.arange(x_min, x_max + bin_width, bin_width)
+    y_bins = np.arange(y_min, y_max + bin_width, bin_width)
+
+    # estimate vmax on mean count of elements across bins
+    counts, _, _ = np.histogram2d(
+        pf_node["Vm"],
+        pf_node["Vm_pred_corrected"],
+        bins=[x_bins, y_bins],
+    )
+
+    counts[counts == 0] = np.nan
+    means = np.nanmean(counts)
+    std = np.nanstd(counts)
+    vmax = means + 3 * std
+
+    # Pearson correlations
+    corr_vm, _ = pearsonr(pf_node["Vm"], pf_node["Vm_pred_corrected"])
+
+    # Create figure with shared x-axis
+    fig, ax1 = plt.subplots(figsize=(8, 6), dpi=400)
+
+    # --- GridFM Mass Correlation ---
+    h1 = ax1.hist2d(
+        pf_node["Vm"],
+        pf_node["Vm_pred_corrected"],
+        bins=[x_bins, y_bins],
+        norm=LogNorm(vmin=vmin, vmax=vmax),
+        cmap="inferno",
+    )
+    ax1.axvline(x_min + 0.05, color="black", linestyle="--", linewidth=2.0)
+    ax1.axhline(y_min + 0.05, color="black", linestyle="--", linewidth=2.0)
+    ax1.axvline(x_max - 0.05, color="black", linestyle="--", linewidth=2.0)
+    ax1.axhline(y_max - 0.05, color="black", linestyle="--", linewidth=2.0)
+
+    ax1.plot([0, 5], [0, 5], "k--", linewidth=0.5)
+    ax1.set_xlabel("True Voltage Magnitude", fontsize=12)
+    ax1.set_ylabel("Predicted Voltage magnitude", fontsize=12)
+    ax1.set_title(label_plot, fontsize=14)
+    ax1.text(
+        0.5,
+        0.95,
+        f"r = {corr_vm:.5f}",
+        transform=ax1.transAxes,
+        fontsize=13,
+        weight="bold",
+        ha="center",
+        va="top",
+    )
+
+    # Colorbar
+    cbar = fig.colorbar(h1[3], ax=ax1, pad=0.02)
+    cbar.set_label("Number of samples", fontsize=10)
+
+    # Style adjustments
+    ax1.set_xlim(x_min, x_max)
+    ax1.set_ylim(y_min, y_max)
+    ax1.grid(True, linewidth=0.3)
+    ax1.tick_params(axis="both", labelsize=10)
+
+    plt.tight_layout()
+    plt.savefig(
+        f"mass_correlation_density_voltage_{prediction_dir}.png",
+        bbox_inches="tight",
+    )
+    plt.show()

{gridfm_graphkit-0.0.4 → gridfm_graphkit-0.0.6}/gridfm_graphkit.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: gridfm-graphkit
-Version: 0.0.4
+Version: 0.0.6
 Summary: Grid Foundation Model
 Author-email: Matteo Mazzonelli <matteo.mazzonelli1@ibm.com>, Alban Puech <apuech@seas.harvard.edu>, Tamara Govindasamy <tamara.govindasamy@ibm.com>, Mangaliso Mngomezulu <mngomezulum@ibm.com>, Etienne Vos <etienne.vos@ibm.com>, Celia Cintas <celia.cintas@ibm.com>, Jonas Weiss <jwe@zurich.ibm.com>
 Maintainer-email: Matteo Mazzonelli <matteo.mazzonelli1@ibm.com>
@@ -14,30 +14,39 @@ Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
 Requires-Python: <3.13,>=3.10
 Description-Content-Type: text/markdown
 License-File: LICENSE
-Requires-Dist: mlflow>=3.1.0
-Requires-Dist: nbformat>=5.10.4
-Requires-Dist: networkx>=3.4.2
-Requires-Dist: numpy>=2.2.6
-Requires-Dist: pandas>=2.3.0
-Requires-Dist: plotly>=6.1.2
-Requires-Dist: pyyaml>=6.0.2
-Requires-Dist: torch>=2.7.1
-Requires-Dist: torch-geometric>=2.6.1
-Requires-Dist: torchaudio>=2.7.1
-Requires-Dist: torchvision>=0.22.1
+Requires-Dist: torch>2.0
+Requires-Dist: torch-geometric
+Requires-Dist: mlflow
+Requires-Dist: nbformat
+Requires-Dist: networkx
+Requires-Dist: numpy
+Requires-Dist: pandas
+Requires-Dist: plotly
+Requires-Dist: pyyaml
 Requires-Dist: lightning
+Requires-Dist: seaborn
 Provides-Extra: dev
 Requires-Dist: mkdocs-material; extra == "dev"
 Requires-Dist: mkdocstrings[python]; extra == "dev"
-Requires-Dist: pre-commit>=4.2.0; extra == "dev"
-Requires-Dist: bandit>=1.8.5; extra == "dev"
+Requires-Dist: pre-commit; extra == "dev"
+Requires-Dist: bandit; extra == "dev"
 Requires-Dist: build; extra == "dev"
 Provides-Extra: test
 Requires-Dist: pytest; extra == "test"
 Requires-Dist: pytest-cov; extra == "test"
 Dynamic: license-file
 
-# gridfm-graphkit
+<p align="center">
+  <img src="https://raw.githubusercontent.com/gridfm/gridfm-graphkit/refs/heads/main/docs/figs/KIT.png" alt="GridFM logo" style="width: 40%; height: auto;"/>
+  <br/>
+</p>
+
+<p align="center" style="font-size: 25px;">
+    <b>gridfm-graphkit</b>
+</p>
+
+
+[![DOI](https://zenodo.org/badge/1007159095.svg)](https://doi.org/10.5281/zenodo.17016737)
 [![Docs](https://img.shields.io/badge/docs-available-brightgreen)](https://gridfm.github.io/gridfm-graphkit/)
 ![Coverage](https://img.shields.io/badge/coverage-83%25-yellowgreen)
 ![Python](https://img.shields.io/badge/python-3.10%20%E2%80%93%203.12-blue)
@@ -47,11 +56,6 @@ This library is brought to you by the GridFM team to train, finetune and interac
 
 ---
 
-<p align="center">
-  <img src="https://raw.githubusercontent.com/gridfm/gridfm-graphkit/refs/heads/main/docs/figs/pre_training.png" alt="GridFM logo"/>
-  <br/>
-</p>
-
 # Installation
 
 You can install `gridfm-graphkit` directly from PyPI:

{gridfm_graphkit-0.0.4 → gridfm_graphkit-0.0.6}/gridfm_graphkit.egg-info/SOURCES.txt

@@ -13,6 +13,7 @@ gridfm_graphkit.egg-info/top_level.txt
 gridfm_graphkit/datasets/__init__.py
 gridfm_graphkit/datasets/globals.py
 gridfm_graphkit/datasets/normalizers.py
+gridfm_graphkit/datasets/postprocessing.py
 gridfm_graphkit/datasets/powergrid_datamodule.py
 gridfm_graphkit/datasets/powergrid_dataset.py
 gridfm_graphkit/datasets/transforms.py
@@ -29,6 +30,7 @@ gridfm_graphkit/training/__init__.py
 gridfm_graphkit/training/callbacks.py
 gridfm_graphkit/training/loss.py
 gridfm_graphkit/utils/__init__.py
+gridfm_graphkit/utils/utils.py
 gridfm_graphkit/utils/visualization.py
 tests/test_data_module.py
 tests/test_full_pipeline.py

gridfm_graphkit-0.0.6/gridfm_graphkit.egg-info/requires.txt

@@ -0,0 +1,22 @@
+torch>2.0
+torch-geometric
+mlflow
+nbformat
+networkx
+numpy
+pandas
+plotly
+pyyaml
+lightning
+seaborn
+
+[dev]
+mkdocs-material
+mkdocstrings[python]
+pre-commit
+bandit
+build
+
+[test]
+pytest
+pytest-cov

{gridfm_graphkit-0.0.4 → gridfm_graphkit-0.0.6}/pyproject.toml

@@ -9,7 +9,7 @@ namespaces = false
 [project]
 name = "gridfm-graphkit"
 description = "Grid Foundation Model"
-version = "0.0.4"
+version = "0.0.6"
 readme = "README.md"
 license = "Apache-2.0"
 requires-python = ">=3.10,<3.13"
@@ -40,26 +40,25 @@ classifiers = [
 
 
 dependencies = [
-    "mlflow>=3.1.0",
-    "nbformat>=5.10.4",
-    "networkx>=3.4.2",
-    "numpy>=2.2.6",
-    "pandas>=2.3.0",
-    "plotly>=6.1.2",
-    "pyyaml>=6.0.2",
-    "torch>=2.7.1",
-    "torch-geometric>=2.6.1",
-    "torchaudio>=2.7.1",
-    "torchvision>=0.22.1",
+    "torch>2.0",
+    "torch-geometric",
+    "mlflow",
+    "nbformat",
+    "networkx",
+    "numpy",
+    "pandas",
+    "plotly",
+    "pyyaml",
     "lightning",
+    "seaborn",
 ]
 
 [project.optional-dependencies]
 dev = [
   "mkdocs-material",
   "mkdocstrings[python]",
-  "pre-commit>=4.2.0",
-  "bandit>=1.8.5",
+  "pre-commit",
+  "bandit",
   "build"
 ]
 

gridfm_graphkit-0.0.4/gridfm_graphkit/utils/visualization.py

@@ -1,99 +0,0 @@
-import networkx as nx
-from gridfm_graphkit.training.loss import PBELoss
-from gridfm_graphkit.datasets.globals import PQ, PV, REF
-import matplotlib.pyplot as plt
-
-
-def visualize_error(data_point, output, node_normalizer):
-    loss = PBELoss(visualization=True)
-
-    loss_dict = loss(
-        output,
-        data_point.y,
-        data_point.edge_index,
-        data_point.edge_attr,
-        data_point.mask,
-    )
-    active_loss = loss_dict["Nodal Active Power Loss in p.u."]
-    active_loss = active_loss.cpu() * node_normalizer.baseMVA
-
-    # Create a graph
-    G = nx.Graph()
-    edges = [
-        (u, v)
-        for u, v in zip(
-            data_point.edge_index[0].tolist(),
-            data_point.edge_index[1].tolist(),
-        )
-        if u != v
-    ]
-    G.add_edges_from(edges)
-
-    # Assign labels based on node type
-    node_shapes = {"REF": "s", "PV": "H", "PQ": "o"}
-    num_nodes = data_point.x.shape[0]
-    mask_PQ = data_point.x[:, PQ] == 1
-    mask_PV = data_point.x[:, PV] == 1
-    mask_REF = data_point.x[:, REF] == 1
-    node_labels = {}
-    for i in range(num_nodes):
-        if mask_REF[i]:
-            node_labels[i] = "REF"
-        elif mask_PV[i]:
-            node_labels[i] = "PV"
-        elif mask_PQ[i]:
-            node_labels[i] = "PQ"
-
-    # Set node positions
-    pos = nx.spring_layout(G, seed=42)
-
-    # Define colormap
-    cmap = plt.cm.viridis
-    vmin = min(active_loss)
-    vmax = max(active_loss)
-    norm = plt.Normalize(vmin=vmin, vmax=vmax)
-
-    # Create a figure and axis
-    fig, ax = plt.subplots(figsize=(13, 7))
-
-    # Draw nodes with heatmap coloring
-    for node_type, shape in node_shapes.items():
-        nodes = [i for i in node_labels if node_labels[i] == node_type]
-        nx.draw_networkx_nodes(
-            G,
-            pos,
-            nodelist=nodes,
-            node_color=[active_loss[i] for i in nodes],
-            cmap=cmap,
-            node_size=800,
-            ax=ax,
-            vmin=vmin,
-            vmax=vmax,
-            node_shape=shape,
-        )
-
-    # Draw edges
-    nx.draw_networkx_edges(G, pos, edge_color="gray", alpha=0.5, ax=ax)
-
-    # Draw labels (node types)
-    nx.draw_networkx_labels(
-        G,
-        pos,
-        labels=node_labels,
-        font_size=10,
-        font_color="white",
-        font_weight="bold",
-        ax=ax,
-    )
-
-    # Add colorbar
-    cbar = plt.colorbar(plt.cm.ScalarMappable(cmap=cmap, norm=norm), ax=ax)
-    cbar.set_label("Active Power Residuals (MW)", fontsize=12)
-    cbar.ax.tick_params(labelsize=12)
-
-    for spine in ax.spines.values():
-        spine.set_linewidth(2)  # Adjust thickness here (e.g., 2 or any value)
-
-    # Show plot
-    plt.title("Nodal Active Power Residuals", fontsize=14, fontweight="bold")
-    plt.show()

gridfm_graphkit-0.0.4/gridfm_graphkit.egg-info/requires.txt

@@ -1,23 +0,0 @@
-mlflow>=3.1.0
-nbformat>=5.10.4
-networkx>=3.4.2
-numpy>=2.2.6
-pandas>=2.3.0
-plotly>=6.1.2
-pyyaml>=6.0.2
-torch>=2.7.1
-torch-geometric>=2.6.1
-torchaudio>=2.7.1
-torchvision>=0.22.1
-lightning
-
-[dev]
-mkdocs-material
-mkdocstrings[python]
-pre-commit>=4.2.0
-bandit>=1.8.5
-build
-
-[test]
-pytest
-pytest-cov