CellBudgetFile2.py

# CellBudgetFile2.py
import sys
import flopy.utils.binaryfile as bf

class CellBudgetFile2(bf.CellBudgetFile):

    def __init__(self, filename, precision="auto", verbose=False, **kwargs):
        super().__init__(filename, precision, verbose, **kwargs)
        # Modflow6DataSource properties
        self.hasSpecificDischargeData = bytes("{:>16}".format("DATA-SPDIS"), 'utf-8') in self.textlist
        self.numberOfBudgetItems = len(self.textlist)

    @staticmethod
    def IsModelFeature(flowType):
        if flowType == None:
            return False
        if "MVR" in flowType:
            return False
        if "UZF-GWD" in flowType:
            return False
        if "UZF-GWET" in flowType:
            return False
        if "LAK" in flowType:
            return False
        if "FLOW-JA-FACE" in flowType:
            return False
        if "DATA-SPDIS" in flowType:
            return False
        ##  This was used to fix Model Viewer 1.0.0
        ##    if "CSUB" in flowType:
        ##        return False
        return True

    def CountBudgetAndFeatures(self):
        dict = {}    #    keys = m_BudgetText    values = vtk_cell_count
        idx = 0
        TXT2ID2 = "{:<16}".format("CSUB")  # Expected when imeth==6 and (flowtype=="CSUB-ELASTIC" or flowtype=="CSUB-INELASTIC")
        for rec in self.recordarray:
            if isinstance(rec, bytes):
                rec = rec.decode()
            flowtype = rec[2].decode()     # Record 1 TEXT
            imeth = rec[6]                 # Record 2 IMETH
            if imeth == 6 and self.IsModelFeature(flowtype):
                text = rec[13].decode()    # Record 6 TXT2ID2
                ##  
                ##  special cases
                if "CSUB-ELASTIC" in flowtype:
                    assert text == TXT2ID2, f"Expected '{TXT2ID2}' got: '{text}'"
                    text = f"{flowtype.strip() : <16}"
                if "CSUB-INELASTIC" in flowtype:
                    assert text == TXT2ID2, f"Expected '{TXT2ID2}' got: '{text}'"
                    text = f"{flowtype.strip() : <16}"
                if not text in dict:
                    dict[text] = 0
                nlist = (self.get_data(idx)[0].shape)[0]
                dict[text] = max(dict[text], nlist)
            idx += 1
        return dict

def main(argv) -> int:
    try:
        if (len(argv) != 2) :
            print(f"Usage: {argv[0]} budgetfile")
            print("   ie")
            print(f"       {argv[0]} ex-gwf-csub-p04.cbc")
            return 1
        filename = argv[1]
        cbf2 = CellBudgetFile2(filename, precision="double", verbose=False)
        print(f"{filename}:")
        dct = cbf2.CountBudgetAndFeatures()
        print(dct)
        mfas = 0
        for key in dct:
            mfas += dct[key] + 1
        print(f"m_ModelFeatureArraySize = {mfas}")
        ##print(f"m_NumberOfBudgetItems = {cbf2.numberOfBudgetItems}")
        ##print(f"m_HasSpecificDischargeData = {cbf2.hasSpecificDischargeData}")
    except Exception as exc:
        print("{0}".format(exc))
        return 2
    return 0

if __name__ == "__main__":
    sys.exit(main(sys.argv))