Contributions to #34

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))