Contributions to #34 (#57)

I added alternative functions for calculation of toa irradiance and
border masks using similar input and output as in the existing ones.
Also removed need for pyproj by reformulating projections using cartopy.

create_forcings.py

@@ -16,7 +16,6 @@
 import xarray as xr
 from affine import Affine
 from anemoi.datasets.grids import cutout_mask
-from pyproj import Proj, Transformer
 from rasterio.features import rasterize
 from scipy.ndimage import binary_dilation, center_of_mass
 
@@ -181,6 +180,142 @@ def create_boundary_mask(
     return interior_mask, boundary_mask
 
 
+
+def create_boundary_mask_alt(
+    lonlat_state,
+    lonlat_boundary,
+    crs_state,
+    boundary_min_distance,
+    boundary_max_distance,
+    overlap=False
+):
+    """
+    Returns a mask such that lonlat_boundary[msk] gives the boundary points
+    satisfying the requirements to be within a distance of 
+    boundary_max_distance from the state area while also not being closer to
+    it than boundary_min_distance.
+    If overap = True then also the points within the state area are included. 
+    In this case boundary_min_distance is set to zero.
+
+    Args:
+        lonlat_state: A NumPy array of shape (2, sy, sx)
+        lonlat_boundary: A NumPy array of shape (2, by, bx)
+        crs_state: Instance of cartopy class CRS
+        boundary_min_distance: Minimum distance [m]
+        boundary_max_distance: Maximum distance [m]
+        overlap: Boolean
+
+    Returns:
+        msk_interior: A boolean (integer) xarray of shape (y, x) indicating 
+        the boundary points within the interior.
+        msk_boundary: A boolean (integer) xarray of shape (y, x) indicating
+        the queried boundary points.
+    """
+
+    import cartopy.crs as ccrs, numpy as np
+
+    def distance_points_to_line(points, line_start, line_end):
+        """
+        Calculates distances from points in an array to a line in 2D.
+
+        Args:
+            points: A NumPy array of shape (N, 2) representing N points (x, y).
+            line_start: Array with start point (x1, y1) of the line.
+            line_end: An array with the end point (x2, y2) of the line.
+
+        Returns:
+            A NumPy array of shape (N,) containing the distances for each point.
+        """
+
+        # Line direction vector
+        line_direction = line_end - line_start
+
+        # Vector from each point to line start
+        point_to_start = points - line_start
+
+        # Calculate projection scalars for all points
+        t = np.dot(point_to_start, line_direction) / np.sum(line_direction**2)
+
+        # constrain to line segment
+        t = np.fmax(0, np.fmin(1, t))
+
+        # Project all points onto the line
+        projections = line_start + t[:, np.newaxis] * line_direction
+
+        # Distances between points and their projections
+        distances = np.sqrt(np.sum((points - projections)**2, axis=1))
+
+        return distances
+
+
+    by, bx = lonlat_boundary[0].shape
+    lon_state, lat_state = lonlat_state[0].ravel(), lonlat_state[1].ravel()
+    lon_boundary, lat_boundary = lonlat_boundary[0].ravel(), lonlat_boundary[1].ravel()
+
+    # regular state grid coordinates
+    crs_boundary = ccrs.Geodetic() 
+    coords_state = crs_state.transform_points(crs_boundary, lon_state, lat_state)
+
+    # state corner points
+    xm_state, xM_state = np.amin(coords_state[:,0]), np.amax(coords_state[:,0])
+    ym_state, yM_state = np.amin(coords_state[:,1]), np.amax(coords_state[:,1])
+    x_ll_state, y_ll_state = xm_state, ym_state
+    x_ur_state, y_ur_state = xM_state, yM_state
+    # ll ul ur lr
+    pts_corners = np.array([[x_ll_state,y_ll_state], [x_ll_state,y_ur_state], [x_ur_state,y_ur_state], [x_ur_state,y_ll_state]])
+
+    # boundary points in state area projection
+    coords_boundary = crs_state.transform_points(crs_boundary, lon_boundary, lat_boundary)
+    # find boundary points outside the state area
+    msk_outside = ~((coords_boundary[:,0] > xm_state) & (coords_boundary[:,0] < xM_state) & (coords_boundary[:,1] > ym_state) & (coords_boundary[:,1] < yM_state))
+
+    # boundary points outside
+    pts_outside = np.vstack((coords_boundary[msk_outside,0], coords_boundary[msk_outside,1])).T
+
+    # distances from points outside to left, top, right and bottom borders
+    d_l = distance_points_to_line(pts_outside, pts_corners[0], pts_corners[1])
+    d_t = distance_points_to_line(pts_outside, pts_corners[1], pts_corners[2])
+    d_r = distance_points_to_line(pts_outside, pts_corners[2], pts_corners[3])
+    d_b = distance_points_to_line(pts_outside, pts_corners[3], pts_corners[0])
+
+    if overlap:
+        boundary_min_distance = 0
+
+    min_d, max_d = boundary_min_distance, boundary_max_distance
+    msk_min = (d_l > min_d) & (d_t > min_d) & (d_r > min_d) & (d_b > min_d) 
+    msk_max = (d_l < max_d) | (d_t < max_d) | (d_r < max_d) | (d_b < max_d)
+
+    msk_border = msk_min & msk_max
+
+    # construct mask to apply to input boundary data
+    msk_boundary = np.zeros(lon_boundary.shape, dtype='bool')
+    msk_boundary[np.where(msk_outside)[0][msk_border]] = True
+
+    msk_interior = np.zeros(lon_boundary.shape, dtype='bool')
+    if overlap:
+        # add boundary points inside
+        msk_interior = ((coords_boundary[:,0] >= xm_state) & (coords_boundary[:,0] <= xM_state) & (coords_boundary[:,1] >= ym_state) & (coords_boundary[:,1] <= yM_state))
+        msk_boundary[msk_interior] = True
+
+    msk_interior = xr.DataArray(
+        msk_interior.reshape((by,bx)).astype(int),
+        dims=("y", "x"),
+        coords={
+            "lat": (("y", "x"), lat_boundary.reshape((by,bx))),
+            "lon": (("y", "x"), lon_boundary.reshape((by,bx))),
+        },
+    )
+    msk_boundary = xr.DataArray(
+        msk_boundary.reshape((by,bx)).astype(int),
+        dims=("y", "x"),
+        coords={
+            "lat": (("y", "x"), lat_boundary.reshape((by,bx))),
+            "lon": (("y", "x"), lon_boundary.reshape((by,bx))),
+        },
+    )
+    return msk_interior, msk_boundary
+
+
 def generate_toa_radiation_forcing(ds, lonlat):
     """
     Pre-compute all static features related to the grid nodes
@@ -214,6 +349,41 @@ def generate_toa_radiation_forcing(ds, lonlat):
     return toa_radiation
 
 
+def generate_toa_radiation_forcing_approximation(ds, lonlat):
+    """
+    Calculates approximate TOA irradiance (instantaneous values [W*m**-2])
+    """
+
+    # solar constant
+    E0 = 1366
+
+    dt = ds.time.dt
+    lon_m = lonlat[0][np.newaxis, :]
+    lat_m = lonlat[1][np.newaxis, :]
+    day = dt.dayofyear.values[:, np.newaxis, np.newaxis]
+    hr_utc = dt.hour.values[:, np.newaxis, np.newaxis]
+    # Eq. 1.6.1a in Solar Engineering of Thermal Processes 4th ed.
+    dec = np.pi/180 * 23.45 * np.sin(2*np.pi * (284+day) / 365) 
+    hr_lst = hr_utc + lon_m / 15
+    hr_angle = 15 * (hr_lst - 12)
+    # Eq. 1.6.2 with beta=0 in Solar Engineering of Thermal Processes 4th ed.
+    cos_sza = np.sin(lat_m * np.pi / 180) * np.sin(dec) + np.cos(lat_m * np.pi / 180) * np.cos(dec) * np.cos(hr_angle * np.pi / 180)
+    toa_radiation = np.fmax(0, E0 * cos_sza) 
+    toa_radiation = toa_radiation.transpose(0, 2, 1)
+
+    toa_radiation = xr.DataArray(
+        toa_radiation,
+        dims=("time", "x", "y"),
+        coords={
+            "time": ds.time,
+            "lat": (("y", "x"), lonlat[1]),
+            "lon": (("y", "x"), lonlat[0]),
+        },
+    )
+    return toa_radiation
+
+
+
 def main():
     """Main function for creating the datetime forcing and boundary mask."""
     parser = argparse.ArgumentParser()
@@ -224,6 +394,8 @@ def main():
     parser.add_argument("--tempdir", type=str, default="./shps")
     parser.add_argument("--seconds_in_year", type=int, default=31536000)
     parser.add_argument("--boundary_thickness", type=int, default=30)
+    parser.add_argument("--boundary_min_distance", type=int, default=15000)
+    parser.add_argument("--boundary_max_distance", type=int, default=150000)
     parser.add_argument("--overlap", type=bool, default=False)
     parser.add_argument("--plot", type=bool, default=True)
     args = parser.parse_args()
@@ -238,24 +410,23 @@ def main():
         "proj": "lcc",
     }
 
+    globe = ccrs.Globe(semimajor_axis=lambert_proj_params["a"], semiminor_axis=lambert_proj_params["b"])
     lambert_proj = ccrs.LambertConformal(
         central_longitude=lambert_proj_params["lon_0"],
         central_latitude=lambert_proj_params["lat_0"],
         standard_parallels=(
             lambert_proj_params["lat_1"],
             lambert_proj_params["lat_2"],
         ),
+        globe=globe,
     )
 
+
     # Make sure the example data is available
     xy_state = np.load("data/meps_example/static/nwp_xy.npy")
-    latlong_proj = Proj(proj="latlong", datum="WGS84")
-    transformer = Transformer.from_proj(lambert_proj, latlong_proj)
-    x_state = xy_state[0]
-    y_state = xy_state[1]
-    # pylint: disable=unpacking-non-sequence
-    lon_state, lat_state = transformer.transform(x_state, y_state)
-    lonlat_state = np.stack((lon_state, lat_state), axis=0)
+    latlong_proj = ccrs.Geodetic() 
+    lonlatz_state = latlong_proj.transform_points(lambert_proj, xy_state[0], xy_state[1])
+    lonlat_state = np.stack((lonlatz_state[:,:,0], lonlatz_state[:,:,1]), axis=0)
 
     ds_meps = np.load(
         "data/meps_example/samples/train/nwp_2022040100_mbr000.npy"
@@ -314,6 +485,15 @@ def main():
         f"TOA radiation saved to {args.zarr_path} and boundary_{args.zarr_path}"
     )
 
+    # alternative boundary masks
+    interior_mask, boundary_mask = create_boundary_mask_alt(
+        lonlat_state, lonlat_boundary, lambert_proj, args.boundary_min_distance, args.boundary_max_distance, args.overlap
+    )
+
+    # alternative TOA
+    toa_radtiation_boundary = generate_toa_radiation_forcing_approximation(ds_boundary, lonlat_boundary)
+
+
     if args.plot:
 
         _, ax = plt.subplots(2, 2, figsize=(10, 8))