feat(vtk): export vectors to vtk

* Implement the possibility to export vectors (e.g. hydraulic gradient, groundwater velocity) to vtk
* Add tests for it

autotest/t050_test.py

@@ -1,4 +1,5 @@
 import shutil
+import numpy as np
 import os
 import flopy
 from flopy.export import vtk
@@ -341,6 +342,38 @@ def test_vtk_cbc():
 
     return
 
+def test_vtk_vector():
+    from flopy.utils import postprocessing as pp
+    # test mf 2005 freyberg
+    mpth = os.path.join('..', 'examples', 'data',
+                        'freyberg_multilayer_transient')
+    namfile = 'freyberg.nam'
+    cbcfile = os.path.join(mpth, 'freyberg.cbc')
+    m = flopy.modflow.Modflow.load(namfile, model_ws=mpth, verbose=False,
+                                   load_only=['dis', 'bas6'])
+    q = pp.get_specific_discharge(m, cbcfile=cbcfile)
+    output_dir = os.path.join(cpth, 'freyberg_vector')
+    filenametocheck = 'discharge.vtu'
+
+    # export and check with point scalar
+    vtk.export_vector(m, q, output_dir, 'discharge', point_scalars=True)
+    filetocheck = os.path.join(output_dir, filenametocheck)
+    totalbytes = os.path.getsize(filetocheck)
+    # assert(totalbytes==2249214)
+    nlines = count_lines_in_file(filetocheck)
+    assert(nlines==10605)
+
+    # with point scalars and binary
+    vtk.export_vector(m, q, output_dir + '_bin', 'discharge', point_scalars=True,
+                      binary=True)
+    filetocheck = os.path.join(output_dir + '_bin', filenametocheck)
+    totalbytes1 = os.path.getsize(filetocheck)
+    # assert(totalbytes1==917033)
+    nlines1 = count_lines_in_file(filetocheck, binary=True)
+    assert(nlines1==2725)
+
+    return
+
 def test_vtk_vti():
     # test mf 2005 ibs2k
     mpth = os.path.join('..', 'examples', 'data', 'mf2005_test')
@@ -391,6 +424,25 @@ def test_vtk_vti():
     nlines4 = count_lines_in_file(filetocheck)
     assert(nlines4==993)
 
+    # vector
+    filenametocheck = 'T.vti'
+    T = (m.bcf6.tran.array, m.bcf6.tran.array, np.zeros(m.modelgrid.shape))
+    vtk.export_vector(m, T, output_dir, 'T', point_scalars=True)
+    filetocheck = os.path.join(output_dir, filenametocheck)
+    totalbytes5 = os.path.getsize(filetocheck)
+    # assert(totalbytes5==12621)
+    nlines5 = count_lines_in_file(filetocheck)
+    assert(nlines5==20)
+
+    # vector binary
+    vtk.export_vector(m, T, output_dir + '_bin', 'T', point_scalars=True,
+                      binary=True)
+    filetocheck = os.path.join(output_dir + '_bin', filenametocheck)
+    totalbytes6 = os.path.getsize(filetocheck)
+    # assert(totalbytes6==16716)
+    nlines6 = count_lines_in_file(filetocheck, binary=True)
+    assert(nlines6==18)
+
     return
 
 def test_vtk_vtr():
@@ -444,5 +496,6 @@ def test_vtk_vtr():
     test_vtk_mf6()
     test_vtk_binary_head_export()
     test_vtk_cbc()
+    test_vtk_vector()
     test_vtk_vti()
     test_vtk_vtr()

flopy/export/vtk.py

