Adding a model restore and copy tool

[avdudchenko]

Fixes/Resolves:

This is a granular alternative to pyomo model.clone tool (which does not work on large WaterTAP IDAES models).

The tool enable storage of model variable states (value, and fixed state), constraint state (activated/deactivated), and variable scaling values, and then mediates their re-aplication to the orignal model or alternative model with same structure at any point in code. Example of usage are below:

from watertap.tools.model_state_tool import modelStateStorage
model=waterTapModel()
model.init() # Initialize model state
model.init_store_point=modelStateStorage(model) 
# do stuff to model 
model.init_store_point.restore_state() # restore model state to state after initilization 
# do other stuff 

Copying model state from one flow sheet to many (should work but have not tested).

modelSet=waterTapModelSet(many blocks of same model) 
modelSet[0].init() # init first flow sheet in block
modelSet[0].base_state=modelStateStorage(model) 
for sub_model in modelSet:
      modelSet[0].base_state.restore_state(alternative_model=sub_model )

#complete initialization do other stuff

This PR is really meant to help with parameter sweep tool and failed initilization when looping. If a model fails to intilize, it can leave behind unfixed variables/constraints, and changed scaling values which will prevent the model from solving/initiating again (leading to issues such as DOF changing after failed init, or auto-scaling messing up scaling values) The objective here is to reset the model to previous good state that is known to solve, or state after model was just built.

Additionally, this tool should make it easy to copy block states from one to another to accelerate initiation.

User can specify which variable to store (Default everything is stored and restore).

model.prior_state = modelStateStorage( model, restorVars=False, restorVarFixedState=False, restorConstraints=False, restoreScaling=True, )

Legal Acknowledgement

By contributing to this software project, I agree to the following terms and conditions for my contribution:

1. I agree my contributions are submitted under the license terms described in the LICENSE.txt file at the top level of this directory.
2. I represent I am authorized to make the contributions and grant the license. If my employer has rights to intellectual property that includes these contributions, I represent that I have received permission to make contributions and grant the required license on behalf of that employer.

[andrewlee94]

@avdudchenko Are you aware of the IDAES tools for saving a model to json and restoring the information? This looks like it might be replicating that functionality.

from idaes.core.util import model_serializer as ms

[codecov]

Codecov Report

Merging #691 (3ae518d) into main (a1478fa) will increase coverage by 0.08%.
The diff coverage is 93.61%.

@@            Coverage Diff             @@
##             main     #691      +/-   ##
==========================================
+ Coverage   95.43%   95.52%   +0.08%     
==========================================
  Files         257      258       +1     
  Lines       24762    24853      +91     
==========================================
+ Hits        23632    23740     +108     
+ Misses       1130     1113      -17     

Impacted Files	Coverage Δ
watertap/property_models/NaCl_prop_pack.py	97.88% <ø> (ø)
watertap/core/membrane_channel1d.py	88.88% <50.00%> (-2.92%)	⬇️
watertap/unit_models/reverse_osmosis_base.py	96.02% <50.00%> (+0.01%)	⬆️
watertap/tools/model_state_tool.py	96.59% <96.59%> (ø)
watertap/edb/db_api.py	85.00% <0.00%> (+8.46%)	⬆️

Help us with your feedback. Take ten seconds to tell us how you rate us. Have a feature suggestion? Share it here.

[avdudchenko]

@andrewlee94: Hmm, I was not aware of this tool...I will test it out and see if it works the same. If it does then we can junk this PR unless anyone think this tool has advantage over the IDEAS tool. Looking at IDAES tool, it is not clear to me if it also copies scaling, (although seems as it does create a complete copy), so it might be a better tool.

[adam-a-a]

Given what @andrewlee94 said, my question is what gap does this fill that the other IDAES tool is missing, if any? Perhaps this could be answered with a comparative test or two.

[avdudchenko]

Given what @andrewlee94 said, my question is what gap does this fill that the other IDAES tool is missing, if any? Perhaps this could be answered with a comparative test or two.

@andrewlee94 @adam-a-a

So I started comparing the two and at the moment it seems that the json utility does not work in manner I expect. I specifically have an issue, where if your initialization fails with NF_DSPE model it will result with extra fixed variables or constraints, not sure, (leading to -X DOFs). The PR here in fixes these DOFs and the model can be reinitialized/solved with out issues.

Using the json tool does not reset the model fixed variables/constraints in same manner as the PR leading to failed solutions. I am not sure exactly why, as it seems the json dict keeps track of all states (variables, constraints, etc). (Fixed, unfixed etc) and I assume it would reset these states upon reload of the model.

I did simple tests to see if it properly resets fixed/active states for vars and constraints and it worked, but when running the model it did not fix the DOF issue, which implies that it fails to record states for some variables/constraints (although I have not tried to figure out which), while the PR here manages to capture them.

The basic code was (with no passed options it records absolutely everything it seems!)

m=model()
model_copy=to_json(m,return_dict=True)
for i in tests:
from_json(m,model_copy)
#try to initalize

If this issue is simply me missing an option in the tool, than the only gap this PR can serve is a cleaner interface and ability to stare model history on the model it self.

[avdudchenko]

To add to the comment above. I think that I will either update this tool to use json utility (in conjunction with my PR, or leave as is and add more features to it. The class based implementation is really nice. As you can attach it to the model and store the flowsheet state at different points. A use case I am using now with parameter sweep looks like this, it also ensure consistency across all of the code since I don't have to call to_json and store the dict in every script and pass it between functions (although you can store it on disk, this becomes an issue for param sweep as with multiple instances you will need to keep unique file names for each node....)

def build():
  build model
  model.build_state=modelStateStorage(m)

def init():
  model.build_state.restore_state() # ensure fresh model is used at every init. 
  # update init guess values 
  # do initi steps 
  model.init_state.store_state() # Store the initialized state for future use

[dangunter]

I think a corollary to the question @andrewlee94 is asking is whether this utility should be incorporated into the IDAES codebase. This could of course proceed in parallel with merging it in watertap for practical reasons. However I will also note that checkpoint/restart (this is application level checkpoint/restart) is a tricky topic and may have other sharp edges in the context of a workflow tool.

[ksbeattie]

@avdudchenko, moving this to the March release. Can you try to chase down some reviewers?