Add type hints Issue #113 (Verified this time)

trajpy/auxiliar_functions.py

@@ -1,6 +1,7 @@
 import numpy as np
+from typing import Union
 
-def moment_(trajectory, order=2, l_size=np.array([0, 0]), periodic=False):
+def moment_(trajectory: np.ndarray, order: int = 2, l_size: np.ndarray = np.array([0, 0]), periodic: bool = False) -> float:
     """
     Calculates the n-th statistical moment of the trajectory r.
 
@@ -27,7 +28,7 @@ def moment_(trajectory, order=2, l_size=np.array([0, 0]), periodic=False):
         moment[n] = np.sum(np.power(dr, order))
     return np.sum(moment) / (n_points - 1)
 
-def einstein_diffusion_probability(r, D, t):
+def einstein_diffusion_probability(r: Union[float, np.ndarray], D: float, t: float) -> Union[float, np.ndarray]:
     """
     Calculates the probability of a Brownian particle with
     diffusivity D arriving in the position r after a period of time t.
@@ -46,7 +47,7 @@ def einstein_diffusion_probability(r, D, t):
     return probability
 
 
-def unfold (r_old, r, box):
+def unfold (r_old: np.ndarray, r: np.ndarray, box: Union[float, np.ndarray]) -> np.ndarray:
     """
     Removes effects of periodic boundaries on particle trajectories.
     r_old is the configuration at the previous step 

trajpy/gui.py

@@ -12,14 +12,15 @@
 import trajpy.trajpy as tj
 import webbrowser
 import trajpy
+from typing import List, Union, Dict, Optional
 
 print("Tcl Version: {}".format(tk.Tcl().eval('info patchlevel')))
 print("TrajPy Version: {}".format(trajpy.__version__))
 print("Loading TrajPy GUI...")
 
 class trajpy_gui:
 
-    def __init__(self, master):
+    def __init__(self, master: ThemedTk) -> None:
         self.app = master
         self.init_window()
         self.app.resizable(False, False)
@@ -51,7 +52,7 @@ def __init__(self, master):
 
         self.placement()
 
-    def init_window(self):
+    def init_window(self) -> None:
         self.app.title('TrajPy GUI')
         self.app.geometry('600x600')
 
@@ -77,7 +78,7 @@ def init_window(self):
     def client_exit(self):
         exit()
 
-    def About(self):
+    def About(self) -> None:
         self.newWindow = tk.Toplevel(self.app)
         self.newWindow.resizable(False, False)
         self.img = Image.open(os.path.dirname(os.path.realpath(__file__))+'/logo.png')
@@ -100,7 +101,7 @@ def About(self):
         self.email.pack()
 
 
-    def placement(self):
+    def placement(self) -> None:
         self.title.place(x=250, y=10)
         self._version.place(x=300,y=50)
         self.entry.place(x=80, y=100)
@@ -114,7 +115,7 @@ def placement(self):
             self.feats_[button].configure(command=partial(self.select_feat, self.feats_[button]))
             self.feats_[button].place(x=20, y=220 + n * 20)
 
-    def open(self, kind):
+    def open(self, kind: str) -> None:
 
         if kind=="file":
             self.trajectory_list.append(tj.Trajectory(self.path, skip_header=1, delimiter=','))
@@ -128,7 +129,7 @@ def open(self, kind):
                 self.trajectory_list.append(tj.Trajectory(self.path + "/"+ file, skip_header=1, delimiter=','))
 
 
-    def compute_selected(self):
+    def compute_selected(self) -> None:
         '''
             compute the selected features
         '''
@@ -211,7 +212,7 @@ def compute_selected(self):
             self.results.delete(0, 'end')
             self.results.insert(0, ','.join(map(str, [*self.data[0].values()]))+'\n')
 
-    def compute(self):
+    def compute(self) -> None:
         f = tk.filedialog.asksaveasfile(mode='w', defaultextension=".csv")
         for n, trajectory in enumerate(self.trajectory_list):
             self.r = trajectory
@@ -222,7 +223,7 @@ def compute(self):
 
         self.results.insert(0, 'Results saved to {}'.format(f.name))
 
-    def get_file(self, kind):
+    def get_file(self, kind: str) -> None:
         '''
             get a file if `kind`=="file" or file list from a given directory if `kind`=="dir"
         '''
@@ -253,7 +254,7 @@ def get_file(self, kind):
 
         self.open(kind)
 
-    def save_file(self):
+    def save_file(self) -> None:
         f = tk.filedialog.asksaveasfile(mode='w', defaultextension=".csv")
         if f is None:
             return
@@ -262,14 +263,14 @@ def save_file(self):
         f.close()
 
 
-    def show_plot(self):
+    def show_plot(self) -> None:
         self._fig = Figure(figsize=(3, 3), dpi=100)
         self._canvas = FigureCanvasTkAgg(self._fig, master=self.app)
         self._fig.add_subplot(111).plot(self.r._t, self.r._r, ls='-.')
         self._canvas.draw()
         self._canvas.get_tk_widget().place(x=200, y=200)
 
-    def select_feat(self, button):
+    def select_feat(self, button: tk.Widget) -> None:
         if any(button.cget('text') in feature for feature in self.selected_features):
             self.selected_features.remove(button.cget('text'))
             button.deselect()
@@ -285,7 +286,7 @@ def select_feat(self, button):
             self.selected_features.append(button.cget('text'))
             button.select()
 
-    def select_all(self):
+    def select_all(self) -> None:
         self.selected_features = []
 
         for n, button in enumerate(self.feats_):

trajpy/traj_generator.py

@@ -1,7 +1,7 @@
 import numpy as np
+from typing import Union, Tuple
 
-
-def weierstrass_mandelbrot(t, n_displacements, alpha):
+def weierstrass_mandelbrot(t: float, n_displacements: int, alpha: float) -> float:
     """
     Calculates the weierstrass mandelbrot function
 
