TSM tests are all passing!

examples/dev_sandbox/tsm_debugger.py

@@ -0,0 +1,38 @@
+"""A script to allow for debugging of the TSM module."""
+import clearwater_modules
+import numpy as np
+
+
+def main():
+    # define starting state values
+    state_i = {
+        'water_temp_c': 20.0,
+        'surface_area': 1.0,
+        'volume': 1.0,
+    }
+
+    # instantiate the TSM module
+    tsm = clearwater_modules.tsm.EnergyBudget(
+        initial_state_values=state_i,
+    )
+    print(tsm.static_variable_values)
+
+    input_dataset = tsm.dataset.isel(time_step=0).copy()
+
+    inputs = map(
+        lambda x: clearwater_modules.utils._prep_inputs(input_dataset, x),
+        tsm.computation_order,
+    )
+    dims = input_dataset.dims
+
+    for name, func, arrays in inputs:
+        print(name)
+        array: np.ndarray = func(*arrays)
+        input_dataset[name] = (dims, array)
+        print(array)
+        print()
+    return input_dataset
+
+
+if __name__ == '__main__':
+    main()

src/clearwater_modules/tsm/constants.py

@@ -38,22 +38,22 @@ class Meteorological(TypedDict):
 
 
 DEFAULT_METEOROLOGICAL = Meteorological(
-    air_temp_c=20,
-    q_solar=400,
+    air_temp_c=20.0,
+    q_solar=400.0,
     sed_temp_c=5.0,
     eair_mb=1.0,
     pressure_mb=1013.0,
     cloudiness=0.1,
     wind_speed=3.0,
     wind_a=0.3,
     wind_b=1.5,
-    wind_c=1.0,
+    wind_c=3.0,
     wind_kh_kw=1.0,
 )
 
 DEFAULT_TEMPERATURE = Temperature(
     stefan_boltzmann=5.67e-8,
-    cp_air=1005,
+    cp_air=1005.0,
     emissivity_water=0.97,
     gravity=-9.806,
     a0=6984.505294,

src/clearwater_modules/tsm/processes.py

@@ -81,6 +81,7 @@ def emissivity_air(
 
 @numba.njit
 def wind_function(
+    ri_function: xr.DataArray,
     wind_a: xr.DataArray,
     wind_b: xr.DataArray,
     wind_c: xr.DataArray,
@@ -89,17 +90,23 @@ def wind_function(
     """Calculate wind function (unitless) for latent and sensible heat.
 
     Args:
+        ri_function: Richardson function (unitless)
         wind_a: Wind function coefficient (unitless)
         wind_b: Wind function coefficient (unitless)
         wind_c: Wind function coefficient (unitless)
         wind_speed: Wind speed (m/s)
     """
-    return wind_a / 1000000.0 + wind_b / 1000000.0 * wind_speed**wind_c
+    return (
+        ri_function * (
+            (wind_a / 1000000.0) +
+            (wind_b / 1000000.0) *
+            (wind_speed**wind_c)
+        )
+    )
 
 
 @numba.njit
 def q_latent(
-    ri_function: xr.DataArray,
     pressure_mb: xr.DataArray,
     density_water: xr.DataArray,
     lv: xr.DataArray,
@@ -118,9 +125,10 @@ def q_latent(
         eair_mb: Vapour pressure of air (mb)
     """
     return (
-        ri_function *
         (0.622 / pressure_mb) *
-        lv * density_water * wind_function *
+        lv *
+        density_water *
+        wind_function *
         (esat_mb - eair_mb)
     )
 
@@ -355,28 +363,28 @@ def ri_function(ri_number: xr.DataArray) -> np.ndarray:
     )
     out: np.ndarray = np.select(
         condlist=[
-            (ri_number_bounded < 0.0) & (ri_number_bounded >= -0.01), # neutral
+            (ri_number_bounded < 0.0) & (ri_number_bounded >= -0.01),  # neutral
             (ri_number_bounded < 0.0) & (ri_number_bounded < -0.01),  # unstable
-            (ri_number_bounded >= 0.0) & (ri_number_bounded <= 0.01), # neutral
+            (ri_number_bounded >= 0.0) & (ri_number_bounded <= 0.01),  # neutral
             (ri_number_bounded >= 0.0) & (ri_number_bounded > 0.01),  # stable
         ],
         choicelist=[
             1.0,                                         # neutral
             (1.0 - 22.0 * ri_number_bounded) ** 0.80,    # unstable
             1.0,                                         # neutral
-            (1.0 + 34.0 * ri_number_bounded) ** (-0.80), # stable
+            (1.0 + 34.0 * ri_number_bounded) ** (-0.80),  # stable
         ],
     )
     warnings.filterwarnings("default", category=RuntimeWarning)
     return out
 
     # old method with where, same logic
-    #da: xr.DataArray = xr.where(ri_number > 2.0, (1.0 + 34.0*2.0)**(-0.80),
+    # da: xr.DataArray = xr.where(ri_number > 2.0, (1.0 + 34.0*2.0)**(-0.80),
     #                   xr.where(ri_number < -1.0,  (1.0 - 22.0*-1.0)**(-0.80),
     #                   xr.where((ri_number < 0.0) & (ri_number >= -0.01), 1.0,
     #                   xr.where(ri_number < -0.01, (1.0 - 22.0 * ri_number)**0.80,
     #                   xr.where((ri_number >= 0.0) & (ri_number <= 0.01), 1.0, (1.0 + 34.0 * ri_number)**(-0.80)
-    #)))))
+    # )))))
 
 
 @numba.njit
@@ -447,7 +455,7 @@ def mf_cp_water(water_temp_c: xr.DataArray) -> xr.DataArray:
     )
 
     # previous code
-    #return xr.where(water_temp_c <= 0.0, 4218.0,
+    # return xr.where(water_temp_c <= 0.0, 4218.0,
     #    xr.where(water_temp_c <= 5.0, 4202.0,
     #       xr.where(water_temp_c <= 10.0, 4192.0,
     #  xr.where(water_temp_c <= 15.0, 4186.0,
@@ -476,12 +484,12 @@ def q_net(
         q_sediment: Sediment heat flux (W/m^2)
     """
     return (
-        q_sensible -
-        q_latent -
-        q_longwave_up -
-        q_longwave_down +
+        q_sensible +
         q_solar +
-        q_sediment
+        q_sediment +
+        q_longwave_down -
+        q_longwave_up -
+        q_latent
     )