Custom interpolation functions in embedded datasets

The cartesian yt dataset returned by yt_xarray.transformations.build_interpolated_cartesian_ds consists of a cartesian grid (or set of refined grids) that will sample the underyling xarray dataset as needed. In the final step of this process, the native coordinates of the cartesian grid points are calculated (using the supplied Transformer object – see previous notebooks) and then the underlying xarray dataset is sampled at those native coordinates. Because the exact values of requested native coordinates is unlikely to exist, an interpolation is required and there are a number of ways to control the interpolation:

The keyword argument, interp_method, can be set to 'nearest' (the default) or 'interpolate'. If 'nearest', then the nearest value to the request point is used and if 'interpolate', a linear interpolation will be used. In both cases, underlying xarray methods are used. For 'nearest', xr.DataArray.data.sel(interp_dict, method="nearest") and for 'interpolate', xr.DataArray.interp is used (which relies on the scipy N-D linear interpolation method).

Additionally, you can supply your own interpolation function to build_interpolated_cartesian_ds. The function will be called with interp_func(data=data_array, coords=eval_coords), where data_array is an xarray DataArray and eval_coords is a list of 1d np.ndarray ordered by the native_coords of the supplied transformer. The function must return an np.ndarray of the same shape as the eval_coords.

As an example, we’ll build a custom interpolation function similar to the implementation when you set interp_method='nearest'.

Let’s first import all the modules we need then initialize a transformer:

import numpy as np
import numpy.typing as npt
import xarray as xr
import yt_xarray
from yt_xarray.transformations import GeocentricCartesian, build_interpolated_cartesian_ds
from yt_xarray.sample_data import load_random_xr_data
from typing import List
import yt

gc = GeocentricCartesian(radial_type='altitude', r_o=6371.)
gc.native_coords

('altitude', 'latitude', 'longitude')

The transformer native_coords tuple above indicates the order in which coordinate arrays will be supplied.

To write an interpolation function that samples an xarray.DataArray at its nearest points:

def my_interp(data: xr.DataArray = None,
              coords: List[npt.NDArray] =  None) -> npt.NDArray:
    print("hello from my_interp!")
    c0 = coords[0] # altitude
    c1 = coords[1] # latitude
    c2 = coords[2] # longitude

    interp_dict = {
        'altitude':  xr.DataArray(c0, dims='points'),
        'latitude':  xr.DataArray(c1, dims='points'),
        'longitude':  xr.DataArray(c2, dims='points'),
    }
    vals = data.sel(interp_dict, method='nearest')
    return vals.values

Now let’s get a sample dataset to use:

ds = load_random_xr_data({'field1':('altitude', 'latitude', 'longitude')},
                         {'altitude': (0, 2000, 10),
                          'latitude': (0, 20, 15),
                          'longitude': (10, 20, 12)},
                          length_unit='m')

ds

<xarray.Dataset>
Dimensions:    (altitude: 10, latitude: 15, longitude: 12)
Coordinates:
  * altitude   (altitude) float64 0.0 222.2 444.4 ... 1.556e+03 1.778e+03 2e+03
  * latitude   (latitude) float64 0.0 1.429 2.857 4.286 ... 17.14 18.57 20.0
  * longitude  (longitude) float64 10.0 10.91 11.82 12.73 ... 18.18 19.09 20.0
Data variables:
    field1     (altitude, latitude, longitude) float64 0.9765 0.4326 ... 0.01994
Attributes:
    geospatial_vertical_units:  m

xarray.Dataset

• Dimensions:
  • altitude: 10
  • latitude: 15
  • longitude: 12
• Coordinates: (3)
  • altitude
    (altitude)
    float64
    0.0 222.2 444.4 ... 1.778e+03 2e+03

    array([   0.      ,  222.222222,  444.444444,  666.666667,  888.888889,
           1111.111111, 1333.333333, 1555.555556, 1777.777778, 2000.      ])

  • latitude
    (latitude)
    float64
    0.0 1.429 2.857 ... 18.57 20.0

    array([ 0.      ,  1.428571,  2.857143,  4.285714,  5.714286,  7.142857,
            8.571429, 10.      , 11.428571, 12.857143, 14.285714, 15.714286,
           17.142857, 18.571429, 20.      ])

  • longitude
    (longitude)
    float64
    10.0 10.91 11.82 ... 19.09 20.0

    array([10.      , 10.909091, 11.818182, 12.727273, 13.636364, 14.545455,
           15.454545, 16.363636, 17.272727, 18.181818, 19.090909, 20.      ])