@@ -25,7 +25,7 @@ def weierstrass_mandelbrot(t, n_displacements, alpha):
     return wsm
 
 
-def anomalous_diffusion(n_steps, n_samples, time_step, alpha):
+def anomalous_diffusion(n_steps: int, n_samples: int, time_step: float, alpha: float) -> Tuple[np.ndarray, np.ndarray]:
     """
     Generates an ensemble of anomalous trajectories.
 
@@ -51,7 +51,7 @@ def anomalous_diffusion(n_steps, n_samples, time_step, alpha):
     return x, y
 
 
-def normal_distribution(u, D, dt):
+def normal_distribution(u: float, D: float, dt: float) -> float:
     """
     This is the steplength probability density function for normal diffusion.
 
@@ -66,7 +66,7 @@ def normal_distribution(u, D, dt):
     return pdf
 
 
-def normal_diffusion(n_steps, n_samples, dx, y0, D, dt):
+def normal_diffusion(n_steps: int, n_samples: int, dx: float, y0: float, D: float, dt: float) -> Tuple[np.ndarray, np.ndarray]:
     """
     Generates an ensemble of normal diffusion trajectories.
 
@@ -99,7 +99,7 @@ def normal_diffusion(n_steps, n_samples, dx, y0, D, dt):
     return x, y
 
 
-def confined_diffusion(radius, n_steps, n_samples, dx, y0, D, dt):
+def confined_diffusion(radius: float, n_steps: int, n_samples: int, dx: float, y0: float, D: float, dt: float) -> Tuple[np.ndarray, np.ndarray]:
     """
     Generates trajectories under confinement.
 
@@ -131,7 +131,7 @@ def confined_diffusion(radius, n_steps, n_samples, dx, y0, D, dt):
     return x, y
 
 
-def superdiffusion(velocity, n_steps, n_samples, y0, dt):
+def superdiffusion(velocity: float, n_steps: int, n_samples: int, y0: float, dt: float) -> Tuple[np.ndarray, np.ndarray]:
     """
     Generates direct diffusion trajectories. 
     Combine pairwise with normal diffusion components.
@@ -156,7 +156,7 @@ def superdiffusion(velocity, n_steps, n_samples, y0, dt):
     return x, y
 
 
-def save_to_file(y, param, path):
+def save_to_file(y: np.ndarray, param: Union[int, float, str], path: str) -> None:
     """
     Saves the trajectories to a file.
 

