Suggestions for vectorizing a for loop?

[ahuang11]

I am looping through some 3D coordinates (time, lat, lon), find the nearest grid point and extract a 6x6 degree values from the gridded field, add it to an initialized with zeros array, and take the mean.

However, I was wondering if there's a vectorized way to do this rather than a loop method because it takes really long for many coordinate points?

import xarray as xr
import pandas as pd
import numpy as np

# generate example data
ds = xr.tutorial.open_dataset('air_temperature').sortby('lat')
df = pd.DataFrame({
    'time': ds['time'][:10],
    'lat': np.arange(40, 45, 0.5),
    'lon': np.arange(220, 225, 0.5)
})

# create a placeholder to add to
ds_mean = xr.zeros_like(ds).isel(time=0, drop=True)

# loop through each coordinate and add changes within 6x6 degrees 
for time, lat, lon in zip(df['time'], df['lat'], df['lon']):
    ds_mean = ds_mean + ds.sel(
        time=time, method='nearest'
    ).sel(
        lat=slice(lat - 3, lat + 3),
        lon=slice(lon - 3, lon + 3)
    ).reindex_like(ds_mean).fillna(0)

# take the average
ds_mean = ds_mean / len(df)

# plot
ds_mean['air'].plot()

I could interpolate it like this, but doesn't satisfy the 6x6 degree requirement.

ds_mean_vector = (
    ds
    .interp(**df.to_xarray())  # interpolate to points
    .to_dataframe().set_index(['time', 'lat', 'lon']).to_xarray()  # somewhat hacky method of making a 10x10x10 array
    ['air'].mean('time')  # take the mean across time
    .reindex_like(ds).fillna(0)  # make it look like the original grid
)
ds_mean_vector.plot()

[ahuang11]

This vectorizes the interpolation, but only slightly faster since it still loops to create the base coordinates.

import xarray as xr
import pandas as pd
import numpy as np

# generate example data
ds = xr.tutorial.open_dataset('air_temperature').sortby('lat')
df = pd.DataFrame({
    'time': ds['time'][:10],
    'lat': np.arange(40, 45, 0.5),
    'lon': np.arange(220, 225, 0.5)
})
df = pd.concat([df] * 1000)  # just to test speed

df_coords = pd.concat((
    xr.Dataset(coords={
        'lat': np.arange(row['lat'] - 3, row['lat'] + 3.5, 0.5),
        'lon': np.arange(row['lon'] - 3, row['lon'] + 3.5, 0.5),
        'time': time
    }).to_dataframe().reset_index() for r, row in df.iterrows()
))

ds_mean_vector = (
    ds
    .interp(**df_coords.to_xarray())  # interpolate to points
    .to_dataframe().set_index(['time', 'lat', 'lon'])
    .sum(level=['lat', 'lon']).to_xarray()  # somewhat hacky method of making a 10x10x10 array
    ['air']
    .reindex_like(ds).fillna(0)  # make it look like the original grid
) / len(df)
ds_mean_vector.plot()

[ahuang11]

Okay instead of looping through each row, looping through the offsets and adding it to the lat/lon in a vectorized fashion is much faster from 25 seconds down to ~2 seconds

import xarray as xr
import pandas as pd
import numpy as np

# generate example data
ds = xr.tutorial.open_dataset('air_temperature').sortby('lat')
df = pd.DataFrame({
    'time': ds['time'][:10],
    'lat': np.arange(40, 45, 0.5),
    'lon': np.arange(220, 225, 0.5)
})
df = pd.concat([df] * 1000)

df_list = []
df = df.set_index('time')
for lon_offset in np.arange(-3, 3.5, 0.5):
    for lat_offset in np.arange(-3, 3.5, 0.5):
        df_copy = df.copy()
        df_copy['lon'] += lon_offset
        df_copy['lat'] += lat_offset
        df_list.append(df_copy)
df_coords = pd.concat(df_list).reset_index()

ds_mean_vector2 = (
    ds
    .interp(**df_coords.to_xarray())  # interpolate to points
    .to_dataframe().set_index(['time', 'lat', 'lon'])
    .sum(level=['lat', 'lon']).to_xarray()  # somewhat hacky method of making a 10x10x10 array
    ['air']
    .reindex_like(ds).fillna(0)  # make it look like the original grid
) / len(df)
ds_mean_vector2.plot()