Skip to content
Snippets Groups Projects
Select Git revision
  • b95ac77699f0a46022095f0fee8ebdbd44b90cd6
  • master default protected
  • use_flight_altitude
  • distribute
4 results

lon_lat_grid.py

Blame
  • user avatar
    tomrink authored
    9628db45
    History
    lon_lat_grid.py 12.70 KiB
    
    import numpy as np
    
    from util.setup import home_dir
    from scipy.interpolate import interp2d
    from scipy import linalg
    from scipy.spatial import KDTree
    from sklearn.neighbors import BallTree
    from util.util import haversine_np
    from scipy.ndimage import gaussian_filter
    
    
    class MyGenericException(Exception):
        def __init__(self, message):
            self.message = message
    
    
    # Author: T.Rink
    # For nearest neighbor/interpolation of a function whose domain is a geostationary satellite fixed grid.
    # This is a utility wrapper for sklearn.neighbors.BallTree adapted for 2D manifold (Longitude, Latitude).
    class LonLatGrid:
        # grd_lons, grd_lats: the longitude, latitude of each grid point (must be 2D grids), must have same shape
        # Incoming longitude must be in range: 0 - 360 degrees
        # Can have NaN for off Earth grid points (these are handled internally).
        # closeness_threshold (m): if < distance of located_point to target, return off grid
        # reduce: resolution blow down factor
        # leaf_size: Number of points at which to switch to brute-force, default=40
        def __init__(self, grd_lons, grd_lats, closeness_threshold=2000, reduce=1, leaf_size=40):
            if grd_lons.shape != grd_lats.shape:
                raise MyGenericException('incoming lons,lats must have same shape')
    
            self.grd_lons = grd_lons
            self.grd_lats = grd_lats
            self.grd_lons = grd_lons[::reduce, ::reduce]
            self.grd_lats = grd_lats[::reduce, ::reduce]
            self.shape = self.grd_lons.shape
            self.lenx = self.shape[1]
            self.leny = self.shape[0]
            self.fnc_dct = {}
            self.x = np.zeros(9, np.float32)
            self.y = np.zeros(9, np.float32)
            self.z = np.zeros(9, np.float32)
    
            flons = self.grd_lons.flatten()
            flats = self.grd_lats.flatten()
    
            keep = np.invert(np.isnan(flons))
    
            self.map_indexes = (np.arange(self.grd_lons.size))[keep]
    
            flons = flons[keep]
            flats = flats[keep]
            self.lon_lo, self.lon_hi = flons.min(), flons.max()
            self.lat_lo, self.lat_hi = flats.min(), flats.max()
    
            points = np.stack([flons, flats], axis=1)
    
            self.kd = BallTree(np.deg2rad(points), leaf_size=leaf_size, metric='haversine')
            self.closeness_threshold = closeness_threshold * reduce
            self.reduce = reduce
    
        # locate nearest neighbor for incoming target in earth coordinates (Should not have NaNs)
        # lons, lats can be flat or 2D
        # incoming longitude must be in range: 0 - 360 degrees
        # returns NN indexes relative to the grid determined in the ctr.
        def value_to_index(self, lons, lats):
            if lons.shape != lats.shape:
                raise MyGenericException('incoming lons,lats must have same shape')
    
            if len(lons) == 0 or len(lats) == 0:
                return None, None
    
            lons = lons.flatten()
            lats = lats.flatten()
    
            xy = np.stack([lons, lats], axis=1)
    
            # Assumes nn for a query pont returned in order of increasing distance
            dist, indices = self.kd.query(np.deg2rad(xy), k=1)
            dist = dist[:, 0]
            indices = indices[:, 0]
            dist *= 6370000  # convert unit radius to meters
    
            valid = (indices < self.map_indexes.size) & (dist < self.closeness_threshold)
            indices = indices[valid]
    
            return self.map_indexes[indices], indices
    
        def earth_to_lc_s(self, lons, lats):
            if lons.shape != lats.shape:
                raise MyGenericException('incoming lons,lats must have same shape')
    
            if len(lons) == 0 or len(lats) == 0:
                return None, None
    
            lons = np.where(lons < 0, lons + 360, lons)
    
            lons = lons.flatten()
            lats = lats.flatten()
    
            xy = np.stack([lons, lats], axis=1)
    
            # Assumes nn for a query pont returned in order of increasing distance
            dist, indices = self.kd.query(np.deg2rad(xy), k=1)
            dist = dist[:, 0]
            indices = indices[:, 0]
            dist *= 6370000  # convert unit radius to meters
    
            valid = (indices < self.map_indexes.size) & (dist < self.closeness_threshold)
            indices = np.where(valid, indices, -1)
            ll = np.where(indices >= 0, (indices / self.lenx).astype(np.int32), indices)
            cc = np.where(indices >= 0, indices % self.lenx, indices)
    
            cc *= self.reduce
            ll *= self.reduce
    
            return cc, ll
    
        # bi-linear interpolation of a function on the grid, the range of which is grd_zvals for target lons,lats.
        # returns the interpolated values at each target point.
        # target lons, lats should not have NaNs, and lons must be in range 0 - 360 degrees
        def interp(self, grid_zvals, lons, lats):
            if grid_zvals.shape != self.shape:
                raise MyGenericException('range and domain must have same 2D shape')
    
            if lons.shape != lats.shape:
                raise MyGenericException('incoming lons, lats must have same shape')
    
            abs_idxs, rel_idxs = self.value_to_index(lons, lats)
    
            x = np.zeros(9, np.float32)
            y = np.zeros(9, np.float32)
            z = np.zeros(9, np.float32)
    
            lons = lons.flatten()
            lats = lats.flatten()
    
            outv = np.zeros(lons.size, np.float32)
            outv.fill(np.nan)
    
            for k in range(len(abs_idxs)):
                idx = abs_idxs[k]
                lon = lons[k]
                lat = lats[k]
    
                fnc = self.fnc_dct.get(idx)
                if fnc is None:
                    gj = int(idx / self.lenx)
                    gi = int(idx % self.lenx)
    
                    L = gi-1
                    R = gi+1
                    D = gj-1
                    U = gj+1
    
                    x[0] = self.grd_lons[D, L]
                    y[0] = self.grd_lats[D, L]
                    z[0] = grid_zvals[D, L]
                    x[1] = self.grd_lons[D, gi]
                    y[1] = self.grd_lats[D, gi]
                    z[1] = grid_zvals[D, gi]
                    x[2] = self.grd_lons[D, R]
                    y[2] = self.grd_lats[D, R]
                    z[2] = grid_zvals[D, R]
                    x[3] = self.grd_lons[gj, L]
                    y[3] = self.grd_lats[gj, L]
                    z[3] = grid_zvals[gj, L]
                    x[4] = self.grd_lons[gj, gi]
                    y[4] = self.grd_lats[gj, gi]
                    z[4] = grid_zvals[gj, gi]
                    x[5] = self.grd_lons[gj, R]
                    y[5] = self.grd_lats[gj, R]
                    z[5] = grid_zvals[gj, R]
                    x[6] = self.grd_lons[U, L]
                    y[6] = self.grd_lats[U, L]
                    z[6] = grid_zvals[U, L]
                    x[7] = self.grd_lons[U, gi]
                    y[7] = self.grd_lats[U, gi]
                    z[7] = grid_zvals[U, gi]
                    x[8] = self.grd_lons[U, R]
                    y[8] = self.grd_lats[U, R]
                    z[8] = grid_zvals[U, R]
    
                    if np.isnan(x).sum() > 0:
                        continue
                    if np.isnan(y).sum() > 0:
                        continue
    
                    fnc = interp2d(x, y, z, copy=False)
                    self.fnc_dct[idx] = fnc
    
                outv[k] = fnc(lon, lat)
    
            return outv
    
        def check_inside(self, lon, lat):
            if lon < 0:
                lon += 360
            return (self.lon_lo < lon < self.lon_hi) & (self.lat_lo < lat < self.lat_hi)
    
    
    DEG_TO_RAD = np.pi/180.0;
    RAD_TO_DEG = 180.0/np.pi;
    
    
    r_pol = 6356.5838  # km
    r_eq = 6378.1690  # km
    # AHI -------------------------------
    # h = 42164.0  # barycentric height, km
    # sub_lon = 140.7
    # sub_lon *= DEG_TO_RAD
    # scan_geom = 'GEOS'
    # signs modified from static nav file to work with ABI FGF code
    # CFAC = 5.588799E-05
    # LFAC = -5.588799E-05
    # COFF = -0.1537199
    # LOFF = 0.1537199
    
    # computed from static nav file (lon, lat) grid
    # Note 0-based to the pixel center
    # CFAC = 5.58924125e-05
    # LFAC = -5.58810490e-05
    # COFF = -1.53678977e-01
    # LOFF = 1.53644345e-01
    
    # GOES ------------------------------
    h = 42164.16 # barycentric height, km
    scan_geom = 'GOES'
    sub_lon = -75.0
    sub_lon *= DEG_TO_RAD
    
    # official for FD
    CFAC = 5.6E-05
    LFAC = -5.6E-05
    COFF = -0.151844
    LOFF = 0.151844
    
    # official for CONUS
    #CFAC = 5.6E-05
    #COFF = -0.101332
    #LFAC = -5.6E-05
    #LOFF = 0.128212
    
    # computed for CLAVRx FD
    #CFAC = 5.60016368e-05
    #LFAC = -5.59941969e-05
    #COFF = -1.51780260e-01
    #LOFF = 1.51773560e-01
    
    # 65536 = 2^16
    
    
    def earth_to_sat(geographic_lon, geographic_lat):
    
        geographic_lon *= DEG_TO_RAD
        geographic_lat *= DEG_TO_RAD
    
        geocentric_lat = np.arctan(((r_pol*r_pol)/(r_eq*r_eq))*np.tan(geographic_lat))
    
        r_earth = r_pol / np.sqrt(1.0 - ((r_eq * r_eq - r_pol * r_pol) / (r_eq * r_eq)) * np.cos(geocentric_lat) * np.cos(geocentric_lat))
    
        r_1 = h - r_earth * np.cos(geocentric_lat) * np.cos(geographic_lon - sub_lon)
        r_2 = -r_earth * np.cos(geocentric_lat) * np.sin(geographic_lon - sub_lon)
        r_3 = r_earth * np.sin(geocentric_lat)
    
        if r_1 > h:
            return np.nan, np.nan
    
        if scan_geom == 'GEOS':
            lamda_sat = np.arctan(-r_2/r_1)
            theta_sat = np.arcsin(r_3/np.sqrt(r_1*r_1 + r_2*r_2 + r_3*r_3))
        elif scan_geom == 'GOES':
            lamda_sat = np.arcsin(-r_2/np.sqrt(r_1*r_1 + r_2*r_2 + r_3*r_3))
            theta_sat = np.arctan(r_3/r_1)
    
        return lamda_sat, theta_sat
    
    
    def earth_to_sat_s(geographic_lon, geographic_lat):
    
        geographic_lon = geographic_lon * DEG_TO_RAD
        geographic_lat = geographic_lat * DEG_TO_RAD
    
        geocentric_lat = np.arctan(((r_pol*r_pol)/(r_eq*r_eq))*np.tan(geographic_lat))
    
        r_earth = r_pol / np.sqrt(1.0 - ((r_eq * r_eq - r_pol * r_pol) / (r_eq * r_eq)) * (np.cos(geocentric_lat))**2)
    
        r_1 = h - r_earth * np.cos(geocentric_lat) * np.cos(geographic_lon - sub_lon)
        r_2 = -r_earth * np.cos(geocentric_lat) * np.sin(geographic_lon - sub_lon)
        r_3 = r_earth * np.sin(geocentric_lat)
    
        if scan_geom == 'GEOS':
            lamda_sat = np.arctan(-r_2/r_1)
            theta_sat = np.arcsin(r_3/np.sqrt(r_1*r_1 + r_2*r_2 + r_3*r_3))
        elif scan_geom == 'GOES':
            lamda_sat = np.arcsin(-r_2/np.sqrt(r_1*r_1 + r_2*r_2 + r_3*r_3))
            theta_sat = np.arctan(r_3/r_1)
    
        np.where(r_1 > h, lamda_sat == np.nan, lamda_sat)
        np.where(r_1 > h, theta_sat == np.nan, theta_sat)
    
        return lamda_sat, theta_sat
    
    
    # def earth_to_sat_s(geographic_lons, geographic_lats):
    #     lambdas = []
    #     thetas = []
    #
    #     for i in range(len(geographic_lons)):
    #         l, t = earth_to_sat(geographic_lons[i], geographic_lats[i])
    #         lambdas.append(l)
    #         thetas.append(t)
    #
    #     return lambdas, thetas
    
    
    def sat_to_lc(lamda, theta):
        # c = COFF + np.rint(lamda * (1/np.power(2, 16)) * CFAC)
        # l = LOFF + np.rint(theta * (1/np.power(2, 16)) * LFAC)
    
        # (float(2 ** 16) * (float(lc) - off)) / float(fac) * DEG_TO_RAD
    
        cc = (lamda - COFF) / CFAC
        ll = (theta - LOFF) / LFAC
    
        cc = np.floor(cc + 0.5)
        ll = np.floor(ll + 0.5)
    
        # cc = int(cc)
        # ll = int(ll)
    
        cc = cc.astype(np.int32)
        ll = ll.astype(np.int32)
    
        return cc, ll
    
    
    def test(flons, flats):
        fflons = flons.flatten()
        fflats = flats.flatten()
        num = len(fflons)
    
        for i in range(0, num-1, 50000):
            lon = fflons[i]
            lat = fflats[i]
            if lon == -999.0:
                continue
            c, l = earth_to_lc(lon, lat)
            print(lon, flons[l,c])
            print(lat, flats[l,c])
            print(haversine_np(lon, lat, flons[l,c], flats[l,c]))
            print('-----------------------')
    
    
    def earth_to_lc(lon, lat):
        lamda, theta = earth_to_sat(lon, lat)
        if np.isnan(lamda):
            return None, None
        cc, ll = sat_to_lc(lamda, theta)
        return cc, ll
    
    
    # def earth_to_lc_s(lons, lats):
    #     num = lons.shape[0]
    #     cc_s = []
    #     ll_s = []
    #     for i in range(num):
    #         cc, ll = earth_to_lc(lons[i], lats[i])
    #         if cc is None:
    #             cc_s.append(-1)
    #             ll_s.append(-1)
    #         else:
    #             cc_s.append(cc)
    #             ll_s.append(ll)
    #
    #     return cc_s, ll_s
    
    
    def earth_to_lc_s(lons, lats):
        lamda, theta = earth_to_sat_s(lons, lats)
        cc, ll = sat_to_lc(lamda, theta)
        np.where(np.isnan(cc), cc == -1, cc)
        np.where(np.isnan(ll), ll == -1, ll)
        return cc, ll
    
    
    def earth_to_indexs(lons, lats, len_x):
        num = lons.shape[0]
        idxs = []
        for i in range(num):
            cc, ll = earth_to_lc(lons[i], lats[i])
            if cc is None:
                idxs.append(-1)
            else:
                idx = ll * len_x + cc
                idxs.append(idx)
        idxs = np.array(idxs)
    
        return idxs
    
    
    def compute_scale_offset(lon_a, lat_a, col_a, line_a, lon_b, lat_b, col_b, line_b):
        lamda_a, theta_a = earth_to_sat(lon_a, lat_a)
        lamda_b, theta_b = earth_to_sat(lon_b, lat_b)
    
        # To setup the problem below
        # (CFAC * col_a) + COFF = lamda_a
        # (CFAC * col_b) + COFF = lamda_b
    
        # (LFAC * line_a) + LOFF = theta_a
        # (LFAC * line_b) + LOFF = theta_b
    
        a = np.array([[col_a, 1], [col_b, 1]])
        b = np.array([lamda_a, lamda_b])
        x = linalg.solve(a, b)
        print(x)
    
        a = np.array([[line_a, 1], [line_b, 1]])
        b = np.array([theta_a, theta_b])
        x = linalg.solve(a, b)
        print(x)