1. WNTR basic tutorial.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "tags": []
   },
   "source": [
    "# WNTR Basic Tutorial\n",
    "The following tutorial illustrates basic use of WNTR, including use of the `WaterNetworkModel` object, the ability to read/write model files to other formats, run hydraulic and water quality simulations, compute resilience metrics, define and use fragility curves, skeletonize water network models, identify network segments associated with isolation valves, and assign geospatial data to junctions and pipes."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "source": [
    "## Imports\n",
    "Install and import WNTR and additional Python packages that are needed for the tutorial\n",
    "- Numpy is required to define comparison operators (i.e., np.greater) in queries\n",
    "- Scipy is required to define lognormal fragility curves\n",
    "- NetworkX is used to compute topographic metrics\n",
    "- Geopandas is used to load geospatial data\n",
    "- Matplotlib is required to create graphics"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install wntr numpy scipy networkx geopandas==0.14.4 matplotlib "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "from scipy.stats import lognorm\n",
    "import networkx as nx\n",
    "import geopandas as gpd\n",
    "import matplotlib.pylab as plt\n",
    "import wntr"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Create directories to hold results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "if not os.path.isdir('jsons'):\n",
    "    os.makedirs('jsons')\n",
    "if not os.path.isdir('geojsons'):\n",
    "    os.makedirs('geojsons')\n",
    "if not os.path.isdir('shapefiles'):\n",
    "    os.makedirs('shapefiles')\n",
    "if not os.path.isdir('analysis'):\n",
    "    os.makedirs('analysis')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Units\n",
    "WNTR uses SI (International System) units (length in meters, time in seconds, mass in kilograms).  See https://usepa.github.io/WNTR/units.html for more details."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "tags": []
   },
   "source": [
    "# Water Network Model"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The `WaterNetworkModel` object defines the water distribution system and simulation options. The object can be created from an EPANET INP file."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create a WaterNetworkModel from an EPANET INP file\n",
    "wn = wntr.network.WaterNetworkModel('networks/Net3.inp')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Print a basic description of the model.  The level can be 0, 1, or 2 and defines the level of detail included in the description.\n",
    "wn.describe(level=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# List properties and methods associated with the WaterNetworkModel (omitting private underscore names)\n",
    "[name for name in dir(wn) if not name.startswith('_')]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Plot a basic network graphic\n",
    "ax = wntr.graphics.plot_network(wn)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "tags": []
   },
   "source": [
    "## Nodes\n",
    "Nodes define junctions, tanks, and reservoirs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Print the names of all junctions, tanks, and reservoirs\n",
    "print(\"Node names\", wn.node_name_list)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Print the names of just tanks\n",
    "print(\"Tank names\", wn.tank_name_list)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Get a tank object\n",
    "tank = wn.get_node('1')\n",
    "print(type(tank))\n",
    "tank"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# List properties and methods associated with the tank (omitting private underscore names)\n",
    "[name for name in dir(tank) if not name.startswith('_')]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Change the max level of a tank\n",
    "print(\"Original max level\", tank.max_level)\n",
    "tank.max_level = 10\n",
    "print(\"New max level\", tank.max_level)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Add a junction to the WaterNetworkModel\n",
    "wn.add_junction('new_junction', base_demand=0.0, demand_pattern=None, elevation=0.0, coordinates=None, demand_category=None)\n",
    "print(wn.junction_name_list)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Remove a junction from the WaterNetworkModel\n",
    "wn.remove_node('new_junction')\n",
    "print(wn.junction_name_list)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Links\n",
    "Links define pipes, pumps, and valves"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Print the names of all pipes, pumps, and valves\n",
    "print(\"Link names\", wn.link_name_list)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Print the names of just head pumps\n",
    "print(\"Head pump names\", wn.head_pump_name_list)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Get the name of links connected to a specific node\n",
    "connected_links = wn.get_links_for_node('229')\n",
    "print('Links connected to node 229 =', connected_links)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Get a pipe object\n",
    "pipe = wn.get_link('105')\n",
    "print(type(pipe))\n",
    "pipe"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# List properties and methods associated with the pipe (omitting private underscore names)\n",
    "[name for name in dir(pipe) if not name.startswith('_')]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Change the diameter of a pipe\n",
    "print(\"Original diameter\", pipe.diameter)\n",
    "pipe.diameter = 10\n",
    "print(\"New diameter\", pipe.diameter)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Add a pipe to the WaterNetworkModel\n",
    "wn.add_pipe(name=\"new_pipe\", start_node_name=\"10\", end_node_name=\"123\", length=304.8, diameter=0.3048, roughness=100, minor_loss=0.0, initial_status='OPEN', check_valve=False)\n",
    "print(wn.pipe_name_list)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Remove a pipe from the WaterNetworkModel\n",
    "wn.remove_link(\"new_pipe\")\n",
    "print(wn.pipe_name_list)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "tags": []
   },
   "source": [
    "## Demands and Patterns\n",
    "Junctions can have multiple demands which are stored as Timeseries objects in a `demand_timeseries_list`. Each Timeseries contains a base value, pattern, and category.  Patterns contain multipliers and the pattern timestep.  \n",
    "\n",
    "The following example illustrates how to\n",
    "* Compute expected demand (which accounts for base demand, demand patterns, and demand multiplier)\n",
    "* Compute average expected demand (average value for a 24 hour period -  accounts for base demand, demand patterns, and demand multiplier)\n",
    "* Add demands to a junction\n",
    "* Modify demand base value and pattern\n",
    "* Remove demands from a junction\n",
    "* Plot expected and simulated demands"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Compute expected demand\n",
    "expected_demand = wntr.metrics.expected_demand(wn)\n",
    "expected_demand.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Compute and plot average expected demand \n",
    "AED = wntr.metrics.average_expected_demand(wn)\n",
    "print(AED.head())\n",
    "ax = wntr.graphics.plot_network(wn, node_attribute=AED, node_range=(0,0.025), title='Average expected demand (m$^3$/s)')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Identify junctions with zero demand\n",
    "zero_demand = AED[AED == 0].index\n",
    "print(zero_demand)\n",
    "ax = wntr.graphics.plot_network(wn, node_attribute=list(zero_demand), title='Zero demand junctions')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Get the demands on Junction 15\n",
    "junction = wn.get_node('15')\n",
    "junction.demand_timeseries_list"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Get the pattern associated with the demand\n",
    "pattern = wn.get_pattern(junction.demand_timeseries_list[0].pattern_name)\n",
    "pattern"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Modify the base value of the demand\n",
    "junction.demand_timeseries_list[0].base_value = 0.005\n",
    "\n",
    "# Add a new pattern to the model\n",
    "wn.add_pattern('New', [1,1,1,0,0,0,1,0,0.5,0.5,0.5,1])\n",
    "\n",
    "# Use the new pattern to modify the junction demand\n",
    "junction.demand_timeseries_list[0].pattern_name = \"New\"\n",
    "print(junction.demand_timeseries_list)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Add a demand to Junction 15\n",
    "junction.add_demand(base=0.015, pattern_name='1')\n",
    "print(junction.demand_timeseries_list)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Plot original and modified expected demands\n",
    "new_expected_demand = wntr.metrics.expected_demand(wn) \n",
    "\n",
    "plt.figure()\n",
    "ax = expected_demand.loc[0:48*3600, \"15\"].plot()\n",
    "new_expected_demand.loc[0:48*3600, \"15\"].plot(ax=ax)\n",
    "tmp = ax.set_xlabel('Time (s)')\n",
    "tmp = ax.set_ylabel('Expected demand (m$^3$/s)')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "tags": []
   },
   "source": [
    "## Curves\n",
    "Curves define pump head curves, tank volume curves, and pump efficiency curves.  The following example illustrates how to work with pump head curves and tank volume curves."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Get a head pump object and plot the head pump curve\n",
    "pump = wn.get_link('10')\n",
    "print(type(pump))\n",
    "ax = wntr.graphics.plot_pump_curve(pump)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Get the head curve and print the points\n",
    "pump_curve_name = pump.pump_curve_name\n",
    "curve = wn.get_curve(pump_curve_name)\n",
    "curve.points"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Modify the curve points and re-plot the pump curve\n",
    "curve.points = [(0.10, 20)]\n",
    "ax = wntr.graphics.plot_pump_curve(pump)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Add a tank volume curve to the model and assign it to a tank\n",
    "wn.add_curve('new_tank_curve', 'VOLUME', [\n",
    "   (1,  0),\n",
    "   (2,  60),\n",
    "   (3,  188),\n",
    "   (4,  372),\n",
    "   (5,  596),\n",
    "   (6,  848),\n",
    "   (7,  1114),\n",
    "   (8,  1379),\n",
    "   (9,  1631),\n",
    "   (10, 1856),\n",
    "   (11, 2039),\n",
    "   (12, 2168),\n",
    "   (13, 2228)])\n",
    "tank = wn.get_node('2')\n",
    "tank.vol_curve_name = 'new_tank_curve'\n",
    "ax = wntr.graphics.plot_tank_volume_curve(tank)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Controls\n",
    "\n",
    "Controls define conditions and actions that operate pipes, pumps, and valves.  WNTR includes support for EPANET controls and rules (note that both are stored as WNTR controls). As with EPANET, controls are evaluated after each simulation timestep, while rules are evaluated after each rule timestep (see `wn.options.time`). The method `convert_controls_to_rules` can be used to convert controls to rules, which can help avoid unintended behavior when controls and rules are both used in complex simulations."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Get a list of control names\n",
    "wn.control_name_list"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Print all controls\n",
    "for name, controls in wn.controls():\n",
    "    print(name, controls)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Get a specific control object\n",
    "control = wn.get_control('control 18')\n",
    "print(control)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Modify the control priority\n",
    "control.update_priority(1)\n",
    "print(control)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Add a time based pump control\n",
    "pump = wn.get_link('10')\n",
    "action = wntr.network.controls.ControlAction(pump, 'status', 1)\n",
    "condition = wntr.network.controls.SimTimeCondition(wn, '=', '121:00:00')\n",
    "control = wntr.network.controls.Control(condition, action, name='new_control')\n",
    "wn.add_control('NewControl', control)\n",
    "print(wn.control_name_list)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Remove a control\n",
    "wn.remove_control('NewControl')\n",
    "print(wn.control_name_list)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Convert controls to rules.  This can help avoid unintended behavior when controls and rules are both used in complex simulations.\n",
    "wn.convert_controls_to_rules()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Queries\n",
    "Queries return attributes of nodes and links.  Comparison operations (like >, =) can be used to return a subset of attributes that meet specific criteria."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Return all pipe diameters (no comparison operator used in the query) \n",
    "all_pipe_diameters = wn.query_link_attribute('diameter')\n",
    "all_pipe_diameters.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Return pipes diameters > 12 inches\n",
    "large_pipe_diameters = wn.query_link_attribute('diameter', np.greater, 12*0.0254)\n",
    "print(\"Number of pipes:\", len(all_pipe_diameters))\n",
    "print(\"Number of pipes > 12 inches:\", len(large_pipe_diameters))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Plot pipes diameters > 12 inches\n",
    "ax = wntr.graphics.plot_network(wn, link_attribute=large_pipe_diameters, node_size=0, link_width=2, title=\"Pipes with diameter > 12 inches\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "tags": []
   },
   "source": [
    "## Coordinates\n",
    "Node coordinates can be obtained using a node query.  Node coordinates can also be modified using functions in `wntr.morph`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Get node coordinates\n",
    "coords = wn.query_node_attribute('coordinates')\n",
    "coords"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Rotate node coordinates counterclockwise by 30 degrees\n",
    "wn_rotated = wntr.morph.rotate_node_coordinates(wn, 30)\n",
    "ax = wntr.graphics.plot_network(wn_rotated)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "tags": []
   },
   "source": [
    "## Loops and generators\n",
    "Loops and generators are commonly used to modify network components or run stochastic simulations"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Loop over tank names and objects with a generator\n",
    "for name, tank in wn.tanks():\n",
    "    print(\"Max level for tank\", name, \"=\", tank.max_level)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Loop over tank names and then get the associated tank object\n",
    "for name in wn.tank_name_list:\n",
    "    tank = wn.get_node(name)\n",
    "    print(\"Max level for tank\", name, \"=\", tank.max_level)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Pipe breaks and leaks\n",
    "Pipes can be split (adding 1 junction to the model) or broken (adding two junctions to the model) using the `split_pipe` and `break_pipe` functions.  While a split pipe retains the network connectivity, a broken pipe does not connect across the break.  By default these functions return a copy of the WaterNetworkModel."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Split pipe 123 and add a leak to the new node which starts at hour 2 and ends at hour 12\n",
    "wn = wntr.morph.split_pipe(wn, pipe_name_to_split='123', new_pipe_name='123_B', new_junction_name='123_node')\n",
    "leak_node = wn.get_node('123_node')\n",
    "leak_node.add_leak(wn, area=0.05, start_time=2*3600, end_time=12*3600)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Break pipe 121\n",
    "wn = wntr.morph.break_pipe(wn, pipe_name_to_split='121', new_pipe_name='121_B', \n",
    "                           new_junction_name_old_pipe='121_node', new_junction_name_new_pipe='121B_node')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Model I/O\n",
    "A `WaterNetworkModel` can be converted to and from the following data formats and file types: \n",
    "* EPANET INP file\n",
    "* dictionary\n",
    "* JSON file\n",
    "* NetworkX graph\n",
    "* Geopandas GeoDataFrame\n",
    "* GeoJSON and Shapefile file"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create a WaterNetworkModel from an EPANET INP file\n",
    "wn = wntr.network.WaterNetworkModel('networks/Net3.inp')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## EPANET INP files\n",
    "WaterNetworkModel objects are commonly built from EPANET INP files.  WaterNetworkModel objects can also be saved as an EPANET INP file. Note that model attributes that are not EPANET compatible will not be saved in the INP file (i.e., leak attributes)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create an EPANET INP file from a WaterNetworkModel\n",
    "wntr.network.write_inpfile(wn, 'networks/Net3_LPS.inp', units='LPS')\n",
    "\n",
    "# Create a WaterNetworkModel from an EPANET INP file.  Note, this is equivalent to running `wn = wntr.network.WaterNetworkModel('Net3_LPS.inp')`\n",
    "wn2 = wntr.network.read_inpfile('networks/Net3_LPS.inp') "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Dictionaries\n",
    "Dictionaries offer a convenient Python format to store all the information in a WaterNetworkModel object. The dictionary can be saved to a file (typically a JSON file) to save the model. Unlike an EPANET INP file, dictionaries can contain custom model attributes."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Convert the WaterNetworkModel to a dictionary\n",
    "wn_dict = wn.to_dict()\n",
    "print(wn_dict.keys())\n",
    "\n",
    "# Create a WaterNetworkModel from a dictionary\n",
    "wn2 = wntr.network.from_dict(wn_dict)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## JSON files\n",
    "JSON files can be created directly from a WaterNetworkModel object and hold the same information as the dictionary representation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create a JSON file from the WaterNetworkModel\n",
    "wntr.network.write_json(wn, 'jsons/Net3.json')\n",
    "\n",
    "# Create a WaterNetworkModel from a JSON file\n",
    "wn2 = wntr.network.read_json('jsons/Net3.json')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## NetworkX graphs\n",
    "Graphs facilitate topographic analysis using NetworkX. WaterNetworkModel objects are represented as a MultiDiGraph, which can have multiple edges between nodes and are directed (from start node to end node). Note that WaterNetworkModel objects cannot currently be created from NetworkX graphs."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Convert the WaterNetworkModel to a MultiDiGraph\n",
    "G = wntr.network.to_graph(wn)\n",
    "print(G)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## GeoPandas GeoDataFrames\n",
    "GeoDataFrames store network attributes and geospatial geometry of junctions, tanks, reservoirs, pipes, pumps, and valves.  GeoDataFrames can be used in geospatial analysis, such as spap and intersection with other geospatial data. Note that WaterNetworkModels created from a collection of GeoDataFrames will not contain patterns, curves, rules, controls, or sources. The GeoDataFrames can be saved to file (typically GeoJSON or Shapefile) and loaded into GIS software platforms for further analysis."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Convert the WaterNetworkModel to a collection of GeoDataFrames\n",
    "wn_gis = wntr.network.to_gis(wn)\n",
    "print(wn_gis.junctions.head())\n",
    "#print(wn_gis.tanks.head())\n",
    "#print(wn_gis.reservoirs.head())\n",
    "#print(wn_gis.pipes.head())\n",
    "#print(wn_gis.pumps.head())\n",
    "#print(wn_gis.valves.head())\n",
    "\n",
    "# Create a WaterNetworkModel from a collection of GeoDataFrames\n",
    "wn2 = wntr.network.from_gis(wn_gis)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## GeoJSON files and ESRI Shapefile files\n",
    "GeoJSON and ESRI Shapefile files can be created directly from a WaterNetworkModel object.  The files can be loaded into GIS software platforms for further analysis. Note that column names longer than 10 characters will be truncated when saved to Shapefile. \n",
    "\n",
    "WaterNetworkModels can also be created from GeoJSON files or Shapefiles. A specific set of column names are required to define junctions, tanks, reservoirs, pipes, pumps, and valves (see the use of `valid_gis_names` below). Model attributes including controls, patterns, curves, and options need to be added separately."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create GeoJSON files from the WaterNetworkModel\n",
    "wntr.network.write_geojson(wn, 'geojsons/Net3')\n",
    "\n",
    "# Create a WaterNetworkModel from GeoJSON files\n",
    "geojson_files = {'junctions': 'geojsons/Net3_junctions.geojson',\n",
    "                 'tanks': 'geojsons/Net3_tanks.geojson',\n",
    "                 'reservoirs': 'geojsons/Net3_reservoirs.geojson',\n",
    "                 'pipes': 'geojsons/Net3_pipes.geojson',\n",
    "                 'pumps': 'geojsons/Net3_pumps.geojson'}\n",
    "wn2 = wntr.network.read_geojson(geojson_files)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Compare model attributes of the original model with the model built from Shapefiles (note the absence of patterns and controls)\n",
    "print(wn.describe(level=1))\n",
    "print(wn2.describe(level=1))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create Shapefiles from the WaterNetworkModel.  The following code will produce UserWarnings: \"Column names longer than 10 characters will be truncated when saved to ESRI Shapefile.\"\n",
    "wntr.network.write_shapefile(wn, 'shapefiles/Net3')\n",
    "\n",
    "# Create a WaterNetworkModel from Shapefiles\n",
    "shapefile_dirs = {'junctions': 'shapefiles/Net3_junctions',\n",
    "                  'tanks': 'shapefiles/Net3_tanks',\n",
    "                  'reservoirs': 'shapefiles/Net3_reservoirs',\n",
    "                  'pipes': 'shapefiles/Net3_pipes',\n",
    "                  'pumps': 'shapefiles/Net3_pumps'}\n",
    "wn2 = wntr.network.read_shapefile(shapefile_dirs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Compare model attributes of the original model with the model built from Shapefiles (note the absence of patterns and controls)\n",
    "print(wn.describe(level=1))\n",
    "print(wn2.describe(level=1))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Print valid GeoJSON or Shapefiles column names required to build a model\n",
    "column_names = wntr.network.io.valid_gis_names()\n",
    "print(\"Junction column names\", column_names['junctions'])\n",
    "print()\n",
    "print(\"Tank column names\", column_names['tanks'])\n",
    "print()\n",
    "print(\"Reservoir column names\", column_names['reservoirs'])\n",
    "print()\n",
    "print(\"Pipe column names\", column_names['pipes'])\n",
    "print()\n",
    "print(\"Pump column names\", column_names['pumps'])\n",
    "print()\n",
    "print(\"Valve column names\", column_names['valves'])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "tags": []
   },
   "source": [
    "# Hydraulic and Water Quality Simulations"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "WNTR includes two simulators: the `EpanetSimulator` and the `WNTRSimulator`.  Both include the ability to run pressure dependent demand (PDD) or demand-driven (DD) hydraulic simulation.  Only the EpanetSimulator runs water quality simulations."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Create a WaterNetworkModel from an EPANET INP file\n",
    "wn = wntr.network.WaterNetworkModel('networks/Net3.inp')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "tags": []
   },
   "source": [
    "## Simulation options\n",
    "Simulation options include options related to simulation time, hydraulics, water quality, reactions, energy calculations, reporting, graphics, and user/custom options."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Print the WaterNetworkModel options\n",
    "wn.options"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Change the simulation duration to 4 days\n",
    "wn.options.time.duration = 4*24*3600 # seconds\n",
    "print(wn.options.time)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Change the simulation to use pressure dependent hydraulic analysis\n",
    "wn.options.hydraulic.demand_model = 'PDD'\n",
    "wn.options.hydraulic.required_pressure = 20 # m\n",
    "wn.options.hydraulic.minimum_pressure = 2 # m\n",
    "print(wn.options.hydraulic)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## EPANET and WNTR Simulators"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Simulate hydraulics using EPANET\n",
    "sim = wntr.sim.EpanetSimulator(wn)\n",
    "results_EPANET = sim.run_sim()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Simulate hydraulics using the WNTRSimulator\n",
    "sim = wntr.sim.WNTRSimulator(wn)\n",
    "results_WNTR = sim.run_sim()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "tags": []
   },
   "source": [
    "## Simulation results\n",
    "Simulation results are stored in an object which includes a dictionary of DataFrames for nodes and a dictionary of DataFrames for links.  Each DataFrame is indexed by time (in seconds) and the columns are node or link names."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Print available node results\n",
    "results_EPANET.node.keys()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Print available link results\n",
    "results_EPANET.link.keys()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# View EpanetSimulator pressure results\n",
    "results_EPANET.node['pressure'].head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Compare EpanetSimulator and WNTRSimulator pressure results\n",
    "diff = results_EPANET.node['pressure'] - results_WNTR.node['pressure']\n",
    "ax = diff.max(axis=1).plot(title='Max difference in pressure')\n",
    "ax.set_xlabel('Time (s)')\n",
    "ax.set_ylabel('Pressure difference (m)')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Plot timeseries of tank levels\n",
    "tank_levels = results_EPANET.node['pressure'].loc[:,wn.tank_name_list]\n",
    "ax = tank_levels.plot(title='Tank level')\n",
    "ax.set_xlabel('Time (s)')\n",
    "ax.set_ylabel('Tank Level (m)')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Plot timeseries of pump flowrates\n",
    "pump_flowrates = results_EPANET.link['flowrate'].loc[:,wn.pump_name_list]\n",
    "ax = pump_flowrates.plot(title='Pump flowrate')\n",
    "ax.set_xlabel('Time (s)')\n",
    "ax.set_ylabel('Pump flowrate (m$^3$/s)')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Plot pressure at hour 5 on the network\n",
    "pressure_at_5hr = results_EPANET.node['pressure'].loc[5*3600, :]\n",
    "ax = wntr.graphics.plot_network(wn, node_attribute=pressure_at_5hr, node_size=30, title='Pressure at 5 hours')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "tags": []
   },
   "source": [
    "## Reset initial conditions\n",
    "Reset initial values, including simulation time, tank head, reservoir head, pipe status, pump status, and valve status.  This is required when using the WNTRSimulator in multiple simulations. Note, the EPANETSimulator is automatically reset."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "wn.reset_initial_values()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "tags": []
   },
   "source": [
    "# Resilience Metrics"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "tags": []
   },
   "source": [
    "WNTR includes a wide range of metrics that can be used to compute resilience and related properties."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Create a WaterNetworkModel from an EPANET INP file\n",
    "wn = wntr.network.WaterNetworkModel('networks/Net3.inp')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Topographic\n",
    "Topographic metrics describe the physical layout of the system."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Convert the WaterNetworkModel to a MultiDiGraph\n",
    "G = wn.to_graph() # directed multigraph"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Some topographic metrics require an undirected graph or a graph with a single edge between two nodes\n",
    "uG = G.to_undirected() # undirected multigraph\n",
    "sG = nx.Graph(uG) # undirected simple graph (single edge between two nodes)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Articulation points\n",
    "articulation_points = list(nx.articulation_points(uG))\n",
    "ax = wntr.graphics.plot_network(wn, node_attribute=articulation_points, title=\"Articulation points\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Compute betweenness centrality\n",
    "betweenness_centrality = nx.betweenness_centrality(G)\n",
    "ax = wntr.graphics.plot_network(wn, node_attribute=betweenness_centrality, title=\"Betweenness centrality\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "tags": []
   },
   "source": [
    "## Hydraulic\n",
    "Hydraulic metrics are based on flow, demand, and/or pressure"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Run a pressure dependent hydraulic simulation\n",
    "wn.options.hydraulic.demand_model = 'PDD'\n",
    "wn.options.hydraulic.required_pressure = 50 # m, The required pressure is set to create a scenario where not all demands are met\n",
    "\n",
    "sim = wntr.sim.EpanetSimulator(wn)\n",
    "results = sim.run_sim()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Compute water service availability, defined as the ratio of delivered demand to the expected demand. \n",
    "expected_demand = wntr.metrics.expected_demand(wn)\n",
    "demand = results.node['demand'].loc[:,wn.junction_name_list]\n",
    "wsa = wntr.metrics.water_service_availability(expected_demand.sum(axis=0), demand.sum(axis=0))\n",
    "ax = wntr.graphics.plot_network(wn, node_attribute=wsa, title='Water service availability')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "tags": []
   },
   "source": [
    "## Water Quality\n",
    "Water quality metrics are based on the concentration or water age."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Compute water age using the last 48 hours of a water quality simulation\n",
    "wn.options.quality.parameter = 'AGE'\n",
    "sim = wntr.sim.EpanetSimulator(wn)\n",
    "results = sim.run_sim()\n",
    "\n",
    "age = results.node['quality']\n",
    "age_last_48h = age.loc[age.index[-1]-48*3600:age.index[-1]]\n",
    "average_age = age_last_48h.mean()/3600 # convert to hours for the plot\n",
    "ax = wntr.graphics.plot_network(wn, node_attribute=average_age, title=\"Average water age (hr)\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Compute the population that is impacted by water age greater than 24 hours\n",
    "pop = wntr.metrics.population(wn)\n",
    "threshold = 24 # hours\n",
    "pop_impacted = wntr.metrics.population_impacted(pop, average_age, np.greater, threshold)\n",
    "ax = wntr.graphics.plot_network(wn, node_attribute=pop_impacted, title=\"Population impacted by water age > 24 hours\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "tags": []
   },
   "source": [
    "# Fragility Curves"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Fragility curves define the probability of exceeding a damage state as a function of environmental condition.  Fragility curves are commonly used in earthquake analysis, but can be defined for other scenarios."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Create a WaterNetworkModel from an EPANET INP file\n",
    "wn = wntr.network.WaterNetworkModel('networks/Net3.inp')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "FC = wntr.scenario.FragilityCurve()\n",
    "FC.add_state('Minor', 1, {'Default': lognorm(0.5,scale=0.3)})\n",
    "FC.add_state('Major', 2, {'Default': lognorm(0.5,scale=0.7)})\n",
    "ax = wntr.graphics.plot_fragility_curve(FC, xlabel='Peak Ground Acceleration (g)')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Model peak ground acceleration from a magnitude 6.5 earthquake, assigning values to pipes.\n",
    "wn = wntr.morph.scale_node_coordinates(wn, 1000)\n",
    "epicenter = (32000,15000) # x,y location\n",
    "magnitude = 6.5 # Richter scale\n",
    "depth = 10000 # m, shallow depth\n",
    "earthquake = wntr.scenario.Earthquake(epicenter, magnitude, depth)\n",
    "distance = earthquake.distance_to_epicenter(wn, element_type=wntr.network.Pipe)\n",
    "pga = earthquake.pga_attenuation_model(distance)\n",
    "ax = wntr.graphics.plot_network(wn, link_attribute=pga, node_size=0, link_width=2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Sample the failure probability and damage states for each pipe\n",
    "failure_probability = FC.cdf_probability(pga)\n",
    "damage_state = FC.sample_damage_state(failure_probability)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Plot damage state (converted to numeric values) on the network\n",
    "priority_map = FC.get_priority_map()\n",
    "damage_value = damage_state.map(priority_map)\n",
    "custom_cmp = wntr.graphics.custom_colormap(3, ['grey', 'royalblue', 'darkorange'])\n",
    "ax = wntr.graphics.plot_network(wn, link_attribute=damage_value,\n",
    "    node_size=0, link_width=2, link_cmap=custom_cmp,\n",
    "    title='Damage state: 0=None, 1=Minor, 2=Major')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "tags": []
   },
   "source": [
    "# Network Skeletonization"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Network skeletonization reduces the size of a WaterNetworkModel with the goal of having minimal impact on hydraulics. \n",
    "\n",
    "The skeletonization process retains all tanks, reservoirs, valves, and pumps, along with all junctions and pipes that are associated with controls. Junction demands and demand patterns are retained in the skeletonized model. Merged pipes are assigned equivalent properties for diameter, length, and roughness to approximate the updated system behavior. Pipes that are less than or equal to a user-defined pipe diameter threshold are candidates for removal based on branch trimming, series and parallel pipe merge."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Create a WaterNetworkModel from an EPANET INP file\n",
    "wn = wntr.network.WaterNetworkModel('networks/Net6.inp')\n",
    "wn.describe(level=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Skeletonize the network using a 12 inch pipe diameter threshold\n",
    "skel_wn = wntr.morph.skeletonize(wn, 12*0.0254)\n",
    "skel_wn.describe(level=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Plot the original and skeletonized networks\n",
    "ax = wntr.graphics.plot_network(wn, node_size=0, title='Original')\n",
    "ax = wntr.graphics.plot_network(skel_wn, node_size=0, title='Skeletonized')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Simulate hydraulics on the original and skeletonized models and plot average pressure\n",
    "sim = wntr.sim.EpanetSimulator(wn)\n",
    "results_original = sim.run_sim()\n",
    "\n",
    "sim = wntr.sim.EpanetSimulator(skel_wn)\n",
    "results_skel = sim.run_sim()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Plot average pressure at junctions that exist in both the original and skeletonized model\n",
    "skel_junctions = skel_wn.junction_name_list\n",
    "pressure_orig = results_original.node['pressure'].loc[:,skel_junctions]\n",
    "pressure_skel = results_skel.node['pressure'].loc[:,skel_junctions]\n",
    "\n",
    "ax = pressure_orig.mean(axis=1).plot(label='Original')\n",
    "ax = pressure_skel.mean(axis=1).plot(ax=ax, label='Skeletonized')\n",
    "plt.title('Average pressure')\n",
    "ax.set_xlabel('Time (s)')\n",
    "ax.set_ylabel('Pressure (m)')\n",
    "plt.legend()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Valve Segmentation"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Valve segmentation groups links and nodes into segments based on the location of isolation valves. Unlike valves that are part of a WaterNetworkModel and used in hydraulic simulations (i.e. PRV, FCV), isolation valves are not included in the WaterNetworkModel and are defined as a separate data layer."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Create a WaterNetworkModel from an EPANET INP file\n",
    "wn = wntr.network.WaterNetworkModel('networks/Net3.inp')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Create a N-2 strategic valve layer.  Note that the user can create strategic or random valve placements, or use real valve data.\n",
    "valve_layer = wntr.network.generate_valve_layer(wn, 'strategic', 2)\n",
    "ax = wntr.graphics.plot_valve_layer(wn, valve_layer, add_colorbar=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Identify nodes and links that are in each valve segment\n",
    "G = wn.to_graph()\n",
    "node_segments, link_segments, seg_sizes = wntr.metrics.topographic.valve_segments(G, valve_layer)\n",
    "seg_sizes.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Plot segments\n",
    "N = seg_sizes.shape[0] # number of segments\n",
    "cmap = wntr.graphics.random_colormap(N) # random color map helps visualize segments\n",
    "ax = wntr.graphics.plot_network(wn, link_attribute=link_segments, node_size=0, link_width=2, link_range=[0,N],  link_cmap=cmap, title='Valve segment ID')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "tags": []
   },
   "source": [
    "# Geospatial Capabilities\n",
    "Geospatial data can be used within WNTR to build a WaterNetworkModel, associate geospatial data with nodes and links, and save simulation results to GIS compatible files.\n",
    "\n",
    "**Note,  this example assumes the coordinate reference system (CRS) for the WaterNetworkModel and GIS data are in EPSG:4326 (lat/long). Geographic CRS are not suitable for measuring distance and the example will result in UserWarnings.  The example will be updated to use a projected CRS at a later date.**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Create a WaterNetworkModel from an EPANET INP file\n",
    "wn = wntr.network.WaterNetworkModel('networks/Net1.inp')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Convert a WaterNetworkModel to a collection of GeoDataFrames in EPSG:4326 coordinates\n",
    "wn_gis = wntr.network.to_gis(wn, crs='EPSG:4326')\n",
    "wn_gis.pipes.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Load hydrant data\n",
    "hydrant_data = gpd.read_file('data/Net1_hydrant_data.geojson') \n",
    "print(hydrant_data.crs)\n",
    "hydrant_data.head()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Snap data\n",
    "The `snap` function is used to find the nearest point or line to a set of points. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Snap hydrants to junctions\n",
    "snapped_to_junctions = wntr.gis.snap(hydrant_data, wn_gis.junctions, tolerance=5.0)\n",
    "snapped_to_junctions.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Plot results\n",
    "ax = hydrant_data.plot()\n",
    "ax.set_aspect('equal', adjustable='box')\n",
    "ax = wntr.graphics.plot_network(wn, node_attribute=snapped_to_junctions['node'].to_list(), ax=ax)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Intersect data\n",
    "The `intersect` function is used to find the intersection between geometries."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Load demographic data associated with census block groups\n",
    "demographic_data = gpd.read_file('data/Net1_demographic_data.geojson') \n",
    "demographic_data.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Intersect junctions with census block groups, extract mean income\n",
    "junction_demographics = wntr.gis.intersect(wn_gis.junctions, demographic_data, 'mean_income')\n",
    "junction_demographics.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Intersect pipes with census block groups, extract mean income\n",
    "pipe_demographics = wntr.gis.intersect(wn_gis.pipes, demographic_data, 'mean_income')\n",
    "pipe_demographics.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Plot results\n",
    "ax = demographic_data.plot(column='mean_income', alpha=0.5, cmap='bone', vmin=10000, vmax=100000)\n",
    "ax.set_aspect('equal', adjustable='box')\n",
    "ax = wntr.graphics.plot_network(wn, \n",
    "                                node_attribute=junction_demographics['mean'], \n",
    "                                link_attribute=pipe_demographics['weighted_mean'], \n",
    "                                link_width=1.5, ax=ax, title='Mean income ($)')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Write analysis results to GIS files\n",
    "In addition to the node and link attributes stored in the `WaterNetworkGIS` object, additional analysis results can be stored and saved to GIS compatible files."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Compute betweenness centrality and add it to the WaterNetworkGIS object\n",
    "G = wn.to_graph() # directed multigraph\n",
    "betweenness_centrality = nx.betweenness_centrality(G)\n",
    "\n",
    "wn_gis.add_node_attributes(betweenness_centrality, 'betweenness_centrality')\n",
    "wn_gis.junctions.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Run a hydraulic simulation and store average flowrate to the WaterNetworkGIS object\n",
    "sim = wntr.sim.EpanetSimulator(wn)\n",
    "results = sim.run_sim()\n",
    "flowrate = results.link['flowrate'].mean()\n",
    "\n",
    "wn_gis.add_link_attributes(flowrate, 'flowrate')\n",
    "wn_gis.pipes.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Write the model and analysis results to GIS compatible files.  These files can be loaded into GIS platforms for further analysis.\n",
    "wn_gis.write_geojson('geojsons/Net3_analysis')"
   ]
  }
 ],
 "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.9.13"
  },
  "toc-autonumbering": true,
  "vscode": {
   "interpreter": {
    "hash": "45e5fc2add644b17ca9d771f06d4d0791b1baf8fc0ad9b13d75afaf563ce7ce4"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}