Make Profile work with Snap physical units

plonk/analysis/profile.py

@@ -12,7 +12,9 @@
 from numpy import ndarray
 from pandas import DataFrame
 
+from .. import Quantity
 from ..snap import Snap
+from ..utils.math import norm
 from .particles import eccentricity as _eccentricity
 from .particles import specific_angular_momentum
 
@@ -102,11 +104,6 @@ def __init__(
         spacing: str = 'linear',
         ignore_accreted: bool = True,
     ):
-        if snap._physical_units:
-            raise ValueError(
-                'Profile does not (currently) work with Snap in physical units.\n'
-                'Unset physical units (snap.physical_units(unset=True)) and try again.'
-            )
 
         self.snap = snap
         self.ndim = ndim
@@ -122,11 +119,21 @@ def __init__(
 
         self.bin_edges, self['size'] = self._setup_bins()
         self.bin_centers = 0.5 * (self.bin_edges[:-1] + self.bin_edges[1:])
-        self._particle_bin = np.digitize(self._x, self.bin_edges)
+        if self.snap._physical_units:
+            self._particle_bin = np.digitize(
+                self._x.magnitude, self.bin_edges.magnitude
+            )
+        else:
+            self._particle_bin = np.digitize(self._x, self.bin_edges)
         self.bin_indicies = self._set_particle_bin_indicies()
 
         self._profiles['radius'] = self.bin_centers
-        self._profiles['number'] = np.histogram(self._x, self.bin_edges)[0]
+        if self.snap._physical_units:
+            self._profiles['number'] = np.histogram(
+                self._x.magnitude, self.bin_edges.magnitude
+            )[0]
+        else:
+            self._profiles['number'] = np.histogram(self._x, self.bin_edges)[0]
 
     def _check_spacing(self, spacing):
         if spacing.lower() in ('lin', 'linear'):
@@ -152,6 +159,14 @@ def _calculate_x(self) -> ndarray:
             return np.sqrt(pos[:, 0] ** 2 + pos[:, 1] ** 2 + pos[:, 2] ** 2)
 
     def _set_range(
+        self, radius_min: Optional[Any], radius_max: Optional[Any]
+    ) -> Tuple[float, float]:
+        if self.snap._physical_units:
+            return self._set_range_physical_units(radius_min, radius_max)
+        else:
+            return self._set_range_code_units(radius_min, radius_max)
+
+    def _set_range_code_units(
         self, radius_min: Optional[float], radius_max: Optional[float]
     ) -> Tuple[float, float]:
         if radius_min is None:
@@ -165,7 +180,53 @@ def _set_range(
 
         return rmin, rmax
 
+    def _set_range_physical_units(
+        self, radius_min: Optional[Any], radius_max: Optional[Any]
+    ) -> Tuple[float, float]:
+        if radius_min is None:
+            rmin = self._x.min()
+        else:
+            rmin = radius_min
+        if radius_max is None:
+            rmax = np.percentile(self._x.magnitude, 99, axis=0) * self._x.units
+        else:
+            rmax = radius_max
+
+        return rmin, rmax
+
     def _setup_bins(self) -> ndarray:
+        if self.snap._physical_units:
+            bin_edges = self._bin_edges_physical_units()
+        else:
+            bin_edges = self._bin_edges_code_units()
+        if self.ndim == 2:
+            bin_sizes = np.pi * (bin_edges[1:] ** 2 - bin_edges[:-1] ** 2)
+        elif self.ndim == 3:
+            bin_sizes = 4 / 3 * np.pi * (bin_edges[1:] ** 3 - bin_edges[:-1] ** 3)
+        return bin_edges, bin_sizes
+
+    def _bin_edges_physical_units(self):
+        if self.spacing == 'linear':
+            bin_edges = (
+                np.linspace(
+                    self.range[0].magnitude, self.range[1].magnitude, self.n_bins + 1
+                )
+                * self.range[0].units
+            )
+        elif self.spacing == 'log':
+            bin_edges = (
+                np.logspace(
+                    np.log10(self.range[0].magnitude),
+                    np.log10(self.range[1].magnitude),
+                    self.n_bins + 1,
+                )
+                * self.range[0].units
+            )
+        else:
+            raise ValueError('Cannot determine spacing to setup bins')
+        return bin_edges
+
+    def _bin_edges_code_units(self):
         if self.spacing == 'linear':
             bin_edges = np.linspace(self.range[0], self.range[1], self.n_bins + 1)
         elif self.spacing == 'log':
@@ -174,11 +235,7 @@ def _setup_bins(self) -> ndarray:
             )
         else:
             raise ValueError('Cannot determine spacing to setup bins')
-        if self.ndim == 2:
-            bin_sizes = np.pi * (bin_edges[1:] ** 2 - bin_edges[:-1] ** 2)
-        elif self.ndim == 3:
-            bin_sizes = 4 / 3 * np.pi * (bin_edges[1:] ** 3 - bin_edges[:-1] ** 3)
-        return bin_edges, bin_sizes
+        return bin_edges
 
     def _set_particle_bin_indicies(self) -> List[ndarray]:
         sortind = self._particle_bin.argsort()
@@ -212,8 +269,8 @@ def __getitem__(self, name: str) -> ndarray:
 
     def __setitem__(self, name: str, item: ndarray):
         """Set the profile directly."""
-        if not isinstance(item, ndarray):
-            raise ValueError('"item" must be ndarray')
+        if not isinstance(item, (ndarray, Quantity)):
+            raise ValueError('"item" must be ndarray or pint Quantity')
         if item.shape[0] != self.n_bins:
             raise ValueError('Length of array does not match number of bins')
         if name in self.loaded_keys():
@@ -315,9 +372,13 @@ def particles_to_binned_quantity(self, function, *args) -> ndarray:
         """
         binned_quantity = np.zeros(self.n_bins)
         for idx, bin_ind in enumerate(self.bin_indicies):
-            binned_quantity[idx] = function(
-                tuple([arg[self._mask][bin_ind] for arg in args])
-            )
+            val = function(*tuple([arg[self._mask][bin_ind] for arg in args]))
+            if self.snap._physical_units:
+                binned_quantity[idx] = val.magnitude
+            else:
+                binned_quantity[idx] = val
+        if self.snap._physical_units:
+            binned_quantity *= val.units
         return binned_quantity
 
 
@@ -363,7 +424,7 @@ def aspect_ratio(prof) -> ndarray:
 def angmom_mag(prof) -> ndarray:
     """Magnitude of specific angular momentum profile."""
     angmom = specific_angular_momentum(snap=prof.snap)
-    angmom_mag = np.linalg.norm(angmom, axis=1)
+    angmom_mag = norm(angmom, axis=1)
 
     return prof.particles_to_binned_quantity(np.mean, angmom_mag)
 
@@ -417,7 +478,7 @@ def eccentricity(prof) -> ndarray:
     """Orbital eccentricity profile.
 
     This profile assumes the central mass is at (0, 0, 0) and that the
-    gravitational parameter (G*M) is set.
+    gravitational parameter (mu = G M) is set.
     """
     try:
         gravitational_parameter = prof.properties['gravitational_parameter']