[WIP] Implement Einasto and Core Einasto profiles

src/gala/potential/potential/builtin/builtin_potentials.h

@@ -145,6 +145,14 @@ extern double burkert_value(double t, double *pars, double *q, int n_dim, void *
 extern void burkert_gradient(double t, double *pars, double *q, int n_dim, size_t N, double *grad, void *state);
 extern double burkert_density(double t, double *pars, double *q, int n_dim, void *state);
 
+extern double einasto_value(double t, double *pars, double *q, int n_dim, void *state);
+extern void einasto_gradient(double t, double *pars, double *q, int n_dim, size_t N, double *grad, void *state);
+extern double einasto_density(double t, double *pars, double *q, int n_dim, void *state);
+
+extern double cEinasto_value(double t, double *pars, double *q, int n_dim, void *state);
+extern void cEinasto_gradient(double t, double *pars, double *q, int n_dim, size_t N, double *grad, void *state);
+extern double cEinasto_density(double t, double *pars, double *q, int n_dim, void *state);
+
 // Spherical spline interpolated potentials
 extern double spherical_spline_density_value(double t, double *pars, double *q, int n_dim, void *state);
 extern void spherical_spline_density_gradient(double t, double *pars, double *q, int n_dim, size_t N, double *grad, void *state);

src/gala/potential/potential/builtin/core.py

@@ -18,7 +18,9 @@
 from gala.potential.common import PotentialParameter
 from gala.potential.potential.builtin.cybuiltin import (
     BurkertWrapper,
+    CoreEinastoWrapper,
     CylSplineWrapper,
+    EinastoWrapper,
     FlattenedNFWWrapper,
     HenonHeilesWrapper,
     HernquistWrapper,
@@ -53,8 +55,10 @@
 
 __all__ = [
     "BurkertPotential",
+    "CoreEinastoPotential",
     "CylSplinePotential",
     "EXPPotential",
+    "EinastoPotential",
     "HenonHeilesPotential",
     "HernquistPotential",
     "IsochronePotential",
@@ -449,6 +453,55 @@ def from_r0(cls, r0, units=None):
         return cls(rho=rho_d0, r0=r0, units=units)
 
 
+@format_doc(common_doc=_potential_docstring)
+class EinastoPotential(CPotentialBase):
+    r"""The Einasto model.
+
+    Parameters
+    ----------
+    rho_m2 : :class:`~astropy.units.Quantity`, numeric [mass density]
+        Density at the radius where the logarithmic slope of the density profile is -2.
+    r_m2 : :class:`~astropy.units.Quantity`, numeric [length]
+        Radius where the logarithmic slope of the density profile is -2.
+    alpha :
+        Shape parameter.
+    {common_doc}
+    """
+
+    rho_m2 = PotentialParameter("rho_m2", physical_type="mass density")
+    r_m2 = PotentialParameter("r_m2", physical_type="length")
+    alpha = PotentialParameter("alpha", physical_type="dimensionless")
+
+    Wrapper = EinastoWrapper
+    _symmetry = SphericalSymmetry()
+
+
+@format_doc(common_doc=_potential_docstring)
+class CoreEinastoPotential(CPotentialBase):
+    r"""The core Einasto model.
+
+    Parameters
+    ----------
+    rho_s : :class:`~astropy.units.Quantity`, numeric [mass density]
+        Scale density (analogous to rho_-2 in the standard Einasto profile).
+    r_s : :class:`~astropy.units.Quantity`, numeric [length]
+        Scale radius (analogous to r_-2 in the standard Einasto profile).
+    alpha :
+        Shape parameter.
+    r_c : :class:`~astropy.units.Quantity`, numeric [length]
+        The core radius.
+    {common_doc}
+    """
+
+    rho_s = PotentialParameter("rho_s", physical_type="mass density")
+    r_s = PotentialParameter("r_s", physical_type="length")
+    alpha = PotentialParameter("alpha", physical_type="dimensionless")
+    r_c = PotentialParameter("r_c", physical_type="length")
+
+    Wrapper = CoreEinastoWrapper
+    _symmetry = SphericalSymmetry()
+
+
 # ============================================================================
 # Flattened, axisymmetric models
 #

src/gala/potential/potential/builtin/cybuiltin.pxd

@@ -100,3 +100,11 @@ cdef extern from "potential/potential/builtin/builtin_potentials.h":
     double burkert_value(double t, double *pars, double *q, int n_dim, void *state) nogil
     void burkert_gradient(double t, double *pars, double *q, int n_dim, size_t N, double *grad, void *state) nogil
     double burkert_density(double t, double *pars, double *q, int n_dim, void *state) nogil
+
+    double einasto_value(double t, double *pars, double *q, int n_dim, void *state) nogil
+    void einasto_gradient(double t, double *pars, double *q, int n_dim, size_t N, double *grad, void *state) nogil
+    double einasto_density(double t, double *pars, double *q, int n_dim, void *state) nogil
+
+    double cEinasto_value(double t, double *pars, double *q, int n_dim, void *state) nogil
+    void cEinasto_gradient(double t, double *pars, double *q, int n_dim, size_t N, double *grad, void *state) nogil
+    double cEinasto_density(double t, double *pars, double *q, int n_dim, void *state) nogil

src/gala/potential/potential/builtin/cybuiltin.pyx

@@ -223,6 +223,7 @@ cdef class PowerLawCutoffWrapper(CPotentialWrapper):
             self.cpotential.gradient[0] = <gradientfunc>(powerlawcutoff_gradient)
             self.cpotential.hessian[0] = <hessianfunc>(powerlawcutoff_hessian)
 
+
 cdef class BurkertWrapper(CPotentialWrapper):
 
     def __init__(self, G, parameters, q0, R):
@@ -234,6 +235,28 @@ cdef class BurkertWrapper(CPotentialWrapper):
         self.cpotential.gradient[0] = <gradientfunc>(burkert_gradient)
 
 
+cdef class EinastoWrapper(CPotentialWrapper):
+
+    def __init__(self, G, parameters, q0, R):
+        self.init([G] + list(parameters),
+                  np.ascontiguousarray(q0),
+                  np.ascontiguousarray(R))
+        self.cpotential.value[0] = <energyfunc>(einasto_value)
+        self.cpotential.density[0] = <densityfunc>(einasto_density)
+        self.cpotential.gradient[0] = <gradientfunc>(einasto_gradient)
+
+
+cdef class CoreEinastoWrapper(CPotentialWrapper):
+
+    def __init__(self, G, parameters, q0, R):
+        self.init([G] + list(parameters),
+                  np.ascontiguousarray(q0),
+                  np.ascontiguousarray(R))
+        self.cpotential.value[0] = <energyfunc>(cEinasto_value)
+        self.cpotential.density[0] = <densityfunc>(cEinasto_density)
+        self.cpotential.gradient[0] = <gradientfunc>(cEinasto_gradient)
+
+
 # ============================================================================
 # Flattened, axisymmetric models
 #

src/gala/potential/potential/builtin/potential_helpers.h

@@ -0,0 +1,35 @@
+#include <math.h>
+#include "src/vectorization.h"
+
+static inline double norm2_sq(const double *q) {
+    return q[0]*q[0] + q[1]*q[1];
+}
+
+static inline double norm2(const double *q) {
+    return sqrt(norm2_sq(q));
+}
+
+static inline double norm3_sq(const double *q) {
+    return q[0]*q[0] + q[1]*q[1] + q[2]*q[2];
+}
+
+static inline double norm3(const double *q) {
+    return sqrt(norm3_sq(q));
+}
+
+static inline double norm3_flat_z(const double *q, const double qz) {
+    return sqrt(q[0]*q[0] + q[1]*q[1] + q[2]*q[2]/(qz*qz));
+}
+
+// Helper functions for computing norms with double6ptr
+static inline double norm3_sq(const double6ptr& q) {
+    return (*q.x) * (*q.x) + (*q.y) * (*q.y) + (*q.z) * (*q.z);
+}
+
+static inline double norm3(const double6ptr& q) {
+    return sqrt(norm3_sq(q));
+}
+
+static inline double norm3_flat_z(const double6ptr& q, const double qz) {
+    return sqrt((*q.x) * (*q.x) + (*q.y) * (*q.y) + (*q.z) * (*q.z) / (qz * qz));
+}

tests/potential/potential/potential_helpers.py

@@ -6,11 +6,12 @@
 import matplotlib.pyplot as plt
 import numpy as np
 import pytest
+from astropy.constants import G
 
 from gala._compat_utils import SCIPY_LT_1_15
 from gala._optional_deps import HAS_SYMPY
 from gala.dynamics import PhaseSpacePosition
-from gala.potential import Hamiltonian, StaticFrame
+from gala.potential import Hamiltonian, SphericalSymmetry, StaticFrame
 from gala.potential.potential.io import load
 from gala.units import DimensionlessUnitSystem, UnitSystem
 
@@ -65,6 +66,9 @@ class PotentialTestBase:
     check_zero_at_infinity = True
     rotation = False
 
+    skip_hessian = False
+    skip_hessian_density = False
+
     def setup_method(self):
         # set up hamiltonian
         if self.frame is None:
@@ -159,6 +163,9 @@ def test_gradient(self):
             )
 
     def test_hessian(self):
+        if self.skip_hessian:
+            pytest.skip("Hessian not implemented for this potential")
+
         for arr, shp in zip(self.w0s, self._hess_return_shapes):
             g = self.potential.hessian(arr[: self.ndim])
             assert g.shape == shp
@@ -322,6 +329,80 @@ def compute_energy(xyz):
         ).T
         np.testing.assert_allclose(grad, num_grad, rtol=self.tol, atol=1e-8)
 
+    def test_numerical_density_vs_density(self):
+        """
+        Compare a numerically estimated Laplacian (trace of Hessian) times 4*pi*G
+        to the implemented density function via Poisson's equation:
+            Laplacian(Phi) = 4*pi*G * rho
+        """
+        if not isinstance(self.potential._symmetry, SphericalSymmetry):
+            pytest.skip("only for spherical potentials")
+
+        dx = 1e-3 * np.sqrt(np.sum(self.w0[: self.w0.size // 2] ** 2))
+        max_x = np.sqrt(np.sum([x**2 for x in self.w0[: self.w0.size // 2]]))
+
+        # use a slightly smaller grid to limit cost of nested derivatives
+        grid = np.linspace(-max_x, max_x, 6)
+        grid = grid[grid != 0.0]
+        grids = [grid for i in range(self.w0.size // 2)]
+        xyz = np.ascontiguousarray(
+            np.vstack(list(map(np.ravel, np.meshgrid(*grids)))).T
+        )
+
+        def compute_energy(xyz):
+            xyz = np.ascontiguousarray(np.atleast_2d(xyz))
+            return np.squeeze(self.potential._energy(xyz, t=np.array([0.0])))
+
+        def numeric_laplacian_at(pt):
+            # sum of second partial derivatives via nested numerical differentiation
+            s = 0.0
+            ndim = self.w0.size // 2
+            for dim in range(ndim):
+
+                def first_deriv(x):
+                    return partial_derivative(
+                        compute_energy,
+                        x,
+                        dim_ix=dim,  # noqa: B023
+                        dx=dx,
+                    )
+
+                s += partial_derivative(first_deriv, pt, dim_ix=dim, dx=dx)
+            return s
+
+        num_lap = np.zeros(xyz.shape[0])
+        for i in range(xyz.shape[0]):
+            num_lap[i] = numeric_laplacian_at(xyz[i])
+
+        num_rho = num_lap / (4.0 * np.pi * self.potential.G)
+
+        try:
+            dens = np.squeeze(
+                self.potential._density(np.ascontiguousarray(xyz), t=np.array([0.0]))
+            )
+        except Exception:
+            pytest.skip("density not implemented for this potential")
+
+        np.testing.assert_allclose(dens, num_rho, rtol=self.tol, atol=1e-6)
+
+    def test_hessian_density_consistency(self):
+        """
+        Check that the trace of the Hessian matches the density via Poisson's equation
+        """
+        if self.skip_hessian or self.skip_hessian_density:
+            pytest.skip("Hessian not implemented for this potential")
+
+        _G = G if not isinstance(self.potential.units, DimensionlessUnitSystem) else 1.0
+
+        for arr in self.w0s:
+            hess = self.potential.hessian(arr[: self.ndim])
+            lap = np.sum(np.diagonal(hess, axis1=0, axis2=1), axis=-1)
+
+            dens = self.potential.density(arr[: self.ndim])
+            rho_from_hess = lap / (4.0 * np.pi * _G)
+            print("YO", dens, rho_from_hess, np.diagonal(hess, axis1=0, axis2=1))
+            assert u.allclose(dens, rho_from_hess, rtol=self.tol)
+
     def test_orbit_integration(self, t1=0.0, t2=1000.0, nsteps=10000):
         """
         Make sure we can integrate an orbit in this potential