Skip to content

Commit

Permalink
WIP
Browse files Browse the repository at this point in the history 
Signed-off-by: nstarman <[email protected]>
  • Loading branch information
@nstarman
nstarman committed Jan 28, 2024
1 parent ed4c377 commit c4ed857
Showing 1 changed file with 21 additions and 39 deletions.
60 changes: 21 additions & 39 deletions docs/potential/define-new-potential.rst
View file
Original file line number Diff line number Diff line change
Expand Up @@ -7,28 +7,20 @@ Defining your own potential class
Introduction
============

There are two ways to define a new potential class: with pure-Python, or with C
and Cython. The advantage to writing a new class in Cython is that the
computations can execute with C-like speeds, however only certain integrators
support using this functionality (Leapfrog and DOP853) and it is a bit more
complicated to set up the code to build the C+Cython code properly. If you are
not familiar with Cython, you probably want to stick to a pure Python class for
initial testing. If there is a potential class that you think should be
included as a built-in Cython potential, feel free to suggest the new addition
as a `GitHub issue <https://github.com/adrn/gala/issues>`_!

For the examples below the following imports have already been executed::

>>> import numpy as np
>>> import gala.potential as gp
>>> import gala.dynamics as gd
>>> 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:`~gala.potential.potential.PotentialBase` and defining functions that
: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
<http://en.wikipedia.org/wiki/H%C3%A9non-Heiles_System>`_.
Expand All @@ -52,36 +44,26 @@ 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:`~gala.potential.potential.PotentialBase` superclass methods
:meth:`~gala.potential.potential.PotentialBase.energy` and
:meth:`~gala.potential.potential.PotentialBase.gradient`. The superclass methods
: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:`~gala.dynamics.PhaseSpacePosition` objects, or :class:`~numpy.ndarray`.
: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.PotentialBase):
... A = gp.PotentialParameter("A")
... ndim = 2
>>> class CustomHenonHeilesPotential(gp.AbstractPotential):
... A: gp.AbstractParameter = gp.ParameterField(dimensions="dimensionless")
...
... def _energy(self, xy, t):
... A = self.parameters['A'].value
... x,y = xy.T
... 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)
...
... 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

The internal energy and gradient methods compute the numerical value and
gradient of the potential. The ``__init__`` method must take a single argument,
Expand All @@ -103,9 +85,9 @@ 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(pos=[0., 0.3],
... vel=[0.38, 0.])
>>> orbit = gp.Hamiltonian(pot).integrate_orbit(w0, dt=0.05, n_steps=10000)
>>> 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::
Expand All @@ -115,8 +97,8 @@ specify four coordinate values)::

import matplotlib.pyplot as pl
import numpy as np
import gala.dynamics as gd
import gala.potential as gp
import galax.dynamics as gd
import galax.potential as gp

class CustomHenonHeilesPotential(gp.PotentialBase):
A = gp.PotentialParameter("A")
Expand Down Expand Up @@ -178,8 +160,8 @@ 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:`~gala.potential.potential.CPotentialBase`, subclassing
:class:`~gala.potential.potential.CPotentialWrapper`, and defining C functions
: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).

0 comments on commit c4ed857

Please sign in to comment.