{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "6b44f80a-374e-4150-a07c-154fc191a151",
   "metadata": {},
   "source": [
    "# Visualization Example"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d143eafa-0289-4dcb-9015-9c6f7091ba5d",
   "metadata": {},
   "outputs": [],
   "source": [
    "from IPython.display import display, HTML\n",
    "display(HTML('<a target=\\\"_blank\\\" href=\\\"https://colab.research.google.com/github/WaterFutures/EPyT-Flow/blob/main/docs/examples/visualization.ipynb\\\"><img src=\\\"https://colab.research.google.com/assets/colab-badge.svg\\\" alt=\\\"Open In Colab\\\"/></a>'))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "02cda241-7733-407a-9d72-21890039b6d2",
   "metadata": {},
   "source": [
    "This example demonstrates the usage of the visualization class in order to plot network topology and simulation data per hydraulic component."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f8dbec64-b868-45ad-8bc7-4a5be527e720",
   "metadata": {},
   "outputs": [],
   "source": [
    "%pip install epyt-flow --quiet"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fe774c30-559d-4f87-a35a-f75e66fcb015",
   "metadata": {},
   "outputs": [],
   "source": [
    "from epyt_flow.data.networks import load_anytown\n",
    "from epyt_flow.simulation import ScenarioSimulator\n",
    "from epyt_flow.visualization import ScenarioVisualizer"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "81394382-c707-4f49-aefc-b2c2cde47ebe",
   "metadata": {},
   "source": [
    "Load Anytown network by calling [load_anytown()](https://epyt-flow.readthedocs.io/en/stable/epyt_flow.data.html#epyt_flow.data.networks.load_anytown):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b0449e45-1c49-4bd1-96e5-7a613b1063b4",
   "metadata": {},
   "outputs": [],
   "source": [
    "network_config = load_anytown(verbose=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ccfdc00b-36ca-4161-8690-63534a42afa8",
   "metadata": {},
   "source": [
    "Create a dummy [ScenarioSimulator](https://epyt-flow.readthedocs.io/en/stable/epyt_flow.simulation.html#epyt_flow.simulation.scenario_simulator.ScenarioSimulator):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d52f9a4a-3649-4e86-9491-ee98520a79bd",
   "metadata": {},
   "outputs": [],
   "source": [
    "wdn = ScenarioSimulator(scenario_config=network_config)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "aeb5270c-33fe-48b5-b5fb-11633e4c43ca",
   "metadata": {},
   "source": [
    "Plot the network by creating a new [ScenarioVisualizer](https://epyt-flow.readthedocs.io/en/stable/epyt_flow.visualization.html#module-epyt_flow.visualization.scenario_visualizer.ScenarioVisualizer) object and calling [show_plot()](https://epyt-flow.readthedocs.io/en/stable/epyt_flow.visualization.html#epyt_flow.visualization.scenario_visualizer.ScenarioVisualizer.show_plot)\n",
    ":"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3fab0452-199e-47bc-8f6f-deeb330aeedd",
   "metadata": {},
   "outputs": [],
   "source": [
    "vis = ScenarioVisualizer(wdn)\n",
    "vis.show_plot()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2a6e1c7b-17c5-438d-ba3f-d1bb25b98c01",
   "metadata": {},
   "source": [
    "The hydraulic components can be colored according to the simulation results by calling the corresponding functions: [color_nodes()]((https://epyt-flow.readthedocs.io/en/stable/epyt_flow.visualization.html#epyt_flow.visualization.scenario_visualizer.ScenarioVisualizer.color_nodes)\n",
    "), [color_links()](https://epyt-flow.readthedocs.io/en/stable/epyt_flow.visualization.html#epyt_flow.visualization.scenario_visualizer.ScenarioVisualizer.color_links)\n",
    ", [color_pumps()](https://epyt-flow.readthedocs.io/en/stable/epyt_flow.visualization.html#epyt_flow.visualization.scenario_visualizer.ScenarioVisualizer.color_pumps)\n",
    ", [color_tanks()](https://epyt-flow.readthedocs.io/en/stable/epyt_flow.visualization.html#epyt_flow.visualization.scenario_visualizer.ScenarioVisualizer.color_tanks)\n",
    " and [color_valves()](https://epyt-flow.readthedocs.io/en/stable/epyt_flow.visualization.html#epyt_flow.visualization.scenario_visualizer.ScenarioVisualizer.color_valves)\n",
    ".\n",
    "The parameters (e.g. pressure, diameter) are component specific and can be found in the function documentation. The statistic options are the same for the different components and include mean, min, max and time_step. When using the values at a certain timestep, a point in time must be given through the pit parameter.\n",
    "\n",
    "This example shows the nodes in this function colored according to the pressure at timestep 8."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0fef82d3-8b4a-43b0-bf47-f984a3d8497a",
   "metadata": {},
   "outputs": [],
   "source": [
    "vis = ScenarioVisualizer(wdn)\n",
    "vis.color_nodes(parameter='pressure', statistic='time_step', pit=8, \n",
    "                colormap='autumn', show_colorbar=True)\n",
    "vis.show_plot()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "81710b19-b981-4af4-b0de-48d542f196fe",
   "metadata": {},
   "source": [
    "Additionally to the [color_links()](https://epyt-flow.readthedocs.io/en/stable/epyt_flow.visualization.html#epyt_flow.visualization.scenario_visualizer.ScenarioVisualizer.color_links)\n",
    " method, it is possible to resize the links according to a parameter. Multiple manipulations can be applied to the same graph at once.\n",
    "If several calls are made to color the same hydraulic component, only the last call is valid."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0df61c33-ef9f-47c6-aa82-34f1c5b2eec8",
   "metadata": {},
   "outputs": [],
   "source": [
    "vis = ScenarioVisualizer(wdn)\n",
    "vis.color_links(parameter='flow_rate', statistic='mean', show_colorbar=True)\n",
    "vis.resize_links(parameter='flow_rate', statistic='mean')\n",
    "vis.hide_nodes()\n",
    "vis.show_plot()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9da62c3c-04c0-4fb7-a8fa-adc0955ead0e",
   "metadata": {},
   "source": [
    "In order to apply a certain color scheme do all hydraulic components that are not affected by method calls, import the color scheme and apply it when creating the ScenarioVisualizer."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f611cedb-5d31-4381-8b3f-c291c7f9fc57",
   "metadata": {},
   "outputs": [],
   "source": [
    "from epyt_flow.visualization import epanet_colors, black_colors\n",
    "\n",
    "vis = ScenarioVisualizer(wdn, color_scheme=epanet_colors)\n",
    "vis.show_plot()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0771406c-3056-42ae-a83e-f20d8efc6a20",
   "metadata": {},
   "source": [
    "The colors applied to a component don't have to be continuous. It is possible to either supply an amount of intervals the components should be grouped into by using intervals=3 or to custom divide them by supplying the borders through intervals=[-100, 0, 100].\n",
    "\n",
    "Conversions to different units can also be applied by using a dictionary containing the required [epyt-flow](https://epyt-flow.readthedocs.io/en/stable/epyt_flow.simulation.scada.html#epyt_flow.simulation.scada.scada_data.ScadaData.convert_units) conversion units."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c6b082b2-ab1f-4798-ac88-b8c459d9fd65",
   "metadata": {},
   "outputs": [],
   "source": [
    "vis = ScenarioVisualizer(wdn, color_scheme=black_colors)\n",
    "vis.color_nodes(parameter='demand', statistic='time_step', pit=0, \n",
    "                colormap='autumn', intervals=3, conversion={\"flow_unit\": 8})\n",
    "vis.show_plot()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9d1cc19a-17c0-458e-adcb-fb501a2a1864",
   "metadata": {},
   "outputs": [],
   "source": [
    "vis = ScenarioVisualizer(wdn, color_scheme=black_colors)\n",
    "vis.color_links(parameter='flow_rate', statistic='max', colormap='autumn', \n",
    "                intervals=[-100, 0, 100, 200])\n",
    "vis.hide_nodes()\n",
    "vis.show_plot()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "77ae8d76-cfc4-442a-97dc-eb26cfbcad9b",
   "metadata": {},
   "source": [
    "The colormap can be adjusted by providing a matplotlib colormap as a parameter."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3cd5f617-4a02-4a70-948f-880452854157",
   "metadata": {},
   "outputs": [],
   "source": [
    "vis = ScenarioVisualizer(wdn, color_scheme=black_colors)\n",
    "vis.color_links(parameter='diameter', colormap='Blues')\n",
    "vis.resize_links(parameter='diameter', line_widths=(1, 3))\n",
    "vis.hide_nodes()\n",
    "vis.show_plot()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "25d0b91b-daf6-41f1-9caf-96fa924877be",
   "metadata": {},
   "source": [
    "Components can be labeld by calling [add_labels()](https://epyt-flow.readthedocs.io/en/stable/epyt_flow.visualization.html#epyt_flow.visualization.scenario_visualizer.ScenarioVisualizer.add_labels). The labels can be added for the sensor_config (see below), for all components via ```'all'``` or for certain components by using a list of the components names which should be labeled: ```['nodes', 'tanks', 'reservoirs', 'pipes', 'valves', 'pumps']```."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cc084795-8e2e-4499-b8ab-652f6bc3aa43",
   "metadata": {},
   "outputs": [],
   "source": [
    "from epyt_flow.data.networks import load_hanoi\n",
    "from epyt_flow.visualization import epyt_flow_colors\n",
    "\n",
    "network_config_sc = load_hanoi(verbose=False, \n",
    "                               include_default_sensor_placement=True)\n",
    "wdn_sc = ScenarioSimulator(scenario_config=network_config_sc)\n",
    "vis = ScenarioVisualizer(wdn_sc, color_scheme=epyt_flow_colors)\n",
    "vis.highlight_sensor_config()\n",
    "vis.add_labels('sensor_config', font_size=10)\n",
    "vis.show_plot()\n",
    "wdn_sc.close()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c3efea66-864e-46f9-94a9-3c3dad956598",
   "metadata": {},
   "source": [
    "Given a set of sensor applied to a water distribution network, it is possible to use the sensor data (instead of the raw simulation data) for visualization. Sensors may differ in values to the raw data due to sensor faults or noise.\n",
    "To do this, set the `use_sensor_data flag` to `True`, when coloring a component:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4c556a11-ba86-4c30-ac43-8d50948ba500",
   "metadata": {},
   "outputs": [],
   "source": [
    "network_config = load_hanoi(verbose=False, include_default_sensor_placement=False)\n",
    "wdn_sc = ScenarioSimulator(scenario_config=network_config)\n",
    "wdn_sc.set_pressure_sensors(['5', '25', '17', '20', '11'])\n",
    "vis = ScenarioVisualizer(wdn_sc, color_scheme=black_colors)\n",
    "vis.color_nodes(parameter=\"pressure\", use_sensor_data=True)\n",
    "vis.add_labels('sensor_config', font_size=10)\n",
    "vis.show_plot()\n",
    "wdn_sc.close()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4302ed97-b19f-44ed-b6ac-9e1569073acd",
   "metadata": {},
   "source": [
    "Note that in this case, the pressure can only be shown for nodes that have a pressure sensor attached to them."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "233c3212-757d-432d-87cf-1e0b3050b2e0",
   "metadata": {},
   "source": [
    "## Animations"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9b4a55ca-8858-4bca-8272-977e73db5373",
   "metadata": {},
   "source": [
    "It is possible to do an animation of the timestep values over an interval in time. For this, the point in time (pit) parameter must be provided as a tuple of the start and end values.\n",
    "Further, instead of calling [show_plot()](https://epyt-flow.readthedocs.io/en/stable/epyt_flow.visualization.html#epyt_flow.visualization.scenario_visualizer.ScenarioVisualizer.show_plot)\n",
    ", [show_animation()](https://epyt-flow.readthedocs.io/en/stable/epyt_flow.visualization.html#epyt_flow.visualization.scenario_visualizer.ScenarioVisualizer.show_animation)\n",
    " must be selected. If a file_path (anim.gif) is provided through the export_to_file parameter, the resulting animation is saved."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7e20ae24-39f2-475b-85ca-d6facb4e3880",
   "metadata": {},
   "outputs": [],
   "source": [
    "from IPython.display import HTML\n",
    "\n",
    "vis = ScenarioVisualizer(wdn, color_scheme=black_colors)\n",
    "vis.color_links(parameter='flow_rate', statistic='time_step', pit=(0, 40))\n",
    "vis.resize_links(parameter='flow_rate', statistic='time_step', pit=(0, 40), \n",
    "                 line_widths=(1, 3))\n",
    "vis.hide_nodes()\n",
    "anim = vis.show_animation(return_animation=True)\n",
    "HTML(anim.to_jshtml())"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "073e61ce-0448-4a63-8c00-a4cfed770d02",
   "metadata": {},
   "source": [
    "It is also possible to use custom data, instead of generated scada data, to visualize behaviour. In this example, periodic flow rates are animated over time. The necessary dimensions for the custom data are: time_steps x links. For other hydraulic components, it would be: time_steps x hydraulic_components."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "400bd1b2-1d87-4238-8dbe-bcb6cccbb43f",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "\n",
    "# Generate custom data for better demonstration\n",
    "timesteps = 50\n",
    "links = 41\n",
    "t = np.linspace(0, 2 * np.pi, timesteps)\n",
    "\n",
    "frequencies = np.linspace(1, 3, links)\n",
    "phases = np.linspace(0, np.pi, links)\n",
    "amplitudes = np.linspace(0.5, 1.5, links)\n",
    "\n",
    "custom_data_table = np.array([a * np.sin(f * t + p) for f, p, a in \n",
    "                              zip(frequencies, phases, amplitudes)]).T\n",
    "\n",
    "# Create new ScenarioVisualizer with black color scheme\n",
    "vis = ScenarioVisualizer(wdn, color_scheme=black_colors)\n",
    "\n",
    "# Visualize the custom data from time step 0 to the last time step (-1) as \n",
    "# link color and link size\n",
    "vis.color_links(data=custom_data_table, parameter='custom_data', \n",
    "                statistic='time_step', pit=(0, -1))\n",
    "vis.resize_links(data=custom_data_table, parameter='custom_data', \n",
    "                 statistic='time_step', pit=(0, -1), line_widths=(1, 3))\n",
    "\n",
    "# Hide the nodes such that only the links remain visible\n",
    "vis.hide_nodes()\n",
    "\n",
    "anim = vis.show_animation(return_animation=True)\n",
    "HTML(anim.to_jshtml(fps=15))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "40e1d4f4-6e5b-4bbd-b6d9-1840ed7a93f6",
   "metadata": {},
   "source": [
    "Please close the simulator after use:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0ae71123-397c-4e01-bd86-6b3429acf0c2",
   "metadata": {},
   "outputs": [],
   "source": [
    "wdn.close()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}