.. include:: ../references.txt

.. module:: gala.potential

*************************************************
Gravitational potentials (`gala.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 <define-new-potential.html>`_
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 `~gala.integrate` and
`~gala.dynamics` subpackages.

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
`~gala.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 numpy as np
    >>> import gala.potential as gp
    >>> from gala.units import galactic, solarsystem, dimensionless

Getting Started: Built-in Methods of Potential Classes
======================================================

Potential classes are initialized by passing parameter values as
:class:`~astropy.units.Quantity` objects or as numeric values with a
specified unit system. You must also specify a `~gala.units.UnitSystem`
— a set of non-reducible units defining length, mass, time, and angle units.
Common unit systems are built in (e.g., ``galactic``, ``solarsystem``,
``dimensionless``).

For example, to create a Kepler potential (point mass) with mass = 1 solar mass::

    >>> ptmass = gp.KeplerPotential(m=1.0 * u.Msun, units=solarsystem)
    >>> ptmass
    <KeplerPotential: m=1.00 (AU,yr,solMass,rad)>

Parameters with different units are automatically converted to the specified
unit system::

    >>> gp.KeplerPotential(m=1047.6115 * u.Mjup, units=solarsystem)
    <KeplerPotential: m=1.00 (AU,yr,solMass,rad)>

Parameters without units are assumed to be in the specified unit system::

    >>> gp.KeplerPotential(m=1.0, units=solarsystem)
    <KeplerPotential: m=1.00 (AU,yr,solMass,rad)>

To work without units, use `~gala.units.DimensionlessUnitSystem` or pass
``None`` (this is the default, so you can also omit the units argument)::

    >>> gp.KeplerPotential(m=1.0, units=None)
    <KeplerPotential: m=1.00 (dimensionless)>

Unit systems can also be specified by passing a string name::

    >>> gp.KeplerPotential(m=1.0 * u.Msun, units='solarsystem')
    <KeplerPotential: m=1.00 (AU,yr,solMass,rad)>

All built-in potential objects have methods to evaluate the potential energy
and gradient/acceleration at given positions. For example, to evaluate the
potential energy at ``(x, y, z) = (1, -1, 0) AU``::

    >>> ptmass.energy([1.0, -1.0, 0.0] * u.au)
    <Quantity [-27.91440236] AU2 / yr2>

These functions accept both :class:`~astropy.units.Quantity` objects and
plain array-like objects (assumed to be in the potential's unit system)::

    >>> ptmass.energy([1.0, -1.0, 0.0])
    <Quantity [-27.91440236] AU2 / yr2>

For multiple positions, pass a 2D array where each column is a position::

    >>> pos = np.array([[1.0, -1.0, 0], [2.0, 3.0, 0]]).T
    >>> ptmass.energy(pos * u.au)
    <Quantity [-27.91440236, -10.94892941] AU2 / yr2>

We can also compute the gradient or acceleration::

    >>> ptmass.gradient([1.0, -1.0, 0] * u.au)  # doctest: +FLOAT_CMP
    <Quantity [[ 13.95720118],
               [-13.95720118],
               [  0.        ]] AU / yr2>
    >>> ptmass.acceleration([1.0, -1.0, 0] * u.au)  # doctest: +FLOAT_CMP
    <Quantity [[-13.95720118],
               [ 13.95720118],
               [ -0.        ]] AU / yr2>

Using Symmetry Coordinates
--------------------------

Many potentials have symmetries that make certain coordinate systems more natural or
concise to work with. For example, spherically-symmetric potentials only depend on the
radius :math:`r`, and axisymmetric potentials only depend on the cylindrical radius
:math:`R` and height :math:`z`.

For potentials with these symmetries, you can use **symmetry coordinates** as a
shorthand instead of full 3D Cartesian coordinates. Internally, gala operates in
Cartesian coordinates, so the symmetry coordinates are simply a front-end convenience.

Spherical Potentials
~~~~~~~~~~~~~~~~~~~~~

For spherically-symmetric potentials (like :class:`~gala.potential.HernquistPotential`,
:class:`~gala.potential.PlummerPotential`, :class:`~gala.potential.KeplerPotential`,
etc.), you can pass just the radius using ``r=``::

    >>> pot = gp.HernquistPotential(m=1e10 * u.Msun, c=5 * u.kpc, units=galactic)
    >>> r = np.array([1.0, 5.0, 10.0]) * u.kpc
    >>> pot.energy(r=r)  # doctest: +FLOAT_CMP
    <Quantity [-0.06674208, -0.03367003, -0.01792115] kpc2 / Myr2>

This is equivalent to passing Cartesian coordinates ``[r, 0, 0]``, but much cleaner. All
potential methods support symmetry coordinates::

    >>> pot.gradient(r=r)  # Note: Still returns gradient in Cartesian  # doctest: +SKIP
    >>> pot.density(r=r)  # doctest: +SKIP
    >>> pot.mass_enclosed(r=r)  # doctest: +SKIP
    >>> pot.circular_velocity(r=r)  # doctest: +SKIP

Note that gradients and accelerations are always returned in Cartesian coordinates, even
when using symmetry inputs. This ensures consistency and compatibility with orbit
integration.

Cylindrical (Axisymmetric) Potentials
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

For axisymmetric potentials (like :class:`~gala.potential.MiyamotoNagaiPotential`),
you can use cylindrical coordinates with ``R=`` and ``z=``::

    >>> pot = gp.MiyamotoNagaiPotential(m=1e11 * u.Msun, a=3 * u.kpc, b=0.3 * u.kpc, units=galactic)
    >>> R = np.linspace(4, 12, 100) * u.kpc
    >>> z = np.linspace(-1, 1, 100) * u.kpc
    >>> pot.energy(R=R, z=z)  # doctest: +SKIP

The ``z`` coordinate defaults to zero, making midplane calculations particularly
convenient::

    >>> pot.energy(R=R)  # Evaluate in the midplane (z=0)  # doctest: +SKIP

This is especially useful for plotting rotation curves or computing circular
velocities in the disk plane.

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:`~gala.potential.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.0 * u.kpc, units=galactic)
    >>> pot.density([1.0, -1.0, 0] * u.kpc)  # doctest: +FLOAT_CMP
    <Quantity [7997938.82200887] solMass / kpc3>
    >>> pot.hessian([1.0, -1.0, 0] * u.kpc)  # doctest: +SKIP
    <Quantity [[[ -4.68318131e-05],
                [  5.92743432e-04],
                [  0.00000000e+00]],

               [[  5.92743432e-04],
                [ -4.68318131e-05],
                [  0.00000000e+00]],

               [[  0.00000000e+00],
                [  0.00000000e+00],
                [  5.45911619e-04]]] 1 / Myr2>

