rule_fcst.py

Go to the documentation of this file.

# -*- coding: utf-8 -*-
"""
Created on Mon May 18 19:48:52 2020
 
@author: cd
"""
 
 
import abc
import numpy as np
from efo.rule import RuleBase, RuleRlsSpecified
from efo.junction import ReservoirJunction
# import sys
 
# TODO: Numpy array variable should be passed as reference or pointer
# TODO: Function arguments should be labeled "ref" or "ptr" at beginning for clarification
# TODO: Function arguments that are array copies should be labeled "cpy"
# TODO: Should forecast rules also return a release schedule?
 
 
class RuleFcstBase(RuleBase):
    def __init__(self, name, timeFcst):
        # Call super class constructor
        super().__init__(name, timeFcst)
        # self.T = timeFcst
        self.releaserelease = np.full(self.T.Tcont.nSteps, np.nan)
        self.tsRlstsRls = np.full(self.T.Tcont.nSteps, int(0))
        self.releaseSchedreleaseSched = np.empty(self.T.nSteps)
 
    def calc_release(self, rlsProposed, rlsUnCtrl, stor, rlsPrev=None, qIn=None):
        rlsProposed = super().calc_release(
            rlsProposed, rlsPrev=rlsPrev, rlsUnCtrl=rlsUnCtrl, stor=stor, qIn=qIn)
        # TODO: Maybe create a property in TimeFcst that sets the current forecast datetime
        # ...indexing to the current continuous time step like done here
        # TODO: Maybe use the fcstDate property of TimeFcst here instead of dateTime[1]
        # if self.T.Tcont.curDT == self.T.tFcst[self.T.Tcont.step].dateTime[1]:
        # if self.T.Tcont.curDT == self.T.fcstDate[self.T.Tcont.step]:
        # TODO: Mabe add a function in TimeFcst that checks the current forecast date
        if np.any(self.T.fcstDate == self.T.Tcont.curDT):
            self.releaseSchedreleaseSched = self.calc_release_fcstcalc_release_fcst(rlsProposed, rlsUnCtrl, stor, qIn)
            return max(rlsProposed, self.releaseSchedreleaseSched[1].item())
        else:
            # iCurTsFcst = self.T.curFcstT.dateTime == self.T.Tcont.curDT
            # return self.releaseSched[iCurTsFcst][0] if np.any(iCurTsFcst[1:]) else 0
            iFcstTs = self.T.get_fcst_step()
            rlsCur = self.releaseSchedreleaseSched[iFcstTs].item() if iFcstTs else 0.
            return max(rlsProposed, rlsCur)
 
    # TODO: Should this include rlsPrev?
    # TODO: Should this pass a proposed release schedule?
    @abc.abstractmethod
    def calc_release_fcst(self, rlsProposed, rlsUnCtrl, stor, qIn):
        # TODO: This should return a tuple that has the release for the current time step
        # but also the release schedule and possilby the max and min?
        # TODO: In the world of OOP this should be a protected method because we only
        # want to call this from other forecast rules, but in python not sure.
        pass
 
    @classmethod
    def __subclasshook__(cls, C):
        if cls is RuleFcstBase:
            attrs = set(dir(C))
            if set(cls.__abstractmethods__) <= attrs:
                return True
        return NotImplemented
 
 
