2
$\begingroup$

I am working on the ORAS5 dataset which comes as an archive of netCDF files. Currently, I am using python and netcdf4 to process it.

My goal is to get surface level ocean currents across the globe (almost, 70°N to 70°S). The dataset itself has several variables. These are relevant for my problem:

  • depth 1D (75,) array describing increasing depths
  • lat 2D (1021, 1442) array describing latitude
  • lon 2D (1021, 1442) array describing longitude
  • vozocrte/vomecrtn 3D (75, 1021, 1442) array describing zonal or meridional velocity

Currently, I am creating a lat-lon grid with desired resolution and then loop over this grid to extract vozocrte or vomecrtn. In principle it looks like this:

import numpy as np
from netCDF4 import Dataset

def mask_map(vals: np.ndarray, range: tuple[int, int], res=0.25) -> dict[float, np.ndarray]:
    return {
        d: (vals >= d - res / 2) & (vals <= d + res / 2)
        for d in np.arange(range[0], range[1] + res, res)
    }

ds = Dataset(file, "r", format="NETCDF4")
depths = ds.variables["depth"][:]  # (75,) depth in m
lats = ds.variables["nav_lat"][:]  # (1021, 1442) lat in degrees
lons = ds.variables["nav_lon"][:]  # (1021, 1442) lon in degrees
vals = ds.variables["vozocrte"][:]  # (76, 1021, 1442) in m/s
vals = vals[0, depths < 5]
lon_2_mask = mask_map(vals=lons, range=lon_range)
lat_2_mask = mask_map(vals=lats, range=lat_range)

records = []
for lon, lons_mask in lon_2_mask.items():
    for lat, lats_mask in lat_2_mask.items():
        masked = vals[..., lats_mask & lons_mask]
        if not masked.mask.all():
            records.append({"component": masked.mean(), "lat": lat, "lon": lon})

It works but it is very slow.

I am unfamiliar with this type of dataset (mapping depth-lon-lat to z-y-x, then using z-y-x to get some values). Is there a more efficient way of doing this?

Edit

I realized I should explain the data better. The variables depth, nav_lat, nav_lon, vozocrte/vomecrtn mentioned above all contain masked numpy arrays of different shape.

  • vozocrte/vomecrtn describe a velocity in m/s as velocity(z, y, x).
  • depth describes depth in m as depth(z).
  • nav_lon describes longitude in degree as longitude(y, x).
  • nav_lat describes latitude in degree as latitude(y, x).

E.g. nav_lat can contain a masked array of [[87.021, 86.721, 86.101, ...], ...]. And nav_lon can contain a masked array of [[178.131, 177.809, 176.980, ...], ...]. depth a masked array of e.g. [0.023, 1.201, ...].

My naive approach above follows the logic of first getting a mask for z, y, x from the variables depth, nav_lat, nav_lon, then using them to get velocities. E.g. if I wanted 10°-11°N, 15°-16°E, less than 5m deep:

depth_mask = depths < 5
lon_mask = (lons >= 15) & (lons < 16)
lat_mask = (lats >= 10) & (lats < 11)
vals[depth_mask, ..., lon_mask & lat_mask]
$\endgroup$

2 Answers 2

3
$\begingroup$

First of all: NEVER, EVER, EVER (I repeat, EVER) loop in Python over an array to subset some values :)

xarray was created exactly to avoid having to write something like this by exploiting the native coordinates/dimensions indexing. For a brief introduction on how to subset an array

List item

have a look here

Regarding your specific issue, it is not clear what you are trying to achieve because you don't state what you want to "extract" in the description. Judging from the code, it seems you're trying to downsample the grid to a specific resolution (res). In pure xarray going from e.g. 0.25 to 0.5 degrees resolution would be as simple as doing

import xarray as xr

ds = xr.open_dataset("test.nc")

ds.sel(lon=ds.lon[::2], lat=ds.lat[::2])

There are many other ways of subsetting data to extract a single point or set of points using the sel operator, for example

ds.sel(lon=43, lat=0, method='nearest') # single point

There are many functions that are written to do operations on geospatial multi-dimensional arrays (interpolation, subsetting,...). As a general rule of thumb, if you're writing your own low-level function to access element one by one in Python...you're doing it wrong :D There's almost surely another function that does that in a vectorized way!

Please take some time to learn how to use xarray to work with Netcdf files and subset data, you'll save a lot of time in the future!

P.S. Even when using xarray, Python remains quite slow with large dataset. If you need to do operations on large files I'd suggest you to pre-process the file BEFORE loading it into Python using something like CDO

$\endgroup$
2
  • $\begingroup$ As mentioned above, I want to extract ocean current zonal and meridional vectors which are variables vozocrte and vomecrtn in this dataset. Your selection statement gives me a subset of lons/lats. This will not yield a regular grid. I understand this is not clear from seeing the code alone. That's why I mentioned the dataset and type of dataset itself. $\endgroup$ Commented Feb 5 at 14:13
  • 1
    $\begingroup$ So you basically want to subset in lat/lon, right? Why not use xr.where? This will return a regular grid in both cases, unless you specify drop=True. You just have to write your condition where you want to keep the data (it is still not clear to me what this condition is :)). But honestly, for your own sake, try to avoid loops as much as possible. Also, try to make a self-contained and self-explanatory MWE next time, it would be easier for people to respond. $\endgroup$
    – Droid
    Commented Feb 5 at 17:14
0
$\begingroup$

There is a very simple way to do this. The point here is that these arrays are all ordered consistently. The dataset needs only to be subset by depth if all lon,lat pairs should be extracted. Lat and lon values can still be corrected later on to create a specific resolution.

What I ended up doing is in essence this:

import numpy as np
import pandas as pd
from netCDF4 import Dataset


def _rm_mask(masked: np.ma.MaskedArray, fill=np.nan) -> np.ndarray:
    return np.where(masked.mask, fill, masked.data).astype(masked.dtype)


def _digitize(df: pd.DataFrame, var: str, interval: tuple[float, float], res: float):
    for target in np.arange(interval[0], interval[1] + res, res):
        mask = (df[var] >= target - res / 2) & (df[var] < target + res / 2)
        df.loc[mask, var] = target


ds = Dataset(file, "r", format="NETCDF4")
depths = _rm_mask(ds.variables["depth"][:])  # (75,) depth in m
lats = _rm_mask(ds.variables["nav_lat"][:])  # (1021, 1442) lat in degrees
lons = _rm_mask(ds.variables["nav_lon"][:])  # (1021, 1442) lon in degrees
vals = _rm_mask(ds.variables["vozocrte"][:])  # (76, 1021, 1442) in m/s
vals = vals[0, depths < 5]

dfdict = {"lat": lats.flatten().tolist(), "lon": lons.flatten().tolist()}
for layer_i in range(vals.shape[0]):
    dfdict[f"l{layer_i}"] = vals[layer_i].flatten().tolist()

df = pd.DataFrame(dfdict)
_digitize(df=df, var="lon", interval=(-180, 180), res=0.25)
_digitize(df=df, var="lat", interval=(-70, 70), res=0.25)
df = df.groupby(["lon", "lat"]).mean()
df.sort_index(inplace=True)

Here, most heavy lifting is done by numpy and is multithreaded. _digitize can also be written using np.digitize but it looks a bit confusing (mapping to bin indices) and doesn't make such a huge difference. Overall this computes ~36 times faster (now just a few seconds).

However, as Droid mentioned xarray might be even faster.

$\endgroup$

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service and acknowledge you have read our privacy policy.

Not the answer you're looking for? Browse other questions tagged or ask your own question.