Another useful method is
:meth:`~gala.potential.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(<r)\approx\frac{r^2}{G} \frac{d \Phi}{d r}`. This function can be used
to compute, for example, a mass profile. For spherical potentials, you can use
the ``r=`` symmetry coordinate::

    >>> pot = gp.NFWPotential(m=1e11 * u.Msun, r_s=20.0 * u.kpc, units=galactic)
    >>> r = np.logspace(np.log10(20.0 / 100.0), np.log10(20 * 100.0), 100) * u.kpc
    >>> m_profile = pot.mass_enclosed(r=r)
    >>> plt.loglog(r, m_profile, marker="")  # doctest: +SKIP
    >>> plt.xlabel("$r$ [{}]".format(r.unit.to_string(format="latex")))  # doctest: +SKIP
    >>> plt.ylabel("$M(<r)$ [{}]".format(m_profile.unit.to_string(format="latex")))  # doctest: +SKIP

.. plot::
    :align: center
    :context: close-figs
    :width: 60%

    import astropy.units as u
    import numpy as np
    import gala.potential as gp
    from gala.units import galactic, solarsystem
    import matplotlib.pyplot as plt

    pot = gp.NFWPotential(m=1e11 * u.Msun, r_s=20.0 * u.kpc, units=galactic)
    r = np.logspace(np.log10(20.0 / 100.0), np.log10(20 * 100.0), 100) * u.kpc
    m_profile = pot.mass_enclosed(r=r)

    plt.figure()
    plt.loglog(r, m_profile, marker="")  # doctest: +SKIP
    plt.xlabel("$r$ [{}]".format(r.unit.to_string(format="latex")))
    plt.ylabel("$M(<r)$ [{}]".format(m_profile.unit.to_string(format="latex")))
    plt.tight_layout()

Plotting Equipotential and Isodensity contours
==============================================

Potential objects provide methods for visualizing isopotential
(:meth:`~gala.potential.potential.PotentialBase.plot_contours`) and isodensity
(:meth:`~gala.potential.potential.PotentialBase.plot_density_contours`) contours.
These methods create either 1D slices or 2D contour plots.

