diff --git a/docs/conventions.rst b/docs/conventions.rst new file mode 100644 index 00000000..bac0ec86 --- /dev/null +++ b/docs/conventions.rst @@ -0,0 +1,45 @@ + +.. _conventions: + +*********** +Conventions +*********** + +.. _name-conventions: + +Common variable names +===================== + +Unless otherwise stated (in function or class docstrings), this package tries to +adhere to using standard variable names for function and class arguments, and in +example code. For shorthand, the variable ``w`` is used to represent arrays of +phase-space coordinates (e.g., positions _and_ velocities). When representing +only positions, the variable ``q`` is used. For just velocities or momenta, the +variables ``p`` or ``v`` are used. + +.. _shape-conventions: + +Array shapes +============ + +The arrays and :class:`~astropy.units.Quantity` objects expected as input and +returned as output throughout :mod:`galax` have shapes that follow a particular +convention, unless otherwise specified in function or class docstrings. + +For arrays containing coordinate or kinematic information, ``axis=-1`` is assumed +to be the coordinate dimension. For example, for representing 128 different 3D +Cartesian positions, the object should have shape ``(128, 3)``. + +For collections of orbits, arrays or array-valued objects have three axes: As +above, ``axis=2`` is assumed to be the coordinate dimension, ``axis=1`` is +interpreted as the time axis, and ``axis=0`` are the different orbits. + +.. _energy-momentum: + +Energy and momentum +=================== + +The :mod:`galax` documentation and functions often refer to energy and angular +momentum and the respective quantities *per unit mass* interchangeably. Unless +otherwise specified, all such quantities -- e.g., energy, angular momentum, +momenta, conjugate momenta -- are in fact used and returned *per unit mass*. diff --git a/docs/potential/compositepotential.rst b/docs/potential/compositepotential.rst new file mode 100644 index 00000000..ae94a3f1 --- /dev/null +++ b/docs/potential/compositepotential.rst @@ -0,0 +1,74 @@ +.. _galax-compositepotential: + +************************************************* +Creating a composite (multi-component ) potential +************************************************* + +Potential objects can be combined into more complex *composite* potentials using +:class:`~galax.potential.CompositePotential`. This class operates like a +Python dictionary in that each component potential must be named, and the +potentials can either be passed in to the initializer or added after the +composite potential container is already created. + +With either class, interaction with the class (e.g., by calling methods) is +identical to the individual potential classes. To compose potentials with unique +but arbitrary names, you can also simply add pre-defined potential class +instances:: + + >>> import jax.numpy as jnp + >>> import galax.potential as gp + >>> from galax.units import galactic + >>> disk = gp.MiyamotoNagaiPotential(m=1E11, a=6.5, b=0.27, units=galactic) + >>> bulge = gp.HernquistPotential(m=3E10, c=0.7, units=galactic) + >>> pot = disk + bulge + >>> print(pot.__class__.__name__) + CompositePotential + >>> list(pot.keys()) # doctest: +SKIP + ['c655f07d-a1fe-4905-bdb2-e8a202d15c81', + '8098cb0b-ebad-4388-b685-2f93a874296e'] + +The two components are assigned unique names and composed into a +:class:`~galax.potential.CompositePotential` instance. + +Alternatively, the potentials can be composed directly into the object by +treating it like an immutable dictionary. This allows you to specify the keys or +names of the components in the resulting +:class:`~galax.potential.CompositePotential` instance:: + + >>> disk = gp.MiyamotoNagaiPotential(m=1E11, a=6.5, b=0.27, units=galactic) + >>> bulge = gp.HernquistPotential(m=3E10, c=0.7, units=galactic) + >>> pot = gp.CompositePotential(disk=disk, bulge=bulge) + >>> list(pot.keys()) + ['disk', 'bulge'] + +The order of insertion is preserved, and sets the order that the potentials are +called. In the above example, the disk potential would always be called first +and the bulge would always be called second. + +The resulting potential object has all of the same properties as individual +potential objects:: + + >>> q = jnp.asarray([1., -1., 0.]) + >>> pot.potential_energy(q, t=0) + Array(-0.12887588, dtype=float64) + >>> pot.acceleration(q, t=0) + Array([-0.02270876, 0.02270876, -0. ], dtype=float64) + +.. >>> grid = jnp.linspace(-3., 3., 100) +.. >>> fig = pot.plot_contours(grid=(grid, 0, grid)) + +.. plot:: + :align: center + :width: 60% + + import jax.numpy as jnp + import gala.dynamics as gd + import galax.potential as gp + from galax.units import galactic + + disk = gp.MiyamotoNagaiPotential(m=1E11, a=6.5, b=0.27, units=galactic) + bulge = gp.HernquistPotential(m=3E10, c=0.7, units=galactic) + pot = gp.CompositePotential(disk=disk, bulge=bulge) + + grid = jnp.linspace(-3.,3.,100) + fig = pot.plot_contours(grid=(grid,0,grid)) diff --git a/docs/potential/define-new-potential.rst b/docs/potential/define-new-potential.rst new file mode 100644 index 00000000..b000652f --- /dev/null +++ b/docs/potential/define-new-potential.rst @@ -0,0 +1,167 @@ +.. _define-new-potential: + +********************************* +Defining your own potential class +********************************* + +Introduction +============ + +For the examples below the following imports have already been executed:: + + >>> import numpy as np + >>> import galax.potential as gp + >>> import galax.dynamics as gd + >>> from galax.typing import FloatOrIntScalarLike, FloatScalar + >>> from jaxtyping import Array, Float + +======================================== +Implementing a new potential with Python +======================================== + +New Python potentials are implemented by subclassing +:class:`~galax.potential.potential.PotentialBase` and defining functions that +compute (at minimum) the energy and gradient of the potential. We will work +through an example below for adding the `Henon-Heiles potential +`_. + +The expression for the potential is: + +.. math:: + + \Phi(x,y) = \frac{1}{2}(x^2 + y^2) + A\,(x^2 y - \frac{y^3}{3}) + +With this parametrization, there is only one free parameter (``A``), and the +potential is two-dimensional. + +At minimum, the subclass must implement the following methods: + +- ``__init__()`` +- ``_energy()`` +- ``_gradient()`` + +The ``_energy()`` method should compute the potential energy at a given position +and time. The ``_gradient()`` method should compute the gradient of the +potential. Both of these methods must accept two arguments: a position, and a +time. These internal methods are then called by the +:class:`~galax.potential.potential.PotentialBase` superclass methods +:meth:`~galax.potential.potential.PotentialBase.energy` and +:meth:`~galax.potential.potential.PotentialBase.gradient`. The superclass methods +convert the input position to an array in the unit system of the potential for +fast evaluation. The input to these superclass methods can be +:class:`~astropy.units.Quantity` objects, +:class:`~galax.dynamics.PhaseSpacePosition` objects, or :class:`~numpy.ndarray`. + +Because this potential has a parameter, the ``__init__`` method must accept +a parameter argument and store this in the ``parameters`` dictionary attribute +(a required attribute of any subclass). Let's write it out, then work through +what each piece means in detail:: + + >>> class CustomHenonHeilesPotential(gp.AbstractPotential): + ... A: gp.AbstractParameter = gp.ParameterField(dimensions="dimensionless") + ... + ... def _potential_energy(self, q: Float[Array, "3"], t: FloatOrIntScalarLike) -> FloatScalar: + ... A = self.A(t=t) + ... x, y = xy + ... return 0.5*(x**2 + y**2) + A*(x**2*y - y**3/3) + +The internal energy and gradient methods compute the numerical value and +gradient of the potential. The ``__init__`` method must take a single argument, +``A``, and store this to a parameter dictionary. The expected shape of the +position array (``xy``) passed to the internal ``_energy()`` and ``_gradient()`` +methods is always 2-dimensional with shape ``(n_points, n_dim)`` where +``n_points >= 1`` and ``n_dim`` must match the dimensionality of the potential +specified in the initializer. Note that this is different from the shape +expected when calling the public methods ``energy()`` and ``gradient()``! + +Let's now create an instance of the class and see how it works. For now, let's +pass in ``None`` for the unit system to designate that we'll work with +dimensionless quantities:: + + >>> pot = CustomHenonHeilesPotential(A=1., units=None) + +That's it! We now have a potential object with all of the same functionality as +the built-in potential classes. For example, we can integrate an orbit in this +potential (but note that this potential is two-dimensional, so we only have to +specify four coordinate values):: + + >>> w0 = gd.PhaseSpacePosition(q=[0., 0.3], p=[0.38, 0.]) + >>> t = jnp.arange(0, 500, step=0.05) + >>> orbit = pot.integrate_orbit(w0, t=t) + >>> fig = orbit.plot(marker=',', linestyle='none', alpha=0.5) # doctest: +SKIP + +.. plot:: + :align: center + :context: close-figs + :width: 60% + + import matplotlib.pyplot as pl + import numpy as np + import galax.dynamics as gd + import galax.potential as gp + + class CustomHenonHeilesPotential(gp.PotentialBase): + A = gp.PotentialParameter("A") + ndim = 2 + def _energy(self, xy, t): + A = self.parameters['A'].value + x,y = xy.T + return 0.5*(x**2 + y**2) + A*(x**2*y - y**3/3) + def _gradient(self, xy, t): + A = self.parameters['A'].value + x,y = xy.T + grad = np.zeros_like(xy) + grad[:,0] = x + 2*A*x*y + grad[:,1] = y + A*(x**2 - y**2) + return grad + + pot = CustomHenonHeilesPotential(A=1., units=None) + w0 = gd.PhaseSpacePosition(pos=[0.,0.3], + vel=[0.38,0.]) + orbit = gp.Hamiltonian(pot).integrate_orbit(w0, dt=0.05, n_steps=10000) + fig = orbit.plot(marker=',', linestyle='none', alpha=0.5) + +We could also, for example, create a contour plot of equipotentials:: + + >>> grid = np.linspace(-1., 1., 100) + >>> from matplotlib import colors + >>> import matplotlib.pyplot as plt + >>> fig, ax = plt.subplots(1, 1, figsize=(5,5)) + >>> fig = pot.plot_contours(grid=(grid, grid), + ... levels=np.logspace(-3, 1, 10), + ... norm=colors.LogNorm(), + ... cmap='Blues', ax=ax) + +.. plot:: + :align: center + :context: close-figs + :width: 60% + + from matplotlib import colors + import matplotlib.pyplot as plt + + grid = np.linspace(-1., 1., 100) + fig, ax = plt.subplots(1, 1, figsize=(5,5)) + fig = pot.plot_contours(grid=(grid,grid), cmap='Blues', + levels=np.logspace(-3, 1, 10), + norm=colors.LogNorm(), ax=ax) + +===================================== +Adding a custom potential with Cython +===================================== + +Adding a new Cython potential class is a little more involved as it requires +writing C-code and setting it up properly to compile when the code is built. +For this example, we'll work through how to define a new C-implemented potential +class representation of a Keplerian (point-mass) potential. Because this example +requires using Cython to build code, we provide a separate +`demo GitHub repository `_ with an +implementation of this potential with a demonstration of a build system that +successfully sets up the code. + +New Cython potentials are implemented by subclassing +:class:`~galax.potential.potential.CPotentialBase`, subclassing +:class:`~galax.potential.potential.CPotentialWrapper`, and defining C functions +that compute (at minimum) the energy and gradient of the potential. This +requires creating (at minimum) a Cython file (.pyx), a C header file (.h), and a +C source file (.c). diff --git a/docs/potential/index.rst b/docs/potential/index.rst new file mode 100644 index 00000000..82f984c2 --- /dev/null +++ b/docs/potential/index.rst @@ -0,0 +1,322 @@ +.. include:: ../references.txt + +.. module:: galax.potential + +******************************************** +Gravitational potentials (`galax.potential`) +******************************************** + +Introduction +============ + +This subpackage provides a number of classes for working with parametric models +of gravitational potentials. There are a number of built-in potentials +implemented in C and Cython (for speed), and there are base classes that allow +for easy creation of `new custom potential classes `_ +in pure Python or by writing custom C/Cython extensions. The ``Potential`` +objects have convenience methods for computing common dynamical quantities, for +example: potential energy, spatial gradient, density, or mass profiles. These +are particularly useful in combination with the :mod:`~galax.integrate` and +:mod:`~galax.dynamics` subpackages. + +.. TODO: uncomment +.. Also defined in this subpackage are a set of reference frames which can be used +.. for numerical integration of orbits in non-static reference frames. See the page +.. on :ref:`hamiltonian-reference-frames` for more information. ``Potential`` +.. objects can be combined with a reference frame and stored in a +.. :class:`~galax.potential.hamiltonian.Hamiltonian` object that provides an easy +.. interface to numerical orbit integration. + +For the examples below the following imports have already been executed:: + + >>> import astropy.units as u + >>> import matplotlib.pyplot as plt + >>> import jax.numpy as jnp + >>> import galax.potential as gp + >>> from galax.units import galactic, solarsystem, dimensionless + + +Getting Started: Built-in Methods of Potential Classes +====================================================== + +Any of the built-in ``Potential`` classes are initialized by passing in keyword +argument parameter values as :class:`~astropy.units.Quantity` objects or as +numeric values in a specified unit system. To see what parameters are available +for a given potential, check the documentation for the individual classes below. +You must also specify a :class:`~galax.units.UnitSystem` when initializing a +potential. A unit system is a set of non-reducible units that define (at +minimum) the length, mass, time, and angle units. A few common unit systems are +built in to the package (e.g., ``galactic``, ``solarsystem``, +``dimensionless``). For example, to create an object to represent a Kepler +potential (point mass) at the origin with mass = 1 solar mass, we would +instantiate a :class:`~gala.potential.KeplerPotential` object::: + + >>> ptmass = gp.KeplerPotential(m=1.*u.Msun, units=solarsystem) + >>> ptmass + KeplerPotential( + units=UnitSystem(AU, yr, solMass, rad), + m=ConstantParameter(unit=Unit("solMass"), value=f64[]) + ) + +If you pass in parameters with different units, they will be converted to the +specified unit system:: + + >>> gp.KeplerPotential(m=1047.6115*u.Mjup, units=solarsystem) + KeplerPotential( + units=UnitSystem(AU, yr, solMass, rad), + m=ConstantParameter(unit=Unit("solMass"), value=f64[]) + ) + +If no units are specified for a parameter (i.e. a parameter value is passed in +as a Python numeric value or array), it is assumed to be in the specified +:class:`~galax.units.UnitSystem`:: + + >>> gp.KeplerPotential(m=1., units=solarsystem) + KeplerPotential( + units=UnitSystem(AU, yr, solMass, rad), + m=ConstantParameter(unit=Unit("solMass"), value=f64[]) + ) + +The potential classes work well with the :mod:`astropy.units` framework, but to +ignore units you can use the :class:`~galax.units.DimensionlessUnitSystem` or +pass :obj:`~None` as the unit system:: + + >>> gp.KeplerPotential(m=1., units=None) + KeplerPotential( + units=DimensionlessUnitSystem(), + m=ConstantParameter(unit=Unit(dimensionless), value=f64[]) + ) + +All of the built-in potential objects have defined methods to evaluate the +potential energy and the gradient/acceleration at a given position or array of +positions. For example, to evaluate the potential energy at the 3D position +``(x, y, z) = (1, -1, 0) AU``:: + + >>> ptmass.potential_energy([1., -1., 0.] * u.au, t=0) + Array(-27.91440236, dtype=float64) + +These functions also accept both :class:`~astropy.units.Quantity` objects or +plain :class:`~numpy.ndarray`-like objects (in which case the position is +assumed to be in the unit system of the potential):: + + >>> ptmass.potential_energy(jnp.asarray([1., -1., 0.]), t=0) + Array(-27.91440236, dtype=float64) + +This also works for multiple positions by passing in a 2D position (but see +:ref:`conventions` for a description of the interpretation of different axes):: + + >>> pos = jnp.array([[1., -1. ,0], + ... [2., 3., 0]]) + >>> ptmass.potential_energy(pos * u.au, t=0) + Array([-27.91440236, -10.94892941], dtype=float64) + +We can also compute the gradient or acceleration:: + + >>> ptmass.gradient([1., -1., 0] * u.au, t=0) + Array([ 13.95720118, -13.95720118, 0. ], dtype=float64) + >>> ptmass.acceleration([1., -1., 0] * u.au, t=0) + Array([-13.95720118, 13.95720118, -0. ], dtype=float64) + +Most of the potential objects also have methods implemented for computing the +corresponding mass density and the Hessian of the potential (the matrix of 2nd +derivatives) at given locations. For example, with the +:class:`~galax.potential.HernquistPotential`, we can evaluate both the mass +density and Hessian at the position ``(x, y, z) = (1, -1, 0) kpc``:: + + >>> pot = gp.HernquistPotential(m=1E9*u.Msun, c=1.*u.kpc, units=galactic) + >>> pot.density([1., -1., 0] * u.kpc, t=0) + Array(7997938.82200887, dtype=float64) + >>> pot.hessian([1., -1., 0] * u.kpc, t=0) + Array([[-4.68187913e-05, 5.92578618e-04, 0.00000000e+00], + [ 5.92578618e-04, -4.68187913e-05, 0.00000000e+00], + [ 0.00000000e+00, 0.00000000e+00, 5.45759827e-04]], dtype=float64) + +.. Another useful method is :meth:`~galax.potential.PotentialBase.mass_enclosed`, +.. which numerically estimates the mass enclosed within a spherical shell defined +.. by the specified position. This numerically estimates :math:`\frac{d \Phi}{d r}` +.. along the vector pointing at the specified position and estimates the enclosed +.. mass simply as :math:`M(>> pot = gp.NFWPotential(m=1E11*u.Msun, r_s=20.*u.kpc, units=galactic) +.. >>> pos = jnp.zeros((3,100)) * u.kpc +.. >>> pos[0] = jnp.logspace(jnp.log10(20./100.), jnp.log10(20*100.), pos.shape[1]) * u.kpc +.. >>> m_profile = pot.mass_enclosed(pos) +.. >>> plt.loglog(pos[0], m_profile, marker='') +.. >>> plt.xlabel("$r$ [{}]".format(pos.unit.to_string(format='latex'))) +.. >>> plt.ylabel("$M(>> p = gp.MiyamotoNagaiPotential(m=1E11, a=6.5, b=0.27, units=galactic) +.. >>> fig, ax = plt.subplots() +.. >>> plot_contours(p, grid=(jnp.linspace(-15,15,100), 0., 1.), marker='', ax=ax) +.. >>> E_unit = p.units['energy'] / p.units['mass'] +.. >>> ax.set_xlabel("$x$ [{}]".format(p.units['length'].to_string(format='latex'))) +.. >>> ax.set_ylabel("$\Phi(x,0,1)$ [{}]".format(E_unit.to_string(format='latex'))) + +.. .. plot:: +.. :align: center +.. :context: close-figs +.. :width: 90% + +.. pot = gp.MiyamotoNagaiPotential(m=1E11, a=6.5, b=0.27, units=galactic) + +.. fig, ax = plt.subplots(1,1) +.. pot.plot_contours(grid=(jnp.linspace(-15,15,100), 0., 1.), marker='', ax=ax) +.. E_unit = pot.units['energy'] / pot.units['mass'] +.. ax.set_xlabel("$x$ [{}]".format(pot.units['length'].to_string(format='latex'))) +.. ax.set_ylabel("$\Phi(x,0,1)$ [{}]".format(E_unit.to_string(format='latex'))) +.. fig.tight_layout() + +.. To instead make a 2D contour plot over :math:`x` and :math:`z` along with +.. :math:`y=0`, pass in a 1D grid of values for :math:`x` and a 1D grid of values +.. for :math:`z` (the meshgridding is taken care of internally). Here, we choose to +.. draw on a pre-defined ``matplotlib`` axes object so we can set the labels and +.. aspect ratio of the plot:: + +.. >>> fig,ax = plt.subplots(1, 1, figsize=(12, 4)) +.. >>> x = jnp.linspace(-15, 15, 100) +.. >>> z = jnp.linspace(-5, 5, 100) +.. >>> p.plot_contours(grid=(x, 1., z), ax=ax) +.. >>> ax.set_xlabel("$x$ [kpc]") +.. >>> ax.set_ylabel("$z$ [kpc]") + +.. .. plot:: +.. :align: center +.. :context: close-figs +.. :width: 60% + +.. x = jnp.linspace(-15,15,100) +.. z = jnp.linspace(-5,5,100) + +.. fig,ax = plt.subplots(1, 1, figsize=(12,4)) +.. pot.plot_contours(grid=(x, 1., z), ax=ax) +.. ax.set_xlabel("$x$ [{}]".format(pot.units['length'].to_string(format='latex'))) +.. ax.set_ylabel("$z$ [{}]".format(pot.units['length'].to_string(format='latex'))) +.. fig.tight_layout() + +.. Saving / loading potential objects +.. ================================== + +.. Potential objects can be `pickled +.. `_ and can therefore be stored +.. for later use. However, pickles are saved as binary files. It may be useful to +.. save to or load from text-based specifications of Potential objects. This can be +.. done with the :meth:`~galax.potential.PotentialBase.save` method and the +.. :func:`~galax.potential.load` function, which serialize and de-serialize +.. (respectively) a ``Potential`` object to a `YAML `_ file:: + +.. >>> from gala.potential import load +.. >>> pot = gp.NFWPotential(m=6E11*u.Msun, r_s=20.*u.kpc, +.. ... units=galactic) +.. >>> pot.save("potential.yml") +.. >>> load("potential.yml") +.. + +.. Exporting potentials as ``sympy`` expressions +.. ============================================= + +.. Most of the potential classes can be exported to a :mod:`~sympy` expression that can +.. be used to manipulate or evaluate the form of the potential. To access this +.. functionality, the potential classes have a +.. :meth:`~galax.potential.PotentialBase.to_sympy` classmethod (note: this +.. requires :mod:`sympy` to be installed): + +.. .. doctest-requires:: sympy + +.. >>> expr, vars_, pars = gp.LogarithmicPotential.to_sympy() +.. >>> str(expr) +.. '0.5*v_c**2*log(r_h**2 + z**2/q3**2 + y**2/q2**2 + x**2/q1**2)' + +.. This method also returns a dictionary containing the coordinate variables used +.. in the expression as ``sympy`` symbols, here defined as ``vars_``: + +.. .. doctest-requires:: sympy + +.. >>> vars_ +.. {'x': x, 'y': y, 'z': z} + +.. A second dictionary containing the potential parameters as :mod:`sympy` symbols is +.. also returned, here defined as ``pars``: + +.. .. doctest-requires:: sympy + +.. >>> pars +.. {'v_c': v_c, 'r_h': r_h, 'q1': q1, 'q2': q2, 'q3': q3, 'phi': phi, 'G': G} + +.. The expressions and variables returned can be used to perform operations on the +.. potential expression. For example, to create a :mod:`sympy` expression for the +.. gradient of the potential: + +.. .. doctest-requires:: sympy + +.. >>> import sympy as sy +.. >>> grad = sy.derive_by_array(expr, list(vars_.values())) +.. >>> grad[0] # dPhi/dx +.. 1.0*v_c**2*x/(q1**2*(r_h**2 + z**2/q3**2 + y**2/q2**2 + x**2/q1**2)) + + +Using gala.potential +==================== +More details are provided in the linked pages below: + +.. toctree:: + :maxdepth: 1 + + define-new-potential + compositepotential + +.. .. toctree:: +.. :maxdepth: 1 + +.. define-new-potential +.. compositepotential +.. origin-rotation +.. hamiltonian-reference-frames +.. scf + + +API +=== + +.. automodapi:: galax.potential + +.. TODO: uncomment +.. .. automodapi:: gala.potential.frame.builtin + +.. TODO: uncomment +.. .. automodapi:: gala.potential.hamiltonian diff --git a/docs/user_guide.rst b/docs/user_guide.rst index 0dbe2b4e..8029ea65 100644 --- a/docs/user_guide.rst +++ b/docs/user_guide.rst @@ -7,15 +7,17 @@ User Guide ********** The user guide contains exhaustive descriptions of all of the functions and -classes available in :mod:`gala`, with some inline narrative descriptions and +classes available in :mod:`galax`, with some inline narrative descriptions and demonstrations of functionality. This portion of the documentation should be -treated like reference material, whereas the :ref:`gala-tutorials` show how -pieces of :mod:`gala` work together to meet more realistic research needs. +treated like reference material, whereas the :ref:`galax-tutorials` show how +pieces of :mod:`galax` work together to meet more realistic research needs. -.. TODO -.. .. toctree:: -.. :maxdepth: 1 +.. toctree:: + :maxdepth: 1 + + potential/index +.. TODO .. conventions .. coordinates/index .. integrate/index