Skip to content

update imod.visualize.streamfunction for general case

In imod.visualize.streamfunction, a more general option for:

    # _meshcoords returns mesh edges, but useful here to allow 1D/2D
    # dimensions in da
    # go back to cell centers in 2D
    X = np.vstack(da.shape[0] * [da.s])
    Y = 0.5 * da.top + 0.5 * da.bottom
    if Y.ndim == 1:
        # promote to 2D
        Y2 = xr.concat(da.shape[1] * [Y], dim=da.s)

would be:

if len(da["top"].dims) == 1:
    top = da["top"] * xr.ones_like(da["s"])
else: 
    top = da["top"]

if len(da["bottom"].dims) == 1:
    bottom = da["bottom"] * xr.ones_like(da["s"])
else:
    bottom = da["bottom"] 

X = xr.ones_like(da["layer"]) * da["s"]
Y = 0.5 * top + 0.5 * bottom