From c4ed857a9ac47e389ab4c1be5006354c56c8fed0 Mon Sep 17 00:00:00 2001 From: nstarman Date: Sun, 28 Jan 2024 02:24:46 -0500 Subject: [PATCH] WIP Signed-off-by: nstarman --- docs/potential/define-new-potential.rst | 60 +++++++++---------------- 1 file changed, 21 insertions(+), 39 deletions(-) diff --git a/docs/potential/define-new-potential.rst b/docs/potential/define-new-potential.rst index 4514f338..b000652f 100644 --- a/docs/potential/define-new-potential.rst +++ b/docs/potential/define-new-potential.rst @@ -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 `_! - 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 `_. @@ -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, @@ -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:: @@ -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") @@ -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).