diff --git a/docs/tutorials/contingency_analysis.md b/docs/tutorials/contingency_analysis.md
deleted file mode 100644
index f894230..0000000
--- a/docs/tutorials/contingency_analysis.md
+++ /dev/null
@@ -1,325 +0,0 @@
-👉 [Link](https://colab.research.google.com/github/gridfm/gridfm-graphkit/blob/main/examples/notebooks/Tutorial_contingency_analisys.ipynb) to the tutorial on Google Colab
-
----
-Contingency analysis is a critical process in power system operations used to assess the impact of potential failures (e.g., line outages) on grid stability and reliability. It helps operators prepare for unexpected events by simulating scenarios such as N-1 or N-2 contingencies, where one or more components are removed from service. This analysis ensures that the grid can continue to operate within safe limits even under stressed conditions.
-
----
-
-## Dataset Generation and Model Evaluation
-
-The dataset used in this study originates from the Texas transmission grid, which includes approximately 2,000 nodes. Using the contingency mode of the `gridfm-datakit`, we simulated N-2 contingencies by removing up to two transmission lines at a time. For each scenario, we first solved the optimal power flow (OPF) problem to determine the generation dispatch. Then, we applied the contingency by removing lines and re-solved the power flow to observe the resulting grid state.
-
-This process generated around 100,000 unique scenarios. Our model, **GridFMv0.1**, was fine-tuned on this dataset to predict power flow outcomes. For demonstration purposes, we selected a subsample of 10 scenarios. The `gridfm-datakit` also computed DC power flow results, enabling a comparison between GridFMv0.1 predictions and traditional DC power flow estimates, specifically in terms of line loading accuracy.
-
-All predictions are benchmarked against the ground truth obtained from AC power flow simulations. Additionally, we analyze bus voltage violations, which GridFM can predict but are not captured by the DC solver, highlighting GridFM’s enhanced capabilities in modeling grid behavior.
-
-```python
-import sys
-
-if "google.colab" in sys.modules:
- try:
- !git clone https://github.com/gridfm/gridfm-graphkit.git
- %cd /content/gridfm-graphkit
- !pip install .
- %cd examples/notebooks/
- except Exception as e:
- print(f"Failed to start Google Collab setup, due to {e}")
-```
-
-```python
-from gridfm_graphkit.datasets.postprocessing import (
- compute_branch_currents_kA,
- compute_loading,
-)
-from gridfm_graphkit.datasets.postprocessing import create_admittance_matrix
-from gridfm_graphkit.utils.utils import compute_cm_metrics
-from gridfm_graphkit.utils.visualization import (
- plot_mass_correlation_density,
- plot_cm,
- plot_loading_predictions,
- plot_mass_correlation_density_voltage,
-)
-
-import os
-from tqdm import tqdm
-import matplotlib.pyplot as plt
-from sklearn.metrics import f1_score
-import numpy as np
-import pandas as pd
-```
-
-## Load Data
-
-We load both the ground truth and predicted values of the power flow solution. The predictions are generated using the `gridfm-graphkit` CLI:
-
-```bash
-gridfm-graphkit predict ...
-```
-
-We then merge the datasets using `scenario` and `bus` as keys, allowing us to align the predicted and actual values for each grid state and bus.
-
-```python
-root_pred_folder = "../data/contingency_texas/"
-prediction_dir = "prediction_gridfm01"
-label_plot = "GridFM_v0.1 Fine-tuned"
-
-pf_node_GT = pd.read_csv(os.path.join(root_pred_folder, "pf_node_10_examples.csv"))
-pg_node_predicted = pd.read_csv(
- os.path.join(root_pred_folder, "predictions_10_examples.csv")
-)
-
-branch_idx_removed = pd.read_csv("{}branch_idx_removed.csv".format(root_pred_folder))
-edge_params = pd.read_csv("{}edge_params.csv".format(root_pred_folder))
-bus_params = pd.read_csv("{}bus_params.csv".format(root_pred_folder))
-
-pf_node = pg_node_predicted.merge(pf_node_GT, on=["scenario", "bus"], how="left")
-```
-
-## Create Admittance matrix
-
-```python
-sn_mva = 100
-Yf, Yt, Vf_base_kV, Vt_base_kV = create_admittance_matrix(
- bus_params, edge_params, sn_mva
-)
-rate_a = edge_params["rate_a"]
-```
-
-## Correct voltage predictions for GridFM and DC
-
-```python
-pf_node["Vm_pred_corrected"] = pf_node["VM_pred"]
-pf_node["Va_pred_corrected"] = pf_node["VA_pred"]
-
-pf_node.loc[pf_node.PV == 1, "Vm_pred_corrected"] = pf_node.loc[pf_node.PV == 1, "Vm"]
-pf_node.loc[pf_node.REF == 1, "Va_pred_corrected"] = pf_node.loc[pf_node.REF == 1, "Va"]
-
-pf_node["Vm_dc_corrected"] = pf_node["Vm_dc"]
-pf_node["Va_dc_corrected"] = pf_node["Va_dc"]
-
-pf_node.loc[pf_node.PV == 1, "Vm_dc_corrected"] = pf_node.loc[pf_node.PV == 1, "Vm"]
-pf_node.loc[pf_node.REF == 1, "Va_dc_corrected"] = pf_node.loc[pf_node.REF == 1, "Va"]
-```
-
-## Compute branch current and line loading
-
-```python
-loadings = []
-loadings_pred = []
-loadings_dc = []
-
-for scenario_idx in tqdm(pf_node.scenario.unique()):
- pf_node_scenario = pf_node[pf_node.scenario == scenario_idx]
- branch_idx_removed_scenario = (
- branch_idx_removed[branch_idx_removed.scenario == scenario_idx]
- .iloc[:, 1:]
- .values
- )
- # remove nan
- branch_idx_removed_scenario = branch_idx_removed_scenario[
- ~np.isnan(branch_idx_removed_scenario)
- ].astype(np.int32)
- V_true = pf_node_scenario["Vm"].values * np.exp(
- 1j * pf_node_scenario["Va"].values * np.pi / 180
- )
- V_pred = pf_node_scenario["Vm_pred_corrected"].values * np.exp(
- 1j * pf_node_scenario["Va_pred_corrected"].values * np.pi / 180
- )
- V_dc = pf_node_scenario["Vm_dc_corrected"].values * np.exp(
- 1j * pf_node_scenario["Va_dc_corrected"].values * np.pi / 180
- )
- If_true, It_true = compute_branch_currents_kA(
- Yf, Yt, V_true, Vf_base_kV, Vt_base_kV, sn_mva
- )
- If_pred, It_pred = compute_branch_currents_kA(
- Yf, Yt, V_pred, Vf_base_kV, Vt_base_kV, sn_mva
- )
- If_dc, It_dc = compute_branch_currents_kA(
- Yf, Yt, V_dc, Vf_base_kV, Vt_base_kV, sn_mva
- )
-
- loading_true = compute_loading(If_true, It_true, Vf_base_kV, Vt_base_kV, rate_a)
- loading_pred = compute_loading(If_pred, It_pred, Vf_base_kV, Vt_base_kV, rate_a)
- loading_dc = compute_loading(If_dc, It_dc, Vf_base_kV, Vt_base_kV, rate_a)
-
- # remove the branches that are removed from loading
- loading_true[branch_idx_removed_scenario] = -1
- loading_pred[branch_idx_removed_scenario] = -1
- loading_dc[branch_idx_removed_scenario] = -1
-
- loadings.append(loading_true)
- loadings_pred.append(loading_pred)
- loadings_dc.append(loading_dc)
-
-
-loadings = np.array(loadings)
-loadings_pred = np.array(loadings_pred)
-loadings_dc = np.array(loadings_dc)
-removed_lines = loadings == -1
-removed_lines_pred = loadings_pred == -1
-removed_lines_dc = loadings_dc == -1
-
-
-# assert the same lines are removed
-assert (removed_lines == removed_lines_pred).all()
-assert (removed_lines == removed_lines_dc).all()
-
-# assert the same number of lines are removed
-assert removed_lines.sum() == removed_lines_pred.sum()
-assert removed_lines.sum() == removed_lines_dc.sum()
-
-overloadings = loadings[removed_lines == False] > 1.0
-overloadings_pred = loadings_pred[removed_lines == False] > 1.0
-overloadings_dc = loadings_dc[removed_lines == False] > 1.0
-```
-
-## Histogram of true line loadings
-
-```python
-plt.hist(loadings[removed_lines == False], bins=100)
-plt.xlabel("Line Loadings")
-plt.ylabel("Frequency")
-# log scale
-plt.savefig(f"loadings_histogram_{prediction_dir}.png")
-plt.show()
-```
-
- 
-
-
-
-
-## Predicted vs True line loading
-
-```python
-# Valid lines
-valid_mask = removed_lines == False
-
-# Extract valid values
-true_vals = loadings[valid_mask]
-gfm_vals = loadings_pred[valid_mask]
-dc_vals = loadings_dc[valid_mask]
-```
-
-```python
-plot_mass_correlation_density(true_vals, gfm_vals, prediction_dir, label_plot)
-```
-
- 
-
-
-
-```python
-plot_mass_correlation_density(true_vals, dc_vals, "DC", "DC Solver")
-```
-
- 
-
-
-
-```python
-plot_cm(TN_gridfm, FP_gridfm, FN_gridfm, TP_gridfm, prediction_dir, label_plot)
-```
-
- 
-
-
-
-```python
-plot_cm(TN_dc, FP_dc, FN_dc, TP_dc, "DC", "DC Solver")
-```
-
- 
-
-
-
-
-```python
-# Histograms of loadings
-plot_loading_predictions(
- loadings_pred[removed_lines == False],
- loadings_dc[removed_lines == False],
- loadings[removed_lines == False],
- prediction_dir,
- label_plot,
-)
-```
-
- 
-
-
-
-
-```python
-# create df from loadings
-loadings_df = pd.DataFrame(loadings)
-loadings_df.columns = [f"branch_{i}" for i in range(loadings_df.shape[1])]
-
-loadings_pred_df = pd.DataFrame(loadings_pred)
-loadings_pred_df.columns = [f"branch_{i}" for i in range(loadings_pred_df.shape[1])]
-
-loadings_df["scenario"] = pf_node["scenario"].unique()
-loadings_pred_df["scenario"] = pf_node["scenario"].unique()
-
-# make bar plot of wrongly classified loadings for different bins
-bins = np.arange(0, 2.2, 0.2)
-mse_pred = []
-mse_dc = []
-for i in range(len(bins) - 1):
- idx_in_bins = (loadings[removed_lines == False] > bins[i]) & (
- loadings[removed_lines == False] < bins[i + 1]
- )
- mse_pred.append(
- np.mean(
- (
- loadings_pred[removed_lines == False][idx_in_bins]
- - loadings[removed_lines == False][idx_in_bins]
- )
- ** 2
- )
- )
- mse_dc.append(
- np.mean(
- (
- loadings_dc[removed_lines == False][idx_in_bins]
- - loadings[removed_lines == False][idx_in_bins]
- )
- ** 2
- )
- )
-
-
-# labels
-labels = [f"{bins[i]:.1f}-{bins[i + 1]:.1f}" for i in range(len(bins) - 1)]
-plt.bar(labels, mse_pred, label=label_plot, alpha=0.5)
-plt.bar(labels, mse_dc, label="DC", alpha=0.5)
-plt.legend()
-plt.xlabel("Loadings")
-plt.ylabel("MSE")
-# y log scale
-plt.yscale("log")
-# rotate x labels
-plt.xticks(rotation=45)
-plt.savefig(f"loading_mse_{prediction_dir}.png")
-plt.show()
-```
-
-
- 
-
-
-
-
-## Voltage violations
-
-```python
-# merge bus_params["vmax"] and bus_params["vmin"] with pf_node on bus_idx
-pf_node = pd.merge(pf_node, bus_params[["bus", "vmax", "vmin"]], on="bus", how="left")
-
-plot_mass_correlation_density_voltage(pf_node, prediction_dir, label_plot)
-```
-
-
- 
-
-
diff --git a/docs/tutorials/feature_reconstruction.md b/docs/tutorials/feature_reconstruction.md
deleted file mode 100644
index 01a502a..0000000
--- a/docs/tutorials/feature_reconstruction.md
+++ /dev/null
@@ -1,200 +0,0 @@
-👉 [Link](https://colab.research.google.com/github/gridfm/gridfm-graphkit/blob/main/examples/notebooks/Tutorial_reconstruction_visualization.ipynb) to the tutorial on Google Colab
-
----
-This notebook demonstrates the state reconstruction capabilities of **GridFM-v0.2**, a graph-based neural network model for transmission grids. We focus on the IEEE case30 network, a standard benchmark with 30 buses, chosen for its compact size and suitability for visualization.
-
-The dataset includes **1,023 load scenarios**, each representing a different operating condition of the grid. For each scenario, the model reconstructs the six key features of the power flow solution that are masked:
-
-- Active Power Demand (MW)
-- Reactive Power Demand (MVar)
-- Active Power Generated (MW)
-- Reactive Power Generated (MVar)
-- Voltage Magnitude (p.u.)
-- Voltage Angle (degrees)
-
-```python
-import sys
-
-if "google.colab" in sys.modules:
- try:
- !git clone https://github.com/gridfm/gridfm-graphkit.git
- %cd /content/gridfm-graphkit
- !pip install .
- %cd examples/notebooks/
- except Exception as e:
- print(f"Failed to start Google Collab setup, due to {e}")
-```
-
-```python
-from gridfm_graphkit.datasets.powergrid_datamodule import LitGridDataModule
-from gridfm_graphkit.io.param_handler import NestedNamespace
-from gridfm_graphkit.tasks.feature_reconstruction_task import FeatureReconstructionTask
-from gridfm_graphkit.utils.visualization import visualize_error, visualize_quantity_heatmap
-from gridfm_graphkit.datasets.globals import PD, QD, PG, QG, VM, VA
-
-import yaml
-import torch
-import numpy as np
-import random
-```
-
-## Load YAML configuration file
-
-```python
-config_path = "../config/case30_ieee_base.yaml"
-with open(config_path) as f:
- config_dict = yaml.safe_load(f)
-
-config_args = NestedNamespace(**config_dict)
-torch.manual_seed(config_args.seed)
-random.seed(config_args.seed)
-np.random.seed(config_args.seed)
-```
-
-## Initialize the DataModule
-
-```python
-data_module = LitGridDataModule(config_args, "../data")
-data_module.setup("test")
-test_loader = data_module.test_dataloader()
-```
-
-## Load the pre-trained model GridFM-v0.2
-
-```python
-model = FeatureReconstructionTask(
- config_args, data_module.node_normalizers, data_module.edge_normalizers
-)
-state_dict = torch.load("../models/GridFM_v0_2.pth")
-model.load_state_dict(state_dict)
-```
-
-## Perform inference, batch size is equal to 1 for further visualization purposes
-
-```python
-batch = next(iter(test_loader[0]))
-
-model.eval()
-with torch.no_grad():
- output = model(
- x=batch.x,
- pe=batch.pe,
- edge_index=batch.edge_index,
- edge_attr=batch.edge_attr,
- batch=batch.batch,
- mask=batch.mask,
- )
-```
-
-## Visualize Nodal Active Power Residuals
-
-```python
-visualize_error(batch, output, data_module.node_normalizers[0])
-```
-
-
- 
-
-
-
-
-## Visualize the state reconstruction capability of gridFM-v0.2 for each feature:
-- Active Power Demand (MW)
-- Reactive Power Demand (MVar)
-- Active Power Generated (MW)
-- Reactive Power Generated (MVar)
-- Voltage Magnitude (p.u.)
-- Voltage Angle (degrees)
-
-```python
-visualize_quantity_heatmap(
- batch,
- output,
- PD,
- "Active Power Demand",
- "MW",
- data_module.node_normalizers[0],
-)
-```
-
- 
-
-
-
-
-```python
-visualize_quantity_heatmap(
- batch,
- output,
- QD,
- "Reactive Power Demand",
- "MVar",
- data_module.node_normalizers[0],
-)
-```
-
- 
-
-
-
-
-
-```python
-visualize_quantity_heatmap(
- batch,
- output,
- PG,
- "Active Power Generated",
- "MW",
- data_module.node_normalizers[0],
-)
-```
-
- 
-
-
-
-```python
-visualize_quantity_heatmap(
- batch,
- output,
- QG,
- "Reactive Power Generated",
- "MVar",
- data_module.node_normalizers[0],
-)
-```
-
- 
-
-
-
-```python
-visualize_quantity_heatmap(
- batch,
- output,
- VM,
- "Voltage magnitude",
- "p.u.",
- data_module.node_normalizers[0],
-)
-```
-
- 
-
-
-
-```python
-visualize_quantity_heatmap(
- batch,
- output,
- VA,
- "Voltage Angle",
- "degrees",
- data_module.node_normalizers[0],
-)
-```
-
- 
-
-