Skip to content

Commit

Permalink
add function for computing angle
Browse files Browse the repository at this point in the history
  • Loading branch information
simonbyrne committed Sep 28, 2024
1 parent 1487746 commit 3b66299
Showing 1 changed file with 21 additions and 6 deletions.
27 changes: 21 additions & 6 deletions earth2grid/lcc.py
Original file line number Diff line number Diff line change
Expand Up @@ -61,16 +61,27 @@ def __init__(self, lat0: float, lon0: float, lat1: float, lat2: float, radius: f
def _rho(self, lat):
return self.RF / np.power(np.tan(np.pi / 4 + np.deg2rad(lat) / 2), self.n)

def _theta(self, lon):
"""
Angle of deviation (in radians) of the projected grid from the regular grid,
for a given longitude (in degrees).
To convert to U and V on the projected grid to easterly / northerly components:
UN = cos(theta) * U + sin(theta) * V
VN = - sin(theta) * U + cos(theta) * V
"""
# center about reference longitude
delta_lon = lon - self.lon0
delta_lon = delta_lon - np.round(delta_lon/360) * 360 # convert to [-180, 180]
return self.n * np.deg2rad(delta_lon)


def proj(self, lat, lon):
"""
Compute the projected x,y from lat,lon.
"""
rho = self._rho(lat)

delta_lon = lon - self.lon0
delta_lon = delta_lon - np.round(delta_lon/360) * 360 # convert to [-180, 180]
theta = self.n * np.deg2rad(delta_lon)
theta = self._theta(lon)

x = rho * np.sin(theta)
y = self.rho0 - rho * np.cos(theta)
Expand Down Expand Up @@ -170,15 +181,19 @@ def hrrr_conus_grid(ix0 = 0, iy0 = 0, nx = 1799, ny = 1059):


class _RegridFromLCC(torch.nn.Module):
"""Regrid from lat-lon to unstructured grid with bilinear interpolation"""
"""Regrid from LambertConformalConicGrid to unstructured grid with bilinear interpolation"""

def __init__(self, src: LambertConformalConicGrid, lat: np.ndarray, lon: np.ndarray):
super().__init__()

x, y = src.projection.proj(lat, lon)

self.shape = lat.shape
self._bilinear = BilinearInterpolator(torch.from_numpy(src.x), torch.from_numpy(src.y), y_query= torch.from_numpy(y.ravel()), x_query=torch.from_numpy(x.ravel()))
self._bilinear = BilinearInterpolator(
x_coords = torch.from_numpy(src.x),
y_coords = torch.from_numpy(src.y),
x_query = torch.from_numpy(x.ravel()),
y_query = torch.from_numpy(y.ravel()))

def forward(self, x: torch.Tensor) -> torch.Tensor:
out = self._bilinear(x)
Expand Down

0 comments on commit 3b66299

Please sign in to comment.