From 9bb38517f1ac92599c4a1861c0abf5cce5bc8614 Mon Sep 17 00:00:00 2001 From: Angus Gibson Date: Wed, 6 Sep 2023 01:19:35 +0000 Subject: [PATCH] Implement quadrilateral_areas using broadcasting We need to define a vecdot function (which is available in the Array API but not the main numpy API, sadly), which will compute the dot product of the last axis of two arrays. --- regional_mom6/regional_mom6.py | 50 ++++++++++++++-------------------- regional_mom6/utils.py | 4 +++ tests/test_grid_generation.py | 28 +++++++++---------- 3 files changed, 39 insertions(+), 43 deletions(-) create mode 100644 regional_mom6/utils.py diff --git a/regional_mom6/regional_mom6.py b/regional_mom6/regional_mom6.py index 0cf55ac7..1154b755 100644 --- a/regional_mom6/regional_mom6.py +++ b/regional_mom6/regional_mom6.py @@ -5,7 +5,7 @@ import dask.bag as db import numpy as np import xarray as xr -import xesmf as xe +# import xesmf as xe import subprocess from scipy.ndimage import binary_fill_holes import netCDF4 @@ -13,6 +13,7 @@ from dask.diagnostics import ProgressBar import datetime as dt import warnings +from .utils import vecdot warnings.filterwarnings("ignore") @@ -22,8 +23,8 @@ "dz", "angle_between", "latlon_to_cartesian", - "quadilateral_area", - "quadilateral_areas", + "quadrilateral_area", + "quadrilateral_areas", "rectangular_hgrid", "experiment", "segment", @@ -286,28 +287,28 @@ def angle_between(v1, v2, v3): v1xv2 = np.cross(v1, v2) v1xv3 = np.cross(v1, v3) - cosangle = np.dot(v1xv2, v1xv3) / np.sqrt( - np.dot(v1xv2, v1xv2) * np.dot(v1xv3, v1xv3) + cosangle = vecdot(v1xv2, v1xv3) / np.sqrt( + vecdot(v1xv2, v1xv2) * vecdot(v1xv3, v1xv3) ) return np.arccos(cosangle) -def quadilateral_area(v1, v2, v3, v4): +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.isclose(np.dot(v1, v1), np.dot(v2, v2)) - & np.isclose(np.dot(v1, v1), np.dot(v2, v2)) - & np.isclose(np.dot(v1, v1), np.dot(v3, v3)) - & np.isclose(np.dot(v1, v1), np.dot(v4, v4)) + 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") - R = np.sqrt(np.dot(v1, v1)) + R = np.sqrt(vecdot(v1, v1)) a1 = angle_between(v1, v2, v4) a2 = angle_between(v2, v3, v1) @@ -328,7 +329,7 @@ def latlon_to_cartesian(lat, lon, R=1): return x, y, z -def quadilateral_areas(lat, lon, R=1): +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. @@ -342,23 +343,14 @@ def quadilateral_areas(lat, lon, R=1): then `areas` is `(m-1) x (n-1)`. """ - x, y, z = latlon_to_cartesian(lat, lon, R) + coords = np.dstack(latlon_to_cartesian(lat, lon, R)) - nx, ny = np.shape(lat) - - areas = np.zeros((nx - 1, ny - 1)) - - for j in range(ny - 1): - for i in range(nx - 1): - # construct the 4 3-vectors v1, v2, v3, v4 that point to the vertices of the quadrilaterals - v1 = [x[i, j], y[i, j], z[i, j]] - v2 = [x[i, j + 1], y[i, j + 1], z[i, j + 1]] - v3 = [x[i + 1, j + 1], y[i + 1, j + 1], z[i + 1, j + 1]] - v4 = [x[i + 1, j], y[i + 1, j], z[i + 1, j]] - - areas[i, j] = quadilateral_area(v1, v2, v3, v4) - - return areas + return quadrilateral_area( + coords[:-1,:-1,:], + coords[:-1,1:,:], + coords[1:,1:,:], + coords[1:,:-1,:] + ) def rectangular_hgrid(λ, φ): @@ -398,7 +390,7 @@ def rectangular_hgrid(λ, φ): lon, lat = np.meshgrid(λ, φ) - area = quadilateral_areas(lat, lon, R) + 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..465edafe --- /dev/null +++ b/regional_mom6/utils.py @@ -0,0 +1,4 @@ +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 0be9f4a8..80fc3e91 100644 --- a/tests/test_grid_generation.py +++ b/tests/test_grid_generation.py @@ -2,8 +2,8 @@ import pytest from regional_mom6 import angle_between from regional_mom6 import latlon_to_cartesian -from regional_mom6 import quadilateral_area -from regional_mom6 import quadilateral_areas +from regional_mom6 import quadrilateral_area +from regional_mom6 import quadrilateral_areas from regional_mom6 import rectangular_hgrid import xarray as xr @@ -38,23 +38,23 @@ def test_angle_between(v1, v2, v3, true_angle): ("v1", "v2", "v3", "v4", "true_area"), [ ( - latlon_to_cartesian(0, 0), - latlon_to_cartesian(0, 90), - latlon_to_cartesian(90, 0), - latlon_to_cartesian(0, -90), + 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, ), ( - latlon_to_cartesian(0, 0), - latlon_to_cartesian(90, 0), - latlon_to_cartesian(0, 90), - latlon_to_cartesian(-90, 0), + 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_quadilateral_area(v1, v2, v3, v4, true_area): - assert np.isclose(quadilateral_area(v1, v2, v3, v4), true_area) +def test_quadrilateral_area(v1, v2, v3, v4, true_area): + assert np.isclose(quadrilateral_area(v1, v2, v3, v4), true_area) v1 = latlon_to_cartesian(0, 0, R=2) @@ -86,8 +86,8 @@ def test_quadilateral_area_exception(): (lat2, lon2, area2), ], ) -def test_quadilateral_areas(lat, lon, true_area): - assert np.isclose(np.sum(quadilateral_areas(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