From 6e21a5f468ef85dbc35c6b14ad342f460c486d57 Mon Sep 17 00:00:00 2001 From: Jo Bovy Date: Thu, 20 Jul 2023 14:46:18 -0400 Subject: [PATCH] Allow Xsun/Zsun/vsun to be arrays in x,y,z,vx,vy,vz -> galcen conversions --- galpy/util/coords.py | 139 ++++++++++++++++---------------- tests/test_coords.py | 188 +++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 256 insertions(+), 71 deletions(-) diff --git a/galpy/util/coords.py b/galpy/util/coords.py index 67ebd32d5..524237e34 100644 --- a/galpy/util/coords.py +++ b/galpy/util/coords.py @@ -1094,9 +1094,9 @@ def XYZ_to_galcenrect(X, Y, Z, Xsun=1.0, Zsun=0.0, _extra_rot=True): Z - Z - Xsun - cylindrical distance to the GC + Xsun - cylindrical distance to the GC (can be array of same length as R) - Zsun - Sun's height above the midplane + Zsun - Sun's height above the midplane (can be array of same length as R) _extra_rot= (True) if True, perform an extra tiny rotation to align the Galactocentric coordinate frame with astropy's definition @@ -1112,14 +1112,21 @@ def XYZ_to_galcenrect(X, Y, Z, Xsun=1.0, Zsun=0.0, _extra_rot=True): 2018-04-18 - Tweaked to be consistent with astropy's Galactocentric frame - Bovy (UofT) + 2023-07-23 - Allowed Xsun/Zsun to be arrays - Bovy (UofT) + """ if _extra_rot: X, Y, Z = numpy.dot(galcen_extra_rot, numpy.array([X, Y, Z])) dgc = numpy.sqrt(Xsun**2.0 + Zsun**2.0) costheta, sintheta = Xsun / dgc, Zsun / dgc - return numpy.dot( + zero = numpy.zeros_like(costheta) + one = numpy.ones_like(costheta) + return numpy.einsum( + "ijk,jk->ik" + if isinstance(costheta, numpy.ndarray) and costheta.ndim > 0 + else "ij,jk->ik", numpy.array( - [[costheta, 0.0, -sintheta], [0.0, 1.0, 0.0], [sintheta, 0.0, costheta]] + [[costheta, zero, -sintheta], [zero, one, zero], [sintheta, zero, costheta]] ), numpy.array([-X + dgc, Y, numpy.sign(Xsun) * Z]), ).T @@ -1163,10 +1170,13 @@ def galcenrect_to_XYZ(X, Y, Z, Xsun=1.0, Zsun=0.0, _extra_rot=True): """ dgc = numpy.sqrt(Xsun**2.0 + Zsun**2.0) costheta, sintheta = Xsun / dgc, Zsun / dgc - if isinstance(Xsun, numpy.ndarray): - zero = numpy.zeros(len(Xsun)) - one = numpy.ones(len(Xsun)) - Carr = numpy.rollaxis( + zero = numpy.zeros_like(costheta) + one = numpy.ones_like(costheta) + out = ( + numpy.einsum( + "ijk,jk->ik" + if isinstance(costheta, numpy.ndarray) and costheta.ndim > 0 + else "ij,jk->ik", numpy.array( [ [-costheta, zero, -sintheta], @@ -1174,22 +1184,10 @@ def galcenrect_to_XYZ(X, Y, Z, Xsun=1.0, Zsun=0.0, _extra_rot=True): [-numpy.sign(Xsun) * sintheta, zero, numpy.sign(Xsun) * costheta], ] ), - 2, - ) - out = (Carr * numpy.array([[X, X, X], [Y, Y, Y], [Z, Z, Z]]).T).sum( - -1 - ) + numpy.array([dgc, zero, zero]).T - else: - out = numpy.dot( - numpy.array( - [ - [-costheta, 0.0, -sintheta], - [0.0, 1.0, 0.0], - [-numpy.sign(Xsun) * sintheta, 0.0, numpy.sign(Xsun) * costheta], - ] - ), numpy.array([X, Y, Z]), - ).T + numpy.array([dgc, 0.0, 0.0]) + ).T + + numpy.array([dgc, zero, zero]).T + ) if _extra_rot: return numpy.dot(galcen_extra_rot.T, out.T).T else: @@ -1325,9 +1323,9 @@ def XYZ_to_galcencyl(X, Y, Z, Xsun=1.0, Zsun=0.0, _extra_rot=True): Z - Z - Xsun - cylindrical distance to the GC + Xsun - cylindrical distance to the GC (can be array of same length as R) - Zsun - Sun's height above the midplane + Zsun - Sun's height above the midplane (can be array of same length as R) _extra_rot= (True) if True, perform an extra tiny rotation to align the Galactocentric coordinate frame with astropy's definition @@ -1339,6 +1337,8 @@ def XYZ_to_galcencyl(X, Y, Z, Xsun=1.0, Zsun=0.0, _extra_rot=True): 2010-09-24 - Written - Bovy (NYU) + 2023-07-19 - Allowed Xsun/Zsun to be arrays - Bovy (UofT) + """ XYZ = numpy.atleast_2d( XYZ_to_galcenrect(X, Y, Z, Xsun=Xsun, Zsun=Zsun, _extra_rot=_extra_rot) @@ -1403,11 +1403,11 @@ def vxvyvz_to_galcenrect( vz - W - vsun - velocity of the sun in the GC frame ndarray[3] + vsun - velocity of the sun in the GC frame ndarray[3] (can also be array of same length as vx; shape [3,N]) - Xsun - cylindrical distance to the GC + Xsun - cylindrical distance to the GC (can be array of same length as vXg) - Zsun - Sun's height above the midplane + Zsun - Sun's height above the midplane (can be array of same length as vXg) _extra_rot= (True) if True, perform an extra tiny rotation to align the Galactocentric coordinate frame with astropy's definition @@ -1423,17 +1423,31 @@ def vxvyvz_to_galcenrect( 2018-04-18 - Tweaked to be consistent with astropy's Galactocentric frame - Bovy (UofT) + 2023-07-23 - Allowed Xsun/Zsun/vsun to be arrays- Bovy (UofT) + """ if _extra_rot: vx, vy, vz = numpy.dot(galcen_extra_rot, numpy.array([vx, vy, vz])) dgc = numpy.sqrt(Xsun**2.0 + Zsun**2.0) costheta, sintheta = Xsun / dgc, Zsun / dgc - return numpy.dot( - numpy.array( - [[costheta, 0.0, -sintheta], [0.0, 1.0, 0.0], [sintheta, 0.0, costheta]] - ), - numpy.array([-vx, vy, numpy.sign(Xsun) * vz]), - ).T + numpy.array(vsun) + zero = numpy.zeros_like(costheta) + one = numpy.ones_like(costheta) + return ( + numpy.einsum( + "ijk,jk->ik" + if isinstance(costheta, numpy.ndarray) and costheta.ndim > 0 + else "ij,jk->ik", + numpy.array( + [ + [costheta, zero, -sintheta], + [zero, one, zero], + [sintheta, zero, costheta], + ] + ), + numpy.array([-vx, vy, numpy.sign(Xsun) * vz]), + ).T + + numpy.array(vsun).T + ) @scalarDecorator @@ -1473,11 +1487,11 @@ def vxvyvz_to_galcencyl( Z - Z in Galactocentric rectangular coordinates - vsun - velocity of the sun in the GC frame ndarray[3] + vsun - velocity of the sun in the GC frame ndarray[3] (can also be array of same length as vx; shape [3,N]) - Xsun - cylindrical distance to the GC + Xsun - cylindrical distance to the GC (can be array of same length as vXg) - Zsun - Sun's height above the midplane + Zsun - Sun's height above the midplane (can be array of same length as vXg) galcen - if True, then X,Y,Z are in cylindrical Galactocentric coordinates rather than rectangular coordinates @@ -1491,6 +1505,8 @@ def vxvyvz_to_galcencyl( 2010-09-24 - Written - Bovy (NYU) + 2023-07-19 - Allowed Xsun/Zsun/vsun to be arrays- Bovy (UofT) + """ vxyz = vxvyvz_to_galcenrect( vx, vy, vz, vsun=vsun, Xsun=Xsun, Zsun=Zsun, _extra_rot=_extra_rot @@ -1546,40 +1562,21 @@ def galcenrect_to_vxvyvz( """ dgc = numpy.sqrt(Xsun**2.0 + Zsun**2.0) costheta, sintheta = Xsun / dgc, Zsun / dgc - if isinstance(Xsun, numpy.ndarray): - zero = numpy.zeros(len(Xsun)) - one = numpy.ones(len(Xsun)) - Carr = numpy.rollaxis( - numpy.array( - [ - [-costheta, zero, -sintheta], - [zero, one, zero], - [-numpy.sign(Xsun) * sintheta, zero, numpy.sign(Xsun) * costheta], - ] - ), - 2, - ) - out = ( - Carr - * numpy.array( - [ - [vXg - vsun[0], vXg - vsun[0], vXg - vsun[0]], - [vYg - vsun[1], vYg - vsun[1], vYg - vsun[1]], - [vZg - vsun[2], vZg - vsun[2], vZg - vsun[2]], - ] - ).T - ).sum(-1) - else: - out = numpy.dot( - numpy.array( - [ - [-costheta, 0.0, -sintheta], - [0.0, 1.0, 0.0], - [-numpy.sign(Xsun) * sintheta, 0.0, numpy.sign(Xsun) * costheta], - ] - ), - numpy.array([vXg - vsun[0], vYg - vsun[1], vZg - vsun[2]]), - ).T + zero = numpy.zeros_like(costheta) + one = numpy.ones_like(costheta) + out = numpy.einsum( + "ijk,jk->ik" + if isinstance(costheta, numpy.ndarray) and costheta.ndim > 0 + else "ij,jk->ik", + numpy.array( + [ + [-costheta, zero, -sintheta], + [zero, one, zero], + [-numpy.sign(Xsun) * sintheta, zero, numpy.sign(Xsun) * costheta], + ] + ), + numpy.array([vXg - vsun[0], vYg - vsun[1], vZg - vsun[2]]), + ).T if _extra_rot: return numpy.dot(galcen_extra_rot.T, out.T).T else: diff --git a/tests/test_coords.py b/tests/test_coords.py index 91251b46f..11aa1239c 100644 --- a/tests/test_coords.py +++ b/tests/test_coords.py @@ -765,6 +765,27 @@ def test_XYZ_to_galcenrect_negXsun(): ), "XYZ_to_galcenrect conversion did not work as expected for negative Xsun" +def test_XYZ_to_galcenrect_vecSun(): + X, Y, Z = ( + numpy.array([1.0, 2.0]), + numpy.array([3.0, 3.0]), + numpy.array([-2.0, -2.0]), + ) + gcXYZ = coords.XYZ_to_galcenrect( + X, Y, Z, Xsun=numpy.array([1.0, -2.0]), Zsun=numpy.array([0.0, 0.0]) + ) + assert numpy.all( + numpy.fabs(gcXYZ[:, 0]) < 10.0**-5.0 + ), "XYZ_to_galcenrect conversion did not work as expected" + assert numpy.all( + numpy.fabs(gcXYZ[:, 1] - 3.0) < 10.0**-5.0 + ), "XYZ_to_galcenrect conversion did not work as expected" + assert numpy.all( + numpy.fabs(gcXYZ[:, 2] + 2.0) < 10.0**-5.0 + ), "XYZ_to_galcenrect conversion did not work as expected" + return None + + def test_lbd_to_galcenrect_galpyvsastropy(): # Test that galpy's transformations agree with astropy's import astropy.units as u @@ -990,6 +1011,27 @@ def test_XYZ_to_galcencyl(): return None +def test_XYZ_to_galcencyl_vecSun(): + X, Y, Z = ( + numpy.array([5.0, 4.0]), + numpy.array([4.0, 4.0]), + numpy.array([-2.0, -2.0]), + ) + gcRpZ = coords.XYZ_to_galcencyl( + X, Y, Z, Xsun=numpy.array([8.0, 7.0]), Zsun=numpy.array([0.0, 0.0]) + ) + assert numpy.all( + numpy.fabs(gcRpZ[:, 0] - 5.0) < 10.0**-5.0 + ), "XYZ_to_galcencyl conversion did not work as expected" + assert numpy.all( + numpy.fabs(gcRpZ[:, 1] - numpy.arctan(4.0 / 3.0)) < 10.0**-5.0 + ), "XYZ_to_galcencyl conversion did not work as expected" + assert numpy.all( + numpy.fabs(gcRpZ[:, 2] + 2.0) < 10.0**-4.8 + ), "XYZ_to_galcencyl conversion did not work as expected" + return None + + def test_galcencyl_to_XYZ(): gcR, gcp, gcZ = 5.0, numpy.arctan(4.0 / 3.0), 2.0 XYZ = coords.galcencyl_to_XYZ(gcR, gcp, gcZ, Xsun=8.0, Zsun=0.0) @@ -1138,6 +1180,79 @@ def test_vxvyvz_to_galcenrect_negXsun(): return None +def test_vxvyvz_to_galcenrect_vecXsun(): + vx, vy, vz = ( + numpy.array([10.0, 10.0]), + numpy.array([-20.0, -20.0]), + numpy.array([30.0, 30.0]), + ) + vgc = coords.vxvyvz_to_galcenrect( + vx, + vy, + vz, + vsun=[-5.0, 10.0, 5.0], + Xsun=numpy.array([1.1, 1.0]), + Zsun=numpy.array([0.0, 0.0]), + ) + assert numpy.all( + numpy.fabs(vgc[:, 0] + 15.0) < 10.0**-4.0 + ), "vxvyvz_to_galcenrect conversion did not work as expected" + assert numpy.all( + numpy.fabs(vgc[:, 1] + 10.0) < 10.0**-4.0 + ), "vxvyvz_to_galcenrect conversion did not work as expected" + assert numpy.all( + numpy.fabs(vgc[:, 2] - 35.0) < 10.0**-4.0 + ), "vxvyvz_to_galcenrect conversion did not work as expected" + return None + + +def test_vxvyvz_to_galcenrect_vecvsun(): + vx, vy, vz = ( + numpy.array([10.0, 5.0]), + numpy.array([-20.0, -10.0]), + numpy.array([30.0, 25.0]), + ) + vgc = coords.vxvyvz_to_galcenrect( + vx, vy, vz, vsun=numpy.array([[-5.0, 10.0, 5.0], [-10.0, 0.0, 10.0]]).T + ) + assert numpy.all( + numpy.fabs(vgc[:, 0] + 15.0) < 10.0**-4.0 + ), "vxvyvz_to_galcenrect conversion did not work as expected" + assert numpy.all( + numpy.fabs(vgc[:, 1] + 10.0) < 10.0**-4.0 + ), "vxvyvz_to_galcenrect conversion did not work as expected" + assert numpy.all( + numpy.fabs(vgc[:, 2] - 35.0) < 10.0**-4.0 + ), "vxvyvz_to_galcenrect conversion did not work as expected" + return None + + +def test_vxvyvz_to_galcenrect_vecXsunvsun(): + vx, vy, vz = ( + numpy.array([10.0, 5.0]), + numpy.array([-20.0, -10.0]), + numpy.array([30.0, 25.0]), + ) + vgc = coords.vxvyvz_to_galcenrect( + vx, + vy, + vz, + Xsun=numpy.array([1.1, 1.0]), + Zsun=numpy.array([0.0, 0.0]), + vsun=numpy.array([[-5.0, 10.0, 5.0], [-10.0, 0.0, 10.0]]).T, + ) + assert numpy.all( + numpy.fabs(vgc[:, 0] + 15.0) < 10.0**-4.0 + ), "vxvyvz_to_galcenrect conversion did not work as expected" + assert numpy.all( + numpy.fabs(vgc[:, 1] + 10.0) < 10.0**-4.0 + ), "vxvyvz_to_galcenrect conversion did not work as expected" + assert numpy.all( + numpy.fabs(vgc[:, 2] - 35.0) < 10.0**-4.0 + ), "vxvyvz_to_galcenrect conversion did not work as expected" + return None + + def test_vrpmllpmbb_to_galcenrect_galpyvsastropy(): # Only run this for astropy>3 if not _APY3: @@ -1438,6 +1553,79 @@ def test_galcenrect_to_vxvyvz_asInverse(): return None +def test_galcenrect_to_vxvyvz_vecXsun(): + vxg, vyg, vzg = ( + numpy.array([-15.0, -15.0]), + numpy.array([-10.0, -10.0]), + numpy.array([35.0, 35.0]), + ) + vxyz = coords.galcenrect_to_vxvyvz( + vxg, + vyg, + vzg, + vsun=[-5.0, 10.0, 5.0], + Xsun=numpy.array([1.1, 1.0]), + Zsun=numpy.array([0.0, 0.0]), + ) + assert numpy.all( + numpy.fabs(vxyz[:, 0] - 10.0) < 10.0**-4.0 + ), "galcenrect_to_vxvyvz conversion did not work as expected" + assert numpy.all( + numpy.fabs(vxyz[:, 1] + 20.0) < 10.0**-4.0 + ), "galcenrect_to_vxvyvz conversion did not work as expected" + assert numpy.all( + numpy.fabs(vxyz[:, 2] - 30.0) < 10.0**-4.0 + ), "galcenrect_to_vxvyvz conversion did not work as expected" + return None + + +def test_galcenrect_to_vxvyvz_vecvsun(): + vxg, vyg, vzg = ( + numpy.array([-15.0, -5.0]), + numpy.array([-10.0, -20.0]), + numpy.array([35.0, 32.5]), + ) + vxyz = coords.galcenrect_to_vxvyvz( + vxg, vyg, vzg, vsun=numpy.array([[-5.0, 10.0, 5.0], [5.0, 0.0, 2.5]]).T + ) + assert numpy.all( + numpy.fabs(vxyz[:, 0] - 10.0) < 10.0**-4.0 + ), "galcenrect_to_vxvyvz conversion did not work as expected" + assert numpy.all( + numpy.fabs(vxyz[:, 1] + 20.0) < 10.0**-4.0 + ), "galcenrect_to_vxvyvz conversion did not work as expected" + assert numpy.all( + numpy.fabs(vxyz[:, 2] - 30.0) < 10.0**-4.0 + ), "galcenrect_to_vxvyvz conversion did not work as expected" + return None + + +def test_galcenrect_to_vxvyvz_vecXsunvsun(): + vxg, vyg, vzg = ( + numpy.array([-15.0, -5.0]), + numpy.array([-10.0, -20.0]), + numpy.array([35.0, 32.5]), + ) + vxyz = coords.galcenrect_to_vxvyvz( + vxg, + vyg, + vzg, + vsun=numpy.array([[-5.0, 10.0, 5.0], [5.0, 0.0, 2.5]]).T, + Xsun=numpy.array([1.1, 1.0]), + Zsun=numpy.array([0.0, 0.0]), + ) + assert numpy.all( + numpy.fabs(vxyz[:, 0] - 10.0) < 10.0**-4.0 + ), "galcenrect_to_vxvyvz conversion did not work as expected" + assert numpy.all( + numpy.fabs(vxyz[:, 1] + 20.0) < 10.0**-4.0 + ), "galcenrect_to_vxvyvz conversion did not work as expected" + assert numpy.all( + numpy.fabs(vxyz[:, 2] - 30.0) < 10.0**-4.0 + ), "galcenrect_to_vxvyvz conversion did not work as expected" + return None + + def test_galcencyl_to_vxvyvz(): vr, vp, vz = -17.0, 6.0, 35.0 phi = numpy.arctan(4.0 / 3.0)