From fd7d5f315a60593f9d6ecc2359aa41d369c62ea9 Mon Sep 17 00:00:00 2001 From: chavlin Date: Thu, 2 May 2024 11:57:29 -0500 Subject: [PATCH 1/3] BUG: geographic coords conver_to_cartesian with units --- .../coordinates/geographic_coordinates.py | 19 ++++++++++-- .../tests/test_geographic_coordinates.py | 29 +++++++++++++++++++ 2 files changed, 46 insertions(+), 2 deletions(-) diff --git a/yt/geometry/coordinates/geographic_coordinates.py b/yt/geometry/coordinates/geographic_coordinates.py index 3c95349f05..414a613ffd 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,12 +319,13 @@ 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 + theta = _lat_to_theta(coord[:, lat]) phi = coord[:, lon] * np.pi / 180 nc = np.zeros_like(coord) # r, theta, phi @@ -332,7 +334,7 @@ def convert_to_cartesian(self, coord): nc[:, rad] = np.cos(theta) * r else: a, b, c = coord - theta = (90.0 - b) * np.pi / 180 + theta = _lat_to_theta(b) phi = a * np.pi / 180 r = factor * c + offset 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 _lat_to_theta(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 dd45927e6c..12c0fc6db2 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 @@ -118,3 +119,31 @@ 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 From 1cea43b695c9ec75f41de0184254300525ccb228 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Thu, 2 May 2024 17:31:11 +0000 Subject: [PATCH 2/3] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- yt/geometry/coordinates/tests/test_geographic_coordinates.py | 1 - 1 file changed, 1 deletion(-) diff --git a/yt/geometry/coordinates/tests/test_geographic_coordinates.py b/yt/geometry/coordinates/tests/test_geographic_coordinates.py index 12c0fc6db2..566b94f0d8 100644 --- a/yt/geometry/coordinates/tests/test_geographic_coordinates.py +++ b/yt/geometry/coordinates/tests/test_geographic_coordinates.py @@ -123,7 +123,6 @@ def test_geographic_conversions(geometry): @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' From 433c81fd2e0405176efae615affd1ef7c1b67ba3 Mon Sep 17 00:00:00 2001 From: chrishavlin Date: Mon, 6 May 2024 15:14:07 -0500 Subject: [PATCH 3/3] rename theta to colatitude in geo coord handler --- .../coordinates/geographic_coordinates.py | 18 +++++++++--------- 1 file changed, 9 insertions(+), 9 deletions(-) diff --git a/yt/geometry/coordinates/geographic_coordinates.py b/yt/geometry/coordinates/geographic_coordinates.py index 414a613ffd..50cb4841f6 100644 --- a/yt/geometry/coordinates/geographic_coordinates.py +++ b/yt/geometry/coordinates/geographic_coordinates.py @@ -325,22 +325,22 @@ def convert_to_cartesian(self, coord): lon = self.axis_id["longitude"] lat = self.axis_id["latitude"] r = factor * coord[:, rad] + offset - theta = _lat_to_theta(coord[:, lat]) + 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 = _lat_to_theta(b) + 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 @@ -572,7 +572,7 @@ def sanitize_width(self, axis, width, depth): return width -def _lat_to_theta(lat_vals): +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