{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Arsenic Contamination Example"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "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/arsenic_contamination.ipynb\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>'))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This example demonstrates how to add a simple arsenic contamination event to a scenario."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "[EPyT-Flow](https://github.com/WaterFutures/EPyT-Flow) is available on [PyPI](https://pypi.org/project/epyt-flow/) and can be installed via `pip install epyt-flow`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%pip install epyt-flow --quiet"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "from epyt_flow.data.benchmarks import load_leakdb_scenarios\n",
    "from epyt_flow.simulation import ScenarioSimulator, EpanetConstants, ScenarioConfig\n",
    "from epyt_flow.simulation.events import SpeciesInjectionEvent\n",
    "from epyt_flow.utils import to_seconds"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Create a new scenario based on the first Net1 [LeakDB](https://epyt-flow.readthedocs.io/en/stable/epyt_flow.data.benchmarks.html#module-epyt_flow.data.benchmarks.leakdb) scenario -- we also add an additional EPANET-MSX configuration file `arsenic_contamination.msx`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "config, = load_leakdb_scenarios(scenarios_id=[\"1\"], use_net1=True, verbose=False)\n",
    "config = ScenarioConfig(scenario_config=config,\n",
    "                        f_msx_in=\"arsenic_contamination.msx\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sim = ScenarioSimulator(scenario_config=config)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Set simulation duration to 21 days by calling [set_general_parameters()](https://epyt-flow.readthedocs.io/en/stable/epyt_flow.simulation.html#epyt_flow.simulation.scenario_simulator.ScenarioSimulator.set_general_parameters) -- note that all time durations have to be stated in seconds, the function [to_seconds()](https://epyt-flow.readthedocs.io/en/stable/epyt_flow.html#epyt_flow.utils.to_seconds) can convert minutes/hours/days to seconds:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sim.set_general_parameters(simulation_duration=to_seconds(days=21))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Place some chlorine sensors and also keep track of the contaminant:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cl_sensor_locations = [\"10\", \"11\", \"12\", \"13\", \"21\", \"22\", \"23\", \"31\", \"32\"]\n",
    "all_nodes = sim.sensor_config.nodes\n",
    "\n",
    "sim.set_bulk_species_node_sensors({\"Chlorine\": cl_sensor_locations,\n",
    "                                   # Also: Keep track of the contaminant\n",
    "                                   \"AsIII\": all_nodes})   # Arsenite"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Create a 1-day contamination event of an Arsenite injection (100 mg/day) at node \"22\" by creating a new instance of the [SpeciesInjectionEvent](https://epyt-flow.readthedocs.io/en/stable/epyt_flow.simulation.events.html#epyt_flow.simulation.events.quality_events.SpeciesInjectionEvent) class and calling [add_system_event()](https://epyt-flow.readthedocs.io/en/stable/epyt_flow.simulation.html#epyt_flow.simulation.scenario_simulator.ScenarioSimulator.add_system_event):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "contamination_event = SpeciesInjectionEvent(species_id=\"AsIII\", node_id=\"22\",\n",
    "                                            profile=np.array([100]),\n",
    "                                            source_type=EpanetConstants.EN_MASS,\n",
    "                                            start_time=to_seconds(days=3),\n",
    "                                            end_time=to_seconds(days=4))\n",
    "\n",
    "sim.add_system_event(contamination_event)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Run the entire simulation by calling [run_simulation()](https://epyt-flow.readthedocs.io/en/stable/epyt_flow.simulation.html#epyt_flow.simulation.scenario_simulator.ScenarioSimulator.run_simulation):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "scada_data = sim.run_simulation()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Inspect simulation results -- i.e. plot sensor readings over time by calling the [plot_bulk_species_node_concentration](https://epyt-flow.readthedocs.io/en/stable/epyt_flow.simulation.scada.html#epyt_flow.simulation.scada.scada_data.ScadaData.plot_bulk_species_node_concentration) function:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "scada_data.plot_bulk_species_node_concentration({\"Chlorine\": cl_sensor_locations})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "scada_data.plot_bulk_species_node_concentration({\"AsIII\": all_nodes})"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Do not forget to close the simulation by calling [close()](https://epyt-flow.readthedocs.io/en/stable/epyt_flow.simulation.html#epyt_flow.simulation.scenario_simulator.ScenarioSimulator.close):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sim.close()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "epytflow2",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "name": "python",
   "version": "3.9.20"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}