vtk_export_seawater_intrusion.py

import os, re
import numpy as np
from osgeo import gdal
import flopy
from flopy.export import vtk as fv
import vtk
import flopy.utils.binaryfile as fpu
import flopy.utils.binaryfile as bf
import matplotlib.pyplot as plt
from workingFunctions import Functions

zobj = flopy.utils.CellBudgetFile('model1.zta')
kstpkper_list = zobj.get_kstpkper()
zeta = []
zeta.append(zobj.get_data(kstpkper=(0,0), text='ZETASRF  1')[0])

mf1 = flopy.modflow.Modflow.load('model1.nam', verbose=False,check=False)
dis = flopy.modflow.ModflowDis.load('model1.dis', mf1)
bot = dis.botm.array
top = dis.top.array
zs = np.zeros((zobj.nrow, zobj.ncol), np.float)
zall = []
zall.append(top)
eps = 1.0e-3
for i in range (0, zobj.nlay):
    zall.append(bot[i])
for irow in range(0, zobj.nrow):
    for icol in range (0, zobj.ncol):
        for klay in range(0, zobj.nlay):
            zt = zeta[0][klay, irow, icol]
            if (zall[klay][irow, icol] - zt) > eps:
                zs[irow, icol] = zt
for irow in range(0, zobj.nrow):
    for icol in range (0, zobj.ncol):
        if zs[irow, icol] == 0:
            zs[irow, icol] = 0.63
        #if zs[irow, icol] == zall[zobj.nlay][irow, icol]:
            #zs[irow, icol] = np.nan
plt.figure()
plt.imshow(zs,cmap='jet')
plt.colorbar()
plt.show()

def GetExtent(gt,cols,rows):
    ext=[]
    xarr=[0,cols]
    yarr=[0,rows]

    for px in xarr:
        for py in yarr:
            x=gt[0]+(px*gt[1])+(py*gt[2])
            y=gt[3]+(px*gt[4])+(py*gt[5])
            ext.append([x,y])
            print(x,y)
        yarr.reverse()
    return ext
mf1 = flopy.modflow.Modflow.load('model1.nam', verbose=False,check=False, load_only=['upw', 'dis'])
hk = mf1.upw.hk

raster=r'dem.tif'
ds=gdal.Open(raster)
gt=ds.GetGeoTransform()
cols = ds.RasterXSize
rows = ds.RasterYSize
ext=GetExtent(gt,cols,rows)

# change directory to the script path
os.chdir('H:/Users/gauvain/DEM')  # use your own path

# open the DIS, BAS and FHD and DRN files
disLines = open('model1.dis').readlines()  # discretization data
basLines = open('model1.bas').readlines()  # active / inactive data
hds = fpu.HeadFile('model1.hds')
times = hds.get_times()
kstpkper = hds.get_kstpkper()

textoVtk = open('VTU_Saltwater_Freshwater_Interface.pvd', 'w')
textoVtk.write('<VTKFile type="Collection" version="0.1">\n')
textoVtk.write('  <Collection>\n')
for time_step in range (0,731):
    textoVtk.write('    <DataSet timestep="' + str(time_step) + '" part="0" file="VTU_Saltwater_Freshwater_Interface_'+str(time_step)+'.vtu" />\n')
textoVtk.write('  </Collection>\n')
textoVtk.write('</VTKFile>\n')
textoVtk.close()

