diff --git a/yt/geometry/coordinates/geographic_coordinates.py b/yt/geometry/coordinates/geographic_coordinates.py index 3c95349f05..50cb4841f6 100644 --- a/yt/geometry/coordinates/geographic_coordinates.py +++ b/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 diff --git a/yt/geometry/coordinates/tests/test_geographic_coordinates.py b/yt/geometry/coordinates/tests/test_geographic_coordinates.py index 20dde8c9ed..5a200d487b 100644 --- a/yt/geometry/coordinates/tests/test_geographic_coordinates.py +++ b/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