Merge branch 'performance-pre-init' of https://github.com/Ecohydrolog…

…yTeam/ClearWater-modules into performance-pre-init

src/clearwater_modules/__init__.py

@@ -1,3 +1,4 @@
 __version__ = "0.2.0"
 from clearwater_modules import shared
 from clearwater_modules import tsm
+from clearwater_modules import nsm1

src/clearwater_modules/nsm1/CBOD/dynamic_variables.py

@@ -1,32 +1,34 @@
-# TODO: figure out imports...
+"""
+File contains dynamic variables related to the CBOD module
+"""
 
 import clearwater_modules.shared.processes as shared_processes
 from clearwater_modules import base
-from clearwater_modules.tsm.model import EnergyBudget
-from clearwater_modules.nsm1.CBOD import processes
+from clearwater_modules.nsm1.model import NutrientBudget
+import clearwater_modules.nsm1.CBOD.processes as processes
 
 
-@base.register_variable(models=EnergyBudget)
+@base.register_variable(models=NutrientBudget)
 class Variable(base.Variable):
     ...
 
 
 Variable(
-    name='kbod_T',
+    name='kbod_tc',
     long_name='Temperature adjusted oxidation rate',
     units='1/d',
     description='Temperature adjusted oxidation rate',
     use='dynamic',
-    process=processes.kbod_T
+    process=processes.kbod_tc
 )
 
 Variable(
-    name='ksbod_T',
+    name='ksbod_tc',
     long_name='Temperature adjusted sedimentation rate',
     units='m/d',
     description='Temperature adjusted sedimentation rate',
     use='dynamic',
-    process=processes.ksbod_T
+    process=processes.ksbod_tc
 )
 
 Variable(

src/clearwater_modules/nsm1/CBOD/processes.py

@@ -1,52 +1,49 @@
-import numpy as np
+"""
+File contains process to calculate new CBOD concentration and associated dependent variables
+"""
+
 import numba
 import xarray as xr
-from clearwater_modules.shared.processes import (
-    arrhenius_correction,
-)
-
+from clearwater_modules.shared.processes import arrhenius_correction
+import math
 
 @numba.njit
-def kbod_T(
-    water_temp_c: xr.DataArray,
+def kbod_tc(
+    TwaterC: xr.DataArray,
     kbod_20: xr.DataArray,
-    theta: xr.DataArray
 ) -> xr.DataArray:
     """Calculate the temperature adjusted CBOD oxidation rate (1/d)
 
     Args:
-        water_temp_c: water temperature in Celsius
+        TwaterC: water temperature in Celsius
         kbod_20: CBOD oxidation rate at 20 degrees Celsius (1/d)
-        theta: Arrhenius coefficient
     """
 
-    kbod_T = arrhenius_correction(water_temp_c, kbod_20, theta)
-    return kbod_T
+    kbod_tc = arrhenius_correction(TwaterC, kbod_20, 1.047)
+    return kbod_tc
 
 
 @numba.njit
-def ksbod_T(
-    water_temp_c: xr.DataArray,
+def ksbod_tc(
+    TwaterC: xr.DataArray,
     ksbod_20: xr.DataArray,
-    theta: xr.DataArray
 ) -> xr.DataArray:
     """Calculate the temperature adjusted CBOD sedimentation rate (m/d)
 
     Args:
-        water_temp_c: water temperature in Celsius
+        TwaterC: water temperature in Celsius
         ksbod_20: CBOD sedimentation rate at 20 degrees Celsius (m/d)
-        theta: Arrhenius coefficient
     """
 
-    ksbod_T = arrhenius_correction(water_temp_c, ksbod_20, theta)
-    return ksbod_T
+    ksbod_tc = arrhenius_correction(TwaterC, ksbod_20, 1.024)
+    return ksbod_tc
 
 
 
 def CBOD_oxidation(
     DOX: xr.DataArray,
     CBOD: xr.DataArray,
-    kbod_T: xr.DataArray,
+    kbod_tc: xr.DataArray,
     KsOxbod: xr.DataArray,
     use_DOX: xr.DataArray
 ) -> xr.DataArray:
@@ -55,28 +52,28 @@ def CBOD_oxidation(
     Args:
         DOX: Dissolved oxygen concentration (mg-O2/L)
         CBOD: Carbonaceous biochemical oxygen demand (mg-O2/L)
-        kbod_T: Temperature adjusted CBOD oxidation rate (1/d)
+        kbod_tc: Temperature adjusted CBOD oxidation rate (1/d)
         KsOxbod: Half-saturation oxygen attenuation for CBOD oxidation (mg-O2/L)
         use_DOX: Option to consider DOX concentration in calculation of CBOD oxidation
     """
-    da: xr.DataArray = xr.where(use_DOX == True, (DOX / (KsOxbod + DOX)) * kbod_T * CBOD, kbod_T * CBOD)
+    da: xr.DataArray = xr.where(use_DOX == True, (DOX / (KsOxbod + DOX)) * kbod_tc * CBOD, kbod_tc * CBOD)
 
     return da
 
 
 @numba.njit
 def CBOD_sedimentation(
     CBOD: xr.DataArray,
-    ksbod_T: xr.DataArray
+    ksbod_tc: xr.DataArray
 ) -> xr.DataArray:
     """Calculates CBOD sedimentation for each group
 
     Args:
         CBOD: CBOD concentration (mg-O2/L)
-        ksbod_T: Temperature adjusted sedimentation rate (m/d)
+        ksbod_tc: Temperature adjusted sedimentation rate (m/d)
     """
 
-    CBOD_sedimentation = CBOD * ksbod_T
+    CBOD_sedimentation = CBOD * ksbod_tc
     return CBOD_sedimentation
 
 
@@ -95,7 +92,7 @@ def dCBODdt(
 
 
 @numba.njit
-def CBOD_new(
+def CBOD(
     CBOD: xr.DataArray,
     dCBODdt: xr.DataArray,
     timestep: xr.DataArray

src/clearwater_modules/nsm1/CBOD/static_variables.py

@@ -1,14 +1,16 @@
-# TODO: figure out what 'model' to import
+"""
+File contains static variables related to the CBOD module
+"""
+
 import clearwater_modules.base as base
-from clearwater_modules.tsm.model import EnergyBudget
+from clearwater_modules.nsm1.model import NutrientBudget
 
 
-@base.register_variable(models=EnergyBudget)
+@base.register_variable(models=NutrientBudget)
 class Variable(base.Variable):
     ...
 
 
-# CBOD variables for each CBOD group - array
 Variable(
     name='kbod_20',
     long_name='CBOD oxidation rate at 20C',
@@ -26,9 +28,10 @@ class Variable(base.Variable):
 )
 
 Variable(
-    name='ksOxbod',
+    name='KsOxbod',
     long_name='Half saturation oxygen attenuation constant for CBOD oxidation',
     units='mg-O2/L',
     description='Half saturation oxygen attenuation constant for CBOD oxidation',
     use='static'
 )
+

src/clearwater_modules/nsm1/DOX/dynamic_variables.py

@@ -2,11 +2,11 @@
 
 import clearwater_modules.shared.processes as shared_processes
 from clearwater_modules import base
-from clearwater_modules.tsm.model import EnergyBudget
-from clearwater_modules.nsm1.DOX import processes
+from clearwater_modules.nsm1.model import NutrientBudget
+import clearwater_modules.nsm1.DOX.processes as processes
 
 
-@base.register_variable(models=EnergyBudget)
+@base.register_variable(models=NutrientBudget)
 class Variable(base.Variable):
     ...
 
@@ -38,51 +38,6 @@ class Variable(base.Variable):
     process=processes.DOs_atm_alpha
 )
 
-Variable(
-    name='kah_20',
-    long_name='Hydraulic oxygen reaeration rate adjusted for hydraulics',
-    units='1/d',
-    description='Hydraulic oxygen reaeration rate adjusted for hydraulic parameters according to XX lit',
-    use='dynamic',
-    process=processes.kah_20
-)
-
-Variable(
-    name='kah_T',
-    long_name='Hydraulic oxygen reaeration rate adjusted for temperature',
-    units='1/d',
-    description='Hydraulic oxygen reaeration rate adjusted for temperature',
-    use='dynamic',
-    process=processes.kah_T
-)
-
-Variable(
-    name='kaw_20',
-    long_name='Wind oxygen reaeration velocity adjusted for hydraulics',
-    units='m/d',
-    description='Wind oxygen reaeration velocity adjusted for hydraulic parameters according to XX lit',
-    use='dynamic',
-    process=processes.kaw_20
-)
-
-Variable(
-    name='kaw_T',
-    long_name='Wind oxygen reaeration velocity adjusted for temperature',
-    units='m/d',
-    description='Wind oxygen reaeration velocity adjusted for temperature',
-    use='dynamic',
-    process=processes.kaw_T
-)
-
-Variable(
-    name='ka_T',
-    long_name='Oxygen reaeration rate',
-    units='1/d',
-    description='Oxygen reaeration rate',
-    use='dynamic',
-    process=processes.ka_T
-)
-
 Variable(
     name='Atm_O2_reaeration',
     long_name='Atmospheric oxygen reaeration',
@@ -157,14 +112,6 @@ class Variable(base.Variable):
     process=processes.DOX_AbRespiration
 )
 
-Variable(
-    name='SOD_tc',
-    long_name='Sediment oxygen demand corrected by temperature and dissolved oxygen concentration',
-    units='g/m^3/d',
-    description='Sediment oxygen demand corrected by temperature and dissolved oxygen concentration',
-    use='dynamic',
-    process=processes.SOD_tc
-)
 
 Variable(
     name='DOX_SOD',