From c05f0b47ad6ed7b1147d3c6c88246df8ba38fd98 Mon Sep 17 00:00:00 2001
From: elbeejay <jhariharan@usgs.gov>
Date: Mon, 17 Jul 2023 15:39:49 -0400
Subject: [PATCH] start setting up runoff functions

---
 docs/source/conf.py             |  1 +
 docs/source/reference/index.rst |  7 +++
 hyswap/__init__.py              |  1 +
 hyswap/runoff.py                | 80 +++++++++++++++++++++++++++++++++
 tests/test_runoff.py            | 17 +++++++
 5 files changed, 106 insertions(+)
 create mode 100644 hyswap/runoff.py
 create mode 100644 tests/test_runoff.py

diff --git a/docs/source/conf.py b/docs/source/conf.py
index 09a26a8..8d58389 100644
--- a/docs/source/conf.py
+++ b/docs/source/conf.py
@@ -92,6 +92,7 @@
 from hyswap import percentiles
 from hyswap import cumulative
 from hyswap import plots
+from hyswap import runoff
 import numpy as np
 import pandas as pd
 import matplotlib.pyplot as plt
diff --git a/docs/source/reference/index.rst b/docs/source/reference/index.rst
index 150e118..0ae5053 100644
--- a/docs/source/reference/index.rst
+++ b/docs/source/reference/index.rst
@@ -55,3 +55,10 @@ Plotting Functions
 .. automodule:: hyswap.plots
     :members:
     :special-members:
+
+Runoff Calculation Functions
+-----------------------------
+
+.. automodule:: hyswap.runoff
+    :members:
+    :special-members:
diff --git a/hyswap/__init__.py b/hyswap/__init__.py
index 8c66d9b..67833d8 100644
--- a/hyswap/__init__.py
+++ b/hyswap/__init__.py
@@ -6,6 +6,7 @@
 from hyswap.percentiles import *  # noqa
 from hyswap.cumulative import *  # noqa
 from hyswap.plots import *  # noqa
+from hyswap.runoff import *  # noqa
 
 try:
     __version__ = version('hyswap')
diff --git a/hyswap/runoff.py b/hyswap/runoff.py
new file mode 100644
index 0000000..634d37a
--- /dev/null
+++ b/hyswap/runoff.py
@@ -0,0 +1,80 @@
+"""Runoff functions for hyswap."""
+
+
+def convert_cfs_to_mmyr(cfs, drainage_area):
+    """Convert cfs to mm/yr.
+
+    Parameters
+    ----------
+    cfs : float
+        Flow in cubic feet per second.
+
+    drainage_area : float
+        Drainage area in km2.
+
+    Returns
+    -------
+    float
+        Runoff in mm/yr.
+
+    Examples
+    --------
+    Convert 14 cfs to mm/yr for a 250 km2 drainage area.
+
+    .. doctest::
+
+        >>> mmyr = runoff.convert_cfs_to_mmyr(14, 250)
+        >>> np.round(mmyr)
+        50.0
+    """
+    # convert cfs to cubic feet per year
+    cpy = cfs * 60 * 60 * 24 * 365.25
+    # convert cubic feet per year to cubic meters per year
+    cpy = cpy * (0.3048 ** 3)
+    # convert drainage area km2 to m2
+    drainage_area = drainage_area * 1000 * 1000
+    # convert cubic meters per year to mm per year
+    mmyr = cpy / drainage_area * 1000
+    return mmyr
+
+
+def streamflow_to_runoff(df, data_col, drainage_area):
+    """Convert streamflow to runoff for a given drainage area.
+
+    For a given gage/dataframe, convert streamflow to runoff using the
+    drainage area and the convert_cfs_to_mmyr function.
+
+    Parameters
+    ----------
+    df : pandas.DataFrame
+        DataFrame containing streamflow data.
+
+    data_col : str
+        Column name containing streamflow data, assumed to be in cfs.
+
+    drainage_area : float
+        Drainage area in km2.
+
+    Returns
+    -------
+    pandas.DataFrame
+        DataFrame containing runoff data in a column named 'runoff'.
+
+    Examples
+    --------
+    Convert streamflow to runoff for a given drainage area.
+
+    .. doctest::
+
+        >>> df = pd.DataFrame({'streamflow': [14, 15, 16]})
+        >>> runoff_df = runoff.streamflow_to_runoff(df, 'streamflow', 250)
+        >>> print(runoff_df['runoff'].round(1))
+        0    50.0
+        1    53.6
+        2    57.2
+        Name: runoff, dtype: float64
+    """
+    df['runoff'] = df[data_col].apply(
+        lambda x: convert_cfs_to_mmyr(x, drainage_area)
+    )
+    return df
diff --git a/tests/test_runoff.py b/tests/test_runoff.py
new file mode 100644
index 0000000..797801f
--- /dev/null
+++ b/tests/test_runoff.py
@@ -0,0 +1,17 @@
+"""Tests for the runoff.py module."""
+import pytest
+from hyswap import runoff
+import pandas as pd
+
+
+def test_convert_cfs_to_mmyr():
+    """Test the convert_cfs_to_mmyr function."""
+    mmyr = runoff.convert_cfs_to_mmyr(14, 250)
+    assert pytest.approx(mmyr, 0.1) == 50.0
+
+
+def test_streamflow_to_runoff():
+    """Test the streamflow_to_runoff function."""
+    df = pd.DataFrame({"streamflow": [14, 15, 16]})
+    runoff_df = runoff.streamflow_to_runoff(df, "streamflow", 250)
+    assert pytest.approx(runoff_df["runoff"].round(1)) == [50.0, 53.6, 57.2]