refac: Fix plot with bool values

Fixes #11
langflow-ai · Jul 30, 2022 · 9c3ce88 · 9c3ce88
1 parent 7083743
commit 9c3ce88
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 1 deletion.
diff --git a/src/wavy/panel.py b/src/wavy/panel.py
@@ -174,7 +174,8 @@ def _copy_attrs(self, df):
  def _constructor(self):
  def f(*args, **kw):
 
- if len(args[0].axes[1]) == self.num_timesteps:
+ index = [a for a in args[0].axes if isinstance(a, pd.MultiIndex)][0]
+ if len(index) == self.num_timesteps:
  return pd.DataFrame(*args, **kw)
 
  df = Panel(*args, **kw)

diff --git a/src/wavy/plot.py b/src/wavy/plot.py
@@ -143,6 +143,10 @@ def plot(panel, use_timestep=False, add_annotation=False, **kwargs):
 
  panel = panel.row_panel(n=0)
 
+ for col in panel.columns:
+ if panel[col].dtype == bool:
+ panel[col] = panel[col].astype(int)
+
  if use_timestep:
  panel = panel.droplevel(0, axis=0)
  else: