Merge branch 'main' into move_fre_execution

regional_mom6/regional_mom6.py

@@ -13,6 +13,7 @@
 from dask.diagnostics import ProgressBar
 import datetime as dt
 import warnings
+from .utils import vecdot
 
 warnings.filterwarnings("ignore")
 
@@ -21,7 +22,9 @@
     "motu_requests",
     "dz",
     "angle_between",
-    "quadilateral_area",
+    "latlon_to_cartesian",
+    "quadrilateral_area",
+    "quadrilateral_areas",
     "rectangular_hgrid",
     "experiment",
     "segment",
@@ -261,7 +264,6 @@ def dz(npoints, ratio, target_depth, min_dz=0.0001, tolerance=1):
 
     Returns:
         numpy.array: An array containing the thickness profile.
-
     """
 
     profile = min_dz + 0.5 * (np.abs(ratio) * min_dz - min_dz) * (
@@ -279,47 +281,73 @@ def dz(npoints, ratio, target_depth, min_dz=0.0001, tolerance=1):
     return dz(npoints, ratio, target_depth, min_dz * err_ratio)
 
 
-# Borrowed from grid tools (GFDL)
 def angle_between(v1, v2, v3):
-    """Returns angle v2-v1-v3 i.e betweeen v1-v2 and v1-v3."""
+    """Returns the angle v2-v1-v3 (in radians). That is the angle between vectors v1-v2 and v1-v3."""
+
+    v1xv2 = np.cross(v1, v2)
+    v1xv3 = np.cross(v1, v3)
+
+    cosangle = vecdot(v1xv2, v1xv3) / np.sqrt(
+        vecdot(v1xv2, v1xv2) * vecdot(v1xv3, v1xv3)
+    )
+
+    return np.arccos(cosangle)
+
+
+def quadrilateral_area(v1, v2, v3, v4):
+    """Returns area of a spherical quadrilateral on the unit sphere that
+    has vertices on 3-vectors `v1`, `v2`, `v3`, `v4` (counter-clockwise
+    orientation is implied). The area is computed via the excess of the
+    sum of the spherical angles of the quadrilateral from 2π."""
+
+    if not (
+        np.all(np.isclose(vecdot(v1, v1), vecdot(v2, v2)))
+        & np.all(np.isclose(vecdot(v1, v1), vecdot(v2, v2)))
+        & np.all(np.isclose(vecdot(v1, v1), vecdot(v3, v3)))
+        & np.all(np.isclose(vecdot(v1, v1), vecdot(v4, v4)))
+    ):
+        raise ValueError("vectors provided must have the same length")
 
-    # vector product between v1 and v2
-    px = v1[1] * v2[2] - v1[2] * v2[1]
-    py = v1[2] * v2[0] - v1[0] * v2[2]
-    pz = v1[0] * v2[1] - v1[1] * v2[0]
-    # vector product between v1 and v3
-    qx = v1[1] * v3[2] - v1[2] * v3[1]
-    qy = v1[2] * v3[0] - v1[0] * v3[2]
-    qz = v1[0] * v3[1] - v1[1] * v3[0]
+    R = np.sqrt(vecdot(v1, v1))
 
-    ddd = (px * px + py * py + pz * pz) * (qx * qx + qy * qy + qz * qz)
-    ddd = (px * qx + py * qy + pz * qz) / np.sqrt(ddd)
-    angle = np.arccos(ddd)
+    a1 = angle_between(v1, v2, v4)
+    a2 = angle_between(v2, v3, v1)
+    a3 = angle_between(v3, v4, v2)
+    a4 = angle_between(v4, v1, v3)
 
-    return angle
+    return (a1 + a2 + a3 + a4 - 2 * np.pi) * R**2
 
 
-# Borrowed from grid tools (GFDL)
-def quadilateral_area(lat, lon):
-    """Returns area of spherical quadrilaterals on the unit sphere that are formed
-    by constant latitude and longitude lines on the `lat`-`lon` grid provided."""
+def latlon_to_cartesian(lat, lon, R=1):
+    """Convert latitude-longitude (in degrees) to Cartesian coordinates on a sphere of radius `R`.
+    By default `R = 1`."""
 
-    # x, y, z are 3D coordinates on the unit sphere
-    x = np.cos(np.deg2rad(lat)) * np.cos(np.deg2rad(lon))
-    y = np.cos(np.deg2rad(lat)) * np.sin(np.deg2rad(lon))
-    z = np.sin(np.deg2rad(lat))
+    x = R * np.cos(np.deg2rad(lat)) * np.cos(np.deg2rad(lon))
+    y = R * np.cos(np.deg2rad(lat)) * np.sin(np.deg2rad(lon))
+    z = R * np.sin(np.deg2rad(lat))
 
-    c1 = (x[:-1, :-1], y[:-1, :-1], z[:-1, :-1])
-    c2 = (x[:-1, 1:], y[:-1, 1:], z[:-1, 1:])
-    c3 = (x[1:, 1:], y[1:, 1:], z[1:, 1:])
-    c4 = (x[1:, :-1], y[1:, :-1], z[1:, :-1])
+    return x, y, z
 
-    a1 = angle_between(c2, c1, c3)
-    a2 = angle_between(c3, c2, c4)
-    a3 = angle_between(c4, c3, c1)
-    a4 = angle_between(c1, c4, c2)
 
-    return a1 + a2 + a3 + a4 - 2.0 * np.pi
+def quadrilateral_areas(lat, lon, R=1):
+    """Returns area of spherical quadrilaterals on a sphere of radius `R`. By default, `R = 1`.
+    The quadrilaterals are formed by constant latitude and longitude lines on the `lat`-`lon` grid provided.
+
+    Args:
+        lat (array): Array of latitude points (in degrees)
+        lon (array): Array of longitude points (in degrees)
+
+    Returns:
+        areas (array): Array with the areas of the quadrilaterals defined by the
+                       `lat`-`lon` grid provided. If the `lat`-`lon` are `m x n`
+                       then `areas` is `(m-1) x (n-1)`.
+    """
+
+    coords = np.dstack(latlon_to_cartesian(lat, lon, R))
+
+    return quadrilateral_area(
+        coords[:-1, :-1, :], coords[:-1, 1:, :], coords[1:, 1:, :], coords[1:, :-1, :]
+    )
 
 
 def rectangular_hgrid(λ, φ):
@@ -346,18 +374,20 @@ def rectangular_hgrid(λ, φ):
 
     dλ = λ[1] - λ[0]  # assuming that longitude is uniformly spaced
 
-    # dx = R * cos(φ) * np.deg2rad(dλ) / 2. Note, we divide dy by 2 because we're on the supergrid
+    # dx = R * cos(φ) * np.deg2rad(dλ) / 2
+    # Note: division by 2 because we're on the supergrid
     dx = np.broadcast_to(
         R * np.cos(np.deg2rad(φ)) * np.deg2rad(dλ) / 2,
         (λ.shape[0] - 1, φ.shape[0]),
     ).T
 
-    # dy = R * np.deg2rad(dφ) / 2. Note, we divide dy by 2 because we're on the supergrid
+    # dy = R * np.deg2rad(dφ) / 2
+    # Note: division by 2 because we're on the supergrid
     dy = np.broadcast_to(R * np.deg2rad(np.diff(φ)) / 2, (λ.shape[0], φ.shape[0] - 1)).T
 
     lon, lat = np.meshgrid(λ, φ)
 
-    area = quadilateral_area(lat, lon) * R**2
+    area = quadrilateral_areas(lat, lon, R)
 
     attrs = {
         "tile": {

regional_mom6/utils.py

@@ -0,0 +1,5 @@
+import numpy as np
+
+
+def vecdot(v1, v2):
+    return np.sum(v1 * v2, axis=-1)

tests/test_grid_generation.py

@@ -1,38 +1,93 @@
 import numpy as np
 import pytest
 from regional_mom6 import angle_between
-from regional_mom6 import quadilateral_area
+from regional_mom6 import latlon_to_cartesian
+from regional_mom6 import quadrilateral_area
+from regional_mom6 import quadrilateral_areas
 from regional_mom6 import rectangular_hgrid
 import xarray as xr
 
 
+@pytest.mark.parametrize(
+    ("lat", "lon", "true_xyz"),
+    [
+        (0, 0, (1, 0, 0)),
+        (90, 0, (0, 0, 1)),
+        (0, 90, (0, 1, 0)),
+        (-90, 0, (0, 0, -1)),
+    ],
+)
+def test_latlon_to_cartesian(lat, lon, true_xyz):
+    assert np.isclose(latlon_to_cartesian(lat, lon), true_xyz).all()
+
+
 @pytest.mark.parametrize(
     ("v1", "v2", "v3", "true_angle"),
     [
         ([1, 0, 0], [0, 1, 0], [0, 0, 1], np.pi / 2),
         ([1, 0, 0], [1, 1, 0], [0, 1, 1], np.pi / 4),
+        ([1, 0, 0], [1, 1, 1], [0, 0, 1], np.pi / 4),
+        ([1, 1, 1], [1, 1, 0], [0, 1, 1], 2 * np.pi / 3),
     ],
 )
 def test_angle_between(v1, v2, v3, true_angle):
     assert np.isclose(angle_between(v1, v2, v3), true_angle)
 
 
+@pytest.mark.parametrize(
+    ("v1", "v2", "v3", "v4", "true_area"),
+    [
+        (
+            np.dstack(latlon_to_cartesian(0, 0)),
+            np.dstack(latlon_to_cartesian(0, 90)),
+            np.dstack(latlon_to_cartesian(90, 0)),
+            np.dstack(latlon_to_cartesian(0, -90)),
+            np.pi,
+        ),
+        (
+            np.dstack(latlon_to_cartesian(0, 0)),
+            np.dstack(latlon_to_cartesian(90, 0)),
+            np.dstack(latlon_to_cartesian(0, 90)),
+            np.dstack(latlon_to_cartesian(-90, 0)),
+            np.pi,
+        ),
+    ],
+)
+def test_quadrilateral_area(v1, v2, v3, v4, true_area):
+    assert np.isclose(quadrilateral_area(v1, v2, v3, v4), true_area)
+
+
+v1 = np.dstack(latlon_to_cartesian(0, 0, R=2))
+v2 = np.dstack(latlon_to_cartesian(90, 0, R=2))
+v3 = np.dstack(latlon_to_cartesian(0, 90, R=2))
+v4 = np.dstack(latlon_to_cartesian(-90, 0, R=2.1))
+
+
+def test_quadrilateral_area_exception():
+    with pytest.raises(ValueError) as excinfo:
+        quadrilateral_area(v1, v2, v3, v4)
+
+    assert str(excinfo.value) == "vectors provided must have the same length"
+
+
 # create a lat-lon mesh that covers 1/4 of the North Hemisphere
 lon1, lat1 = np.meshgrid(np.linspace(0, 90, 5), np.linspace(0, 90, 5))
+area1 = 1 / 8 * (4 * np.pi)
 
 # create a lat-lon mesh that covers 1/4 of the whole globe
 lon2, lat2 = np.meshgrid(np.linspace(-45, 45, 5), np.linspace(-90, 90, 5))
+area2 = 1 / 4 * (4 * np.pi)
 
 
 @pytest.mark.parametrize(
     ("lat", "lon", "true_area"),
     [
-        (lat1, lon1, 0.5 * np.pi),
-        (lat2, lon2, np.pi),
+        (lat1, lon1, area1),
+        (lat2, lon2, area2),
     ],
 )
-def test_quadilateral_area(lat, lon, true_area):
-    assert np.isclose(np.sum(quadilateral_area(lat, lon)), true_area)
+def test_quadrilateral_areas(lat, lon, true_area):
+    assert np.isclose(np.sum(quadrilateral_areas(lat, lon)), true_area)
 
 
 # a simple test that rectangular_hgrid runs without erroring