correlator class cleaned up

pyerrors/correlators.py

@@ -1,15 +1,12 @@
 import numpy as np
 import autograd.numpy as anp
-#from scipy.special.orthogonal import _IntegerType
+import scipy.linalg
-from .pyerrors import *
+from .pyerrors import Obs, dump_object
 from .fits import standard_fit
-from .linalg import *
+from .linalg import eigh, mat_mat_op
 from .roots import find_root
-from matplotlib import pyplot as plt
+import matplotlib.pyplot as plt
-from matplotlib.ticker import NullFormatter  # useful for `logit` scale
+
-from matplotlib.backends.backend_tkagg import FigureCanvasTkAgg
-#import PySimpleGUI as sg
-import matplotlib
 
 class Corr:
     """The class for a correlator (time dependent sequence of pe.Obs).
@@ -37,7 +34,7 @@ class Corr:
             self.N = 1  # number of smearings
 
         # data_input in the form [np.array(Obs,NxN)]
-        elif all([isinstance(item,np.ndarray) or item==None for item in data_input]) and any([isinstance(item,np.ndarray)for item in data_input]):
+        elif all([isinstance(item, np.ndarray) or item is None for item in data_input]) and any([isinstance(item, np.ndarray) for item in data_input]):
             self.content = data_input
 
             noNull = [a for a in self.content if not (a is None)]  # To check if the matrices are correct for all undefined elements
@@ -105,6 +102,7 @@ class Corr:
 
     def sum(self):
         return np.sqrt(self.N) * self.projected(np.ones(self.N))
+
     # For purposes of debugging and verification, one might want to see a single smearing level. smearing will return a Corr at the specified i,j. where both are integers 0<=i,j<N.
     def smearing(self, i, j):
         if self.N == 1:
@@ -112,15 +110,14 @@ class Corr:
         newcontent = [None if(item is None) else item[i, j] for item in self.content]
         return Corr(newcontent)
 
-
     # Obs and Matplotlib do not play nicely
     # We often want to retrieve x,y,y_err as lists to pass them to something like pyplot.errorbar
     def plottable(self):
         if self.N != 1:
             raise Exception("Can only make Corr[N=1] plottable")  # We could also autoproject to the groundstate or expect vectors, but this is supposed to be a super simple function.
         x_list = [x for x in range(self.T) if (not self.content[x] is None)]
-        y_list = [y[0].value for y in self.content if (not y is None)]
+        y_list = [y[0].value for y in self.content if (y is not None)]
-        y_err_list = [y[0].dvalue for y in self.content if (not y is None)]
+        y_err_list = [y[0].dvalue for y in self.content if (y is not None)]
 
         return x_list, y_list, y_err_list
 
@@ -142,7 +139,6 @@ class Corr:
             raise Exception("Corr could not be symmetrized: No redundant values")
         return Corr(newcontent, prange=self.prange)
 
-
     def anti_symmetric(self):
 
         if self.T % 2 != 0:
@@ -158,7 +154,6 @@ class Corr:
             raise Exception("Corr could not be symmetrized: No redundant values")
         return Corr(newcontent, prange=self.prange)
 
-
     # This method will symmetrice the matrices and therefore make them positive definit.
     def smearing_symmetric(self):
         if self.N > 1:
@@ -167,7 +162,6 @@ class Corr:
         if self.N == 1:
             raise Exception("Trying to symmetrize a smearing matrix, that already has N=1.")
 
-
     # We also include a simple GEVP method based on Scipy.linalg
     def GEVP(self, t0, ts, state=1):
         if (self.content[t0] is None) or (self.content[ts] is None):
@@ -179,11 +173,10 @@ class Corr:
                 Gt[i, j] = self.content[ts][i, j].value
 
         sp_val, sp_vec = scipy.linalg.eig(Gt, G0)
-        sp_vec=sp_vec[:,np.argsort(sp_val)[-state]] #we only want the eigenvector belonging to the selected state
+        sp_vec = sp_vec[:, np.argsort(sp_val)[-state]]  # We only want the eigenvector belonging to the selected state
         sp_vec = sp_vec / np.sqrt(sp_vec @ sp_vec)
         return sp_vec
 
-
     def Eigenvalue(self, t0, state=1):
         G = self.smearing_symmetric()
         G0 = G.content[t0]
@@ -196,16 +189,13 @@ class Corr:
             Gt = G.content[t]
             M = Li @ Gt @ LTi
             eigenvalues = eigh(M)[0]
-            #print(eigenvalues)
             eigenvalue = eigenvalues[-state]
             newcontent.append(eigenvalue)
         return Corr(newcontent)
 
-
     def roll(self, dt):
         return Corr(list(np.roll(np.array(self.content, dtype=object), dt)))
 
-
     def deriv(self, symmetric=True):  # Defaults to symmetric derivative
         if not symmetric:
             newcontent = []
@@ -228,7 +218,6 @@ class Corr:
                 raise Exception('Derivative is undefined at all timeslices')
             return Corr(newcontent, padding_back=1, padding_front=1)
 
-
     def second_deriv(self):
         newcontent = []
         for t in range(1, self.T - 1):
@@ -240,7 +229,6 @@ class Corr:
             raise Exception("Derivative is undefined at all timeslices")
         return Corr(newcontent, padding_back=1, padding_front=1)
 
-
     def m_eff(self, variant='log', guess=1.0):
         """Returns the effective mass of the correlator as correlator object
 