class RuleFcstEfo(RuleFcstBase):
    def __init__(self, name, timeFcst, efoResrvr, storTh, riskTol, constants, *,
                  efoNetwork=None, ruleRlsSpecified=None, efoRlsSched=None,
                  maxItr=100, storMax=None, nHrsLead=None):
        # Call super class constructor
        super().__init__(name, timeFcst)
        # Set sub class properties
        self.efoNetworkefoNetwork = efoNetwork
        self.efoResrvrefoResrvr = efoResrvr
        self.ruleRlsSpecifiedruleRlsSpecified = ruleRlsSpecified
        self.efoRlsSchedefoRlsSched = efoRlsSched
        self.maxItrmaxItr = maxItr
        self.nMembersnMembers = len(efoResrvr)
        # TODO: This shoud be a storage guide curve rule for flexibility
        self.storThstorTh = float(storTh)
        # TODO: Why is riskTol a 2D array?
        # TODO: Check to make sure that have enough risk tol lead for forecast lead
        self.riskTolriskTol = riskTol
        self.constantsconstants = constants
        self.mbrRlsmbrRls = np.full(self.T.Tcont.nSteps, int(0))
        # TODO: This should probably be nHoursLead then calc the number of steps
        self.nStepsLeadnStepsLead = int(nHrsLead/timeFcst.nHrs) + 1 if nHrsLead else timeFcst.nSteps
        # TODO: This value could also be pulled from the efoResrvr
        self.storMaxstorMax = storMax if storMax else storTh
        # if __name__ == "__main__": pool = mp.Pool(mp.cpu_count())
        # self.pool = mp.Pool(mp.cpu_count())
 
    def calc_release_fcst(self, rlsProposed, rlsUnCtrl, stor, qIn=None):
        # Set initial conditions of the ensemble networks to the current observed
        # TODO: Would it make sense to initialize the big class arrays with the constructor for performance?
        # TODO: Need to create efo network that has the primary network and the ensembles
        # TODO: Should this have max and min release properties for easier access for results?
        # TODO: This could be replaced with the observer pattern to update forecast network
        if self.efoNetworkefoNetwork is not None: 
            for net in self.efoNetworkefoNetwork: net.set_init_cond()
        # Initialize variables
        t = self.T.get_fcst_time()
        # self.storFcst = np.empty(t.nSteps, self.nMembers, t.nSteps - 1)
        # self.storFcst[:] = np.nan
        # TODO: These should be set as references to the enemble reservoirs
        self.storFcststorFcst = np.full((t.nSteps, self.nMembersnMembers, t.nSteps-1), np.nan)
        # self.storFcst[0, :, :] = [curRes.stor[0] for curRes in self.efoResrvr]
        self.storFcststorFcst[0, :, :] = self.efoResrvrefoResrvr[0].stor[0]
        self.storFcst_preRlsstorFcst_preRls = self.storFcststorFcst.copy()
        # self.rlsFcst = np.full((t.nSteps, self.nMembers, t.nSteps - 1), np.nan)
        # TODO: If your efo reservoir does not have an uncontrolled outlet then you don't need these variables
        # TODO: These should be set as references to the enemble reservoirs
        self.rlsUnCtrlFcstrlsUnCtrlFcst = np.full((t.nSteps, self.nMembersnMembers, t.nSteps-1), np.nan)
        if self.efoResrvrefoResrvr[0].outletUnCtrl:
            self.rlsUnCtrlFcstrlsUnCtrlFcst[0, :, :] = self.efoResrvrefoResrvr[0].rlsUnCtrl[0]
        else:
            self.rlsUnCtrlFcstrlsUnCtrlFcst[0, :, :] = 0.
        self.rlsUnCtrlFcst_preRlsrlsUnCtrlFcst_preRls = self.rlsUnCtrlFcstrlsUnCtrlFcst.copy()
        self.rlsFcstrlsFcst = np.zeros((t.nSteps, self.nMembersnMembers, t.nSteps-1))
        for i in range(t.nSteps-1):
            self.rlsFcstrlsFcst[i+2:, 0:self.nMembersnMembers, i] = np.nan
        # self.rlsFcst[0, :, :] = [curRes.rlsCtrl[0] for curRes in self.efoResrvr]
        # TODO: Maybe use the passed rlsPrev here?
        self.rlsFcstrlsFcst[0, :, :] = self.efoResrvrefoResrvr[0].rlsCtrl[0]
        self.rlsMaxrlsMax = np.full((t.nSteps, self.nMembersnMembers, t.nSteps-1), np.nan)
        self.rlsNoConstraintrlsNoConstraint = np.copy(self.rlsFcstrlsFcst)
        self.prExcThprExcTh = np.full(t.nSteps, np.nan)
        # self.prExcTh_preRls = np.copy(self.prExcTh)
        self.prExcTh_preRlsprExcTh_preRls = np.full((t.nSteps, t.nSteps-1), np.nan)
        self.iRiskyMbrsiRiskyMbrs = np.full((t.nSteps, self.nMembersnMembers), False)
        riskySelect = np.arange(0, self.nMembersnMembers, dtype=int)
        riskyMbr = np.copy(riskySelect)
        iRisky = np.full(self.nMembersnMembers, False)
        iRiskyChk = np.full(self.nMembersnMembers, True)
        volAbvThresh = np.zeros([t.nSteps, self.nMembersnMembers])
        vol2Rls = np.zeros([t.nSteps, self.nMembersnMembers])
        rlsFcst = np.full(t.nSteps, np.nan)
        rlsFcst[0] = self.efoResrvrefoResrvr[0].rlsCtrl[0]
        rlsMax = np.full(t.nSteps, np.nan)
        rlsToday = np.empty((t.nSteps-1, self.nMembersnMembers))
        rlsSched = np.zeros(t.nSteps)
        for tsHoriz in t.steps[1:self.nStepsLeadnStepsLead]:
            self.T.step = tsHoriz
            if tsHoriz > 1:
                self.storFcststorFcst[:, :, tsHoriz-1] = self.storFcststorFcst[:, :, tsHoriz-2]
                self.rlsUnCtrlFcstrlsUnCtrlFcst[:, :, tsHoriz-1] = self.rlsUnCtrlFcstrlsUnCtrlFcst[:, :, tsHoriz-2]
            # for j in range(self.nMembers):
            for curMbr in riskySelect:
                # TODO: Maybe shouldn't set the override to 0 here?
                if issubclass(type(self.ruleRlsSpecifiedruleRlsSpecified[curMbr]), RuleRlsSpecified):
                    self.ruleRlsSpecifiedruleRlsSpecified[curMbr].set_release(tsHoriz, 0.)
                self.efoResrvrefoResrvr[curMbr].calc_qout()
                # TODO: In a network this should just process up to the efoReseroir...
                # TODO: and not points downstream to save compute time
                # self.storFcst[tsHoriz, curMbr, tsHoriz-1] = np.copy(self.efoResrvr[curMbr].stor[tsHoriz])
                self.storFcststorFcst[tsHoriz, curMbr, tsHoriz-1] = self.efoResrvrefoResrvr[curMbr].stor[tsHoriz].item()
                if self.efoResrvrefoResrvr[curMbr].outletUnCtrl:
                    # self.rlsUnCtrlFcst[tsHoriz, curMbr, tsHoriz-1] = np.copy(self.efoResrvr[curMbr].rlsUnCtrl[tsHoriz])
                    self.rlsUnCtrlFcstrlsUnCtrlFcst[tsHoriz, curMbr, tsHoriz-1] = self.efoResrvrefoResrvr[curMbr].rlsUnCtrl[tsHoriz].item()
                else:
                    self.rlsUnCtrlFcstrlsUnCtrlFcst[tsHoriz, curMbr, tsHoriz-1] = 0.
            self.storFcst_preRlsstorFcst_preRls[:, :, tsHoriz-1] = self.storFcststorFcst[:, :, tsHoriz-1].copy()
            self.rlsUnCtrlFcst_preRlsrlsUnCtrlFcst_preRls[:, :, tsHoriz-1] = self.rlsUnCtrlFcstrlsUnCtrlFcst[:, :, tsHoriz-1].copy()
            risky = True
            while risky:
                storPlusUnCtrl = self.storFcststorFcst[tsHoriz, :, tsHoriz-1]\
                    + self.rlsUnCtrlFcstrlsUnCtrlFcst[tsHoriz, :, tsHoriz-1]*self.constantsconstants.cfs2af
                iRisky = (storPlusUnCtrl > self.storThstorTh) & iRiskyChk
                self.prExcThprExcTh[tsHoriz] = np.sum(iRisky)/self.nMembersnMembers
                if self.prExcThprExcTh[tsHoriz] - self.riskTolriskTol[tsHoriz] > 1/self.nMembersnMembers/2:
                    # if self.T.Tcont.step == 240: 
                    #     print(tsHoriz)
                    #     sys.exit("ts = " + str(tsHoriz))
                    self.prExcTh_preRlsprExcTh_preRls[:, tsHoriz-1] = self.prExcThprExcTh.copy()
                    iRiskyChk = np.copy(iRisky)
                    # riskyMbr = np.where(self.storFcst[tsHoriz, :, tsHoriz-1] > self.storTh)[0]
                    riskyMbr = np.where(iRisky)[0]
                    nMbrs2Reduce = int(len(riskyMbr) - np.floor(self.riskTolriskTol[tsHoriz]*self.nMembersnMembers))
                    volAbvThresh[tsHoriz, riskyMbr] = \
                        self.storFcststorFcst[tsHoriz, riskyMbr, tsHoriz-1]\
                            + sum(self.rlsUnCtrlFcstrlsUnCtrlFcst[1:tsHoriz+1, riskyMbr, tsHoriz-1])*self.constantsconstants.cfs2af\
                                - self.storThstorTh
                    mbrSorted = np.argsort(volAbvThresh[tsHoriz, riskyMbr])
                    riskySelect = riskyMbr[mbrSorted[0:nMbrs2Reduce]]
                    self.iRiskyMbrsiRiskyMbrs[tsHoriz, riskySelect] = True
                    # TODO: Should add a while loop so you can account for evaportion iterations
                    # TODO: This percent increase (1.01) should be a property
                    # vol2Rls[tsHoriz, riskySelect] = \
                    #     1.01*(self.storFcst[tsHoriz, riskySelect, tsHoriz-1] - self.storTh)
                    vol2Rls[tsHoriz, riskySelect] = 1.01*volAbvThresh[tsHoriz, riskySelect]
                    tasks = []
                    # for j in range(len(riskySelect)):
                    for curMbr in riskySelect:
                        # print((tsHoriz, curMbr))
                        rlsFcst[1:tsHoriz+1] = np.sum(
                                vol2Rls[:, curMbr])*self.constantsconstants.af2cfs/tsHoriz
                        self.rlsNoConstraintrlsNoConstraint[1:tsHoriz+1, curMbr, tsHoriz-1] = np.copy(rlsFcst[1:tsHoriz+1])
                        rlsFcst[1:tsHoriz+1], rlsMax[1:tsHoriz+1] = self._calc_release_sched_calc_release_sched(
                            rlsFcst[1:tsHoriz+1], tsHoriz, curMbr)
                        self.rlsMaxrlsMax[1:tsHoriz+1, curMbr, tsHoriz-1] = np.copy(rlsMax[1:tsHoriz+1])
                        if np.sum(vol2Rls[:, curMbr]) - np.sum(rlsFcst[1:tsHoriz+1])*self.constantsconstants.cfs2af  > 0.01:
                            iRiskyChk[curMbr] = False
                            # vol2Rls[tsHoriz, curMbr] = \
                            #     vol2Rls[tsHoriz, curMbr] - \
                            #     (np.sum(vol2Rls[:, curMbr]) - np.sum(rlsFcst[1:tsHoriz+1])*self.constants.cfs2af)
                            vol2Rls[tsHoriz, curMbr] -= \
                                (np.sum(vol2Rls[:, curMbr]) - np.sum(rlsFcst[1:tsHoriz+1])*self.constantsconstants.cfs2af)\
                                    - np.sum(self.rlsUnCtrlFcstrlsUnCtrlFcst[1:tsHoriz+1, curMbr, tsHoriz-1])
                        self.rlsFcstrlsFcst[1:tsHoriz+1, curMbr, tsHoriz-1] = np.copy(rlsFcst[1:tsHoriz+1])
                        # print(tsHoriz)
                    risky = True
                else:
                    if np.isnan(self.prExcTh_preRlsprExcTh_preRls[tsHoriz, tsHoriz-1]):
                        self.prExcTh_preRlsprExcTh_preRls[:, tsHoriz-1] = self.prExcThprExcTh.copy()
                    # TODO: Create a function that initializes these variables, would have to return...
                    # TODO: a tuple of multiple variables. Not sure this would save any code redundancy
                    riskySelect = np.arange(0, self.nMembersnMembers, dtype=int)
                    riskyMbr = np.copy(riskySelect)
                    iRiskyChk = np.full(self.nMembersnMembers, True)
                    risky = False
        # rlsToday[:, :] = self.rlsFcst[1, :, :].T
        # rlsToday[:, :] = np.nansum(self.rlsFcst[1:int(self.T.fcstFreq/self.T.nHrs)+1, :, :], axis=0).T
        rlsToday[:, :] = np.nanmean(self.rlsFcstrlsFcst[1:int(self.T.fcstFreq/self.T.nHrs)+1, :, :], axis=0).T
        if np.any(rlsToday[:] > 0.):
            # TODO: You could have a count of the number of ens mbrs and ts that have the same day 1 release
            # TODO: This could be used to provide multiple release schedules
            # TODO: Needs to check if there is a spill in the 24 hour horizon and if so
            # ...go with the max release to prevent this spill
            tsMax, mbrMax = np.where(rlsToday==np.max(rlsToday))
            # TODO: The time step of release should be the max of the maxes
            # TODO: to do this you have to look out to the min horiz of all of the tsMax...
            # and then determine which of those has the max total release. Tricky!
            # TODO: Check if tsRls considers 0 hr - it shouldn't and if so fix it
            # TODO: the selected release schedule should secondarily priority which...
            # ...release has the highest time step average
            self.tsRlstsRls[self.T.Tcont.step] = int(tsMax[0])
            self.mbrRlsmbrRls[self.T.Tcont.step] = int(mbrMax[0])
            # TODO: Does this need to be set since you return the release schedule and it is reset by super class?
            self.releaserelease[self.T.Tcont.step] = self.rlsFcstrlsFcst[1, mbrMax[0], tsMax[0]]
            rlsSched = self.rlsFcstrlsFcst[:, mbrMax[0], tsMax[0]]
            # Fill in the period to next Fcst if needed
            if np.any(np.isnan(rlsSched[:int(self.T.fcstFreq/self.T.nHrs)+1])): 
                tsLastRls = np.where(np.isnan(rlsSched))[0][0]
                rlsSched[tsLastRls:int(self.T.fcstFreq/self.T.nHrs)+1] = rlsSched[tsLastRls - 1]
        else:
            # TODO: Does this need to be set since you return the release schedule and it is reset by super class?
            self.releaserelease[self.T.Tcont.step] = 0.
            # rlsSched = np.zeros(t.nSteps)
        return rlsSched
    
    def _calc_release_sched(self, rlsFcst, tsHoriz, ensMbr):
        # TODO: This should start at ts 0 instead of 1 in case there is dependency on ts 0 for future applications
        # pctConvg = 0.005
        # rlsTot = np.nansum(rlsFcst)
        # maxItr = 10
        curItr = 0.
        # Create reference variables to ensemble reservoir objects
        rlsMax = np.copy(self.efoResrvrefoResrvr[ensMbr].rlsMax[1:tsHoriz+1])
        rlsCtrlApplied = self.efoResrvrefoResrvr[ensMbr].rlsCtrl[1:tsHoriz+1]
        storFcst = self.storFcststorFcst[1:tsHoriz+1, ensMbr, tsHoriz-1]
        rlsUnCtrlFcst = self.rlsUnCtrlFcstrlsUnCtrlFcst[1:tsHoriz+1, ensMbr, tsHoriz-1]
        # rlsMaxPrev = np.zeros(len(rlsMax))
        rlsFcstInit = np.copy(rlsFcst)
        rlsFcstPrev = np.array([np.copy(rlsFcst)])
        while True:
            # Recaculate storage
            for i in range(1, tsHoriz+1):
                if issubclass(type(self.efoResrvrefoResrvr[ensMbr]), ReservoirJunction):
                    self.T.step = i
                    if issubclass(type(self.ruleRlsSpecifiedruleRlsSpecified[ensMbr]), RuleRlsSpecified):
                        # self.ruleRlsSpecified[ensMbr].set_release(i, np.copy(rlsFcst[i-1]))
                        self.ruleRlsSpecifiedruleRlsSpecified[ensMbr].set_release(i, rlsFcst[i-1].item())
                    # TODO: In a network this should just process up to the efoReseroir...
                    # TODO: and not points downstream to save compute time
                    # Set break at i==7
                    self.efoNetworkefoNetwork[ensMbr].process_fcst_junctions()
                    # TODO: Try to use the locally defined reference variables instead
                    rlsMax[i-1] = self.efoResrvrefoResrvr[ensMbr].rlsMax[i].item()
                    storFcst[i-1] = self.efoResrvrefoResrvr[ensMbr].stor[i].item()
                    # rlsFcst = np.copy(rlsCtrlApplied)
                    rlsFcst[i-1] = self.efoResrvrefoResrvr[ensMbr].rlsCtrl[i].item()
                    if self.efoResrvrefoResrvr[ensMbr].outletUnCtrl:
                        rlsUnCtrlFcst[i-1] = self.efoResrvrefoResrvr[ensMbr].rlsUnCtrl[i].item()
                    else:
                        rlsUnCtrlFcst[i-1] = 0.
            # print(self.T.Tcont.step, tsHoriz, curItr)
            if curItr == 0: rlsMaxPrev = np.copy(rlsMax)
            if np.all(rlsFcstInit == rlsCtrlApplied, axis=0) or curItr >= self.maxItrmaxItr:
                break
            elif np.any(np.all(rlsFcstPrev == rlsCtrlApplied, axis=1)):
                # rlsFcst = np.copy(rlsCtrlApplied)
                # rlsMaxPrev = np.copy(rlsMax)
                break
            # Account for storage above spillway
            # We're reducing the max release by the amount that exceeds our storage target
            # so the release above max can be redistributed to another fcst timestep
            # TODO: Not sure this totally makes sense because for a low storage threshold w/ multi-obj...
            # ...you might frequently exceed the storTh but this would force you rlsMax to 0
            # TODO: This concept works of you exceed you storMax in lead time, but doesn't make sense...
            #... if you exceed it on lead 0
            # TODO: In that case you need a storTh for management and a storMax to adjusting rlsMax
            # TODO: Should this be accounted for if we are not considering uncontrolled releases?
            # TODO: Check this with Lake Mendocino for really big events
            # TODO: May want a separate variable that is the storage above the storMax that get's...
            # ...redistributed differently
            rlsMax2Spill = np.fmax(np.zeros(rlsFcst.size), 
                rlsCtrlApplied - (storFcst - self.storMaxstorMax)*self.constantsconstants.af2cfs - rlsUnCtrlFcst)
            rlsMax = np.fmin(rlsMax, rlsMax2Spill)
            # Get equally distributed release
            rlsFcstPrev = np.vstack((rlsFcstPrev, np.copy(rlsCtrlApplied)))
            rlsFcst = self._get_adjusted_release_eqldist_get_adjusted_release_eqldist(rlsFcstInit, rlsMax)
            # zChk = np.array((storFcst, rlsFcst, rlsMax, rlsCtrlApplied)).T
            rlsMaxPrev = np.copy(rlsMax)
            curItr += 1
        # return rlsFcst, rlsMaxPrev
        return rlsFcst, rlsMax
    
    def old_calc_release_sched(self, rlsFcst, tsHoriz, ensMbr):
        # TODO: This should start at ts 0 instead of 1 in case there is dependency on ts 0 for future applications
        # TODO: Should this first calculate the zero release max to help reduce iterations?
        pctConvg = 0.005
        rlsTot = np.nansum(rlsFcst)
        maxItr = 100
        curItr = 0
        # Create reference variables to ensemble reservoir objects
        rlsMax = self.efoResrvrefoResrvr[ensMbr].rlsMax[1:tsHoriz+1]
        rlsCtrlApplied = self.efoResrvrefoResrvr[ensMbr].rlsCtrl[1:tsHoriz+1]
        storFcst = self.storFcststorFcst[1:tsHoriz+1, ensMbr, tsHoriz-1]
        rlsUnCtrlFcst = self.rlsUnCtrlFcstrlsUnCtrlFcst[1:tsHoriz+1, ensMbr, tsHoriz-1]
        rlsMaxPrev = np.zeros(len(rlsMax))
        rlsFcstInit = np.copy(rlsFcst)
        rlsFcstPrev = np.array([np.copy(rlsFcst)])
        # rlsFcstPrev = np.empty([maxItr, rlsFcst.size])
        # idxZero = 1
        while True:
            # Recaculate storage
            for i in range(1, tsHoriz+1):
                if issubclass(type(self.efoResrvrefoResrvr[ensMbr]), ReservoirJunction):
                    self.T.step = i
                    if issubclass(type(self.ruleRlsSpecifiedruleRlsSpecified[ensMbr]), RuleRlsSpecified):
                        self.ruleRlsSpecifiedruleRlsSpecified[ensMbr].set_release(i, np.copy(rlsFcst[i-1]))
                    # TODO: In a network this should just process up to the efoReseroir...
                    # TODO: and not points downstream to save compute time
                    self.efoNetworkefoNetwork[ensMbr].process_fcst_junctions()
                    # TODO: Use the locally defined reference variables instead
                    storFcst[i-1] = np.copy(self.efoResrvrefoResrvr[ensMbr].stor[i])
                    rlsUnCtrlFcst[i-1] = np.copy(self.efoResrvrefoResrvr[ensMbr].rlsUnCtrl[i])
            # TODO: If you are maintaining forecasted storage levels then all future timesteps...
            # TODO: should be at the max release schedule
            if curItr == 0: rlsMaxPrev = np.copy(rlsMax)
            if (rlsTot - np.sum(rlsCtrlApplied) < pctConvg*rlsTot and np.sum(rlsUnCtrlFcst)==0.)\
                or curItr >= maxItr:
                break
            elif np.any(np.all(rlsFcstPrev == rlsCtrlApplied, axis=1)):
                rlsFcst = np.copy(rlsCtrlApplied)
                rlsMaxPrev = np.copy(rlsMax)
                break
            # zChk = np.array((rlsFcst, self.efoResrvr[ensMbr].rlsCtrl, self.storFcst[:, ensMbr, tsHoriz-1], rlsMax)).T
            # Roll back releases above target
            iAbvTarget = storFcst[:] > self.storThstorTh
            iUnCtrl = rlsUnCtrlFcst[:] > 0.
            # volAbvTarget = sum(self.storFcst[1:tsHoriz+1, ensMbr, tsHoriz-1] - self.storTh) + \
            #     sum(self.rlsUnCtrlFcst[1:tsHoriz+1, ensMbr, tsHoriz-1])*self.constants.cfs2af
            rlsMax[iAbvTarget | iUnCtrl] = np.fmax(0., rlsCtrlApplied[iAbvTarget | iUnCtrl] \
                - (storFcst[iAbvTarget | iUnCtrl] - self.storThstorTh)*self.constantsconstants.af2cfs \
                - rlsUnCtrlFcst[iAbvTarget | iUnCtrl])
            # Get equally distributed release
            rlsFcstPrev = np.vstack((rlsFcstPrev, np.copy(rlsCtrlApplied)))
            rlsFcst = self._get_adjusted_release_eqldist_get_adjusted_release_eqldist(rlsFcstInit, rlsMax)
            rlsMaxPrev = np.copy(rlsMax)
            curItr += 1
        return rlsFcst, rlsMaxPrev
    
    def _get_adjusted_release_eqldist(self, rlsFcstInit, rlsMax):
        rlsFcst = np.copy(rlsFcstInit)
        rlsFcstCheck = np.zeros(len(rlsFcst))
        iVal = ~np.isnan(rlsFcst) & ~np.isnan(rlsMax)
        iMaxed = np.full(len(rlsFcst), False)
        # iMaxed[0] = True
        while any(iMaxed==False) and any(rlsFcst[iVal] - rlsFcstCheck[iVal] != 0.):
            rlsFcstCheck = np.copy(rlsFcst)
            iGrtrMax = rlsFcst > rlsMax
            iMaxed = iMaxed | iGrtrMax
            rlsDiffVol = np.sum(rlsFcst[iGrtrMax] - rlsMax[iGrtrMax])
            rlsFcst[iGrtrMax] = rlsMax[iGrtrMax]
            rlsAdd = rlsDiffVol/max(1,np.sum(~iMaxed & iVal))
            rlsFcst[~iMaxed] += rlsAdd
        return rlsFcst
    
    def _get_adjusted_release_rollback(self, rlsFcstInit, tsHoriz, rlsMax):
        rlsFcst = np.copy(rlsFcstInit)
        rlsFcstCheck = np.copy(rlsFcst)
        carryOver = 0.
        for i in range(tsHoriz, 0, -1):
            rlsFcst[i] += carryOver
            rlsFcstCheck[i] = rlsFcst[i]
            if rlsFcst[i] > rlsMax[i]:
                rlsFcst[i] = rlsMax[i]
            # Readjust carryover
            carryOver += (rlsFcstCheck[i] -  rlsFcst[i])/(i-1) if i > 1 else 0.
        return rlsFcst
    
 
