Add method to set gravitational parameter on Snap

Also, add keplerian_frequency analysis function. Plus add Toomre Q
profile.

plonk/__init__.py

@@ -82,10 +82,3 @@
 
 # Canonical version number
 __version__ = '0.4.0'
-
-# Add units
-units.define('solar_mass = 1.9891e33 g')
-units.define('solar_radius = 6.959500e10 cm')
-units.define('earth_mass = 5.979e27 g')
-units.define('earth_radius = 6.371315e8 cm')
-units.define('jupiter_mass = 1.89813e30 g')

plonk/analysis/particles.py

@@ -183,6 +183,53 @@ def specific_kinetic_energy(snap: SnapLike, ignore_accreted: bool = False) -> nd
     return 1 / 2 * norm(vel, axis=1) ** 2
 
 
+def keplerian_frequency(
+    snap: SnapLike,
+    gravitational_parameter: Any,
+    origin: Union[ndarray, Tuple[float, float, float]] = (0.0, 0.0, 0.0),
+    ignore_accreted: bool = False,
+) -> ndarray:
+    """Calculate the Keplerian orbital frequency.
+
+    The Keplerian orbital frequency of particles around a mass specified
+    by gravitational parameter with an optional to specify the position
+    of the mass.
+
+    Parameters
+    ----------
+    snap
+        The Snap object.
+    gravitational_parameter
+        The gravitational parameter (mu = G M). Can be a float or a Pint
+        quantity.
+    origin : optional
+        The origin around which to compute the angular momentum as a
+        ndarray or tuple (x, y, z). Default is (0, 0, 0).
+    ignore_accreted : optional
+        Ignore accreted particles. Default is False.
+
+    Returns
+    -------
+    ndarray
+        The eccentricity on the particles.
+    """
+    if ignore_accreted:
+        h: ndarray = snap['smoothing_length']
+        pos: ndarray = snap['position'][h > 0]
+    else:
+        pos = snap['position']
+
+    origin = np.array(origin)
+    pos = pos - origin
+
+    mu = gravitational_parameter
+    if not isinstance(pos, Quantity):
+        mu = (mu * snap.units['time'] ** 2 / snap.units['length'] ** 3).magnitude
+
+    radius = norm(pos, axis=1)
+    return np.sqrt(mu / radius ** 3)
+
+
 def semi_major_axis(
     snap: SnapLike,
     gravitational_parameter: Any,
@@ -224,10 +271,9 @@ def semi_major_axis(
     origin = np.array(origin)
     pos = pos - origin
 
-    if isinstance(pos, Quantity):
-        mu = gravitational_parameter.to_reduced_units()
-    else:
-        mu = gravitational_parameter
+    mu = gravitational_parameter
+    if not isinstance(pos, Quantity):
+        mu = (mu * snap.units['time'] ** 2 / snap.units['length'] ** 3).magnitude
 
     radius = norm(pos, axis=1)
 
@@ -246,8 +292,6 @@ def semi_major_axis(
         mu * (1 - eccentricity ** 2)
     )
 
-    if isinstance(pos, Quantity):
-        return semi_major_axis.magnitude
     return semi_major_axis
 
 
@@ -292,10 +336,9 @@ def eccentricity(
     origin = np.array(origin)
     pos = pos - origin
 
-    if isinstance(pos, Quantity):
-        mu = gravitational_parameter.to_reduced_units()
-    else:
-        mu = gravitational_parameter
+    mu = gravitational_parameter
+    if not isinstance(pos, Quantity):
+        mu = (mu * snap.units['time'] ** 2 / snap.units['length'] ** 3).magnitude
 
     radius = norm(pos, axis=1)
 
@@ -309,8 +352,6 @@ def eccentricity(
     term = specific_energy * (specific_angular_momentum_magnitude / mu) ** 2
     eccentricity = np.sqrt(1 + 2 * term)
 
-    if isinstance(pos, Quantity):
-        return eccentricity.magnitude
     return eccentricity
 
 

plonk/analysis/profile.py

@@ -13,7 +13,8 @@
 from pandas import DataFrame
 
 from .. import Quantity
-from ..snap import Snap
+from .. import units as plonk_units
+from ..snap import Snap, gravitational_constant_in_code_units
 from ..utils.math import norm
 
 
@@ -311,7 +312,7 @@ def __getitem__(self, name: str) -> ndarray:
                 array = norm(self.snap[name_root], axis=1)
             elif name_suffix == 'total':
                 array = np.sum(self.snap[name_root], axis=1)
-            elif str_is_int(name_suffix):
+            elif _str_is_int(name_suffix):
                 array = self.snap[name_root][:, int(name_suffix) - 1]
                 name = name_root + f'_{int(name_suffix):03}'
             else:
@@ -540,7 +541,21 @@ def angular_momentum_phi(prof) -> ndarray:
     return np.arctan2(angular_momentum_y, angular_momentum_x)
 
 
-def str_is_int(string):
+@Profile.profile_property
+def toomre_Q(prof) -> ndarray:
+    """Toomre Q parameter."""
+    if not prof.snap._physical_units:
+        G = gravitational_constant_in_code_units(prof.snap)
+    else:
+        G = (1 * plonk_units.newtonian_constant_of_gravitation).to_base_units()
+    return (
+        prof['sound_speed']
+        * prof['keplerian_frequency']
+        / (np.pi * G * prof['density'])
+    )
+
+
+def _str_is_int(string):
     try:
         int(string)
         return True

plonk/snap/__init__.py

@@ -6,5 +6,6 @@
 
 from .readers import load_snap
 from .snap import Snap, SnapLike
+from .units import gravitational_constant_in_code_units
 
-__all__ = ['Snap', 'SnapLike', 'load_snap']
+__all__ = ['Snap', 'SnapLike', 'gravitational_constant_in_code_units', 'load_snap']

plonk/snap/snap.py

@@ -16,6 +16,7 @@
 from scipy.spatial.transform import Rotation
 
 from .. import Quantity
+from .. import units as plonk_units
 from ..analysis import particles
 
 
@@ -485,6 +486,28 @@ def translate(self, translation: ndarray) -> Snap:
 
         return self
 
+    def set_gravitational_parameter(self, sink_idx: Union[int, List[int]]):
+        """Set standard gravitational parameter.
+
+        Calculate the standard gravitational parameter (G M) given a
+        sink index, or list of sink indices for multiple systems, etc.
+
+        Parameters
+        ----------
+        sink_idx
+            The sink index or list of indices.
+        """
+        G = plonk_units.newtonian_constant_of_gravitation
+        if isinstance(sink_idx, (int, list)):
+            M = self.sinks['mass'][sink_idx]
+        else:
+            raise ValueError('Cannot determine gravitational parameter')
+        if self._physical_units:
+            G = (G * M).to_base_units()
+        else:
+            G = (G * self.units['mass'] * M).to_base_units()
+        self.properties['gravitational_parameter'] = G
+
     def to_dataframe(
         self, columns: Union[Tuple[str, ...], List[str]] = None
     ) -> DataFrame:
@@ -785,6 +808,21 @@ def specific_kinetic_energy(snap) -> ndarray:
         """Specific kinetic energy."""
         return particles.specific_kinetic_energy(snap=snap)
 
+    @Snap.add_array(unit='frequency')
+    def keplerian_frequency(snap) -> ndarray:
+        """Keplerian orbital frequency."""
+        try:
+            gravitational_parameter = snap.properties['gravitational_parameter']
+        except KeyError:
+            raise ValueError(
+                'To get Keplerian frequency, first set the gravitational parameter\n'
+                'via snap.set_gravitational_parameter.'
+            )
+        origin = snap.translation if snap.translation is not None else (0.0, 0.0, 0.0)
+        return particles.keplerian_frequency(
+            snap=snap, gravitational_parameter=gravitational_parameter, origin=origin
+        )
+
     @Snap.add_array(unit='length')
     def semi_major_axis(snap) -> ndarray:
         """Semi-major axis."""
@@ -793,7 +831,7 @@ def semi_major_axis(snap) -> ndarray:
         except KeyError:
             raise ValueError(
                 'To get semi-major axis, first set the gravitational parameter\n'
-                'via snap.properties["gravitational_parameter"].'
+                'via snap.set_gravitational_parameter.'
             )
         origin = snap.translation if snap.translation is not None else (0.0, 0.0, 0.0)
         return particles.semi_major_axis(
@@ -808,7 +846,7 @@ def eccentricity(snap) -> ndarray:
         except KeyError:
             raise ValueError(
                 'To get eccentricity, first set the gravitational parameter\n'
-                'via snap.properties["gravitational_parameter"].'
+                'via snap.set_gravitational_parameter.'
             )
         origin = snap.translation if snap.translation is not None else (0.0, 0.0, 0.0)
         return particles.eccentricity(

plonk/snap/units.py

@@ -1,5 +1,15 @@
 """Physical units on snapshots."""
 
+from .. import units as plonk_units
+from .snap import SnapLike
+
+# Add units
+plonk_units.define('solar_mass = 1.9891e33 g')
+plonk_units.define('solar_radius = 6.959500e10 cm')
+plonk_units.define('earth_mass = 5.979e27 g')
+plonk_units.define('earth_radius = 6.371315e8 cm')
+plonk_units.define('jupiter_mass = 1.89813e30 g')
+
 
 def generate_units_dictionary(length, mass, time, magnetic_field):
     """Generate units dictionary.
@@ -40,3 +50,22 @@ def generate_units_dictionary(length, mass, time, magnetic_field):
     units['pressure'] = mass / time ** 2 / length
 
     return units
+
+
+def gravitational_constant_in_code_units(snap: SnapLike) -> float:
+    """Gravitational constant in code units.
+
+    Parameters
+    ----------
+    snap
+        The Snap object.
+
+    Returns
+    -------
+    float
+        The gravitational constant in code units.
+    """
+    G = plonk_units.newtonian_constant_of_gravitation
+    G_units = snap.units['length'] ** 3 / snap.units['mass'] / snap.units['time'] ** 2
+    G = (G / G_units).to_base_units().magnitude
+    return G