Out of memory error during initialization of hash table - CREG12 curvilinear grid

[noemieplanat]

Parcels version

main
python 3.11

Description

Parcels fails to evaluate U (any) field due to a memory error. Initial grid of the netcdf file is 75 (depth), 1580 (x) and 1801 (y). Initial xarray (ds_fields) has the same dimensions. Converted fset (ds_fset) has 1580 along x and x_center dimensions (idem for y) plus has a x variable that has (y,x) dimensions.

Fails on fset.U.eval(...)

Code sample

The data is stored in @VeckoTheGecko 's home folder on Lorenz. We can probably trim this example down even further using the instructions on sharing datasets (which would indicate that this is a dimension/metadata/data transformation problem and not a coordinates issue) but he is yet to pursue that.

from parcels import FieldSet, ParticleSet, Variable, convert, ParticleFile, StatusCode # type: ignore
import xarray as xr
from datetime import timedelta as delta
import numpy as np
import datetime
import warnings
import cftime as cf


listeU = ['data/CREG12.L75-REF12_y2001m02.1d_gridU.nc']
listeV = ['data/CREG12.L75-REF12_y2001m02.1d_gridV.nc']
listeW = ['data/CREG12.L75-REF12_y2001m02.1d_gridW.nc']
mesh_file_z = 'data/CREG12.L75-REF09_domain_mask_lat_lon_updated_z.nc'


ds_fields = xr.open_mfdataset(
    listeU +listeV+listeW,
    data_vars={"vozocrtx","vomecrty", "vovecrtz"},
    coords="minimal",
    compat="override",
)
ds_fields = ds_fields.drop({"depthu_bounds", \
                           "time_centered_bounds",\
                          "time_counter_bounds",\
                           "e3u","e3v", "sozotaux",\
                           "utau_atmoce", "utau_iceoce",\
                           "uwspd10", "ahtu_3d", "depthv_bounds",\
                           "sometauy", "vtau_atmoce",\
                           "vtau_iceoce","vwspd10",\
                           "ahtv_3d", "depthw_bounds",\
                           "votkeavt","e3w"}
                          )
ds_coords = xr.open_dataset(mesh_file_z, decode_times=False)
ds_fset = convert.nemo_to_sgrid(fields=dict(U=ds_fields["vozocrtx"], V=ds_fields["vomecrty"], W = ds_fields["vovecrtz"]), coords=ds_coords)

ds_fset = ds_fset.drop(['V','W']).isel(time=[0])
fset = FieldSet.from_sgrid_conventions(ds_fset, ds_coords)

fset.U.eval(np.array([0.]), np.array([50.]),np.array([89.]), np.array([120.]))

error

Traceback (most recent call last):
  File "/storage/home/hodgs004/parcels4-debug-naomi/min-example.py", line 44, in <module>
    fset.U.eval(np.array([0.]), np.array([50.]),np.array([89.]), np.array([120.]))
    ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/storage/home/hodgs004/parcels4-debug-naomi/.pixi/envs/default/lib/python3.11/site-packages/parcels/_core/field.py", line 220, in eval
    particle_positions, grid_positions = _get_positions(self, time, z, y, x, particles, _ei)
                                         ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/storage/home/hodgs004/parcels4-debug-naomi/.pixi/envs/default/lib/python3.11/site-packages/parcels/_core/field.py", line 469, in _get_positions
    grid_positions.update(field.grid.search(z, y, x, ei=_ei))
                          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/storage/home/hodgs004/parcels4-debug-naomi/.pixi/envs/default/lib/python3.11/site-packages/parcels/_core/xgrid.py", line 313, in search
    yi, eta, xi, xsi = _search_indices_curvilinear_2d(self, y, x, yi, xi)
                       ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/storage/home/hodgs004/parcels4-debug-naomi/.pixi/envs/default/lib/python3.11/site-packages/parcels/_core/index_search.py", line 192, in _search_indices_curvilinear_2d
    yi_q, xi_q, coords_q = grid.get_spatial_hash().query(y_check, x_check)
                           ^^^^^^^^^^^^^^^^^^^^^^^
  File "/storage/home/hodgs004/parcels4-debug-naomi/.pixi/envs/default/lib/python3.11/site-packages/parcels/_core/basegrid.py", line 210, in get_spatial_hash
    self._spatialhash = SpatialHash(self)
                        ^^^^^^^^^^^^^^^^^
  File "/storage/home/hodgs004/parcels4-debug-naomi/.pixi/envs/default/lib/python3.11/site-packages/parcels/_core/spatialhash.py", line 177, in __init__
    self._hash_table = self._initialize_hash_table()
                       ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/storage/home/hodgs004/parcels4-debug-naomi/.pixi/envs/default/lib/python3.11/site-packages/parcels/_core/spatialhash.py", line 224, in _initialize_hash_table
    morton_codes = np.zeros(total_hash_entries, dtype=np.uint32)
                   ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
