Skip to content

Commit

Permalink
Docstring cleanup (#70)
Browse files Browse the repository at this point in the history 
* Docstring cleanup
  • Loading branch information
@bwohlberg
bwohlberg authored Oct 29, 2021
1 parent d2b1b5e commit b1443ea
Show file tree
Hide file tree
Showing 3 changed files with 132 additions and 88 deletions.
162 changes: 98 additions & 64 deletions scico/linop/optics.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -7,9 +7,9 @@

"""Optical propagator classes.
This module provides classes that model the propagation of a monochromatic waveform between two
parallel planes in a homogeneous medium.
This module provides classes that model the propagation of a
monochromatic waveform between two parallel planes in a homogeneous
medium.
"""


Expand Down Expand Up @@ -37,18 +37,21 @@
def radial_transverse_frequency(
input_shape: Shape, dx: Union[float, Tuple[float, ...]]
) -> np.ndarray:
"""Construct radial Fourier coordinate system
"""Construct radial Fourier coordinate system.
Args:
input_shape: Tuple of length 1 or 2 containing the number of samples per dimension
dx: Sampling interval at source plane. If a float and `len(input_shape)==2`
the same sampling interval is applied to both dimensions. If `dx` is a
tuple, must have same length as `input_shape`.
input_shape: Tuple of length 1 or 2 containing the number of
samples per dimension
dx: Sampling interval at source plane. If a float and
`len(input_shape)==2` the same sampling interval is applied
to both dimensions. If `dx` is a tuple, must have same
length as `input_shape`.
Returns:
If `len(input_shape)==1` returns an ndarray containing corresponding
Fourier coordinates. If `len(input_shape) == 2`, returns an ndarray
containing the radial Fourier coordinates; sqrt(kx^2 + ky^2)
If `len(input_shape)==1` returns an ndarray containing
corresponding Fourier coordinates. If `len(input_shape) == 2`,
returns an ndarray containing the radial Fourier coordinates;
sqrt(kx^2 + ky^2).
"""

ndim = len(input_shape) # 1 or 2 dimensions
Expand All @@ -60,7 +63,8 @@ def radial_transverse_frequency(
else:
if len(dx) != ndim:
raise ValueError(
f"dx must be a scalar or have len(dx) == len(input_shape); got len(dx)={len(dx)}, len(input_shape)={ndim}"
"dx must be a scalar or have len(dx) == len(input_shape); "
f"got len(dx)={len(dx)}, len(input_shape)={ndim}"
)

if ndim == 1:
Expand All @@ -86,10 +90,12 @@ def __init__(
r"""Base class for AngularSpectrum and Fresnel propagator
Args:
input_shape: Shape of input array. Can be a tuple of length 1 or 2.
dx: Sampling interval at source plane. If a float and `len(input_shape)==2`
the same sampling interval is applied to both dimensions. If `dx` is a
tuple, must have same length as `input_shape`.
input_shape: Shape of input array. Can be a tuple of length
1 or 2.
dx: Sampling interval at source plane. If a float and
`len(input_shape)==2` the same sampling interval is
applied to both dimensions. If `dx` is a tuple, must
have same length as `input_shape`.
k0 : Illumination wavenumber; 2π/wavelength
z : Propagation distance
pad_factor: Amount of padding to apply in DFT step
Expand All @@ -104,7 +110,8 @@ def __init__(
else:
if len(dx) != ndim:
raise ValueError(
f"dx must be a scalar or have len(dx) == len(input_shape); got len(dx)={len(dx)}, len(input_shape)={ndim}"
"dx must be a scalar or have len(dx) == len(input_shape); "
f"got len(dx)={len(dx)}, len(input_shape)={ndim}"
)

#: Illumination wavenumber; 2π/wavelength
Expand Down Expand Up @@ -158,25 +165,32 @@ class AngularSpectrumPropagator(Propagator):
Notes:
Propagates a source field with coordinates :math:`(x, y, z_0)` to a destination plane at a distance :math:`z`
with coordinates :math:`(x, y, z_0 + z)`.
Propagates a source field with coordinates :math:`(x, y, z_0)`
to a destination plane at a distance :math:`z` with coordinates
:math:`(x, y, z_0 + z)`.
The action of this linear operator is given by (Eq 3.74, :cite:`goodman-2005-fourier`):
The action of this linear operator is given by
(Eq 3.74, :cite:`goodman-2005-fourier`):
.. math ::
(A \mb{u})(x, y) = \iint_{-\infty}^{\infty} \mb{\hat{u}}(k_x, k_y) e^{j \sqrt{k_0^2 - k_x^2 - k_y^2} \abs{z}} e^{j (x k_x + y k_y) } d k_x \ d k_y
(A \mb{u})(x, y) = \iint_{-\infty}^{\infty}
\mb{\hat{u}}(k_x, k_y) e^{j \sqrt{k_0^2 - k_x^2 - k_y^2}
\abs{z}} e^{j (x k_x + y k_y) } d k_x \ d k_y
where the :math:`\mb{\hat{u}}` is the Fourier transform of the field in the plane :math:`z=0`, given by
where the :math:`\mb{\hat{u}}` is the Fourier transform of the
field in the plane :math:`z=0`, given by
.. math ::
\mb{\hat{u}}(k_x, k_y) = \iint_{-\infty}^{\infty} \mb{u}(x, y) e^{- j (x k_x + y k_y)} d k_x \ d k_y.
\mb{\hat{u}}(k_x, k_y) = \iint_{-\infty}^{\infty}
\mb{u}(x, y) e^{- j (x k_x + y k_y)} d k_x \ d k_y.
The angular spectrum propagator can be written
.. math ::
A\mb{u} = F^{-1} D F \mb{u}
where :math:`F` is the Fourier transform with respect to :math:`(x, y)` and
where :math:`F` is the Fourier transform with respect to
:math:`(x, y)` and
.. math ::
D = \mathrm{diag}\left(\exp \left\{ j
Expand All @@ -188,7 +202,8 @@ class AngularSpectrumPropagator(Propagator):
The propagator is adequately sampled when :cite:`voelz-2009-digital`
.. math ::
(\Delta x)^2 \geq \frac{\pi}{k_0 N} \sqrt{ (\Delta x)^2 N^2 + 4 z^2}
(\Delta x)^2 \geq \frac{\pi}{k_0 N} \sqrt{ (\Delta x)^2 N^2
+ 4 z^2}
"""

Expand All @@ -205,14 +220,15 @@ def __init__(
r"""Angular Spectrum init method.
Args:
input_shape: Shape of input array. Can be a tuple of length 2 or 3.
input_shape: Shape of input array. Can be a tuple of length
2 or 3.
dx : Spatial sampling rate.
k0 : Illumination wavenumber
z : Propagation distance
pad_factor: Amount of padding to apply in DFT step
jit: If ``True``, call :meth:`.jit()` on this LinearOperator to jit the forward,
adjoint, and gram functions. Same as calling :meth:`.jit` after
the LinearOperator is created.
k0 : Illumination wavenumber.
z : Propagation distance.
pad_factor: Amount of padding to apply in DFT step.
jit: If ``True``, call :meth:`.jit()` on this LinearOperator
to jit the forward, adjoint, and gram functions. Same as
calling :meth:`.jit` after the LinearOperator is created.
"""

Expand Down Expand Up @@ -241,7 +257,8 @@ def check_sampling(self):
Returns:
True if the angular spectrum kernel is adequately sampled, False otherwise.
True if the angular spectrum kernel is adequately sampled,
False otherwise.
"""
tmp = []
for d, N in zip(self.dx, self.padded_shape):
Expand All @@ -260,30 +277,35 @@ def pinv(self, y):
class FresnelPropagator(Propagator):
r"""Fresnel Propagator
Propagates a source field with coordinates :math:`(x, y, z_0)` to a destination
plane at a distance :math:`z` with coordinates :math:`(x, y, z_0 + z)`. The
action of this linear operator is given by (Eq 4.20, :cite:`goodman-2005-fourier`):
Propagates a source field with coordinates :math:`(x, y, z_0)` to a
destination plane at a distance :math:`z` with coordinates
:math:`(x, y, z_0 + z)`. The action of this linear operator is given
by (Eq 4.20, :cite:`goodman-2005-fourier`):
.. math ::
(A \mb{u})(x, y) = e^{j k_0 z} \iint_{-\infty}^{\infty} \mb{\hat{u}}(k_x, k_y) e^{-j \frac{z}{2 k_0} (k_x^2 + k_y^2) }
(A \mb{u})(x, y) = e^{j k_0 z} \iint_{-\infty}^{\infty}
\mb{\hat{u}}(k_x, k_y) e^{-j \frac{z}{2 k_0} (k_x^2 + k_y^2) }
e^{j (x k_x + y k_y) } d k_x \ d k_y,
where the :math:`\mb{\hat{u}}` is the Fourier transform of the field in the source plane, given by
where the :math:`\mb{\hat{u}}` is the Fourier transform of the field
in the source plane, given by
.. math ::
\mb{\hat{u}}(k_x, k_y) = \iint_{-\infty}^{\infty} \mb{u}(x, y) e^{- j (x k_x + y k_y)} d k_x \ d k_y.
\mb{\hat{u}}(k_x, k_y) = \iint_{-\infty}^{\infty} \mb{u}(x, y)
e^{- j (x k_x + y k_y)} d k_x \ d k_y.
The Fresnel propagator can be written
.. math ::
A\mb{u} = F^{-1} D F \mb{u}
where :math:`F` is the Fourier transform with respect to :math:`(x, y)` and
where :math:`F` is the Fourier transform with respect to
:math:`(x, y)` and
.. math ::
D = \mathrm{diag}\left(\exp \left\{ j
\sqrt{k_0^2 - k_x^2 - k_y^2} \abs{z}
\right\} \right)
\right\} \right) \;.
"""

Expand Down Expand Up @@ -321,7 +343,8 @@ def check_sampling(self):
Returns:
True if the Fresnel propagation kernel is adequately sampled, False otherwise.
True if the Fresnel propagation kernel is adequately sampled,
False otherwise.
"""
tmp = []
for d, N in zip(self.dx, self.padded_shape):
Expand All @@ -337,15 +360,19 @@ class FraunhoferPropagator(LinearOperator):
Notes:
Propagates a source field with coordinates :math:`(x_S, y_S)` to a destination plane at a distance :math:`z`
with coordinates :math:`(x_D, y_D)`.
Propagates a source field with coordinates :math:`(x_S, y_S)` to
a destination plane at a distance :math:`z` with coordinates
:math:`(x_D, y_D)`.
The action of this linear operator is given by (Eq 4.25, :cite:`goodman-2005-fourier`):
The action of this linear operator is given by
(Eq 4.25, :cite:`goodman-2005-fourier`):
.. math ::
(A \mb{u})(x_D, y_D) = \underbrace{\frac{k_0}{2 \pi} \frac{e^{j k_0 z}}{j z} \mathrm{exp} \left\{ j \frac{k_0}{2 z}
(A \mb{u})(x_D, y_D) = \underbrace{\frac{k_0}{2 \pi}
\frac{e^{j k_0 z}}{j z} \mathrm{exp} \left\{ j \frac{k_0}{2 z}
(x_D^2 + y_D^2) \right\}}_{\triangleq P(x_D, y_D)}
\int \mb{u}(x_S, y_S) e^{-j \frac{k_0}{z} (x_D x_S + y_D y_S) } dx_S \ dy_S.
\int \mb{u}(x_S, y_S) e^{-j \frac{k_0}{z} (x_D x_S + y_D y_S)
} dx_S \ dy_S.
Writing the Fourier transform of the field :math:`\mb{u}` as
Expand All @@ -355,22 +382,25 @@ class FraunhoferPropagator(LinearOperator):
the action of this linear operator can be written
.. math ::
(A \mb{u})(x_D, y_D) = P(x_D, y_D) \ \hat{\mb{u}}\left({\frac{k_0}{z} x_D, \frac{k_0}{z} y_D}\right).
Ignoring multiplicative prefactors, the Fraunhofer propagated field is the Fourier transform of the source field,
evaluated at coordinates :math:`(k_x, k_y) = (\frac{k_0}{z} x_D, \frac{k_0}{z} y_D)`.
(A \mb{u})(x_D, y_D) = P(x_D, y_D) \ \hat{\mb{u}}
\left({\frac{k_0}{z} x_D, \frac{k_0}{z} y_D}\right).
Ignoring multiplicative prefactors, the Fraunhofer propagated
field is the Fourier transform of the source field, evaluated at
coordinates :math:`(k_x, k_y) = (\frac{k_0}{z} x_D, \frac{k_0}{z} y_D)`.
In general, the sampling intervals (and thus plane lengths) differ between source and destination planes.
In particular, :cite:`voelz-2011-computational`
In general, the sampling intervals (and thus plane lengths)
differ between source and destination planes. In particular,
:cite:`voelz-2011-computational`
.. math ::
\begin{aligned}
\Delta x_D &= \frac{2 \pi z}{k_0 L_S } \\
L_D &= \frac{2 \pi z}{k_0 \Delta x_S }
\end{aligned}
The sampling intervals and plane lengths coincide in the case of critical sampling:
The sampling intervals and plane lengths coincide in the case of
critical sampling:
.. math ::
\Delta x_S = \sqrt{\frac{2 \pi z}{N k_0}}
Expand All @@ -392,15 +422,16 @@ def __init__(
):
"""
Args:
input_shape: Shape of input plane. Must be length 1 or 2. The output plane will
have the same number of samples.
dx: Sampling interval at source plane. If a float and `len(input_shape)==2`
the same sampling interval is applied to both dimensions. If `dx` is a
tuple, must have same length as `input_shape`.
input_shape: Shape of input plane. Must be length 1 or 2.
The output plane will have the same number of samples.
dx: Sampling interval at source plane. If a float and
`len(input_shape)==2` the same sampling interval is
applied to both dimensions. If `dx` is a tuple, must have
same length as `input_shape`.
k0 : Illumination wavenumber; 2π/wavelength
z : Propagation distance
jit: If ``True``, jit the evaluation, adjoint, and gram functions of this LinearOperator.
Default: True.
jit: If ``True``, jit the evaluation, adjoint, and gram
functions of this LinearOperator. Default: True.
"""

ndim = len(input_shape) # 1 or 2 dimensions
Expand All @@ -412,7 +443,8 @@ def __init__(
else:
if len(dx) != ndim:
raise ValueError(
f"dx must be a scalar or have len(dx) == len(input_shape); got len(dx)={len(dx)}, len(input_shape)={ndim}"
"dx must be a scalar or have len(dx) == len(input_shape); "
f"got len(dx)={len(dx)}, len(input_shape)={ndim}"
)

L: Tuple[float, ...] = tuple(s * d for s, d in zip(input_shape, dx))
Expand Down Expand Up @@ -477,13 +509,15 @@ def _eval(self, x):
def check_sampling(self):
r"""Verify the Fraunhofer propagation kernel is not aliased.
Checks the condition for adequate sampling, :cite:`voelz-2011-computational`
Checks the condition for adequate sampling,
:cite:`voelz-2011-computational`
.. math ::
\Delta x^2 \geq \frac{2 \pi z }{k_0 N}
Returns:
True if the Fresnel propagation kernel is adequately sampled, False otherwise.
True if the Fresnel propagation kernel is adequately sampled,
False otherwise.
"""
tmp = []
for d, N in zip(self.dx, self.padded_shape):
Expand Down
Loading

0 comments on commit b1443ea

Please sign in to comment.