From 227b43e889d4b4884a26edc48fdb011d93b7ec66 Mon Sep 17 00:00:00 2001
From: Simon Byrne <sbyrne@nvidia.com>
Date: Fri, 27 Sep 2024 16:03:19 -0700
Subject: [PATCH 1/7] add HRRR CONUS grid

---
 earth2grid/__init__.py |   7 +-
 earth2grid/lcc.py      | 185 +++++++++++++++++++++++++++++++++++++++++
 tests/test_lcc.py      |  56 +++++++++++++
 3 files changed, 246 insertions(+), 2 deletions(-)
 create mode 100644 earth2grid/lcc.py
 create mode 100644 tests/test_lcc.py

diff --git a/earth2grid/__init__.py b/earth2grid/__init__.py
index 6d52f27..30a0eca 100644
--- a/earth2grid/__init__.py
+++ b/earth2grid/__init__.py
@@ -14,13 +14,14 @@
 # limitations under the License.
 import torch
 
-from earth2grid import base, healpix, latlon
+from earth2grid import base, healpix, latlon, lcc
 from earth2grid._regrid import BilinearInterpolator, Identity, KNNS2Interpolator, Regridder
 
 __all__ = [
     "base",
     "healpix",
     "latlon",
+    "lcc",
     "get_regridder",
     "BilinearInterpolator",
     "KNNS2Interpolator",
@@ -32,10 +33,12 @@ def get_regridder(src: base.Grid, dest: base.Grid) -> torch.nn.Module:
     """Get a regridder from `src` to `dest`"""
     if src == dest:
         return Identity()
-    elif isinstance(src, latlon.LatLonGrid) and isinstance(dest, latlon.LatLonGrid):
+    elif isinstance(src, latlon.LatLonGrid)
         return src.get_bilinear_regridder_to(dest.lat, dest.lon)
     elif isinstance(src, latlon.LatLonGrid) and isinstance(dest, healpix.Grid):
         return src.get_bilinear_regridder_to(dest.lat, dest.lon)
+    elif isinstance(src, lcc.LambertConformalConicGrid)
+        return src.get_bilinear_regridder_to(dest.lat, dest.lon)
     elif isinstance(src, healpix.Grid):
         return src.get_bilinear_regridder_to(dest.lat, dest.lon)
     elif isinstance(dest, healpix.Grid):
diff --git a/earth2grid/lcc.py b/earth2grid/lcc.py
new file mode 100644
index 0000000..8168d6d
--- /dev/null
+++ b/earth2grid/lcc.py
@@ -0,0 +1,185 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+import numpy as np
+import torch
+
+from earth2grid import base
+from earth2grid._regrid import BilinearInterpolator
+
+try:
+    import pyvista as pv
+except ImportError:
+    pv = None
+
+
+
+class LambertConformalConicProjection:
+    def __init__(self, lat0: float, lon0: float, lat1: float, lat2: float, radius: float):
+        """
+        
+        Args:
+            lat0: latitude of origin (degrees)
+            lon0: longitude of origin (degrees)
+            lat1: first standard parallel (degrees)
+            lat2: second standard parallel (degrees)
+            radius: radius of sphere (m)
+
+        """
+
+        self.lon0 = lon0
+        self.lat0 = lat0
+        self.lat1 = lat1
+        self.lat2 = lat2
+        self.radius = radius
+
+
+        c1 = np.cos(np.deg2rad(lat1))
+        c2 = np.cos(np.deg2rad(lat2))
+        t1 = np.tan(np.pi / 4 + np.deg2rad(lat1) / 2)
+        t2 = np.tan(np.pi / 4 + np.deg2rad(lat2) / 2)
+
+        if np.abs(lat1 - lat2) < 1e-8:
+            self.n = np.sin(np.deg2rad(lat1))
+        else:            
+            self.n = np.log(c1/c2) / np.log(t2/t1)
+
+        self.RF = radius * c1 * np.power(t1, self.n) / self.n
+        self.rho0 = self._rho(lat0)
+
+    def _rho(self, lat):
+        return self.RF / np.power(np.tan(np.pi / 4 + np.deg2rad(lat) / 2), self.n)
+
+
+    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)
+    
+        x = rho * np.sin(theta)
+        y = self.rho0 - rho * np.cos(theta)
+        return x, y
+
+    def inv(self, x, y):
+        """
+        Compute the lat,lon from the projected x,y.
+        """
+        rho = np.hypot(x, self.rho0 - y)
+        theta = np.arctan2(x, self.rho0 - y)
+        
+        lat = np.rad2deg(2 * np.arctan(np.power(self.RF/rho, 1/self.n))) - 90
+        lon = self.lon0 + np.rad2deg(theta / self.n)
+        return lat, lon
+
+# Projection used by HRRR CONUS (Continental US) data
+HRRR_CONUS_PROJECTION = LambertConformalConicProjection(
+    lon0 = -97.5,
+    lat0 = 38.5,
+    lat1 = 38.5,
+    lat2 = 38.5,
+    radius = 6371229.0
+)
+
+
+class LambertConformalConicGrid(base.Grid):
+    def __init__(self, projection: LambertConformalConicProjection, x, y):
+        """        
+        Args:
+            projection: LambertConformalConicProjection object
+            x: range of x values
+            y: range of y values
+
+        """
+        self.projection = projection
+
+        self.x = np.array(x)
+        self.y = np.array(y)
+
+    @property
+    def lat_lon(self):
+        mesh_x, mesh_y = np.meshgrid(self.x, self.y)
+        return self.projection.inv(mesh_x, mesh_y)
+
+    @property
+    def lat(self):
+        return self.lat_lon[0]
+
+    @property
+    def lon(self):
+        return self.lat_lon[1]
+
+    @property
+    def shape(self):
+        return (len(self.y), len(self.x))
+
+    def __getitem__(self, idxs):
+        yidxs, xidxs = idxs
+        return LambertConformalConicGrid(self.projection, x=self.x[xidxs], y=self.y[yidxs])
+
+    def get_bilinear_regridder_to(self, lat: np.ndarray, lon: np.ndarray):
+        """Get regridder to the specified lat and lon points"""
+        return _RegridFromLCC(self, lat, lon)
+
+    def visualize(self, data):
+        raise NotImplementedError()
+
+    def to_pyvista(self):
+        if pv is None:
+            raise ImportError("Need to install pyvista")
+
+        lat, lon = self.lat_lon
+        y = np.cos(np.deg2rad(lat)) * np.sin(np.deg2rad(lon))
+        x = np.cos(np.deg2rad(lat)) * np.cos(np.deg2rad(lon))
+        z = np.sin(np.deg2rad(lat))
+        grid = pv.StructuredGrid(x, y, z)
+        return grid
+
+def hrrr_conus_grid(ix0 = 0, iy0 = 0, nx = 1799, ny = 1059):
+    # coordinates of point in top-left corner
+    lat0 = 21.138123
+    lon0 = 237.280472
+    # grid length (m)
+    scale = 3000.0
+    # coordinates on projected space
+    x0, y0 = HRRR_CONUS_PROJECTION.proj(lat0, lon0)
+
+    x = [x0 + i * scale for i in range(ix0, ix0+nx)]
+    y = [y0 + i * scale for i in range(iy0, iy0+ny)]
+
+    return LambertConformalConicGrid(HRRR_CONUS_PROJECTION, x, y)
+
+# Grid used by HRRR CONUS (Continental US) data
+HRRR_CONUS_GRID = hrrr_conus_grid()
+
+
+class _RegridFromLCC(torch.nn.Module):
+    """Regrid from lat-lon 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()))
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        out = self._bilinear(x)
+        return out.view(out.shape[:-1] + self.shape)
+
diff --git a/tests/test_lcc.py b/tests/test_lcc.py
new file mode 100644
index 0000000..8f0d5bf
--- /dev/null
+++ b/tests/test_lcc.py
@@ -0,0 +1,56 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+#%%
+from earth2grid.lcc import HRRR_CONUS_GRID
+import numpy as np
+import torch
+from pytest import approx
+
+def test_grid_shape():
+    assert HRRR_CONUS_GRID.lat.shape == HRRR_CONUS_GRID.shape 
+    assert HRRR_CONUS_GRID.lon.shape == HRRR_CONUS_GRID.shape
+
+lats = np.array([
+        [21.138123, 21.801926, 22.393631, 22.911015],
+        [23.636763, 24.328228, 24.944668, 25.48374 ],
+        [26.155672, 26.875362, 27.517046, 28.078257],
+        [28.69017 , 29.438608, 30.106009, 30.68978 ]])
+
+lons = np.array([
+        [-122.71953 , -120.03195 , -117.304596, -114.54146 ],
+        [-123.491356, -120.72898 , -117.92319 , -115.07828 ],
+        [-124.310524, -121.469505, -118.58098 , -115.649574],
+        [-125.181404, -122.25762 , -119.28173 , -116.25871 ]])
+
+
+def test_grid_vals():
+    assert HRRR_CONUS_GRID.lat[0:400:100,0:400:100] == approx(lats)
+    assert HRRR_CONUS_GRID.lon[0:400:100,0:400:100] == approx(lons)
+
+def test_grid_slice():
+    slice_grid = HRRR_CONUS_GRID[0:400:100,0:400:100]
+    assert slice_grid.lat == approx(lats)
+    assert slice_grid.lon == approx(lons)
+
+def test_grid_slice():
+
+    src = HRRR_CONUS_GRID
+    dest = HRRR_CONUS_GRID[1:-1, 1:-1]
+    regrid = src.get_bilinear_regridder_to(dest.lat, dest.lon)
+    lat = torch.broadcast_to(torch.tensor(src.lat), src.shape)
+    out = regrid(lat)
+
+    assert torch.allclose(out, lat[1:-1, 1:-1])

From 009f8a05434809a4a5b5da5fec91a67e512fcf38 Mon Sep 17 00:00:00 2001
From: Simon Byrne <sbyrne@nvidia.com>
Date: Fri, 27 Sep 2024 16:10:30 -0700
Subject: [PATCH 2/7] add comment

---
 earth2grid/lcc.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/earth2grid/lcc.py b/earth2grid/lcc.py
index 8168d6d..ce60bf1 100644
--- a/earth2grid/lcc.py
+++ b/earth2grid/lcc.py
@@ -98,6 +98,7 @@ def inv(self, x, y):
 
 
 class LambertConformalConicGrid(base.Grid):
+    # nothing here is specific to the projection, so could be shared by any projected rectilinear grid
     def __init__(self, projection: LambertConformalConicProjection, x, y):
         """        
         Args:

From 1487746929687e4c62c766e911af74372a861792 Mon Sep 17 00:00:00 2001
From: Simon Byrne <simonbyrne@gmail.com>
Date: Fri, 27 Sep 2024 23:45:48 +0000
Subject: [PATCH 3/7] fix import, add 1d and 2d tests

---
 earth2grid/__init__.py |  4 ++--
 tests/test_lcc.py      | 22 ++++++++++++++++------
 2 files changed, 18 insertions(+), 8 deletions(-)

diff --git a/earth2grid/__init__.py b/earth2grid/__init__.py
index 30a0eca..3fe995c 100644
--- a/earth2grid/__init__.py
+++ b/earth2grid/__init__.py
@@ -33,11 +33,11 @@ def get_regridder(src: base.Grid, dest: base.Grid) -> torch.nn.Module:
     """Get a regridder from `src` to `dest`"""
     if src == dest:
         return Identity()
-    elif isinstance(src, latlon.LatLonGrid)
+    elif isinstance(src, latlon.LatLonGrid) and isinstance(dest, latlon.LatLonGrid):
         return src.get_bilinear_regridder_to(dest.lat, dest.lon)
     elif isinstance(src, latlon.LatLonGrid) and isinstance(dest, healpix.Grid):
         return src.get_bilinear_regridder_to(dest.lat, dest.lon)
-    elif isinstance(src, lcc.LambertConformalConicGrid)
+    elif isinstance(src, lcc.LambertConformalConicGrid):
         return src.get_bilinear_regridder_to(dest.lat, dest.lon)
     elif isinstance(src, healpix.Grid):
         return src.get_bilinear_regridder_to(dest.lat, dest.lon)
diff --git a/tests/test_lcc.py b/tests/test_lcc.py
index 8f0d5bf..3f9764e 100644
--- a/tests/test_lcc.py
+++ b/tests/test_lcc.py
@@ -45,12 +45,22 @@ def test_grid_slice():
     assert slice_grid.lat == approx(lats)
     assert slice_grid.lon == approx(lons)
 
-def test_grid_slice():
+def test_regrid_1d():
+    src = HRRR_CONUS_GRID
+    dest_lat = np.linspace(25.0, 33.0, 10)
+    dest_lon = np.linspace(-123, -98, 10)
+    regrid = src.get_bilinear_regridder_to(dest_lat, dest_lon)
+    src_lat = torch.broadcast_to(torch.tensor(src.lat), src.shape)
+    out_lat = regrid(src_lat)
 
+    assert torch.allclose(out_lat, torch.tensor(dest_lat))
+
+def test_regrid_2d():
     src = HRRR_CONUS_GRID
-    dest = HRRR_CONUS_GRID[1:-1, 1:-1]
-    regrid = src.get_bilinear_regridder_to(dest.lat, dest.lon)
-    lat = torch.broadcast_to(torch.tensor(src.lat), src.shape)
-    out = regrid(lat)
+    dest_lat, dest_lon = np.meshgrid(np.linspace(25.0, 33.0, 10), np.linspace(-123, -98, 10))
+    regrid = src.get_bilinear_regridder_to(dest_lat, dest_lon)
+    src_lat = torch.broadcast_to(torch.tensor(src.lat), src.shape)
+    out_lat = regrid(src_lat)
+
+    assert torch.allclose(out_lat, torch.tensor(dest_lat))
 
-    assert torch.allclose(out, lat[1:-1, 1:-1])

From 3b662997ccfcc0932c686b78493cff5fecb73cb8 Mon Sep 17 00:00:00 2001
From: Simon Byrne <sbyrne@nvidia.com>
Date: Fri, 27 Sep 2024 18:28:46 -0700
Subject: [PATCH 4/7] add function for computing angle

---
 earth2grid/lcc.py | 27 +++++++++++++++++++++------
 1 file changed, 21 insertions(+), 6 deletions(-)

diff --git a/earth2grid/lcc.py b/earth2grid/lcc.py
index ce60bf1..2448e51 100644
--- a/earth2grid/lcc.py
+++ b/earth2grid/lcc.py
@@ -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)
@@ -170,7 +181,7 @@ 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__()
@@ -178,7 +189,11 @@ def __init__(self, src: LambertConformalConicGrid, lat: np.ndarray, lon: np.ndar
         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)

From 1ccdb7ec528108267277efc7e918df77558725cc Mon Sep 17 00:00:00 2001
From: Simon Byrne <simonbyrne@gmail.com>
Date: Mon, 30 Sep 2024 20:29:19 +0000
Subject: [PATCH 5/7] reviewer comments

---
 earth2grid/lcc.py | 17 ++++++++++++-----
 tests/test_lcc.py | 10 +++++-----
 2 files changed, 17 insertions(+), 10 deletions(-)

diff --git a/earth2grid/lcc.py b/earth2grid/lcc.py
index 2448e51..fbba5fe 100644
--- a/earth2grid/lcc.py
+++ b/earth2grid/lcc.py
@@ -23,6 +23,12 @@
 except ImportError:
     pv = None
 
+__all__ = [
+    "LambertConformalConicProjection",
+    "LambertConformalConicGrid",
+    "HRRR_CONUS_PROJECTION",
+    "HRRR_CONUS_GRID",
+]
 
 
 class LambertConformalConicProjection:
@@ -76,7 +82,7 @@ def _theta(self, lon):
         return self.n * np.deg2rad(delta_lon)
 
 
-    def proj(self, lat, lon):
+    def project(self, lat, lon):
         """
         Compute the projected x,y from lat,lon.
         """
@@ -87,7 +93,7 @@ def proj(self, lat, lon):
         y = self.rho0 - rho * np.cos(theta)
         return x, y
 
-    def inv(self, x, y):
+    def inverse_project(self, x, y):
         """
         Compute the lat,lon from the projected x,y.
         """
@@ -99,6 +105,7 @@ def inv(self, x, y):
         return lat, lon
 
 # Projection used by HRRR CONUS (Continental US) data
+# https://rapidrefresh.noaa.gov/hrrr/HRRR_conus.domain.txt
 HRRR_CONUS_PROJECTION = LambertConformalConicProjection(
     lon0 = -97.5,
     lat0 = 38.5,
@@ -126,7 +133,7 @@ def __init__(self, projection: LambertConformalConicProjection, x, y):
     @property
     def lat_lon(self):
         mesh_x, mesh_y = np.meshgrid(self.x, self.y)
-        return self.projection.inv(mesh_x, mesh_y)
+        return self.projection.inverse_project(mesh_x, mesh_y)
 
     @property
     def lat(self):
@@ -169,7 +176,7 @@ def hrrr_conus_grid(ix0 = 0, iy0 = 0, nx = 1799, ny = 1059):
     # grid length (m)
     scale = 3000.0
     # coordinates on projected space
-    x0, y0 = HRRR_CONUS_PROJECTION.proj(lat0, lon0)
+    x0, y0 = HRRR_CONUS_PROJECTION.project(lat0, lon0)
 
     x = [x0 + i * scale for i in range(ix0, ix0+nx)]
     y = [y0 + i * scale for i in range(iy0, iy0+ny)]
@@ -186,7 +193,7 @@ class _RegridFromLCC(torch.nn.Module):
     def __init__(self, src: LambertConformalConicGrid, lat: np.ndarray, lon: np.ndarray):
         super().__init__()
 
-        x, y = src.projection.proj(lat, lon)
+        x, y = src.projection.project(lat, lon)
         
         self.shape = lat.shape
         self._bilinear = BilinearInterpolator(
diff --git a/tests/test_lcc.py b/tests/test_lcc.py
index 3f9764e..64207ba 100644
--- a/tests/test_lcc.py
+++ b/tests/test_lcc.py
@@ -17,7 +17,7 @@
 from earth2grid.lcc import HRRR_CONUS_GRID
 import numpy as np
 import torch
-from pytest import approx
+import pytest
 
 def test_grid_shape():
     assert HRRR_CONUS_GRID.lat.shape == HRRR_CONUS_GRID.shape 
@@ -37,13 +37,13 @@ def test_grid_shape():
 
 
 def test_grid_vals():
-    assert HRRR_CONUS_GRID.lat[0:400:100,0:400:100] == approx(lats)
-    assert HRRR_CONUS_GRID.lon[0:400:100,0:400:100] == approx(lons)
+    assert HRRR_CONUS_GRID.lat[0:400:100,0:400:100] == pytest.approx(lats)
+    assert HRRR_CONUS_GRID.lon[0:400:100,0:400:100] == pytest.approx(lons)
 
 def test_grid_slice():
     slice_grid = HRRR_CONUS_GRID[0:400:100,0:400:100]
-    assert slice_grid.lat == approx(lats)
-    assert slice_grid.lon == approx(lons)
+    assert slice_grid.lat == pytest.approx(lats)
+    assert slice_grid.lon == pytest.approx(lons)
 
 def test_regrid_1d():
     src = HRRR_CONUS_GRID

From bcf86421ce35aaa99157ea3925f80ebd89342f5b Mon Sep 17 00:00:00 2001
From: Simon Byrne <simonbyrne@gmail.com>
Date: Mon, 30 Sep 2024 21:32:36 +0000
Subject: [PATCH 6/7] make BilinearInterpolator work with multidimensional
 destinations

---
 earth2grid/_regrid.py |  4 ++--
 earth2grid/lcc.py     | 30 ++++++++----------------------
 tests/test_lcc.py     |  3 +--
 3 files changed, 11 insertions(+), 26 deletions(-)

diff --git a/earth2grid/_regrid.py b/earth2grid/_regrid.py
index 5e9e5e8..564a7cd 100644
--- a/earth2grid/_regrid.py
+++ b/earth2grid/_regrid.py
@@ -177,8 +177,8 @@ def forward(self, z: torch.Tensor):
         interpolated = torch.full(
             [self.mask.numel(), zrs.shape[1]], fill_value=self.fill_value, dtype=z.dtype, device=z.device
         )
-        interpolated.masked_scatter_(self.mask.unsqueeze(-1), output)
-        interpolated = interpolated.T.view(*shape, self.mask.numel())
+        interpolated.masked_scatter_(self.mask.view(-1,1), output)
+        interpolated = interpolated.T.view(*shape, *self.mask.shape)
         return interpolated
 
 
diff --git a/earth2grid/lcc.py b/earth2grid/lcc.py
index fbba5fe..abac978 100644
--- a/earth2grid/lcc.py
+++ b/earth2grid/lcc.py
@@ -153,7 +153,14 @@ def __getitem__(self, idxs):
 
     def get_bilinear_regridder_to(self, lat: np.ndarray, lon: np.ndarray):
         """Get regridder to the specified lat and lon points"""
-        return _RegridFromLCC(self, lat, lon)
+
+        x, y = self.projection.project(lat, lon)
+        
+        return BilinearInterpolator(
+            x_coords = torch.from_numpy(self.x), 
+            y_coords = torch.from_numpy(self.y), 
+            x_query = torch.from_numpy(x),
+            y_query = torch.from_numpy(y))
 
     def visualize(self, data):
         raise NotImplementedError()
@@ -185,24 +192,3 @@ def hrrr_conus_grid(ix0 = 0, iy0 = 0, nx = 1799, ny = 1059):
 
 # Grid used by HRRR CONUS (Continental US) data
 HRRR_CONUS_GRID = hrrr_conus_grid()
-
-
-class _RegridFromLCC(torch.nn.Module):
-    """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.project(lat, lon)
-        
-        self.shape = lat.shape
-        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)
-        return out.view(out.shape[:-1] + self.shape)
-
diff --git a/tests/test_lcc.py b/tests/test_lcc.py
index 64207ba..73dcacb 100644
--- a/tests/test_lcc.py
+++ b/tests/test_lcc.py
@@ -57,10 +57,9 @@ def test_regrid_1d():
 
 def test_regrid_2d():
     src = HRRR_CONUS_GRID
-    dest_lat, dest_lon = np.meshgrid(np.linspace(25.0, 33.0, 10), np.linspace(-123, -98, 10))
+    dest_lat, dest_lon = np.meshgrid(np.linspace(25.0, 33.0, 10), np.linspace(-123, -98, 12))
     regrid = src.get_bilinear_regridder_to(dest_lat, dest_lon)
     src_lat = torch.broadcast_to(torch.tensor(src.lat), src.shape)
     out_lat = regrid(src_lat)
 
     assert torch.allclose(out_lat, torch.tensor(dest_lat))
-

From 643cedc8ea359263cd8ca9526941e2cdfc0d5fd2 Mon Sep 17 00:00:00 2001
From: Simon Byrne <simonbyrne@gmail.com>
Date: Mon, 30 Sep 2024 21:35:29 +0000
Subject: [PATCH 7/7] apply linter

---
 earth2grid/_regrid.py |  6 +++---
 earth2grid/lcc.py     | 46 ++++++++++++++++++++-----------------------
 tests/test_lcc.py     | 42 +++++++++++++++++++++++++--------------
 3 files changed, 51 insertions(+), 43 deletions(-)

diff --git a/earth2grid/_regrid.py b/earth2grid/_regrid.py
index 564a7cd..daba8bb 100644
--- a/earth2grid/_regrid.py
+++ b/earth2grid/_regrid.py
@@ -44,7 +44,7 @@ def forward(self, z):
         weight = self.weight.view(-1, p)
 
         # using embedding bag is 2x faster on cpu and 4x on gpu.
-        output = torch.nn.functional.embedding_bag(index, zrs, per_sample_weights=weight, mode='sum')
+        output = torch.nn.functional.embedding_bag(index, zrs, per_sample_weights=weight, mode="sum")
         output = output.T.view(*shape, -1)
         return output.reshape(list(shape) + output_shape)
 
@@ -173,11 +173,11 @@ def forward(self, z: torch.Tensor):
         *shape, y, x = z.shape
         zrs = z.view(-1, y * x).T
         # using embedding bag is 2x faster on cpu and 4x on gpu.
-        output = torch.nn.functional.embedding_bag(self.index, zrs, per_sample_weights=self.weights, mode='sum')
+        output = torch.nn.functional.embedding_bag(self.index, zrs, per_sample_weights=self.weights, mode="sum")
         interpolated = torch.full(
             [self.mask.numel(), zrs.shape[1]], fill_value=self.fill_value, dtype=z.dtype, device=z.device
         )
-        interpolated.masked_scatter_(self.mask.view(-1,1), output)
+        interpolated.masked_scatter_(self.mask.view(-1, 1), output)
         interpolated = interpolated.T.view(*shape, *self.mask.shape)
         return interpolated
 
diff --git a/earth2grid/lcc.py b/earth2grid/lcc.py
index abac978..55146d0 100644
--- a/earth2grid/lcc.py
+++ b/earth2grid/lcc.py
@@ -34,7 +34,7 @@
 class LambertConformalConicProjection:
     def __init__(self, lat0: float, lon0: float, lat1: float, lat2: float, radius: float):
         """
-        
+
         Args:
             lat0: latitude of origin (degrees)
             lon0: longitude of origin (degrees)
@@ -50,7 +50,6 @@ def __init__(self, lat0: float, lon0: float, lat1: float, lat2: float, radius: f
         self.lat2 = lat2
         self.radius = radius
 
-
         c1 = np.cos(np.deg2rad(lat1))
         c2 = np.cos(np.deg2rad(lat2))
         t1 = np.tan(np.pi / 4 + np.deg2rad(lat1) / 2)
@@ -58,8 +57,8 @@ def __init__(self, lat0: float, lon0: float, lat1: float, lat2: float, radius: f
 
         if np.abs(lat1 - lat2) < 1e-8:
             self.n = np.sin(np.deg2rad(lat1))
-        else:            
-            self.n = np.log(c1/c2) / np.log(t2/t1)
+        else:
+            self.n = np.log(c1 / c2) / np.log(t2 / t1)
 
         self.RF = radius * c1 * np.power(t1, self.n) / self.n
         self.rho0 = self._rho(lat0)
@@ -78,17 +77,16 @@ def _theta(self, lon):
         """
         # center about reference longitude
         delta_lon = lon - self.lon0
-        delta_lon = delta_lon - np.round(delta_lon/360) * 360 # convert to [-180, 180]
+        delta_lon = delta_lon - np.round(delta_lon / 360) * 360  # convert to [-180, 180]
         return self.n * np.deg2rad(delta_lon)
 
-
     def project(self, lat, lon):
         """
         Compute the projected x,y from lat,lon.
         """
         rho = self._rho(lat)
         theta = self._theta(lon)
-    
+
         x = rho * np.sin(theta)
         y = self.rho0 - rho * np.cos(theta)
         return x, y
@@ -99,26 +97,21 @@ def inverse_project(self, x, y):
         """
         rho = np.hypot(x, self.rho0 - y)
         theta = np.arctan2(x, self.rho0 - y)
-        
-        lat = np.rad2deg(2 * np.arctan(np.power(self.RF/rho, 1/self.n))) - 90
+
+        lat = np.rad2deg(2 * np.arctan(np.power(self.RF / rho, 1 / self.n))) - 90
         lon = self.lon0 + np.rad2deg(theta / self.n)
         return lat, lon
 
+
 # Projection used by HRRR CONUS (Continental US) data
 # https://rapidrefresh.noaa.gov/hrrr/HRRR_conus.domain.txt
-HRRR_CONUS_PROJECTION = LambertConformalConicProjection(
-    lon0 = -97.5,
-    lat0 = 38.5,
-    lat1 = 38.5,
-    lat2 = 38.5,
-    radius = 6371229.0
-)
+HRRR_CONUS_PROJECTION = LambertConformalConicProjection(lon0=-97.5, lat0=38.5, lat1=38.5, lat2=38.5, radius=6371229.0)
 
 
 class LambertConformalConicGrid(base.Grid):
     # nothing here is specific to the projection, so could be shared by any projected rectilinear grid
     def __init__(self, projection: LambertConformalConicProjection, x, y):
-        """        
+        """
         Args:
             projection: LambertConformalConicProjection object
             x: range of x values
@@ -155,12 +148,13 @@ def get_bilinear_regridder_to(self, lat: np.ndarray, lon: np.ndarray):
         """Get regridder to the specified lat and lon points"""
 
         x, y = self.projection.project(lat, lon)
-        
+
         return BilinearInterpolator(
-            x_coords = torch.from_numpy(self.x), 
-            y_coords = torch.from_numpy(self.y), 
-            x_query = torch.from_numpy(x),
-            y_query = torch.from_numpy(y))
+            x_coords=torch.from_numpy(self.x),
+            y_coords=torch.from_numpy(self.y),
+            x_query=torch.from_numpy(x),
+            y_query=torch.from_numpy(y),
+        )
 
     def visualize(self, data):
         raise NotImplementedError()
@@ -176,7 +170,8 @@ def to_pyvista(self):
         grid = pv.StructuredGrid(x, y, z)
         return grid
 
-def hrrr_conus_grid(ix0 = 0, iy0 = 0, nx = 1799, ny = 1059):
+
+def hrrr_conus_grid(ix0=0, iy0=0, nx=1799, ny=1059):
     # coordinates of point in top-left corner
     lat0 = 21.138123
     lon0 = 237.280472
@@ -185,10 +180,11 @@ def hrrr_conus_grid(ix0 = 0, iy0 = 0, nx = 1799, ny = 1059):
     # coordinates on projected space
     x0, y0 = HRRR_CONUS_PROJECTION.project(lat0, lon0)
 
-    x = [x0 + i * scale for i in range(ix0, ix0+nx)]
-    y = [y0 + i * scale for i in range(iy0, iy0+ny)]
+    x = [x0 + i * scale for i in range(ix0, ix0 + nx)]
+    y = [y0 + i * scale for i in range(iy0, iy0 + ny)]
 
     return LambertConformalConicGrid(HRRR_CONUS_PROJECTION, x, y)
 
+
 # Grid used by HRRR CONUS (Continental US) data
 HRRR_CONUS_GRID = hrrr_conus_grid()
diff --git a/tests/test_lcc.py b/tests/test_lcc.py
index 73dcacb..536068a 100644
--- a/tests/test_lcc.py
+++ b/tests/test_lcc.py
@@ -13,38 +13,49 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-#%%
-from earth2grid.lcc import HRRR_CONUS_GRID
+# %%
 import numpy as np
-import torch
 import pytest
+import torch
+
+from earth2grid.lcc import HRRR_CONUS_GRID
+
 
 def test_grid_shape():
-    assert HRRR_CONUS_GRID.lat.shape == HRRR_CONUS_GRID.shape 
+    assert HRRR_CONUS_GRID.lat.shape == HRRR_CONUS_GRID.shape
     assert HRRR_CONUS_GRID.lon.shape == HRRR_CONUS_GRID.shape
 
-lats = np.array([
+
+lats = np.array(
+    [
         [21.138123, 21.801926, 22.393631, 22.911015],
-        [23.636763, 24.328228, 24.944668, 25.48374 ],
+        [23.636763, 24.328228, 24.944668, 25.48374],
         [26.155672, 26.875362, 27.517046, 28.078257],
-        [28.69017 , 29.438608, 30.106009, 30.68978 ]])
+        [28.69017, 29.438608, 30.106009, 30.68978],
+    ]
+)
 
-lons = np.array([
-        [-122.71953 , -120.03195 , -117.304596, -114.54146 ],
-        [-123.491356, -120.72898 , -117.92319 , -115.07828 ],
-        [-124.310524, -121.469505, -118.58098 , -115.649574],
-        [-125.181404, -122.25762 , -119.28173 , -116.25871 ]])
+lons = np.array(
+    [
+        [-122.71953, -120.03195, -117.304596, -114.54146],
+        [-123.491356, -120.72898, -117.92319, -115.07828],
+        [-124.310524, -121.469505, -118.58098, -115.649574],
+        [-125.181404, -122.25762, -119.28173, -116.25871],
+    ]
+)
 
 
 def test_grid_vals():
-    assert HRRR_CONUS_GRID.lat[0:400:100,0:400:100] == pytest.approx(lats)
-    assert HRRR_CONUS_GRID.lon[0:400:100,0:400:100] == pytest.approx(lons)
+    assert HRRR_CONUS_GRID.lat[0:400:100, 0:400:100] == pytest.approx(lats)
+    assert HRRR_CONUS_GRID.lon[0:400:100, 0:400:100] == pytest.approx(lons)
+
 
 def test_grid_slice():
-    slice_grid = HRRR_CONUS_GRID[0:400:100,0:400:100]
+    slice_grid = HRRR_CONUS_GRID[0:400:100, 0:400:100]
     assert slice_grid.lat == pytest.approx(lats)
     assert slice_grid.lon == pytest.approx(lons)
 
+
 def test_regrid_1d():
     src = HRRR_CONUS_GRID
     dest_lat = np.linspace(25.0, 33.0, 10)
@@ -55,6 +66,7 @@ def test_regrid_1d():
 
     assert torch.allclose(out_lat, torch.tensor(dest_lat))
 
+
 def test_regrid_2d():
     src = HRRR_CONUS_GRID
     dest_lat, dest_lon = np.meshgrid(np.linspace(25.0, 33.0, 10), np.linspace(-123, -98, 12))