Set version 1.0.0

exptool/models/__init__.py

@@ -1,4 +1,2 @@
 
-
-# this is EXPtool
-
+__all__ = ['mndisc','logpot','hernquist','nfw','plummer']

exptool/models/arbitrarypotential.py

@@ -0,0 +1,19 @@
+"""
+Provides functions to compute the density from an arbitrary potential using JAX.
+"""
+
+import jax
+import jax.numpy as jnp
+
+def calculate_density(potential, gravitational_constant):
+    # Define a function to compute the Laplacian of the potential
+    def laplacian(potential):
+        return jnp.sum(jax.hessian(jnp.sum)(potential))
+
+    # Calculate the Laplacian of the potential
+    laplacian_potential = jax.vmap(laplacian)(potential)
+
+    # Calculate the density using Poisson's equation
+    density = laplacian_potential / (4 * jnp.pi * gravitational_constant)
+
+    return density

exptool/models/hernquist.py

@@ -8,6 +8,7 @@
 29 Oct 2017  First construction
 11 Jul 2021  Much more usable (read:correct) now!
 25 Jan 2022  Revise, add force, comment
+24 Sep 2023  New docstrings
 
 
 example
@@ -26,85 +27,112 @@
 # general python imports
 import numpy as np
 
-
 class Hernquist():
-    """The Hernquist (1990) model
+    '''
+    The Hernquist (1990) model
 
+    The Hernquist model represents a spherically symmetric mass distribution characterized by a scale radius (rscl) and an overall mass (M).
 
-    notes
-    ----------
-    The Hernquist profile has a gravitational radii of 6rscl.
-    """
-    def __init__(self,rscl=1.0,G=1,M=1):
-        """
+    Attributes:
+        rscl (float): The scale radius of the Hernquist model.
+        G (float): The gravitational constant.
+        M (float): The overall mass of the model.
+        rho0 (float): The central density of the Hernquist model.
+
+    Methods:
+        get_rho0():
+            Calculate the central density of the Hernquist model.
+
+        get_mass(r):
+            Calculate the mass enclosed within a given radius (r).
+
+        get_dens(r):
+            Calculate the density at a given radius (r).
+
+        get_pot(r):
+            Calculate the gravitational potential at a given radius (r).
+
+        get_dphi_dr(r):
+            Calculate the radial force at a given radius (r).
 
-        inputs
-        ----------------
-        rscl  : the scale radius, in units of length
-        G     : the gravitational constant, in units of (length^3) * (mass) / (time^2)
-                 (an astronomical value is 0.0000043009125, in the equivalent (km/s)^2 * kpc / Msun)
-        M     : the overall mass of the model, in units of mass
+    Example usage:
+        H = Hernquist(rscl=1.0, G=0.0000043009125, M=1.0)
+        print(H.get_mass(2.0))
+    '''
 
+    def __init__(self, rscl=1.0, G=1, M=1):
         """
+        Initialize the Hernquist model.
+
+        Args:
+            rscl (float, optional): The scale radius (rscl) of the model. Defaults to 1.0.
+            G (float, optional): The gravitational constant. Defaults to 1.
+            M (float, optional): The overall mass of the model. Defaults to 1.
+
+        Attributes:
+            rscl (float): The scale radius of the Hernquist model.
+            G (float): The gravitational constant.
+            M (float): The overall mass of the model.
+            rho0 (float): The central density of the Hernquist model.
+        """
+
         self.rscl = rscl
-        self.G    = G
-        self.M    = M
+        self.G = G
+        self.M = M
         self.rho0 = self.get_rho0()
 
     def get_rho0(self):
-        """solving BT08, eq. 2.66 at r==1000a
-
-        this is a reasonable approximation because Hernquist is finite in mass.
-        exercise to the coder: what fraction of the mass are we excluding?
+        """Calculate the central density of the Hernquist model.
 
-        returned in units of density, mass/(length^3)
+        Returns:
+            float: The central density in units of mass / (length^3).
         """
-        rs = 1000.
-        return self.M*(2*(1+rs)*(1+rs))/(rs*rs)/(4*np.pi*np.power(self.rscl,3))
+        rs = 1000.  # Evaluate at r = 1000 * rscl (reasonable approximation)
+        return self.M * (2 * (1 + rs) * (1 + rs)) / (rs * rs) / (4 * np.pi * np.power(self.rscl, 3))
 
-    def get_mass(self,r):
-        """Hernquist mass enclosed: BT08, eq. 2.66
+    def get_mass(self, r):
+        """Calculate the mass enclosed within a given radius (r).
 
-        inputs
-        ---------------
-        r    : radius to compute enclosed mass for
+        Args:
+            r (float): The radius at which to compute the enclosed mass.
+
+        Returns:
+            float: The enclosed mass in units of mass.
         """
-        rs = r/self.rscl
-        return 4*np.pi*self.rho0*np.power(self.rscl,3)*(rs*rs)/(2*(1+rs)*(1+rs))
+        rs = r / self.rscl
+        return 4 * np.pi * self.rho0 * np.power(self.rscl, 3) * (rs * rs) / (2 * (1 + rs) * (1 + rs))
+
+    def get_dens(self, r):
+        """Calculate the density at a given radius (r).
 
-    def get_dens(self,r):
-        """Hernquist density: BT08, eq. 2.64, with alpha=1, beta=4
+        Args:
+            r (float): The radius at which to compute the density.
 
-        inputs
-        ---------------
-        r    : radius to compute density at
+        Returns:
+            float: The density in units of mass / (length^3).
         """
         alpha = 1
-        beta  = 4
-        return self.rho0 * (r/self.rscl)**(-alpha) * (1. + r/self.rscl)**(-beta+alpha)
-
-    def get_pot(self,r):
-        """Hernquist potential: BT08, eq. 2.67
+        beta = 4
+        return self.rho0 * (r / self.rscl) ** (-alpha) * (1. + r / self.rscl) ** (-beta + alpha)
 
-        inputs
-        ---------------
-        r    : radius to compute potential at
+    def get_pot(self, r):
+        """Calculate the gravitational potential at a given radius (r).
 
-        Note: this is a more complicated generalisation of the Hernquist potential,
-        \Phi = -GM/(r+rscl)
-        Feel free to inject that formula instead!
+        Args:
+            r (float): The radius at which to compute the potential.
 
-        returned in units of (length^2)*(mass^2)/(time^2)
+        Returns:
+            float: The gravitational potential in units of (length^2) * (mass^2) / (time^2).
         """
-        return -4*np.pi*self.G*self.rho0 *self.rscl*self.rscl * ( (2*(1.+r/self.rscl))**-1.)
+        return -4 * np.pi * self.G * self.rho0 * self.rscl * self.rscl * ((2 * (1. + r / self.rscl)) ** -1.)
 
-    def get_dphi_dr(self,r):
-        """Hernquist radial force: differentiate -GM/(r+a) using Wolfram Alpha
+    def get_dphi_dr(self, r):
+        """Calculate the radial force at a given radius (r).
 
-        inputs
-        ---------------
-        r    : radius to compute radial force at
+        Args:
+            r (float): The radius at which to compute the radial force.
 
-        returned in units of length * mass / (time^2)
+        Returns:
+            float: The radial force in units of (length * mass) / (time^2).
         """
-        return self.G*self.M/(self.rscl+r)**2
+        return self.G * self.M / (self.rscl + r) ** 2

exptool/models/icosphere.py

@@ -0,0 +1,86 @@
+"""
+Provides the Icosphere class for generating evenly distributed points on a sphere
+using recursive subdivision of an icosahedron.
+"""
+
+import numpy as np
+
+class Icosphere:
+    def __init__(self, subdivisions):
+        self.subdivisions = subdivisions
+        self.phi = (1.0 + np.sqrt(5.0)) / 2.0  # Golden ratio
+
+        # Define the vertices of an icosahedron
+        self.vertices = np.array([
+            [-1, 0, self.phi],
+            [1, 0, self.phi],
+            [-1, 0, -self.phi],
+            [1, 0, -self.phi],
+            [0, self.phi, 1],
+            [0, self.phi, -1],
+            [0, -self.phi, 1],
+            [0, -self.phi, -1],
+            [self.phi, 1, 0],
+            [self.phi, -1, 0],
+            [-self.phi, 1, 0],
+            [-self.phi, -1, 0]
+        ])
+
+        # Define the faces of the icosahedron
+        self.faces = np.array([
+            [0, 11, 5],
+            [0, 5, 1],
+            [0, 1, 7],
+            [0, 7, 10],
+            [0, 10, 11],
+            [1, 5, 9],
+            [5, 11, 4],
+            [11, 10, 2],
+            [10, 7, 6],
+            [7, 1, 8],
+            [3, 9, 4],
+            [3, 4, 2],
+            [3, 2, 6],
+            [3, 6, 8],
+            [3, 8, 9],
+            [4, 9, 5],
+            [2, 4, 11],
+            [6, 2, 10],
+            [8, 6, 7],
+            [9, 8, 1]
+        ], dtype=int)
+
+        # Initialize the list of points
+        self.points = []
+
+        # Divide each triangle face into smaller triangles
+        for face in self.faces:
+            a, b, c = self.vertices[face]
+            self._divide_triangle(a, b, c, self.subdivisions)
+
+        # Convert points to a NumPy array
+        self.points = np.array(self.points)
+
+        # verify that no anomalous points have been found
+        radii = np.linalg.norm(self.points,axis=1)
+        self.points = self.points[radii<=1.0]
+
+    def _normalize(self, v):
+        """Normalize a vector to have unit length."""
+        norm = np.linalg.norm(v)
+        if norm == 0:
+            return v
+        return v / norm
+
+    def _divide_triangle(self, a, b, c, depth):
+        """Recursively divide a triangle into smaller triangles."""
+        if depth == 0:
+            self.points.extend([a, b, c])
+        else:
+            ab = self._normalize((a + b) / 2)
+            ac = self._normalize((a + c) / 2)
+            bc = self._normalize((b + c) / 2)
+            self._divide_triangle(a, ab, ac, depth - 1)
+            self._divide_triangle(ab, b, bc, depth - 1)
+            self._divide_triangle(ac, bc, c, depth - 1)
+            self._divide_triangle(ab, bc, ac, depth - 1)