Set cgs units

dmentipl · Mar 10, 2020 · 05f5dd8 · 05f5dd8
1 parent eda5d2f
commit 05f5dd8
Show file tree

Hide file tree

Showing 2 changed files with 13 additions and 11 deletions.
diff --git a/plonk/__init__.py b/plonk/__init__.py
@@ -63,7 +63,7 @@
 
 import pint
 
-units = pint.UnitRegistry()
+units = pint.UnitRegistry(system='cgs')
 Quantity = units.Quantity
 
 from . import analysis, simulation, snap, utils, visualize
@@ -84,8 +84,8 @@
 __version__ = '0.3.1'
 
 # Add units
-units.define('solarm = 1.9891e33 g')
-units.define('solarr = 6.959500e10 cm')
-units.define('earthm = 5.979e27 g')
-units.define('earthr = 6.371315e8 cm')
-units.define('jupiterm = 1.89813e30 g')
+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')
diff --git a/plonk/analysis/particles.py b/plonk/analysis/particles.py
@@ -3,7 +3,7 @@
 Calculate various quantities on the particles.
 """
 
-from typing import Tuple, Union
+from typing import Any, Tuple, Union
 
 import numpy as np
 from numpy import ndarray
@@ -180,7 +180,7 @@ def specific_kinetic_energy(snap: SnapLike, ignore_accreted: bool = False) -> nd
 
 def semi_major_axis(
     snap: SnapLike,
-    gravitational_parameter: float,
+    gravitational_parameter: Any,
     origin: Union[ndarray, Tuple[float, float, float]] = (0.0, 0.0, 0.0),
     ignore_accreted: bool = False,
 ) -> ndarray:
@@ -195,7 +195,8 @@ def semi_major_axis(
     snap
         The Snap object.
     gravitational_parameter
-        The gravitational parameter (G*M).
+        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).
@@ -241,7 +242,7 @@ def semi_major_axis(
 
 def eccentricity(
     snap: SnapLike,
-    gravitational_parameter: float,
+    gravitational_parameter: Any,
     origin: Union[ndarray, Tuple[float, float, float]] = (0.0, 0.0, 0.0),
     ignore_accreted: bool = False,
 ) -> ndarray:
@@ -256,7 +257,8 @@ def eccentricity(
     snap
         The Snap object.
     gravitational_parameter
-        The gravitational parameter (G*M).
+        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).