for time_step in range (0,len(kstpkper)
                        ):
    print(time_step)
    # create a empty dictionay to store the model features
    modDis = {}
    modFhd = {}

    modDis["vertexXmin"] = float(ext[0][0])
    modDis["vertexYmin"] = float(ext[2][1])
    modDis["vertexXmax"] = float(ext[2][0])
    modDis["vertexYmax"] = float(ext[0][1])
    # get the number of layers, rows, columns, cell and vertex numbers
    linelaycolrow = disLines[1].split()
    modDis["cellLays"] = int(linelaycolrow[0])
    modDis["cellRows"] = int(linelaycolrow[1])
    modDis["cellCols"] = int(linelaycolrow[2])
    modDis["vertexLays"] = modDis["cellLays"] + 1
    modDis["vertexRows"] = modDis["cellRows"] + 1
    modDis["vertexCols"] = modDis["cellCols"] + 1
    modDis["vertexperlay"] = modDis["vertexRows"] * modDis["vertexCols"]
    modDis["cellsperlay"] = modDis["cellRows"] * modDis["cellCols"]
    # ### Get the DIS Breakers
    modDis['disBreakers'] = {}
    breakerValues = ["INTERNAL", "CONSTANT"]
    vertexLay = 0
    for item in breakerValues:
        for line in disLines:
            if item in line:
                if 'delr' in line:  # DELR is cell width along rows
                    modDis['disBreakers']['DELR'] = disLines.index(line)
                elif 'delc' in line:  # DELC is cell width along columns
                    modDis['disBreakers']['DELC'] = disLines.index(line)
                else:
                    modDis['disBreakers']['vertexLay' + str(vertexLay)] = disLines.index(line)
                    vertexLay += 1
    modDis['DELR'] = Functions.getListFromDEL(modDis['disBreakers']['DELR'], disLines, modDis['cellCols'])
    modDis['DELC'] = Functions.getListFromDEL(modDis['disBreakers']['DELC'], disLines, modDis['cellRows'])
    modDis['cellCentroidZList'] = {}
    for lay in range(modDis['vertexLays']):
        # add auxiliar variables to identify breakers
        lineaBreaker = modDis['disBreakers']['vertexLay' + str(lay)]
        # two cases in breaker line
        if 'INTERNAL' in disLines[lineaBreaker]:
            lista = Functions.getListFromBreaker(lineaBreaker, modDis, disLines)
            modDis['cellCentroidZList']['lay' + str(lay)] = lista
        elif 'CONSTANT' in disLines[lineaBreaker]:
            constElevation = float(disLines[lineaBreaker].split()[1])
            modDis['cellCentroidZList']['lay' + str(lay)] = [constElevation for x in range(modDis["cellsperlay"])]
        else:
            pass
    modDis['vertexEasting'] = np.array(
        [modDis['vertexXmin'] + np.sum(modDis['DELR'][:col]) for col in range(modDis['vertexCols'])])
    modDis['vertexNorthing'] = np.array(
        [modDis['vertexYmax'] - np.sum(modDis['DELC'][:row]) for row in range(modDis['vertexRows'])])
    modDis['cellEasting'] = np.array(
        [modDis['vertexXmin'] + np.sum(modDis['DELR'][:col]) + modDis['DELR'][col] / 2 for col in
         range(modDis['cellCols'])])
    modDis['cellNorthing'] = np.array(
        [modDis['vertexYmax'] - np.sum(modDis['DELC'][:row]) - modDis['DELC'][row] / 2 for row in
         range(modDis['cellRows'])])

    modFhd['cellHeadGrid'] = {}
    lay = 0
    head = hds.get_data(kstpkper=kstpkper[time_step])
    for i in range(0, head.shape[0]):
        modFhd['cellHeadGrid']['lay' + str(lay)] = zs
        lay += 1

    listLayerQuadSequence = []

    # definition of hexahedrons cell coordinates
    for row in range(modDis['cellRows']):
        for col in range(modDis['cellCols']):
            pt0 = modDis['vertexCols'] * (row + 1) + col
            pt1 = modDis['vertexCols'] * (row + 1) + col + 1
            pt2 = modDis['vertexCols'] * (row) + col + 1
            pt3 = modDis['vertexCols'] * (row) + col
            anyList = [pt0, pt1, pt2, pt3]
            listLayerQuadSequence.append(anyList)

    vertexHeadGridCentroid = {}
    # arrange to hace positive heads in all vertex of an active cell
    for lay in range(modDis['cellLays']):
        matrix = np.zeros([modDis['vertexRows'], modDis['vertexCols']])
        for row in range(modDis['cellRows']):
            for col in range(modDis['cellCols']):
                headLay = zs
                neighcartesianlist = [headLay[row, col], headLay[row, col], headLay[row, col],
                                      headLay[row, col]]
                headList = []
                for item in neighcartesianlist:
                    if item > -1e+30:
                        headList.append(item)
                if len(headList) > 0:
                    headMean = sum(headList) / len(headList)
                else:
                    headMean = -1e+30

                matrix[row, col] = headMean

        matrix[-1, :-1] = zs[-1, :]
        matrix[:-1, -1] = zs[:, -1]
        matrix[-1, -1] = zs[-1, -1]

        vertexHeadGridCentroid['lay' + str(lay)] = matrix

    # empty temporal dictionary to store transformed heads
    vertexHKGridCentroid = {}

    # arrange to hace positive heads in all vertex of an active cell
    for lay in range(modDis['cellLays']):
        matrix = np.zeros([modDis['vertexRows'], modDis['vertexCols']])
        for row in range(modDis['cellRows']):
            for col in range(modDis['cellCols']):
                headLay = hk.array[lay]
                neighcartesianlist = [headLay[row, col], headLay[row, col], headLay[row, col],
                                      headLay[row, col]]
                headList = []
                for item in neighcartesianlist:
                    if item > -1e+30:
                        headList.append(item)
                if len(headList) > 0:
                    headMean = sum(headList) / len(headList)
                else:
                    headMean = -1e+30

                matrix[row, col] = headMean

        matrix[-1, :-1] = modFhd['cellHeadGrid']['lay' + str(lay)][-1, :]
        matrix[:-1, -1] = modFhd['cellHeadGrid']['lay' + str(lay)][:, -1]
        matrix[-1, -1] = modFhd['cellHeadGrid']['lay' + str(lay)][-1, -1]

        vertexHKGridCentroid['lay' + str(lay)] = matrix

    # In[15]:

    modFhd['vertexHeadGrid'] = {}
    for lay in range(modDis['vertexLays']):
        anyGrid = vertexHeadGridCentroid
        if lay == modDis['cellLays']:
            modFhd['vertexHeadGrid']['lay' + str(lay)] = anyGrid['lay' + str(lay - 1)]
        elif lay == 0:
            modFhd['vertexHeadGrid']['lay0'] = anyGrid['lay0']
        else:

            value = np.where(anyGrid['lay' + str(lay)] > -1e+30,
                             anyGrid['lay' + str(lay)],
                             (anyGrid['lay' + str(lay - 1)] + anyGrid['lay' + str(lay)]) / 2
                             )
            modFhd['vertexHeadGrid']['lay' + str(lay)] = value

    # empty numpy array for the water table
    waterTableVertexGrid = np.zeros((modDis['vertexRows'], modDis['vertexCols']))
    # obtain the first positive or real head from the head array
    for row in range(modDis['vertexRows']):
        for col in range(modDis['vertexCols']):
            anyList = []
            for lay in range(modDis['cellLays']):
                anyList.append(modFhd['vertexHeadGrid']['lay' + str(lay)][row, col])
            a = np.asarray(anyList)
            if list(a[a > -1e+10]) != []:  # just in case there are some inactive zones
                waterTableVertexGrid[row, col] = a[a > -1e+10][0]
            else:
                waterTableVertexGrid[row, col] = -1e+10

    # empty list to store all vertex Water Table XYZ
    vertexWaterTableXYZPoints = []
    # definition of xyz points for all vertex
    for row in range(modDis['vertexRows']):
        for col in range(modDis['vertexCols']):
            if waterTableVertexGrid[row, col] > -1e+10:
                waterTable = waterTableVertexGrid[row, col]
            else:
                waterTable = 1000
            xyz = [
                modDis['vertexEasting'][col],
                modDis['vertexNorthing'][row],
                waterTable
            ]
            vertexWaterTableXYZPoints.append(xyz)

    waterTableCellGrid = np.zeros((modDis['cellRows'], modDis['cellCols']))

    # obtain the first positive or real head from the head array
    for row in range(modDis['cellRows']):
        for col in range(modDis['cellCols']):
            anyList = []
            for lay in range(modDis['cellLays']):
                anyList.append(modFhd['cellHeadGrid']['lay' + str(lay)][row, col])
            a = np.asarray(anyList)
            if list(a[a > -1e+10]) != []:  # just in case there are some inactive zones
                waterTableCellGrid[row, col] = a[a > -1e+10][0]
            else:
                waterTableCellGrid[row, col] = -1e+10

    listWaterTableCell = list(waterTableCellGrid.flatten())

    listWaterTableQuadSequenceDef = []
    listWaterTableCellDef = []
    listDrawdownCellDef = []
    for item in range(len(listWaterTableCell)):
        if listWaterTableCell[item] > -1e10:
            listWaterTableQuadSequenceDef.append(listLayerQuadSequence[item])
            listWaterTableCellDef.append(listWaterTableCell[item])
    for item in range(len(listWaterTableCellDef)):
        drawdown = modDis['cellCentroidZList']['lay0'][item] - listWaterTableCellDef[item]
        listDrawdownCellDef.append(drawdown)

    textoVtk = open('VTU_Saltwater_Freshwater_Interface_'+str(time_step)+'.vtu', 'w')
    # add header
    textoVtk.write('<VTKFile type="UnstructuredGrid" version="1.0" byte_order="LittleEndian" header_type="UInt64">\n')
    textoVtk.write('  <UnstructuredGrid>\n')
    textoVtk.write('    <Piece NumberOfPoints="' + str(len(vertexWaterTableXYZPoints)) + '" NumberOfCells="' +
                   str(len(listWaterTableCellDef)) + '">\n')
    # cell data
    textoVtk.write('      <CellData Scalars="Water Table">\n')
    textoVtk.write('        <DataArray type="Float64" Name="Altitude" format="ascii">\n')
    for item in range(len(listWaterTableCellDef)):
        textvalue = str(listWaterTableCellDef[item])
        if item == 0:
            textoVtk.write('          ' + textvalue + ' ')
        elif item % 20 == 0:
            textoVtk.write(textvalue + '\n          ')
        else:
            textoVtk.write(textvalue + ' ')
    textoVtk.write('\n')
    textoVtk.write('        </DataArray>\n')


    textoVtk.write('      </CellData>\n')
    # points definition
    textoVtk.write('      <Points>\n')
    textoVtk.write('        <DataArray type="Float64" Name="Points" NumberOfComponents="3" format="ascii">\n')
    for item in range(len(vertexWaterTableXYZPoints)):
        tuplevalue = tuple(vertexWaterTableXYZPoints[item])
        if item == 0:
            textoVtk.write("          %.2f %.2f %.2f " % tuplevalue)
        elif item % 4 == 0:
            textoVtk.write('%.2f %.2f %.2f \n          ' % tuplevalue)
        elif item == len(vertexWaterTableXYZPoints) - 1:
            textoVtk.write("%.2f %.2f %.2f \n" % tuplevalue)
        else:
            textoVtk.write("%.2f %.2f %.2f " % tuplevalue)
    textoVtk.write('        </DataArray>\n')
    textoVtk.write('      </Points>\n')
    # cell connectivity
    textoVtk.write('      <Cells>\n')
    textoVtk.write('        <DataArray type="Int64" Name="connectivity" format="ascii">\n')
    for item in range(len(listWaterTableQuadSequenceDef)):
        textoVtk.write('          ')
        textoVtk.write('%s %s %s %s \n' % tuple(listWaterTableQuadSequenceDef[item]))
    textoVtk.write('        </DataArray>\n')
    # cell offsets
    textoVtk.write('        <DataArray type="Int64" Name="offsets" format="ascii">\n')
    for item in range(len(listWaterTableQuadSequenceDef)):
        offset = str((item + 1) * 4)
        if item == 0:
            textoVtk.write('          ' + offset + ' ')
        elif item % 20 == 0:
            textoVtk.write(offset + ' \n          ')
        elif item == len(listWaterTableQuadSequenceDef) - 1:
            textoVtk.write(offset + ' \n')
        else:
            textoVtk.write(offset + ' ')
    textoVtk.write('        </DataArray>\n')
    # cell types
    textoVtk.write('        <DataArray type="UInt8" Name="types" format="ascii">\n')
    for item in range(len(listWaterTableQuadSequenceDef)):
        if item == 0:
            textoVtk.write('          ' + '9 ')
        elif item % 20 == 0:
            textoVtk.write('9 \n          ')
        elif item == len(listWaterTableQuadSequenceDef) - 1:
            textoVtk.write('9 \n')
        else:
            textoVtk.write('9 ')
    textoVtk.write('        </DataArray>\n')
    textoVtk.write('      </Cells>\n')
    # footer
    textoVtk.write('    </Piece>\n')
    textoVtk.write('  </UnstructuredGrid>\n')
    textoVtk.write('</VTKFile>\n')

    textoVtk.close()