minor updates

hanjinliu · Feb 22, 2024 · c459a36 · c459a36
1 parent 5448413
commit c459a36
Show file tree

Hide file tree

Showing 2 changed files with 17 additions and 9 deletions.
diff --git a/whitecanvas/canvas/dataframe/_feature_cat.py b/whitecanvas/canvas/dataframe/_feature_cat.py
@@ -215,12 +215,12 @@ def add_hist2d(
         self,
         *,
         name: str | None = None,
-        color: str | None = None,
+        color: str | Sequence[str] | None = None,
         bins: HistBinType | tuple[HistBinType, HistBinType] = "auto",
         rangex: tuple[float, float] | None = None,
         rangey: tuple[float, float] | None = None,
         cmap=None,  # deprecated
-    ) -> _lt.DFHeatmap[_DF]:
+    ) -> _lt.DFMultiHeatmap[_DF]:
         """
         Add 2-D histogram of given x/y columns.
 
@@ -229,10 +229,10 @@ def add_hist2d(
 
         Parameters
         ----------
-        cmap : colormap-like, default "inferno"
-            Colormap to use for the heatmap.
         name : str, optional
             Name of the layer.
+        color : str, optional
+            Column name(s) for coloring the histogram.
         bins : int, array, str or tuple of them, default "auto"
             If int, the number of bins for both x and y. If tuple, the number of bins
             for x and y respectively.
@@ -243,7 +243,7 @@ def add_hist2d(
 
         Returns
         -------
-        DFHeatmap
+        DFMultiHeatmap
             Dataframe bound heatmap layer.
         """
         canvas = self._canvas()
@@ -260,7 +260,7 @@ def add_kde2d(
         name: str | None = None,
         color: str | None = None,
         band_width: KdeBandWidthType = "scott",
-    ) -> _lt.DFHeatmap[_DF]:
+    ) -> _lt.DFMultiHeatmap[_DF]:
         """
         Add 2-D kernel density estimation of given x/y columns.
 
@@ -269,16 +269,16 @@ def add_kde2d(
 
         Parameters
         ----------
-        cmap : colormap-like, default "inferno"
-            Colormap to use for the heatmap.
         name : str, optional
             Name of the layer.
+        color : str, optional
+            Column name(s) for coloring the densities.
         band_width : float, default None
             Bandwidth of the kernel density estimation. If None, use Scott's rule.
 
         Returns
         -------
-        DFHeatmap
+        DFMultiHeatmap
             Dataframe bound heatmap layer.
         """
         canvas = self._canvas()

diff --git a/whitecanvas/canvas/dataframe/_one_cat.py b/whitecanvas/canvas/dataframe/_one_cat.py
@@ -104,6 +104,14 @@ def __init__(
                     "Category column(s) must be specified by a string or a sequence "
                     f"of strings, got {offset!r}."
                 )
+        # check dtype
+        for col in offset:
+            arr = self._df[col]
+            if arr.dtype.kind in "fc":
+                raise ValueError(
+                    f"Column {col!r} cannot be interpreted as a categorical column "
+                    f"(dtype={arr.dtype!r})."
+                )
         self._offset: tuple[str, ...] = offset
         self._cat_iter = CatIterator(self._df, offset)
         self._value = value