Merge branch 'main' into test-domain-generation

README.md

@@ -2,7 +2,7 @@
 
 *Python package for automatic generation of regional configurations for the [Modular Ocean Model 6](https://github.com/mom-ocean/MOM6).*
 
-[![Repo status](https://www.repostatus.org/badges/latest/active.svg?style=flat-square)](https://www.repostatus.org/#active) [![License](https://img.shields.io/badge/License-MIT-blue.svg?style=flat-square)](https://mit-license.org) [![codecov](https://codecov.io/gh/COSIMA/regional-mom6/branch/master/graph/badge.svg?token=7OEZ1UZRY4)](https://codecov.io/gh/COSIMA/regional-mom6) [![Documentation Status](https://readthedocs.org/projects/regional-mom6/badge/?version=latest)](https://regional-mom6.readthedocs.io/en/latest/?badge=latest) [![Code style: black](https://img.shields.io/badge/code%20style-black-000000.svg)](https://github.com/psf/black)
+[![Repo status](https://www.repostatus.org/badges/latest/active.svg?style=flat-square)](https://www.repostatus.org/#active) [![License](https://img.shields.io/badge/License-MIT-blue.svg?style=flat-square)](https://mit-license.org) [![codecov](https://codecov.io/gh/COSIMA/regional-mom6/branch/main/graph/badge.svg?token=7OEZ1UZRY4)](https://codecov.io/gh/COSIMA/regional-mom6) [![Documentation Status](https://readthedocs.org/projects/regional-mom6/badge/?version=latest)](https://regional-mom6.readthedocs.io/en/latest/?badge=latest) [![Code style: black](https://img.shields.io/badge/code%20style-black-000000.svg)](https://github.com/psf/black)
 
 Users just need to provide some information about where, when, and how big their domain is and also where raw input forcing files are. The package sorts out all the boring details and creates a set of MOM6-friendly input files along with setup directories ready to go! 
 

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