• Data variables: (1)
  • field1
    (altitude, latitude, longitude)
    float64
    0.9765 0.4326 ... 0.5896 0.01994

    array([[[0.97648284, 0.43259147, 0.49484548, ..., 0.2100358 ,
             0.22365835, 0.87192953],
            [0.21698072, 0.17765705, 0.44666109, ..., 0.43118691,
             0.81980718, 0.60155836],
            [0.0496882 , 0.52415221, 0.81631591, ..., 0.89868623,
             0.1196812 , 0.92476923],
            ...,
            [0.85168786, 0.18861585, 0.1010581 , ..., 0.65116744,
             0.80913434, 0.99554168],
            [0.94994778, 0.73501508, 0.33130486, ..., 0.09214206,
             0.79772242, 0.4401901 ],
            [0.6513003 , 0.56971011, 0.88293915, ..., 0.80956066,
             0.06383272, 0.01737054]],
    
           [[0.31845955, 0.84642622, 0.36578082, ..., 0.2169022 ,
             0.14139395, 0.53262278],
            [0.16335444, 0.32337149, 0.48930071, ..., 0.08371515,
             0.07075334, 0.30153669],
            [0.38996034, 0.42410577, 0.42855397, ..., 0.72421923,
             0.68320305, 0.55654341],
    ...
            [0.99251558, 0.09839139, 0.78490873, ..., 0.62758996,
             0.34293408, 0.6483621 ],
            [0.3781159 , 0.18710431, 0.37240958, ..., 0.368602  ,
             0.33157903, 0.22529801],
            [0.80450525, 0.6103617 , 0.31176877, ..., 0.15239449,
             0.1567234 , 0.47431269]],
    
           [[0.43128443, 0.1044148 , 0.57556023, ..., 0.9485631 ,
             0.69214876, 0.30355404],
            [0.95951226, 0.10288168, 0.55770819, ..., 0.98880661,
             0.64012401, 0.92253858],
            [0.00666931, 0.28934773, 0.54202978, ..., 0.13282967,
             0.86856248, 0.17350749],
            ...,
            [0.12687199, 0.68315539, 0.01124447, ..., 0.81112687,
             0.11683871, 0.05323632],
            [0.64666392, 0.77458101, 0.53552976, ..., 0.91914587,
             0.42077197, 0.20270676],
            [0.53949917, 0.368608  , 0.77486222, ..., 0.58419937,
             0.58956882, 0.01993657]]])

• Indexes: (3)
  • altitude
    PandasIndex

    PandasIndex(Index([               0.0, 222.22222222222223, 444.44444444444446,
            666.6666666666667,  888.8888888888889,  1111.111111111111,
           1333.3333333333335, 1555.5555555555557, 1777.7777777777778,
                       2000.0],
          dtype='float64', name='altitude'))

  • latitude
    PandasIndex

    PandasIndex(Index([               0.0, 1.4285714285714286,  2.857142857142857,
            4.285714285714286,  5.714285714285714,  7.142857142857143,
            8.571428571428571,               10.0, 11.428571428571429,
           12.857142857142858, 14.285714285714286, 15.714285714285715,
           17.142857142857142, 18.571428571428573,               20.0],
          dtype='float64', name='latitude'))

  • longitude
    PandasIndex

    PandasIndex(Index([              10.0, 10.909090909090908, 11.818181818181818,
           12.727272727272727, 13.636363636363637, 14.545454545454545,
           15.454545454545453, 16.363636363636363, 17.272727272727273,
            18.18181818181818,  19.09090909090909,               20.0],
          dtype='float64', name='longitude'))

• Attributes: (1)

  geospatial_vertical_units :

  m

Finally, we supply the function handle of our my_interp function when building the interpolated dataset:

ds_yt = build_interpolated_cartesian_ds(ds,
                                gc,
                                grid_resolution=(128,128,128),
                                length_unit='m',
                                interp_func=my_interp
                                )

yt_xarray : [INFO ] 2025-11-05 16:25:34,716:  Interpolation function provided, switching interp_method to 'interpolate'.
yt : [INFO     ] 2025-11-05 16:25:34,855 Parameters: current_time              = 0.0
yt : [INFO     ] 2025-11-05 16:25:34,856 Parameters: domain_dimensions         = [128 128 128]
yt : [INFO     ] 2025-11-05 16:25:34,856 Parameters: domain_left_edge          = [5625. 1039.    0.]
yt : [INFO     ] 2025-11-05 16:25:34,857 Parameters: domain_right_edge         = [8244. 2864. 2864.]
yt : [INFO     ] 2025-11-05 16:25:34,858 Parameters: cosmological_simulation   = 0

yt.SlicePlot(ds_yt, 'z', 'field1')

hello from my_interp!

yt : [INFO     ] 2025-11-05 16:25:36,230 xlim = 5625.000000 8244.000000
yt : [INFO     ] 2025-11-05 16:25:36,230 ylim = 1039.000000 2864.000000
yt : [INFO     ] 2025-11-05 16:25:36,232 xlim = 5625.000000 8244.000000
yt : [INFO     ] 2025-11-05 16:25:36,233 ylim = 1039.000000 2864.000000
yt : [INFO     ] 2025-11-05 16:25:36,239 Making a fixed resolution buffer of (('stream', 'field1')) 800 by 800

the custom interpolation function can be as complex as you need.