Xarray Office Hours November 15, 2022

[dcherian]

High-level discussion: #7216

Announcement

Registration: https://www.eventbrite.com/e/xarray-virtual-office-hours-2022-tickets-452159089007

When: 11/15/2022 9-11 PST (calendar)

Where: Online (Zoom link to be provided to those registered).

Cost: Free!

Who: Will cap this first session at 20 people so that we can ensure small group interactions and utilize breakout rooms for screen sharing and pair programming

What: Get to know developers and users of Xarray and have an opportunity to troubleshoot issues and learn new skills! We are hosting a virtual office hour session via Zoom. Anyone is welcome to join, for all or part of the session. We encourage you to come with a specific question or piece of code to work on, for example, is any Xarray documentation (https://docs.xarray.dev/) or tutorial (https://tutorial.xarray.dev) material challenging to understand? Do you have questions about adopting examples to your own data analysis?
Please also consider posting specific questions on Xarray’s Discussions page! https://github.com/pydata/xarray/discussions.

Let’s make progress and grow the fantastic Xarray community together!

Please review our Code of Conduct before the event! (https://github.com/pydata/xarray/blob/main/CODE_OF_CONDUCT.md)

[dcherian]

@pydata/xarray (or general community) Please comment in a reply here if you can help out with this session.

Members of the community are also very welcome to participate as facilitators!

[TomNicholas]

I can attend! (but now only for the second hour)

[keewis]

I would be available as well!

[JessicaS11]

I'll be there

[scottyhq]

i'll be there too!

[mgrover1]

I can attend!

[dcherian]

Looks like we have 10 signups so it'll be pretty close to a 1-on-1 setup if everyone who signed up can still make it.

[rsignell-usgs]

I would be happy to help out also if desired. But I didn't sign up (registration was full).

Hmmm, and thinking some more about this....
... even if you have enough instructors, would it be possible to be a fly on the wall? I would like to see how it goes to inform similar initiatives in the future (I'm intending to spend most of 2023 helping with training and mentoring).

[TomNicholas]

I can't attend the first half of the xarray office hours anymore (have a meeting about SGRID conventions instead at 12pm EST xgcm/xgcm#109). So you could replace me if you wanted @rsignell-usgs . But it also sounds like there are plenty of devs free.

[dcherian]

I've emailed you the info, Rich.

Tom, the SGRID meeting sounds interesting. Too bad I can't make it, please make the notes public if you can.

I'm intending to spend most of 2023 helping with training and mentoring

Thank you for your service! :)

[andrewdnolan]

Thanks hosting this event, it was huge help!

I've created a gist with my problem description, some notes, and the solution we settled on:

• https://gist.github.com/andrewdnolan/3214d59375e667aaf573d3e0e2264570

I've also put this info in the hackmd file.

Thanks again!

[rsignell-usgs]

@andrewdnolan , thanks so much for posting this!
Were you able to get your zarr dataset to complete writing?

[andrewdnolan]

Thanks for the recommendation!

I've updated the sample data array to be a Dataset and slightly more similar in size to the real files I'm working with:

import numpy as np
import xarray as xr 
import pandas as pd
import dask.array
from dask.distributed import Client

client = Client()

# hard code the array sizes
NT = 20000  # number of timesteps
NN = 3552   # number of nodes
NX = 222    # number nodes in the x direction
NZ = 16     # number nodes in the z direction


# create dask array
array = dask.array.ones((NT, NN), chunks=(-1, -1))

ds = xr.Dataset({
    f'variable_{i}':(["time", "nMesh_node"], array)  for i in range(20)
    },
    coords={"time": np.linspace(0, 2000, NT)},
)


# reshape by construct.coarsen 
coarsen = ds.coarsen(nMesh_node=NX).construct(nMesh_node=("coord_2", "coord_1")) 


# reshape by unstack 
coord_2 = np.arange(0, NZ)
coord_1 = np.arange(0, NX)
spatial = pd.MultiIndex.from_product((coord_2, coord_1),names=('coord_2','coord_1'))

unstack = ds.assign_coords(spatial=('nMesh_node', spatial)).swap_dims({'nMesh_node':'spatial'}).unstack('spatial')

Then test the time it takes to write to disk for both approaches:

It seems as though the unstack method is slightly faster, does this makes sense?

I can try both methods with an actual data file and see if this holds up.

[rsignell-usgs]

The coarsen method seems easier to understand (and remember!). Even though it wasn't faster, it wasn't much slower either. Did it use less memory?

[keewis]

For reference, I'd probably shorten the unstack example to:

ds.assign_coords(nMesh_node=spatial).unstack("nMesh_node")

(when initially proposing it I wanted to avoid overwriting any existing coordinate, but that turned out not to be necessary)

It seems I didn't pay enough attention to the recent PRs and missed #7192, where the coarsen method was added to the documentation. This still feels a bit like magic, though, because you only specify the size of one of the dimensions. In that regard, something like Variable.unstack for Dataset / DataArray would be both fully explicit and efficient (it directly translates to a reshape):

var = xr.Variable("z", np.linspace(0, 1, 12))
var.unstack(z={"x": 3, "y": 4})

we'd only need to figure out what to do with indexes along z (probably drop them)

unstack will create at least one copy IIRC

Right, so the copy is here:

xarray/xarray/core/variable.py

Lines 1816 to 1826 in df909b4

	data = np.full_like(
	self.data,
	fill_value=fill_value,
	shape=new_shape,
	dtype=dtype,
	)
	# Indexer is a list of lists of locations. Each list is the locations
	# on the new dimension. This is robust to the data being sparse; in that
	# case the destinations will be NaN / zero.
	data[(..., *indexer)] = reordered

I wonder if we can avoid that in case np.prod(new_shape) == self.data.size (i.e. the unstack does not change the number of elements)? In that case we would be able to just reshape the variable.

[dcherian]

something like Variable.unstack for Dataset / DataArray would be both fully explicit and efficient (it directly translates to a reshape):

stack/unstack are pretty general and quite tied in to the coordinate labels. We could consider adding a reshape method if we can figure out the syntax.

# stack dims a,b into c
# unstack d to e,f
ds.reshape({("a", "b"): "c", "d": {"e": 10, "f":20}})

The coarsen method seems easier to understand (and remember!).

Yes that's why I like it.

I wonder if we can avoid that in case np.prod(new_shape) == self.data.size (i.e. the unstack does not change the number of elements)? In that case we would be able to just reshape the variable.

The order of elements also needs to match for this to work.

[andrewdnolan]

Did it use less memory?

I'm not totally sure the best way to profile this, but using this code snippet:

from distributed.diagnostics.memory_sampler import MemorySampler

ms = MemorySampler()

for method in ["coarsen", "unstack"]:
    with ms.sample(method, interval=0.01):
        eval(method).to_zarr(f"{method}.zarr", mode = 'w')

ms.plot(align=True)

So they still seem pretty comparable.

This still feels a bit like magic, though, because you only specify the size of one of the dimensions.

As a relatively new user, I agree with this. The coarsen method feels pretty un-intuitive, especially since the size of only one dimension is specified. I hadn't realized there was a "Reshaping via coarsen" section in the documentation. The code I was using to do my reshaping predated this section of the documentation, so I must have missed it.

Given that coarsen seems to be recommended way, I'll use this going forward. The performance difference seem to be minor, and not consequential for anything I'll be doing.

Thanks again for all the help!