plot: upgrade nice plot simple for colors

siria/plot/__init__.py

@@ -6,7 +6,8 @@ from pyrocko.plot import AutoScaler
 def nice_plot_simple(
         ax, xlabel=None, ylabel=None,
         title=None,
-        spine_pos=['left', 'bottom']):
+        spine_pos=['left', 'bottom'],
+        color=None):
 
     if xlabel:
         ax.set_xlabel(xlabel)
@@ -14,7 +15,13 @@ def nice_plot_simple(
     if ylabel:
         ax.set_ylabel(ylabel)
 
-    adjust_spines(ax, spine_pos)
+    if color is not None:
+        ax.xaxis.label.set_color(color)
+        ax.yaxis.label.set_color(color)
+        ax.tick_params(axis='x', colors=color)
+        ax.tick_params(axis='y', colors=color)
+
+    adjust_spines(ax, spine_pos, color=color)
 
     if title:
         ax.set_title(title)
@@ -22,7 +29,7 @@ def nice_plot_simple(
     return ax
 
 
-def adjust_spines(ax, spines):
+def adjust_spines(ax, spines, color=None):
     xlim = ax.get_xlim()
     ylim = ax.get_ylim()
 
@@ -55,6 +62,9 @@ def adjust_spines(ax, spines):
     ax.set_yticks(yticks[(yticks > ylim[0]) & (yticks < ylim[1])])
 
     for loc, spine in ax.spines.items():
+        if color is not None:
+            spine.set_color(color)
+
         if loc not in spines:
             spine.set_visible(False)  # don't draw spine