Refactored and greatly speed up ErrorSleeve

examples/derivativesAndErrorsleeves.py

@@ -11,171 +11,157 @@ from scipy.optimize import curve_fit as fit
 import matplotlib.pyplot as plt
 
 import pyVis
+
 pyVis.styling.defaultStyling()
 
 np.random.seed(42)
 
-'''
+"""
 ===============================================================================
 The 2 derivatives:
     - numDeriv() is just a wrapper for np.gradient(). In numDerivError() there
       is a method implemented that gives the corresponding pointwise error for
       noisy data (just by linear error propagation).
-    - TVRegDiff() is the total-variation regularized derivative. It has much 
+    - TVRegDiff() is the total-variation regularized derivative. It has much
       better properties for noisy data. Only caveat: Due to the global
       optimization involved there is no pointwise error one could plot.
-    
+
     This example is strongly inspired by the example of TVRegDiff.
-    
+
     For demonstration purposes, we create some fake data an find the derivative.
 ===============================================================================
-'''
+"""
 
-dx=.05
-x=np.arange(0,2*np.pi,dx)
-y=np.sin(x)
-noise=np.random.normal(0.0, 0.004, x.shape)
-yNoisy=y+noise
-yPrime=np.cos(x)
+dx = 0.05
+x = np.arange(0, 2 * np.pi, dx)
+y = np.sin(x)
+noise = np.random.normal(0.0, 0.004, x.shape)
+yNoisy = y + noise
+yPrime = np.cos(x)
 
 # Show what we've done.
 # One should note that with this tiny noise, the noisy data is hardly
 # distinguishable from the original data.
-f,a=pyVis.styling.newFig()
-a.plot(x,y,label='orig.',zorder=10)
-a.plot(x,yNoisy,
-           linestyle='None',marker='o',
-           label='noisy')
+f, a = pyVis.styling.newFig()
+a.plot(x, y, label="orig.", zorder=10)
+a.plot(x, yNoisy, linestyle="None", marker="o", label="noisy")
 a.legend()
 
 # Calculate the 2 derivatives.
-tv=pyVis.derivatives.TVRegDiff(yNoisy, 1, 5e-2, dx=dx,
-                               ep=1e-1, scale='small', plotflag=0, diagflag=0)
-naive=pyVis.derivatives.numDeriv(yNoisy,x)
-naiveError=pyVis.derivatives.numDerivError(noise,x)
+tv = pyVis.derivatives.TVRegDiff(
+    yNoisy, 1, 5e-2, dx=dx, ep=1e-1, scale="small", plotflag=0, diagflag=0
+)
+naive = pyVis.derivatives.numDeriv(yNoisy, x)
+naiveError = pyVis.derivatives.numDerivError(noise, x)
 
 # Show the derivatives.
 # While the TV approach struggles only at the edges, the naive derivative is
 # noticeably spread.
-f,a=pyVis.styling.newFig()
-a.plot(x,yPrime,label='exact',zorder=10)
-a.plot(x[1:]-dx/2,tv[1:-1],label='TV',zorder=20)
-a.errorbar(x,naive,
-           naiveError,
-           withLines=False,
-           color=pyVis.styling.colors[2],
-           label='naive')
+f, a = pyVis.styling.newFig()
+a.plot(x, yPrime, label="exact", zorder=10)
+a.plot(x[1:] - dx / 2, tv[1:-1], label="TV", zorder=20)
+a.errorbar(
+    x, naive, naiveError, withLines=False, color=pyVis.styling.colors[2], label="naive"
+)
 
 a.legend()
 
 # Now, we take a short glance at the limit dx --> 0:
 # TV gets even better with smaller artifact zones at the edges while
 # the naive approach is almost useless, now.
-dx2=.01
-x2=np.arange(0,2*np.pi,dx2)
-y2=np.sin(x2)
-noise2=np.random.normal(0.0, 0.004, x2.shape)
-yNoisy2=y2+noise2
-yPrime2=np.cos(x2)
-
-tv2=pyVis.derivatives.TVRegDiff(yNoisy2, 1, 5e-2, dx=dx2,
-                               ep=1e-1, scale='small', plotflag=0, diagflag=0)
-naive2=pyVis.derivatives.numDeriv(yNoisy2,x2)
-naive2Error=pyVis.derivatives.numDerivError(noise2,x2)
-
-f,a=pyVis.styling.newFig()
-a.plot(x2,yPrime2,label='exact',zorder=10)
-a.plot(x2[1:]-dx/2,tv2[1:-1],label='TV',zorder=20)
-a.errorbar(x2,naive2,
-           naive2Error,
-           linestyle='None',marker='o',
-           color=pyVis.styling.colors[2],
-           label='naive')
-
-'''
+dx2 = 0.01
+x2 = np.arange(0, 2 * np.pi, dx2)
+y2 = np.sin(x2)
+noise2 = np.random.normal(0.0, 0.004, x2.shape)
+yNoisy2 = y2 + noise2
+yPrime2 = np.cos(x2)
+
+tv2 = pyVis.derivatives.TVRegDiff(
+    yNoisy2, 1, 5e-2, dx=dx2, ep=1e-1, scale="small", plotflag=0, diagflag=0
+)
+naive2 = pyVis.derivatives.numDeriv(yNoisy2, x2)
+naive2Error = pyVis.derivatives.numDerivError(noise2, x2)
+
+f, a = pyVis.styling.newFig()
+a.plot(x2, yPrime2, label="exact", zorder=10)
+a.plot(x2[1:] - dx / 2, tv2[1:-1], label="TV", zorder=20)
+a.errorbar(
+    x2,
+    naive2,
+    naive2Error,
+    linestyle="None",
+    marker="o",
+    color=pyVis.styling.colors[2],
+    label="naive",
+)
+
+"""
 ===============================================================================
 Errorsleves:
     The current implementation samples the distribution of the function values
     assuming a multi-variate normal distribution of the parameters as given
     by the covariance matrix of the fit parameters.
