add where

Author: lucascolley

src/pint_array/__init__.py

@@ -356,6 +356,31 @@ def func(x, /, *args, func_str=func_str, **kwargs):
 
         setattr(mod, func_str, func)
 
+    # def searchsorted(x1, x2, /, *, side='left', sorter=None):
+    #     if sorter is not None:
+    #         x1 = take(x1, sorter)
+
+    #     mask_count = xp.cumulative_sum(xp.astype(x1.mask, xp.int64))
+    #     x1_compressed = x1.data[~x1.mask]
+    #     count = xp.zeros(x1_compressed.size+1, dtype=xp.int64)
+    #     count[:-1] = mask_count[~x1.mask]
+    #     count[-1] = count[-2]
+    #     i = xp.searchsorted(x1_compressed, x2.data, side=side)
+    #     j = i + xp.take(count, i)
+    #     return MArray(j, mask=x2.mask)
+
+    # ignore units of condition, convert x2 to units of x1
+    def where(condition, x1, x2, /):
+        condition = asarray(condition)
+        x1 = asarray(x1)
+        x2 = asarray(x2)
+        units = x1.units
+        magnitude = xp.where(condition.magnitude, x1.magnitude, x2.m_as(units))
+        return ArrayUnitQuantity(magnitude, units)
+
+    # mod.searchsorted = searchsorted
+    mod.where = where
+
     # strip_unit_input_output_ufuncs = ["isnan", "isinf", "isfinite", "signbit", "sign"]
     # matching_input_bare_output_ufuncs = [
     #     "equal",