diff --git a/regional_mom6/regional_mom6.py b/regional_mom6/regional_mom6.py index fbeb362d..a3dd6ad0 100644 --- a/regional_mom6/regional_mom6.py +++ b/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": { diff --git a/regional_mom6/utils.py b/regional_mom6/utils.py new file mode 100644 index 00000000..36749635 --- /dev/null +++ b/regional_mom6/utils.py @@ -0,0 +1,5 @@ +import numpy as np + + +def vecdot(v1, v2): + return np.sum(v1 * v2, axis=-1) diff --git a/tests/test_grid_generation.py b/tests/test_grid_generation.py index d39bed9e..ea0f7662 100644 --- a/tests/test_grid_generation.py +++ b/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