How to work with large data models

[deltamarnix]

One of the bigger questions that we still have is how we are going to handle large datasets, also depending on the type of data (grid based / array based).

We should investigate how we can work with large data models where the data is chunked, preferably on disk. We see the benefits of using a package like Dask. Package that build on top of that are numpy, xarray, and uxarray. We are wondering if a package like uxarray also supports our more complex DISU model (fully unstructured 3d grid).

We have to consider how we are going to let users provide data to the flopy package. That can be via indexed arrays, but that can also be with a grid mask where NaNs are filled in when no data is provided. Depending on the sparsity of the data that is typically used, we can make a decision on the type of data we want users to put in via the API.

import xarray as xa
import uxarray as uxa

class ModflowGwfriv(mfpackage.MFPackage):
    def __init__(self, stress_period_data: Union[xa.Grid, uxa.UnstructuredGrid, ...], ...):
        # If the river package requires information on all cells, it's likely best to accept a full grid
        pass

import numpy as np

class  ModflowGwfriv(mfpackage.MFPackage):
    def __init__(self, stress_period_data: np.ArrayLike, ...):
        # This option would not give the user much information. The array is a 1-d array where the order of the data is important, like the current implementation. [cellid, stage, cond, rbot, aux, boundname]
        pass

import numpy as np

class  ModflowGwfriv(mfpackage.MFPackage):
    def __init__(self, cellids: np.ArrayLike, stage: np.ArrayLike, cond: np.ArrayLike, rbot: np.ArrayLike, aux: np.ArrayLike, boundname: np.ArrayLike, ...):
        # As an alternative, the user could supply separate arrays and we combine it for the user and check that the lengths of these arrays are equal. It gives the user more information on what needs to go in, but the documentation has to be very clear on what is considered equal for the user.
        pass

import pandas as pa

class  ModflowGwfriv(mfpackage.MFPackage):
    def __init__(self, stress_period_data: pa.DataFrame, ...):
        # Similar to a single np.ArrayLike, we could use pandas DataFrame so that the data is labelled and the order of the data is not relevant anymore.
        pass

Packages of interest:

• Xarray
• UXArray
• Zarr
• XUGrid
• Dask
• Numpy

XArray seems to support on disk formats for zarr and netcdf: https://docs.xarray.dev/en/stable/api.html#dataset-methods. Letting users open these data formats would only let us reference to it, perhaps write in it if needed. At the moment that the data needs to go to the MF6 input format, we can convert at that moment and not earlier. Making it more efficient.
If MF6 would support netcdf formats directly, there would be no conversion needed and it would save even more time.

[wpbonelli]

Just to note, MF6 distinguishes array and list data.

Arrays are always homogeneous and often (but not always) represent some variable living on features of the grid. I guess xarray would work only where this is true?

Lists can be homogeneous (elements all having the same record type) or heterogeneous (elements can be any of several record types). The former are amenable to numpy, pandas, providing columns separately, etc. The latter not so. Period data can be heterogeneous for some packages and so will not always be representable in a tabular format.

Also, for homogeneous arrays, maybe we could consider type hinting input context class attributes with e.g. numpy.typing.NDArray[np.int_] (and potentially providing a shape parameter in future, once numpy supports that). This would provide more info than simply ArrayLike, while the init method and other APIs could stay lenient and have an ArrayLike hint, and try to coerce the argument to an ndarray (or subclass / numpy-interoperable array class) of the proper type/shape. Kind of a best-of-both-worlds approach?

[deltamarnix]

Maybe we can work with DataTree? https://docs.xarray.dev/en/latest/user-guide/data-structures.html#datatree

[wpbonelli]

That is nice- especially the dimension inheritance. And it even has hierarchical addressing. It could represent the entire simulation context, at least all the space- or time-indexed variables, I guess?

Reminds me, we've talked about distinguishing "configuration" from "data" variables before, but I didn't have a clear idea where the line was. I am thinking now that if a variable is alignable with xarray it can be considered data? I'm not sure how useful a distinction this is, but maybe it has consequences, like which variables we can store in and/or present to user in a DataTree.

[wpbonelli]

some more notes on DataTree

• per the docs, x, y, z, time etc would be dimensional coords, while kstp and kper could be non-dimensional coordinates
• we could use custom accessors for e.g. plotting, layers, any information or functionality we want to attach to a dataset

