diff --git a/.travis.yml b/.travis.yml index ccad088b6..67b423c07 100644 --- a/.travis.yml +++ b/.travis.yml @@ -23,9 +23,11 @@ env: # overidden underneath. They are defined here in order to save having # to repeat them for all configurations. - NUMPY_VERSION=stable - - ASTROPY_VERSION=stable + # - ASTROPY_VERSION=stable # HACK: to use mhvk's branch - SETUP_CMD='test' - CONDA_DEPENDENCIES='cython jinja2 scipy matplotlib pyyaml h5py sympy qt' + # HACK: to use mhvk's branch + - PIP_DEPENDENCIES='git+https://github.com/mhvk/astropy.git@representation-differentials' matrix: # Make sure that egg_info works without dependencies diff --git a/docs/coordinates/index.rst b/docs/coordinates/index.rst index 940331db2..92f9e52c2 100644 --- a/docs/coordinates/index.rst +++ b/docs/coordinates/index.rst @@ -29,7 +29,8 @@ The core functions in this subpackage provide support to: - Convert proper motions and radial velocities to Galactocentric, cartesian velocities. - Convert proper motions from/to ICRS to/from Galactic. -- Convert radial velocities from/to the Galactic Standard of Rest (GSR) to/from a barycentric frame. +- Convert radial velocities from/to the Galactic Standard of Rest (GSR) to/from + a barycentric frame. These functions work naturally with the :mod:`astropy.units` and :mod:`astropy.coordinates` subpackages. Handling positional transformations @@ -39,38 +40,45 @@ currently no support for transforming velocities in Astropy. The functions below attempt to bridge that gap as a temporary solution until support is added (planned for v1.2). -For example, to convert a spherical, heliocentric velocity (proper motion and radial -velocity) in an ICRS frame to a Galactocentric, cartesian velocity, we first have -to define an Astropy coordinate to specify the position of the object:: +For example, to convert a spherical, heliocentric velocity (proper motion and +radial velocity) in an ICRS frame to a Galactocentric, cartesian velocity, we +first have to define an Astropy coordinate to specify the position of the +object:: - >>> c = coord.SkyCoord(ra=100.68458*u.deg, dec=41.26917*u.deg, distance=1.1*u.kpc) + >>> c = coord.SkyCoord(ra=100.68458*u.deg, dec=41.26917*u.deg, + ... distance=1.1*u.kpc) Then pass this object in to the heliocentric to galactocentric conversion function, :func:`~gala.coordinates.vhel_to_gal`:: - >>> pm = [1.5, -1.7] * u.mas/u.yr - >>> rv = 151.1 * u.km/u.s - >>> gc.vhel_to_gal(c.icrs, pm=pm, rv=rv) - - -Because the input coordinate is given in the ICRS frame, the function assumes that -the proper motion is also in this frame, e.g., that the proper motion components are -:math:`(\mu_\alpha\cos\delta, \mu_\delta)`. If we instead passed in a coordinate in -the Galactic frame, the components are assumed to be :math:`(\mu_l\cos b, \mu_b)` :: - - >>> gc.vhel_to_gal(c.galactic, pm=pm, rv=rv) - - -The velocity transformation functions allow specifying the circular velocity at the Sun -(``vcirc``) and a 3-vector specifying the Sun's velocity with respect to the local -standard of rest (``vlsr``). Further customization of the Sun's location can be made via -the :class:`~astropy.coordinates.Galactocentric` frame attributes and passed in with the -keyword argument ``galactocentric_frame`` :: + >>> vhel = coord.SphericalDifferential(d_lon=1.5 * u.mas/u.yr, + ... d_lat=-1.7 * u.mas/u.yr, + ... d_distance=151.1 * u.km/u.s) + >>> gc.vhel_to_gal(c.icrs, vhel) + + +Because the input coordinate is given in the ICRS frame, the function assumes +that the proper motion is also in this frame, e.g., that the proper motion +components are :math:`(\mu_\alpha, \mu_\delta)`. If we instead passed in a +coordinate in the Galactic frame, the components are assumed to be +:math:`(\mu_l, \mu_b)` :: + + >>> gc.vhel_to_gal(c.galactic, vhel) + + +The velocity transformation functions allow specifying the circular velocity at +the Sun (``vcirc``) and a 3-vector specifying the Sun's velocity with respect to +the local standard of rest (``vlsr``). Further customization of the Sun's +location can be made via the :class:`~astropy.coordinates.Galactocentric` frame +attributes and passed in with the keyword argument ``galactocentric_frame`` :: >>> frame = coord.Galactocentric(z_sun=10.*u.pc, galcen_distance=8.3*u.kpc) - >>> gc.vhel_to_gal(c.icrs, pm=pm, rv=rv, galactocentric_frame=frame, + >>> gc.vhel_to_gal(c.icrs, vhel, galactocentric_frame=frame, ... vcirc=218*u.km/u.s, vlsr=[0.,0.,0.]*u.km/u.s) - + The inverse transformations are also available, with the function :func:`~gala.coordinates.vgal_to_hel`. Here, because the input coordinate is passed @@ -79,27 +87,21 @@ given in the ICRS frame :math:`(\mu_\alpha\cos\delta, \mu_\delta)`:: >>> xyz = coord.Galactocentric([11., 15, 25] * u.kpc) >>> vxyz = [121., 150., -75.] * u.km/u.s - >>> pm_ra,pm_dec,vr = gc.vgal_to_hel(xyz.transform_to(coord.ICRS), vxyz) - >>> pm_ra # doctest: +FLOAT_CMP - - >>> pm_dec # doctest: +FLOAT_CMP - - >>> vr # doctest: +FLOAT_CMP - + >>> c = xyz.transform_to(coord.ICRS) + >>> gc.vgal_to_hel(c, vxyz) # doctest: +FLOAT_CMP + Passing in coordinates in the Galactic frame means that the output proper motions will -instead be :math:`(\mu_l\cos b, \mu_b)` :: - - >>> pm_l,pm_b,vr = gc.vgal_to_hel(xyz.transform_to(coord.Galactic), vxyz) - >>> pm_l # doctest: +FLOAT_CMP - +instead be :math:`(\mu_l, \mu_b)` :: -All of these functions also work on arrays of coordinates and velocities, e.g.:: + >>> c = xyz.transform_to(coord.Galactic) + >>> gc.vgal_to_hel(c, vxyz) # doctest: +FLOAT_CMP + - >>> xyz = coord.Galactocentric(np.random.uniform(-20,20,size=(3,10)) * u.kpc) - >>> vxyz = np.random.uniform(-150,150,size=(3,10)) * u.km/u.s - >>> gc.vgal_to_hel(xyz.transform_to(coord.ICRS), vxyz) # doctest: +SKIP - ... +All of these functions also work on arrays of many coordinates and velocities as +well. Using gala.coordinates ====================== diff --git a/docs/dynamics/actionangle.rst b/docs/dynamics/actionangle.rst index a4f07e327..efa89c845 100644 --- a/docs/dynamics/actionangle.rst +++ b/docs/dynamics/actionangle.rst @@ -37,12 +37,15 @@ been executed:: >>> import astropy.coordinates as coord >>> import astropy.units as u - >>> import matplotlib.pyplot as pl + >>> import matplotlib.pyplot as plt >>> import numpy as np >>> import gala.potential as gp >>> import gala.dynamics as gd >>> from gala.units import galactic +For many more options for action calculation, see +`tact `_. + .. _tube-axisymmetric: A tube orbit in an axisymmetric potential @@ -74,25 +77,21 @@ We first create a potential and set up our initial conditions:: >>> pot = gp.LogarithmicPotential(v_c=150*u.km/u.s, q1=1., q2=1., q3=0.9, r_h=0, ... units=galactic) - >>> w0 = gd.CartesianPhaseSpacePosition(pos=[8, 0, 0.]*u.kpc, - ... vel=[75, 150, 50.]*u.km/u.s) + >>> w0 = gd.PhaseSpacePosition(pos=[8, 0, 0.]*u.kpc, + ... vel=[75, 150, 50.]*u.km/u.s) We will now integrate the orbit and plot it in the meridional plane:: >>> w = pot.integrate_orbit(w0, dt=0.5, n_steps=10000) - >>> cyl_pos, cyl_vel = w.represent_as(coord.CylindricalRepresentation) - >>> fig,ax = pl.subplots(1,1,figsize=(6,6)) - >>> ax.plot(cyl_pos.rho.to(u.kpc), cyl_pos.z.to(u.kpc), - ... marker=None, linestyle='-') # doctest: +SKIP - >>> ax.set_xlabel("R [kpc]") # doctest: +SKIP - >>> ax.set_ylabel("z [kpc]") # doctest: +SKIP + >>> cyl = w.represent_as('cylindrical') + >>> fig = cyl.plot(['rho', 'z'], linestyle='-') .. plot:: :align: center import astropy.coordinates as coord import astropy.units as u - import matplotlib.pyplot as pl + import matplotlib.pyplot as plt import numpy as np import gala.potential as gp import gala.dynamics as gd @@ -100,16 +99,12 @@ We will now integrate the orbit and plot it in the meridional plane:: pot = gp.LogarithmicPotential(v_c=150*u.km/u.s, q1=1., q2=1., q3=0.9, r_h=0, units=galactic) - w0 = gd.CartesianPhaseSpacePosition(pos=[8, 0, 0.]*u.kpc, - vel=[75, 150, 50.]*u.km/u.s) + w0 = gd.PhaseSpacePosition(pos=[8, 0, 0.]*u.kpc, + vel=[75, 150, 50.]*u.km/u.s) w = pot.integrate_orbit(w0, dt=0.5, n_steps=10000) - cyl_pos, cyl_vel = w.represent_as(coord.CylindricalRepresentation) - fig,ax = pl.subplots(1,1,figsize=(6,6)) - ax.plot(cyl_pos.rho.to(u.kpc).value, cyl_pos.z.to(u.kpc).value, - marker=None, linestyle='-') - ax.set_xlabel("R [kpc]") - ax.set_ylabel("z [kpc]") + cyl = w.represent_as('cylindrical') + cyl.plot(['rho', 'z'], linestyle='-') To solve for the actions in the true potential, we first compute the actions in a "toy" potential -- a potential in which we can compute the actions and angles @@ -132,7 +127,7 @@ angles would be perfectly straight lines with slope equal to the frequencies. Instead, the orbit is wobbly in the toy potential angles:: >>> toy_actions,toy_angles,toy_freqs = toy_potential.action_angle(w) - >>> fig,ax = pl.subplots(1,1,figsize=(5,5)) + >>> fig,ax = plt.subplots(1,1,figsize=(5,5)) >>> ax.plot(toy_angles[0], toy_angles[2], linestyle='none', marker=',') # doctest: +SKIP >>> ax.set_xlim(0,2*np.pi) # doctest: +SKIP >>> ax.set_ylim(0,2*np.pi) # doctest: +SKIP @@ -144,7 +139,7 @@ Instead, the orbit is wobbly in the toy potential angles:: import astropy.coordinates as coord import astropy.units as u - import matplotlib.pyplot as pl + import matplotlib.pyplot as plt import numpy as np import gala.potential as gp import gala.dynamics as gd @@ -152,13 +147,13 @@ Instead, the orbit is wobbly in the toy potential angles:: pot = gp.LogarithmicPotential(v_c=150*u.km/u.s, q1=1., q2=1., q3=0.9, r_h=0, units=galactic) - w0 = gd.CartesianPhaseSpacePosition(pos=[8, 0, 0.]*u.kpc, - vel=[75, 150, 50.]*u.km/u.s) + w0 = gd.PhaseSpacePosition(pos=[8, 0, 0.]*u.kpc, + vel=[75, 150, 50.]*u.km/u.s) w = pot.integrate_orbit(w0, dt=0.5, n_steps=10000) toy_potential = gd.fit_isochrone(w) actions,angles,freqs = toy_potential.action_angle(w) - fig,ax = pl.subplots(1,1,figsize=(5,5)) + fig,ax = plt.subplots(1,1,figsize=(5,5)) ax.plot(angles[0], angles[2], linestyle='none', marker=',') ax.set_xlim(0,2*np.pi) ax.set_ylim(0,2*np.pi) @@ -169,7 +164,7 @@ Instead, the orbit is wobbly in the toy potential angles:: This can also be seen in the value of the action variables, which are not time-independent in the toy potential:: - >>> fig,ax = pl.subplots(1,1) + >>> fig,ax = plt.subplots(1,1) >>> ax.plot(w.t, toy_actions[0], marker=None) # doctest: +SKIP >>> ax.set_xlabel(r"$t$ [Myr]") # doctest: +SKIP >>> ax.set_ylabel(r"$J_1$ [rad]") # doctest: +SKIP @@ -179,7 +174,7 @@ time-independent in the toy potential:: import astropy.coordinates as coord import astropy.units as u - import matplotlib.pyplot as pl + import matplotlib.pyplot as plt import numpy as np import gala.potential as gp import gala.dynamics as gd @@ -187,13 +182,13 @@ time-independent in the toy potential:: pot = gp.LogarithmicPotential(v_c=150*u.km/u.s, q1=1., q2=1., q3=0.9, r_h=0, units=galactic) - w0 = gd.CartesianPhaseSpacePosition(pos=[8, 0, 0.]*u.kpc, - vel=[75, 150, 50.]*u.km/u.s) + w0 = gd.PhaseSpacePosition(pos=[8, 0, 0.]*u.kpc, + vel=[75, 150, 50.]*u.km/u.s) w = pot.integrate_orbit(w0, dt=0.5, n_steps=10000) toy_potential = gd.fit_isochrone(w) actions,angles,freqs = toy_potential.action_angle(w) - fig,ax = pl.subplots(1,1) + fig,ax = plt.subplots(1,1) ax.plot(w.t, actions[0].to(u.km/u.s*u.kpc*u.Msun), marker=None) ax.set_xlabel(r"$t$ [Myr]") ax.set_ylabel(r"$J_1$ [kpc ${\rm M}_\odot$ km/s]") @@ -222,7 +217,7 @@ actions computed using this machinery:: >>> act_correction = nvecs.T[...,None] * result['Sn'][None,:,None] * np.cos(nvecs.dot(toy_angles))[None] >>> action_approx = toy_actions - 2*np.sum(act_correction, axis=1)*u.kpc**2/u.Myr*u.Msun >>> - >>> fig,ax = pl.subplots(1,1) + >>> fig,ax = plt.subplots(1,1) >>> ax.plot(w.t, toy_actions[0].to(u.km/u.s*u.kpc*u.Msun), marker=None, label='$J_1$') # doctest: +SKIP >>> ax.plot(w.t, action_approx[0].to(u.km/u.s*u.kpc*u.Msun), marker=None, label="$J_1'$") # doctest: +SKIP >>> ax.set_xlabel(r"$t$ [Myr]") # doctest: +SKIP @@ -234,7 +229,7 @@ actions computed using this machinery:: import astropy.coordinates as coord import astropy.units as u - import matplotlib.pyplot as pl + import matplotlib.pyplot as plt import numpy as np import gala.potential as gp import gala.dynamics as gd @@ -242,8 +237,8 @@ actions computed using this machinery:: pot = gp.LogarithmicPotential(v_c=150*u.km/u.s, q1=1., q2=1., q3=0.9, r_h=0, units=galactic) - w0 = gd.CartesianPhaseSpacePosition(pos=[8, 0, 0.]*u.kpc, - vel=[75, 150, 50.]*u.km/u.s) + w0 = gd.PhaseSpacePosition(pos=[8, 0, 0.]*u.kpc, + vel=[75, 150, 50.]*u.km/u.s) w = pot.integrate_orbit(w0, dt=0.5, n_steps=10000) toy_potential = gd.fit_isochrone(w) @@ -252,7 +247,7 @@ actions computed using this machinery:: nvecs = gd.generate_n_vectors(8, dx=1, dy=2, dz=2) act_correction = nvecs.T[...,None] * result['Sn'][None,:,None] * np.cos(nvecs.dot(toy_angles))[None] action_approx = toy_actions - 2*np.sum(act_correction, axis=1)*u.kpc**2/u.Myr*u.Msun - fig,ax = pl.subplots(1,1) + fig,ax = plt.subplots(1,1) ax.plot(w.t, toy_actions[0].to(u.km/u.s*u.kpc*u.Msun), marker=None, label='$J_1$') ax.plot(w.t, action_approx[0].to(u.km/u.s*u.kpc*u.Msun), marker=None, label="$J_1'$") ax.set_xlabel(r"$t$ [Myr]") @@ -298,7 +293,7 @@ and the same initial conditions as above: import astropy.coordinates as coord import astropy.units as u - import matplotlib.pyplot as pl + import matplotlib.pyplot as plt import numpy as np import gala.potential as gp import gala.dynamics as gd @@ -309,8 +304,8 @@ and the same initial conditions as above: units=galactic) # define initial conditions - w0 = gd.CartesianPhaseSpacePosition(pos=[8, 0, 0.]*u.kpc, - vel=[75, 150, 50.]*u.km/u.s) + w0 = gd.PhaseSpacePosition(pos=[8, 0, 0.]*u.kpc, + vel=[75, 150, 50.]*u.km/u.s) # integrate orbit w = pot.integrate_orbit(w0, dt=0.5, n_steps=10000) @@ -330,7 +325,7 @@ and the same initial conditions as above: act_correction = nvecs.T[...,None] * result['Sn'][None,:,None] * np.cos(nvecs.dot(toy_angles))[None] action_approx = toy_actions - 2*np.sum(act_correction, axis=1)*u.kpc**2/u.Myr*u.Msun - fig,axes = pl.subplots(3,1,figsize=(6,14)) + fig,axes = plt.subplots(3,1,figsize=(6,14)) for i,ax in enumerate(axes): ax.plot(w.t, toy_actions[i].to(u.km/u.s*u.kpc*u.Msun), marker=None, label='$J_{}$'.format(i+1)) diff --git a/docs/dynamics/index.rst b/docs/dynamics/index.rst index 327596efb..220b585cc 100644 --- a/docs/dynamics/index.rst +++ b/docs/dynamics/index.rst @@ -6,21 +6,8 @@ Dynamics (`gala.dynamics`) ******************************** -Introduction -============ - -This subpackage contains functions and classes useful for gravitational -dynamics. The fundamental objects used by many of the functions and utilities -in this and other subpackages are the `~gala.dynamics.PhaseSpacePosition` and -`~gala.dynamics.Orbit` subclasses. - -There are utilities for transforming orbits in phase-space to action-angle -coordinates, tools for visualizing and computing dynamical quantities from -orbits, tools to generate mock stellar streams, and tools useful for nonlinear -dynamics such as Lyapunov exponent estimation. - -For code blocks below and any pages linked below, I assume the following -imports have already been excuted:: +For code blocks below and any pages linked below, we'll assume the following +imports have already been executed:: >>> import astropy.units as u >>> import numpy as np @@ -28,25 +15,40 @@ imports have already been excuted:: >>> import gala.dynamics as gd >>> from gala.units import galactic +Introduction +============ + +This subpackage contains functions and classes useful for gravitational +dynamics. There are utilities for transforming orbits in phase-space to +action-angle coordinates, tools for visualizing and computing dynamical +quantities from orbits, tools to generate mock stellar streams, and tools useful +for nonlinear dynamics such as Lyapunov exponent estimation. + +The fundamental objects used by many of the functions and utilities in this and +other subpackages are the |psp| and |orb| classes. + Getting started: Working with orbits ==================================== -Some simple tools are provided for inspecting and plotting orbits. For example, -we'll start by integrating an orbit in Cartesian coordinates using the -:mod:`gala.potential` and :mod:`gala.integrate` subpackages:: +As a demonstration of how to use these objects, we'll start by integrating an +orbit using the :mod:`gala.potential` and :mod:`gala.integrate` subpackages:: >>> pot = gp.MiyamotoNagaiPotential(m=2.5E11*u.Msun, a=6.5*u.kpc, ... b=0.26*u.kpc, units=galactic) - >>> w0 = gd.CartesianPhaseSpacePosition(pos=[11., 0., 0.2]*u.kpc, - ... vel=[0., 200, 100]*u.km/u.s) + >>> w0 = gd.PhaseSpacePosition(pos=[11., 0., 0.2]*u.kpc, + ... vel=[0., 200, 100]*u.km/u.s) >>> orbit = pot.integrate_orbit(w0, dt=1., n_steps=1000) -This will integrate an orbit from the specified initial conditions (``w0``) and -return an orbit object. There are many useful methods of the -`~gala.dynamics.Orbit` subclasses and many functions that accept -`~gala.dynamics.Orbit` objects. For example, we can easily visualize the orbit -by plotting the time series in all projections using the -:meth:`~gala.dynamics.CartesianOrbit.plot` method:: +This numerically integrates an orbit from the specified initial conditions, +``w0``, and returns an |orb| object. By default, the position and velocity are +assumed to be Cartesian coordinates but other coordinate systems are supported +(see the :ref:`orbits-in-detail` and :ref:`nd-representations` pages for more +information). + +The |orb| object that is returned contains many useful methods, and can be +passed to many of the analysis functions implemented in Gala. For example, we +can easily visualize the orbit by plotting the time series in all Cartesian +projections using the :meth:`~gala.dynamics.Orbit.plot` method:: >>> fig = orbit.plot() @@ -58,35 +60,73 @@ by plotting the time series in all projections using the import gala.dynamics as gd from gala.units import galactic pot = gp.MiyamotoNagaiPotential(m=2.5E11, a=6.5, b=0.26, units=galactic) - w0 = gd.CartesianPhaseSpacePosition(pos=[11., 0., 0.2]*u.kpc, - vel=[0., 200, 100]*u.km/u.s) + w0 = gd.PhaseSpacePosition(pos=[11., 0., 0.2]*u.kpc, + vel=[0., 200, 100]*u.km/u.s) orbit = pot.integrate_orbit(w0, dt=1., n_steps=1000) fig = orbit.plot() -From this object, we can easily compute dynamical quantities such as the energy -or angular momentum (I take the 0th element because these functions return the -quantities computed at every timestep):: +Or, we can visualize the orbit in just one projection of some transformed +coordinate representation, for example, cylindrical radius :math:`\rho` and +:math:`z`:: + + >>> fig = orbit.represent_as('cylindrical').plot(['rho', 'z']) - >>> orbit.energy()[0] # doctest: +FLOAT_CMP - +.. plot:: + :align: center + + import astropy.units as u + import gala.potential as gp + import gala.dynamics as gd + from gala.units import galactic + pot = gp.MiyamotoNagaiPotential(m=2.5E11, a=6.5, b=0.26, units=galactic) + w0 = gd.PhaseSpacePosition(pos=[11., 0., 0.2]*u.kpc, + vel=[0., 200, 100]*u.km/u.s) + orbit = pot.integrate_orbit(w0, dt=1., n_steps=1000) + _ = orbit.represent_as('cylindrical').plot(['rho', 'z']) + +From the |orb| object, we can also easily compute dynamical quantities such as +the energy or angular momentum (we take the 0th element because these functions +return the quantities computed at every timestep):: + + >>> E = orbit.energy() + >>> E[0] # doctest: +SKIP + Let's see how well the integrator conserves energy and the ``z`` component of angular momentum:: - >>> E = orbit.energy() >>> Lz = orbit.angular_momentum()[2] - >>> np.std(E), np.std(Lz) + >>> np.std(E), np.std(Lz) # doctest: +FLOAT_CMP (, ) +We can access the position and velocity components of the orbit separately using +the attributes ``.pos`` and ``.vel``. These objects are +`~astropy.coordinates.BaseRepresentation` and +`~astropy.coordinates.BaseDifferential` subclasses. By default, as in this +example, both are Cartesian (`~astropy.coordinates.CartesianRepresentation` and +`~astropy.coordinates.CartesianDifferential`), so to access the individual +components, e.g., ``x``, we use:: + + >>> orbit.pos.x # doctest: +FLOAT_CMP + + >>> orbit.vel.d_x # doctest: +FLOAT_CMP + + +Continue to the :ref:`orbits-in-detail` page for more information. + Using gala.dynamics =================== + More details are provided in the linked pages below: .. toctree:: :maxdepth: 2 orbits-in-detail + nd-representations actionangle mockstreams nonlinear diff --git a/docs/dynamics/mockstreams.rst b/docs/dynamics/mockstreams.rst index 71d4be47b..f1867d631 100644 --- a/docs/dynamics/mockstreams.rst +++ b/docs/dynamics/mockstreams.rst @@ -88,8 +88,8 @@ initial conditions that place the progenitor on a mildly eccentric orbit: >>> pot = gp.SphericalNFWPotential(v_c=175*u.km/u.s, r_s=10*u.kpc, ... units=galactic) >>> prog_mass = 1E4*u.Msun - >>> prog_w0 = gd.CartesianPhaseSpacePosition(pos=[15, 0, 0.]*u.kpc, - ... vel=[75, 150, 30.]*u.km/u.s) + >>> prog_w0 = gd.PhaseSpacePosition(pos=[15, 0, 0.]*u.kpc, + ... vel=[75, 150, 30.]*u.km/u.s) >>> prog_orbit = pot.integrate_orbit(prog_w0, dt=0.5, n_steps=4000) >>> fig = prog_orbit.plot() @@ -105,8 +105,8 @@ initial conditions that place the progenitor on a mildly eccentric orbit: pot = gp.SphericalNFWPotential(v_c=175*u.km/u.s, r_s=10*u.kpc, units=galactic) prog_mass = 1E4*u.Msun - prog_w0 = gd.CartesianPhaseSpacePosition(pos=[15, 0, 0.]*u.kpc, - vel=[75, 150, 0.]*u.km/u.s) + prog_w0 = gd.PhaseSpacePosition(pos=[15, 0, 0.]*u.kpc, + vel=[75, 150, 0.]*u.km/u.s) prog_orbit = pot.integrate_orbit(prog_w0, dt=0.5, n_steps=4000) fig = prog_orbit.plot() @@ -138,8 +138,8 @@ from both Lagrange points by setting ``release_every=1``: pot = gp.SphericalNFWPotential(v_c=175*u.km/u.s, r_s=10*u.kpc, units=galactic) prog_mass = 1E4*u.Msun - prog_w0 = gd.CartesianPhaseSpacePosition(pos=[15, 0, 0.]*u.kpc, - vel=[75, 150, 0.]*u.km/u.s) + prog_w0 = gd.PhaseSpacePosition(pos=[15, 0, 0.]*u.kpc, + vel=[75, 150, 0.]*u.km/u.s) prog_orbit = pot.integrate_orbit(prog_w0, dt=0.5, n_steps=4000) k_mean = [1., 0, 0, 0, 1., 0] k_disp = np.zeros(6) @@ -163,8 +163,8 @@ Or, zooming in around the progenitor: pot = gp.SphericalNFWPotential(v_c=175*u.km/u.s, r_s=10*u.kpc, units=galactic) prog_mass = 1E4*u.Msun - prog_w0 = gd.CartesianPhaseSpacePosition(pos=[15, 0, 0.]*u.kpc, - vel=[75, 150, 30.]*u.km/u.s) + prog_w0 = gd.PhaseSpacePosition(pos=[15, 0, 0.]*u.kpc, + vel=[75, 150, 30.]*u.km/u.s) prog_orbit = pot.integrate_orbit(prog_w0, dt=0.5, n_steps=4000) k_mean = [1., 0, 0, 0, 1., 0] k_disp = np.zeros(6) @@ -211,8 +211,8 @@ stream using this method: pot = gp.SphericalNFWPotential(v_c=175*u.km/u.s, r_s=10*u.kpc, units=galactic) prog_mass = 1E4*u.Msun - prog_w0 = gd.CartesianPhaseSpacePosition(pos=[15, 0, 0.]*u.kpc, - vel=[75, 150, 0.]*u.km/u.s) + prog_w0 = gd.PhaseSpacePosition(pos=[15, 0, 0.]*u.kpc, + vel=[75, 150, 0.]*u.km/u.s) prog_orbit = pot.integrate_orbit(prog_w0, dt=0.5, n_steps=4000) k_mean = [2., 0, 0, 0, 0.3, 0] k_disp = [0.5, 0, 0.5, 0, 0.5, 0.5] @@ -236,8 +236,8 @@ Or, again, zooming in around the progenitor: pot = gp.SphericalNFWPotential(v_c=175*u.km/u.s, r_s=10*u.kpc, units=galactic) prog_mass = 1E4*u.Msun - prog_w0 = gd.CartesianPhaseSpacePosition(pos=[15, 0, 0.]*u.kpc, - vel=[75, 150, 30.]*u.km/u.s) + prog_w0 = gd.PhaseSpacePosition(pos=[15, 0, 0.]*u.kpc, + vel=[75, 150, 30.]*u.km/u.s) prog_orbit = pot.integrate_orbit(prog_w0, dt=0.5, n_steps=4000) k_mean = [2., 0, 0, 0, 0.3, 0] k_disp = [0.5, 0, 0.5, 0, 0.5, 0.5] diff --git a/docs/dynamics/nd-representations.rst b/docs/dynamics/nd-representations.rst new file mode 100644 index 000000000..a2169643e --- /dev/null +++ b/docs/dynamics/nd-representations.rst @@ -0,0 +1,69 @@ +.. include:: references.txt + +.. _nd-representations: + +************************************ +N-dimensional representation classes +************************************ + +For code blocks below and any pages linked below, we'll assume the following +imports have already been executed:: + + >>> import astropy.units as u + >>> import numpy as np + >>> import gala.dynamics as gd + +Introduction +============ + +The Astropy |astropyrep|_ presently only support 3D positions and differential +objects. The `~gala.dynamics.NDCartesianRepresentation` and +`~gala.dynamics.NDCartesianDifferential` classes add Cartesian representation +classes that can handle arbitrary numbers of dimensions. For example, 2D +coordinates:: + + >>> xy = np.arange(16).reshape(2, 8) * u.kpc + >>> rep = gd.NDCartesianRepresentation(xy) + >>> rep + + +4D coordinates:: + + >>> x = np.arange(16).reshape(4, 4) * u.kpc + >>> rep = gd.NDCartesianRepresentation(x) + >>> rep + + +These can be passed in to the |psp| or |orb| classes as with any of the Astropy +core representation objects:: + + >>> xy = np.arange(16).reshape(2, 8) * u.kpc + >>> vxy = np.arange(16).reshape(2, 8) / 10. * u.kpc / u.Myr + >>> w = gd.PhaseSpacePosition(pos=xy, vel=vxy) + >>> fig = w.plot() + +.. plot:: + :align: center + + import astropy.units as u + import numpy as np + import gala.dynamics as gd + xy = np.arange(16).reshape(2, 8) * u.kpc + vxy = np.arange(16).reshape(2, 8) / 10. * u.kpc / u.Myr + w = gd.PhaseSpacePosition(pos=xy, vel=vxy) + fig = w.plot() + +However, certain functionality such as representation transformations, dynamical +quantity calculation, and coordinate frame transformations are disabled when the +number of dimensions is not 3 (i.e. when not using the Astropy core +representation classes). + +N-dimensional representations API +--------------------------------- +.. automodapi:: gala.dynamics.representation_nd + :no-heading: + :headings: ^^ diff --git a/docs/dynamics/nonlinear.rst b/docs/dynamics/nonlinear.rst index 3305841d6..1d7e78f6e 100644 --- a/docs/dynamics/nonlinear.rst +++ b/docs/dynamics/nonlinear.rst @@ -35,14 +35,14 @@ a set of initial conditions:: >>> pot = gp.LogarithmicPotential(v_c=150*u.km/u.s, r_h=0.1*u.kpc, ... q1=1., q2=0.8, q3=0.6, units=galactic) - >>> w0 = gd.CartesianPhaseSpacePosition(pos=[5.5,0.,5.5]*u.kpc, - ... vel=[0.,100.,0]*u.km/u.s) + >>> w0 = gd.PhaseSpacePosition(pos=[5.5,0.,5.5]*u.kpc, + ... vel=[0.,100.,0]*u.km/u.s) >>> lyap,orbit = gd.fast_lyapunov_max(w0, pot, dt=2., n_steps=100000) # doctest: +SKIP This returns two objects: an `~astropy.units.Quantity` object that contains the maximum Lyapunov exponent estimate for each offset orbit, (we can control the number of offset orbits with the ``noffset_orbits`` -argument) and an `~gala.dynamics.CartesianOrbit` object that contains +argument) and an `~gala.dynamics.Orbit` object that contains the parent orbit and each offset orbit. Let's plot the parent orbit:: >>> fig = orbit[:,0].plot(marker=',', alpha=0.1, linestyle='none') # doctest: +SKIP @@ -58,8 +58,8 @@ the parent orbit and each offset orbit. Let's plot the parent orbit:: pot = gp.LogarithmicPotential(v_c=150*u.km/u.s, r_h=0.1*u.kpc, q1=1., q2=0.8, q3=0.6, units=galactic) - w0 = gd.CartesianPhaseSpacePosition(pos=[5.5,0.,5.5]*u.kpc, - vel=[0.,100.,0]*u.km/u.s) + w0 = gd.PhaseSpacePosition(pos=[5.5,0.,5.5]*u.kpc, + vel=[0.,100.,0]*u.km/u.s) lyap,orbit = gd.fast_lyapunov_max(w0, pot, dt=2., n_steps=100000) fig = orbit[:,0].plot(marker=',', linestyle='none', alpha=0.1) @@ -87,8 +87,8 @@ array. This plots one line per offset orbit:: pot = gp.LogarithmicPotential(v_c=150*u.km/u.s, r_h=0.1*u.kpc, q1=1., q2=0.8, q3=0.6, units=galactic) - w0 = gd.CartesianPhaseSpacePosition(pos=[5.5,0.,5.5]*u.kpc, - vel=[0.,100.,0]*u.km/u.s) + w0 = gd.PhaseSpacePosition(pos=[5.5,0.,5.5]*u.kpc, + vel=[0.,100.,0]*u.km/u.s) lyap,orbit = gd.fast_lyapunov_max(w0, pot, dt=2., n_steps=100000) pl.figure() @@ -104,8 +104,8 @@ Regular orbit To compare, we will compute the estimate for a regular orbit as well:: - >>> w0 = gd.CartesianPhaseSpacePosition(pos=[5.5,0.,0.]*u.kpc, - ... vel=[0.,140.,25]*u.km/u.s) + >>> w0 = gd.PhaseSpacePosition(pos=[5.5,0.,0.]*u.kpc, + ... vel=[0.,140.,25]*u.km/u.s) >>> lyap,orbit = gd.fast_lyapunov_max(w0, pot, dt=2., n_steps=100000) # doctest: +SKIP >>> fig = orbit[:,0].plot(marker=',', alpha=0.1, linestyle='none') # doctest: +SKIP @@ -120,8 +120,8 @@ To compare, we will compute the estimate for a regular orbit as well:: pot = gp.LogarithmicPotential(v_c=150*u.km/u.s, r_h=0.1*u.kpc, q1=1., q2=0.8, q3=0.6, units=galactic) - w0 = gd.CartesianPhaseSpacePosition(pos=[5.5,0.,0.]*u.kpc, - vel=[0.,140.,25]*u.km/u.s) + w0 = gd.PhaseSpacePosition(pos=[5.5,0.,0.]*u.kpc, + vel=[0.,140.,25]*u.km/u.s) lyap,orbit = gd.fast_lyapunov_max(w0, pot, dt=1., n_steps=200000) fig = orbit[:,0].plot(marker=',', linestyle='none', alpha=0.1) @@ -145,8 +145,8 @@ following a characteristic power-law (a straight line in a log-log plot):: pot = gp.LogarithmicPotential(v_c=150*u.km/u.s, r_h=0.1*u.kpc, q1=1., q2=0.8, q3=0.6, units=galactic) - w0 = gd.CartesianPhaseSpacePosition(pos=[5.5,0.,0.]*u.kpc, - vel=[0.,140.,25]*u.km/u.s) + w0 = gd.PhaseSpacePosition(pos=[5.5,0.,0.]*u.kpc, + vel=[0.,140.,25]*u.km/u.s) lyap,orbit = gd.fast_lyapunov_max(w0, pot, dt=1., n_steps=200000) pl.figure() diff --git a/docs/dynamics/orbits-in-detail.rst b/docs/dynamics/orbits-in-detail.rst index ccd116b00..60c601ce1 100644 --- a/docs/dynamics/orbits-in-detail.rst +++ b/docs/dynamics/orbits-in-detail.rst @@ -1,125 +1,164 @@ +.. include:: references.txt + .. _orbits-in-detail: ***************************************************** Orbit and phase-space position objects in more detail ***************************************************** -Introduction -============ - -The `astropy.units` subpackage is excellent for working with numbers and -associated units, however, dynamical quantities often contain many -quantities with mixed units. An example is a position in phase-space, which -may contain some quantities with length units and some quantities with -velocity units. The `~gala.dynamics.PhaseSpacePosition` and -`~gala.dynamics.Orbit` subclasses are designed to work with these structures -- -click these shortcuts to jump to a section below: - - * :ref:`phase-space-position` - * :ref:`orbit` - -Some imports needed for the code below:: +We'll assume the following imports have already been executed:: >>> import astropy.units as u >>> import numpy as np >>> import gala.potential as gp >>> import gala.dynamics as gd + >>> from astropy.coordinates import (CylindricalRepresentation, + ... CylindricalDifferential) >>> from gala.units import galactic >>> np.random.seed(42) +Introduction +============ + +The `astropy.units` subpackage is excellent for working with numbers and +associated units, but dynamical quantities often contain many quantities with +mixed units. An example is a position in phase-space, which may contain some +quantities with length units and some quantities with velocity or momentum +units. The |psp| and |orb| classes are designed to work with these data +structures and provide a consistent API for visualizing and computing further +dynamical quantities. Click these shortcuts to jump to a section below, or start +reading below: + + * :ref:`phase-space-position` + * :ref:`orbit` + .. _phase-space-position: Phase-space positions ===================== -It is often useful to represent full phase-space positions quantities jointly. -For example, if you need to transform the velocities to a new coordinate -representation or frame, the positions often enter into the transformations. -The `~gala.dynamics.PhaseSpacePosition` subclasses provide an interface for -handling these numbers. At present, only the -`~gala.dynamics.CartesianPhaseSpacePosition` is fully implemented. - -To create a `~gala.dynamics.CartesianPhaseSpacePosition` object, pass in a -cartesian position and velocity to the initializer:: +The |psp| class provides an interface for representing full phase-space +positions--coordinate positions and momenta (velocities). This class is useful +as a container for initial conditions and for transforming phase-space positions +to new coordinate representations or reference frames. - >>> gd.CartesianPhaseSpacePosition(pos=[4.,8.,15.]*u.kpc, - ... vel=[-150.,50.,15.]*u.km/u.s) - +The easiest way to create a |psp| object is to pass in a pair of +`~astropy.units.Quantity` objects that represent the Cartesian position and +velocity vectors:: -Of course, this works with arrays of positions and velocities as well:: + >>> gd.PhaseSpacePosition(pos=[4.,8.,15.]*u.kpc, + ... vel=[-150.,50.,15.]*u.km/u.s) + - >>> x = np.random.uniform(-10,10,size=(3,128)) - >>> v = np.random.uniform(-200,200,size=(3,128)) - >>> gd.CartesianPhaseSpacePosition(pos=x*u.kpc, - ... vel=v*u.km/u.s) - +By default, passing in `~astropy.units.Quantity`'s are interpreted as Cartesian +coordinates and velocities. This works with arrays of positions and velocities +as well:: -This works for arbitrary numbers of dimensions, e.g., we define a position:: - - >>> w = gd.CartesianPhaseSpacePosition(pos=[4.,8.]*u.kpc, - ... vel=[-150,45.]*u.km/u.s) + >>> x = np.arange(24).reshape(3,8) + >>> v = np.arange(24).reshape(3,8) + >>> w = gd.PhaseSpacePosition(pos=x * u.kpc, + ... vel=v * u.km/u.s) + >>> w + + +This is interpreted as 8, 6-dimensional phase-space positions. + +The class internally stores the positions and velocities as +`~astropy.coordinates.BaseRepresentation` and +`~astropy.coordinates.BaseDifferential` subclasses; in this case, +`~astropy.coordinates.CartesianRepresentation` and +`~astropy.coordinates.CartesianDifferential`:: + + >>> w.pos + + >>> w.vel + + +All of the components of these classes are added as attributes of the +phase-space position class for convenience. For example, to access the ``x`` +component of the position and the ``d_x`` component of the velocity:: + + >>> w.x, w.d_x # doctest: +FLOAT_CMP + (, + ) + +The default representation is Cartesian, but the class can also be instantiated +with representation objects instead of `~astropy.units.Quantity`'s:: + + >>> pos = CylindricalRepresentation(rho=np.linspace(1., 4, 4) * u.kpc, + ... phi=np.linspace(0, np.pi, 4) * u.rad, + ... z=np.linspace(-1, 1., 4) * u.kpc) + >>> vel = CylindricalDifferential(d_rho=np.linspace(100, 150, 4) * u.km/u.s, + ... d_phi=np.linspace(-1, 1, 4) * u.rad/u.Myr, + ... d_z=np.linspace(-15, 15., 4) * u.km/u.s) + >>> w = gd.PhaseSpacePosition(pos=pos, vel=vel) >>> w - - -We can check the dimensionality using the `~gala.dynamics.CartesianPhaseSpacePosition.ndim` -attribute:: - - >>> w.ndim - 2 - -For objects with ``ndim=3``, we can also easily transform the full -phase-space vector to new representations or coordinate frames. These -transformations use the :mod:`astropy.coordinates` framework and the -velocity transforms from `gala.coordinates`. - - >>> from astropy.coordinates import CylindricalRepresentation - >>> x = np.random.uniform(-10,10,size=(3,128)) - >>> v = np.random.uniform(-200,200,size=(3,128)) - >>> w = gd.CartesianPhaseSpacePosition(pos=x*u.kpc, - ... vel=v*u.km/u.s) - >>> cyl_pos, cyl_vel = w.represent_as(CylindricalRepresentation) - -The `~gala.dynamics.CartesianPhaseSpacePosition.represent_as` method returns two -objects that contain the position in the new representation, and the velocity -in the new representation. The position is returned as a -:class:`~astropy.coordinates.BaseRepresentation` subclass. The velocity is -presently returned as a single :class:`~astropy.units.Quantity` object with -the velocity components represented in velocity units (not in angular velocity -units!) but this will change in the future when velocity support is added -to Astropy:: - - >>> cyl_pos # doctest: +SKIP - >> cyl_vel # doctest: +SKIP - + >>> w.rho + + +We can easily transform the full phase-space vector to new representations or +coordinate frames. These transformations use the :mod:`astropy.coordinates` +|astropyrep|_:: + + >>> cart = w.represent_as('cartesian') + >>> cart.x + + >>> sph = w.represent_as('spherical') + >>> sph.distance + + +There is also support for transforming the positions and velocities (assumed to +be in a `~astropy.coordinates.Galactocentric` frame) to any of the other +coordinate frames. The transformation returns two objects: an +initialized coordinate frame for the position, and a ``Differential`` class +instance for the velocity (usually +`~astropy.coordinates.SphericalDifferential`). + +The velocities are represented in angular velocities conjugate to the angle +variables in the output coordinate frame. For example, in the below +transformation to :class:`~astropy.coordinates.Galactic` coordinates, the +returned velocity object is a tuple with angular velocities and radial velocity +in the :class:`~astropy.coordinates.Galactic` frame:: >>> from astropy.coordinates import Galactic >>> gal_c, gal_v = w.to_coord_frame(Galactic) - >>> gal_c # doctest: +SKIP + >>> gal_c # doctest: +FLOAT_CMP >> gal_v[0].unit, gal_v[1].unit, gal_v[2].unit - (Unit("mas / yr"), Unit("mas / yr"), Unit("km / s")) - -We can easily plot projections of the positions using the -`~gala.dynamics.CartesianPhaseSpacePosition.plot` method:: + [( 4.42801092e-05, -6.11537341, 9.35649038), + ( 1.05488650e+01, -1.99824507, 9.46673245), + ( 2.09134381e+01, 2.58371838, 7.28582479), + ( 7.26282965e-05, 12.9365465 , 4.40866775)]> + >>> gal_v # doctest: +FLOAT_CMP + + +It's important to note that the longitudinal angular velocity component in the +`~astropy.coordinates.SphericalDifferential` class does *not* include the +:math:`\cos{\rm lat}` term. For this example, to get :math:`\mu_l\,cos{b}` you +would need to do:: + + >>> pm_l = gal_v.d_lon * np.cos(gal_c.b) + >>> pm_l + + +We can easily plot projections of the phase-space positions using the +`~gala.dynamics.PhaseSpacePosition.plot` method:: + >>> np.random.seed(42) + >>> x = np.random.uniform(-10, 10, size=(3,128)) + >>> v = np.random.uniform(-200, 200, size=(3,128)) + >>> w = gd.PhaseSpacePosition(pos=x * u.kpc, + ... vel=v * u.km/u.s) >>> fig = w.plot() .. plot:: @@ -131,14 +170,15 @@ We can easily plot projections of the positions using the np.random.seed(42) x = np.random.uniform(-10,10,size=(3,128)) v = np.random.uniform(-200,200,size=(3,128)) - w = gd.CartesianPhaseSpacePosition(pos=x*u.kpc, - vel=v*u.km/u.s) + w = gd.PhaseSpacePosition(pos=x*u.kpc, + vel=v*u.km/u.s) fig = w.plot() -This is a thin wrapper around the `~gala.dynamics.three_panel` +This is a thin wrapper around the `~gala.dynamics.plot_projections` function and any keyword arguments are passed through to that function:: - >>> fig = w.plot(marker='o', s=40, alpha=0.5) + >>> fig = w.plot(components=['x', 'd_z'], color='r', + ... facecolor='', marker='o', s=20, alpha=0.5) .. plot:: :align: center @@ -149,101 +189,88 @@ function and any keyword arguments are passed through to that function:: np.random.seed(42) x = np.random.uniform(-10,10,size=(3,128)) v = np.random.uniform(-200,200,size=(3,128)) - w = gd.CartesianPhaseSpacePosition(pos=x*u.kpc, - vel=v*u.km/u.s) - fig = w.plot(marker='o', s=40, alpha=0.5) + w = gd.PhaseSpacePosition(pos=x*u.kpc, + vel=v*u.km/u.s) + fig = w.plot(components=['x', 'd_z'], color='r', + facecolor='', marker='o', s=20, alpha=0.5) Phase-space position API ------------------------ .. automodapi:: gala.dynamics.core :no-heading: :headings: ^^ + :skip: CartesianPhaseSpacePosition .. _orbit: Orbits ====== -The `~gala.dynamics.Orbit` subclasses all inherit the functionality described -above from `~gala.dynamics.PhaseSpacePosition`, but similarly, at present only the -`~gala.dynamics.CartesianOrbit` is fully implemented. There are some differences -between the methods and some functionality that is particular to the orbit classes. - -A `~gala.dynamics.CartesianOrbit` is initialized much like the -`~gala.dynamics.CartesianPhaseSpacePosition`. `~gala.dynamics.CartesianOrbit`s can be -created with just position and velocity information, however now the -interpretation of the input object shapes is different. Whereas an input position with -shape ``(2,128)`` to a `~gala.dynamics.CartesianPhaseSpacePosition` represents -128, 2D positions, for an orbit it would represent a single orbit's positions -at 128 timesteps:: - - >>> t = np.linspace(0,10,128) - >>> pos,vel = np.zeros((2,128)),np.zeros((2,128)) - >>> pos[0] = np.cos(t) - >>> pos[1] = np.sin(t) - >>> vel[0] = -np.sin(t) - >>> vel[1] = np.cos(t) - >>> orbit = gd.CartesianOrbit(pos=pos*u.kpc, vel=vel*u.km/u.s) +The |orb| class inherits much of the functionality from |psp| (described above) +and adds some additional features that are useful for time-series orbits. + +An |orb| instance is initialized like the |psp|--with arrays of positions and +velocities-- but usually also requires specifying a time array as well. Also, +the extra axes in these arrays hold special meaning for the |orb| class. The +position and velocity arrays passed to |psp| can have arbitrary numbers of +dimensions as long as the 0th axis specifies the dimensionality. For the |orb| +class, the 0th axis remains the axis of dimensionality, but the 1st axis now is +always assumed to be the time axis. For example, an input position with shape +``(2,128)`` to a |psp| represents 128 independent 2D positions, but to a |orb| +it represents a single orbit's positions at 128 times:: + + >>> t = np.linspace(0, 100, 128) * u.Myr + >>> Om = 1E-1 * u.rad / u.Myr + >>> pos = np.vstack((5*np.cos(Om*t), np.sin(Om*t))).value * u.kpc + >>> vel = np.vstack((-5*np.sin(Om*t), np.cos(Om*t))).value * u.kpc/u.Myr + >>> orbit = gd.Orbit(pos=pos, vel=vel) >>> orbit - - -To create a single object that contains multiple orbits, the input position object -should have 3 axes. The last axis (``axis=2``) contains each orbit. So, an input -position with shape ``(2,128,16)`` would represent 16, 2D orbits with 128 timesteps:: - - >>> t = np.linspace(0,10,128) - >>> pos,vel = np.zeros((2,128,16)),np.zeros((2,128,16)) - >>> Omega = np.random.uniform(size=16) - >>> pos[0] = np.cos(Omega[np.newaxis]*t[:,np.newaxis]) - >>> pos[1] = np.sin(Omega[np.newaxis]*t[:,np.newaxis]) - >>> vel[0] = -np.sin(Omega[np.newaxis]*t[:,np.newaxis]) - >>> vel[1] = np.cos(Omega[np.newaxis]*t[:,np.newaxis]) - >>> orbit = gd.CartesianOrbit(pos=pos*u.kpc, vel=vel*u.km/u.s) + + +To create a single object that contains multiple orbits, the input position +object should have 3 axes. The last axis (``axis=2``) specifies the number of +orbits. So, an input position with shape ``(2,128,16)`` would represent 16, 2D +orbits, each with the same 128 times:: + + >>> t = np.linspace(0, 100, 128) * u.Myr + >>> Om = np.random.uniform(size=16) * u.rad / u.Myr + >>> angle = Om[None] * t[:,None] + >>> pos = np.stack((5*np.cos(angle), np.sin(angle))).value * u.kpc + >>> vel = np.stack((-5*np.sin(angle), np.cos(angle))).value * u.kpc/u.Myr + >>> orbit = gd.Orbit(pos=pos, vel=vel) >>> orbit - + To make full use of the orbit functionality, you must also pass in an array with -the time values and an instance of a `~gala.potential.PotentialBase` subclass that -represents the potential that the orbit was integrated in:: +the time values and an instance of a `~gala.potential.PotentialBase` subclass +that represents the potential that the orbit was integrated in:: >>> pot = gp.PlummerPotential(m=1E10, b=1., units=galactic) - >>> orbit = gd.CartesianOrbit(pos=pos*u.kpc, vel=vel*u.km/u.s, - ... t=t*u.Myr, potential=pot) + >>> orbit = gd.Orbit(pos=pos*u.kpc, vel=vel*u.km/u.s, + ... t=t*u.Myr, potential=pot) (note, in this case ``pos`` and ``vel`` were not generated from integrating -an orbit in the potential ``pot``!) Orbit objects -are returned by the `~gala.potential.PotentialBase.integrate_orbit` method -of potential objects that already have the ``time`` and ``potential`` set:: - - >>> pot = gp.PlummerPotential(m=1E10, b=1., units=galactic) - >>> w0 = gd.CartesianPhaseSpacePosition(pos=[10.,0,0]*u.kpc, - ... vel=[0.,75,0]*u.km/u.s) - >>> orbit = pot.integrate_orbit(w0, dt=1., n_steps=500) +an orbit in the potential ``pot``!). However, most of the time you won't need to +create |orb| objects from scratch! They are returned from any of the numerical +intergration routines provided in `gala`. For example, they are returned by the +`~gala.potential.PotentialBase.integrate_orbit` method of potential objects and +will automatically contain the ``time`` array and ``potential`` object. For +example:: + + >>> pot = gp.PlummerPotential(m=1E10 * u.Msun, b=1. * u.kpc, units=galactic) + >>> w0 = gd.PhaseSpacePosition(pos=[10.,0,0] * u.kpc, + ... vel=[0.,75,0] * u.km/u.s) + >>> orbit = pot.integrate_orbit(w0, dt=1., n_steps=5000) >>> orbit - - >>> orbit.t # doctest: +SKIP - + >>> orbit.t + >>> orbit.potential -From an Orbit object, we can quickly compute quantities like the angular momentum, -and estimates for the pericenter, apocenter, eccentricity of the orbit. Estimates -for the latter few get better with smaller timesteps:: - - >>> orbit = pot.integrate_orbit(w0, dt=0.1, n_steps=100000) - >>> np.mean(orbit.angular_momentum(), axis=1) # doctest: +FLOAT_CMP - - >>> orbit.eccentricity() # doctest: +FLOAT_CMP - - >>> orbit.pericenter() # doctest: +FLOAT_CMP - - >>> orbit.apocenter() # doctest: +FLOAT_CMP - - -Just like above, we can quickly visualize an orbit using the -`~gala.dynamics.CartesianOrbit.plot` method:: +Just like for |psp|, we can quickly visualize an orbit using the +`~gala.dynamics.Orbit.plot` method:: >>> fig = orbit.plot() @@ -255,18 +282,16 @@ Just like above, we can quickly visualize an orbit using the import gala.potential as gp from gala.units import galactic - pot = gp.PlummerPotential(m=1E10, b=1., units=galactic) - w0 = gd.CartesianPhaseSpacePosition(pos=[1.,0,0]*u.kpc, - vel=[0.,50,0]*u.km/u.s) - orbit = pot.integrate_orbit(w0, dt=1., n_steps=500) + pot = gp.PlummerPotential(m=1E10 * u.Msun, b=1. * u.kpc, units=galactic) + w0 = gd.PhaseSpacePosition(pos=[10.,0,0] * u.kpc, + vel=[0.,75,0] * u.km/u.s) + orbit = pot.integrate_orbit(w0, dt=1., n_steps=5000) fig = orbit.plot() -This is a thin wrapper around the `~gala.dynamics.plot_orbits` +Again, this is a thin wrapper around the `~gala.dynamics.plot_projections` function and any keyword arguments are passed through to that function:: >>> fig = orbit.plot(linewidth=4., alpha=0.5, color='r') - >>> fig.axes[0].set_xlim(-1.5,1.5) # doctest: +SKIP - >>> fig.axes[0].set_ylim(-1.5,1.5) # doctest: +SKIP .. plot:: :align: center @@ -276,17 +301,41 @@ function and any keyword arguments are passed through to that function:: import gala.potential as gp from gala.units import galactic - pot = gp.PlummerPotential(m=1E10, b=1., units=galactic) - w0 = gd.CartesianPhaseSpacePosition(pos=[1.,0,0]*u.kpc, - vel=[0.,50,0]*u.km/u.s) - orbit = pot.integrate_orbit(w0, dt=1., n_steps=500) + pot = gp.PlummerPotential(m=1E10 * u.Msun, b=1. * u.kpc, units=galactic) + w0 = gd.PhaseSpacePosition(pos=[10.,0,0] * u.kpc, + vel=[0.,75,0] * u.km/u.s) + orbit = pot.integrate_orbit(w0, dt=1., n_steps=5000) fig = orbit.plot(linewidth=4., alpha=0.5, color='r') - fig.axes[0].set_xlim(-1.5,1.5) - fig.axes[0].set_ylim(-1.5,1.5) +We can also quickly compute quantities like the angular momentum, and estimates +for the pericenter, apocenter, eccentricity of the orbit. Estimates for the +latter few get better with smaller timesteps:: + + >>> orbit = pot.integrate_orbit(w0, dt=0.1, n_steps=100000) + >>> np.mean(orbit.angular_momentum(), axis=1) # doctest: +FLOAT_CMP + + >>> orbit.eccentricity() # doctest: +FLOAT_CMP + + >>> orbit.pericenter() # doctest: +FLOAT_CMP + + >>> orbit.apocenter() # doctest: +FLOAT_CMP + Orbit API --------- .. automodapi:: gala.dynamics.orbit :no-heading: :headings: ^^ + :skip: CartesianOrbit + +More information +================ + +Internally, both of the above classes rely on the Astropy representation +transformation framework (i.e. the subclasses of +`~astropy.coordinates.BaseRepresentation` and +`~astropy.coordinates.BaseDifferential`). However, at present these classes only +support 3D positions and differentials (velocities). The |psp| and |orb| classes +both support arbitrary numbers of dimensions and, when relevant, rely on custom +subclasses of the representation classes to handle such cases. See the +:ref:`nd-representations` page for more information about these classes. diff --git a/docs/dynamics/references.txt b/docs/dynamics/references.txt new file mode 100644 index 000000000..7a444969c --- /dev/null +++ b/docs/dynamics/references.txt @@ -0,0 +1,4 @@ +.. |psp| replace:: `~gala.dynamics.PhaseSpacePosition` +.. |orb| replace:: `~gala.dynamics.Orbit` +.. |astropyrep| replace:: representations framework +.. _astropyrep: http://docs.astropy.org/en/latest/coordinates/skycoord.html#astropy-skycoord-representations diff --git a/docs/examples/integrate-potential-example.rst b/docs/examples/integrate-potential-example.rst index da896992d..79d3ec2c6 100644 --- a/docs/examples/integrate-potential-example.rst +++ b/docs/examples/integrate-potential-example.rst @@ -12,27 +12,31 @@ We first need to import some relevant packages:: >>> import gala.potential as gp >>> from gala.units import galactic -In the examples below, we will work use the ``galactic`` `~gala.units.UnitSystem`: -as I define it, this is: :math:`{\rm kpc}`, :math:`{\rm Myr}`, :math:`{\rm M}_\odot`. +In the examples below, we will work use the ``galactic`` +`~gala.units.UnitSystem`: as I define it, this is: :math:`{\rm kpc}`, +:math:`{\rm Myr}`, :math:`{\rm M}_\odot`. We first create a potential object to work with. For this example, we'll use a spherical NFW potential, parametrized by a scale radius and the circular velocity at the scale radius:: - >>> pot = gp.SphericalNFWPotential(v_c=200*u.km/u.s, r_s=10.*u.kpc, units=galactic) + >>> pot = gp.SphericalNFWPotential(v_c=200*u.km/u.s, r_s=10.*u.kpc, + ... units=galactic) -As a demonstration, we're going to first integrate a single orbit in this potential. +As a demonstration, we're going to first integrate a single orbit in this +potential. -The easiest way to do this is to use the `~gala.potential.PotentialBase.integrate_orbit` -method of the potential object, which accepts a set of initial conditions and -a specification for the time-stepping. We'll define the initial conditions as a -`~gala.dynamics.CartesianPhaseSpacePosition` object:: +The easiest way to do this is to use the +`~gala.potential.PotentialBase.integrate_orbit` method of the potential object, +which accepts a set of initial conditions and a specification for the +time-stepping. We'll define the initial conditions as a +`~gala.dynamics.PhaseSpacePosition` object:: - >>> ics = gd.CartesianPhaseSpacePosition(pos=[10,0,0.]*u.kpc, - ... vel=[0,175,0]*u.km/u.s) + >>> ics = gd.PhaseSpacePosition(pos=[10,0,0.] * u.kpc, + ... vel=[0,175,0] * u.km/u.s) >>> orbit = pot.integrate_orbit(ics, dt=2., n_steps=2000) -This method returns a `~gala.dynamics.CartesianOrbit` object that contains an +This method returns a `~gala.dynamics.Orbit` object that contains an array of times and the (6D) position at each time-step. By default, this method uses Leapfrog integration to compute the orbit (:class:`~gala.integrate.LeapfrogIntegrator`), but you can optionally specify @@ -47,9 +51,11 @@ conditions by sampling from a Gaussian around the initial orbit (with a positional scale of 100 pc, and a velocity scale of 1 km/s):: >>> norbits = 128 - >>> new_pos = np.random.normal(ics.pos.to(u.pc).value, 100., size=(3,norbits))*u.pc - >>> new_vel = np.random.normal(ics.vel.to(u.km/u.s).value, 1, size=(3,norbits))*u.km/u.s - >>> new_ics = gd.CartesianPhaseSpacePosition(pos=new_pos, vel=new_vel) + >>> new_pos = np.random.normal(ics.pos.xyz.to(u.pc).value, 100., + ... size=(norbits,3)).T * u.pc + >>> new_vel = np.random.normal(ics.vel.d_xyz.to(u.km/u.s).value, 1., + ... size=(norbits,3)).T * u.km/u.s + >>> new_ics = gd.PhaseSpacePosition(pos=new_pos, vel=new_vel) >>> orbits = pot.integrate_orbit(new_ics, dt=2., n_steps=2000) We'll now plot the final positions of these orbits over isopotential contours. @@ -78,14 +84,14 @@ the orbit points:: pot = gp.SphericalNFWPotential(v_c=200*u.km/u.s, r_s=10.*u.kpc, units=galactic) - ics = gd.CartesianPhaseSpacePosition(pos=[10,0,0.]*u.kpc, + ics = gd.PhaseSpacePosition(pos=[10,0,0.]*u.kpc, vel=[0,175,0]*u.km/u.s) orbit = pot.integrate_orbit(ics, dt=2., n_steps=2000) norbits = 1024 new_pos = np.random.normal(ics.pos.to(u.pc).value, 100., size=(3,norbits))*u.pc new_vel = np.random.normal(ics.vel.to(u.km/u.s).value, 1., size=(3,norbits))*u.km/u.s - new_ics = gd.CartesianPhaseSpacePosition(pos=new_pos, vel=new_vel) + new_ics = gd.PhaseSpacePosition(pos=new_pos, vel=new_vel) orbits = pot.integrate_orbit(new_ics, dt=2., n_steps=2000) grid = np.linspace(-15,15,64) diff --git a/docs/integrate/index.rst b/docs/integrate/index.rst index a8f141f5f..b747c6821 100644 --- a/docs/integrate/index.rst +++ b/docs/integrate/index.rst @@ -26,31 +26,32 @@ imports have already been executed:: Getting Started =============== -All of the integrator classes have the same basic call structure. To create -an integrator object, you pass in a function that evaluates derivatives of, -for example, phase-space coordinates, then you call the `~gala.integrate.Integrator.run` -method while specifying timestep information. - -The integration function must accept, at minimum, two arguments: the current time, ``t``, -and the current position in phase-space, ``w``. The time is a single floating-point number -and the position will have shape ``(ndim, norbits)`` where ``ndim`` is the full dimensionality -of the phase-space (e.g., 6 for a 3D coordinate system) and ``norbits`` is the number of -orbits. These inputs will *not* have units associated with them (e.g., they are not -:class:`astropy.units.Quantity` objects). An example of such a function (that represents -a simple harmonic oscillator) is:: +All of the integrator classes have the same basic call structure. To create an +integrator object, you pass in a function that evaluates derivatives of, for +example, phase-space coordinates, then you call the +`~gala.integrate.Integrator.run` method while specifying timestep information. + +The integration function must accept, at minimum, two arguments: the current +time, ``t``, and the current position in phase-space, ``w``. The time is a +single floating-point number and the position will have shape ``(ndim, +norbits)`` where ``ndim`` is the full dimensionality of the phase-space (e.g., 6 +for a 3D coordinate system) and ``norbits`` is the number of orbits. These +inputs will *not* have units associated with them (e.g., they are not +:class:`astropy.units.Quantity` objects). An example of such a function (that + represents a simple harmonic oscillator) is:: >>> def F(t, w): ... x,x_dot = w ... return np.array([x_dot, -x]) -Even though time does not explicitly enter into the equation, the function must still -accept a time argument. We will now create an instance of `~gala.integrate.LeapfrogIntegrator` -to integrate an orbit:: +Even though time does not explicitly enter into the equation, the function must +still accept a time argument. We will now create an instance of +`~gala.integrate.LeapfrogIntegrator` to integrate an orbit:: >>> integrator = gi.LeapfrogIntegrator(F) -To actually run the integrator, we need to specify a set of initial conditions. The -simplest way to do this is to specify an array:: +To actually run the integrator, we need to specify a set of initial conditions. +The simplest way to do this is to specify an array:: >>> w0 = np.array([1.,0.]) @@ -62,18 +63,18 @@ steps:: >>> orbit = integrator.run(w0, dt=0.5, n_steps=100) >>> orbit.t.unit Unit(dimensionless) - >>> orbit.pos.unit + >>> orbit.pos.xyz.unit Unit(dimensionless) We can instead specify the unit system that the function (``F``) expects, and -then pass in a `~gala.dynamics.CartesianPhaseSpacePosition` object with arbitrary -units in as initial conditions:: +then pass in a `~gala.dynamics.PhaseSpacePosition` object with arbitrary units +in as initial conditions:: >>> import gala.dynamics as gd >>> from gala.units import UnitSystem >>> usys = UnitSystem(u.m, u.s, u.kg, u.radian) >>> integrator = gi.LeapfrogIntegrator(F, func_units=usys) - >>> w0 = gd.CartesianPhaseSpacePosition(pos=100.*u.cm, vel=0*u.cm/u.yr) + >>> w0 = gd.PhaseSpacePosition(pos=[100.]*u.cm, vel=[0]*u.cm/u.yr) >>> orbit = integrator.run(w0, dt=0.5, n_steps=100) >>> orbit.t.unit Unit("s") @@ -107,25 +108,27 @@ Our derivative function is then:: ... p_dot = -np.sin(q) + A*omega_D*np.cos(omega_D*t) ... return np.array([q_dot, p_dot]) -This function has two arguments -- :math:`A` (``A``), the amplitude of the forcing, -and :math:`\omega_D` (``omega_D``), the driving frequency. We define an integrator -object by specifying this function, along with values for the function arguments:: +This function has two arguments -- :math:`A` (``A``), the amplitude of the +forcing, and :math:`\omega_D` (``omega_D``), the driving frequency. We define an +integrator object by specifying this function, along with values for the +function arguments:: >>> integrator = gi.DOPRI853Integrator(F, func_args=(0.07, 0.75)) -To integrate an orbit, we use the `~gala.integrate.Integrator.run` method. -We have to specify the initial conditions along with information about how long -to integrate and with what stepsize. There are several options for how to specify -the time step information. We could pre-generate an array of times and pass that in, -or pass in an initial time, end time, and timestep. Or, we could simply pass in the -number of steps to run for and a timestep. For this example, we will use the last -option. See the API below under *"Other Parameters"* for more information.:: +To integrate an orbit, we use the `~gala.integrate.Integrator.run` method. We +have to specify the initial conditions along with information about how long to +integrate and with what stepsize. There are several options for how to specify +the time step information. We could pre-generate an array of times and pass that +in, or pass in an initial time, end time, and timestep. Or, we could simply pass +in the number of steps to run for and a timestep. For this example, we will use +the last option. See the API below under *"Other Parameters"* for more +information.:: >>> orbit = integrator.run([3.,0.], dt=0.1, n_steps=10000) We can plot the integrated (chaotic) orbit:: - >>> fig = orbit.plot() + >>> fig = orbit.plot(subplots_kwargs=dict(figsize=(8,4))) .. plot:: :align: center @@ -143,7 +146,7 @@ We can plot the integrated (chaotic) orbit:: integrator = gi.DOPRI853Integrator(F, func_args=(0.07, 0.75)) orbit = integrator.run([3.,0.], dt=0.1, n_steps=10000) - fig = orbit.plot() + fig = orbit.plot(subplots_kwargs=dict(figsize=(8,4))) Example: Lorenz equations ------------------------- diff --git a/docs/potential/define-new-potential.rst b/docs/potential/define-new-potential.rst index 756718795..439e8797c 100644 --- a/docs/potential/define-new-potential.rst +++ b/docs/potential/define-new-potential.rst @@ -104,8 +104,8 @@ dimensionless quantities:: That's it! Now we have a fully-fledged potential object. For example, we can integrate an orbit in this potential:: - >>> w0 = gd.CartesianPhaseSpacePosition(pos=[0.,0.3], - ... vel=[0.38,0.]) + >>> w0 = gd.PhaseSpacePosition(pos=[0.,0.3], + ... vel=[0.38,0.]) >>> orbit = pot.integrate_orbit(w0, dt=0.05, n_steps=10000) >>> fig = orbit.plot(marker=',', linestyle='none', alpha=0.5) @@ -140,8 +140,8 @@ can integrate an orbit in this potential:: return grad pot = HenonHeilesPotential(A=1., units=None) - w0 = gd.CartesianPhaseSpacePosition(pos=[0.,0.3], - vel=[0.38,0.]) + w0 = gd.PhaseSpacePosition(pos=[0.,0.3], + vel=[0.38,0.]) orbit = pot.integrate_orbit(w0, dt=0.05, n_steps=10000) fig = orbit.plot(marker=',', linestyle='none', alpha=0.5) diff --git a/docs/potential/hamiltonian-reference-frames.rst b/docs/potential/hamiltonian-reference-frames.rst index 5455acd95..af23f5c63 100644 --- a/docs/potential/hamiltonian-reference-frames.rst +++ b/docs/potential/hamiltonian-reference-frames.rst @@ -23,8 +23,8 @@ When :ref:`integrating orbits using the potential classes directly >>> pot = gp.HernquistPotential(m=1E10*u.Msun, c=1.*u.kpc, ... units=galactic) - >>> w0 = gd.CartesianPhaseSpacePosition(pos=[5.,0,0]*u.kpc, - ... vel=[0,0,50.]*u.km/u.s) + >>> w0 = gd.PhaseSpacePosition(pos=[5.,0,0]*u.kpc, + ... vel=[0,0,50.]*u.km/u.s) >>> orbit = pot.integrate_orbit(w0, dt=0.5, n_steps=1000) it is implicitly assumed that the initial conditions and orbit are in an @@ -68,8 +68,8 @@ object and a :class:`~gala.potential.frame.Staticframe` instance:: ... units=galactic) >>> frame = gp.StaticFrame(units=galactic) >>> H = gp.Hamiltonian(potential=pot, frame=frame) - >>> w0 = gd.CartesianPhaseSpacePosition(pos=[5.,0,0]*u.kpc, - ... vel=[0,0,50.]*u.km/u.s) + >>> w0 = gd.PhaseSpacePosition(pos=[5.,0,0]*u.kpc, + ... vel=[0,0,50.]*u.km/u.s) >>> orbit = H.integrate_orbit(w0, dt=0.5, n_steps=1000) In this case, the ``orbit`` object returned from integration knows what @@ -102,8 +102,8 @@ rotation around that axis:: units=galactic) frame = gp.StaticFrame(units=galactic) H = gp.Hamiltonian(potential=pot, frame=frame) - w0 = gd.CartesianPhaseSpacePosition(pos=[5.,0,0]*u.kpc, - vel=[0,0,50.]*u.km/u.s) + w0 = gd.PhaseSpacePosition(pos=[5.,0,0]*u.kpc, + vel=[0,0,50.]*u.km/u.s) orbit = H.integrate_orbit(w0, dt=0.5, n_steps=1000) rotation_axis = np.array([8.2, -1.44, 3.25]) @@ -143,8 +143,8 @@ frame:: pot = gp.HernquistPotential(m=1E10*u.Msun, c=1.*u.kpc, units=galactic) - w0 = gd.CartesianPhaseSpacePosition(pos=[5.,0,0]*u.kpc, - vel=[0,0,50.]*u.km/u.s) + w0 = gd.PhaseSpacePosition(pos=[5.,0,0]*u.kpc, + vel=[0,0,50.]*u.km/u.s) rotation_axis = np.array([8.2, -1.44, 3.25]) rotation_axis /= np.linalg.norm(rotation_axis) # make a unit vector @@ -243,8 +243,8 @@ radius:: >>> r_corot = res.x * u.kpc >>> v_circ = Om_bar * r_corot * u.kpc >>> - >>> w0 = gd.CartesianPhaseSpacePosition(pos=[r_corot.value, 0, 0] * r_corot.unit, - ... vel=[0, v_circ.value, 0.] * v_circ.unit) + >>> w0 = gd.PhaseSpacePosition(pos=[r_corot.value, 0, 0] * r_corot.unit, + ... vel=[0, v_circ.value, 0.] * v_circ.unit) >>> orbit = H.integrate_orbit(w0, dt=0.1, n_steps=40000, ... Integrator=gi.DOPRI853Integrator) >>> fig = orbit.plot(marker=',', linestyle='none', alpha=0.5) # doctest: +SKIP @@ -287,8 +287,8 @@ radius:: r_corot = res.x * u.kpc v_circ = Om_bar * r_corot - w0 = gd.CartesianPhaseSpacePosition(pos=[r_corot.value, 0, 0] * r_corot.unit, - vel=[0,v_circ.value, 0.] * v_circ.unit) + w0 = gd.PhaseSpacePosition(pos=[r_corot.value, 0, 0] * r_corot.unit, + vel=[0,v_circ.value, 0.] * v_circ.unit) orbit = H.integrate_orbit(w0, dt=0.1, n_steps=40000, Integrator=gi.DOPRI853Integrator) @@ -342,8 +342,8 @@ orbit looks like in an inertial frame:: r_corot = res.x * u.kpc v_circ = Om_bar * r_corot - w0 = gd.CartesianPhaseSpacePosition(pos=[r_corot.value, 0, 0] * r_corot.unit, - vel=[0,v_circ.value, 0.] * v_circ.unit) + w0 = gd.PhaseSpacePosition(pos=[r_corot.value, 0, 0] * r_corot.unit, + vel=[0,v_circ.value, 0.] * v_circ.unit) orbit = H.integrate_orbit(w0, dt=0.1, n_steps=40000, Integrator=gi.DOPRI853Integrator) diff --git a/docs/why.rst b/docs/why.rst index eee565103..7bdf6e04b 100644 --- a/docs/why.rst +++ b/docs/why.rst @@ -55,8 +55,8 @@ implemented in C, enabling extremely fast orbit integration for single or composite potentials: >>> pot = gp.IsochronePotential(m=1E10*u.Msun, b=15.*u.kpc, units=galactic) - >>> w0 = gd.CartesianPhaseSpacePosition(pos=[7.,0,0]*u.kpc, - ... vel=[0.,50.,0]*u.km/u.s) + >>> w0 = gd.PhaseSpacePosition(pos=[7.,0,0]*u.kpc, + ... vel=[0.,50.,0]*u.km/u.s) >>> import timeit >>> timeit.timeit(lambda: pot.integrate_orbit(w0, dt=0.5, n_steps=10000), number=100) / 100. # doctest: +SKIP 0.0028513244865462184 @@ -83,7 +83,7 @@ Easy visualization ================== Numerically integrated orbits can be easily visualized using the -`~gala.dynamics.CartesianOrbit.plot()` method: +`~gala.dynamics.Orbit.plot()` method: >>> orbit.plot() # doctest: +SKIP @@ -101,8 +101,8 @@ Numerically integrated orbits can be easily visualized using the disk = gp.MiyamotoNagaiPotential(m=6E10*u.Msun, a=3*u.kpc, b=0.26*u.kpc, units=galactic) pot = gp.CCompositePotential(bulge=bulge, disk=disk) - w0 = gd.CartesianPhaseSpacePosition(pos=[7.,0,0]*u.kpc, - vel=[0.,50.,0]*u.km/u.s) + w0 = gd.PhaseSpacePosition(pos=[7.,0,0]*u.kpc, + vel=[0.,50.,0]*u.km/u.s) orbit = pot.integrate_orbit(w0, dt=0.5, n_steps=10000, Integrator=gi.DOPRI853Integrator) diff --git a/gala/coordinates/__init__.py b/gala/coordinates/__init__.py index ea5468828..814994e71 100644 --- a/gala/coordinates/__init__.py +++ b/gala/coordinates/__init__.py @@ -2,7 +2,6 @@ from .orphan import * from .gd1 import * from .propermotion import * -from .velocity_coord_transforms import * from .velocity_frame_transforms import * from .poincarepolar import * from .quaternion import * diff --git a/gala/coordinates/tests/test_velocity_frame_transforms.py b/gala/coordinates/tests/test_velocity_frame_transforms.py index 1d4337fa2..118d5d98f 100644 --- a/gala/coordinates/tests/test_velocity_frame_transforms.py +++ b/gala/coordinates/tests/test_velocity_frame_transforms.py @@ -7,18 +7,19 @@ # Standard library -import os -import pytest -import numpy as np import tempfile # Third-party import astropy.coordinates as coord import astropy.units as u from astropy.utils.data import get_pkg_data_filename +from astropy.tests.helper import quantity_allclose +import pytest +import numpy as np # This package -from ..velocity_frame_transforms import * +from ..velocity_frame_transforms import (vgal_to_hel, vhel_to_gal, + vgsr_to_vhel, vhel_to_vgsr) def test_vgsr_to_vhel(): filename = get_pkg_data_filename('idl_vgsr_vhel.txt') @@ -130,186 +131,214 @@ def setup(self): temp.seek(0) self.data = np.genfromtxt(temp, names=True, skip_header=1) + # This should make the transformations more compatible + g = coord.Galactic(l=0*u.deg, b=0*u.deg).transform_to(coord.ICRS) + self.galcen_frame = coord.Galactocentric(galcen_ra=g.ra, + galcen_dec=g.dec, + z_sun=0*u.kpc) + def test_vhel_to_gal_single(self): - # test a single entry - row = self.data[0] - c = coord.SkyCoord(ra=row['ra']*u.deg, dec=row['dec']*u.deg, distance=row['dist']*u.pc) - pm = [row['pml'], row['pmb']]*u.mas/u.yr - rv = row['rv']*u.km/u.s + for row in self.data: # test one entry at a time + c = coord.SkyCoord(ra=row['ra']*u.deg, dec=row['dec']*u.deg, + distance=row['dist']*u.pc) + icrs = c.icrs + gal = c.galactic + v_hel = coord.SphericalDifferential(d_lon=row['pml']/np.cos(gal.b)*u.mas/u.yr, + d_lat=row['pmb']*u.mas/u.yr, + d_distance=row['rv']*u.km/u.s) - # stupid check - vxyz_i = vhel_to_gal(c.icrs, pm=pm, rv=rv, - vcirc=0*u.km/u.s, - vlsr=[0.,0,0]*u.km/u.s) + # stupid check + vxyz_i = vhel_to_gal(icrs, v_hel, + vcirc=0*u.km/u.s, + vlsr=[0.,0,0]*u.km/u.s) - vxyz = vhel_to_gal(c.galactic, pm=pm, rv=rv, - vcirc=0*u.km/u.s, - vlsr=[0.,0,0]*u.km/u.s) + vxyz = vhel_to_gal(gal, v_hel, + vcirc=0*u.km/u.s, + vlsr=[0.,0,0]*u.km/u.s) - assert vxyz_i.shape == vxyz.shape + assert vxyz_i.shape == vxyz.shape - true_UVW = [row['U'],row['V'],row['W']]*u.km/u.s - found_UVW = vxyz - np.testing.assert_allclose(true_UVW.value, found_UVW.value, atol=1.) + true_UVW = np.array([row['U'], row['V'], row['W']]) + UVW = vxyz.d_xyz.to(u.km/u.s).value - # some sanity checks - first, some convenience definitions - g = coord.Galactic(l=0*u.deg, b=0*u.deg).transform_to(coord.ICRS) - galcen_frame = coord.Galactocentric(galcen_ra=g.ra, - galcen_dec=g.dec, - z_sun=0*u.kpc) + # catalog values are rounded + assert np.allclose(UVW, true_UVW, rtol=1E-2, atol=0.1) # -------------------------------------------------------------------- # l = 0 # without LSR and circular velocity - c = coord.SkyCoord(ra=galcen_frame.galcen_ra, dec=galcen_frame.galcen_dec, distance=2*u.kpc) - pm = [0., 0]*u.mas/u.yr - rv = 20*u.km/u.s - vxyz = vhel_to_gal(c.galactic, pm=pm, rv=rv, + c = coord.SkyCoord(ra=self.galcen_frame.galcen_ra, + dec=self.galcen_frame.galcen_dec, + distance=2*u.kpc) + v_hel = coord.SphericalDifferential(d_lon=0*u.mas/u.yr, + d_lat=0*u.mas/u.yr, + d_distance=20*u.km/u.s) + vxyz = vhel_to_gal(c.galactic, v_hel, vcirc=0*u.km/u.s, vlsr=[0.,0,0]*u.km/u.s, - galactocentric_frame=galcen_frame) - np.testing.assert_allclose(vxyz.to(u.km/u.s).value, [20,0,0.], atol=1E-12) + galactocentric_frame=self.galcen_frame) + assert np.allclose(vxyz.d_xyz.to(u.km/u.s).value, [20,0,0.], + atol=1E-12) # with LSR and circular velocity - c = coord.SkyCoord(ra=galcen_frame.galcen_ra, dec=galcen_frame.galcen_dec, distance=2*u.kpc) - pm = [0., 0]*u.mas/u.yr - rv = 20*u.km/u.s - vxyz = vhel_to_gal(c.galactic, pm=pm, rv=rv, + c = coord.SkyCoord(ra=self.galcen_frame.galcen_ra, + dec=self.galcen_frame.galcen_dec, + distance=2*u.kpc) + vxyz = vhel_to_gal(c.galactic, v_hel, vcirc=200*u.km/u.s, vlsr=[-20.,0,10]*u.km/u.s, - galactocentric_frame=galcen_frame) - np.testing.assert_allclose(vxyz.to(u.km/u.s).value, [0,200,10], atol=1E-12) + galactocentric_frame=self.galcen_frame) + assert np.allclose(vxyz.d_xyz.to(u.km/u.s).value, [0,200,10], + atol=1E-12) # l = 90 # with LSR and circular velocity - c = coord.SkyCoord(l=90*u.deg, b=0*u.deg, distance=2*u.kpc, frame=coord.Galactic) - pm = [0., 0]*u.mas/u.yr - rv = 20*u.km/u.s - vxyz = vhel_to_gal(c.galactic, pm=pm, rv=rv, + c = coord.SkyCoord(l=90*u.deg, b=0*u.deg, + distance=2*u.kpc, frame=coord.Galactic) + vxyz = vhel_to_gal(c.galactic, v_hel, vcirc=200*u.km/u.s, vlsr=[-20.,0,10]*u.km/u.s, - galactocentric_frame=galcen_frame) - np.testing.assert_allclose(vxyz.to(u.km/u.s).value, [-20,220,10], atol=1E-5) + galactocentric_frame=self.galcen_frame) + assert np.allclose(vxyz.d_xyz.to(u.km/u.s).value, [-20,220,10], + atol=1E-5) # l = 180 # with LSR and circular velocity - c = coord.SkyCoord(l=180*u.deg, b=0*u.deg, distance=2*u.kpc, frame=coord.Galactic) - pm = [0., 0]*u.mas/u.yr - rv = 20*u.km/u.s - vxyz = vhel_to_gal(c.galactic, pm=pm, rv=rv, + c = coord.SkyCoord(l=180*u.deg, b=0*u.deg, + distance=2*u.kpc, frame=coord.Galactic) + vxyz = vhel_to_gal(c.galactic, v_hel, vcirc=200*u.km/u.s, vlsr=[-20.,0,10]*u.km/u.s, - galactocentric_frame=galcen_frame) - np.testing.assert_allclose(vxyz.to(u.km/u.s).value, [-40,200,10], atol=1E-12) + galactocentric_frame=self.galcen_frame) + assert np.allclose(vxyz.d_xyz.to(u.km/u.s).value, [-40,200,10], + atol=1E-12) # l = 270 # with LSR and circular velocity - c = coord.SkyCoord(l=270*u.deg, b=0*u.deg, distance=2*u.kpc, frame=coord.Galactic) - pm = [0., 0]*u.mas/u.yr - rv = 20*u.km/u.s - vxyz = vhel_to_gal(c.galactic, pm=pm, rv=rv, + c = coord.SkyCoord(l=270*u.deg, b=0*u.deg, + distance=2*u.kpc, frame=coord.Galactic) + vxyz = vhel_to_gal(c.galactic, v_hel, vcirc=200*u.km/u.s, vlsr=[-20.,0,10]*u.km/u.s, - galactocentric_frame=galcen_frame) - np.testing.assert_allclose(vxyz.to(u.km/u.s).value, [-20,180,10], atol=1E-5) + galactocentric_frame=self.galcen_frame) + assert np.allclose(vxyz.d_xyz.to(u.km/u.s).value, [-20,180,10], + atol=1E-5) def test_vhel_to_gal_array(self): # test all together d = self.data - c = coord.SkyCoord(ra=d['ra']*u.deg, dec=d['dec']*u.deg, distance=d['dist']*u.pc) - pm = np.vstack((d['pml'], d['pmb']))*u.mas/u.yr - rv = d['rv']*u.km/u.s + c = coord.SkyCoord(ra=d['ra']*u.deg, dec=d['dec']*u.deg, + distance=d['dist']*u.pc) + icrs = c.icrs + gal = c.galactic + v_hel = coord.SphericalDifferential(d_lon=d['pml']/np.cos(gal.b)*u.mas/u.yr, + d_lat=d['pmb']*u.mas/u.yr, + d_distance=d['rv']*u.km/u.s) # stupid check - vxyz_i = vhel_to_gal(c.icrs, pm=pm, rv=rv, + vxyz_i = vhel_to_gal(icrs, v_hel, vcirc=0*u.km/u.s, vlsr=[0.,0,0]*u.km/u.s) - vxyz = vhel_to_gal(c.galactic, pm=pm, rv=rv, + vxyz = vhel_to_gal(gal, v_hel, vcirc=0*u.km/u.s, vlsr=[0.,0,0]*u.km/u.s) assert vxyz_i.shape == vxyz.shape # check values - true_UVW = np.vstack((d['U'],d['V'],d['W']))*u.km/u.s - found_UVW = vxyz - np.testing.assert_allclose(true_UVW.value, found_UVW.value, atol=1.) + true_UVW = np.array([d['U'], d['V'], d['W']]) + UVW = vxyz.d_xyz.to(u.km/u.s).value - def test_vgal_to_hel_single(self): - - # test a single entry - row = self.data[0] - c = coord.SkyCoord(ra=row['ra']*u.deg, dec=row['dec']*u.deg, distance=row['dist']*u.pc) - pm = [row['pml'],row['pmb']]*u.mas/u.yr - rv = row['rv']*u.km/u.s + # catalog values are rounded + assert np.allclose(UVW, true_UVW, rtol=1E-2, atol=0.1) - true_pmrv = (pm[0], pm[1], rv) - vxyz = [row['U'],row['V'],row['W']]*u.km/u.s - pmrv = vgal_to_hel(c.galactic, vxyz=vxyz, - vcirc=0.*u.km/u.s, - vlsr=[0.,0,0]*u.km/u.s) - - for i in range(3): - np.testing.assert_allclose(pmrv[i].to(true_pmrv[i].unit).value, - true_pmrv[i].value, - atol=1.) + def test_vgal_to_hel_single(self): - # some sanity checks - first, some convenience definitions - g = coord.Galactic(l=0*u.deg, b=0*u.deg).transform_to(coord.ICRS) - frargs = dict(galcen_ra=g.ra, - galcen_dec=g.dec, - z_sun=0*u.kpc, - galcen_distance=8*u.kpc) - galcen_frame = coord.Galactocentric(**frargs) + for row in self.data: # test one entry at a time + c = coord.SkyCoord(ra=row['ra']*u.deg, dec=row['dec']*u.deg, + distance=row['dist']*u.pc) + gal = c.galactic + vxyz = [row['U'], row['V'], row['W']] * u.km/u.s + + vhel = vgal_to_hel(gal, vxyz, + vcirc=0.*u.km/u.s, + vlsr=[0.,0,0]*u.km/u.s, + galactocentric_frame=self.galcen_frame) + + # tolerance set by the catalog rounded numbers + assert quantity_allclose(vhel.d_lon * np.cos(gal.b), + row['pml'] * u.mas/u.yr, rtol=1E-2) + assert quantity_allclose(vhel.d_lat, row['pmb'] * u.mas/u.yr, + rtol=1E-2) + assert quantity_allclose(vhel.d_distance, row['rv'] * u.km/u.s, + rtol=1E-2) # -------------------------------------------------------------------- # l = 0 # without LSR and circular velocity - # c = coord.Galactocentric([6,0,0]*u.kpc,**frargs) - c = coord.SkyCoord(l=0*u.deg, b=0*u.deg, distance=2*u.kpc, frame=coord.Galactic) - vxyz = [20.,0,0]*u.km/u.s - pmv = vgal_to_hel(c.galactic, vxyz, - vcirc=0*u.km/u.s, - vlsr=[0.,0,0]*u.km/u.s, - galactocentric_frame=galcen_frame) - np.testing.assert_allclose(pmv[0].to(u.mas/u.yr).value, 0., atol=1E-12) - np.testing.assert_allclose(pmv[1].to(u.mas/u.yr).value, 0., atol=1E-12) - np.testing.assert_allclose(pmv[2].to(u.km/u.s).value, 20., atol=1E-12) + c = coord.SkyCoord(l=0*u.deg, b=0*u.deg, distance=2*u.kpc, + frame=coord.Galactic) + vxyz = [20., 0, 0]*u.km/u.s + vhel = vgal_to_hel(c.galactic, vxyz, + vcirc=0*u.km/u.s, + vlsr=[0.,0,0]*u.km/u.s, + galactocentric_frame=self.galcen_frame) + + assert np.allclose(vhel.d_lon.value, 0., atol=1E-12) + assert np.allclose(vhel.d_lat.value, 0., atol=1E-12) + assert np.allclose(vhel.d_distance.to(u.km/u.s).value, 20., atol=1E-12) + + vxyz = [20., 0, 50]*u.km/u.s + vhel = vgal_to_hel(c.galactic, vxyz, + vcirc=0*u.km/u.s, + vlsr=[0.,0,0]*u.km/u.s, + galactocentric_frame=self.galcen_frame) + + assert np.allclose(vhel.d_lon.value, 0., atol=1E-5) # TODO: astropy precision issues + with u.set_enabled_equivalencies(u.dimensionless_angles()): + assert quantity_allclose(vhel.d_lat, 50*u.km/u.s / (2*u.kpc), + atol=1E-10*u.mas/u.yr) + assert quantity_allclose(vhel.d_distance.to(u.km/u.s), vxyz[0], + atol=1E-10*u.km/u.s) # with LSR and circular velocity - c = coord.SkyCoord(l=0*u.deg, b=0*u.deg, distance=2*u.kpc, frame=coord.Galactic) - vxyz = [20.,0,0]*u.km/u.s - pmv = vgal_to_hel(c.galactic, vxyz, - vcirc=-200*u.km/u.s, - vlsr=[0.,0,10]*u.km/u.s, - galactocentric_frame=galcen_frame) + vxyz = [20., 0, 50]*u.km/u.s + vhel = vgal_to_hel(c.galactic, vxyz, + vcirc=-200*u.km/u.s, + vlsr=[0., 0, 10]*u.km/u.s, + galactocentric_frame=self.galcen_frame) with u.set_enabled_equivalencies(u.dimensionless_angles()): - np.testing.assert_allclose(pmv[0].to(u.mas/u.yr).value, - ((200.*u.km/u.s)/(2*u.kpc)).to(u.mas/u.yr).value, - atol=1E-12) - np.testing.assert_allclose(pmv[1].to(u.mas/u.yr).value, - ((-10.*u.km/u.s)/(2*u.kpc)).to(u.mas/u.yr).value, - atol=1E-4) - np.testing.assert_allclose(pmv[2].to(u.km/u.s).value, 20., atol=1E-12) + assert quantity_allclose(vhel.d_lon, + (200.*u.km/u.s) / (2*u.kpc), + atol=1E-10*u.mas/u.yr) + assert quantity_allclose(vhel.d_lat, + (40.*u.km/u.s) / (2*u.kpc), + atol=1E-6*u.mas/u.yr) + + assert quantity_allclose(vhel.d_distance, 20.*u.km/u.s, + atol=1E-10*u.km/u.s) def test_vgal_to_hel_array(self): # test all together d = self.data - c = coord.SkyCoord(ra=d['ra']*u.deg, dec=d['dec']*u.deg, distance=d['dist']*u.pc) + c = coord.SkyCoord(ra=d['ra']*u.deg, dec=d['dec']*u.deg, + distance=d['dist']*u.pc) + pm = np.vstack([d['pml'],d['pmb']])*u.mas/u.yr rv = d['rv']*u.km/u.s - true_pmrv = (pm[0], pm[1], rv) - vxyz = np.vstack((d['U'],d['V'],d['W']))*u.km/u.s - pmrv = vgal_to_hel(c.galactic, vxyz=vxyz, + vxyz = np.vstack((d['U'], d['V'], d['W']))*u.km/u.s + vhel = vgal_to_hel(c.galactic, vxyz, vcirc=0.*u.km/u.s, - vlsr=[0.,0,0]*u.km/u.s) + vlsr=[0.,0,0]*u.km/u.s, + galactocentric_frame=self.galcen_frame) - for i in range(3): - np.testing.assert_allclose(pmrv[i].to(true_pmrv[i].unit).value, - true_pmrv[i].value, - atol=1.) + # tolerance set by the catalog rounded numbers + assert quantity_allclose(vhel.d_lon * np.cos(c.galactic.b), pm[0], + rtol=1E-2) + assert quantity_allclose(vhel.d_lat, pm[1], rtol=1E-2) + assert quantity_allclose(vhel.d_distance, rv, rtol=5E-3) def test_roundtrip_icrs(self): np.random.seed(42) @@ -322,20 +351,18 @@ def test_roundtrip_icrs(self): pm = np.random.uniform(-20,20,size=(2,n)) * u.mas/u.yr vr = np.random.normal(0., 75., size=n)*u.km/u.s - mua,mud = pm # initial + vhel = coord.SphericalDifferential(d_lon=pm[0], d_lat=pm[1], + d_distance=vr) # first to galactocentric - vxyz = vhel_to_gal(c.icrs, pm=pm, rv=vr) + vxyz = vhel_to_gal(c.icrs, vhel) # then back again, wooo - pmv = vgal_to_hel(c.icrs, vxyz=vxyz) + vhel2 = vgal_to_hel(c.icrs, vxyz) - mua2,mud2 = pmv[:2] - vr2 = pmv[2] - - np.testing.assert_allclose(mua.to(u.mas/u.yr).value, mua2.to(u.mas/u.yr).value, atol=1e-12) - np.testing.assert_allclose(mud.to(u.mas/u.yr).value, mud2.to(u.mas/u.yr).value, atol=1e-12) - np.testing.assert_allclose(vr.to(u.km/u.s).value, vr2.to(u.km/u.s).value, atol=1e-12) + for c in vhel.components: + assert quantity_allclose(getattr(vhel, c), getattr(vhel2, c), + rtol=1e-12) def test_roundtrip_gal(self): np.random.seed(42) @@ -348,18 +375,16 @@ def test_roundtrip_gal(self): pm = np.random.uniform(-20,20,size=(2,n)) * u.mas/u.yr vr = np.random.normal(0., 75., size=n)*u.km/u.s - mul,mub = pm # initial + vhel = coord.SphericalDifferential(d_lon=pm[0], d_lat=pm[1], + d_distance=vr) # first to galactocentric - vxyz = vhel_to_gal(c.galactic, pm=pm, rv=vr) + vxyz = vhel_to_gal(c.galactic, vhel) # then back again, wooo - pmv = vgal_to_hel(c.galactic, vxyz=vxyz) - - mul2,mub2 = pmv[:2] - vr2 = pmv[2] + vhel2 = vgal_to_hel(c.galactic, vxyz) - # TODO: why such bad roundtripping? - np.testing.assert_allclose(mul.to(u.mas/u.yr).value, mul2.to(u.mas/u.yr).value, rtol=1E-4, atol=1e-12) - np.testing.assert_allclose(mub.to(u.mas/u.yr).value, mub2.to(u.mas/u.yr).value, rtol=1E-4, atol=1e-12) - np.testing.assert_allclose(vr.to(u.km/u.s).value, vr2.to(u.km/u.s).value, rtol=1E-4, atol=1e-12) + # TODO: why such bad roundtripping??? + for c in vhel.components: + assert quantity_allclose(getattr(vhel, c), getattr(vhel2, c), + rtol=1e-4) diff --git a/gala/coordinates/tests/test_velocity_transforms.py b/gala/coordinates/tests/test_velocity_transforms.py deleted file mode 100644 index 883c6b052..000000000 --- a/gala/coordinates/tests/test_velocity_transforms.py +++ /dev/null @@ -1,89 +0,0 @@ -# coding: utf-8 - -""" Test velocity transformations. """ - -from __future__ import division, print_function - - -# Standard library -import logging - -# Third-party -import numpy as np -from astropy import log as logger -import astropy.coordinates as coord -import astropy.units as u - -# This project -from ..velocity_coord_transforms import * - -logger.setLevel(logging.DEBUG) - -class TestTransforms(object): - - def setup(self): - self.pos = ([[10.,0,0], [10,0,0], [4,4,0]] * u.kpc).T - self.vel = ([[0.,100,500], [0,-100,-500], [100.,100.,-500]] * u.km/u.s).T - - self.pos_repr = coord.CartesianRepresentation(self.pos) - - def test_cartesian_to(self): - for i,func in enumerate([cartesian_to_spherical, - cartesian_to_physicsspherical, - cartesian_to_cylindrical]): - - # dimensionless - vsph1 = func(self.pos.value * u.dimensionless_unscaled, - self.vel.value * u.dimensionless_unscaled) - assert vsph1.unit == u.dimensionless_unscaled - - # astropy coordinates - cpos = coord.SkyCoord(self.pos_repr) - vsph2 = func(cpos, self.vel) - assert vsph2.unit == u.km/u.s - - # astropy representation - vsph3 = func(self.pos_repr, self.vel) - assert vsph3.unit == u.km/u.s - - np.testing.assert_allclose(vsph1.value, vsph2.value, atol=1E-10) - np.testing.assert_allclose(vsph1.value, vsph3.value, atol=1E-10) - - true_v = self.vel.copy() - if func == cartesian_to_physicsspherical: - true_v[2] *= -1. - - np.testing.assert_allclose(vsph1[:,:2].value, true_v[:,:2].value) - - assert vsph1[0,2] > 0. # vr - assert np.sign(vsph1[2,2]) == np.sign(true_v[2,2]) - np.testing.assert_allclose(np.sqrt(np.sum(vsph1**2,axis=0)).value, - np.sqrt(np.sum(true_v**2,axis=0)).value) - - def test_to_cartesian(self): - for i,func in enumerate([spherical_to_cartesian, - physicsspherical_to_cartesian, - cylindrical_to_cartesian]): - true_v = ([[0.,100,500], [0,-100,-500], [0.,np.sqrt(2)*100,-500]] * u.km/u.s).T - - # dimensionless - vsph1 = func(self.pos.value * u.dimensionless_unscaled, - self.vel.value * u.dimensionless_unscaled) - assert vsph1.unit == u.dimensionless_unscaled - - # astropy coordinates - cpos = coord.SkyCoord(self.pos_repr) - vsph2 = func(cpos, self.vel) - assert vsph2.unit == u.km/u.s - - # astropy representation - vsph3 = func(self.pos_repr, self.vel) - assert vsph3.unit == u.km/u.s - - np.testing.assert_allclose(vsph1.value, vsph2.value, atol=1E-10) - np.testing.assert_allclose(vsph1.value, vsph3.value, atol=1E-10) - - if func == physicsspherical_to_cartesian: - true_v[2] *= -1. - - np.testing.assert_allclose(vsph1.value, true_v.value, atol=1E-10) diff --git a/gala/coordinates/velocity_frame_transforms.py b/gala/coordinates/velocity_frame_transforms.py index 01fd4df78..a0c2e389d 100644 --- a/gala/coordinates/velocity_frame_transforms.py +++ b/gala/coordinates/velocity_frame_transforms.py @@ -13,10 +13,11 @@ import astropy.units as u from astropy.coordinates.builtin_frames.galactocentric import _ROLL0 as ROLL0 try: - from astropy.coordinates.matrix_utilities import rotation_matrix, matrix_transpose, matrix_product + from astropy.coordinates.matrix_utilities import (rotation_matrix, + matrix_product) ASTROPY_1_3 = True except ImportError: - from .matrix_utilities import rotation_matrix, matrix_transpose, matrix_product + from .matrix_utilities import rotation_matrix, matrix_product ASTROPY_1_3 = False if not ASTROPY_1_3: @@ -32,8 +33,8 @@ # This is the default circular velocity and LSR peculiar velocity of the Sun # TODO: make this a config item? -VCIRC = 220.*u.km/u.s -VLSR = [10., 5.25, 7.17]*u.km/u.s +VCIRC = 220. * u.km/u.s +VLSR = [10., 5.25, 7.17] * u.km/u.s def _icrs_gctc_velocity_matrix(galactocentric_frame): """ @@ -56,69 +57,86 @@ def _icrs_gctc_velocity_matrix(galactocentric_frame): return matrix_product(M4, M3, M1, M2) # this is right: 4,3,1,2 -def vgal_to_hel(coordinate, vxyz, vcirc=VCIRC, vlsr=VLSR, galactocentric_frame=None): +def vgal_to_hel(coordinate, velocity, vcirc=None, vlsr=None, + galactocentric_frame=None): r""" - Convert a Galactocentric, cartesian velocity to a Heliocentric velocity in - spherical coordinates (e.g., proper motion and radial velocity). + Convert a Galactocentric velocity to a Heliocentric velocity. - The frame of the input coordinate determines the output frame of the proper motions. - For example, if the input coordinate is in the ICRS frame, the proper motions - returned will be :math:`(\mu_\alpha\cos\delta,\mu_\delta)`. This function also - handles array inputs (see examples below). + The frame of the input coordinate determines the output frame of the + heliocentric velocity. For example, if the input coordinate is in the ICRS + frame, heliocentric velocity will also be in the ICRS. Parameters ---------- coordinate : :class:`~astropy.coordinates.SkyCoord`, :class:`~astropy.coordinates.BaseCoordinateFrame` - This is most commonly a :class:`~astropy.coordinates.SkyCoord` object, but - alternatively, it can be any coordinate frame object that is transformable to the - Galactocentric frame. - vxyz : :class:`~astropy.units.Quantity`, iterable - Cartesian velocity components :math:`(v_x,v_y,v_z)`. This should either be a single - :class:`~astropy.units.Quantity` object with shape (3,N), or an iterable - object with 3 :class:`~astropy.units.Quantity` objects as elements. + This is most commonly a :class:`~astropy.coordinates.SkyCoord` object, + but alternatively, it can be any coordinate frame object that is + transformable to the Galactocentric frame. + velocity : :class:`~astropy.coordinates.BaseDifferential`, :class:`~astropy.units.Quantity`, iterable + If not provided as a Differential instance, the velocity components are + assumed to be Cartesian :math:`(v_x,v_y,v_z)` and should either + be a single :class:`~astropy.units.Quantity` object with shape ``(3,N)`` + or an iterable object with 3 :class:`~astropy.units.Quantity` objects as + elements. vcirc : :class:`~astropy.units.Quantity` (optional) Circular velocity of the Sun. vlsr : :class:`~astropy.units.Quantity` (optional) Velocity of the Sun relative to the local standard of rest (LSR). galactocentric_frame : :class:`~astropy.coordinates.Galactocentric` (optional) - An instantiated :class:`~astropy.coordinates.Galactocentric` frame object with - custom parameters for the Galactocentric coordinates. For example, if you want - to set your own position of the Galactic center, you can pass in a frame with - custom `galcen_ra` and `galcen_dec`. + An instantiated :class:`~astropy.coordinates.Galactocentric` frame + object with custom parameters for the Galactocentric coordinates. For + example, if you want to set your own position of the Galactic center, + you can pass in a frame with custom `galcen_ra` and `galcen_dec`. Returns ------- - pmv : tuple - A tuple containing the proper motions (in Galactic coordinates) and - radial velocity, all as :class:`~astropy.units.Quantity` objects. + helio_velocity : tuple + The computed heliocentric velocity. Examples -------- >>> import astropy.units as u >>> import astropy.coordinates as coord - >>> c = coord.Galactocentric(x=15.*u.kpc, y=13.*u.kpc, z=2.*u.kpc) + >>> c = coord.Galactocentric([15., 13, 2]*u.kpc) >>> vxyz = [-115., 100., 95.]*u.km/u.s >>> icrs = c.transform_to(coord.ICRS) >>> vgal_to_hel(icrs, vxyz) # doctest: +FLOAT_CMP - (, , ) + - >>> c = coord.Galactocentric([[15.,11.],[13,21.],[2.,-7]]*u.kpc) - >>> vxyz = [[-115.,11.], [100.,-21.], [95.,103]]*u.km/u.s + >>> c = coord.Galactocentric([[15.,11.], [13,21.], [2.,-7]]*u.kpc) + >>> vxyz = [[-115.,11.], [100.,-21.], [95.,103]] * u.km/u.s >>> icrs = c.transform_to(coord.ICRS) - >>> vgal_to_hel(icrs, vxyz) # doctest: +FLOAT_CMP - (, , ) + >>> vhel = vgal_to_hel(icrs, vxyz) + >>> vhel # doctest: +FLOAT_CMP + + >>> vhel.d_lon # doctest: +FLOAT_CMP + """ + if vcirc is None: + vcirc = VCIRC + + if vlsr is None: + vlsr = VLSR + if galactocentric_frame is None: galactocentric_frame = coord.Galactocentric - # so I don't accidentally modify in place - vxyz = vxyz.copy() - c = coordinate + v = velocity + if isinstance(v, coord.BaseDifferential): + pos = c.represent_as(v.base_representation) + vxyz = v.represent_as(coord.CartesianDifferential, base=pos).d_xyz + + else: + vxyz = v + R = _icrs_gctc_velocity_matrix(galactocentric_frame) # remove circular and LSR velocities @@ -137,14 +155,18 @@ def vgal_to_hel(coordinate, vxyz, vcirc=VCIRC, vlsr=VLSR, galactocentric_frame=N vr = np.sum(x_icrs * v_icrs, axis=0) / d - mua = ((x_icrs[0]*v_icrs[1] - v_icrs[0]*x_icrs[1]) / dxy**2).to(u.mas/u.yr, equivalencies=u.dimensionless_angles()) - mua_cosd = (mua * dxy / d).to(u.mas/u.yr, equivalencies=u.dimensionless_angles()) - mud = (-(x_icrs[2]*(x_icrs[0]*v_icrs[0] + x_icrs[1]*v_icrs[1]) - dxy**2*v_icrs[2]) / d**2 / dxy).to(u.mas/u.yr, equivalencies=u.dimensionless_angles()) + mua = ((x_icrs[0]*v_icrs[1] - v_icrs[0]*x_icrs[1]) / dxy**2) + mua_cosd = (mua * dxy / d).to(u.mas/u.yr, + equivalencies=u.dimensionless_angles()) + + mud = (-(x_icrs[2]*(x_icrs[0]*v_icrs[0] + x_icrs[1]*v_icrs[1]) - + dxy**2*v_icrs[2]) / d**2 / dxy) + mud = mud.to(u.mas/u.yr, equivalencies=u.dimensionless_angles()) pm_radec = (mua_cosd, mud) if c.name == 'icrs': - pm = u.Quantity(map(np.atleast_1d,pm_radec)) + pm = u.Quantity(map(np.atleast_1d, pm_radec)) else: # transform ICRS proper motions to whatever frame @@ -154,47 +176,41 @@ def vgal_to_hel(coordinate, vxyz, vcirc=VCIRC, vlsr=VLSR, galactocentric_frame=N vr = vr.reshape(()) pm = (pm[0].reshape(()), pm[1].reshape(())) - return tuple(pm) + (vr,) + # NOTE: remember that d_lon does not include the cos(lat) term! + return coord.SphericalDifferential(d_lon=pm[0] / np.cos(c.spherical.lat), + d_lat=pm[1], + d_distance=vr) -def vhel_to_gal(coordinate, pm, rv, vcirc=VCIRC, vlsr=VLSR, galactocentric_frame=None): +def vhel_to_gal(coordinate, velocity, vcirc=None, vlsr=None, + galactocentric_frame=None): r""" - Convert a Heliocentric velocity in spherical coordinates (e.g., proper motion - and radial velocity) in the ICRS or Galactic frame to a Galactocentric, cartesian - velocity. + Convert a Heliocentric velocity to a Galactocentric velocity. The frame of the input coordinate determines how to interpret the given - proper motions. For example, if the input coordinate is in the ICRS frame, the - proper motions are assumed to be :math:`(\mu_\alpha\cos\delta,\mu_\delta)`. This - function also handles array inputs (see examples below). - - TODO: Roundtrip using galactic coordinates only maintains relative precision - of ~1E-5. Why? + proper motions. For example, if the input coordinate is in the ICRS frame, + the input velocity is assumed to be as well. Parameters ---------- coordinate : :class:`~astropy.coordinates.SkyCoord`, :class:`~astropy.coordinates.BaseCoordinateFrame` - This is most commonly a :class:`~astropy.coordinates.SkyCoord` object, but - alternatively, it can be any coordinate frame object that is transformable to the - Galactocentric frame. - pm : :class:`~astropy.units.Quantity` or iterable of :class:`~astropy.units.Quantity` objects - Proper motion in the same frame as the coordinate. For example, if your input - coordinate is in :class:`~astropy.coordinates.ICRS`, then the proper motion is - assumed to be in this frame as well. The order of elements should always be - proper motion in (longitude, latitude), and should have shape (2,N). The longitude - component is assumed to have the cosine of the latitude already multiplied in, so - that in ICRS, for example, this would be :math:`\mu_\alpha\cos\delta`. - rv : :class:`~astropy.units.Quantity` - Barycentric radial velocity. Should have shape (1,N) or (N,). + This is most commonly a :class:`~astropy.coordinates.SkyCoord` object, + but alternatively, it can be any coordinate frame object that is + transformable to the Galactocentric frame. + velocity : :class:`~astropy.coordinates.BaseDifferential`, iterable + If not provided as a Differential instance, the velocity input is + assumed to be a length-3 iterable containing proper motion components + and radial velocity. The proper motion in longitude should *not* contain + the cos(latitude) term. vcirc : :class:`~astropy.units.Quantity` (optional) Circular velocity of the Sun. vlsr : :class:`~astropy.units.Quantity` (optional) Velocity of the Sun relative to the local standard of rest (LSR). galactocentric_frame : :class:`~astropy.coordinates.Galactocentric` (optional) - An instantiated :class:`~astropy.coordinates.Galactocentric` frame object with - custom parameters for the Galactocentric coordinates. For example, if you want - to set your own position of the Galactic center, you can pass in a frame with - custom `galcen_ra` and `galcen_dec`. + An instantiated :class:`~astropy.coordinates.Galactocentric` frame + object with custom parameters for the Galactocentric coordinates. For + example, if you want to set your own position of the Galactic center, + you can pass in a frame with custom `galcen_ra` and `galcen_dec`. Returns ------- @@ -207,33 +223,59 @@ def vhel_to_gal(coordinate, pm, rv, vcirc=VCIRC, vlsr=VLSR, galactocentric_frame >>> import astropy.units as u >>> import astropy.coordinates as coord - >>> c = coord.SkyCoord(ra=196.5*u.degree, dec=-10.33*u.deg, distance=16.2*u.kpc) - >>> pm = [-1.53, 3.5]*u.mas/u.yr - >>> rv = 161.4*u.km/u.s - >>> vhel_to_gal(c, pm=pm, rv=rv) # doctest: +FLOAT_CMP - - - >>> c = coord.SkyCoord(ra=[196.5,51.3]*u.degree, dec=[-10.33,2.1]*u.deg, distance=[16.2,11.]*u.kpc) - >>> pm = [[-1.53,4.5], [3.5,10.9]]*u.mas/u.yr - >>> rv = [161.4,-210.2]*u.km/u.s - >>> vhel_to_gal(c, pm=pm, rv=rv) # doctest: +FLOAT_CMP - + >>> c = coord.SkyCoord(ra=196.5*u.degree, dec=-10.33*u.deg, + ... distance=16.2*u.kpc) + >>> vhel = coord.SphericalDifferential(d_lon=-1.53*u.mas/u.yr, + ... d_lat=3.5*u.mas/u.yr, + ... d_distance=161.4*u.km/u.s) + >>> vgal = vhel_to_gal(c, vhel) + >>> vgal # doctest: +FLOAT_CMP + + + >>> c = coord.SkyCoord(ra=[196.5,51.3]*u.degree, dec=[-10.33,2.1]*u.deg, + ... distance=[16.2,11.]*u.kpc) + >>> vhel = coord.SphericalDifferential(d_lon=[-1.53, 4.5]*u.mas/u.yr, + ... d_lat=[3.5, 10.9]*u.mas/u.yr, + ... d_distance=[161.4, -210.2]*u.km/u.s) + >>> vgal = vhel_to_gal(c, vhel) + >>> vgal # doctest: +FLOAT_CMP + """ + if vcirc is None: + vcirc = VCIRC + + if vlsr is None: + vlsr = VLSR + if galactocentric_frame is None: galactocentric_frame = coord.Galactocentric - # make sure this is a coordinate and get the frame for later use c = coordinate + v = velocity + if isinstance(v, coord.BaseDifferential): + pos = c.represent_as(v.base_representation) + vsph = v.represent_as(coord.SphericalDifferential, base=pos) + vsph = [vsph.d_lon, vsph.d_lat, vsph.d_distance] + + else: + vsph = v + + pm = vsph[:2] + rv = vsph[2] if c.name == 'icrs': - pm_radec = u.Quantity(map(np.atleast_1d,pm)) + pm_radec = u.Quantity([np.atleast_1d(pm[0]) * np.cos(c.dec), + np.atleast_1d(pm[1])]) icrs = c else: + pm[0] = pm[0] * np.cos(c.spherical.lat) + # pmra *includes* cos(dec) term! pm_radec = transform_proper_motion(c, pm, coord.ICRS) icrs = c.transform_to(coord.ICRS) @@ -261,14 +303,13 @@ def vhel_to_gal(coordinate, pm, rv, vcirc=VCIRC, vlsr=VLSR, galactocentric_frame v_gc[i] = v_gc[i] + vlsr[i] if c.isscalar: - return v_gc.reshape((3,)) - else: - return v_gc + v_gc = v_gc.reshape((3,)) + return coord.CartesianDifferential(*v_gc) # ----------------------------------------------------------------------------- -def vgsr_to_vhel(coordinate, vgsr, vcirc=VCIRC, vlsr=VLSR): +def vgsr_to_vhel(coordinate, vgsr, vcirc=None, vlsr=None): """ Convert a radial velocity in the Galactic standard of rest (GSR) to a barycentric radial velocity. @@ -293,6 +334,12 @@ def vgsr_to_vhel(coordinate, vgsr, vcirc=VCIRC, vlsr=VLSR): """ + if vcirc is None: + vcirc = VCIRC + + if vlsr is None: + vlsr = VLSR + c = coord.SkyCoord(coordinate) g = c.galactic l,b = g.l, g.b @@ -311,7 +358,7 @@ def vgsr_to_vhel(coordinate, vgsr, vcirc=VCIRC, vlsr=VLSR): return vhel -def vhel_to_vgsr(coordinate, vhel, vcirc=VCIRC, vlsr=VLSR): +def vhel_to_vgsr(coordinate, vhel, vcirc=None, vlsr=None): """ Convert a velocity from a heliocentric radial velocity to the Galactic standard of rest (GSR). @@ -336,6 +383,12 @@ def vhel_to_vgsr(coordinate, vhel, vcirc=VCIRC, vlsr=VLSR): """ + if vcirc is None: + vcirc = VCIRC + + if vlsr is None: + vlsr = VLSR + c = coord.SkyCoord(coordinate) g = c.galactic l,b = g.l, g.b diff --git a/gala/dynamics/__init__.py b/gala/dynamics/__init__.py index 3499ad538..22f139fe8 100644 --- a/gala/dynamics/__init__.py +++ b/gala/dynamics/__init__.py @@ -5,3 +5,4 @@ from .nonlinear import * from .plot import * from .util import * +from .representation_nd import * diff --git a/gala/dynamics/actionangle.py b/gala/dynamics/actionangle.py index d0ed58c62..2e3f39f00 100644 --- a/gala/dynamics/actionangle.py +++ b/gala/dynamics/actionangle.py @@ -88,7 +88,7 @@ def fit_isochrone(orbit, m0=2E11, b0=1.): Parameters ---------- - orbit : `~gala.dynamics.CartesianOrbit` + orbit : `~gala.dynamics.Orbit` m0 : numeric (optional) Initial mass guess. b0 : numeric (optional) @@ -141,7 +141,7 @@ def fit_harmonic_oscillator(orbit, omega0=[1.,1.,1.]): Parameters ---------- - orbit : `~gala.dynamics.CartesianOrbit` + orbit : `~gala.dynamics.Orbit` omega0 : array_like (optional) Initial frequency guess. @@ -185,7 +185,7 @@ def fit_toy_potential(orbit, force_harmonic_oscillator=False): Parameters ---------- - orbit : `~gala.dynamics.CartesianOrbit` + orbit : `~gala.dynamics.Orbit` force_harmonic_oscillator : bool (optional) Force using the harmonic oscillator potential as the toy potential. @@ -430,7 +430,7 @@ def _single_orbit_find_actions(orbit, N_max, toy_potential=None, Parameters ---------- - orbit : `~gala.dynamics.CartesianOrbit` + orbit : `~gala.dynamics.Orbit` N_max : int Maximum integer Fourier mode vector length, |n|. toy_potential : Potential (optional) @@ -517,7 +517,7 @@ def find_actions(orbit, N_max, force_harmonic_oscillator=False, toy_potential=No Parameters ---------- - orbit : `~gala.dynamics.CartesianOrbit` + orbit : `~gala.dynamics.Orbit` N_max : int Maximum integer Fourier mode vector length, :math:`|\boldsymbol{n}|`. force_harmonic_oscillator : bool (optional) diff --git a/gala/dynamics/analyticactionangle.py b/gala/dynamics/analyticactionangle.py index 804597d13..3b7826f65 100644 --- a/gala/dynamics/analyticactionangle.py +++ b/gala/dynamics/analyticactionangle.py @@ -16,7 +16,6 @@ # Project from ..potential import (Hamiltonian, PotentialBase, IsochronePotential, HarmonicOscillatorPotential) -from ..coordinates import physicsspherical_to_cartesian from ..util import atleast_2d __all__ = ['isochrone_to_aa', 'harmonic_oscillator_to_aa'] @@ -41,7 +40,7 @@ def isochrone_to_aa(w, potential): Parameters ---------- - w : :class:`gala.dynamics.CartesianPhaseSpacePosition`, :class:`gala.dynamics.CartesianOrbit` + w : :class:`gala.dynamics.PhaseSpacePosition`, :class:`gala.dynamics.Orbit` potential : :class:`gala.potential.IsochronePotential`, dict An instance of the potential to use for computing the transformation to angle-action coordinates. Or, a dictionary of parameters used to @@ -70,9 +69,17 @@ def isochrone_to_aa(w, potential): raise ValueError("Unbound particle. (E = {})".format(E)) # convert position, velocity to spherical polar coordinates - sph,vsph = w.represent_as(coord.PhysicsSphericalRepresentation) - r,phi,theta = map(np.squeeze, [sph.r.value, sph.phi.value, sph.theta.value]) - vr,vphi,vtheta = map(np.squeeze, vsph.value) + w_sph = w.represent_as(coord.PhysicsSphericalRepresentation) + r,phi,theta = map(np.squeeze, [w_sph.r.decompose(usys).value, + w_sph.phi.radian, + w_sph.theta.radian]) + + ang_unit = u.radian/usys['time'] + vr,phi_dot,theta_dot = map(np.squeeze, [w_sph.d_r.decompose(usys).value, + w_sph.d_phi.to(ang_unit).value, + w_sph.d_theta.to(ang_unit).value]) + vphi = r*np.sin(theta) * phi_dot + vtheta = r*theta_dot # ---------------------------- # Compute the actions @@ -268,6 +275,9 @@ def F(x, y): vtheta = -L*sini*np.cos(psi)/np.sin(theta)/r vphi = Lz / (r*np.sin(theta)) + d_phi = vphi / (r*np.sin(theta)) + d_theta = vtheta / r + # phi sinu = np.sin(psi)*cosi/np.sin(theta) @@ -282,9 +292,13 @@ def F(x, y): # We now need to convert from spherical polar coord to cart. coord. pos = coord.PhysicsSphericalRepresentation(r=r*u.dimensionless_unscaled, phi=phi*u.rad, theta=theta*u.rad) - x = pos.represent_as(coord.CartesianRepresentation).xyz.value - v = physicsspherical_to_cartesian(pos, [vr,vphi,vtheta]*u.dimensionless_unscaled).value - return x,v + pos = pos.represent_as(coord.CartesianRepresentation) + x = pos.xyz.value + + vel = coord.PhysicsSphericalDifferential(d_phi=d_phi, d_theta=d_theta, d_r=vr) + v = vel.represent_as(coord.CartesianDifferential, base=pos).d_xyz.value + + return x, v def harmonic_oscillator_to_aa(w, potential): """ @@ -304,17 +318,17 @@ def harmonic_oscillator_to_aa(w, potential): Parameters ---------- - w : :class:`gala.dynamics.CartesianPhaseSpacePosition`, :class:`gala.dynamics.CartesianOrbit` + w : :class:`gala.dynamics.PhaseSpacePosition`, :class:`gala.dynamics.Orbit` potential : Potential """ usys = potential.units if usys is not None: - x = w.pos.decompose(usys).value - v = w.vel.decompose(usys).value + x = w.pos.xyz.decompose(usys).value + v = w.vel.d_xyz.decompose(usys).value else: - x = w.pos.value - v = w.vel.value + x = w.pos.xyz.value + v = w.vel.d_xyz.value _new_omega_shape = (3,) + tuple([1]*(len(x.shape)-1)) # compute actions -- just energy (hamiltonian) over frequency diff --git a/gala/dynamics/core.py b/gala/dynamics/core.py index cabd07808..0cb10a1fa 100644 --- a/gala/dynamics/core.py +++ b/gala/dynamics/core.py @@ -2,157 +2,124 @@ from __future__ import division, print_function - # Standard library import warnings import inspect # Third-party -from astropy import log as logger import astropy.coordinates as coord import astropy.units as u -uno = u.dimensionless_unscaled import numpy as np +from six import string_types # Project -from .plot import three_panel -from ..coordinates import velocity_coord_transforms as vtrans +from . import representation_nd as rep_nd +from .plot import plot_projections from ..coordinates import vgal_to_hel -from ..units import UnitSystem, DimensionlessUnitSystem +from ..units import UnitSystem, DimensionlessUnitSystem, _greek_letters from ..util import atleast_2d -__all__ = ['CartesianPhaseSpacePosition', 'combine'] +__all__ = ['PhaseSpacePosition', 'CartesianPhaseSpacePosition'] class PhaseSpacePosition(object): - - # ------------------------------------------------------------------------ - # Display - # - def __repr__(self): - return "<{} N={}, shape={}>".format(self.__class__.__name__, self.ndim, self.shape) - - def __str__(self): - # TODO: should show some arrays instead -- get inspiration from CartesianRepresentation - return "{} {}D, {}".format(self.__class__.__name__, self.ndim, self.shape) - - # ------------------------------------------------------------------------ - # Shape and size - # - @property - def ndim(self): - """ - Number of coordinate dimensions. 1/2 of the phase-space dimensionality. - - .. warning:: - - This is *not* the number of axes in the position or velocity - arrays. That is accessed by doing ``{}.pos.ndim``. - - Returns - ------- - n : int - - """.format(self.__class__.__name__) - return self.pos.shape[0] - - @property - def shape(self): - """ - .. warning:: - - This is *not* the shape of the position or velocity - arrays. That is accessed by doing ``{}.pos.shape``. - - Returns - ------- - shp : tuple - - """.format(self.__class__.__name__) - return self.pos.shape[1:] - -class CartesianPhaseSpacePosition(PhaseSpacePosition): """ - Represents phase-space positions in Cartesian coordinates, e.g., - positions and conjugate momenta (velocities). + Represents phase-space positions, i.e. positions and conjugate momenta + (velocities). - .. warning:: + The class can be instantiated with Astropy representation objects (e.g., + :class:`~astropy.coordinates.CartesianRepresentation`), Astropy + :class:`~astropy.units.Quantity` objects, or plain Numpy arrays. - This is an experimental class. The API may change in a future release! + If passing in representation objects, the default representation is taken to + be the class that is passed in. - The position and velocity quantities (arrays) can have an arbitrary - number of dimensions, but the first axis (0) has special meaning:: + If passing in Quantity or Numpy array instances for both position and + velocity, they are assumed to be Cartesian. Array inputs are interpreted as + dimensionless quantities. The input position and velocity objects can have + an arbitrary number of (broadcastable) dimensions. For Quantity or array + inputs, the first axis (0) has special meaning:: - - `axis=0` is the coordinate dimension (e.g., x, y, z) + - `axis=0` is the coordinate dimension (e.g., x, y, z for Cartesian) So if the input position array, `pos`, has shape `pos.shape = (3, 100)`, this would represent 100 3D positions (`pos[0]` is `x`, `pos[1]` is `y`, - etc.). The same is true for velocity. The position and velocity arrays - must have the same shape. - - If the input position and velocity are arrays or array-like rather than - :class:`~astropy.units.Quantity` objects, they are internally stored with - dimensionles units. + etc.). The same is true for velocity. Parameters ---------- - pos : :class:`~astropy.units.Quantity`, array_like + pos : :class:`~astropy.coordinates.BaseRepresentation`, :class:`~astropy.units.Quantity`, array_like Positions. If a numpy array (e.g., has no units), this will be stored as a dimensionless :class:`~astropy.units.Quantity`. See the note above about the assumed meaning of the axes of this object. - vel : :class:`~astropy.units.Quantity`, array_like + vel : :class:`~astropy.coordinates.BaseDifferential`, :class:`~astropy.units.Quantity`, array_like Velocities. If a numpy array (e.g., has no units), this will be stored as a dimensionless :class:`~astropy.units.Quantity`. See the note above about the assumed meaning of the axes of this object. frame : :class:`~gala.potential.FrameBase` (optional) The reference frame of the input phase-space positions. + """ def __init__(self, pos, vel, frame=None): - # make sure position and velocity input are 2D - pos = atleast_2d(pos, insert_axis=1) - vel = atleast_2d(vel, insert_axis=1) + if not isinstance(pos, coord.BaseRepresentation): + # assume Cartesian if not specified + if not hasattr(pos, 'unit'): + pos = pos * u.one + + # 3D coordinates get special treatment + ndim = pos.shape[0] + if ndim == 3: + # TODO: HACK: until this stuff is in astropy core + if isinstance(pos, coord.BaseRepresentation): + kw = [(k,getattr(pos,k)) for k in pos.components] + pos = getattr(coord, pos.__class__.__name__)(**kw) + + else: + pos = coord.CartesianRepresentation(pos) + + else: + pos = rep_nd.NDCartesianRepresentation(pos) - # make sure position and velocity have at least a dimensionless unit! - if not hasattr(pos, 'unit'): - pos = pos * uno + else: + ndim = 3 - if not hasattr(vel, 'unit'): - vel = vel * uno + if not isinstance(vel, coord.BaseDifferential): + # assume representation is same as pos if not specified + if not hasattr(vel, 'unit'): + vel = vel * u.one - if (pos.unit == uno and vel.unit != uno): - logger.warning("Position unit is dimensionless but velocity unit is not." - "Are you sure that's what you want?") - elif (vel.unit == uno and pos.unit != uno): - logger.warning("Velocity unit is dimensionless but position unit is not." - "Are you sure that's what you want?") + if ndim == 3: + name = pos.__class__.get_name() + Diff = coord.representation.DIFFERENTIAL_CLASSES[name] + vel = Diff(*vel) + else: + Diff = rep_nd.NDCartesianDifferential + vel = Diff(vel) # make sure shape is the same - for i in range(pos.ndim): - if pos.shape[i] != vel.shape[i]: - raise ValueError("Position and velocity must have the same shape " - "{} vs {}".format(pos.shape, vel.shape)) + if pos.shape != vel.shape: + raise ValueError("Position and velocity must have the same shape " + "{} vs {}".format(pos.shape, vel.shape)) from ..potential.frame import FrameBase if frame is not None and not isinstance(frame, FrameBase): - raise TypeError("Input reference frame must be a gala.potential.FrameBase " + raise TypeError("Input reference frame must be a FrameBase " "subclass instance.") self.pos = pos self.vel = vel self.frame = frame + self.ndim = ndim - def __getitem__(self, slyce): - if isinstance(slyce, np.ndarray) or isinstance(slyce, list): - _slyce = np.array(slyce) - _slyce = (slice(None),) + (slyce,) - else: - try: - _slyce = (slice(None),) + tuple(slyce) - except TypeError: - _slyce = (slice(None),) + (slyce,) + for name in pos.components: + setattr(self, name, getattr(pos,name)) + + for name in vel.components: + setattr(self, name, getattr(vel,name)) - return self.__class__(pos=self.pos[_slyce], - vel=self.vel[_slyce], + def __getitem__(self, slyce): + return self.__class__(pos=self.pos[slyce], + vel=self.vel[slyce], frame=self.frame) # ------------------------------------------------------------------------ @@ -171,28 +138,28 @@ def represent_as(self, Representation): Returns ------- - pos : `~astropy.coordinates.BaseRepresentation` - vel : `~astropy.units.Quantity` - The velocity in the new representation. All components - have units of velocity -- e.g., to get an angular velocity - in spherical representations, you'll need to divide by the radius. + new_psp : `gala.dynamics.PhaseSpacePosition` """ + if self.ndim != 3: - raise ValueError("Representation changes require a 3D (ndim=3) " - "position and velocity.") + raise ValueError("Can only change representation for " + "ndim=3 instances.") # get the name of the desired representation - rep_name = Representation.get_name() + if isinstance(Representation, string_types): + name = Representation + else: + name = Representation.get_name() - # transform the position - car_pos = coord.CartesianRepresentation(self.pos) - new_pos = car_pos.represent_as(Representation) + Representation = coord.representation.REPRESENTATION_CLASSES[name] + Differential = coord.representation.DIFFERENTIAL_CLASSES[name] - # transform the velocity - vfunc = getattr(vtrans, "cartesian_to_{}".format(rep_name)) - new_vel = vfunc(car_pos, self.vel) + new_pos = self.pos.represent_as(Representation) + new_vel = self.vel.represent_as(Differential, self.pos) - return new_pos, new_vel + return self.__class__(pos=new_pos, + vel=new_vel, + frame=self.frame) def to_frame(self, frame, current_frame=None, **kwargs): """ @@ -206,7 +173,8 @@ def to_frame(self, frame, current_frame=None, **kwargs): The current frame the phase-space position is in. **kwargs Any additional arguments are passed through to the individual frame - transformation functions (see: `~gala.potential.frame.builtin.transformations`). + transformation functions (see: + `~gala.potential.frame.builtin.transformations`). Returns ------- @@ -215,21 +183,23 @@ def to_frame(self, frame, current_frame=None, **kwargs): """ - from ..potential.frame.builtin import StaticFrame, ConstantRotatingFrame from ..potential.frame.builtin import transformations as frame_trans if ((inspect.isclass(frame) and issubclass(frame, coord.BaseCoordinateFrame)) or - isinstance(frame, coord.BaseCoordinateFrame)): + isinstance(frame, coord.BaseCoordinateFrame)): import warnings - warnings.warn("This function now expects a `gala.potential.FrameBase` instance. To " - " transform to an Astropy coordinate frame, use the `.to_coord_frame()`" - " method instead.", DeprecationWarning) + warnings.warn("This function now expects a " + "`gala.potential.FrameBase` instance. To transform to" + " an Astropy coordinate frame, use the " + "`.to_coord_frame()` method instead.", + DeprecationWarning) return self.to_coord_frame(frame=frame, **kwargs) if self.frame is None and current_frame is None: - raise ValueError("If no frame was specified when this {} was initialized, you must " - "pass the current frame in via the current_frame argument to " - "transform to a new frame.") + raise ValueError("If no frame was specified when this {} was " + "initialized, you must pass the current frame in " + "via the current_frame argument to transform to a " + "new frame.") elif self.frame is not None and current_frame is None: current_frame = self.frame @@ -244,11 +214,11 @@ def to_frame(self, frame, current_frame=None, **kwargs): else: trans_func = getattr(frame_trans, func_name) - pos,vel = trans_func(current_frame, frame, self, **kwargs) - return CartesianPhaseSpacePosition(pos=pos, vel=vel, frame=frame) + pos, vel = trans_func(current_frame, frame, self, **kwargs) + return PhaseSpacePosition(pos=pos, vel=vel, frame=frame) - def to_coord_frame(self, frame, galactocentric_frame=coord.Galactocentric(), - vcirc=None, vlsr=None): + def to_coord_frame(self, frame, + galactocentric_frame=None, vcirc=None, vlsr=None): """ Transform the orbit from Galactocentric, cartesian coordinates to Heliocentric coordinates in the specified Astropy coordinate frame. @@ -267,34 +237,46 @@ def to_coord_frame(self, frame, galactocentric_frame=coord.Galactocentric(), ------- c : :class:`~astropy.coordinates.BaseCoordinateFrame` An instantiated coordinate frame. - v : tuple - A tuple of velocities represented as - :class:`~astropy.units.Quantity` objects. + v : :class:`~astropy.coordinates.BaseDifferential` + An astropy differential class containing the velocity components. """ if self.ndim != 3: - raise ValueError("Frame transformations require a 3D (ndim=3) " - "position and velocity.") + raise ValueError("Can only change representation for " + "ndim=3 instances.") - car_pos = coord.CartesianRepresentation(self.pos) - gc_c = galactocentric_frame.realize_frame(car_pos) - c = gc_c.transform_to(frame) + if galactocentric_frame is None: + galactocentric_frame = coord.Galactocentric() kw = dict() kw['galactocentric_frame'] = galactocentric_frame - if vcirc is not None: - kw['vcirc'] = vcirc - if vlsr is not None: - kw['vlsr'] = vlsr + kw['vcirc'] = vcirc + kw['vlsr'] = vlsr + + # first we need to turn the position into a Galactocentric instance + gc_c = galactocentric_frame.realize_frame(self.pos) + c = gc_c.transform_to(frame) + rep = c.represent_as(c.representation) + + # get the corresponding differential class + new_Diff = coord.representation.DIFFERENTIAL_CLASSES[rep.get_name()] + vxyz = self.vel.represent_as(coord.CartesianDifferential, + base=self.pos).d_xyz + v = vgal_to_hel(c, vxyz, galactocentric_frame=galactocentric_frame) + v = v.represent_as(new_Diff, + base=self.pos.represent_as(v.base_representation)) - v = vgal_to_hel(c, self.vel, **kw) return c, v # Convenience attributes @property def cartesian(self): - return self.pos, self.vel + return self.represent_as(coord.CartesianRepresentation) + + @property + def spherical(self): + return self.represent_as(coord.SphericalRepresentation) @property def spherical(self): @@ -322,15 +304,30 @@ def w(self, units=None): Will have shape ``(2*ndim,...)``. """ - if (units is None or isinstance(units, DimensionlessUnitSystem)) \ - and (self.pos.unit == uno and self.vel.unit == uno): + if self.ndim == 3: + cart = self.cartesian + else: + cart = self + + xyz = cart.pos.xyz + d_xyz = cart.vel.d_xyz + + x_unit = xyz.unit + v_unit = d_xyz.unit + if ((units is None or isinstance(units, DimensionlessUnitSystem)) and + (x_unit == u.one and v_unit == u.one)): units = DimensionlessUnitSystem() elif units is None: raise ValueError("A UnitSystem must be provided.") - x = self.pos.decompose(units).value - v = self.vel.decompose(units).value + x = xyz.decompose(units).value + if x.ndim < 2: + x = atleast_2d(x, insert_axis=1) + + v = d_xyz.decompose(units).value + if v.ndim < 2: + v = atleast_2d(v, insert_axis=1) return np.vstack((x,v)) @@ -357,15 +354,18 @@ def from_w(cls, w, units=None, **kwargs): """.format(name=cls.__name__) + w = np.array(w) + ndim = w.shape[0]//2 pos = w[:ndim] vel = w[ndim:] - # TODO: this is bad form - UnitSystem should know what to do with a Dimensionless + # TODO: this is bad form - UnitSystem should know what to do with a + # Dimensionless if units is not None and not isinstance(units, DimensionlessUnitSystem): units = UnitSystem(units) pos = pos*units['length'] - vel = vel*units['length']/units['time'] # velocity in w is from _core_units + vel = vel*units['length']/units['time'] # from _core_units return cls(pos=pos, vel=vel, **kwargs) @@ -385,7 +385,7 @@ def kinetic_energy(self): E : :class:`~astropy.units.Quantity` The kinetic energy. """ - return 0.5*np.sum(self.vel**2, axis=0) + return 0.5 * self.vel.norm()**2 def potential_energy(self, potential): r""" @@ -405,7 +405,8 @@ def potential_energy(self, potential): E : :class:`~astropy.units.Quantity` The potential energy. """ - return potential.value(self.pos) + # TODO: check that potential ndim is consistent with here + return potential.value(self) def energy(self, hamiltonian): r""" @@ -425,10 +426,11 @@ def energy(self, hamiltonian): if isinstance(hamiltonian, PotentialBase): from ..potential import Hamiltonian - warnings.warn("This function now expects a `Hamiltonian` instance instead of " - "a `PotentialBase` subclass instance. If you are using a " - "static reference frame, you just need to pass your " - "potential object in to the Hamiltonian constructor to use, e.g., " + warnings.warn("This function now expects a `Hamiltonian` instance " + "instead of a `PotentialBase` subclass instance. If " + "you are using a static reference frame, you just " + "need to pass your potential object in to the " + "Hamiltonian constructor to use, e.g., " "Hamiltonian(potential).", DeprecationWarning) hamiltonian = Hamiltonian(hamiltonian) @@ -459,42 +461,104 @@ def angular_momentum(self): >>> import astropy.units as u >>> pos = np.array([1., 0, 0]) * u.au >>> vel = np.array([0, 2*np.pi, 0]) * u.au/u.yr - >>> w = {}(pos, vel) + >>> w = PhaseSpacePosition(pos, vel) >>> w.angular_momentum() - """.format(self.__class__.__name__) - return np.cross(self.pos.value, self.vel.value, axis=0) * self.pos.unit * self.vel.unit + """ + cart = self.represent_as(coord.CartesianRepresentation) + return cart.pos.cross(cart.vel).xyz # ------------------------------------------------------------------------ # Misc. useful methods # - def plot(self, **kwargs): + def _plot_prepare(self, components, units): + """ + Prepare the ``PhaseSpacePosition`` or subclass for passing to a plotting + routine to plot all projections of the object. """ - Plot the positions in all projections. This is a thin wrapper around - `~gala.dynamics.three_panel` -- the docstring for this function is - included here. - .. warning:: + # components to plot + if components is None: + components = self.pos.components + n_comps = len(components) + + # if units not specified, get units from the components + if units is not None: + if isinstance(units, u.UnitBase): + units = [units]*n_comps # global unit + + elif len(units) != n_comps: + raise ValueError('You must specify a unit for each axis, or a ' + 'single unit for all axes.') + + labels = [] + x = [] + for i,name in enumerate(components): + val = getattr(self, name) - This will currently fail for orbits with fewer than 3 dimensions. + if units is not None: + val = val.to(units[i]) + unit = units[i] + else: + unit = val.unit + + if val.unit != u.one: + uu = unit.to_string(format='latex_inline') + unit_str = ' [{}]'.format(uu) + else: + unit_str = '' + + # Figure out how to fancy display the component name + if name.startswith('d_'): + dot = True + name = name[2:] + else: + dot = False + + if name in _greek_letters: + name = r"\{}".format(name) + + if dot: + name = "\dot{{{}}}".format(name) + + labels.append('${}$'.format(name) + unit_str) + x.append(val.value) + + return x, labels + + def plot(self, components=None, units=None, **kwargs): + """ + Plot the positions in all projections. This is a wrapper around + `~gala.dynamics.plot_projections` for fast access and quick + visualization. All extra keyword arguments are passed to that function + (the docstring for this function is included here for convenience). Parameters ---------- + components : iterable (optional) + A list of component names (strings) to plot. By default, this is the + Cartesian positions ``['x', 'y', 'z']``. To plot Cartesian + velocities, pass in the velocity component names + ``['d_x', 'd_y', 'd_z']``. + units : `~astropy.units.UnitBase`, iterable (optional) + A single unit or list of units to display the components in. relative_to : bool (optional) Plot the values relative to this value or values. autolim : bool (optional) Automatically set the plot limits to be something sensible. axes : array_like (optional) Array of matplotlib Axes objects. - triangle : bool (optional) - Make a triangle plot instead of plotting all projections in a single row. subplots_kwargs : dict (optional) Dictionary of kwargs passed to :func:`~matplotlib.pyplot.subplots`. labels : iterable (optional) - List or iterable of axis labels as strings. They should correspond to the - dimensions of the input orbit. + List or iterable of axis labels as strings. They should correspond to + the dimensions of the input orbit. + plot_function : callable (optional) + The ``matplotlib`` plot function to use. By default, this is + :func:`~matplotlib.pyplot.scatter`, but can also be, e.g., + :func:`~matplotlib.pyplot.plot`. **kwargs - All other keyword arguments are passed to :func:`~matplotlib.pyplot.scatter`. + All other keyword arguments are passed to the ``plot_function``. You can pass in any of the usual style kwargs like ``color=...``, ``marker=...``, etc. @@ -503,65 +567,77 @@ def plot(self, **kwargs): fig : `~matplotlib.Figure` """ - _label_unit = '' - if self.pos.unit != u.dimensionless_unscaled: - _label_unit = ' [{}]'.format(self.pos.unit) + + try: + import matplotlib.pyplot as plt + except ImportError: + msg = 'matplotlib is required for visualization.' + raise ImportError(msg) + + if components is None: + components = self.pos.components + + x,labels = self._plot_prepare(components=components, + units=units) default_kwargs = { 'marker': '.', - 'color': 'k', - 'labels': ('$x${}'.format(_label_unit), - '$y${}'.format(_label_unit), - '$z${}'.format(_label_unit)) + 'labels': labels, + 'plot_function': plt.scatter, + 'autolim': False } for k,v in default_kwargs.items(): kwargs[k] = kwargs.get(k, v) - return three_panel(self.pos.value, **kwargs) + fig = plot_projections(x, **kwargs) -def combine(args): - """ - Combine the input `PhaseSpacePosition` objects into a single object. + if self.pos.get_name() == 'cartesian' and \ + all([not c.startswith('d_') for c in components]): + for ax in fig.axes: + ax.set(aspect='equal', adjustable='datalim') - Parameters - ---------- - args : iterable - Multiple instances of `PhaseSpacePosition`. + return fig - Returns - ------- - obj : `~gala.dynamics.PhaseSpacePosition` - A single objct with positions and velocities stacked along the last axis. - """ + # ------------------------------------------------------------------------ + # Display + # + def __repr__(self): + return "<{} {}, dim={}, shape={}>".format(self.__class__.__name__, + self.pos.get_name(), + self.ndim, + self.pos.shape) - ndim = None - pos_unit = None - vel_unit = None - all_pos = [] - all_vel = [] - for x in args: - if ndim is None: - ndim = x.ndim - pos_unit = x.pos.unit - vel_unit = x.vel.unit - else: - if x.ndim != ndim: - raise ValueError("All objects must have the same dimensionality.") + def __str__(self): + return "pos={}\nvel={}".format(self.pos, self.vel) + + # ------------------------------------------------------------------------ + # Shape and size + # + + @property + def shape(self): + """ + .. warning:: + This is *not* the shape of the position or velocity + arrays. That is accessed by doing ``obj.pos.shape``. - pos = x.pos - if pos.ndim < 2: - pos = pos[...,np.newaxis] + Returns + ------- + shp : tuple - vel = x.vel - if vel.ndim < 2: - vel = vel[...,np.newaxis] + """ + return self.pos.shape - all_pos.append(pos.to(pos_unit).value) - all_vel.append(vel.to(vel_unit).value) +class CartesianPhaseSpacePosition(PhaseSpacePosition): - all_pos = np.hstack(all_pos)*pos_unit - all_vel = np.hstack(all_vel)*vel_unit + def __init__(self, pos, vel, frame=None): + """ + Deprecated. + """ - return CartesianPhaseSpacePosition(pos=all_pos, vel=all_vel) + warnings.warn("This class is now deprecated! Use the general interface " + "provided by PhaseSpacePosition instead.", + DeprecationWarning) + super(CartesianPhaseSpacePosition, self).__init__(pos, vel, frame=frame) diff --git a/gala/dynamics/mockstream/core.py b/gala/dynamics/mockstream/core.py index 88a1a0e71..a72de68d2 100644 --- a/gala/dynamics/mockstream/core.py +++ b/gala/dynamics/mockstream/core.py @@ -11,7 +11,7 @@ import numpy as np # Project -from .. import CartesianPhaseSpacePosition, Orbit +from .. import PhaseSpacePosition, Orbit from ...potential import Hamiltonian, CPotentialBase from ...integrate import DOPRI853Integrator, LeapfrogIntegrator from ._mockstream import _mock_stream_dop853, _mock_stream_leapfrog, _mock_stream_animate @@ -59,7 +59,7 @@ def mock_stream(hamiltonian, prog_orbit, prog_mass, k_mean, k_disp, Returns ------- - stream : `~gala.dynamics.CartesianPhaseSpacePosition` + stream : `~gala.dynamics.PhaseSpacePosition` """ @@ -148,7 +148,7 @@ def mock_stream(hamiltonian, prog_orbit, prog_mass, k_mean, k_disp, else: raise RuntimeError("Should never get here...") - return CartesianPhaseSpacePosition.from_w(w=stream_w.T, units=hamiltonian.units) + return PhaseSpacePosition.from_w(w=stream_w.T, units=hamiltonian.units) def streakline_stream(hamiltonian, prog_orbit, prog_mass, release_every=1, Integrator=DOPRI853Integrator, Integrator_kwargs=dict(), @@ -183,7 +183,7 @@ def streakline_stream(hamiltonian, prog_orbit, prog_mass, release_every=1, Returns ------- - stream : `~gala.dynamics.CartesianPhaseSpacePosition` + stream : `~gala.dynamics.PhaseSpacePosition` """ k_mean = np.zeros(6) @@ -245,7 +245,7 @@ def fardal_stream(hamiltonian, prog_orbit, prog_mass, release_every=1, Returns ------- - stream : `~gala.dynamics.CartesianPhaseSpacePosition` + stream : `~gala.dynamics.PhaseSpacePosition` """ k_mean = np.zeros(6) @@ -311,7 +311,7 @@ def dissolved_fardal_stream(hamiltonian, prog_orbit, prog_mass, t_disrupt, relea Returns ------- - stream : `~gala.dynamics.CartesianPhaseSpacePosition` + stream : `~gala.dynamics.PhaseSpacePosition` """ diff --git a/gala/dynamics/mockstream/tests/test_mockstream.py b/gala/dynamics/mockstream/tests/test_mockstream.py index 514a9ecd4..7b1fa0e5d 100644 --- a/gala/dynamics/mockstream/tests/test_mockstream.py +++ b/gala/dynamics/mockstream/tests/test_mockstream.py @@ -12,7 +12,7 @@ # Custom from ....potential import Hamiltonian, NFWPotential, HernquistPotential -from ....dynamics import CartesianPhaseSpacePosition +from ....dynamics import PhaseSpacePosition from ....integrate import DOPRI853Integrator, LeapfrogIntegrator, RK5Integrator from ....units import galactic @@ -27,8 +27,8 @@ def test_mock_stream(Integrator, kwargs): units=galactic) ham = Hamiltonian(potential) - w0 = CartesianPhaseSpacePosition(pos=[0.,15.,0]*u.kpc, - vel=[-0.13,0,0]*u.kpc/u.Myr) + w0 = PhaseSpacePosition(pos=[0.,15.,0]*u.kpc, + vel=[-0.13,0,0]*u.kpc/u.Myr) prog = ham.integrate_orbit(w0, dt=-2., n_steps=1023) prog = prog[::-1] @@ -46,9 +46,9 @@ def test_mock_stream(Integrator, kwargs): # pl.show() # return - assert stream.pos.shape == (3,2048) # two particles per step + assert stream.pos.shape == (2048,) # two particles per step - diff = np.abs(stream[-2:].pos - prog[-1].pos) + diff = np.abs(stream[-2:].pos.xyz - prog[-1:].pos.xyz) assert np.allclose(diff[0].value, 0.) assert np.allclose(diff[1,0].value, diff[1,1].value) assert np.allclose(diff[2].value, 0.) @@ -61,8 +61,8 @@ def test_each_type(mock_func, extra_kwargs): units=galactic) ham = Hamiltonian(potential) - w0 = CartesianPhaseSpacePosition(pos=[0.,15.,0]*u.kpc, - vel=[-0.13,0,0]*u.kpc/u.Myr) + w0 = PhaseSpacePosition(pos=[0.,15.,0]*u.kpc, + vel=[-0.13,0,0]*u.kpc/u.Myr) prog = ham.integrate_orbit(w0, dt=-2., n_steps=1023) prog = prog[::-1] @@ -76,7 +76,7 @@ def test_each_type(mock_func, extra_kwargs): # plt.show() assert prog.t.shape == (1024,) - assert stream.pos.shape == (3,2048) # two particles per step + assert stream.pos.shape == (2048,) # two particles per step # ----------------------- # Test expected failures: @@ -103,7 +103,7 @@ def test_animate(tmpdir): np.random.seed(42) pot = HernquistPotential(m=1E11, c=1., units=galactic) - w0 = CartesianPhaseSpacePosition(pos=[5.,0,0]*u.kpc, vel=[0,0.1,0]*u.kpc/u.Myr) + w0 = PhaseSpacePosition(pos=[5.,0,0]*u.kpc, vel=[0,0.1,0]*u.kpc/u.Myr) orbit = pot.integrate_orbit(w0, dt=1., n_steps=1000, Integrator=DOPRI853Integrator) @@ -123,5 +123,5 @@ def test_animate(tmpdir): assert np.allclose(t, orbit.t.value) for idx in range(pos.shape[2]): - assert np.allclose(pos[:,-1,idx], stream.pos.value[:,idx], rtol=1E-4) - assert np.allclose(vel[:,-1,idx], stream.vel.value[:,idx], rtol=1E-4) + assert np.allclose(pos[:,-1,idx], stream.pos.xyz.value[:,idx], rtol=1E-4) + assert np.allclose(vel[:,-1,idx], stream.vel.d_xyz.value[:,idx], rtol=1E-4) diff --git a/gala/dynamics/nonlinear.py b/gala/dynamics/nonlinear.py index c31f5e90e..7ac15df1c 100644 --- a/gala/dynamics/nonlinear.py +++ b/gala/dynamics/nonlinear.py @@ -9,7 +9,7 @@ from scipy.signal import argrelmin # Project -from . import CartesianPhaseSpacePosition, CartesianOrbit +from . import PhaseSpacePosition, Orbit from ..potential.potential import PotentialBase __all__ = ['fast_lyapunov_max', 'lyapunov_max', 'surface_of_section'] @@ -45,7 +45,7 @@ def fast_lyapunov_max(w0, hamiltonian, dt, n_steps, d0=1e-5, ------- LEs : :class:`~astropy.units.Quantity` Lyapunov exponents calculated from each offset / deviation orbit. - orbit : `~gala.dynamics.CartesianOrbit` (optional) + orbit : `~gala.dynamics.Orbit` (optional) """ @@ -60,11 +60,11 @@ def fast_lyapunov_max(w0, hamiltonian, dt, n_steps, d0=1e-5, raise TypeError("Input Hamiltonian must contain a C-implemented " "potential and frame.") - if not isinstance(w0, CartesianPhaseSpacePosition): + if not isinstance(w0, PhaseSpacePosition): w0 = np.asarray(w0) ndim = w0.shape[0]//2 - w0 = CartesianPhaseSpacePosition(pos=w0[:ndim], - vel=w0[ndim:]) + w0 = PhaseSpacePosition(pos=w0[:ndim], + vel=w0[ndim:]) _w0 = np.squeeze(w0.w(hamiltonian.units)) if _w0.ndim > 1: @@ -82,8 +82,8 @@ def fast_lyapunov_max(w0, hamiltonian, dt, n_steps, d0=1e-5, except (TypeError, AttributeError): tunit = u.dimensionless_unscaled - orbit = CartesianOrbit.from_w(w=w, units=hamiltonian.units, - t=t*tunit, hamiltonian=hamiltonian) + orbit = Orbit.from_w(w=w, units=hamiltonian.units, + t=t*tunit, hamiltonian=hamiltonian) return l/tunit, orbit else: @@ -134,7 +134,7 @@ def lyapunov_max(w0, integrator, dt, n_steps, d0=1e-5, n_steps_per_pullback=10, ------- LEs : :class:`~astropy.units.Quantity` Lyapunov exponents calculated from each offset / deviation orbit. - orbit : `~gala.dynamics.CartesianOrbit` + orbit : `~gala.dynamics.Orbit` """ if units is not None: @@ -144,11 +144,11 @@ def lyapunov_max(w0, integrator, dt, n_steps, d0=1e-5, n_steps_per_pullback=10, pos_unit = u.dimensionless_unscaled vel_unit = u.dimensionless_unscaled - if not isinstance(w0, CartesianPhaseSpacePosition): + if not isinstance(w0, PhaseSpacePosition): w0 = np.asarray(w0) ndim = w0.shape[0]//2 - w0 = CartesianPhaseSpacePosition(pos=w0[:ndim]*pos_unit, - vel=w0[ndim:]*vel_unit) + w0 = PhaseSpacePosition(pos=w0[:ndim]*pos_unit, + vel=w0[ndim:]*vel_unit) _w0 = w0.w(units) ndim = 2*w0.ndim @@ -205,8 +205,8 @@ def lyapunov_max(w0, integrator, dt, n_steps, d0=1e-5, n_steps_per_pullback=10, except (TypeError, AttributeError): t_unit = u.dimensionless_unscaled - orbit = CartesianOrbit.from_w(w=full_w[:,:total_steps_taken], - units=units, t=full_ts[:total_steps_taken]*t_unit) + orbit = Orbit.from_w(w=full_w[:,:total_steps_taken], + units=units, t=full_ts[:total_steps_taken]*t_unit) return LEs/t_unit, orbit def surface_of_section(orbit, plane_ix, interpolate=False): @@ -219,7 +219,7 @@ def surface_of_section(orbit, plane_ix, interpolate=False): Parameters ---------- - orbit : `~gala.dynamics.CartesianOrbit` + orbit : `~gala.dynamics.Orbit` plane_ix : int Integer that represents the coordinate to record crossings in. For example, for a 2D Hamiltonian where you want to make a SoS in diff --git a/gala/dynamics/orbit.py b/gala/dynamics/orbit.py index d8717cf25..c37040cf4 100644 --- a/gala/dynamics/orbit.py +++ b/gala/dynamics/orbit.py @@ -7,41 +7,280 @@ import warnings # Third-party +import astropy.coordinates as coord from astropy import log as logger import astropy.units as u -uno = u.dimensionless_unscaled import numpy as np from scipy.signal import argrelmin, argrelmax from scipy.interpolate import InterpolatedUnivariateSpline from scipy.optimize import minimize +from six import string_types # Project -from .core import CartesianPhaseSpacePosition +from .core import PhaseSpacePosition from .util import peak_to_peak_period -from .plot import plot_orbits +from .plot import plot_projections from ..util import atleast_2d +from ..units import dimensionless -__all__ = ['CartesianOrbit', 'combine'] +__all__ = ['Orbit', 'CartesianOrbit'] -class Orbit(object): +class Orbit(PhaseSpacePosition): + """ + Represents an orbit: positions and velocities (conjugate momenta) as a + function of time. + + The class can be instantiated with Astropy representation objects (e.g., + :class:`~astropy.coordinates.CartesianRepresentation`), Astropy + :class:`~astropy.units.Quantity` objects, or plain Numpy arrays. + + If passing in Quantity or Numpy array instances for both position and + velocity, they are assumed to be Cartesian. Array inputs are interpreted as + dimensionless quantities. The input position and velocity objects can have + an arbitrary number of (broadcastable) dimensions. For Quantity or array + inputs, the first axes have special meaning:: + + - `axis=0` is the coordinate dimension (e.g., x, y, z) + - `axis=1` is the time dimension + + So if the input position array, `pos`, has shape `pos.shape = (3, 100)`, + this would be a 3D orbit at 100 times (`pos[0]` is `x`, `pos[1]` is `y`, + etc.). For representing multiple orbits, the position array could have 3 + axes, e.g., it might have shape `pos.shape = (3, 100, 8)`, where this is + interpreted as a 3D position at 100 times for 8 different orbits. The same + is true for velocity. The position and velocity arrays must have the same + shape. + + If a time argument is specified, the position and velocity arrays must have + the same number of timesteps as the length of the time object:: + + len(t) == pos.shape[1] + + Parameters + ---------- + pos : :class:`~astropy.coordinates.BaseRepresentation`, :class:`~astropy.units.Quantity`, array_like + Positions. If a numpy array (e.g., has no units), this will be + stored as a dimensionless :class:`~astropy.units.Quantity`. See + the note above about the assumed meaning of the axes of this object. + vel : :class:`~astropy.coordinates.BaseDifferential`, :class:`~astropy.units.Quantity`, array_like + Velocities. If a numpy array (e.g., has no units), this will be + stored as a dimensionless :class:`~astropy.units.Quantity`. See + the note above about the assumed meaning of the axes of this object. + t : array_like, :class:`~astropy.units.Quantity` (optional) + Array of times. If a numpy array (e.g., has no units), this will be + stored as a dimensionless :class:`~astropy.units.Quantity`. + hamiltonian : `~gala.potential.Hamiltonian` (optional) + The Hamiltonian that the orbit was integrated in. + + """ + def __init__(self, pos, vel, t=None, + hamiltonian=None, potential=None, frame=None): + + super(Orbit, self).__init__(pos=pos, vel=vel) + + if self.pos.ndim < 1: + self.pos = self.pos.reshape(1) + self.vel = self.vel.reshape(1) + + # TODO: check that Hamiltonian ndim is consistent with here + + if t is not None: + t = np.atleast_1d(t) + if self.pos.shape[0] != len(t): + raise ValueError("Position and velocity must have the same " + "length along axis=1 as the length of the " + "time array {} vs {}" + .format(len(t), self.pos.shape[0])) + + if not hasattr(t, 'unit'): + t = t * u.one + + self.t = t + + if hamiltonian is not None: + self.potential = hamiltonian.potential + self.frame = hamiltonian.frame + + else: + self.potential = potential + self.frame = frame + + def __getitem__(self, slice_): + + if isinstance(slice_, np.ndarray) or isinstance(slice_, list): + slice_ = (slice_,) + + try: + slice_ = tuple(slice_) + except TypeError: + slice_ = (slice_,) + + kw = dict() + if self.t is not None: + kw['t'] = self.t[slice_[0]] + + pos = self.pos[slice_] + vel = self.vel[slice_] + + # if one time is sliced out, return a phasespaceposition + if isinstance(slice_[0], int): + return PhaseSpacePosition(pos=pos, vel=vel, frame=self.frame) + + else: + return self.__class__(pos=pos, vel=vel, + potential=self.potential, + frame=self.frame, **kw) + + @property + def hamiltonian(self): + if self.potential is None or self.frame is None: + return None + + try: + return self._hamiltonian + except AttributeError: + from ..potential import Hamiltonian + self._hamiltonian = Hamiltonian(potential=self.potential, + frame=self.frame) + + return self._hamiltonian + + def w(self, units=None): + """ + This returns a single array containing the phase-space positions. + + Parameters + ---------- + units : `~gala.units.UnitSystem` (optional) + The unit system to represent the position and velocity in + before combining into the full array. + + Returns + ------- + w : `~numpy.ndarray` + A numpy array of all positions and velocities, without units. + Will have shape ``(2*ndim,...)``. + + """ + + if units is None: + if self.hamiltonian is None: + units = dimensionless + else: + units = self.hamiltonian.units + + return super(Orbit, self).w(units=units) + + # ------------------------------------------------------------------------ + # Convert from Cartesian to other representations + # + def represent_as(self, Representation): + """ + Represent the position and velocity of the orbit in an alternate + coordinate system. Supports any of the Astropy coordinates + representation classes. + + Parameters + ---------- + Representation : :class:`~astropy.coordinates.BaseRepresentation` + The class for the desired representation. + + Returns + ------- + new_orbit : `gala.dynamics.Orbit` + """ + + if self.ndim != 3: + raise ValueError("Can only change representation for " + "ndim=3 instances.") + + # get the name of the desired representation + if isinstance(Representation, string_types): + name = Representation + else: + name = Representation.get_name() + + Representation = coord.representation.REPRESENTATION_CLASSES[name] + Differential = coord.representation.DIFFERENTIAL_CLASSES[name] + + new_pos = self.pos.represent_as(Representation) + new_vel = self.vel.represent_as(Differential, self.pos) + + return self.__class__(pos=new_pos, + vel=new_vel, + hamiltonian=self.hamiltonian) # ------------------------------------------------------------------------ # Shape and size # ------------------------------------------------------------------------ @property def ntimes(self): - return self.pos.shape[1] + return self.shape[0] @property def norbits(self): - if self.pos.ndim < 3: + if self.pos.xyz.ndim < 3: return 1 else: - return self.pos.shape[2] + return self.shape[2] # ------------------------------------------------------------------------ # Computed dynamical quantities - # ------------------------------------------------------------------------ + # + + def potential_energy(self, potential=None): + r""" + The potential energy *per unit mass*: + + .. math:: + + E_\Phi = \Phi(\boldsymbol{q}) + + Returns + ------- + E : :class:`~astropy.units.Quantity` + The potential energy. + """ + if self.hamiltonian is None and potential is None: + raise ValueError("To compute the potential energy, a potential" + " object must be provided!") + if potential is None: + potential = self.hamiltonian.potential + + return super(Orbit,self).potential_energy(potential) + + def energy(self, hamiltonian=None): + r""" + The total energy *per unit mass*: + + Returns + ------- + E : :class:`~astropy.units.Quantity` + The total energy. + """ + + if self.hamiltonian is None and hamiltonian is None: + raise ValueError("To compute the total energy, a hamiltonian" + " object must be provided!") + + from ..potential import PotentialBase + if isinstance(hamiltonian, PotentialBase): + from ..potential import Hamiltonian + + warnings.warn("This function now expects a `Hamiltonian` instance " + "instead of a `PotentialBase` subclass instance. If " + "you are using a static reference frame, you just " + "need to pass your potential object in to the " + "Hamiltonian constructor to use, e.g., Hamiltonian" + "(potential).", DeprecationWarning) + + hamiltonian = Hamiltonian(hamiltonian) + + if hamiltonian is None: + hamiltonian = self.hamiltonian + + return hamiltonian(self) + def pericenter(self, return_times=False, func=np.mean, interp_kwargs=None, minimize_kwargs=None): """ @@ -92,11 +331,11 @@ def pericenter(self, return_times=False, func=np.mean, # default scipy function kwargs interp_kwargs.setdefault('k', 3) - interp_kwargs.setdefault('ext', 3) # don't extrapolate, return boundary val + interp_kwargs.setdefault('ext', 3) # don't extrapolate, return boundary minimize_kwargs.setdefault('method', 'powell') # orbital radius - r = self.r.value + r = self.spherical.pos.r.value t = self.t.value _ix = argrelmin(r, mode='wrap')[0] @@ -111,7 +350,7 @@ def pericenter(self, return_times=False, func=np.mean, res = minimize(interp_func, t[j], **minimize_kwargs) refined_times[i] = res.x - peri = interp_func(refined_times) * self.r.unit + peri = interp_func(refined_times) * self.cartesian.pos.x.unit if return_times: return peri, refined_times @@ -169,11 +408,11 @@ def apocenter(self, return_times=False, func=np.mean, # default scipy function kwargs interp_kwargs.setdefault('k', 3) - interp_kwargs.setdefault('ext', 3) # don't extrapolate, return boundary val + interp_kwargs.setdefault('ext', 3) # don't extrapolate, return boundary minimize_kwargs.setdefault('method', 'powell') # orbital radius - r = self.r.value + r = self.spherical.pos.r.value t = self.t.value _ix = argrelmax(r, mode='wrap')[0] @@ -188,7 +427,7 @@ def apocenter(self, return_times=False, func=np.mean, res = minimize(lambda x: -interp_func(x), t[ix], **minimize_kwargs) refined_times[i] = res.x - apo = interp_func(refined_times) * self.r.unit + apo = interp_func(refined_times) * self.cartesian.pos.x.unit if return_times: return apo, refined_times @@ -239,12 +478,13 @@ def estimate_period(self, radial=True): " Specify a time array when creating this object.") if radial: - r = self.r.value + r = self.spherical.pos.r.value if self.norbits == 1: T = peak_to_peak_period(self.t.value, r) T = T * self.t.unit else: - T = [peak_to_peak_period(self.t.value, r[:,n]) for n in range(r.shape[1])] + T = [peak_to_peak_period(self.t.value, r[:,n]) + for n in range(r.shape[1])] T = T * self.t.unit else: @@ -252,247 +492,6 @@ def estimate_period(self, radial=True): return T -class CartesianOrbit(CartesianPhaseSpacePosition, Orbit): - """ - Represents an orbit in Cartesian coordinates -- positions - and velocities (conjugate momenta) at different times. - - .. warning:: - - This is an experimental class. The API may change in a future release! - - The position and velocity quantities (arrays) can have an arbitrary - number of dimensions, but the first two axes (0, 1) have - special meaning:: - - - `axis=0` is the coordinate dimension (e.g., x, y, z) - - `axis=1` is the time dimension - - So if the input position array, `pos`, has shape `pos.shape = (3, 100)`, this - would be a 3D orbit (`pos[0]` is `x`, `pos[1]` is `y`, etc.). For representing - multiple orbits, the position array could have 3 axes, e.g., it might have shape - `pos.shape = (3, 100, 8)`, where this is interpreted as a 3D position at 100 times - for 8 different orbits. The same is true for velocity. The position and velocity - arrays must have the same shape. - - If a time argument is specified, the position and velocity arrays must have - the same number of timesteps as the length of the time object:: - - len(t) == pos.shape[1] - - Parameters - ---------- - pos : :class:`~astropy.units.Quantity`, array_like - Positions. If a numpy array (e.g., has no units), this will be - stored as a dimensionless :class:`~astropy.units.Quantity`. See - the note above about the assumed meaning of the axes of this object. - vel : :class:`~astropy.units.Quantity`, array_like - Velocities. If a numpy array (e.g., has no units), this will be - stored as a dimensionless :class:`~astropy.units.Quantity`. See - the note above about the assumed meaning of the axes of this object. - t : array_like, :class:`~astropy.units.Quantity` (optional) - Array of times. If a numpy array (e.g., has no units), this will be - stored as a dimensionless :class:`~astropy.units.Quantity`. - hamiltonian : `~gala.potential.Hamiltonian` (optional) - The Hamiltonian that the orbit was integrated in. - - """ - def __init__(self, pos, vel, t=None, hamiltonian=None, potential=None, frame=None): - - super(CartesianOrbit, self).__init__(pos=pos, vel=vel) - - if t is not None: - t = np.atleast_1d(t) - if self.pos.shape[1] != len(t): - raise ValueError("Position and velocity must have the same length " - "along axis=1 as the length of the time array " - "{} vs {}".format(len(t), self.pos.shape[1])) - - if not hasattr(t, 'unit'): - t = t * uno - - self.t = t - - if hamiltonian is not None: - self.potential = hamiltonian.potential - self.frame = hamiltonian.frame - - else: - self.potential = potential - self.frame = frame - - def __getitem__(self, slyce): - if isinstance(slyce, np.ndarray) or isinstance(slyce, list): - _slyce = np.array(slyce) - _slyce = (slice(None),) + (slyce,) - else: - try: - _slyce = (slice(None),) + tuple(slyce) - except TypeError: - _slyce = (slice(None),) + (slyce,) - - kw = dict() - if self.t is not None: - kw['t'] = self.t[_slyce[1]] - - pos = self.pos[_slyce] - vel = self.vel[_slyce] - - if isinstance(_slyce[1], int) or pos.ndim == 1: - return CartesianPhaseSpacePosition(pos=pos, vel=vel, frame=self.frame) - - else: - return self.__class__(pos=pos, vel=vel, - potential=self.potential, - frame=self.frame, **kw) - - def w(self, units=None): - """ - This returns a single array containing the phase-space positions. - - Parameters - ---------- - units : `~gala.units.UnitSystem` (optional) - The unit system to represent the position and velocity in - before combining into the full array. - - Returns - ------- - w : `~numpy.ndarray` - A numpy array of all positions and velocities, without units. - Will have shape ``(2*ndim,...)``. - - """ - if self.pos.unit == uno and self.vel.unit == uno and units is None: - units = [uno] - - else: - if units is None and self.hamiltonian is None: - raise ValueError("If no UnitSystem is specified, the orbit must have " - "an associated hamiltonian object.") - - if units is None and self.hamiltonian.units is None: - raise ValueError("If no UnitSystem is specified, the hamiltonian object " - "associated with this orbit must have a UnitSystem defined.") - - if units is None: - units = self.hamiltonian.units - - x = self.pos.decompose(units).value - v = self.vel.decompose(units).value - w = np.vstack((x,v)) - if w.ndim < 3: - w = w[...,np.newaxis] # one orbit - return w - - @property - def hamiltonian(self): - if self.potential is None or self.frame is None: - return None - - try: - return self._hamiltonian - except AttributeError: - from ..potential import Hamiltonian - self._hamiltonian = Hamiltonian(potential=self.potential, frame=self.frame) - - return self._hamiltonian - - # ------------------------------------------------------------------------ - # Computed dynamical quantities - # - @property - def r(self): - """ - The orbital radius. - """ - return np.sqrt(np.sum(self.pos**2, axis=0)) - - def potential_energy(self, potential=None): - r""" - The potential energy *per unit mass*: - - .. math:: - - E_\Phi = \Phi(\boldsymbol{q}) - - Returns - ------- - E : :class:`~astropy.units.Quantity` - The potential energy. - """ - if self.hamiltonian is None and potential is None: - raise ValueError("To compute the potential energy, a potential" - " object must be provided!") - if potential is None: - potential = self.hamiltonian.potential - - return super(CartesianOrbit,self).potential_energy(potential) - - def energy(self, hamiltonian=None): - r""" - The total energy *per unit mass*: - - Returns - ------- - E : :class:`~astropy.units.Quantity` - The total energy. - """ - - if self.hamiltonian is None and hamiltonian is None: - raise ValueError("To compute the total energy, a hamiltonian" - " object must be provided!") - - from ..potential import PotentialBase - if isinstance(hamiltonian, PotentialBase): - from ..potential import Hamiltonian - - warnings.warn("This function now expects a `Hamiltonian` instance instead of " - "a `PotentialBase` subclass instance. If you are using a " - "static reference frame, you just need to pass your " - "potential object in to the Hamiltonian constructor to use, e.g., " - "Hamiltonian(potential).", DeprecationWarning) - - hamiltonian = Hamiltonian(hamiltonian) - - if hamiltonian is None: - hamiltonian = self.hamiltonian - - return hamiltonian(self) - - def to_frame(self, frame, current_frame=None, **kwargs): - """ - TODO: - - Parameters - ---------- - frame : `gala.potential.CFrameBase` - The frame to transform to. - current_frame : `gala.potential.CFrameBase` (optional) - If the Orbit has no associated Hamiltonian, this specifies the - current frame of the orbit. - - Returns - ------- - orbit : `gala.dynamics.CartesianOrbit` - The orbit in the new reference frame. - - """ - - kw = kwargs.copy() - - # TODO: need a better way to do this! - from ..potential.frame.builtin import ConstantRotatingFrame - for fr in [frame, current_frame, self.frame]: - if isinstance(fr, ConstantRotatingFrame): - if 't' not in kw: - kw['t'] = self.t - - psp = super(CartesianOrbit, self).to_frame(frame, current_frame, **kw) - - return CartesianOrbit(pos=psp.pos, vel=psp.vel, t=self.t, - frame=frame, potential=self.potential) - # ------------------------------------------------------------------------ # Misc. useful methods # ------------------------------------------------------------------------ @@ -517,7 +516,8 @@ def circulation(self): circulation : :class:`numpy.ndarray` An array that specifies whether there is circulation about any of the axes of the input orbit. For a single orbit, will return a - 1D array, but for multiple orbits, the shape will be ``(3, norbits)``. + 1D array, but for multiple orbits, the shape will be + ``(3, norbits)``. """ L = self.angular_momentum() @@ -556,31 +556,37 @@ def align_circulation_with_z(self, circulation=None): Parameters ---------- circulation : array_like (optional) - Array of bits that specify the axis about which the orbit circulates. If - not provided, will compute this using - :meth:`~gala.dynamics.CartesianOrbit.circulation`. See that method for - more information. + Array of bits that specify the axis about which the orbit + circulates. If not provided, will compute this using + :meth:`~gala.dynamics.Orbit.circulation`. See that method for more + information. Returns ------- - orb : :class:`~gala.dynamics.CartesianOrbit` - A copy of the original orbit object with circulation aligned with the z axis. + orb : :class:`~gala.dynamics.Orbit` + A copy of the original orbit object with circulation aligned with + the z axis. """ if circulation is None: circulation = self.circulation() circulation = atleast_2d(circulation, insert_axis=1) - if self.pos.ndim < 3: - pos = self.pos[...,np.newaxis] - vel = self.vel[...,np.newaxis] - else: - pos = self.pos - vel = self.vel + cart = self.cartesian + pos = cart.pos.xyz + vel = np.vstack((cart.vel.d_x.value[None], + cart.vel.d_y.value[None], + cart.vel.d_z.value[None])) * cart.vel.d_x.unit + + if pos.ndim < 3: + pos = pos[...,np.newaxis] + vel = vel[...,np.newaxis] - if circulation.shape[0] != self.ndim or circulation.shape[1] != pos.shape[2]: - raise ValueError("Shape of 'circulation' array should match the shape" - " of the position/velocity (minus the time axis).") + if (circulation.shape[0] != self.ndim or + circulation.shape[1] != pos.shape[2]): + raise ValueError("Shape of 'circulation' array should match the " + "shape of the position/velocity (minus the time " + "axis).") new_pos = pos.copy() new_vel = vel.copy() @@ -606,35 +612,42 @@ def align_circulation_with_z(self, circulation=None): new_vel[circ,:,n] = vel[2,:,n] new_vel[2,:,n] = vel[circ,:,n] - return self.__class__(pos=new_pos, vel=new_vel, t=self.t, hamiltonian=self.hamiltonian) + return self.__class__(pos=new_pos, vel=new_vel, t=self.t, + hamiltonian=self.hamiltonian) - def plot(self, **kwargs): + def plot(self, components=None, units=None, **kwargs): """ - Plot the orbit in all projections. This is a thin wrapper around - `~gala.dynamics.plot_orbits` -- the docstring for this function is - included here. - - .. warning:: - - This will currently fail for orbits with fewer than 3 dimensions. + Plot the positions in all projections. This is a wrapper around + `~gala.dynamics.plot_projections` for fast access and quick + visualization. All extra keyword arguments are passed to that function + (the docstring for this function is included here for convenience). Parameters ---------- - ix : int, array_like (optional) - Index or array of indices of orbits to plot. For example, if `x` is an - array of shape ``(3,1024,32)`` - 1024 timesteps for 32 orbits in 3D - positions -- `ix` would specify which of the 32 orbits to plot. + components : iterable (optional) + A list of component names (strings) to plot. By default, this is the + Cartesian positions ``['x', 'y', 'z']``. To plot Cartesian + velocities, pass in the velocity component names + ``['d_x', 'd_y', 'd_z']``. + units : `~astropy.units.UnitBase`, iterable (optional) + A single unit or list of units to display the components in. + relative_to : bool (optional) + Plot the values relative to this value or values. + autolim : bool (optional) + Automatically set the plot limits to be something sensible. axes : array_like (optional) Array of matplotlib Axes objects. - triangle : bool (optional) - Make a triangle plot instead of plotting all projections in a single row. subplots_kwargs : dict (optional) Dictionary of kwargs passed to :func:`~matplotlib.pyplot.subplots`. labels : iterable (optional) - List or iterable of axis labels as strings. They should correspond to the - dimensions of the input orbit. + List or iterable of axis labels as strings. They should correspond to + the dimensions of the input orbit. + plot_function : callable (optional) + The ``matplotlib`` plot function to use. By default, this is + :func:`~matplotlib.pyplot.scatter`, but can also be, e.g., + :func:`~matplotlib.pyplot.plot`. **kwargs - All other keyword arguments are passed to :func:`~matplotlib.pyplot.plot`. + All other keyword arguments are passed to the ``plot_function``. You can pass in any of the usual style kwargs like ``color=...``, ``marker=...``, etc. @@ -643,115 +656,87 @@ def plot(self, **kwargs): fig : `~matplotlib.Figure` """ - _label_unit = '' - if self.pos.unit != u.dimensionless_unscaled: - _label_unit = ' [{}]'.format(self.pos.unit) + + try: + import matplotlib.pyplot as plt + except ImportError: + msg = 'matplotlib is required for visualization.' + raise ImportError(msg) + + if components is None: + if self.ndim == 1: # only a 1D orbit, so just plot time series + components = ['t', self.pos.components[0]] + else: + components = self.pos.components + + x,labels = self._plot_prepare(components=components, + units=units) default_kwargs = { - 'marker': None, + 'marker': '', 'linestyle': '-', - 'labels': ('$x${}'.format(_label_unit), - '$y${}'.format(_label_unit), - '$z${}'.format(_label_unit)) + 'labels': labels, + 'plot_function': plt.plot } for k,v in default_kwargs.items(): kwargs[k] = kwargs.get(k, v) - return plot_orbits(self.pos.value, **kwargs) - -def combine(args, along_time_axis=False): - """ - Combine the input `Orbit` objects into a single object. + fig = plot_projections(x, **kwargs) - The `Orbits` must all have the same hamiltonian and time array. + if self.pos.get_name() == 'cartesian' and \ + all([not c.startswith('d_') for c in components]) and \ + 't' not in components: + for ax in fig.axes: + ax.set(aspect='equal', adjustable='datalim') - Parameters - ---------- - args : iterable - Multiple instances of `Orbit` objects. - along_time_axis : bool (optional) - If True, will combine single orbits along the time axis. + return fig - Returns - ------- - obj : orbit_like - A single objct with positions and velocities stacked along the last axis. + def to_frame(self, frame, current_frame=None, **kwargs): + """ + TODO: - """ + Parameters + ---------- + frame : `gala.potential.CFrameBase` + The frame to transform to. + current_frame : `gala.potential.CFrameBase` (optional) + If the Orbit has no associated Hamiltonian, this specifies the + current frame of the orbit. - ndim = None - time = None - ham = None - pos_unit = None - vel_unit = None - cls = None - - all_pos = [] - all_vel = [] - all_time = [] - for x in args: - if ndim is None: - ndim = x.ndim - pos_unit = x.pos.unit - vel_unit = x.vel.unit - time = x.t - if time is not None: - t_unit = time.unit - else: - t_unit = None - ham = x.hamiltonian - cls = x.__class__ - else: - if x.__class__.__name__ != cls.__name__: - raise ValueError("All objects must have the same class.") + Returns + ------- + orbit : `gala.dynamics.Orbit` + The orbit in the new reference frame. - if x.ndim != ndim: - raise ValueError("All objects must have the same dimensionality.") + """ - if not along_time_axis: - if time is not None: - if x.t is None or len(x.t) != len(time) or np.any(x.t.to(time.unit).value != time.value): - raise ValueError("All orbits must have the same time array.") + kw = kwargs.copy() - if x.hamiltonian != ham: - raise ValueError("All orbits must have the same Hamiltonian object.") + # TODO: need a better way to do this! + from ..potential.frame.builtin import ConstantRotatingFrame + for fr in [frame, current_frame, self.frame]: + if isinstance(fr, ConstantRotatingFrame): + if 't' not in kw: + kw['t'] = self.t - pos = x.pos - if pos.ndim < 3: - pos = pos[...,np.newaxis] + psp = super(Orbit, self).to_frame(frame, current_frame, **kw) - vel = x.vel - if vel.ndim < 3: - vel = vel[...,np.newaxis] + return Orbit(pos=psp.pos, vel=psp.vel, t=self.t, + frame=frame, potential=self.potential) - all_pos.append(pos.to(pos_unit).value) - all_vel.append(vel.to(vel_unit).value) - if time is not None: - all_time.append(x.t.to(t_unit).value) - - norbits = np.array([pos.shape[-1] for pos in all_pos] + [pos.shape[-1] for pos in all_vel]) - if along_time_axis: - if np.all(norbits == norbits[0]): - all_pos = np.hstack(all_pos)*pos_unit - all_vel = np.hstack(all_vel)*vel_unit - if len(all_time) > 0: - all_time = np.concatenate(all_time)*t_unit - else: - all_time = None - else: - raise ValueError("To combine along time axis, all orbit objects must have " - "the same number of orbits.") - if args[0].pos.ndim == 2: - all_pos = all_pos[...,0] - all_vel = all_vel[...,0] - - else: - all_pos = np.dstack(all_pos)*pos_unit - all_vel = np.dstack(all_vel)*vel_unit - if len(all_time) > 0: - all_time = all_time[0]*t_unit - else: - all_time = None +class CartesianOrbit(Orbit): - return cls(pos=all_pos, vel=all_vel, t=all_time, hamiltonian=args[0].hamiltonian) + def __init__(self, pos, vel, t=None, + hamiltonian=None, potential=None, frame=None): + """ + Deprecated. + """ + warnings.warn("This class is now deprecated! Use the general interface " + "provided by Orbit instead.", + DeprecationWarning) + + super(CartesianOrbit, self).__init__(pos, vel, t=t, + hamiltonian=hamiltonian, + potential=potential, + frame=frame) diff --git a/gala/dynamics/plot.py b/gala/dynamics/plot.py index 5453f30ae..1766760a4 100644 --- a/gala/dynamics/plot.py +++ b/gala/dynamics/plot.py @@ -8,262 +8,120 @@ # Third-party import numpy as np -# Project -from ..util import atleast_2d +__all__ = ['plot_projections'] -__all__ = ['plot_orbits', 'three_panel'] - -def _get_axes(dim, axes=None, triangle=False, subplots_kwargs=dict()): +def _get_axes(dim, subplots_kwargs=dict()): """ Parameters ---------- dim : int Dimensionality of the orbit. - axes : array_like (optional) - Array of matplotlib Axes objects. - triangle : bool (optional) - Make a triangle plot instead of plotting all projections in a single row. subplots_kwargs : dict (optional) Dictionary of kwargs passed to :func:`~matplotlib.pyplot.subplots`. """ import matplotlib.pyplot as plt - if dim == 3: - if triangle and axes is None: - figsize = subplots_kwargs.pop('figsize', (8,8)) - sharex = subplots_kwargs.pop('sharex', True) - sharey = subplots_kwargs.pop('sharey', True) - fig,axes = plt.subplots(2,2,figsize=figsize, sharex=sharex, sharey=sharey, - **subplots_kwargs) - axes[0,1].set_visible(False) - axes = axes.flat - axes = [axes[0],axes[2],axes[3]] - - elif triangle and axes is not None: - try: - axes = axes.flat - except: - pass - - if len(axes) == 4: - axes = [axes[0],axes[2],axes[3]] - - elif not triangle and axes is None: - figsize = subplots_kwargs.pop('figsize', (10,3.5)) - fig,axes = plt.subplots(1, 3, figsize=figsize, **subplots_kwargs) - - elif dim <= 2: - if axes is not None: - try: - len(axes) - except TypeError: # single axes object - axes = [axes] - - else: - if dim ==1: - figsize = subplots_kwargs.pop('figsize', (14,8)) - elif dim == 2: - figsize = subplots_kwargs.pop('figsize', (8,8)) + if dim > 1: + n_panels = int(dim * (dim - 1) / 2) + else: + n_panels = 1 + figsize = subplots_kwargs.pop('figsize', (4*n_panels, 4)) + fig, axes = plt.subplots(1, n_panels, figsize=figsize, + **subplots_kwargs) - fig,axes = plt.subplots(1, 1, figsize=figsize, **subplots_kwargs) - axes = [axes] + if n_panels == 1: + axes = [axes] else: - raise ValueError("Orbit must have dimensions <= 3.") + axes = axes.flat return axes -def plot_orbits(x, t=None, ix=None, axes=None, triangle=False, - subplots_kwargs=dict(), labels=("$x$", "$y$", "$z$"), **kwargs): +def plot_projections(x, relative_to=None, autolim=True, axes=None, + subplots_kwargs=dict(), labels=None, plot_function=None, + **kwargs): """ - Given time series of positions, `x`, make nice plots of the orbit in - cartesian projections. + Given N-dimensional quantity, ``x``, make a figure containing 2D projections + of all combinations of the axes. Parameters ---------- x : array_like - Array of positions. ``axis=0`` is assumed to be the dimensionality, - ``axis=1`` is the time axis. See :ref:`shape-conventions` for more information. - t : array_like (optional) - Array of times. Only used if the input orbit is 1-dimensional. - ix : int, array_like (optional) - Index or array of indices of orbits to plot. For example, if `x` is an - array of shape ``(3,1024,32)`` - 1024 timesteps for 32 orbits in 3D - positions -- `ix` would specify which of the 32 orbits to plot. - axes : array_like (optional) - Array of matplotlib Axes objects. - triangle : bool (optional) - Make a triangle plot instead of plotting all projections in a single row. - subplots_kwargs : dict (optional) - Dictionary of kwargs passed to :func:`~matplotlib.pyplot.subplots`. - labels : iterable (optional) - List or iterable of axis labels as strings. They should correspond to the - dimensions of the input orbit. - **kwargs - All other keyword arguments are passed to :func:`~matplotlib.pyplot.plot`. - You can pass in any of the usual style kwargs like ``color=...``, - ``marker=...``, etc. - - Returns - ------- - fig : `~matplotlib.Figure` - """ - - x = atleast_2d(x, insert_axis=1) - if x.ndim == 2: - x = x[...,np.newaxis] - - # dimensionality of input orbit - dim = x.shape[0] - if dim > 3: - # if orbit has more than 3 dimensions, only use the first 3 - dim = 3 - - # hack in some defaults to subplots kwargs so by default share x and y axes - if 'sharex' not in subplots_kwargs: - subplots_kwargs['sharex'] = True - - if 'sharey' not in subplots_kwargs: - subplots_kwargs['sharey'] = True - - axes = _get_axes(dim=dim, axes=axes, triangle=triangle, - subplots_kwargs=subplots_kwargs) - - if ix is not None: - ixs = np.atleast_1d(ix) - else: - ixs = range(x.shape[-1]) - - if dim == 3: - for ii in ixs: - axes[0].plot(x[0,:,ii], x[1,:,ii], **kwargs) - axes[1].plot(x[0,:,ii], x[2,:,ii], **kwargs) - axes[2].plot(x[1,:,ii], x[2,:,ii], **kwargs) - - if triangle: - # HACK: until matplotlib 1.4 comes out, need this - axes[0].set_ylim(axes[0].get_xlim()) - axes[2].set_xlim(axes[0].get_ylim()) - - axes[0].set_ylabel(labels[1]) - axes[1].set_xlabel(labels[0]) - axes[1].set_ylabel(labels[2]) - axes[2].set_xlabel(labels[1]) - - else: - axes[0].set_xlabel(labels[0]) - axes[0].set_ylabel(labels[1]) - - axes[1].set_xlabel(labels[0]) - axes[1].set_ylabel(labels[2]) - - axes[2].set_xlabel(labels[1]) - axes[2].set_ylabel(labels[2]) - - if not triangle: - axes[0].figure.tight_layout() - - elif dim == 2: - for ii in ixs: - axes[0].plot(x[0,:,ii], x[1,:,ii], **kwargs) - - axes[0].set_xlabel(labels[0]) - axes[0].set_ylabel(labels[1]) - axes[0].figure.tight_layout() - - elif dim == 1: - if t is None: - t = np.arange(x.shape[1]) - - for ii in ixs: - axes[0].plot(t, x[0,:,ii], **kwargs) - - axes[0].set_xlabel("$t$") - axes[0].set_ylabel(labels[0]) - axes[0].figure.tight_layout() - - return axes[0].figure - -def three_panel(q, relative_to=None, autolim=True, axes=None, - triangle=False, subplots_kwargs=dict(), labels=None, **kwargs): - """ - Given 3D quantities, ``q``, make a nice three-panel or triangle plot - of projections of the values. - - Parameters - ---------- - q : array_like Array of values. ``axis=0`` is assumed to be the dimensionality, - ``axis=1`` is the time axis. See :ref:`shape-conventions` for more information. + ``axis=1`` is the time axis. See :ref:`shape-conventions` for more + information. relative_to : bool (optional) Plot the values relative to this value or values. autolim : bool (optional) Automatically set the plot limits to be something sensible. axes : array_like (optional) Array of matplotlib Axes objects. - triangle : bool (optional) - Make a triangle plot instead of plotting all projections in a single row. subplots_kwargs : dict (optional) Dictionary of kwargs passed to :func:`~matplotlib.pyplot.subplots`. labels : iterable (optional) - List or iterable of axis labels as strings. They should correspond to the - dimensions of the input orbit. + List or iterable of axis labels as strings. They should correspond to + the dimensions of the input orbit. + plot_function : callable (optional) + The ``matplotlib`` plot function to use. By default, this is + :func:`~matplotlib.pyplot.scatter`, but can also be, e.g., + :func:`~matplotlib.pyplot.plot`. **kwargs - All other keyword arguments are passed to :func:`~matplotlib.pyplot.scatter`. + All other keyword arguments are passed to the ``plot_function``. You can pass in any of the usual style kwargs like ``color=...``, ``marker=...``, etc. Returns ------- fig : `~matplotlib.Figure` + """ # don't propagate changes back... - q = q.copy() + x = np.array(x, copy=True) + ndim = x.shape[0] # get axes object from arguments - axes = _get_axes(dim=3, axes=axes, triangle=triangle, subplots_kwargs=subplots_kwargs) + if axes is None: + axes = _get_axes(dim=ndim, subplots_kwargs=subplots_kwargs) # if the quantities are relative if relative_to is not None: - q -= relative_to + x -= relative_to - axes[0].scatter(q[0], q[1], **kwargs) - axes[1].scatter(q[0], q[2], **kwargs) - axes[2].scatter(q[1], q[2], **kwargs) + # name of the plotting function + plot_fn_name = plot_function.__name__ - if labels is not None: - if triangle: - axes[0].set_ylabel(labels[1]) - axes[1].set_xlabel(labels[0]) - axes[1].set_ylabel(labels[2]) - axes[2].set_xlabel(labels[0]) + # automatically determine limits + if autolim: + lims = [] + for i in range(ndim): + max_,min_ = np.max(x[i]), np.min(x[i]) + delta = max_ - min_ - else: - axes[0].set_xlabel(labels[0]) - axes[0].set_ylabel(labels[1]) + if delta == 0.: + delta = 1. - axes[1].set_xlabel(labels[0]) - axes[1].set_ylabel(labels[2]) + lims.append([min_ - delta*0.02, max_ + delta*0.02]) - axes[2].set_xlabel(labels[1]) - axes[2].set_ylabel(labels[2]) + k = 0 + for i in range(ndim): + for j in range(ndim): + if i >= j: + continue # skip diagonal, upper triangle - if autolim: - lims = [] - for i in range(3): - mx,mi = q[i].max(), q[i].min() - delta = mx-mi - lims.append((mi-delta*0.05, mx+delta*0.05)) + plot_func = getattr(axes[k], plot_fn_name) + plot_func(x[i], x[j], **kwargs) + + if labels is not None: + axes[k].set_xlabel(labels[i]) + axes[k].set_ylabel(labels[j]) + + if autolim: + axes[k].set_xlim(lims[i]) + axes[k].set_ylim(lims[j]) - axes[0].set_xlim(lims[0]) - axes[0].set_ylim(lims[1]) - axes[1].set_xlim(lims[0]) - axes[1].set_ylim(lims[2]) - axes[2].set_xlim(lims[1]) - axes[2].set_ylim(lims[2]) + k += 1 - if not triangle: - axes[0].figure.tight_layout() + axes[0].figure.tight_layout() return axes[0].figure diff --git a/gala/dynamics/representation_nd.py b/gala/dynamics/representation_nd.py new file mode 100644 index 000000000..f164a861f --- /dev/null +++ b/gala/dynamics/representation_nd.py @@ -0,0 +1,220 @@ +# coding: utf-8 + +from __future__ import division, print_function + +# Standard library +from collections import OrderedDict +import operator + +# Third-party +import astropy.coordinates as coord +import astropy.units as u +import numpy as np + +__all__ = ['NDCartesianRepresentation', 'NDCartesianDifferential'] + +class NDMixin(object): + + def _apply(self, method, *args, **kwargs): + """Create a new representation with ``method`` applied to the arrays. + + In typical usage, the method is any of the shape-changing methods for + `~numpy.ndarray` (``reshape``, ``swapaxes``, etc.), as well as those + picking particular elements (``__getitem__``, ``take``, etc.), which + are all defined in `~astropy.utils.misc.ShapedLikeNDArray`. It will be + applied to the underlying arrays (e.g., ``x``, ``y``, and ``z`` for + `~astropy.coordinates.CartesianRepresentation`), with the results used + to create a new instance. + + Internally, it is also used to apply functions to the components + (in particular, `~numpy.broadcast_to`). + + Parameters + ---------- + method : str or callable + If str, it is the name of a method that is applied to the internal + ``components``. If callable, the function is applied. + args : tuple + Any positional arguments for ``method``. + kwargs : dict + Any keyword arguments for ``method``. + """ + if callable(method): + apply_method = lambda array: method(array, *args, **kwargs) + else: + apply_method = operator.methodcaller(method, *args, **kwargs) + return self.__class__([apply_method(getattr(self, component)) + for component in self.components], copy=False) + +class NDCartesianRepresentation(NDMixin, coord.CartesianRepresentation): + """ + Representation of points in ND cartesian coordinates. + + Parameters + ---------- + x : `~astropy.units.Quantity` or array + The Cartesian coordinates of the point(s). If not quantity, + ``unit`` should be set. + unit : `~astropy.units.Unit` or str + If given, the coordinates will be converted to this unit (or taken to + be in this unit if not given. + copy : bool, optional + If `True` (default), arrays will be copied rather than referenced. + """ + + attr_classes = OrderedDict() + + def __init__(self, x, unit=None, copy=True): + + if unit is None: + if not hasattr(x[0], 'unit'): + unit = u.one + else: + unit = x[0].unit + + x = u.Quantity(x, unit, copy=copy, subok=True) + copy = False + + self.attr_classes = OrderedDict([('x'+str(i), u.Quantity) + for i in range(1, len(x)+1)]) + + super(coord.CartesianRepresentation, self).__init__(*x, copy=copy) + + ptype = None + for name,_ in self.attr_classes.items(): + if ptype is None: + ptype = getattr(self, '_'+name).unit.physical_type + + else: + if getattr(self, '_'+name).unit.physical_type != ptype: + raise u.UnitsError("All components should have matching " + "physical types") + + cls = self.__class__ + if not hasattr(cls, name): + setattr(cls, name, + property(coord.representation._make_getter(name), + doc=("The '{0}' component of the points(s)." + .format(name)))) + + def get_xyz(self, xyz_axis=0): + """Return a vector array of the x, y, and z coordinates. + + Parameters + ---------- + xyz_axis : int, optional + The axis in the final array along which the x, y, z components + should be stored (default: 0). + + Returns + ------- + xs : `~astropy.units.Quantity` + With dimension 3 along ``xyz_axis``. + """ + # Add new axis in x, y, z so one can concatenate them around it. + # NOTE: just use np.stack once our minimum numpy version is 1.10. + result_ndim = self.ndim + 1 + if not -result_ndim <= xyz_axis < result_ndim: + raise IndexError('xyz_axis {0} out of bounds [-{1}, {1})' + .format(xyz_axis, result_ndim)) + if xyz_axis < 0: + xyz_axis += result_ndim + + # Get components to the same units (very fast for identical units) + # since np.concatenate cannot deal with quantity. + unit = self._x1.unit + + sh = self.shape + sh = sh[:xyz_axis] + (1,) + sh[xyz_axis:] + components = [getattr(self, '_'+name).reshape(sh).to(unit).value + for name in self.attr_classes] + xs_value = np.concatenate(components, axis=xyz_axis) + return u.Quantity(xs_value, unit=unit, copy=False) + + xyz = property(get_xyz) + +class NDCartesianDifferential(NDMixin, coord.CartesianDifferential): + """Differentials in of points in ND cartesian coordinates. + + Parameters + ---------- + *d_x : `~astropy.units.Quantity` or array + The Cartesian coordinates of the differentials. If not quantity, + ``unit`` should be set. + unit : `~astropy.units.Unit` or str + If given, the differentials will be converted to this unit (or taken to + be in this unit if not given. + copy : bool, optional + If `True` (default), arrays will be copied rather than referenced. + """ + base_representation = NDCartesianRepresentation + attr_classes = OrderedDict() + + def __init__(self, d_x, unit=None, copy=True): + + if unit is None: + if not hasattr(d_x[0], 'unit'): + unit = u.one + else: + unit = d_x[0].unit + + d_x = u.Quantity(d_x, unit, copy=copy, subok=True) + copy = False + + self.attr_classes = OrderedDict([('d_x'+str(i), u.Quantity) + for i in range(1, len(d_x)+1)]) + + super(coord.CartesianDifferential, self).__init__(*d_x, copy=copy) + + ptype = None + for name,_ in self.attr_classes.items(): + if ptype is None: + ptype = getattr(self, '_'+name).unit.physical_type + + else: + if getattr(self, '_'+name).unit.physical_type != ptype: + raise u.UnitsError("All components should have matching " + "physical types") + + cls = self.__class__ + if not hasattr(cls, name): + setattr(cls, name, + property(coord.representation._make_getter(name), + doc=("The '{0}' component of the points(s)." + .format(name)))) + + def get_d_xyz(self, xyz_axis=0): + """Return a vector array of the x, y, and z coordinates. + + Parameters + ---------- + xyz_axis : int, optional + The axis in the final array along which the x, y, z components + should be stored (default: 0). + + Returns + ------- + d_xs : `~astropy.units.Quantity` + With dimension 3 along ``xyz_axis``. + """ + # Add new axis in x, y, z so one can concatenate them around it. + # NOTE: just use np.stack once our minimum numpy version is 1.10. + result_ndim = self.ndim + 1 + if not -result_ndim <= xyz_axis < result_ndim: + raise IndexError('xyz_axis {0} out of bounds [-{1}, {1})' + .format(xyz_axis, result_ndim)) + if xyz_axis < 0: + xyz_axis += result_ndim + + # Get components to the same units (very fast for identical units) + # since np.concatenate cannot deal with quantity. + unit = self._d_x1.unit + + sh = self.shape + sh = sh[:xyz_axis] + (1,) + sh[xyz_axis:] + components = [getattr(self, '_'+name).reshape(sh).to(unit).value + for name in self.components] + xs_value = np.concatenate(components, axis=xyz_axis) + return u.Quantity(xs_value, unit=unit, copy=False) + + d_xyz = property(get_d_xyz) diff --git a/gala/dynamics/tests/helpers.py b/gala/dynamics/tests/helpers.py index 69c6ebfa1..de1caa6c9 100644 --- a/gala/dynamics/tests/helpers.py +++ b/gala/dynamics/tests/helpers.py @@ -14,7 +14,6 @@ # Project # from ..actionangle import classify_orbit from ...units import galactic -from ...coordinates import physicsspherical_to_cartesian from ...potential import HarmonicOscillatorPotential, IsochronePotential from .._genfunc import genfunc_3d, solver, toy_potentials @@ -123,9 +122,16 @@ def isotropic_w0(N=100): theta=theta*u.radian) x = rep.represent_as(coord.CartesianRepresentation).xyz.T.value - vr = vr.decompose(galactic).value - vphi = vphi.decompose(galactic).value - vtheta = vtheta.decompose(galactic).value - v = physicsspherical_to_cartesian(rep, [vr,vphi,vtheta]*u.dimensionless_unscaled).T.value + vr = vr.decompose(galactic).value * u.one + vphi = vphi.decompose(galactic).value * u.one + vtheta = vtheta.decompose(galactic).value * u.one + + vsph = coord.PhysicsSphericalDifferential(d_phi=vphi/(d*np.sin(theta)), + d_theta=vtheta/d, + d_r=vr) + + with u.set_enabled_equivalencies(u.dimensionless_angles()): + v = vsph.represent_as(coord.CartesianDifferential, + base=rep).d_xyz.value.T return np.hstack((x,v)).T diff --git a/gala/dynamics/tests/test_analyticactionangle.py b/gala/dynamics/tests/test_analyticactionangle.py index d502d6eb7..a5ec481c8 100644 --- a/gala/dynamics/tests/test_analyticactionangle.py +++ b/gala/dynamics/tests/test_analyticactionangle.py @@ -44,8 +44,8 @@ def test(self): assert np.allclose(actions[i,1:], actions[i,0], rtol=1E-5) # Compare to genfunc - x = self.w.pos.value[...,n] - v = self.w.vel.value[...,n] + x = self.w.pos.xyz.value[...,n] + v = self.w.vel.d_xyz.value[...,n] s_v = (v*u.kpc/u.Myr).to(u.km/u.s).value s_w = np.vstack((x,s_v)) m = self.potential.parameters['m'].value / 1E11 @@ -93,8 +93,8 @@ def test(self): assert np.allclose(actions[i,1:], actions[i,0], rtol=1E-5) # Compare to genfunc - x = self.w.pos.value[...,n] - v = self.w.vel.value[...,n] + x = self.w.pos.xyz.value[...,n] + v = self.w.vel.d_xyz.value[...,n] s_w = np.vstack((x,v)) omega = self.potential.parameters['omega'].value aa = np.array([toy_potentials.angact_ho(s_w[:,i].T, omega=omega) for i in range(s_w.shape[1])]) diff --git a/gala/dynamics/tests/test_core.py b/gala/dynamics/tests/test_core.py index 146f6e754..46b4c80fb 100644 --- a/gala/dynamics/tests/test_core.py +++ b/gala/dynamics/tests/test_core.py @@ -9,13 +9,15 @@ # Third-party import astropy.units as u -from astropy.coordinates import SphericalRepresentation, Galactic +from astropy.coordinates import (Galactic, CartesianRepresentation, + SphericalRepresentation, CartesianDifferential, + SphericalDifferential) from astropy.tests.helper import quantity_allclose import numpy as np import pytest # Project -from ..core import * +from ..core import PhaseSpacePosition from ...potential import Hamiltonian, HernquistPotential from ...potential.frame import StaticFrame, ConstantRotatingFrame from ...units import galactic, solarsystem @@ -25,67 +27,96 @@ def test_initialize(): with pytest.raises(ValueError): x = np.random.random(size=(3,10)) v = np.random.random(size=(3,8)) - CartesianPhaseSpacePosition(pos=x, vel=v) + PhaseSpacePosition(pos=x, vel=v) x = np.random.random(size=(3,10)) v = np.random.random(size=(3,10)) - o = CartesianPhaseSpacePosition(pos=x, vel=v) - assert o.ndim == 3 + o = PhaseSpacePosition(pos=x, vel=v) + assert o.shape == (10,) x = np.random.random(size=(3,10))*u.kpc v = np.random.random(size=(3,10))*u.km/u.s - o = CartesianPhaseSpacePosition(pos=x, vel=v) - assert o.pos.unit == u.kpc - assert o.vel.unit == u.km/u.s + o = PhaseSpacePosition(pos=x, vel=v) + assert o.pos.xyz.unit == u.kpc + assert o.vel.d_x.unit == u.km/u.s + # Not 3D x = np.random.random(size=(2,10)) v = np.random.random(size=(2,10)) - o = CartesianPhaseSpacePosition(pos=x, vel=v) + o = PhaseSpacePosition(pos=x, vel=v) assert o.ndim == 2 - o = CartesianPhaseSpacePosition(pos=x, vel=v, frame=StaticFrame(galactic)) + o = PhaseSpacePosition(pos=x, vel=v, frame=StaticFrame(galactic)) assert o.ndim == 2 assert o.frame is not None assert isinstance(o.frame, StaticFrame) + x = np.random.random(size=(4,10)) + v = np.random.random(size=(4,10)) + o = PhaseSpacePosition(pos=x, vel=v) + assert o.ndim == 4 + + # back to 3D + pos = CartesianRepresentation(np.random.random(size=(3,10))*u.one) + vel = CartesianDifferential(np.random.random(size=(3,10))*u.one) + o = PhaseSpacePosition(pos=pos, vel=vel) + assert hasattr(o, 'x') + assert hasattr(o, 'y') + assert hasattr(o, 'z') + assert hasattr(o, 'd_x') + assert hasattr(o, 'd_y') + assert hasattr(o, 'd_z') + + o = o.represent_as(SphericalRepresentation) + assert hasattr(o, 'distance') + assert hasattr(o, 'lat') + assert hasattr(o, 'lon') + assert hasattr(o, 'd_distance') + assert hasattr(o, 'd_lat') + assert hasattr(o, 'd_lon') + with pytest.raises(TypeError): - o = CartesianPhaseSpacePosition(pos=x, vel=v, frame="blah blah blah") + o = PhaseSpacePosition(pos=x, vel=v, frame="blah blah blah") + + # check that old class raises deprecation warning + from ..core import CartesianPhaseSpacePosition + warnings.simplefilter('always') + with pytest.warns(DeprecationWarning): + o = CartesianPhaseSpacePosition(pos=x, vel=v) def test_from_w(): w = np.random.random(size=(6,10)) - o = CartesianPhaseSpacePosition.from_w(w, galactic) - assert o.pos.unit == u.kpc - assert o.vel.unit == u.kpc/u.Myr + o = PhaseSpacePosition.from_w(w, galactic) + assert o.pos.x.unit == u.kpc + assert o.vel.d_x.unit == u.kpc/u.Myr + assert o.shape == (10,) def test_slice(): # simple x = np.random.random(size=(3,10)) v = np.random.random(size=(3,10)) - o = CartesianPhaseSpacePosition(pos=x, vel=v) + o = PhaseSpacePosition(pos=x, vel=v) new_o = o[:5] - assert new_o.pos.shape == (3,5) - assert new_o.vel.shape == (3,5) + assert new_o.shape == (5,) # 1d slice on 3d - x = np.random.random(size=(3,100,8)) - v = np.random.random(size=(3,100,8)) - o = CartesianPhaseSpacePosition(pos=x, vel=v) + x = np.random.random(size=(3,10,8)) + v = np.random.random(size=(3,10,8)) + o = PhaseSpacePosition(pos=x, vel=v) new_o = o[:5] - assert new_o.pos.shape == (3,5,8) - assert new_o.vel.shape == (3,5,8) + assert new_o.shape == (5,8) # 3d slice on 3d - o = CartesianPhaseSpacePosition(pos=x, vel=v) + o = PhaseSpacePosition(pos=x, vel=v) new_o = o[:5,:4] - assert new_o.pos.shape == (3,5,4) - assert new_o.vel.shape == (3,5,4) + assert new_o.shape == (5,4) # boolean array x = np.random.random(size=(3,10)) v = np.random.random(size=(3,10)) - o = CartesianPhaseSpacePosition(pos=x, vel=v) + o = PhaseSpacePosition(pos=x, vel=v) ix = np.array([0,0,0,0,0,1,1,1,1,1]).astype(bool) new_o = o[ix] assert new_o.shape == (sum(ix),) @@ -93,7 +124,7 @@ def test_slice(): # integer array x = np.random.random(size=(3,10)) v = np.random.random(size=(3,10)) - o = CartesianPhaseSpacePosition(pos=x, vel=v) + o = PhaseSpacePosition(pos=x, vel=v) ix = np.array([0,3,5]) new_o = o[ix] assert new_o.shape == (len(ix),) @@ -106,26 +137,32 @@ def test_represent_as(): # simple / unitless x = np.random.random(size=(3,10)) v = np.random.random(size=(3,10)) - o = CartesianPhaseSpacePosition(pos=x, vel=v) - sph_pos, sph_vel = o.represent_as(SphericalRepresentation) + o = PhaseSpacePosition(pos=x, vel=v) + new_o = o.represent_as(SphericalRepresentation) o.spherical o.cylindrical o.cartesian - assert sph_pos.distance.unit == u.dimensionless_unscaled - assert sph_vel.unit == u.dimensionless_unscaled + assert new_o.pos.distance.unit == u.one + assert new_o.vel.d_distance.unit == u.one # simple / with units x = np.random.random(size=(3,10))*u.kpc v = np.random.normal(0.,100.,size=(3,10))*u.km/u.s - o = CartesianPhaseSpacePosition(pos=x, vel=v) - sph_pos, sph_vel = o.represent_as(SphericalRepresentation) - assert sph_pos.distance.unit == u.kpc + o = PhaseSpacePosition(pos=x, vel=v) + sph = o.represent_as(SphericalRepresentation) + assert sph.pos.distance.unit == u.kpc + + sph2 = o.represent_as('spherical') + for c in sph.pos.components: + assert quantity_allclose(getattr(sph.pos, c), + getattr(sph2.pos, c), + rtol=1E-12) # doesn't work for 2D x = np.random.random(size=(2,10)) v = np.random.random(size=(2,10)) - o = CartesianPhaseSpacePosition(pos=x, vel=v) + o = PhaseSpacePosition(pos=x, vel=v) with pytest.raises(ValueError): o.represent_as(SphericalRepresentation) @@ -133,7 +170,7 @@ def test_to_coord_frame(): # simple / unitless x = np.random.random(size=(3,10)) v = np.random.random(size=(3,10)) - o = CartesianPhaseSpacePosition(pos=x, vel=v) + o = PhaseSpacePosition(pos=x, vel=v) with pytest.raises(u.UnitConversionError): o.to_coord_frame(Galactic) @@ -141,7 +178,7 @@ def test_to_coord_frame(): # simple / with units x = np.random.random(size=(3,10))*u.kpc v = np.random.normal(0.,100.,size=(3,10))*u.km/u.s - o = CartesianPhaseSpacePosition(pos=x, vel=v) + o = PhaseSpacePosition(pos=x, vel=v) coo,vel = o.to_coord_frame(Galactic) assert coo.name == 'galactic' @@ -152,7 +189,7 @@ def test_to_coord_frame(): # doesn't work for 2D x = np.random.random(size=(2,10))*u.kpc v = np.random.normal(0.,100.,size=(2,10))*u.km/u.s - o = CartesianPhaseSpacePosition(pos=x, vel=v) + o = PhaseSpacePosition(pos=x, vel=v) with pytest.raises(ValueError): o.to_coord_frame(Galactic) @@ -160,14 +197,33 @@ def test_w(): # simple / unitless x = np.random.random(size=(3,10)) v = np.random.random(size=(3,10)) - o = CartesianPhaseSpacePosition(pos=x, vel=v) + o = PhaseSpacePosition(pos=x, vel=v) w = o.w() assert w.shape == (6,10) + x = np.random.random(size=3) + v = np.random.random(size=3) + o = PhaseSpacePosition(pos=x, vel=v) + w = o.w() + assert w.shape == (6,1) + + # simple / unitless, 2D + x = np.random.random(size=(2,10)) + v = np.random.random(size=(2,10)) + o = PhaseSpacePosition(pos=x, vel=v) + w = o.w() + assert w.shape == (4,10) + + x = np.random.random(size=2) + v = np.random.random(size=2) + o = PhaseSpacePosition(pos=x, vel=v) + w = o.w() + assert w.shape == (4,1) + # simple / with units x = np.random.random(size=(3,10))*u.kpc v = np.random.normal(0.,100.,size=(3,10))*u.km/u.s - o = CartesianPhaseSpacePosition(pos=x, vel=v) + o = PhaseSpacePosition(pos=x, vel=v) with pytest.raises(ValueError): o.w() w = o.w(units=galactic) @@ -178,7 +234,7 @@ def test_w(): p = HernquistPotential(units=galactic, m=1E11, c=0.25) x = np.random.random(size=(3,10))*u.kpc v = np.random.normal(0.,100.,size=(3,10))*u.km/u.s - o = CartesianPhaseSpacePosition(pos=x, vel=v) + o = PhaseSpacePosition(pos=x, vel=v) w = o.w(p.units) assert np.allclose(x.value, w[:3]) assert np.allclose(v.value, (w[3:]*u.kpc/u.Myr).to(u.km/u.s).value) @@ -194,17 +250,17 @@ def test_energy(): # with units x = np.random.random(size=(3,10))*u.kpc v = np.random.normal(0.,100.,size=(3,10))*u.km/u.s - o = CartesianPhaseSpacePosition(pos=x, vel=v) + o = PhaseSpacePosition(pos=x, vel=v) KE = o.kinetic_energy() - assert KE.unit == (o.vel.unit)**2 - assert KE.shape == o.pos.shape[1:] + assert KE.unit == (o.vel.d_x.unit)**2 + assert KE.shape == o.shape # with units and potential p = HernquistPotential(units=galactic, m=1E11, c=0.25) H = Hamiltonian(p) x = np.random.random(size=(3,10))*u.kpc v = np.random.normal(0.,100.,size=(3,10))*u.km/u.s - o = CartesianPhaseSpacePosition(pos=x, vel=v) + o = PhaseSpacePosition(pos=x, vel=v) PE = o.potential_energy(p) E = o.energy(H) @@ -214,59 +270,32 @@ def test_energy(): def test_angular_momentum(): - w = CartesianPhaseSpacePosition([1.,0.,0.], [0.,0.,1.]) - assert np.allclose(np.squeeze(w.angular_momentum()), [0., -1, 0]) + w = PhaseSpacePosition([1.,0.,0.], [0.,0.,1.]) + assert quantity_allclose(np.squeeze(w.angular_momentum()), [0.,-1,0]*u.one) - w = CartesianPhaseSpacePosition([1.,0.,0.], [0.,1.,0.]) - assert np.allclose(np.squeeze(w.angular_momentum()), [0., 0, 1]) + w = PhaseSpacePosition([1.,0.,0.], [0.,1.,0.]) + assert quantity_allclose(np.squeeze(w.angular_momentum()), [0.,0, 1]*u.one) - w = CartesianPhaseSpacePosition([0.,1.,0.],[0.,0.,1.]) - assert np.allclose(np.squeeze(w.angular_momentum()), [1., 0, 0]) + w = PhaseSpacePosition([0.,1.,0.],[0.,0.,1.]) + assert quantity_allclose(np.squeeze(w.angular_momentum()), [1., 0, 0]*u.one) - w = CartesianPhaseSpacePosition([1.,0,0]*u.kpc, [0.,200.,0]*u.pc/u.Myr) + w = PhaseSpacePosition([1.,0,0]*u.kpc, [0.,200.,0]*u.pc/u.Myr) assert quantity_allclose(np.squeeze(w.angular_momentum()), [0,0,0.2]*u.kpc**2/u.Myr) # multiple - known q = np.array([[1.,0.,0.],[1.,0.,0.],[0,1.,0.]]).T p = np.array([[0,0,1.],[0,1.,0.],[0,0,1]]).T - L = CartesianPhaseSpacePosition(q, p).angular_momentum() - true_L = np.array([[0., -1, 0],[0., 0, 1],[1., 0, 0]]).T + L = PhaseSpacePosition(q, p).angular_momentum() + true_L = np.array([[0., -1, 0],[0., 0, 1],[1., 0, 0]]).T * u.one assert L.shape == (3,3) - assert np.allclose(L, true_L) + assert quantity_allclose(L, true_L) # multiple - random q = np.random.uniform(size=(3,128)) p = np.random.uniform(size=(3,128)) - L = CartesianPhaseSpacePosition(q, p).angular_momentum() + L = PhaseSpacePosition(q, p).angular_momentum() assert L.shape == (3,128) -def test_combine(): - - o1 = CartesianPhaseSpacePosition.from_w(np.random.random(size=6), units=galactic) - o2 = CartesianPhaseSpacePosition.from_w(np.random.random(size=6), units=galactic) - o = combine((o1, o2)) - assert o.pos.shape == (3,2) - assert o.vel.shape == (3,2) - - o1 = CartesianPhaseSpacePosition.from_w(np.random.random(size=6), units=galactic) - o2 = CartesianPhaseSpacePosition.from_w(np.random.random(size=(6,10)), units=galactic) - o = combine((o1, o2)) - assert o.pos.shape == (3,11) - assert o.vel.shape == (3,11) - - o1 = CartesianPhaseSpacePosition.from_w(np.random.random(size=6), units=galactic) - o2 = CartesianPhaseSpacePosition.from_w(np.random.random(size=6), units=solarsystem) - o = combine((o1, o2)) - assert o.pos.unit == galactic['length'] - assert o.vel.unit == galactic['length']/galactic['time'] - - o1 = CartesianPhaseSpacePosition.from_w(np.random.random(size=6), units=galactic) - x = np.random.random(size=(2,10))*u.kpc - v = np.random.normal(0.,100.,size=(2,10))*u.km/u.s - o2 = CartesianPhaseSpacePosition(pos=x, vel=v) - with pytest.raises(ValueError): - o = combine((o1, o2)) - def test_frame_transform(): static = StaticFrame(galactic) rotating = ConstantRotatingFrame(Omega=[0.53,1.241,0.9394]*u.rad/u.Myr, units=galactic) @@ -275,12 +304,12 @@ def test_frame_transform(): v = np.array([[0.0034,0.2,0.0014],[0.0001,0.002532,-0.2]]).T * u.kpc/u.Myr # no frame specified at init - psp = CartesianPhaseSpacePosition(pos=x, vel=v) + psp = PhaseSpacePosition(pos=x, vel=v) with pytest.raises(ValueError): psp.to_frame(rotating) psp.to_frame(rotating, current_frame=static, t=0.4*u.Myr) # frame specified at init - psp = CartesianPhaseSpacePosition(pos=x, vel=v, frame=static) + psp = PhaseSpacePosition(pos=x, vel=v, frame=static) psp.to_frame(rotating, t=0.4*u.Myr) diff --git a/gala/dynamics/tests/test_nonlinear.py b/gala/dynamics/tests/test_nonlinear.py index 7d9848001..8996ff60f 100644 --- a/gala/dynamics/tests/test_nonlinear.py +++ b/gala/dynamics/tests/test_nonlinear.py @@ -110,12 +110,7 @@ def setup(self): def test_integrate_orbit(self, tmpdir): integrator = DOPRI853Integrator(self.F_max, func_args=self.par) - - # pl.clf() orbit = integrator.run(self.w0, dt=self.dt, n_steps=self.n_steps) - # pl.plot(w[:,0,0], w[:,0,1], marker=None) - - # pl.savefig(os.path.join(str(tmpdir),"hh_orbit_{}.png".format(self.__class__.__name__))) def test_lyapunov_max(self, tmpdir): n_steps_per_pullback = 10 diff --git a/gala/dynamics/tests/test_orbit.py b/gala/dynamics/tests/test_orbit.py index cd6677dd5..dd31c0a22 100644 --- a/gala/dynamics/tests/test_orbit.py +++ b/gala/dynamics/tests/test_orbit.py @@ -1,9 +1,9 @@ # coding: utf-8 -""" Test ... """ - from __future__ import division, print_function +# standard library +import warnings # Third-party from astropy.coordinates import SphericalRepresentation, Galactic @@ -13,135 +13,45 @@ import scipy.optimize as so # Project -from ..core import CartesianPhaseSpacePosition -from ..orbit import * +from ..core import PhaseSpacePosition +from ..orbit import Orbit from ...integrate import DOPRI853Integrator from ...potential import Hamiltonian, HernquistPotential, LogarithmicPotential, KeplerPotential from ...potential.frame import StaticFrame, ConstantRotatingFrame from ...units import galactic, solarsystem -def make_known_orbit(tmpdir, x, vx, potential, name): - # See Binney & Tremaine (2008) Figure 3.8 and 3.9 - E = -0.337 - y = 0. - vy = np.sqrt(2*(E - potential.value([x,y,0.]).value))[0] - - w = [x,y,0.,vx,vy,0.] - ham = Hamiltonian(potential) - orbit = ham.integrate_orbit(w, dt=0.05, n_steps=10000) - - # fig,ax = pl.subplots(1,1) - # ax.plot(ws[0], ws[1]) - # fig = plot_orbits(ws, linestyle='none', alpha=0.1) - # fig.savefig(os.path.join(str(tmpdir), "{}.png".format(name))) - # logger.debug(os.path.join(str(tmpdir), "{}.png".format(name))) - - return orbit - -def test_circulation(tmpdir): - - potential = LogarithmicPotential(v_c=1., r_h=0.14, q1=1., q2=0.9, q3=1., - units=galactic) - - # individual - w1 = make_known_orbit(tmpdir, 0.5, 0., potential, "loop") - circ = w1.circulation() - assert circ.shape == (3,) - assert circ.sum() == 1 - - w2 = make_known_orbit(tmpdir, 0., 1.5, potential, "box") - circ = w2.circulation() - assert circ.shape == (3,) - assert circ.sum() == 0 - - # try also for both, together - w3 = combine((w1,w2)) - circ = w3.circulation() - assert circ.shape == (3,2) - assert np.allclose(circ.sum(axis=0), [1,0]) - -def test_align_circulation(): - - t = np.linspace(0,100,1024) - w = np.zeros((6,1024,4)) - - # loop around x axis - w[1,:,0] = np.cos(t) - w[2,:,0] = np.sin(t) - w[4,:,0] = -np.sin(t) - w[5,:,0] = np.cos(t) - - # loop around y axis - w[0,:,1] = -np.cos(t) - w[2,:,1] = np.sin(t) - w[3,:,1] = np.sin(t) - w[5,:,1] = np.cos(t) - - # loop around z axis - w[0,:,2] = np.cos(t) - w[1,:,2] = np.sin(t) - w[3,:,2] = -np.sin(t) - w[4,:,2] = np.cos(t) - - # box - w[0,:,3] = np.cos(t) - w[1,:,3] = -np.cos(0.5*t) - w[2,:,3] = np.cos(0.25*t) - w[3,:,3] = -np.sin(t) - w[4,:,3] = 0.5*np.sin(0.5*t) - w[5,:,3] = -0.25*np.sin(0.25*t) - - # First, individually - for i in range(w.shape[2]): - orb = CartesianOrbit.from_w(w[...,i], units=galactic) - new_orb = orb.align_circulation_with_z() - circ = new_orb.circulation() - - if i == 3: - assert np.sum(circ) == 0 - else: - assert circ[2] == 1. - - # all together now - orb = CartesianOrbit.from_w(w, units=galactic) - circ = orb.circulation() - assert circ.shape == (3,4) - - new_orb = orb.align_circulation_with_z() - new_circ = new_orb.circulation() - assert np.all(new_circ[2,:3] == 1.) - assert np.all(new_circ[:,3] == 0.) - -# Tests below should be fixed a bit... +# Tests below should be cleaned up a bit... def test_initialize(): with pytest.raises(ValueError): x = np.random.random(size=(3,10)) v = np.random.random(size=(3,8)) - CartesianOrbit(pos=x, vel=v) + Orbit(pos=x, vel=v) with pytest.raises(ValueError): x = np.random.random(size=(3,10)) v = np.random.random(size=(3,10)) t = np.arange(8) - CartesianOrbit(pos=x, vel=v, t=t) + Orbit(pos=x, vel=v, t=t) - x = np.random.random(size=(3,10)) - v = np.random.random(size=(3,10)) - o = CartesianOrbit(pos=x, vel=v) - assert o.ndim == 3 + # TODO: always? + # x = np.random.random(size=(3,10)) + # v = np.random.random(size=(3,10)) + # o = Orbit(pos=x, vel=v) + # assert o.ndim == 3 x = np.random.random(size=(3,10))*u.kpc v = np.random.random(size=(3,10))*u.km/u.s - o = CartesianOrbit(pos=x, vel=v) - assert o.pos.unit == u.kpc - assert o.vel.unit == u.km/u.s - - x = np.random.random(size=(2,10)) - v = np.random.random(size=(2,10)) - o = CartesianOrbit(pos=x, vel=v) - assert o.ndim == 2 - assert o.hamiltonian is None + o = Orbit(pos=x, vel=v) + assert o.pos.xyz.unit == u.kpc + assert o.vel.d_x.unit == u.km/u.s + + # TODO: don't support < 3 dim? + # x = np.random.random(size=(2,10)) + # v = np.random.random(size=(2,10)) + # o = Orbit(pos=x, vel=v) + # assert o.ndim == 2 + # assert o.hamiltonian is None # Check that passing in frame and potential or Hamiltonian works x = np.random.random(size=(3,10))*u.kpc @@ -150,57 +60,81 @@ def test_initialize(): potential = LogarithmicPotential(v_c=1., r_h=0.14, q1=1., q2=0.9, q3=1., units=galactic) - o = CartesianOrbit(pos=x, vel=v, frame=frame) + o = Orbit(pos=x, vel=v, frame=frame) assert o.hamiltonian is None assert o.potential is None - o = CartesianOrbit(pos=x, vel=v, potential=potential) + o = Orbit(pos=x, vel=v, potential=potential) assert o.hamiltonian is None assert o.frame is None - o = CartesianOrbit(pos=x, vel=v, potential=potential, frame=frame) - o = CartesianOrbit(pos=x, vel=v, - hamiltonian=Hamiltonian(potential, frame=frame)) + o = Orbit(pos=x, vel=v, potential=potential, frame=frame) + o = Orbit(pos=x, vel=v, + hamiltonian=Hamiltonian(potential, frame=frame)) assert isinstance(o.hamiltonian, Hamiltonian) assert isinstance(o.potential, LogarithmicPotential) assert isinstance(o.frame, StaticFrame) + # check that old class raises deprecation warning + from ..orbit import CartesianOrbit + warnings.simplefilter('always') + with pytest.warns(DeprecationWarning): + o = CartesianOrbit(pos=x, vel=v) + def test_from_w(): w = np.random.random(size=(6,10)) - o = CartesianOrbit.from_w(w, galactic) - assert o.pos.unit == u.kpc - assert o.vel.unit == u.kpc/u.Myr + o = Orbit.from_w(w, galactic) + assert o.pos.xyz.unit == u.kpc + assert o.vel.d_x.unit == u.kpc/u.Myr def test_slice(): # simple x = np.random.random(size=(3,10)) v = np.random.random(size=(3,10)) - o = CartesianOrbit(pos=x, vel=v) + o = Orbit(pos=x, vel=v) + new_o = o[:5] + assert new_o.shape == (5,) + + x = np.random.random(size=(3,10)) + v = np.random.random(size=(3,10)) + t = np.linspace(0,10,10) + o = Orbit(pos=x, vel=v, t=t) new_o = o[:5] - assert new_o.pos.shape == (3,5) + assert new_o.shape == (5,) # 1d slice on 3d - x = np.random.random(size=(3,100,8)) - v = np.random.random(size=(3,100,8)) + x = np.random.random(size=(3,10,8)) + v = np.random.random(size=(3,10,8)) t = np.arange(x.shape[1]) - o = CartesianOrbit(pos=x, vel=v, t=t) + o = Orbit(pos=x, vel=v, t=t) new_o = o[:5] - assert new_o.pos.shape == (3,5,8) + assert new_o.shape == (5,8) assert new_o.t.shape == (5,) + # pick a single orbit + new_o = o[:,0] + assert isinstance(new_o, Orbit) + assert new_o.shape == (10,) + assert new_o.t.shape == (10,) + + # pick a single time + new_o = o[3] + assert isinstance(new_o, PhaseSpacePosition) + assert new_o.shape == (8,) + # 3d slice on 3d - o = CartesianOrbit(pos=x, vel=v, t=t) + o = Orbit(pos=x, vel=v, t=t) new_o = o[:5,:4] - assert new_o.pos.shape == (3,5,4) + assert new_o.shape == (5,4) assert new_o.t.shape == (5,) # boolean array x = np.random.random(size=(3,10)) v = np.random.random(size=(3,10)) t = np.arange(x.shape[1]) - o = CartesianOrbit(pos=x, vel=v, t=t) + o = Orbit(pos=x, vel=v, t=t) ix = np.array([0,0,0,0,0,1,1,1,1,1]).astype(bool) new_o = o[ix] assert new_o.shape == (sum(ix),) @@ -210,7 +144,7 @@ def test_slice(): x = np.random.random(size=(3,10,4)) v = np.random.random(size=(3,10,4)) t = np.arange(x.shape[1]) - o = CartesianOrbit(pos=x, vel=v, t=t) + o = Orbit(pos=x, vel=v, t=t) ix = np.array([0,0,0,0,0,1,1,1,1,1]).astype(bool) new_o = o[ix] assert new_o.shape == (sum(ix),x.shape[-1]) @@ -220,7 +154,7 @@ def test_slice(): x = np.random.random(size=(3,10)) v = np.random.random(size=(3,10)) t = np.arange(x.shape[1]) - o = CartesianOrbit(pos=x, vel=v, t=t) + o = Orbit(pos=x, vel=v, t=t) ix = np.array([0,3,5]) new_o = o[ix] assert new_o.shape == (len(ix),) @@ -231,24 +165,25 @@ def test_represent_as(): # simple / unitless x = np.random.random(size=(3,10)) v = np.random.random(size=(3,10)) - o = CartesianOrbit(pos=x, vel=v) - sph_pos, sph_vel = o.represent_as(SphericalRepresentation) + o = Orbit(pos=x, vel=v) + sph = o.represent_as(SphericalRepresentation) - assert sph_pos.distance.unit == u.dimensionless_unscaled - assert sph_vel.unit == u.dimensionless_unscaled + assert sph.pos.distance.unit == u.one + assert sph.vel.d_distance.unit == u.one # simple / with units x = np.random.random(size=(3,10))*u.kpc v = np.random.normal(0.,100.,size=(3,10))*u.km/u.s - o = CartesianOrbit(pos=x, vel=v) - sph_pos, sph_vel = o.represent_as(SphericalRepresentation) - assert sph_pos.distance.unit == u.kpc + o = Orbit(pos=x, vel=v) + sph = o.represent_as(SphericalRepresentation) + assert sph.pos.distance.unit == u.kpc + assert sph.vel.d_distance.unit == u.km/u.s def test_to_coord_frame(): # simple / unitless x = np.random.random(size=(3,10)) v = np.random.random(size=(3,10)) - o = CartesianOrbit(pos=x, vel=v) + o = Orbit(pos=x, vel=v) with pytest.raises(u.UnitConversionError): o.to_coord_frame(Galactic) @@ -256,14 +191,14 @@ def test_to_coord_frame(): # simple / with units x = np.random.random(size=(3,10))*u.kpc v = np.random.normal(0.,100.,size=(3,10))*u.km/u.s - o = CartesianOrbit(pos=x, vel=v) + o = Orbit(pos=x, vel=v) coo,vel = o.to_coord_frame(Galactic) assert coo.name == 'galactic' # simple / with units and time x = np.random.random(size=(3,128,10))*u.kpc v = np.random.normal(0.,100.,size=(3,128,10))*u.km/u.s - o = CartesianOrbit(pos=x, vel=v) + o = Orbit(pos=x, vel=v) coo,vel = o.to_coord_frame(Galactic) assert coo.name == 'galactic' @@ -271,47 +206,47 @@ def test_w(): # simple / unitless x = np.random.random(size=(3,10)) v = np.random.random(size=(3,10)) - o = CartesianOrbit(pos=x, vel=v) + o = Orbit(pos=x, vel=v) w = o.w() - assert w.shape == (6,10,1) + assert w.shape == (6,10) # simple / with units x = np.random.random(size=(3,10))*u.kpc v = np.random.normal(0.,100.,size=(3,10))*u.km/u.s - o = CartesianOrbit(pos=x, vel=v) + o = Orbit(pos=x, vel=v) with pytest.raises(ValueError): o.w() w = o.w(units=galactic) - assert np.allclose(x.value, w[:3,:,0]) - assert np.allclose(v.value, (w[3:,:,0]*u.kpc/u.Myr).to(u.km/u.s).value) + assert np.allclose(x.value, w[:3,:]) + assert np.allclose(v.value, (w[3:,:]*u.kpc/u.Myr).to(u.km/u.s).value) # simple / with units and potential p = HernquistPotential(units=galactic, m=1E11, c=0.25) x = np.random.random(size=(3,10))*u.kpc v = np.random.normal(0.,100.,size=(3,10))*u.km/u.s - o = CartesianOrbit(pos=x, vel=v, potential=p, frame=StaticFrame(galactic)) + o = Orbit(pos=x, vel=v, potential=p, frame=StaticFrame(galactic)) w = o.w() - assert np.allclose(x.value, w[:3,:,0]) - assert np.allclose(v.value, (w[3:,:,0]*u.kpc/u.Myr).to(u.km/u.s).value) + assert np.allclose(x.value, w[:3,:]) + assert np.allclose(v.value, (w[3:,:]*u.kpc/u.Myr).to(u.km/u.s).value) w = o.w(units=solarsystem) - assert np.allclose(x.value, (w[:3,:,0]*u.au).to(u.kpc).value) - assert np.allclose(v.value, (w[3:,:,0]*u.au/u.yr).to(u.km/u.s).value) + assert np.allclose(x.value, (w[:3,:]*u.au).to(u.kpc).value) + assert np.allclose(v.value, (w[3:,:]*u.au/u.yr).to(u.km/u.s).value) def test_energy(): # with units x = np.random.random(size=(3,10))*u.kpc v = np.random.normal(0.,100.,size=(3,10))*u.km/u.s - o = CartesianOrbit(pos=x, vel=v) + o = Orbit(pos=x, vel=v) KE = o.kinetic_energy() - assert KE.unit == (o.vel.unit)**2 - assert KE.shape == o.pos.shape[1:] + assert KE.unit == (o.vel.d_x.unit)**2 + assert KE.shape == o.pos.shape # with units and potential p = HernquistPotential(units=galactic, m=1E11, c=0.25) x = np.random.random(size=(3,10))*u.kpc v = np.random.normal(0.,100.,size=(3,10))*u.km/u.s - o = CartesianOrbit(pos=x, vel=v, potential=p, frame=StaticFrame(galactic)) + o = Orbit(pos=x, vel=v, potential=p, frame=StaticFrame(galactic)) PE = o.potential_energy() E = o.energy() @@ -319,15 +254,15 @@ def test_angular_momentum(): # with units x = np.random.random(size=(3,10))*u.kpc v = np.random.normal(0.,100.,size=(3,10))*u.km/u.s - o = CartesianOrbit(pos=x, vel=v) + o = Orbit(pos=x, vel=v) L = o.angular_momentum() - assert L.unit == (o.vel.unit*o.pos.unit) - assert L.shape == o.pos.shape + assert L.unit == (o.vel.d_x.unit*o.pos.x.unit) + assert L.shape == ((3,) + o.shape) def test_eccentricity(): pot = KeplerPotential(m=1., units=solarsystem) - w0 = CartesianPhaseSpacePosition(pos=[1,0,0.]*u.au, - vel=[0.,2*np.pi,0.]*u.au/u.yr) + w0 = PhaseSpacePosition(pos=[1,0,0.]*u.au, + vel=[0.,2*np.pi,0.]*u.au/u.yr) ham = Hamiltonian(pot) w = ham.integrate_orbit(w0, dt=0.01, n_steps=10000, Integrator=DOPRI853Integrator) e = w.eccentricity() @@ -335,8 +270,8 @@ def test_eccentricity(): def test_apocenter_pericenter(): pot = KeplerPotential(m=1., units=solarsystem) - w0 = CartesianPhaseSpacePosition(pos=[1,0,0.]*u.au, - vel=[0.,1.5*np.pi,0.]*u.au/u.yr) + w0 = PhaseSpacePosition(pos=[1,0,0.]*u.au, + vel=[0.,1.5*np.pi,0.]*u.au/u.yr) ham = Hamiltonian(pot) w = ham.integrate_orbit(w0, dt=0.01, n_steps=10000, Integrator=DOPRI853Integrator) @@ -383,48 +318,102 @@ def test_estimate_period(): pos[1] = R*np.sin(phi) vel = np.zeros_like(pos) - orb = CartesianOrbit(pos*u.kpc, vel*u.kpc/u.Myr, t=t*u.Gyr) + orb = Orbit(pos*u.kpc, vel*u.kpc/u.Myr, t=t*u.Gyr) T = orb.estimate_period() assert np.allclose(T.value, true_T_R, rtol=1E-3) -def test_combine(): +def make_known_orbits(tmpdir, xs, vxs, potential, names): + # See Binney & Tremaine (2008) Figure 3.8 and 3.9 + E = -0.337 + y = 0. + + ws = [] + for x,vx,name in zip(xs, vxs, names): + vy = np.sqrt(2*(E - potential.value([x,y,0.]).value))[0] + w = [x,y,0.,vx,vy,0.] + ws.append(w) + ws = np.array(ws).T - o1 = CartesianOrbit.from_w(np.random.random(size=6), units=galactic) - o2 = CartesianOrbit.from_w(np.random.random(size=6), units=galactic) - o = combine((o1, o2)) - assert o.pos.shape == (3,1,2) - assert o.vel.shape == (3,1,2) + ham = Hamiltonian(potential) + orbit = ham.integrate_orbit(ws, dt=0.05, n_steps=10000) - o1 = CartesianOrbit.from_w(np.random.random(size=(6,11,1)), units=galactic) - o2 = CartesianOrbit.from_w(np.random.random(size=(6,11,10)), units=galactic) - o = combine((o1, o2)) - assert o.pos.shape == (3,11,11) - assert o.vel.shape == (3,11,11) + return orbit - o1 = CartesianOrbit.from_w(np.random.random(size=(6,20,1)), units=galactic) - o2 = CartesianOrbit.from_w(np.random.random(size=(6,10,1)), units=galactic) - o = combine((o1, o2), along_time_axis=True) - assert o.pos.shape == (3,30,1) - assert o.vel.shape == (3,30,1) +def test_circulation(tmpdir): - with pytest.raises(ValueError): - o1 = CartesianOrbit.from_w(np.random.random(size=(6,10,2)), units=galactic) - o2 = CartesianOrbit.from_w(np.random.random(size=(6,10,1)), units=galactic) - o = combine((o1, o2), along_time_axis=True) + potential = LogarithmicPotential(v_c=1., r_h=0.14, q1=1., q2=0.9, q3=1., + units=galactic) - # With time - t1 = np.linspace(0.,1,20) - t2 = np.linspace(0.,1,10) - o1 = CartesianOrbit.from_w(np.random.random(size=(6,20,1)), units=galactic, t=t1*u.Myr) - o2 = CartesianOrbit.from_w(np.random.random(size=(6,10,1)), units=galactic, t=t2*u.Myr) + # individual + ws = make_known_orbits(tmpdir, [0.5, 0], [0., 1.5], + potential, ["loop", "box"]) - with pytest.raises(ValueError): - o = combine((o1, o2)) + w1 = ws[:,0] + circ = w1.circulation() + assert circ.shape == (3,) + assert circ.sum() == 1 + + w2 = ws[:,1] + circ = w2.circulation() + assert circ.shape == (3,) + assert circ.sum() == 0 + + # try also for both, together + circ = ws.circulation() + assert circ.shape == (3,2) + assert np.allclose(circ.sum(axis=0), [1,0]) + +def test_align_circulation(): + + t = np.linspace(0,100,1024) + w = np.zeros((6,1024,4)) + + # loop around x axis + w[1,:,0] = np.cos(t) + w[2,:,0] = np.sin(t) + w[4,:,0] = -np.sin(t) + w[5,:,0] = np.cos(t) + + # loop around y axis + w[0,:,1] = -np.cos(t) + w[2,:,1] = np.sin(t) + w[3,:,1] = np.sin(t) + w[5,:,1] = np.cos(t) + + # loop around z axis + w[0,:,2] = np.cos(t) + w[1,:,2] = np.sin(t) + w[3,:,2] = -np.sin(t) + w[4,:,2] = np.cos(t) + + # box + w[0,:,3] = np.cos(t) + w[1,:,3] = -np.cos(0.5*t) + w[2,:,3] = np.cos(0.25*t) + w[3,:,3] = -np.sin(t) + w[4,:,3] = 0.5*np.sin(0.5*t) + w[5,:,3] = -0.25*np.sin(0.25*t) + + # First, individually + for i in range(w.shape[2]): + orb = Orbit.from_w(w[...,i], units=galactic) + new_orb = orb.align_circulation_with_z() + circ = new_orb.circulation() + + if i == 3: + assert np.sum(circ) == 0 + else: + assert circ[2] == 1. + + # all together now + orb = Orbit.from_w(w, units=galactic) + circ = orb.circulation() + assert circ.shape == (3,4) - o = combine((o1, o2), along_time_axis=True) - assert o.pos.shape == (3,30,1) - assert o.vel.shape == (3,30,1) - assert o.t.shape == (30,) + new_orb = orb.align_circulation_with_z() + new_circ = new_orb.circulation() + assert np.all(new_circ[2,:3] == 1.) + assert np.all(new_circ[:,3] == 0.) def test_frame_transform(): static = StaticFrame(galactic) @@ -435,7 +424,7 @@ def test_frame_transform(): t = np.linspace(0,1,10)*u.Myr # no frame specified at init - o = CartesianOrbit(pos=x, vel=v, t=t) + o = Orbit(pos=x, vel=v, t=t) with pytest.raises(ValueError): o.to_frame(rotating) @@ -443,8 +432,8 @@ def test_frame_transform(): o.to_frame(rotating, current_frame=static) # frame specified at init - o = CartesianOrbit(pos=x, vel=v, t=t, - frame=static, - potential=HernquistPotential(m=1E10, c=0.5, units=galactic)) + o = Orbit(pos=x, vel=v, t=t, + frame=static, + potential=HernquistPotential(m=1E10, c=0.5, units=galactic)) o.to_frame(rotating) o.to_frame(rotating, t=o.t) diff --git a/gala/dynamics/tests/test_plot.py b/gala/dynamics/tests/test_plot.py index 140423584..753a5e5b0 100644 --- a/gala/dynamics/tests/test_plot.py +++ b/gala/dynamics/tests/test_plot.py @@ -1,90 +1,130 @@ # coding: utf-8 -""" Test dynamics plotting """ +""" Test dynamics plotting functions """ from __future__ import division, print_function - -# Standard library -import os -import logging - # Third-party import numpy as np -from astropy import log as logger -import gala.dynamics as gd +import astropy.units as u # Project -from ..plot import * +from ..core import PhaseSpacePosition +from ..orbit import Orbit +from ..plot import plot_projections + +class TestPlotting(object): + + def setup(self): + + psps = [] + psps.append(PhaseSpacePosition(pos=np.random.random(size=3), + vel=np.random.random(size=3))) + psps.append(PhaseSpacePosition(pos=np.random.random(size=(3,16)), + vel=np.random.random(size=(3,16)))) + psps.append(PhaseSpacePosition(pos=np.random.random(size=(3,16))*u.kpc, + vel=np.random.random(size=(3,16))*u.km/u.s)) + + orbits = [] + orbits.append(Orbit(pos=np.random.random(size=(3,16)), + vel=np.random.random(size=(3,16)), + t=np.linspace(0,1,16))) + orbits.append(Orbit(pos=np.random.random(size=(3,16,2)), + vel=np.random.random(size=(3,16,2)), + t=np.linspace(0,1,16))) + orbits.append(Orbit(pos=np.random.random(size=(3,16))*u.kpc, + vel=np.random.random(size=(3,16))*u.km/u.s, + t=np.linspace(0,1,16)*u.Myr)) + + self.psps = psps + self.orbits = orbits + + def test_plot_projections(self): + import matplotlib.pyplot as plt + + # TODO: need major test improvements here + # let the function create the figure + for p in self.psps: + _ = p.plot() + + for o in self.orbits: + _ = o.plot() + + x = self.psps[0].pos.xyz.value + fig,axes = plt.subplots(1,2) + fig = plot_projections(x[:2], autolim=True, axes=axes, + subplots_kwargs=dict(sharex=True), + labels=['x', 'y'], + plot_function=plt.plot, + marker='o', linestyle='--', color='r') -logger.setLevel(logging.DEBUG) -def test_orbits(tmpdir): +# def test_orbits(tmpdir): - # generate an "orbit" - n = 8 - t = np.linspace(0, 100, 1000).reshape(1000,1) - x = np.cos(np.random.uniform(1.,8.,size=(1,n))*t).T[None] - y = np.sin(np.random.uniform(1.,8.,size=(1,n))*t).T[None] - z = np.cos(np.random.uniform(1.,8.,size=(1,n))*t).T[None] - vx = vy = vz = np.zeros_like(x) - w = np.rollaxis(np.vstack((x,y,z,vx,vy,vz)), -1, 1) +# # generate an "orbit" +# n = 8 +# t = np.linspace(0, 100, 1000).reshape(1000,1) +# x = np.cos(np.random.uniform(1.,8.,size=(1,n))*t).T[None] +# y = np.sin(np.random.uniform(1.,8.,size=(1,n))*t).T[None] +# z = np.cos(np.random.uniform(1.,8.,size=(1,n))*t).T[None] +# vx = vy = vz = np.zeros_like(x) +# w = np.rollaxis(np.vstack((x,y,z,vx,vy,vz)), -1, 1) - fig = plot_orbits(w, linestyle='none', marker='.', alpha=0.25) - fig.savefig(os.path.join(str(tmpdir), "all_orbits.png")) +# fig = plot_orbits(w, linestyle='none', marker='.', alpha=0.25) +# fig.savefig(os.path.join(str(tmpdir), "all_orbits.png")) - fig = plot_orbits(w, ix=0, linestyle='none', marker='.', alpha=0.25) - fig.savefig(os.path.join(str(tmpdir), "one_orbit.png")) +# fig = plot_orbits(w, ix=0, linestyle='none', marker='.', alpha=0.25) +# fig.savefig(os.path.join(str(tmpdir), "one_orbit.png")) - fig = plot_orbits(w, ix=0, linestyle='none', marker='.', alpha=0.25, - labels=("herp","derp","merp")) - fig.savefig(os.path.join(str(tmpdir), "one_orbit_labels.png")) +# fig = plot_orbits(w, ix=0, linestyle='none', marker='.', alpha=0.25, +# labels=("herp","derp","merp")) +# fig.savefig(os.path.join(str(tmpdir), "one_orbit_labels.png")) - fig = plot_orbits(w, triangle=True, linestyle='-', marker=None) - fig.savefig(os.path.join(str(tmpdir), "all_orbits_triangle.png")) +# fig = plot_orbits(w, triangle=True, linestyle='-', marker=None) +# fig.savefig(os.path.join(str(tmpdir), "all_orbits_triangle.png")) - fig = plot_orbits(w, ix=0, triangle=True, linestyle='-', marker=None) - fig.savefig(os.path.join(str(tmpdir), "one_orbit_triangle.png")) +# fig = plot_orbits(w, ix=0, triangle=True, linestyle='-', marker=None) +# fig.savefig(os.path.join(str(tmpdir), "one_orbit_triangle.png")) - fig = plot_orbits(w, ix=0, triangle=True, linestyle='-', marker=None, - labels=("herp","derp","merp")) - fig.savefig(os.path.join(str(tmpdir), "one_orbit_triangle_labels.png")) +# fig = plot_orbits(w, ix=0, triangle=True, linestyle='-', marker=None, +# labels=("herp","derp","merp")) +# fig.savefig(os.path.join(str(tmpdir), "one_orbit_triangle_labels.png")) -def test_three_panel(tmpdir): +# def test_three_panel(tmpdir): - q = np.random.uniform(0.,10.,size=(1000,3)) - q0 = np.array([5,5,5]) +# q = np.random.uniform(0.,10.,size=(1000,3)) +# q0 = np.array([5,5,5]) - fig = three_panel(q) - fig.savefig(os.path.join(str(tmpdir), "three-panel-random.png")) +# fig = three_panel(q) +# fig.savefig(os.path.join(str(tmpdir), "three-panel-random.png")) - fig = three_panel(q, triangle=True) - fig.savefig(os.path.join(str(tmpdir), "three-panel-random_triangle.png")) +# fig = three_panel(q, triangle=True) +# fig.savefig(os.path.join(str(tmpdir), "three-panel-random_triangle.png")) - fig = three_panel(q, relative_to=q0, labels=[r'$\Omega_1$',r'$\Omega_2$',r'$\Omega_3$']) - fig.savefig(os.path.join(str(tmpdir), "three-panel-random-relative.png")) +# fig = three_panel(q, relative_to=q0, labels=[r'$\Omega_1$',r'$\Omega_2$',r'$\Omega_3$']) +# fig.savefig(os.path.join(str(tmpdir), "three-panel-random-relative.png")) - fig = three_panel(q, relative_to=q0, triangle=True, labels=[r'$\Omega_1$',r'$\Omega_2$',r'$\Omega_3$']) - fig.savefig(os.path.join(str(tmpdir), "three-panel-random-relative_triangle.png")) +# fig = three_panel(q, relative_to=q0, triangle=True, labels=[r'$\Omega_1$',r'$\Omega_2$',r'$\Omega_3$']) +# fig.savefig(os.path.join(str(tmpdir), "three-panel-random-relative_triangle.png")) -def test_1d(tmpdir): +# def test_1d(tmpdir): - t = np.linspace(0,100.,1000) - q = np.cos(2*np.pi*t/10.)[None] +# t = np.linspace(0,100.,1000) +# q = np.cos(2*np.pi*t/10.)[None] - fig = plot_orbits(q, labels=(r"$\theta$",)) - fig.savefig(os.path.join(str(tmpdir), "1d-orbit-labels.png")) +# fig = plot_orbits(q, labels=(r"$\theta$",)) +# fig.savefig(os.path.join(str(tmpdir), "1d-orbit-labels.png")) - fig = plot_orbits(q, t=t, labels=(r"$\theta$",)) - fig.savefig(os.path.join(str(tmpdir), "1d-orbit-labels-time.png")) +# fig = plot_orbits(q, t=t, labels=(r"$\theta$",)) +# fig.savefig(os.path.join(str(tmpdir), "1d-orbit-labels-time.png")) -def test_2d(tmpdir): +# def test_2d(tmpdir): - t = np.linspace(0,100.,1000) +# t = np.linspace(0,100.,1000) - q = np.zeros((2,len(t))) - q[0,:] = np.cos(2*np.pi*t/10.) - q[1,:] = np.sin(2*np.pi*t/5.5) +# q = np.zeros((2,len(t))) +# q[0,:] = np.cos(2*np.pi*t/10.) +# q[1,:] = np.sin(2*np.pi*t/5.5) - fig = plot_orbits(q, labels=(r"$\theta$",r"$\omega$")) - fig.savefig(os.path.join(str(tmpdir), "2d-orbit-labels.png")) +# fig = plot_orbits(q, labels=(r"$\theta$",r"$\omega$")) +# fig.savefig(os.path.join(str(tmpdir), "2d-orbit-labels.png")) diff --git a/gala/dynamics/tests/test_representation_nd.py b/gala/dynamics/tests/test_representation_nd.py new file mode 100644 index 000000000..e4e09bedb --- /dev/null +++ b/gala/dynamics/tests/test_representation_nd.py @@ -0,0 +1,74 @@ +# coding: utf-8 + +from __future__ import division, print_function + +# Third-party +import astropy.units as u +from astropy.coordinates import Galactic, CartesianRepresentation +from astropy.tests.helper import quantity_allclose +import numpy as np +import pytest + +# Project +from ..core import PhaseSpacePosition +from ...potential import Hamiltonian, HernquistPotential +from ...potential.frame import StaticFrame, ConstantRotatingFrame +from ...units import galactic, solarsystem +from ..representation_nd import (NDCartesianRepresentation, + NDCartesianDifferential) + +def test_init_repr(): + + # Passing in x1, x2 + rep = NDCartesianRepresentation([1., 1.]) + assert rep.xyz.shape == (2,) + + # Passing in x1, x2 + rep = NDCartesianRepresentation(np.random.random(size=(2, 8))) + assert rep.xyz.shape == (2,8) + rep[:1] + + for n in range(1, 6+1): + print('N: '+str(n)) + + xs = np.random.uniform(size=(n, 16)) * u.one + rep = NDCartesianRepresentation(xs) + for i in range(1, n+1): + assert hasattr(rep, 'x'+str(i)) + + xs2 = rep.xyz + assert quantity_allclose(xs, xs2) + + rep2 = rep[:8] + + assert rep.shape == (16,) + assert rep2.shape == (8,) + +def test_init_diff(): + + # Passing in x1, x2 + rep = NDCartesianDifferential([1., 1.]) + assert rep.d_xyz.shape == (2,) + with pytest.raises(TypeError): + rep[:1] + + # Passing in x1, x2 + rep = NDCartesianDifferential(np.random.random(size=(2, 8))) + assert rep.d_xyz.shape == (2,8) + rep[:1] + + for n in range(1, 6+1): + print('N: '+str(n)) + + xs = np.random.uniform(size=(n, 16)) * u.one + rep = NDCartesianDifferential(xs) + for i in range(1, n+1): + assert hasattr(rep, 'd_x'+str(i)) + + xs2 = rep.d_xyz + assert quantity_allclose(xs, xs2) + + rep2 = rep[:8] + + assert rep.shape == (16,) + assert rep2.shape == (8,) diff --git a/gala/dynamics/util.py b/gala/dynamics/util.py index 442fa7e8f..c9a99b4a8 100644 --- a/gala/dynamics/util.py +++ b/gala/dynamics/util.py @@ -12,7 +12,7 @@ from scipy.signal import argrelmax, argrelmin # This package -from .core import CartesianPhaseSpacePosition +from .core import PhaseSpacePosition from ..integrate import LeapfrogIntegrator __all__ = ['peak_to_peak_period', 'estimate_dt_n_steps'] @@ -70,9 +70,11 @@ def peak_to_peak_period(t, f, amplitude_threshold=1E-2): # then take the mean of these two return np.mean([T_max, T_min]) * t_unit -def _autodetermine_initial_dt(w0, potential, dE_threshold=1E-9, Integrator=LeapfrogIntegrator): - if w0.shape[0] > 1: - raise ValueError("Only one set of initial conditions may be passed in at a time.") +def _autodetermine_initial_dt(w0, potential, dE_threshold=1E-9, + Integrator=LeapfrogIntegrator): + if w0.shape and w0.shape[0] > 1: + raise ValueError("Only one set of initial conditions may be passed " + "in at a time.") if dE_threshold is None: return 1. @@ -82,7 +84,8 @@ def _autodetermine_initial_dt(w0, potential, dE_threshold=1E-9, Integrator=Leapf for dt in dts: n_steps = int(round(_base_n_steps / dt)) - orbit = potential.integrate_orbit(w0, dt=dt, n_steps=n_steps, Integrator=Integrator) + orbit = potential.integrate_orbit(w0, dt=dt, n_steps=n_steps, + Integrator=Integrator) E = orbit.energy() dE = np.abs((E[-1] - E[0]) / E[0]).value @@ -91,8 +94,8 @@ def _autodetermine_initial_dt(w0, potential, dE_threshold=1E-9, Integrator=Leapf return dt -def estimate_dt_n_steps(w0, potential, n_periods, n_steps_per_period, dE_threshold=1E-9, - func=np.nanmax): +def estimate_dt_n_steps(w0, potential, n_periods, n_steps_per_period, + dE_threshold=1E-9, func=np.nanmax): """ Estimate the timestep and number of steps to integrate an orbit for given its initial conditions and a potential object. @@ -108,12 +111,12 @@ def estimate_dt_n_steps(w0, potential, n_periods, n_steps_per_period, dE_thresho n_steps_per_period : int Number of steps to take per (max) orbital period. dE_threshold : numeric (optional) - Maximum fractional energy difference -- used to determine initial timestep. - Set to ``None`` to ignore this. + Maximum fractional energy difference -- used to determine initial + timestep. Set to ``None`` to ignore this. func : callable (optional) - Determines which period to use. By default, this takes the maximum period using - :func:`~numpy.nanmax`. Other options could be :func:`~numpy.nanmin`, - :func:`~numpy.nanmean`, :func:`~numpy.nanmedian`. + Determines which period to use. By default, this takes the maximum + period using :func:`~numpy.nanmax`. Other options could be + :func:`~numpy.nanmin`, :func:`~numpy.nanmean`, :func:`~numpy.nanmedian`. Returns ------- @@ -123,9 +126,9 @@ def estimate_dt_n_steps(w0, potential, n_periods, n_steps_per_period, dE_thresho The number of timesteps to integrate for. """ - if not isinstance(w0, CartesianPhaseSpacePosition): + if not isinstance(w0, PhaseSpacePosition): w0 = np.asarray(w0) - w0 = CartesianPhaseSpacePosition.from_w(w0, units=potential.units) + w0 = PhaseSpacePosition.from_w(w0, units=potential.units) # integrate orbit dt = _autodetermine_initial_dt(w0, potential, dE_threshold=dE_threshold) @@ -136,17 +139,19 @@ def estimate_dt_n_steps(w0, potential, n_periods, n_steps_per_period, dE_thresho circ = orbit.circulation() if np.any(circ): orbit = orbit.align_circulation_with_z(circulation=circ) - cyl,_ = orbit.represent_as(coord.CylindricalRepresentation) # ignore velocity return + cyl = orbit.represent_as(coord.CylindricalRepresentation) # convert to cylindrical coordinates R = cyl.rho.value phi = cyl.phi.value z = cyl.z.value - T = np.array([peak_to_peak_period(orbit.t, f).value for f in [R, phi, z]])*orbit.t.unit + T = np.array([peak_to_peak_period(orbit.t, f).value + for f in [R, phi, z]])*orbit.t.unit else: - T = np.array([peak_to_peak_period(orbit.t, f).value for f in orbit.pos])*orbit.t.unit + T = np.array([peak_to_peak_period(orbit.t, f).value + for f in orbit.pos])*orbit.t.unit # timestep from number of steps per period T = func(T) diff --git a/gala/integrate/core.py b/gala/integrate/core.py index 9354163e2..801c05090 100644 --- a/gala/integrate/core.py +++ b/gala/integrate/core.py @@ -39,14 +39,12 @@ def _prepare_ws(self, w0, mmap, n_steps): particularly useful for integrating a large number of orbits or integrating a large number of time steps. """ - from ..dynamics import CartesianPhaseSpacePosition - if not isinstance(w0, CartesianPhaseSpacePosition): - w0 = np.asarray(w0) - ndim = w0.shape[0]//2 - w0 = CartesianPhaseSpacePosition(pos=w0[:ndim], - vel=w0[ndim:]) + from ..dynamics import PhaseSpacePosition + if not isinstance(w0, PhaseSpacePosition): + w0 = PhaseSpacePosition.from_w(w0) arr_w0 = w0.w(self._func_units) + self.ndim,self.norbits = arr_w0.shape self.ndim = self.ndim//2 @@ -78,17 +76,17 @@ def _handle_output(self, w0, t, w): t_unit = self._func_units['time'] vel_unit = pos_unit / t_unit - from ..dynamics import CartesianOrbit - orbit = CartesianOrbit(pos=w[:self.ndim]*pos_unit, - vel=w[self.ndim:]*vel_unit, - t=t*t_unit) # HACK: BADDDD + from ..dynamics import Orbit + orbit = Orbit(pos=w[:self.ndim]*pos_unit, + vel=w[self.ndim:]*vel_unit, + t=t*t_unit) # HACK: BADDDD return orbit def run(self): """ Run the integrator starting from the specified phase-space position. - The initial conditions ``w0`` should be a `~gala.dynamics.CartesianPhaseSpacePosition` + The initial conditions ``w0`` should be a `~gala.dynamics.PhaseSpacePosition` instance. There are a few combinations of keyword arguments accepted for @@ -105,7 +103,7 @@ def run(self): .. warning:: - Right now, this always returns a `~gala.dynamics.CartesianOrbit` instance. + Right now, this always returns a `~gala.dynamics.Orbit` instance. This is wrong and will change! .. todo:: @@ -114,7 +112,7 @@ def run(self): Parameters ---------- - w0 : `~gala.dynamics.CartesianPhaseSpacePosition` + w0 : `~gala.dynamics.PhaseSpacePosition` Initial conditions. **time_spec Timestep information passed to @@ -122,7 +120,7 @@ def run(self): Returns ------- - orbit : `~gala.dynamics.CartesianOrbit` + orbit : `~gala.dynamics.Orbit` """ pass diff --git a/gala/observation/__init__.py b/gala/observation/__init__.py deleted file mode 100644 index bb67a43fa..000000000 --- a/gala/observation/__init__.py +++ /dev/null @@ -1 +0,0 @@ -from .core import * diff --git a/gala/observation/core.py b/gala/observation/core.py deleted file mode 100644 index 0da4716e5..000000000 --- a/gala/observation/core.py +++ /dev/null @@ -1,89 +0,0 @@ -# coding: utf-8 - -""" Convenience functions for observational data """ - -from __future__ import division, print_function - - -# Third-party -import numpy as np -import astropy.units as u - -__all__ = ["apparent_magnitude", "absolute_magnitude", - "distance_modulus", "distance"] - -def apparent_magnitude(mag_abs, distance): - """ - Compute the apparent magnitude of a source given an absolute magnitude - and a distance. - - Parameters - ---------- - mag_abs : numeric or iterable - Absolute magnitude. - distance : :class:`~astropy.units.Quantity` - The distance to the source as a Quantity object. - - Returns - ------- - mag_app : :class:`~numpy.ndarray` - The apparent magnitude. - """ - - # Compute the apparent magnitude - return mag_abs + distance_modulus(distance) - -def absolute_magnitude(mag_app, distance): - """ - Compute the absolute magnitude of a source given an apparent magnitude - and a distance. - - Parameters - ---------- - mag_app : numeric or iterable - Apparent magnitude. - distance : :class:`~astropy.units.Quantity` - The distance to the source as a Quantity object. - - Returns - ------- - mag_abs : :class:`~numpy.ndarray` - The absolute magnitude. - """ - - # Compute the absolute magnitude - return mag_app - distance_modulus(distance) - -def distance_modulus(distance): - """ - Compute the distance modulus given a distance. - - Parameters - ---------- - distance : astropy.units.Quantity - The distance as a Quantity object. - - Returns - ------- - distance_mod : :class:`~numpy.ndarray` - The distance modulus. - """ - - return 5.*(np.log10(distance.to(u.pc).value) - 1.) - -def distance(distance_mod): - """ - Compute the distance from the distance modulus. - - Parameters - ---------- - distance_mod : astropy.units.Quantity - The distance modulus. - - Returns - ------- - distance : :class:`~astropy.units.Quantity` - Distance. - - """ - return 10**(distance_mod/5. + 1) * u.pc diff --git a/gala/observation/tests/__init__.py b/gala/observation/tests/__init__.py deleted file mode 100644 index e69de29bb..000000000 diff --git a/gala/observation/tests/test_core.py b/gala/observation/tests/test_core.py deleted file mode 100644 index edd2647f1..000000000 --- a/gala/observation/tests/test_core.py +++ /dev/null @@ -1,18 +0,0 @@ -# coding: utf-8 -""" - Test the core observational utilities. -""" - -from __future__ import absolute_import, unicode_literals, division, print_function - - -# Standard library -import os, sys - -# Third-party -import astropy.units as u -import numpy as np -import pytest - -from ..core import * - diff --git a/gala/potential/frame/builtin/transformations.py b/gala/potential/frame/builtin/transformations.py index cd66d8b9a..cb94ae700 100644 --- a/gala/potential/frame/builtin/transformations.py +++ b/gala/potential/frame/builtin/transformations.py @@ -27,8 +27,8 @@ def rodrigues_axis_angle_rotate(x, vec, theta): vec = np.array(vec).T theta = np.array(theta).T[...,None] - out = np.cos(theta)*x + np.sin(theta)*np.cross(vec, x) \ - + (1 - np.cos(theta)) * (vec * x).sum(axis=-1)[...,None] * vec + out = np.cos(theta)*x + np.sin(theta)*np.cross(vec, x) + \ + (1 - np.cos(theta)) * (vec * x).sum(axis=-1)[...,None] * vec return out.T @@ -53,6 +53,7 @@ def z_angle_rotate(xy, theta): return out.T def _constantrotating_static_helper(frame_r, frame_i, w, t=None, sign=1.): + # TODO: use representation arithmetic instead Omega = -frame_r.parameters['Omega'].decompose(frame_i.units).value if not isinstance(w, Orbit) and t is None: @@ -65,10 +66,23 @@ def _constantrotating_static_helper(frame_r, frame_i, w, t=None, sign=1.): elif not hasattr(t, 'unit'): t = t * frame_i.units['time'] + if t is None: + raise ValueError('Time must be supplied either through the input ' + 'Orbit class instance or through the t argument.') t = t.decompose(frame_i.units).value - pos = w.pos.decompose(frame_i.units).value - vel = w.vel.decompose(frame_i.units).value + # HACK: this is a little bit crazy...this makes it so that !=3D + # representations will work here + if hasattr(w.pos, 'xyz'): + pos = w.pos + vel = w.vel + else: + cart = w.cartesian + pos = cart.pos + vel = cart.vel + + pos = pos.xyz.decompose(frame_i.units).value + vel = vel.d_xyz.decompose(frame_i.units).value # get rotation angle, axis vs. time if isiterable(Omega): # 3D @@ -95,7 +109,7 @@ def static_to_constantrotating(frame_i, frame_r, w, t=None): ---------- frame_i : `~gala.potential.StaticFrame` frame_r : `~gala.potential.ConstantRotatingFrame` - w : `~gala.dynamics.CartesianPhaseSpacePosition`, `~gala.dynamics.CartesianOrbit` + w : `~gala.dynamics.PhaseSpacePosition`, `~gala.dynamics.Orbit` t : quantity_like (optional) Required if input coordinates are just a phase-space position. @@ -117,7 +131,7 @@ def constantrotating_to_static(frame_r, frame_i, w, t=None): ---------- frame_i : `~gala.potential.StaticFrame` frame_r : `~gala.potential.ConstantRotatingFrame` - w : `~gala.dynamics.CartesianPhaseSpacePosition`, `~gala.dynamics.CartesianOrbit` + w : `~gala.dynamics.PhaseSpacePosition`, `~gala.dynamics.Orbit` t : quantity_like (optional) Required if input coordinates are just a phase-space position. diff --git a/gala/potential/frame/cframe.pyx b/gala/potential/frame/cframe.pyx index b5abc3122..5b63b3555 100644 --- a/gala/potential/frame/cframe.pyx +++ b/gala/potential/frame/cframe.pyx @@ -15,7 +15,7 @@ np.import_array() from .core import FrameBase from ..potential.cpotential import _validate_pos_arr from ..potential.cpotential cimport energyfunc, gradientfunc, hessianfunc -from ...dynamics import CartesianPhaseSpacePosition +from ...dynamics import PhaseSpacePosition cdef class CFrameWrapper: """ Wrapper class for C implementation of reference frames. """ diff --git a/gala/potential/frame/tests/test_transformations.py b/gala/potential/frame/tests/test_transformations.py index d712c1365..e6518278b 100644 --- a/gala/potential/frame/tests/test_transformations.py +++ b/gala/potential/frame/tests/test_transformations.py @@ -11,7 +11,7 @@ from ..builtin.transformations import (static_to_constantrotating, constantrotating_to_static, rodrigues_axis_angle_rotate) -from ....dynamics import CartesianOrbit, CartesianPhaseSpacePosition +from ....dynamics import Orbit, PhaseSpacePosition from ....units import galactic, dimensionless def test_axis_angle_rotate(): @@ -35,36 +35,42 @@ def test_axis_angle_rotate(): def _helper(fi, fr, w, t=None): pos_r,vel_r = static_to_constantrotating(fi, fr, w, t=t) - w2 = CartesianOrbit(pos=pos_r, vel=vel_r, t=t) + if isinstance(w, Orbit): + w2 = Orbit(pos=pos_r, vel=vel_r, t=t) + else: + w2 = PhaseSpacePosition(pos=pos_r, vel=vel_r) pos_i,vel_i = constantrotating_to_static(fr, fi, w2, t=t) - assert quantity_allclose(pos_i, w.pos) - assert quantity_allclose(vel_i, w.vel) + assert quantity_allclose(pos_i, w.pos.xyz) + assert quantity_allclose(vel_i, w.vel.d_xyz) pos_i,vel_i = constantrotating_to_static(fr, fi, w, t=t) - w2 = CartesianOrbit(pos=pos_i, vel=vel_i, t=t) + if isinstance(w, Orbit): + w2 = Orbit(pos=pos_i, vel=vel_i, t=t) + else: + w2 = PhaseSpacePosition(pos=pos_i, vel=vel_i) pos_r,vel_r = static_to_constantrotating(fi, fr, w2, t=t) - assert quantity_allclose(pos_r, w.pos) - assert quantity_allclose(vel_r, w.vel) + assert quantity_allclose(pos_r, w.pos.xyz) + assert quantity_allclose(vel_r, w.vel.d_xyz) def test_frame_transforms_3d(): frame_i = StaticFrame(units=galactic) frame_r = ConstantRotatingFrame(Omega=[0.112, 1.235, 0.8656]*u.rad/u.Myr, units=galactic) - w = CartesianOrbit(pos=np.random.random(size=(3,32))*u.kpc, - vel=np.random.random(size=(3,32))*u.kpc/u.Myr, - t=np.linspace(0,1,32)*u.Myr) + w = Orbit(pos=np.random.random(size=(3,32))*u.kpc, + vel=np.random.random(size=(3,32))*u.kpc/u.Myr, + t=np.linspace(0,1,32)*u.Myr) _helper(frame_i, frame_r, w, t=w.t) - w = CartesianOrbit(pos=np.random.random(size=(3,32,8))*u.kpc, - vel=np.random.random(size=(3,32,8))*u.kpc/u.Myr, - t=np.linspace(0,1,32)*u.Myr) + w = Orbit(pos=np.random.random(size=(3,32,8))*u.kpc, + vel=np.random.random(size=(3,32,8))*u.kpc/u.Myr, + t=np.linspace(0,1,32)*u.Myr) _helper(frame_i, frame_r, w, t=w.t) - w = CartesianPhaseSpacePosition(pos=np.random.random(size=3)*u.kpc, - vel=np.random.random(size=3)*u.kpc/u.Myr) + w = PhaseSpacePosition(pos=np.random.random(size=3)*u.kpc, + vel=np.random.random(size=3)*u.kpc/u.Myr) with pytest.raises(ValueError): _helper(frame_i, frame_r, w) _helper(frame_i, frame_r, w, t=0.*u.Myr) @@ -75,18 +81,18 @@ def test_frame_transforms_2d(): frame_r = ConstantRotatingFrame(Omega=0.529*u.rad/u.Myr, units=galactic) - w = CartesianOrbit(pos=np.random.random(size=(2,32))*u.kpc, - vel=np.random.random(size=(2,32))*u.kpc/u.Myr, - t=np.linspace(0,1,32)*u.Myr) + w = Orbit(pos=np.random.random(size=(2,32))*u.kpc, + vel=np.random.random(size=(2,32))*u.kpc/u.Myr, + t=np.linspace(0,1,32)*u.Myr) _helper(frame_i, frame_r, w, t=w.t) - w = CartesianOrbit(pos=np.random.random(size=(2,32,8))*u.kpc, - vel=np.random.random(size=(2,32,8))*u.kpc/u.Myr, - t=np.linspace(0,1,32)*u.Myr) + w = Orbit(pos=np.random.random(size=(2,32,8))*u.kpc, + vel=np.random.random(size=(2,32,8))*u.kpc/u.Myr, + t=np.linspace(0,1,32)*u.Myr) _helper(frame_i, frame_r, w, t=w.t) - w = CartesianPhaseSpacePosition(pos=np.random.random(size=2)*u.kpc, - vel=np.random.random(size=2)*u.kpc/u.Myr) + w = PhaseSpacePosition(pos=np.random.random(size=2)*u.kpc, + vel=np.random.random(size=2)*u.kpc/u.Myr) with pytest.raises(ValueError): _helper(frame_i, frame_r, w) _helper(frame_i, frame_r, w, t=0.*u.Myr) diff --git a/gala/potential/hamiltonian/chamiltonian.pyx b/gala/potential/hamiltonian/chamiltonian.pyx index 24dc7029d..4d9ce8b6f 100644 --- a/gala/potential/hamiltonian/chamiltonian.pyx +++ b/gala/potential/hamiltonian/chamiltonian.pyx @@ -11,7 +11,7 @@ from ..common import CommonBase from ..potential import PotentialBase, CPotentialBase from ..frame import FrameBase, CFrameBase, StaticFrame from ...integrate import LeapfrogIntegrator, DOPRI853Integrator -from ...dynamics import PhaseSpacePosition, CartesianOrbit, CartesianPhaseSpacePosition +from ...dynamics import PhaseSpacePosition, Orbit __all__ = ["Hamiltonian"] @@ -254,7 +254,7 @@ class Hamiltonian(CommonBase): Returns ------- - orbit : `~gala.dynamics.CartesianOrbit` + orbit : `~gala.dynamics.Orbit` """ @@ -266,11 +266,10 @@ class Hamiltonian(CommonBase): # use the Integrator provided pass - if not isinstance(w0, CartesianPhaseSpacePosition): + if not isinstance(w0, PhaseSpacePosition): w0 = np.asarray(w0) ndim = w0.shape[0]//2 - w0 = CartesianPhaseSpacePosition(pos=w0[:ndim], - vel=w0[ndim:]) + w0 = PhaseSpacePosition(pos=w0[:ndim], vel=w0[ndim:]) ndim = w0.ndim arr_w0 = w0.w(self.units) @@ -315,7 +314,7 @@ class Hamiltonian(CommonBase): tunit = self.units['time'] except (TypeError, AttributeError): tunit = u.dimensionless_unscaled - return CartesianOrbit.from_w(w=w, units=self.units, t=t*tunit, hamiltonian=self) + return Orbit.from_w(w=w, units=self.units, t=t*tunit, hamiltonian=self) def save(self, f): """ diff --git a/gala/potential/hamiltonian/tests/helpers.py b/gala/potential/hamiltonian/tests/helpers.py index ebb864658..a081f7f41 100644 --- a/gala/potential/hamiltonian/tests/helpers.py +++ b/gala/potential/hamiltonian/tests/helpers.py @@ -7,10 +7,10 @@ import numpy as np # Project -from ....dynamics import CartesianPhaseSpacePosition, CartesianOrbit +from ....dynamics import PhaseSpacePosition, Orbit from ....units import galactic -PSP = CartesianPhaseSpacePosition -ORB = CartesianOrbit +PSP = PhaseSpacePosition +ORB = Orbit class _TestBase(object): use_half_ndim = False @@ -65,9 +65,9 @@ def setup_class(cls): vel=np.random.random(size=(ndim//2,ntimes)))) cls.w0s.append(ORB(pos=np.random.random(size=(ndim//2,ntimes))*u.kpc, vel=np.random.random(size=(ndim//2,ntimes))*u.km/u.s)) - cls.energy_return_shapes += [(ntimes,1)]*2 - cls.gradient_return_shapes += [(r_ndim,ntimes,1)]*2 - cls.hessian_return_shapes += [(r_ndim,r_ndim,ntimes,1)]*2 + cls.energy_return_shapes += [(ntimes,)]*2 + cls.gradient_return_shapes += [(r_ndim,ntimes,)]*2 + cls.hessian_return_shapes += [(r_ndim,r_ndim,ntimes,)]*2 # 3D - orbit cls.w0s.append(ORB(pos=np.random.random(size=(ndim//2,ntimes,norbits)), @@ -91,7 +91,7 @@ def setup_class(cls): def test_energy(self): for arr,shp in zip(self.w0s, self.energy_return_shapes): if self.E_unit.is_equivalent(u.one) and hasattr(arr, 'pos') and \ - not arr.pos.unit.is_equivalent(u.one): + not arr.pos.xyz.unit.is_equivalent(u.one): continue v = self.obj.energy(arr) @@ -106,7 +106,7 @@ def test_energy(self): def test_gradient(self): for arr,shp in zip(self.w0s, self.gradient_return_shapes): if self.E_unit.is_equivalent(u.one) and hasattr(arr, 'pos') and \ - not arr.pos.unit.is_equivalent(u.one): + not arr.pos.xyz.unit.is_equivalent(u.one): continue v = self.obj.gradient(arr) @@ -121,7 +121,7 @@ def test_gradient(self): def test_hessian(self): for arr,shp in zip(self.w0s, self.hessian_return_shapes): if self.E_unit.is_equivalent(u.one) and hasattr(arr, 'pos') and \ - not arr.pos.unit.is_equivalent(u.one): + not arr.pos.xyz.unit.is_equivalent(u.one): continue g = self.obj.hessian(arr) diff --git a/gala/potential/hamiltonian/tests/test_with_frame_potential.py b/gala/potential/hamiltonian/tests/test_with_frame_potential.py index c20267353..e0139386c 100644 --- a/gala/potential/hamiltonian/tests/test_with_frame_potential.py +++ b/gala/potential/hamiltonian/tests/test_with_frame_potential.py @@ -2,7 +2,9 @@ # Third-party import astropy.units as u +from astropy.tests.helper import quantity_allclose import pytest +import numpy as np # Project from .helpers import _TestBase @@ -10,15 +12,11 @@ from ...potential.builtin import NFWPotential, KeplerPotential, HernquistPotential from ...frame.builtin import StaticFrame, ConstantRotatingFrame from ....units import galactic, dimensionless -from ....dynamics import CartesianPhaseSpacePosition +from ....dynamics import PhaseSpacePosition, Orbit from ....integrate import DOPRI853Integrator # ---------------------------------------------------------------------------- -import astropy.units as u -import numpy as np -from gala.dynamics import PhaseSpacePosition, Orbit - def to_rotating_frame(omega, w, t=None): """ TODO: figure out units shit for omega and t @@ -56,12 +54,12 @@ def to_rotating_frame(omega, w, t=None): if isinstance(w, PhaseSpacePosition) or isinstance(w, Orbit): Cls = w.__class__ - x_shape = w.pos.shape - x_unit = w.pos.unit - v_unit = w.vel.unit + x_shape = w.pos.xyz.shape + x_unit = w.pos.x.unit + v_unit = w.vel.d_x.unit - x = w.pos.reshape(3,-1).value - v = w.vel.reshape(3,-1).value + x = w.pos.xyz.reshape(3,-1).value + v = w.vel.d_xyz.reshape(3,-1).value else: Cls = None @@ -122,15 +120,15 @@ def test_hessian(self): def test_integrate(self): - w0 = CartesianPhaseSpacePosition(pos=[1.,0,0.], vel=[0,1.,0.]) + w0 = PhaseSpacePosition(pos=[1.,0,0.], vel=[0,1.,0.]) for bl in [True, False]: orbit = self.obj.integrate_orbit(w0, dt=1., n_steps=1000, cython_if_possible=bl, Integrator=DOPRI853Integrator) - assert np.allclose(orbit.pos.value[0], 1., atol=1E-7) - assert np.allclose(orbit.pos.value[1:], 0., atol=1E-7) + assert np.allclose(orbit.pos.x.value, 1., atol=1E-7) + assert np.allclose(orbit.pos.xyz.value[1:], 0., atol=1E-7) class TestKepler2RotatingFrame(_TestBase): Omega = [1.,1.,1.]*u.one @@ -147,7 +145,7 @@ def test_integrate(self): # -------------------------------------------------------------- # when Omega is off from orbital frequency # - w0 = CartesianPhaseSpacePosition(pos=[1.,0,0.], vel=[0,1.1,0.]) + w0 = PhaseSpacePosition(pos=[1.,0,0.], vel=[0,1.1,0.]) for bl in [True, False]: orbit = self.obj.integrate_orbit(w0, dt=0.1, n_steps=10000, @@ -155,7 +153,7 @@ def test_integrate(self): Integrator=DOPRI853Integrator) L = orbit.angular_momentum() - C = orbit.energy()[:,0] - np.sum(self.Omega[:,None] * L, axis=0) + C = orbit.energy() - np.sum(self.Omega[:,None] * L, axis=0) dC = np.abs((C[1:]-C[0])/C[0]) assert np.all(dC < 1E-9) # conserve Jacobi constant @@ -171,8 +169,8 @@ def test_velocity_rot_frame(name, Omega, tol): r0 = 1.245246 potential = HernquistPotential(m=1., c=0.2, units=dimensionless) vc = potential.circular_velocity([r0,0,0]).value[0] - w0 = CartesianPhaseSpacePosition(pos=[r0, 0, 0.], - vel=[0, vc, 0.]) + w0 = PhaseSpacePosition(pos=[r0, 0, 0.], + vel=[0, vc, 0.]) Omega = Omega * [1., 1., vc/r0] H_r = Hamiltonian(potential, ConstantRotatingFrame(Omega=Omega, units=dimensionless)) @@ -184,15 +182,10 @@ def test_velocity_rot_frame(name, Omega, tol): orbit_i2r = orbit_i.to_frame(ConstantRotatingFrame(Omega=Omega, units=dimensionless)) orbit_r2i = orbit_r.to_frame(StaticFrame(units=dimensionless)) - dx = lambda x1,x2: np.sqrt(np.sum((x1-x2)**2, axis=0)) + assert quantity_allclose(orbit_i.pos.xyz, orbit_r2i.pos.xyz, atol=tol) + assert quantity_allclose(orbit_i.vel.d_xyz, orbit_r2i.vel.d_xyz, atol=tol) - # import matplotlib.pyplot as plt - # plt.plot(orbit_i.pos.value[0], orbit_i.pos.value[1]) - # plt.plot(orbit_r2i.pos.value[0], orbit_r2i.pos.value[1]) - # plt.show() + assert quantity_allclose(orbit_r.pos.xyz, orbit_i2r.pos.xyz, atol=tol) + assert quantity_allclose(orbit_r.vel.d_xyz, orbit_i2r.vel.d_xyz, atol=tol) - assert np.all(dx(orbit_i.pos.value, orbit_r2i.pos.value) < tol) - assert np.all(dx(orbit_i.vel.value, orbit_r2i.vel.value) < tol) - assert np.all(dx(orbit_r.pos.value, orbit_i2r.pos.value) < tol) - assert np.all(dx(orbit_r.vel.value, orbit_i2r.vel.value) < tol) diff --git a/gala/potential/potential/builtin/cybuiltin.pyx b/gala/potential/potential/builtin/cybuiltin.pyx index fb00fcaf2..5183c1f19 100644 --- a/gala/potential/potential/builtin/cybuiltin.pyx +++ b/gala/potential/potential/builtin/cybuiltin.pyx @@ -236,7 +236,7 @@ class IsochronePotential(CPotentialBase): Parameters ---------- - w : :class:`gala.dynamics.CartesianPhaseSpacePosition`, :class:`gala.dynamics.CartesianOrbit` + w : :class:`gala.dynamics.PhaseSpacePosition`, :class:`gala.dynamics.Orbit` The positions or orbit to compute the actions, angles, and frequencies at. """ from ....dynamics.analyticactionangle import isochrone_to_aa diff --git a/gala/potential/potential/builtin/pybuiltin.py b/gala/potential/potential/builtin/pybuiltin.py index 708798af9..3633de070 100644 --- a/gala/potential/potential/builtin/pybuiltin.py +++ b/gala/potential/potential/builtin/pybuiltin.py @@ -63,7 +63,7 @@ def action_angle(self, w): Parameters ---------- - w : :class:`gala.dynamics.CartesianPhaseSpacePosition`, :class:`gala.dynamics.CartesianOrbit` + w : :class:`gala.dynamics.PhaseSpacePosition`, :class:`gala.dynamics.Orbit` The positions or orbit to compute the actions, angles, and frequencies at. """ from ....dynamics.analyticactionangle import harmonic_oscillator_to_aa diff --git a/gala/potential/potential/core.py b/gala/potential/potential/core.py index 3e93b578c..328acb8b3 100644 --- a/gala/potential/potential/core.py +++ b/gala/potential/potential/core.py @@ -81,7 +81,7 @@ def _remove_units_prepare_shape(self, x): x = x.decompose(self.units).value elif isinstance(x, PhaseSpacePosition): - x = x.pos.decompose(self.units).value + x = x.cartesian.pos.xyz.decompose(self.units).value x = atleast_2d(x, insert_axis=1).astype(np.float64) return x @@ -581,7 +581,7 @@ def integrate_orbit(self, *args, **kwargs): Returns ------- - orbit : `~gala.dynamics.CartesianOrbit` + orbit : `~gala.dynamics.Orbit` """ diff --git a/gala/potential/potential/tests/helpers.py b/gala/potential/potential/tests/helpers.py index 967732258..a506a724a 100644 --- a/gala/potential/potential/tests/helpers.py +++ b/gala/potential/potential/tests/helpers.py @@ -17,7 +17,7 @@ from ..io import load from ..core import CompositePotential from ....units import UnitSystem, DimensionlessUnitSystem -from ....dynamics import CartesianPhaseSpacePosition +from ....dynamics import PhaseSpacePosition from ....integrate import LeapfrogIntegrator def partial_derivative(func, point, dim_ix=0, **kwargs): @@ -234,8 +234,8 @@ def test_orbit_integration(self): plt.close(f) us = self.potential.units - w0 = CartesianPhaseSpacePosition(pos=w0[:self.ndim]*us['length'], - vel=w0[self.ndim:]*us['length']/us['time']) + w0 = PhaseSpacePosition(pos=w0[:self.ndim]*us['length'], + vel=w0[self.ndim:]*us['length']/us['time']) orbit = self.potential.integrate_orbit(w0, dt=1., n_steps=10000, Integrator=LeapfrogIntegrator) diff --git a/gala/potential/potential/tests/test_composite.py b/gala/potential/potential/tests/test_composite.py index 01cf19a12..f22aae618 100644 --- a/gala/potential/potential/tests/test_composite.py +++ b/gala/potential/potential/tests/test_composite.py @@ -76,8 +76,8 @@ def test_integrate(self): w_py = potential.integrate_orbit([1.,0,0, 0,2*np.pi,0], dt=0.01, n_steps=1000, Integrator=Integrator, cython_if_possible=False) - for i in range(3): - np.testing.assert_allclose(w_cy.pos[i].value, w_cy.pos[i].value) + assert np.allclose(w_cy.pos.xyz.value, w_py.pos.xyz.value) + assert np.allclose(w_cy.vel.d_xyz.value, w_py.vel.d_xyz.value) # ------------------------------------------------------------------------ diff --git a/gala/tests/coveragerc b/gala/tests/coveragerc index 97d01018b..fead606c3 100644 --- a/gala/tests/coveragerc +++ b/gala/tests/coveragerc @@ -14,6 +14,7 @@ omit = {packagename}/dynamics/_genfunc/* {packagename}/coordinates/quaternion.py {packagename}/coordinates/poincarepolar.py + {packagename}/dynamics/extern/* [report] exclude_lines = diff --git a/gala/units.py b/gala/units.py index ff1c4608f..b688eeb82 100644 --- a/gala/units.py +++ b/gala/units.py @@ -8,6 +8,11 @@ from astropy.units.physical import _physical_unit_mapping import astropy.constants as const +_greek_letters = ["alpha", "beta", "gamma", "delta", "epsilon", "zeta", "eta", + "theta", "iota", "kappa", "lambda", "mu", "nu", "xi", "pi", + "o", "rho", "sigma", "tau", "upsilon", "phi", "chi", "psi", + "omega"] + class UnitSystem(object): """ Represents a system of units. At minimum, this consists of a set of