Use explicit xarray's indexes: Delete UgridDataArray, UgridDataset

[Huite]

With the changes so far, this snippet already works!

import xarray as xr
import xugrid as xu

# %%
# Grab some data with UGRID conventions, run ugrid.initialize to create the UgridIndex:

ds = xu.data.elevation_nl(xarray=True)
uds = ds.ugrid.initialize()

# %%
# Grab one variable (would've been a "UgridDataArray" before) and try plotting it:
uda = uds["elevation"]
uda.ugrid.plot()

# %%
# Get it directly as a DataArray with a UgridIndex:
uda = xu.data.elevation_nl()
uda.ugrid.plot()
# %%

# Get rid of roughly half the data
masked = uda.where(uda["mesh2d_face_x"] < 150_000)
masked.ugrid.plot()
# %%

# Fill it back in with Laplace interpolation
filled = masked.ugrid.laplace_interpolate()
filled.ugrid.plot()

[deltamarnix]

Amazing if we can pull this off and can get rid of custom classes.

[deltamarnix]

Does this automatically fix isinstance(x, UgridDataset) so we don't need to create classes with deprecation warnings?

[Huite]

No, this is just so that UgridDataArray and UgridDataset are defined, so I can still import xugrid while I haven't purged UgridDataArray and UgridDataset from the code base. To be fair, this was a thinko, I should have just defined UgridDataArray = None.

[deltamarnix]

Are these classes part of the API? Perhaps add documentation and type hints

[Huite]

Yes, they will be; this is the bare minimum to get anything working. We'll implement some of of the xarray Index base class methods here: xarray.Index

At least: copy, equals, rename
Less priority, arguable: sel, concat, join, reindex_like, maybe create_variables
Won't do: isel, roll, stack, to_pandas_index, unstack

[deltamarnix]

How would you eventually envision going for fully unstructured grids as a program like flopy would need?

[Huite]

The first step is to define a Ugrid3d topology, which is where all the hard work lies. For Ugrid1d's the fundamental connectivity is edge_node_connectivity, and for Ugrid2d it's the face_node_connectivity. For Ugrid3d, we need a volume_node_connectivity (and we'll require a z coordinate in the node_coordinates).

The 3D indexes would then look like:

class UgridIndex3d
     def __init__(
          self,
          node_x,
          node_y,
          node_z,
          volume_node_connectivity,
     ):

For the specific code here, we'd check the ugrid dimension of the topology and dispatch.

So something like this

UGRID_INDEXES = {
     1: UgridIndex1d,
     2: UgridIndex2d,
     3: UgridIndex3d,
}
for topology in self.obj.ugrid_roles.topology:
     ugrid_vars = self.obj.ugrid_roles[topology]
     topology_dimension = self[topology].attrs["topology_dimension"]   # should be added to the UgridRoles accessor instead
     index_type = UGRID_INDEXES[topology_dimension]
     # Collect variables, maybe through a static method?
     variables = index_type._collect_variables(self.obj, ugrid_vars)
     new = new.set_coords(variables).set_xindex(variables, index_type)

And to pick a nit, Flopy doesn't need "fully" unstructured, since the MODFLOW grids are always prismatic (horizontal and vertical faces). This has major benefits, since we won't need the more complicated 3D geometry algorithms, but we can rely on 2D (x, y) operations instead and apply the z operation afterwards. This should also be more efficient since it's more specialized.

The only argument against it is that prismatic grids are a bit odd and they might be MODFLOW 6 only as I'm not aware of any simulation codes using it, they would generally use tetrahedrons. But that's an acceptable idiosyncracy to me.

[Huite]

The actual implementation will look a bit more complicated due to the optional or "optional required" attributes that the UGRID conventions prescribe, such as the edge_node_connectivity for a 2D topology (not always relevant), but this is just a case of providing some optional arguments to the init and including them as coordinates -- so basically the number of names provided in variables will not be constant.