{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# PGM Calculation Engine Examples\n",
    "\n",
    "These examples show how to interact with the [power-grid-model](https://github.com/PowerGridModel/power-grid-model) calculation engine to perform power flow calculations. For a detailed documentation on the calculation engine please refer to [power-grid-model docs](https://power-grid-model.readthedocs.io/en/stable/).\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Single Power Flow example\n",
    "\n",
    "We will demonstrate how to use the PGMCoreInterface to perform Power Flow calculations on a `Grid` object.\n",
    "In the examples the `RadialGridGenerator` is used to create randomised input networks.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "from power_grid_model_ds import Grid, PowerGridModelInterface\n",
    "from power_grid_model_ds.generators import RadialGridGenerator"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "grid_generator = RadialGridGenerator(grid_class=Grid, nr_nodes=5, nr_sources=1, nr_nops=0)\n",
    "grid = grid_generator.run(seed=0)\n",
    "\n",
    "print(\"Created Grid:\")\n",
    "print(grid.branches)\n",
    "\n",
    "core_interface = PowerGridModelInterface(grid=grid)\n",
    "\n",
    "# Create input from grid and calculate power flow\n",
    "core_interface.create_input_from_grid()\n",
    "output = core_interface.calculate_power_flow()\n",
    "\n",
    "print(\"Power Flow Results:\")\n",
    "display(output[\"node\"])\n",
    "display(output[\"line\"])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Updating the Grid object\n",
    "\n",
    "If you want to store specific outputs of the calculation in the `Grid` object you can specify the columns in extended arrays. For a more detailed how to on extending arrays there is a seperate example.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "from dataclasses import dataclass\n",
    "\n",
    "import numpy as np\n",
    "from numpy.typing import NDArray\n",
    "\n",
    "from power_grid_model_ds import GraphContainer, Grid\n",
    "from power_grid_model_ds.arrays import LineArray, NodeArray\n",
    "\n",
    "\n",
    "class ExtendedNodeArray(NodeArray):\n",
    "    \"\"\"Extends the node array with an output value\"\"\"\n",
    "\n",
    "    _defaults = {\"u\": 0}\n",
    "\n",
    "    u: NDArray[np.float64]\n",
    "\n",
    "\n",
    "class ExtendedLineArray(LineArray):\n",
    "    \"\"\"Extends the line array with an output value\"\"\"\n",
    "\n",
    "    _defaults = {\"i_from\": 0}\n",
    "\n",
    "    i_from: NDArray[np.float64]\n",
    "\n",
    "\n",
    "@dataclass\n",
    "class ExtendedGrid(Grid):\n",
    "    \"\"\"\n",
    "    This is my own grid to extend.\n",
    "    \"\"\"\n",
    "\n",
    "    node: ExtendedNodeArray\n",
    "    line: ExtendedLineArray\n",
    "    graphs: GraphContainer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "grid_generator = RadialGridGenerator(grid_class=ExtendedGrid, nr_nodes=5, nr_sources=1, nr_nops=0)\n",
    "grid = grid_generator.run(seed=0)\n",
    "\n",
    "core_interface = PowerGridModelInterface(grid=grid)\n",
    "core_interface.create_input_from_grid()\n",
    "core_interface.calculate_power_flow()\n",
    "core_interface.update_grid()\n",
    "\n",
    "print(grid.node)\n",
    "print(grid.line)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Batch Load Flow Example\n",
    "\n",
    "To use the `PGM` batch calculation functionality you can supply a dictionary with update data to the `calculate_power_flow` function.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from power_grid_model import initialize_array\n",
    "\n",
    "grid_generator = RadialGridGenerator(grid_class=Grid, nr_nodes=5, nr_sources=1, nr_nops=0)\n",
    "grid = grid_generator.run(seed=0)\n",
    "core_interface = PowerGridModelInterface(grid=grid)\n",
    "\n",
    "update_sym_load = initialize_array(\"update\", \"sym_load\", (10, len(grid.sym_load)))\n",
    "update_sym_load[\"id\"] = [grid.sym_load.id.tolist()]\n",
    "update_sym_load[\"p_specified\"] = [grid.sym_load.p_specified.tolist()] * np.linspace(0, 1, 10).reshape(-1, 1)\n",
    "update_sym_load[\"q_specified\"] = [grid.sym_load.q_specified.tolist()] * np.linspace(0, 1, 10).reshape(-1, 1)\n",
    "update_data = {\n",
    "    \"sym_load\": update_sym_load,\n",
    "}\n",
    "output = core_interface.calculate_power_flow(update_data=update_data)\n",
    "\n",
    "# Results have been calculated for all 10 scenarios\n",
    "display(output[\"line\"])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Create grid from input data example\n",
    "\n",
    "To create a `Grid` object from `PGM` input data, the `create_grid_from_input_data()` method can be used."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "input_data = core_interface.create_input_from_grid()\n",
    "\n",
    "core_interface = PowerGridModelInterface(input_data=input_data)\n",
    "output = core_interface.create_grid_from_input_data()\n",
    "\n",
    "print(input_data[\"node\"])\n",
    "print(output.node)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": ".venv",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}