improved type annotations and doc strings

Author: ptbrown1729

Diff for: mcsim/analysis/camera_noise.py

@@ -154,7 +154,7 @@ def _compute(arr):
 def get_gain_map(means: np.ndarray,
                  vars: np.ndarray,
                  vars_sd: Optional[np.ndarray] = None,
-                 max_mean: float = np.inf) -> (np.ndarray, np.ndarray, np.ndarray):
+                 max_mean: float = np.inf) -> tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]:
     """
     Compute camera pixel gains by finding the best-fit line describing the variance as a function of the mean
 
@@ -193,14 +193,15 @@ def plot_camera_noise_results(gains: np.ndarray,
                               means_err: Optional[np.ndarray] = None,
                               vars_err: Optional[np.ndarray] = None,
                               nbins: int = 600,
-                              **kwargs) -> (list[Figure], list[str]):
+                              **kwargs) -> tuple[list[Figure], list[str]]:
     """
     Display gain curve for single pixel vs. illumination data.
     Additional keyword arguments are passed through to plt.figure()
 
     :param gains: ny x nx
     :param means: nill x ny x nx
     :param vars: nill x ny x nx
+    :param gains_err:
     :param gains: estimated uncertainty (standard dev) of each gain value
     :param means_err: estimated uncertainty (standard dev) of each mean value
     :param vars_err: estimated uncertainty (standard dev) of each variance value

Diff for: mcsim/analysis/dmd_patterns.py

@@ -40,7 +40,7 @@ def get_sim_pattern(dmd_size: Sequence[int, int],
                     vec_b: array,
                     nphases: int,
                     phase_index: int,
-                    geometry: str = "orthogonal") -> (array, array):
+                    geometry: str = "orthogonal") -> tuple[array, array]:
     """
     Convenience function for generating SIM patterns from the tile_patterns() function.
     This function can run either on the CPU or the GPU. If vec_a is a NumPy array, it will
@@ -233,7 +233,7 @@ def double_cell(cell: array,
                 vec_a: array,
                 vec_b: array,
                 na: int = 1,
-                nb: int = 0) -> (np.ndarray, np.ndarray, np.ndarray):
+                nb: int = 0) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
     """
     Create new unit cell by doubling the original one by a factor of na along vec_a and nb along vec_b
 
@@ -290,7 +290,7 @@ def double_cell(cell: array,
 def get_sim_unit_cell(vec_a: array,
                       vec_b: array,
                       nphases: int,
-                      phase_index: int = 0) -> (array, array, array):
+                      phase_index: int = 0) -> tuple[array, array, array]:
     """
     Get unit cell, which can be repeated to form SIM pattern.
 
@@ -339,7 +339,7 @@ def get_sim_unit_cell(vec_a: array,
 
 
 def get_unit_cell(vec_a: array,
-                  vec_b: array) -> (array, array, array):
+                  vec_b: array) -> tuple[array, array, array]:
     """
     Generate a mask which represents one unit cell of a pattern for given vectors.
     This mask is a square array with NaNs at positions outside the unit cell, and
@@ -470,7 +470,7 @@ def line(x, p1, p2): return ((p2[1] - p1[1]) * x + p1[1] * p2[0] - p1[0] * p2[1]
 
 def reduce2cell(point: array,
                 va: array,
-                vb: array) -> (array, array, array):
+                vb: array) -> tuple[array, array, array]:
     """
     Given a vector, reduce it to coordinates within the unit cell.
 
@@ -507,7 +507,7 @@ def convert_cell(cell1: array,
                  va1: array,
                  vb1: array,
                  va2: array,
-                 vb2: array) -> (array, array, array):
+                 vb2: array) -> tuple[array, array, array]:
     """
     Given a unit cell described by vectors va1 and vb2, convert to equivalent description
     in terms of va2, vb2
@@ -550,7 +550,7 @@ def get_minimal_cell(cell: array,
                      x: array,
                      y: array,
                      va: array,
-                     vb: array) -> (array, array, array, array):
+                     vb: array) -> tuple[array, array, array, array, array]:
     """
     Convert to cell using the smallest lattice vectors
 
@@ -636,7 +636,7 @@ def show_cell(v1: np.ndarray,
 
 # determine parameters of SIM patterns
 def get_reciprocal_vects(vec_a: array,
-                         vec_b: array) -> (array, array):
+                         vec_b: array) -> tuple[array, array]:
     """
     Compute the reciprocal vectors for (real-space) lattice vectors vec_a and vec_b.
     If we call the lattice vectors a_i and the reciprocal vectors b_j, then these should be defined such that