@@ -399,6 +399,7 @@ def __init__(self, model, verbose=None, nanval=-1e+20, smooth=False,
 
         self.cell_type = 11
         self.arrays = {}
+        self.vectors = {}
 
         self.smooth = smooth
         self.point_scalars = point_scalars
@@ -499,9 +500,9 @@ def _vtk_grid_type(self, vtk_grid_type='auto'):
         # return vtk grid type and file extension
         return (vtk_grid_type, file_extension)
 
-    def add_array(self, name, a, array2d=False):
+    def _format_array(self, a, array2d=False):
         """
-        Adds an array to the vtk object
+        Format array for vtk output
 
         Parameters
         ----------
@@ -512,6 +513,10 @@ def add_array(self, name, a, array2d=False):
             the array to be added to the vtk object
         array2d : bool
             True if the array is 2d
+
+        Return
+        ------
+        Formatted array (a copy is made)
         """
         # if array is 2d reformat to 3d array
         if array2d:
@@ -530,8 +535,55 @@ def add_array(self, name, a, array2d=False):
         where_to_nan = np.logical_or(a==self.nanval, self.ibound==0)
         a = np.where(where_to_nan, np.nan, a)
 
+        return a
+
+    def add_array(self, name, a, array2d=False):
+        """
+        Adds an array to the vtk object
+
+        Parameters
+        ----------
+
+        name : str
+            name of the array
+        a : flopy array
+            the array to be added to the vtk object
+        array2d : bool
+            True if the array is 2d
+        """
+        # format array
+        a = self._format_array(a, array2d)
+
         # add to self.arrays
-        self.arrays[name] = a
+        if a is not None:
+            self.arrays[name] = a
+
+        return
+
+    def add_vector(self, name, v, array2d=False):
+        """
+        Adds a vector (i.e., a tuple of arrays) to the vtk object
+
+        Parameters
+        ----------
+
+        name : str
+            name of the vector
+        v : tuple of arrays
+            the vector to be added to the vtk object
+        array2d : bool
+            True if the vector components are 2d arrays
+        """
+        # format each component of the vector
+        vf = ()
+        for vcomp in v:
+            vcomp = self._format_array(vcomp, array2d=array2d)
+            if vcomp is None:
+                return
+            vf = vf + (vcomp,)
+
+        # add to self.vectors
+        self.vectors[name] = vf
 
         return
 
@@ -687,6 +739,23 @@ def write(self, output_file, timeval=None):
                 a = np.flip(a, axis=1)
                 xml.write_array(a, Name=name, NumberOfComponents='1')
 
+        # loop through stored vectors
+        for name, v in self.vectors.items():
+            ncomp = len(v)
+            v_as_array = np.moveaxis(np.array(v), 0, -1)
+            if self.vtk_grid_type == 'UnstructuredGrid':
+                shape4d = actwcells3d.shape + (ncomp,)
+                actwcells4d = actwcells3d.reshape(actwcells3d.shape + (1,))
+                actwcells4d = np.broadcast_to(actwcells4d, shape4d)
+                xml.write_array(v_as_array, actwcells=actwcells4d, Name=name,
+                                NumberOfComponents=ncomp)
+            else:
+                # flip "v" so coordinates increase along with indices as in vtk
+                v_as_array = np.flip(v_as_array, axis=0)
+                v_as_array = np.flip(v_as_array, axis=1)
+                xml.write_array(v_as_array, Name=name,
+                                NumberOfComponents=ncomp)
+
         # end cell data
         xml.close_element('CellData')
 
@@ -707,8 +776,32 @@ def write(self, output_file, timeval=None):
                     # vtk
                     a = np.flip(a, axis=0)
                     a = np.flip(a, axis=1)
+                # write to file
                 xml.write_array(a, Name=name, NumberOfComponents='1')
 
+            # loop through stored vectors
+            for name, v in self.vectors.items():
+                # get the vector values onto vertices
+                v_ext = ()
+                for vcomp in v:
+                    if self.vtk_grid_type == 'UnstructuredGrid':
+                        _, _, zverts = self.get_3d_vertex_connectivity(
+                                     actwcells=actwcells3d, zvalues=vcomp)
+                        vcomp = np.array(list(zverts.values()))
+                    else:
+                        vcomp = self.extendedDataArray(vcomp)
+                    v_ext = v_ext + (vcomp,)
+                # write to file
+                ncomp = len(v_ext)
+                v_ext_as_array = np.moveaxis(np.array(v_ext), 0, -1)
+                if self.vtk_grid_type != 'UnstructuredGrid':
+                    # flip "v" so coordinates increase along with indices as in
+                    # vtk
+                    v_ext_as_array = np.flip(v_ext_as_array, axis=0)
+                    v_ext_as_array = np.flip(v_ext_as_array, axis=1)
+                xml.write_array(v_ext_as_array, Name=name,
+                                NumberOfComponents=ncomp)
+
             # end point data
             xml.close_element('PointData')
 
@@ -723,6 +816,7 @@ def write(self, output_file, timeval=None):
 
         # clear arrays
         self.arrays.clear()
+        self.vectors.clear()
 
     def _configure_data_arrays(self):
         """
@@ -745,6 +839,16 @@ def _configure_data_arrays(self):
             # set the active array to 1
             ot_idx_array[idxs] = 1
 
+        # loop through vectors
+        for name in self.vectors:
+            for vcomp in self.vectors[name]:
+                # make array 1d
+                a = vcomp.ravel()
+                # get the indexes where there is data
+                idxs = np.argwhere(np.logical_not(np.isnan(a)))
+                # set the active array to 1
+                ot_idx_array[idxs] = 1
+
         # reset the shape of the active data array
         ot_idx_array = ot_idx_array.reshape(self.shape)
 
@@ -1203,6 +1307,60 @@ def export_array(model, array, output_folder, name, nanval=-1e+20,
     return
 
 
+def export_vector(model, vector, output_folder, name, nanval=-1e+20,
+                 array2d=False, smooth=False, point_scalars=False,
+                 vtk_grid_type='auto', binary=False):
+
+    """
+
+    Export vector (i.e., a tuple of arrays) to vtk
+
+    Parameters
+    ----------
+
+    model : flopy model instance
+        the flopy model instance
+    vector : tuple of arrays
+        vector to be exported
+    output_folder : str
+        output folder to write the data
+    name : str
+        name of vector
+    nanval : scalar
+        no data value, default value is -1e20
+    array2d : bool
+            True if the vector components are 2d arrays, default is False
+    smooth : bool
+        if True, will create smooth layer elevations, default is False
+    point_scalars : bool
+        if True, will also output array values at cell vertices, default is
+        False; note this automatically sets smooth to True
+    vtk_grid_type : str
+            Specific vtk_grid_type or 'auto'. Possible specific values are
+            'ImageData', 'RectilinearGrid', and 'UnstructuredGrid'.
+            If 'auto', the grid type is automatically determined. Namely:
+                * A regular grid (in all three directions) will be saved as an
+                  'ImageData'.
+                * A rectilinear (in all three directions), non-regular grid
+                  will be saved as a 'RectilinearGrid'.
+                * Other grids will be saved as 'UnstructuredGrid'.
+    binary : bool
+        if True the output file will be binary, default is False
+
+    """
+
+    if not os.path.exists(output_folder):
+        os.mkdir(output_folder)
+
+    vtk = Vtk(model, nanval=nanval, smooth=smooth, point_scalars=point_scalars,
+              vtk_grid_type=vtk_grid_type, binary=binary)
+    vtk.add_vector(name, vector, array2d=array2d)
+    otfile = os.path.join(output_folder, '{}'.format(name))
+    vtk.write(otfile)
+
+    return
+
+
 def export_transient(model, array, output_folder, name, nanval=-1e+20,
                      array2d=False, smooth=False, point_scalars=False,
                      vtk_grid_type='auto', binary=False):