Nearest-neighbour indexing using xoak

Contents

Nearest-neighbour indexing using xoak#

This notebook experiments subsetting datasets using TLONG, TLAT, ULONG, ULAT by making use of the xoak package

import matplotlib as mpl
import numpy as np
import xarray as xr
import xoak

import pop_tools

---------------------------------------------------------------------------
ModuleNotFoundError                       Traceback (most recent call last)
Cell In[1], line 4
      2 import numpy as np
      3 import xarray as xr
----> 4 import xoak
      6 import pop_tools

ModuleNotFoundError: No module named 'xoak'

grid, xds = pop_tools.to_xgcm_grid_dataset(ds)

Set the “index”#

This is what allows the indexing magic.

xds.xoak.set_index(['TLAT', 'TLONG'], 'scipy_kdtree')

Extracting sections#

# have to know that 0.1 is a reasonable choice
lons = xr.Variable("points", np.arange(150, 270, 0.1))
lats = xr.zeros_like(lons)

eqsection = xds.xoak.sel(TLONG=lons, TLAT=lats)

# plot
eqsection.TEMP.sel(z_t=slice(50000)).plot(y="z_t", cmap=mpl.cm.Spectral_r, yincrease=False)

<matplotlib.collections.QuadMesh at 0x7f18c3a888e0>

../_images/c460949869e5adf06420dc70664292157c162e45c1a3c696f004fafb50170168.png

Extracting multiple points#

moorings = xds.xoak.sel(
    TLONG=xr.Variable("moor", [360 - 140, 360 - 110]),
    TLAT=xr.Variable("moor", [0, 0]),
)

moorings.TEMP.plot(hue="TLONG", y="z_t", yincrease=False)

[<matplotlib.lines.Line2D at 0x7f18c3aa57f0>,
 <matplotlib.lines.Line2D at 0x7f18c3aa58e0>]

../_images/cb5e4671578f786ce2c9e6048ee86dc9092847e23689fccb8196812a9e9392bd.png

Limitations#

[ ] cannot simply index by point
[ ] Can only set one index at a time (so only, TLONG, TLAT or ULONG, ULAT)
[ ] indexes are not propagated so xds.TEMP.xoak will only work after xds.TEMP.set_index is called

xds.TEMP.xoak.sel(TLONG=220, TLAT=0)

---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
<ipython-input-13-550ba18bf81d> in <module>
----> 1 xds.TEMP.xoak.sel(TLONG=220, TLAT=0)

~/miniconda3/envs/dcpy/lib/python3.8/site-packages/xoak/accessor.py in sel(self, indexers, **indexers_kwargs)
    246         """
    247         if not getattr(self, '_index', False):
--> 248             raise ValueError(
    249                 'The index(es) has/have not been built yet. Call `.xoak.set_index()` first'
    250             )

ValueError: The index(es) has/have not been built yet. Call `.xoak.set_index()` first

xds.xoak.sel(TLONG=220, TLAT=0)

---------------------------------------------------------------------------
AttributeError                            Traceback (most recent call last)
<ipython-input-12-227cdaa281aa> in <module>
----> 1 xds.xoak.sel(TLONG=220, TLAT=0)

~/miniconda3/envs/dcpy/lib/python3.8/site-packages/xoak/accessor.py in sel(self, indexers, **indexers_kwargs)
    251 
    252         indexers = either_dict_or_kwargs(indexers, indexers_kwargs, 'xoak.sel')
--> 253         indices = self._query(indexers)
    254 
    255         if not isinstance(indices, np.ndarray):

~/miniconda3/envs/dcpy/lib/python3.8/site-packages/xoak/accessor.py in _query(self, indexers)
    135 
    136     def _query(self, indexers):
--> 137         X = coords_to_point_array([indexers[c] for c in self._index_coords])
    138 
    139         if isinstance(X, np.ndarray) and isinstance(self._index, XoakIndexWrapper):

~/miniconda3/envs/dcpy/lib/python3.8/site-packages/xoak/accessor.py in coords_to_point_array(coords)
     18 
     19     """
---> 20     c_chunks = [c.chunks for c in coords]
     21 
     22     if any([chunks is None for chunks in c_chunks]):

~/miniconda3/envs/dcpy/lib/python3.8/site-packages/xoak/accessor.py in <listcomp>(.0)
     18 
     19     """
---> 20     c_chunks = [c.chunks for c in coords]
     21 
     22     if any([chunks is None for chunks in c_chunks]):

AttributeError: 'int' object has no attribute 'chunks'

%load_ext watermark
%watermark -d -iv -m -g

Compiler    : GCC 9.3.0
OS          : Linux
Release     : 5.8.0-44-generic
Machine     : x86_64
Processor   : x86_64
CPU cores   : 8
Architecture: 64bit

Git hash: bd1236ca615b32595c43cfa689e85fc9a112eb9f

numpy     : 1.20.1
xoak      : 0.1.0
xarray    : 0.17.1.dev3+g48378c4b1
matplotlib: 3.3.4
pop_tools : 2020.12.15.post6+dirty