docs/examples/arrow_example.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Arrow example\n",
    "\n",
    "This example illustrates how to convert Arrow data that contains power-grid-model data to NumPy structured arrays, which the power-grid-model requests.\n",
    "\n",
    "It is by no means intended to provide complete documentation on the topic, but only to show how such conversions could be done.\n",
    "\n",
    "In particular, this example restricts itself to PyArrow Tables, but for more advanced cases, RecordBatches obviously are a better solution.\n",
    "\n",
    "**NOTE:** To run this example, the optional `examples` dependencies are required:\n",
    "\n",
    "```sh\n",
    "pip install .[examples]\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%capture cap --no-stderr\n",
    "from IPython.display import display\n",
    "\n",
    "from power_grid_model import PowerGridModel, initialize_array, CalculationMethod\n",
    "import pyarrow as pa\n",
    "import pandas as pd\n",
    "import numpy as np"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Model\n",
    "\n",
    "For clarity, a simple network is created. More complex cases work similarly and can be found in the other examples:\n",
    "\n",
    "```\n",
    "node 1 ---- line 4 ---- node 2 ----line 5 ---- node 3\n",
    "   |                       |                      |\n",
    "source 6               sym_load 7             sym_load 8\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Single symmetric calculations\n",
    "\n",
    "Construct the input data for the model and construct the actual model.\n",
    "\n",
    "Arrow uses a columnar data format while the power-grid-model uses a row-based data format with continuous memory.\n",
    "Because of that, at least one copy is required.\n",
    "\n",
    "### List the power-grid-model data types\n",
    "\n",
    "See which attributes exist for a given component and which data types are used"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "node: {'names': ['id', 'u_rated'], 'formats': ['<i4', '<f8'], 'offsets': [0, 8], 'itemsize': 16, 'aligned': True}\n",
      "line: {'names': ['id', 'from_node', 'to_node', 'from_status', 'to_status', 'r1', 'x1', 'c1', 'tan1', 'r0', 'x0', 'c0', 'tan0', 'i_n'], 'formats': ['<i4', '<i4', '<i4', 'i1', 'i1', '<f8', '<f8', '<f8', '<f8', '<f8', '<f8', '<f8', '<f8', '<f8'], 'offsets': [0, 4, 8, 12, 13, 16, 24, 32, 40, 48, 56, 64, 72, 80], 'itemsize': 88, 'aligned': True}\n",
      "source: {'names': ['id', 'node', 'status', 'u_ref', 'u_ref_angle', 'sk', 'rx_ratio', 'z01_ratio'], 'formats': ['<i4', '<i4', 'i1', '<f8', '<f8', '<f8', '<f8', '<f8'], 'offsets': [0, 4, 8, 16, 24, 32, 40, 48], 'itemsize': 56, 'aligned': True}\n",
      "asym_load: {'names': ['id', 'node', 'status', 'type', 'p_specified', 'q_specified'], 'formats': ['<i4', '<i4', 'i1', 'i1', ('<f8', (3,)), ('<f8', (3,))], 'offsets': [0, 4, 8, 12, 16, 40], 'itemsize': 64, 'aligned': True}\n"
     ]
    }
   ],
   "source": [
    "node_input_dtype = initialize_array(\"input\", \"node\", 0).dtype\n",
    "line_input_dtype = initialize_array(\"input\", \"line\", 0).dtype\n",
    "source_input_dtype = initialize_array(\"input\", \"source\", 0).dtype\n",
    "asym_load_input_dtype = initialize_array(\"input\", \"asym_load\", 0).dtype\n",
    "print(\"node:\", node_input_dtype)\n",
    "print(\"line:\", line_input_dtype)\n",
    "print(\"source:\", source_input_dtype)\n",
    "print(\"asym_load:\", asym_load_input_dtype)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Create the grid using Arrow tables\n",
    "\n",
    "The [power-grid-model documentation on Components](https://power-grid-model.readthedocs.io/en/stable/user_manual/components.html) provides documentation on which components are required and which ones are optional.\n",
    "\n",
    "Construct the Arrow data as a table with the correct headers and data types."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "pyarrow.Table\n",
       "id: int32\n",
       "u_rated: double\n",
       "----\n",
       "id: [[1,2,3]]\n",
       "u_rated: [[10500,10500,10500]]"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "nodes = pa.table(\n",
    "    [\n",
    "        pa.array([1, 2, 3], type=pa.int32()),  # id\n",
    "        pa.array([10500.0, 10500.0, 10500.0], type=pa.float64()),\n",
    "    ],\n",
    "    names=(\"id\", \"u_rated\"),\n",
    ")\n",
    "lines = pa.table(\n",
    "    [\n",
    "        pa.array([4, 5], type=pa.int32()),  # id\n",
    "        pa.array([1, 2], type=pa.int32()),  # from_node\n",
    "        pa.array([2, 3], type=pa.int32()),  # to_node\n",
    "        pa.array([1, 1], type=pa.int8()),  # from_status\n",
    "        pa.array([1, 1], type=pa.int8()),  # to_status\n",
    "        pa.array([0.11, 0.15], type=pa.float64()),  # r1\n",
    "        pa.array([0.12, 0.16], type=pa.float64()),  # x1\n",
    "        pa.array([4.1e-05, 5.4e-05], type=pa.float64()),  # c1\n",
    "        pa.array([0.1, 0.1], type=pa.float64()),  # tan1\n",
    "        pa.array([0.01, 0.05], type=pa.float64()),  # r0\n",
    "        pa.array([0.22, 0.06], type=pa.float64()),  # x0\n",
    "        pa.array([4.1e-05, 5.4e-05], type=pa.float64()),  # c0\n",
    "        pa.array([0.4, 0.1], type=pa.float64()),  # tan0\n",
    "    ],\n",
    "    names=(\"id\", \"from_node\", \"to_node\", \"from_status\", \"to_status\", \"r1\", \"x1\", \"c1\", \"tan1\", \"r0\", \"x0\", \"c0\", \"tan0\"),\n",
    ")\n",
    "sources = pa.table(\n",
    "    [\n",
    "        pa.array([6], type=pa.int32()),  # id\n",
    "        pa.array([1], type=pa.int32()),  # node\n",
    "        pa.array([1], type=pa.int8()),  # status\n",
    "        pa.array([1.0], type=pa.float64()),  # u_ref\n",
    "    ],\n",
    "    names=(\"id\", \"node\", \"status\", \"u_ref\"),\n",
    ")\n",
    "sym_loads = pa.table(\n",
    "    [\n",
    "        pa.array([7, 8], type=pa.int32()),  # id\n",
    "        pa.array([2, 3], type=pa.int32()),  # node\n",
    "        pa.array([1, 1], type=pa.int8()),  # status\n",
    "        pa.array([0, 0], type=pa.int8()),  # type\n",
    "        pa.array([1.0, 2.0], type=pa.float64()),  # p_specified\n",
    "        pa.array([0.5, 1.5], type=pa.float64()),  # q_specified\n",
    "    ],\n",
    "    names=(\"id\", \"node\", \"status\", \"type\", \"p_specified\", \"q_specified\"),\n",
    ")\n",
    "\n",
    "nodes\n",
    "# the tables of the other components can be printed similarly"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Convert the Arrow data to power-grid-model input data\n",
    "\n",
    "No direct conversion from Arrow Tables to NumPy exists and a copy is always required.\n",
    "\n",
    "To ensure support for optional arguments and to prevent version lock, it is recommended to create an empty power-grid-model data set using `initialize_array` and then fill it with the Arrow data."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([(1, 10500.), (2, 10500.), (3, 10500.)],\n",
       "      dtype={'names': ['id', 'u_rated'], 'formats': ['<i4', '<f8'], 'offsets': [0, 8], 'itemsize': 16, 'aligned': True})"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "def arrow_to_numpy(data: pa.lib.Table, data_type: str, component: str) -> np.ndarray:\n",
    "    \"\"\"Convert Arrow data to NumPy data.\"\"\"\n",
    "    result = initialize_array(data_type, component, len(data))\n",
    "    for name, column in zip(data.column_names, data.columns):\n",
    "        if name in result.dtype.names:\n",
    "            result[name] = column.to_numpy()\n",
    "    return result\n",
    "\n",
    "\n",
    "node_input = arrow_to_numpy(nodes, \"input\", \"node\")\n",
    "line_input = arrow_to_numpy(lines, \"input\", \"line\")\n",
    "source_input = arrow_to_numpy(sources, \"input\", \"source\")\n",
    "sym_load_input = arrow_to_numpy(sym_loads, \"input\", \"sym_load\")\n",
    "\n",
    "node_input"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Construct the complete input data structure"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'node': array([(1, 10500.), (2, 10500.), (3, 10500.)],\n",
       "       dtype={'names': ['id', 'u_rated'], 'formats': ['<i4', '<f8'], 'offsets': [0, 8], 'itemsize': 16, 'aligned': True}),\n",
       " 'line': array([(4, 1, 2, 1, 1, 0.11, 0.12, 4.1e-05, 0.1, 0.01, 0.22, 4.1e-05, 0.4, nan),\n",
       "        (5, 2, 3, 1, 1, 0.15, 0.16, 5.4e-05, 0.1, 0.05, 0.06, 5.4e-05, 0.1, nan)],\n",
       "       dtype={'names': ['id', 'from_node', 'to_node', 'from_status', 'to_status', 'r1', 'x1', 'c1', 'tan1', 'r0', 'x0', 'c0', 'tan0', 'i_n'], 'formats': ['<i4', '<i4', '<i4', 'i1', 'i1', '<f8', '<f8', '<f8', '<f8', '<f8', '<f8', '<f8', '<f8', '<f8'], 'offsets': [0, 4, 8, 12, 13, 16, 24, 32, 40, 48, 56, 64, 72, 80], 'itemsize': 88, 'aligned': True}),\n",
       " 'source': array([(6, 1, 1, 1., nan, nan, nan, nan)],\n",
       "       dtype={'names': ['id', 'node', 'status', 'u_ref', 'u_ref_angle', 'sk', 'rx_ratio', 'z01_ratio'], 'formats': ['<i4', '<i4', 'i1', '<f8', '<f8', '<f8', '<f8', '<f8'], 'offsets': [0, 4, 8, 16, 24, 32, 40, 48], 'itemsize': 56, 'aligned': True}),\n",
       " 'sym_load': array([(7, 2, 1, 0, 1., 0.5), (8, 3, 1, 0, 2., 1.5)],\n",
       "       dtype={'names': ['id', 'node', 'status', 'type', 'p_specified', 'q_specified'], 'formats': ['<i4', '<i4', 'i1', 'i1', '<f8', '<f8'], 'offsets': [0, 4, 8, 12, 16, 24], 'itemsize': 32, 'aligned': True})}"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "input_data = {\n",
    "    \"node\": node_input,\n",
    "    \"line\": line_input,\n",
    "    \"source\": source_input,\n",
    "    \"sym_load\": sym_load_input,\n",
    "}\n",
    "\n",
    "input_data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Optional: validate the input data\n",
    "from power_grid_model.validation import validate_input_data\n",
    "\n",
    "validate_input_data(input_data)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Use the power-grid-model\n",
    "\n",
    "For more extensive examples, visit the [power-grid-model documentation](https://power-grid-model.readthedocs.io/en/stable/index.html)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>id</th>\n",
       "      <th>energized</th>\n",
       "      <th>u_pu</th>\n",
       "      <th>u</th>\n",
       "      <th>u_angle</th>\n",
       "      <th>p</th>\n",
       "      <th>q</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>1</td>\n",
       "      <td>1</td>\n",
       "      <td>1.000325</td>\n",
       "      <td>10503.410670</td>\n",
       "      <td>-0.000067</td>\n",
       "      <td>338777.246279</td>\n",
       "      <td>-3.299419e+06</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>2</td>\n",
       "      <td>1</td>\n",
       "      <td>1.002879</td>\n",
       "      <td>10530.228073</td>\n",
       "      <td>-0.002932</td>\n",
       "      <td>-1.000000</td>\n",
       "      <td>-5.000000e-01</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>3</td>\n",
       "      <td>1</td>\n",
       "      <td>1.004113</td>\n",
       "      <td>10543.184974</td>\n",
       "      <td>-0.004342</td>\n",
       "      <td>-2.000000</td>\n",
       "      <td>-1.500000e+00</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "   id  energized      u_pu             u   u_angle              p  \\\n",
       "0   1          1  1.000325  10503.410670 -0.000067  338777.246279   \n",
       "1   2          1  1.002879  10530.228073 -0.002932      -1.000000   \n",
       "2   3          1  1.004113  10543.184974 -0.004342      -2.000000   \n",
       "\n",
       "              q  \n",
       "0 -3.299419e+06  \n",
       "1 -5.000000e-01  \n",
       "2 -1.500000e+00  "
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# construct the model\n",
    "model = PowerGridModel(input_data=input_data, system_frequency=50)\n",
    "\n",
    "# run the calculation\n",
    "sym_result = model.calculate_power_flow()\n",
    "\n",
    "# use pandas to tabulate and display the results\n",
    "pd_sym_node_result = pd.DataFrame(sym_result[\"node\"])\n",
    "display(pd_sym_node_result)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Convert power-grid-model output data to Arrow output data\n",
    "\n",
    "Using Pandas DataFrames as an intermediate type, constructing Arrow data formats is straightfoward"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "pyarrow.Table\n",
       "id: int32\n",
       "energized: int8\n",
       "u_pu: double\n",
       "u: double\n",
       "u_angle: double\n",
       "p: double\n",
       "q: double\n",
       "----\n",
       "id: [[1,2,3]]\n",
       "energized: [[1,1,1]]\n",
       "u_pu: [[1.000324825742982,1.0028788641128947,1.004112854674026]]\n",
       "u: [[10503.410670301311,10530.228073185395,10543.184974077272]]\n",
       "u_angle: [[-0.00006651843181519333,-0.0029317915196014274,-0.004341587216862399]]\n",
       "p: [[338777.2462788447,-1.0000001549705169,-1.9999999440349978]]\n",
       "q: [[-3299418.6613065186,-0.4999999565008232,-1.4999999075367236]]"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "pa_sym_node_result = pa.table(pd_sym_node_result)\n",
    "\n",
    "# and similar for other components\n",
    "\n",
    "pa_sym_node_result"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Single asymmetric calculations\n",
    "\n",
    "Asymmetric calculations have a tuple of values for some of the attributes and are not easily convertible to pandas data frames.\n",
    "Instead, one can have a look at the individual components of those attributes and/or flatten the arrays to access all components.\n",
    "\n",
    "### Asymmetric input\n",
    "\n",
    "To illustrate the conversion, let's consider a similar grid but with asymmetric loads.\n",
    "\n",
    "```\n",
    "node 1 ---- line 4 ---- node 2 ----line 5 ---- node 3\n",
    "   |                       |                      |\n",
    "source 6              asym_load 7            asym_load 8\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "asym_load: {'names': ['id', 'node', 'status', 'type', 'p_specified', 'q_specified'], 'formats': ['<i4', '<i4', 'i1', 'i1', ('<f8', (3,)), ('<f8', (3,))], 'offsets': [0, 4, 8, 12, 16, 40], 'itemsize': 64, 'aligned': True}\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "pyarrow.Table\n",
       "id: int32\n",
       "node: int32\n",
       "status: int8\n",
       "type: int8\n",
       "p_specified_a: double\n",
       "p_specified_b: double\n",
       "p_specified_c: double\n",
       "q_specified_a: double\n",
       "q_specified_b: double\n",
       "q_specified_c: double\n",
       "----\n",
       "id: [[7,8]]\n",
       "node: [[2,3]]\n",
       "status: [[1,1]]\n",
       "type: [[0,0]]\n",
       "p_specified_a: [[1,2]]\n",
       "p_specified_b: [[0.01,2.5]]\n",
       "p_specified_c: [[0.011,450]]\n",
       "q_specified_a: [[0.5,1.5]]\n",
       "q_specified_b: [[1500,2.5]]\n",
       "q_specified_c: [[0.1,1500]]"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "asym_load_input_dtype = initialize_array(\"input\", \"asym_load\", 0).dtype\n",
    "print(\"asym_load:\", asym_load_input_dtype)\n",
    "\n",
    "asym_loads = pa.table(\n",
    "    [\n",
    "        pa.array([7, 8], type=pa.int32()),  # id\n",
    "        pa.array([2, 3], type=pa.int32()),  # node\n",
    "        pa.array([1, 1], type=pa.int8()),  # status\n",
    "        pa.array([0, 0], type=pa.int8()),  # type\n",
    "        pa.array([1.0, 2.0], type=pa.float64()),  # p_specified_a\n",
    "        pa.array([1.0e-2, 2.5], type=pa.float64()),  # p_specified_b\n",
    "        pa.array([1.1e-2, 4.5e2], type=pa.float64()),  # p_specified_c\n",
    "        pa.array([0.5, 1.5], type=pa.float64()),  # q_specified_a\n",
    "        pa.array([1.5e3, 2.5], type=pa.float64()),  # q_specified_b\n",
    "        pa.array([0.1, 1.5e3], type=pa.float64()),  # q_specified_c\n",
    "    ],\n",
    "    names=(\"id\", \"node\", \"status\", \"type\", \"p_specified_a\", \"p_specified_b\", \"p_specified_c\", \"q_specified_a\", \"q_specified_b\", \"q_specified_c\"),\n",
    ")\n",
    "\n",
    "asym_loads"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([(7, 2, 1, 0, [1.0e+00, 1.0e-02, 1.1e-02], [5.0e-01, 1.5e+03, 1.0e-01]),\n",
       "       (8, 3, 1, 0, [2.0e+00, 2.5e+00, 4.5e+02], [1.5e+00, 2.5e+00, 1.5e+03])],\n",
       "      dtype={'names': ['id', 'node', 'status', 'type', 'p_specified', 'q_specified'], 'formats': ['<i4', '<i4', 'i1', 'i1', ('<f8', (3,)), ('<f8', (3,))], 'offsets': [0, 4, 8, 12, 16, 40], 'itemsize': 64, 'aligned': True})"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "def arrow_to_numpy_asym(data: pa.lib.Table, data_type: str, component: str) -> np.ndarray:\n",
    "    \"\"\"Convert asymmetric Arrow data to NumPy data.\n",
    "    \n",
    "    This function is similar to the arrow_to_numpy function, but also supports asymmetric data.\"\"\"\n",
    "    result = initialize_array(data_type, component, len(data))\n",
    "    phases = (\"a\", \"b\", \"c\")\n",
    "    for name, (dtype, _) in result.dtype.fields.items():\n",
    "        if len(dtype.shape) == 0:\n",
    "            # simple or symmetric data type\n",
    "            if name in data.column_names:\n",
    "                result[name] = data.column(name).to_numpy()\n",
    "        else:\n",
    "            # asymmetric data type\n",
    "            for phase_index, phase in enumerate(phases):\n",
    "                phase_name = f\"{name}_{phase}\"\n",
    "\n",
    "                if phase_name in data.column_names:\n",
    "                    result[name][:, phase_index] = data.column(phase_name).to_numpy()\n",
    "\n",
    "    return result\n",
    "\n",
    "asym_load_input = arrow_to_numpy_asym(asym_loads, \"input\", \"asym_load\")\n",
    "\n",
    "asym_load_input\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Use the power-grid-model in asymmetric calculations"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>0</th>\n",
       "      <th>1</th>\n",
       "      <th>2</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>-0.000067</td>\n",
       "      <td>-2.094462</td>\n",
       "      <td>2.094328</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>-0.002930</td>\n",
       "      <td>-2.097322</td>\n",
       "      <td>2.091464</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>-0.004338</td>\n",
       "      <td>-2.098733</td>\n",
       "      <td>2.090057</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "          0         1         2\n",
       "0 -0.000067 -2.094462  2.094328\n",
       "1 -0.002930 -2.097322  2.091464\n",
       "2 -0.004338 -2.098733  2.090057"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "asym_input_data = {\n",
    "    \"node\": node_input,\n",
    "    \"line\": line_input,\n",
    "    \"source\": source_input,\n",
    "    \"asym_load\": asym_load_input,\n",
    "}\n",
    "\n",
    "validate_input_data(asym_input_data, symmetric=False)\n",
    "\n",
    "# construct the model\n",
    "asym_model = PowerGridModel(input_data=asym_input_data, system_frequency=50)\n",
    "\n",
    "# run the calculation\n",
    "asym_result = asym_model.calculate_power_flow(symmetric=False)\n",
    "\n",
    "# use pandas to display the results, but beware the data types\n",
    "pd.DataFrame(asym_result[\"node\"][\"u_angle\"])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Convert asymmetric power-grid-model output data to Arrow output data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "pyarrow.Table\n",
       "id: int32\n",
       "energized: int8\n",
       "u_pu_a: double\n",
       "u_pu_b: double\n",
       "u_pu_c: double\n",
       "u_a: double\n",
       "u_b: double\n",
       "u_c: double\n",
       "u_angle_a: double\n",
       "u_angle_b: double\n",
       "u_angle_c: double\n",
       "p_a: double\n",
       "p_b: double\n",
       "p_c: double\n",
       "q_a: double\n",
       "q_b: double\n",
       "q_c: double\n",
       "----\n",
       "id: [[1,2,3]]\n",
       "energized: [[1,1,1]]\n",
       "u_pu_a: [[1.0003248257977395,1.0028803762176168,1.004114300817404]]\n",
       "u_pu_b: [[1.000324376948685,1.0028710993140397,1.0041033583077168]]\n",
       "u_pu_c: [[1.00032436416241,1.002873078902152,1.0041004935738533]]\n",
       "u_a: [[6064.146978239599,6079.639179329459,6087.119449677851]]\n",
       "u_b: [[6064.144257236812,6079.582941090295,6087.053114238259]]\n",
       "u_c: [[6064.1441797241405,6079.594941705455,6087.035747712152]]\n",
       "u_angle_a: [[-0.00006651848125692708,-0.0029298831864833634,-0.004337685507209539]]\n",
       "u_angle_b: [[-2.0944615736658134,-2.0973219974462594,-2.098732840554144]]\n",
       "..."
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "def numpy_to_arrow(data: np.ndarray) -> pa.lib.table:\n",
    "    \"\"\"Convert NumPy data to Arrow data.\"\"\"\n",
    "    simple_data_types = []\n",
    "    multi_value_data_types = []\n",
    "\n",
    "    for name, (dtype, _) in data.dtype.fields.items():\n",
    "        if len(dtype.shape) == 0:\n",
    "            simple_data_types.append(name)\n",
    "        else:\n",
    "            multi_value_data_types.append(name)\n",
    "\n",
    "    result = pa.table(pd.DataFrame(data[simple_data_types]))\n",
    "\n",
    "    phases = (\"a\", \"b\", \"c\")\n",
    "    for name in multi_value_data_types:\n",
    "        column = data[name]\n",
    "\n",
    "        assert column.shape[1] == len(phases), \"Asymmetric data has 3 phase output\"\n",
    "\n",
    "        for phase_index, phase in enumerate(phases):\n",
    "            sub_column = column[:, phase_index]\n",
    "            result = result.append_column(f\"{name}_{phase}\", [pd.Series(sub_column)])\n",
    "\n",
    "    return result\n",
    "\n",
    "\n",
    "pa_asym_node_result = numpy_to_arrow(asym_result[\"node\"])\n",
    "\n",
    "pa_asym_node_result"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Batch data\n",
    "\n",
    "power-grid-model supports batch calculations by providing an `update_data` argument, as shown in [this example](https://power-grid-model.readthedocs.io/en/stable/examples/Power%20Flow%20Example.html#batch-calculation).\n",
    "\n",
    "Both the `update_data` and the output result are similar to the `input_data` and output data in the above, except that they have another dimension representing the batch index: the first index in the NumPy structured arrays.\n",
    "\n",
    "This extra index can be represented in Arrow using a [`RecordBatch`](https://arrow.apache.org/docs/cpp/api/table.html#two-dimensional-datasets) or using any other multi-index data format."
   ]
  }
 ],
 "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.11.2"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}