Make xhistogram dependency optional

pymc_experimental/distributions/histogram_utils.py

@@ -17,7 +17,6 @@
 
 import numpy as np
 import pymc as pm
-import xhistogram.core
 from numpy.typing import ArrayLike
 
 try:
@@ -26,13 +25,21 @@
 except ImportError:
     dask = None
 
+try:
+    import xhistogram.core
+except ImportError:
+    xhistogram = None
+
 
 __all__ = ["quantile_histogram", "discrete_histogram", "histogram_approximation"]
 
 
 def quantile_histogram(
     data: ArrayLike, n_quantiles=1000, zero_inflation=False
 ) -> Dict[str, ArrayLike]:
+    if xhistogram is None:
+        raise RuntimeError("quantile_histogram requires xhistogram package")
+
     if dask and isinstance(data, (dask.dataframe.Series, dask.dataframe.DataFrame)):
         data = data.to_dask_array(lengths=True)
     if zero_inflation:
@@ -67,6 +74,9 @@ def quantile_histogram(
 
 
 def discrete_histogram(data: ArrayLike, min_count=None) -> Dict[str, ArrayLike]:
+    if xhistogram is None:
+        raise RuntimeError("discrete_histogram requires xhistogram package")
+
     if dask and isinstance(data, (dask.dataframe.Series, dask.dataframe.DataFrame)):
         data = data.to_dask_array(lengths=True)
     mid, count_uniq = np.unique(data, return_counts=True)

requirements.txt

@@ -1,2 +1 @@
 pymc>=4.3.0
-xhistogram

setup.py

@@ -81,5 +81,9 @@ def get_version():
         classifiers=classifiers,
         python_requires=">=3.8",
         install_requires=install_reqs,
-        extras_requires=dict(dask=["dask[all]"]),
+        extras_requires=dict(
+            dask=["dask[all]"], 
+            histogram=["xhistogram"], 
+            complete=["dask[all]", "xhistogram"]
+        ),
     )