For a 1D slice, pass a grid for one dimension and fixed values for others.
For example, to plot potential vs. :math:`x` at :math:`y=0, z=1`::

    >>> p = gp.MiyamotoNagaiPotential(m=1e11, a=6.5, b=0.27, units=galactic)
    >>> fig, ax = plt.subplots()  # doctest: +SKIP
    >>> p.plot_contours(
    ...     grid=(np.linspace(-15, 15, 100), 0.0, 1.0), marker="", ax=ax
    ... )  # doctest: +SKIP

.. 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)  # doctest: +SKIP
    pot.plot_contours(
        grid=(np.linspace(-15, 15, 100), 0.0, 1.0), marker="", ax=ax
    )  # doctest: +SKIP
    E_unit = pot.units["energy"] / pot.units["mass"]
    ax.set_xlabel("$x$ [{}]".format(pot.units["length"].to_string(format="latex")))  # doctest: +SKIP
    ax.set_ylabel("$\Phi(x,0,1)$ [{}]".format(E_unit.to_string(format="latex")))  # doctest: +SKIP
    fig.tight_layout()

For 2D contour plots, pass 1D grids for two dimensions and a fixed value
for the third. For example, to plot :math:`x`-:math:`z` contours at :math:`y=0`::

    >>> fig, ax = plt.subplots(1, 1, figsize=(12, 4))
    >>> x = np.linspace(-15, 15, 100)
    >>> z = np.linspace(-5, 5, 100)
    >>> p.plot_contours(grid=(x, 1.0, z), ax=ax)  # doctest: +SKIP

.. plot::
    :align: center
    :context: close-figs
    :width: 60%

    x = np.linspace(-15, 15, 100)
    z = np.linspace(-5, 5, 100)

    fig, ax = plt.subplots(1, 1, figsize=(12, 4))
    pot.plot_contours(grid=(x, 1.0, 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 saved to and loaded from YAML files using the
:meth:`~gala.potential.potential.PotentialBase.save` method and
:func:`~gala.potential.potential.load` function::

    >>> from gala.potential import load
    >>> pot = gp.NFWPotential(m=6e11 * u.Msun, r_s=20.0 * u.kpc, units=galactic)
    >>> pot.save("potential.yml")
    >>> load("potential.yml")
    <NFWPotential: m=6.00e+11, r_s=20.00, a=1.00, b=1.00, c=1.00 (kpc,Myr,solMass,rad)>

Exporting potentials as ``sympy`` expressions
=============================================

Most of the potential classes can be exported to a `sympy` expression that can
be used to manipulate or evaluate the form of the potential. To access this
functionality, the potential classes have a
`~gala.potential.potential.PotentialBase.to_sympy` classmethod (note: this
requires `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 `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 `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))

Time-varying Potentials
=======================

For modeling time-dependent gravitational potentials, Gala provides
:class:`~gala.potential.potential.TimeInterpolatedPotential`.
This class can wrap any potential class and interpolate its parameters, origin, and/or
rotation over time.
This is useful for modeling, for example, time-evolving masses, such as from mass-loss
in a mock stream simulation, or for rotating components like galactic bars.

The class uses GSL spline interpolation to smoothly interpolate parameter values between
specified time knots. For example, to model a black hole that grows in mass over time::

    >>> import astropy.units as u
    >>> import numpy as np
    >>> times = np.linspace(0, 1, 11) * u.Gyr
    >>> masses = np.linspace(1e6, 1e8, 11) * u.Msun  # Grows linearly by a factor of 100
    >>> pot = gp.TimeInterpolatedPotential(
    ...     gp.KeplerPotential,
    ...     time_knots=times,
    ...     m=masses,
    ...     units=galactic
    ... )
    >>> pot.energy([10., 0, 0] * u.pc, t=0 * u.Gyr)  # doctest: +SKIP
    <Quantity [-0.0044985] kpc2 / Myr2>
    >>> pot.energy([10., 0, 0] * u.pc, t=1 * u.Gyr)  # doctest: +SKIP
    <Quantity [-0.44985021] kpc2 / Myr2>

Parameters can be constant (passed as scalars) or time-varying (passed as arrays
matching the number of time knots). You can also specify time-varying origin offsets and
rotation matrices. See :doc:`time-interpolated` for more details and examples.

Using gala.potential
====================
More details are provided in the linked pages below:

.. toctree::
   :maxdepth: 1

   symmetry-coordinates
   define-new-potential
   compositepotential
   origin-rotation
   time-interpolated
   hamiltonian-reference-frames
   spherical-spline
   scf

API
===

.. automodapi:: gala.potential.potential

.. automodapi:: gala.potential.frame.builtin

.. automodapi:: gala.potential.hamiltonian