How to get back to a UxDataset?

[Chrismarsh]

This is probably a dumb question, but I'm struggling to figure out how to get back to a UxDataset. I don't think this is an AB problem, but I will write what I am trying just in case.

I have a UxDataset and I want to subset it. The .subset method only seems to be applicable to a single variable, so I end up in a situation like this

swe = ds.swe.subset.bounding_box(lon, lat)   # Snow
t = ds.t.subset.bounding_box(lon, lat)  # air temp
elev = ds.Elevation.subset.bounding_box(lon, lat)  # elevation

Now, I have my subset, I want to do something like "find all snow at an elevation > 1500 m". So I need to merge all these UxDataArrays back together. This works:

all = xr.merge([t.to_xarray(), elev.to_xarray()])

as .merge does not seem to exist in the uxarray api -- correct?

So I can now answer my query, but suppose I want to get back to the goodness of uxarray -- further subsetting, plotting, etc.

How do I do that? I thought it would be as easy as

all = xr.merge([t.to_xarray(), elev.to_xarray()])
ux.UxDataset().from_xarray(all)

but it is not: RuntimeError: Could not recognize dataset format.

I've tried the various from_ methods, but none seem to work. I'm sure this is a PEBKAC, but I'm struggling to figure out how I'm violating the assumption in the ux api.

Would appreciate any insight!

[Chrismarsh]

@philipc2 I'm wondering if you had any thoughts on this?

[philipc2]

Thanks for the ping, I was out of office when you originally posted this.

The easiest way (if you already have a Grid object) is to use the default constructor and pass in the grid that matches the data.

# Grid object 
grid = ... 

# Xarray data that resides on the grid above
ds = ... 

uxds = ux.UxDataset(ds, uxgrid=uxgrid)

[philipc2]

as .merge does not seem to exist in the uxarray api -- correct?

While it does exist via inheritance, we haven't implemented our own version, so there may be some inconsistencies.

[Chrismarsh]

And I guess that uxgrid should be the subset, e.g.,

t=ds.t.isel(time=slice(0,4)).subset.bounding_box([-115.4,-155.1],[50.7,50.9])
dem=ds.Elevation.subset.bounding_box([-115.4,-155.1],[50.7,50.9])
all = xr.merge([t.to_xarray(), dem.to_xarray()])
all = ux.UxDataset(all, uxgrid=dem.uxgrid)

I have had problems in the past where xr.merge calls required an xr.align to enforce the lat/lon coords to be identical despite any numerical noise. Is ux.UxDataset(all, uxgrid=dem.uxgrid) enforcing that all of all is defined on the grid in dem.uxgrid?

[philipc2]

Now, I have my subset, I want to do something like "find all snow at an elevation > 1500 m"

If you want to mask out all values over the threshold with np.nan and then do something like plot:

elevation.where(elevation > 2000).plot()

If you want to remove the elements that are below the threshold and create a new Grid:

import numpy as np
indices_over_threshold = np.arange(elevation.uxgrid.n_face)[elevation > 2000]

elevation.isel(n_face=indices_over_threshold).plot()

The main difference here is that the first example preserves the original grid, but masks values with np.nan. The second one filters and reduces the data along the new Grid.

[Chrismarsh]

Great thanks very much!