Use function to strip units instead of try..except

arpes/utilities/conversion/kx_ky_conversion.py

@@ -2,16 +2,17 @@
 
 Broadly, this covers cases where we are not performing photon energy scans.
 """
-import numpy as np
+import math
+from typing import Any, Callable, Dict, List
 
 import numba
-import math
+import numpy as np
+import pint
+import xarray as xr
 
 import arpes.constants
-import xarray as xr
-from typing import Any, Callable, Dict, List
 
-from .base import CoordinateConverter, K_SPACE_BORDER, MOMENTUM_BREAKPOINTS
+from .base import K_SPACE_BORDER, MOMENTUM_BREAKPOINTS, CoordinateConverter
 from .bounds_calculations import calculate_kp_bounds, calculate_kx_ky_bounds
 
 __all__ = ["ConvertKp", "ConvertKxKy"]
@@ -74,6 +75,11 @@ def _safe_compute_k_tot(hv, work_function, binding_energy):
 
     return k_tot
 
+def _strip_units(val):
+    if isinstance(val, pint.Quantity):
+        return val.magnitude
+    return val
+
 
 class ConvertKp(CoordinateConverter):
     """A momentum converter for single ARPES (kp) cuts."""
@@ -143,12 +149,9 @@ def kspace_to_phi(
 
         offset = self.arr.S.phi_offset + parallel_angle
 
-        try:
-            polar_angle = polar_angle.magnitude
-            parallel_angle = parallel_angle.magnitude
-            offset = offset.magnitude
-        except:
-            pass
+        polar_angle = _strip_units(polar_angle)
+        parallel_angle = _strip_units(parallel_angle)
+        offset = _strip_units(offset)
 
         if self.k_tot is None:
             self.compute_k_tot(binding_energy)
@@ -318,11 +321,7 @@ def rkx_rky(self, kx, ky):
             return self.rkx, self.rky
 
         chi = self.arr.S.lookup_offset_coord("chi")
-
-        try:
-            chi = chi.magnitude
-        except:
-            pass
+        chi = _strip_units(chi)
 
         self.rkx = np.zeros_like(kx)
         self.rky = np.zeros_like(ky)
@@ -349,10 +348,7 @@ def kspace_to_phi(
         scan_angle = self.direct_angles[1]
         self.phi = np.zeros_like(ky)
         offset = self.arr.S.phi_offset + self.arr.S.lookup_offset_coord(self.parallel_angles[0])
-        try:
-            offset = offset.magnitude
-        except:
-            pass
+        offset = _strip_units(offset)
 
         par_tot = isinstance(self.k_tot, np.ndarray) and len(self.k_tot) != 1
         assert len(self.k_tot) == len(self.phi) or len(self.k_tot) == 1
@@ -404,37 +400,25 @@ def kspace_to_perp_angle(
                 offset = self.arr.S.psi_offset - self.arr.S.lookup_offset_coord(
                     self.parallel_angles[1]
                 )
-                try:
-                    offset = offset.magnitude
-                except:
-                    pass
+                offset = _strip_units(offset)
                 _small_angle_arcsin(kx, self.k_tot, self.perp_angle, offset, par_tot, True)
             else:
                 offset = self.arr.S.psi_offset + self.arr.S.lookup_offset_coord(
                     self.parallel_angles[1]
                 )
-                try:
-                    offset = offset.magnitude
-                except:
-                    pass
+                offset = _strip_units(offset)
                 _small_angle_arcsin(ky, self.k_tot, self.perp_angle, offset, par_tot, False)
         elif scan_angle == "beta":
             offset = self.arr.S.beta_offset + self.arr.S.lookup_offset_coord(
                 self.parallel_angles[1]
             )
-            try:
-                offset = offset.magnitude
-            except:
-                pass
+            offset = _strip_units(offset)
             _exact_arcsin(ky, kx, self.k_tot, self.perp_angle, offset, par_tot, True)
         elif scan_angle == "theta":
             offset = self.arr.S.theta_offset - self.arr.S.lookup_offset_coord(
                 self.parallel_angles[1]
             )
-            try:
-                offset = offset.magnitude
-            except:
-                pass
+            offset = _strip_units(offset)
             _exact_arcsin(kx, ky, self.k_tot, self.perp_angle, offset, par_tot, True)
         else:
             raise ValueError(