Find the closest grid point to a set of points

map_tunnel_coords.py

import numpy as np
import math
from scipy.spatial import cKDTree as KDTree

WGS84_A = 6378137.0  # semi major axis in metres

def tunnel_coords(lats, lons, elev=0.0, rearth=WGS84_A):
    """
    Transforms lat/lon coordinates to tunnel coordinates.

    Args:
        lats: array, latitude coordinates.
        lons: array, longitude coordinates.

    Returns: tuple, transformed coordinates.
    """
    rlats, rlons = np.deg2rad(lats), np.deg2rad(lons)
    x = (elev + rearth) * np.cos(rlats) * np.cos(rlons)
    y = (elev + rearth) * np.cos(rlats) * np.sin(rlons)
    z = (elev + rearth) * np.sin(rlats)
    return x, y, z
  
def map_closest(a,b,coord_trans_func=tunnel_coords):
    """ Map points a (NxD) to points b(MxD)
    a or b are just coordinate sets, can be sites or grids it doesn't matter.
    N,M number of points, D number of spatial dimensions.
    Args:
       a: points to map (NxD)
       b: points to map to (MxD)
       coord_func: function to transform

    Returns: distances between points in a and b, indices of mapped-to points
    """
    xi,yi,zi = coord_func(a[:,0],a[:,1])
    xo,yo,zo = coord_func(b[:,0],b[:,1])
    btrans = np.dstack([xo,yo,zo]).squeeze()
    atrans = np.dstack([xi,yi,zi]).squeeze()
    kt = KDTree(btrans)
    dist,ind = kt.query(atrans)
    return np.atleast_1d(dist),np.atleast_1d(ind)