class RuleFcstPerfectEfoRisk(RuleFcstBase):
    def __init__(self, name, rulePfo, ruleEfoNoRls):
         # Call super class constructor
        super().__init__(name, rulePfo.T)
        self.rulePforulePfo = rulePfo
        self.ruleEfoNoRlsruleEfoNoRls = ruleEfoNoRls
        self.riskEforiskEfo = np.full((self.T.Tcont.nSteps, self.T.nSteps), 0.)
        
    def calc_release_fcst(self, rlsProposed, rlsUnCtrl, stor, qIn=None):
        zeroRls = self.ruleEfoNoRlsruleEfoNoRls.calc_release_fcst(rlsProposed, rlsUnCtrl, stor, qIn)
        self.riskEforiskEfo[self.T.Tcont.step,1:]  = self.ruleEfoNoRlsruleEfoNoRls.prExcTh[1:]
        rlsSched = self.rulePforulePfo.calc_release_fcst(rlsProposed, rlsUnCtrl, stor, qIn)
        # TODO: Does this need to be set since you return the release schedule and it is reset by super class?
        self.releaserelease[self.T.Tcont.step] = rlsSched[1]
        return rlsSched
    
    
class RuleFcstEfoPlotResults(RuleFcstBase):
    def __init__(self, name, ruleEfo, ruleEfo_NoRls, tsArchive=0, 
                 fileName='efoFcstResults', tsRlsOvrd=None, saveAllRlsSched=False):
         # Call super class constructor
        super().__init__(name, ruleEfo.T)
        self.ruleEforuleEfo = ruleEfo
        self.ruleEfo_NoRlsruleEfo_NoRls = ruleEfo_NoRls
        self.tsArchivetsArchive = tsArchive
        self.fileNamefileName = fileName
        self.tsRlsOvrdtsRlsOvrd = tsRlsOvrd
        self.SaveAllRlsSchedSaveAllRlsSched = saveAllRlsSched
        
    def calc_release_fcst(self, rlsProposed, rlsUnCtrl, stor, qIn=None):
        rlsSched = self.ruleEforuleEfo.calc_release_fcst(rlsProposed, rlsUnCtrl, stor, qIn)
        self.releaserelease[self.T.Tcont.step] = rlsSched[1]
        # TODO: This should take multiple time steps and archive to a 3d array of results
        if self.T.Tcont.step == self.tsArchivetsArchive:
            zeroRls = self.ruleEfo_NoRlsruleEfo_NoRls.calc_release_fcst(rlsProposed, rlsUnCtrl, stor, qIn)
            t = self.T.get_fcst_time()
            self.tFcsttFcst = t
            # self.tsRls = int(self.ruleEfo.tsRls[self.T.Tcont.step]) if self.tsRlsOvrd is None else int(self.tsRlsOvrd)
            if self.tsRlsOvrdtsRlsOvrd:
                self.tsRlstsRlstsRls = int(self.tsRlsOvrdtsRlsOvrd)
                rlsToday = np.nanmean(
                    self.ruleEforuleEfo.rlsFcst[1:int(self.T.fcstFreq/self.T.nHrs)+1, :, :self.tsRlstsRlstsRls], axis=0).T
                self.mbrRlsmbrRls = int(np.where(rlsToday==np.max(rlsToday))[1][0])
            else:
                self.tsRlstsRlstsRls = int(self.ruleEforuleEfo.tsRls[self.T.Tcont.step])
                self.mbrRlsmbrRls = int(self.ruleEforuleEfo.mbrRls[self.T.Tcont.step])
            self.storFcst_NoRlsstorFcst_NoRls = self.ruleEfo_NoRlsruleEfo_NoRls.storFcst[:, : , t.end-1]
            self.rlsUnCtrlFcst_NoRlsrlsUnCtrlFcst_NoRls = self.ruleEfo_NoRlsruleEfo_NoRls.rlsUnCtrlFcst[:, : , t.end-1]
            self.prExcTh_NoRlsprExcTh_NoRls = self.ruleEfo_NoRlsruleEfo_NoRls.prExcTh
            self.riskTolriskTol = self.ruleEforuleEfo.riskTol
            self.storFcststorFcst = self.ruleEforuleEfo.storFcst[:, : , t.end-1]
            self.rlsUnCtrlFcstrlsUnCtrlFcst = self.ruleEforuleEfo.rlsUnCtrlFcst[:, : , t.end-1]
            # self.rlsFcst = self.ruleEfo.rlsFcst[:, : , t.end-1]
            self.releaseSchedreleaseSchedreleaseSched = rlsSched
            self.prExcThprExcTh = self.ruleEforuleEfo.prExcTh
            self.iRiskyMbrsAlliRiskyMbrsAll = np.any(self.ruleEforuleEfo.iRiskyMbrs, axis=0)
            # Time step of release variables
            # TODO: The NoRls variables are temporary to create the explanation figures, kinda silly to have
            self.storFcst_NoRls_tsRlsstorFcst_NoRls_tsRls = self.ruleEfo_NoRlsruleEfo_NoRls.storFcst[:, : , self.tsRlstsRlstsRls]
            self.rlsUnCtrlFcst_NoRls_tsRlsrlsUnCtrlFcst_NoRls_tsRls = self.ruleEfo_NoRlsruleEfo_NoRls.rlsUnCtrlFcst[:, : , self.tsRlstsRlstsRls]
            self.rlsFcst_NoRls_tsRlsrlsFcst_NoRls_tsRls = self.ruleEfo_NoRlsruleEfo_NoRls.rlsFcst[:, : , self.tsRlstsRlstsRls]
            self.storFcst_preRlsstorFcst_preRls = self.ruleEforuleEfo.storFcst_preRls[:, : , self.tsRlstsRlstsRls]
            self.rlsUnCtrlFcst_preRlsrlsUnCtrlFcst_preRls = self.ruleEforuleEfo.rlsUnCtrlFcst_preRls[:, : , self.tsRlstsRlstsRls]
            self.storFcst_tsRlsstorFcst_tsRls = self.ruleEforuleEfo.storFcst[:, :, self.tsRlstsRlstsRls]
            self.rlsUnCtrlFcst_tsRlsrlsUnCtrlFcst_tsRls = self.ruleEforuleEfo.rlsUnCtrlFcst[:, : , self.tsRlstsRlstsRls]
            self.iRiskyMbrs_AlliRiskyMbrs_All = np.any(self.ruleEforuleEfo.iRiskyMbrs, axis=0)
            # TODO: Check if tsRls considers 0 hr
            self.iRiskyMbrs_tsRlsiRiskyMbrs_tsRls = self.ruleEforuleEfo.iRiskyMbrs[self.tsRlstsRlstsRls+1, :].flatten()
            self.iRiskyMbrs_tsRls_AllPreviRiskyMbrs_tsRls_AllPrev = np.any(self.ruleEforuleEfo.iRiskyMbrs[:self.tsRlstsRlstsRls+1, :], axis=1)
            self.prExcTh_preRlsprExcTh_preRls = self.ruleEforuleEfo.prExcTh_preRls[:, self.tsRlstsRlstsRls]
            # self.rlsFcst_tsRls = self.ruleEfo.rlsFcst[:, self.mbrRls, self.tsRls]
            self.rlsFcst_tsRlsrlsFcst_tsRls = self.ruleEforuleEfo.rlsFcst[:, :, self.tsRlstsRlstsRls]
            self.rlsMax_tsRlsrlsMax_tsRls = self.ruleEforuleEfo.rlsMax[:, self.mbrRlsmbrRls, self.tsRlstsRlstsRls]
            # Temporary for presentation
            if self.SaveAllRlsSchedSaveAllRlsSched:
                if self.tsRlsOvrdtsRlsOvrd:
                    self.rlsFcstAllrlsFcstAll = self.ruleEforuleEfo.rlsFcst[:, :, :self.tsRlsOvrdtsRlsOvrd].reshape(self.ruleEforuleEfo.rlsFcst.shape[0], -1) 
                else:
                    self.rlsFcstAllrlsFcstAll = self.ruleEforuleEfo.rlsFcst.reshape(self.ruleEforuleEfo.rlsFcst.shape[0], -1)  
            else: None
            # Save to a zipped numpy file
            np.savez(self.fileNamefileName+'_'+self.T.Tcont.curDT.strftime('%Y-%m-%d-%H'),
                     tFcst=t.dateTime,
                     tsRls = self.tsRlstsRlstsRls,
                     mbrRls = self.mbrRlsmbrRls,
                     releaseSched = self.releaseSchedreleaseSchedreleaseSched, 
                     storFcst_NoRls=self.storFcst_NoRlsstorFcst_NoRls,
                     rlsUnCtrlFcst_NoRls = self.rlsUnCtrlFcst_NoRlsrlsUnCtrlFcst_NoRls,
                     storFcst_preRls = self.storFcst_preRlsstorFcst_preRls,
                     rlsUnCtrlFcst_preRls = self.rlsUnCtrlFcst_preRlsrlsUnCtrlFcst_preRls,
                     prExcTh_NoRls=self.prExcTh_NoRlsprExcTh_NoRls,
                     riskTol=self.riskTolriskTol,
                     storFcst=self.storFcststorFcst,
                     rlsUnCtrlFcst = self.rlsUnCtrlFcstrlsUnCtrlFcst,
                     prExcTh=self.prExcThprExcTh,
                     iRiskyMbrsAll=self.iRiskyMbrsAlliRiskyMbrsAll,
                     storFcst_NoRls_tsRls = self.storFcst_NoRls_tsRlsstorFcst_NoRls_tsRls,
                     rlsUnCtrlFcst_NoRls_tsRls = self.rlsUnCtrlFcst_NoRls_tsRlsrlsUnCtrlFcst_NoRls_tsRls,
                     rlsFcst_NoRls_tsRls = self.rlsFcst_NoRls_tsRlsrlsFcst_NoRls_tsRls,
                     storFcst_tsRls = self.storFcst_tsRlsstorFcst_tsRls,
                     rlsUnCtrlFcst_tsRls = self.rlsUnCtrlFcst_tsRlsrlsUnCtrlFcst_tsRls,
                     iRiskyMbrs_tsRls = self.iRiskyMbrs_tsRlsiRiskyMbrs_tsRls,
                     rlsFcst_tsRls = self.rlsFcst_tsRlsrlsFcst_tsRls,
                     rlsMax_tsRls = self.rlsMax_tsRlsrlsMax_tsRls,
                     rlsFcstAll = self.rlsFcstAllrlsFcstAll,
                     iRiskyMbrs_All = self.iRiskyMbrs_AlliRiskyMbrs_All,
                     prExcTh_preRls = self.prExcTh_preRlsprExcTh_preRls
                     )
        return rlsSched
 
 
# TODO: This should keep track of the mapping between the fcst reservoir and the network.
class RuleFcstRlsSched(RuleFcstBase):
    def __init__(self, name, timeFcst, fcstNetwork):
        # Call super class constructor
        super().__init__(name, timeFcst)
        self.fcstNetworkfcstNetwork = fcstNetwork
 
    def calc_release_fcst(self, rlsProposed, rlsUnCtrl, stor):
        pass