Add scaling for cbar in correlation_heatmap (#143)

* Add scaling for cbar * Update CHANGELOG
tinkoff-ai · Oct 7, 2021 · b92d6bb · b92d6bb
1 parent f7c7ed3
commit b92d6bb
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 0 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -28,6 +28,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - Delete offset from WindowStatisticsTransform ([#111](https://github.com/tinkoff-ai/etna-ts/pull/111))
 - Add Pipeline example in Get started notebook ([#115](https://github.com/tinkoff-ai/etna-ts/pull/115))
 - Internal implementation of BinsegTrendTransform ([#141](https://github.com/tinkoff-ai/etna-ts/pull/141))
+- Colorebar scaling in Correlation heatmap plotter ([#143](https://github.com/tinkoff-ai/etna-ts/pull/143))
 
 ### Fixed
 - Add more obvious Exception Error for forecasting with unfitted model ([#102](https://github.com/tinkoff-ai/etna-ts/pull/102))

diff --git a/etna/analysis/plotters.py b/etna/analysis/plotters.py
@@ -240,6 +240,10 @@ def plot_correlation_matrix(
     """
     if segments is None:
         segments = sorted(ts.segments)
+    if "vmin" not in heatmap_kwargs:
+        heatmap_kwargs["vmin"] = -1
+    if "vmax" not in heatmap_kwargs:
+        heatmap_kwargs["vmax"] = 1
 
     correlation_matrix = get_correlation_matrix(ts, segments, method)
     ax = sns.heatmap(correlation_matrix, annot=True, fmt=".1g", square=True, **heatmap_kwargs)