STY: Plottools

Added docstring for resax and fixed minor formatting errors

Author: igresh

tools/plottools.py

@@ -12,7 +12,7 @@ def plot_ellipsdata(
     xaxis="aoi",
     plot_labels=True,
     legend=True,
-    resAx=None
+    resax=None
 ):
     """
     Plots delta and psi values as a function of wavelength or angle of incidence.
@@ -34,31 +34,33 @@ def plot_ellipsdata(
         The default is None.
     objective : refellips.objectiveSE.ObjectiveSE, optional
         Objective (containing model and data) to plot. If the objective is provided,
-        neither model or data should be provided. The default is None.
+        neither model nor data should be provided. The default is None.
     xaxis : String, optional
         Either 'aoi' or 'wavelength'. The default is 'aoi'.
     plot_labels : Bool, optional
         Whether to plot axis labels. The default is True.
     legend : Bool, optional
         Whether to plot the legend. The default is True.
-
+    resax : matplotlib.axes._subplots.AxesSubplot, optional
+        Axis object on which delta and psi residuals. If not provided,
+        no residual will be plotted.
     Returns
     -------
     None.
 
     """
 
-    if objective != None:
+    if objective is not None:
         assert (
-            data == None and model == None
+                data is None and model is None
         ), "If objective is supplied, model and data should not be passed"
         data = objective.data
         model = objective.model
-    elif model != None:
-        assert data != None, "If you supply a model, you must also supply data"
+    elif model is not None:
+        assert data is not None, "If you supply a model, you must also supply data"
     else:
         assert (
-            data != None
+                data is not None
         ), "must supply at least one of data, model or objective"
 
     assert (
@@ -73,7 +75,7 @@ def plot_ellipsdata(
         x = data.aoi
         xlab = "AOI (°)"
 
-        if model != None:
+        if model is not None:
             for wav in unique_wavs:
                 psis, deltas = model(np.c_[np.ones_like(aois) * wav, aois])
                 ax.plot(aois, psis, color="r")
@@ -85,8 +87,9 @@ def plot_ellipsdata(
             np.min(data.wavelength) - 50, np.max(data.wavelength) + 50
         )
         x = data.wavelength
+        xlab = "Wavelength (nm)"
 
-        if model != None:
+        if model is not None:
             wavelength_aois = np.c_[wavs, data.aoi[0]*np.ones_like(wavs)]
             psi, delta = model(wavelength_aois)
             ax.plot(wavs, psi, color="r")
@@ -100,16 +103,19 @@ def plot_ellipsdata(
 #                 ax.plot(np.ones_like(psi) * wavelength, psi, color="r")
 #                 axt.plot(np.ones_like(delta) * wavelength, delta, color="b")
 
-            xlab = "Wavelength (nm)"
+
+    else:
+        assert False, "xaxis must be 'aoi' or 'wavelength'"
 
     # Plot data
     p = ax.scatter(x, data.psi, color="r", alpha=0.5)
     d = axt.scatter(x, data.delta, color="b", alpha=0.5)
 
-    # ax.legend(handles=[p, d], labels=["Psi", "Delta"], loc='center right')
+    if legend:
+        ax.legend(handles=[p, d], labels=["Psi", "Delta"], loc='center right')
 
-    if resAx != None:
-        assert objective != None, 'To plot residuals you must supply an objective'
+    if resax is not None:
+        assert objective is not None, 'To plot residuals you must supply an objective'
         res = objective.residuals()
         numdp = int(len(res)/2)
         psires = res[:numdp]
@@ -125,7 +131,7 @@ def plot_ellipsdata(
         ax.set_xlabel(xlab)
 
         axt.set_ylabel("Delta", color='blue')
-        if resAx != None:
+        if resAx is not None:
             resAx.set(ylabel='error')
 
 def plot_structure(
@@ -151,7 +157,7 @@ def plot_structure(
     objective : refellips.objectiveSE.ObjectiveSE, optional
         Objective (containing model and data) to plot. If the objective is provided,
         structure should not be provided. The default is None.
-    structure : refnx.reflect.structure.Structure
+    structure : refnx.reflect.structure.Structure.
         Structure (which represents the interface) to be plotted. If structure is provided,
         objective should not be provided. The default is None.
     reverse_structure : bool
@@ -164,15 +170,15 @@ def plot_structure(
     None.
 
     """
-    if objective != None:
+    if objective is not None:
         assert (
-            structure == None
+                structure is None
         ), "you must supply either an objective or structure, not both"
         structure = objective.model.structure
         wavelengths = np.unique(objective.data.wavelength)
     else:
         assert (
-            structure != None
+                structure is not None
         ), "you must supply either an objective or structure"
         wavelengths = [658]