Merge pull request #270 from VlachosGroup/cmap-patch

fix import of get_cmap
VlachosGroup · May 16, 2024 · e170435 · e170435
2 parents efdfc4b + 305b89b
commit e170435
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 2 deletions.
diff --git a/aimsim/__init__.py b/aimsim/__init__.py
@@ -7,4 +7,4 @@
 except ImportError:
     pass  # aimsim_core does not include this
 
-__version__ = "2.2.0"
+__version__ = "2.2.1"
diff --git a/aimsim/tasks/cluster_data.py b/aimsim/tasks/cluster_data.py
@@ -3,6 +3,10 @@
 from os.path import dirname
 
 import matplotlib.pyplot as plt
+try:
+    from matplotlib.cm import get_cmap
+except:  # noqa: E722
+    from matplotlib.pyplot import get_cmap
 from matplotlib.colors import rgb2hex
 import yaml
 
@@ -54,7 +58,7 @@ def _extract_configs(self):
         self.plot_settings = dict()
         self.plot_settings["cluster_plot"] = {
             "cluster_colors": [
-                rgb2hex(plt.cm.get_cmap("tab20", self.n_clusters)(cluster_id))
+                rgb2hex(get_cmap("tab20", self.n_clusters)(cluster_id))
                 for cluster_id in range(self.n_clusters)],
             "response": "Response",
         }