Add Stokes number

plonk/analysis/particles.py

@@ -202,8 +202,7 @@ def keplerian_frequency(
     snap
         The Snap object.
     gravitational_parameter
-        The gravitational parameter (mu = G M). Can be a float or a Pint
-        quantity.
+        The gravitational parameter (mu = G M) as 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).
@@ -213,7 +212,7 @@ def keplerian_frequency(
     Returns
     -------
     ndarray
-        The eccentricity on the particles.
+        The Keplerian frequency on the particles.
     """
     if ignore_accreted:
         h: ndarray = snap['smoothing_length']
@@ -232,6 +231,53 @@ def keplerian_frequency(
     return np.sqrt(mu / radius ** 3)
 
 
+def stokes_number(
+    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 Stokes number.
+
+    Parameters
+    ----------
+    snap
+        The Snap object.
+    gravitational_parameter
+        The gravitational parameter (mu = G M) as 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 Stokes number on the particles.
+    """
+    if ignore_accreted:
+        h: ndarray = snap['smoothing_length']
+        pos: ndarray = snap['position'][h > 0]
+        t_s: ndarray = snap['stopping_time'][h > 0]
+    else:
+        pos = snap['position']
+        t_s = snap['stopping_time']
+
+    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)
+    Omega_k = np.sqrt(mu / radius ** 3)
+
+    Stokes = t_s * Omega_k[:, np.newaxis]
+    return Stokes
+
+
 def semi_major_axis(
     snap: SnapLike,
     gravitational_parameter: Any,
@@ -249,8 +295,7 @@ def semi_major_axis(
     snap
         The Snap object.
     gravitational_parameter
-        The gravitational parameter (mu = G M). Can be a float or a Pint
-        quantity.
+        The gravitational parameter (mu = G M) as 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).
@@ -314,8 +359,7 @@ def eccentricity(
     snap
         The Snap object.
     gravitational_parameter
-        The gravitational parameter (mu = G M). Can be a float or a Pint
-        quantity.
+        The gravitational parameter (mu = G M) as 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).

plonk/analysis/profile.py

@@ -293,7 +293,10 @@ def __getitem__(self, name: str) -> ndarray:
                     del self.snap[name]
                 return self._profiles[name]
             else:
-                if name.split('_')[0] == 'dust' or name in ('stopping_time',):
+                if name.split('_')[0] == 'dust' or name in (
+                    'stopping_time',
+                    'stokes_number',
+                ):
                     raise ValueError(
                         'To access dust profiles try, for example, '
                         'prof["stopping_time_001"] or\nprof["dust_mass_total"]'

plonk/snap/snap.py

@@ -947,3 +947,22 @@ def gas_density(snap) -> ndarray:
         def dust_density(snap) -> ndarray:
             """Dust density."""
             return particles.dust_density(snap=snap)
+
+        @Snap.add_array(unit='dimensionless')
+        def stokes_number(snap) -> ndarray:
+            """Stokes number."""
+            try:
+                gravitational_parameter = snap.properties['gravitational_parameter']
+            except KeyError:
+                raise ValueError(
+                    'To get eccentricity, 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.stokes_number(
+                snap=snap,
+                gravitational_parameter=gravitational_parameter,
+                origin=origin,
+            )