fix: backport "finally fix doublecrystalball for good" (#1377)

* actually check for apply_where import

* check for ModuleNotFoundError too

* missing return statement here

* missing normalization and import typo

Author: ikrommyd

coffea/lookup_tools/doublecrystalball.py

@@ -2,7 +2,7 @@
 
 _old_style_where = False
 try:
-    import scipy._lib.array_api_extra as xpx
+    from scipy._lib.array_api_extra import apply_where
     from scipy.stats._continuous_distns import (
         _norm_cdf,
         _norm_pdf_C,
@@ -74,11 +74,11 @@ def lhs(x, betaL, betaH, mL, mH):
         def rhs(x, betaL, betaH, mL, mH):
             if _old_style_where:
                 return _lazywhere(x < betaH, (-x, betaH, mH), f=core, f2=tail)
-            return xpx.apply_where(x < betaH, (-x, betaH, mH), core, tail)
+            return apply_where(x < betaH, (-x, betaH, mH), core, tail)
 
         if _old_style_where:
-            N * _lazywhere(x > -betaL, (x, betaL, betaH, mL, mH), f=rhs, f2=lhs)
-        return N * xpx.apply_where(x > -betaL, (x, betaL, betaH, mL, mH), rhs, lhs)
+            return N * _lazywhere(x > -betaL, (x, betaL, betaH, mL, mH), f=rhs, f2=lhs)
+        return N * apply_where(x > -betaL, (x, betaL, betaH, mL, mH), rhs, lhs)
 
     def _logpdf(self, x, betaL, betaH, mL, mH):
         """
@@ -106,15 +106,13 @@ def lhs(x, betaL, betaH, mL, mH):
         def rhs(x, betaL, betaH, mL, mH):
             if _old_style_where:
                 return _lazywhere(x < betaH, (-x, betaH, mH), f=core, f2=tail)
-            return xpx.apply_where(x < betaH, (-x, betaH, mH), core, tail)
+            return apply_where(x < betaH, (-x, betaH, mH), core, tail)
 
         if _old_style_where:
             return np.log(N) + _lazywhere(
                 x > -betaL, (x, betaL, betaH, mL, mH), f=rhs, f2=lhs
             )
-        return np.log(N) + xpx.apply_where(
-            x > -betaL, (x, betaL, betaH, mL, mH), rhs, lhs
-        )
+        return np.log(N) + apply_where(x > -betaL, (x, betaL, betaH, mL, mH), rhs, lhs)
 
     def _cdf(self, x, betaL, betaH, mL, mH):
         """
@@ -159,11 +157,11 @@ def rhs(x, betaL, betaH, mL, mH):
                 return _lazywhere(
                     x < betaH, (x, betaL, betaH, mL, mH), f=core, f2=hightail
                 )
-            return xpx.apply_where(x < betaH, (x, betaL, betaH, mL, mH), core, hightail)
+            return apply_where(x < betaH, (x, betaL, betaH, mL, mH), core, hightail)
 
         if _old_style_where:
-            return _lazywhere(x > -betaL, (x, betaL, betaH, mL, mH), f=rhs, f2=lhs)
-        return N * xpx.apply_where(x > -betaL, (x, betaL, betaH, mL, mH), rhs, lhs)
+            return N * _lazywhere(x > -betaL, (x, betaL, betaH, mL, mH), f=rhs, f2=lhs)
+        return N * apply_where(x > -betaL, (x, betaL, betaH, mL, mH), rhs, lhs)
 
     def _ppf(self, p, betaL, betaH, mL, mH):
         """
@@ -215,19 +213,15 @@ def ppf_greater(p, betaL, betaH, mL, mH):
                 return _lazywhere(
                     p > pbetaH, (p, betaL, betaH, mL, mH), f=hightail, f2=core
                 )
-            return xpx.apply_where(
-                p > pbetaH, (p, betaL, betaH, mL, mH), hightail, core
-            )
+            return apply_where(p > pbetaH, (p, betaL, betaH, mL, mH), hightail, core)
 
         N = 1.0 / (inttail(betaL, mL) + intcore(betaL, betaH) + inttail(betaH, mH))
         pbetaL = N * (mL / betaL) * np.exp(-0.5 * betaL * betaL) / (mL - 1)
         if _old_style_where:
             return _lazywhere(
                 p < pbetaL, (p, betaL, betaH, mL, mH), f=lowtail, f2=ppf_greater
             )
-        return xpx.apply_where(
-            p < pbetaL, (p, betaL, betaH, mL, mH), lowtail, ppf_greater
-        )
+        return apply_where(p < pbetaL, (p, betaL, betaH, mL, mH), lowtail, ppf_greater)
 
     def _munp(self, n, betaL, betaH, mL, mH):
         """