trajpy/trajpy.py

@@ -1,6 +1,7 @@
 import numpy as np
 from scipy.stats import linregress
 import trajpy.auxiliar_functions as aux
+from typing import Union, List, Dict, Tuple
 import warnings
 
 class Trajectory(object):
@@ -9,7 +10,7 @@ class Trajectory(object):
     as a dummy object for calling its functions or you can initialize
     it with a trajectory array or csv file.
     """
-    def __init__(self, trajectory=np.zeros((1, 2)), box_length=None, **params):
+    def __init__(self, trajectory: np.ndarray=np.zeros((1, 2)), box_length: int = None, **params) -> None:
         """
         Initialization function that can be left blank for using staticmethods.
         It can be initialized with an array with shape (N, dim)
@@ -59,7 +60,7 @@ def __init__(self, trajectory=np.zeros((1, 2)), box_length=None, **params):
         self.velocity_description = None
         self.frequency_spectrum = None
 
-    def compute_features(self):
+    def compute_features(self) -> str:
         """
         Compute every feature for the trajectory saved in self._r.
 
@@ -104,7 +105,7 @@ def compute_features(self):
 
 
     @staticmethod
-    def msd_time_averaged_(spatial_components, tau):
+    def msd_time_averaged_(spatial_components: np.ndarray, tau: Union[np.ndarray, int]) -> np.ndarray:
         """
         calculates the time-averaged mean squared displacement
 
@@ -137,7 +138,7 @@ def msd_time_averaged_(spatial_components, tau):
         return msd
 
     @staticmethod
-    def msd_ensemble_averaged_(spatial_components):
+    def msd_ensemble_averaged_(spatial_components: np.ndarray) -> np.ndarray:
         """
         calculates the ensemble-averaged mean squared displacement
 
@@ -158,7 +159,7 @@ def msd_ensemble_averaged_(spatial_components):
         return msd
 
     @staticmethod
-    def msd_ratio_(msd_ta, n1, n2):
+    def msd_ratio_(msd_ta: np.ndarray, n1: int, n2: int) -> float:
         """
         Ratio of the ensemble averaged mean squared displacements.
 
@@ -178,7 +179,7 @@ def msd_ratio_(msd_ta, n1, n2):
         return msd_ratio
 
     @staticmethod
-    def anomalous_exponent_(msd, time_lag):
+    def anomalous_exponent_(msd: np.ndarray, time_lag: np.ndarray) -> float:
         """
         Calculates the diffusion anomalous exponent
 
@@ -202,7 +203,7 @@ def anomalous_exponent_(msd, time_lag):
         return anomalous_exponent
 
     @staticmethod
-    def fractal_dimension_(trajectory):
+    def fractal_dimension_(trajectory: np.ndarray) -> Tuple[float, float]:
         """
         Estimates the fractal dimension of the trajectory
 
@@ -236,7 +237,7 @@ def fractal_dimension_(trajectory):
         return fractal_dimension, d_max
 
     @staticmethod
-    def gyration_radius_(trajectory):
+    def gyration_radius_(trajectory: np.ndarray) -> Dict[str, Union[np.ndarray, float]]:
         """
         Calculates the gyration radius tensor of the trajectory
 
@@ -271,7 +272,7 @@ def gyration_radius_(trajectory):
         return gyration_radius #dictionary 
 
     @staticmethod
-    def asymmetry_(eigenvalues):
+    def asymmetry_(eigenvalues: np.ndarray) -> float:
         """
         Takes the eigenvalues of the gyration radius tensor 
         to estimate the asymmetry between axis.
@@ -294,7 +295,7 @@ def asymmetry_(eigenvalues):
         return asymmetry
 
     @staticmethod