@@ -711,7 +711,7 @@ def get_sim_period(vec_a: array,
 
 
 def get_sim_frqs(vec_a: array,
-                 vec_b: array) -> (float, float):
+                 vec_b: array) -> tuple[float, float]:
     """
     Get spatial frequency of SIM pattern constructed from periodicity vectors.
 
@@ -772,7 +772,7 @@ def get_sim_phase(vec_a: array,
 
 
 def ldftfreq(vec_a: np.ndarray,
-             vec_b: np.ndarray) -> (array, array, array):
+             vec_b: np.ndarray) -> tuple[array, array, array]:
     """
     Get the Fourier frequencies for the lattice fourier transform (LDFT). See ldft2() for more details.
     This function is an analog of fftfreq().
@@ -862,7 +862,7 @@ def get_pattern_fourier_component(unit_cell: array,
                                   nphases: int = 3,
                                   phase_index: int = 0,
                                   use_fft_origin: bool = True,
-                                  dmd_size: Optional[Sequence[int]] = None) -> (np.ndarray, np.ndarray):
+                                  dmd_size: Optional[Sequence[int]] = None) -> tuple[np.ndarray, np.ndarray]:
     """
     Get fourier component at f = n * recp_vec_a + m * recp_vec_b.
 
@@ -1038,7 +1038,7 @@ def get_intensity_fourier_components(unit_cell: np.ndarray,
                                      use_fft_origin: bool = True,
                                      include_blaze_correction: bool = True,
                                      dmd_params: Optional[dict] = None) -> \
-        (np.ndarray, np.ndarray, np.ndarray, np.ndarray, np.ndarray):
+        tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray, np.ndarray]:
     """
     Utility function for computing many electric field and intensity components of the Fourier pattern, including the
     effect of the Blaze angle and system numerical aperture
@@ -1501,7 +1501,7 @@ def show_fourier_components(vec_a: np.ndarray,
 
 # Lagrange-Gauss basis reduction
 def reduce_basis(va: array,
-                 vb: array) -> (array, array):
+                 vb: array) -> tuple[array, array]:
     """
     Find the "smallest" set of basis vectors using Lagrange-Gauss basis reduction.
 
@@ -1544,7 +1544,7 @@ def reduce_basis(va: array,
 
 
 def reduce_recp_basis(va: array,
-                      vb: array) -> (array, array):
+                      vb: array) -> tuple[array, array]:
     """
     Compute the shortest pair of reciprocal basis vectors. These vectors may not be dual to the lattice vectors
     in the sense that vi * rsj = delta_{ij}, but they do form a basis for the reciprocal lattice vectors.
@@ -1559,7 +1559,7 @@ def reduce_recp_basis(va: array,
 
 def get_closest_lattice_vec(point: np.ndarray,
                             va: np.ndarray,
-                            vb: np.ndarray) -> (np.ndarray, int, int):
+                            vb: np.ndarray) -> tuple[np.ndarray, int, int]:
     """
     Find the closest lattice vector to point
 
@@ -1610,7 +1610,7 @@ def get_closest_lattice_vec(point: np.ndarray,
 
 def get_closest_recip_vec(recp_point: np.ndarray,
                           va: np.ndarray,
-                          vb: np.ndarray) -> (np.ndarray, int, int):
+                          vb: np.ndarray) -> tuple[np.ndarray, int, int]:
     """
     Find the closest reciprocal lattive vector, f = na * rva + nb * rvb, to a given point in reciprocal space,
     recp_point.
@@ -1819,7 +1819,7 @@ def find_closest_multicolor_set(period: float,
                                 angle_sep_tol: float = 5*np.pi/180,
                                 max_solutions_to_search: int = 20,
                                 pitch: float = 7560.,
-                                minimize_leakage: bool = True) -> (np.ndarray, np.ndarray):
+                                minimize_leakage: bool = True) -> tuple[np.ndarray, np.ndarray]:
     """
     Generate set of SIM patterns for multiple colors with period close to specified value and maximizing distance
     between angles. The patterns are determined such that the diffracted orders will pass through the same positions
@@ -2173,7 +2173,7 @@ def find_rational_approx_angle(angle: float,
 
 def find_allowed_periods(angle: float,
                          nphases: int,
-                         nmax: int) -> (list[float], list[int], list):
+                         nmax: int) -> tuple[array, array, array]:
     """
     Given a DMD pattern with fixed angle, get allowed pattern periods which allow perfect phase shifting for nphases
 
@@ -2295,7 +2295,7 @@ def vects2pattern_data(dmd_size: Sequence[int, int],
                        invert: bool = False,
                        pitch: float = 7560,
                        generate_patterns: bool = True,
-                       geometry: str = "orthogonal") -> (np.ndarray, dict):
+                       geometry: str = "orthogonal") -> tuple[np.ndarray, dict]:
     """
     Generate pattern and useful data (angles, phases, frequencies, reciprocal vectors, ...) from the lattice
     vectors for a given pattern set.
@@ -2504,7 +2504,7 @@ def export_pattern_set(dmd_size: Sequence[int, int],
                        wavelength: float = 1.,
                        save_dir: Union[str, Path] = 'sim_patterns',
                        plot_results: bool = False,
-                       geometry: str = "orthogonal") -> (np.ndarray, dict, Figure):
+                       geometry: str = "orthogonal") -> tuple[np.ndarray, dict, Figure]:
     """
     Export a single set of SIM patterns, i.e. single wavelength, single period
 
@@ -2854,7 +2854,7 @@ def export_otf_test_set(dmd_size: Sequence[int],
                         bvec_max_size: int = 40,
                         phase_index: int = 0,
                         save_dir: Optional[Union[str, Path]] = None,
-                        verbose: bool = True) -> (np.ndarray, dict):
+                        verbose: bool = True) -> tuple[np.ndarray, dict]:
     """
     Generate many SIM-like patterns at a variety of different periods and angles
 

Diff for: mcsim/analysis/fft.py

@@ -2,7 +2,7 @@
 CPU/GPU agnostic FFT functions using our preferred idioms
 """
 
-from typing import Union, Optional
+from typing import Union, Optional, Any
 from collections.abc import Sequence
 import numpy as np
 import numpy.fft as fft_cpu
@@ -26,7 +26,7 @@ def _noshift(arr: array, **kwargs) -> array:
 
 def ft2(m: array,
         axes: tuple[int, int] = (-1, -2),
-        plan: Optional = None,
+        plan: Optional[Any] = None,
         no_cache: bool = False,
         shift: bool = True,
         adjoint: bool = False) -> array:
