Merge pull request #4893 from chrishavlin/geographic_coords_convert_u…

…nits

yt/geometry/coordinates/geographic_coordinates.py

@@ -1,4 +1,5 @@
 import numpy as np
+import unyt
 
 from yt.utilities.lib.pixelization_routines import pixelize_cartesian, pixelize_cylinder
 
@@ -318,27 +319,28 @@ def convert_from_cartesian(self, coord):
 
     def convert_to_cartesian(self, coord):
         offset, factor = self._retrieve_radial_offset()
+
         if isinstance(coord, np.ndarray) and len(coord.shape) > 1:
             rad = self.axis_id[self.radial_axis]
             lon = self.axis_id["longitude"]
             lat = self.axis_id["latitude"]
             r = factor * coord[:, rad] + offset
-            theta = (90.0 - coord[:, lat]) * np.pi / 180
+            colatitude = _latitude_to_colatitude(coord[:, lat])
             phi = coord[:, lon] * np.pi / 180
             nc = np.zeros_like(coord)
             # r, theta, phi
-            nc[:, lat] = np.cos(phi) * np.sin(theta) * r
-            nc[:, lon] = np.sin(phi) * np.sin(theta) * r
-            nc[:, rad] = np.cos(theta) * r
+            nc[:, lat] = np.cos(phi) * np.sin(colatitude) * r
+            nc[:, lon] = np.sin(phi) * np.sin(colatitude) * r
+            nc[:, rad] = np.cos(colatitude) * r
         else:
             a, b, c = coord
-            theta = (90.0 - b) * np.pi / 180
+            colatitude = _latitude_to_colatitude(b)
             phi = a * np.pi / 180
             r = factor * c + offset
             nc = (
-                np.cos(phi) * np.sin(theta) * r,
-                np.sin(phi) * np.sin(theta) * r,
-                np.cos(theta) * r,
+                np.cos(phi) * np.sin(colatitude) * r,
+                np.sin(phi) * np.sin(colatitude) * r,
+                np.cos(colatitude) * r,
             )
         return nc
 
@@ -568,3 +570,16 @@ def sanitize_width(self, axis, width, depth):
             outermost = factor * self.ds.domain_left_edge[ri] + offset
             width = [outermost, 2.0 * outermost]
         return width
+
+
+def _latitude_to_colatitude(lat_vals):
+    # convert latitude to theta, accounting for units,
+    # including the case where the units are code_length
+    # due to how yt stores the domain_center units for
+    # geographic coordinates.
+    if isinstance(lat_vals, unyt.unyt_array):
+        if lat_vals.units.dimensions == unyt.dimensions.length:
+            return (90.0 - lat_vals.d) * np.pi / 180.0
+        ninety = unyt.unyt_quantity(90.0, "degree")
+        return (ninety - lat_vals).to("radian")
+    return (90 - lat_vals) * np.pi / 180.0

yt/geometry/coordinates/tests/test_geographic_coordinates.py

@@ -2,6 +2,7 @@
 
 import numpy as np
 import pytest
+import unyt
 from numpy.testing import assert_equal
 
 from yt.testing import assert_rel_equal, fake_amr_ds
@@ -117,3 +118,30 @@ def test_geographic_conversions(geometry):
     z = xyz[:, 2]
     assert z[0] == r_val
     assert z[1] == -r_val
+
+
+@pytest.mark.parametrize("geometry", ("geographic", "internal_geographic"))
+def test_geographic_conversions_with_units(geometry):
+    ds = fake_amr_ds(geometry=geometry)
+
+    # _sanitize_center will give all values in 'code_length'
+    coords = ds.arr(np.zeros((2, 3)), "code_length")
+    xyz_u = ds.coordinates.convert_to_cartesian(coords)
+    xyz = ds.coordinates.convert_to_cartesian(coords.d)
+    assert_equal(xyz, xyz_u)
+
+    coords = ds.arr(np.zeros((3,)), "code_length")
+    xyz_u = ds.coordinates.convert_to_cartesian(coords)
+    xyz = ds.coordinates.convert_to_cartesian(coords.d)
+    assert_equal(xyz, xyz_u)
+
+    # also check that if correct units are supplied, the
+    # result has dimensions of length.
+    coords = [
+        ds.arr(np.zeros((10,)), "degree"),
+        ds.arr(np.zeros((10,)), "degree"),
+        ds.arr(np.linspace(0, 100, 10), "code_length"),
+    ]
+    xyz = ds.coordinates.convert_to_cartesian(coords)
+    for dim in xyz:
+        assert dim.units.dimensions == unyt.dimensions.length