diff --git a/HSP2/HYDR.py b/HSP2/HYDR.py index 89f8f264..1aa56f2c 100644 --- 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 index d3cef16f..f18238f3 100644 --- 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 \ No newline at end of file diff --git a/HSP2/main.py b/HSP2/main.py index c9264e83..30d69ad0 100644 --- 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: diff --git a/HSP2/state.py b/HSP2/state.py new file mode 100644 index 00000000..074f653c --- /dev/null +++ b/HSP2/state.py @@ -0,0 +1,200 @@ +''' General routines for SPECL ''' + +import numpy as np +import time +from pandas import DataFrame, date_range +from pandas.tseries.offsets import Minute +from numba.typed import Dict +from numpy import zeros +from numba import int8, float32, njit, types, typed # import the types +import os +import importlib.util +import sys + +def init_state_dicts(): + """ + This contains the base dictionaries used to pass model state amongst modules and custom code plugins + """ + state = {} # shared state Dictionary, contains numba-ready Dicts + state_paths = Dict.empty(key_type=types.unicode_type, value_type=types.int64) + state_ix = Dict.empty(key_type=types.int64, value_type=types.float64) + dict_ix = Dict.empty(key_type=types.int64, value_type=types.float64[:,:]) + ts_ix = Dict.empty(key_type=types.int64, value_type=types.float64[:]) + # initialize state for hydr + # now put all of these Dicts into the state Dict + state['state_paths'], state['state_ix'], state['dict_ix'], state['ts_ix'] = state_paths, state_ix, dict_ix, ts_ix + return state + + +def find_state_path(state_paths, parent_path, varname): + """ + We should get really good at using docstrings... + """ + # this is a bandaid, we should have an object routine that searches the parent for variables or inputs + var_path = parent_path + "/states/" + str(varname) + return var_path + +def op_path_name(operation, id): + """ + Used to generate hdf5 operation name in a central fashion to avoid naming convention slip-ups + """ + tid = str(id).zfill(3) + path_name = f'{operation}_{operation[0]}{tid}' + return path_name + +def get_op_state_path(operation, id, activity = ''): + """ + Used to generate hdf5 paths in a central fashion to avoid naming convention slip-ups + """ + op_name = op_path_name(operation, id) + if activity == '': + op_path = f'/STATE/{op_name}' + else: + op_path = f'/STATE/{op_name}/{activity}' + return op_path + + +def get_state_ix(state_ix, state_paths, var_path): + """ + Find the integer key of a variable name in state_ix + """ + if not (var_path in list(state_paths.keys())): + # we need to add this to the state + return False # should throw an error + var_ix = state_paths[var_path] + return var_ix + + +def get_ix_path(state_paths, var_ix): + """ + Find the integer key of a variable name in state_ix + """ + for spath, ix in state_paths.items(): + if var_ix == ix: + # we need to add this to the state + return spath + return False + +def set_state(state_ix, state_paths, var_path, default_value = 0.0, debug = False): + """ + Given an hdf5 style path to a variable, set the value + If the variable does not yet exist, create it. + Returns the integer key of the variable in the state_ix Dict + """ + if not (var_path in state_paths.keys()): + # we need to add this to the state + state_paths[var_path] = append_state(state_ix, default_value) + var_ix = get_state_ix(state_ix, state_paths, var_path) + if (debug == True): + print("Setting state_ix[", var_ix, "], to", default_value) + state_ix[var_ix] = default_value + return var_ix + + +def set_dict_state(state_ix, dict_ix, state_paths, var_path, default_value = {}): + """ + Given an hdf5 style path to a variable, set the value in the dict + If the variable does not yet exist, create it. + Returns the integer key of the variable in the state_ix Dict + """ + if not (var_path in state_paths.keys()): + # we need to add this to the state + state_paths[var_path] = append_state(state_ix, default_value) + var_ix = get_state_ix(state_ix, state_paths, var_path) + return var_ix + + +def append_state(state_ix, var_value): + """ + Add a new variable on the end of the state_ix Dict + Return the key of this new variable + """ + if (len(state_ix) == 0): + val_ix = 1 + else: + val_ix = max(state_ix.keys()) + 1 # next ix value + state_ix[val_ix] = var_value + return val_ix + +def state_context_hsp2(state, operation, segment, activity): + state['operation'] = operation + state['segment'] = segment # + state['activity'] = activity + # give shortcut to state path for the upcoming function + state['domain'] = "/STATE/" + operation + "_" + segment + "/" + activity + +def state_siminfo_hsp2(uci_obj, siminfo): + # Add crucial simulation info for dynamic operation support + delt = uci_obj.opseq.INDELT_minutes[0] # get initial value for STATE objects + siminfo['delt'] = delt + siminfo['tindex'] = date_range(siminfo['start'], siminfo['stop'], freq=Minute(delt))[1:] + siminfo['steps'] = len(siminfo['tindex']) + +def state_load_dynamics_hsp2(state, io_manager, siminfo): + # Load any dynamic components if present, and store variables on objects + hsp2_local_py = load_dynamics(io_manager, siminfo) + # if a local file with state_step_hydr() was found in load_dynamics(), we add it to state + state['state_step_hydr'] = siminfo['state_step_hydr'] # enabled or disabled + state['hsp2_local_py'] = hsp2_local_py # Stores the actual function in state + +def hydr_init_ix(state_ix, state_paths, domain): + # get a list of keys for all hydr state variables + hydr_state = ["DEP","IVOL","O1","O2","O3","OVOL1","OVOL2","OVOL3","PRSUPY","RO","ROVOL","SAREA","TAU","USTAR","VOL","VOLEV"] + hydr_ix = Dict.empty(key_type=types.unicode_type, value_type=types.int64) + for i in hydr_state: + #var_path = f'{domain}/{i}' + var_path = domain + "/" + i + hydr_ix[i] = set_state(state_ix, state_paths, var_path, 0.0) + return hydr_ix + +@njit +def hydr_get_ix(state_ix, state_paths, domain): + # get a list of keys for all hydr state variables + hydr_state = ["DEP","IVOL","O1","O2","O3","OVOL1","OVOL2","OVOL3","PRSUPY","RO","ROVOL","SAREA","TAU","USTAR","VOL","VOLEV"] + hydr_ix = Dict.empty(key_type=types.unicode_type, value_type=types.int64) + for i in hydr_state: + #var_path = f'{domain}/{i}' + var_path = domain + "/" + i + hydr_ix[i] = state_paths[var_path] + return hydr_ix + +# function to dynamically load module, based on "Using imp module" in https://www.tutorialspoint.com/How-I-can-dynamically-import-Python-module# +#def dynamic_module_import(module_name, class_name): +def dynamic_module_import(local_name, local_path, module_name): + # find_module() is used to find the module in current directory + # it gets the pointer, path and description of the module + module = False + local_spec = False + try: + print ("Looking for local_name, local_path", local_name, local_path) + local_spec = importlib.util.spec_from_file_location(local_name, local_path) + except ImportError: + print ("Imported module {} not found".format(local_name)) + try: + # load_module dynamically loads the module + # the parameters are pointer, path and description of the module + if (local_spec != False): + module = importlib.util.module_from_spec(local_spec) + sys.modules[local_spec.name] = module + sys.modules[module_name] = module + local_spec.loader.exec_module(module) + except Exception as e: + print(e) + return module + + +def load_dynamics(io_manager, siminfo): + local_path = os.getcwd() + # try this + hdf5_path = io_manager._input.file_path + (fbase, fext) = os.path.splitext(hdf5_path) + # see if there is a code module with custom python + print("Looking for custom python code ", (fbase + ".py")) + hsp2_local_py = dynamic_module_import(fbase, local_path + "/" + fbase + ".py", "hsp2_local_py") + siminfo['state_step_hydr'] = 'disabled' + if 'state_step_hydr' in dir(hsp2_local_py): + siminfo['state_step_hydr'] = 'enabled' + else: + print("state_step_hydr function not defined. Using default") + return False + return hsp2_local_py \ No newline at end of file diff --git a/HSP2/state_fn_defaults.py b/HSP2/state_fn_defaults.py new file mode 100644 index 00000000..f4b4035b --- /dev/null +++ b/HSP2/state_fn_defaults.py @@ -0,0 +1,6 @@ +# null function to be loaded when not supplied by user +from numba import int8, float32, njit, types, typed # import the types + +@njit +def state_step_hydr(state_info, state_paths, state_ix, dict_ix, ts_ix, hydr_ix, step): + return \ No newline at end of file diff --git a/environment.yml b/environment.yml index 057cd9a3..807dd2c7 100644 --- a/environment.yml +++ b/environment.yml @@ -22,7 +22,7 @@ dependencies: - mando # for Python CLI app in HSP2tools/HSP2_CLI.py; hasn't been updated since 2017 and doesn't support Python 3.9: https://github.com/rubik/mando # Interactivity & Visualization via Jupyter Notebooks (optional, but required for tutorials) - - jupyterlab =3.0.* # also installs classic Jupyter notbook + - jupyterlab ==3.0.* # also installs classic Jupyter notbook - ipympl # jupyter-matplotlib, https://github.com/matplotlib/ipympl - nodejs # required for many JupyterLab extensions - nb_conda # Conda environment & package access extension from within Jupyter @@ -42,4 +42,4 @@ dependencies: # Optional, but recommended for tutorials # - lckr-jupyterlab-variableinspector # https://github.com/lckr/jupyterlab-variableInspector # - jupyterlab_hdf # Explore HDF5 files in JupyterLab. Requires an additional step to install. - # Installation instructions: https://github.com/jupyterlab/jupyterlab-hdf5#installation \ No newline at end of file + # Installation instructions: https://github.com/jupyterlab/jupyterlab-hdf5#installation diff --git a/tests/test10/HSP2results/test10.py b/tests/test10/HSP2results/test10.py new file mode 100644 index 00000000..47d10924 --- /dev/null +++ b/tests/test10/HSP2results/test10.py @@ -0,0 +1,25 @@ +import sys +from numba import int8, float32, njit, types, typed # import the types + +print("Loaded a set of HSP2 code!") + +@njit +def state_step_hydr(state_info, state_paths, state_ix, dict_ix, ts_ix, hydr_ix, step): + if (step <= 1): + print("Custom state_step_hydr() called for ", state_info['segment']) + print("operation", state_info['operation']) + print("state at start", state_ix) + print("domain info", state_info) + print("state_paths", state_paths) + # Do a simple withdrawal of 10 MGGD from R001 + if (state_info['segment'] == 'R001') and (state_info['activity'] == 'HYDR'): + state_ix[hydr_ix['O1']] = 10.0 * 1.547 + # Route point source return from segment R001 demand to R005 inflow (IVOL) + # For demo purposes this will only use the last state_ix value for R001 demand + # Realistic approach would run all segments simultaneously or use value from ts_ix (ts_ix loading TBD) + if (state_info['segment'] == 'R005') and (state_info['activity'] == 'HYDR'): + state_ix[hydr_ix['IVOL']] += 0.85 * state_ix[state_paths['/STATE/RCHRES_R001/HYDR/O1']] + if (step <= 1): + print("IVOL after", state_ix[hydr_ix['IVOL']]) + return +