@@ -252,7 +240,7 @@ class Corr:
         """
         if self.N != 1:
             raise Exception('Correlator must be projected before getting m_eff')
-        if variant is 'log':
+        if variant == 'log':
             newcontent = []
             for t in range(self.T - 1):
                 if (self.content[t] is None) or (self.content[t + 1] is None):
@@ -264,7 +252,7 @@ class Corr:
 
             return np.log(Corr(newcontent, padding_back=1))
 
-        elif variant is 'periodic':
+        elif variant == 'periodic':
             newcontent = []
             for t in range(self.T - 1):
                 if (self.content[t] is None) or (self.content[t + 1] is None):
@@ -276,7 +264,8 @@ class Corr:
                 raise Exception('m_eff is undefined at all timeslices')
 
             return Corr(newcontent, padding_back=1)
-        elif variant is 'arccosh':
+
+        elif variant == 'arccosh':
             newcontent = []
             for t in range(1, self.T - 1):
                 if (self.content[t] is None) or (self.content[t + 1] is None) or (self.content[t - 1] is None):
@@ -300,7 +289,6 @@ class Corr:
         # if none is provided, use the range of the corr
         # if this is also not set, use the whole length of the corr (This could come with a warning!)
 
-
         if fitrange is None:
             if self.prange:
                 fitrange = self.prange
@@ -318,8 +306,6 @@ class Corr:
             raise Exception('Unexpected fit result.')
         return result
 
-
-    #we want to quickly get a plateau
     def plateau(self, plateau_range=None, method="fit"):
         if not plateau_range:
             if self.prange:
@@ -335,7 +321,7 @@ class Corr:
                 return a[0]  # At some point pe.standard fit had an issue with single parameter fits. Being careful does not hurt
             return self.fit(const_func, plateau_range)[0]
         elif method in ["avg", "average", "mean"]:
-            returnvalue= np.mean([item[0] for item in self.content[plateau_range[0]:plateau_range[1]+1] if not item is None])
+            returnvalue = np.mean([item[0] for item in self.content[plateau_range[0]:plateau_range[1] + 1] if item is not None])
             returnvalue.gamma_method()
             return returnvalue
 
@@ -408,8 +394,8 @@ class Corr:
         if fit_res:
             x_samples = np.arange(x_range[0], x_range[1] + 1, 0.05)
             ax1.plot(x_samples,
-                    fit_res['fit_function']([o.value for o in fit_res['fit_parameters']], x_samples)
+                     fit_res['fit_function']([o.value for o in fit_res['fit_parameters']], x_samples),
-                    , ls='-', marker=',', lw=2)
+                     ls='-', marker=',', lw=2)
 
         ax1.set_xlabel(r'$x_0 / a$')
         if ylabel:
@@ -418,7 +404,7 @@ class Corr:
 
         handles, labels = ax1.get_legend_handles_labels()
         if labels:
-            legend = ax1.legend()
+            ax1.legend()
         plt.draw()
 
         if save:
@@ -449,9 +435,9 @@ class Corr:
                     content_string += '\t' + element.__repr__()[4:-1]
                 content_string += '\n'
         return content_string
+
     def __str__(self):
         return self.__repr__()
-        #return ("Corr[T="+str(self.T)+" , N="+str(self.N)+" , content="+str([o[0] for o in [o for o in self.content]])+"]")
 
     # We define the basic operations, that can be performed with correlators.
     # While */+- get defined here, they only work for Corr*Obs and not Obs*Corr.
@@ -524,14 +510,11 @@ class Corr:
 
             if all([item is None for item in newcontent]):
                 raise Exception("Division returns completely undefined correlator")
-
-
-
             return Corr(newcontent)
 
         elif isinstance(y, Obs):
             if y.value == 0:
-                raise Exception("Division by Zero will return undefined correlator")
+                raise Exception('Division by zero will return undefined correlator')
             newcontent = []
             for t in range(self.T):
                 if (self.content[t] is None):
@@ -542,16 +525,16 @@ class Corr:
 
         elif isinstance(y, int) or isinstance(y, float):
             if y == 0:
-                raise Exception("Division by Zero will return undefined correlator")
+                raise Exception('Division by zero will return undefined correlator')
             newcontent = []
             for t in range(self.T):
                 if (self.content[t] is None):
                     newcontent.append(None)
                 else:
                     newcontent.append(self.content[t] / y)
-            return Corr(newcontent,prange= self.prange if hasattr(self,"prange") else None)
+            return Corr(newcontent, prange=self.prange)
         else:
-            raise TypeError("Corr / wrong type")
+            raise TypeError('Corr / wrong type')
 
     def __neg__(self):
         newcontent = [None if (item is None) else -1. * item for item in self.content]
@@ -565,11 +548,11 @@ class Corr:
             newcontent = [None if (item is None) else item**y for item in self.content]
             return Corr(newcontent, prange=self.prange)
         else:
-            raise TypeError("type of exponent not supported")
+            raise TypeError('Type of exponent not supported')
 
     def __abs__(self):
         newcontent = [None if (item is None) else np.abs(item) for item in self.content]
-        return Corr(newcontent,prange= self.prange if hasattr(self,"prange") else None)
+        return Corr(newcontent, prange=self.prange)
 
     # The numpy functions:
     def sqrt(self):
@@ -577,123 +560,65 @@ class Corr:
 
     def log(self):
         newcontent = [None if (item is None) else np.log(item) for item in self.content]
-        return Corr(newcontent,prange= self.prange if hasattr(self,"prange") else None)
+        return Corr(newcontent, prange=self.prange)
 
     def exp(self):
         newcontent = [None if (item is None) else np.exp(item) for item in self.content]
-        return Corr(newcontent,prange= self.prange if hasattr(self,"prange") else None)
+        return Corr(newcontent, prange=self.prange)
+
+    def _apply_func_to_corr(self, func):
+        newcontent = [None if (item is None) else func(item) for item in self.content]
+        for t in range(self.T):
+            if newcontent[t] is None:
+                continue
+            if np.isnan(np.sum(newcontent[t]).value):
+                newcontent[t] = None
+        if all([item is None for item in newcontent]):
+            raise Exception('Operation returns undefined correlator')
+        return Corr(newcontent)
 
     def sin(self):
-        newcontent=[None if (item is None) else np.sin(item) for item in self.content]
+        return self._apply_func_to_corr(np.sin)
-        return Corr(newcontent)
 
     def cos(self):
-        newcontent=[None if (item is None) else np.cos(item) for item in self.content]
+        return self._apply_func_to_corr(np.cos)
-        return Corr(newcontent)
 
     def tan(self):
-        newcontent=[None if (item is None) else np.tan(item) for item in self.content]
+        return self._apply_func_to_corr(np.tan)
-
-        for t in range(self.T):
-                if newcontent[t] is None:
-                    continue
-                if np.isnan(np.sum(newcontent[t]).value):
-                    newcontent[t]=None
-        if all([item is None for item in newcontent]):
-                raise Exception("Operation returns completely undefined correlator")
-
-        return Corr(newcontent)
 
     def sinh(self):
-        newcontent=[None if (item is None) else np.sinh(item) for item in self.content]
+        return self._apply_func_to_corr(np.sinh)
-        return Corr(newcontent)
 
     def cosh(self):
-        newcontent=[None if (item is None) else np.cosh(item) for item in self.content]
+        return self._apply_func_to_corr(np.cosh)
-        return Corr(newcontent)
 
     def tanh(self):
-        newcontent=[None if (item is None) else np.tanh(item) for item in self.content]
+        return self._apply_func_to_corr(np.tanh)
-        for t in range(self.T):
-                if newcontent[t] is None:
-                    continue
-                if np.isnan(np.sum(newcontent[t]).value):
-                    newcontent[t]=None
-        if all([item is None for item in newcontent]):
-                raise Exception("Operation returns completely undefined correlator")
-        return Corr(newcontent)
 
     def arcsin(self):
-        newcontent=[None if (item is None) else np.arcsin(item) for item in self.content]
+        return self._apply_func_to_corr(np.arcsin)
-        for t in range(self.T):
-                if newcontent[t] is None:
-                    continue
-                if np.isnan(np.sum(newcontent[t]).value):
-                    newcontent[t]=None
-        if all([item is None for item in newcontent]):
-                raise Exception("Operation returns completely undefined correlator")
-        return Corr(newcontent)
 
     def arccos(self):
-        newcontent=[None if (item is None) else np.arccos(item) for item in self.content]
+        return self._apply_func_to_corr(np.arccos)
-        for t in range(self.T):
-                if newcontent[t] is None:
-                    continue
-                if np.isnan(np.sum(newcontent[t]).value):
-                    newcontent[t]=None
-        if all([item is None for item in newcontent]):
-                raise Exception("Operation returns completely undefined correlator")
-        return Corr(newcontent)
 
     def arctan(self):
-        newcontent=[None if (item is None) else np.arctan(item) for item in self.content]
+        return self._apply_func_to_corr(np.arctan)
-        for t in range(self.T):
-                if newcontent[t] is None:
-                    continue
-                if np.isnan(np.sum(newcontent[t]).value):
-                    newcontent[t]=None
-        if all([item is None for item in newcontent]):
-                raise Exception("Operation returns completely undefined correlator")
-        return Corr(newcontent)
 
     def arcsinh(self):
-        newcontent=[None if (item is None) else np.arcsinh(item) for item in self.content]
+        return self._apply_func_to_corr(np.arcsinh)
-        for t in range(self.T):
-                if newcontent[t] is None:
-                    continue
-                if np.isnan(np.sum(newcontent[t]).value):
-                    newcontent[t]=None
-        if all([item is None for item in newcontent]):
-                raise Exception("Operation returns completely undefined correlator")
-        return Corr(newcontent)
 
     def arccosh(self):
-        newcontent=[None if (item is None) else np.arccosh(item) for item in self.content]
+        return self._apply_func_to_corr(np.arccosh)
-        for t in range(self.T):
-                if newcontent[t] is None:
-                    continue
-                if np.isnan(np.sum(newcontent[t]).value):
-                    newcontent[t]=None
-        if all([item is None for item in newcontent]):
-                raise Exception("Operation returns completely undefined correlator")
-        return Corr(newcontent)
 
     def arctanh(self):
-        newcontent=[None if (item is None) else np.arctanh(item) for item in self.content]
+        return self._apply_func_to_corr(np.arctanh)
-        for t in range(self.T):
-                if newcontent[t] is None:
-                    continue
-                if np.isnan(np.sum(newcontent[t]).value):
-                    newcontent[t]=None
-        if all([item is None for item in newcontent]):
-                raise Exception("Operation returns completely undefined correlator")
-        return Corr(newcontent)
 
-
+    # Right hand side operations (require tweak in main module to work)
-    #right hand side operations (require tweak in main module to work)
     def __rsub__(self, y):
         return -self + y
+
     def __rmul__(self, y):
         return self * y
+
     def __radd__(self, y):
         return self + y