-    def anisotropy_(eigenvalues):
+    def anisotropy_(eigenvalues: np.ndarray) -> float:
         """
         Calculates the trajectory anisotropy using the eigenvalues of the gyration radius tensor.
 
@@ -315,7 +316,7 @@ def anisotropy_(eigenvalues):
         return anisotropy
 
     @staticmethod
-    def straightness_(trajectory):
+    def straightness_(trajectory: np.ndarray) -> float:
         """
         Estimates how much straight is the trajectory
 
@@ -336,7 +337,7 @@ def straightness_(trajectory):
         return straightness
 
     @staticmethod
-    def kurtosis_(trajectory, eigenvector):
+    def kurtosis_(trajectory: np.ndarray, eigenvector: np.ndarray) -> float:
         """
         We obtain the kurtosis by projecting each position of the trajectory along the main principal eigenvector of the radius of gyration tensor
          :math:`r_i^p = \\mathbf{r} \\cdot \\hat{e}_1` and then calculating the quartic moment
@@ -363,7 +364,7 @@ def kurtosis_(trajectory, eigenvector):
         return kurtosis
 
     @staticmethod
-    def gaussianity_(trajectory):
+    def gaussianity_(trajectory: np.ndarray) -> float:
         """
         measure of how close to a gaussian distribution is the trajectory.
 
@@ -380,7 +381,7 @@ def gaussianity_(trajectory):
         return gaussianity
 
     @staticmethod
-    def confinement_probability_(r0, D, t, N=100):
+    def confinement_probability_(r0: int, D: float, t: float, N: int = 100) -> float:
         """ new
         Estimate the probability of Brownian particle with
         diffusivity :math:`D` being trapped in the interval :math:`[-r0, +r0]` after a period of time t.
@@ -402,7 +403,7 @@ def confinement_probability_(r0, D, t, N=100):
         return 1-probability
 
     @staticmethod
-    def efficiency_(trajectory):
+    def efficiency_(trajectory: np.ndarray) -> float:
         """
             Calculates the efficiency of the movement, a measure that is related to
             the straightness.
@@ -424,7 +425,7 @@ def efficiency_(trajectory):
         return efficiency
 
     @staticmethod
-    def velocity_(position,time):
+    def velocity_(position: np.ndarray, time: np.ndarray) -> np.ndarray:
         """
             Computes the velocity associated with the trajectory
         """
@@ -434,7 +435,7 @@ def velocity_(position,time):
         return velocity
 
     @staticmethod
-    def stationary_velocity_correlation_(velocity,t,taus):
+    def stationary_velocity_correlation_(velocity: np.ndarray, t: np.ndarray, taus: np.ndarray) -> np.ndarray:
         """
             Computes the stationary velocity autocorrelation function by time average
             .. math:
@@ -455,7 +456,7 @@ def stationary_velocity_correlation_(velocity,t,taus):
         return time_averaged_corr_velocity
 
     @staticmethod
-    def green_kubo_(velocity,t,vacf):
+    def green_kubo_(velocity: np.ndarray, t:np.ndarray, vacf:np.ndarray) -> float:
         """
             Computes the generalised Green-Kubo's diffusion constant
             :return diffusivity: diffusion constant obtained by the Green-Kubo relation 
@@ -468,7 +469,7 @@ def green_kubo_(velocity,t,vacf):
         return diffusivity
 
     @staticmethod
-    def velocity_description_(velocity):
+    def velocity_description_(velocity: np.ndarray) -> Dict[str, Union[np.ndarray, float, str] ]:
         """
             Computes the main features of the velocity distribuition: mean, median, mode, variance,
             standard deviation, range, skewness and kurtosis
@@ -504,7 +505,7 @@ def velocity_description_(velocity):
         return velocity_description
 
     @staticmethod
-    def frequency_spectrum_(position, time):
+    def frequency_spectrum_(position: np.ndarray, time: np.ndarray) -> Dict[str, Union[np.ndarray, float] ]:
         '''
             Computes the frequency spectrum for each spatial coordinate by using the Fast Fourier Transform algorithm
             param position: spatial coordinates