`develop-state-class`: Reorganize Files and create STATE `jit` class

[rburghol]

Developmental Code: https://github.com/HARPgroup/HSPsquared/blob/develop-state-class/src/hsp2/hsp2/

This branch will perform some code reorganization and add class-based state data structures, I think it will enhance readability and organization. Specifically:

• Move src/hsp2/hsp2/state* to src/hsp2/state/
• Create a class ModelSTATE which contains all state related entities (state_ix, dict_ix, state_paths, ...) as attributes
• Reduce the number of arguments passed to functions
• Organize high-performance functions under the object to make it explicit when a function is jit-ed.
  • i.e., if a function is called as state.func_name() it is guaranteed to be jit compatible
• See example of jit-able object with these properties in https://github.com/HARPgroup/HSPsquared/blob/develop-hydrclass/src/hsp2/hsp2/classes/model.test.class.py

Goals

• Re-organize code for STATE model classes and function files to reside in a sub-directory

Info

• Data Model for STATE: Data Model for Shared STATE (supports model modularity) RFC #112
• Current code base in fork, will PR when ready: https://github.com/HARPgroup/HSPsquared/tree/develop-state-tidy/src/hsp2/

Tasks

• Document Needed functions and classes to implement How to Implement New Module in STATE system #127
• Prototype full workflow for using shared STATE and all helper object-classes/runtime exec routines from OM
• Migrate code from OM into state sub-module for ModelObject, ModelConstant, Timer, and SpecialAction
• Create sub-directory module with state support and object model component code
• migrate import statements
• Create base class for state and helper function to make class attributes
• Migrate old references to new object:
  • init_state_dicts
  • state_siminfo_hsp2(parameter_obj, siminfo, io_manager, state)
  • Separate model_object_cache from state and go into om_operations or just be separate argument (these objects cannot be jitted)
    • in om is fairly simple, as om_operations can be created and passed around
    • in ModelObject and children classes, this can be more challenging. Will need to pass in at creation, or modify the state that is passed in... maybe pass in om_operations instead of state, or as a container that includes state?
      • If we make a state container that combines om_operations, this should be fine and propagate all changes up to the actual state object?
  • model_root_object
    • Where to store instead of on state? In om_operations
  • state_load_dynamics_hsp2(state, io_manager, siminfo)
  • model_data - put into om_operations
  • state_init_hsp2(state, opseq, activities)
  • model_input_dependencies - this function is actually unused, but maye will be used for timeseries? Add timeseries/other DSN access to STATE/specl HARPgroup/HSPsquared#60 (comment)
  • state_context_hsp2:
    • includes operation, segment activity, and domain
  • siminfo['tindex']
    • This is passed into state, creatin a complex object (dict), so not possible in jit
    • Only tindex is actually needed, and that could be obtained elsewhere
    • However, thinking forward, the SimTimer object needs to have its tindex updated each time that the time step switches, so we should tackle all of this at once if possible.
    • (YES it is in dict_ix) Is tindex gotten from state in simtimer anywhere?
    • Update ModelLinkage to Load SimTimer from model_object_cache and get, or, just grab from the dict_ix of SimTimer
  • om_init_state(state) # set up operational model specific state entries -- now In om_operations`
  • specactions: this was in state only to pass to the specl parsing routine at startup, and really, can be passed as an argument anyway. Easy fix.
    • specl_load_state(state, specactions) # traditional special actions
    • specl_load_om
  • Migrate model_exec_list from inside functions to prep routine. Stored as domain_action specific array
    • In HYDR
    • In SEDTRN
  • state_load_dynamics_om(state, io_manager, siminfo ) # operational model for custom python
  • state_om_model_run_prep(state, io_manager, siminfo)
  • add timer to state, which eliminates the need to implement dict_ix at this time which is most likely going to be far better as an object with runtime declared np array
  • HYDR
  • SEDMNT
  • SEDTRN
  • RQUAL_class
  • om
  • om_model_object
  • om_special_action
  • om_equation
  • om_model_linkage
  • om_sim_timer

Module Declaration / Outline

• THIS IS OLD INFORMATION - MAY NO LONGER BE RELEVANT
• via: https://docs.python.org/3/tutorial/modules.html#packages
  • Note, in the above outlines how to support from state.state import *

state / 
        __init__.py
        state.py
        specl.py
        om.py
        model_object.py
        actions.py

SEDTRN, HYDR adaptation to new state object class

