Currently, to_raster() is extremely slow for high resolution grids, even when only looking at a small subset of the grid.
Background details / exploration
Using a CESM-HR grid with 8.64M faces (the same one as in @erogluorhan's notebook from the recent usclivar meeting, filename "02-usclivar-so-sst-trends.ipynb"), calling to_raster() by itself simply crashes the kernel after a little while, probably due to running out of memory.
To understand the problem better, I tried slicing to a smaller section of the grid. Running arr.isel(n_faces=slice(None, 1000)).to_raster() takes roughly 5 seconds. Running arr.isel(n_faces=slice(None, 10000)).to_raster() takes roughly 37 seconds. The scaling itself isn't terrible (10x more faces but less than 10x more time), but the base amount of time this takes is too large to be practical. Even if the kernel didn't crash, at this speed it would take roughly 1 hour per million faces, i.e. 8 hours to build the raster for this grid. A factor of 10 speedup would start to make this maybe feasible at all, while a factor of 100 would bring it in line with its holoviews counterpart (where rendering an image of data across this full grid took only roughly 3 mins).
Comparing those timings with a very low-resolution grid is a bit surprising:
import cartopy.crs as ccrs
import matplotlib.pyplot as plt
import uxarray as ux
fig, ax = plt.subplots(subplot_kw={"projection": ccrs.Robinson()})
ax.set_global()
example_arr = ux.tutorial.open_dataset("outCSne30-vortex")['psi']
raster = example_arr.to_raster(ax=ax)
ax.imshow(raster, origin="lower", extent=ax.get_xlim() + ax.get_ylim())This grid has 5400 faces, but this cell takes only roughly 0.5 seconds, far faster than the 1000-faces slice of the high resolution grid. Something about the underlying methods cares very much about the resolution, rather than the total number of faces.
Underlying methods / thoughts on how to proceed
I explored the underlying methods, inlining their logic until I reached a numba method. I saw that the vast majority of the time of to_raster() is coming from Grid.get_faces_containing_point(); the vast majority of the time there comes from _point_in_face_query(); and the vast majority of the time there comes from _batch_point_in_face(). That method is a numba method, which is difficult to profile internally.
But, I can see that it is internally calling _get_faces_containing_point() for each point in the raster (using default parameters in the example above, there are 159k points in the raster, coming from shape=(369, 496)). Each call to that function creates a numpy array:
hit_buf = np.empty(candidate_indices.shape[0], dtype=INT_DTYPE)
and the length of candidate_indices will be much larger for grids with higher resolution. This is like #1648 but worse; the arrays might not all be tiny. hit_buf gets returned, then consumed immediately inside _batch_point_in_face(), assigning results[i, j] = hits[j] for all hits, then forgotten. A simple rewrite should be able to remove the need for ever creating it in the first place. (For example, it might be sufficient to just pass results into _get_faces_containing_point and set its values there directly.)
The _get_faces_containing_point() method also calls _get_cartesian_face_edge_nodes() for each candidate; this repeats work when candidates appear multiple times across different points, which is much more common for high resolution grids. In the slice(None, 10000) example above, of the 7414 points in the raster with at least 1 candidate, the median number of candidates is 1333, so this represents a lot of repeated work.
The _get_cartesian_face_edge_nodes() method itself constructs a tiny array for each face, of shape (n_edges, 2, 3). Its docstring claims that it "Computes the Cartesian Coordinates of the edge nodes that make up a given face," but really there is no computation going on here at all, just indexing using the face_node_connectivity array to get a result with a convenient shape. Here, the result is passed directly to _face_contains_point as face_edges, where it is consumed to produce a boolean ("does the face contain the point, or not?") and then forgotten. Thus, it is definitely possible to avoid constructing these tiny arrays, even if it means needing to inline the relevant indexing logic directly inside _face_contains_point.
Finally, _face_contains_point() uses numpy methods like np.allclose and np.linalg.norm (which only construct new arrays internally if the inputs aren't numpy arrays yet, but this is easy to fix, see e.g. _numba_norm3), along with point_within_gca() and _small_angle_of_2_vectors(). The latter has already been optimized to avoid creating tiny numpy arrays (see #1674). The former (i.e., point_within_gca()) should be easy to optimize; it only calls _angle_of_2_vectors() and simple numpy helper functions for handling 3-vectors, and _angle_of_2_vectors() similarly only calls simple numpy helper functions.
A full solution to this issue should optimize all the methods mentioned here, but probably won't need to touch any other parts of the code. Based on speedups seen in other sub-issues of #1648, and considering all the different inefficiencies mentioned above, I'm hopeful that this might lead to a huge speedup for to_raster().
Self-assigning because I'm really interested to solve this one, after doing such a deep dive on it (Claude helped me hone in on _batch_point_in_face() as the relevant method to consider, but everything after that point I did without AI tools)!
Currently, to_raster() is extremely slow for high resolution grids, even when only looking at a small subset of the grid.
Background details / exploration
Using a CESM-HR grid with 8.64M faces (the same one as in @erogluorhan's notebook from the recent usclivar meeting, filename "02-usclivar-so-sst-trends.ipynb"), calling to_raster() by itself simply crashes the kernel after a little while, probably due to running out of memory.
To understand the problem better, I tried slicing to a smaller section of the grid. Running
arr.isel(n_faces=slice(None, 1000)).to_raster()takes roughly 5 seconds. Runningarr.isel(n_faces=slice(None, 10000)).to_raster()takes roughly 37 seconds. The scaling itself isn't terrible (10x more faces but less than 10x more time), but the base amount of time this takes is too large to be practical. Even if the kernel didn't crash, at this speed it would take roughly 1 hour per million faces, i.e. 8 hours to build the raster for this grid. A factor of 10 speedup would start to make this maybe feasible at all, while a factor of 100 would bring it in line with its holoviews counterpart (where rendering an image of data across this full grid took only roughly 3 mins).Comparing those timings with a very low-resolution grid is a bit surprising:
This grid has 5400 faces, but this cell takes only roughly 0.5 seconds, far faster than the 1000-faces slice of the high resolution grid. Something about the underlying methods cares very much about the resolution, rather than the total number of faces.
Underlying methods / thoughts on how to proceed
I explored the underlying methods, inlining their logic until I reached a numba method. I saw that the vast majority of the time of
to_raster()is coming fromGrid.get_faces_containing_point(); the vast majority of the time there comes from_point_in_face_query(); and the vast majority of the time there comes from_batch_point_in_face(). That method is a numba method, which is difficult to profile internally.But, I can see that it is internally calling
_get_faces_containing_point()for each point in the raster (using default parameters in the example above, there are 159k points in the raster, coming from shape=(369, 496)). Each call to that function creates a numpy array:and the length of candidate_indices will be much larger for grids with higher resolution. This is like #1648 but worse; the arrays might not all be tiny.
hit_bufgets returned, then consumed immediately inside_batch_point_in_face(), assigningresults[i, j] = hits[j]for all hits, then forgotten. A simple rewrite should be able to remove the need for ever creating it in the first place. (For example, it might be sufficient to just passresultsinto_get_faces_containing_pointand set its values there directly.)The
_get_faces_containing_point()method also calls_get_cartesian_face_edge_nodes()for each candidate; this repeats work when candidates appear multiple times across different points, which is much more common for high resolution grids. In the slice(None, 10000) example above, of the 7414 points in the raster with at least 1 candidate, the median number of candidates is 1333, so this represents a lot of repeated work.The
_get_cartesian_face_edge_nodes()method itself constructs a tiny array for each face, of shape (n_edges, 2, 3). Its docstring claims that it "Computes the Cartesian Coordinates of the edge nodes that make up a given face," but really there is no computation going on here at all, just indexing using the face_node_connectivity array to get a result with a convenient shape. Here, the result is passed directly to_face_contains_pointas face_edges, where it is consumed to produce a boolean ("does the face contain the point, or not?") and then forgotten. Thus, it is definitely possible to avoid constructing these tiny arrays, even if it means needing to inline the relevant indexing logic directly inside_face_contains_point.Finally,
_face_contains_point()uses numpy methods like np.allclose and np.linalg.norm (which only construct new arrays internally if the inputs aren't numpy arrays yet, but this is easy to fix, see e.g._numba_norm3), along withpoint_within_gca()and_small_angle_of_2_vectors(). The latter has already been optimized to avoid creating tiny numpy arrays (see #1674). The former (i.e.,point_within_gca()) should be easy to optimize; it only calls_angle_of_2_vectors()and simple numpy helper functions for handling 3-vectors, and_angle_of_2_vectors()similarly only calls simple numpy helper functions.A full solution to this issue should optimize all the methods mentioned here, but probably won't need to touch any other parts of the code. Based on speedups seen in other sub-issues of #1648, and considering all the different inefficiencies mentioned above, I'm hopeful that this might lead to a huge speedup for to_raster().
Self-assigning because I'm really interested to solve this one, after doing such a deep dive on it (Claude helped me hone in on
_batch_point_in_face()as the relevant method to consider, but everything after that point I did without AI tools)!