numpy._core._exceptions._ArrayMemoryError: Unable to allocate 791. GiB for an array with shape (212287160821,) and data type uint32

[VeckoTheGecko]

@fluidnumerics-joe just a quick pulsecheck, any idea why this is happening? I can provide ncdump and see if we can trim this example down further next week

[fluidnumerics-joe]

What type of grid is this? Is it a global tri-pole grid? Typically, if there are locations in the grid where the cells become small (near poles) this can cause a high number of hits in the hashing. In the grids we've worked on, we got this sorted out. It'd be good to have a visualisation of the grid here. Of course, access to data would be even better

[noemieplanat]

I documented another issue with that same grid (ORCA, no discontinuity in y/x directions at the north pole) here : #1979 you can have a look at the grid there.

[fluidnumerics-joe]

It has less to do with discontinuities and more to do with the number of physical grid cells per hash cell

[fluidnumerics-joe]

What is the colormap representative of in those plots ?

[noemieplanat]

I am plotting the x and y indices. I have no idea what size is the hash cell, but no strong change in the number of grid cells at the north pole.

[fluidnumerics-joe]

Ok, but what are the grid cell sizes for your original grid ?Just looking at the shape (and not knowing what grid generation method you're using), it looks like theres some high curvature in the mesh over canada. To resolve that, I suspect there's intense crowding of grid cells there. Can you plot the cell area of each grid cell for your grid, or make a plot showing grid lines of your curvilinear grid ?

[fluidnumerics-joe]

Maybe I'm mistaken, but the call to create the fieldset doesn't seem right

fset = FieldSet.from_sgrid_conventions(ds_fset, ds_coords)

The API is

    def from_sgrid_conventions(
        cls, ds: xr.Dataset, mesh: Mesh | None = None
    ):  # TODO: Update mesh to be discovered from the dataset metadata
        """Create a FieldSet from a dataset using SGRID convention metadata.

        This is the primary ingestion method in Parcels for structured grid datasets.

        Assumes that U, V, (and optionally W) variables are named 'U', 'V', and 'W' in the dataset.

        Parameters
        ----------
        ds : xarray.Dataset
            xarray.Dataset with SGRID convention metadata.
        mesh : str
            String indicating the type of mesh coordinates used during
            velocity interpolation. Options are "spherical" or "flat"

@VeckoTheGecko - am I missing something here ?

Assuming the mesh type resolves to spherical, the one adjustment I could see on hash grid creation is to change the bounding box for spherical mesh types. Right now, the assumption we've been working with is that spherical mesh types are global grids - under the hood, the hash grid ends up fairly coarse with respect to your current grid (I think).

This can translate to a large number of physical cells per hash cell. Is there an easy way to get this dataset into the SurfDrive @VeckoTheGecko ? I'm happy to work on a patch for this in the hash table initialization to see if this resolves the issue

[VeckoTheGecko]

fset = FieldSet.from_sgrid_conventions(ds_fset, ds_coords)

yes, you're right. Normally this is flagged by our code, but there is a bug in xarray causing this guard to be skipped.

Anyway, if you do

-fset = FieldSet.from_sgrid_conventions(ds_fset, ds_coords)
+fset = FieldSet.from_sgrid_conventions(ds_fset, "spherical")

the error still reproduces.

---

Sounds great that you see a way forward here! The files are in the SURF at https://surfdrive.surf.nl/index.php/s/7xlfdOFaUGDEmpD

[fluidnumerics-joe]

Yeah, my suspicion is that we can improve the hashing by adjusting the bounding box for the polar cap. Further its likely there is intense crowding of the grid cells around some of the high curvature areas . I'll take a look at the files this afternoon.

@noemieplanat , have you been able to make plots of the grid lines and the cell areas of the mesh you created?