-    
+
     We try to fit the data from above to get back the cosine.
 ===============================================================================
-'''
+"""
+
+
+def cos(x, a, b, c):
+    return a * np.cos(b * x) + c
 
-def cos(x,a,b,c):
-    return a*np.cos(b*x)+c
 
 # Only use every tenth and not take the errors into account to get a bad fit.
-p,pe=fit(cos,x2[::10],naive2[::10],#sigma=naive2Error[::10]
-         )
+p, pe = fit(cos, x2[::10], naive2[::10],)  # sigma=naive2Error[::10]
 
 # First, print the found parameter values.
-pyVis.printPar(p,pe,names=['Ampl.','Freq.','Offs.'],title='Fit of naive derivative')
+pyVis.printPar(
+    p, pe, names=["Ampl.", "Freq.", "Offs."], title="Fit of naive derivative"
+)
 
 # Plot the result.
-f,a=pyVis.styling.newFig()
-a.plot(x2,yPrime2,label='exact',zorder=10)
-a.plot(x2,cos(x2,*p),label='fit',zorder=10)
+f, a = pyVis.styling.newFig()
+a.plot(x2, yPrime2, label="exact", zorder=10)
+a.plot(x2, cos(x2, *p), label="fit", zorder=10)
 
 # Calculate the error sleeve via bootstrap.
-pyVis.plotting.errorSleeve(cos,p,pe,
-                           method='bootstrap',ax=a,color=pyVis.styling.colors[2],
-                           # plot the sleeve but not the fit itself
-                           plot=1,
-                           # use the following grid for the plot:
-                           grid=x2,
-                           # The random numbers from the bootstrap create noisy
-                           # edges. Smooth them by smearing (i.e. averaging)
-                           # over neighboring points.
-                           smear=20,
-                           # Specify the number of random samples.
-                           N=200,
-                           # Specifiy the label for plotting.
-                           label='bootstrap')
-
-a.set_ylim(-1.1,1.1)
+pyVis.plotting.ErrorSleeve(cos, p, pe).plot(a, x2)
+
+a.set_ylim(-1.1, 1.1)
 
 # Let's do it again to see how it handles constraints.
 # Use a really simple fit and contrain it from below.
-def cos(x,c):
-    return np.cos(x)+c
+def cos(x, c):
+    return np.cos(x) + c
 
-bounds=((0,),(np.inf,))
 
-p,pe=fit(cos,x2[::10],naive2[::10],#sigma=naive2Error[::10],
-         bounds=bounds)
+bounds = ((0,), (np.inf,))
+
+p, pe = fit(cos, x2[::10], naive2[::10], bounds=bounds)  # sigma=naive2Error[::10],
 
 # First, print the found parameter value.
-pyVis.printPar(p,pe,names=['Offs.'],title='Constrained Fit')
+pyVis.printPar(p, pe, names=["Offs."], title="Constrained Fit")
 
 # Plot the result.
-f,a=pyVis.styling.newFig()
-a.plot(x2,yPrime2,label='exact',zorder=10)
-a.plot(x2,cos(x2,*p),label='fit',zorder=10)
+f, a = pyVis.styling.newFig()
+a.plot(x2, yPrime2, label="exact", zorder=10)
+a.plot(x2, cos(x2, *p), label="fit", zorder=10)
 
 # Calculate the error sleeve via bootstrap.
-pyVis.plotting.errorSleeve(cos,p,pe,
-                           method='bootstrap',ax=a,color=pyVis.styling.colors[2],
-                           plot=1,
-                           grid=x2,
-                           rel=.5,
-                           smear=20,
-                           N=1000,
-                           bounds=bounds,
-                           label='bootstrap')
+
+pyVis.plotting.ErrorSleeve(cos, p, pe, bounds=bounds,).plot(a, x2)
 
 # One can see that the cut-off in the distribution is correctly reproduced by
 # errorSleeve().
 
 a.legend()
-a.set_ylim(-1.1,1.1)
+a.set_ylim(-1.1, 1.1)
 
 
-'''
+"""
 ===============================================================================
 Finalize
 ===============================================================================
-'''
+"""
 pyVis.styling.tightLayout()
-plt.show()
\ No newline at end of file
+plt.show()