• State-Mutable Variables Existing SEDTRN, HYDR need to load the pointers to state-mutable variables from state
  • SEDTRN: RSED4, RSED5 and RSED6
  • HYDR:
  • Future version will use an object based SEDTRN which will:
    • Store pointers to state as properties of the object in order to facilitate simultaneous simulations, but for now, we want to support the old model, but with using state class methods to obtain the indices.
    • domain will similarly be obsolete since the domain will be associated with the SEDTRN object state_path property
• Functional domain: For now, we pass domain in as a property of the state object, but this will also be obsolete as we move to an object based system for functional areas.
• Dynamic Code: experimental support for dynamically loaded python code and the operational models is supported in the state_info dictionary as:
  • model_exec_list has been replaced as an array on the state object, keyed by ix value for each domain + function, like /STATE/RCHRES_R001/HYDR , so model_exec_list = state.op_exec_lists[op_ix] will contain the list of dynamically executed operations for that function area. In the future this could
  • the variable state_step_hydr is a string enabled or disabled, and determines if a function called state_step_hydro() has been defined in a local code file (in the h5 directory)-- this is now located on state.state_step_hydr. The string variable state_step_hydr is actually somewhat superfluous as we could maybe call the stub function state_step_hydr() which does nothing. But we will keep it on state for now as it costs nothing and will keep code changes smaller. Also, it may be that the process of using that functon will require a recompile every time regardless of whether or not the function is new, just the location of the include forces it in the jited function.
  • state.state_step_hydr is also passed in to hydr() as part of the old state dicionary, and this contains the actual function passed into the non-jitted parent function hydr() (which calls jitted function _hydr() --

sedtrn_ix = sedtrn_get_ix(state_ix, state_paths, state_info["domain"])
    # these are integer placeholders faster than calling the array look each timestep
    rsed4_ix, rsed5_ix, rsed6_ix = (
        sedtrn_ix["RSED4"],
        sedtrn_ix["RSED5"],
        sedtrn_ix["RSED6"],
    )

• to:

[rburghol]

Enabled a new directory:

I think all you have to do is make an empty __init__.py file in the new directories.  That is how you inform pip to understand that it needs to install the *.py files in that directory.  There is a new way to handle this without the __init__.py, but we are currently using the old-fangled approach.

[rburghol]

Dev notes 12/1/2025:

• Created 3 test cases to compare updates to state to insure working in HYDR before finishing in SEDTRN and RQUAL
  • PL3_5250_0001.uci (no json), PL3_5250_0001wd.json, PL3_5250_0001eq.json
  • See testing code in https://github.com/HARPgroup/HSPsquared/blob/develop-state-class/tests/testcbp/HSP2results/
  • Diff: develop...HARPgroup:HSPsquared:develop-state-class
• om_constant demand:
  • Partial validation of withdrawals with constant (PL3_5250_0001wd.json)
  • O2 = 10.0 all the time (see check_equation.py)
  • Is O3 decreased? Not sure yet. Looks the same.
• om_equation demand
  • Note yet producing results with new state, O2 = 0.0 all the time
  • state_paths when loaded from the h5 (see check_equation.py) do not yield expected variables, only those loaded via JSON, not the other standard hsp variables (like O1, O2, etc.)

om_operations["model_object_cache"].keys()
dict_keys(['/STATE/PL3_5250_0001eq', '/STATE/timer', '/STATE/PL3_5250_0001eq/RCHRES_R001', '/STATE/PL3_5250_0001eq/RCHRES_R001/wd_cfs', '/STATE/PL3_5250_0001eq/RCHRES_R001/wd_cfs/min_value', '/STATE/PL3_5250_0001eq/RCHRES_R001/IVOLwrite', '/STATE/PL3_5250_0001eq/RCHRES_R001/wd_cfs/_op_0_1'])
>>> state.state_paths
DictType[unicode_type,int64]<iv=None>({/STATE/PL3_5250_0001eq/RCHRES_R001: 0, /STATE/PL3_5250_0001eq: 1, /STATE/timer: 2, /STATE/year: 3, /STATE/month: 4, /STATE/day: 5, /STATE/hour: 6, /STATE/minute: 7, /STATE/second: 8, /STATE/weekday: 9, /STATE/dt: 10, /STATE/jday: 11, /STATE/modays: 12, /STATE/dts: 13, /STATE/PL3_5250_0001eq/RCHRES_R001/wd_cfs: 14, /STATE/PL3_5250_0001eq/RCHRES_R001/wd_cfs/min_value: 15, /STATE/PL3_5250_0001eq/RCHRES_R001/IVOLwrite: 16, /STATE/PL3_5250_0001eq/RCHRES_R001/wd_cfs/_op_0_1: 17})