{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Grid extension Examples\n",
    "\n",
    "You can extend the grid by inheriting from the power-grid-model-ds. In this way you can profit from existing functionality but also add your own.\n",
    "Here we show how to extend the Grid with output variables in the node and line arrays.\n",
    "\n",
    "## Adding columns to the Grid arrays\n",
    "\n",
    "Output variables are not present in the basic Grid as defined in pgm-ds, since these might or might not be useful to specific projects.\n",
    "To add these as you please you can extend the definitions of the arrays in your own project.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "from numpy.typing import NDArray\n",
    "\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]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Adding the new arrays to the Grid\n",
    "\n",
    "When these arrays have been defined they should also be provided to the Grid object.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from dataclasses import dataclass\n",
    "\n",
    "from power_grid_model_ds import GraphContainer, Grid\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": "markdown",
   "metadata": {},
   "source": [
    "## Create a grid with this array container\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ExtendedGrid.empty()"
   ]
  }
 ],
 "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"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}