[wpbonelli]

We've discussed/prototyped a modflow-specific duck array type, supporting e.g. an optional multiplication factor, convenient layer access, and optimized storage for constant arrays. I wonder now if xarray can do at least the first two: layer as a labeled dimension and/or coordinate, factor stored as metadata and accessed/applied at write time.

[wpbonelli]

xarray works with duck arrays so any extra duck array types we provide would work with it

[wpbonelli]

Some questions we might be able to answer with prototypes in the near term:

• can xarray (via uxarray etc) represent DISV/DISU satisfactorily?
  • imod-python has disv https://deltares.github.io/imod-python/examples/mf6/circle.html#create-ugriddataarray
  • disu graph representation Disu as graph #58
• what kind of improvements would xarray offer for large datasets?*
• why is flopy3 IO slow? is it at the write layer or further up the stack?*
• can we use DataTree as an internal data store for the entire simulation context?
  • mentioned in How to work with large data models #49 (reply in thread)
  • prototyped in demo xarray with attrs #57
• understand how flopy creates/maintains grids and dis/v/u packages

*These will require some benchmarking/profiling/performance testing, comparing flopy3 legacy IO with numpy/pandas/xarray IO

[deltamarnix]

I've researched some of the packages that seem to support UGRID. It's hard to find a library that supports 3D unstructured grids. I found the following so far:

• UGridPy, a Deltares product. Only supports 1D and 2D.
• gridded is a package that only supports 2D grids (lat and lon). And it seems to store everything in memory.
• XUGrid, a Deltares product by Huite. Only supports 2D, and can extend with adding an extra XArray dimension for layer in order to support DISV.
• UXArray, currently being worked on by a masters student at NCAR. Supports 2D with XArray, so would likely also be able to add an extra dimension layer for DISV.

I haven't found any DISU compatible libraries. And I was thinking that it might be wise to keep the FloPy design as close to the MODFLOW6 input format as possible, like it is right now. So perhaps it's best as @mjreno suggested, to use numpy-like arrays so we can support dask where possible as input formats. Preferable as lazy arrays. And the moment that we need to write the MF6 files, we compute() the values and write them directly to the file format that is needed.

I will try to make a little prototype that shows how that works.

[deltamarnix]

I'm still playing in my head with the fact that DISU is actually just a graph, and should be treated that way. The graph itself might fit into memory, but the data on the nodes and edges might not.

There area some core packages in python that represent graphs: python graphblas, networkx, scipy. And on top of that there is a package that can translate graphs from one to another, or call algorithms from one on another type: metagraph (metagraph was last released in 2021, so not as interesting anymore).

So far we want to stay as close as possible with the packages to the input of modflow, but we might want to use these packages in generator helper functions networkx -> gwfdisu. Or we might accept a networkx-graph as input if it is set up correctly and convert it to our own representation internally.

For the output we can do something similar. By default we return a single list of data points based on the cell index, but if there are use cases for traversal, we can add a translation layer on top to return a networkx or graphblas object with the data linked to the correct nodes and edges.

[deltamarnix]

I'm diving a little deeper into what XArray supports for representing a graph (DISU) in its data format. XArray supports numpy sparse arrays.

NetworkX has the option to transform its representation to numpy arrays to_numpy_array. This information can be directly stored in an XArray object. And the other way around, a network representation in XArray can be transformed back to a NetworkX object with from_numpy_array.

[deltamarnix]

I made a little prototype in #58 to show how we could store networkx graph via a sparse scipy matrix into an XArray Dataset, and get the structure back into networkx too.

[wpbonelli]

a question: should we send all array data through xarray/dask and do lazy/chunked operations on it, or operate on data below a certain size in memory?

[deltamarnix]

We were discussing how XArray and CAttrs would work together like in this PR: #62 .

cattrs will always return a dictionary of some sort, so that it can be passed into packages like pyyaml or tomlkit. cattrs provides the conversion for datetimes and encoding. But cattrs is not made for file-io, because when you would want to convert everything into one big string, that would all be passed into memory in a single instance.

It would be better to use functionality like jinja for file writing and just pass it the actual python instance of simulation / model / package. We can write filter functions that can handle the dask chunking while writing large datasets from a netcdf file / lazy dask array. This is something we still need to prototype.