Polar plotting is really slow

[wwieder]

Hello,

I'm reaching out to see if you have suggestions for accelerating the creation of polar plots using uxarray on the CESM ne30 grid.

The larger script I'm calling is here,
which specifically calls this plotting function and the ux.to_raster function.

This script is slow when using regular lat-lon grids and xarray, but it's even slower when using unstructured grids. Any suggestions on how I could try gaining efficiency here?

[erogluorhan]

Hi Will,

Thanks for your feedback! We'll be looking into this and get back to you soon.

[erogluorhan]

Hi @wwieder , I will have time to look into it now, but we have made changes in the last couple weeks to relevant functions that improved the performance significantly. For me to make sure first whether those improvements might have gotten things improved on your end, can you give me:

• some numbers about the performance you were seeing
• in what environment you were running
• pointer to CESM ne30 data you used

[wwieder]

Great to hear you've made progress on this, @erogluorhan

• Sample data I'm trying to plot are here
  /glade/derecho/scratch/wwieder/ADF_unstruct/b.e30_alpha08a.B1850C_LTso.ne30_t232_wgx3.283i/climo

• I was plotting this with ADF /glade/u/home/wwieder/python/adf

  • specifically the call to polar plotting (https://github.com/NCAR/ADF/blob/clm-diags/scripts/plotting/polar_map.py)
  • using the cupid-analysis environment (https://github.com/NCAR/CUPiD/blob/main/environments/cupid-analysis.yml)

• I don't have timing numbers other than this is much slower with uxarray that it is for regular grids using xarray / matplotlib.

  • If you're interested sample data for this are here /glade/derecho/scratch/wwieder/ADF_unstruct/ADF_unstruct/ctsm5.4.CMIP7_ciso_ctsm5.3.075_f09_124_HIST/climo

[erogluorhan]

Thanks @wwieder , let me look into it and get back to you with what I have.

[erogluorhan]

Hi @wwieder , I took a stab on this without much luck, and I would really appreciate it if we could minimize the scope of the problem because:

• Currently, understanding how polar_map.py works (e.g. how to setup ADF configs, what the adfobj should be, etc.) and then running and debugging it to pick an iteration or two from all those loops & if-checks to get an xarray vs uxarray comparison would be well above a minimal recreatable example to focus on a performance comparison.
• I attempted to open unstructured grid data files from /glade/derecho/scratch/wwieder/ADF_unstruct/b.e30_alpha08a.B1850C_LTso.ne30_t232_wgx3.283i/climo but they don't seem to have grid definition in them. I checked its attributes as well but wasn't sure if any of those file locations pointed in the attributes is the grid file either. Can you point me to the grid file for them?
• I assume using the same variable from the unstructured grid and lat-lon grid datasets from those directories you provided would give me almost an apples to apples comparison between xarray and uxarray (because lat-lon is 192x288 and #cells is 48600, pretty close). Please let me know if this is valid.

Once I can open unstructured grid datasets, I think I'd be able to start some minimal uxarray vs xarray comparisons without focusing too much on ADF internals.

[wwieder]

Yeah, this is kind of a complicated example. Sorry I don't have the more basic plotting script it was developed from any more.

The mesh file is here /glade/campaign/cesm/cesmdata/inputdata/share/meshes/ne30pg3_ESMFmesh_cdf5_c20211018.nc. Yes, you can pick any variable on those files. Some thing like GPP or ELAI should be share on all the files.

This that what you needed?

[erogluorhan]

Thanks @wwieder , yes this was what exactly I needed.

I put together a simple notebook (issue1449-polar-plot.ipynb) to have a comparison with minimal code.

Apologies for how long the rest of this message will be, but I will first share my observations on it and then discuss what could be done with UXarray executions.

Observations

For similar-size GPP data on lat-lon and USG, xarray's contourf() plot (I used that since it is what was being used in the ADF script) takes only a few of ten miliseconds while uxarray's to_raster() takes around 1.5 seconds on the same environment (processing power, conda env, etc.). So, yes, this is significant difference especially when you want to do time-series plotting for many cases etc. Thoughts:

• For simplicity, my simple code does not treat cyclic points etc. in the xarray case while our uxarray plotting natively does that; hence, the xarray case might be looking a bit better than how it should actually be, but probably not much of a difference it would introduce.
• I don't have much information about how xarray's contourf() works internally, but it might be something extremely optimized for lat-lon grids.
• Our to_raster() needs to rasterize the grid first (as its name suggests), employing a nearest-neighbor resampling approach, using features from a third party package, cartopy.GeoAxes. For every pixel in the GeoAxes, it finds the face of the unstructured grid that contains the pixel's geographic coordinate and colors that pixel with the face's data value. While we've done a number of optimizations so far on this, there can be things that could further optimize the performance.
• However, these plotting cases are really different than each other in the way of computation.

What you could do additionally with UXarray

Actually, to_raster() has a number of options that you should definitely try for better performance. Checking out to_raster()'s documentation along with looking into the examples would be helpful for you to do that, but briefly, to_raster() has optional arguments:

• pixel_ratio, if you didn't specify, it defaults to 1.0, but you can try lower values, say 0.5 for better performance and see if you like the end result (quality of the plot will be diminished). I got around 600 ms of execution time with 0.5 for instance.
• Especially when you are using the same grid geometry over plotting many time-series etc. in a for-loop or so, you should consider using the pair of arguments return_pixel_mapping and pixel_mapping. Namely, if the very first call, you will run to_raster() with the former and assigned the returned pixel_mapping output to an additional variable of yours, and then use that one through the latter through the rest of your plots.
  • This could significantly improve your execution times as the pixel mapping calculation is the main bottleneck in this plotting.

Please let me know if these help, and always reach out if you need anything else.