Force compilation of computing functions for round and ellipse 'precise' shaped objects.
-
-
verbose_report : bool, optional
- Flag for printing out the elapsed for compilation time. The default is False.
+
Clean locally cached files created by numba after compilation of computation methods from the 'calc_psfs_numba' module.
Returns
-- None or tuple with 2 used ZernPSF instances depending on 'numba_installed' flag.
+- bool: Flag if local cache is assumed (at least 6 cached files successfully removed) to be cleaned.
diff --git a/pyproject.toml b/pyproject.toml
index b9dc177..bc21129 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -2,26 +2,31 @@
name = "zernpy"
dynamic = ["version"] # Discovering of the version by the setuptools (see below: in the 'file' or 'attr')
authors = [
- {name = "Sergei Klykov"},
- {email = "sergej.klykow@gmail.com"}
+ { name = "Sergei Klykov" },
+ { email = "sergej.klykow@gmail.com" }
]
description = "Calculation of real Zernike polynomials values, associated PSFs, plotting of their profiles in polar coordinates"
readme = "README.md"
-# license = {file = "LICENSE"} # includes the whole text in METADATA, maybe not so convienient
requires-python = ">=3.8"
classifiers = [
"Programming Language :: Python :: 3",
"Operating System :: OS Independent",
]
+
dependencies = [
"numpy",
"matplotlib",
- "scipy",
+ "scipy",
]
+[project.optional-dependencies]
+performance = ["numba>=0.57.1"]
+dev = ["ruff", "mypy", "pytest"]
+all = ["numba>=0.57.1"]
+
keywords = ["zernike-polynomials", "zernike-psf"]
[project.urls] # METADATA in wheel file will represent the data below by using pip show zernpy -v
-# But the the links themselves are mapped / parced by the PyPi website
+# But the links themselves are mapped / parsed by the PyPi website
"Homepage" = "https://pypi.org/project/zernpy/"
"Repository" = "https://github.com/sklykov/zernpy/"
"Bug Tracker" = "https://github.com/sklykov/zernpy/issues/"
@@ -31,14 +36,59 @@ keywords = ["zernike-polynomials", "zernike-psf"]
[tool.setuptools.packages.find] # Manifest.in file is only required for adding some package-data
where = ["src"]
include = ["zernpy*"]
-exclude = ["tests"]
+exclude = ["zernpy.tests*"] # only packages should be specified, folders - exclude in MANIFEST.in
+namespaces = false # forbid treating a folder as a module without explicit __init__.py presence
-[tool.setuptools.exclude-package-data]
-zernpy = ["*.png"] # should exclude image files from distribution
+[tool.setuptools]
+include-package-data = false # forbid to make a confusion on the src/zernpy/tests folder presence
+# confusion is produced because it is an importable module
[tool.setuptools.dynamic]
-version = {attr = "zernpy.__version__"} # Variable set in the __init__.py
+version = { attr = "zernpy.__version__" } # Variable set in the __init__.py
[build-system]
requires = ["setuptools>=75.0"]
build-backend = "setuptools.build_meta"
+
+[tool.mypy]
+python_version = "3.11"
+warn_unused_configs = true # mypy will warn if any of options below not valid
+warn_redundant_casts = true # warn about unnecessary cast(): from typing import cast
+warn_unused_ignores = true # warn about unused comments '# type: ignore'
+show_error_codes = true # enables error code in a report
+pretty = true # better formatted report in a console
+strict_equality = true # checks proper usage of ==, !=, in
+no_implicit_optional = true # throws error on: def f(x: int = None): ...
+# Note: flags below introduces too strict type checks
+# disallow_untyped_defs = true # all parameters and return types must be fully annotated
+check_untyped_defs = true # mypy checks untyped expressions and tries to infer types
+# warn_return_any = true # Any couldn't be used as return type
+# disallow_any_generics = true # forbid bare generics like list, dict, tuple... -> instead List[float]
+exclude = [
+ '^tests/',
+ '^build/',
+ '^dist/',
+ '^src/zernpy/tests/'
+]
+[[tool.mypy.overrides]] # for ignoring warning about not checked imported libraries
+module = ["numba", "joblib"]
+ignore_missing_imports = true
+# [[tool.mypy.overrides]] # ignore any errors in specific test files - retained for ref.
+# module = ["zernpy.run_zernpsf_as_main"]
+# ignore_errors = true
+
+[tool.ruff]
+line-length = 145 # my own preference
+indent-width = 4
+target-version = "py311"
+[tool.ruff.lint]
+select = ["E4", "E7", "E9", "F", "I", "B", "SIM"]
+ignore = ["E702", "E402", # used in the IDE
+ "SIM108", # many ternany operator replacements in a code
+ "SIM113", # potentially can break logic in a loop
+]
+[tool.ruff.format]
+quote-style = "double"
+indent-style = "space"
+line-ending = "auto"
+docstring-code-format = true
diff --git a/src/zernpy/__init__.py b/src/zernpy/__init__.py
index 02f60ba..a537efa 100644
--- a/src/zernpy/__init__.py
+++ b/src/zernpy/__init__.py
@@ -8,21 +8,11 @@
"""
-__version__ = "0.1.0" # Straightforward way of specifying package version and including it to the package attributes
+__version__ = "0.1.1" # Straightforward way of specifying package version and including it to the package attributes
-if __name__ == "__main__":
- # use absolute imports for importing as module
- __all__ = ['zernikepol', 'zernpsf'] # for specifying from zernpy import * if package imported from some script
-elif __name__ == "zernpy":
- pass # do not add module "zernikepol" to __all__ attribute, because it demands to construct explicit path
+# Univesal logic for making all main classes and functions available after calling 'from project import *'
+from .zernikepol import ZernPol, fit_polynomials, fit_polynomials_vectors, generate_phases_image, generate_polynomials, generate_random_phases
+from .zernpsf import ZernPSF, clean_zernpy_cache, force_get_psf_compilation
-# Automatically bring the main class and some methods to the name space when one of import command is used commands:
-# 1) from zernpy import ZernPol, ... functions; 2) from zernpy import *
-if __name__ != "__main__" and __name__ != "__mp_main__":
- from .zernikepol import ZernPol # main class auto export on the import call of the package
- # functions auto export - when everything imported from the module
- from .zernikepol import generate_polynomials, fit_polynomials, generate_random_phases, fit_polynomials_vectors, generate_phases_image
- from .zernpsf import ZernPSF # class for ZernPSF auto export on the import call of the package
- from .zernpsf import force_get_psf_compilation # function for precompile functions by numba library
- __all__ = ["ZernPol", "generate_polynomials", "fit_polynomials", "generate_random_phases", "fit_polynomials_vectors",
- "generate_phases_image", "ZernPSF", "force_get_psf_compilation"]
+__all__ = ["ZernPol", "generate_polynomials", "fit_polynomials", "generate_random_phases", "fit_polynomials_vectors",
+ "generate_phases_image", "ZernPSF", "force_get_psf_compilation", "clean_zernpy_cache"]
diff --git a/src/zernpy/calculations/calc_psfs.py b/src/zernpy/calculations/calc_psfs.py
index 406d067..eb9745b 100644
--- a/src/zernpy/calculations/calc_psfs.py
+++ b/src/zernpy/calculations/calc_psfs.py
@@ -7,19 +7,18 @@
"""
# %% Global imports
-import numpy as np
-import matplotlib.pyplot as plt
-from scipy.special import jv
-from pathlib import Path
-import warnings
-from math import sqrt, pi, e
-import time
-from scipy.ndimage import convolve
import json
-from matplotlib.patches import Circle
+import time
+import warnings
from functools import partial
-from typing import Union
-# import time
+from math import e, pi, sqrt
+from pathlib import Path
+from typing import Optional, Union
+
+import matplotlib.pyplot as plt
+import numpy as np
+from scipy.ndimage import convolve
+from scipy.special import jv
# Testing the parallelization with joblib. Native Pool.map() tested and results transferred to the collection_numCalc repo
try:
@@ -29,21 +28,17 @@
joblib_installed = False
# %% Local (package-scoped) imports
-if __name__ == "__main__" or __name__ == Path(__file__).stem or __name__ == "__mp_main__":
- from calc_zernike_pol import define_orders
-else:
- from .calc_zernike_pol import define_orders
+from .calc_zernike_pol import define_orders
# %% Module parameters
__docformat__ = "numpydoc"
um_char = "\u00B5" # Unicode char code for micrometers
lambda_char = "\u03BB" # Unicode char code for lambda (wavelength)
pi_char = "\u03C0" # Unicode char code for pi
-# print(um_char, lambda_char, pi_char) # check the codes for characters from above
# %% Reference - Airy profile for Z(0, 0)
-def airy_ref_pattern(r: float):
+def airy_ref_pattern(r: float) -> float:
"""
Return Airy pattern radial function J1(r)/r.
@@ -59,17 +54,15 @@ def airy_ref_pattern(r: float):
"""
r = round(r, 12)
- if r == 0.0:
- ratio = jv(1, 1E-11)/1E-11
- else:
- ratio = jv(1, r)/r
+ ratio = jv(1, 1E-11)/1E-11 if r == 0.0 else jv(1, r)/r
# NOTE that the values produced by exact equation below is off with the direct computation of diffraction integral
# apporoximate coefficient is 0.986711 for a central point. Most likely, the difference because of numerical integration
return 4.0*pow(ratio, 2)
# %% PSF pixel value calc.
-def diffraction_integral_r(zernike_pol, alpha: float, phi: float, p: Union[float, np.ndarray], theta: float, r: float) -> np.array:
+def diffraction_integral_r(zernike_pol, alpha: float, phi: float, p: Union[float, np.ndarray], theta: float,
+ r: float) -> Union[complex, np.ndarray]:
"""
Diffraction integral function for the formed image point (see the references as the sources of the equation).
@@ -263,17 +256,16 @@ def get_psf_point_r_parallel(zernike_pol, r: float, theta: float, alpha: float,
# Vectorized or parallelized form of for loop for even and odd phi-s
if not joblib_installed or paralleljobs is None:
- even_sums = [radial_integral(zernike_pol, r, theta, i*h_phi, alpha, n_int_r_points) for i in range(2, n_int_phi_points-2, 2)]
- even_sums = np.asarray(even_sums); even_sum = np.sum(even_sums)
- odd_sums = [radial_integral(zernike_pol, r, theta, i*h_phi, alpha, n_int_r_points) for i in range(1, n_int_phi_points-1, 2)]
- odd_sums = np.asarray(odd_sums); odd_sum = np.sum(odd_sums)
+ even_sums = np.asarray([radial_integral(zernike_pol, r, theta, i*h_phi, alpha, n_int_r_points)
+ for i in range(2, n_int_phi_points-2, 2)])
+ odd_sums = np.asarray([radial_integral(zernike_pol, r, theta, i*h_phi, alpha, n_int_r_points)
+ for i in range(1, n_int_phi_points-1, 2)])
+ even_sum = np.sum(even_sums); odd_sum = np.sum(odd_sums)
else:
if paralleljobs is not None and isinstance(paralleljobs, Parallel):
- even_sums = paralleljobs(delayed(radial_integral_fixed_args)(i*h_phi) for i in range(2, n_int_phi_points-2, 2))
- odd_sums = paralleljobs(delayed(radial_integral_fixed_args)(i*h_phi) for i in range(1, n_int_phi_points-1, 2))
- # even_sum = sum(even_sums, start=even_sum); odd_sum = sum(even_sums, start=odd_sum)
- even_sums = np.asarray(even_sums); even_sum = np.sum(even_sums); odd_sums = np.asarray(odd_sums); odd_sum = np.sum(odd_sums)
-
+ even_sums = np.asarray(paralleljobs(delayed(radial_integral_fixed_args)(i*h_phi) for i in range(2, n_int_phi_points-2, 2)))
+ odd_sums = np.asarray(paralleljobs(delayed(radial_integral_fixed_args)(i*h_phi) for i in range(1, n_int_phi_points-1, 2)))
+ even_sum = np.sum(even_sums); odd_sum = np.sum(odd_sums)
# Simpson integration rule implementation
yA = radial_integral(zernike_pol, r, theta, 0.0, alpha, n_int_r_points)
yB = radial_integral(zernike_pol, r, theta, 2.0*pi, alpha, n_int_r_points)
@@ -306,13 +298,13 @@ def get_kernel_size(zernike_pol, len2pixels: float, alpha: float, wavelength: fl
Estimated kernel size.
"""
- (m, n) = define_orders(zernike_pol) # get polynomial orders
+ m, n = define_orders(zernike_pol) # get polynomial orders
size_ext = 0 # additional size depending on some parameters below
- if m == 0 and n == 0: # Airy
+ if m == 0 and n == 0: # Airy profile
if 0.25 < NA < 1.0:
- multiplier = 5.0*(1.0 - NA) + 1.5
+ multiplier = 5.0*(1.0 - NA) + 1.5 + alpha
else:
- multiplier = 4.5 + 2.5*sqrt(1.0 / NA)
+ multiplier = 4.5 + 2.5*sqrt(1.0 / NA) + 1.25*alpha
else:
multiplier = 1.25*sqrt(n) # Enlarge kernel size according to the provided radial order n
if abs(m) > 0 and n % 2 != 0: # Enlarge kernel size for the not symmetrical orders
@@ -336,7 +328,7 @@ def get_kernel_size(zernike_pol, len2pixels: float, alpha: float, wavelength: fl
def get_psf_kernel(zernike_pol, len2pixels: float, alpha: float, wavelength: float, NA: float, n_int_r_points: int = 320,
- n_int_phi_points: int = 300, show_kernel: bool = False, fig_title: str = None, normalize_values: bool = False,
+ n_int_phi_points: int = 300, show_kernel: bool = False, fig_title: Optional[str] = None, normalize_values: bool = False,
airy_pattern: bool = False, kernel_size: int = 0, test_parallel: bool = False, fig_id: str = "",
test_vectorized: bool = False, suppress_warns: bool = False, verbose: bool = False) -> np.ndarray:
"""
@@ -386,7 +378,7 @@ def get_psf_kernel(zernike_pol, len2pixels: float, alpha: float, wavelength: flo
Matrix with PSF values.
"""
- (m, n) = define_orders(zernike_pol) # get polynomial orders
+ m, n = define_orders(zernike_pol) # get polynomial orders
# Convert provided absolute value of Zernike expansion coefficient (in um) into fraction of wavelength
alpha /= wavelength; k = 2.0*pi/wavelength # Calculate angular frequency (k)
# Empirical estimation of the sufficient size for the kernel
@@ -410,9 +402,9 @@ def get_psf_kernel(zernike_pol, len2pixels: float, alpha: float, wavelength: flo
# Check that the calibration coefficient is sufficient for calculation
pixel_size_nyquist = 0.5*0.61*wavelength/NA
if len2pixels > pixel_size_nyquist and not suppress_warns:
- __warn_message = f"Provided calibration coefficient {len2pixels} {um_char}/pixels isn't sufficient enough"
+ __warn_message = f"\nProvided calibration coefficient {len2pixels} {um_char}/pixels isn't sufficient enough"
__warn_message += f" (defined by the relation between Nyquist freq. and the optical resolution: 0.61{lambda_char}/NA)"
- warnings.warn(__warn_message)
+ warnings.warn(__warn_message, stacklevel=2)
# Calculate the PSF kernel for usage in convolution operation
# mean_time_integration = 0.0; n = 0
if not joblib_installed or not test_parallel:
@@ -475,7 +467,7 @@ def get_psf_kernel(zernike_pol, len2pixels: float, alpha: float, wavelength: flo
if kernel_border_max > np.max(kernel)/20.0 and not suppress_warns:
__warn_message = (f"\nThe calculated size for PSF ({size}) isn't sufficient for its proper representation, "
+ "because the maximum value on the kernel border is bigger than 5% of maximum overall kernel")
- warnings.warn(__warn_message)
+ warnings.warn(__warn_message, stacklevel=2)
# Plotting the calculated kernel
if show_kernel:
if airy_pattern:
@@ -484,7 +476,7 @@ def get_psf_kernel(zernike_pol, len2pixels: float, alpha: float, wavelength: flo
plt.figure(fig_title, figsize=(6, 6))
else:
plt.figure(f"{(m, n)} {zernike_pol.get_polynomial_name(True)}: {round(alpha, 2)}*wavelength {fig_id}", figsize=(6, 6))
- plt.imshow(kernel, cmap=plt.cm.viridis, origin='upper'); plt.tight_layout()
+ plt.imshow(kernel, cmap=plt.colormaps["viridis"], origin='upper'); plt.tight_layout()
return kernel
@@ -592,10 +584,11 @@ def get_psf_point_r_pols(polynomials, amplitudes: np.ndarray, r: float, theta: f
"""
h_phi = 2.0*pi/n_int_phi_points; even_sum = 0.0j; odd_sum = 0.0j
# Vectorized or parallelized form of for loop for even and odd phi-s
- even_sums = [radial_integral_pols(polynomials, amplitudes, r, theta, i*h_phi, n_int_r_points) for i in range(2, n_int_phi_points-2, 2)]
- even_sums = np.asarray(even_sums); even_sum = np.sum(even_sums)
- odd_sums = [radial_integral_pols(polynomials, amplitudes, r, theta, i*h_phi, n_int_r_points) for i in range(1, n_int_phi_points-1, 2)]
- odd_sums = np.asarray(odd_sums); odd_sum = np.sum(odd_sums)
+ even_sums = np.asarray([radial_integral_pols(polynomials, amplitudes, r, theta, i*h_phi, n_int_r_points)
+ for i in range(2, n_int_phi_points-2, 2)])
+ odd_sums = np.asarray([radial_integral_pols(polynomials, amplitudes, r, theta, i*h_phi, n_int_r_points)
+ for i in range(1, n_int_phi_points-1, 2)])
+ even_sum = np.sum(even_sums); odd_sum = np.sum(odd_sums)
# Simpson integration rule implementation
yA = radial_integral_pols(polynomials, amplitudes, r, theta, 0.0, n_int_r_points)
yB = radial_integral_pols(polynomials, amplitudes, r, theta, 2.0*pi, n_int_r_points)
@@ -661,7 +654,7 @@ def get_psf_kernel_zerns(polynomials, amplitudes: np.ndarray, len2pixels: float,
if len2pixels > pixel_size_nyquist and not suppress_warns:
__warn_message = f"\nProvided calibration coefficient {len2pixels} {um_char}/pixels isn't sufficient enough"
__warn_message += f" (defined by the relation between Nyquist freq. and the optical resolution: 0.61{lambda_char}/NA)"
- warnings.warn(__warn_message)
+ warnings.warn(__warn_message, stacklevel=2)
# Verbose info preparation
if verbose:
calculated_points = 0 # for explicit showing of performance
@@ -701,7 +694,7 @@ def get_psf_kernel_zerns(polynomials, amplitudes: np.ndarray, len2pixels: float,
if kernel_border_max > np.max(kernel)/20.0 and not suppress_warns:
__warn_message = (f"\nThe calculated size for PSF ({size}) isn't sufficient for its proper representation, "
+ "because the maximum value on the kernel border is bigger than 5% of maximum overall kernel")
- warnings.warn(__warn_message)
+ warnings.warn(__warn_message, stacklevel=2)
return kernel
@@ -723,8 +716,9 @@ def convolute_img_psf(img: np.ndarray, psf_kernel: np.ndarray, scale2original: b
Result of convolution (used scipy.ndimage.convolve).
"""
- img_type = img.dtype; img = np.copy(img) # get the image type and copy its content to a new container
- convolved_img = convolve(np.float64(img), psf_kernel, mode='reflect'); conv_coeff = np.sum(psf_kernel) # convolution using scipy
+ img_type = img.dtype; img = np.copy(img).astype(np.float64) # get the image type and copy its content to a new container
+ convolved_img: np.ndarray
+ convolved_img = convolve(img, psf_kernel, mode='reflect'); conv_coeff = np.sum(psf_kernel) # convolution using scipy
if conv_coeff > 0.0:
convolved_img /= conv_coeff # correct the convolution result by dividing to the kernel sum
if scale2original:
@@ -736,8 +730,8 @@ def convolute_img_psf(img: np.ndarray, psf_kernel: np.ndarray, scale2original: b
# %% Save and read the calculated PSF matrices
def save_psf(psf_kernel: np.ndarray, NA: float, wavelength: float, pixel_size: float, expansion_coefficient: Union[float, np.ndarray],
- kernel_size: int, n_int_points_r: int, n_int_points_phi: int, zernike_pol, folder_path: str = None,
- overwrite: bool = True, additional_file_name: str = None) -> str:
+ kernel_size: int, n_int_points_r: int, n_int_points_phi: int, zernike_pol, folder_path: Optional[str] = None,
+ overwrite: bool = True, additional_file_name: Optional[str] = None) -> str:
"""
Save the calculated PSF kernel along with the used for the calculation parameters.
@@ -775,8 +769,10 @@ def save_psf(psf_kernel: np.ndarray, NA: float, wavelength: float, pixel_size: f
"""
single_pol_used = False; osa_indices = [] # flag for saving different parameters
+ if isinstance(expansion_coefficient, np.ndarray):
+ ampls = expansion_coefficient.copy().tolist()
if not hasattr(zernike_pol, '__len__'):
- (m, n) = define_orders(zernike_pol); single_pol_used = True # get polynomial orders
+ m, n = define_orders(zernike_pol); single_pol_used = True # get polynomial orders
else:
for zernpol in zernike_pol:
osa_indices.append(zernpol.get_indices()[1])
@@ -794,12 +790,12 @@ def save_psf(psf_kernel: np.ndarray, NA: float, wavelength: float, pixel_size: f
if single_pol_used:
json_file_path = saved_psfs_folder.joinpath(f"psf_{(m, n)}_{additional_file_name}_{expansion_coefficient}.json")
else:
- json_file_path = saved_psfs_folder.joinpath(f"psf_{osa_indices}_{additional_file_name}_{expansion_coefficient}.json")
+ json_file_path = saved_psfs_folder.joinpath(f"psf_{osa_indices}_{additional_file_name}_{ampls}.json")
else:
if single_pol_used:
json_file_path = saved_psfs_folder.joinpath(f"psf_{(m, n)}_{expansion_coefficient}.json")
else:
- json_file_path = saved_psfs_folder.joinpath(f"psf_{osa_indices}_{expansion_coefficient}.json")
+ json_file_path = saved_psfs_folder.joinpath(f"psf_{osa_indices}_{ampls}.json")
# Data composing for recording
data4serialization = {} # python dictionary is similar to the JSON file structure and can be dumped directly there
data4serialization['PSF Kernel'] = psf_kernel.tolist(); data4serialization['NA'] = NA; data4serialization['Wavelength'] = wavelength
@@ -808,17 +804,18 @@ def save_psf(psf_kernel: np.ndarray, NA: float, wavelength: float, pixel_size: f
if single_pol_used:
data4serialization["Expansion Coefficient"] = expansion_coefficient; data4serialization["Polynomial"] = zernike_pol.get_indices()[1]
else:
- data4serialization["Amplitudes"] = expansion_coefficient.tolist(); data4serialization["Polynomials"] = osa_indices
+ data4serialization["Amplitudes"] = ampls; data4serialization["Polynomials"] = osa_indices
# File presence check and recording
if json_file_path.exists() and not overwrite:
- _warn_message = "The file already exists, the content won't be overwritten."; warnings.warn(_warn_message)
+ _warn_message = "\nThe file already exists, the content won't be overwritten."
+ warnings.warn(_warn_message, stacklevel=2)
if not json_file_path.exists() or (json_file_path.exists() and overwrite):
with open(json_file_path, 'w') as json_write_file:
json.dump(data4serialization, json_write_file)
return str(json_file_path.absolute())
-def read_psf(file_path: str) -> dict:
+def read_psf(file_path: str) -> Optional[dict]:
"""
Read the saved PSF data from the *json file.
@@ -891,87 +888,3 @@ def get_bumped_circle(radius: float, max_intensity: int = 255) -> np.ndarray:
# %% Define standard exports from this module
__all__ = ['get_psf_kernel', 'save_psf', 'read_psf', 'convolute_img_psf', 'radial_integral_s', 'get_kernel_size', 'radial_integral',
'get_kernel_size', 'um_char', 'lambda_char', 'get_bumped_circle', 'get_psf_kernel_zerns']
-
-# %% Tests
-if __name__ == '__main__':
- from zernpy import ZernPol
- plt.ion(); plt.close('all') # close all plots before plotting new ones
- # Physical parameters of a system (an objective)
- wavelength = 0.55 # in micrometers
- NA = 0.95 # objective property, ultimately NA = d/2*f, there d - aperture diameter, f - distance to the object (focal length for an image)
- # Note that ideal Airy pattern will be (2*J1(x)/x)^2, there x = k*NA*r, there r - radius in the polar coordinates on the image
- resolution = 0.61*wavelength/NA # ultimate theoretical physical resolution of an objective
- pixel_size_nyquist = 0.5*resolution # Nyquist's resolution needed for using theoretical physical resolution above
- pixel_size = 0.95*pixel_size_nyquist # the relation between um / pixels for calculating the coordinate in physical units for each pixel
-
- # Flags for performing tests
- check_zero_case = True # checking that integral equation is corresponding to the Airy pattern (zero case)
- check_sign_coeff = False # checking the same amplitude applied for the same polynomial (trefoil)
- check_performance_optimizations = False # checking optimization of calculations
- check_various_pols = False # checking the shape of some Zernike polynomials for comparing with the link below
- check_warnings = False # flag for checking the warning producing
- check_io = False # check save / read kernels
- show_convolution_results = False # check result of convolution of several kernel with the disk
-
- # Definition of some Zernike polynomials for further tests
- pol1 = (0, 0); pol2 = (-1, 1); pol3 = (0, 2); pol4 = (-2, 2); pol5 = (-3, 3); pol6 = (2, 2); pol7 = (-1, 3); pol8 = (0, 4)
- pol9 = (-4, 4); pol7z = ZernPol(m=pol7[0], n=pol7[1])
- pol1z = ZernPol(m=pol1[0], n=pol1[1]); pol2z = ZernPol(m=pol2[0], n=pol2[1]); pol3z = ZernPol(m=pol3[0], n=pol3[1])
-
- if check_zero_case:
- kern_zc = get_psf_kernel(pol1z, len2pixels=wavelength/5.5, alpha=-0.4, wavelength=0.55, NA=0.65,
- show_kernel=True, normalize_values=True)
- kern_zc_ref = get_psf_kernel(pol1z, len2pixels=wavelength/5.5, alpha=-0.4, wavelength=0.55, NA=0.65, airy_pattern=True,
- show_kernel=True, normalize_values=True)
- diff = kern_zc_ref - kern_zc; plt.figure("Difference Airy and Piston", figsize=(6, 6))
- plt.imshow(diff, cmap=plt.cm.viridis, origin='upper')
-
- if check_sign_coeff:
- kern_sign_n = get_psf_kernel(pol5, len2pixels=pixel_size, alpha=-0.5, wavelength=wavelength, NA=NA, normalize_values=False)
- kern_sign_p = get_psf_kernel(pol5, len2pixels=pixel_size, alpha=0.5, wavelength=wavelength, NA=NA, normalize_values=False)
-
- if check_performance_optimizations:
- t1 = time.perf_counter()
- kern_sign = get_psf_kernel(ZernPol(m=-3, n=3), len2pixels=pixel_size, alpha=0.2, wavelength=wavelength,
- NA=NA, normalize_values=True, show_kernel=False)
- print("Calc. time for 'for loops' form ms:", int(round(1000*(time.perf_counter() - t1), 0)))
- t1 = time.perf_counter()
- kern_sign2 = get_psf_kernel(ZernPol(m=-3, n=3), len2pixels=pixel_size, alpha=0.2, wavelength=wavelength,
- NA=NA, normalize_values=True, show_kernel=False, test_vectorized=True, fig_id="Vector. Form")
- print("Calc. time for 'vectorized' form ms:", int(round(1000*(time.perf_counter() - t1), 0)))
- t1 = time.perf_counter()
- kern_sign3 = get_psf_kernel(ZernPol(m=-3, n=3), len2pixels=pixel_size, alpha=0.2, wavelength=wavelength,
- NA=NA, normalize_values=True, show_kernel=False, test_parallel=True, fig_id="Parallel. Form")
- print("Calc. time for 'parallelized' form ms:", int(round(1000*(time.perf_counter() - t1), 0)))
- kern_diff1 = np.round(kern_sign - kern_sign2, 9); kern_diff2 = np.round(kern_sign - kern_sign3, 9)
-
- if check_various_pols:
- kern_def = get_psf_kernel(pol3z, len2pixels=pixel_size, alpha=0.5, wavelength=wavelength, NA=NA, normalize_values=True,
- show_kernel=True)
- kern_ast = get_psf_kernel(pol6, len2pixels=pixel_size, alpha=0.5, wavelength=wavelength, NA=NA, normalize_values=True,
- show_kernel=True)
- kern_coma = get_psf_kernel(pol7, len2pixels=pixel_size, alpha=0.5, wavelength=wavelength, NA=NA, normalize_values=True,
- show_kernel=True)
- kern_spher = get_psf_kernel(pol8, len2pixels=pixel_size, alpha=0.5, wavelength=wavelength,
- NA=NA, normalize_values=True, show_kernel=True)
- kern_4foil = get_psf_kernel(pol9, len2pixels=pixel_size, alpha=0.5, wavelength=wavelength,
- NA=NA, normalize_values=True, show_kernel=True)
-
- if check_warnings:
- kern_def = get_psf_kernel(pol3z, len2pixels=1.0, alpha=0.5, wavelength=wavelength, NA=NA, normalize_values=True)
- if show_convolution_results:
- # Generate the ideal centralized circle with the blurred edges
- radius = 4.0; sample = get_bumped_circle(radius); plt.figure("Sample disk"); m_center, n_center = sample.shape
- m_center = m_center // 2; n_center = n_center // 2; axes_img = plt.imshow(sample, cmap=plt.cm.viridis, origin='upper')
- plt.tight_layout(); axes_img.axes.add_patch(Circle((n_center, m_center), radius, edgecolor='red', facecolor='none'))
- # Visualize results of convolution with various PSFs
- kern_def = get_psf_kernel(pol3z, len2pixels=pixel_size, alpha=0.7, wavelength=wavelength,
- NA=NA, normalize_values=True, show_kernel=True)
- conv_def = convolute_img_psf(img=sample, psf_kernel=kern_def, scale2original=True)
- plt.figure("Defocused"); axes_img2 = plt.imshow(conv_def, cmap=plt.cm.viridis, origin='upper'); plt.tight_layout()
- axes_img2.axes.add_patch(Circle((n_center, m_center), radius, edgecolor='red', facecolor='none'))
- kern_coma = get_psf_kernel(pol7z, len2pixels=pixel_size, alpha=-0.45, wavelength=wavelength,
- NA=NA, normalize_values=True, show_kernel=True)
- conv_coma = convolute_img_psf(img=sample, psf_kernel=kern_coma, scale2original=True)
- plt.figure("*Coma"); axes_img3 = plt.imshow(conv_coma, cmap=plt.cm.viridis, origin='upper'); plt.tight_layout()
- axes_img3.axes.add_patch(Circle((n_center, m_center), radius, edgecolor='red', facecolor='none'))
diff --git a/src/zernpy/calculations/calc_psfs_check.py b/src/zernpy/calculations/calc_psfs_check.py
deleted file mode 100644
index 79f68bb..0000000
--- a/src/zernpy/calculations/calc_psfs_check.py
+++ /dev/null
@@ -1,707 +0,0 @@
-# -*- coding: utf-8 -*-
-"""
-Check calculation and plotting of associated with polynomials PSFs.
-
-@author: Sergei Klykov
-@licence: MIT
-
-"""
-# %% Global imports
-import numpy as np
-import matplotlib.pyplot as plt
-from pathlib import Path
-from scipy.special import jv
-import warnings
-from math import cos, pi
-from zernpy import ZernPol
-import time
-import os
-from skimage import io
-from skimage.util import img_as_ubyte
-
-# %% Local (package-scoped) imports
-if __name__ == "__main__" or __name__ == Path(__file__).stem or __name__ == "__mp_main__":
- from calc_zernike_pol import define_orders
- from calc_psfs import convolute_img_psf, save_psf, read_psf
-else:
- from .calc_zernike_pol import define_orders
- from .calc_psfs import convolute_img_psf, save_psf, read_psf
-
-# %% Module parameters
-__docformat__ = "numpydoc"
-n_phi_points = 300; n_p_points = 320
-# Physical parameters
-wavelength = 0.55 # in micrometers
-k = 2.0*np.pi/wavelength # angular frequency
-NA = 0.95 # microobjective property, ultimately NA = d/2*f, there d - aperture diameter, f - distance to the object (focal length for an image)
-# Note that ideal Airy pattern will be (2*J1(x)/x)^2, there x = k*NA*r, there r - radius in the polar coordinates on the image
-pixel_size = 0.14 # in micrometers, physical length in pixels (um / pixels)
-pixel2um_coeff = k*NA*pixel_size # coefficient used for relate pixels to physical units
-pixel2um_coeff_plot = k*NA*(pixel_size/10.0) # coefficient used for better plotting with the reduced pixel size for preventing pixelated
-
-
-# %% Integral Functions - integration by trapezoidal rule first on radius (normalized on the pupil), after - on angle
-def diffraction_integral_r(zernike_pol: ZernPol, alpha: float, phi: float, p: np.array, theta: float, r: float) -> np.array:
- """
- Diffraction integral function (see the references for the source of the equation).
-
- Parameters
- ----------
- zernike_pol : ZernPol
- Zernike polynomial definition.
- alpha : float
- Amplitude of the polynomial.
- phi : float
- Angle on the pupil coordinates.
- p : np.ndarray
- Integration interval on the pupil coordinates.
- theta : float
- Angle on the image coordinates.
- r : float
- Radius on the image coordinates.
-
- References
- ----------
- [1] https://nijboerzernike.nl/_downloads/Thesis_Nijboer.pdf
- [2] https://opg.optica.org/ao/abstract.cfm?uri=ao-52-10-2062 (DOI: 10.1364/AO.52.002062)
-
- Returns
- -------
- numpy.ndarray
- Values of the diffraction integral.
-
- """
- # Multiplication by phi (in r*p*np.cos(...) only guides to scaling of the resulting PSF
- phase_arg = (alpha*zernike_pol.polynomial_value(p, phi) - r*p*np.cos(phi - theta))*1j
- return np.exp(phase_arg)*p
-
-
-def radial_integral(zernike_pol: ZernPol, r: float, theta: float, phi: float, alpha: float) -> complex:
- """
- Make integration of the diffraction integral on the radius.
-
- Parameters
- ----------
- zernike_pol : ZernPol
- Zernike polynomial definition.
- r : float
- Radius on the image coordinates.
- theta : float
- Angle on the image coordinates.
- phi : float
- Angle on the pupil coordinates.
- alpha : float
- Amplitude of the polynomial.
-
- Returns
- -------
- complex
- Complex amplitude of the field as result of integration on the pupil radius coordinate.
-
- """
- # Integration on the pupil angle. Vectorized form of the trapezoidal rule
- h_p = 1.0/n_p_points; p = np.arange(start=h_p, stop=1.0, step=h_p)
- fa = diffraction_integral_r(zernike_pol, alpha, phi, 0.0, theta, r); fb = diffraction_integral_r(zernike_pol, alpha, phi, 1.0, theta, r)
- ang_int = np.sum(diffraction_integral_r(zernike_pol, alpha, phi, p, theta, r)) + 0.5*(fa + fb)
- return h_p*ang_int
-
-
-def get_psf_point_r(zernike_pol, r: float, theta: float, alpha: float = 1.0) -> float:
- """
- Get the point for calculation of PSF.
-
- Parameters
- ----------
- zernike_pol : tuple (m, n) or ZernPol
- Zernike polynomial definition.
- r : float
- Radius on the image coordinates.
- theta : float
- Angle on the image coordinates.
- alpha : float, optional
- Amplitude of the polynomial. The default is 1.0.
-
- Returns
- -------
- float
- |U|**2 - the module and square of the amplitude, intensity as the PSF value.
-
- """
- # Check and initialize Zernike polynomial if provided only orders
- (m, n) = define_orders(zernike_pol) # get polynomial orders
- if not isinstance(zernike_pol, ZernPol):
- zernike_pol = ZernPol(m=m, n=n)
- # Integration on the pupil radius using Simpson equation
- n_integral_points = n_phi_points; h_phi = 2.0*pi/n_integral_points
- even_sum = 0.0j; odd_sum = 0.0j
- for i in range(2, n_integral_points-2, 2):
- phi = i*h_phi; even_sum += radial_integral(zernike_pol, r, theta, phi, alpha)
- for i in range(1, n_integral_points-1, 2):
- phi = i*h_phi; odd_sum += radial_integral(zernike_pol, r, theta, phi, alpha)
- yA = radial_integral(zernike_pol, r, theta, 0.0, alpha); yB = radial_integral(zernike_pol, r, theta, 2.0*pi, alpha)
- integral_sum = (h_phi/3.0)*(yA + yB + 2.0*even_sum + 4.0*odd_sum); integral_normalization = 1.0/(pi*pi)
- return np.power(np.abs(integral_sum), 2)*integral_normalization
-
-
-# %% Integral Functions - integration on the radial external, angular - internal by the trapezoidal rule
-def diffraction_integral_ang(zernike_pol: ZernPol, alpha: float, phi: np.array, p: float, theta: float, r: float) -> np.array:
- """
- Diffraction integral function (see the references for the source of the equation).
-
- Parameters
- ----------
- zernike_pol : ZernPol
- Zernike polynomial definition.
- alpha : float
- Amplitude of the polynomial.
- phi : float
- Integration interval of angles on the pupil coordinates.
- p : float
- Radius on the pupil coordinates.
- theta : float
- Angle on the image coordinates.
- r : float
- Radius on the image coordinates.
-
- References
- ----------
- [1] https://nijboerzernike.nl/_downloads/Thesis_Nijboer.pdf
- [2] https://opg.optica.org/ao/abstract.cfm?uri=ao-52-10-2062 (DOI: 10.1364/AO.52.002062)
-
- Returns
- -------
- numpy.ndarray
- Values of the diffraction integral.
-
- """
- phase_arg = (alpha*zernike_pol.polynomial_value(p, phi) - r*p*np.cos(phi - theta))*1j
- return np.exp(phase_arg)*p
-
-
-def angular_integral(zernike_pol: ZernPol, r: float, theta: float, p: float, alpha: float) -> complex:
- """
- Make integration of the diffraction integral on the angle.
-
- Parameters
- ----------
- zernike_pol : ZernPol
- Zernike polynomial definition.
- r : float
- Radius on the image coordinates.
- theta : float
- Angle on the image coordinates.
- p : float
- Radius on the pupil coordinates.
- alpha : float
- Amplitude of the polynomial.
-
- Returns
- -------
- complex
- Complex amplitude of the field as result of integration on the pupil radius coordinate.
-
- """
- # Integration on the pupil angle. Vectorized form of the trapezoidal rule
- h_phi = pi/n_phi_points; phi = np.arange(start=h_phi, stop=2.0*pi, step=h_phi)
- fa = diffraction_integral_ang(zernike_pol, alpha, 0.0, p, theta, r); fb = diffraction_integral_ang(zernike_pol, alpha, 2.0*pi, p, theta, r)
- ang_int = np.sum(diffraction_integral_ang(zernike_pol, alpha, phi, p, theta, r)) + 0.5*(fa + fb)
- return h_phi*ang_int
-
-
-def get_psf_point_ang(zernike_pol, r: float, theta: float, alpha: float = 1.0) -> float:
- """
- Get the point for calculation of PSF.
-
- Parameters
- ----------
- zernike_pol : tuple (m, n) or ZernPol
- Zernike polynomial definition.
- r : float
- Radius on the image coordinates.
- theta : float
- Angle on the image coordinates.
- alpha : float, optional
- Amplitude of the polynomial. The default is 1.0.
-
- Returns
- -------
- float
- |U|**2 - the module and square of the amplitude, intensity as the PSF value.
-
- """
- # Check and initialize Zernike polynomial if provided only orders
- (m, n) = define_orders(zernike_pol) # get polynomial orders
- if not isinstance(zernike_pol, ZernPol):
- zernike_pol = ZernPol(m=m, n=n)
- # Integration on the pupil radius using Simpson equation
- n_integral_points = n_p_points; h_p = 1.0/n_integral_points
- even_sum = 0.0j; odd_sum = 0.0j
- for i in range(2, n_integral_points-2, 2):
- p = i*h_p; even_sum += angular_integral(zernike_pol, r, theta, p, alpha)
- for i in range(1, n_integral_points-1, 2):
- p = i*h_p; odd_sum += angular_integral(zernike_pol, r, theta, p, alpha)
- yA = angular_integral(zernike_pol, r, theta, 0.0, alpha); yB = angular_integral(zernike_pol, r, theta, 1.0, alpha)
- integral_sum = (h_p/3.0)*(yA + yB + 2.0*even_sum + 4.0*odd_sum)
- return np.power(np.abs(integral_sum), 2)/(4.0*pi*pi)
-
-
-# %% Attempt to speed up calculations by using provided in the Born & Wolf's handbook equation
-def diffraction_int_approx_sum(zernike_pol: ZernPol, s: int, alpha: float, theta, r):
- m, n = define_orders(zernike_pol)
- if s == 0:
- c = 2.0
- else:
- c = 4.0
- coefficient = c*pow(1j, (m-1)*s)*np.cos(m*s*theta)
- h_p = 1.0/n_p_points; p_arr = np.arange(start=h_p, stop=1.0, step=h_p)
- eps_zero = 1E-11 # substitution of the exact 0.0 value
- if r <= 1E-12:
- r = eps_zero
- if alpha <= 1E-12:
- alpha = 1E-11
- fa = jv(s, alpha*zernike_pol.radial(eps_zero))*jv(m*s, eps_zero*r)*eps_zero # f(0.0)
- fb = jv(s, alpha*zernike_pol.radial(1.0))*jv(m*s, r); fAB = 0.5*(fa + fb)
- integral_sum = h_p*np.sum(jv(s, alpha*zernike_pol.radial(p_arr))*jv(m*s, r*p_arr)*p_arr)
- return coefficient*(fAB + integral_sum)
-
-
-def get_psf_point_bwap(zernike_pol, r: float, theta: float, alpha: float = 1.0) -> float:
- m, n = define_orders(zernike_pol)
- if not isinstance(zernike_pol, ZernPol):
- zernike_pol = ZernPol(m=m, n=n)
- s_max = 30; phases_sum = 0.0j
- for s_i in range(s_max):
- phases_sum += diffraction_int_approx_sum(zernike_pol, s_i, alpha, theta, r)
- integral_normalization = 1.0/(pi*pi)
- return integral_normalization*np.power(np.abs(phases_sum), 2)
-
-
-# %% Nijboer's Thesis Functions
-def radial_func(n: int, r: float) -> float:
- """
- Define the radial function Jn(r)/r, there Jn - Bessel function of 1st kind with order n.
-
- Parameters
- ----------
- n : int
- Order of the Bessel function.
- r : float
- Radial parameter r.
-
- Returns
- -------
- float
- Evaluated function Jn(r)/r.
-
- """
- # Defining pure radial function (angular independent) - Jv(r)/r
- if isinstance(r, int): # Convert int to float explicitly
- r = float(r)
- radial = 0.0 # default value
- # Calculate value only for input r as the float number
- if isinstance(r, float):
- if abs(round(r, 12)) == 0.0: # check that the argument provided with 0 value
- radial = round(2.0*pow(jv(n, 1E-11)/1E-11, 2), 11) # approximation of the limit for the special condition jv(x)/x, where x -> 0
- else:
- radial = round(2.0*pow(jv(n, r)/r, 2), 11)
- return radial
-
-
-def radial_ampl_func(n: int, r: np.ndarray) -> np.ndarray:
- """
- Calculate the radial amplitude functions used in the Nijboer's thesis as Jv(r)/r.
-
- Parameters
- ----------
- n : int
- Radial order of the Zernike polynomial.
- r : np.ndarray
- Radii on the pupil coordinates.
- Note that the array can only start with the special zero point (r[0] = 0.0).
-
- Reference
- ---------
- [1] https://nijboerzernike.nl/_downloads/Thesis_Nijboer.pdf
-
- Returns
- -------
- radial_ampl_f : np.ndarray
- Jv(r)/r function values.
-
- """
- # Expanse the calculation for using the numpy array, assuming it starting with 0.0
- radial_ampl_f = None
- print(type(r))
- if isinstance(r, np.ndarray):
- r = np.round(r, 12)
- if r[0] == 0.0:
- radial_ampl_f = np.zeros(shape=(r.shape[0]))
- r1 = r[1:]; radial_ampl_f[1:] = jv(n, r1)/r1
- radial_ampl_f[0] = jv(n, 1E-11)/1E-11
- return radial_ampl_f
-
-
-def get_aberrated_psf(zernike_pol, r: float, theta: float, alpha: float = 1.0) -> float:
- """
- Get the precalculated expansion of diffraction integral based on the Nijboer's thesis.
-
- Parameters
- ----------
- zernike_pol : tuple (m, n) or ZernPol
- Zernike polynomial definition.
- r : float
- Radius on the image coordinates.
- theta : float
- Angle on the image coordinates.
- alpha : float, optional
- Amplitude of the polynomial. The default is 1.0.
-
- Reference
- ---------
- [1] https://nijboerzernike.nl/_downloads/Thesis_Nijboer.pdf
-
- Returns
- -------
- float
- PSF point as |U|**2.
-
- """
- (m, n) = define_orders(zernike_pol) # get polynomial orders
- # The analytical equation could be found in the Nijboer's thesis
- if m == 0 and n == 0: # piston value
- x1 = radial_func(1, r)
- return x1*x1*(1.0 + alpha*alpha)
- # Approximation of the (-1, 1) polynomial (Y tilt). Note that each polynomial should be somehow evaluated or overall equation used
- elif m == -1 and n == 1:
- alpha = round(alpha, 4) # rounding with extended precision, commonly alpha up to 0.001 useful
- x1 = radial_func(1, r); x2 = radial_func(2, r); x3 = radial_func(3, r); x4 = radial_func(4, r)
- c1 = -(alpha*alpha)/8.0; c2 = (alpha*alpha)/4.0; c4 = pow(alpha, 3)/24.0; c5 = pow(alpha, 3)/132.0
- u = x1 + alpha*x2*cos(theta) + c1*(x1 - x3) + c2*x3*cos(2.0*theta) + c4*x4*cos(3.0*theta) + c5*(x4 - 2.0*x1)
- if alpha > 1.0:
- c6 = pow(alpha, 4)/192.0; x5 = radial_func(5, r)
- u += c6*(x1 + -1.5*x3 + 0.5*x5 - cos(2.0*theta)*(3.0*x3 - x5) + cos(4.0*theta)*x5)
- if alpha > 2.0:
- __warn_message = f"For such big aberration amplitude ({alpha}) not precise equation was found, use it with caution!"
- warnings.warn(__warn_message)
- return u*u
- # Approximation of the (-2, 2) polynomial (defocus)
- elif m == 0 and n == 2:
- x1 = radial_func(1, r); x3 = radial_func(3, r); x5 = radial_func(5, r); x7 = radial_func(7, r)
- c1 = -(alpha*alpha)/2.0; c2 = pow(alpha, 3)/6.0
- a = x1 + c1*((4.0/3.0)*x1 - x3 + (2.0/3.0)*x5); b = alpha*x3 + c2*(-7.0*x1 + 4.8*x3 + 2.0*x5 - 0.8*x7)
- if alpha < 0.0:
- return a*a - b*b
- else:
- return b*b - a*a
- else:
- return 0.0
-
-
-# %% Another attempt to use Nijboer's Thesis Functions
-# Tested, this approximation somehow doesn't work (maybe, the amplitudes of polynomial aren't small enough)
-def radial_integral_nijboer(zernike_pol: ZernPol, r: float, order: int, power: int) -> float:
- """
- Calculate the radial integrals used in the Nijboer's thesis.
-
- Parameters
- ----------
- zernike_pol : tuple (m, n) or ZernPol
- Zernike polynomial definition.
- r : float
- Radius on the image coordinates.
- order : int
- Of the Bessel function Jv(r).
- power : int
- Power of the polynomial radial component.
-
- Reference
- ---------
- [1] https://nijboerzernike.nl/_downloads/Thesis_Nijboer.pdf
-
- Returns
- -------
- float
- Result of integration.
-
- """
- integral_sum = 0.0; (m, n) = define_orders(zernike_pol) # get polynomial orders
- # Integration on the pupil radius. Vectorized form of simple integration equation
- h_p = 1.0/1000; p = np.arange(start=h_p, stop=1.0-h_p, step=h_p)
- fa = 0.0; fb = np.power(zernike_pol.radial(1.0), power)*jv(order, r)
- integral_sum = h_p*(np.sum(p*np.power(zernike_pol.radial(p), power)*jv(order, p*r)) + 0.5*(fa + fb))
- return integral_sum
-
-
-def psf_point_approx_sum(zernike_pol: ZernPol, r: float, theta: float, alpha: float):
- """
- Get the point for calculation of PSF as the expansion of the diffraction integral acquired in the Nijboer's thesis.
-
- Parameters
- ----------
- zernike_pol : tuple (m, n) or ZernPol
- Zernike polynomial definition.
- r : float
- Radius on the image coordinates.
- theta : float
- Angle on the image coordinates.
- alpha : float, optional
- Amplitude of the polynomial. The default is 1.0.
-
- Reference
- ---------
- [1] https://nijboerzernike.nl/_downloads/Thesis_Nijboer.pdf
-
-
- Returns
- -------
- float
- |U|**2 - the module and square of the amplitude, intensity as the PSF value.
-
- """
- (m, n) = define_orders(zernike_pol) # get polynomial orders
- if not isinstance(zernike_pol, ZernPol):
- zernike_pol = ZernPol(m=m, n=n)
- # Integration on the pupil radius. Vectorized form of simple integration equation
- r = round(r, 11)
- if r == 0.0:
- s1 = 2.0*(jv(1, 1E-11)/1E-11); s2 = -2.0*alpha*pow(1j, m+n+1)*np.cos(m*theta)*(jv(n+1, 1E-11)/1E-11)
- else:
- s1 = 2.0*(jv(1, r)/r); s2 = -2.0*alpha*pow(1j, m+n+1)*np.cos(m*theta)*(jv(n+1, r)/r)
- s3_1 = radial_integral_nijboer(zernike_pol, r, order=0, power=2); s3_2 = radial_integral_nijboer(zernike_pol, r, order=2*m, power=2)
- s3 = -(pow(alpha, 2)/2.0)*(s3_1 + pow(1j, 2*m)*np.cos(2*m*theta)*s3_2)
- s4_1 = radial_integral_nijboer(zernike_pol, r, order=m, power=3); s4_2 = radial_integral_nijboer(zernike_pol, r, order=3*m, power=3)
- s4 = (pow(alpha, 3)/12.0)*(3.0*pow(1j, m+1)*np.cos(m*theta)*s4_1 + pow(1j, 3*m+1)*np.cos(3*m*theta)*s4_2)
- s5_1 = 3.0*radial_integral_nijboer(zernike_pol, r, order=0, power=4)
- s5_2 = 4.0*pow(1j, 2*m)*np.cos(2*m*theta)*radial_integral_nijboer(zernike_pol, r, order=2*m, power=4)
- s5_3 = pow(1j, 4*m)*np.cos(4*m*theta)*radial_integral_nijboer(zernike_pol, r, order=4*m, power=4)
- s5 = (pow(alpha, 4)/96.0)*(s5_1 + s5_2 + s5_3)
- return np.power(np.abs(s1 + s2 + s3 + s4 + s5), 2)
-
-
-# %% Plotting of PSFs
-def get_psf_kernel(zernike_pol, calibration_coefficient: float, alpha: float, unified_kernel_size: bool = False) -> np.ndarray:
- """
- Calculate centralized matrix with PSF mask.
-
- Parameters
- ----------
- zernike_pol : tuple (m, n) or instance of the ZernPol class.
- Required for calculation Zernike polynomial.
-
- calibration_coefficient : float
- Relation between pixels and distance (physical).
-
- unified_kernel_size : bool
- Flag for adjusting or not the kernel size to the provided absolute value of alpha. The default is False.
-
- Returns
- -------
- kernel : numpy.ndarray
- Matrix with PSF values.
- """
- (m, n) = define_orders(zernike_pol) # get polynomial orders
- # Define the kernel size just empirically, based on made experimental plots
- if not unified_kernel_size:
- if abs(alpha) < 1.0:
- multiplier = 20.0
- elif abs(alpha) < 1.5:
- multiplier = 24.0
- elif abs(alpha) <= 2.0:
- multiplier = 26.0
- else:
- multiplier = 30.0
- else:
- multiplier = 26.0
- max_size = int(round(multiplier*(1.0/calibration_coefficient), 0)) + 1 # Note that with the amplitude growth, it requires to grow as well
- size = max_size # Just use the maximum estimation
- # Auto definition of the required PSF size is failed for the different PSFs forms (e.g., vertical coma with different amplitudes)
- # Make kernel with odd sizes for precisely centering the kernel
- if size % 2 == 0:
- size += 1
- kernel = np.zeros(shape=(size, size))
- i_center = size//2; j_center = size//2
- # Calculate the PSF kernel for usage in convolution operation
- for i in range(size):
- for j in range(size):
- pixel_dist = np.sqrt(np.power((i - i_center), 2) + np.power((j - j_center), 2))
- # The PSF also has the angular dependency, not only the radial one
- theta = np.arctan2((i - i_center), (j - j_center))
- theta += np.pi # shift angles to the range [0, 2pi]
- # The scaling below is not needed because the Zernike polynomial is scaled as the RMS values
- kernel[i, j] = get_psf_point_r(zernike_pol, pixel_dist*calibration_coefficient, theta, alpha)
- return kernel
-
-
-def show_ideal_psf(zernike_pol, size: int, calibration_coefficient: float, alpha: float, title: str = None, test_alt: bool = False):
- """
- Plot the intensity distribution on the image with WxH: (size, size) and using coefficient between pixel and physical distance.
-
- Note the color map is viridis.
-
- Parameters
- ----------
- size : int
- Size of picture for plotting.
- calibration_coefficient : float
- Relation between distance in pixels and um (see parameters at the start lines of the script).
- title : str, optional
- Title for the plotted figure. The default is None.
-
- Returns
- -------
- None.
-
- """
- if size % 2 == 0:
- size += 1 # make the image with odd sizes
- img = np.zeros((size, size), dtype=float)
- i_center = size//2; j_center = size//2
- for i in range(size):
- for j in range(size):
- pixel_dist = np.sqrt(np.power((i - i_center), 2) + np.power((j - j_center), 2))
- # The PSF also has the angular dependency, not only the radial one
- theta = np.arctan2((i - i_center), (j - j_center))
- theta += np.pi # shift angles to the range [0, 2pi]
- if not test_alt:
- img[i, j] = get_psf_point_r(zernike_pol, pixel_dist*calibration_coefficient, theta, alpha)
- else:
- # Implementation for testing and comparing with the general equation
- theta -= np.pi/2.0
- img[i, j] = get_psf_point_bwap(zernike_pol, pixel_dist*calibration_coefficient, theta, alpha)
- if img[0, 0] > np.max(img)/100:
- __warn_message = f"The provided size for plotting PSF ({size}) isn't sufficient for proper representation"
- warnings.warn(__warn_message)
- if title is not None and len(title) > 0:
- plt.figure(title, figsize=(6, 6))
- else:
- plt.figure(figsize=(6, 6))
- plt.imshow(img, cmap=plt.cm.viridis, origin='upper'); plt.tight_layout()
- return img
-
-
-def plot_correlation(zernike_pol, size: int, calibration_coefficient: float, alpha: float, title: str = None,
- show_original: bool = True, show_psf: bool = True, R: int = 1):
- """
- Plot result of correlation with the object shown as the pixelated circle.
-
- Parameters
- ----------
- zernike_pol : tuple (m, n) or ZernPol
- Zernike polynomial definition.
- size : int
- Size of the picture with the object.
- calibration_coefficient : float
- Coefficient between physical values and pixels.
- alpha : float
- Amplitude of the polynomial.
- title : str, optional
- Title of plots. The default is None.
- show_original : bool, optional
- Flag for plotting the original object. The default is True.
- show_psf : bool, optional
- Flag for plotting the used PSF. The default is True.
- R : int, optional
- Radius of the circle. The default is 1.
-
- Returns
- -------
- None.
-
- """
- if size % 2 == 0:
- size += 1 # make the image with odd sizes
- img = np.zeros((size, size), dtype=float)
- i_center = size//2; j_center = size//2; R = 1
- for i in range(size):
- for j in range(size):
- pixel_dist = np.sqrt(np.power((i - i_center), 2) + np.power((j - j_center), 2))/R
- if pixel_dist < 1.0:
- img[i, j] = 1.0
- else:
- continue
- if show_original:
- plt.figure("Original object", figsize=(6, 6)); plt.imshow(img, cmap=plt.cm.viridis, extent=(0, size, 0, size))
- plt.tight_layout()
- psf_kernel = get_psf_kernel(zernike_pol, calibration_coefficient, alpha)
- if show_psf:
- plt.figure(f"PSF for {title}", figsize=(6, 6)); plt.imshow(psf_kernel, cmap=plt.cm.viridis)
- conv_img = convolute_img_psf(img, psf_kernel, scale2original=True)
- plt.figure(f"Convolved with {title} image"); plt.imshow(conv_img, cmap=plt.cm.viridis, extent=(0, size, 0, size)); plt.tight_layout()
-
-
-def plot_correlation_photo(zernike_pol, calibration_coefficient: float, alpha: float, title: str = None,
- show_original: bool = True, show_psf: bool = False):
- """
- Plot result of convolution of PSF with the sample image.
-
- Parameters
- ----------
- zernike_pol : tuple (m, n) or ZernPol
- Zernike polynomial definition.
- size : int
- Size of the picture with the object.
- calibration_coefficient : float
- Coefficient between physical values and pixels.
- alpha : float
- Amplitude of the polynomial.
- title : str, optional
- Title of plots. The default is None.
- show_original : bool, optional
- Flag for plotting the original object. The default is True.
- show_psf : bool, optional
- Flag for plotting the used PSF. The default is True.
-
- Returns
- -------
- None.
-
- """
- try:
- sample = img_as_ubyte(io.imread(os.path.join(os.getcwd(), "nesvizh_grey.jpg"), as_gray=True))
- if show_original:
- plt.figure("Original Photo"); plt.imshow(sample, cmap=plt.cm.gray); plt.axis('off'); plt.tight_layout()
- psf_kernel = get_psf_kernel(zernike_pol, calibration_coefficient, alpha)
- if show_psf:
- plt.figure(f"PSF for {title}", figsize=(6, 6)); plt.imshow(psf_kernel, cmap=plt.cm.viridis)
- conv_img = convolute_img_psf(sample, psf_kernel, scale2original=True)
- plt.figure(f"Convolved with {title} image"); plt.imshow(conv_img, cmap=plt.cm.gray); plt.axis('off'); plt.tight_layout()
- except FileNotFoundError:
- _warn_message = "The sample photo is ignored in the repository, load it to the folder with this script from repo 'collection_numCalc'"
- warnings.warn(_warn_message)
-
-
-# %% Tests
-if __name__ == '__main__':
- orders1 = (0, 2); orders2 = (0, 0); orders3 = (-1, 1); orders4 = (-3, 3); plot_pure_psfs = True; plot_photo_convolution = False
- plot_photo_convolution_row = False; figsizes = (6.5, 6.5); test_write_read_psf = False
-
- # Plotting
- plt.ion(); plt.close('all'); conv_pic_size = 14; detailed_plots_sizes = 22; calibration_coeff = pixel2um_coeff; alpha = 2.0
- if plot_pure_psfs:
- t1 = time.time(); p_img = show_ideal_psf(orders2, detailed_plots_sizes-10, calibration_coeff, alpha, "Piston")
- print(f"Calculation (steps on phi/p {n_phi_points}/{n_p_points}) of Piston takes s:", round((time.time() - t1), 3)); t1 = time.time()
-
- # Testing the kernel calculation
- v_coma_m = get_psf_kernel((-1, 3), calibration_coeff, -1.0)
- plt.figure(figsize=figsizes); plt.imshow(v_coma_m, cmap=plt.cm.viridis); plt.tight_layout()
-
- # Plot convolution of the sample photo with the various psfs
- if plot_photo_convolution:
- plot_correlation_photo(orders2, pixel2um_coeff/1.5, 1.0, "Piston", show_psf=True, show_original=True)
- plot_correlation_photo(orders3, pixel2um_coeff/1.5, 1.0, "Y Tilt", show_psf=True, show_original=True)
- plot_correlation_photo(orders1, pixel2um_coeff/1.5, 1.0, "Defocus", show_psf=True, show_original=True)
- plot_correlation_photo(orders4, pixel2um_coeff/1.5, 1.0, "Vertical Trefoil", show_psf=True, show_original=True)
-
- # Testing of the convolution results with different amplitudes and signs
- if plot_photo_convolution_row:
- amplitudes = [-2.0, 0.0, 2.0]; vert_coma = (-1, 3)
- for amplitude in amplitudes:
- plot_correlation_photo(vert_coma, pixel2um_coeff/1.5, amplitude, f"Y Coma {amplitude}", show_psf=False, show_original=False)
-
- # Testing saving / reading
- if test_write_read_psf:
- ampl = -1.0; psf_kernel = get_psf_kernel(orders4, calibration_coeff, ampl)
- file_path = save_psf("test", psf_kernel, NA, wavelength, calibration_coeff, ampl, orders4)
- psf_stored_data = read_psf(file_path)
- read_psf_kernel = np.asarray(psf_stored_data['PSF kernel'])
- plt.figure(figsize=figsizes); plt.imshow(read_psf_kernel, cmap=plt.cm.viridis); plt.tight_layout()
diff --git a/src/zernpy/calculations/calc_psfs_numba.py b/src/zernpy/calculations/calc_psfs_numba.py
index c569a67..2e14516 100644
--- a/src/zernpy/calculations/calc_psfs_numba.py
+++ b/src/zernpy/calculations/calc_psfs_numba.py
@@ -7,17 +7,19 @@
"""
# %% Global imports
-import numpy as np
-import matplotlib.pyplot as plt
+import logging
+import time
import warnings
from math import pi
-import time
-from typing import Union
+from typing import Optional, Union
+
+import matplotlib.pyplot as plt
+import numpy as np
# %% Checking and import the numba library for speeding up the calculation
-methods_compiled = False # flag for storing if the methods compiled
try:
from numba import njit
+ logging.getLogger('numba').setLevel(logging.WARNING) # disable many DEBUG level logs caused by numba during compilation
except ModuleNotFoundError:
pass
@@ -34,8 +36,8 @@
# %% Airy profile for Z(0, 0) ('airy_ref_pattern' cannot be compiled, deleted)
# %% Exchange ZernPol class call to calculation functions
-@njit
-def zernpol_value(orders: tuple, r: Union[float, np.ndarray], theta: Union[float, np.ndarray]) -> np.ndarray:
+@njit(cache=True)
+def zernpol_value(orders: tuple, r: Union[float, np.ndarray], theta: Union[float, np.ndarray]) -> Union[float, np.ndarray]:
"""
Provide composed Zernike polynomial value calculation for compilation by numba.
@@ -50,8 +52,8 @@ def zernpol_value(orders: tuple, r: Union[float, np.ndarray], theta: Union[float
Returns
-------
- np.ndarray
- Polynomial value(-s).
+ Union[float, np.ndarray] (float | np.ndarray)
+ Polynomial value(-s) depending on used only float values for r and therat or arrays as inputs.
"""
m, n = orders # transfer definition of a polynomial
@@ -71,7 +73,7 @@ def zernpol_value(orders: tuple, r: Union[float, np.ndarray], theta: Union[float
radial = np.power(r, 0)
# 1st order
elif ((m == -1) and (n == 1)) or ((m == 1) and (n == 1)):
- radial = r
+ radial = np.power(r, 1) # explicit for avoiding mypy confusing
# 2nd order
elif ((m == -2) and (n == 2)) or ((m == 2) and (n == 2)):
radial = np.power(r, 2) # r^2
@@ -154,13 +156,14 @@ def zernpol_value(orders: tuple, r: Union[float, np.ndarray], theta: Union[float
return np.power(r, n)
elif n > 10 and abs(m) == n-2: # equation for high order polynomials (orders with abs(m) == n-2)
return float(n)*np.power(r, n) - float(n-1)*np.power(r, n-2)
- # Polynomial value as the multiplication of calculated above components
+ # Polynomial value as the multiplication of 3 parts: normalization, radial, triangular
return norm*triangular*radial
# %% PSF pixel value calc.
-@njit
-def diffraction_integral_r_comp(orders: tuple, alpha: float, phi: float, p: Union[float, np.ndarray], theta: float, r: float) -> np.ndarray:
+@njit(cache=True)
+def diffraction_integral_r_comp(orders: tuple, alpha: float, phi: float, p: Union[float, np.ndarray], theta: float,
+ r: float) -> Union[float, np.ndarray]:
"""
Diffraction integral function for the formed image point (see the references as the sources of the equation).
@@ -186,15 +189,15 @@ def diffraction_integral_r_comp(orders: tuple, alpha: float, phi: float, p: Unio
Returns
-------
- numpy.ndarray
- Values of the diffraction integral.
+ Union[float, np.ndarray] (float | np.ndarray)
+ Value(-s) of the diffraction integral.
"""
phase_arg = (alpha*zernpol_value(orders, p, phi) - r*p*np.cos(phi - theta))*1j
return np.exp(phase_arg)*p
-@njit
+@njit(cache=True)
def radial_integral_comp(orders: tuple, r: float, theta: float, phi: float, alpha: float, n_int_r_points: int) -> complex:
"""
Make integration of the diffraction integral on the radius of the entrance pupil.
@@ -229,7 +232,7 @@ def radial_integral_comp(orders: tuple, r: float, theta: float, phi: float, alph
# %% Testing various speeding up calculation approaches
-@njit
+@njit(cache=True)
def get_psf_point_r_comp(orders: tuple, r: float, theta: float, alpha: float, n_int_r_points: int, n_int_phi_points: int) -> float:
"""
Calculate PSF point for the kernel using Parallel class from the joblib library.
@@ -258,10 +261,9 @@ def get_psf_point_r_comp(orders: tuple, r: float, theta: float, alpha: float, n_
"""
h_phi = 2.0*pi/n_int_phi_points; even_sum = 0.0j; odd_sum = 0.0j
- even_sums = [radial_integral_comp(orders, r, theta, i*h_phi, alpha, n_int_r_points) for i in range(2, n_int_phi_points-2, 2)]
- even_sums = np.asarray(even_sums); even_sum = np.sum(even_sums)
- odd_sums = [radial_integral_comp(orders, r, theta, i*h_phi, alpha, n_int_r_points) for i in range(1, n_int_phi_points-1, 2)]
- odd_sums = np.asarray(odd_sums); odd_sum = np.sum(odd_sums)
+ even_sums = np.asarray([radial_integral_comp(orders, r, theta, i*h_phi, alpha, n_int_r_points) for i in range(2, n_int_phi_points-2, 2)])
+ odd_sums = np.asarray([radial_integral_comp(orders, r, theta, i*h_phi, alpha, n_int_r_points) for i in range(1, n_int_phi_points-1, 2)])
+ even_sum = np.sum(even_sums); odd_sum = np.sum(odd_sums)
# Simpson integration rule implementation
yA = radial_integral_comp(orders, r, theta, 0.0, alpha, n_int_r_points)
yB = radial_integral_comp(orders, r, theta, 2.0*pi, alpha, n_int_r_points)
@@ -271,7 +273,7 @@ def get_psf_point_r_comp(orders: tuple, r: float, theta: float, alpha: float, n_
# %% PSF kernel calc. (accelerated)
def get_psf_kernel_comp(zernike_pol, len2pixels: float, alpha: Union[float, np.ndarray], wavelength: float, NA: float, n_int_r_points: int = 320,
- n_int_phi_points: int = 300, show_kernel: bool = False, fig_title: str = None, normalize_values: bool = False,
+ n_int_phi_points: int = 300, show_kernel: bool = False, fig_title: Optional[str] = None, normalize_values: bool = False,
kernel_size: int = 3, fig_id: str = "", suppress_warns: bool = False, verbose: bool = False) -> np.ndarray:
"""
Calculate centralized matrix (kernel) with the PSF mask values.
@@ -327,7 +329,7 @@ def get_psf_kernel_comp(zernike_pol, len2pixels: float, alpha: Union[float, np.n
# Provide performance tip if the provided kernel size is quite big for calculations
if size > 85 and not suppress_warns:
__warn_message = f"\nCalculation of provided kernel size ({size}x{size}) may take more than 20 seconds"
- warnings.warn(__warn_message); __warn_message = ""
+ warnings.warn(__warn_message, stacklevel=2); __warn_message = ""
kernel = np.zeros(shape=(size, size)); i_center = size//2; j_center = size//2
# Get the orders of polynomial and check if the equation for compilation was implemented
single_polynomial_provided = False # flag for using single polynomial functions
@@ -339,18 +341,18 @@ def get_psf_kernel_comp(zernike_pol, len2pixels: float, alpha: Union[float, np.n
else:
if not isinstance(alpha, np.ndarray):
alpha = np.asarray(alpha)
- polynomials_orders = [] # for checking and providing for further compilation polynomials in a tuple
+ polynomials_orders_l = [] # for checking and providing for further compilation polynomials in a tuple
for pol in zernike_pol:
- m, n = pol.get_mn_orders(); polynomials_orders.append((m, n))
+ m, n = pol.get_mn_orders(); polynomials_orders_l.append((m, n))
if n > 10 and (abs(m) != n or abs(m) != n-2):
raise ValueError(f"The calculation PSF function isn't implemented for these orders: {m, n}")
- polynomials_orders = tuple(polynomials_orders) # convert list to tuple
+ polynomials_orders = tuple(polynomials_orders_l) # convert list to tuple
# Check that the calibration coefficient is sufficient for calculation
pixel_size_nyquist = 0.5*0.61*wavelength/NA
if len2pixels > pixel_size_nyquist and not suppress_warns:
__warn_message = f"\nProvided calibration coefficient {len2pixels} {um_char}/pixels isn't sufficient enough"
__warn_message += f" (defined by the relation between Nyquist freq. and the optical resolution: 0.61{lambda_char}/NA)"
- warnings.warn(__warn_message); __warn_message = ""
+ warnings.warn(__warn_message, stacklevel=2); __warn_message = ""
# Calculate the PSF kernel for usage in convolution operation
if verbose:
calculated_points = 0 # for explicit showing of performance
@@ -391,23 +393,23 @@ def get_psf_kernel_comp(zernike_pol, len2pixels: float, alpha: Union[float, np.n
if kernel_border_max > np.max(kernel)/20.0 and not suppress_warns:
__warn_message = (f"\nThe calculated size for PSF ({size}) isn't sufficient for its proper representation, "
+ "because the maximum value on the kernel border is bigger than 5% of maximum overall kernel")
- warnings.warn(__warn_message)
+ warnings.warn(__warn_message, stacklevel=2)
# Plotting the calculated kernel
if show_kernel:
if fig_title is not None and len(fig_title) > 0:
plt.figure(fig_title, figsize=(6, 6))
else:
if not hasattr(zernike_pol, "__len__"):
- plt.figure(f"{zernike_pol.get_mn_orders()} {zernike_pol.get_polynomial_name(True)}: {round(alpha, 2)}*wavelength {fig_id}",
+ plt.figure(f"{zernike_pol.get_mn_orders()} {zernike_pol.get_polynomial_name(True)}: {np.round(alpha, 2)}*wavelength {fig_id}",
figsize=(6, 6))
else:
plt.figure(f"Sum of provided #{len(zernike_pol)} of polynomials {fig_id}", figsize=(6, 6))
- plt.imshow(kernel, cmap=plt.cm.viridis, origin='upper'); plt.tight_layout()
+ plt.imshow(kernel, cmap=plt.colormaps["viridis"], origin='upper'); plt.tight_layout()
return kernel
# %% PSF calc. for several polynomials
-@njit
+@njit(cache=True)
def pol_sums(polynomials_orders: tuple, amplitudes: np.ndarray, p: float, phi: float) -> float:
"""
Wrap calculation of polynomials values sum for compilation.
@@ -437,9 +439,9 @@ def pol_sums(polynomials_orders: tuple, amplitudes: np.ndarray, p: float, phi: f
return sum_pols
-@njit
+@njit(cache=True)
def diffraction_integral_r_pols_comp(polynomials_orders: tuple, amplitudes: np.ndarray, phi: float,
- p: float, theta: float, r: float) -> float:
+ p: float, theta: float, r: float) -> np.ndarray:
"""
Diffraction integral function for the formed image point (see the references as the sources of the equation).
@@ -473,7 +475,7 @@ def diffraction_integral_r_pols_comp(polynomials_orders: tuple, amplitudes: np.n
return np.exp(phase_arg)*p
-@njit
+@njit(cache=True)
def radial_integral_pols_comp(polynomials_orders: tuple, amplitudes: np.ndarray, r: float, theta: float,
phi: float, n_int_r_points: int) -> complex:
"""
@@ -508,7 +510,7 @@ def radial_integral_pols_comp(polynomials_orders: tuple, amplitudes: np.ndarray,
return h_p*ang_int
-@njit
+@njit(cache=True)
def get_psf_point_r_pols_comp(polynomials_orders: tuple, amplitudes: np.ndarray, r: float, theta: float,
n_int_r_points: int, n_int_phi_points: int) -> float:
"""
@@ -540,12 +542,11 @@ def get_psf_point_r_pols_comp(polynomials_orders: tuple, amplitudes: np.ndarray,
"""
h_phi = 2.0*pi/n_int_phi_points; even_sum = 0.0j; odd_sum = 0.0j
- even_sums = [radial_integral_pols_comp(polynomials_orders, amplitudes, r, theta, i*h_phi, n_int_r_points)
- for i in range(2, n_int_phi_points-2, 2)]
- even_sums = np.asarray(even_sums); even_sum = np.sum(even_sums)
- odd_sums = [radial_integral_pols_comp(polynomials_orders, amplitudes, r, theta, i*h_phi, n_int_r_points)
- for i in range(1, n_int_phi_points-1, 2)]
- odd_sums = np.asarray(odd_sums); odd_sum = np.sum(odd_sums)
+ even_sums = np.asarray([radial_integral_pols_comp(polynomials_orders, amplitudes, r, theta, i*h_phi, n_int_r_points)
+ for i in range(2, n_int_phi_points-2, 2)])
+ odd_sums = np.asarray([radial_integral_pols_comp(polynomials_orders, amplitudes, r, theta, i*h_phi, n_int_r_points)
+ for i in range(1, n_int_phi_points-1, 2)])
+ even_sum = np.sum(even_sums); odd_sum = np.sum(odd_sums)
# Simpson integration rule implementation
yA = radial_integral_pols_comp(polynomials_orders, amplitudes, r, theta, 0.0, n_int_r_points)
yB = radial_integral_pols_comp(polynomials_orders, amplitudes, r, theta, 2.0*pi, n_int_r_points)
@@ -553,52 +554,5 @@ def get_psf_point_r_pols_comp(polynomials_orders: tuple, amplitudes: np.ndarray,
return np.power(np.abs(integral_sum), 2)*integral_normalization
-# %% Utility functions
-def set_methods_compiled():
- """
- Reset the flag by external call.
-
- Returns
- -------
- None.
-
- """
- global methods_compiled; methods_compiled = True
-
-
# %% Define standard exports from this module
-__all__ = ['get_psf_kernel_comp', 'methods_compiled', 'set_methods_compiled']
-
-# %% Tests
-if __name__ == '__main__':
- from zernpy import ZernPol # for polynomials initialization
- plt.ion(); plt.close('all') # close all plots before plotting new ones
- # Physical parameters of a system (an objective)
- wavelength = 0.55 # in micrometers
- NA = 0.95 # objective property, ultimately NA = d/2*f, there d - aperture diameter, f - distance to the object (focal length for an image)
- # Note that ideal Airy pattern will be (2*J1(x)/x)^2, there x = k*NA*r, there r - radius in the polar coordinates on the image
- resolution = 0.61*wavelength/NA # ultimate theoretical physical resolution of an objective
- pixel_size_nyquist = 0.5*resolution # Nyquist's resolution needed for using theoretical physical resolution above
- pixel_size = 0.95*pixel_size_nyquist # the relation between um / pixels for calculating the coordinate in physical units for each pixel
-
- # Flags for testing various scenarios
- test_single_polynomial = False; test_few_polynomials = True
-
- # Testing implemented calculations for single polynomial
- if test_single_polynomial:
- get_psf_kernel_comp(ZernPol(m=0, n=0), pixel_size*0.7, alpha=1.5, wavelength=wavelength, NA=NA, normalize_values=True,
- show_kernel=True, kernel_size=11, verbose=True)
- print("*******************************************")
- get_psf_kernel_comp(ZernPol(m=0, n=4), pixel_size*0.7, alpha=0.75, wavelength=wavelength, NA=NA, normalize_values=True,
- show_kernel=True, kernel_size=21, verbose=True)
- # For several (calculated their sum)
- if test_few_polynomials:
- # First, checking below sequentially all functions to be compilable
- # pols = ((-2, 2), (1, 3)); ampls = np.asarray([-0.4, 0.6])
- # diffraction_integral_r_pols_comp(pols, ampls, phi=0.2, p=0.1, theta=0.3, r=0.5)
- # radial_integral_pols_comp(pols, ampls, r=0.5, theta=0.3, phi=1.01, n_int_r_points=300)
- # get_psf_point_r_pols_comp(pols, ampls, r=0.5, theta=0.3, n_int_r_points=250, n_int_phi_points=320)
- # Second, test all at once for calling the function
- zp1 = ZernPol(m=-2, n=2); zp2 = ZernPol(m=0, n=2); zp3 = ZernPol(m=2, n=2); pols = (zp1, zp2, zp3); coeffs = (-0.86, 0.4, 0.7)
- get_psf_kernel_comp(pols, pixel_size*0.7, alpha=coeffs, wavelength=wavelength, NA=NA, normalize_values=True,
- show_kernel=True, kernel_size=24, verbose=True)
+__all__ = ['get_psf_kernel_comp']
diff --git a/src/zernpy/calculations/calc_zernike_pol.py b/src/zernpy/calculations/calc_zernike_pol.py
index 65f2b67..2db1bb3 100644
--- a/src/zernpy/calculations/calc_zernike_pol.py
+++ b/src/zernpy/calculations/calc_zernike_pol.py
@@ -7,17 +7,13 @@
"""
# %% Global imports
-import numpy as np
import math
import time
-from pathlib import Path
-# from decimal import Decimal
+from typing import List, Union
+
+import numpy as np
-# %% Local (package-scoped) imports
-if __name__ == "__main__" or __name__ == Path(__file__).stem or __name__ == "__mp_main__":
- from find_pols_coeffs import find_coeffs_orders, find_coeffs_orders_dr
-else:
- from .find_pols_coeffs import find_coeffs_orders, find_coeffs_orders_dr
+from .find_pols_coeffs import find_coeffs_orders, find_coeffs_orders_dr
# %% Module parameters
__docformat__ = "numpydoc"
@@ -47,9 +43,9 @@ def define_orders(zernike_pol) -> tuple:
"""
# Get azimuthal and radial orders
if isinstance(zernike_pol, tuple):
- (m, n) = zernike_pol
+ m, n = zernike_pol
else:
- (m, n) = zernike_pol.get_mn_orders()
+ m, n = zernike_pol.get_mn_orders()
return m, n
@@ -72,7 +68,7 @@ def normalization_factor(zernike_pol) -> float:
Normalization factor, depending only on Zernike type.
"""
- (m, n) = define_orders(zernike_pol) # get orders
+ m, n = define_orders(zernike_pol) # get orders
# Calculation of the value according to Ref.
if m == 0:
return np.sqrt(n + 1)
@@ -80,7 +76,7 @@ def normalization_factor(zernike_pol) -> float:
return np.sqrt(2*(n + 1))
-def radial_polynomial(zernike_pol, r):
+def radial_polynomial(zernike_pol, r: Union[float, np.ndarray]) -> Union[float, np.ndarray]:
"""
Calculate radial polynomial R(m, n) value for input r lying in the range [0, 1].
@@ -104,7 +100,7 @@ def radial_polynomial(zernike_pol, r):
Depending on the type of theta, return float or np.ndarray with calculated values of radial polynomial.
"""
- (m, n) = define_orders(zernike_pol) # get orders
+ m, n = define_orders(zernike_pol) # get orders
# Radial polynomials defined as analytical equations for up to 10th order (check tables from [2])
# 0th order
if (m == 0) and (n == 0):
@@ -205,7 +201,7 @@ def radial_polynomial(zernike_pol, r):
- radial_polynomial((m, n-2), r)) # general recurrence formula from [1]
-def radial_derivative(zernike_pol, r):
+def radial_derivative(zernike_pol, r: Union[float, np.ndarray]) -> Union[float, np.ndarray]:
"""
Calculate the derivative of radial polynomial dR(m, n)/dr value for input r lying in the range [0, 1].
@@ -222,7 +218,7 @@ def radial_derivative(zernike_pol, r):
Depending on the type of r, return float or np.ndarray with calculated values of radial polynomial derivative.
"""
- (m, n) = define_orders(zernike_pol) # get orders
+ m, n = define_orders(zernike_pol) # get orders
# Derivatives of radial polynomials defined as analytical equations for up to 8th order (check function above)
# 0th order
if (m == 0) and (n == 0):
@@ -327,7 +323,7 @@ def radial_derivative(zernike_pol, r):
- radial_derivative((m, n-2), r))
-def radial_polynomial_eq(zernike_pol, r):
+def radial_polynomial_eq(zernike_pol, r: Union[float, np.ndarray]) -> Union[float, np.ndarray]:
"""
Calculate the radial polynomial R(m, n) using exact equation from the Reference below.
@@ -348,11 +344,12 @@ def radial_polynomial_eq(zernike_pol, r):
Depending on the type of theta, return float or numpy.ndarray with calculated values of radial polynomial.
"""
+ value: Union[float, np.ndarray] # fix for mypy checks
if isinstance(r, float):
value = 0.0
elif isinstance(r, np.ndarray):
value = np.zeros(r.shape)
- (m, n) = define_orders(zernike_pol) # get orders
+ m, n = define_orders(zernike_pol) # get orders
for k in range(0, ((n - abs(m))//2) + 1):
a = (n + m)//2; b = (n - m)//2
value += ((-1)**k)*(((math.factorial(n-k))/(math.factorial(k)
@@ -362,7 +359,7 @@ def radial_polynomial_eq(zernike_pol, r):
return value
-def radial_derivative_eq(zernike_pol, r):
+def radial_derivative_eq(zernike_pol, r: Union[float, np.ndarray]) -> Union[float, np.ndarray]:
"""
Calculate the derivative of radial polynomial R(m, n) on r (eq. dR(m, n)/dr).
@@ -379,11 +376,12 @@ def radial_derivative_eq(zernike_pol, r):
Depending on the type of theta, return float or numpy.ndarray with calculated values of radial polynomial derivative.
"""
+ value: Union[float, np.ndarray] # fix for mypy checks
if isinstance(r, float):
value = 0.0
elif isinstance(r, np.ndarray):
value = np.zeros(r.shape)
- (m, n) = define_orders(zernike_pol) # get orders
+ m, n = define_orders(zernike_pol) # get orders
if n > 0:
for k in range(0, ((n - abs(m))//2) + 1):
a = (n + m)//2; b = (n - m)//2
@@ -395,7 +393,7 @@ def radial_derivative_eq(zernike_pol, r):
return value
-def triangular_function(zernike_pol, theta):
+def triangular_function(zernike_pol, theta: Union[float, np.ndarray]) -> Union[float, np.ndarray]:
"""
Return triangular component of the Zernike polynomial.
@@ -418,7 +416,7 @@ def triangular_function(zernike_pol, theta):
Depending on the type of theta, return float or numpy.ndarray with calculated values of triangular function.
"""
- (m, n) = define_orders(zernike_pol) # get orders
+ m, n = define_orders(zernike_pol) # get orders
# Calculation of the value according to Refs.
if m >= 0:
return np.cos(m*theta)
@@ -426,7 +424,7 @@ def triangular_function(zernike_pol, theta):
return -np.sin(m*theta)
-def triangular_derivative(zernike_pol, theta):
+def triangular_derivative(zernike_pol, theta: Union[float, np.ndarray]) -> Union[float, np.ndarray]:
"""
Return derivative of triangular function of the Zernike polynomial.
@@ -444,7 +442,7 @@ def triangular_derivative(zernike_pol, theta):
Depending on the type of theta, return float or numpy.ndarray with calculated values of derivative of triangular function.
"""
- (m_order, n) = define_orders(zernike_pol) # get orders
+ m_order, n = define_orders(zernike_pol) # get orders
# Calculation of the value according to the analytical derivative of the equation above
m = float(m_order)
if m_order >= 0:
@@ -453,7 +451,7 @@ def triangular_derivative(zernike_pol, theta):
return -m*np.cos(m*theta)
-def radial_polynomial_coeffs(zernike_pol, r):
+def radial_polynomial_coeffs(zernike_pol, r: Union[float, np.ndarray]) -> Union[float, np.ndarray]:
"""
Calculate radial polynomial using recursive finding algorithm of each coefficient for radial component (e.g. R^6).
@@ -481,6 +479,7 @@ def radial_polynomial_coeffs(zernike_pol, r):
if n <= MAX_RADIAL_ORDER_COEFFS:
pols_coeffs = find_coeffs_orders((m, n))
# Initial value for sum calculation
+ r_sum: Union[float, np.ndarray]
if isinstance(r, float):
r_sum = 0.0
elif isinstance(r, np.ndarray):
@@ -506,7 +505,7 @@ def radial_polynomial_coeffs(zernike_pol, r):
return r_sum
-def radial_polynomial_coeffs_dr(zernike_pol, r):
+def radial_polynomial_coeffs_dr(zernike_pol, r: Union[float, np.ndarray]) -> Union[float, np.ndarray]:
"""
Calculate radial polynomial derivative using recursive finding algorithm of each coefficient for radial component.
@@ -531,6 +530,7 @@ def radial_polynomial_coeffs_dr(zernike_pol, r):
if n <= MAX_RADIAL_ORDER_COEFFS_dR:
pols_coeffs = find_coeffs_orders_dr((m, n))
# Initial value for sum calculation
+ r_sum: Union[float, np.ndarray]
if isinstance(r, float):
r_sum = 0.0
elif isinstance(r, np.ndarray):
@@ -582,30 +582,21 @@ def compare_radial_calculations(max_order: int) -> np.ndarray:
if not isinstance(max_order, int) and max_order < 2:
print("NOTE that max_order by default set to 2")
max_order = 2
-
# Generating Zernike orders in OSA/ANSI indexing scheme
orders_list = [(0, 0)]
for order in range(1, max_order):
m = -order; n = order
orders_list.append((m, n))
- for n_azimuthals in range(0, order):
- m += 2
- orders_list.append((m, n))
-
- # Generation numpy array with radii
- n_points = 21
- test_r = np.zeros(shape=(n_points, ))
- for i in range(n_points):
- test_r[i] = i/(n_points-1)
- test_r = np.round(test_r, 4)
+ for _ in range(0, order):
+ m += 2; orders_list.append((m, n))
+ n_points = 21; test_r = np.linspace(start=0.0, stop=1.0, num=n_points) # radii points for testing
# Testing that exact calculation and implementation of tabular / recursive are the same
diff = np.ones(shape=(len(orders_list), n_points))
- for i, order in enumerate(orders_list):
- diff[i, :] = radial_polynomial(order, test_r) - radial_polynomial_eq(order, test_r)
+ for i, orders in enumerate(orders_list):
+ diff[i, :] = radial_polynomial(orders, test_r) - radial_polynomial_eq(orders, test_r)
diff = np.round(diff, 9)
- assert np.max(np.abs(diff)) < 1E-9, (f"Order {order} has inconsistency between tabular/recursion"
+ assert np.max(np.abs(diff)) < 1E-9, ("Computation of Rmn has inconsistency between tabular/recursion"
+ f" and exact implementations, diff: {np.max(np.abs(diff))}")
- print("Difference between analytical and implemented equations for Zernike pol-s is negligible, test passed")
return diff
@@ -631,30 +622,21 @@ def compare_radial_derivatives(max_order: int) -> np.ndarray:
if not isinstance(max_order, int) and max_order < 2:
print("NOTE that max_order by default set to 2")
max_order = 2
-
# Generating Zernike orders in OSA/ANSI indexing scheme
orders_list = [(0, 0)]
for order in range(1, max_order):
m = -order; n = order
orders_list.append((m, n))
- for n_azimuthals in range(0, order):
- m += 2
- orders_list.append((m, n))
-
- # Generation numpy array with radii
- n_points = 21
- test_r = np.zeros(shape=(n_points, ))
- for i in range(n_points):
- test_r[i] = i/(n_points-1)
- test_r = np.round(test_r, 4)
+ for _ in range(0, order):
+ m += 2; orders_list.append((m, n))
+ n_points = 21; test_r = np.linspace(start=0.0, stop=1.0, num=n_points) # radii points for testing
# Testing that exact calculation and implementation of tabular / recursive are the same
diff = np.ones(shape=(len(orders_list), n_points))
- for i, order in enumerate(orders_list):
- diff[i, :] = radial_derivative(order, test_r) - radial_derivative_eq(order, test_r)
+ for i, orders in enumerate(orders_list):
+ diff[i, :] = radial_derivative(orders, test_r) - radial_derivative_eq(orders, test_r)
diff = np.round(diff, 9)
- assert np.max(np.abs(diff)) < 1E-9, (f"Order {order} has inconsistency between tabular/recursion"
+ assert np.max(np.abs(diff)) < 1E-9, ("Computation of dRmn/dr has inconsistency between tabular/recursion"
+ f" and exact implementations, diff: {np.max(np.abs(diff))}")
- print("Difference between analytical and implemented equations for derivatives of Zernike pol-s is negligible, test passed")
return diff
@@ -677,27 +659,16 @@ def compare_recursive_coeffs_radials() -> np.ndarray:
for order in range(15, MAX_RADIAL_ORDER_COEFFS+2):
m = -order; n = order
orders_list.append((m, n))
- for n_azimuthals in range(0, order):
- m += 2
- orders_list.append((m, n))
- # Generation numpy array with radii
- n_points = 101
- test_r = np.zeros(shape=(n_points, ))
- for i in range(n_points):
- test_r[i] = i/(n_points-1)
- test_r = np.round(test_r, 4)
+ for _ in range(0, order):
+ m += 2; orders_list.append((m, n))
+ n_points = 101; test_r = np.linspace(start=0.0, stop=1.0, num=n_points) # radii points for testing
# Testing that exact calculation and implementation of tabular / recursive are the same
diff = np.ones(shape=(len(orders_list), n_points))
- for i, order in enumerate(orders_list):
- diff[i, :] = radial_polynomial_eq(order, test_r) - radial_polynomial_coeffs(order, test_r)
- # diff[i, :] = radial_polynomial_coeffs(order, test_r)
- # diff[i, :] = radial_polynomial_eq(order, test_r)
- # if np.max(np.abs(diff)) > 1.0:
- # print(order, np.max(np.abs(diff)))
+ for i, orders in enumerate(orders_list):
+ diff[i, :] = radial_polynomial_eq(orders, test_r) - radial_polynomial_coeffs(orders, test_r)
diff = np.round(diff, 9)
- assert np.max(np.abs(diff)) < 2E-2, (f"Order {order} has inconsistency between tabular/recursion"
+ assert np.max(np.abs(diff)) < 2E-2, ("Computation of Rmn has inconsistency between tabular/recursion"
+ f" and exact implementations, diff: {np.max(np.abs(diff))}")
- print("Difference between exact equation and pols. coeffs. finding algorithm is negligible, test passed")
return diff
@@ -720,25 +691,16 @@ def compare_recursive_coeffs_radials_dr() -> np.ndarray:
for order in range(15, MAX_RADIAL_ORDER_COEFFS_dR+2):
m = -order; n = order
orders_list.append((m, n))
- for n_azimuthals in range(0, order):
- m += 2
- orders_list.append((m, n))
- # Generation numpy array with radii
- n_points = 101
- test_r = np.zeros(shape=(n_points, ))
- for i in range(n_points):
- test_r[i] = i/(n_points-1)
- test_r = np.round(test_r, 4)
+ for _ in range(0, order):
+ m += 2; orders_list.append((m, n))
+ n_points = 101; test_r = np.linspace(start=0.0, stop=1.0, num=n_points) # radii points for testing
# Testing that exact calculation and implementation of tabular / recursive are the same
diff = np.ones(shape=(len(orders_list), n_points))
- for i, order in enumerate(orders_list):
- diff[i, :] = radial_derivative_eq(order, test_r) - radial_polynomial_coeffs_dr(order, test_r)
- # if np.max(np.abs(diff)) > 1.0:
- # print(order, np.max(np.abs(diff)))
+ for i, orders in enumerate(orders_list):
+ diff[i, :] = radial_derivative_eq(orders, test_r) - radial_polynomial_coeffs_dr(orders, test_r)
diff = np.round(diff, 9)
- assert np.max(np.abs(diff)) < 5E-2, (f"Order {order} has inconsistency between tabular/recursion"
+ assert np.max(np.abs(diff)) < 5E-2, ("Computation of Rmn has inconsistency between tabular/recursion"
+ f" and exact implementations, diff: {np.max(np.abs(diff))}")
- print("Difference between exact equation and pols. coeffs. finding algorithm is negligible, test passed")
return diff
@@ -757,22 +719,15 @@ def check_high_orders_recursion() -> np.ndarray:
for order in range(MAX_RADIAL_ORDER_COEFFS+1, MAX_RADIAL_ORDER_COEFFS+6):
m = -order; n = order
orders_list.append((m, n))
- for n_azimuthals in range(0, order):
- m += 2
- orders_list.append((m, n))
- # Generation numpy array with radii
- n_points = 51
- test_r = np.zeros(shape=(n_points, ))
- for i in range(n_points):
- test_r[i] = i/(n_points-1)
- test_r = np.round(test_r, 4)
+ for _ in range(0, order):
+ m += 2; orders_list.append((m, n))
+ n_points = 51; test_r = np.linspace(start=0.0, stop=1.0, num=n_points) # radii points for testing
# Testing that exact calculation and implementation of tabular / recursive are the same
diff = np.ones(shape=(len(orders_list), n_points))
- for i, order in enumerate(orders_list):
- diff[i, :] = radial_polynomial_coeffs(order, test_r)
+ for i, orders in enumerate(orders_list):
+ diff[i, :] = radial_polynomial_coeffs(orders, test_r)
diff = np.round(diff, 6)
- assert np.max(np.abs(diff)) <= 1.0, f"Order {order} has inconsistency in max abs value: {np.max(np.abs(diff))}"
- print("Abs max in calculated recursively radial polynomials <= 1.0, test passed")
+ assert np.max(np.abs(diff)) <= 1.0, f"Rmn has inconsistency in max abs value (should be <= 1.0): {np.max(np.abs(diff))}"
return diff
@@ -785,7 +740,7 @@ def time_radial_pols():
None.
"""
- calc_times_ms = [] # for storing calculation times
+ calc_times_ms: List[float] = [] # for storing calculation times
zp1 = (2, 16); zp2 = (0, 18); zp3 = (-2, 20); zp4 = (6, 22); zp5 = (-4, 24); r = 0.425
zpols = [zp1, zp2, zp3, zp4, zp5]
for zp in zpols:
@@ -826,19 +781,3 @@ def time_radial_pols():
calc_times_ms.append(round(1000*(t2-t1), 3))
print("Timed calc. using exact equation rad. pol.", zpols, ":", calc_times_ms)
-
-# %% Tests
-if __name__ == '__main__':
- R = [0.0, 0.2, 0.4, 0.6, 0.8, 1.0]; R = np.asarray(R)
- Theta = [i*np.pi/3 for i in range(6)]; Theta = np.asarray(Theta)
- orders = (-2, 2); r = 0.5
- ZR = radial_polynomial(orders, R); ZR1 = radial_polynomial(orders, r)
- TR = triangular_function(orders, Theta)
- diff = compare_radial_calculations(max_order=20)
- diff_deriv = compare_radial_derivatives(max_order=18)
- # Test specific implementations
- orders = (-8, 52); r = 0.95; val = radial_polynomial_coeffs(orders, r)
- time_radial_pols() # initial estimation of performance of calculations
- diff_coeffs = compare_recursive_coeffs_radials()
- diff_deriv_coeffs = compare_recursive_coeffs_radials_dr()
- high_R = check_high_orders_recursion()
diff --git a/src/zernpy/calculations/find_pols_coeffs.py b/src/zernpy/calculations/find_pols_coeffs.py
index 02f7753..1b10f0c 100644
--- a/src/zernpy/calculations/find_pols_coeffs.py
+++ b/src/zernpy/calculations/find_pols_coeffs.py
@@ -44,7 +44,7 @@ def make_orders_coeffs(defined_coeff: dict, max_order: int, minus: bool = False)
# Setting provided already defined coefficients to the initialized dictionary
if len(defined_coeff.keys()) > 0:
for key, value in defined_coeff.items():
- if key in coefficients.keys():
+ if key in coefficients:
if not minus:
coefficients[key] = value
else:
@@ -185,7 +185,7 @@ def check_special_orders(orders: tuple) -> dict:
Polynomials coefficients with radial order: value pairs.
"""
- special_coefficients = None
+ special_coefficients = {}
m, n = orders
if abs(m) == n:
special_coefficients = make_orders_coeffs({n: 1}, n)
@@ -209,7 +209,7 @@ def check_special_orders_dr(orders: tuple) -> dict:
Derivatives of polynomials coefficients with radial order: value pairs.
"""
- special_coefficients = None
+ special_coefficients = {}
m, n = orders
if abs(m) == n:
special_coefficients = make_orders_coeffs({n-1: n}, n-1)
@@ -248,16 +248,16 @@ def find_coeffs_orders(orders: tuple, use_test_dict: bool = False) -> dict:
initial_coefficients = initial_coefficients_test
else:
initial_coefficients = precalculated_initial_coeffs
- if orders in initial_coefficients.keys():
+ if orders in initial_coefficients:
# print("Found in initial dict.: ", initial_coefficients[orders])
return initial_coefficients[orders] # return stored in the dictionary value
- elif check_special_orders(orders) is not None:
+ elif check_special_orders(orders):
# Cashing already calculated coefficients in global dictionary specified above
if not use_test_dict:
- if orders not in precalculated_initial_coeffs.keys():
+ if orders not in precalculated_initial_coeffs:
precalculated_initial_coeffs[orders] = check_special_orders(orders)
else:
- if orders not in initial_coefficients_test.keys():
+ if orders not in initial_coefficients_test:
initial_coefficients_test[orders] = check_special_orders(orders)
return check_special_orders(orders) # some special shorthanded cases for polynomials values calculation
else:
@@ -269,14 +269,14 @@ def find_coeffs_orders(orders: tuple, use_test_dict: bool = False) -> dict:
max_order=n-2, minus=True))
# Cashing already calculated coefficients in global dictionary specified above
if not use_test_dict:
- if (abs(m-1), n-1) not in precalculated_initial_coeffs.keys():
+ if (abs(m-1), n-1) not in precalculated_initial_coeffs:
precalculated_initial_coeffs[(abs(m-1), n-1)] = polm1n1
- if (m+1, n-1) not in precalculated_initial_coeffs.keys():
+ if (m+1, n-1) not in precalculated_initial_coeffs:
precalculated_initial_coeffs[(m+1, n-1)] = polmP1n1
else:
- if (abs(m-1), n-1) not in initial_coefficients_test.keys():
+ if (abs(m-1), n-1) not in initial_coefficients_test:
initial_coefficients_test[(abs(m-1), n-1)] = polm1n1
- if (m+1, n-1) not in initial_coefficients_test.keys():
+ if (m+1, n-1) not in initial_coefficients_test:
initial_coefficients_test[(m+1, n-1)] = polmP1n1
return sum_dict_coeffs
@@ -311,16 +311,16 @@ def find_coeffs_orders_dr(orders: tuple, use_test_dict: bool = False) -> dict:
initial_coefficients_dr = initial_coefficients_test_dr
else:
initial_coefficients_dr = precalculated_initial_coeffs_dr
- if orders in initial_coefficients_dr.keys():
+ if orders in initial_coefficients_dr:
# print("Found in initial dict.: ", initial_coefficients[orders])
return initial_coefficients_dr[orders] # return stored in the dictionary value
- elif check_special_orders_dr(orders) is not None:
+ elif check_special_orders_dr(orders):
# Cashing already calculated coefficients in global dictionary specified above
if not use_test_dict:
- if orders not in precalculated_initial_coeffs_dr.keys():
+ if orders not in precalculated_initial_coeffs_dr:
precalculated_initial_coeffs_dr[orders] = check_special_orders_dr(orders)
else:
- if orders not in initial_coefficients_test_dr.keys():
+ if orders not in initial_coefficients_test_dr:
initial_coefficients_test_dr[orders] = check_special_orders_dr(orders)
return check_special_orders_dr(orders) # some special shorthanded cases for derivatives values calculation
else:
@@ -335,22 +335,22 @@ def find_coeffs_orders_dr(orders: tuple, use_test_dict: bool = False) -> dict:
max_order=n-2, minus=True))
# Cashing already calculated coefficients in global dictionary specified above
if not use_test_dict:
- if (abs(m-1), n-1) not in precalculated_initial_coeffs.keys():
+ if (abs(m-1), n-1) not in precalculated_initial_coeffs:
precalculated_initial_coeffs[(abs(m-1), n-1)] = polm1n1
- if (m+1, n-1) not in precalculated_initial_coeffs.keys():
+ if (m+1, n-1) not in precalculated_initial_coeffs:
precalculated_initial_coeffs[(m+1, n-1)] = polmP1n1
- if (abs(m-1), n-1) not in precalculated_initial_coeffs_dr.keys():
+ if (abs(m-1), n-1) not in precalculated_initial_coeffs_dr:
precalculated_initial_coeffs_dr[(abs(m-1), n-1)] = polm1n1_dr
- if (m+1, n-1) not in precalculated_initial_coeffs_dr.keys():
+ if (m+1, n-1) not in precalculated_initial_coeffs_dr:
precalculated_initial_coeffs_dr[(m+1, n-1)] = polmP1n1_dr
else:
- if (abs(m-1), n-1) not in initial_coefficients_test.keys():
+ if (abs(m-1), n-1) not in initial_coefficients_test:
initial_coefficients_test[(abs(m-1), n-1)] = polm1n1
- if (m+1, n-1) not in initial_coefficients_test.keys():
+ if (m+1, n-1) not in initial_coefficients_test:
initial_coefficients_test[(m+1, n-1)] = polmP1n1
- if (abs(m-1), n-1) not in initial_coefficients_test_dr.keys():
+ if (abs(m-1), n-1) not in initial_coefficients_test_dr:
initial_coefficients_test_dr[(abs(m-1), n-1)] = polm1n1_dr
- if (m+1, n-1) not in initial_coefficients_test_dr.keys():
+ if (m+1, n-1) not in initial_coefficients_test_dr:
initial_coefficients_test_dr[(m+1, n-1)] = polmP1n1_dr
return sum_dict_coeffs
diff --git a/src/zernpy/calculations/fit_zernike_pols.py b/src/zernpy/calculations/fit_zernike_pols.py
index 77144a7..8dfe7ef 100644
--- a/src/zernpy/calculations/fit_zernike_pols.py
+++ b/src/zernpy/calculations/fit_zernike_pols.py
@@ -7,11 +7,9 @@
"""
# %% Global imports
-import numpy as np
import warnings
-# %% Local imports
-
+import numpy as np
# %% Module parameters
__docformat__ = "numpydoc"
@@ -73,17 +71,17 @@ def crop_phases_img(phases_image: np.ndarray, crop_radius: float = 1.0, suppress
cropped_radii_vector = np.zeros(shape=(rows*cols,)); cropped_thetas_vector = np.zeros(shape=(rows*cols,))
# Check input image shape
if rows != cols:
- __warn_message = "Phases image isn't square, results of fitting could be ambiguous"
+ __warn_message = "\nPhase profile (image) isn't square, results of fitting could be ambiguous"
img_min_size = min(rows, cols); img_max_size = max(rows, cols)
if not suppress_warns:
- warnings.warn(__warn_message)
+ warnings.warn(__warn_message, stacklevel=2)
else:
img_min_size = rows; img_max_size = rows
if rows % 2 == 0 or cols % 2 == 0:
- __warn_message = ("Phases image provided with even rows or columns, "
+ __warn_message = ("\nPhase profile (image) provided with even rows or columns, "
+ "it's error prone to define exact image center")
if not suppress_warns:
- warnings.warn(__warn_message)
+ warnings.warn(__warn_message, stacklevel=2)
if img_min_size < 3:
raise ValueError("Provided image is too small (min size < 3) for producing any meaningful result")
# Calculate center of the input image - depending on its sizes
@@ -114,7 +112,7 @@ def crop_phases_img(phases_image: np.ndarray, crop_radius: float = 1.0, suppress
+ " or with 1 pixel difference between them. \n"
+ "Note that additional border pixels are included to crop.")
if not suppress_warns:
- warnings.warn(__warn_message)
+ warnings.warn(__warn_message, stacklevel=2)
# Calculate polar coordinates of pixels
recalibrate_radii = False; recalibration_coeff = 1.0; vector_index = 0
# Cropping phases, collecting data on cropping process
@@ -185,7 +183,7 @@ def fit_zernikes(phases_coordinates_vectors: tuple, polynomials: tuple) -> np.nd
# Calculate polynomials values in the unit circle defined by polar coordinates
zernike_values = np.zeros(shape=(vector_length, len(polynomials)))
# Checking that all polynomials are unique
- if not len(polynomials) == 0:
+ if len(polynomials) != 0:
provided_orders = []
for polynomial in polynomials:
provided_orders.append(polynomial.get_mn_orders())
@@ -212,28 +210,3 @@ def fit_zernikes(phases_coordinates_vectors: tuple, polynomials: tuple) -> np.nd
else:
raise ValueError("There isn't any polynomials provided")
return zernike_coefficients
-
-
-# %% Basic tests
-if __name__ == "__main__":
- phases_sample = 4*np.ones(shape=(3, 4), dtype="uint8")
- crop_deforms1, polar_coordinates1 = crop_phases_img(phases_sample)
- phases_sample = np.ones(shape=(4, 4), dtype="int16")
- crop_deforms2, polar_coordinates2 = crop_phases_img(phases_sample)
- crop_deforms2a, polar_coordinates2a = crop_phases_img(phases_sample, strict_border=True)
- phases_sample = np.ones(shape=(6, 6))
- crop_deforms3, polar_coordinates = crop_phases_img(phases_sample)
- phases_sample = np.ones(shape=(3, 3))
- crop_deforms4, polar_coordinates = crop_phases_img(phases_sample)
- phases_sample = np.ones(shape=(5, 5))
- crop_deforms5, polar_coordinates = crop_phases_img(phases_sample)
- phases_sample = np.ones(shape=(4, 6))
- crop_deforms6, polar_coordinates6 = crop_phases_img(phases_sample)
- crop_deforms6a, polar_coordinates6a = crop_phases_img(phases_sample, strict_border=True)
- phases_sample = np.ones(shape=(7, 6))
- crop_deforms7, polar_coordinates7 = crop_phases_img(phases_sample)
- crop_deforms7a, polar_coordinates7a = crop_phases_img(phases_sample, strict_border=True)
- phases_sample = np.ones(shape=(5, 5))
- crop_deforms81, polar_coordinates = crop_phases_img(phases_sample)
- crop_deforms82, polar_coordinates = crop_phases_img(phases_sample, crop_radius=0.5)
- thetas_grads = polar_coordinates[1]*(180.0/np.pi)
diff --git a/src/zernpy/examples/define_pv_values.py b/src/zernpy/examples/define_pv_values.py
index 31d1a70..b25fd3f 100644
--- a/src/zernpy/examples/define_pv_values.py
+++ b/src/zernpy/examples/define_pv_values.py
@@ -7,6 +7,7 @@
"""
# %% Global imports
import numpy as np
+
try:
from zernpy import generate_polynomials
zernpy_installed = True
diff --git a/src/zernpy/plotting/plot_zerns.py b/src/zernpy/plotting/plot_zerns.py
index 01183b4..d681948 100644
--- a/src/zernpy/plotting/plot_zerns.py
+++ b/src/zernpy/plotting/plot_zerns.py
@@ -7,8 +7,8 @@
"""
# %% Global imports
-import numpy as np
import matplotlib.pyplot as plt
+import numpy as np
# %% Module parameters
__docformat__ = "numpydoc"
diff --git a/src/zernpy/tests/test_calc_psf.py b/src/zernpy/tests/test_calc_psf.py
index 4864b87..969db34 100644
--- a/src/zernpy/tests/test_calc_psf.py
+++ b/src/zernpy/tests/test_calc_psf.py
@@ -9,21 +9,26 @@
@licence: MIT
"""
-import numpy as np
-from pathlib import Path
import os
+from pathlib import Path
+
+import numpy as np
-# Importing the written in the modules test functions for letting pytest library their automatic exploration
-if __name__ != "__main__":
- from ..calculations.calc_psfs import (get_psf_kernel)
- from ..zernpsf import ZernPSF, force_get_psf_compilation
- from ..zernikepol import ZernPol
-else:
- from zernpy import ZernPol
+# Relative import from the modules - designed to be used only by pytest runs
+from ..calculations.calc_psfs import get_psf_kernel
+from ..zernikepol import ZernPol
+from ..zernpsf import ZernPSF, clean_zernpy_cache, force_get_psf_compilation
# Testing functions
def test_psf_kernel_calc():
+ """
+ Test common calculation of PSF kernel + comparison with the result of exact equation - Airy pattern.
+
+ Returns
+ -------
+ None
+ """
NA = 0.35; wavelength = 0.55; pixel_size = wavelength / 3.05; ampl = -0.28 # Common physical properties
# Basic test - calculating kernel by numerical integration and by Airy pattern exact equation and compare them
piston = ZernPol(m=0, n=0)
@@ -52,6 +57,13 @@ def test_psf_kernel_calc():
def test_zernpsf_usage():
+ """
+ Test wrong initialization parameters for ZernPSF, at the end - normal one.
+
+ Returns
+ -------
+ None
+ """
NA = 0.95; wavelength = 0.55; pixel_size = wavelength / 5.0; ampl = 0.55 # Common physical properties
zp1 = ZernPol(m=0, n=2) # defocus
test_init = True # flag, if set to False, the test failed (e.g., expected Error not caught)
@@ -101,6 +113,13 @@ def test_zernpsf_usage():
def test_save_load_zernpsf():
+ """
+ Test saving and loading of PSF kernel stored in a json file.
+
+ Returns
+ -------
+ None
+ """
zp2 = ZernPol(m=1, n=3); zpsf2 = ZernPSF(zp2) # horizontal coma
zp3 = ZernPol(osa=21); zpsf3 = ZernPSF(zp3) # additional classes for testing reading and reassigning values
NA = 0.4; wavelength = 0.4; pixel_size = wavelength / 3.0; ampl = 0.11 # Common physical properties
@@ -118,7 +137,14 @@ def test_save_load_zernpsf():
def test_numba_compilation():
- force_get_psf_compilation() # force compilation of computation methods
+ """
+ Test compilation by numba of the calculation functions.
+
+ Returns
+ -------
+ None
+ """
+ force_get_psf_compilation() # force compilational of computation methods
# Test the difference between accelerated and not accelerated calculation methods
NA = 0.95; wavelength = 0.55; pixel_size = wavelength / 4.25; ampl = -0.12
zp6 = ZernPol(m=0, n=2); zpsf6 = ZernPSF(zp6); zpsf7 = ZernPSF(zp6) # defocus
@@ -126,4 +152,5 @@ def test_numba_compilation():
zpsf7.set_physical_props(NA=NA, wavelength=wavelength, expansion_coeff=ampl, pixel_physical_size=pixel_size)
kernel_acc = zpsf6.calculate_psf_kernel(normalized=True, accelerated=True) # accelerated by numba compilation
kernel_norm = zpsf7.calculate_psf_kernel(normalized=True) # normal calculation
- assert np.max(np.abs(kernel_acc - kernel_norm) < 1E-6), "Accelerated and not one calculation of kernel methods have significant differences"
+ assert np.max(np.abs(kernel_acc - kernel_norm) < 1E-6), "Accelerated and not computational of a kernel methods have significant differences"
+ assert clean_zernpy_cache(), "Local cache with compiled files not cleaned"
diff --git a/src/zernpy/tests/test_calculations.py b/src/zernpy/tests/test_calculations.py
index a8d76c6..78b1d26 100644
--- a/src/zernpy/tests/test_calculations.py
+++ b/src/zernpy/tests/test_calculations.py
@@ -1,6 +1,6 @@
# -*- coding: utf-8 -*-
"""
-Test the implemented calculation functions for module calc_zernike_pol and ZernPol static methods by using pytest library.
+Test the implemented calculation functions for the module 'calc_zernike_pol' and the static methods from 'ZernPol' by using pytest library.
The pytest library available on: https://docs.pytest.org/en/latest/contents.html
For running collected here tests, it's enough to run the command "pytest" from the repository location in the command line.
@@ -10,45 +10,75 @@
"""
import math
+
import numpy as np
-# Importing the written in the modules test functions for letting pytest library their automatic exploration
-if __name__ != "__main__":
- from ..calculations.calc_zernike_pol import (compare_radial_calculations, compare_radial_derivatives,
- compare_recursive_coeffs_radials, compare_recursive_coeffs_radials_dr,
- check_high_orders_recursion)
- from ..zernikepol import ZernPol
+from ..calculations.calc_zernike_pol import (
+ check_high_orders_recursion,
+ compare_radial_calculations,
+ compare_radial_derivatives,
+ compare_recursive_coeffs_radials,
+ compare_recursive_coeffs_radials_dr,
+)
+from ..zernikepol import ZernPol
-# Testing implemented equations for tabular R(m, n) from References by comparing with the exact ones (with factorials)
+# %% Func-s to be run by pytest
def test_tabular_orders():
- compare_radial_calculations(max_order=10)
+ """
+ # Test implemented equations for tabular R(m, n) from the References by comparing with the exact ones (with factorials usage).
+
+ Returns
+ -------
+ None
+ """
+ compare_radial_calculations(max_order=17)
-# Testing implemented tabular and recursive equations for R(m, n) by comparing with the exact ones (with factorials)
-# Also, testing implemented recursive scheme for finding the polynomials coefficients for each order and comparison
-# with the exact equations (with factorials)
def test_recursive_orders():
- compare_radial_calculations(max_order=17)
+ """
+ Test implemented recursive scheme for finding the polynomials coefficients for each order and comparison with the exact equations.
+
+ Exact equations are based on factorials usage for coefficients calculation.
+
+ Returns
+ -------
+ None
+ """
compare_recursive_coeffs_radials()
check_high_orders_recursion()
-# Testing derived equations for derivatives dR(m, n)/dr by comparing with the exact ones (with factorials)
def test_tabular_derivatives():
- compare_radial_derivatives(max_order=10)
+ """
+ Test derived equations for derivatives dR(m, n)/dr by comparing with the exact ones (with factorials).
+
+ Returns
+ -------
+ None
+ """
+ compare_radial_derivatives(max_order=17)
-# Testing recursive and derived equations for derivatives dR(m, n)/dr by comparing with the exact ones (with factorials)
-# Also, testing implemented recursive scheme for finding the polynomials coefficients for each order and comparison
-# with the exact equations (with factorials) - derivative cases
def test_recursive_derivatives():
- compare_radial_derivatives(max_order=17)
+ """
+ Test implemented recursive scheme for finding derivatives dR(m, n)/dr for each order and comparison with the exact equations.
+
+ Returns
+ -------
+ None
+ """
compare_recursive_coeffs_radials_dr()
-# Testing sum of Zernike polynomials
def test_sum_zernikes():
+ """
+ Test summing of several Zernike polynomials, generation of phase profile and plotting method.
+
+ Returns
+ -------
+ None
+ """
zp1 = ZernPol(osa_index=1); zp2 = ZernPol(noll_index=2); ampls = [-0.5, 0.5]; theta = math.pi/3; r = 1.0
sum_pols = ZernPol.sum_zernikes(coefficients=ampls, polynomials=[zp1, zp2], r=r, theta=theta)
sum_pols_manual = math.cos(theta) - math.sin(theta) # manual calculation of specified above sum
@@ -59,14 +89,6 @@ def test_sum_zernikes():
assert isinstance(zern_surface, tuple) and isinstance(zern_surface.ZernSurf, np.ndarray), ("Check gen_zernikes_surface()"
+ " method output (tuple len=3)")
assert len(zern_surface.ZernSurf.shape) == 2, "Check gen_zernikes_surface() method for output matrix shape"
- try:
- from matplotlib.figure import Figure
- fig = Figure()
- plotted_fig = ZernPol.plot_sum_zernikes_on_fig(coefficients=ampls, polynomials=[zp1, zp2],
- figure=fig, zernikes_sum_surface=zern_surface)
- assert isinstance(plotted_fig, Figure) and plotted_fig.tight_layout, "Something wrong with the plotting function"
- except ModuleNotFoundError:
- assert False, "Install matplotlib for passing the test"
# Test difference between direct (naive) implementation and using meshgrids implementation of sum of Zernikes
pols = [ZernPol(osa=2), ZernPol(osa=4), ZernPol(osa=7), ZernPol(osa=10), ZernPol(osa=15)]
ampls = [-0.85, 0.85, 0.24, -0.37, 1.0]; radii = np.arange(start=0.0, stop=1.0 + 0.05, step=0.05)
@@ -79,8 +101,14 @@ def test_sum_zernikes():
+ f" and have abs(min) = {abs(np.min(sum_pols_d - sum_pols))}")
-# Testing the calculation of polynomial values for edge cases
def test_pol_values_edge_cases():
+ """
+ Test edge cases for proper denial / reporting.
+
+ Returns
+ -------
+ None
+ """
zp = ZernPol(osa=8) # test polynomial, could be any
# Combination list + float
r = [0, 0.1, 0.2, 0.5]; theta = 0.2; values = zp.polynomial_value(r, theta)
diff --git a/src/zernpy/tests/test_fitting.py b/src/zernpy/tests/test_fitting.py
index 4fef08a..15c6be5 100644
--- a/src/zernpy/tests/test_fitting.py
+++ b/src/zernpy/tests/test_fitting.py
@@ -10,12 +10,11 @@
"""
# %% Global imports
-import numpy as np
import random
-# %% Imports from modules
-if __name__ != "__main__":
- from ..zernikepol import generate_random_phases, fit_polynomials, ZernPol, fit_polynomials_vectors, generate_phases_image
+import numpy as np
+
+from ..zernikepol import ZernPol, fit_polynomials, fit_polynomials_vectors, generate_phases_image, generate_random_phases
# %% Test functions
@@ -62,7 +61,7 @@ def test_random_fitting():
assert rmse_percentage <= (mdp//2) + 1, ("RMSE between fitted and randomly generated polynomials:"
+ f" {rmse_percentage} > assumed value {(mdp//2) + 1}")
# Test the sign of fitted and randomly generated amplitudes
- for i in range(2): # run tests
+ for _ in range(2): # run tests
random_phases_image, random_amplitudes, polynomials_tuple = generate_random_phases(img_height=101, img_width=101)
fitted_amplitudes, _ = fit_polynomials(random_phases_image, polynomials_tuple)
max_fit_ampl = np.max(fitted_amplitudes); max_src_ampl = np.max(random_amplitudes)
@@ -73,9 +72,7 @@ def test_random_fitting():
else:
src_ampl = max_src_ampl; fit_ampl = max_fit_ampl
same_sign = False
- if src_ampl < 0.0 and fit_ampl < 0.0:
- same_sign = True
- elif src_ampl > 0.0 and fit_ampl > 0.0:
+ if src_ampl < 0.0 and fit_ampl < 0.0 or src_ampl > 0.0 and fit_ampl > 0.0:
same_sign = True
assert same_sign, ("\n Fitted and source amplitudes have abs. maximum values with different"
+ f" signs: source ampl.: {src_ampl}, fitted ampl.: {fit_ampl}")
diff --git a/src/zernpy/tests/test_polynomials_initialization.py b/src/zernpy/tests/test_polynomials_initialization.py
index 9dd12c6..0b3685a 100644
--- a/src/zernpy/tests/test_polynomials_initialization.py
+++ b/src/zernpy/tests/test_polynomials_initialization.py
@@ -11,18 +11,95 @@
"""
import math
-# Importing the written in the modules test functions for letting pytest library their automatic exploration
-if __name__ != "__main__":
- from ..zernikepol import check_conformity, ZernPol
+import numpy as np
+
+from ..zernikepol import ZernPol
-# Testing initialization of Zernike polynomials, for details, see the zernikepol module
def test_polynomials_initialization():
- check_conformity()
+ """
+ Test various initialization ways of Zernike polynomials, for details, see the imported function.
+
+ Returns
+ -------
+ None
+ """
+ zp = ZernPol(m=-2, n=2) # Initialization with orders
+ (m1, n1), osa_i, noll_i, fringe_i = zp.get_indices()
+ assert (osa_i == 3 and noll_i == 5 and fringe_i == 6), (f"Check consistency of Z{(m1, n1)} indices: "
+ + f"OSA: {osa_i}, Noll: {noll_i}, Fringe: {fringe_i}")
+ zp = ZernPol(l=-3, n=5)
+ (m2, n2), osa_i, noll_i, fringe_i = zp.get_indices()
+ assert (osa_i == 16 and noll_i == 19 and fringe_i == 20), (f"Check consistency of Z{(m2, n2)} indices: "
+ + f"OSA: {osa_i}, Noll: {noll_i}, Fringe: {fringe_i}")
+ assert len(zp.get_polynomial_name(short=True)) > 0, f"Short name for Z{(m2, n2)} is zero length"
+ zp = ZernPol(azimuthal_order=-1, radial_order=5)
+ (m3, n3), osa_i, noll_i, fringe_i = zp.get_indices()
+ assert (osa_i == 17 and noll_i == 17 and fringe_i == 15), (f"Check consistency of Z{(m3, n3)} indices: "
+ + f"OSA: {osa_i}, Noll: {noll_i}, Fringe: {fringe_i}")
+ assert len(zp.get_polynomial_name()) > 0, f"Name for Z{(m3, n3)} is zero length"
+ m4, n4 = zp.get_mn_orders()
+ assert m4 == m3 and n3 == n4, f"Check method get_mn_orders() for Z{(m3, n3)}"
+ print(f"Initialization of polynomials Z{(m1, n1)}, Z{(m2, n2)}, Z{(m3, n3)} tested")
+ osa_i = 12; zp = ZernPol(osa_index=osa_i) # Initialization with OSA index
+ m, n = zp.get_mn_orders()
+ assert (m == 0 and n == 4), f"Check consistency of Z[OSA index = {osa_i}] orders {m, n}"
+ assert zp.get_fringe_index(m, n) == 9, f"Check consistency of Z[OSA index = {osa_i}] Fringe index"
+ assert zp.get_noll_index(m, n) == 11, f"Check consistency of Z[OSA index = {osa_i}] Noll index"
+ print(f"Initialization of polynomial Z[OSA index = {osa_i}] tested")
+ noll_i = 10 # Testing static methods
+ assert ZernPol.noll2osa(noll_i) == 9, f"Check consistency of Noll index {noll_i} conversion to OSA index"
+ assert ZernPol.osa2fringe(ZernPol.noll2osa(noll_i)) == 10, ("Check consistency of Noll "
+ + f"index {noll_i} conversion to OSA index")
+ print(f"Conversion of Noll index {noll_i} to OSA and Fringe indices tested")
+ # Test for not proper initialization
+ try:
+ m_f = 2; n_f = -2
+ zp = ZernPol(m=m_f, n=n_f)
+ asserting_value = False
+ except ValueError:
+ print(f"Polynomial Z{(m_f, n_f)} haven't been initialized, test passed")
+ asserting_value = True
+ assert asserting_value, f"Polynomial Z{(m_f, n_f)} initialized with wrong orders assignment"
+ # Testing input parameters for calculation
+ zp = ZernPol(m=0, n=2); r = 0.0; theta = math.pi
+ assert abs(zp.polynomial_value(r, theta) + math.sqrt(3)) < 1E-6, f"Check value of Z[{m}, {n}]({r}, {theta})"
+ zp = ZernPol(m=-1, n=1); r = 0.5; theta = math.pi/2
+ assert abs(zp.polynomial_value(r, theta) - 1.0) < 1E-6, f"Check value of Z[{m}, {n}]({r}, {theta})"
+ print("Simple values of Zernike polynomials tested successfully")
+ try:
+ r = 'd'; theta = [1, 2]
+ zp.polynomial_value(r, theta)
+ asserting_value = False
+ except ValueError:
+ print("Input as string is not allowed for calculation of polynomial value, tested successfully")
+ asserting_value = True
+ assert asserting_value, "Wrong parameter passed (string) for calculation of polynomial value"
+ try:
+ r = [0.1, 0.2, 1.0+1E-9]; theta = math.pi
+ zp.polynomial_value(r, theta)
+ asserting_value = False
+ except ValueError:
+ print("Radius more than 1.0 is not allowed, tested successfully")
+ asserting_value = True
+ assert asserting_value, "Wrong parameter passed (r > 1.0) for calculation of polynomial value"
+ # Compare two implementations of Zernike pol-s sum calculation: direct and using meshgrid
+ pols = [ZernPol(osa=2), ZernPol(osa=4), ZernPol(osa=7), ZernPol(osa=10), ZernPol(osa=15),
+ ZernPol(osa=3), ZernPol(osa=9), ZernPol(osa=12), ZernPol(osa=16), ZernPol(osa=19)]
+ ampls = [-0.85, 0.85, 0.24, -0.37, 1.0, 0.1, -1.0, -0.05, 1.1, 0.41]
+ radii = np.arange(start=0.0, stop=1.0 + 0.001, step=0.001); thetas = np.arange(start=0.0, stop=2.0*np.pi + np.pi/180, step=np.pi/180)
+ ZernPol.sum_zernikes(ampls, pols, radii, thetas, get_surface=True)
+ ZernPol._sum_zernikes_meshgrid(ampls, pols, radii, thetas)
-# Explicit testing initialization of Zernike polynomials
def test_explicit_initialization():
+ """
+ Test particular initialization scenarios of Zernike polynomials.
+
+ Returns
+ -------
+ None
+ """
# Testing the ordinary, normal initialization of polynomials
m = 0; n = 2; zp = ZernPol(l=m, n=n)
assert abs(zp.radial_dr(0.25) - 1.0) < 1E-9, f"Radial derivative calculated with error for Z{(m, n)}"
@@ -167,7 +244,7 @@ def test_explicit_initialization():
assert zp1 == zp2, "Implemented method '==' isn't correct"
zp1 = ZernPol(fringe=21); zp2 = ZernPol(noll=8)
- assert not zp1 == zp2, "Implemented method '==' isn't correct"
+ assert zp1 != zp2, "Implemented method '==' isn't correct"
zp1 = ZernPol(fringe=17); zp2 = ZernPol(osa=14)
assert zp1 == zp2, "Implemented method '==' isn't correct"
diff --git a/src/zernpy/utils/intmproc.py b/src/zernpy/utils/intmproc.py
index 84107fe..7d62963 100644
--- a/src/zernpy/utils/intmproc.py
+++ b/src/zernpy/utils/intmproc.py
@@ -7,11 +7,13 @@
"""
# %% Global imports
-from multiprocessing import Process, Event, Queue
import os
-import warnings
import time
-from typing import Callable, Any, Union # Callable[..., Any] - should check for any callable instance accepting / returning any types
+import warnings
+from multiprocessing import Event, Process, Queue
+from typing import Any, Callable, List, Optional, Union
+
+# Callable[..., Any] - providing any callable instance accepting / returning any types
# from collections.abc import Sequence # for more broad typing check - Sequence[Any] | np.ndarray - accept list, set, tuple, str or np.ndarray
@@ -19,11 +21,13 @@
class DispenserManager():
"""Manager for the distributing calculation job on several Processes."""
- __MAX_WORKERS: int = os.cpu_count(); workers_number: int = __MAX_WORKERS // 2; __warn_message: str = ""
- __workers_pool: list = []; __results: list = []; __triggers: list = []; __initialized: bool = False
- __global_live_trigger: Event = None; __queues: list = []; __parameters_vector: list = []
+ __cpu_count = os.cpu_count()
+ __MAX_WORKERS: int = __cpu_count if __cpu_count is not None else 1
+ workers_number: int = __MAX_WORKERS // 2; __warn_message: str = ""
+ __workers_pool: list = []; __results: list = []; __triggers: List[Any] = []; __initialized: bool = False
+ __global_live_trigger: Any = None; __queues: list = []; __parameters_vector: list = []
- def __init__(self, compute_func: Callable[..., Any], params_list: list, n_workers: int = None, verbose_info: bool = False):
+ def __init__(self, compute_func: Callable[..., Any], params_list: list, n_workers: Optional[int] = None, verbose_info: bool = False):
"""
Initialize main handling class along with provided (automatically calculated) number of workers.
@@ -61,13 +65,13 @@ def __init__(self, compute_func: Callable[..., Any], params_list: list, n_worker
if n_workers > self.__MAX_WORKERS:
self.__warn_message = (f"Provided more workers than available detected by os.cpu_count() method: {self.__MAX_WORKERS}. "
+ f"By default, number of workers is set to: {self.workers_number}")
- warnings.warn(self.__warn_message)
+ warnings.warn(self.__warn_message, stacklevel=2)
else:
self.workers_number = n_workers
elif n_workers is not None and n_workers < 2:
self.__warn_message = ("Provided number of workers is meaningless for spreading jobs on workers. "
+ f"Default number of workers will be used: {self.workers_number}")
- warnings.warn(self.__warn_message)
+ warnings.warn(self.__warn_message, stacklevel=2)
# Checking input parameters - compute function / method
if not callable(compute_func):
raise TypeError("Provided computing function isn't callable (not callable(compute_func) - True)")
@@ -81,8 +85,8 @@ def __init__(self, compute_func: Callable[..., Any], params_list: list, n_worker
# Initializing the pool with Processes
self.__global_live_trigger = Event(); self.__global_live_trigger.set(); time.sleep(0.01)
# print(f"Initializing {self.workers_number} Processes")
- for i in range(self.workers_number):
- trigger_event = Event(); queue = Queue(); self.__triggers.append(trigger_event)
+ for _ in range(self.workers_number):
+ trigger_event = Event(); queue: Queue = Queue(); self.__triggers.append(trigger_event)
worker = IndiWorker(keep_run_trigger=self.__global_live_trigger, trigger=trigger_event, data_queue=queue, compute_func=compute_func)
worker.start(); received_confirmation = False # for waiting of initializaton of the Process
while not received_confirmation:
@@ -92,7 +96,7 @@ def __init__(self, compute_func: Callable[..., Any], params_list: list, n_worker
received_confirmation = True; break
else:
self.__warn_message = f"Not recognized confirmation message '{confirmation}', contact the developer"
- warnings.warn(self.__warn_message); break
+ warnings.warn(self.__warn_message, stacklevel=2); break
else:
time.sleep(0.005)
self.__workers_pool.append(worker); self.__queues.append(queue)
@@ -125,8 +129,9 @@ def compute(self):
"""
if self.__initialized:
computed_tasks = 0; init_step = True; indices2process = [i for i in range(len(self.__parameters_vector))]
- processing_indices = [None]*len(self.__workers_pool); task_assigned = [False]*len(self.__workers_pool); self.done_jobs_percentage = 0
- while computed_tasks < len(self.__results):
+ processing_indices = [0]*len(self.__workers_pool); task_assigned = [False]*len(self.__workers_pool)
+ self.done_jobs_percentage = 0
+ while computed_tasks < len(self.__results): # noqa: SIM113
if init_step:
current_param_index = 0
for proc_i in range(len(self.__workers_pool)):
@@ -166,8 +171,8 @@ def compute(self):
self.done_jobs_percentage = done_jobs_percent
return self.__results
else:
- self.__warn_message = ("All workers have been released, no computation has been performed")
- warnings.warn(self.__warn_message); return None
+ self.__warn_message = ("\nAll workers have been released, no computation has been performed")
+ warnings.warn(self.__warn_message, stacklevel=2); return None
def close(self):
"""
@@ -236,9 +241,9 @@ def update_params(self, params_list: list):
class IndiWorker(Process):
"""Individual Worker based on Process() class for performing computations using provided function."""
- __keep_event: Event = None; __trigger: Event; __data_queue: Queue
+ __keep_event: Any = None; __trigger: Any; __data_queue: Any
- def __init__(self, keep_run_trigger: Event, trigger: Event, data_queue: Queue, compute_func):
+ def __init__(self, keep_run_trigger: Any, trigger: Any, data_queue: Any, compute_func: Callable[..., Any]):
"""
Process from multiprocessing based class.
@@ -275,11 +280,6 @@ def run(self):
while self.__keep_event.is_set():
self.__trigger.wait() # wait that the parameter is transferred
if self.__trigger.is_set() and self.__keep_event.is_set():
- # attempts = 0
- # if not self.__data_queue.empty() and attempts < 11:
- # parameter = self.__data_queue.get_nowait()
- # else:
- # time.sleep(0.001); attempts += 1
parameter = self.__data_queue.get()
# print(f"Start computing for parameter {parameter}", flush=True)
result = self.computation(parameter) # computation work
@@ -335,7 +335,7 @@ def test_minus(x: Union[float, int]) -> Union[float, int]:
if __name__ == "__main__":
params = [10*(i+1) for i in range(177)] # ... computation points
# Direct computation - for loop for performance check
- t1 = time.perf_counter(); results_for = [None]*len(params)
+ t1 = time.perf_counter(); results_for = [0.0]*len(params)
for i, par in enumerate(params):
results_for[i] = test_plus(par)
elapsed_ms = int(round(1000.0*(time.perf_counter() - t1), 0))
diff --git a/src/zernpy/zernikepol.py b/src/zernpy/zernikepol.py
index f436eac..5931df8 100644
--- a/src/zernpy/zernikepol.py
+++ b/src/zernpy/zernikepol.py
@@ -4,40 +4,40 @@
Also, provides a few functions useful for fitting set of Zernike polynomials to an image with phases.
-@author: Sergei Klykov, @year: 2026, @licence: MIT \n
+@author: Sergei Klykov, @year: 2026, @license: MIT \n
"""
# %% Global imports
-import numpy as np
-from pathlib import Path
+import random
import warnings
-import math
from collections import namedtuple
+from typing import Optional, Sequence, Tuple, Union
+
import matplotlib.pyplot as plt
-import random
-import time
-from typing import Union, Sequence, Tuple
+import numpy as np
# %% Local (package-scoped) imports
-if __name__ == "__main__" or __name__ == Path(__file__).stem or __name__ == "__mp_main__":
- from calculations.calc_zernike_pol import (normalization_factor, radial_polynomial, triangular_function, triangular_derivative,
- radial_derivative, radial_polynomial_eq, radial_derivative_eq, radial_polynomial_coeffs,
- radial_polynomial_coeffs_dr, MAX_RADIAL_ORDER_COEFFS, MAX_RADIAL_ORDER_COEFFS_dR)
- from plotting.plot_zerns import plot_sum_fig, subplot_sum_on_fig, plot_sum_fig_3d, subplot_sum_on_fig_3d
- from calculations.fit_zernike_pols import crop_phases_img, fit_zernikes
- from props.properties import polynomial_names, short_polynomial_names, warn_mess_r_long, warn_mess_dr_long, warn_mess_slow_calc
-else:
- from .calculations.calc_zernike_pol import (normalization_factor, radial_polynomial, triangular_function, triangular_derivative,
- radial_derivative, radial_polynomial_eq, radial_derivative_eq, radial_polynomial_coeffs,
- radial_polynomial_coeffs_dr, MAX_RADIAL_ORDER_COEFFS, MAX_RADIAL_ORDER_COEFFS_dR)
- from .plotting.plot_zerns import plot_sum_fig, subplot_sum_on_fig, plot_sum_fig_3d, subplot_sum_on_fig_3d
- from .calculations.fit_zernike_pols import crop_phases_img, fit_zernikes
- from .props.properties import polynomial_names, short_polynomial_names, warn_mess_r_long, warn_mess_dr_long, warn_mess_slow_calc
+from .calculations.calc_zernike_pol import (
+ MAX_RADIAL_ORDER_COEFFS,
+ MAX_RADIAL_ORDER_COEFFS_dR,
+ normalization_factor,
+ radial_derivative,
+ radial_derivative_eq,
+ radial_polynomial,
+ radial_polynomial_coeffs,
+ radial_polynomial_coeffs_dr,
+ radial_polynomial_eq,
+ triangular_derivative,
+ triangular_function,
+)
+from .calculations.fit_zernike_pols import crop_phases_img, fit_zernikes
+from .plotting.plot_zerns import plot_sum_fig, plot_sum_fig_3d, subplot_sum_on_fig, subplot_sum_on_fig_3d
+from .props.properties import polynomial_names, short_polynomial_names, warn_mess_dr_long, warn_mess_r_long, warn_mess_slow_calc
# %% Module parameters
__docformat__ = "numpydoc"
-polar_vectors = namedtuple("PolarVectors", "R Theta") # re-used below as the return type from the method
-zernikes_surface = namedtuple("ZernikesSurface", "ZernSurf R Theta") # used as the input type
+polar_vectors = namedtuple("polar_vectors", "R Theta") # re-used below as the return type from the method
+zernikes_surface = namedtuple("zernikes_surface", "ZernSurf R Theta") # used as the input type
# %% Zernike Pol. class
@@ -86,30 +86,29 @@ def __init__(self, **kwargs):
key = ""
# Zernike polynomial specified with key arguments m, n - check firstly for these parameters
if len(kwargs.keys()) == 2:
- if "n" in kwargs.keys() or "radial_order" in kwargs.keys():
+ if "n" in kwargs or "radial_order" in kwargs:
if self.__initialized:
raise ValueError("The polynomial has been already initialized, but radial_order/n/m parsed")
else:
# get the actual name of the key for radial order
- if "n" in kwargs.keys():
+ if "n" in kwargs:
key = "n"
else:
key = "radial_order"
if isinstance(kwargs.get(key), int):
- self.__n = kwargs.get(key) # radial order acknowledged
+ self.__n = int(kwargs[key]) # radial order acknowledged
# Below each time key arguments are checked in the list of keys
- if ("m" in kwargs.keys() or "l" in kwargs.keys() or "azimuthal_order" in kwargs.keys()
- or "angular_frequency" in kwargs.keys()):
- if "m" in kwargs.keys():
+ if ("m" in kwargs or "l" in kwargs or "azimuthal_order" in kwargs or "angular_frequency" in kwargs):
+ if "m" in kwargs:
key = "m"
- elif "l" in kwargs.keys():
+ elif "l" in kwargs:
key = "l"
- elif "azimuthal_order" in kwargs.keys():
+ elif "azimuthal_order" in kwargs:
key = "azimuthal_order"
else:
key = "angular_frequency"
if isinstance(kwargs.get(key), int):
- self.__m = kwargs.get(key) # azimuthal order acknowledged
+ self.__m = int(kwargs[key]) # azimuthal order acknowledged
# Checking that the provided orders are reasonable
if not (self.__n - abs(self.__m)) % 2 == 0: # see [1]
raise ValueError("Failed sanity check: n - |m| == even number")
@@ -120,7 +119,7 @@ def __init__(self, **kwargs):
elif self.__n < 0:
raise ValueError("Failed sanity check: order n less than 0")
elif self.__n > 54:
- raise ValueError("Initialization of Zernike with radial order higher than 54"
+ raise ValueError("\nInitialization of Zernike with radial order higher than 54"
+ " is meaningless because of very slow calculation performance")
# m and n specified correctly, calculate other properties - various indices
else:
@@ -138,25 +137,24 @@ def __init__(self, **kwargs):
raise ValueError("Radial order n provided not as an integer")
elif len(kwargs.keys()) == 1:
# OSA / ANSI index used for Zernike polynomial initialization
- if ("osa_index" in kwargs.keys() or "osa" in kwargs.keys() or "ansi_index" in kwargs.keys()
- or "ansi" in kwargs.keys()):
+ if ("osa_index" in kwargs or "osa" in kwargs or "ansi_index" in kwargs or "ansi" in kwargs):
if self.__initialized:
raise ValueError("The polynomial has been already initialized, but osa_index/osa... parsed")
else:
- if "osa_index" in kwargs.keys():
+ if "osa_index" in kwargs:
key = "osa_index"
- elif "osa" in kwargs.keys():
+ elif "osa" in kwargs:
key = "osa"
- elif "ansi_index" in kwargs.keys():
+ elif "ansi_index" in kwargs:
key = "ansi_index"
- elif "ansi" in kwargs.keys():
+ elif "ansi" in kwargs:
key = "ansi"
if isinstance(kwargs.get(key), int):
- osa_i = kwargs.get(key)
+ osa_i = int(kwargs[key])
if osa_i < 0:
raise ValueError("OSA / ANSI index should be non-negative integer")
elif osa_i > 1539:
- ValueError("Initialization of Zernike with OSA index higher than 1539"
+ ValueError("\nInitialization of Zernike with OSA index higher than 1539"
+ " is meaningless because of very slow calculation performance")
else:
self.__osa_index = osa_i; self.__initialized = True
@@ -166,20 +164,20 @@ def __init__(self, **kwargs):
else:
raise ValueError("OSA / ANSI index provided not as an integer")
# Noll index used for Zernike polynomial initialization
- elif "noll_index" in kwargs.keys() or "noll" in kwargs.keys():
+ elif "noll_index" in kwargs or "noll" in kwargs:
if self.__initialized:
raise ValueError("The polynomial has been already initialized, but noll_index/noll parsed")
else:
- if "noll_index" in kwargs.keys():
+ if "noll_index" in kwargs:
key = "noll_index"
- elif "noll" in kwargs.keys():
+ elif "noll" in kwargs:
key = "noll"
if isinstance(kwargs.get(key), int):
- noll_i = kwargs.get(key)
+ noll_i = int(kwargs[key])
if noll_i < 1:
raise ValueError("Noll index should be not less than 1 integer")
elif noll_i > 1540:
- ValueError("Initialization of Zernike with Noll index higher than 1540"
+ ValueError("\nInitialization of Zernike with Noll index higher than 1540"
+ " is meaningless because of very slow calculation performance")
else:
self.__noll_index = noll_i; self.__initialized = True
@@ -189,23 +187,23 @@ def __init__(self, **kwargs):
else:
raise ValueError("Noll index provided not as an integer")
# Fringe / Univ. of Arizona index used for Zernike polynomial initialization
- elif "fringe_index" in kwargs.keys() or "fringe" in kwargs.keys():
+ elif "fringe_index" in kwargs or "fringe" in kwargs:
if self.__initialized:
raise ValueError("The polynomial has been already initialized, but fringe_index/fringe parsed")
else:
- if "fringe_index" in kwargs.keys():
+ if "fringe_index" in kwargs:
key = "fringe_index"
- elif "fringe" in kwargs.keys():
+ elif "fringe" in kwargs:
key = "fringe"
if isinstance(kwargs.get(key), int):
- fringe_i = kwargs.get(key)
+ fringe_i = int(kwargs[key])
if fringe_i < 1:
raise ValueError("Fringe index should be not less than 1 integer")
else:
self.__fringe_index = fringe_i; self.__initialized = True
self.__m, self.__n = ZernPol.index2orders(fringe_index=self.__fringe_index)
if self.__n > 54:
- raise ValueError("Initialization of Zernike with radial order higher than 54"
+ raise ValueError("\nInitialization of Zernike with radial order higher than 54"
+ " is meaningless because of very slow calculation performance")
self.__osa_index = ZernPol.get_osa_index(self.__m, self.__n)
self.__noll_index = ZernPol.get_noll_index(self.__m, self.__n)
@@ -218,11 +216,11 @@ def __init__(self, **kwargs):
raise ValueError("The initialization parameters for Zernike polynomial hasn't been parsed / recognized")
# Generate warning messages for possible re-usage
if self.__n > MAX_RADIAL_ORDER_COEFFS:
- self.__warn_mes_r = f"Call for radial order {self.__n} is higher than {MAX_RADIAL_ORDER_COEFFS}"
+ self.__warn_mes_r = f"\nCall for radial order {self.__n} is higher than {MAX_RADIAL_ORDER_COEFFS}"
else:
self.warn_mes_r = ""
if self.__n > MAX_RADIAL_ORDER_COEFFS_dR:
- self.__warn_mes_dr = f"Call for derivative of radial order {self.__n} is > than {MAX_RADIAL_ORDER_COEFFS_dR}"
+ self.__warn_mes_dr = f"\nCall for derivative of radial order {self.__n} is > than {MAX_RADIAL_ORDER_COEFFS_dR}"
else:
self.warn_mes_dr = ""
@@ -273,10 +271,10 @@ def get_polynomial_name(self, short: bool = False) -> str:
"""
name = ""
if short:
- if (self.__m, self.__n) in short_polynomial_names.keys():
+ if (self.__m, self.__n) in short_polynomial_names:
name = short_polynomial_names[(self.__m, self.__n)]
else:
- if (self.__m, self.__n) in polynomial_names.keys():
+ if (self.__m, self.__n) in polynomial_names:
name = polynomial_names[(self.__m, self.__n)]
return name
@@ -400,9 +398,8 @@ def polynomial_value(self, r: Union[float, np.ndarray], theta: Union[float, np.n
r = ZernPol._check_radii(r) # Check radii type and that they are not lying outside range [0.0, 1.0] - unit circle
theta = ZernPol._check_angles(theta) # Checking that angles lie in the range [0, 2*pi] and their type
# Checking coincidence of shapes if theta and r are arrays
- if isinstance(r, type(np.zeros(1))) and isinstance(theta, type(np.zeros(1))):
- if r.shape != theta.shape:
- raise ValueError("Shape of input arrays r and theta is not equal")
+ if isinstance(r, type(np.zeros(1))) and isinstance(theta, type(np.zeros(1))) and r.shape != theta.shape:
+ raise ValueError("Shape of input arrays r and theta is not equal")
# Calculation using imported function from submodule depending on radial order, use different eq.
nTr = normalization_factor(self)*triangular_function(self, theta)
if not use_exact_eq:
@@ -411,14 +408,14 @@ def polynomial_value(self, r: Union[float, np.ndarray], theta: Union[float, np.n
else:
# Raise warning about slow calculations for orders more than 50, only once
if self.__n > 50 and not self.__show_slow_calc_warn:
- warn_mess = f"ZernPol(m={self.__m}, n={self.__n})" + warn_mess_slow_calc
- warnings.warn(warn_mess)
+ warn_mess = f"\nZernPol(m={self.__m}, n={self.__n})" + warn_mess_slow_calc
+ warnings.warn(warn_mess, stacklevel=2)
self.__show_slow_calc_warn = True
# Returning values using recursive scheme for finding the coefficients for all radial orders (e.g. R^6)
return nTr*radial_polynomial_coeffs(self, r)
else:
if self.__n > MAX_RADIAL_ORDER_COEFFS:
- warnings.warn(self.__warn_mes_r + warn_mess_r_long)
+ warnings.warn(self.__warn_mes_r + warn_mess_r_long, stacklevel=2)
if isinstance(r, float):
return 0.0
elif isinstance(r, np.ndarray):
@@ -478,14 +475,14 @@ def radial(self, r: Union[float, np.ndarray], use_exact_eq: bool = False) -> Uni
else:
# Raise warning about slow calculations for orders more than 50, only once
if self.__n > 50 and not self.__show_slow_calc_warn:
- warn_mess = f"ZernPol(m={self.__m}, n={self.__n})" + warn_mess_slow_calc
- warnings.warn(warn_mess)
+ warn_mess = f"\nZernPol(m={self.__m}, n={self.__n})" + warn_mess_slow_calc
+ warnings.warn(warn_mess, stacklevel=2)
self.__show_slow_calc_warn = True
# Returning values using recursive scheme for finding the coefficients for all radial orders (e.g. R^6)
return radial_polynomial_coeffs(self, r)
else:
if self.__n > MAX_RADIAL_ORDER_COEFFS:
- warnings.warn(self.__warn_mes_r + warn_mess_r_long)
+ warnings.warn(self.__warn_mes_r + warn_mess_r_long, stacklevel=2)
if isinstance(r, float):
return 0.0
elif isinstance(r, np.ndarray):
@@ -561,8 +558,7 @@ def radial_dr(self, r: Union[float, np.ndarray], use_exact_eq: bool = False) ->
Calculated derivative of Zernike radial function value(-s) on provided float values / arrays of radiuses.
"""
- # Checking input parameters for avoiding errors and unexpectable values
- r = ZernPol._check_radii(r)
+ r = ZernPol._check_radii(r) # Checking input parameters for avoiding errors and unexpectable values
# Calculation using imported function from submodule depending on radial order, use different eq.
if not use_exact_eq:
if self.__n <= 12: # condition to switch from direct recurrence equation to finding of coeffs. algorithm
@@ -570,14 +566,14 @@ def radial_dr(self, r: Union[float, np.ndarray], use_exact_eq: bool = False) ->
else:
# Raise warning about slow calculations for orders more than 50, only once
if self.__n > 48 and not self.__show_slow_calc_warn:
- warn_mess = f"ZernPol(m={self.__m}, n={self.__n})" + warn_mess_slow_calc
- warnings.warn(warn_mess)
+ warn_mess = f"\nZernPol(m={self.__m}, n={self.__n})" + warn_mess_slow_calc
+ warnings.warn(warn_mess, stacklevel=2)
self.__show_slow_calc_warn = True
# Returning values using recursive scheme for finding the coefficients for all radial orders (e.g. R^6)
return radial_polynomial_coeffs_dr(self, r)
else:
if self.__n > MAX_RADIAL_ORDER_COEFFS_dR:
- warnings.warn(self.__warn_mes_dr + warn_mess_dr_long)
+ warnings.warn(self.__warn_mes_dr + warn_mess_dr_long, stacklevel=2)
if isinstance(r, float):
return 0.0
elif isinstance(r, np.ndarray):
@@ -680,15 +676,10 @@ def get_noll_index(m: int, n: int) -> int:
"""
add_n = 1
if m > 0:
- if (n % 4) == 0:
- add_n = 0
- elif ((n - 1) % 4) == 0:
- add_n = 0
- elif m < 0:
- if ((n - 2) % 4) == 0:
- add_n = 0
- elif ((n - 3) % 4) == 0:
+ if (n % 4) == 0 or ((n - 1) % 4) == 0:
add_n = 0
+ elif m < 0 and (((n - 2) % 4) == 0 or ((n - 3) % 4) == 0):
+ add_n = 0
return (n*(n + 1))//2 + abs(m) + add_n
@staticmethod
@@ -731,30 +722,29 @@ def index2orders(**kwargs) -> Tuple[int, int]:
(m, n) - contains azimuthal and radial orders as integers.
"""
- osa_index = -1; noll_index = -1; fringe_index = -1; m = -1; n = -1
+ osa_index: int = -1; noll_index: int = -1; fringe_index: int = -1; m: int = -1; n: int = -1
highest_order = 56 # default value, limited by the allowed maximal radial order + 2
- if "osa_index" in kwargs.keys():
- osa_index = kwargs.get("osa_index")
- elif "noll_index" in kwargs.keys():
- noll_index = kwargs.get("noll_index")
- elif "fringe_index" in kwargs.keys():
- fringe_index = kwargs.get("fringe_index")
+ if "osa_index" in kwargs:
+ osa_index = int(kwargs["osa_index"])
+ elif "noll_index" in kwargs:
+ noll_index = int(kwargs["noll_index"])
+ elif "fringe_index" in kwargs:
+ fringe_index = int(kwargs["fringe_index"])
highest_order = 250 # to guarantee that right Fringe index found
# Define m, n orders up to the specified above highest order (radial)
stop_search = False
for order in range(0, highest_order):
m = -order # azimuthal order
n = order # radial order
- for polynomial in range(0, order+1):
+ for _ in range(0, order+1):
if osa_index >= 0:
if osa_index == ZernPol.get_osa_index(m, n):
stop_search = True; break
elif noll_index >= 0:
if noll_index == ZernPol.get_noll_index(m, n):
stop_search = True; break
- elif fringe_index >= 0:
- if fringe_index == ZernPol.get_fringe_index(m, n):
- stop_search = True; break
+ elif fringe_index >= 0 and fringe_index == ZernPol.get_fringe_index(m, n):
+ stop_search = True; break
m += 2
if stop_search:
break
@@ -905,9 +895,9 @@ def _sum_zernikes_meshgrid(coefficients: Sequence[float], polynomials: Sequence,
raise ValueError("Lengths of coefficients and polynomials aren't equal")
else:
if not isinstance(r, np.ndarray) or not isinstance(theta, np.ndarray):
- warnings.warn("Requested calculation of surface (mesh) values with"
+ warnings.warn("\nRequested calculation of surface (mesh) values with"
+ " provided r or theta as not numpy.ndarray.\n"
- + "The surface will be generated automatically.")
+ + "The surface will be generated automatically.", stacklevel=2)
else:
r_size = np.size(r, 0); theta_size = np.size(theta, 0)
theta_grid, r_grid = np.meshgrid(theta, r); S = np.zeros(shape=(r_size, theta_size))
@@ -954,7 +944,7 @@ def sum_zernikes(coefficients: Sequence[float], polynomials: Sequence, r: Union[
Depending on the input values and parameter get_surface - can be: float, 1D or 2D numpy.ndarrays.
"""
- S = 0.0 # default value - sum
+ S: Union[float, np.ndarray]
if len(coefficients) != len(polynomials):
raise ValueError("Lengths of lists with polynomials and their amplitudes aren't equal")
elif len(coefficients) == 0:
@@ -971,12 +961,10 @@ def sum_zernikes(coefficients: Sequence[float], polynomials: Sequence, r: Union[
S += coefficient*polynomials[i].polynomial_value(r, theta)
elif get_surface:
if not isinstance(r, np.ndarray) or not isinstance(theta, np.ndarray):
- warnings.warn("Requested calculation of surface (mesh) values with"
- + " provided r or theta as not numpy.ndarray.\n"
- + "The surface will be generated automatically.")
+ warnings.warn("\nRequested calculation of surface (mesh) values with provided r or theta as not numpy.ndarray.\n"
+ + "The surface will be generated automatically.", stacklevel=2)
else:
- r_size = np.size(r, 0); theta_size = np.size(theta, 0)
- S = np.zeros(shape=(r_size, theta_size))
+ r_size = np.size(r, 0); theta_size = np.size(theta, 0); S = np.zeros(shape=(r_size, theta_size))
for i, coefficient in enumerate(coefficients):
if not isinstance(polynomials[i], ZernPol):
raise ValueError(f"Variable {polynomials[i]} isn't an instance of ZernPol class")
@@ -992,7 +980,7 @@ def sum_zernikes(coefficients: Sequence[float], polynomials: Sequence, r: Union[
@staticmethod
def gen_polar_coordinates(r_step: float = 0.01, theta_rad_step: float = round(np.pi/240, 7)) -> polar_vectors:
"""
- Generate the named tuple "PolarVectors" with R and Theta - vectors with polar coordinates for an entire unit circle.
+ Generate the named tuple "polar_vectors" with R and Theta - vectors with polar coordinates for an entire unit circle.
Note that R and Theta are generated as the numpy.ndarrays vectors (shape like (n elements, )). Their shapes are
defined by the specification of r_step and theta_rad_step parameters.
@@ -1012,7 +1000,7 @@ def gen_polar_coordinates(r_step: float = 0.01, theta_rad_step: float = round(np
Returns
-------
polar_vectors
- namedtuple("PolarVectors", "R Theta"), where R - vector with radiuses values [0.0, r_step, ... 1.0],
+ namedtuple("polar_vectors", "R Theta"), where R - vector with radiuses values [0.0, r_step, ... 1.0],
Theta - vector with theta angles values [0.0, theta_rad_step, ... 2*pi].
"""
@@ -1031,7 +1019,7 @@ def gen_polar_coordinates(r_step: float = 0.01, theta_rad_step: float = round(np
@staticmethod
def gen_equal_polar_mesh(n_points: int = 250) -> polar_vectors:
"""
- Generate the named tuple "PolarVectors" with R and Theta - vectors with polar coordinates for an entire unit circle.
+ Generate the named tuple "polar_vectors" with R and Theta - vectors with polar coordinates for an entire unit circle.
Note that R and Theta are generated as the numpy.ndarrays vectors (shape like (n elements, )). Their shapes are
equal and defined by the parameter n_points.
@@ -1044,7 +1032,7 @@ def gen_equal_polar_mesh(n_points: int = 250) -> polar_vectors:
Returns
-------
polar_vectors
- namedtuple("PolarVectors", "R Theta"), where R - vector with radiuses values [0.0, r_step, ... 1.0],
+ namedtuple("polar_vectors", "R Theta"), where R - vector with radiuses values [0.0, r_step, ... 1.0],
Theta - vector with theta angles values [0.0, theta_rad_step, ... 2*pi].
"""
@@ -1055,7 +1043,7 @@ def gen_equal_polar_mesh(n_points: int = 250) -> polar_vectors:
@staticmethod
def plot_profile(polynomial, color_map: str = "coolwarm", show_title: bool = True, use_defaults: bool = True,
- projection: str = "2d", polar_coordinates: polar_vectors = ()):
+ projection: str = "2d", polar_coordinates: Optional[polar_vectors] = None):
"""
Plot the provided Zernike polynomial (instance of ZernPol class) on the matplotlib figure.
@@ -1077,7 +1065,7 @@ def plot_profile(polynomial, color_map: str = "coolwarm", show_title: bool = Tru
Either "2d" ("2D") - for 2D profile plot, or "3d" ("3D") - for 3D surface plot. The default is "2d".
polar_coordinates : polar_vectors, optional
If the flag 'use_defaults' is False, this named tuple is used for accessing polar coordinates for plotting.
- The default is ().
+ The default is None.
Raises
------
@@ -1090,8 +1078,8 @@ def plot_profile(polynomial, color_map: str = "coolwarm", show_title: bool = Tru
"""
if isinstance(polynomial, ZernPol):
- if not use_defaults and len(polar_coordinates) != 2:
- raise ValueError("Polar coordinates isn't provided as a tuple with values R, Theta")
+ if not use_defaults and polar_coordinates is None:
+ raise ValueError("\nPolar coordinates isn't provided as a tuple with values R, Theta")
# Get polar coordinates
if use_defaults:
if projection == "3d" or projection == "3D":
@@ -1099,10 +1087,13 @@ def plot_profile(polynomial, color_map: str = "coolwarm", show_title: bool = Tru
else:
r, theta = ZernPol.gen_polar_coordinates()
else:
- r, theta = polar_coordinates.R, polar_coordinates.Theta
+ if polar_coordinates is not None:
+ r, theta = polar_coordinates.R, polar_coordinates.Theta
+ else:
+ raise ValueError("\nPolar coordinates isn't provided as a tuple with values R, Theta")
# Get profile or surface to plot
amplitudes = [1.0]; zernikes = [polynomial] # for reusing the sum function of polynomials
- zern_surface = ZernPol.sum_zernikes(amplitudes, zernikes, r, theta, get_surface=True)
+ zern_surface = np.asarray(ZernPol.sum_zernikes(amplitudes, zernikes, r, theta, get_surface=True))
# Select plotting function between 2D and 3D plotting functions
if projection == "3d" or projection == "3D":
plot_sum_fig_3d(zern_surface, r, theta, color_map)
@@ -1115,8 +1106,8 @@ def plot_profile(polynomial, color_map: str = "coolwarm", show_title: bool = Tru
@staticmethod
def gen_zernikes_surface(coefficients: Sequence[float], polynomials: Sequence, r_step: float = 0.01,
- theta_rad_step: float = round(np.pi/180, 7),
- equal_n_coordinates: bool = False, n_points: int = 250) -> zernikes_surface:
+ theta_rad_step: float = round(np.pi/180, 7), equal_n_coordinates: bool = False,
+ n_points: int = 250) -> zernikes_surface:
"""
Generate surface of provided Zernike polynomials on the generated polar coordinates used steps.
@@ -1141,7 +1132,7 @@ def gen_zernikes_surface(coefficients: Sequence[float], polynomials: Sequence, r
Returns
-------
zernikes_surface
- namedtuple("ZernikesSurface", "ZernSurf R Theta") - tuple for storing mesh values for polar coordinates.
+ namedtuple("zernikes_surface", "ZernSurf R Theta") - tuple for storing mesh values for polar coordinates.
ZernSurf variable is 2D matrix with the sum of the input polynomials on generated polar coordinates (R, Theta).
"""
@@ -1149,14 +1140,13 @@ def gen_zernikes_surface(coefficients: Sequence[float], polynomials: Sequence, r
polar_vectors = ZernPol.gen_polar_coordinates(r_step, theta_rad_step)
else:
polar_vectors = ZernPol.gen_equal_polar_mesh(n_points)
- zernikes_sum = ZernPol.sum_zernikes(coefficients, polynomials, polar_vectors.R,
- polar_vectors.Theta, get_surface=True)
+ zernikes_sum = ZernPol.sum_zernikes(coefficients, polynomials, polar_vectors.R, polar_vectors.Theta, get_surface=True)
return zernikes_surface(zernikes_sum, polar_vectors.R, polar_vectors.Theta)
@staticmethod
def plot_sum_zernikes_on_fig(figure: plt.Figure, coefficients: Sequence[float] = (), polynomials: Sequence = (), use_defaults: bool = True,
- zernikes_sum_surface: zernikes_surface = (), show_range: bool = True, color_map: str = "coolwarm",
- projection: str = "2d") -> plt.Figure:
+ zernikes_sum_surface: Optional[zernikes_surface] = None,
+ show_range: bool = True, color_map: str = "coolwarm", projection: str = "2d") -> plt.Figure:
"""
Plot a sum of the specified Zernike polynomials by input lists (see function parameters) on the provided figure.
@@ -1174,10 +1164,10 @@ def plot_sum_zernikes_on_fig(figure: plt.Figure, coefficients: Sequence[float] =
Coefficients of Zernike polynomials for calculation of their sum. The default is ().
polynomials : Sequence[ZernPol], optional
Initialized polynomials as class instances of ZernPol class specified in this module. The default is ().
- zernikes_sum_surface : namedtuple("ZernikesSurface", "ZernSurf R Theta") , optional
+ zernikes_sum_surface : namedtuple("zernikes_surface", "ZernSurf R Theta"), optional
This tuple should contain the ZernSurf calculated on a mesh of polar coordinates R, Theta.
This tuple could be generated by the call of the static method gen_zernikes_surface().
- Check the method signature for details. The default is ().
+ Check the method signature for details. The default is None.
show_range : bool, optional
Flag for showing range of provided values as the colorbar on the figure. The default is True.
color_map : str, optional
@@ -1199,16 +1189,15 @@ def plot_sum_zernikes_on_fig(figure: plt.Figure, coefficients: Sequence[float] =
"""
if use_defaults and len(coefficients) == 0 and len(polynomials) == 0:
- raise ValueError("Input Sequence with coefficients or with polynomials is empty along with the flag 'use_defaults' - True")
- if not use_defaults and len(zernikes_sum_surface) != 3:
- raise ValueError("Zernike surface isn't provided as a tuple with values Sum surface, R, Theta")
+ raise ValueError("\nInput Sequence with coefficients or with polynomials is empty along with the flag 'use_defaults' - True")
+ if not use_defaults and zernikes_sum_surface is None:
+ raise ValueError("\nZernike surface isn't provided as a tuple with values Sum surface, R, Theta")
if use_defaults:
if projection == "3d" or projection == "3D":
polar_vectors = ZernPol.gen_equal_polar_mesh()
else:
polar_vectors = ZernPol.gen_polar_coordinates()
- zernikes_sum = ZernPol.sum_zernikes(coefficients, polynomials, polar_vectors.R,
- polar_vectors.Theta, get_surface=True)
+ zernikes_sum = np.asarray(ZernPol.sum_zernikes(coefficients, polynomials, polar_vectors.R, polar_vectors.Theta, get_surface=True))
if projection == "3d" or projection == "3D":
figure = subplot_sum_on_fig_3d(figure, zernikes_sum, polar_vectors.R, polar_vectors.Theta,
show_range_colorbar=show_range, color_map=color_map)
@@ -1216,20 +1205,21 @@ def plot_sum_zernikes_on_fig(figure: plt.Figure, coefficients: Sequence[float] =
figure = subplot_sum_on_fig(figure, zernikes_sum, polar_vectors.R, polar_vectors.Theta,
show_range_colorbar=show_range, color_map=color_map)
else:
- if projection == "3d" or projection == "3D":
- figure = subplot_sum_on_fig_3d(figure, zernikes_sum_surface.ZernSurf, zernikes_sum_surface.R,
- zernikes_sum_surface.Theta, show_range_colorbar=show_range,
- color_map=color_map)
+ if zernikes_sum_surface is not None:
+ if projection == "3d" or projection == "3D":
+ figure = subplot_sum_on_fig_3d(figure, zernikes_sum_surface.ZernSurf, zernikes_sum_surface.R,
+ zernikes_sum_surface.Theta, show_range_colorbar=show_range, color_map=color_map)
+ else:
+ figure = subplot_sum_on_fig(figure, zernikes_sum_surface.ZernSurf, zernikes_sum_surface.R,
+ zernikes_sum_surface.Theta, show_range_colorbar=show_range, color_map=color_map)
else:
- figure = subplot_sum_on_fig(figure, zernikes_sum_surface.ZernSurf, zernikes_sum_surface.R,
- zernikes_sum_surface.Theta, show_range_colorbar=show_range,
- color_map=color_map)
+ raise ValueError("\nZernike surface isn't provided as a tuple with values Sum surface, R, Theta")
return figure
@staticmethod
def _plot_zernikes_half_pyramid():
"""
- Generate halb-pyramid with Zernikes polynomials.
+ Generate half-pyramid with Zernikes polynomials.
Returns
-------
@@ -1244,10 +1234,8 @@ def _plot_zernikes_half_pyramid():
for j in range(len(axes[0])):
axes[i, j].grid(False) # demanded by pcolormesh function, if not called - deprecation warning
if j > ignored_column-1:
- zps = [ZernPol(osa=k)]
- zernike_surface, r, theta = ZernPol.gen_zernikes_surface(coefficients=ampls, polynomials=zps)
- axes[i, j].pcolormesh(theta, r, zernike_surface, cmap=plt.cm.coolwarm, shading='nearest')
- k += 1
+ zps = [ZernPol(osa=k)]; zernike_surface, r, theta = ZernPol.gen_zernikes_surface(coefficients=ampls, polynomials=zps)
+ axes[i, j].pcolormesh(theta, r, zernike_surface, cmap=plt.colormaps['coolwarm'], shading='nearest'); k += 1
axes[i, j].axis('off') # off polar coordinate axes
ignored_column -= 1
fig.subplots_adjust(left=0, bottom=0, right=1, top=1); fig.tight_layout()
@@ -1276,17 +1264,16 @@ def _check_radii(radii: Union[list, tuple, float, int, np.ndarray]) -> Union[flo
"""
# Trying to convert known (list, tuple) data types into numpy, if they provided as input
if not isinstance(radii, np.ndarray) and not isinstance(radii, float):
- if isinstance(radii, list) or isinstance(radii, tuple) or isinstance(radii, set):
+ if isinstance(radii, (list, tuple, set)):
radii = np.asarray(radii) # convert list or tuple to np.array
else:
radii = float(radii) # attempt to convert r to float number, will raise ValueError if it's impossible
# Checking that radii or radius lie in the range [0.0, 1.0]
if isinstance(radii, np.ndarray):
if np.min(radii) < 0.0 or np.max(radii) > 1.0:
- raise ValueError("Minimal or maximal value of radii laying outside unit circle [0.0, 1.0]")
- elif isinstance(radii, float):
- if radii > 1.0 or radii < 0.0:
- raise ValueError("Radius laying outside unit circle [0.0, 1.0]")
+ raise ValueError("\nMinimal or maximal value of radii laying outside unit circle [0.0, 1.0]")
+ elif isinstance(radii, float) and (radii > 1.0 or radii < 0.0):
+ raise ValueError("Radius laying outside unit circle [0.0, 1.0]")
return radii
@staticmethod
@@ -1315,27 +1302,26 @@ def _check_angles(angles: Union[list, tuple, float, int, np.ndarray]) -> Union[f
"""
# Check input parameter type and attempt to convert to acceptable types
if not isinstance(angles, np.ndarray) and not isinstance(angles, float):
- if isinstance(angles, list) or isinstance(angles, tuple) or isinstance(angles, set):
+ if isinstance(angles, (list, tuple, set)):
angles = np.asarray(angles) # convert list or tuple to np.array
else:
angles = float(angles) # attempt to convert to float number, will raise ValueError if it's impossible
# Checking that angles lie in the range [0, 2*pi]
if isinstance(angles, np.ndarray):
if np.max(angles) - np.min(angles) > 2.0*np.pi:
- _warn_message_ = "Theta angles defined in range outside of interval [0.0, 2.0*pi]"
- warnings.warn(_warn_message_)
+ _warn_message_ = "\nTheta angles defined in range outside of interval [0.0, 2.0*pi]"
+ warnings.warn(_warn_message_, stacklevel=2)
elif np.max(angles) > 2.0*np.pi or np.min(angles) < 0.0:
- _warn_message_ = "Max or min of theta angles lies outside of interval [0.0, 2.0*pi]"
- warnings.warn(_warn_message_)
- elif isinstance(angles, float):
- if angles < 0.0 or angles > 2.0*np.pi:
- _warn_message_ = "Max or min of theta angles lies outside of interval [0.0, 2.0*pi]"
- warnings.warn(_warn_message_)
+ _warn_message_ = "\nMax or min of theta angles lies outside of interval [0.0, 2.0*pi]"
+ warnings.warn(_warn_message_, stacklevel=2)
+ elif isinstance(angles, float) and angles < 0.0 or angles > 2.0*np.pi:
+ _warn_message_ = "\nMax or min of theta angles lies outside of interval [0.0, 2.0*pi]"
+ warnings.warn(_warn_message_, stacklevel=2)
return angles
# %% Independent functions defs.
-def generate_polynomials(max_order: int = 10) -> Tuple[ZernPol]:
+def generate_polynomials(max_order: int = 10) -> Tuple[ZernPol, ...]:
"""
Generate tuple with ZernPol instances (ultimately, representing polynomials) indexed using OSA scheme, starting with Piston(m=0,n=0).
@@ -1357,14 +1343,14 @@ def generate_polynomials(max_order: int = 10) -> Tuple[ZernPol]:
"""
# Sanity check of max_order parameter
if not isinstance(max_order, int):
- __warning_mess = "The parameter max_order provided not as integer, there will be attempt to convert it to int"
- warnings.warn(__warning_mess)
+ __warning_mess = "\nThe parameter max_order provided not as integer, there will be attempt to convert it to int"
+ warnings.warn(__warning_mess, stacklevel=2)
max_order = int(max_order)
if max_order < 0:
raise ValueError("The maximum order should be not less than 0")
if max_order > 30:
- __warning_mess = "Calculation polynomial values with orders higher than 30 is really slow"
- warnings.warn(__warning_mess)
+ __warning_mess = "\nCalculation polynomial values with orders higher than 30 is really slow"
+ warnings.warn(__warning_mess, stacklevel=2)
polynomials_list = [ZernPol(m=0, n=0)] # list starting with piston
for order in range(1, max_order + 1): # going through all specified orders
m = -order # azimuthal order
@@ -1376,7 +1362,7 @@ def generate_polynomials(max_order: int = 10) -> Tuple[ZernPol]:
def generate_random_phases(max_order: int = 4, img_width: int = 513, img_height: int = 513,
- round_digits: int = 4) -> Tuple[np.ndarray, np.ndarray, Tuple[ZernPol]]:
+ round_digits: int = 4) -> Tuple[np.ndarray, np.ndarray, Tuple[ZernPol, ...]]:
"""
Generate phases image (profile) for random set of polynomials with randomly selected amplitudes.
@@ -1437,8 +1423,7 @@ def generate_random_phases(max_order: int = 4, img_width: int = 513, img_height:
for j in range(img_width):
euclidean_dist = np.linalg.norm(center - np.asarray([i, j]))
if euclidean_dist <= img_radius:
- r[j] = euclidean_dist / img_radius
- theta[j] = np.arctan2(row_center - i, j - cols_center)
+ r[j] = euclidean_dist / img_radius; theta[j] = np.arctan2(row_center - i, j - cols_center)
if theta[j] < 0.0:
theta[j] += 2.0*np.pi
# Speed up calculations by using vectors (r, theta) as the input parameters
@@ -1449,13 +1434,11 @@ def generate_random_phases(max_order: int = 4, img_width: int = 513, img_height:
euclidean_dist = np.linalg.norm(center - np.asarray([i, j]))
if euclidean_dist > img_radius:
phases_image[i, j] = 0.0
- # Final conversion
- polynomials_list = tuple(polynomials_list)
- return phases_image, polynomials_amplitudes, polynomials_list
+ return phases_image, polynomials_amplitudes, tuple(polynomials_list)
-def generate_phases_image(polynomials: tuple = (), polynomials_amplitudes: tuple = (),
- img_width: int = 513, img_height: int = 513) -> np.ndarray:
+def generate_phases_image(polynomials: tuple = (), polynomials_amplitudes: tuple = (), img_width: int = 513,
+ img_height: int = 513) -> np.ndarray:
"""
Generate phases image (profile) for provided set of polynomials with provided coefficients (amplitudes).
@@ -1488,8 +1471,7 @@ def generate_phases_image(polynomials: tuple = (), polynomials_amplitudes: tuple
for j in range(img_width):
euclidean_dist = np.linalg.norm(center - np.asarray([i, j]))
if euclidean_dist <= img_radius:
- r[j] = euclidean_dist / img_radius
- theta[j] = np.arctan2(row_center - i, j - cols_center)
+ r[j] = euclidean_dist / img_radius; theta[j] = np.arctan2(row_center - i, j - cols_center)
if theta[j] < 0.0:
theta[j] += 2.0*np.pi
phases_image[i, :] = ZernPol.sum_zernikes(coefficients=list(polynomials_amplitudes), polynomials=list(polynomials), r=r, theta=theta)
@@ -1542,16 +1524,13 @@ def fit_polynomials(phases_image: np.ndarray, polynomials: tuple, crop_radius: f
if it is False, the following tuple will be returned: zernike_coefficients, None - 1st with the same
meaning and type as explained before.
"""
- zernike_coefficients = np.zeros(shape=(len(polynomials), ))
+ zernike_coefficients = np.zeros(shape=(len(polynomials), )); cropped_image = None
logic_mask, cropped_phases_coordinates = crop_phases_img(phases_image, crop_radius, suppress_warnings, strict_circle_border)
if return_cropped_image:
cropped_image = logic_mask*phases_image # for debugging
zernike_coefficients = fit_zernikes(cropped_phases_coordinates, polynomials)
zernike_coefficients = np.round(zernike_coefficients, round_digits)
- if return_cropped_image:
- return zernike_coefficients, cropped_image
- else:
- return zernike_coefficients, None
+ return zernike_coefficients, cropped_image
def fit_polynomials_vectors(polynomials: tuple, phases_vector: np.ndarray, radii_vector: np.ndarray,
@@ -1589,260 +1568,10 @@ def fit_polynomials_vectors(polynomials: tuple, phases_vector: np.ndarray, radii
# Checking input data consistency
if len(phases_vector.shape) > 1 or len(radii_vector.shape) > 1 or len(thetas_vector.shape) > 1:
raise TypeError("Some of provided vector is not 1D, check the call len(...shape)")
- # Call to fitting procedure
zernike_coefficients = fit_zernikes((phases_vector, radii_vector, thetas_vector), polynomials)
zernike_coefficients = np.round(zernike_coefficients, round_digits)
return zernike_coefficients
-def compare_performances(min_order: int, max_order: int) -> Tuple[int, int, str]:
- """
- Compare performances of radial polynomials calculation by using recursive and exact equations.
-
- Comparison achieved by simple measuring of time in msec needed for calculation of all radial
- polynomials from minimal radial order up to maximum radial order returned as tuple.
-
- Parameters
- ----------
- min_order : int
- Minimum radial order of used polynomials (n).
- max_order : int
- Maximum radial order of used polynomials (n).
-
- Returns
- -------
- tuple
- Composed by time for recursive calculation and time for exact calculation.
-
- """
- # Generation of orders in OSA/ANSI indexing scheme for initializing Zernike polynomials
- zernpols = []
- for order in range(min_order, max_order+1):
- m = -order; n = order
- zernpols.append(ZernPol(m=m, n=n))
- for n_azimuthals in range(0, order):
- m += 2
- zernpols.append(ZernPol(m=m, n=n))
- # Generation numpy array with radii
- n_points = 251
- test_r = np.zeros(shape=(n_points, ))
- for i in range(n_points):
- test_r[i] = i/(n_points-1)
- test_r = np.round(test_r, 6)
- # Measuring performance of radial polynomials calculations using recursive implementation
- t1 = time.perf_counter()
- for i, polynomial in enumerate(zernpols):
- polynomial.radial(test_r) # calculate radial polynomials over vector of radii
- t2 = time.perf_counter()
- t_recursive_ms = round(1000*(t2-t1), 3)
- # Measuring performance of radial polynomials calculations using exact implementation
- t1 = time.perf_counter()
- for i, polynomial in enumerate(zernpols):
- polynomial.radial(test_r, use_exact_eq=True) # calculate radial polynomials over vector of radii
- t2 = time.perf_counter()
- t_exact_ms = round(1000*(t2-t1), 3)
- return t_recursive_ms, t_exact_ms, f"Used polynomials: {i}", f"Radii: {n_points}"
-
-
-def _estimate_high_order_calc_times():
- """
- Estimate slowing down of radial polynomial calculation with increasing of radial order.
-
- Returns
- -------
- None.
-
- """
- r = 0.55
- high_order_pols = [ZernPol(m=-2, n=46), ZernPol(m=1, n=47), ZernPol(m=2, n=48), ZernPol(m=-1, n=49),
- ZernPol(m=2, n=50), ZernPol(m=1, n=51), ZernPol(m=-2, n=52)]
- times = []
- # Single polynomials values
- for pol in high_order_pols:
- # calculate radial polynomials over vector of radii
- t1 = time.perf_counter(); pol.radial(r); t2 = time.perf_counter()
- times.append(f"{pol.get_mn_orders()}: {int(round(1000*(t2-t1), 0))} ms")
- # Delete polynomials with too high orders for derivatives calculates
- high_order_pols.pop(len(high_order_pols)-1); high_order_pols.pop(len(high_order_pols)-1)
- high_order_pols.pop(len(high_order_pols)-1)
- print(times); times = []
- # Single derivative polynomials values
- for pol in high_order_pols:
- # calculate derivatives of radial polynomials over vector of radii
- t1 = time.perf_counter(); pol.radial_dr(r); t2 = time.perf_counter()
- times.append(f"Deriv. {pol.get_mn_orders()}: {int(round(1000*(t2-t1), 0))} ms")
- print(times)
-
-
-# %% Test functions for the external call
-def check_conformity():
- """
- Test initialization parameters and transform between indices consistency.
-
- Returns
- -------
- None.
-
- """
- zp = ZernPol(m=-2, n=2) # Initialization with orders
- (m1, n1), osa_i, noll_i, fringe_i = zp.get_indices()
- assert (osa_i == 3 and noll_i == 5 and fringe_i == 6), (f"Check consistency of Z{(m1, n1)} indices: "
- + f"OSA: {osa_i}, Noll: {noll_i}, Fringe: {fringe_i}")
- zp = ZernPol(l=-3, n=5)
- (m2, n2), osa_i, noll_i, fringe_i = zp.get_indices()
- assert (osa_i == 16 and noll_i == 19 and fringe_i == 20), (f"Check consistency of Z{(m2, n2)} indices: "
- + f"OSA: {osa_i}, Noll: {noll_i}, Fringe: {fringe_i}")
- assert len(zp.get_polynomial_name(short=True)) > 0, f"Short name for Z{(m2, n2)} is zero length"
- zp = ZernPol(azimuthal_order=-1, radial_order=5)
- (m3, n3), osa_i, noll_i, fringe_i = zp.get_indices()
- assert (osa_i == 17 and noll_i == 17 and fringe_i == 15), (f"Check consistency of Z{(m3, n3)} indices: "
- + f"OSA: {osa_i}, Noll: {noll_i}, Fringe: {fringe_i}")
- assert len(zp.get_polynomial_name()) > 0, f"Name for Z{(m3, n3)} is zero length"
- m4, n4 = zp.get_mn_orders()
- assert m4 == m3 and n3 == n4, f"Check method get_mn_orders() for Z{(m3, n3)}"
- print(f"Initialization of polynomials Z{(m1, n1)}, Z{(m2, n2)}, Z{(m3, n3)} tested")
- osa_i = 12; zp = ZernPol(osa_index=osa_i) # Initialization with OSA index
- m, n = zp.get_mn_orders()
- assert (m == 0 and n == 4), f"Check consistency of Z[OSA index = {osa_i}] orders {m, n}"
- assert zp.get_fringe_index(m, n) == 9, f"Check consistency of Z[OSA index = {osa_i}] Fringe index"
- assert zp.get_noll_index(m, n) == 11, f"Check consistency of Z[OSA index = {osa_i}] Noll index"
- print(f"Initialization of polynomial Z[OSA index = {osa_i}] tested")
- noll_i = 10 # Testing static methods
- assert ZernPol.noll2osa(noll_i) == 9, f"Check consistency of Noll index {noll_i} conversion to OSA index"
- assert ZernPol.osa2fringe(ZernPol.noll2osa(noll_i)) == 10, ("Check consistency of Noll "
- + f"index {noll_i} conversion to OSA index")
- print(f"Conversion of Noll index {noll_i} to OSA and Fringe indices tested")
- # Test for not proper initialization
- try:
- m_f = 2; n_f = -2
- zp = ZernPol(m=m_f, n=n_f)
- asserting_value = False
- except ValueError:
- print(f"Polynomial Z{(m_f, n_f)} haven't been initialized, test passed")
- asserting_value = True
- assert asserting_value, f"Polynomial Z{(m_f, n_f)} initialized with wrong orders assignment"
- # Testing input parameters for calculation
- zp = ZernPol(m=0, n=2); r = 0.0; theta = math.pi
- assert abs(zp.polynomial_value(r, theta) + math.sqrt(3)) < 1E-6, f"Check value of Z[{m}, {n}]({r}, {theta})"
- zp = ZernPol(m=-1, n=1); r = 0.5; theta = math.pi/2
- assert abs(zp.polynomial_value(r, theta) - 1.0) < 1E-6, f"Check value of Z[{m}, {n}]({r}, {theta})"
- print("Simple values of Zernike polynomials tested successfully")
- try:
- r = 'd'; theta = [1, 2]
- zp.polynomial_value(r, theta)
- asserting_value = False
- except ValueError:
- print("Input as string is not allowed for calculation of polynomial value, tested successfully")
- asserting_value = True
- assert asserting_value, "Wrong parameter passed (string) for calculation of polynomial value"
- try:
- r = [0.1, 0.2, 1.0+1E-9]; theta = math.pi
- zp.polynomial_value(r, theta)
- asserting_value = False
- except ValueError:
- print("Radius more than 1.0 is not allowed, tested successfully")
- asserting_value = True
- assert asserting_value, "Wrong parameter passed (r > 1.0) for calculation of polynomial value"
- # Compare two implementations of Zernike pol-s sum calculation: direct and using meshgrid
- pols = [ZernPol(osa=2), ZernPol(osa=4), ZernPol(osa=7), ZernPol(osa=10), ZernPol(osa=15),
- ZernPol(osa=3), ZernPol(osa=9), ZernPol(osa=12), ZernPol(osa=16), ZernPol(osa=19)]
- ampls = [-0.85, 0.85, 0.24, -0.37, 1.0, 0.1, -1.0, -0.05, 1.1, 0.41]
- radii = np.arange(start=0.0, stop=1.0 + 0.001, step=0.001)
- thetas = np.arange(start=0.0, stop=2.0*np.pi + np.pi/180, step=np.pi/180)
- t1 = time.perf_counter(); ZernPol.sum_zernikes(ampls, pols, radii, thetas, get_surface=True)
- t_direct = int(round(1000*(time.perf_counter() - t1), 0)); t1 = time.perf_counter()
- ZernPol._sum_zernikes_meshgrid(ampls, pols, radii, thetas); t_meshgr = int(round(1000*(time.perf_counter() - t1), 0))
- print(f"Diff. calc. time b/t direct ({t_direct} ms) and meshgrid ({t_meshgr} ms) sums: {t_direct - t_meshgr} ms")
- print("ALL TEST PASSED")
-
-
-# %% Define default export classes and methods used with import * statement (import * from zernikepol)
-__all__ = ['ZernPol', 'fit_polynomials_vectors', 'fit_polynomials', 'generate_phases_image',
- 'generate_random_phases', 'generate_polynomials']
-
-# %% Tests
-if __name__ == "__main__":
- _test_plots = False # regulates testing of plotting various plots
- _test_calculations = False # regulates tests below concerning calculations
- check_conformity() # testing initialization
-
- # Testing plotting, the plots will be opened in the additional pop-up windows
- if _test_plots:
- plt.close("all") # close all previously opened plots
- t1 = time.perf_counter(); zp = ZernPol(m=0, n=2); ZernPol.plot_profile(zp, color_map="jet", show_title=True) # basic plot
- t2 = time.perf_counter(); print("Plotting of 1 non-zero polynomial takes ms: ", int(round(1000*(t2-t1), 0)))
- coordinates = ZernPol.gen_polar_coordinates(r_step=0.005)
- zp = ZernPol(m=-10, n=30); ZernPol.plot_profile(zp, color_map="jet", show_title=False, polar_coordinates=coordinates) # high order plot
- zp = ZernPol(m=0, n=0); ZernPol.plot_profile(zp, color_map="turbo", show_title=True) # plot of piston polynomial
-
- # Testing 3D surface plotting
- ZernPol.plot_profile(ZernPol(m=0, n=2), color_map="viridis", projection="3d")
-
- # Testing 3D figure plotting on the externally initialized Figure class
- fig3d = plt.figure(figsize=(6.8, 6.8))
- zern_surface = ZernPol.gen_zernikes_surface([1.0], [ZernPol(m=0, n=2)], equal_n_coordinates=True, n_points=400)
- ZernPol.plot_sum_zernikes_on_fig(figure=fig3d, use_defaults=False, zernikes_sum_surface=zern_surface,
- show_range=True, color_map="magma", projection="3D")
- fig3d2 = plt.figure(figsize=(5.8, 5.8))
- ZernPol.plot_sum_zernikes_on_fig(figure=fig3d2, coefficients=[1.0], polynomials=[ZernPol(m=0, n=2)],
- show_range=True, color_map="bwr", projection="3D")
-
- # Testing accelerated plotting / sum calculation
- fig3 = plt.figure(figsize=(3, 3))
- t1 = time.perf_counter(); n_pols = 31; polynomials = []; coefficients = [0.0]*n_pols
- for i in range(n_pols):
- polynomials.append(ZernPol(osa=58+i))
- coefficients[0] = 1.0 # only 1st polynomial will be plotted
- fig3 = ZernPol.plot_sum_zernikes_on_fig(figure=fig3, coefficients=coefficients, polynomials=polynomials,
- show_range=False, color_map="turbo")
- fig3.subplots_adjust(0, 0, 1, 1)
- t2 = time.perf_counter(); print("Plotting of 1 non-zero and 30 zero pol-s takes ms: ", int(round(1000*(t2-t1), 0)))
-
- # Tests with generation / restoring Zernike profiles (phases images)
- phases_image, polynomials_ampls, polynomials = generate_random_phases(img_height=301, img_width=321)
- plt.figure(); plt.axis("off"); plt.imshow(phases_image, cmap="jet"); plt.tight_layout(); plt.subplots_adjust(0, 0, 1, 1)
- polynomials_amplitudes, cropped_img = fit_polynomials(phases_image, polynomials, return_cropped_image=True,
- strict_circle_border=False, crop_radius=1.0)
- plt.figure(); plt.axis("off"); plt.imshow(cropped_img, cmap="jet")
- plt.tight_layout(); plt.subplots_adjust(0, 0, 1, 1)
-
- # Updated test of fitting including piston polynomial
- height = 500; width = 481; crop_r = 1.0; strict_border = True; pols_coeffs = [-0.75, 0.86, 0.41]; fig4 = plt.figure(figsize=(4, 4))
- polynomials = [ZernPol(osa=0), ZernPol(m=0, n=2), ZernPol(m=-3, n=3)]; rs, angles = ZernPol.gen_polar_coordinates()
- phase_profile = ZernPol.sum_zernikes(coefficients=pols_coeffs, polynomials=polynomials, r=rs, theta=angles, get_surface=True)
- # Below - plotting specified polynomials on the polar coordinates
- ZernPol.plot_sum_zernikes_on_fig(figure=fig4, use_defaults=False, zernikes_sum_surface=zernikes_surface(phase_profile, rs, angles),
- color_map="jet")
- # Below - generate phases image with the cartesian coordinates
- phases_image2 = generate_phases_image(polynomials=tuple(polynomials), polynomials_amplitudes=tuple(pols_coeffs),
- img_height=height, img_width=width)
- plt.figure(); plt.axis("off"); im = plt.imshow(phases_image2, cmap="jet"); plt.tight_layout(); plt.subplots_adjust(0, 0, 1, 1)
- plt.colorbar(mappable=im)
- # Below - fitting procedure on the provided phases image
- polynomials_amplitudes2, cropped_img2 = fit_polynomials(phases_image2, polynomials, return_cropped_image=True,
- strict_circle_border=strict_border, crop_radius=crop_r)
- print("Difference between used amplitudes and fitted ones:", np.asarray(pols_coeffs) - polynomials_amplitudes2)
- plt.figure(); plt.axis("off"); im = plt.imshow(cropped_img2, cmap="jet"); plt.tight_layout(); plt.subplots_adjust(0, 0, 1, 1)
- plt.colorbar(mappable=im)
-
- plt.show() # show all images created by plt.figure() calls
-
- # Testing calculations and their performance comparison
- if _test_calculations:
- # Simple test of two concepts of calculations - exact and recursive equations
- z = ZernPol(n=30, m=-2); print("Diff. between recursive and exact equations:",
- round(z.radial(0.85) - z.radial(0.85, use_exact_eq=True), 9))
- z = ZernPol(n=32, l=0); print("Diff. between recursive and exact equations:",
- round(z.radial(0.35) - z.radial(0.35, use_exact_eq=True), 9))
- r = 0.955; theta = np.pi/8; z = ZernPol(osa=55)
- print("Diff. between recursive and exact equations:",
- round(z.polynomial_value(r, theta) - z.polynomial_value(r, theta, use_exact_eq=True), 9))
- z = ZernPol(n=35, l=-1); print("Diff. between recursive & exact eq-s for derivatives:",
- round(z.radial_dr(0.78) - z.radial_dr(0.78, use_exact_eq=True), 9))
- z = ZernPol(n=38, m=-2); print("Diff. between recursive & exact eq-s for derivatives:",
- round(z.radial_dr(0.9) - z.radial_dr(0.9, use_exact_eq=True), 9))
- # Compare performances
- print("Tabular (10th order) / exact calc. times:", compare_performances(1, 10))
- print("Recursive / exact calc. times for high orders:", compare_performances(12, 40))
- # Statement below producing expected warnings, it's used for performance estimation
- _estimate_high_order_calc_times()
+# %% Define exported class and functions
+__all__ = ['ZernPol', 'fit_polynomials_vectors', 'fit_polynomials', 'generate_phases_image', 'generate_random_phases', 'generate_polynomials']
diff --git a/src/zernpy/zernpsf.py b/src/zernpy/zernpsf.py
index d6b5071..b557d76 100644
--- a/src/zernpy/zernpsf.py
+++ b/src/zernpy/zernpsf.py
@@ -2,18 +2,19 @@
"""
PSF class definition based on Zernike polynomial for computation of its kernel for convolution / deconvolution.
-@author: Sergei Klykov, @year: 2025, @licence: MIT \n
+@author: Sergei Klykov, @year: 2025, @license: MIT \n
"""
# %% Global imports
-import numpy as np
-from pathlib import Path
-import warnings
-import matplotlib.pyplot as plt
-from math import pi
import time
-from typing import Union, Sequence
+import warnings
from importlib.metadata import version
+from math import pi
+from pathlib import Path
+from typing import Optional, Sequence, Tuple, Union
+
+import matplotlib.pyplot as plt
+import numpy as np
# Check if numba library installed for importing compilable methods
numba_installed = False # default value for checking if 'numba' library is installed
@@ -25,26 +26,23 @@
pass
# %% Local (package-scoped) imports
-if __name__ == "__main__" or __name__ == Path(__file__).stem or __name__ == "__mp_main__":
- from calculations.calc_psfs import (get_psf_kernel, lambda_char, radial_integral_s, radial_integral, get_kernel_size,
- convolute_img_psf, get_bumped_circle, save_psf, read_psf, get_psf_kernel_zerns)
- from zernikepol import ZernPol
- from utils.intmproc import DispenserManager
- if numba_installed:
- from calculations.calc_psfs_numba import get_psf_kernel_comp, methods_compiled, set_methods_compiled
- num_methods_compiled = methods_compiled
- else:
- methods_compiled = False
-else:
- from .calculations.calc_psfs import (get_psf_kernel, lambda_char, radial_integral_s, radial_integral, get_kernel_size,
- convolute_img_psf, get_bumped_circle, save_psf, read_psf, get_psf_kernel_zerns)
- from .zernikepol import ZernPol
- from .utils.intmproc import DispenserManager
- if numba_installed:
- from .calculations.calc_psfs_numba import get_psf_kernel_comp, methods_compiled, set_methods_compiled
- num_methods_compiled = methods_compiled
- else:
- methods_compiled = False
+from .calculations.calc_psfs import (
+ convolute_img_psf,
+ get_bumped_circle,
+ get_kernel_size,
+ get_psf_kernel,
+ get_psf_kernel_zerns,
+ lambda_char,
+ radial_integral,
+ radial_integral_s,
+ read_psf,
+ save_psf,
+)
+from .utils.intmproc import DispenserManager
+from .zernikepol import ZernPol
+
+if numba_installed:
+ from .calculations.calc_psfs_numba import get_psf_kernel_comp
# %% Module parameters
__docformat__ = "numpydoc"
@@ -72,19 +70,18 @@ class ZernPSF:
# Class predefined values along with their types for providing reference how the default values computed
kernel: np.ndarray = np.ones(shape=(1, 1)); kernel_size: int = 1 # by default, single point as the 2D matrix [[1.0]]
NA: float = 1.0; wavelength: float = 0.532; expansion_coeff: float = 0.5 # in um (or other physical unit)
- zernpol: ZernPol = ZernPol(m=0, n=0) # by default, Piston as the zero case (Airy pattern)
- polynomials = () # empty tuple - holder for provided Sequence with instances of ZernPol
+ zernpol: Optional[ZernPol] = None; polynomials: tuple = () # several provided polynomials will be stored in the last var.
__physical_props_set: bool = False # flag for saving that physical properties have been provided
__warn_message: str = ""; pixel_size_nyquist: float = 0.5*0.5*wavelength # based on the Abbe limit
pixel_size_nyquist_eStr = "{:.3e}".format(pixel_size_nyquist)
pixel_size = 0.98*pixel_size_nyquist # default value based on the limit above
alpha: float = expansion_coeff / wavelength # the role of amplitude for PSF calculation
airy: bool = False # flag if the Piston is provided as the Zernike polynomial
- __ParallelCalc: DispenserManager = None; __integration_params: list = [] # for speeding up the calculations using several Processes
+ __ParallelCalc: Optional[DispenserManager] = None; __integration_params: list = [] # for speeding up the calculations using Processes
n_int_r_points: int = 320; n_int_phi_points: int = 300 # integration parameters on the unit radius and angle - polar coordinates
k: float = 2.0*pi/wavelength # angular frequency
json_file_path: str = "" # shifted to the __init__ method to prevent putting path to API doc (by pydoc)
- coefficients: np.ndarray = None; amplitudes: np.ndarray = None # for storing amplitudes of polynomials
+ coefficients: Optional[np.ndarray] = None; amplitudes: Optional[np.ndarray] = None # for storing amplitudes of polynomials
# Dev. Note: always carefully check the definition of variables above, since the error in their definition may cause hard traceable bug
def __init__(self, zernpol: Union[ZernPol, Sequence[ZernPol]]):
@@ -110,13 +107,13 @@ def __init__(self, zernpol: Union[ZernPol, Sequence[ZernPol]]):
"""
self.json_file_path = str(Path(__file__).parent.absolute()) # initialize default path as the root folder containing the script
- if not hasattr(zernpol, '__len__') and isinstance(zernpol, ZernPol):
+ if isinstance(zernpol, ZernPol):
self.zernpol = zernpol; m, n = self.zernpol.get_mn_orders()
if m == 0 and n == 0:
self.airy = True
else:
self.airy = False
- elif hasattr(zernpol, '__len__'):
+ else:
seq_length = len(zernpol)
# Empty sequence as the input not acceptable
if seq_length == 0:
@@ -140,15 +137,13 @@ def __init__(self, zernpol: Union[ZernPol, Sequence[ZernPol]]):
if m == 0 and n == 0:
self.airy = True # check if Airy pattern is among provided polynomials
if all_are_polls:
- self.polynomials = zernpol; self.zernpol = None
+ self.polynomials = tuple(zernpol); self.zernpol = None
unique_osa_orders = set(osa_orders) # check that all OSA orders are unique
if len(osa_orders) > len(unique_osa_orders):
raise ValueError("There might be some repeated polynomials provided, "
+ f"# of provided: {len(osa_orders)} and # of unique: {len(unique_osa_orders)}")
else:
raise ValueError("Not all objects in Sequence are instances of the 'ZernPol' class")
- else:
- ValueError("ZernPSF class requires single 'ZernPol' instance or Sequence (class with __len__ attr.) with 'ZernPol' instances")
# %% Set properties
def set_physical_props(self, NA: float, wavelength: float, expansion_coeff: Union[float, Sequence[float]], pixel_physical_size: float):
@@ -195,6 +190,7 @@ def set_physical_props(self, NA: float, wavelength: float, expansion_coeff: Unio
raise ValueError("Wavelength should be positive real number")
self.k = 2.0*pi/self.wavelength # Calculate angular frequency (k)
self.NA = NA; self.wavelength = wavelength # save as the class properties
+ self.amplitudes: np.ndarray # set explicitly expected data type
# Sanity check of provided wavelength, pixel physical size (Nyquist criteria)
self.pixel_size_nyquist = 0.5*0.5*wavelength/NA # based on half of the Abbe resolution limit, see references in the docstring
self.pixel_size_nyquist_eStr = "{:.3e}".format(self.pixel_size_nyquist) # formatting calculated pixel size in scientific notation
@@ -205,7 +201,7 @@ def set_physical_props(self, NA: float, wavelength: float, expansion_coeff: Unio
+ f" computed from the Abbe's resolution limit (0.5*{lambda_char}/NA)")
self.pixel_size = pixel_physical_size
# Check which type is provided as the expansion_coeff parameter
- if hasattr(expansion_coeff, '__len__'):
+ if not isinstance(expansion_coeff, float):
coeffs_len = len(expansion_coeff) # for checking how many amplitudes provided
if coeffs_len != len(self.polynomials):
if coeffs_len == 1 and len(self.polynomials) == 0: # polynomials maybe provided also as a sequence with 1 element
@@ -213,7 +209,7 @@ def set_physical_props(self, NA: float, wavelength: float, expansion_coeff: Unio
# Sanity check for the expansion coefficient of the polynomial
self.__warn_message = _sanity_check_expansion_coefficient(abs(expansion_coeff) / wavelength)
if len(self.__warn_message) > 0:
- warnings.warn(self.__warn_message); self.__warn_message = ""
+ warnings.warn(self.__warn_message, stacklevel=2); self.__warn_message = ""
self.expansion_coeff = expansion_coeff; self.alpha = self.expansion_coeff / self.wavelength
else:
raise ValueError(f"Length of provided coefficients ({coeffs_len}) is not equal to stored number "
@@ -224,7 +220,7 @@ def set_physical_props(self, NA: float, wavelength: float, expansion_coeff: Unio
max_module_coeff = max(np.max(self.coefficients), abs(np.min(self.coefficients)))
self.__warn_message = _sanity_check_expansion_coefficient(abs(max_module_coeff) / wavelength, max_coeff_check=True)
if len(self.__warn_message) > 0:
- warnings.warn(self.__warn_message); self.__warn_message = ""
+ warnings.warn(self.__warn_message, stacklevel=2); self.__warn_message = ""
self.amplitudes = self.coefficients / self.wavelength
else:
# Check consistency of provided type of polynomials and coefficients
@@ -237,7 +233,7 @@ def set_physical_props(self, NA: float, wavelength: float, expansion_coeff: Unio
# Sanity check for the expansion coefficient of the polynomial
self.__warn_message = _sanity_check_expansion_coefficient(abs(expansion_coeff) / wavelength)
if len(self.__warn_message) > 0:
- warnings.warn(self.__warn_message); self.__warn_message = ""
+ warnings.warn(self.__warn_message, stacklevel=2); self.__warn_message = ""
self.expansion_coeff = expansion_coeff; self.alpha = self.expansion_coeff / self.wavelength
# Kernel size estimation (could be changed explicitly in the method 'set_calculation_props'). Redefine it for each call
if self.zernpol is not None: # for single polynomial
@@ -265,7 +261,7 @@ def set_physical_props(self, NA: float, wavelength: float, expansion_coeff: Unio
self.kernel_size += 1 # kernel size should odd
self.__physical_props_set = True # set internal flag True if no ValueError raised
- def set_calculation_props(self, kernel_size: int, n_integration_points_r: int, n_integration_points_phi: int) -> None:
+ def set_calculation_props(self, kernel_size: int, n_integration_points_r: int, n_integration_points_phi: int):
"""
Set calculation properties: kernel size, number of integration points on polar coordinates.
@@ -301,12 +297,12 @@ def set_calculation_props(self, kernel_size: int, n_integration_points_r: int, n
kernel_size += 1; print("Note that kernel_size should be odd integer for centering PSF kernel")
if self.__physical_props_set:
if kernel_size < self.kernel_size:
- self.__warn_message = (f"Empirically estimated kernel size = {self.kernel_size} based on the physical properties "
+ self.__warn_message = (f"\nEmpirically estimated kernel size = {self.kernel_size} based on the physical properties "
+ f"is larger than provided size = {kernel_size}. PSF may be truncated.")
- warnings.warn(self.__warn_message); self.__warn_message = ""
+ warnings.warn(self.__warn_message, stacklevel=2); self.__warn_message = ""
else:
- self.__warn_message = "Please set the physical properties first for estimation of kernel size and after call this method"
- warnings.warn(self.__warn_message); self.__warn_message = ""
+ self.__warn_message = "\nPlease set the physical properties first for estimation of kernel size and after call this method"
+ warnings.warn(self.__warn_message, stacklevel=2); self.__warn_message = ""
self.kernel_size = kernel_size # overwrite stored property, show before all warnings if kernel size is inconsistent
# Sanity checks for n_integration_points_r and n_integration_points_phi
if n_integration_points_r < 20:
@@ -317,13 +313,13 @@ def set_calculation_props(self, kernel_size: int, n_integration_points_r: int, n
# Associated with slow calculations warnings
if n_integration_points_r > 500 and n_integration_points_phi > 400:
self.__warn_message = "Selected integration precision may be unnecessary for PSF calculation and slow down it significantly"
- warnings.warn(self.__warn_message); self.__warn_message = ""
+ warnings.warn(self.__warn_message, stacklevel=2); self.__warn_message = ""
if abs(n_integration_points_phi - 300) < 40 and abs(n_integration_points_r - 320) < 50 and self.kernel_size > 25:
print(f"Note that the approximate calc. time: {int(round(self.kernel_size*self.kernel_size*38.5/1000, 0))} sec.")
# %% Calculation
def calculate_psf_kernel(self, suppress_warnings: bool = False, normalized: bool = True, verbose_info: bool = False,
- accelerated: bool = None) -> np.ndarray:
+ accelerated: Optional[bool] = None) -> np.ndarray:
"""
Calculate PSF kernel using the specified or default calculation parameters and physical values.
@@ -370,10 +366,10 @@ def calculate_psf_kernel(self, suppress_warnings: bool = False, normalized: bool
"""
if len(self.__warn_message) > 0 and not suppress_warnings:
- warnings.warn(self.__warn_message)
+ warnings.warn(self.__warn_message, stacklevel=2)
if not self.__physical_props_set and not suppress_warnings:
self.kernel_size = 19; self.__warn_message = "\nPhysical properties haven't been set before, the default ones will be used"
- warnings.warn(self.__warn_message); self.__warn_message = ""
+ warnings.warn(self.__warn_message, stacklevel=2); self.__warn_message = ""
global numba_installed # access the flag
# Check provided flag
if numba_installed and accelerated is None:
@@ -383,13 +379,13 @@ def calculate_psf_kernel(self, suppress_warnings: bool = False, normalized: bool
# Check if accelerated flag set to True but no numba installed
if accelerated and not numba_installed and not suppress_warnings:
self.__warn_message = "\nAcceleration isn't possible because 'numba' library not installed in the current environment"
- warnings.warn(self.__warn_message); self.__warn_message = ""
+ warnings.warn(self.__warn_message, stacklevel=2); self.__warn_message = ""
# For providing performance verbose report
if verbose_info:
t1 = time.perf_counter() # for explicit showing of performance
if self.kernel_size*self.kernel_size >= 301:
print("Kernel calculation started...")
- # Calculation using the vectorised form
+ # Calculation using the vectorized form
if self.zernpol is not None:
if not accelerated or (accelerated and not numba_installed):
self.kernel = get_psf_kernel(zernike_pol=self.zernpol, len2pixels=self.pixel_size, alpha=self.expansion_coeff,
@@ -403,6 +399,7 @@ def calculate_psf_kernel(self, suppress_warnings: bool = False, normalized: bool
kernel_size=self.kernel_size, n_int_r_points=self.n_int_r_points,
suppress_warns=suppress_warnings, n_int_phi_points=self.n_int_phi_points)
elif len(self.polynomials) > 0:
+ self.amplitudes: np.ndarray # set explicitly expected data type
if not accelerated or (accelerated and not numba_installed):
self.kernel = get_psf_kernel_zerns(polynomials=self.polynomials, amplitudes=self.amplitudes, len2pixels=self.pixel_size,
wavelength=self.wavelength, NA=self.NA, normalize_values=normalized, verbose=verbose_info,
@@ -451,7 +448,8 @@ def plot_kernel(self, id_str: str = ""):
fig_title = f"Composed kernel for #{len(self.polynomials)} polynomials: {orders} {id_str}"
if not plt.isinteractive():
plt.ion()
- plt.figure(fig_title, figsize=(6, 6)); plt.imshow(self.kernel, cmap=plt.cm.viridis, origin='upper'); plt.colorbar(); plt.tight_layout()
+ plt.figure(fig_title, figsize=(6, 6)); plt.imshow(self.kernel, cmap=plt.colormaps["viridis"], origin='upper')
+ plt.colorbar(); plt.tight_layout()
# %% Utilities
def convolute_img(self, image: np.ndarray, scale2original: bool = True) -> np.ndarray:
@@ -461,7 +459,7 @@ def convolute_img(self, image: np.ndarray, scale2original: bool = True) -> np.nd
Parameters
----------
image : numpy.ndarray
- Sample image, not colour.
+ Sample gray-scaled (2D array) image.
scale2original : bool
Convolution resulting image will be rescaled to the max intensity of the provided image if True. The default is True.
@@ -492,10 +490,10 @@ def visualize_convolution(self, radius: float = 4.0, max_intensity: int = 255):
target_disk = get_bumped_circle(radius, max_intensity) # get the sample image of the even centered circle with blurred edges
if not plt.isinteractive():
plt.ion()
- plt.figure("Sample Image: Disk", figsize=(6, 6)); plt.imshow(target_disk, cmap=plt.cm.viridis, origin='upper')
+ plt.figure("Sample Image: Disk", figsize=(6, 6)); plt.imshow(target_disk, cmap=plt.colormaps["viridis"], origin='upper')
plt.axis('off'); plt.tight_layout()
convolved_img = self.convolute_img(image=target_disk); plt.figure("Convolved PSF and Disk", figsize=(6, 6))
- plt.imshow(convolved_img, cmap=plt.cm.viridis, origin='upper'); plt.axis('off'); plt.tight_layout()
+ plt.imshow(convolved_img, cmap=plt.colormaps["viridis"], origin='upper'); plt.axis('off'); plt.tight_layout()
def crop_kernel(self, min_part_of_max: float = 0.01):
"""
@@ -544,7 +542,7 @@ def crop_kernel(self, min_part_of_max: float = 0.01):
self.kernel_size = self.kernel.shape[0]
# %% I/O methods
- def save_json(self, abs_path: Union[str, Path] = None, overwrite: bool = False):
+ def save_json(self, abs_path: Union[str, Path, None] = None, overwrite: bool = False):
"""
Save class attributes (PSF kernel, physical properties, etc.) in the JSON file for further reloading and avoiding long computation.
@@ -563,13 +561,14 @@ def save_json(self, abs_path: Union[str, Path] = None, overwrite: bool = False):
"""
if abs_path is None:
abs_path = Path(self.json_file_path).joinpath("saved_psfs") # default folder for storing saved JSON files (root for the package)
- if not Path.exists(abs_path):
- abs_path = "" # the folder "saved_psfs" will be created in the root of the package
if isinstance(abs_path, Path):
- abs_path = str(abs_path) # convert to the expected string format
+ if not Path.exists(abs_path):
+ abs_path = "" # the folder "saved_psfs" will be created in the root of the package
+ else:
+ abs_path = str(abs_path) # convert to the expected string format
if self.kernel_size == 1:
- self.__warn_message = "Kernel most likely hasn't been calculated, the kernel size == 1 - default value"
- warnings.warn(self.__warn_message); self.__warn_message = ""
+ self.__warn_message = "\nKernel most likely hasn't been calculated, the kernel size == 1 - default value"
+ warnings.warn(self.__warn_message, stacklevel=2); self.__warn_message = ""
if self.zernpol is not None: # single polynomial saving
self.json_file_path = save_psf(psf_kernel=self.kernel, NA=self.NA, wavelength=self.wavelength,
expansion_coefficient=self.expansion_coeff, pixel_size=self.pixel_size,
@@ -577,13 +576,14 @@ def save_json(self, abs_path: Union[str, Path] = None, overwrite: bool = False):
n_int_points_phi=self.n_int_phi_points, zernike_pol=self.zernpol,
overwrite=overwrite, folder_path=abs_path)
elif len(self.polynomials) > 0: # saving several polynomials
+ self.coefficients: np.ndarray # set explicitly data type for mypy
self.json_file_path = save_psf(psf_kernel=self.kernel, NA=self.NA, wavelength=self.wavelength,
expansion_coefficient=self.coefficients, pixel_size=self.pixel_size,
kernel_size=self.kernel_size, n_int_points_r=self.n_int_r_points,
n_int_points_phi=self.n_int_phi_points, zernike_pol=self.polynomials,
overwrite=overwrite, folder_path=abs_path)
- def read_json(self, abs_path: Union[str, Path] = None):
+ def read_json(self, abs_path: Union[str, Path, None] = None):
"""
Read the JSON file with the saved attributes and setting it for the class.
@@ -603,7 +603,7 @@ def read_json(self, abs_path: Union[str, Path] = None):
abs_path = str(abs_path)
json_data = read_psf(abs_path) # raw parsed data from a file
if json_data is not None:
- wavelen: float; na: float; a: float; pols: list; ps: float; read_props = 0
+ wavelen: float; na: float; a: Union[float, Sequence[float]]; pols: list; ps: float; read_props = 0
for key, item in json_data.items():
# Calculation properties + calculated kernel
if key == "PSF Kernel":
@@ -619,10 +619,8 @@ def read_json(self, abs_path: Union[str, Path] = None):
na = item; read_props += 1
elif key == "Wavelength":
wavelen = item; read_props += 1
- elif key == "Expansion Coefficient":
+ elif key == "Expansion Coefficient" or key == "Amplitudes":
a = item; read_props += 1
- elif key == "Amplitudes":
- a = np.asarray(item); read_props += 1
elif key == "Pixel Size":
ps = item; read_props += 1
# Getting and reassigning used polynomial(-s)
@@ -639,23 +637,23 @@ def read_json(self, abs_path: Union[str, Path] = None):
else:
for index in osa_indices:
pols.append(ZernPol(osa=index))
- self.polynomials = pols; self.zernpol = None; self.airy = False
+ self.polynomials = tuple(pols); self.zernpol = None; self.airy = False
# Assign read physical properties
if read_props == 8:
self.set_physical_props(NA=na, wavelength=wavelen, expansion_coeff=a, pixel_physical_size=ps)
else:
- self.__warn_message = "Provided json file doesn't contain all necessary keys for physical / calculation properties"
- warnings.warn(self.__warn_message); self.__warn_message = ""
+ self.__warn_message = "\nProvided json file doesn't contain all necessary keys for physical / calculation properties"
+ warnings.warn(self.__warn_message, stacklevel=2); self.__warn_message = ""
else:
- self.__warn_message = "Provided path doesn't contain valid JSON data"
- warnings.warn(self.__warn_message); self.__warn_message = ""
+ self.__warn_message = "\nProvided path doesn't contain valid JSON data"
+ warnings.warn(self.__warn_message, stacklevel=2); self.__warn_message = ""
# %% Parallelized computing methods
- def initialize_parallel_workers(self):
+ def __initialize_parallel_workers(self):
"""
Initialize 4 Processes() for performing integration.
- See intmproc.py script (utils module) for implementation details.
+ See intmproc.py script (utils module) for implementation details.\n
Tests showed that this way doesn't provide performance gain.
Returns
@@ -668,7 +666,7 @@ def initialize_parallel_workers(self):
self.__ParallelCalc = DispenserManager(compute_func=radial_integral_s, params_list=self.__integration_params,
n_workers=4, verbose_info=False)
- def get_psf_point_r_parallel(self, r: float, theta: float) -> float:
+ def __get_psf_point_r_parallel(self, r: float, theta: float) -> float:
"""
Parallel implementation of numerical integration.
@@ -704,7 +702,7 @@ def get_psf_point_r_parallel(self, r: float, theta: float) -> float:
integral_sum = (h_phi/3.0)*(yA + yB + 2.0*even_sum + 4.0*odd_sum); integral_normalization = 1.0/(pi*pi)
return np.power(np.abs(integral_sum), 2)*integral_normalization
- def get_kernel_parallel(self, normalize_values: bool = True) -> np.ndarray:
+ def __get_kernel_parallel(self, normalize_values: bool = True) -> np.ndarray:
"""
Parallelized implementation of PSF kernel calculation.
@@ -723,7 +721,7 @@ def get_kernel_parallel(self, normalize_values: bool = True) -> np.ndarray:
"""
if self.kernel_size < 3 and self.zernpol is not None:
self.kernel_size = get_kernel_size(zernike_pol=self.zernpol, len2pixels=self.pixel_size, alpha=self.alpha,
- wavelength=self.wavelength)
+ wavelength=self.wavelength, NA=self.NA)
if self.zernpol is None:
self.kernel_size = 3
self.kernel = np.zeros(shape=(self.kernel_size, self.kernel_size)); i_center = self.kernel_size//2; j_center = self.kernel_size//2
@@ -736,7 +734,7 @@ def get_kernel_parallel(self, normalize_values: bool = True) -> np.ndarray:
theta = np.arctan2((i - i_center), (j - j_center)) # The PSF also has the angular dependency, not only the radial one
theta += np.pi # shift angles to the range [0, 2pi]
if self.zernpol is not None:
- self.kernel[i, j] = self.get_psf_point_r_parallel(r=distance, theta=theta); calculated_points += 1
+ self.kernel[i, j] = self.__get_psf_point_r_parallel(r=distance, theta=theta); calculated_points += 1
# print(f"Calculated point {[i, j]} from {[self.kernel_size-1, self.kernel_size-1]}")
print(f"Calculated point #{calculated_points} from {self.kernel_size*self.kernel_size}, takes ms: ",
int(round(1000*(time.perf_counter() - t1), 0)))
@@ -744,7 +742,7 @@ def get_kernel_parallel(self, normalize_values: bool = True) -> np.ndarray:
self.kernel /= np.max(self.kernel)
return self.kernel
- def deinitialize_workers(self):
+ def __deinitialize_workers(self):
"""
Release initialized before Processes for performing parallel computation.
@@ -786,7 +784,7 @@ def _sanity_check_expansion_coefficient(normalized_coefficient: float, max_coeff
return ""
-def force_get_psf_compilation(verbose_report: bool = False) -> Union[tuple, None]:
+def force_get_psf_compilation(verbose_report: bool = False) -> Optional[Tuple[ZernPSF, ZernPSF]]:
"""
Force compilation of computing functions for round and ellipse 'precise' shaped objects.
@@ -800,15 +798,14 @@ def force_get_psf_compilation(verbose_report: bool = False) -> Union[tuple, None
"""
if numba_installed:
if verbose_report:
- print("Precompilation started..."); t1 = time.perf_counter() # for explicit showing of performance
+ print("Precompilation of 'zernpy' started...", flush=True); t1 = time.perf_counter() # for explicit showing of performance
# Check the version of numba for guarantee working of compilation calls
try:
- if int(numba_ver_n[0]) == 0:
- if int(numba_ver_n[1]) <= 57:
- if int(numba_ver_n[2]) < 1:
- __warn_message = "\nRecommended numba version for compilation >= 0.57.1"; warnings.warn(__warn_message)
+ if int(numba_ver_n[0]) == 0 and int(numba_ver_n[1]) <= 57 and int(numba_ver_n[2]) < 1:
+ __warn_message = "\nRecommended numba version for compilation >= 0.57.1"; warnings.warn(__warn_message, stacklevel=2)
except ValueError:
- __warn_message = "\nCannot parse 'numba' version, expected in format 'x.yy.zz' - all integers"; warnings.warn(__warn_message)
+ __warn_message = "\nCannot parse 'numba' version, expected in format 'x.yy.zz' - all integers"
+ warnings.warn(__warn_message, stacklevel=2)
# Compilation of methods for single polynomial PSF calculation
NA = 0.2; wavelength = 0.55; pixel_size = wavelength / 0.82; ampl = 0.042
zp_prc = ZernPol(m=-1, n=1); zpsf_prc = ZernPSF(zp_prc) # Airy pattern - for compilation methods for single polynomial
@@ -822,189 +819,45 @@ def force_get_psf_compilation(verbose_report: bool = False) -> Union[tuple, None
zpsf_mpc.kernel_size = 3 # force to calculate only small kernel size for saving sometime
# print("Precompilation kernel size for several pol.:", zpsf_mpc.kernel_size)
zpsf_mpc.calculate_psf_kernel(normalized=True, accelerated=True, suppress_warnings=True, verbose_info=False)
- set_methods_compiled() # setting the flag implicitly in the module
- global num_methods_compiled; num_methods_compiled = True # setting copied variable in this script
if verbose_report:
passed_time_ms = int(round(1000.0*(time.perf_counter() - t1), 0))
if passed_time_ms > 1000:
- print(f"Precompilation took: {round(passed_time_ms/1000.0, 2)} sec.")
+ print(f"Precompilation took: {round(passed_time_ms/1000.0, 2)} sec.", flush=True)
else:
- print(f"Precompilation took: {passed_time_ms} ms.")
- print("--------------------------------------------")
+ print(f"Precompilation took: {passed_time_ms} ms.", flush=True)
+ print("--------------------------------------------", flush=True)
return zpsf_prc, zpsf_mpc
else:
__warn_message = "\nAcceleration isn't possible because 'numba' library not installed in the current environment"
- warnings.warn(__warn_message)
+ warnings.warn(__warn_message, stacklevel=2)
return None
+def clean_zernpy_cache() -> bool:
+ """
+ Clean locally cached files created by numba after compilation of computation methods from the 'calc_psfs_numba' module.
+
+ Returns
+ -------
+ bool \n
+ Flag if local cache is assumed (at least 6 cached files successfully removed) to be cleaned.
+
+ """
+ _local_cache_cleaned = False # flag to show that local numba cache found and has been cleaned
+ if numba_installed:
+ root = Path(__file__).resolve().parent; cache_folder = root.joinpath("calculations").joinpath("__pycache__")
+ if cache_folder.exists() and cache_folder.is_dir():
+ _n_cleaned = 0
+ for obj in cache_folder.iterdir():
+ if obj.is_file() and (obj.suffix == ".nbc" or obj.suffix == ".nbi"):
+ try:
+ obj.unlink(missing_ok=True); _n_cleaned += 1
+ except (PermissionError, OSError):
+ _n_cleaned = 0; break # automatically breaks loop assuming that invalid permission is universal
+ if _n_cleaned >= 6: # 8 - minimal count of caches created by numba, assume that cache is cleaned if no error encountered
+ _local_cache_cleaned = True
+ return _local_cache_cleaned
+
+
# %% Define default export classes and methods used with import * statement (import * from zernpsf)
-__all__ = ['ZernPSF', 'force_get_psf_compilation']
-
-
-# %% Test as the main script
-if __name__ == "__main__":
- plt.close("all") # close all opened before figures
- wavelength_um = 0.55 # used below in a several calls
- check_other_pols = False; check_small_na_wl = False # flag for checking some other polynomials PSFs
- check_airy = False; check_common_psf = False; check_io_kernel = False; check_parallel_calculation = False; check_test = False
- check_faster_airy = False; check_test_conditions = False; check_test_conditions2 = False; check_several_pols = False
- check_edge_conditions = False; test_acceleration_single_pol = False; test_acceleration_few_pol = False
- prepare_pic_readme = False # for plotting the sum of polynomials produced profile
- test_io_few_pols = False; standard_path = Path.home().joinpath("Desktop") # for saving json on the Desktop
- check_acceleration_flag = False # for testing fallback calculation with the wrong flag for calculate PSF kernel
- check_init_several_pols = False # check calculation several pol-s: Airy (Piston) + some of polynomial
- check_airy_patterns = False # check the difference between calculated by equation for Airy pattern and diffraction integral
- check_precompilation = False # checks precompilation
- check_cropping = False # checks how kernel is cropped
-
- # Common PSF for testing
- if check_common_psf:
- force_get_psf_compilation(True)
- zpsf1 = ZernPSF(ZernPol(m=1, n=1)); zpsf2 = ZernPSF(ZernPol(m=-3, n=3)); ampl = -0.43
- zpsf2.set_physical_props(NA=0.95, wavelength=wavelength_um, expansion_coeff=ampl, pixel_physical_size=wavelength_um/5.05)
- zpsf1.set_physical_props(NA=0.5, wavelength=0.6, expansion_coeff=0.25, pixel_physical_size=wavelength_um/5.5)
- zpsf2.set_calculation_props(kernel_size=zpsf2.kernel_size, n_integration_points_r=250, n_integration_points_phi=320)
- zpsf1.set_calculation_props(kernel_size=zpsf2.kernel_size, n_integration_points_r=220, n_integration_points_phi=360)
- kernel3 = zpsf2.calculate_psf_kernel(suppress_warnings=False, verbose_info=True); zpsf2.plot_kernel("for loop")
- zpsf2.visualize_convolution() # Visualize convolution on the disk image
- if check_io_kernel:
- # default_path = Path(__file__).parent.joinpath("saved_psfs").absolute() # default path for storing JSON files
- zpsf2.save_json(overwrite=True, abs_path=standard_path)
- zpsf1.read_json(zpsf2.json_file_path) # for testing reading and assigning values to ZernPSF class (substitution)
- if check_parallel_calculation:
- zpsf2.initialize_parallel_workers(); kernel2 = zpsf2.get_kernel_parallel()
- zpsf2.plot_kernel("parallel"); zpsf2.deinitialize_workers()
-
- # Another Zernike polynomial, big NA
- if check_other_pols:
- zpsf3 = ZernPSF(ZernPol(m=0, n=4))
- zpsf3.set_physical_props(NA=1.25, wavelength=wavelength_um, expansion_coeff=0.4, pixel_physical_size=wavelength_um/5.25)
- zpsf3.calculate_psf_kernel(suppress_warnings=False, verbose_info=True); zpsf3.plot_kernel()
-
- # Another Zernike polynomial, average to small NA and wavelength
- if check_small_na_wl:
- NA = 0.45; wavelength = 0.4; pixel_size = wavelength / 5.25; ampl = 0.18
- zp2 = ZernPol(m=1, n=3); zpsf2 = ZernPSF(zp2) # horizontal coma
- zpsf2.set_physical_props(NA=NA, wavelength=wavelength, expansion_coeff=ampl, pixel_physical_size=pixel_size)
- zpsf2.calculate_psf_kernel(normalized=True); zpsf2.plot_kernel()
-
- if check_airy:
- NA = 0.12; wavelength = 0.8; pixel_size = wavelength / 4.0; ampl = 1.25
- zp4 = ZernPol(m=0, n=0); zpsf4 = ZernPSF(zp4) # piston for the Airy pattern
- zpsf4.set_physical_props(NA=NA, wavelength=wavelength, expansion_coeff=ampl, pixel_physical_size=pixel_size)
- zpsf4.calculate_psf_kernel(normalized=True); zpsf4.plot_kernel("Plus")
- zpsf4.set_physical_props(NA=NA, wavelength=wavelength, expansion_coeff=-ampl, pixel_physical_size=pixel_size)
- zpsf4.calculate_psf_kernel(normalized=True); zpsf4.plot_kernel("Minus")
-
- if check_faster_airy: # For set the test for pytest library
- NA = 0.35; wavelength = 0.55; pixel_size = wavelength / 3.05; ampl = -0.4
- zp4 = ZernPol(m=0, n=0); zpsf4 = ZernPSF(zp4) # piston for the Airy pattern
- zpsf4.set_physical_props(NA=NA, wavelength=wavelength, expansion_coeff=ampl, pixel_physical_size=pixel_size)
- zpsf4.calculate_psf_kernel(normalized=True); zpsf4.plot_kernel()
-
- if check_test_conditions:
- NA = 0.95; wavelength = 0.55; pixel_size = wavelength / 5.0; ampl = 0.55
- zp6 = ZernPol(m=0, n=2); zpsf6 = ZernPSF(zp6) # defocus
- zpsf6.set_physical_props(NA=NA, wavelength=wavelength, expansion_coeff=ampl, pixel_physical_size=pixel_size) # normal assignment
- zpsf6.set_calculation_props(kernel_size=zpsf6.kernel_size, n_integration_points_r=250, n_integration_points_phi=300)
- zpsf6.calculate_psf_kernel(normalized=True); zpsf6.plot_kernel()
- if check_test_conditions2:
- zp7 = ZernPol(m=1, n=3); zpsf7 = ZernPSF(zp7) # horizontal coma
- NA = 0.4; wavelength = 0.4; pixel_size = wavelength / 3.2; ampl = 0.185 # Common physical properties
- zpsf7.set_physical_props(NA=NA, wavelength=wavelength, expansion_coeff=ampl, pixel_physical_size=pixel_size)
- zpsf7.calculate_psf_kernel(normalized=True); zpsf7.plot_kernel()
-
- # Test calculation of a PSF for several polynomials and I/O operations (see flags)
- if check_several_pols:
- zp1 = ZernPol(m=-2, n=2); zp2 = ZernPol(m=0, n=2); zp3 = ZernPol(m=2, n=2); pols = (zp1, zp2, zp3); coeffs = (-0.36, 0.25, 0.4)
- zpsf8 = ZernPSF(pols); zpsf8.set_physical_props(NA=0.95, wavelength=0.5, expansion_coeff=coeffs, pixel_physical_size=0.5/4.5)
- composed_kernel = zpsf8.calculate_psf_kernel(normalized=True, verbose_info=True); zpsf8.plot_kernel()
- if test_io_few_pols:
- zpsf14 = ZernPSF(ZernPol(osa=19)); zpsf14.set_physical_props(NA=0.1, wavelength=0.4, expansion_coeff=0.82,
- pixel_physical_size=0.05)
- zpsf8.save_json(overwrite=True, abs_path=standard_path); zpsf14.read_json(zpsf8.json_file_path) # save / read calculated kernel
-
- # Test some edge conditions - e.g., specifying 1 polynomial in a list with huge coefficient
- if check_edge_conditions:
- zp4 = ZernPol(m=3, n=3); pols2 = [zp4]; coeff = 5.1
- zpsf9 = ZernPSF(pols2); zpsf9.set_physical_props(NA=0.95, wavelength=0.5, expansion_coeff=coeff, pixel_physical_size=0.5/4.75)
- zpsf9.set_physical_props(NA=0.95, wavelength=0.5, expansion_coeff=(-0.67), pixel_physical_size=0.5/4.75)
- zpsf9.calculate_psf_kernel(normalized=True, verbose_info=True); zpsf9.plot_kernel()
- # Test acceleration by using numba library utilities
- if test_acceleration_single_pol:
- force_get_psf_compilation(); NA = 0.95; wavelength = 0.55; pixel_size = wavelength / 4.6; ampl = -0.16
- zp6 = ZernPol(m=0, n=2); zpsf6 = ZernPSF(zp6) # defocus
- zpsf6.set_physical_props(NA=NA, wavelength=wavelength, expansion_coeff=ampl, pixel_physical_size=pixel_size)
- kernel_acc = zpsf6.calculate_psf_kernel(normalized=True, accelerated=True, verbose_info=True); zpsf6.plot_kernel("Accelerated")
- kernel_norm = zpsf6.calculate_psf_kernel(normalized=True); zpsf6.plot_kernel("Normal")
- if test_acceleration_few_pol:
- force_get_psf_compilation()
- zp1 = ZernPol(m=-2, n=2); zp2 = ZernPol(m=0, n=2); zp3 = ZernPol(m=2, n=2); pols = (zp1, zp2, zp3); coeffs = (-0.12, 0.15, 0.1)
- zpsf8 = ZernPSF(pols); zpsf8.set_physical_props(NA=0.35, wavelength=0.5, expansion_coeff=coeffs, pixel_physical_size=0.5/1.5)
- composed_kernel = zpsf8.calculate_psf_kernel(normalized=True, verbose_info=True); zpsf8.plot_kernel("Normal")
- composed_kernel_acc = zpsf8.calculate_psf_kernel(normalized=True, verbose_info=True, accelerated=True)
- zpsf8.plot_kernel("Accelerated")
- if prepare_pic_readme:
- force_get_psf_compilation(verbose_report=True)
- zp1 = ZernPol(m=-1, n=3); zp2 = ZernPol(m=2, n=4); zp3 = ZernPol(m=0, n=4); pols = (zp1, zp2, zp3); coeffs = (0.5, 0.21, 0.15)
- zpsf_pic = ZernPSF(pols); zpsf_pic.set_physical_props(NA=0.65, wavelength=0.6, expansion_coeff=coeffs, pixel_physical_size=0.6/5.0)
- zpsf_pic.calculate_psf_kernel(normalized=True, verbose_info=True, accelerated=True)
- zpsf_pic.plot_kernel("Vert. Coma Vert. 2nd Astigmatism Spherical")
- if check_acceleration_flag:
- zp16 = ZernPol(m=-1, n=3); zp18 = ZernPol(m=2, n=4); pols10 = (zp16, zp18); coeffs10 = (-0.1, 0.13); zpsf_acc = ZernPSF(pols10)
- zpsf_norm = ZernPSF(pols10); zpsf_acc.set_physical_props(NA=0.43, wavelength=0.6, expansion_coeff=coeffs10,
- pixel_physical_size=0.6/3.0)
- zpsf_norm.set_physical_props(NA=0.43, wavelength=0.6, expansion_coeff=coeffs10, pixel_physical_size=0.6/3.0)
- kern_acc = zpsf_acc.calculate_psf_kernel(normalized=True, accelerated=True)
- kern_norm = zpsf_norm.calculate_psf_kernel(normalized=True)
- kern_diff = np.round(kern_acc - kern_norm, 9) # for checking the difference in calculations
-
- if check_test:
- pols = (ZernPol(osa=10), ZernPol(osa=15)); coeffs = (0.28, -0.33); NA = 0.35; wavelength = 0.55
- zpsf = ZernPSF(pols); zpsf.set_physical_props(NA, wavelength, expansion_coeff=coeffs, pixel_physical_size=wavelength / 3.5)
- zpsf.set_calculation_props(kernel_size=25, n_integration_points_r=200, n_integration_points_phi=180)
- psf_kernel = zpsf.calculate_psf_kernel(normalized=False); zpsf.plot_kernel()
-
- if check_init_several_pols:
- force_get_psf_compilation(True)
- zpsf30 = ZernPSF(zernpol=(ZernPol(osa=1), ZernPol(osa=5)))
- try:
- zpsf30.set_physical_props(NA=0.95, wavelength=0.5, expansion_coeff=0.5, pixel_physical_size=0.5/5.0)
- except ValueError:
- print("Check for 1 ampl. and 2 pol-s passed") # as expected, transfer to test function
- zpsf31 = ZernPSF(zernpol=(ZernPol(osa=0), ZernPol(m=-1, n=3)))
- zpsf31.set_physical_props(NA=0.95, wavelength=0.5, expansion_coeff=[1.5, 0.24], pixel_physical_size=0.5/3.8)
- # zpsf31.calculate_psf_kernel(verbose_info=True, accelerated=False); zpsf31.plot_kernel("Not accelerated 2 pol-s")
- # kernel_not_acc = np.copy(zpsf31.kernel)
- zpsf31.calculate_psf_kernel(verbose_info=True, accelerated=True); zpsf31.plot_kernel("Airy 1.5")
- kernel_acc = np.copy(zpsf31.kernel)
- # zpsf31.kernel = np.abs(kernel_acc - kernel_not_acc); zpsf31.plot_kernel("Diff. 2 pol-s")
- zpsf31.set_physical_props(NA=0.95, wavelength=0.5, expansion_coeff=[-1.5, 0.24], pixel_physical_size=0.5/3.8)
- zpsf31.calculate_psf_kernel(verbose_info=True, accelerated=True); zpsf31.plot_kernel("Airy -1.5")
- kernel_neg_acc = np.copy(zpsf31.kernel)
- zpsf31.kernel = np.abs(kernel_acc - kernel_neg_acc); zpsf31.plot_kernel("Diff. neg. pos. Airy + Coma")
-
- if check_airy_patterns:
- force_get_psf_compilation(verbose_report=True)
- zpsf40 = ZernPSF(zernpol=ZernPol(osa=0))
- zpsf40.set_physical_props(NA=0.95, wavelength=0.5, expansion_coeff=3.0, pixel_physical_size=0.5/4.0)
- zpsf40.calculate_psf_kernel(verbose_info=True, accelerated=False, normalized=False); zpsf40.plot_kernel("Not accelerated Airy")
- kernel_not_acc = np.copy(zpsf40.kernel)
- zpsf40.calculate_psf_kernel(verbose_info=True, accelerated=True, normalized=False); zpsf40.plot_kernel("Accelerated Airy")
- kernel_acc = np.copy(zpsf40.kernel)
- zpsf40.kernel = np.abs(kernel_acc - kernel_not_acc); zpsf40.plot_kernel("Diff. Airy")
-
- if check_precompilation:
- force_get_psf_compilation(True)
- zpsf50 = ZernPSF(zernpol=(ZernPol(m=1, n=3)))
- zpsf50.set_physical_props(NA=1.25, wavelength=0.52, expansion_coeff=0.5, pixel_physical_size=0.5/4.85)
- zpsf50.calculate_psf_kernel(verbose_info=True, accelerated=True, normalized=True); zpsf50.plot_kernel()
-
- if check_cropping:
- zpsf60 = ZernPSF(zernpol=(ZernPol(m=0, n=4)))
- zpsf60.set_physical_props(NA=1.25, wavelength=0.5, expansion_coeff=0.47, pixel_physical_size=0.5/5.0)
- zpsf60.calculate_psf_kernel(accelerated=True, verbose_info=True); zpsf60.plot_kernel("Not Cropped")
- original_kernel = np.copy(zpsf60.kernel)
- zpsf60.crop_kernel(min_part_of_max=0.025); zpsf60.plot_kernel("Cropped"); cropped_kernel = np.copy(zpsf60.kernel)
- print("Original kernel shape:", original_kernel.shape, "\nCropped kernel shape:", cropped_kernel.shape)
+__all__ = ['ZernPSF', 'force_get_psf_compilation', 'clean_zernpy_cache']
diff --git a/tests/run_zernikepol_as_main.py b/tests/run_zernikepol_as_main.py
new file mode 100644
index 0000000..95a52d5
--- /dev/null
+++ b/tests/run_zernikepol_as_main.py
@@ -0,0 +1,194 @@
+# -*- coding: utf-8 -*-
+"""
+Run a few methods from 'zernikepol' as the main one for performing tests in IDE.
+
+@author: Sergei Klykov, '@sklykov' on GitHub
+
+@licence: MIT, @year: 2026
+
+"""
+# %% Imports and checking its validity
+import importlib
+import time
+from contextlib import suppress
+from typing import Tuple
+
+import matplotlib
+import numpy as np
+
+# Explicit backend assignment for matplotlib - for compatibility between running configurations in Spyder and PyCharm IDEs
+with suppress(ImportError): # will be thrown in the environment doesn't contain Qt-like library
+ matplotlib.use('Qt5Agg')
+import matplotlib.pyplot as plt
+
+# Import of developed version of package, install in the editable mode: pip install -e .
+import zernpy.zernikepol
+
+importlib.reload(zernpy.zernikepol) # trick to guarantee local changes to be always used for run
+from zernpy.zernikepol import ZernPol, fit_polynomials, generate_phases_image, generate_random_phases, zernikes_surface
+
+
+# %% Test functions for the external call
+def compare_performances(min_order: int, max_order: int) -> Tuple[int, int, str]:
+ """
+ Compare performances of radial polynomials calculation by using recursive and exact equations.
+
+ Comparison achieved by simple measuring of time in msec needed for calculation of all radial
+ polynomials from minimal radial order up to maximum radial order returned as tuple.
+
+ Parameters
+ ----------
+ min_order : int
+ Minimum radial order of used polynomials (n).
+ max_order : int
+ Maximum radial order of used polynomials (n).
+
+ Returns
+ -------
+ tuple
+ Composed by time for recursive calculation and time for exact calculation.
+
+ """
+ # Generation of orders in OSA/ANSI indexing scheme for initializing Zernike polynomials
+ zernpols = []
+ for order in range(min_order, max_order+1):
+ m = -order; n = order
+ zernpols.append(ZernPol(m=m, n=n))
+ for _ in range(0, order):
+ m += 2; zernpols.append(ZernPol(m=m, n=n))
+ # Generation numpy array with radii
+ n_points = 251
+ test_r = np.zeros(shape=(n_points, ))
+ for i in range(n_points):
+ test_r[i] = i/(n_points-1)
+ test_r = np.round(test_r, 6)
+ # Measuring performance of radial polynomials calculations using recursive implementation
+ t1 = time.perf_counter()
+ for polynomial in zernpols:
+ polynomial.radial(test_r) # calculate radial polynomials over vector of radii
+ t2 = time.perf_counter()
+ t_recursive_ms = round(1000*(t2-t1), 3)
+ # Measuring performance of radial polynomials calculations using exact implementation
+ t1 = time.perf_counter()
+ for polynomial in zernpols:
+ polynomial.radial(test_r, use_exact_eq=True) # calculate radial polynomials over vector of radii
+ t2 = time.perf_counter()
+ t_exact_ms = round(1000*(t2-t1), 3)
+ return t_recursive_ms, t_exact_ms, f"Used polynomials: {i}", f"Radii: {n_points}"
+
+
+def _estimate_high_order_calc_times():
+ """
+ Estimate slowing down of radial polynomial calculation with increasing of radial order.
+
+ Returns
+ -------
+ None.
+
+ """
+ r = 0.55
+ high_order_pols = [ZernPol(m=-2, n=46), ZernPol(m=1, n=47), ZernPol(m=2, n=48), ZernPol(m=-1, n=49),
+ ZernPol(m=2, n=50), ZernPol(m=1, n=51), ZernPol(m=-2, n=52)]
+ times = []
+ # Single polynomials values
+ for pol in high_order_pols:
+ # calculate radial polynomials over vector of radii
+ t1 = time.perf_counter(); pol.radial(r); t2 = time.perf_counter()
+ times.append(f"{pol.get_mn_orders()}: {int(round(1000*(t2-t1), 0))} ms")
+ # Delete polynomials with too high orders for derivatives calculates
+ high_order_pols.pop(len(high_order_pols)-1); high_order_pols.pop(len(high_order_pols)-1)
+ high_order_pols.pop(len(high_order_pols)-1)
+ print(times); times = []
+ # Single derivative polynomials values
+ for pol in high_order_pols:
+ # calculate derivatives of radial polynomials over vector of radii
+ t1 = time.perf_counter(); pol.radial_dr(r); t2 = time.perf_counter()
+ times.append(f"Deriv. {pol.get_mn_orders()}: {int(round(1000*(t2-t1), 0))} ms")
+ print(times)
+
+
+# %% Tests
+if __name__ == "__main__":
+ _test_plots = True # regulates testing of plotting various plots
+ _test_calculations = False # regulates tests below concerning calculations
+
+ # Testing plotting, the plots will be opened in the additional pop-up windows
+ if _test_plots:
+ plt.close("all") # close all previously opened plots
+ t1 = time.perf_counter(); zp = ZernPol(m=0, n=2); ZernPol.plot_profile(zp, color_map="jet", show_title=True) # basic plot
+ t2 = time.perf_counter(); print("Plotting of 1 non-zero polynomial takes ms: ", int(round(1000*(t2-t1), 0)))
+ coordinates = ZernPol.gen_polar_coordinates(r_step=0.005)
+ zp = ZernPol(m=-10, n=30); ZernPol.plot_profile(zp, color_map="jet", show_title=False, polar_coordinates=coordinates) # high order plot
+ zp = ZernPol(m=0, n=0); ZernPol.plot_profile(zp, color_map="turbo", show_title=True) # plot of piston polynomial
+
+ # Testing 3D surface plotting
+ ZernPol.plot_profile(ZernPol(m=0, n=2), color_map="viridis", projection="3d")
+
+ # Testing 3D figure plotting on the externally initialized Figure class
+ fig3d = plt.figure(figsize=(6.8, 6.8))
+ zern_surface = ZernPol.gen_zernikes_surface([1.0], [ZernPol(m=0, n=2)], equal_n_coordinates=True, n_points=400)
+ ZernPol.plot_sum_zernikes_on_fig(figure=fig3d, use_defaults=False, zernikes_sum_surface=zern_surface,
+ show_range=True, color_map="magma", projection="3D")
+ fig3d2 = plt.figure(figsize=(5.8, 5.8))
+ ZernPol.plot_sum_zernikes_on_fig(figure=fig3d2, coefficients=[1.0], polynomials=[ZernPol(m=0, n=2)],
+ show_range=True, color_map="bwr", projection="3D")
+
+ # Testing accelerated plotting / sum calculation
+ fig3 = plt.figure(figsize=(3, 3))
+ t1 = time.perf_counter(); n_pols = 31; polynomials = []; coefficients = [0.0]*n_pols
+ for i in range(n_pols):
+ polynomials.append(ZernPol(osa=58+i))
+ coefficients[0] = 1.0 # only 1st polynomial will be plotted
+ fig3 = ZernPol.plot_sum_zernikes_on_fig(figure=fig3, coefficients=coefficients, polynomials=polynomials,
+ show_range=False, color_map="turbo")
+ fig3.subplots_adjust(0, 0, 1, 1)
+ t2 = time.perf_counter(); print("Plotting of 1 non-zero and 30 zero pol-s takes ms: ", int(round(1000*(t2-t1), 0)))
+
+ # Tests with generation / restoring Zernike profiles (phases images)
+ phases_image, polynomials_ampls, polynomials = generate_random_phases(img_height=301, img_width=321)
+ plt.figure(); plt.axis("off"); plt.imshow(phases_image, cmap="jet"); plt.tight_layout(); plt.subplots_adjust(0, 0, 1, 1)
+ polynomials_amplitudes, cropped_img = fit_polynomials(phases_image, polynomials, return_cropped_image=True,
+ strict_circle_border=False, crop_radius=1.0)
+ plt.figure(); plt.axis("off"); plt.imshow(cropped_img, cmap="jet")
+ plt.tight_layout(); plt.subplots_adjust(0, 0, 1, 1)
+
+ # Updated test of fitting including piston polynomial
+ height = 500; width = 481; crop_r = 1.0; strict_border = True; pols_coeffs = [-0.75, 0.86, 0.41]; fig4 = plt.figure(figsize=(4, 4))
+ polynomials = [ZernPol(osa=0), ZernPol(m=0, n=2), ZernPol(m=-3, n=3)]; rs, angles = ZernPol.gen_polar_coordinates()
+ phase_profile = ZernPol.sum_zernikes(coefficients=pols_coeffs, polynomials=polynomials, r=rs, theta=angles, get_surface=True)
+ # Below - plotting specified polynomials on the polar coordinates
+ ZernPol.plot_sum_zernikes_on_fig(figure=fig4, use_defaults=False, zernikes_sum_surface=zernikes_surface(phase_profile, rs, angles),
+ color_map="jet")
+ # Below - generate phases image with the cartesian coordinates
+ phases_image2 = generate_phases_image(polynomials=tuple(polynomials), polynomials_amplitudes=tuple(pols_coeffs),
+ img_height=height, img_width=width)
+ plt.figure(); plt.axis("off"); im = plt.imshow(phases_image2, cmap="jet"); plt.tight_layout(); plt.subplots_adjust(0, 0, 1, 1)
+ plt.colorbar(mappable=im)
+ # Below - fitting procedure on the provided phases image
+ polynomials_amplitudes2, cropped_img2 = fit_polynomials(phases_image2, polynomials, return_cropped_image=True,
+ strict_circle_border=strict_border, crop_radius=crop_r)
+ print("Difference between used amplitudes and fitted ones:", np.asarray(pols_coeffs) - polynomials_amplitudes2)
+ plt.figure(); plt.axis("off"); im = plt.imshow(cropped_img2, cmap="jet"); plt.tight_layout(); plt.subplots_adjust(0, 0, 1, 1)
+ plt.colorbar(mappable=im)
+
+ plt.show() # show all images created by plt.figure() calls
+
+ # Testing calculations and their performance comparison
+ if _test_calculations:
+ # Simple test of two concepts of calculations - exact and recursive equations
+ z = ZernPol(n=30, m=-2); print("Diff. between recursive and exact equations:",
+ round(z.radial(0.85) - z.radial(0.85, use_exact_eq=True), 9))
+ z = ZernPol(n=32, l=0); print("Diff. between recursive and exact equations:",
+ round(z.radial(0.35) - z.radial(0.35, use_exact_eq=True), 9))
+ r = 0.955; theta = np.pi/8; z = ZernPol(osa=55)
+ print("Diff. between recursive and exact equations:",
+ round(z.polynomial_value(r, theta) - z.polynomial_value(r, theta, use_exact_eq=True), 9))
+ z = ZernPol(n=35, l=-1); print("Diff. between recursive & exact eq-s for derivatives:",
+ round(z.radial_dr(0.78) - z.radial_dr(0.78, use_exact_eq=True), 9))
+ z = ZernPol(n=38, m=-2); print("Diff. between recursive & exact eq-s for derivatives:",
+ round(z.radial_dr(0.9) - z.radial_dr(0.9, use_exact_eq=True), 9))
+ # Compare performances
+ print("Tabular (10th order) / exact calc. times:", compare_performances(1, 10))
+ print("Recursive / exact calc. times for high orders:", compare_performances(12, 40))
+ # Statement below producing expected warnings, it's used for performance estimation
+ _estimate_high_order_calc_times()
diff --git a/tests/run_zernpsf_as_main.py b/tests/run_zernpsf_as_main.py
new file mode 100644
index 0000000..da8af1e
--- /dev/null
+++ b/tests/run_zernpsf_as_main.py
@@ -0,0 +1,197 @@
+# -*- coding: utf-8 -*-
+"""
+Run a few methods from 'zernpsf' as the main one for performing tests in IDE.
+
+@author: Sergei Klykov, '@sklykov' on GitHub
+
+@licence: MIT, @year: 2026
+
+"""
+import importlib
+from contextlib import suppress
+from pathlib import Path
+
+import matplotlib
+import numpy as np
+
+# Explicit backend assignment for matplotlib - for compatibility between running configurations in Spyder and PyCharm IDEs
+with suppress(ImportError):
+ matplotlib.use('Qt5Agg')
+import matplotlib.pyplot as plt
+
+# Import of developed version of package, install in the editable mode: pip install -e .
+import zernpy.zernikepol
+import zernpy.zernpsf
+
+importlib.reload(zernpy.zernikepol) # trick to guarantee local changes to be always used for run
+importlib.reload(zernpy.zernpsf)
+
+from zernpy.zernikepol import ZernPol
+from zernpy.zernpsf import ZernPSF, force_get_psf_compilation
+
+# %% Test as the main script
+if __name__ == "__main__":
+ plt.close("all") # close all opened before figures
+ wavelength_um = 0.55 # used below in a several calls
+ check_other_pols = False; check_small_na_wl = False # flag for checking some other polynomials PSFs
+ check_airy = False; check_common_psf = True; check_io_kernel = False; check_parallel_calculation = False; check_test = False
+ check_faster_airy = True; check_test_conditions = False; check_test_conditions2 = False; check_several_pols = False
+ check_edge_conditions = False; test_acceleration_single_pol = False; test_acceleration_few_pol = False
+ prepare_pic_readme = False # for plotting the sum of polynomials produced profile
+ test_io_few_pols = False; standard_path = Path.home().joinpath("Desktop") # for saving json on the Desktop
+ check_acceleration_flag = False # for testing fallback calculation with the wrong flag for calculate PSF kernel
+ check_init_several_pols = False # check calculation several pol-s: Airy (Piston) + some of polynomial
+ check_airy_patterns = False # check the difference between calculated by equation for Airy pattern and diffraction integral
+ check_precompilation = False # checks precompilation
+ check_cropping = False # checks how kernel is cropped
+
+ # Common PSF for testing
+ if check_common_psf:
+ force_get_psf_compilation(True)
+ zpsf1 = ZernPSF(ZernPol(m=1, n=1)); zpsf2 = ZernPSF(ZernPol(m=-3, n=3)); ampl = -0.43
+ zpsf2.set_physical_props(NA=0.95, wavelength=wavelength_um, expansion_coeff=ampl, pixel_physical_size=wavelength_um/5.05)
+ zpsf1.set_physical_props(NA=0.5, wavelength=0.6, expansion_coeff=0.25, pixel_physical_size=wavelength_um/5.5)
+ zpsf2.set_calculation_props(kernel_size=zpsf2.kernel_size, n_integration_points_r=250, n_integration_points_phi=320)
+ zpsf1.set_calculation_props(kernel_size=zpsf2.kernel_size, n_integration_points_r=220, n_integration_points_phi=360)
+ kernel3 = zpsf2.calculate_psf_kernel(suppress_warnings=False, verbose_info=True); zpsf2.plot_kernel("for loop")
+ zpsf2.visualize_convolution() # Visualize convolution on the disk image
+ if check_io_kernel:
+ # default_path = Path(__file__).parent.joinpath("saved_psfs").absolute() # default path for storing JSON files
+ zpsf2.save_json(overwrite=True, abs_path=standard_path)
+ zpsf1.read_json(zpsf2.json_file_path) # for testing reading and assigning values to ZernPSF class (substitution)
+ if check_parallel_calculation:
+ zpsf2.initialize_parallel_workers(); kernel2 = zpsf2.get_kernel_parallel()
+ zpsf2.plot_kernel("parallel"); zpsf2.deinitialize_workers()
+
+ # Another Zernike polynomial, big NA
+ if check_other_pols:
+ zpsf3 = ZernPSF(ZernPol(m=0, n=4))
+ zpsf3.set_physical_props(NA=1.25, wavelength=wavelength_um, expansion_coeff=0.4, pixel_physical_size=wavelength_um/5.25)
+ zpsf3.calculate_psf_kernel(suppress_warnings=False, verbose_info=True); zpsf3.plot_kernel()
+
+ # Another Zernike polynomial, average to small NA and wavelength
+ if check_small_na_wl:
+ NA = 0.45; wavelength = 0.4; pixel_size = wavelength / 5.25; ampl = 0.18
+ zp2 = ZernPol(m=1, n=3); zpsf2 = ZernPSF(zp2) # horizontal coma
+ zpsf2.set_physical_props(NA=NA, wavelength=wavelength, expansion_coeff=ampl, pixel_physical_size=pixel_size)
+ zpsf2.calculate_psf_kernel(normalized=True); zpsf2.plot_kernel()
+
+ if check_airy:
+ NA = 0.12; wavelength = 0.8; pixel_size = wavelength / 4.0; ampl = 1.25
+ zp4 = ZernPol(m=0, n=0); zpsf4 = ZernPSF(zp4) # piston for the Airy pattern
+ zpsf4.set_physical_props(NA=NA, wavelength=wavelength, expansion_coeff=ampl, pixel_physical_size=pixel_size)
+ zpsf4.calculate_psf_kernel(normalized=True); zpsf4.plot_kernel("Plus")
+ zpsf4.set_physical_props(NA=NA, wavelength=wavelength, expansion_coeff=-ampl, pixel_physical_size=pixel_size)
+ zpsf4.calculate_psf_kernel(normalized=True); zpsf4.plot_kernel("Minus")
+
+ if check_faster_airy: # For set the test for pytest library
+ NA = 0.35; wavelength = 0.55; pixel_size = wavelength / 3.05; ampl = -0.4
+ zp4 = ZernPol(m=0, n=0); zpsf4 = ZernPSF(zp4) # piston for the Airy pattern
+ zpsf4.set_physical_props(NA=NA, wavelength=wavelength, expansion_coeff=ampl, pixel_physical_size=pixel_size)
+ zpsf4.calculate_psf_kernel(normalized=True); zpsf4.plot_kernel()
+
+ if check_test_conditions:
+ NA = 0.95; wavelength = 0.55; pixel_size = wavelength / 5.0; ampl = 0.55
+ zp6 = ZernPol(m=0, n=2); zpsf6 = ZernPSF(zp6) # defocus
+ zpsf6.set_physical_props(NA=NA, wavelength=wavelength, expansion_coeff=ampl, pixel_physical_size=pixel_size) # normal assignment
+ zpsf6.set_calculation_props(kernel_size=zpsf6.kernel_size, n_integration_points_r=250, n_integration_points_phi=300)
+ zpsf6.calculate_psf_kernel(normalized=True); zpsf6.plot_kernel()
+ if check_test_conditions2:
+ zp7 = ZernPol(m=1, n=3); zpsf7 = ZernPSF(zp7) # horizontal coma
+ NA = 0.4; wavelength = 0.4; pixel_size = wavelength / 3.2; ampl = 0.185 # Common physical properties
+ zpsf7.set_physical_props(NA=NA, wavelength=wavelength, expansion_coeff=ampl, pixel_physical_size=pixel_size)
+ zpsf7.calculate_psf_kernel(normalized=True); zpsf7.plot_kernel()
+
+ # Test calculation of a PSF for several polynomials and I/O operations (see flags)
+ if check_several_pols:
+ zp1 = ZernPol(m=-2, n=2); zp2 = ZernPol(m=0, n=2); zp3 = ZernPol(m=2, n=2); pols = (zp1, zp2, zp3); coeffs = (-0.36, 0.25, 0.4)
+ zpsf8 = ZernPSF(pols); zpsf8.set_physical_props(NA=0.95, wavelength=0.5, expansion_coeff=coeffs, pixel_physical_size=0.5/4.5)
+ composed_kernel = zpsf8.calculate_psf_kernel(normalized=True, verbose_info=True); zpsf8.plot_kernel()
+ if test_io_few_pols:
+ zpsf14 = ZernPSF(ZernPol(osa=19)); zpsf14.set_physical_props(NA=0.1, wavelength=0.4, expansion_coeff=0.82,
+ pixel_physical_size=0.05)
+ zpsf8.save_json(overwrite=True, abs_path=standard_path); zpsf14.read_json(zpsf8.json_file_path) # save / read calculated kernel
+
+ # Test some edge conditions - e.g., specifying 1 polynomial in a list with huge coefficient
+ if check_edge_conditions:
+ zp4 = ZernPol(m=3, n=3); pols2 = [zp4]; coeff = 5.1
+ zpsf9 = ZernPSF(pols2); zpsf9.set_physical_props(NA=0.95, wavelength=0.5, expansion_coeff=coeff, pixel_physical_size=0.5/4.75)
+ zpsf9.set_physical_props(NA=0.95, wavelength=0.5, expansion_coeff=(-0.67), pixel_physical_size=0.5/4.75)
+ zpsf9.calculate_psf_kernel(normalized=True, verbose_info=True); zpsf9.plot_kernel()
+ # Test acceleration by using numba library utilities
+ if test_acceleration_single_pol:
+ force_get_psf_compilation(); NA = 0.95; wavelength = 0.55; pixel_size = wavelength / 4.6; ampl = -0.16
+ zp6 = ZernPol(m=0, n=2); zpsf6 = ZernPSF(zp6) # defocus
+ zpsf6.set_physical_props(NA=NA, wavelength=wavelength, expansion_coeff=ampl, pixel_physical_size=pixel_size)
+ kernel_acc = zpsf6.calculate_psf_kernel(normalized=True, accelerated=True, verbose_info=True); zpsf6.plot_kernel("Accelerated")
+ kernel_norm = zpsf6.calculate_psf_kernel(normalized=True); zpsf6.plot_kernel("Normal")
+ if test_acceleration_few_pol:
+ force_get_psf_compilation()
+ zp1 = ZernPol(m=-2, n=2); zp2 = ZernPol(m=0, n=2); zp3 = ZernPol(m=2, n=2); pols = (zp1, zp2, zp3); coeffs = (-0.12, 0.15, 0.1)
+ zpsf8 = ZernPSF(pols); zpsf8.set_physical_props(NA=0.35, wavelength=0.5, expansion_coeff=coeffs, pixel_physical_size=0.5/1.5)
+ composed_kernel = zpsf8.calculate_psf_kernel(normalized=True, verbose_info=True); zpsf8.plot_kernel("Normal")
+ composed_kernel_acc = zpsf8.calculate_psf_kernel(normalized=True, verbose_info=True, accelerated=True)
+ zpsf8.plot_kernel("Accelerated")
+ if prepare_pic_readme:
+ force_get_psf_compilation(verbose_report=True)
+ zp1 = ZernPol(m=-1, n=3); zp2 = ZernPol(m=2, n=4); zp3 = ZernPol(m=0, n=4); pols = (zp1, zp2, zp3); coeffs = (0.5, 0.21, 0.15)
+ zpsf_pic = ZernPSF(pols); zpsf_pic.set_physical_props(NA=0.65, wavelength=0.6, expansion_coeff=coeffs, pixel_physical_size=0.6/5.0)
+ zpsf_pic.calculate_psf_kernel(normalized=True, verbose_info=True, accelerated=True)
+ zpsf_pic.plot_kernel("Vert. Coma Vert. 2nd Astigmatism Spherical")
+ if check_acceleration_flag:
+ zp16 = ZernPol(m=-1, n=3); zp18 = ZernPol(m=2, n=4); pols10 = (zp16, zp18); coeffs10 = (-0.1, 0.13); zpsf_acc = ZernPSF(pols10)
+ zpsf_norm = ZernPSF(pols10); zpsf_acc.set_physical_props(NA=0.43, wavelength=0.6, expansion_coeff=coeffs10,
+ pixel_physical_size=0.6/3.0)
+ zpsf_norm.set_physical_props(NA=0.43, wavelength=0.6, expansion_coeff=coeffs10, pixel_physical_size=0.6/3.0)
+ kern_acc = zpsf_acc.calculate_psf_kernel(normalized=True, accelerated=True)
+ kern_norm = zpsf_norm.calculate_psf_kernel(normalized=True)
+ kern_diff = np.round(kern_acc - kern_norm, 9) # for checking the difference in calculations
+
+ if check_test:
+ pols = (ZernPol(osa=10), ZernPol(osa=15)); coeffs = (0.28, -0.33); NA = 0.35; wavelength = 0.55
+ zpsf = ZernPSF(pols); zpsf.set_physical_props(NA, wavelength, expansion_coeff=coeffs, pixel_physical_size=wavelength / 3.5)
+ zpsf.set_calculation_props(kernel_size=25, n_integration_points_r=200, n_integration_points_phi=180)
+ psf_kernel = zpsf.calculate_psf_kernel(normalized=False); zpsf.plot_kernel()
+
+ if check_init_several_pols:
+ force_get_psf_compilation(True)
+ zpsf30 = ZernPSF(zernpol=(ZernPol(osa=1), ZernPol(osa=5)))
+ try:
+ zpsf30.set_physical_props(NA=0.95, wavelength=0.5, expansion_coeff=0.5, pixel_physical_size=0.5/5.0)
+ except ValueError:
+ print("Check for 1 ampl. and 2 pol-s passed") # as expected, transfer to test function
+ zpsf31 = ZernPSF(zernpol=(ZernPol(osa=0), ZernPol(m=-1, n=3)))
+ zpsf31.set_physical_props(NA=0.95, wavelength=0.5, expansion_coeff=[1.5, 0.24], pixel_physical_size=0.5/3.8)
+ # zpsf31.calculate_psf_kernel(verbose_info=True, accelerated=False); zpsf31.plot_kernel("Not accelerated 2 pol-s")
+ # kernel_not_acc = np.copy(zpsf31.kernel)
+ zpsf31.calculate_psf_kernel(verbose_info=True, accelerated=True); zpsf31.plot_kernel("Airy 1.5")
+ kernel_acc = np.copy(zpsf31.kernel)
+ # zpsf31.kernel = np.abs(kernel_acc - kernel_not_acc); zpsf31.plot_kernel("Diff. 2 pol-s")
+ zpsf31.set_physical_props(NA=0.95, wavelength=0.5, expansion_coeff=[-1.5, 0.24], pixel_physical_size=0.5/3.8)
+ zpsf31.calculate_psf_kernel(verbose_info=True, accelerated=True); zpsf31.plot_kernel("Airy -1.5")
+ kernel_neg_acc = np.copy(zpsf31.kernel)
+ zpsf31.kernel = np.abs(kernel_acc - kernel_neg_acc); zpsf31.plot_kernel("Diff. neg. pos. Airy + Coma")
+
+ if check_airy_patterns:
+ force_get_psf_compilation(verbose_report=True)
+ zpsf40 = ZernPSF(zernpol=ZernPol(osa=0))
+ zpsf40.set_physical_props(NA=0.95, wavelength=0.5, expansion_coeff=3.0, pixel_physical_size=0.5/4.0)
+ zpsf40.calculate_psf_kernel(verbose_info=True, accelerated=False, normalized=False); zpsf40.plot_kernel("Not accelerated Airy")
+ kernel_not_acc = np.copy(zpsf40.kernel)
+ zpsf40.calculate_psf_kernel(verbose_info=True, accelerated=True, normalized=False); zpsf40.plot_kernel("Accelerated Airy")
+ kernel_acc = np.copy(zpsf40.kernel)
+ zpsf40.kernel = np.abs(kernel_acc - kernel_not_acc); zpsf40.plot_kernel("Diff. Airy")
+
+ if check_precompilation:
+ force_get_psf_compilation(True)
+ zpsf50 = ZernPSF(zernpol=(ZernPol(m=1, n=3)))
+ zpsf50.set_physical_props(NA=1.25, wavelength=0.52, expansion_coeff=0.5, pixel_physical_size=0.5/4.85)
+ zpsf50.calculate_psf_kernel(verbose_info=True, accelerated=True, normalized=True); zpsf50.plot_kernel()
+
+ if check_cropping:
+ zpsf60 = ZernPSF(zernpol=(ZernPol(m=0, n=4)))
+ zpsf60.set_physical_props(NA=1.25, wavelength=0.5, expansion_coeff=0.47, pixel_physical_size=0.5/5.0)
+ zpsf60.calculate_psf_kernel(accelerated=True, verbose_info=True); zpsf60.plot_kernel("Not Cropped")
+ original_kernel = np.copy(zpsf60.kernel)
+ zpsf60.crop_kernel(min_part_of_max=0.025); zpsf60.plot_kernel("Cropped"); cropped_kernel = np.copy(zpsf60.kernel)
+ print("Original kernel shape:", original_kernel.shape, "\nCropped kernel shape:", cropped_kernel.shape)
diff --git a/tests/test_package_import.py b/tests/test_package_import.py
index 177974f..a307404 100644
--- a/tests/test_package_import.py
+++ b/tests/test_package_import.py
@@ -9,7 +9,7 @@
def test_initialization():
try:
- from zernpy import ZernPol, generate_polynomials, ZernPSF, force_get_psf_compilation
+ from zernpy import ZernPol, ZernPSF, force_get_psf_compilation, generate_polynomials
try:
import numba
if numba is not None: