respec · rburghol · Jun 21, 2023 · Jun 21, 2023 · Jun 21, 2023 · Jun 22, 2023
diff --git a/HSP2/HYDR.py b/HSP2/HYDR.py
@@ -12,13 +12,14 @@
 '''
 
 
-from numpy import zeros, any, full, nan, array, int64
+from numpy import zeros, any, full, nan, array, int64, arange
 from pandas import DataFrame
 from math import sqrt, log10
 from numba import njit
 from numba.typed import List
 from HSP2.utilities import initm, make_numba_dict
-from HSP2.SPECL import specl, _specl_
+from HSP2.state import *
+from HSP2.SPECL import specl
 
 
 ERRMSGS =('HYDR: SOLVE equations are indeterminate',             #ERRMSG0
@@ -31,7 +32,7 @@
 MAXLOOPS  = 100     # newton method exit tolerance
 
 
-def hydr(io_manager, siminfo, uci, ts, ftables, specactions):
+def hydr(io_manager, siminfo, uci, ts, ftables, state):
     ''' find the state of the reach/reservoir at the end of the time interval
     and the outflows during the interval
 
@@ -40,7 +41,9 @@ def hydr(io_manager, siminfo, uci, ts, ftables, specactions):
        general is a dictionary with simulation level infor (OP_SEQUENCE for example)
        ui is a dictionary with RID specific HSPF UCI like data
        ts is a dictionary with RID specific timeseries
-       specactions is a dictionary with all special actions'''
+       state is a dictionary that contains all dynamic code dictionaries such as: 
+       - specactions is a dictionary with all special actions
+    '''
 
     steps   = siminfo['steps']                # number of simulation points
     uunits  = siminfo['units']
@@ -72,33 +75,11 @@ def hydr(io_manager, siminfo, uci, ts, ftables, specactions):
 
     # OUTDGT timeseries might come in as OUTDGT, OUTDGT0, etc. otherwise UCI default
     names = list(sorted([n for n in ts if n.startswith('OUTDG')], reverse=True))
-    print(names)
     df = DataFrame()
     for i,c in enumerate(ODGTF):
         df[i] = ts[names.pop()][0:steps] if c > 0 else zeros(steps)
     OUTDGT = df.to_numpy()
 
-    print(OUTDGT)
-
-    # OUTDGT replacement testing
-    # OUTDGT = ts[names][0:steps]
-    xkeys = ['OUTDGT1', 'OUTDGT2']
-    # xvalues = list(map(ts.get, xkeys))
-    xvalues = array(ts['OUTDGT1'],ts['OUTDGT2'])
-    # xvalues = array(map(ts.get, xkeys))
-    # xvalues = list(map(ts.get, names))
-    # print(names)
-    print(xvalues)
-    # print(xvalues.shape)
-    # print(xvalues[0][0])
-    # xvalues[0][0] = xvalues[0][0] * 2
-    # print("ts['OUTDGT1']", ts['OUTDGT1'][0])
-    # print("ts['OUTDGT2']", ts['OUTDGT2'][0])
-
-    # List all names in ts, for jk testing purposes only
-    # ts_names = list(sorted([n for n in ts], reverse=True))
-    # print(ts_names)
-
     # generic SAVE table doesn't know nexits for output flows and rates
     if nexits > 1:
         u = uci['SAVE']
@@ -145,11 +126,33 @@ def hydr(io_manager, siminfo, uci, ts, ftables, specactions):
     for i in range(nexits):
         Olabels.append(f'O{i+1}')
         OVOLlabels.append(f'OVOL{i+1}')
-
-    specactions = make_numba_dict(specactions) # Note: all values coverted to float automatically
+
+    # must split dicts out of state Dict since numba cannot handle mixed-type nested Dicts
+    # state_info is some generic things about the simulation 
+    state_info = Dict.empty(key_type=types.unicode_type, value_type=types.unicode_type)
+    state_info['operation'], state_info['segment'], state_info['activity'] = state['operation'], state['segment'], state['activity']
+    state_info['domain'], state_info['state_step_hydr'] = state['domain'], state['state_step_hydr']
+    hsp2_local_py = state['hsp2_local_py']
+    # It appears necessary to load this here, instead of from main.py, otherwise,
+    # _hydr_() does not recognize the function state_step_hydr()? 
+    if (hsp2_local_py != False):
+        from hsp2_local_py import state_step_hydr
+    else:
+        from HSP2.state_fn_defaults import state_step_hydr
+    state_ix, dict_ix, ts_ix = state['state_ix'], state['dict_ix'], state['ts_ix']
+    state_paths = state['state_paths']
+    # initialize the hydr paths in case they don't already reside here
+    hydr_init_ix(state_ix, state_paths, state['domain'])
+
+    ###########################################################################
+    # specactions - special actions code TBD
+    ###########################################################################
+    specactions = make_numba_dict(state['specactions']) # Note: all values coverted to float automatically
+
+    ###########################################################################
+    # Do the simulation with _hydr_()
     ###########################################################################
