poor performance of subsetting

[uturuncoglu]

I am trying to subset a global dataset belongs to unstructured ADCIRC ocean model output (under /glade/work/turuncu/COASTAL/s104). I have two version of code to subset the dataset to a specified region.

• directly uses Xarray

combined_msl_file_path = "GSTOFSv5.6.5_msl_2017-2019_avg_20241619.nc"
combined_msl_file = xr.open_dataset(combined_msl_file_path)

min_lon, min_lat = -74.10, 40.575
max_lon, max_lat = -74.025, 40.65

cmsl_lon = combined_msl_file['x'].values
cmsl_lat = combined_msl_file['y'].values
cmsl_idx = np.where(
    (cmsl_lon >= min_lon) & (cmsl_lon <= max_lon) &
    (cmsl_lat >= min_lat) & (cmsl_lat <= max_lat)
)[0]
subset_cmsl = combined_msl_file['mean_sea_level'].isel(node=cmsl_idx)

• uxarray

import uxarray as ux

mesh_file = '/glade/work/turuncu/COASTAL/s104/mesh/esmf_mesh.nc'
uxds = ux.open_dataset(
    mesh_file, 
    combined_msl_file_path, 
    chunks='auto',
    drop_variables=['nvel', 'nvell', 'adcirc_mesh', 'max_nvell', 'ibtype', 'nbvv', 'depth']
)

lon_bounds = [ -74.100, -74.025 ]
lat_bounds = [  40.575,  40.650 ]

subset = uxds['mean_sea_level'].subset.bounding_box(
    lon_bounds,
    lat_bounds,
)

The UXarray way is much slower than the first one and I wonder why? If I measure time with %%time under JupyterHub, the first one is taking CPU times: user 27.1 ms, sys: 32 ms, total: 59.1 ms / Wall time: 92.5 ms and the UXarray way is taking CPU times: user 8min 32s, sys: 8.84 s, total: 8min 41s / Wall time: 3min 5s. So, it seems it is the difference is huge like 27 ms vs. 8min 32s. Am I missing something in here? Any suggestion will help at this point.

[philipc2]

Hi @uturuncoglu

Which version of UXarray are you on?

[uturuncoglu]

I have 2025.03.0-pyhd8ed1ab_0

[philipc2]

Can you try updating to the latest version? We had some large optimizations in 2025.04.0

For reference, I just ran the your code on Derecho:

%%time
subset = uxgrid.subset.bounding_box(
    lon_bounds,
    lat_bounds,
)

Initial Computation

CPU times: user 6min 21s, sys: 1.98 s, total: 6min 23s
Wall time: 4.6 s

Subsequent Computations

After the bounds are computed internally (which takes up most of the execution time), the results are comparable.

CPU times: user 87.8 ms, sys: 12.1 ms, total: 99.8 ms
Wall time: 98.5 ms

[philipc2]

Here's what my result looks like (only the grid since I couldn't access the data)

[uturuncoglu]

Okay. There is uxarray-2025.05.0 on conda. Let me try.

[uturuncoglu]

Thanks for your help. I have two scripts under the same folder: test.py is the UXarray one and test2.py is the Xarray one.

With the new version (uxarray-2025.05.0) I got following warning in the first run

/glade/work/turuncu/COASTAL/s104/myenv/lib/python3.12/site-packages/uxarray/grid/grid.py:1432: RuntimeWarning: Necessary functions for computing the bounds of each face are not yet compiled with Numba. This initial execution will be significantly longer.
  warn(` .

but I am not seeing any difference between initial and subsequent one.

run1

(myenv) turuncu@derecho6/glade/work/turuncu/COASTAL/s104 $ time python test.py

real	1m14.327s
user	7m11.615s
sys	0m7.224s

run2

(myenv) turuncu@derecho6/glade/work/turuncu/COASTAL/s104 $ time python test.py

real	1m17.632s
user	7m32.598s
sys	0m6.957s

test2.py still way more faster for me.

(myenv) turuncu@derecho6/glade/work/turuncu/COASTAL/s104 $ time python test2.py 

real	0m2.654s
user	0m0.982s
sys	0m1.373s

[uturuncoglu]

The data file is in /glade/work/turuncu/COASTAL/s104/GSTOFSv5.6.5_msl_2017-2019_avg_20241619.nc. This has only one time step and used to correct sea level of actual data file (/glade/work/turuncu/COASTAL/s104/stofs_2d_glo.t06z.fields.cwl.nc).

BTW, I have some code to create ESMF mesh file from UXarray grid and store it use create ESMF mesh object under ESMpy to do the interpolation. As I see there is no any capability like that at this point but it would be nice to have it.

[philipc2]

BTW, I have some code to create ESMF mesh file from UXarray grid and store it use create ESMF mesh object under ESMpy to do the interpolation. As I see there is no any capability like that at this point but it would be nice to have it.

Might be related to #172

This does sound like something that would be good to implement. I'm currently refactoring and optimizing our remapping functionality in #1244, so we can consider adding something like that in the near future.

Also pinging @rajeeja since he has done some work with ESMPy

[rajeeja]

BTW, I have some code to create ESMF mesh file from UXarray grid and store it use create ESMF mesh object under ESMpy to do the interpolation. As I see there is no any capability like that at this point but it would be nice to have it.

Might be related to #172

This does sound like something that would be good to implement. I'm currently refactoring and optimizing our remapping functionality in #1244, so we can consider adding something like that in the near future.

Also pinging @rajeeja since he has done some work with ESMPy

Yes, I'd be happy to implement this

[rajeeja]

BTW, I have some code to create ESMF mesh file from UXarray grid and store it use create ESMF mesh object under ESMpy to do the interpolation. As I see there is no any capability like that at this point but it would be nice to have it.

can you share this code? I can access glade

[uturuncoglu]

@rajeeja You could find the actual function that creates the mesh file is in the comment (https://gist.github.com/uturuncoglu/0df264a060668e370f4604b57e7c20d5?permalink_comment_id=5580195#gistcomment-5580195). The same code also exits in the notebook but it was failing for other case. So, I fixed but put it into comment. The usage of the mesh with ESMPy can be seen in the actual notebook (it basically interpolates global unstructured data to regional lat-lon after subsetting the global data with UXarray). I need to create mesh file (for subsetted region) since ESMPy uses that mesh file to create ESMF mesh object (I think it can be done with the API but I go with simple approach and read the mesh file after creating with the function). BTW, the global ADCIRC dataset can be used with UXarray but that also requires ESMF mesh file and I created that one externally with some Fortran code. Let me know if you need any other information.

[uturuncoglu]

@philipc2 BTW, I tried to run in same notebook and all the time it prints /glade/work/turuncu/COASTAL/s104/myenv/lib/python3.12/site-packages/uxarray/grid/grid.py:1454: RuntimeWarning: Necessary functions for computing the bounds of each face are not yet compiled with Numba. This initial execution will be significantly longer. bounds_normalized = (bounds_deg + 180.0) % 360.0 - 180.0. So, I am not sure why but it seems caching not working in this case (I think). Or maybe this message is something else and I need to configure the python environment in a spacial way.

[philipc2]

The caching here is for Numba, which we use to accelerate our numerical computations.

Numba translate Python code to machine code, which can add a couple seconds from a cold install of UXarray. However, once you've run the function at least once, it uses the cached machine code.

This is different than how how we cache the bounds array. The array is computed once per Grid when it is required, either for subsetting or zonal averaging.

I'm not sure why it keeps printing, as it should only show up the first time you run the code after a fresh install.