Fix bug in Visualization in setting units

dmentipl · Mar 10, 2020 · 0965e3f · 0965e3f
1 parent 7e3f435
commit 0965e3f
Show file tree

Hide file tree

Showing 2 changed files with 15 additions and 7 deletions.
diff --git a/plonk/visualize/__init__.py b/plonk/visualize/__init__.py
@@ -9,16 +9,16 @@
 
 >>> viz = plonk.visualize.plot(
 ...     snap=snap,
-...     x=snap['x'],
-...     y=snap['y'],
+...     x='x',
+...     y='y',
 ...     extent=(-100, 100, -100, 100),
 ... )
 
 Render the surface density in xz-plane.
 
 >>> viz = plonk.visualize.plot(
 ...     snap=snap,
-...     quantity=snap['density'],
+...     quantity='density',
 ...     x='x',
 ...     y='z',
 ...     extent=(-100, 100, -25, 25),

diff --git a/plonk/visualize/visualization.py b/plonk/visualize/visualization.py
@@ -219,6 +219,7 @@ def plot(
             if axis is not None:
                 raise ValueError('Trying to change existing axis attribute')
 
+        quantity_str: Optional[str] = quantity if isinstance(quantity, str) else None
         quantity, x, y, z, kind = _check_input(
             snap=self.snap, quantity=quantity, x=x, y=y, z=z, kind=kind
         )
@@ -242,8 +243,15 @@ def plot(
                 **_kwargs,
             )
             if units is not None:
+                if quantity_str is None:
+                    raise ValueError(
+                        'Cannot set units when passing in arrays. '
+                        'Instead, use strings to access\n'
+                        'quantities on Snap. E.g. plot(..., quantity="density", ...).'
+                    )
                 interpolated_data, extent = _convert_units(
                     snap=self.snap,
+                    quantity_str=quantity_str,
                     interpolated_data=interpolated_data,
                     extent=extent,
                     units=units,
@@ -321,22 +329,22 @@ def _check_input(*, snap, quantity, x, y, z, kind):
 def _convert_units(
     *,
     snap: SnapLike,
+    quantity_str: str,
     interpolated_data: ndarray,
     extent: Extent,
     units: Dict[str, Any],
     interp: str,
 ):
 
+    _quantity_unit = snap.get_array_unit(quantity_str)
     if interp == 'projection':
         data = (
-            (interpolated_data * snap.units['density'] * snap.units['length'])
+            (interpolated_data * _quantity_unit * snap.units['length'])
             .to(units['quantity'] * units['projection'])
             .magnitude
         )
     elif interp == 'cross_section':
-        data = (
-            (interpolated_data * snap.units['density']).to(units['quantity']).magnitude
-        )
+        data = (interpolated_data * _quantity_unit).to(units['quantity']).magnitude
 
     new_extent = tuple((extent * snap.units['length']).to(units['extent']).magnitude)