-    errors = _hydr_(ui, ts, COLIND, OUTDGT, rchtab, funct, Olabels, OVOLlabels, specactions)                  # run reaches simulation code
-#    errors = _hydr_(ui, ts, COLIND, OUTDGT, rchtab, funct, Olabels, OVOLlabels)                  # run reaches simulation code
+    errors = _hydr_(ui, ts, COLIND, OUTDGT, rchtab, funct, Olabels, OVOLlabels, state_info, state_paths, state_ix, dict_ix, ts_ix, specactions, state_step_hydr) # run reaches simulation code
     ###########################################################################
 
     if 'O'    in ts:  del ts['O']
@@ -164,7 +167,7 @@ def hydr(io_manager, siminfo, uci, ts, ftables, specactions):
 
 
 @njit(cache=True)
-def _hydr_(ui, ts, COLIND, OUTDGT, rowsFT, funct, Olabels, OVOLlabels, specactions):
+def _hydr_(ui, ts, COLIND, OUTDGT, rowsFT, funct, Olabels, OVOLlabels, state_info, state_paths, state_ix, dict_ix, ts_ix, specactions, state_step_hydr):
     errors = zeros(int(ui['errlen'])).astype(int64)
 
     steps  = int(ui['steps'])            # number of simulation steps
@@ -290,38 +293,51 @@ def _hydr_(ui, ts, COLIND, OUTDGT, rowsFT, funct, Olabels, OVOLlabels, specactio
     for index in range(nexits):
         ui['OS' + str(index + 1)] = o[index]
 
+    # other initial vars
+    rovol = 0.0
+    volev = 0.0
+    IVOL0   = ts['IVOL']                   # the actual inflow in simulation native units 
+    # prepare for dynamic state
+    hydr_ix = hydr_get_ix(state_ix, state_paths, state_info['domain'])
+    # these are integer placeholders faster than calling the array look each timestep
+    o1_ix, o2_ix, o3_ix, ivol_ix = hydr_ix['O1'], hydr_ix['O2'], hydr_ix['O3'], hydr_ix['IVOL']
+    ro_ix, rovol_ix, volev_ix, vol_ix = hydr_ix['RO'], hydr_ix['ROVOL'], hydr_ix['VOLEV'], hydr_ix['VOL']
+    # handle varying length outdgt
+    out_ix = arange(nexits)
+    if nexits > 0:
+        out_ix[0] = o1_ix
+    if nexits > 1:
+        out_ix[1] = o2_ix
+    if nexits > 2:
+        out_ix[2] = o3_ix
     # HYDR (except where noted)
     for step in range(steps):
-        # print('\n', 'step: ', step, ' of: ', steps, ' steps')
-
-        # ------------------------------------------------------------------------
-        # print('Trying specl')   
-        # OUTDGT2_save = ts['OUTDGT2'][step - 1] # save before calling specl()
-        # OUTDGT1_save = ts['OUTDGT1'][step - 1] 
-        # print("OUTDGT[step, :]", OUTDGT[step, :])
-        # print("ro", ro)
-
         # call specl
-        specl(ui, ts, step, specactions)
-        # print("ts['OUTDGT2'][step]", ts['OUTDGT2'][step])
-        # print("ts['OUTDGT1'][step]", ts['OUTDGT1'][step])
-        # print("OUTDGT[step, :]", OUTDGT[step, :])
-
-        # set OUTDGT using the values in the ts object which were set inside specl()
-        OUTDGT[step, :] = [ts['OUTDGT1'][step], ts['OUTDGT2'][step], 0.0]
-        # print("OUTDGT[step, :]", OUTDGT[step, :])
-        # ------------------------------------------------------------------------
-
+        specl(ui, ts, step, state_info, state_paths, state_ix, specactions)
         convf  = CONVF[step]
         outdgt[:] = OUTDGT[step, :]
         colind[:] = COLIND[step, :]
         roseff = ro
         oseff[:] = o[:]
-
-        # jk test prints
-        # print("roseff", roseff)
-        # print("outdgt[:]", outdgt[:])
-
+        # set state_ix with value of local state variables and/or needed vars
+        # Note: we pass IVOL0, not IVOL here since IVOL has been converted to different units
+        state_ix[ro_ix], state_ix[rovol_ix] = ro, rovol
+        di = 0
+        for oi in range(nexits):
+            state_ix[out_ix[oi]] = outdgt[oi] 
+        state_ix[vol_ix], state_ix[ivol_ix] = vol, IVOL0[step]
+        state_ix[volev_ix] = volev
+        # Execute dynamic code if enabled
+        if (state_info['state_step_hydr'] == 'enabled'):
+            state_step_hydr(state_info, state_paths, state_ix, dict_ix, ts_ix, hydr_ix, step)
+            # Do write-backs for editable STATE variables
+            # OUTDGT is writeable
+            for oi in range(nexits):
+                outdgt[oi] = state_ix[out_ix[oi]]
+            # IVOL is writeable. 
+            # Note: we must convert IVOL to the units expected in _hydr_ 
+            # maybe routines should do this, and this is not needed (but pass VFACT in state)
+            IVOL[step] = state_ix[ivol_ix] * VFACT
         # vols, sas variables and their initializations  not needed.
         if irexit >= 0:             # irrigation exit is set, zero based number
             if rirwdl > 0.0:  # equivalent to OVOL for the irrigation exit

diff --git a/HSP2/SPECL.py b/HSP2/SPECL.py
@@ -12,34 +12,7 @@
 from numba import njit
 
 @njit
-# def specl(io_manager, siminfo, uci, ts, step, specl_actions):
-def specl(ui, ts, step, specactions):
-
-    # print('Made it to specl()')
-    ts = _specl_(ui, ts, step, specactions)
-
-    # return errors, ERRMSGS
-    # return ts
-
-
-
-# def _specl_(ui, ts, COLIND, OUTDGT, rowsFT, funct, Olabels, OVOLlabels):
-@njit
-def _specl_(ui, ts, step, specactions):
-
-    # print('Made it to _specl_()')
-    # ts['VOL'][step - 1] = ts['VOL'][step - 1] * 5.0
-    # ts['VOL'][step - 1] = ts['VOL'][step - 1] - specactions['test_wd']
-
-    # ts['OUTDGT'][step - 1] = ts['OUTDGT'][step - 1]  - specactions['test_wd']
-    # ts['OUTDGT2'][step - 1] = 99
-    # ts['OUTDGT2'][step - 1] = ts['OUTDGT2'][step - 1]
-    # ts['OUTDGT2'][step] = ts['OUTDGT2'][step - 1] + 99
-
-    ts['OUTDGT2'][step] = 99
-    # ts['OUTDGT2'][step, :] = 99 # this resulted in errors
-
-    # print(specactions['outdgt'])
-    # return errors
-    # return ts
+def specl(ui, ts, step, state_info, state_paths, state_ix, specactions):
+    # ther eis no need for _specl_ because this code must already be njit
+    return
 
diff --git a/HSP2/main.py b/HSP2/main.py
@@ -11,6 +11,7 @@
 import os
 from HSP2.utilities import versions, get_timeseries, expand_timeseries_names, save_timeseries, get_gener_timeseries
 from HSP2.configuration import activities, noop, expand_masslinks
+from HSP2.state import *
 
 from HSP2IO.io import IOManager, SupportsReadTS, Category
 
@@ -56,6 +57,19 @@ def main(io_manager:IOManager, saveall:bool=False, jupyterlab:bool=True) -> None
     copy_instances = {}
     gener_instances = {}
 
+    #######################################################################################
+    # initilize STATE dicts
+    #######################################################################################
+    # Set up Things in state that will be used in all modular activitis like SPECL
+    state = init_state_dicts()
+    state_siminfo_hsp2(uci_obj, siminfo)
+    # Add support for dynamic functins to operate on STATE
+    # - Load any dynamic components if present, and store variables on objects 
+    state_load_dynamics_hsp2(state, io_manager, siminfo)
+    # - finally stash specactions in state, not domain (segment) dependent so do it once
+    state['specactions'] = specactions # stash the specaction dict in state
+    #######################################################################################
+
     # main processing loop
     msg(1, f'Simulation Start: {start}, Stop: {stop}')
     tscat = {}
@@ -99,7 +113,9 @@ def main(io_manager:IOManager, saveall:bool=False, jupyterlab:bool=True) -> None
                     continue
 
                 msg(3, f'{activity}')
-
+                # Set context for dynamic executables.
+                state_context_hsp2(state, operation, segment, activity)
+
                 ui = uci[(operation, activity, segment)]   # ui is a dictionary
                 if operation == 'PERLND' and activity == 'SEDMNT':
                     # special exception here to make CSNOFG available
@@ -209,7 +225,7 @@ def main(io_manager:IOManager, saveall:bool=False, jupyterlab:bool=True) -> None
                 ############ calls activity function like snow() ##############
                 if operation not in ['COPY','GENER']:
                     if (activity == 'HYDR'):
-                        errors, errmessages = function(io_manager, siminfo, ui, ts, ftables, specactions)
+                        errors, errmessages = function(io_manager, siminfo, ui, ts, ftables, state)
                     elif (activity != 'RQUAL'):
                         errors, errmessages = function(io_manager, siminfo, ui, ts)
                     else: