rendering multiple variables from an xarray dataset?

[Thomas-Moore-Creative]

Thanks for all this work and sharing it openly - it's really appreciated!

I have a very basic question which I've been unable to answer by reading the docs and trying things out.

Is it possible to render more than one variable in an xarray dataset? For example, something like this which would ideally animate 12 monthly frames each with two line plots on them, one for each variable reduced along a single latitude:

import xarray as XR
from streamjoy import stream

# load xarray object
# a dataset with two 3D variables, var1 & var 2 along dimensions month, latitude, & longitude
ds = xr.open_dataset(my_file.nc)
ds_reduced = ds.sel(latitude=0,method='nearest') # reduce the dataset to select 2D data along the equator

stream(ds[['var1','var2']], uri="test.gif",dim='month')

[ahuang11]

Thanks for your interest in StreamJoy!

To answer your question, I'm not exactly sure I follow. ds_reduced still has longitude correct? What should x be and what should y be?

If lat/lon is reduced to 0-dim, I think you could create a pandas dataframe and melt, and finally groupby.

import xarray as xr
from streamjoy import stream


ds = xr.tutorial.open_dataset("air_temperature")
ds["airx2"] = ds.air * 2
ds = ds.mean(["lon", "lat"])


df = ds[["air", "airx2"]].to_dataframe().reset_index().melt(["time"])
stream(df, groupby="variable").write()

[ahuang11]

Oh you might mean something like this?

[ahuang11]

A bit convoluted, but it gets the job done! The benefit of this is you can also customize how you want the plot.

import xarray as xr
import hvplot.xarray
import matplotlib.pyplot as plt
from streamjoy import stream, wrap_matplotlib

@wrap_matplotlib()
def renderer(da):
    fig = plt.figure()
    ax = fig.add_subplot(111)
    da.plot.line(x="lon", ax=ax, ylim=(250, 600))
    return fig


ds = xr.tutorial.open_dataset("air_temperature").isel(time=slice(0, 100))
ds["airx2"] = ds.air * 2
ds = ds.mean(["lat"])

stream(ds.to_array(), dim="time", renderer=renderer).write()

Underneath, the renderer essentially just does

ds.plot()

So what I did was first explore how to get that plot through basic methods

ds.to_array().isel(time=0).plot.line(x="lon", ylim=(250, 600))

And then migrate into a renderer function.

However, in future versions, I'll try to find a way to make this easy

#40

[Thomas-Moore-Creative]

Andrew! Much appreciated - seeing this done helps me better understand how to build renderer functions. And I also appreciate if this can be built into a method. I think people will appreciate it.

[Thomas-Moore-Creative]

@ahuang11 - I'm trying to follow your renderer example code but this time apply it to an animation with two subplots in each frame, each plotting a different data array in a single dataset.

something like:

ds_combined = xr.Dataset({'u': u, 'v': v})
@wrap_matplotlib()
def renderer(ds):
    fig = plt.figure(figsize=(16, 9))  # Set the figure size to be wider
    
    fig.suptitle('Currents N/S & E/W', fontsize=16)  # Add main title
    
    ax1 = fig.add_subplot(121)  # First subplot
    ax2 = fig.add_subplot(122)  # Second subplot
    
    # Create a colormap and set the 'bad' color for NaN values
    cmap = plt.cm.RdBu.copy()
    cmap.set_bad(color='grey')  # Set NaN color to black
    
    # Plot the data on the first subplot
    ds['v'].plot(ax=ax1, vmin=-0.7, vmax=0.7, cmap=cmap, cbar_kwargs={'extend': 'neither'})
    
    # Plot the data on the second subplot
    ds['u'].plot(ax=ax2, vmin=-0.7, vmax=0.7, cmap=cmap, cbar_kwargs={'extend': 'neither'})
    ax2.set_title('')  # Remove the title from the second plot
    ax2.set_ylabel('')  # Remove the x axis label from the second plot
    
    return fig

    
stream(ds_combined.to_array(),renderer=renderer,dim='Time',max_frames=50,
                client=client,
                uri="currents.gif",
      threads_per_worker=1)

but I get an error like:

ValueError: Can only handle 3D arrays; 4D array found

any thoughts on how to get a multi-variable multi-subplot to animate? Thanks.

[ahuang11]

Hmm I think it's an issue I need to fix / release.

import matplotlib.pyplot as plt
import numpy as np
import xarray as xr
from streamjoy import stream, wrap_matplotlib


@wrap_matplotlib()
def renderer(ds):
    fig = plt.figure(figsize=(16, 9))  # Set the figure size to be wider

    fig.suptitle("Currents N/S & E/W", fontsize=16)  # Add main title

    ax1 = fig.add_subplot(121)  # First subplot
    ax2 = fig.add_subplot(122)  # Second subplot

    # Create a colormap and set the 'bad' color for NaN values
    cmap = plt.cm.RdBu.copy()
    cmap.set_bad(color="grey")  # Set NaN color to black

    # Plot the data on the first subplot
    ds.sel(variable="air").plot(ax=ax1, cmap=cmap, cbar_kwargs={"extend": "neither"})

    # Plot the data on the second subplot
    ds.sel(variable="airx2").plot(ax=ax2, cmap=cmap, cbar_kwargs={"extend": "neither"})
    ax2.set_title("")  # Remove the title from the second plot
    ax2.set_ylabel("")  # Remove the x axis label from the second plot

    return fig


ds = xr.tutorial.open_dataset("air_temperature")
ds["airx2"] = ds["air"] * 2
stream(
    ds.to_array(),
    renderer=renderer,
    dim="time",
    uri="currents.gif",
    threads_per_worker=1,
)

[ahuang11]

I just released v0.0.10 which should fix this I think.

[Thomas-Moore-Creative]

Thanks, will test ASAP.

[Thomas-Moore-Creative]

This works! I no longer have to generate subplots for variables individually and paste them together with ffmpeg. Thanks @ahuang11