From 4daf40c6b09f8907b4e099e4e0d6b142902edbb7 Mon Sep 17 00:00:00 2001
From: Noe Brucy <noe.brucy@cea.fr>
Date: Tue, 20 Jul 2021 10:46:24 +0200
Subject: [PATCH] [plotter] add new options to plot map

---
 plotter.py | 68 ++++++++++++++++++++++++++++++++++++++++--------------
 1 file changed, 51 insertions(+), 17 deletions(-)

diff --git a/plotter.py b/plotter.py
index bbd5c12..4ba9428 100644
--- a/plotter.py
+++ b/plotter.py
@@ -11,6 +11,8 @@ This is the plotter module.
 import os
 from functools import partial
 import matplotlib as mpl
+from mpl_toolkits.axes_grid1.anchored_artists import AnchoredSizeBar
+
 import numpy as np
 import tables
 from astrophysix.simdm.datafiles import Datafile, PlotInfo, PlotType
@@ -570,6 +572,11 @@ class Plotter(Aggregator, BaseProcessor):
         transform=None,
         vmin=None,
         vmax=None,
+        scalebar=False,
+        scalebar_size=1,
+        axes=True,
+        colorbar=True,
+        tight_cb=False,
         **kwargs,
     ):
         """
@@ -617,28 +624,56 @@ class Plotter(Aggregator, BaseProcessor):
 
         plt.locator_params(axis="both", nbins=self.params.plot.ntick)
 
-        if xlabel is None:
-            xlabel = self._ax_title[ax_h]
-        if ylabel is None:
-            ylabel = self._ax_title[ax_v]
-        if put_units:
-            xlabel = xlabel + unit_str(unit_space)
-            ylabel = ylabel + unit_str(unit_space)
-        plt.xlabel(xlabel)
-        plt.ylabel(ylabel)
+        if scalebar:
+            scalebar = AnchoredSizeBar(
+                plt.gca().transData,
+                scalebar_size,
+                f'{scalebar_size} {unit_str(unit_space)[2:-1]}',
+                'lower left',
+                pad=1,
+                color='white',
+                frameon=False,
+                size_vertical=1
+            )
+            plt.gca().add_artist(scalebar)
 
-        try:
-            cbar = plt.colorbar(im, cax=plt.gca().cax)
-        except AttributeError:
-            cbar = plt.colorbar()
+        if axes:
+            if xlabel is None:
+                xlabel = self._ax_title[ax_h]
+            if ylabel is None:
+                ylabel = self._ax_title[ax_v]
+            if put_units:
+                xlabel = xlabel + unit_str(unit_space)
+                ylabel = ylabel + unit_str(unit_space)
+            plt.xlabel(xlabel)
+            plt.ylabel(ylabel)
+        else:
+            plt.xticks([])
+            plt.yticks([])
+
+        if colorbar:
+            if tight_cb:
+                ax = plt.gca()
+                fig = plt.gcf()
+                cax = fig.add_axes([
+                    ax.get_position().x1,
+                    ax.get_position().y0,
+                    0.03,
+                    ax.get_position().height
+                ])
+                cbar = plt.colorbar(cax=cax)
+            else:
+                try:
+                    cbar = plt.colorbar(im, cax=plt.gca().cax)
+                except AttributeError:
+                    cbar = plt.colorbar()
+            if label is not None:
+                cbar.set_label(label)
 
         if put_title:
             title = self.snapshot_title(run, title, nml_key, put_time, unit_time)
             plt.title(title)
 
-        if label is not None:
-            cbar.set_label(label)
-
         for i, plot_overlay in enumerate(overlays):
             if plot_overlay in self.overlays:
 
@@ -650,7 +685,6 @@ class Plotter(Aggregator, BaseProcessor):
                     )
                 else:
                     plot_overlay = self.overlays[plot_overlay]
-
             try:
                 plot_overlay(ax_los, im_extent, **overlays_kwargs[i])
             except IndexError: