proplot-dev / proplot

🎨 A succinct matplotlib wrapper for making beautiful, publication-quality graphics

Home Page:https://proplot.readthedocs.io

Geek Repo:Geek Repo

Github PK Tool:Github PK Tool

[Feature request]: xarray accessor

savannahferretti opened this issue · comments

I am a PhD student in climate science, and love proplot! But I would love to use it from xarray using an accessor. That might look like this:

So you can toy with the data yourself, I used the Xarray tutorial air temperature dataset for an "example" of sorts. The environment needs both pooch (```conda install pooch``) and Xarray installed.

import xarray as xr

airtemps = xr.tutorial.open_dataset("air_temperature")
airtemp1D = airtemps.mean(dim=['lat','lon']).air # example of 1D (time)
airtemp2D = airtemps.isel(time=0).air # example of 2D (lat,lon)
airtemps3D = airtemps.air # example of 3D (time,lat,lon)

Make plotting simple graphs in proplot an accessor to Xarray:

import proplot as pplt

@xr.register_dataarray_accessor("proplot")
class ProPlotAccessor:
    def __init__(self, xarray_obj):
        self = xarray_obj

    def plot1D(self):
        fig,ax = pplt.subplots(nrows=1,ncols=1)
        ax.plot(self,color='k')
        
    def plot2D(self):
        fig,ax = pplt.subplots(nrows=1,ncols=1)
        ax.contourf(self,cmap='viridis')
        
    def plot(self):
        if len(self.dims) == 1:  # if data is 1D
            return self.plot1D()
        elif len(self.dims) == 2:  # if data is 2D
            return self.plot2D()
        else:
            raise ValueError("Data is not 1D or 2D")

Test the accessor on the Xarray tutorial data with different dimensions:

airtemps1D.proplot.plot()

image

airtemps2D.proplot.plot()

image

I think this would be impactful because it would make it easier for those in the climate science space who use Xarray to get their feet wet with proplot, moving away from matplotlib. I'd be happy to collaborate on this effort if wanted/needed!

@TomNicholas

I think this would be super neat. It would effectively copy how xarray already allows you to plot matplotlib plots, i.e.

da.plot.line()  # matplotlib plot
da.proplot.line()  # proplot plot

Xarray's documentation on accessors is here.

You might also want to follow this issue in case we ever implement it 😅