Fixing the initial notebooks

bezpy/mt/plotting.py

@@ -4,6 +4,7 @@
 
 import numpy as np
 import matplotlib.pyplot as plt
+from matplotlib import container
 from mpl_toolkits.axes_grid1 import make_axes_locatable
 from .utils import apparent_resistivity
 
@@ -16,50 +17,44 @@ def plot_apparent_resistivity(periods, Z, Z_var=None, fig=None, ax_res=None, ax_
 
     resistivity, resistivity_var, phase, phase_var = apparent_resistivity(periods, Z, Z_var)
 
-    # No variance, so this should be just a standard line plot
-    if Z_var is None:
-        ax_res.plot(periods, resistivity[1,:], c='C0', linewidth=1.5, label='xy')
-        ax_res.plot(periods, resistivity[2,:], c='C3', linewidth=1.5, label='yx')
-
-        ax_res.plot(periods, resistivity[0,:], c='C1', linewidth=1.5, label='xx', alpha=alpha)
-        ax_res.plot(periods, resistivity[3,:], c='C2', linewidth=1.5, label='yy', alpha=alpha)
-
-        ax_phase.plot(periods, phase[1,:], c='C0', linewidth=1.5, label='xy')
-        ax_phase.plot(periods, phase[2,:], c='C3', linewidth=1.5, label='yx')
-
-        ax_phase.plot(periods, phase[0,:], c='C1', linewidth=1.5, label='xx', alpha=alpha)
-        ax_phase.plot(periods, phase[3,:], c='C2', linewidth=1.5, label='yy', alpha=alpha)
-
-    # Got the variance, so plot it as scatter points with error bars
-    else:
-        ax_res.errorbar(periods, resistivity[1,:], yerr=2*resistivity_var[1,:], label='xy',
-                    color='C0', marker='o', fmt='o', markersize=markersize, elinewidth=2,
-                    markerfacecolor='none', markeredgecolor='C0', markeredgewidth=2)
-        ax_res.errorbar(periods, resistivity[2,:], yerr=2*resistivity_var[2,:], label='yx',
-                    color='C3', marker='X', fmt='o', markersize=markersize, elinewidth=2)
-
-        ax_res.errorbar(periods, resistivity[0,:], yerr=2*resistivity_var[0,:], label='xx',
-                    color='C1', marker='*', fmt='o', markersize=markersize, elinewidth=2, alpha=alpha)
-        ax_res.errorbar(periods, resistivity[3,:], yerr=2*resistivity_var[3,:], label='yy',
-                    color='C2', marker='^', fmt='o', markersize=markersize, elinewidth=2, alpha=alpha)
-
-        ax_phase.errorbar(periods, phase[1,:], yerr=2*phase_var[1,:],
-                    color='C0', marker='o', fmt='o', markersize=markersize, elinewidth=2,
-                    markerfacecolor='none', markeredgecolor='C0', markeredgewidth=2)
-        ax_phase.errorbar(periods, phase[2,:], yerr=2*phase_var[2,:],
-                    color='C3', marker='X', fmt='o', markersize=markersize, elinewidth=2)
-
-        ax_phase.errorbar(periods, phase[0,:], yerr=2*phase_var[0,:],
-                    color='C1', marker='*', fmt='o', markersize=markersize, elinewidth=2, alpha=alpha)
-        ax_phase.errorbar(periods, phase[3,:], yerr=2*phase_var[3,:],
-                    color='C2', marker='^', fmt='o', markersize=markersize, elinewidth=2, alpha=alpha)
-        # get handles
-        #ax_res.legend(ncol=1, fancybox=False)
-        handles, labels = ax_res.get_legend_handles_labels()
-        # remove the errorbars
-        handles = [h[0] for h in handles]
-        # use them in the legend
-        ax_res.legend(handles, labels, ncol=1, fancybox=False, numpoints=1)
+    # Default entries for each component of the tensor
+    labels = ['xx', 'xy', 'yx', 'yy']
+    colors = ['C1', 'C0', 'C3', 'C2']
+    alphas = [alpha, 1., 1., alpha]
+
+    for i in range(4):
+        if Z_var is None:
+            ax_res.plot(periods, resistivity[i,:], c=colors[i], linewidth=1.5,
+                        label=labels[i], alpha=alphas[i])
+            ax_phase.plot(periods, phase[i,:], c=colors[i], linewidth=1.5,
+                          label=labels[i], alpha=alphas[i])
+        else:
+            # Ignoring nans
+            good_vals = ~np.isnan(resistivity[i,:])
+            x = periods[good_vals]
+            y = resistivity[i,good_vals]
+            yerr = 2*resistivity_var[i,good_vals]
+
+            ax_res.errorbar(x, y, yerr=yerr,
+                            label=labels[i], color=colors[i], marker='o', fmt='o',
+                            markersize=markersize, elinewidth=2, alpha=alphas[i],
+                            markerfacecolor='none', markeredgecolor=colors[i], markeredgewidth=2)
+            y = phase[i,good_vals]
+            yerr = 2*phase_var[i,good_vals]
+            ax_phase.errorbar(x, y, yerr=yerr,
+                              color=colors[i], marker='o', fmt='o',
+                              markersize=markersize, elinewidth=2, alpha=alphas[i],
+                              markerfacecolor='none', markeredgecolor=colors[i], markeredgewidth=2)
+
+    # get handles and remove error bars
+    handles, leg_labels = ax_res.get_legend_handles_labels()
+    handles = [h[0] if isinstance(h, container.ErrorbarContainer) else h for h in handles]
+    # Don't plot line/circle separate if interpolation and points are used
+    # combine the handles together to produce a nicer legend
+    if len(handles) == 8:
+        handles = ((handles[0], handles[4]), (handles[1], handles[5]),
+                   (handles[2], handles[6]), (handles[3], handles[7]))
+    ax_res.legend(handles, leg_labels[:4], ncol=1, fancybox=False, numpoints=1)
 
     ax_res.set_xscale('log')
     ax_res.set_yscale('log')

bezpy/mt/site.py

@@ -44,7 +44,6 @@ def __init__(self, name):
         self.min_period = None
         self.max_period = None
 
-    # FFT of magnetic field time series and the MT site
     def convolve_fft(self, magX, magY, dt=60, extrapolate=0):
         """Convolution in frequency space."""
 
@@ -182,7 +181,7 @@ def plot_apparent_resistivity(self, interp_freqs=None):
                      "\n" +
                      r"Longitude: {lon:4.1f}$^\circ$ ".format(lon=self.longitude) +
                      r"Latitude: {lat:4.1f}$^\circ$".format(lat=self.latitude), size=20)
-        return fig
+        return (fig, ax_res, ax_phase)
 
     def calcZ(self, freqs):
         # extrapolate=1: bandpass filter to only interpolate between data points (no extrapolation)