<jupyter_start><jupyter_text>Indexing / Selecting Large Data**Note: this feature is experimental!**When the dataset have chunked coordinates (dask arrays), Xoak may build the index and/or performs the selection in-parallel.<jupyter_code>import dask
import dask.array as da
import numpy as np
import xarray as xr
import xoak

xr.set_options(display_style='text');<jupyter_output><empty_output><jupyter_text>Let's first create an `xarray.Dataset` of latitude / longitude points located randomly on the sphere, forming a 2-dimensional (x, y) model mesh. The array coordinates are split into 4 chunks of 250x250 elements each.<jupyter_code>shape = (500, 500)
chunk_shape = (250, 250)

lat = da.random.uniform(-90, 90, size=shape, chunks=chunk_shape)
lon = da.random.uniform(-180, 180, size=shape, chunks=chunk_shape)

field = lat + lon

ds_mesh = xr.Dataset(
    coords={'lat': (('x', 'y'), lat), 'lon': (('x', 'y'), lon)},
    data_vars={'field': (('x', 'y'), field)},
)

ds_mesh<jupyter_output><empty_output><jupyter_text>Xoak builds an index structure for each chunk (all coordinates must have the same chunks):<jupyter_code>ds_mesh.xoak.set_index(['lat', 'lon'], 'sklearn_geo_balltree')

# here returns a list of BallTree objects, one for each chunk
ds_mesh.xoak.index<jupyter_output><empty_output><jupyter_text>Let's create some query data points, which may also be chunked (here 2 chunks).<jupyter_code>shape = (100, 10)
chunk_shape = (50, 10)

ds_data = xr.Dataset({
    'lat': (('y', 'x'), da.random.uniform(-90, 90, size=shape, chunks=chunk_shape)),
    'lon': (('y', 'x'), da.random.uniform(-180, 180, size=shape, chunks=chunk_shape))
})

ds_data<jupyter_output><empty_output><jupyter_text>Queries may be perfomed in parallel using Dask. Please note, however, that some indexes may not support multi-threads and/or multi-process parallelism.<jupyter_code>from dask.diagnostics import ProgressBar

with ProgressBar(), dask.config.set(scheduler='processes'):
    ds_selection = ds_mesh.xoak.sel(lat=ds_data.lat, lon=ds_data.lon)

ds_selection<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.13246628131021196, 'fraction_numerical': 0.024084778420038536, 'mean_word_length': 6.2369337979094075, 'pattern_counts': {'":': 0, '<': 21, '<?xml version=': 0, '>': 21, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '5', 'dir': 'jupyter-structured-clean-dedup', 'id': '263866', 'license': 'MIT', 'n_tokens_mistral': 748, 'n_tokens_neox': 706, 'n_words': 199, 'path': 'doc/examples/dask_support.ipynb', 'repo_name': 'raybellwaves/xoak', 'size': '2076'}