Repository: IMMM-SFA/mosartwmpy (language: python)

Issue:
Investigate better parallel processing schemes
Current parallel implementation only scales efficiently to ~8 CPUs due to reliance on `numexpr` vector math. May need to implement true MPI via `mpi4py` or otherwise convert vector based math to scalar cell-by-cell math.


Discussion / hints:

Fix patch:
diff --git a/.github/workflows/build_linux.yml b/.github/workflows/build.yml
similarity index 98%
rename from .github/workflows/build_linux.yml
rename to .github/workflows/build.yml
index 5820793..a54de76 100644
--- a/.github/workflows/build_linux.yml
+++ b/.github/workflows/build.yml
@@ -1,4 +1,4 @@
-name: linux
+name: build
 
 on: [push]
 
diff --git a/.github/workflows/build_osx.yml b/.github/workflows/build_osx.yml
deleted file mode 100644
index 97f2526..0000000
--- a/.github/workflows/build_osx.yml
+++ /dev/null
@@ -1,34 +0,0 @@
-name: osx
-
-on: [pull_request]
-
-jobs:
-  build:
-      runs-on: ${{ matrix.os }}
-
-      strategy:
-        matrix:
-          os: [macos-latest]
-
-      env:
-        OS: ${{ matrix.os }}
-        PYTHON: '3.7'
-
-      steps:
-
-        - uses: actions/checkout@v1
-
-        - name: Set up Python
-          uses: actions/setup-python@master
-          with:
-            python-version: 3.7
-
-        - name: Install dependencies
-          run: |
-            python -m pip install --upgrade pip
-            pip install .
-
-        - name: Test
-          run: |
-            pip install pytest
-            pytest
diff --git a/.github/workflows/build_windows.yml b/.github/workflows/build_windows.yml
deleted file mode 100644
index 61e0962..0000000
--- a/.github/workflows/build_windows.yml
+++ /dev/null
@@ -1,40 +0,0 @@
-name: windows
-
-on: [pull_request]
-
-jobs:
-  build:
-      runs-on: ${{ matrix.os }}
-
-      strategy:
-        matrix:
-          os: [windows-latest]
-
-      env:
-        OS: ${{ matrix.os }}
-        PYTHON: '3.7'
-
-      steps:
-
-        - uses: actions/checkout@v1
-
-        - name: Set up Python
-          uses: actions/setup-python@master
-          with:
-            python-version: 3.7
-
-        - name: Install dependencies
-          run: |
-            python -m pip install --upgrade pip
-            curl -O https://download.lfd.uci.edu/pythonlibs/q4trcu4l/numpy-1.20.3+mkl-cp37-cp37m-win_amd64.whl
-            pip install numpy-1.20.3+mkl-cp37-cp37m-win_amd64.whl
-            curl -O https://download.lfd.uci.edu/pythonlibs/q4trcu4l/GDAL-3.3.0-cp37-cp37m-win_amd64.whl
-            pip install GDAL-3.3.0-cp37-cp37m-win_amd64.whl
-            curl -O https://download.lfd.uci.edu/pythonlibs/q4trcu4l/rasterio-1.2.4-cp37-cp37m-win_amd64.whl
-            pip install rasterio-1.2.4-cp37-cp37m-win_amd64.whl
-            pip install .
-
-        - name: Test
-          run: |
-            pip install pytest
-            pytest
diff --git a/README.md b/README.md
index c653e23..f87d45c 100644
--- a/README.md
+++ b/README.md
@@ -1,7 +1,6 @@
-[![linux](https://github.com/IMMM-SFA/mosartwmpy/actions/workflows/build_linux.yml/badge.svg)](https://github.com/IMMM-SFA/mosartwmpy/actions/workflows/build_linux.yml)
-[![osx](https://github.com/IMMM-SFA/mosartwmpy/actions/workflows/build_osx.yml/badge.svg)](https://github.com/IMMM-SFA/mosartwmpy/actions/workflows/build_osx.yml)
-[![windows](https://github.com/IMMM-SFA/mosartwmpy/actions/workflows/build_windows.yml/badge.svg)](https://github.com/IMMM-SFA/mosartwmpy/actions/workflows/build_windows.yml)
+[![build](https://github.com/IMMM-SFA/mosartwmpy/actions/workflows/build.yml/badge.svg)](https://github.com/IMMM-SFA/mosartwmpy/actions/workflows/build.yml)
 [![codecov](https://codecov.io/gh/IMMM-SFA/mosartwmpy/branch/main/graph/badge.svg?token=IPOY8984MB)](https://codecov.io/gh/IMMM-SFA/mosartwmpy)
+[![DOI](https://joss.theoj.org/papers/10.21105/joss.03221/status.svg)](https://doi.org/10.21105/joss.03221)
 [![DOI](https://zenodo.org/badge/312114600.svg)](https://zenodo.org/badge/latestdoi/312114600)
 
 ## mosartwmpy
diff --git a/build.sh b/build.sh
new file mode 100755
index 0000000..502be10
--- /dev/null
+++ b/build.sh
@@ -0,0 +1,3 @@
+rm -f dist/*
+python -m build
+python -m twine upload dist/*
\ No newline at end of file
diff --git a/config.yaml b/config.yaml
index 8a31624..142f4ae 100644
--- a/config.yaml
+++ b/config.yaml
@@ -6,29 +6,23 @@ simulation:
     # path to a directory for writing output files; will be created if doesn't exist
     output_path: ./output
     # date to begin the simulation
-    start_date: 1981-05-24
+    start_date: 1981-05-01
     # date to end the simulation
-    end_date: 1981-05-26
+    end_date: 1981-05-30
 
-# For example -- to change the input files, update the paths below:
+grid:
+  path: ./input/domains/mosart.nc
+  land:
+    path: ./input/domains/land.nc
 
-# grid:
-#     # path to the grid domain file; can be absolute or relative to the source code root
-#     path: /pic/projects/im3/iwmm/rof/mosart/MOSART_NLDAS_8th_20160426.nc
-#     land:
-#       # path to the land grid domain file; can be absolute or relative to the source code root
-#       path: /pic/projects/im3/iwmm/share/domains/domain.clm/domain.lnd.nldas2_0224x0464_c110415.nc
-# runoff:
-#     # whether or not to read runoff input from file
-#     read_from_file: true
-#     # path to the runoff file; can be absolute or relative to the source code root
-#     path: /pic/projects/im3/iwmm/lnd/dlnd7/NLDAS/Livneh_NLDAS_1980_2011.nc
-# water_management:
-#     # flag to turn this feature on or off
-#     enabled: true
-#     demand:
-#       # path to the demand file; can be absolute or relative to the source code root
-#       path: /pic/projects/im3/iwmm/input_data/demand/GCAM_waterdemand_nc/rcp8.5/RCP8.5_GCAM_water_demand_{year}_{month}.nc
-#     reservoirs:
-#       # path to the reservoir parameter file; can be absolute or relative to the source code root
-#       path: /pic/projects/im3/iwmm/input_data/runoff/US_reservoir_8th_NLDAS3_updated_20200421.nc
+runoff:
+  read_from_file: true
+  path: ./input/runoff/runoff_1981_05.nc
+
+water_management:
+  enabled: true
+  demand:
+    read_from_file: true
+    path: ./input/demand/demand_1981_05.nc
+  reservoirs:
+    path: ./input/reservoirs/reservoirs.nc
diff --git a/docs/_static/api.svg b/docs/_static/api.svg
new file mode 100644
index 0000000..4f23413
--- /dev/null
+++ b/docs/_static/api.svg
@@ -0,0 +1,1 @@
+<svg id="icon" xmlns="http://www.w3.org/2000/svg" viewBox="0 0 32 32"><defs><style>.cls-1{fill:none;}</style></defs><title>api</title><path d="M26,22a3.86,3.86,0,0,0-2,.57l-3.09-3.1a6,6,0,0,0,0-6.94L24,9.43A3.86,3.86,0,0,0,26,10a4,4,0,1,0-4-4,3.86,3.86,0,0,0,.57,2l-3.1,3.09a6,6,0,0,0-6.94,0L9.43,8A3.86,3.86,0,0,0,10,6a4,4,0,1,0-4,4,3.86,3.86,0,0,0,2-.57l3.09,3.1a6,6,0,0,0,0,6.94L8,22.57A3.86,3.86,0,0,0,6,22a4,4,0,1,0,4,4,3.86,3.86,0,0,0-.57-2l3.1-3.09a6,6,0,0,0,6.94,0L22.57,24A3.86,3.86,0,0,0,22,26a4,4,0,1,0,4-4ZM26,4a2,2,0,1,1-2,2A2,2,0,0,1,26,4ZM4,6A2,2,0,1,1,6,8,2,2,0,0,1,4,6ZM6,28a2,2,0,1,1,2-2A2,2,0,0,1,6,28Zm10-8a4,4,0,1,1,4-4A4,4,0,0,1,16,20Zm10,8a2,2,0,1,1,2-2A2,2,0,0,1,26,28Z" transform="translate(0 0)"/><rect class="cls-1" width="32" height="32"/></svg>
\ No newline at end of file
diff --git a/docs/_static/cognitive.svg b/docs/_static/cognitive.svg
new file mode 100644
index 0000000..dc8f670
--- /dev/null
+++ b/docs/_static/cognitive.svg
@@ -0,0 +1,1 @@
+<svg id="icon" xmlns="http://www.w3.org/2000/svg" viewBox="0 0 32 32"><defs><style>.cls-1{fill:none;}</style></defs><title>cognitive</title><path d="M30,13A11,11,0,0,0,19,2H11a9,9,0,0,0-9,9v3a5,5,0,0,0,5,5H8.1A5,5,0,0,0,13,23h1.38l4,7,1.73-1-4-6.89A2,2,0,0,0,14.38,21H13a3,3,0,0,1,0-6h1V13H13a5,5,0,0,0-4.9,4H7a3,3,0,0,1-3-3V12H6A3,3,0,0,0,9,9V8H7V9a1,1,0,0,1-1,1H4.08A7,7,0,0,1,11,4h6V6a1,1,0,0,1-1,1H14V9h2a3,3,0,0,0,3-3V4a9,9,0,0,1,8.05,5H26a3,3,0,0,0-3,3v1h2V12a1,1,0,0,1,1-1h1.77A8.76,8.76,0,0,1,28,13v1a5,5,0,0,1-5,5H20v2h3a7,7,0,0,0,3-.68V21a3,3,0,0,1-3,3H22v2h1a5,5,0,0,0,5-5V18.89A7,7,0,0,0,30,14Z"/><rect id="_Transparent_Rectangle_" data-name="&lt;Transparent Rectangle&gt;" class="cls-1" width="32" height="32"/></svg>
\ No newline at end of file
diff --git a/docs/_static/education.svg b/docs/_static/education.svg
new file mode 100644
index 0000000..5f5c20d
--- /dev/null
+++ b/docs/_static/education.svg
@@ -0,0 +1,1 @@
+<svg id="icon" xmlns="http://www.w3.org/2000/svg" viewBox="0 0 32 32"><defs><style>.cls-1{fill:none;}</style></defs><title>education</title><path d="M26,30H24V27a5.0059,5.0059,0,0,0-5-5H13a5.0059,5.0059,0,0,0-5,5v3H6V27a7.0082,7.0082,0,0,1,7-7h6a7.0082,7.0082,0,0,1,7,7Z"/><path d="M5,6A1,1,0,0,0,4,7v9H6V7A1,1,0,0,0,5,6Z"/><path d="M4,2V4H9v7a7,7,0,0,0,14,0V4h5V2Zm7,2H21V7H11Zm5,12a5,5,0,0,1-5-5V9H21v2A5,5,0,0,1,16,16Z"/><rect id="_Transparent_Rectangle_" data-name="&lt;Transparent Rectangle&gt;" class="cls-1" width="32" height="32"/></svg>
\ No newline at end of file
diff --git a/docs/_static/soccer.svg b/docs/_static/soccer.svg
new file mode 100644
index 0000000..8f5dc8a
--- /dev/null
+++ b/docs/_static/soccer.svg
@@ -0,0 +1,1 @@
+<svg id="icon" xmlns="http://www.w3.org/2000/svg" width="32" height="32" viewBox="0 0 32 32"><defs><style>.cls-1{fill:none;}</style></defs><title>soccer</title><circle cx="17" cy="28" r="2"/><rect x="7.5501" y="17.5001" width="7.8997" height="1.9998" transform="translate(-9.7132 13.5503) rotate(-45)"/><path d="M28,16.584,19.4143,8H6v2H18.5857l3,3L6,28.5859,7.4141,30,23,14.415,26.584,18,23,21.5859,24.4141,23,28,19.416A2.0044,2.0044,0,0,0,28,16.584Z"/><path d="M24.5,9A3.5,3.5,0,1,1,28,5.5,3.504,3.504,0,0,1,24.5,9Zm0-5A1.5,1.5,0,1,0,26,5.5,1.5017,1.5017,0,0,0,24.5,4Z"/><rect id="_Transparent_Rectangle_" data-name="&lt;Transparent Rectangle&gt;" class="cls-1" width="32" height="32"/></svg>
\ No newline at end of file
diff --git a/docs/_static/style.css b/docs/_static/style.css
new file mode 100644
index 0000000..d4da7a5
--- /dev/null
+++ b/docs/_static/style.css
@@ -0,0 +1,48 @@
+/* Override some aspects of the pydata-sphinx-theme */
+
+:root {
+  /* Use softer blue from bootstrap's default info color */
+  --pst-color-info: 23, 162, 184;
+}
+
+/* Main index page overview cards */
+
+.intro-card {
+  background: #fff;
+  border-radius: 0;
+  padding: 30px 10px 20px 10px;
+  margin: 10px 0px;
+}
+
+.intro-card p.card-text {
+  margin: 0px;
+}
+
+.intro-card .card-img-top {
+  margin: 0px;
+  height: 52px;
+}
+
+.intro-card .card-header {
+  border: none;
+  background-color:white;
+  color: #150458 !important;
+  font-size: var(--pst-font-size-h5);
+  font-weight: bold;
+  padding: 0.5rem 0rem 0.5rem 0rem;
+}
+
+.intro-card .card-footer {
+  border: none;
+  background-color:white;
+}
+
+.intro-card .card-footer p.card-text{
+  max-width: 220px;
+  margin-left: auto;
+  margin-right: auto;
+}
+
+.card-body {
+    padding: 0.5rem;
+}
\ No newline at end of file
diff --git a/docs/conf.py b/docs/conf.py
index 0549199..e9ea4aa 100644
--- a/docs/conf.py
+++ b/docs/conf.py
@@ -32,10 +32,11 @@
 # ones.
 extensions = [
     'nbsphinx',
-    'recommonmark',                                                             
-    'sphinx.ext.napoleon',                                                      
-    'sphinx.ext.autodoc',                                                       
-    'sphinx_rtd_theme' 
+    'recommonmark',
+    'sphinx.ext.napoleon',
+    'sphinx.ext.autodoc',
+    'sphinx_panels',
+    'sphinx_rtd_theme'
 ]
 
 # Add any paths that contain templates here, relative to this directory.
@@ -57,4 +58,10 @@
 # Add any paths that contain custom static files (such as style sheets) here,
 # relative to this directory. They are copied after the builtin static files,
 # so a file named "default.css" will overwrite the builtin "default.css".
-html_static_path = ['_static']
\ No newline at end of file
+html_static_path = ['_static']
+html_css_files = ['style.css']
+
+panels_delimiters = (r"^\>{3,}$", r"^\^{3,}$", r"^\+{3,}$")
+
+import mosartwmpy
+version = str(mosartwmpy.__version__)
diff --git a/docs/index.rst b/docs/index.rst
index 656d50d..cbe383a 100644
--- a/docs/index.rst
+++ b/docs/index.rst
@@ -1,30 +1,89 @@
-mosartwmpy
-==========
+:notoc:
 
-Introduction
-------------
-.. toctree::
-   :maxdepth: 1
+.. mosartwmpy documentation master file, created by
 
-   README.md
+.. module:: mosartwmpy
 
-Tutorial
---------
-.. toctree::
-   :maxdepth: 1
+************************
+mosartwmpy documentation
+************************
 
-   tutorial.ipynb
+**Date**: |today| **Version**: |version|
 
-Tips and Tricks
----------------
-.. toctree::
-   :maxdepth: 1
+**Useful links**:
+`Source Repository <https://github.com/immm-sfa/mosartwmpy>`__ |
+`Issues & Ideas <https://github.com/immm-sfa/mosartwmpy/issues>`__
+
+`mosartwmpy` is a Python translation of MOSART-WM, a water routing and reservoir management model written in Fortran.
+
+
+.. panels::
+    :card: + intro-card text-center
+    :column: col-lg-6 col-md-6 col-sm-6 col-xs-12 d-flex
+
+    >>>
+    :img-top: _static/cognitive.svg
+
+    Getting started
+    ^^^
+
+    Get to know the *mosartwmpy* model.
+
+    +++
+    .. link-button:: README
+            :type: ref
+            :text: Getting started
+            :classes: btn-block btn-secondary stretched-link
+
+    >>>
+    :img-top: _static/education.svg
+
+    Tutorial
+    ^^^
+
+    Follow along with this Jupyter notebook to learn the ropes of *mosartwmpy*.
+
+    +++
+    .. link-button:: tutorial
+            :type: ref
+            :text: Tutorial
+            :classes: btn-block btn-secondary stretched-link
+
+    >>>
+    :img-top: _static/soccer.svg
+
+    Tips & tricks
+    ^^^
+
+    Learn about ways to manage Python virtual environments.
+
+    +++
+    .. link-button:: virtualenv
+            :type: ref
+            :text: Virtual environments
+            :classes: btn-block btn-secondary stretched-link
+
+    >>>
+    :img-top: _static/api.svg
+
+    API reference
+    ^^^
+
+    A detailed description of the *mosartwmpy* API.
+
+    +++
+    .. link-button:: mosartwmpy
+            :type: ref
+            :text: API
+            :classes: btn-block btn-secondary stretched-link
 
-   virtualenv
 
-Application Programming Interface
----------------------------------
 .. toctree::
-   :maxdepth: 1
+    :maxdepth: 1
+    :hidden:
+    :titlesonly:
 
-   mosartwmpy
+    README.md
+    tutorial.ipynb
+    virtualenv
+    mosartwmpy
diff --git a/docs/virtualenv.rst b/docs/virtualenv.rst
index 94b18e1..56fccbe 100644
--- a/docs/virtualenv.rst
+++ b/docs/virtualenv.rst
@@ -1,10 +1,11 @@
 Python Virtual Environments
 ===========================
 
-Maintaining different versions of Python can be a chore, especially when ``mosartwmpy`` is forcing you to use Python 3.9+ 😈. There are many tools for managing Python environments, and many more will probably be developed after I write this, but at this moment in time my two recommendations are:
+Maintaining different versions of Python can be a chore. Thankfully, there are many tools for managing Python environments; here are a few recommendations:
 
-* `PyCharm <https://www.jetbrains.com/pycharm/>`_ IDE -- this is great for developing code in Python, and can automatically create virtual environments for a codebase by detecting versions and dependencies from the ``setup.py`` or ``setup.cfg``.
-* `pyenv <https://github.com/pyenv/pyenv>`_ CLI -- this is a shell based tool for installing and switching between different versions of Python and dependencies. I will give a brief tutorial of using ``pyenv`` below, but recognize that the instructions may change over time so the ``pyenv`` documentation is the best place to look.
+* `PyCharm <https://www.jetbrains.com/pycharm/>`_ IDE -- great for developing code in Python, and can automatically create virtual environments for a codebase by detecting versions and dependencies from the ``setup.py`` or ``setup.cfg``.
+* `Conda <https://docs.conda.io>`_ package manager -- a Python package manager focused on scientific computing that can also manage virtual environments.
+* `pyenv <https://github.com/pyenv/pyenv>`_ CLI -- a shell based tool for installing and switching between different versions of Python and dependencies. I will give a brief tutorial of using ``pyenv`` below, but recognize that the instructions may change over time so the ``pyenv`` documentation is the best place to look.
 
 To create a Python 3.9 virtual environment, try the following steps:
 
diff --git a/mosartwmpy/__init__.py b/mosartwmpy/__init__.py
index 3b4d86e..d4caa30 100644
--- a/mosartwmpy/__init__.py
+++ b/mosartwmpy/__init__.py
@@ -1,1 +1,2 @@
 from .model import Model
+from ._version import __version__
diff --git a/mosartwmpy/_version.py b/mosartwmpy/_version.py
new file mode 100644
index 0000000..3dc1f76
--- /dev/null
+++ b/mosartwmpy/_version.py
@@ -0,0 +1,1 @@
+__version__ = "0.1.0"
diff --git a/mosartwmpy/config/parameters.py b/mosartwmpy/config/parameters.py
index 814c30c..f2f11f4 100644
--- a/mosartwmpy/config/parameters.py
+++ b/mosartwmpy/config/parameters.py
@@ -35,21 +35,21 @@ def __init__(self):
         # minimum main channel slope (replaces 0s from grid file)
         self.channel_slope_minimum = 0.0001
         # kinematic wave condition # TODO what is it?
-        self.kinematic_wave_condition = 1.0e6
+        self.kinematic_wave_parameter = 1.0e6
         
         # reservoir parameters # TODO better describe
         self.reservoir_minimum_flow_condition = 0.05
         self.reservoir_flood_control_condition = 1.0
         self.reservoir_small_magnitude_difference = 0.01
         self.reservoir_regulation_release_parameter = 0.85
-        self.reservoir_runoff_capacity_condition = 0.1
+        self.reservoir_runoff_capacity_parameter = 0.1
         self.reservoir_flow_volume_ratio = 0.9
         
         # number of supply iterations
         self.reservoir_supply_iterations = 3
         
         # minimum depth to perform irrigation extraction [m]
-        self.irrigation_extraction_condition = 0.1
+        self.irrigation_extraction_parameter = 0.1
         # maximum fraction of flow that can be extracted from main channel
         self.irrigation_extraction_maximum_fraction = 0.5
         
diff --git a/mosartwmpy/config_defaults.yaml b/mosartwmpy/config_defaults.yaml
index f51f5aa..78074af 100644
--- a/mosartwmpy/config_defaults.yaml
+++ b/mosartwmpy/config_defaults.yaml
@@ -21,14 +21,12 @@ simulation:
     subcycles: 3
     # routing iterations - subsubtimestep will be timestep / subcycles / routing_iterations - then further subdivided within subnetwork and main channel based on convergence estimations
     routing_iterations: 5
-    # routing method to use
-    # TODO only 1 is implemented so far
-    # 1 --> variable storage method from SWAT model; 2 --> Muskingum method?
-    routing_method: 1
     # log level, one of DEBUG, INFO, WARNING, ERROR, CRITICAL
     log_level: INFO
     # whether or not to log to console; slightly more efficient to disable this when running on a cluster
     log_to_std_out: true
+    # whether or not to log to file; slightly more efficient to disable, but useful for keeping a record of experiment
+    log_to_file: true
     # resume simulation from restart file
     restart_file: ~
     # how often to write a restart file; one of 'daily', 'monthly', 'yearly' -- a restart file is also always written at the end of the simulation
@@ -191,12 +189,6 @@ runoff:
 water_management:
     # flag to turn this feature on or off
     enabled: true
-    # flag to enable or disable regulation
-    regulation_enabled: true
-    # flag to enable or disable extraction
-    extraction_enabled: true
-    # extraction from main channel
-    extraction_main_channel_enabled: true
     # Demand
     demand:
         # flag to enable/disable reading demand from file
diff --git a/mosartwmpy/data_manifest.yaml b/mosartwmpy/data_manifest.yaml
index accdbdb..c35818a 100644
--- a/mosartwmpy/data_manifest.yaml
+++ b/mosartwmpy/data_manifest.yaml
@@ -12,7 +12,7 @@ sample_input:
 
 validation:
   description: Result datasets that can be used for validating the model; includes results with and without WM; covers 1981-1982.
-  url: https://zenodo.org/record/4539693/files/mosartwmpy_validation.zip?download=1
+  url: https://zenodo.org/record/5083443/files/mosartwmpy_validation.zip?download=1
   destination: ./validation
 
 # TODO add other weather files, demand files, etc as they become ready
\ No newline at end of file
diff --git a/mosartwmpy/direct_to_ocean/__init__.py b/mosartwmpy/direct_to_ocean/__init__.py
deleted file mode 100644
index e69de29..0000000
diff --git a/mosartwmpy/direct_to_ocean/direct_to_ocean.py b/mosartwmpy/direct_to_ocean/direct_to_ocean.py
deleted file mode 100644
index 70dc3c3..0000000
--- a/mosartwmpy/direct_to_ocean/direct_to_ocean.py
+++ /dev/null
@@ -1,106 +0,0 @@
-import numpy as np
-import numexpr as ne
-import pandas as pd
-
-from benedict.dicts import benedict as Benedict
-
-from mosartwmpy.config.parameters import Parameters
-from mosartwmpy.grid.grid import Grid
-from mosartwmpy.state.state import State
-
-def direct_to_ocean(state: State, grid: Grid, parameters: Parameters, config: Benedict) -> None:
-    """Direct transfer to outlet point; mutates state.
-
-    Args:
-        state (State): the current model state; will be mutated
-        grid (Grid): the model grid
-        parameters (Parameters): the model parameters
-        config (Benedict): the model configuration
-    """
-
-    # direct to ocean
-    # note - in fortran mosart this direct_to_ocean forcing could be provided from LND component, but we don't seem to be using it
-    source_direct = 1.0 * state.direct_to_ocean
-    
-    # wetland runoff
-    wetland_runoff_volume = state.hillslope_wetland_runoff * config.get('simulation.timestep') / config.get('simulation.subcycles')
-    river_volume_minimum = parameters.river_depth_minimum * grid.area
-
-    # if wetland runoff is negative and it would bring main channel storage below the minimum, send it directly to ocean
-    condition = ((state.channel_storage + wetland_runoff_volume) < river_volume_minimum) & (state.hillslope_wetland_runoff < 0)
-    source_direct = np.where(
-        condition,
-        source_direct + state.hillslope_wetland_runoff,
-        source_direct
-    )
-    state.hillslope_wetland_runoff = np.where(
-        condition,
-        0,
-        state.hillslope_wetland_runoff
-    )
-    # remove remaining wetland runoff (negative and positive)
-    source_direct = source_direct + state.hillslope_wetland_runoff
-    state.hillslope_wetland_runoff = 0.0 * state.zeros
-    
-    # runoff from hillslope
-    # remove negative subsurface water
-    condition = state.hillslope_subsurface_runoff < 0
-    source_direct = np.where(
-        condition,
-        source_direct + state.hillslope_subsurface_runoff,
-        source_direct
-    )
-    state.hillslope_subsurface_runoff = np.where(
-        condition,
-        0,
-        state.hillslope_subsurface_runoff
-    )
-    # remove negative surface water
-    condition = state.hillslope_surface_runoff < 0
-    source_direct = np.where(
-        condition,
-        source_direct + state.hillslope_surface_runoff,
-        source_direct
-    )
-    state.hillslope_surface_runoff = np.where(
-        condition,
-        0,
-        state.hillslope_surface_runoff
-    )
-
-    # if ocean cell or ice tracer, remove the rest of the sub and surface water
-    # other cells will be handled by mosart euler
-    condition = (grid.mosart_mask == 0) | (state.tracer == parameters.ICE_TRACER)
-    source_direct = np.where(
-        condition,
-        source_direct + state.hillslope_subsurface_runoff + state.hillslope_surface_runoff,
-        source_direct
-    )
-    state.hillslope_subsurface_runoff = np.where(
-        condition,
-        0,
-        state.hillslope_subsurface_runoff
-    )
-    state.hillslope_surface_runoff = np.where(
-        condition,
-        0,
-        state.hillslope_surface_runoff
-    )
-    
-    state.direct[:] = source_direct
-
-    # send the direct water to outlet for each tracer
-    state.direct[:] = pd.DataFrame(
-        grid.id, columns=['id']
-    ).merge(
-        pd.DataFrame(
-            state.direct, columns=['direct']
-        ).join(
-            pd.DataFrame(
-                grid.outlet_id, columns=['outlet_id']
-            )
-        ).groupby('outlet_id').sum(),
-        how='left',
-        left_on='id',
-        right_index=True
-    ).direct.fillna(0.0).values
\ No newline at end of file
diff --git a/mosartwmpy/flood/__init__.py b/mosartwmpy/flood/__init__.py
deleted file mode 100644
index e69de29..0000000
diff --git a/mosartwmpy/flood/flood.py b/mosartwmpy/flood/flood.py
deleted file mode 100644
index c80d17a..0000000
--- a/mosartwmpy/flood/flood.py
+++ /dev/null
@@ -1,43 +0,0 @@
-import numpy as np
-
-from benedict.dicts import benedict as Benedict
-from mosartwmpy.config.parameters import Parameters
-from mosartwmpy.grid.grid import Grid
-from mosartwmpy.state.state import State
-
-def flood(state: State, grid: Grid, parameters: Parameters, config: Benedict) -> None:
-    """Excess runoff is removed from the available groundwater; mutates state.
-
-    Args:
-        state (State): the current model state; will be mutated
-        grid (Grid): the model grid
-        parameters (Parameters): the model parameters
-        config (Benedict): the model configuration
-    """
-    
-    ###
-    ### Compute Flood
-    ### Remove excess liquid water from land
-    ###
-    ### TODO tcraig leaves a comment here concerning surface_runoff in fortran mosart:
-    ### "This seems like an odd approach, you
-    ### might create negative forcing.  why not take it out of
-    ### the volr directly?  it's also odd to compute this
-    ### at the initial time of the time loop.  why not do
-    ### it at the end or even during the run loop as the
-    ### new volume is computed.  fluxout depends on volr, so
-    ### how this is implemented does impact the solution."
-    ###
-
-    # flux sent back to land
-    state.flood = np.where(
-        (grid.land_mask == 1) & (state.storage > parameters.flood_threshold) & (state.tracer == parameters.LIQUID_TRACER),
-        (state.storage - parameters.flood_threshold) / config.get('simulation.timestep'),
-        0
-    )
-    # remove this flux from the input runoff from land
-    state.hillslope_surface_runoff = np.where(
-        state.tracer == parameters.LIQUID_TRACER,
-        state.hillslope_surface_runoff - state.flood,
-        state.hillslope_surface_runoff
-    )
\ No newline at end of file
diff --git a/mosartwmpy/grid/grid.py b/mosartwmpy/grid/grid.py
index a8eb502..80b5a43 100644
--- a/mosartwmpy/grid/grid.py
+++ b/mosartwmpy/grid/grid.py
@@ -1,4 +1,3 @@
-import logging
 import numpy as np
 import pandas as pd
 from pathlib import Path
@@ -7,14 +6,20 @@
 import xarray as xr
 
 from benedict.dicts import benedict as Benedict
+from numba.core import types
+from numba.typed import Dict
 from xarray import open_dataset
 from zipfile import ZIP_DEFLATED, ZipFile
 
 from mosartwmpy.config.parameters import Parameters
 from mosartwmpy.reservoirs.grid import load_reservoirs
 
+# ignore numpy NaN and invalid warnings
+# (i.e. divide by zero and NaN logicals -- in Pandas/Dask, these simply remain NaN instead of becoming infinite)
+np.seterr(all='ignore')
 
-class Grid():
+
+class Grid:
     """Class to store grid related values that are constant throughout a simulation."""
     
     # initialize all properties
@@ -74,8 +79,11 @@ class Grid():
     reservoir_use_fish_protection: np.ndarray = np.empty(0)
     reservoir_withdrawal: np.ndarray = np.empty(0)
     reservoir_conveyance: np.ndarray = np.empty(0)
-    reservoir_count: np.ndarray = np.empty(0)
     reservoir_to_grid_mapping: pd.DataFrame = pd.DataFrame()
+    reservoir_to_grid_map: Dict = Dict.empty(
+        key_type=types.int64,
+        value_type=types.int64[:],
+    )
     reservoir_streamflow_schedule: xr.DataArray = xr.DataArray()
     reservoir_demand_schedule: xr.DataArray = xr.DataArray()
     reservoir_prerelease_schedule: xr.DataArray = xr.DataArray()
@@ -93,8 +101,6 @@ def __init__(self, config: Benedict = None, parameters: Parameters = None, empty
         if empty:
             return
         
-        logging.debug('Loading grid file.')
-        
         # open dataset
         grid_dataset = open_dataset(config.get('grid.path'))
     
@@ -118,7 +124,6 @@ def __init__(self, config: Benedict = None, parameters: Parameters = None, empty
         grid_dataset.close()
         
         # use ID and dnID field to calculate masks, upstream, downstream, and outlet indices, as well as count of upstream cells
-        logging.debug(' - masks, downstream, upstream, and outlet cell indices')
     
         # ocean/land mask
         # 1 == land
@@ -192,7 +197,6 @@ def __init__(self, config: Benedict = None, parameters: Parameters = None, empty
         
         # recalculate area to fill in missing values
         # assumes grid spacing is in degrees and uniform
-        logging.debug(' - area')
         deg2rad = np.pi / 180.0
         self.area = np.where(
             self.local_drainage_area <= 0,
@@ -230,8 +234,6 @@ def __init__(self, config: Benedict = None, parameters: Parameters = None, empty
         except:
             self.land_fraction = np.full(self.id.size, 1.0)
 
-        logging.debug(' - main channel iterations')
-
         # parameter for calculating number of main channel iterations needed
         # phi_r
         phi_main = np.where(
@@ -249,9 +251,7 @@ def __init__(self, config: Benedict = None, parameters: Parameters = None, empty
                 1 + config.get('simulation.subcycles'),
                 1
             )
-        )
-
-        logging.debug(' - subnetwork substeps')
+        ).astype(np.int64)
 
         # total main channel length [m]
         # rlenTotal
@@ -335,12 +335,11 @@ def __init__(self, config: Benedict = None, parameters: Parameters = None, empty
                 )
             ),
             1
-        )
+        ).astype(np.int64)
 
         # if water management is enabled, load the reservoir parameters and build the grid cell mapping
         # note that reservoir grid is assumed to be the same as the domain grid
         if config.get('water_management.enabled', False):
-            logging.debug(' - reservoirs')
             load_reservoirs(self, config, parameters)
 
     def __getitem__(self, item):
@@ -404,11 +403,11 @@ def to_files(self, path: str) -> None:
             for df in dfs:
                 names.append(f'{df["key"]}.df.nc')
                 paths.append(f'{tmpdir}/{names[-1]}')
-                df['frame'].to_xarray().to_netcdf(paths[-1], engine='scipy')
+                df['frame'].to_xarray().to_netcdf(paths[-1], engine='h5netcdf')
             for ds in xrs:
                 names.append(f'{ds["key"]}.xr.nc')
                 paths.append(f'{tmpdir}/{names[-1]}')
-                ds['data_array'].to_netcdf(paths[-1], engine='scipy')
+                ds['data_array'].to_netcdf(paths[-1], engine='h5netcdf')
             with ZipFile(path, 'w', compression=ZIP_DEFLATED, compresslevel=9) as zip:
                 for i, filename in enumerate(paths):
                     zip.write(filename, names[i])
@@ -441,14 +440,19 @@ def from_files(path: Path) -> 'Grid':
                             setattr(grid, key, npdf[key].values)
                     if filename.endswith('df.nc'):
                         key = filename.split('.')[0]
-                        ds = xr.open_dataset(file, engine='scipy')
+                        ds = xr.open_dataset(file, engine='h5netcdf')
                         df = ds.to_dataframe()
                         setattr(grid, key, df)
                         ds.close()
                     if filename.endswith('xr.nc'):
                         key = filename.split('.')[0]
-                        ds = xr.open_dataarray(file, engine='scipy').load()
+                        ds = xr.open_dataarray(file, engine='h5netcdf').load()
                         setattr(grid, key, ds)
                         ds.close()
+
+        # recreate the numba grid to reservoir map
+        if grid.reservoir_to_grid_mapping.size > 0:
+            for grid_cell_id, group in grid.reservoir_to_grid_mapping.reset_index().groupby('grid_cell_id'):
+                grid.reservoir_to_grid_map[grid_cell_id] = group.reservoir_id.values
         
-        return grid
\ No newline at end of file
+        return grid
diff --git a/mosartwmpy/hillslope/routing.py b/mosartwmpy/hillslope/routing.py
index 45b94e4..b989ae4 100644
--- a/mosartwmpy/hillslope/routing.py
+++ b/mosartwmpy/hillslope/routing.py
@@ -1,83 +1,55 @@
-import numpy as np
-import numexpr as ne
-
-from mosartwmpy.config.parameters import Parameters
-from mosartwmpy.grid.grid import Grid
-from mosartwmpy.state.state import State
-from mosartwmpy.hillslope.state import update_hillslope_state
-
-def hillslope_routing(state: State, grid: Grid, parameters: Parameters, delta_t: float) -> None:
-    """Tracks the storage of runoff water in the hillslope and the flow of runoff water from the hillslope into the channels.
-
-    Args:
-        state (State): the current model state; will be mutated
-        grid (Grid): the model grid
-        parameters (Parameters): the model parameters
-        delta_t (float): the timestep for this subcycle (overall timestep / subcycles)
-    """
-    
-    base_condition = (grid.mosart_mask > 0) & state.euler_mask
-    
-    velocity_hillslope = calculate_velocity_hillslope(base_condition, state.hillslope_depth, grid.hillslope, grid.hillslope_manning)
-    state.hillslope_overland_flow = calculate_base_hillslope_overland_flow(base_condition, state.hillslope_depth, velocity_hillslope, grid.drainage_density, state.hillslope_overland_flow)
-    state.hillslope_overland_flow = calculate_hillslope_overland_flow(base_condition, state.hillslope_overland_flow, state.hillslope_storage, delta_t, state.hillslope_surface_runoff, parameters.tiny_value)
-    state.hillslope_delta_storage = calculate_hillslope_delta_storage(base_condition, state.hillslope_surface_runoff, state.hillslope_overland_flow, state.hillslope_delta_storage)
-    state.hillslope_storage = calculate_hillslope_storage(base_condition, state.hillslope_storage, delta_t, state.hillslope_delta_storage)
-    
-    update_hillslope_state(state, base_condition)
-    
-    state.subnetwork_lateral_inflow = calculate_subnetwork_lateral_inflow(base_condition, state.hillslope_subsurface_runoff, state.hillslope_overland_flow, grid.drainage_fraction, grid.area, state.subnetwork_lateral_inflow)
-
-calculate_velocity_hillslope = ne.NumExpr(
-    'where('
-        'base_condition & (hillslope_depth > 0),'
-        '(hillslope_depth ** (2/3)) * sqrt(hillslope) / hillslope_manning,'
-        '0'
-    ')',
-    (('base_condition', np.bool), ('hillslope_depth', np.float64), ('hillslope', np.float64), ('hillslope_manning', np.float64))
+import numba as nb
+
+
+@nb.jit(
+    "void("
+        "int64, float64,"
+        "int64[:], float64[:], float64[:], float64[:], float64[:], float64[:],"
+        "boolean[:], float64[:], float64[:], float64[:], float64[:], float64[:],"
+        "float64[:], float64[:], float64"
+    ")",
+    nopython=True,
+    nogil=True,
+    cache=True,
 )
-
-calculate_base_hillslope_overland_flow = ne.NumExpr(
-    'where('
-        'base_condition,'
-        '-hillslope_depth * velocity_hillslope * drainage_density,'
-        'hillslope_overland_flow'
-    ')',
-    (('base_condition', np.bool), ('hillslope_depth', np.float64), ('velocity_hillslope', np.float64), ('drainage_density', np.float64), ('hillslope_overland_flow', np.float64))
-)
-
-calculate_hillslope_overland_flow = ne.NumExpr(
-    'where('
-        'base_condition & (hillslope_overland_flow < 0) & ((hillslope_storage + delta_t * (hillslope_surface_runoff + hillslope_overland_flow)) < tiny_value),'
-        '-(hillslope_surface_runoff + hillslope_storage / delta_t),'
-        'hillslope_overland_flow'
-    ')',
-    (('base_condition', np.bool), ('hillslope_overland_flow', np.float64), ('hillslope_storage', np.float64), ('delta_t', np.float64), ('hillslope_surface_runoff', np.float64), ('tiny_value', np.float64))
-)
-
-calculate_hillslope_delta_storage = ne.NumExpr(
-    'where('
-        'base_condition,'
-        'hillslope_surface_runoff + hillslope_overland_flow,'
-        'hillslope_delta_storage'
-    ')',
-    (('base_condition', np.bool), ('hillslope_surface_runoff', np.float64), ('hillslope_overland_flow', np.float64), ('hillslope_delta_storage', np.float64))
-)
-
-calculate_hillslope_storage = ne.NumExpr(
-    'where('
-        'base_condition,'
-        'hillslope_storage + delta_t * hillslope_delta_storage,'
-        'hillslope_storage'
-    ')',
-    (('base_condition', np.bool), ('hillslope_storage', np.float64), ('delta_t', np.float64), ('hillslope_delta_storage', np.float64))
-)
-
-calculate_subnetwork_lateral_inflow = ne.NumExpr(
-    'where('
-        'base_condition,'
-        '(hillslope_subsurface_runoff - hillslope_overland_flow) * drainage_fraction * area,'
-        'subnetwork_lateral_inflow'
-    ')',
-    (('base_condition', np.bool), ('hillslope_subsurface_runoff', np.float64), ('hillslope_overland_flow', np.float64), ('drainage_fraction', np.float64), ('area', np.float64), ('subnetwork_lateral_inflow', np.float64))
-)
\ No newline at end of file
+def hillslope_routing(
+    i,
+    delta_t,
+    mosart_mask,
+    hillslope,
+    hillslope_manning,
+    drainage_density,
+    drainage_fraction,
+    area,
+    euler_mask,
+    hillslope_depth,
+    hillslope_overland_flow,
+    hillslope_storage,
+    hillslope_surface_runoff,
+    hillslope_delta_storage,
+    hillslope_subsurface_runoff,
+    subnetwork_lateral_inflow,
+    tiny_value,
+):
+    """Tracks the storage of runoff water in the hillslope and the flow of runoff water from the hillslope into the channels."""
+
+    if ~euler_mask[i] or ~(mosart_mask[i] > 0):
+        return
+
+    if hillslope_depth[i] > 0.0:
+        velocity_hillslope = (hillslope_depth[i] ** (2.0/3.0)) * (hillslope[i] ** (1.0/2.0)) / hillslope_manning[i]
+    else:
+        velocity_hillslope = 0.0
+
+    hillslope_overland_flow[i] = -hillslope_depth[i] * velocity_hillslope * drainage_density[i]
+
+    if (hillslope_overland_flow[i] < 0.0) and ((hillslope_storage[i] + delta_t * (hillslope_surface_runoff[i] + hillslope_overland_flow[i])) < tiny_value):
+        hillslope_overland_flow[i] = -(hillslope_surface_runoff[i] + hillslope_storage[i] / delta_t)
+
+    hillslope_delta_storage[i] = hillslope_surface_runoff[i] + hillslope_overland_flow[i]
+
+    hillslope_storage[i] = hillslope_storage[i] + delta_t * hillslope_delta_storage[i]
+
+    hillslope_depth[i] = hillslope_storage[i]
+
+    subnetwork_lateral_inflow[i] = (hillslope_subsurface_runoff[i] - hillslope_overland_flow[i]) * drainage_fraction[i] * area[i]
diff --git a/mosartwmpy/hillslope/state.py b/mosartwmpy/hillslope/state.py
deleted file mode 100644
index cca37b5..0000000
--- a/mosartwmpy/hillslope/state.py
+++ /dev/null
@@ -1,23 +0,0 @@
-import numpy as np
-import numexpr as ne
-
-from mosartwmpy.state.state import State
-
-def update_hillslope_state(state: State, base_condition: np.ndarray) -> None:
-    """Updates the depth of water remaining in the hillslope.
-
-    Args:
-        state (State): the current model state; will be mutated
-        base_condition (np.ndarray): a boolean array representing where the update should occur in the state
-    """
-    
-    state.hillslope_depth = calculate_hillslope_depth(base_condition, state.hillslope_storage, state.hillslope_depth)
-
-calculate_hillslope_depth = ne.NumExpr(
-    'where('
-        'base_condition,'
-        'hillslope_storage,'
-        'hillslope_depth'
-    ')',
-    (('base_condition', np.bool), ('hillslope_storage', np.float64), ('hillslope_depth', np.float64))
-)
\ No newline at end of file
diff --git a/mosartwmpy/input/runoff.py b/mosartwmpy/input/runoff.py
index 75d45ef..ccef59c 100644
--- a/mosartwmpy/input/runoff.py
+++ b/mosartwmpy/input/runoff.py
@@ -44,4 +44,4 @@ def load_runoff(state: State, grid: Grid, config: Benedict, current_time: dateti
             runoff[config.get('runoff.variables.wetland_runoff')].sel(sel, method='pad')
         ).flatten()
     
-    runoff.close()
\ No newline at end of file
+    runoff.close()
diff --git a/mosartwmpy/input_output_variables.py b/mosartwmpy/input_output_variables.py
index c44f0a8..758199d 100644
--- a/mosartwmpy/input_output_variables.py
+++ b/mosartwmpy/input_output_variables.py
@@ -94,7 +94,7 @@ def __init__(
             variable_type=FLOAT,
             variable_item_size=FLOAT64,
             variable_class=STATE,
-            units='m3 s-1'
+            units='m3'
         ),
         Variable(
             standard_name='deficit_water_amount',
@@ -102,6 +102,6 @@ def __init__(
             variable_type=FLOAT,
             variable_item_size=FLOAT64,
             variable_class=STATE,
-            units='m3 s-1'
+            units='m3'
         ),
     ]
diff --git a/mosartwmpy/main_channel/irrigation.py b/mosartwmpy/main_channel/irrigation.py
index 460b06f..2213648 100644
--- a/mosartwmpy/main_channel/irrigation.py
+++ b/mosartwmpy/main_channel/irrigation.py
@@ -1,114 +1,79 @@
-import numpy as np
-import numexpr as ne
+import numba as nb
 
-from mosartwmpy.config.parameters import Parameters
-from mosartwmpy.grid.grid import Grid
-from mosartwmpy.state.state import State
 from mosartwmpy.main_channel.state import update_main_channel_state
-from mosartwmpy.utilities.timing import timing
 
-# @timing
-def main_channel_irrigation(state: State, grid: Grid, parameters: Parameters) -> None:
-    """Tracks the supply of water from the main river channel extracted into the grid cells.
 
-    Args:
-        state (State): the current model state; will be mutated
-        grid (Grid): the model grid
-        parameters (Parameters): the model parameters
-    """
-    
-    depth_condition = calculate_depth_condition(grid.mosart_mask, state.euler_mask, state.tracer, parameters.LIQUID_TRACER, state.channel_depth, parameters.irrigation_extraction_condition)
-    
-    tiny_condition = calculate_tiny_condition(state.channel_storage, parameters.tinier_value, state.grid_cell_unmet_demand, grid.channel_length)
-    
-    flow_volume = np.empty_like(state.channel_storage)
-    np.copyto(flow_volume, state.channel_storage)
-    
-    volume_condition = calculate_volume_condition(parameters.irrigation_extraction_maximum_fraction, flow_volume, state.grid_cell_unmet_demand)
-    
-    state.grid_cell_supply = calculate_reservoir_supply(depth_condition, tiny_condition, volume_condition, state.grid_cell_supply, state.grid_cell_unmet_demand, parameters.irrigation_extraction_maximum_fraction, flow_volume)
-    
-    flow_volume = calculate_flow_volume(depth_condition, tiny_condition, volume_condition, flow_volume, state.grid_cell_unmet_demand)
-    
-    state.grid_cell_unmet_demand = calculate_reservoir_demand(depth_condition, tiny_condition, volume_condition, state.grid_cell_unmet_demand, parameters.irrigation_extraction_maximum_fraction, flow_volume)
-    
-    flow_volume = update_flow_volume(depth_condition, tiny_condition, volume_condition, parameters.irrigation_extraction_maximum_fraction, flow_volume)
-    
-    state.channel_storage = calculate_channel_storage(depth_condition, tiny_condition, flow_volume, state.channel_storage)
-    
-    # TODO ? fortran mosart appears to do some more math with temp_erout and TRunoff%erout that looks to me to always be zero
-    
-    update_main_channel_state(state, grid, parameters, depth_condition)
-
-calculate_depth_condition = ne.NumExpr(
-    '(mosart_mask > 0) &'
-    'euler_mask &'
-    '(tracer == LIQUID_TRACER) &'
-    '(channel_depth >= irrigation_extraction_condition)',
-    (('mosart_mask', np.int64), ('euler_mask', np.bool), ('tracer', np.int64), ('LIQUID_TRACER', np.int64), ('channel_depth', np.float64), ('irrigation_extraction_condition', np.float64))
+@nb.jit(
+    "void("
+        "int64, int64[:], float64[:], float64[:], float64[:], float64[:], boolean[:],"
+        "float64[:], float64[:], float64[:], float64[:], float64[:], float64[:], float64[:],"
+        "float64, float64, float64, float64, float64, float64"
+    ")",
+    parallel=True,
+    nopython=True,
+    nogil=True,
+    cache=True,
 )
+def main_channel_irrigation(
+    n,
+    mosart_mask,
+    channel_length,
+    grid_channel_depth,
+    channel_width,
+    channel_floodplain_width,
+    euler_mask,
+    channel_depth,
+    channel_storage,
+    grid_cell_unmet_demand,
+    grid_cell_supply,
+    channel_cross_section_area,
+    channel_wetness_perimeter,
+    channel_hydraulic_radii,
+    tiny_value,
+    tinier_value,
+    slope_1_def,
+    inverse_sin_atan_slope_1_def,
+    irrigation_extraction_parameter,
+    irrigation_extraction_maximum_fraction,
+):
+    """Tracks the supply of water from the main river channel extracted into the grid cells."""
 
-calculate_tiny_condition = ne.NumExpr(
-    '(channel_storage > tinier_value) &'
-    '(reservoir_demand > tinier_value) &'
-    '(channel_length > tinier_value)',
-    (('channel_storage', np.float64), ('tinier_value', np.float64), ('reservoir_demand', np.float64), ('channel_length', np.float64))
-)
+    for i in nb.prange(n):
 
-calculate_volume_condition = ne.NumExpr(
-    'irrigation_extraction_maximum_fraction * flow_volume >= reservoir_demand',
-    (('irrigation_extraction_maximum_fraction', np.float64), ('flow_volume', np.float64), ('reservoir_demand', np.float64))
-)
+        # is the channel deep enough to extract from
+        depth_condition = (mosart_mask[i] > 0) and euler_mask[i] and (channel_depth[i] >= irrigation_extraction_parameter)
 
-calculate_reservoir_supply = ne.NumExpr(
-    'where('
-        'depth_condition & tiny_condition,'
-        'where('
-            'volume_condition,'
-            'reservoir_supply + reservoir_demand,'
-            'reservoir_supply + irrigation_extraction_maximum_fraction * flow_volume'
-        '),'
-        'reservoir_supply'
-    ')',
-    (('depth_condition', np.bool), ('tiny_condition', np.bool), ('volume_condition', np.bool), ('reservoir_supply', np.float64), ('reservoir_demand', np.float64), ('irrigation_extraction_maximum_fraction', np.float64), ('flow_volume', np.float64))
-)
+        # skip if certain parameters are very close to zero
+        tiny_condition = (channel_storage[i] > tinier_value) and (grid_cell_unmet_demand[i] > tinier_value) and (channel_length[i] > tinier_value)
 
-calculate_flow_volume = ne.NumExpr(
-    'where('
-        'depth_condition & tiny_condition & volume_condition,'
-        'flow_volume - reservoir_demand,'
-        'flow_volume'
-    ')',
-    (('depth_condition', np.bool), ('tiny_condition', np.bool), ('volume_condition', np.bool), ('flow_volume', np.float64), ('reservoir_demand', np.float64))
-)
+        flow_volume = 1.0 * channel_storage[i]
 
-calculate_reservoir_demand = ne.NumExpr(
-    'where('
-        'depth_condition & tiny_condition,'
-        'where('
-            'volume_condition,'
-            '0,'
-            'reservoir_demand - irrigation_extraction_maximum_fraction * flow_volume'
-        '),'
-        'reservoir_demand'
-    ')',
-    (('depth_condition', np.bool), ('tiny_condition', np.bool), ('volume_condition', np.bool), ('reservoir_demand', np.float64), ('irrigation_extraction_maximum_fraction', np.float64), ('flow_volume', np.float64))
-)
+        # is there enough water to allow full extraction
+        volume_condition = (irrigation_extraction_maximum_fraction * flow_volume) >= grid_cell_unmet_demand[i]
 
-update_flow_volume = ne.NumExpr(
-    'where('
-        'depth_condition & tiny_condition & (~volume_condition),'
-        '(1.0 - irrigation_extraction_maximum_fraction) * flow_volume,'
-        'flow_volume'
-    ')',
-    (('depth_condition', np.bool), ('tiny_condition', np.bool), ('volume_condition', np.bool), ('irrigation_extraction_maximum_fraction', np.float64), ('flow_volume', np.float64))
-)
+        if depth_condition and tiny_condition:
+            if volume_condition:
+                grid_cell_supply[i] = grid_cell_supply[i] + grid_cell_unmet_demand[i]
+                flow_volume = flow_volume - grid_cell_unmet_demand[i]
+                grid_cell_unmet_demand[i] = 0.0
+            else:
+                grid_cell_supply[i] = grid_cell_supply[i] + irrigation_extraction_maximum_fraction * flow_volume
+                grid_cell_unmet_demand[i] = grid_cell_unmet_demand[i] - irrigation_extraction_maximum_fraction * flow_volume
+                flow_volume = flow_volume - irrigation_extraction_maximum_fraction * flow_volume
 
-calculate_channel_storage = ne.NumExpr(
-    'where('
-        'depth_condition & tiny_condition,'
-        'flow_volume,'
-        'channel_storage'
-    ')',
-    (('depth_condition', np.bool), ('tiny_condition', np.bool), ('flow_volume', np.float64), ('channel_storage', np.float64))
-)
\ No newline at end of file
+            channel_storage[i] = 1.0 * flow_volume
+            update_main_channel_state(
+                i,
+                channel_length,
+                grid_channel_depth,
+                channel_width,
+                channel_floodplain_width,
+                channel_storage,
+                channel_cross_section_area,
+                channel_depth,
+                channel_wetness_perimeter,
+                channel_hydraulic_radii,
+                tiny_value,
+                slope_1_def,
+                inverse_sin_atan_slope_1_def
+            )
diff --git a/mosartwmpy/main_channel/kinematic_wave.py b/mosartwmpy/main_channel/kinematic_wave.py
index b86397b..8d38e43 100644
--- a/mosartwmpy/main_channel/kinematic_wave.py
+++ b/mosartwmpy/main_channel/kinematic_wave.py
@@ -1,131 +1,67 @@
-import numpy as np
-import numexpr as ne
+import numba as nb
 
-from mosartwmpy.config.parameters import Parameters
-from mosartwmpy.grid.grid import Grid
-from mosartwmpy.state.state import State
 
-def kinematic_wave_routing(state: State, grid: Grid, parameters: Parameters, delta_t: float, base_condition: np.ndarray) -> None:
-    """Tracks the storage and flow of water in the main channel using the kinematic wave routing method.
-
-    Args:
-        state (State): the current model state; will be mutated
-        grid (Grid): the model grid
-        parameters (Parameters): the model parameters
-        delta_t (float): the timestep for this subcycle (overall timestep / subcycles)
-        base_condition (np.ndarray): a boolean array representing where the update should occur in the state
-    """
+@nb.jit(
+    "void("
+        "int64, float64, float64[:], float64[:], float64[:], float64[:], float64[:], float64[:],"
+        "float64[:], float64[:], float64[:], float64[:], float64[:], float64[:], float64[:],"
+        "float64[:], float64[:], float64[:], float64[:], float64, float64"
+    ")",
+    nopython=True,
+    nogil=True,
+    cache=True,
+)
+def kinematic_wave_routing(
+    i,
+    delta_t,
+    channel_inflow_upstream,
+    channel_outflow_sum_upstream_instant,
+    channel_length,
+    channel_hydraulic_radii,
+    channel_flow_velocity,
+    channel_slope,
+    channel_manning,
+    total_drainage_area_single,
+    channel_width,
+    channel_outflow_downstream,
+    channel_cross_section_area,
+    channel_lateral_flow_hillslope,
+    channel_storage,
+    hillslope_wetland_runoff,
+    area,
+    drainage_fraction,
+    channel_delta_storage,
+    tiny_value,
+    kinematic_wave_parameter,
+):
+    """Tracks the storage and flow of water in the main channel using the kinematic wave routing method."""
     
     # estimation of inflow
-    state.channel_inflow_upstream = calculate_channel_inflow_upstream(base_condition, state.channel_outflow_sum_upstream_instant, state.channel_inflow_upstream)
+    channel_inflow_upstream[i] = -channel_outflow_sum_upstream_instant[i]
     
     # estimation of outflow
-    state.channel_flow_velocity = calculate_channel_flow_velocity(base_condition, grid.channel_length, state.channel_hydraulic_radii, grid.channel_slope, grid.channel_manning, state.channel_flow_velocity)
-    condition = calculate_kinematic_wave_condition(grid.channel_length, grid.total_drainage_area_single, grid.channel_width, grid.channel_length, parameters.kinematic_wave_condition)
-    state.channel_outflow_downstream = calculate_base_channel_outflow_downstream(base_condition, condition, state.channel_flow_velocity, state.channel_cross_section_area, state.channel_inflow_upstream, state.channel_lateral_flow_hillslope, state.channel_outflow_downstream)
-    condition = calculate_flow_condition(base_condition, condition, state.channel_outflow_downstream, parameters.tiny_value, state.channel_storage, state.channel_lateral_flow_hillslope, state.channel_inflow_upstream, delta_t)
-    state.channel_outflow_downstream = calculate_channel_outflow_downstream(condition, state.channel_lateral_flow_hillslope, state.channel_inflow_upstream, state.channel_storage, delta_t, state.channel_outflow_downstream)
-    state.channel_flow_velocity = update_channel_flow_velocity(condition, state.channel_cross_section_area, state.channel_outflow_downstream, state.channel_flow_velocity)
-    
-    # calculate change in storage, but first round small runoff to zero
-    tmp_delta_runoff = calculate_tmp_delta_runoff(base_condition, state.hillslope_wetland_runoff, grid.area, grid.drainage_fraction)
-    tmp_delta_runoff = update_tmp_delta_runoff(base_condition, tmp_delta_runoff, parameters.tiny_value)
-    state.channel_delta_storage = calculate_channel_delta_storage(base_condition, state.channel_lateral_flow_hillslope, state.channel_inflow_upstream, state.channel_outflow_downstream, tmp_delta_runoff)
-
+    if (channel_length[i] > 0) and (channel_hydraulic_radii[i] > 0):
+        channel_flow_velocity[i] = (channel_hydraulic_radii[i] ** (2.0/3.0)) * (channel_slope[i] ** (1.0/2.0)) / channel_manning[i]
+    else:
+        channel_flow_velocity[i] = 0.0
 
-calculate_channel_inflow_upstream = ne.NumExpr(
-    'where('
-        'base_condition,'
-        '-channel_outflow_sum_upstream_instant,'
-        'channel_inflow_upstream'
-    ')',
-    (('base_condition', np.bool), ('channel_outflow_sum_upstream_instant', np.float64), ('channel_inflow_upstream', np.float64))
-)
+    kinematic_wave_condition = (channel_length[i] > 0) and ((total_drainage_area_single[i] / channel_width[i] / channel_length[i]) <= kinematic_wave_parameter)
 
-calculate_channel_flow_velocity = ne.NumExpr(
-    'where('
-        'base_condition,'
-        'where('
-            '(channel_length > 0) & (channel_hydraulic_radii > 0),'
-            '(channel_hydraulic_radii ** (2/3)) * sqrt(channel_slope) / channel_manning,'
-            '0'
-        '),'
-        'channel_flow_velocity'
-    ')',
-    (('base_condition', np.bool), ('channel_length', np.float64), ('channel_hydraulic_radii', np.float64), ('channel_slope', np.float64), ('channel_manning', np.float64), ('channel_flow_velocity', np.float64))
-)
+    if kinematic_wave_condition:
+        channel_outflow_downstream[i] = -channel_flow_velocity[i] * channel_cross_section_area[i]
+    else:
+        channel_outflow_downstream[i] = -channel_inflow_upstream[i] - channel_lateral_flow_hillslope[i]
 
-calculate_kinematic_wave_condition = ne.NumExpr(
-    '(channel_length > 0) & ((total_drainage_area_single / channel_width / channel_length) <= kinematic_wave_condition)',
-    (('channel_length', np.float64), ('total_drainage_area_single', np.float64), ('channel_width', np.float64), ('channel_length', np.float64), ('kinematic_wave_condition', np.float64))
-)
+    flow_condition = kinematic_wave_condition and (-channel_outflow_downstream[i] > tiny_value) and ((channel_storage[i] + (channel_lateral_flow_hillslope[i] + channel_inflow_upstream[i] + channel_outflow_downstream[i]) * delta_t) < tiny_value)
 
-calculate_base_channel_outflow_downstream = ne.NumExpr(
-    'where('
-        'base_condition,'
-        'where('
-            'kinematic_wave_condition,'
-            '-channel_flow_velocity * channel_cross_section_area,'
-            '-channel_inflow_upstream - channel_lateral_flow_hillslope'
-        '),'
-        'channel_outflow_downstream'
-    ')',
-    (('base_condition', np.bool), ('kinematic_wave_condition', np.bool), ('channel_flow_velocity', np.float64), ('channel_cross_section_area', np.float64), ('channel_inflow_upstream', np.float64), ('channel_lateral_flow_hillslope', np.float64), ('channel_outflow_downstream', np.float64))
-)
+    if flow_condition:
+        channel_outflow_downstream[i] = -(channel_lateral_flow_hillslope[i] + channel_inflow_upstream[i] + channel_storage[i] / delta_t)
 
-calculate_flow_condition = ne.NumExpr(
-    'base_condition &'
-    'kinematic_wave_condition &'
-    '(-channel_outflow_downstream > tiny_value) &'
-    '((channel_storage + (channel_lateral_flow_hillslope + channel_inflow_upstream + channel_outflow_downstream) * delta_t) < tiny_value)',
-    (('base_condition', np.bool), ('kinematic_wave_condition', np.bool), ('channel_outflow_downstream', np.float64), ('tiny_value', np.float64), ('channel_storage', np.float64), ('channel_lateral_flow_hillslope', np.float64), ('channel_inflow_upstream', np.float64), ('delta_t', np.float64))
-)
+    if flow_condition and (channel_cross_section_area[i] > 0):
+        channel_flow_velocity[i] = -channel_outflow_downstream[i] / channel_cross_section_area[i]
 
-calculate_channel_outflow_downstream = ne.NumExpr(
-    'where('
-        'flow_condition,'
-        '-(channel_lateral_flow_hillslope + channel_inflow_upstream + channel_storage / delta_t),'
-        'channel_outflow_downstream'
-    ')',
-    (('flow_condition', np.bool), ('channel_lateral_flow_hillslope', np.float64), ('channel_inflow_upstream', np.float64), ('channel_storage', np.float64), ('delta_t', np.float64), ('channel_outflow_downstream', np.float64))
-)
-
-update_channel_flow_velocity = ne.NumExpr(
-    'where('
-        'flow_condition & (channel_cross_section_area > 0),'
-        '-channel_outflow_downstream / channel_cross_section_area,'
-        'channel_flow_velocity'
-    ')',
-    (('flow_condition', np.bool), ('channel_cross_section_area', np.float64), ('channel_outflow_downstream', np.float64), ('channel_flow_velocity', np.float64))
-)
-
-calculate_tmp_delta_runoff = ne.NumExpr(
-    'where('
-        'base_condition,'
-        'hillslope_wetland_runoff * area * drainage_fraction,'
-        '0'
-    ')',
-    (('base_condition', np.bool), ('hillslope_wetland_runoff', np.float64), ('area', np.float64), ('drainage_fraction', np.float64))
-)
-
-update_tmp_delta_runoff = ne.NumExpr(
-    'where('
-        'base_condition,'
-        'where('
-            'abs(tmp_delta_runoff) <= tiny_value,'
-            '0,'
-            'tmp_delta_runoff'
-        '),'
-        'tmp_delta_runoff'
-    ')',
-    (('base_condition', np.bool), ('tmp_delta_runoff', np.float64), ('tiny_value', np.float64))
-)
-
-calculate_channel_delta_storage = ne.NumExpr(
-    'where('
-        'base_condition,'
-        'channel_lateral_flow_hillslope + channel_inflow_upstream + channel_outflow_downstream + tmp_delta_runoff,'
-        '0'
-    ')',
-    (('base_condition', np.bool), ('channel_lateral_flow_hillslope', np.float64), ('channel_inflow_upstream', np.float64), ('channel_outflow_downstream', np.float64), ('tmp_delta_runoff', np.float64))
-)
\ No newline at end of file
+    # calculate change in storage, but first round small runoff to zero
+    delta_runoff = hillslope_wetland_runoff[i] * area[i] * drainage_fraction[i]
+    if abs(delta_runoff) <= tiny_value:
+        delta_runoff = 0
+    channel_delta_storage[i] = channel_lateral_flow_hillslope[i] + channel_inflow_upstream[i] + channel_outflow_downstream[i] + delta_runoff
diff --git a/mosartwmpy/main_channel/routing.py b/mosartwmpy/main_channel/routing.py
index 3b13635..17ea9af 100644
--- a/mosartwmpy/main_channel/routing.py
+++ b/mosartwmpy/main_channel/routing.py
@@ -1,88 +1,120 @@
-import numpy as np
-import numexpr as ne
+import numba as nb
 
-from benedict.dicts import benedict as Benedict
-
-from mosartwmpy.config.parameters import Parameters
-from mosartwmpy.grid.grid import Grid
-from mosartwmpy.state.state import State
 from mosartwmpy.main_channel.kinematic_wave import kinematic_wave_routing
 from mosartwmpy.main_channel.state import update_main_channel_state
-from mosartwmpy.utilities.timing import timing
 
-# @timing
-def main_channel_routing(state: State, grid: Grid, parameters: Parameters, config: Benedict, delta_t: float) -> None:
-    """Tracks the storage and flow of water in the main river channels.
 
-    Args:
-        state (State): the current model state; will be mutated
-        grid (Grid): the model grid
-        parameters (Parameters): the model parameters
-        config (Benedict): the model configuration
-        delta_t (float): the timestep for this subcycle (overall timestep / subcycles)
-    """
-    
-    tmp_outflow_downstream = 0.0 * state.zeros
-    local_delta_t = (delta_t / config.get('simulation.routing_iterations') / grid.iterations_main_channel)
-    
-    # step through max iterations, masking out the unnecessary cells each time
-    base_condition = (grid.mosart_mask > 0) & state.euler_mask
-    for _ in np.arange(np.nanmax(grid.iterations_main_channel)):
-        iteration_condition = base_condition & (grid.iterations_main_channel > _)
-    
-        # routing
-        routing_method = config.get('simulation.routing_method', 1)
-        if routing_method == 1:
-            kinematic_wave_routing(state, grid, parameters, local_delta_t, iteration_condition)
-        else:
-            raise Exception(f"Error - Routing method {routing_method} not implemented.")
-        
-        # update storage
-        state.channel_storage_previous_timestep = calculate_channel_storage_previous_timestep(iteration_condition, state.channel_storage, state.channel_storage_previous_timestep)
-        state.channel_storage = calculate_channel_storage(iteration_condition, state.channel_storage, state.channel_delta_storage, local_delta_t)
-        
-        # update channel state
-        update_main_channel_state(state, grid, parameters, iteration_condition)
-        
-        # update outflow tracking
-        tmp_outflow_downstream = calculate_tmp_outflow_downstream(iteration_condition, tmp_outflow_downstream, state.channel_outflow_downstream)
-    
-    # update outflow
-    state.channel_outflow_downstream = calculate_channel_outflow_downstream(base_condition, tmp_outflow_downstream, grid.iterations_main_channel, state.channel_outflow_downstream)
+@nb.jit(
+    "void("
+        "int64, float64, int64, int64,"
+        "int64[:], int64[:], float64[:], float64[:], float64[:], float64[:], float64[:], float64[:], float64[:],"
+        "float64[:], float64[:], boolean[:], float64[:], float64[:], float64[:], float64[:], float64[:], float64[:],"
+        "float64[:], float64[:], float64[:], float64[:], float64[:], float64[:], float64[:],"
+        "float64, float64, float64, float64"
+    ")",
+    parallel=True,
+    nopython=True,
+    nogil=True,
+    cache=True,
+)
+def main_channel_routing(
+    n,
+    delta_t,
+    routing_iterations,
+    max_iterations_main_channel,
+    iterations_main_channel,
+    mosart_mask,
+    channel_length,
+    channel_slope,
+    channel_manning,
+    total_drainage_area_single,
+    channel_width,
+    area,
+    drainage_fraction,
+    grid_channel_depth,
+    channel_floodplain_width,
+    euler_mask,
+    channel_inflow_upstream,
+    channel_outflow_sum_upstream_instant,
+    channel_hydraulic_radii,
+    channel_flow_velocity,
+    channel_outflow_downstream,
+    channel_cross_section_area,
+    channel_lateral_flow_hillslope,
+    channel_storage,
+    channel_storage_previous_timestep,
+    hillslope_wetland_runoff,
+    channel_delta_storage,
+    channel_depth,
+    channel_wetness_perimeter,
+    slope_1_def,
+    inverse_sin_atan_slope_1_def,
+    tiny_value,
+    kinematic_wave_parameter,
+):
+    """Tracks the storage and flow of water in the main river channels."""
 
+    for i in nb.prange(n):
 
-calculate_channel_storage_previous_timestep = ne.NumExpr(
-    'where('
-        'iteration_condition,'
-        'channel_storage,'
-        'channel_storage_previous_timestep'
-    ')',
-    (('iteration_condition', np.bool), ('channel_storage', np.float64), ('channel_storage_previous_timestep', np.float64))
-)
+        if ~euler_mask[i] or ~(mosart_mask[i] > 0):
+            continue
 
-calculate_channel_storage = ne.NumExpr(
-    'where('
-        'iteration_condition,'
-        'channel_storage + channel_delta_storage * local_delta_t,'
-        'channel_storage'
-    ')',
-    (('iteration_condition', np.bool), ('channel_storage', np.float64), ('channel_delta_storage', np.float64), ('local_delta_t', np.float64))
-)
+        local_delta_t = (delta_t / routing_iterations) / iterations_main_channel[i]
+        outflow = 0.0
 
-calculate_tmp_outflow_downstream = ne.NumExpr(
-    'where('
-        'iteration_condition,'
-        'tmp_outflow_downstream + channel_outflow_downstream,'
-        'tmp_outflow_downstream'
-    ')',
-    (('iteration_condition', np.bool), ('tmp_outflow_downstream', np.float64), ('channel_outflow_downstream', np.float64))
-)
+        # step through max iterations
+        for _ in nb.prange(max_iterations_main_channel):
+            if ~(iterations_main_channel[i] > _):
+                continue
+
+            # route the water
+            kinematic_wave_routing(
+                i,
+                local_delta_t,
+                channel_inflow_upstream,
+                channel_outflow_sum_upstream_instant,
+                channel_length,
+                channel_hydraulic_radii,
+                channel_flow_velocity,
+                channel_slope,
+                channel_manning,
+                total_drainage_area_single,
+                channel_width,
+                channel_outflow_downstream,
+                channel_cross_section_area,
+                channel_lateral_flow_hillslope,
+                channel_storage,
+                hillslope_wetland_runoff,
+                area,
+                drainage_fraction,
+                channel_delta_storage,
+                tiny_value,
+                kinematic_wave_parameter,
+            )
+
+            # update storage
+            channel_storage_previous_timestep[i] = 1.0 * channel_storage[i]
+            channel_storage[i] = channel_storage[i] + channel_delta_storage[i] * local_delta_t
+
+            # update channel state
+            update_main_channel_state(
+                i,
+                channel_length,
+                grid_channel_depth,
+                channel_width,
+                channel_floodplain_width,
+                channel_storage,
+                channel_cross_section_area,
+                channel_depth,
+                channel_wetness_perimeter,
+                channel_hydraulic_radii,
+                tiny_value,
+                slope_1_def,
+                inverse_sin_atan_slope_1_def
+            )
+
+            # track the outflow
+            outflow = outflow + channel_outflow_downstream[i]
 
-calculate_channel_outflow_downstream = ne.NumExpr(
-    'where('
-        'base_condition,'
-        'tmp_outflow_downstream / iterations_main_channel,'
-        'channel_outflow_downstream'
-    ')',
-    (('base_condition', np.bool), ('tmp_outflow_downstream', np.float64), ('iterations_main_channel', np.float64), ('channel_outflow_downstream', np.float64))
-)
\ No newline at end of file
+        # update the final outflow
+        channel_outflow_downstream[i] = outflow / iterations_main_channel[i]
diff --git a/mosartwmpy/main_channel/state.py b/mosartwmpy/main_channel/state.py
index 488508f..5c16985 100644
--- a/mosartwmpy/main_channel/state.py
+++ b/mosartwmpy/main_channel/state.py
@@ -1,129 +1,76 @@
-import numpy as np
-import numexpr as ne
+import numba as nb
 
-from mosartwmpy.config.parameters import Parameters
-from mosartwmpy.grid.grid import Grid
-from mosartwmpy.state.state import State
 
-def update_main_channel_state(state: State, grid: Grid, parameters: Parameters, base_condition: np.ndarray) -> None:
-    """Updates the physical properties of the main river channel based on current state.
-
-    Args:
-        state (State): the current model state; will be mutated
-        grid (Grid): the model grid
-        parameters (Parameters): the model parameters
-        base_condition (np.ndarray): a boolean array representing where the update should occur in the state
-    """
-    
-    condition = calculate_storage_condition(grid.channel_length, state.channel_storage)
-    state.channel_cross_section_area = calculate_channel_cross_section_area(base_condition, condition, state.channel_storage, grid.channel_length, state.channel_cross_section_area)
-
-    # Function for estimating maximum water depth assuming rectangular channel and tropezoidal flood plain
+@nb.jit(
+    "void("
+        "int64, float64[:], float64[:], float64[:], float64[:], float64[:],"
+        "float64[:], float64[:], float64[:], float64[:], float64, float64, float64"
+    ")",
+    nopython=True,
+    nogil=True,
+    cache=True,
+)
+def update_main_channel_state(
+    i,
+    channel_length,
+    grid_channel_depth,
+    channel_width,
+    channel_floodplain_width,
+    channel_storage,
+    channel_cross_section_area,
+    channel_depth,
+    channel_wetness_perimeter,
+    channel_hydraulic_radii,
+    tiny_value,
+    slope_1_def,
+    inverse_sin_atan_slope_1_def,
+):
+    """Updates the physical properties of the main river channel based on current state."""
+
+    storage_condition = (channel_length[i] > 0.0) and (channel_storage[i] > 0.0)
+
+    if storage_condition:
+        channel_cross_section_area[i] = channel_storage[i] / channel_length[i]
+    else:
+        channel_cross_section_area[i] = 0.0
+
+    not_flooded = (channel_cross_section_area[i] - (grid_channel_depth[i] * channel_width[i])) <= tiny_value
+
+    delta_area = channel_cross_section_area[i] - grid_channel_depth[i] * channel_width[i] - (channel_width[i] + channel_floodplain_width[i]) * slope_1_def * ((channel_floodplain_width[i] - channel_width[i]) / 2.0) / 2.0
+
+    # Function for estimating maximum water depth assuming rectangular channel and trapezoidal flood plain
     # here assuming the channel cross-section consists of three parts, from bottom to up,
     # part 1 is a rectangular with bankfull depth (rdep) and bankfull width (rwid)
-    # part 2 is a tropezoidal, bottom width rwid and top width rwid0, height 0.1*((rwid0-rwid)/2), assuming slope is 0.1
-    # part 3 is a rectagular with the width rwid0
-    not_flooded = calculate_not_flooded_depth_condition(state.channel_cross_section_area, grid.grid_channel_depth, grid.channel_width, parameters.tiny_value)
-    delta_area = calculate_delta_area(state.channel_cross_section_area, grid.grid_channel_depth, grid.channel_width, grid.channel_floodplain_width, parameters.slope_1_def)
-    state.channel_depth = calculate_channel_depth(base_condition, condition, state.channel_cross_section_area, parameters.tiny_value, not_flooded, grid.channel_width, delta_area, grid.grid_channel_depth, parameters.slope_1_def, grid.channel_floodplain_width, state.channel_depth)
+    # part 2 is a trapezoidal, bottom width rwid and top width rwid0, height slope*((rwid0-rwid)/2)
+    # part 3 is a rectangular with the width rwid0
+    if storage_condition and (channel_cross_section_area[i] > tiny_value):
+        if not_flooded:
+            channel_depth[i] = channel_cross_section_area[i] / channel_width[i]
+        else:
+            if delta_area > tiny_value:
+                channel_depth[i] = grid_channel_depth[i] + slope_1_def * ((channel_floodplain_width[i] - channel_width[i]) / 2.0) + delta_area / channel_floodplain_width[i]
+            else:
+                channel_depth[i] = grid_channel_depth[i] + (-channel_width[i] + (((channel_width[i] ** 2.0) + 4.0 * (channel_cross_section_area[i] - grid_channel_depth[i] * channel_width[i]) / slope_1_def) ** (1.0/2.0))) * slope_1_def / 2.0
+    else:
+        channel_depth[i] = 0.0
+
+    not_flooded = (channel_depth[i] <= (grid_channel_depth[i] + tiny_value))
+
+    delta_depth = channel_depth[i] - grid_channel_depth[i] - ((channel_floodplain_width[i] - channel_width[i]) / 2.0) * slope_1_def
 
     # Function for estimating wetness perimeter based on same assumptions as above
-    not_flooded = calculate_not_flooded_wetness_condition(state.channel_depth, grid.grid_channel_depth, parameters.tiny_value)
-    delta_depth = calculate_delta_depth(state.channel_depth, grid.grid_channel_depth, grid.channel_floodplain_width, grid.channel_width, parameters.slope_1_def)
-    state.channel_wetness_perimeter = calculate_channel_wetness_perimeter(base_condition, condition, state.channel_depth, parameters.tiny_value, not_flooded, grid.channel_width, delta_depth, grid.grid_channel_depth, grid.channel_floodplain_width, parameters.slope_1_def, parameters.inverse_sin_atan_slope_1_def, state.channel_wetness_perimeter)
-
-    state.channel_hydraulic_radii = calculate_channel_hydraulic_radii(base_condition, condition, state.channel_wetness_perimeter, parameters.tiny_value, state.channel_cross_section_area, state.channel_hydraulic_radii)
-
-
-calculate_storage_condition = ne.NumExpr(
-    '(channel_length > 0) & (channel_storage > 0)',
-    (('channel_length', np.float64), ('channel_storage', np.float64))
-)
-
-calculate_channel_cross_section_area = ne.NumExpr(
-    'where('
-        'base_condition,'
-        'where('
-            'storage_condition,'
-            'channel_storage / channel_length,'
-            '0'
-        '),'
-        'channel_cross_section_area'
-    ')',
-    (('base_condition', np.bool), ('storage_condition', np.bool), ('channel_storage', np.float64), ('channel_length', np.float64), ('channel_cross_section_area', np.float64))
-)
-
-calculate_not_flooded_depth_condition = ne.NumExpr(
-    '(channel_cross_section_area - (grid_channel_depth * channel_width)) <= tiny_value',
-    (('channel_cross_section_area', np.float64), ('grid_channel_depth', np.float64), ('channel_width', np.float64), ('tiny_value', np.float64))
-)
-
-calculate_delta_area = ne.NumExpr(
-    'channel_cross_section_area - grid_channel_depth  * channel_width - (channel_width + channel_floodplain_width) * slope_1_def * ((channel_floodplain_width - channel_width) / 2.0) / 2.0',
-    (('channel_cross_section_area', np.float64), ('grid_channel_depth', np.float64), ('channel_width', np.float64), ('channel_floodplain_width', np.float64), ('slope_1_def', np.float64))
-)
-
-calculate_channel_depth = ne.NumExpr(
-    'where('
-        'base_condition,'
-        'where('
-            'storage_condition & (channel_cross_section_area > tiny_value),'
-            'where('
-                'not_flooded_condition,'
-                'channel_cross_section_area / channel_width,'
-                'where('
-                    'delta_area > tiny_value,'
-                    'grid_channel_depth + slope_1_def * ((channel_floodplain_width  - channel_width) / 2.0) + delta_area / channel_floodplain_width,'
-                    'grid_channel_depth + (-channel_width + (((channel_width ** 2) + 4 * (channel_cross_section_area  - grid_channel_depth * channel_width) / slope_1_def) ** (1/2))) * slope_1_def / 2.0'
-                ')'
-            '),'
-            '0'
-        '),'
-        'channel_depth'
-    ')',
-    (('base_condition', np.bool), ('storage_condition', np.bool), ('channel_cross_section_area', np.float64), ('tiny_value', np.float64), ('not_flooded_condition', np.bool), ('channel_width', np.float64), ('delta_area', np.float64), ('grid_channel_depth', np.float64), ('slope_1_def', np.float64), ('channel_floodplain_width', np.float64), ('channel_depth', np.float64))
-)
-
-calculate_not_flooded_wetness_condition = ne.NumExpr(
-    'channel_depth <= (grid_channel_depth + tiny_value)',
-    (('channel_depth', np.float64), ('grid_channel_depth', np.float64), ('tiny_value', np.float64))
-)
-
-calculate_delta_depth = ne.NumExpr(
-    'channel_depth - grid_channel_depth - ((channel_floodplain_width -  channel_width) / 2.0) * slope_1_def',
-    (('channel_depth', np.float64), ('grid_channel_depth', np.float64), ('channel_floodplain_width', np.float64), ('channel_width', np.float64), ('slope_1_def', np.float64))
-)
-
-calculate_channel_wetness_perimeter = ne.NumExpr(
-    'where('
-        'base_condition,'
-        'where('
-            'storage_condition & (channel_depth >= tiny_value),'
-            'where('
-                'not_flooded_condition,'
-                'channel_width + 2 * channel_depth,'
-                'where('
-                    'delta_depth > tiny_value,'
-                    'channel_width + 2 * (grid_channel_depth + ((channel_floodplain_width - channel_width) / 2.0) * slope_1_def * inverse_sin_atan_slope_1_def + delta_depth),'
-                    'channel_width + 2 * (grid_channel_depth + (channel_depth - grid_channel_depth) * inverse_sin_atan_slope_1_def)'
-                ')'
-            '),'
-            '0'
-        '),'
-        'channel_wetness_perimeter'
-    ')',
-    (('base_condition', np.bool), ('storage_condition', np.bool), ('channel_depth', np.float64), ('tiny_value', np.float64), ('not_flooded_condition', np.bool), ('channel_width', np.float64), ('delta_depth', np.float64), ('grid_channel_depth', np.float64), ('channel_floodplain_width', np.float64), ('slope_1_def', np.float64), ('inverse_sin_atan_slope_1_def', np.float64), ('channel_wetness_perimeter', np.float64))
-)
-
-calculate_channel_hydraulic_radii = ne.NumExpr(
-    'where('
-        'base_condition,'
-        'where('
-            'storage_condition & (channel_wetness_perimeter > tiny_value),'
-            'channel_cross_section_area / channel_wetness_perimeter,'
-            '0'
-        '),'
-        'channel_hydraulic_radii'
-    ')',
-    (('base_condition', np.bool), ('storage_condition', np.bool), ('channel_wetness_perimeter', np.float64), ('tiny_value', np.float64), ('channel_cross_section_area', np.float64), ('channel_hydraulic_radii', np.float64))
-)
\ No newline at end of file
+    if storage_condition and (channel_depth[i] >= tiny_value):
+        if not_flooded:
+            channel_wetness_perimeter[i] = channel_width[i] + 2.0 * channel_depth[i]
+        else:
+            if delta_depth > tiny_value:
+                channel_wetness_perimeter[i] = channel_width[i] + 2.0 * (grid_channel_depth[i] + ((channel_floodplain_width[i] - channel_width[i]) / 2.0) * slope_1_def * inverse_sin_atan_slope_1_def + delta_depth)
+            else:
+                channel_wetness_perimeter[i] = channel_width[i] + 2.0 * (grid_channel_depth[i] + (channel_depth[i] - grid_channel_depth[i]) * inverse_sin_atan_slope_1_def)
+    else:
+        channel_wetness_perimeter[i] = 0.0
+
+    if storage_condition and (channel_wetness_perimeter[i] > tiny_value):
+        channel_hydraulic_radii[i] = channel_cross_section_area[i] / channel_wetness_perimeter[i]
+    else:
+        channel_hydraulic_radii[i] = 0.0
diff --git a/mosartwmpy/model.py b/mosartwmpy/model.py
index 407223e..1c085a3 100644
--- a/mosartwmpy/model.py
+++ b/mosartwmpy/model.py
@@ -1,7 +1,7 @@
 import logging
 import sys
 
-import numexpr as ne
+import matplotlib.pyplot as plt
 import numpy as np
 import psutil
 import regex as re
@@ -10,11 +10,12 @@
 from benedict import benedict
 from bmipy import Bmi
 from datetime import datetime, time, timedelta
+from numba import get_num_threads, threading_layer
 from pathlib import Path
 from pathvalidate import sanitize_filename
 from timeit import default_timer as timer
 from typing import Tuple
-from xarray import open_dataset
+from xarray import DataArray, open_dataset
 
 from mosartwmpy.config.config import get_config
 from mosartwmpy.config.parameters import Parameters
@@ -22,7 +23,7 @@
 from mosartwmpy.input.runoff import load_runoff
 from mosartwmpy.input.demand import load_demand
 from mosartwmpy.input_output_variables import IO
-from mosartwmpy.output.output import initialize_output, update_output, write_restart
+from mosartwmpy.output.output import initialize_output, update_output
 from mosartwmpy.reservoirs.reservoirs import reservoir_release
 from mosartwmpy.state.state import State
 from mosartwmpy.update.update import update
@@ -57,12 +58,21 @@ def __init__(self):
         self.reservoir_streamflow_schedule = None
         self.reservoir_demand_schedule = None
         self.reservoir_prerelease_schedule = None
-        self.git_hash = None
-        self.git_untracked = None
 
     def __getitem__(self, item):
         return getattr(self, item)
 
+    def plot(self, variable: str):
+        """Display a colormap of a spatial variable at the current timestep."""
+        DataArray(
+            self.get_value_ptr(variable).reshape(self.get_grid_shape()),
+            dims=['latitude', 'longitude'],
+            coords={'latitude': self.get_grid_x(), 'longitude': self.get_grid_y()},
+            name=variable.replace('_', ' ').title(),
+            attrs={'units': self.get_var_units(variable)}
+        ).plot(robust=True, levels=16, cmap='winter_r')
+        plt.show()
+
     def initialize(self, config_file_path: str = None, grid: Grid = None, state: State = None) -> None:
         
         t = timer()
@@ -76,7 +86,9 @@ def initialize(self, config_file_path: str = None, grid: Grid = None, state: Sta
             self.name = sanitize_filename(self.config.get('simulation.name')).replace(" ", "_")
             # setup logging and output directories
             Path(f'{self.config.get("simulation.output_path")}/{self.name}/restart_files').mkdir(parents=True, exist_ok=True)
-            handlers = [logging.FileHandler(Path(f'{self.config.get("simulation.output_path")}/{self.name}/mosartwmpy.log'))]
+            handlers = []
+            if self.config.get('simulation.log_to_file'):
+                handlers.append(logging.FileHandler(Path(f'{self.config.get("simulation.output_path")}/{self.name}/mosartwmpy.log')))
             if self.config.get('simulation.log_to_std_out'):
                 h = logging.StreamHandler(sys.stdout)
                 h.setFormatter(logging.Formatter(""))
@@ -92,23 +104,14 @@ def initialize(self, config_file_path: str = None, grid: Grid = None, state: Sta
             self.config.to_yaml(filepath=f'{self.config.get("simulation.output_path")}/{self.name}/config.yaml')
             if config_file_path is None or config_file_path == '':
                 logging.info("No configuration file provided; initializing with all default values.")
-            try:
-                self.git_hash = subprocess.check_output(['git', 'describe', '--always']).strip().decode('utf-8')
-                self.git_untracked = subprocess.check_output(['git', 'diff', '--name-only']).strip().decode('utf-8').split('\n')
-                logging.debug(f'Version: {self.git_hash}')
-                if len(self.git_untracked) > 0:
-                    logging.debug(f'Uncommitted changes:')
-                    for u in self.git_untracked:
-                        logging.debug(f'  * {u}')
-            except:
-                pass
             # ensure that end date is after start date
             if self.config.get('simulation.end_date') < self.config.get('simulation.start_date'):
                 raise ValueError(f"Configured `end_date` {self.config.get('simulation.end_date')} is prior to configured `start_date` {self.config.get('simulation.start_date')}; please update and try again.")
             # detect available physical cores
             self.cores = psutil.cpu_count(logical=False)
             logging.debug(f'Cores: {self.cores}.')
-            ne.set_num_threads(self.cores)
+            logging.debug(f'Numba threads: {get_num_threads()}.')
+            logging.debug(f'Numba threading layer: {threading_layer()}')
         except Exception as e:
             logging.exception('Failed to configure model; see below for stacktrace.')
             raise e
@@ -126,6 +129,7 @@ def initialize(self, config_file_path: str = None, grid: Grid = None, state: Sta
         # load restart file or initialize state
         if state is not None:
             self.state = state
+            self.current_time = datetime.combine(self.config.get('simulation.start_date'), time.min)
         else:
             try:
                 # restart file
@@ -244,10 +248,6 @@ def finalize(self) -> None:
         logging.shutdown()
         return
 
-    def download_data(self, *args, **kwargs) -> None:
-        """Downloads data related to the model."""
-        download_data(*args, **kwargs)
-
     def get_component_name(self) -> str:
         return f'mosartwmpy ({self.git_hash})'
 
diff --git a/mosartwmpy/output/output.py b/mosartwmpy/output/output.py
index 623becc..903af74 100644
--- a/mosartwmpy/output/output.py
+++ b/mosartwmpy/output/output.py
@@ -9,6 +9,7 @@
 
 from mosartwmpy.utilities.timing import timing
 
+
 def initialize_output(self):
     """Initializes the output buffer."""
     
@@ -22,6 +23,7 @@ def initialize_output(self):
             else:
                 self.output_buffer = self.output_buffer.join(pd.DataFrame(self.state.zeros, columns=[output.get('name')]))
 
+
 # @timing
 def update_output(self):
     """Updates the output buffer based on the current state, and averages and writes to file when appropriate."""
@@ -44,6 +46,7 @@ def update_output(self):
     # check if restart file if need
     check_restart(self)
 
+
 def write_output(self):
     """Writes the output buffer and requested grid variables to a netcdf file."""
     # TODO only daily resolution is currently supported - need to support arbitrary resolutions
@@ -136,6 +139,7 @@ def write_output(self):
     frame = frame.rio.write_crs(4326)
     frame.to_netcdf(filename, unlimited_dims=['time'])
 
+
 def check_restart(self):
     """Checks if a restart file is needed based on the current simulation time."""
     
@@ -156,6 +160,7 @@ def check_restart(self):
     if is_needed:
         write_restart(self)
 
+
 def write_restart(self):
     """Writes the state to a netcdf file, with the current simulation time in the file name."""
 
diff --git a/mosartwmpy/reservoirs/grid.py b/mosartwmpy/reservoirs/grid.py
index 8dc39f6..7fb7835 100644
--- a/mosartwmpy/reservoirs/grid.py
+++ b/mosartwmpy/reservoirs/grid.py
@@ -1,12 +1,14 @@
-import logging
 import numpy as np
 import pandas as pd
 
+from numba.core import types
+from numba.typed import Dict
 from xarray import concat, open_dataset, Dataset
 from benedict.dicts import benedict as Benedict
 
 from mosartwmpy.config.parameters import Parameters
 
+
 def load_reservoirs(self, config: Benedict, parameters: Parameters) -> None:
     """Loads the reservoir information from file onto the grid.
 
@@ -15,8 +17,6 @@ def load_reservoirs(self, config: Benedict, parameters: Parameters) -> None:
         parameters (Parameters): the model parameters
     """
 
-    logging.debug('Loading reservoir file.')
-
     # reservoir parameter file
     reservoirs = open_dataset(config.get('water_management.reservoirs.path'))
 
@@ -48,14 +48,16 @@ def load_reservoirs(self, config: Benedict, parameters: Parameters) -> None:
     # set to integer
     self.reservoir_to_grid_mapping = self.reservoir_to_grid_mapping.astype(int)
 
-    # count of the number of reservoirs that can supply each grid cell
-    self.reservoir_count = np.array(pd.DataFrame(self.id).join(
-        self.reservoir_to_grid_mapping.groupby('grid_cell_id').count().rename(columns={'reservoir_id': 'reservoir_count'}),
-        how='left'
-    ).reservoir_count)
+    # create a numba typed dict with key = <grid cell id> and value = <list of reservoir_ids that feed the cell>
+    self.reservoir_to_grid_map = Dict.empty(
+        key_type=types.int64,
+        value_type=types.int64[:],
+    )
+    for grid_cell_id, group in self.reservoir_to_grid_mapping.groupby('grid_cell_id'):
+        self.reservoir_to_grid_map[grid_cell_id] = group.reservoir_id.values
 
     # index by grid cell
-    self.reservoir_to_grid_mapping = self.reservoir_to_grid_mapping.set_index('grid_cell_id')
+    self.reservoir_to_grid_mapping = self.reservoir_to_grid_mapping.set_index('grid_cell_id').sort_index()
 
     # prepare the month based reservoir schedules mapped to the domain
     prepare_reservoir_schedule(self, config, parameters, reservoirs)
@@ -130,4 +132,4 @@ def prepare_reservoir_schedule(self, config: Benedict, parameters: Parameters, r
             flow_avg + self.reservoir_demand_schedule - demand_avg
         )
     )
-    self.reservoir_prerelease_schedule = prerelease
\ No newline at end of file
+    self.reservoir_prerelease_schedule = prerelease
diff --git a/mosartwmpy/reservoirs/regulation.py b/mosartwmpy/reservoirs/regulation.py
index f936c1c..b9f5701 100644
--- a/mosartwmpy/reservoirs/regulation.py
+++ b/mosartwmpy/reservoirs/regulation.py
@@ -1,105 +1,104 @@
+import numba as nb
 import numpy as np
-import numexpr as ne
-import pandas as pd
-
-from benedict.dicts import benedict as Benedict
-
-from mosartwmpy.config.parameters import Parameters
-from mosartwmpy.grid.grid import Grid
-from mosartwmpy.state.state import State
-from mosartwmpy.utilities.timing import timing
-
-# @timing
-def regulation(state: State, grid: Grid, parameters: Parameters, delta_t: float) -> None:
-    """Regulates the flow across the reservoirs.
-
-    Args:
-        state (State): the current model state; will be mutated
-        grid (Grid): the model grid
-        parameters (Parameters): the model parameters
-        delta_t (float): the timestep for this routing iteration (overall timestep / subcycles / routing iterations)
-    """
-    
-    base_condition = (
-        (grid.mosart_mask > 0) &
-        state.euler_mask &
-        (state.tracer == parameters.LIQUID_TRACER) &
-        np.isfinite(grid.reservoir_id)
-    )
-    
-    flow_volume = -state.channel_outflow_downstream * delta_t
-    
-    flow_reservoir = state.reservoir_release * delta_t
-    
-    evaporation = 1e6 * state.reservoir_potential_evaporation * delta_t * grid.reservoir_surface_area
-    
-    minimum_flow = parameters.reservoir_runoff_capacity_condition * state.reservoir_streamflow * delta_t
-    minimum_storage = parameters.reservoir_runoff_capacity_condition * grid.reservoir_storage_capacity
-    maximum_storage = 1 * grid.reservoir_storage_capacity
-    
-    condition_max = flow_volume + state.reservoir_storage - flow_reservoir - evaporation >= maximum_storage
-    condition_min = flow_volume + state.reservoir_storage - flow_reservoir - evaporation < minimum_storage
-    condition_min_one = flow_reservoir <= flow_volume - evaporation
-    condition_min_two = flow_volume - evaporation >= minimum_flow
-    
-    flow_reservoir = np.where(
-        condition_max,
-        flow_volume + state.reservoir_storage - maximum_storage - evaporation,
-        np.where(
-            condition_min,
-            np.where(
-                condition_min_one,
-                flow_reservoir,
-                np.where(
-                    condition_min_two,
-                    flow_volume - evaporation,
-                    flow_volume
-                )
-            ),
-            flow_reservoir
-        )
-    )
-    
-    state.reservoir_storage = np.where(
-        base_condition,
-        np.where(
-            condition_max,
-            maximum_storage,
-            np.where(
-                condition_min,
-                np.where(
-                    condition_min_one,
-                    state.reservoir_storage + flow_volume - flow_reservoir - evaporation,
-                    np.where(
-                        condition_min_two,
-                        state.reservoir_storage,
-                        np.maximum(0, state.reservoir_storage - flow_reservoir + flow_volume - evaporation)
-                    )
-                ),
-                state.reservoir_storage + flow_volume - flow_reservoir - evaporation
-            )
-        ),
-        state.reservoir_storage
-    )
-    
-    state.channel_outflow_downstream = np.where(
-        base_condition,
-        -flow_reservoir / delta_t,
-        state.channel_outflow_downstream
-    )
-
-# @timing
-def extraction_regulated_flow(state: State, grid: Grid, parameters: Parameters, config: Benedict, delta_t: float) -> None:
-    """Tracks the supply of water extracted from the reservoirs to fulfill demand from dependent grid cells.
-
-    Args:
-        state (State): the current model state; will be mutated
-        grid (Grid): the model grid
-        parameters (Parameters): the model parameters
-        config (Benedict): the model configuration
-        delta_t (float): the timestep for this subcycle (overall timestep / subcycles)
-    """
-    
+
+from numba.typed import List
+
+
+@nb.jit(
+    "void("
+        "int64, float64,"
+        "int64[:], float64[:], float64[:], float64[:],"
+        "boolean[:], float64[:], float64[:], float64[:], float64[:], float64[:], float64"
+    ")",
+    nopython=True,
+    nogil=True,
+    cache=True,
+)
+def regulation(
+    n,
+    delta_t,
+    mosart_mask,
+    reservoir_id,
+    reservoir_surface_area,
+    reservoir_storage_capacity,
+    euler_mask,
+    channel_outflow_downstream,
+    reservoir_release,
+    reservoir_potential_evaporation,
+    reservoir_streamflow,
+    reservoir_storage,
+    reservoir_runoff_capacity_parameter,
+):
+    """Regulates the flow across the reservoirs."""
+
+    for i in nb.prange(n):
+
+        if euler_mask[i] and (mosart_mask[i] > 0) and np.isfinite(reservoir_id[i]):
+
+            flow_volume = -channel_outflow_downstream[i] * delta_t
+
+            flow_reservoir = reservoir_release[i] * delta_t
+
+            evaporation = 1e6 * reservoir_potential_evaporation[i] * delta_t * reservoir_surface_area[i]
+
+            minimum_flow = reservoir_runoff_capacity_parameter * reservoir_streamflow[i] * delta_t
+            minimum_storage = reservoir_runoff_capacity_parameter * reservoir_storage_capacity[i]
+            maximum_storage = reservoir_storage_capacity[i]
+
+            condition_max = (flow_volume + reservoir_storage[i] - flow_reservoir - evaporation) >= maximum_storage
+            condition_min = (flow_volume + reservoir_storage[i] - flow_reservoir - evaporation) < minimum_storage
+            condition_min_one = flow_reservoir <= (flow_volume - evaporation)
+            condition_min_two = (flow_volume - evaporation) >= minimum_flow
+
+            if condition_max:
+                flow_reservoir = flow_volume + reservoir_storage[i] - maximum_storage - evaporation
+                reservoir_storage[i] = maximum_storage
+            else:
+                if condition_min:
+                    if condition_min_one:
+                        reservoir_storage[i] = reservoir_storage[i] + flow_volume - flow_reservoir - evaporation
+                    else:
+                        if condition_min_two:
+                            flow_reservoir = flow_volume - evaporation
+                        else:
+                            flow_reservoir = flow_volume
+                            reservoir_storage[i] = max(0, reservoir_storage[i] - flow_reservoir + flow_volume - evaporation)
+                else:
+                    reservoir_storage[i] = reservoir_storage[i] + flow_volume - flow_reservoir - evaporation
+
+            channel_outflow_downstream[i] = -flow_reservoir / delta_t
+
+
+@nb.jit(
+    "void("
+        "int64, int64, int64, int64[:], float64[:], DictType(int64, int64[:]),"
+        "float64[:], float64[:], float64[:], float64[:], float64[:], float64[:], float64[:],"
+        "int64, float64"
+    ")",
+    parallel=True,
+    nopython=True,
+    nogil=True,
+    cache=True,
+)
+def extraction_regulated_flow(
+    n,
+    n_reservoir,
+    delta_t,
+    grid_id,
+    reservoir_id,
+    reservoir_to_grid_map,
+    channel_outflow_downstream,
+    outflow_before_regulation,
+    outflow_after_regulation,
+    channel_flow,
+    grid_cell_unmet_demand,
+    grid_cell_supply,
+    grid_cell_deficit,
+    reservoir_supply_iterations,
+    reservoir_flow_volume_ratio,
+):
+    """Tracks the supply of water extracted from the reservoirs to fulfill demand from dependent grid cells."""
+
     # notes from fortran mosart:
     # This is an iterative algorithm that converts main channel flow
     # at each dam into gridcell supply based on the demand of each
@@ -143,112 +142,89 @@ def extraction_regulated_flow(state: State, grid: Grid, parameters: Parameters,
     # 3. if any sum of dam fraction < 1.0 for a gridcell, then provide fraction of 
     #    demand to gridcell prorated by the dam fraction of each dam.
     #
-    
-    has_reservoir = np.isfinite(grid.reservoir_id)
-    
-    flow_volume = calculate_flow_volume(has_reservoir, parameters.reservoir_flow_volume_ratio, delta_t, state.channel_outflow_downstream)
-    
-    state.channel_outflow_downstream = remove_flow(has_reservoir, state.channel_outflow_downstream, flow_volume, delta_t)
-    
-    cells = pd.DataFrame({'id': grid.id[state.grid_cell_unmet_demand > 0]}).set_index('id')
-    cells['supply'] = 0
-    
-    # join grid cell demand, then drop where no demand
-    demand = grid.reservoir_to_grid_mapping.join(pd.DataFrame(state.grid_cell_unmet_demand, columns=['grid_cell_demand']))
-    demand = demand[demand.grid_cell_demand.gt(0)]
-    
-    # aggregate demand to each reservoir and join to flow volume
-    reservoir_demand_flow = demand.groupby('reservoir_id')[['grid_cell_demand']].sum().rename(columns={'grid_cell_demand': 'reservoir_demand'}).join(pd.DataFrame({'flow_volume': flow_volume, 'reservoir_id': grid.reservoir_id}).dropna().set_index('reservoir_id'))
-    
-    for _ in np.arange(parameters.reservoir_supply_iterations):
-        
-        if _ == 0:
-            case = reservoir_demand_flow
-        else:
-            # subset reservoir list to speed up calculation
-            case = reservoir_demand_flow[np.isin(reservoir_demand_flow.index.astype(int).values, demand.reservoir_id.unique())]
-            case.loc[:, 'reservoir_demand'] = case.join(demand.groupby('reservoir_id')[['grid_cell_demand']].sum()).grid_cell_demand.fillna(0)
-        
-        # ratio of flow to total demand
-        case.loc[:, 'demand_fraction'] = divide(case.flow_volume.values, case.reservoir_demand.values)
-        
-        # case 1
-        if case.demand_fraction.gt(1).any():
-            case = demand[np.isin(demand.reservoir_id.values, case[case.demand_fraction.gt(1)].index.astype(int).values)]
-            case.loc[:, 'condition_count'] = case.groupby(case.index)['reservoir_id'].transform('count')
-            case.loc[:, 'supply'] = divide(case.grid_cell_demand, case.condition_count)
-            taken_from_reservoir = reservoir_demand_flow.join(case.groupby('reservoir_id').supply.sum()).supply.fillna(0).values
-            reservoir_demand_flow.loc[:, 'reservoir_demand'] -= taken_from_reservoir
-            reservoir_demand_flow.loc[:, 'flow_volume'] -= taken_from_reservoir
-            # all demand was supplied to these cells
-            cells.loc[:, 'supply'] += cells.join(case.groupby(case.index)[['grid_cell_demand']].first()).grid_cell_demand.fillna(0)
-            demand = demand[~demand.index.isin(case.index.unique())]
-        
-        else:
-            # sum demand fraction
-            case = demand.merge(case, how='left', left_on='reservoir_id', right_index=True)
-            case.loc[:, 'demand_fraction_sum'] = case.groupby(case.index).demand_fraction.transform('sum').fillna(0).values
-            
-            # case 2
-            if case.demand_fraction_sum.ge(1).any():
-                case = case[case.demand_fraction_sum.ge(1)]
-                case.loc[:, 'supply'] = case.grid_cell_demand.values  * case.demand_fraction.values / case.demand_fraction_sum.values
-                taken_from_reservoir = reservoir_demand_flow.join(case.groupby('reservoir_id')['supply'].sum()).supply.fillna(0).values
-                reservoir_demand_flow.loc[:, 'reservoir_demand'] = subtract(reservoir_demand_flow.reservoir_demand.values, taken_from_reservoir)
-                reservoir_demand_flow.loc[:, 'flow_volume'] = subtract(reservoir_demand_flow.flow_volume.values, taken_from_reservoir)
-                # all demand was supplied to these cells
-                cells.loc[:, 'supply'] += cells.join(case.groupby(case.index)[['grid_cell_demand']].first()).grid_cell_demand.fillna(0)
-                demand = demand[~demand.index.isin(case.index.unique())]
-                
+
+    # note that these reservoir arrays start at zero index but reservoir IDs start at one index
+
+    reservoir_flow_volume = np.full(n_reservoir, 0.0)
+
+    for i in nb.prange(n):
+
+        outflow_before_regulation[i] = -channel_outflow_downstream[i]
+        channel_flow[i] = channel_flow[i] + channel_outflow_downstream[i]
+
+        has_reservoir = np.isfinite(reservoir_id[i])
+
+        if has_reservoir:
+            reservoir_flow_volume[int(reservoir_id[i]) - 1] = -(reservoir_flow_volume_ratio * delta_t * channel_outflow_downstream[i])
+            channel_outflow_downstream[i] = channel_outflow_downstream[i] + reservoir_flow_volume[int(reservoir_id[i]) - 1] / delta_t
+
+    for _ in np.arange(reservoir_supply_iterations):
+        # total potential water demand on a reservoir
+        reservoir_demand = np.full(n_reservoir, 0.0)
+        # ratio of available water to total potential demand on a reservoir
+        demand_fraction = np.full(n_reservoir, 0.0)
+        for i in nb.prange(n):
+            if grid_cell_unmet_demand[i] > 0:
+                # assign this grid cell's demand to each available reservoir
+                if grid_id[i] in reservoir_to_grid_map:
+                    for r in reservoir_to_grid_map[grid_id[i]]:
+                        reservoir_demand[r - 1] = reservoir_demand[r - 1] + grid_cell_unmet_demand[i]
+        for r in nb.prange(n_reservoir):
+            # ratio of available water to total demand on a reservoir
+            if (reservoir_demand[r] > 0.0) and (reservoir_flow_volume[r] > 0.0):
+                demand_fraction[r] = reservoir_flow_volume[r] / reservoir_demand[r]
             else:
-                case = case[case.demand_fraction_sum.gt(0)]
-                case.loc[:, 'supply'] = case.grid_cell_demand.values * case.demand_fraction.values
-                taken_from_reservoir = reservoir_demand_flow.join(case.groupby('reservoir_id')['supply'].sum()).supply.fillna(0).values
-                reservoir_demand_flow.loc[:, 'reservoir_demand'] -= taken_from_reservoir
-                reservoir_demand_flow.loc[:, 'flow_volume'] -= taken_from_reservoir
-                # not all demand was supplied to these cells
-                supplied = cells[[]].join(case.groupby(case.index)[['supply']].sum()).supply.fillna(0)
-                cells.loc[:, 'supply'] += supplied
-                demand.loc[:, 'grid_cell_demand'] -= demand[[]].join(supplied).fillna(0).supply.values
-    
-    # merge the supply back in and update demand
-    supplied = pd.DataFrame(grid.id).join(cells).supply.fillna(0).values
-    state.grid_cell_supply = add(state.grid_cell_supply, supplied)
-    state.grid_cell_unmet_demand = subtract(state.grid_cell_unmet_demand, supplied)
-    
-    # add the residual flow volume back
-    state.channel_outflow_downstream[:] -= pd.DataFrame(grid.reservoir_id, columns=['reservoir_id']).merge(reservoir_demand_flow.flow_volume, how='left', left_on='reservoir_id', right_index=True).flow_volume.fillna(0).values / delta_t
-
-
-calculate_flow_volume = ne.NumExpr(
-    'where('
-        'has_reservoir,'
-        '-(reservoir_flow_volume_ratio * delta_t * channel_outflow_downstream),'
-        '0'
-    ')',
-    (('has_reservoir', np.bool), ('reservoir_flow_volume_ratio',  np.float64), ('delta_t', np.float64), ('channel_outflow_downstream', np.float64))
-)
+                demand_fraction[r] = 0.0
 
-remove_flow = ne.NumExpr(
-    'where('
-        'has_reservoir,'
-        'channel_outflow_downstream + flow_volume / delta_t,'
-        'channel_outflow_downstream'
-    ')',
-    (('has_reservoir', np.bool), ('channel_outflow_downstream',  np.float64), ('flow_volume', np.float64), ('delta_t', np.float64))
-)
+        if np.max(demand_fraction) >= 1.0:
+            # case 1 - provide all water to grid cell split from all available reservoirs with demand_fraction > 1
+            for i in nb.prange(n):
+                if grid_id[i] in reservoir_to_grid_map:
+                    available_reservoirs = List()
+                    for r in reservoir_to_grid_map[grid_id[i]]:
+                        if demand_fraction[r - 1] >= 1.0:
+                            available_reservoirs.append(r - 1)
+                    # take equally from each reservoir
+                    if len(available_reservoirs) > 0:
+                        for r in available_reservoirs:
+                            reservoir_flow_volume[r] = reservoir_flow_volume[r] - grid_cell_unmet_demand[i] / float(len(available_reservoirs))
+                        grid_cell_supply[i] = grid_cell_supply[i] + grid_cell_unmet_demand[i]
+                        grid_cell_unmet_demand[i] = 0.0
 
-divide = ne.NumExpr(
-    'a / b',
-    (('a', np.float64), ('b', np.float64))
-)
+        else:
+            sum_demand_fraction = np.full(n, 0.0)
+            for i in nb.prange(n):
+                if grid_id[i] in reservoir_to_grid_map:
+                    for r in reservoir_to_grid_map[grid_id[i]]:
+                        sum_demand_fraction[i] = sum_demand_fraction[i] + demand_fraction[r - 1]
 
-subtract = ne.NumExpr(
-    'a - b',
-    (('a', np.float64), ('b', np.float64))
-)
+            if np.any(sum_demand_fraction >= 1.0):
+                # case 2 - provide all water to grid cell prorated from all available reservoirs
+                for i in nb.prange(n):
+                    if sum_demand_fraction[i] >= 1.0:
+                        for r in reservoir_to_grid_map[grid_id[i]]:
+                            reservoir_flow_volume[r - 1] = reservoir_flow_volume[r - 1] - grid_cell_unmet_demand[i] * demand_fraction[r - 1] / sum_demand_fraction[i]
+                        grid_cell_supply[i] = grid_cell_supply[i] + grid_cell_unmet_demand[i]
+                        grid_cell_unmet_demand[i] = 0.0
+
+            else:
+                # case 3 - provide fraction of water to grid cell prorated from all available reservoirs
+                for i in nb.prange(n):
+                    if sum_demand_fraction[i] > 0.0:
+                        total_take = 0.0
+                        for r in reservoir_to_grid_map[grid_id[i]]:
+                            take = min(grid_cell_unmet_demand[i] * demand_fraction[r - 1], grid_cell_unmet_demand[i])
+                            total_take = total_take + take
+                            reservoir_flow_volume[r - 1] = reservoir_flow_volume[r - 1] - take
+                        grid_cell_supply[i] = grid_cell_supply[i] + total_take
+                        grid_cell_unmet_demand[i] = grid_cell_unmet_demand[i] - total_take
 
-add = ne.NumExpr(
-    'a + b',
-    (('a', np.float64), ('b', np.float64))
-)
\ No newline at end of file
+    for i in nb.prange(n):
+        has_reservoir = np.isfinite(reservoir_id[i])
+        if has_reservoir:
+            # add the residual flow volume back to channel
+            channel_outflow_downstream[i] = channel_outflow_downstream[i] - reservoir_flow_volume[int(reservoir_id[i]) - 1] / delta_t
+        outflow_after_regulation[i] = -channel_outflow_downstream[i]
+        channel_flow[i] = channel_flow[i] - channel_outflow_downstream[i]
+        # aggregate deficit
+        grid_cell_deficit[i] = grid_cell_deficit[i] + grid_cell_unmet_demand[i]
diff --git a/mosartwmpy/reservoirs/reservoirs.py b/mosartwmpy/reservoirs/reservoirs.py
index 4c0186c..a8ce751 100644
--- a/mosartwmpy/reservoirs/reservoirs.py
+++ b/mosartwmpy/reservoirs/reservoirs.py
@@ -24,6 +24,7 @@ def reservoir_release(state, grid, config, parameters, current_time):
     
     storage_targets(state, grid, config, parameters, current_time)
 
+
 def regulation_release(state, grid, config, parameters, current_time):
     # compute the expected monthly release based on Biemans (2011)
 
@@ -34,11 +35,11 @@ def regulation_release(state, grid, config, parameters, current_time):
     state.reservoir_release = grid.reservoir_streamflow_schedule.mean(dim=streamflow_time_name).values
     
     # TODO what is k
-    k = state.reservoir_storage_operation_year_start / ( parameters.reservoir_regulation_release_parameter * grid.reservoir_storage_capacity)
+    k = state.reservoir_storage_operation_year_start / (parameters.reservoir_regulation_release_parameter * grid.reservoir_storage_capacity)
     
     # TODO what is factor
     factor = np.where(
-        grid.reservoir_runoff_capacity > parameters.reservoir_runoff_capacity_condition,
+        grid.reservoir_runoff_capacity > parameters.reservoir_runoff_capacity_parameter,
         (2.0 / grid.reservoir_runoff_capacity) ** 2.0,
         0
     )
@@ -58,6 +59,7 @@ def regulation_release(state, grid, config, parameters, current_time):
         )
     )
 
+
 def storage_targets(state: State, grid: Grid, config: Benedict, parameters: Parameters, current_time: datetime) -> None:
     """Define the necessary drop in storage based on the reservoir storage targets at the start of the month.
 
@@ -121,7 +123,7 @@ def storage_targets(state: State, grid: Grid, config: Benedict, parameters: Para
         (month < state.reservoir_month_start_operations)
     )
     state.reservoir_release = np.where(
-        (state.reservoir_release> grid.reservoir_streamflow_schedule.mean(dim=streamflow_time_name).values) & (first_condition | second_condition),
+        (state.reservoir_release > grid.reservoir_streamflow_schedule.mean(dim=streamflow_time_name).values) & (first_condition | second_condition),
         grid.reservoir_streamflow_schedule.mean(dim=streamflow_time_name).values,
         state.reservoir_release
     )
diff --git a/mosartwmpy/reservoirs/state.py b/mosartwmpy/reservoirs/state.py
index 6d305a9..718fe80 100644
--- a/mosartwmpy/reservoirs/state.py
+++ b/mosartwmpy/reservoirs/state.py
@@ -5,7 +5,6 @@
 from mosartwmpy.config.parameters import Parameters
 from mosartwmpy.grid.grid import Grid
 
-# TODO in fortran mosart there is a StorCalibFlag that affects how storage targets are calculated -- code so far is written assuming that it == 0
 
 def initialize_reservoir_state(self, grid: Grid, config: Benedict, parameters: Parameters) -> None:
     """Initializes the reservoir state.
@@ -202,4 +201,4 @@ def initialize_reservoir_start_of_operation_year(self, grid: Grid, config: Bened
 
     self.reservoir_month_start_operations = month_start_operations
     self.reservoir_month_flood_control_start = month_flood_control_start
-    self.reservoir_month_flood_control_end = month_flood_control_end
\ No newline at end of file
+    self.reservoir_month_flood_control_end = month_flood_control_end
diff --git a/mosartwmpy/state/state.py b/mosartwmpy/state/state.py
index 0cb9751..9e945b0 100644
--- a/mosartwmpy/state/state.py
+++ b/mosartwmpy/state/state.py
@@ -1,8 +1,5 @@
-import logging
 import numpy as np
 import pandas as pd
-from datetime import datetime, time
-from xarray import open_dataset
 
 from benedict.dicts import benedict as Benedict
 
@@ -10,7 +7,8 @@
 from mosartwmpy.grid.grid import Grid
 from mosartwmpy.reservoirs.state import initialize_reservoir_state
 
-class State():
+
+class State:
 
     def __init__(self, grid: Grid = None, config: Benedict = None, parameters: Parameters = None, grid_size: int = None, empty: bool = False):
         """Initialize the model state.
@@ -204,8 +202,7 @@ def __init__(self, grid: Grid = None, config: Benedict = None, parameters: Param
         # a column of always all zeros, to use as a utility
         self.zeros: np.ndarray = np.empty(0)
 
-
-        ## Reservoir related state variables
+        # Reservoir related state variables
 
         # reservoir streamflow schedule
         self.reservoir_streamflow: np.ndarray = np.empty(0)
@@ -237,7 +234,6 @@ def __init__(self, grid: Grid = None, config: Benedict = None, parameters: Param
             return
         
         # initialize all the state variables
-        logging.debug('Initializing state variables.')
         for key in [key for key in dir(self) if isinstance(getattr(self, key), np.ndarray)]:
             setattr(self, key, np.zeros(grid_size))
         
@@ -252,7 +248,6 @@ def __init__(self, grid: Grid = None, config: Benedict = None, parameters: Param
         )
     
         if config.get('water_management.enabled', False):
-            logging.debug(' - reservoirs')
             initialize_reservoir_state(self, grid, config, parameters)
 
     def __getitem__(self, item):
@@ -283,4 +278,4 @@ def from_dataframe(df: pd.DataFrame) -> 'State':
         state = State(empty=True)
         for key in df.columns:
             setattr(state, key, df[key].values)
-        return state
\ No newline at end of file
+        return state
diff --git a/mosartwmpy/subnetwork/irrigation.py b/mosartwmpy/subnetwork/irrigation.py
index 1e4cfd1..496cd3f 100644
--- a/mosartwmpy/subnetwork/irrigation.py
+++ b/mosartwmpy/subnetwork/irrigation.py
@@ -1,103 +1,66 @@
-import numpy as np
-import numexpr as ne
+import numba as nb
 
-from mosartwmpy.config.parameters import Parameters
-from mosartwmpy.grid.grid import Grid
-from mosartwmpy.state.state import State
 from mosartwmpy.subnetwork.state import update_subnetwork_state
-from mosartwmpy.utilities.timing import timing
 
-# @timing
-def subnetwork_irrigation(state: State, grid: Grid, parameters: Parameters) -> None:
-    """Tracks the supply of water from the subnetwork river channels extracted into the grid cells.
 
-    Args:
-        state (State): the current model state; will be mutated
-        grid (Grid): the model grid
-        parameters (Parameters): the model parameters
-    """
-    
-    depth_condition = calculate_depth_condition(grid.mosart_mask, state.euler_mask, state.tracer, parameters.LIQUID_TRACER, state.subnetwork_depth, parameters.irrigation_extraction_condition)
-    
-    flow_volume = np.empty_like(state.subnetwork_storage)
-    np.copyto(flow_volume, state.subnetwork_storage)
-    
-    volume_condition = calculate_volume_condition(flow_volume, state.grid_cell_unmet_demand)
-    
-    state.grid_cell_supply = calculate_reservoir_supply(depth_condition, volume_condition, state.grid_cell_supply, state.grid_cell_unmet_demand, flow_volume)
-    
-    flow_volume = calculate_flow_volume(depth_condition, volume_condition, flow_volume, state.grid_cell_unmet_demand)
-    
-    state.grid_cell_unmet_demand = calculate_reservoir_demand(depth_condition, volume_condition, state.grid_cell_unmet_demand, flow_volume)
-    
-    flow_volume = update_flow_volume(depth_condition, volume_condition, flow_volume)
-    
-    state.subnetwork_storage = calculate_subnetwork_storage(depth_condition, flow_volume, state.subnetwork_storage)
-    
-    update_subnetwork_state(state, grid, parameters, depth_condition)
-
-calculate_depth_condition = ne.NumExpr(
-    '(mosart_mask > 0) &'
-    'euler_mask &'
-    '(tracer == LIQUID_TRACER) &'
-    '(subnetwork_depth >= irrigation_extraction_condition)',
-    (('mosart_mask', np.int64), ('euler_mask', np.bool), ('tracer', np.int64), ('LIQUID_TRACER', np.int64), ('subnetwork_depth', np.float64), ('irrigation_extraction_condition', np.float64))
+@nb.jit(
+    "void("
+        "int64, int64[:], float64[:], float64[:], boolean[:],"
+        "float64[:], float64[:], float64[:], float64[:], float64[:], float64[:], float64[:],"
+        "float64, float64"
+    ")",
+    parallel=True,
+    nopython=True,
+    nogil=True,
+    cache=True,
 )
+def subnetwork_irrigation(
+    n,
+    mosart_mask,
+    subnetwork_length,
+    subnetwork_width,
+    euler_mask,
+    subnetwork_depth,
+    subnetwork_storage,
+    grid_cell_unmet_demand,
+    grid_cell_supply,
+    subnetwork_cross_section_area,
+    subnetwork_wetness_perimeter,
+    subnetwork_hydraulic_radii,
+    irrigation_extraction_parameter,
+    tiny_value,
+):
+    """Tracks the supply of water from the subnetwork river channels extracted into the grid cells."""
 
-calculate_volume_condition = ne.NumExpr(
-    'flow_volume >= reservoir_demand',
-    (('flow_volume', np.float64), ('reservoir_demand', np.float64))
-)
+    for i in nb.prange(n):
 
-calculate_reservoir_supply = ne.NumExpr(
-    'where('
-        'depth_condition,'
-        'where('
-            'volume_condition,'
-            'reservoir_supply + reservoir_demand,'
-            'reservoir_supply + flow_volume'
-        '),'
-        'reservoir_supply'
-    ')',
-    (('depth_condition', np.bool), ('volume_condition', np.bool), ('reservoir_supply', np.float64), ('reservoir_demand', np.float64), ('flow_volume', np.float64))
-)
+        # is the channel deep enough to extract from
+        depth_condition = (mosart_mask[i] > 0) and euler_mask[i] and (subnetwork_depth[i] >= irrigation_extraction_parameter)
 
-calculate_flow_volume = ne.NumExpr(
-    'where('
-        'depth_condition & volume_condition,'
-        'flow_volume - reservoir_demand,'
-        'flow_volume'
-    ')',
-    (('depth_condition', np.bool), ('volume_condition', np.bool), ('flow_volume', np.float64), ('reservoir_demand', np.float64))
-)
+        flow_volume = 1.0 * subnetwork_storage[i]
 
-calculate_reservoir_demand = ne.NumExpr(
-    'where('
-        'depth_condition,'
-        'where('
-            'volume_condition,'
-            '0,'
-            'reservoir_demand - flow_volume'
-        '),'
-        'reservoir_demand'
-    ')',
-    (('depth_condition', np.bool), ('volume_condition', np.bool), ('reservoir_demand', np.float64), ('flow_volume', np.float64))
-)
+        # is there enough water to allow full extraction
+        volume_condition = flow_volume >= grid_cell_unmet_demand[i]
 
-update_flow_volume = ne.NumExpr(
-    'where('
-        'depth_condition & (~volume_condition),'
-        '0,'
-        'flow_volume'
-    ')',
-    (('depth_condition', np.bool), ('volume_condition', np.bool), ('flow_volume', np.float64))
-)
+        if depth_condition:
+            if volume_condition:
+                grid_cell_supply[i] = grid_cell_supply[i] + grid_cell_unmet_demand[i]
+                flow_volume = flow_volume - grid_cell_unmet_demand[i]
+                grid_cell_unmet_demand[i] = 0.0
+            else:
+                grid_cell_supply[i] = grid_cell_supply[i] + flow_volume
+                grid_cell_unmet_demand[i] = grid_cell_unmet_demand[i] - flow_volume
+                flow_volume = 0.0
 
-calculate_subnetwork_storage = ne.NumExpr(
-    'where('
-        'depth_condition,'
-        'flow_volume,'
-        'subnetwork_storage'
-    ')',
-    (('depth_condition', np.bool), ('flow_volume', np.float64), ('subnetwork_storage', np.float64))
-)
\ No newline at end of file
+            subnetwork_storage[i] = 1.0 * flow_volume
+            update_subnetwork_state(
+                i,
+                subnetwork_length,
+                subnetwork_width,
+                subnetwork_storage,
+                subnetwork_cross_section_area,
+                subnetwork_depth,
+                subnetwork_wetness_perimeter,
+                subnetwork_hydraulic_radii,
+                tiny_value,
+            )
diff --git a/mosartwmpy/subnetwork/routing.py b/mosartwmpy/subnetwork/routing.py
index b28f5b9..29479ee 100644
--- a/mosartwmpy/subnetwork/routing.py
+++ b/mosartwmpy/subnetwork/routing.py
@@ -1,152 +1,105 @@
-import numpy as np
-import numexpr as ne
+import numba as nb
 
-from benedict.dicts import benedict as Benedict
-
-from mosartwmpy.config.parameters import Parameters
-from mosartwmpy.grid.grid import Grid
-from mosartwmpy.state.state import State
 from mosartwmpy.subnetwork.state import update_subnetwork_state
 from mosartwmpy.utilities.timing import timing
 
-# @timing
-def subnetwork_routing(state: State, grid: Grid, parameters: Parameters, config: Benedict, delta_t: float) -> None:
-    """Tracks the storage and flow of water in the subnetwork river channels.
-
-    Args:
-        state (State): the current model state; will be mutated
-        grid (Grid): the model grid
-        parameters (Parameters): the model parameters
-        config (Benedict): the model configuration
-        delta_t (float): the timestep for this subcycle (overall timestep / subcycles)
-    """
-    
-    state.channel_lateral_flow_hillslope[:] = 0
-    local_delta_t = (delta_t / config.get('simulation.routing_iterations') / grid.iterations_subnetwork)
-    
-    # step through max iterations, masking out the unnecessary cells each time
-    base_condition = (grid.mosart_mask > 0) & state.euler_mask
-    sub_condition = grid.subnetwork_length > grid.hillslope_length # has tributaries
-    
-    for _ in np.arange(np.nanmax(grid.iterations_subnetwork)):
-        iteration_condition = base_condition & (grid.iterations_subnetwork > _)
-
-        state.subnetwork_flow_velocity = calculate_subnetwork_flow_velocity(iteration_condition, sub_condition, state.subnetwork_hydraulic_radii, grid.subnetwork_slope, grid.subnetwork_manning, state.subnetwork_flow_velocity)
-        
-        state.subnetwork_discharge = calculate_subnetwork_discharge(iteration_condition, sub_condition, state.subnetwork_flow_velocity, state.subnetwork_cross_section_area, state.subnetwork_lateral_inflow, state.subnetwork_discharge)
-        
-        discharge_condition = calculate_discharge_condition(iteration_condition, sub_condition, state.subnetwork_storage, state.subnetwork_lateral_inflow, state.subnetwork_discharge, local_delta_t, parameters.tiny_value)
-        
-        state.subnetwork_discharge = update_subnetwork_discharge(discharge_condition, state.subnetwork_lateral_inflow, state.subnetwork_storage, local_delta_t, state.subnetwork_discharge)
-        
-        state.subnetwork_flow_velocity = update_flow_velocity(discharge_condition, state.subnetwork_cross_section_area, state.subnetwork_discharge, state.subnetwork_flow_velocity)
-        
-        state.subnetwork_delta_storage = calculate_subnetwork_delta_storage(iteration_condition, state.subnetwork_lateral_inflow, state.subnetwork_discharge, state.subnetwork_delta_storage)
-        
-        # update storage
-        state.subnetwork_storage_previous_timestep = calculate_subnetwork_storage_previous_timestep(iteration_condition, state.subnetwork_storage, state.subnetwork_storage_previous_timestep)
-        state.subnetwork_storage = calculate_subnetwork_storage(iteration_condition, state.subnetwork_storage, state.subnetwork_delta_storage, local_delta_t)
-        
-        update_subnetwork_state(state, grid, parameters, iteration_condition)
-        
-        state.channel_lateral_flow_hillslope = calculate_channel_lateral_flow_hillslope(iteration_condition, state.channel_lateral_flow_hillslope, state.subnetwork_discharge)
-    
-    # average lateral flow over substeps
-    state.channel_lateral_flow_hillslope = average_channel_lateral_flow_hillslope(base_condition, state.channel_lateral_flow_hillslope, grid.iterations_subnetwork)
-
-
-calculate_subnetwork_flow_velocity = ne.NumExpr(
-    'where('
-        'base_condition & length_condition,'
-        'where('
-            'subnetwork_hydraulic_radii > 0,'
-            '(subnetwork_hydraulic_radii ** (2/3)) * sqrt(subnetwork_slope) / subnetwork_manning,'
-            '0'
-        '),'
-        'subnetwork_flow_velocity'
-    ')',
-    (('base_condition', np.bool), ('length_condition', np.bool), ('subnetwork_hydraulic_radii', np.float64), ('subnetwork_slope', np.float64), ('subnetwork_manning', np.float64), ('subnetwork_flow_velocity', np.float64))
-)
-
-calculate_subnetwork_discharge = ne.NumExpr(
-    'where('
-        'base_condition,'
-        'where('
-            'length_condition,'
-            '-subnetwork_flow_velocity * subnetwork_cross_section_area,'
-            '-subnetwork_lateral_inflow'
-        '),'
-        'subnetwork_discharge'
-    ')',
-    (('base_condition', np.bool), ('length_condition', np.bool), ('subnetwork_flow_velocity', np.float64), ('subnetwork_cross_section_area', np.float64), ('subnetwork_lateral_inflow', np.float64), ('subnetwork_discharge', np.float64))
-)
-
-calculate_discharge_condition = ne.NumExpr(
-    'base_condition &'
-    'length_condition &'
-    '((subnetwork_storage + (subnetwork_lateral_inflow + subnetwork_discharge) * local_delta_t) < tiny_value)',
-    (('base_condition', np.bool), ('length_condition', np.bool), ('subnetwork_storage', np.float64), ('subnetwork_lateral_inflow', np.float64), ('subnetwork_discharge', np.float64), ('local_delta_t', np.float64), ('tiny_value', np.float64))
-)
 
-update_subnetwork_discharge = ne.NumExpr(
-    'where('
-        'discharge_condition,'
-        '-(subnetwork_lateral_inflow + subnetwork_storage / local_delta_t),'
-        'subnetwork_discharge'
-    ')',
-    (('discharge_condition', np.bool), ('subnetwork_lateral_inflow', np.float64), ('subnetwork_storage', np.float64), ('local_delta_t', np.float64), ('subnetwork_discharge', np.float64))
-)
-
-update_flow_velocity = ne.NumExpr(
-    'where('
-        'discharge_condition & (subnetwork_cross_section_area > 0),'
-        '-subnetwork_discharge / subnetwork_cross_section_area,'
-        'subnetwork_flow_velocity'
-    ')',
-    (('discharge_condition', np.bool), ('subnetwork_cross_section_area', np.float64), ('subnetwork_discharge', np.float64), ('subnetwork_flow_velocity', np.float64))
-)
-
-calculate_subnetwork_delta_storage = ne.NumExpr(
-    'where('
-        'base_condition,'
-        'subnetwork_lateral_inflow + subnetwork_discharge,'
-        'subnetwork_delta_storage'
-    ')',
-    (('base_condition', np.bool), ('subnetwork_lateral_inflow', np.float64), ('subnetwork_discharge', np.float64), ('subnetwork_delta_storage', np.float64))
-)
-
-calculate_subnetwork_storage_previous_timestep = ne.NumExpr(
-    'where('
-        'base_condition,'
-        'subnetwork_storage,'
-        'subnetwork_storage_previous_timestep'
-    ')',
-    (('base_condition', np.bool), ('subnetwork_storage', np.float64), ('subnetwork_storage_previous_timestep', np.float64))
-)
-
-calculate_subnetwork_storage = ne.NumExpr(
-    'where('
-        'base_condition,'
-        'subnetwork_storage + subnetwork_delta_storage * local_delta_t,'
-        'subnetwork_storage'
-    ')',
-    (('base_condition', np.bool), ('subnetwork_storage', np.float64), ('subnetwork_delta_storage', np.float64), ('local_delta_t', np.float64))
-)
-
-calculate_channel_lateral_flow_hillslope = ne.NumExpr(
-    'where('
-        'base_condition,'
-        'channel_lateral_flow_hillslope - subnetwork_discharge,'
-        'channel_lateral_flow_hillslope'
-    ')',
-    (('base_condition', np.bool), ('channel_lateral_flow_hillslope', np.float64), ('subnetwork_discharge', np.float64))
+# @timing
+@nb.jit(
+    "void("
+        "int64, float64, int64, int64,"
+        "int64[:], int64[:], float64[:], float64[:], float64[:], float64[:], float64[:],"
+        "boolean[:], float64[:], float64[:], float64[:], float64[:], float64[:],"
+        "float64[:], float64[:], float64[:], float64[:], float64[:], float64[:],"
+        "float64"
+    ")",
+    parallel=True,
+    nopython=True,
+    nogil=True,
+    cache=True,
 )
-
-average_channel_lateral_flow_hillslope = ne.NumExpr(
-    'where('
-        'base_condition,'
-        'channel_lateral_flow_hillslope / iterations_subnetwork,'
-        'channel_lateral_flow_hillslope'
-    ')',
-    (('base_condition', np.bool), ('channel_lateral_flow_hillslope', np.float64), ('iterations_subnetwork', np.float64))
-)
\ No newline at end of file
+def subnetwork_routing(
+    n,
+    delta_t,
+    routing_iterations,
+    max_iterations_subnetwork,
+    iterations_subnetwork,
+    mosart_mask,
+    subnetwork_slope,
+    subnetwork_manning,
+    subnetwork_length,
+    subnetwork_width,
+    hillslope_length,
+    euler_mask,
+    channel_lateral_flow_hillslope,
+    subnetwork_flow_velocity,
+    subnetwork_discharge,
+    subnetwork_lateral_inflow,
+    subnetwork_storage,
+    subnetwork_storage_previous_timestep,
+    subnetwork_delta_storage,
+    subnetwork_depth,
+    subnetwork_cross_section_area,
+    subnetwork_wetness_perimeter,
+    subnetwork_hydraulic_radii,
+    tiny_value,
+):
+    """Tracks the storage and flow of water in the subnetwork river channels."""
+
+    for i in nb.prange(n):
+
+        local_delta_t = (delta_t / routing_iterations) / iterations_subnetwork[i]
+        channel_lateral_flow_hillslope[i] = 0.0
+
+        if ~euler_mask[i] or ~(mosart_mask[i] > 0):
+            continue
+
+        has_tributaries = subnetwork_length[i] > hillslope_length[i]
+
+        # step through max iterations
+        for _ in nb.prange(max_iterations_subnetwork):
+            if ~(iterations_subnetwork[i] > _):
+                continue
+
+            if has_tributaries:
+                if subnetwork_hydraulic_radii[i] > 0.0:
+                    subnetwork_flow_velocity[i] = (subnetwork_hydraulic_radii[i] ** (2.0/3.0)) * (subnetwork_slope[i] ** (1.0/2.0)) / subnetwork_manning[i]
+                else:
+                    subnetwork_flow_velocity[i] = 0.0
+
+            if has_tributaries:
+                subnetwork_discharge[i] = -subnetwork_flow_velocity[i] * subnetwork_cross_section_area[i]
+            else:
+                subnetwork_discharge[i] = -1.0 * subnetwork_lateral_inflow[i]
+
+            discharge_condition = has_tributaries and ((subnetwork_storage[i] + (subnetwork_lateral_inflow[i] + subnetwork_discharge[i]) * local_delta_t) < tiny_value)
+
+            if discharge_condition:
+                subnetwork_discharge[i] = -(subnetwork_lateral_inflow[i] + subnetwork_storage[i] / local_delta_t)
+
+            if discharge_condition and (subnetwork_cross_section_area[i] > 0.0):
+                subnetwork_flow_velocity[i] = -subnetwork_discharge[i] / subnetwork_cross_section_area[i]
+
+            subnetwork_delta_storage[i] = subnetwork_lateral_inflow[i] + subnetwork_discharge[i]
+
+            # update storage
+            subnetwork_storage_previous_timestep[i] = subnetwork_storage[i]
+            subnetwork_storage[i] = subnetwork_storage[i] + subnetwork_delta_storage[i] * local_delta_t
+
+            # update subnetwork state
+            update_subnetwork_state(
+                i,
+                subnetwork_length,
+                subnetwork_width,
+                subnetwork_storage,
+                subnetwork_cross_section_area,
+                subnetwork_depth,
+                subnetwork_wetness_perimeter,
+                subnetwork_hydraulic_radii,
+                tiny_value,
+            )
+
+            channel_lateral_flow_hillslope[i] = channel_lateral_flow_hillslope[i] - subnetwork_discharge[i]
+
+        # average lateral flow over substeps
+        channel_lateral_flow_hillslope[i] = channel_lateral_flow_hillslope[i] / iterations_subnetwork[i]
diff --git a/mosartwmpy/subnetwork/state.py b/mosartwmpy/subnetwork/state.py
index cbff5bb..48e430c 100644
--- a/mosartwmpy/subnetwork/state.py
+++ b/mosartwmpy/subnetwork/state.py
@@ -1,81 +1,46 @@
-import numpy as np
-import numexpr as ne
+import numba as nb
 
-from mosartwmpy.config.parameters import Parameters
-from mosartwmpy.grid.grid import Grid
-from mosartwmpy.state.state import State
 
-def update_subnetwork_state(state: State, grid: Grid, parameters: Parameters, base_condition: np.ndarray) -> None:
-    """Updates the physical properties of the subnetwork river channels based on current state.
-
-    Args:
-        state (State): the current model state; will be mutated
-        grid (Grid): the model grid
-        parameters (Parameters): the model parameters
-        base_condition (np.ndarray): a boolean array representing where the update should occur in the state
-    """
-        
-    # update state variables
-    length_condition = calculate_length_condition(grid.subnetwork_length, state.subnetwork_storage)
-    state.subnetwork_cross_section_area = calculate_subnetwork_cross_section_area(base_condition, length_condition, state.subnetwork_storage, grid.subnetwork_length, state.subnetwork_cross_section_area)
-    state.subnetwork_depth = calculate_subnetwork_depth(base_condition, length_condition, state.subnetwork_cross_section_area, parameters.tiny_value, grid.subnetwork_width, state.subnetwork_depth)
-    state.subnetwork_wetness_perimeter = calculate_subnetwork_wetness_perimeter(base_condition, length_condition, state.subnetwork_depth, parameters.tiny_value, grid.subnetwork_width, state.subnetwork_wetness_perimeter)
-    state.subnetwork_hydraulic_radii = calculate_subnetwork_hydraulic_radii(base_condition, length_condition, state.subnetwork_wetness_perimeter, parameters.tiny_value, state.subnetwork_cross_section_area, state.subnetwork_hydraulic_radii)
-
-
-calculate_length_condition = ne.NumExpr(
-    '(subnetwork_length > 0) & (subnetwork_storage > 0)',
-    (('subnetwork_length', np.float64), ('subnetwork_storage', np.float64))
-)
-
-calculate_subnetwork_cross_section_area = ne.NumExpr(
-    'where('
-        'base_condition,'
-        'where('
-            'length_condition,'
-            'subnetwork_storage / subnetwork_length,'
-            '0'
-        '),'
-        'subnetwork_cross_section_area'
-    ')',
-    (('base_condition', np.bool), ('length_condition', np.bool), ('subnetwork_storage', np.float64), ('subnetwork_length', np.float64), ('subnetwork_cross_section_area', np.float64))
+@nb.jit(
+    "void("
+        "int64, float64[:], float64[:], float64[:], float64[:],"
+        "float64[:], float64[:], float64[:], float64"
+    ")",
+    nopython=True,
+    nogil=True,
+    cache=True,
 )
-
-calculate_subnetwork_depth = ne.NumExpr(
-    'where('
-        'base_condition,'
-        'where('
-            'length_condition & (subnetwork_cross_section_area > tiny_value),'
-            'subnetwork_cross_section_area / subnetwork_width,'
-            '0'
-        '),'
-        'subnetwork_depth'
-    ')',
-    (('base_condition', np.bool), ('length_condition', np.bool), ('subnetwork_cross_section_area', np.float64), ('tiny_value', np.float64), ('subnetwork_width', np.float64), ('subnetwork_depth', np.float64))
-)
-
-calculate_subnetwork_wetness_perimeter = ne.NumExpr(
-    'where('
-        'base_condition,'
-        'where('
-            'length_condition & (subnetwork_depth > tiny_value),'
-            'subnetwork_width + 2 * subnetwork_depth,'
-            '0'
-        '),'
-        'subnetwork_wetness_perimeter'
-    ')',
-    (('base_condition', np.bool), ('length_condition', np.bool), ('subnetwork_depth', np.float64), ('tiny_value', np.float64), ('subnetwork_width', np.float64), ('subnetwork_wetness_perimeter', np.float64))
-)
-
-calculate_subnetwork_hydraulic_radii = ne.NumExpr(
-    'where('
-        'base_condition,'
-        'where('
-            'length_condition & (subnetwork_wetness_perimeter > tiny_value),'
-            'subnetwork_cross_section_area / subnetwork_wetness_perimeter,'
-            '0'
-        '),'
-        'subnetwork_hydraulic_radii'
-    ')',
-    (('base_condition', np.bool), ('length_condition', np.bool),  ('subnetwork_wetness_perimeter', np.float64), ('tiny_value', np.float64), ('subnetwork_cross_section_area', np.float64), ('subnetwork_hydraulic_radii', np.float64))
-)
\ No newline at end of file
+def update_subnetwork_state(
+    i,
+    subnetwork_length,
+    subnetwork_width,
+    subnetwork_storage,
+    subnetwork_cross_section_area,
+    subnetwork_depth,
+    subnetwork_wetness_perimeter,
+    subnetwork_hydraulic_radii,
+    tiny_value,
+):
+    """Updates the physical properties of the subnetwork river channels based on current state."""
+
+    has_water = (subnetwork_length[i] > 0.0) and (subnetwork_storage[i] > 0.0)
+
+    if has_water:
+        subnetwork_cross_section_area[i] = subnetwork_storage[i] / subnetwork_length[i]
+    else:
+        subnetwork_cross_section_area[i] = 0.0
+
+    if has_water and (subnetwork_cross_section_area[i] > tiny_value):
+        subnetwork_depth[i] = subnetwork_cross_section_area[i] / subnetwork_width[i]
+    else:
+        subnetwork_depth[i] = 0.0
+
+    if has_water and (subnetwork_depth[i] > tiny_value):
+        subnetwork_wetness_perimeter[i] = subnetwork_width[i] + 2.0 * subnetwork_depth[i]
+    else:
+        subnetwork_wetness_perimeter[i] = 0.0
+
+    if has_water and (subnetwork_wetness_perimeter[i] > tiny_value):
+        subnetwork_hydraulic_radii[i] = subnetwork_cross_section_area[i] / subnetwork_wetness_perimeter[i]
+    else:
+        subnetwork_hydraulic_radii[i] = 0.0
diff --git a/mosartwmpy/update/update.py b/mosartwmpy/update/update.py
index 7f79380..41c1ff1 100644
--- a/mosartwmpy/update/update.py
+++ b/mosartwmpy/update/update.py
@@ -1,6 +1,6 @@
+import numba as nb
 import numpy as np
 import pandas as pd
-import logging
 
 from benedict.dicts import benedict as Benedict
 
@@ -8,21 +8,13 @@
 from mosartwmpy.grid.grid import Grid
 from mosartwmpy.state.state import State
 
-from mosartwmpy.direct_to_ocean.direct_to_ocean import direct_to_ocean
-from mosartwmpy.flood.flood import flood
 from mosartwmpy.hillslope.routing import hillslope_routing
-from mosartwmpy.main_channel.irrigation import  main_channel_irrigation
+from mosartwmpy.main_channel.irrigation import main_channel_irrigation
 from mosartwmpy.main_channel.routing import main_channel_routing
 from mosartwmpy.reservoirs.regulation import extraction_regulated_flow, regulation
 from mosartwmpy.subnetwork.irrigation import subnetwork_irrigation
 from mosartwmpy.subnetwork.routing import subnetwork_routing
 
-# ignore numpy NaN and invalid warnings
-# (i.e. divide by zero and NaN logicals -- in Pandas/Dask, these simply remain NaN instead of becoming infinite)
-np.seterr(all='ignore')
-# filter pandas chained assignment warnings -- pretty sure handling this correctly
-pd.options.mode.chained_assignment = None
-
 
 def update(state: State, grid: Grid, parameters: Parameters, config: Benedict) -> None:
     """Advance the simulation one timestamp.
@@ -33,202 +25,647 @@ def update(state: State, grid: Grid, parameters: Parameters, config: Benedict) -
         parameters (Parameters): the model parameters
         config (Benedict): the model configuration
     """
-    # perform one timestep
-    
-    # ignore nan, overflow, underflow, and div by 0 warnings, since they are handled correctly
-    np.seterr(all='ignore')
-    
-    ###
-    ### Reset certain state variables
-    ###
-    state.flow[:] = 0
-    state.outflow_downstream_previous_timestep[:] = 0
-    state.outflow_downstream_current_timestep[:] = 0
-    state.outflow_before_regulation[:] = 0
-    state.outflow_after_regulation[:] = 0
-    state.outflow_sum_upstream_average[:] = 0
-    state.lateral_flow_hillslope_average[:] = 0
-    state.runoff[:] = 0
-    state.direct[:] = 0
-    state.flood[:] = 0
-    state.runoff_land[:] = 0
-    state.runoff_ocean[:] = 0
-    state.delta_storage[:] = 0
-    state.delta_storage_land[:] = 0
-    state.delta_storage_ocean[:] = 0
-    if config.get('water_management.enabled', False):
-        state.grid_cell_deficit[:] = 0
-
-    # flood
-    flood(state, grid, parameters, config)
-
-    # direct to ocean
-    direct_to_ocean(state, grid, parameters, config)
-    
-    ###
-    ### Subcycling
-    ###
-    
-    # convert runoff to m/s
-    state.hillslope_surface_runoff = state.hillslope_surface_runoff / grid.area
-    state.hillslope_subsurface_runoff = state.hillslope_subsurface_runoff / grid.area
-    state.hillslope_wetland_runoff = state.hillslope_wetland_runoff / grid.area
-
-    # subcycle timestep
-    delta_t =  config.get('simulation.timestep') / config.get('simulation.subcycles')
-    
+
+    # size of flattened grid
+    n = len(grid.id)
+
+    _prepare(
+        n,
+        config.get('water_management.enabled', False),
+        config.get('simulation.timestep'),
+        grid.land_mask,
+        grid.mosart_mask,
+        grid.area,
+        state.flow,
+        state.direct_to_ocean,
+        state.outflow_downstream_previous_timestep,
+        state.outflow_downstream_current_timestep,
+        state.outflow_before_regulation,
+        state.outflow_after_regulation,
+        state.outflow_sum_upstream_average,
+        state.lateral_flow_hillslope_average,
+        state.runoff,
+        state.direct,
+        state.flood,
+        state.runoff_land,
+        state.runoff_ocean,
+        state.delta_storage,
+        state.delta_storage_land,
+        state.delta_storage_ocean,
+        state.grid_cell_deficit,
+        state.storage,
+        state.channel_storage,
+        state.hillslope_wetland_runoff,
+        state.hillslope_subsurface_runoff,
+        state.hillslope_surface_runoff,
+        parameters.flood_threshold,
+        parameters.river_depth_minimum,
+    )
+    # send the direct water to outlet
+    state.direct[:] = pd.DataFrame(grid.id, columns=['id']).merge(
+        pd.DataFrame(state.direct, columns=['direct']).join(
+            pd.DataFrame(grid.outlet_id, columns=['outlet_id'])
+        ).groupby('outlet_id').sum(),
+        how='left',
+        left_on='id',
+        right_index=True
+    ).direct.fillna(0.0).values
+    subcycle_delta_t = config.get('simulation.timestep') / config.get('simulation.subcycles')
     for _ in np.arange(config.get('simulation.subcycles')):
-        
-        ###
-        ### hillslope routing
-        ###
-        hillslope_routing(state, grid, parameters, delta_t)
-        
-        # zero relevant state variables
-        state.channel_flow[:] = 0
-        state.channel_outflow_downstream_previous_timestep[:] = 0
-        state.channel_outflow_downstream_current_timestep[:] = 0
-        state.channel_outflow_sum_upstream_average[:] = 0
-        state.channel_lateral_flow_hillslope_average[:] = 0
-        
-        # get the demand volume for this substep
-        if config.get('water_management.enabled', False):
-            state.grid_cell_unmet_demand = state.grid_cell_demand_rate * delta_t
-        
-        # iterate substeps for remaining routing
+        _subcycle(
+            n,
+            config.get('water_management.enabled', False),
+            subcycle_delta_t,
+            grid.mosart_mask,
+            grid.hillslope,
+            grid.hillslope_manning,
+            grid.drainage_density,
+            grid.drainage_fraction,
+            grid.area,
+            state.euler_mask,
+            state.hillslope_depth,
+            state.hillslope_overland_flow,
+            state.hillslope_storage,
+            state.hillslope_surface_runoff,
+            state.hillslope_delta_storage,
+            state.hillslope_subsurface_runoff,
+            state.subnetwork_lateral_inflow,
+            state.channel_flow,
+            state.channel_outflow_downstream_previous_timestep,
+            state.channel_outflow_downstream_current_timestep,
+            state.channel_outflow_sum_upstream_average,
+            state.channel_lateral_flow_hillslope_average,
+            state.grid_cell_unmet_demand,
+            state.grid_cell_demand_rate,
+            parameters.tiny_value,
+        )
         for __ in np.arange(config.get('simulation.routing_iterations')):
-        
-            ###
-            ### subnetwork
-            ###
-            if config.get('water_management.enabled', False) and config.get('water_management.extraction_enabled', False):
-                subnetwork_irrigation(state, grid, parameters)
-            subnetwork_routing(state, grid, parameters, config, delta_t)
-            
-            ###
-            ### upstream interactions
-            ###
+            if config.get('water_management.enabled', False):
+                subnetwork_irrigation(
+                    n,
+                    grid.mosart_mask,
+                    grid.subnetwork_length,
+                    grid.subnetwork_width,
+                    state.euler_mask,
+                    state.subnetwork_depth,
+                    state.subnetwork_storage,
+                    state.grid_cell_unmet_demand,
+                    state.grid_cell_supply,
+                    state.subnetwork_cross_section_area,
+                    state.subnetwork_wetness_perimeter,
+                    state.subnetwork_hydraulic_radii,
+                    parameters.irrigation_extraction_parameter,
+                    parameters.tiny_value,
+                )
+            subnetwork_routing(
+                n,
+                subcycle_delta_t,
+                config.get('simulation.routing_iterations'),
+                np.nanmax(grid.iterations_subnetwork),
+                grid.iterations_subnetwork,
+                grid.mosart_mask,
+                grid.subnetwork_slope,
+                grid.subnetwork_manning,
+                grid.subnetwork_length,
+                grid.subnetwork_width,
+                grid.hillslope_length,
+                state.euler_mask,
+                state.channel_lateral_flow_hillslope,
+                state.subnetwork_flow_velocity,
+                state.subnetwork_discharge,
+                state.subnetwork_lateral_inflow,
+                state.subnetwork_storage,
+                state.subnetwork_storage_previous_timestep,
+                state.subnetwork_delta_storage,
+                state.subnetwork_depth,
+                state.subnetwork_cross_section_area,
+                state.subnetwork_wetness_perimeter,
+                state.subnetwork_hydraulic_radii,
+                parameters.tiny_value,
+            )
+
+            # upstream interactions
             state.channel_outflow_downstream_previous_timestep = state.channel_outflow_downstream_previous_timestep - state.channel_outflow_downstream
-            state.channel_outflow_sum_upstream_instant[:] = 0
-            
+            state.channel_outflow_sum_upstream_instant[:] = 0.0
             # send channel downstream outflow to downstream cells
-            state.channel_outflow_sum_upstream_instant[:] = pd.DataFrame(
-                grid.id, columns=['id']
-            ).merge(
-                pd.DataFrame(
-                    state.channel_outflow_downstream, columns=['channel_outflow_downstream']
-                ).join(
-                    pd.DataFrame(
-                        grid.downstream_id, columns=['downstream_id']
-                    )
+            state.channel_outflow_sum_upstream_instant[:] = pd.DataFrame(grid.id, columns=['id']).merge(
+                pd.DataFrame(state.channel_outflow_downstream, columns=['channel_outflow_downstream']).join(
+                    pd.DataFrame(grid.downstream_id, columns=['downstream_id'])
                 ).groupby('downstream_id').sum(),
                 how='left',
                 left_on='id',
                 right_index=True
             ).channel_outflow_downstream.fillna(0.0).values
-            
             state.channel_outflow_sum_upstream_average = state.channel_outflow_sum_upstream_average + state.channel_outflow_sum_upstream_instant
             state.channel_lateral_flow_hillslope_average = state.channel_lateral_flow_hillslope_average + state.channel_lateral_flow_hillslope
-            
-            ###
-            ### main channel
-            ###
-            main_channel_routing(state, grid, parameters, config, delta_t)
-            if config.get('water_management.enabled', False) and config.get('water_management.extraction_enabled', False) and config.get('water_management.extraction_main_channel_enabled', False):
-                main_channel_irrigation(state, grid, parameters)
-            
-            ###
-            ### regulation
-            ###
-            if config.get('water_management.enabled', False) and config.get('water_management.regulation_enabled', False):
-                regulation(state, grid, parameters, delta_t / config.get('simulation.routing_iterations'))
-            
+
+            main_channel_routing(
+                n,
+                subcycle_delta_t,
+                config.get('simulation.routing_iterations'),
+                np.nanmax(grid.iterations_main_channel),
+                grid.iterations_main_channel,
+                grid.mosart_mask,
+                grid.channel_length,
+                grid.channel_slope,
+                grid.channel_manning,
+                grid.total_drainage_area_single,
+                grid.channel_width,
+                grid.area,
+                grid.drainage_fraction,
+                grid.grid_channel_depth,
+                grid.channel_floodplain_width,
+                state.euler_mask,
+                state.channel_inflow_upstream,
+                state.channel_outflow_sum_upstream_instant,
+                state.channel_hydraulic_radii,
+                state.channel_flow_velocity,
+                state.channel_outflow_downstream,
+                state.channel_cross_section_area,
+                state.channel_lateral_flow_hillslope,
+                state.channel_storage,
+                state.channel_storage_previous_timestep,
+                state.hillslope_wetland_runoff,
+                state.channel_delta_storage,
+                state.channel_depth,
+                state.channel_wetness_perimeter,
+                parameters.slope_1_def,
+                parameters.inverse_sin_atan_slope_1_def,
+                parameters.tiny_value,
+                parameters.kinematic_wave_parameter,
+            )
+            if config.get('water_management.enabled', False):
+                main_channel_irrigation(
+                    n,
+                    grid.mosart_mask,
+                    grid.channel_length,
+                    grid.grid_channel_depth,
+                    grid.channel_width,
+                    grid.channel_floodplain_width,
+                    state.euler_mask,
+                    state.channel_depth,
+                    state.channel_storage,
+                    state.grid_cell_unmet_demand,
+                    state.grid_cell_supply,
+                    state.channel_cross_section_area,
+                    state.channel_wetness_perimeter,
+                    state.channel_hydraulic_radii,
+                    parameters.tiny_value,
+                    parameters.tinier_value,
+                    parameters.slope_1_def,
+                    parameters.inverse_sin_atan_slope_1_def,
+                    parameters.irrigation_extraction_parameter,
+                    parameters.irrigation_extraction_maximum_fraction,
+                )
+                regulation(
+                    n,
+                    subcycle_delta_t / config.get('simulation.routing_iterations'),
+                    grid.mosart_mask,
+                    grid.reservoir_id,
+                    grid.reservoir_surface_area,
+                    grid.reservoir_storage_capacity,
+                    state.euler_mask,
+                    state.channel_outflow_downstream,
+                    state.reservoir_release,
+                    state.reservoir_potential_evaporation,
+                    state.reservoir_streamflow,
+                    state.reservoir_storage,
+                    parameters.reservoir_runoff_capacity_parameter,
+                )
+
             # update flow
-            base_condition = (grid.mosart_mask > 0) & state.euler_mask
+            mask = np.logical_and(state.euler_mask, grid.mosart_mask > 0)
             state.channel_outflow_downstream_current_timestep = np.where(
-                base_condition,
+                mask,
                 state.channel_outflow_downstream_current_timestep - state.channel_outflow_downstream,
                 state.channel_outflow_downstream_current_timestep
             )
             state.channel_flow = np.where(
-                base_condition,
+                mask,
                 state.channel_flow - state.channel_outflow_downstream,
                 state.channel_flow
             )
-        
-        # average state values over dlevelh2r
-        state.channel_flow = state.channel_flow / config.get('simulation.routing_iterations')
-        state.channel_outflow_downstream_previous_timestep = state.channel_outflow_downstream_previous_timestep / config.get('simulation.routing_iterations')
-        state.channel_outflow_downstream_current_timestep = state.channel_outflow_downstream_current_timestep / config.get('simulation.routing_iterations')
-        state.channel_outflow_sum_upstream_average = state.channel_outflow_sum_upstream_average / config.get('simulation.routing_iterations')
-        state.channel_lateral_flow_hillslope_average = state.channel_lateral_flow_hillslope_average / config.get('simulation.routing_iterations')
-        
-        # regulation extraction
-        if config.get('water_management.enabled', False) and config.get('water_management.regulation_enabled', False):
-            state.outflow_before_regulation = -state.channel_outflow_downstream
-            state.channel_flow = state.channel_flow + state.channel_outflow_downstream
-            if config.get('water_management.extraction_enabled', False):
-                extraction_regulated_flow(state, grid, parameters, config, delta_t)
-            state.outflow_after_regulation = -state.channel_outflow_downstream
-            state.channel_flow = state.channel_flow - state.channel_outflow_downstream
-            # aggregate deficit
-            state.grid_cell_deficit = state.grid_cell_deficit + state.grid_cell_unmet_demand
-        
-        # accumulate local flow field
-        state.flow = state.flow + state.channel_flow
-        state.outflow_downstream_previous_timestep = state.outflow_downstream_previous_timestep + state.channel_outflow_downstream_previous_timestep
-        state.outflow_downstream_current_timestep = state.outflow_downstream_current_timestep + state.channel_outflow_downstream_current_timestep
-        state.outflow_before_regulation = state.outflow_before_regulation + state.channel_outflow_before_regulation
-        state.outflow_after_regulation = state.outflow_after_regulation + state.channel_outflow_after_regulation
-        state.outflow_sum_upstream_average = state.outflow_sum_upstream_average + state.channel_outflow_sum_upstream_average
-        state.lateral_flow_hillslope_average = state.lateral_flow_hillslope_average + state.channel_lateral_flow_hillslope_average
-    
-    if config.get('water_management.enabled'):
-        # convert supply to flux
-        state.grid_cell_supply = state.grid_cell_supply / config.get('simulation.timestep')
-    
-    # convert runoff back to m3/s for output
-    state.hillslope_surface_runoff = state.hillslope_surface_runoff * grid.area
-    state.hillslope_subsurface_runoff = state.hillslope_subsurface_runoff * grid.area
-    state.hillslope_wetland_runoff = state.hillslope_wetland_runoff * grid.area
-    
-    # average state values over subcycles
-    state.flow = state.flow / config.get('simulation.subcycles')
-    state.outflow_downstream_previous_timestep = state.outflow_downstream_previous_timestep / config.get('simulation.subcycles')
-    state.outflow_downstream_current_timestep = state.outflow_downstream_current_timestep / config.get('simulation.subcycles')
-    state.outflow_before_regulation = state.outflow_before_regulation / config.get('simulation.subcycles')
-    state.outflow_after_regulation = state.outflow_after_regulation / config.get('simulation.subcycles')
-    state.outflow_sum_upstream_average = state.outflow_sum_upstream_average / config.get('simulation.subcycles')
-    state.lateral_flow_hillslope_average = state.lateral_flow_hillslope_average / config.get('simulation.subcycles')
-    
-    # update state values
-    previous_storage = 1.0 * state.storage
-    state.storage = (state.channel_storage + state.subnetwork_storage + state.hillslope_storage * grid.area) * grid.drainage_fraction
-    state.delta_storage = (state.storage - previous_storage) / config.get('simulation.timestep')
-    state.runoff[:] = state.flow
-    state.runoff_total[:] = state.direct
-    state.runoff_land = np.where(
-        grid.land_mask == 1,
-        state.runoff,
-        0
-    )
-    state.delta_storage_land = np.where(
-        grid.land_mask == 1,
+
+        _average_over_routing_iterations(
+            n,
+            config.get('simulation.routing_iterations'),
+            state.channel_flow,
+            state.channel_outflow_downstream_previous_timestep,
+            state.channel_outflow_downstream_current_timestep,
+            state.channel_outflow_sum_upstream_average,
+            state.channel_lateral_flow_hillslope_average
+        )
+
+        if config.get('water_management.enabled', False):
+            extraction_regulated_flow(
+                n,
+                int(np.nanmax(grid.reservoir_id) - 1),
+                subcycle_delta_t,
+                grid.id,
+                grid.reservoir_id,
+                grid.reservoir_to_grid_map,
+                state.channel_outflow_downstream,
+                state.outflow_before_regulation,
+                state.outflow_after_regulation,
+                state.channel_flow,
+                state.grid_cell_unmet_demand,
+                state.grid_cell_supply,
+                state.grid_cell_deficit,
+                parameters.reservoir_supply_iterations,
+                parameters.reservoir_flow_volume_ratio,
+            )
+
+        _accumulate_flow_field(
+            n,
+            state.flow,
+            state.channel_flow,
+            state.outflow_downstream_previous_timestep,
+            state.outflow_downstream_current_timestep,
+            state.channel_outflow_downstream_previous_timestep,
+            state.channel_outflow_downstream_current_timestep,
+            state.outflow_before_regulation,
+            state.outflow_after_regulation,
+            state.channel_outflow_before_regulation,
+            state.channel_outflow_after_regulation,
+            state.outflow_sum_upstream_average,
+            state.channel_outflow_sum_upstream_average,
+            state.lateral_flow_hillslope_average,
+            state.channel_lateral_flow_hillslope_average,
+        )
+
+    _finalize(
+        n,
+        config.get('water_management.enabled', False),
+        config.get('simulation.timestep'),
+        config.get('simulation.subcycles'),
+        grid.land_mask,
+        grid.area,
+        grid.drainage_fraction,
+        state.grid_cell_supply,
+        state.hillslope_surface_runoff,
+        state.hillslope_subsurface_runoff,
+        state.hillslope_wetland_runoff,
+        state.flow,
+        state.outflow_downstream_previous_timestep,
+        state.outflow_downstream_current_timestep,
+        state.outflow_before_regulation,
+        state.outflow_after_regulation,
+        state.outflow_sum_upstream_average,
+        state.lateral_flow_hillslope_average,
+        state.storage,
+        state.channel_storage,
+        state.subnetwork_storage,
+        state.hillslope_storage,
         state.delta_storage,
-        0
-    )
-    state.runoff_ocean = np.where(
-        grid.land_mask >= 2,
+        state.direct,
         state.runoff,
-        0
+        state.runoff_total,
+        state.runoff_land,
+        state.delta_storage_land,
+        state.runoff_ocean,
+        state.delta_storage_ocean,
     )
-    state.runoff_total = np.where(
-        grid.land_mask >= 2,
-        state.runoff_total + state.runoff,
-        state.runoff_total
-    )
-    state.delta_storage_ocean = np.where(
-        grid.land_mask >= 2,
-        state.delta_storage,
-        0
-    )
\ No newline at end of file
+
+
+@nb.jit(
+    "void("
+        "int64, boolean, int64, int64[:], int64[:], float64[:], float64[:], float64[:], float64[:],"
+        "float64[:], float64[:], float64[:], float64[:], float64[:], float64[:], float64[:], float64[:], float64[:],"
+        "float64[:], float64[:], float64[:], float64[:], float64[:], float64[:], float64[:], float64[:], float64[:],"
+        " float64[:], float64, float64"
+    ")",
+    parallel=True,
+    nopython=True,
+    nogil=True,
+    cache=True,
+)
+def _prepare(
+    n,
+    water_management_enabled,
+    timestep,
+    land_mask,
+    mosart_mask,
+    area,
+    flow,
+    direct_to_ocean,
+    outflow_downstream_previous_timestep,
+    outflow_downstream_current_timestep,
+    outflow_before_regulation,
+    outflow_after_regulation,
+    outflow_sum_upstream_average,
+    lateral_flow_hillslope_average,
+    runoff,
+    direct,
+    flood,
+    runoff_land,
+    runoff_ocean,
+    delta_storage,
+    delta_storage_land,
+    delta_storage_ocean,
+    grid_cell_deficit,
+    storage,
+    channel_storage,
+    hillslope_wetland_runoff,
+    hillslope_subsurface_runoff,
+    hillslope_surface_runoff,
+    flood_threshold,
+    river_depth_minimum,
+):
+    for i in nb.prange(n):
+        # Reset certain state variables
+        flow[i] = 0.0
+        outflow_downstream_previous_timestep[i] = 0.0
+        outflow_downstream_current_timestep[i] = 0.0
+        outflow_before_regulation[i] = 0.0
+        outflow_after_regulation[i] = 0.0
+        outflow_sum_upstream_average[i] = 0.0
+        lateral_flow_hillslope_average[i] = 0.0
+        runoff[i] = 0.0
+        direct[i] = 0.0
+        flood[i] = 0.0
+        runoff_land[i] = 0.0
+        runoff_ocean[i] = 0.0
+        delta_storage[i] = 0.0
+        delta_storage_land[i] = 0.0
+        delta_storage_ocean[i] = 0.0
+        if water_management_enabled:
+            grid_cell_deficit[i] = 0.0
+
+        ###
+        ### flood
+        ###
+        if (land_mask[i] == 1) and (storage[i] > flood_threshold):
+            flood[i] = (storage[i] - flood_threshold) / timestep
+            # remove this flux from the input runoff from land
+            hillslope_surface_runoff[i] = hillslope_surface_runoff[i] - flood[i]
+        else:
+            flood[i] = 0.0
+
+        ###
+        ### direct to ocean
+        ###
+        # note - in fortran mosart this direct_to_ocean forcing could be provided from LND component, but we don't seem to be using it
+        source_direct = 1.0 * direct_to_ocean[i]
+
+        # wetland runoff
+        wetland_runoff_volume = hillslope_wetland_runoff[i] * timestep
+        river_volume_minimum = river_depth_minimum * area[i]
+
+        # if wetland runoff is negative and it would bring main channel storage below the minimum, send it directly to ocean
+        if ((channel_storage[i] + wetland_runoff_volume) < river_volume_minimum) and (hillslope_wetland_runoff[i] < 0.0):
+            source_direct = source_direct + hillslope_wetland_runoff[i]
+            hillslope_wetland_runoff[i] = 0.0
+
+        # remove remaining wetland runoff (negative and positive)
+        source_direct = source_direct + hillslope_wetland_runoff[i]
+        hillslope_wetland_runoff[i] = 0.0
+
+        # runoff from hillslope
+        # remove negative subsurface water
+        if hillslope_subsurface_runoff[i] < 0.0:
+            source_direct = source_direct + hillslope_subsurface_runoff[i]
+            hillslope_subsurface_runoff[i] = 0.0
+
+        # remove negative surface water
+        if hillslope_surface_runoff[i] < 0.0:
+            source_direct = source_direct + hillslope_surface_runoff[i]
+            hillslope_surface_runoff[i] = 0.0
+
+        # if ocean cell remove the rest of the sub and surface water
+        # other cells will be handled by mosart euler
+        if mosart_mask[i] == 0:
+            source_direct = source_direct + hillslope_subsurface_runoff[i] + hillslope_surface_runoff[i]
+            hillslope_subsurface_runoff[i] = 0.0
+            hillslope_surface_runoff[i] = 0.0
+
+        direct[i] = source_direct
+
+        # convert runoff from m3/s to m/s
+        hillslope_surface_runoff[i] = hillslope_surface_runoff[i] / area[i]
+        hillslope_subsurface_runoff[i] = hillslope_subsurface_runoff[i] / area[i]
+        hillslope_wetland_runoff[i] = hillslope_wetland_runoff[i] / area[i]
+
+
+@nb.jit(
+    "void("
+        "int64, boolean, float64, int64[:], float64[:], float64[:], float64[:], float64[:], float64[:],"
+        "boolean[:], float64[:], float64[:], float64[:], float64[:], float64[:], float64[:], float64[:],"
+        "float64[:], float64[:], float64[:], float64[:], float64[:], float64[:], float64[:], float64"
+    ")",
+    parallel=True,
+    nopython=True,
+    nogil=True,
+    cache=True,
+)
+def _subcycle(
+    n,
+    water_management_enabled,
+    delta_t,
+    mosart_mask,
+    hillslope,
+    hillslope_manning,
+    drainage_density,
+    drainage_fraction,
+    area,
+    euler_mask,
+    hillslope_depth,
+    hillslope_overland_flow,
+    hillslope_storage,
+    hillslope_surface_runoff,
+    hillslope_delta_storage,
+    hillslope_subsurface_runoff,
+    subnetwork_lateral_inflow,
+    channel_flow,
+    channel_outflow_downstream_previous_timestep,
+    channel_outflow_downstream_current_timestep,
+    channel_outflow_sum_upstream_average,
+    channel_lateral_flow_hillslope_average,
+    grid_cell_unmet_demand,
+    grid_cell_demand_rate,
+    tiny_value,
+):
+    for i in nb.prange(n):
+
+        ###
+        ### hillslope routing
+        ###
+        hillslope_routing(
+            i,
+            delta_t,
+            mosart_mask,
+            hillslope,
+            hillslope_manning,
+            drainage_density,
+            drainage_fraction,
+            area,
+            euler_mask,
+            hillslope_depth,
+            hillslope_overland_flow,
+            hillslope_storage,
+            hillslope_surface_runoff,
+            hillslope_delta_storage,
+            hillslope_subsurface_runoff,
+            subnetwork_lateral_inflow,
+            tiny_value,
+        )
+
+        # zero relevant state variables
+        channel_flow[i] = 0.0
+        channel_outflow_downstream_previous_timestep[i] = 0.0
+        channel_outflow_downstream_current_timestep[i] = 0.0
+        channel_outflow_sum_upstream_average[i] = 0.0
+        channel_lateral_flow_hillslope_average[i] = 0.0
+
+        # get the demand volume for this substep
+        if water_management_enabled:
+            grid_cell_unmet_demand[i] = grid_cell_demand_rate[i] * delta_t
+
+
+@nb.jit(
+    "void("
+        "int64, int64,"
+        "float64[:], float64[:], float64[:], float64[:], float64[:]"
+    ")",
+    parallel=True,
+    nopython=True,
+    nogil=True,
+    cache=True,
+)
+def _average_over_routing_iterations(
+    n,
+    routing_iterations,
+    channel_flow,
+    channel_outflow_downstream_previous_timestep,
+    channel_outflow_downstream_current_timestep,
+    channel_outflow_sum_upstream_average,
+    channel_lateral_flow_hillslope_average,
+):
+    for i in nb.prange(n):
+        # average state values over the routing iterations
+        channel_flow[i] = channel_flow[i] / routing_iterations
+        channel_outflow_downstream_previous_timestep[i] = channel_outflow_downstream_previous_timestep[i] / routing_iterations
+        channel_outflow_downstream_current_timestep[i] = channel_outflow_downstream_current_timestep[i] / routing_iterations
+        channel_outflow_sum_upstream_average[i] = channel_outflow_sum_upstream_average[i] / routing_iterations
+        channel_lateral_flow_hillslope_average[i] = channel_lateral_flow_hillslope_average[i] / routing_iterations
+
+
+@nb.jit(
+    "void("
+        "int64, float64[:], float64[:], float64[:], float64[:], float64[:], float64[:], float64[:],"
+        "float64[:], float64[:], float64[:], float64[:], float64[:], float64[:], float64[:]"
+    ")",
+    parallel=True,
+    nopython=True,
+    nogil=True,
+    cache=True,
+)
+def _accumulate_flow_field(
+    n,
+    flow,
+    channel_flow,
+    outflow_downstream_previous_timestep,
+    outflow_downstream_current_timestep,
+    channel_outflow_downstream_previous_timestep,
+    channel_outflow_downstream_current_timestep,
+    outflow_before_regulation,
+    outflow_after_regulation,
+    channel_outflow_before_regulation,
+    channel_outflow_after_regulation,
+    outflow_sum_upstream_average,
+    channel_outflow_sum_upstream_average,
+    lateral_flow_hillslope_average,
+    channel_lateral_flow_hillslope_average,
+):
+    for i in nb.prange(n):
+        # accumulate local flow field
+        flow[i] = flow[i] + channel_flow[i]
+        outflow_downstream_previous_timestep[i] = outflow_downstream_previous_timestep[i] + channel_outflow_downstream_previous_timestep[i]
+        outflow_downstream_current_timestep[i] = outflow_downstream_current_timestep[i] + channel_outflow_downstream_current_timestep[i]
+        outflow_before_regulation[i] = outflow_before_regulation[i] + channel_outflow_before_regulation[i]
+        outflow_after_regulation[i] = outflow_after_regulation[i] + channel_outflow_after_regulation[i]
+        outflow_sum_upstream_average[i] = outflow_sum_upstream_average[i] + channel_outflow_sum_upstream_average[i]
+        lateral_flow_hillslope_average[i] = lateral_flow_hillslope_average[i] + channel_lateral_flow_hillslope_average[i]
+
+
+@nb.jit(
+    "void("
+        "int64, boolean, int64, int64, int64[:], float64[:], float64[:], float64[:], float64[:],"
+        "float64[:], float64[:], float64[:], float64[:], float64[:], float64[:], float64[:],"
+        "float64[:], float64[:], float64[:], float64[:], float64[:], float64[:], float64[:],"
+        "float64[:], float64[:], float64[:], float64[:], float64[:], float64[:], float64[:],"
+    ")",
+    parallel=True,
+    nopython=True,
+    nogil=True,
+    cache=True,
+)
+def _finalize(
+    n,
+    water_management_enabled,
+    timestep,
+    subcycles,
+    land_mask,
+    area,
+    drainage_fraction,
+    grid_cell_supply,
+    hillslope_surface_runoff,
+    hillslope_subsurface_runoff,
+    hillslope_wetland_runoff,
+    flow,
+    outflow_downstream_previous_timestep,
+    outflow_downstream_current_timestep,
+    outflow_before_regulation,
+    outflow_after_regulation,
+    outflow_sum_upstream_average,
+    lateral_flow_hillslope_average,
+    storage,
+    channel_storage,
+    subnetwork_storage,
+    hillslope_storage,
+    delta_storage,
+    direct,
+    runoff,
+    runoff_total,
+    runoff_land,
+    delta_storage_land,
+    runoff_ocean,
+    delta_storage_ocean,
+):
+    for i in nb.prange(n):
+        if water_management_enabled:
+            # convert supply to flux
+            grid_cell_supply[i] = grid_cell_supply[i] / timestep
+
+        # convert runoff back to m3/s for output
+        hillslope_surface_runoff[i] = hillslope_surface_runoff[i] * area[i]
+        hillslope_subsurface_runoff[i] = hillslope_subsurface_runoff[i] * area[i]
+        hillslope_wetland_runoff[i] = hillslope_wetland_runoff[i] * area[i]
+
+        # average state values over the subcycles
+        flow[i] = flow[i] / subcycles
+        outflow_downstream_previous_timestep[i] = outflow_downstream_previous_timestep[i] / subcycles
+        outflow_downstream_current_timestep[i] = outflow_downstream_current_timestep[i] / subcycles
+        outflow_before_regulation[i] = outflow_before_regulation[i] / subcycles
+        outflow_after_regulation[i] = outflow_after_regulation[i] / subcycles
+        outflow_sum_upstream_average[i] = outflow_sum_upstream_average[i] / subcycles
+        lateral_flow_hillslope_average[i] = lateral_flow_hillslope_average[i] / subcycles
+
+        # update state values
+        previous_storage = 1.0 * storage[i]
+        storage[i] = (channel_storage[i] + subnetwork_storage[i] + hillslope_storage[i] * area[i]) * drainage_fraction[i]
+        delta_storage[i] = (storage[i] - previous_storage) / timestep
+        runoff[i] = 1.0 * flow[i]
+        runoff_total[i] = 1.0 * direct[i]
+        if land_mask[i] == 1:
+            runoff_land[i] = 1.0 * runoff[i]
+            delta_storage_land[i] = 1.0 * delta_storage[i]
+        else:
+            runoff_land[i] = 0.0
+            delta_storage_land[i] = 0.0
+        if land_mask[i] >= 2:
+            runoff_ocean[i] = 1.0 * runoff[i]
+            runoff_total[i] = runoff_total[i] + runoff[i]
+            delta_storage_ocean[i] = 1.0 * delta_storage[i]
+        else:
+            runoff_ocean[i] = 0.0
+            delta_storage_ocean[i] = 0.0
diff --git a/mosartwmpy/utilities/timing.py b/mosartwmpy/utilities/timing.py
index 04434d1..11b65e7 100644
--- a/mosartwmpy/utilities/timing.py
+++ b/mosartwmpy/utilities/timing.py
@@ -5,6 +5,7 @@
 
 from mosartwmpy.utilities.pretty_timer import pretty_timer
 
+
 def timing(f: Callable) -> Callable:
     """Decorator for timing a method and outputting to log.
 
@@ -21,4 +22,16 @@ def wrap(*args, **kw):
         seconds = timer() - t
         logging.info(f'{f.__name__}: {pretty_timer(seconds)}')
         return result
-    return wrap
\ No newline at end of file
+    return wrap
+
+
+class Timer:
+    def __init__(self, name=None):
+        self.name = " '" + name + "'" if name else ''
+
+    def __enter__(self):
+        self.start = timer()
+
+    def __exit__(self, exc_type, exc_value, traceback):
+        self.took = (timer() - self.start) * 1000.0
+        logging.info('Code block' + self.name + ' took: ' + str(self.took) + ' ms')
diff --git a/mosartwmpy/validate.py b/mosartwmpy/validate.py
index 08068f7..366f8f3 100644
--- a/mosartwmpy/validate.py
+++ b/mosartwmpy/validate.py
@@ -1,15 +1,17 @@
+import pdb
+
 import matplotlib.pyplot as plt
 import numpy as np
 import os
 from pathlib import Path
-from xarray import open_mfdataset
+from xarray import open_dataset, open_mfdataset
 
 # TODO accept command line path input as alternative
 # TODO accept command line year input
 # TODO allow easily toggling between scenarios for variables of interest (no-wm, wm, heat, etc)
 
 years = [1981, 1982]
-baseline_data_path = Path('validation/mosartwmpy_validation_wm_1981_1982.nc')
+baseline_data_path = str(Path('validation/mosartwmpy_validation_wm_1981_1982.nc'))
 variables_of_interest = ['STORAGE_LIQ', 'RIVER_DISCHARGE_OVER_LAND_LIQ', 'WRM_STORAGE', 'WRM_SUPPLY']
 physical_dimensions = ['lat', 'lon']
 temporal_dimension = 'time'
@@ -33,7 +35,7 @@
     data = data.sel({temporal_dimension: slice(f"{years[0]}", f"{years[-1]}")})
     timeslice = slice(data[temporal_dimension].values[0], data[temporal_dimension].values[len(data[temporal_dimension].values) - 1])
     baseline_data = open_mfdataset(baseline_data_path)
-    baseline_data = baseline_data.sel({ temporal_dimension: timeslice })
+    baseline_data = baseline_data.sel({temporal_dimension: timeslice})
 except:
     print(f"Either your data or the baseline data does not appear to include years {years}.")
     print("Please double check and try again or update this code to look for a different year.")
diff --git a/setup.py b/setup.py
index 8da9505..59bab65 100644
--- a/setup.py
+++ b/setup.py
@@ -1,3 +1,4 @@
+import re
 from setuptools import setup, find_packages
 
 
@@ -7,9 +8,11 @@ def readme():
         return f.read()
 
 
+version = re.search(r"__version__ = ['\"]([^'\"]*)['\"]", open('mosartwmpy/_version.py').read(), re.M).group(1)
+
 setup(
     name='mosartwmpy',
-    version='0.0.8',
+    version=version,
     packages=find_packages(),
     url='https://github.com/IMMM-SFA/mosartwmpy',
     license='BSD2-Simplified',
@@ -22,33 +25,34 @@ def readme():
     include_package_data=True,
     install_requires=[
         'bmipy==2.0',
-        'dask[complete]==2021.5.1',
-        'ipython==7.24.1',
+        'dask[complete]==2021.6.2',
+        'h5netcdf==0.11.0',
         'matplotlib==3.4.2',
-        'nc-time-axis==1.2.0',
-        'netCDF4==1.5.6',
-        'numexpr==2.7.3',
-        'numpy==1.20.3',
-        'pandas==1.2.4',
-        'pathvalidate==2.3.0',
+        'nc-time-axis==1.3.1',
+        'netCDF4==1.5.7',
+        'numba==0.53.1',
+        'numpy==1.21.0',
+        'pandas==1.3.0',
+        'pathvalidate==2.4.1',
         'psutil==5.8.0',
         'pyarrow==4.0.1',
         'python-benedict==0.24.0',
-        'regex==2021.4.4',
+        'regex==2021.7.6',
         'requests==2.25.1',
-        'rioxarray==0.4.2',
-        'tqdm==4.60.0',
+        'rioxarray==0.5.0',
+        'tqdm==4.61.2',
         'xarray==0.18.2'
     ],
     extras_require={
         'dev': [
-            'build==0.4.0',
-            'nbsphinx==0.8.6',
-            'recommonmark==0.7.1',
-            'setuptools==54.1.2',
-            'sphinx==3.5.2',
-            'sphinx-rtd-theme==0.5.1',
-            'twine==3.4.1'
+            'build~=0.5.1',
+            'nbsphinx~=0.8.6',
+            'recommonmark~=0.7.1',
+            'setuptools~=57.0.0',
+            'sphinx~=4.0.2',
+            'sphinx-panels~=0.6.0',
+            'sphinx-rtd-theme~=0.5.2',
+            'twine~=3.4.1'
         ]
     }
 )


Test patch:
diff --git a/mosartwmpy/tests/grid.zip b/mosartwmpy/tests/grid.zip
index 2b73a98..64cc605 100644
Binary files a/mosartwmpy/tests/grid.zip and b/mosartwmpy/tests/grid.zip differ
diff --git a/notebooks/speed_test_grid.zip b/notebooks/speed_test_grid.zip
new file mode 100644
index 0000000..2b73a98
Binary files /dev/null and b/notebooks/speed_test_grid.zip differ
diff --git a/notebooks/speed_test_state.nc b/notebooks/speed_test_state.nc
new file mode 100644
index 0000000..5c8c9d7
Binary files /dev/null and b/notebooks/speed_test_state.nc differ