mfsubwt.py

"""
mfsubwt module.  Contains the ModflowSubwt class. Note that the user can access
the ModflowSubwt class as `flopy.modflow.ModflowSubwt`.

Additional information for this MODFLOW package can be found at the `Online
MODFLOW Guide
<http://water.usgs.gov/ogw/modflow/MODFLOW-2005-Guide/subwt.htm>`_.

"""
import sys

import numpy as np

from ..pakbase import Package
from ..utils import Util2d, Util3d, read1d


class ModflowSubwt(Package):
    """
    MODFLOW SUBWT Package Class.

    Parameters
    ----------
    model : model object
        The model object (of type :class:`flopy.modflow.mf.Modflow`) to which
        this package will be added.
    ipakcb : int
        A flag that is used to determine if cell-by-cell budget data should be
        saved. If ipakcb is non-zero cell-by-cell budget data will be saved.
        (default is 0).
    iswtoc : int
        isuboc is a flag used to control output of information generated by the
        SUB Package. (default is 0).
    nsystem : int
        nsystem is the number of systems of interbeds
    ithk : int
        ithk is a flag to determine how thicknesses of compressible sediments
         vary in response to changes in saturated thickness (default is -1).
    ivoid : int
        ivoid is a flag to determine how void ratios of compressible sediments
         vary in response to changes in saturated thickness (default is -1).
    istpcs : int
        istpcs is a flag to determine how initial preconsolidation stress
         will be obtained (default is 0)
    icrcc : int
        icrcc is a flag to determine how recompression and compression
        indices will be obtained (default is 1).
    lnwt : array of ints
        lnwt is a one-dimensional array specifying the model-layer
        assignments for each system of interbeds. The array has NSYSTM values
    izcfl : int
        izcfl is a flag to specify whether or not initial calculated values
        of layer-center elevation will be printed (default is -1).
    izcfm : int
        izcfm is a code for the format in which layer-center elevation will
        be printed (default is 0).
    iglfl : int
        is a flag to specify whether or not initial calculated values of
        geostatic stress will be printed (default is -1).
    iglfm : int
        iglfm is a code for the format in which geostatic stress will
        be printed (default is 0).
    iestfl : int
        iestfl is a flag to specify whether or not initial calculated values
         of effective stress will be printed (default is -1)
    iestfm : int
        iestfm is a code for the format in which effective stress will be
        printed (default is 0).
    ipcfl : int
        ipcfl is a flag to specify whether or not initial calculated values
        of preconsolidation stress will be printed (default is -1).
    ipcfm : int
        ipcfm is a code for the format in which preconsolidation stress
        will be printed (default is 0).
    istfl : int
        istfl is a flag to specify whether or not initial equivalent storage
        properties will be printed for each system of interbeds. If ICRCC≠0,
        the equivalent storage properties that can be printed are
        recompression and compression indices (Cr and Cc), which are
        calculated from elastic and inelastic skeletal specific storage
        (Sske and Sskv). If ICRCC=0, equivalent storage properties that can
        be printed are elastic and inelastic skeletal specific storage,
        which are calculated from the recompression and compression indices.
        (default is -1)
    istfm : int
        istfm is a code for the format in which equivalent storage properties
         will be printed (default is 0).

    gl0 : array of floats
        is an array specifying the geostatic stress above model layer 1.
        If the top of model layer 1 is the land surface, enter values
        of zero for this array.(default is 0.0)

    sgm : array of floats
        sgm is an array specifying the specific gravity of moist or
        unsaturated sediments (default is 1.7).
    sgs : array of floats
        sgs is an array specifying the specific gravity of saturated sediments.
        (default is 2.0)
    thick : array of floats
        thick is an array specifying the thickness of compressible sediments.
        Note: pass a list of thick values to have a different value applied
        to each nsystm of interbeds
    sse : array of floats (default is 1.0).
        sse is an array specifying the initial elastic skeletal specific
        storage of compressible beds (default is ). Note: pass a list of sse
        values to have a different value applied to each nsystm of interbeds
    ssv : array of floats
        ssv is an array specifying the initial inelastic skeletal
        specific storage of compressible beds (default is ). Note: pass a
        list of ssv values to have a different value applied to each
        nsystm of interbeds
    cr : array of floats
        cr is an array specifying the recompression index of compressible
        beds (default is 0.01). Note: pass a list of cr values
        to have a different value applied to each nsystm of interbeds
    cc : array of floats
        cc is an array specifying the compression index of compressible beds.
        (default is 0.25). Note: pass a list of cc values
        to have a different value applied to each nsystm of interbeds
    void : array of floats
        void is an array specifying the initial void ratio of compressible
        beds (default is 0.82). Note: pass a list of void values
        to have a different value applied to each nsystm of interbeds
    sub : array of floats
        sub is an array specifying the initial compaction in each system of
        interbeds. Compaction values computed by the package are added to
        values in this array so that printed or stored values of compaction
        and land subsidence may include previous components. Values in this
        array do not affect calculations of storage changes or resulting
        compaction. For simulations in which output values will reflect
        compaction and subsidence since the start of the simulation, enter
        zero values for all elements of this array (default is 0.0).
        Note: pass a list of sub values
        to have a different value applied to each nsystm of interbeds
    pscoff : array of floats
        pscoff is an array specifying the offset from initial effective
        stress to initial preconsolidation stress at the bottom of the
        model layer in units of height of a column of water (default is 0.0).
    pcs : array of floats
        pcs is an array specifying the initial preconsolidation stress,
        in units of height of a column of water, at the bottom of the
        model layer (default is 0.0).
    ds16 : array of ints (26)
        ds16 is required if iswtoc > 0.  controls the writing of subwt output.
        see the subwt users manual for more information regarding this
        dataset (default is None)
    ds17 : array of ints (30)
        ds17 controls the starting and ending stress periods related to the
        printing of output in ds16.  see the subwt users manual for more
        information regarding this dataset (default is None)

    Attributes
    ----------

    Methods
    -------

    See Also
    --------

    Notes
    -----
    Parameters are supported in Flopy only when reading in existing models.
    Parameter values are converted to native values in Flopy and the
    connection to "parameters" is thus nonexistent. Parameters are not
    supported in the SUBWT Package.

    Examples
    --------

    >>> import flopy
    >>> m = flopy.modflow.Modflow()
    >>> sub = flopy.modflow.ModflowSubwt(m)

    """

    def __init__(self, model, ipakcb=None, iswtoc=0, nsystm=1,
                 ithk=0,ivoid=0,istpcs=0,icrcc=0,lnwt=1,izcfl=0,
                 izcfm=0,iglfl=0,iglfm=0,iestfl=0,iestfm=0,
                 ipcsfl=0,ipcsfm=0,istfl=0,istfm=0,gl0=0.0,sgm=1.7,
                 sgs=2.0,thick=1.0,sse=1.0e-4,ssv=1.0e-3,
                 cr=0.01,cc=0.25,void=0.82,sub=0.0,pcsoff=0.0,
                 pcs=0.0,ds16=None,ds17=None,
                 extension='swt', unitnumber=33):
        """
        Package constructor.

        """
        # set default unit number of one is not specified
        if unitnumber is None:
            unitnumber = ModflowSubwt.defaultunit()

        # update external file information with cbc output, if necessary
        if ipakcb is not None:
            pth = model.name + '.' + ModflowSubwt.ftype() + '.cbc'
            model.add_externalbudget(ipakcb, fname=pth)
        else:
            ipakcb = 0

        extensions = [extension]
        name = [ModflowSubwt.ftype()]
        units = [unitnumber]
        extra = ['']

        item16_extensions = ["subsidence.hds", "total_comp.hds",
                             "inter_comp.hds", "vert_disp.hds",
                             "precon_stress.hds", "precon_stress_delta.hds",
                             "geostatic_stress.hds","geostatic_stress_delta.hds",
                             "eff_stress.hds","eff_stress_delta.hds",
                             "void_ratio.hds","thick.hds","lay_center.hds"]
        item16_units = [2052 + i for i in range(len(item16_extensions))]

        if iswtoc > 0:
            extensions.append('subbud')
            name.append('DATA(BINARY)')
            units.append(2051)
            extra.append('REPLACE')
            for e, u in zip(item16_extensions, item16_units):
                extensions.append(e)
                name.append('DATA(BINARY)')
                units.append(u)
                extra.append('REPLACE')


        # Call ancestor's init to set self.parent, extension, name and unit number
        Package.__init__(self, model, extension=extensions, name=name,
                         unit_number=units, extra=extra)

        nrow, ncol, nlay, nper = self.parent.nrow_ncol_nlay_nper
        line = '# Subsidence (SUBWT) package file for ' + \
               '{}, generated by Flopy.'.format(model.version)
        self.heading = line
        self.url = 'subwt.htm'

        self.ipakcb = ipakcb
        self.iswtoc = int(iswtoc)
        if self.iswtoc > 1:
            raise NotImplementedError("too lazy")
        self.nsystm = int(nsystm)
        self.ithk = int(ithk)
        self.ivoid = int(ivoid)
        self.istpcs = int(istpcs)
        self.icrcc = int(icrcc)

        self.izcfl = int(izcfl)
        self.izcfm = int(izcfm)
        self.iglfl = int(iglfl)
        self.iglfm = int(iglfm)
        self.iestfl = int(iestfl)
        self.iestfm = int(iestfm)
        self.ipcsfl = int(ipcsfl)
        self.ipcsfm = int(ipcsfm)
        self.istfl = int(istfl)
        self.istfm = int(istfm)

        self.lnwt = Util2d(model,(self.nsystm,),np.int,lnwt,name="lnwt",
                           locat=self.unit_number[0])
        self.gl0 = Util2d(model,(nrow,ncol),np.float32,gl0,name="gl0",
                          locat=self.unit_number[0])
        self.sgs = Util2d(model,(nrow,ncol),np.float32,sgs,name="sgs",
                          locat=self.unit_number[0])
        self.sgm = Util2d(model,(nrow,ncol),np.float32,sgm,name="sgm",
                          locat=self.unit_number[0])

        if isinstance(thick,list):
            assert len(thick) == self.nsystm
        else:
            thick = [thick] * self.nsystm
        self.thick = [Util2d(model,(nrow,ncol),np.float32,t,name="thick_nsy{0}"\
                             .format(i),locat=self.unit_number[0])\
                      for i,t in enumerate(thick)]

        if self.icrcc == 0:
            if isinstance(cc,list):
                assert len(cc) == self.nsystm
            else:
                cc = [cc] * self.nsystm
            self.cc = [Util2d(model,(nrow,ncol),np.float32,t,name="cc_nsy{0}"\
                             .format(i),locat=self.unit_number[0])\
                       for i,t in enumerate(cc)]
            if isinstance(cr,list):
                assert len(cr) == self.nsystm
            else:
                cr = [cr] * self.nsystm
            self.cr = [Util2d(model,(nrow,ncol),np.float32,t,name="cr_nsy{0}"\
                             .format(i),locat=self.unit_number[0])\
                       for i,t in enumerate(cr)]

        else:
            if isinstance(sse,list):
                assert len(sse) == self.nsystm
            else:
                sse = [sse] * self.nsystm
            self.sse = [Util2d(model,(nrow,ncol),np.float32,t,name="sse_nsy{0}"\
                             .format(i),locat=self.unit_number[0])\
                        for i,t in enumerate(sse)]
            if isinstance(ssv,list):
                assert len(ssv) == self.nsystm
            else:
                ssv = [ssv] * self.nsystm
            self.ssv = [Util2d(model,(nrow,ncol),np.float32,t,name="ssv_nsy{0}"\
                             .format(i),locat=self.unit_number[0])\
                        for i,t in enumerate(ssv)]

        if isinstance(void,list):
            assert len(void) == self.nsystm
        else:
            void = [void] * self.nsystm
        self.void = [Util2d(model,(nrow,ncol),np.float32,t,name="void_nsy{0}"\
                             .format(i),locat=self.unit_number[0])\
                     for i,t in enumerate(void)]

        if isinstance(sub,list):
            assert len(sub) == self.nsystm
        else:
            sub = [sub] * self.nsystm
        self.sub = [Util2d(model,(nrow,ncol),np.float32,t,name="sub_nsy{0}"\
                             .format(i),locat=self.unit_number[0])\
                    for i,t in enumerate(sub)]

        if self.istpcs == 0:
            self.pcs = Util3d(model,(nlay,nrow,ncol),np.float32,pcs,name="pcs",
                              locat=self.unit_number[0])
            self.pcsoff = None
        else:
            self.pcsoff = Util3d(model,(nlay,nrow,ncol),np.float32,pcsoff,
                                 name="pcsoff",locat=self.unit_number[0])

        if self.iswtoc > 0:
            if ds16 is None:
                self.ds16 = np.zeros((26),dtype=np.int)
                ui = 0
                for i in range(1,26,2):
                    self.ds16[i] = item16_units[ui]
                    ui += 1
            else:
                if isinstance(ds16,list):
                    ds16 = np.array(ds16)
                assert len(ds16) == 26
                self.ds16 = ds16

            if ds17 is None:
                ds17 = np.ones((30),dtype=np.int)
                ds17[0] = 0
                ds17[2] = 0
                ds17[1] = 9999
                ds17[3] = 9999
                self.ds17 = ds17
            else:
                if isinstance(ds17,list):
                    ds17 = np.array(ds17)
                assert len(ds17) == 30
                self.ds17 = ds17

        # add package to model
        self.parent.add_package(self)

    def write_file(self, check=False):
        """
        Write the package file.

        Returns
        -------
        None

        """
        if check:
            print("warning: check not implemented for subwt")
        nrow, ncol, nlay, nper = self.parent.nrow_ncol_nlay_nper
        # Open file for writing
        f = open(self.fn_path, 'w')
        # First line: heading
        f.write('{}\n'.format(self.heading))
        # write dataset 1
        f.write('{} {} {} {} {} {} {}\n'.format(self.ipakcb, self.iswtoc,
                                              self.nsystm,self.ithk,
                                              self.ivoid, self.istpcs,
                                              self.icrcc))
        lnwt = self.lnwt.array + 1
        [f.write("{} ".format(i)) for i in lnwt]
        f.write("\n")
        f.write('{} {} {} {} {} {} '.format(self.izcfl, self.izcfm,
                                          self.iglfl, self.iglfm,
                                          self.iestfl,self.iestfm))
        f.write("{} {} {} {}\n".format(self.ipcsfl,self.ipcsfm,self.istfl,
                                       self.istfm))

        f.write(self.gl0.get_file_entry())

        f.write(self.sgm.get_file_entry())

        f.write(self.sgs.get_file_entry())

        for isystm in range(self.nsystm):
            f.write(self.thick[isystm].get_file_entry())
            if self.icrcc == 0:
                f.write(self.cr[isystm].get_file_entry())
                f.write(self.cc[isystm].get_file_entry())
            else:
                f.write(self.sse[isystm].get_file_entry())
                f.write(self.ssv[isystm].get_file_entry())

            f.write(self.void[isystm].get_file_entry())
            f.write(self.sub[isystm].get_file_entry())

        p = self.pcsoff
        if self.istpcs == 0:
            p = self.pcs
        for k in range(nlay):
            f.write(p[k].get_file_entry())

        if self.iswtoc > 0:
            [f.write("{0} ".format(i)) for i in self.ds16]
            f.write('  #dataset 16\n')
            t = self.ds17.copy()
            t[0:4] += 1 #zero-based to one-based
            for i in t:
                f.write('{} '.format(i))
            f.write('  #dataset 17\n')
        # close sub file
        f.close()

    @staticmethod
    def load(f, model, ext_unit_dict=None):
        """
        Load an existing package.

        Parameters
        ----------
        f : filename or file handle
            File to load.
        model : model object
            The model object (of type :class:`flopy.modflow.mf.Modflow`) to
            which this package will be added.
        ext_unit_dict : dictionary, optional
            If the arrays in the file are specified using EXTERNAL,
            or older style array control records, then `f` should be a file
            handle.  In this case ext_unit_dict is required, which can be
            constructed using the function
            :class:`flopy.utils.mfreadnam.parsenamefile`.

        Returns
        -------
        subwt : ModflowSubwt object

        Examples
        --------

        >>> import flopy
        >>> m = flopy.modflow.Modflow()
        >>> sub = flopy.modflow.ModflowSubwt.load('test.subwt', m)

        """

        if model.verbose:
            sys.stdout.write('loading subwt package file...\n')

        if not hasattr(f, 'read'):
            filename = f
            f = open(filename, 'r')
        # dataset 0 -- header
        while True:
            line = f.readline()
            if line[0] != '#':
                break
        # determine problem dimensions
        nrow, ncol, nlay, nper = model.get_nrow_ncol_nlay_nper()

        # read dataset 1
        if model.verbose:
            sys.stdout.write('  loading subwt dataset 1\n')
        t = np.array(line.strip().split(),dtype=int)
        ds1_names = "ISWTCB ISWTOC NSYSTM ITHK IVOID ISTPCS ICRCC".lower().split()
        ds3_names = "IZCFL IZCFM IGLFL IGLFM IESTFL IESTFM IPCSFL IPCSFM ISTFL ISTFM".lower().split()
        kwargs = {n:v for n,v in zip(ds1_names,t)}
        ipakcb = kwargs["ipakcb"]
        if ipakcb > 0:
            ipakcb = 53

        if model.verbose:
            sys.stdout.write('  loading sub dataset 2\n')
        lnwt = np.empty((kwargs["nsystm"]), dtype=np.int)
        lnwt = read1d(f, lnwt) - 1
        t = np.array(f.readline().strip().split(),dtype=np.int)
        kw_temp = {n:i for n,i in zip(ds3_names,t)}
        kwargs.update(kw_temp)
        for k in range(ndb):
            t = Util2d.load(f, model, (nrow, ncol), np.float32,
                            'rnb delay bed {}'.format(k + 1),
                            ext_unit_dict)
            rnb[k] = t
        hc = None
        sfe = None
        sfv = None
        com = None
        if nndb > 0:
            hc = [0] * nndb
            sfe = [0] * nndb
            sfv = [0] * nndb
            com = [0] * nndb
            for k in range(nndb):
                kk = ln[k] + 1
                # hc
                if model.verbose:
                    sys.stdout.write(
                        '  loading sub dataset 5 for layer {}\n'.format(kk))
                t = Util2d.load(f, model, (nrow, ncol), np.float32,
                                'hc layer {}'.format(kk),
                                ext_unit_dict)
                hc[k] = t
                # sfe
                if model.verbose:
                    sys.stdout.write(
                        '  loading sub dataset 6 for layer {}\n'.format(kk))
                t = Util2d.load(f, model, (nrow, ncol), np.float32,
                                'sfe layer {}'.format(kk),
                                ext_unit_dict)
                sfe[k] = t
                # sfv
                if model.verbose:
                    sys.stdout.write(
                        '  loading sub dataset 7 for layer {}\n'.format(kk))
                t = Util2d.load(f, model, (nrow, ncol), np.float32,
                                'sfv layer {}'.format(kk),
                                ext_unit_dict)
                sfv[k] = t
                # com
                if model.verbose:
                    sys.stdout.write(
                        '  loading sub dataset 8 for layer {}\n'.format(kk))
                t = Util2d.load(f, model, (nrow, ncol), np.float32,
                                'com layer {}'.format(kk),
                                ext_unit_dict)
                com[k] = t

        # dp
        dp = None
        if ndb > 0:
            dp = np.zeros((nmz, 3), dtype=np.float32)
            for k in range(nmz):
                if model.verbose:
                    sys.stdout.write(
                        '  loading sub dataset 9 for material zone {}\n'.format(
                            k + 1))
                line = f.readline()
                t = line.strip().split()
                dp[k, :] = float(t[0]), float(t[1]), float(t[2])

        dstart = None
        dhc = None
        dcom = None
        dz = None
        nz = None
        if ndb > 0:
            dstart = [0] * ndb
            dhc = [0] * ndb
            dcom = [0] * ndb
            dz = [0] * ndb
            nz = [0] * ndb
            for k in range(ndb):
                kk = ldn[k] + 1
                # dstart
                if model.verbose:
                    sys.stdout.write(
                        '  loading sub dataset 10 for layer {}\n'.format(kk))
                t = Util2d.load(f, model, (nrow, ncol), np.float32,
                                'dstart layer {}'.format(kk),
                                ext_unit_dict)
                dstart[k] = t
                # dhc
                if model.verbose:
                    sys.stdout.write(
                        '  loading sub dataset 11 for layer {}\n'.format(kk))
                t = Util2d.load(f, model, (nrow, ncol), np.float32,
                                'dhc layer {}'.format(kk),
                                ext_unit_dict)
                dhc[k] = t
                # dcom
                if model.verbose:
                    sys.stdout.write(
                        '  loading sub dataset 12 for layer {}\n'.format(kk))
                t = Util2d.load(f, model, (nrow, ncol), np.float32,
                                'dcom layer {}'.format(kk),
                                ext_unit_dict)
                dcom[k] = t
                # dz
                if model.verbose:
                    sys.stdout.write(
                        '  loading sub dataset 13 for layer {}\n'.format(kk))
                t = Util2d.load(f, model, (nrow, ncol), np.float32,
                                'dz layer {}'.format(kk),
                                ext_unit_dict)
                dz[k] = t
                # nz
                if model.verbose:
                    sys.stdout.write(
                        '  loading sub dataset 14 for layer {}\n'.format(kk))
                t = Util2d.load(f, model, (nrow, ncol), np.int,
                                'nz layer {}'.format(kk),
                                ext_unit_dict)
                nz[k] = t

        ids15 = None
        ids16 = None
        if isuboc > 0:
            # dataset 15
            if model.verbose:
                sys.stdout.write(
                    '  loading sub dataset 15 for layer {}\n'.format(kk))
            ids15 = np.empty(12, dtype=np.int)
            ids15 = read1d(f, ids15)
            iu = 1
            for k in range(1, 12, 2):
                model.add_pop_key_list(ids15[k])
                ids15[k] = 2051 + iu  # all subsidence data sent to unit 2051
                iu += 1
            # dataset 16
            ids16 = [0] * isuboc
            for k in range(isuboc):
                if model.verbose:
                    sys.stdout.write(
                        '  loading sub dataset 16 for isuboc {}\n'.format(
                            k + 1))
                t = np.empty(17, dtype=np.int)
                t = read1d(f, t)
                t[0:4] -= 1
                ids16[k] = t
        model.add_pop_key_list(2051)
        # close file
        f.close()

        # determine specified unit number
        unitnumber = None
        if ext_unit_dict is not None:
            for key, value in ext_unit_dict.items():
                if value.filetype == ModflowSubwt.ftype():
                    unitnumber = key

        # create sub instance
        sub = ModflowSubwt(model, ipakcb=ipakcb, isuboc=isuboc, idsave=idsave,
                         idrest=idrest,
                         nndb=nndb, ndb=ndb, nmz=nmz, nn=nn, ac1=ac1, ac2=ac2,
                         itmin=itmin,
                         ln=ln, ldn=ldn, rnb=rnb,
                         hc=hc, sfe=sfe, sfv=sfv, com=com, dp=dp,
                         dstart=dstart, dhc=dhc, dcom=dcom, dz=dz, nz=nz,
                         ids15=ids15, ids16=ids16, unitnumber=unitnumber)
        # return sub instance
        return sub

    @staticmethod
    def ftype():
        return 'SWT'

    @staticmethod
    def defaultunit():
        return 33