@@ -52,7 +52,7 @@ def ft2(m: array,
 
 def ift2(m: array,
          axes: tuple[int, int] = (-1, -2),
-         plan: Optional = None,
+         plan: Optional[Any] = None,
          no_cache: bool = False,
          shift: bool = True,
          adjoint: bool = False) -> array:
@@ -72,7 +72,7 @@ def ift2(m: array,
 
 def irft2(m: array,
           axes: tuple[int, int] = (-2, -1),
-          plan: Optional = None,
+          plan: Optional[Any] = None,
           no_cache: bool = False,
           shift: bool = True,
           adjoint: bool = False) -> array:
@@ -120,7 +120,7 @@ def irft2(m: array,
 
 def ft3(m: array,
         axes: tuple[int, int, int] = (-1, -2, -3),
-        plan: Optional = None,
+        plan: Optional[Any] = None,
         no_cache: bool = False,
         shift: bool = True,
         adjoint: bool = False) -> array:
@@ -141,7 +141,7 @@ def ft3(m: array,
 
 def ift3(m: array,
          axes: tuple[int, int, int] = (-1, -2, -3),
-         plan: Optional = None,
+         plan: Optional[Any] = None,
          no_cache: bool = False,
          shift: bool = True,
          adjoint: bool = False) -> array:
@@ -161,7 +161,7 @@ def ift3(m: array,
 
 def ftn(m: array,
         axes: tuple[int],
-        plan: Optional = None,
+        plan: Optional[Any] = None,
         no_cache: bool = False,
         shift: bool = True,
         adjoint: bool = False) -> array:
@@ -209,7 +209,7 @@ def ftn(m: array,
 
 def iftn(m: array,
          axes: tuple[int],
-         plan: Optional = None,
+         plan: Optional[Any] = None,
          no_cache: bool = False,
          shift: bool = True,
          adjoint: bool = False) -> array:

Diff for: mcsim/analysis/field_prop.py

@@ -66,7 +66,7 @@ def get_v(n: array,
 
 
 def frqs2angles(frqs: array,
-                magnitude: float = 1.) -> (array, array):
+                magnitude: float = 1.) -> tuple[array, array]:
     """
     Convert from frequency vectors to angle vectors
 
@@ -1505,7 +1505,7 @@ def gradient(self,
         return phi_fwd[..., :-2, :, :, 0]
 
     def get_estart(self,
-                   inds: Optional[Sequence[int]] = None) -> (array, array):
+                   inds: Optional[Sequence[int]] = None) -> tuple[array, array]:
         if inds is None:
             inds = list(range(self.n_samples))
 

Diff for: mcsim/analysis/fit_dmd_affine.py

@@ -138,7 +138,7 @@ def fit_pattern_peaks(img: np.ndarray,
                       chi_squared_relative_max: float = np.inf,
                       max_position_err: float = 0.1,
                       img_sd: Optional[np.ndarray] = None,
-                      debug: bool = False) -> (np.ndarray, np.ndarray, np.ndarray):
+                      debug: bool = False) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
     """
     Fit peaks of fluorescence image corresponding to affine calibration pattern.
 
@@ -385,6 +385,7 @@ def plot_affine_summary(img: np.ndarray,
     :param vmin_percentile:
     :param vmax_percentile:
     :param gamma: gamma used in displaying image
+    :param cmap:
     :param kwargs: passed through to figure
     :return fig: figure handle to summary figure
     """
@@ -598,7 +599,7 @@ def estimate_xform(img: np.ndarray,
                    vmax_percentile: float = 99.,
                    gamma: float = 1.,
                    figsize: Sequence[float, float] = (16., 12.),
-                   **kwargs) -> (dict, Figure):
+                   **kwargs) -> tuple[dict, Figure]:
     """
     Estimate affine transformation from DMD space to camera image space from an image.
 

Diff for: mcsim/analysis/gauss_beam.py

@@ -83,8 +83,11 @@ def get_q(wo: np.ndarray,
     assuming we use the phasor convention assuming phasor convention exp(-iwt). If using the other phasor convention,
     then take the complex conjugate of q
 
+    :param wo:
     :param z:
-    :param p: [wo, zc, wavelength, n]
+    :param wavelength:
+    :param zc:
+    :param n:
     :return:
     """
     zr = np.pi * wo ** 2 / wavelength * n
@@ -100,6 +103,7 @@ def q2beam_params(qz,
 
     :param qz: complex beam parameter
     :param wavelength:
+    :param n:
     :return R(z), w(z)^2, w(o)^2, z, zr:
     """
 

Diff for: mcsim/analysis/mie_fields.py

@@ -227,7 +227,7 @@ def jn(n: int,
 
 
 # helper function
-def spherical_hn(n: int, z: np.ndarray) -> (np.ndarray, np.ndarray):
+def spherical_hn(n: int, z: np.ndarray) -> tuple[np.ndarray, np.ndarray]:
     hn = spherical_jn(n, z) + 1j * spherical_yn(n, z)
     dhn = spherical_jn(n, z, derivative=True) + 1j * spherical_yn(n, z, derivative=True)
     return hn, dhn
@@ -244,7 +244,7 @@ def mie_efield(wavelength: float,
                beam_phi: float = 0.,
                beam_psi: float = 0.,
                use_gpu: bool = False,
-               **kwargs) -> (array, array, array, array):
+               **kwargs) -> tuple[array, array, array, array]:
     """
     Use miepython to compute multipolar coefficients, and use either scipy special functions
     or GPU accelerated custom CuPy kernels to calculate fields. There are many good references

Diff for: mcsim/analysis/odt_patterns.py

@@ -213,7 +213,7 @@ def loss(c, include_line_penalty=True):
 
 def add_position_multiplexing(centers: list[np.ndarray],
                               frqs: np.ndarray,
-                              ) -> (list[np.ndarray], list[np.ndarray]):
+                              ) -> tuple[list[np.ndarray], list[np.ndarray]]:
     """
     Given a set of spot positions (which may or may not include angle multiplexing), generate a set of
     frequency multiplexed patterns. To do this, start from our original set of patterns, and
@@ -304,7 +304,7 @@ def get_odt_patterns(pupil_positions: Sequence[np.ndarray],
                      pupil_radius_mirrors: float,
                      frqs: Sequence[np.ndarray],
                      phase: Union[Sequence[np.ndarray], float] = 0.,
-                     use_off_mirrors: bool = True) -> (np.ndarray, list[dict]):
+                     use_off_mirrors: bool = True) -> tuple[np.ndarray, list[dict]]:
     """
     Generate DMD patterns from a list of center positions
 

Diff for: mcsim/analysis/optimize.py

@@ -182,7 +182,7 @@ def test_gradient(self,
                       x: array,
                       jind: int = 0,
                       inds: Optional[Sequence[int]] = None,
-                      dx: float = 1e-5) -> (array, array):
+                      dx: float = 1e-5) -> tuple[array, array]:
         """
         Numerically test the gradient computation at a single coordinate of x.