Fix mcolor arg (#59)

LSYS · Jan 14, 2024 · b903bd1 · b903bd1
1 parent 17c7dcc
commit b903bd1
Show file tree

Hide file tree

Showing 5 changed files with 437 additions and 184 deletions.
diff --git a/docs/images/group-grouporder-pvalue-sort-colorrows.png b/docs/images/group-grouporder-pvalue-sort-colorrows.png
diff --git a/docs/images/multimodel.png b/docs/images/multimodel.png
diff --git a/examples/readme-examples.ipynb b/examples/readme-examples.ipynb
diff --git a/examples/test-multmodel-sleep.ipynb b/examples/test-multmodel-sleep.ipynb
diff --git a/forestplot/mplot_graph_utils.py b/forestplot/mplot_graph_utils.py
@@ -1,3 +1,4 @@
+import warnings
 from typing import Any, List, Optional, Sequence, Tuple, Union
 
 import matplotlib.pyplot as plt
@@ -158,7 +159,7 @@ def mdraw_est_markers(
         _df = dataframe.query(f'{model_col}=="{modelgroup}"')
         base_y_vector = np.arange(len(_df)) - offset / 2 - (offset / 2) * (n - 2)
         _y = base_y_vector + (ix * offset)
-        ax.scatter(y=_y, x=_df[estimate], marker=msymbols[ix], color=mcolor[ix], s=markersize)
+        ax.scatter(y=_y, x=_df[estimate], marker=msymbols[ix], c=mcolor[ix], s=markersize)
     return ax
 
 
@@ -280,9 +281,13 @@ def mdraw_legend(
     leg_markersize = kwargs.get("leg_markersize", 8)
     leg_artists = []
     for ix, symbol in enumerate(msymbols):
-        leg_artists.append(
-            Line2D([0], [0], marker=symbol, color=mcolor[ix], markersize=leg_markersize)
-        )
+        try:
+            leg_artists.append(
+                Line2D([0], [0], marker=symbol, color=mcolor[ix], markersize=leg_markersize)
+            )
+        except IndexError:
+            warnings.warn("'msymbols' and 'mcolor' have different lengths.")
+            pass
     # Handle position of legend
     # bbox_to_anchor = kwargs.get("bbox_to_anchor", None)
     if len(modellabels) <= 2: