jan/meanas
1
0
Fork 0
Code Issues Pull Requests Releases Wiki Activity

type annotation fixup

Browse Source
This commit is contained in:
Jan Petykiewicz 2020-06-11 19:26:17 -07:00
parent 2fceff9930
commit b7d4bbbd95
15 changed files with 149 additions and 122 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

10
meanas/eigensolvers.py
View File

@ -1,7 +1,7 @@
""" """
Solvers for eigenvalue / eigenvector problems Solvers for eigenvalue / eigenvector problems
""" """
from typing import Tuple, List from typing import Tuple, Callable, Optional, Union
import numpy import numpy
from numpy.linalg import norm from numpy.linalg import norm
from scipy import sparse from scipy import sparse
@ -9,7 +9,7 @@ import scipy.sparse.linalg as spalg
def power_iteration(operator: sparse.spmatrix, def power_iteration(operator: sparse.spmatrix,
guess_vector: numpy.ndarray = None, guess_vector: Optional[numpy.ndarray] = None,
iterations: int = 20, iterations: int = 20,
) -> Tuple[complex, numpy.ndarray]: ) -> Tuple[complex, numpy.ndarray]:
""" """
@ -36,11 +36,11 @@ def power_iteration(operator: sparse.spmatrix,
return lm_eigval, v return lm_eigval, v
def rayleigh_quotient_iteration(operator: sparse.spmatrix or spalg.LinearOperator, def rayleigh_quotient_iteration(operator: Union[sparse.spmatrix, spalg.LinearOperator],
guess_vector: numpy.ndarray, guess_vector: numpy.ndarray,
iterations: int = 40, iterations: int = 40,
tolerance: float = 1e-13, tolerance: float = 1e-13,
solver = None, solver: Optional[Callable[..., numpy.ndarray]] = None,
) -> Tuple[complex, numpy.ndarray]: ) -> Tuple[complex, numpy.ndarray]:
""" """
Use Rayleigh quotient iteration to refine an eigenvector guess. Use Rayleigh quotient iteration to refine an eigenvector guess.
@ -83,7 +83,7 @@ def rayleigh_quotient_iteration(operator: sparse.spmatrix or spalg.LinearOperato
return eigval, v return eigval, v
def signed_eigensolve(operator: sparse.spmatrix or spalg.LinearOperator, def signed_eigensolve(operator: Union[sparse.spmatrix, spalg.LinearOperator],
how_many: int, how_many: int,
negative: bool = False, negative: bool = False,
) -> Tuple[numpy.ndarray, numpy.ndarray]: ) -> Tuple[numpy.ndarray, numpy.ndarray]:

30
meanas/fdfd/operators.py
View File

@ -27,7 +27,7 @@ The following operators are included:
- Cross product matrices - Cross product matrices
""" """
from typing import List, Tuple from typing import Tuple, Optional
import numpy import numpy
import scipy.sparse as sparse import scipy.sparse as sparse
@ -41,9 +41,9 @@ __author__ = 'Jan Petykiewicz'
def e_full(omega: complex, def e_full(omega: complex,
dxes: dx_lists_t, dxes: dx_lists_t,
epsilon: vfdfield_t, epsilon: vfdfield_t,
mu: vfdfield_t = None, mu: Optional[vfdfield_t] = None,
pec: vfdfield_t = None, pec: Optional[vfdfield_t] = None,
pmc: vfdfield_t = None, pmc: Optional[vfdfield_t] = None,
) -> sparse.spmatrix: ) -> sparse.spmatrix:
""" """
Wave operator Wave operator
@ -125,9 +125,9 @@ def e_full_preconditioners(dxes: dx_lists_t
def h_full(omega: complex, def h_full(omega: complex,
dxes: dx_lists_t, dxes: dx_lists_t,
epsilon: vfdfield_t, epsilon: vfdfield_t,
mu: vfdfield_t = None, mu: Optional[vfdfield_t] = None,
pec: vfdfield_t = None, pec: Optional[vfdfield_t] = None,
pmc: vfdfield_t = None, pmc: Optional[vfdfield_t] = None,
) -> sparse.spmatrix: ) -> sparse.spmatrix:
""" """
Wave operator Wave operator
@ -181,9 +181,9 @@ def h_full(omega: complex,
def eh_full(omega: complex, def eh_full(omega: complex,
dxes: dx_lists_t, dxes: dx_lists_t,
epsilon: vfdfield_t, epsilon: vfdfield_t,
mu: vfdfield_t = None, mu: Optional[vfdfield_t] = None,
pec: vfdfield_t = None, pec: Optional[vfdfield_t] = None,
pmc: vfdfield_t = None pmc: Optional[vfdfield_t] = None
) -> sparse.spmatrix: ) -> sparse.spmatrix:
""" """
Wave operator for `[E, H]` field representation. This operator implements Maxwell's Wave operator for `[E, H]` field representation. This operator implements Maxwell's
@ -249,8 +249,8 @@ def eh_full(omega: complex,
def e2h(omega: complex, def e2h(omega: complex,
dxes: dx_lists_t, dxes: dx_lists_t,
mu: vfdfield_t = None, mu: Optional[vfdfield_t] = None,
pmc: vfdfield_t = None, pmc: Optional[vfdfield_t] = None,
) -> sparse.spmatrix: ) -> sparse.spmatrix:
""" """
Utility operator for converting the E field into the H field. Utility operator for converting the E field into the H field.
@ -280,7 +280,7 @@ def e2h(omega: complex,
def m2j(omega: complex, def m2j(omega: complex,
dxes: dx_lists_t, dxes: dx_lists_t,
mu: vfdfield_t = None mu: Optional[vfdfield_t] = None
) -> sparse.spmatrix: ) -> sparse.spmatrix:
""" """
Operator for converting a magnetic current M into an electric current J. Operator for converting a magnetic current M into an electric current J.
@ -362,7 +362,7 @@ def e_tfsf_source(TF_region: vfdfield_t,
omega: complex, omega: complex,
dxes: dx_lists_t, dxes: dx_lists_t,
epsilon: vfdfield_t, epsilon: vfdfield_t,
mu: vfdfield_t = None, mu: Optional[vfdfield_t] = None,
) -> sparse.spmatrix: ) -> sparse.spmatrix:
""" """
Operator that turns a desired E-field distribution into a Operator that turns a desired E-field distribution into a
@ -392,7 +392,7 @@ def e_boundary_source(mask: vfdfield_t,
omega: complex, omega: complex,
dxes: dx_lists_t, dxes: dx_lists_t,
epsilon: vfdfield_t, epsilon: vfdfield_t,
mu: vfdfield_t = None, mu: Optional[vfdfield_t] = None,
periodic_mask_edges: bool = False, periodic_mask_edges: bool = False,
) -> sparse.spmatrix: ) -> sparse.spmatrix:
""" """

20
meanas/fdfd/scpml.py
View File

@ -2,10 +2,10 @@
Functions for creating stretched coordinate perfectly matched layer (PML) absorbers. Functions for creating stretched coordinate perfectly matched layer (PML) absorbers.
""" """
from typing import List, Callable from typing import Sequence, Union, Callable, Optional
import numpy import numpy
from ..fdmath import dx_lists_t from ..fdmath import dx_lists_t, dx_lists_mut
__author__ = 'Jan Petykiewicz' __author__ = 'Jan Petykiewicz'
@ -37,12 +37,12 @@ def prepare_s_function(ln_R: float = -16,
return s_factor return s_factor
def uniform_grid_scpml(shape: numpy.ndarray or List[int], def uniform_grid_scpml(shape: Union[numpy.ndarray, Sequence[int]],
thicknesses: numpy.ndarray or List[int], thicknesses: Union[numpy.ndarray, Sequence[int]],
omega: float, omega: float,
epsilon_effective: float = 1.0, epsilon_effective: float = 1.0,
s_function: s_function_t = None, s_function: Optional[s_function_t] = None,
) -> dx_lists_t: ) -> dx_lists_mut:
""" """
Create dx arrays for a uniform grid with a cell width of 1 and a pml. Create dx arrays for a uniform grid with a cell width of 1 and a pml.
@ -63,7 +63,7 @@ def uniform_grid_scpml(shape: numpy.ndarray or List[int],
Default uses `prepare_s_function()` with no parameters. Default uses `prepare_s_function()` with no parameters.
Returns: Returns:
Complex cell widths (dx_lists_t) as discussed in `meanas.fdmath.types`. Complex cell widths (dx_lists_mut) as discussed in `meanas.fdmath.types`.
""" """
if s_function is None: if s_function is None:
s_function = prepare_s_function() s_function = prepare_s_function()
@ -90,13 +90,13 @@ def uniform_grid_scpml(shape: numpy.ndarray or List[int],
return [dx_a, dx_b] return [dx_a, dx_b]
def stretch_with_scpml(dxes: dx_lists_t, def stretch_with_scpml(dxes: dx_lists_mut,
axis: int, axis: int,
polarity: int, polarity: int,
omega: float, omega: float,
epsilon_effective: float = 1.0, epsilon_effective: float = 1.0,
thickness: int = 10, thickness: int = 10,
s_function: s_function_t = None, s_function: Optional[s_function_t] = None,
) -> dx_lists_t: ) -> dx_lists_t:
""" """
Stretch dxes to contain a stretched-coordinate PML (SCPML) in one direction along one axis. Stretch dxes to contain a stretched-coordinate PML (SCPML) in one direction along one axis.
@ -113,7 +113,7 @@ def stretch_with_scpml(dxes: dx_lists_t,
of pml parameters. Default uses `prepare_s_function()` with no parameters. of pml parameters. Default uses `prepare_s_function()` with no parameters.
Returns: Returns:
Complex cell widths (dx_lists_t) as discussed in `meanas.fdmath.types`. Complex cell widths (dx_lists_mut) as discussed in `meanas.fdmath.types`.
Multiple calls to this function may be necessary if multiple absorpbing boundaries are needed. Multiple calls to this function may be necessary if multiple absorpbing boundaries are needed.
""" """
if s_function is None: if s_function is None:

28
meanas/fdfd/waveguide_2d.py
View File

@ -146,7 +146,7 @@ to account for numerical dispersion if the result is introduced into a space wit
""" """
# TODO update module docs # TODO update module docs
from typing import List, Tuple from typing import List, Tuple, Optional
import numpy import numpy
from numpy.linalg import norm from numpy.linalg import norm
import scipy.sparse as sparse import scipy.sparse as sparse
@ -163,7 +163,7 @@ __author__ = 'Jan Petykiewicz'
def operator_e(omega: complex, def operator_e(omega: complex,
dxes: dx_lists_t, dxes: dx_lists_t,
epsilon: vfdfield_t, epsilon: vfdfield_t,
mu: vfdfield_t = None, mu: Optional[vfdfield_t] = None,
) -> sparse.spmatrix: ) -> sparse.spmatrix:
""" """
Waveguide operator of the form Waveguide operator of the form
@ -229,7 +229,7 @@ def operator_e(omega: complex,
def operator_h(omega: complex, def operator_h(omega: complex,
dxes: dx_lists_t, dxes: dx_lists_t,
epsilon: vfdfield_t, epsilon: vfdfield_t,
mu: vfdfield_t = None, mu: Optional[vfdfield_t] = None,
) -> sparse.spmatrix: ) -> sparse.spmatrix:
""" """
Waveguide operator of the form Waveguide operator of the form
@ -298,7 +298,7 @@ def normalized_fields_e(e_xy: numpy.ndarray,
omega: complex, omega: complex,
dxes: dx_lists_t, dxes: dx_lists_t,
epsilon: vfdfield_t, epsilon: vfdfield_t,
mu: vfdfield_t = None, mu: Optional[vfdfield_t] = None,
prop_phase: float = 0, prop_phase: float = 0,
) -> Tuple[vfdfield_t, vfdfield_t]: ) -> Tuple[vfdfield_t, vfdfield_t]:
""" """
@ -332,7 +332,7 @@ def normalized_fields_h(h_xy: numpy.ndarray,
omega: complex, omega: complex,
dxes: dx_lists_t, dxes: dx_lists_t,
epsilon: vfdfield_t, epsilon: vfdfield_t,
mu: vfdfield_t = None, mu: Optional[vfdfield_t] = None,
prop_phase: float = 0, prop_phase: float = 0,
) -> Tuple[vfdfield_t, vfdfield_t]: ) -> Tuple[vfdfield_t, vfdfield_t]:
""" """
@ -366,7 +366,7 @@ def _normalized_fields(e: numpy.ndarray,
omega: complex, omega: complex,
dxes: dx_lists_t, dxes: dx_lists_t,
epsilon: vfdfield_t, epsilon: vfdfield_t,
mu: vfdfield_t = None, mu: Optional[vfdfield_t] = None,
prop_phase: float = 0, prop_phase: float = 0,
) -> Tuple[vfdfield_t, vfdfield_t]: ) -> Tuple[vfdfield_t, vfdfield_t]:
# TODO documentation # TODO documentation
@ -405,7 +405,7 @@ def exy2h(wavenumber: complex,
omega: complex, omega: complex,
dxes: dx_lists_t, dxes: dx_lists_t,
epsilon: vfdfield_t, epsilon: vfdfield_t,
mu: vfdfield_t = None mu: Optional[vfdfield_t] = None
) -> sparse.spmatrix: ) -> sparse.spmatrix:
""" """
Operator which transforms the vector `e_xy` containing the vectorized E_x and E_y fields, Operator which transforms the vector `e_xy` containing the vectorized E_x and E_y fields,
@ -430,7 +430,7 @@ def hxy2e(wavenumber: complex,
omega: complex, omega: complex,
dxes: dx_lists_t, dxes: dx_lists_t,
epsilon: vfdfield_t, epsilon: vfdfield_t,
mu: vfdfield_t = None mu: Optional[vfdfield_t] = None
) -> sparse.spmatrix: ) -> sparse.spmatrix:
""" """
Operator which transforms the vector `h_xy` containing the vectorized H_x and H_y fields, Operator which transforms the vector `h_xy` containing the vectorized H_x and H_y fields,
@ -453,7 +453,7 @@ def hxy2e(wavenumber: complex,
def hxy2h(wavenumber: complex, def hxy2h(wavenumber: complex,
dxes: dx_lists_t, dxes: dx_lists_t,
mu: vfdfield_t = None mu: Optional[vfdfield_t] = None
) -> sparse.spmatrix: ) -> sparse.spmatrix:
""" """
Operator which transforms the vector `h_xy` containing the vectorized H_x and H_y fields, Operator which transforms the vector `h_xy` containing the vectorized H_x and H_y fields,
@ -520,7 +520,7 @@ def exy2e(wavenumber: complex,
def e2h(wavenumber: complex, def e2h(wavenumber: complex,
omega: complex, omega: complex,
dxes: dx_lists_t, dxes: dx_lists_t,
mu: vfdfield_t = None mu: Optional[vfdfield_t] = None
) -> sparse.spmatrix: ) -> sparse.spmatrix:
""" """
Returns an operator which, when applied to a vectorized E eigenfield, produces Returns an operator which, when applied to a vectorized E eigenfield, produces
@ -676,7 +676,7 @@ def solve_modes(mode_numbers: List[int],
epsilon: vfdfield_t, epsilon: vfdfield_t,
mu: vfdfield_t = None, mu: vfdfield_t = None,
mode_margin: int = 2, mode_margin: int = 2,
) -> Tuple[List[vfdfield_t], List[complex]]: ) -> Tuple[numpy.ndarray, List[complex]]:
""" """
Given a 2D region, attempts to solve for the eigenmode with the specified mode numbers. Given a 2D region, attempts to solve for the eigenmode with the specified mode numbers.
@ -691,7 +691,8 @@ def solve_modes(mode_numbers: List[int],
ability to find the correct mode. Default 2. ability to find the correct mode. Default 2.
Returns: Returns:
(e_xys, wavenumbers) e_xys: list of vfdfield_t specifying fields
wavenumbers: list of wavenumbers
""" """
''' '''
@ -733,5 +734,6 @@ def solve_mode(mode_number: int,
Returns: Returns:
(e_xy, wavenumber) (e_xy, wavenumber)
""" """
e_xys, wavenumbers = solve_modes(mode_numbers=[mode_number], *args, **kwargs) kwargs['mode_numbers'] = [mode_number]
e_xys, wavenumbers = solve_modes(*args, **kwargs)
return e_xys[:, 0], wavenumbers[0] return e_xys[:, 0], wavenumbers[0]

28
meanas/fdfd/waveguide_3d.py
View File

@ -4,7 +4,7 @@ Tools for working with waveguide modes in 3D domains.
This module relies heavily on `waveguide_2d` and mostly just transforms This module relies heavily on `waveguide_2d` and mostly just transforms
its parameters into 2D equivalents and expands the results back into 3D. its parameters into 2D equivalents and expands the results back into 3D.
""" """
from typing import Dict, List, Tuple from typing import Dict, List, Tuple, Optional, Sequence, Union
import numpy import numpy
import scipy.sparse as sparse import scipy.sparse as sparse
@ -17,10 +17,10 @@ def solve_mode(mode_number: int,
dxes: dx_lists_t, dxes: dx_lists_t,
axis: int, axis: int,
polarity: int, polarity: int,
slices: List[slice], slices: Sequence[slice],
epsilon: fdfield_t, epsilon: fdfield_t,
mu: fdfield_t = None, mu: Optional[fdfield_t] = None,
) -> Dict[str, complex or numpy.ndarray]: ) -> Dict[str, Union[complex, numpy.ndarray]]:
""" """
Given a 3D grid, selects a slice from the grid and attempts to Given a 3D grid, selects a slice from the grid and attempts to
solve for an eigenmode propagating through that slice. solve for an eigenmode propagating through that slice.
@ -104,9 +104,9 @@ def compute_source(E: fdfield_t,
dxes: dx_lists_t, dxes: dx_lists_t,
axis: int, axis: int,
polarity: int, polarity: int,
slices: List[slice], slices: Sequence[slice],
epsilon: fdfield_t, epsilon: fdfield_t,
mu: fdfield_t = None, mu: Optional[fdfield_t] = None,
) -> fdfield_t: ) -> fdfield_t:
""" """
Given an eigenmode obtained by `solve_mode`, returns the current source distribution Given an eigenmode obtained by `solve_mode`, returns the current source distribution
@ -148,7 +148,7 @@ def compute_overlap_e(E: fdfield_t,
dxes: dx_lists_t, dxes: dx_lists_t,
axis: int, axis: int,
polarity: int, polarity: int,
slices: List[slice], slices: Sequence[slice],
) -> fdfield_t: # TODO DOCS ) -> fdfield_t: # TODO DOCS
""" """
Given an eigenmode obtained by `solve_mode`, calculates an overlap_e for the Given an eigenmode obtained by `solve_mode`, calculates an overlap_e for the
@ -177,9 +177,9 @@ def compute_overlap_e(E: fdfield_t,
start, stop = sorted((slices[axis].start, slices[axis].start - 2 * polarity)) start, stop = sorted((slices[axis].start, slices[axis].start - 2 * polarity))
slices2 = list(slices) slices2_l = list(slices)
slices2[axis] = slice(start, stop) slices2_l[axis] = slice(start, stop)
slices2 = (slice(None), *slices2) slices2 = (slice(None), *slices2_l)
Etgt = numpy.zeros_like(Ee) Etgt = numpy.zeros_like(Ee)
Etgt[slices2] = Ee[slices2] Etgt[slices2] = Ee[slices2]
@ -193,7 +193,7 @@ def expand_e(E: fdfield_t,
dxes: dx_lists_t, dxes: dx_lists_t,
axis: int, axis: int,
polarity: int, polarity: int,
slices: List[slice], slices: Sequence[slice],
) -> fdfield_t: ) -> fdfield_t:
""" """
Given an eigenmode obtained by `solve_mode`, expands the E-field from the 2D Given an eigenmode obtained by `solve_mode`, expands the E-field from the 2D
@ -226,9 +226,9 @@ def expand_e(E: fdfield_t,
# Expand our slice to the entire grid using the phase factors # Expand our slice to the entire grid using the phase factors
E_expanded = numpy.zeros_like(E) E_expanded = numpy.zeros_like(E)
slices_exp = list(slices) slices_exp_l = list(slices)
slices_exp[axis] = slice(E.shape[axis + 1]) slices_exp_l[axis] = slice(E.shape[axis + 1])
slices_exp = (slice(None), *slices_exp) slices_exp = (slice(None), *slices_exp_l)
slices_in = (slice(None), *slices) slices_in = (slice(None), *slices)

4
meanas/fdfd/waveguide_cyl.py
View File

@ -8,7 +8,7 @@ As the z-dependence is known, all the functions in this file assume a 2D grid
""" """
# TODO update module docs # TODO update module docs
from typing import List, Tuple, Dict from typing import List, Tuple, Dict, Union
import numpy import numpy
from numpy.linalg import norm from numpy.linalg import norm
import scipy.sparse as sparse import scipy.sparse as sparse
@ -85,7 +85,7 @@ def solve_mode(mode_number: int,
dxes: dx_lists_t, dxes: dx_lists_t,
epsilon: vfdfield_t, epsilon: vfdfield_t,
r0: float, r0: float,
) -> Dict[str, complex or fdfield_t]: ) -> Dict[str, Union[complex, fdfield_t]]:
""" """
TODO: fixup TODO: fixup
Given a 2d (r, y) slice of epsilon, attempts to solve for the eigenmode Given a 2d (r, y) slice of epsilon, attempts to solve for the eigenmode

2
meanas/fdmath/__init__.py
View File

@ -741,7 +741,7 @@ the true values can be multiplied back in after the simulation is complete if no
normalized results are needed. normalized results are needed.
""" """
from .types import fdfield_t, vfdfield_t, dx_lists_t, fdfield_updater_t from .types import fdfield_t, vfdfield_t, dx_lists_t, dx_lists_mut, fdfield_updater_t
from .vectorization import vec, unvec from .vectorization import vec, unvec
from . import operators, functional, types, vectorization from . import operators, functional, types, vectorization

28
meanas/fdmath/functional.py
View File

@ -3,13 +3,14 @@ Math functions for finite difference simulations
Basic discrete calculus etc. Basic discrete calculus etc.
""" """
from typing import List, Callable, Tuple, Dict from typing import Sequence, Tuple, Dict, Optional
import numpy import numpy
from .types import fdfield_t, fdfield_updater_t from .types import fdfield_t, fdfield_updater_t
def deriv_forward(dx_e: List[numpy.ndarray] = None) -> fdfield_updater_t: def deriv_forward(dx_e: Optional[Sequence[numpy.ndarray]] = None
) -> Tuple[fdfield_updater_t, fdfield_updater_t, fdfield_updater_t]:
""" """
Utility operators for taking discretized derivatives (backward variant). Utility operators for taking discretized derivatives (backward variant).
@ -21,17 +22,18 @@ def deriv_forward(dx_e: List[numpy.ndarray] = None) -> fdfield_updater_t:
List of functions for taking forward derivatives along each axis. List of functions for taking forward derivatives along each axis.
""" """
if dx_e: if dx_e:
derivs = [lambda f: (numpy.roll(f, -1, axis=0) - f) / dx_e[0][:, None, None], derivs = (lambda f: (numpy.roll(f, -1, axis=0) - f) / dx_e[0][:, None, None],
lambda f: (numpy.roll(f, -1, axis=1) - f) / dx_e[1][None, :, None], lambda f: (numpy.roll(f, -1, axis=1) - f) / dx_e[1][None, :, None],
lambda f: (numpy.roll(f, -1, axis=2) - f) / dx_e[2][None, None, :]] lambda f: (numpy.roll(f, -1, axis=2) - f) / dx_e[2][None, None, :])
else: else:
derivs = [lambda f: numpy.roll(f, -1, axis=0) - f, derivs = (lambda f: numpy.roll(f, -1, axis=0) - f,
lambda f: numpy.roll(f, -1, axis=1) - f, lambda f: numpy.roll(f, -1, axis=1) - f,
lambda f: numpy.roll(f, -1, axis=2) - f] lambda f: numpy.roll(f, -1, axis=2) - f)
return derivs return derivs
def deriv_back(dx_h: List[numpy.ndarray] = None) -> fdfield_updater_t: def deriv_back(dx_h: Optional[Sequence[numpy.ndarray]] = None
) -> Tuple[fdfield_updater_t, fdfield_updater_t, fdfield_updater_t]:
""" """
Utility operators for taking discretized derivatives (forward variant). Utility operators for taking discretized derivatives (forward variant).
@ -43,17 +45,17 @@ def deriv_back(dx_h: List[numpy.ndarray] = None) -> fdfield_updater_t:
List of functions for taking forward derivatives along each axis. List of functions for taking forward derivatives along each axis.
""" """
if dx_h: if dx_h:
derivs = [lambda f: (f - numpy.roll(f, 1, axis=0)) / dx_h[0][:, None, None], derivs = (lambda f: (f - numpy.roll(f, 1, axis=0)) / dx_h[0][:, None, None],
lambda f: (f - numpy.roll(f, 1, axis=1)) / dx_h[1][None, :, None], lambda f: (f - numpy.roll(f, 1, axis=1)) / dx_h[1][None, :, None],
lambda f: (f - numpy.roll(f, 1, axis=2)) / dx_h[2][None, None, :]] lambda f: (f - numpy.roll(f, 1, axis=2)) / dx_h[2][None, None, :])
else: else:
derivs = [lambda f: f - numpy.roll(f, 1, axis=0), derivs = (lambda f: f - numpy.roll(f, 1, axis=0),
lambda f: f - numpy.roll(f, 1, axis=1), lambda f: f - numpy.roll(f, 1, axis=1),
lambda f: f - numpy.roll(f, 1, axis=2)] lambda f: f - numpy.roll(f, 1, axis=2))
return derivs return derivs
def curl_forward(dx_e: List[numpy.ndarray] = None) -> fdfield_updater_t: def curl_forward(dx_e: Optional[Sequence[numpy.ndarray]] = None) -> fdfield_updater_t:
""" """
Curl operator for use with the E field. Curl operator for use with the E field.
@ -80,7 +82,7 @@ def curl_forward(dx_e: List[numpy.ndarray] = None) -> fdfield_updater_t:
return ce_fun return ce_fun
def curl_back(dx_h: List[numpy.ndarray] = None) -> fdfield_updater_t: def curl_back(dx_h: Optional[Sequence[numpy.ndarray]] = None) -> fdfield_updater_t:
""" """
Create a function which takes the backward curl of a field. Create a function which takes the backward curl of a field.

20
meanas/fdmath/operators.py
View File

@ -3,14 +3,14 @@ Matrix operators for finite difference simulations
Basic discrete calculus etc. Basic discrete calculus etc.
""" """
from typing import List, Callable, Tuple, Dict from typing import Sequence, List, Callable, Tuple, Dict
import numpy import numpy
import scipy.sparse as sparse import scipy.sparse as sparse
from .types import fdfield_t, vfdfield_t from .types import fdfield_t, vfdfield_t
def rotation(axis: int, shape: List[int], shift_distance: int=1) -> sparse.spmatrix: def rotation(axis: int, shape: Sequence[int], shift_distance: int=1) -> sparse.spmatrix:
""" """
Utility operator for performing a circular shift along a specified axis by a Utility operator for performing a circular shift along a specified axis by a
specified number of elements. specified number of elements.
@ -46,7 +46,7 @@ def rotation(axis: int, shape: List[int], shift_distance: int=1) -> sparse.spmat
return d return d
def shift_with_mirror(axis: int, shape: List[int], shift_distance: int=1) -> sparse.spmatrix: def shift_with_mirror(axis: int, shape: Sequence[int], shift_distance: int=1) -> sparse.spmatrix:
""" """
Utility operator for performing an n-element shift along a specified axis, with mirror Utility operator for performing an n-element shift along a specified axis, with mirror
boundary conditions applied to the cells beyond the receding edge. boundary conditions applied to the cells beyond the receding edge.
@ -87,7 +87,7 @@ def shift_with_mirror(axis: int, shape: List[int], shift_distance: int=1) -> spa
return d return d
def deriv_forward(dx_e: List[numpy.ndarray]) -> List[sparse.spmatrix]: def deriv_forward(dx_e: Sequence[numpy.ndarray]) -> List[sparse.spmatrix]:
""" """
Utility operators for taking discretized derivatives (forward variant). Utility operators for taking discretized derivatives (forward variant).
@ -112,7 +112,7 @@ def deriv_forward(dx_e: List[numpy.ndarray]) -> List[sparse.spmatrix]:
return Ds return Ds
def deriv_back(dx_h: List[numpy.ndarray]) -> List[sparse.spmatrix]: def deriv_back(dx_h: Sequence[numpy.ndarray]) -> List[sparse.spmatrix]:
""" """
Utility operators for taking discretized derivatives (backward variant). Utility operators for taking discretized derivatives (backward variant).
@ -137,7 +137,7 @@ def deriv_back(dx_h: List[numpy.ndarray]) -> List[sparse.spmatrix]:
return Ds return Ds
def cross(B: List[sparse.spmatrix]) -> sparse.spmatrix: def cross(B: Sequence[sparse.spmatrix]) -> sparse.spmatrix:
""" """
Cross product operator Cross product operator
@ -171,7 +171,7 @@ def vec_cross(b: vfdfield_t) -> sparse.spmatrix:
return cross(B) return cross(B)
def avg_forward(axis: int, shape: List[int]) -> sparse.spmatrix: def avg_forward(axis: int, shape: Sequence[int]) -> sparse.spmatrix:
""" """
Forward average operator `(x4 = (x4 + x5) / 2)` Forward average operator `(x4 = (x4 + x5) / 2)`
@ -189,7 +189,7 @@ def avg_forward(axis: int, shape: List[int]) -> sparse.spmatrix:
return 0.5 * (sparse.eye(n) + rotation(axis, shape)) return 0.5 * (sparse.eye(n) + rotation(axis, shape))
def avg_back(axis: int, shape: List[int]) -> sparse.spmatrix: def avg_back(axis: int, shape: Sequence[int]) -> sparse.spmatrix:
""" """
Backward average operator `(x4 = (x4 + x3) / 2)` Backward average operator `(x4 = (x4 + x3) / 2)`
@ -203,7 +203,7 @@ def avg_back(axis: int, shape: List[int]) -> sparse.spmatrix:
return avg_forward(axis, shape).T return avg_forward(axis, shape).T
def curl_forward(dx_e: List[numpy.ndarray]) -> sparse.spmatrix: def curl_forward(dx_e: Sequence[numpy.ndarray]) -> sparse.spmatrix:
""" """
Curl operator for use with the E field. Curl operator for use with the E field.
@ -217,7 +217,7 @@ def curl_forward(dx_e: List[numpy.ndarray]) -> sparse.spmatrix:
return cross(deriv_forward(dx_e)) return cross(deriv_forward(dx_e))
def curl_back(dx_h: List[numpy.ndarray]) -> sparse.spmatrix: def curl_back(dx_h: Sequence[numpy.ndarray]) -> sparse.spmatrix:
""" """
Curl operator for use with the H field. Curl operator for use with the H field.

9
meanas/fdmath/types.py
View File

@ -2,7 +2,7 @@
Types shared across multiple submodules Types shared across multiple submodules
""" """
import numpy import numpy
from typing import List, Callable from typing import Sequence, Callable, MutableSequence
# Field types # Field types
@ -25,7 +25,7 @@ class vfdfield_t(numpy.ndarray):
pass pass
dx_lists_t = List[List[numpy.ndarray]] dx_lists_t = Sequence[Sequence[numpy.ndarray]]
''' '''
'dxes' datastructure which contains grid cell width information in the following format: 'dxes' datastructure which contains grid cell width information in the following format:
@ -36,6 +36,11 @@ dx_lists_t = List[List[numpy.ndarray]]
and `dy_h[0]` is the y-width of the `y=0` cells, as used when calculating dH/dy, etc. and `dy_h[0]` is the y-width of the `y=0` cells, as used when calculating dH/dy, etc.
''' '''
dx_lists_mut = MutableSequence[MutableSequence[numpy.ndarray]]
'''
Mutable version of `dx_lists_t`
'''
fdfield_updater_t = Callable[..., fdfield_t] fdfield_updater_t = Callable[..., fdfield_t]
''' '''

20
meanas/fdmath/vectorization.py
View File

@ -4,16 +4,20 @@ and a 1D array representation of that field `[f_x0, f_x1, f_x2,... f_y0,... f_z0
Vectorized versions of the field use row-major (ie., C-style) ordering. Vectorized versions of the field use row-major (ie., C-style) ordering.
""" """
from typing import List from typing import Optional, TypeVar, overload, Union, List
import numpy import numpy
from .types import fdfield_t, vfdfield_t from .types import fdfield_t, vfdfield_t
@overload
def vec(f: None) -> None:
pass
__author__ = 'Jan Petykiewicz' @overload
def vec(f: Union[fdfield_t, List[numpy.ndarray]]) -> vfdfield_t:
pass
def vec(f: Optional[Union[fdfield_t, List[numpy.ndarray]]]) -> Optional[vfdfield_t]:
def vec(f: fdfield_t) -> vfdfield_t:
""" """
Create a 1D ndarray from a 3D vector field which spans a 1-3D region. Create a 1D ndarray from a 3D vector field which spans a 1-3D region.
@ -31,7 +35,15 @@ def vec(f: fdfield_t) -> vfdfield_t:
return numpy.ravel(f, order='C') return numpy.ravel(f, order='C')
@overload
def unvec(v: None, shape: numpy.ndarray) -> None:
pass
@overload
def unvec(v: vfdfield_t, shape: numpy.ndarray) -> fdfield_t: def unvec(v: vfdfield_t, shape: numpy.ndarray) -> fdfield_t:
pass
def unvec(v: Optional[vfdfield_t], shape: numpy.ndarray) -> Optional[fdfield_t]:
""" """
Perform the inverse of vec(): take a 1D ndarray and output a 3D field Perform the inverse of vec(): take a 1D ndarray and output a 3D field
of form `[f_x, f_y, f_z]` where each of `f_*` is a len(shape)-dimensional of form `[f_x, f_y, f_z]` where each of `f_*` is a len(shape)-dimensional

6
meanas/fdtd/base.py
View File

@ -3,7 +3,7 @@ Basic FDTD field updates
""" """
from typing import List, Callable, Tuple, Dict from typing import List, Callable, Dict, Union
import numpy import numpy
from ..fdmath import dx_lists_t, fdfield_t, fdfield_updater_t from ..fdmath import dx_lists_t, fdfield_t, fdfield_updater_t
@ -47,7 +47,7 @@ def maxwell_e(dt: float, dxes: dx_lists_t = None) -> fdfield_updater_t:
else: else:
curl_h_fun = curl_back() curl_h_fun = curl_back()
def me_fun(e: fdfield_t, h: fdfield_t, epsilon: fdfield_t): def me_fun(e: fdfield_t, h: fdfield_t, epsilon: Union[fdfield_t, float]) -> fdfield_t:
""" """
Update the E-field. Update the E-field.
@ -100,7 +100,7 @@ def maxwell_h(dt: float, dxes: dx_lists_t = None) -> fdfield_updater_t:
else: else:
curl_e_fun = curl_forward() curl_e_fun = curl_forward()
def mh_fun(e: fdfield_t, h: fdfield_t, mu: fdfield_t = None): def mh_fun(e: fdfield_t, h: fdfield_t, mu: Union[fdfield_t, float, None] = None) -> fdfield_t:
""" """
Update the H-field. Update the H-field.

8
meanas/fdtd/boundaries.py
View File

@ -4,7 +4,7 @@ Boundary conditions
#TODO conducting boundary documentation #TODO conducting boundary documentation
""" """
from typing import List, Callable, Tuple, Dict from typing import Callable, Tuple, Dict, Any, List
import numpy import numpy
from ..fdmath import dx_lists_t, fdfield_t, fdfield_updater_t from ..fdmath import dx_lists_t, fdfield_t, fdfield_updater_t
@ -20,8 +20,8 @@ def conducting_boundary(direction: int,
u, v = dirs u, v = dirs
if polarity < 0: if polarity < 0:
boundary_slice = [slice(None)] * 3 boundary_slice = [slice(None)] * 3 # type: List[Any]
shifted1_slice = [slice(None)] * 3 shifted1_slice = [slice(None)] * 3 # type: List[Any]
boundary_slice[direction] = 0 boundary_slice[direction] = 0
shifted1_slice[direction] = 1 shifted1_slice[direction] = 1
@ -42,7 +42,7 @@ def conducting_boundary(direction: int,
if polarity > 0: if polarity > 0:
boundary_slice = [slice(None)] * 3 boundary_slice = [slice(None)] * 3
shifted1_slice = [slice(None)] * 3 shifted1_slice = [slice(None)] * 3
shifted2_slice = [slice(None)] * 3 shifted2_slice = [slice(None)] * 3 # type: List[Any]
boundary_slice[direction] = -1 boundary_slice[direction] = -1
shifted1_slice[direction] = -2 shifted1_slice[direction] = -2
shifted2_slice[direction] = -3 shifted2_slice[direction] = -3

50
meanas/fdtd/energy.py
View File

@ -1,5 +1,5 @@
# pylint: disable=unsupported-assignment-operation # pylint: disable=unsupported-assignment-operation
from typing import List, Callable, Tuple, Dict from typing import Callable, Tuple, Dict, Optional, Union
import numpy import numpy
from ..fdmath import dx_lists_t, fdfield_t, fdfield_updater_t from ..fdmath import dx_lists_t, fdfield_t, fdfield_updater_t
@ -8,7 +8,7 @@ from ..fdmath.functional import deriv_back, deriv_forward
def poynting(e: fdfield_t, def poynting(e: fdfield_t,
h: fdfield_t, h: fdfield_t,
dxes: dx_lists_t = None, dxes: Optional[dx_lists_t] = None,
) -> fdfield_t: ) -> fdfield_t:
""" """
Calculate the poynting vector Calculate the poynting vector
@ -30,16 +30,19 @@ def poynting(e: fdfield_t,
return s return s
def poynting_divergence(s: fdfield_t = None, def poynting_divergence(s: Optional[fdfield_t] = None,
*, *,
e: fdfield_t = None, e: Optional[fdfield_t] = None,
h: fdfield_t = None, h: Optional[fdfield_t] = None,
dxes: dx_lists_t = None, dxes: Optional[dx_lists_t] = None,
) -> fdfield_t: ) -> fdfield_t:
""" """
Calculate the divergence of the poynting vector Calculate the divergence of the poynting vector
""" """
if s is None: if s is None:
assert(e is not None)
assert(h is not None)
assert(dxes is not None)
s = poynting(e, h, dxes=dxes) s = poynting(e, h, dxes=dxes)
Dx, Dy, Dz = deriv_back() Dx, Dy, Dz = deriv_back()
@ -50,9 +53,9 @@ def poynting_divergence(s: fdfield_t = None,
def energy_hstep(e0: fdfield_t, def energy_hstep(e0: fdfield_t,
h1: fdfield_t, h1: fdfield_t,
e2: fdfield_t, e2: fdfield_t,
epsilon: fdfield_t = None, epsilon: Optional[fdfield_t] = None,
mu: fdfield_t = None, mu: Optional[fdfield_t] = None,
dxes: dx_lists_t = None, dxes: Optional[dx_lists_t] = None,
) -> fdfield_t: ) -> fdfield_t:
u = dxmul(e0 * e2, h1 * h1, epsilon, mu, dxes) u = dxmul(e0 * e2, h1 * h1, epsilon, mu, dxes)
return u return u
@ -61,9 +64,9 @@ def energy_hstep(e0: fdfield_t,
def energy_estep(h0: fdfield_t, def energy_estep(h0: fdfield_t,
e1: fdfield_t, e1: fdfield_t,
h2: fdfield_t, h2: fdfield_t,
epsilon: fdfield_t = None, epsilon: Optional[fdfield_t] = None,
mu: fdfield_t = None, mu: Optional[fdfield_t] = None,
dxes: dx_lists_t = None, dxes: Optional[dx_lists_t] = None,
) -> fdfield_t: ) -> fdfield_t:
u = dxmul(e1 * e1, h0 * h2, epsilon, mu, dxes) u = dxmul(e1 * e1, h0 * h2, epsilon, mu, dxes)
return u return u
@ -74,9 +77,9 @@ def delta_energy_h2e(dt: float,
h1: fdfield_t, h1: fdfield_t,
e2: fdfield_t, e2: fdfield_t,
h3: fdfield_t, h3: fdfield_t,
epsilon: fdfield_t = None, epsilon: Optional[fdfield_t] = None,
mu: fdfield_t = None, mu: Optional[fdfield_t] = None,
dxes: dx_lists_t = None, dxes: Optional[dx_lists_t] = None,
) -> fdfield_t: ) -> fdfield_t:
""" """
This is just from (e2 * e2 + h3 * h1) - (h1 * h1 + e0 * e2) This is just from (e2 * e2 + h3 * h1) - (h1 * h1 + e0 * e2)
@ -92,9 +95,9 @@ def delta_energy_e2h(dt: float,
e1: fdfield_t, e1: fdfield_t,
h2: fdfield_t, h2: fdfield_t,
e3: fdfield_t, e3: fdfield_t,
epsilon: fdfield_t = None, epsilon: Optional[fdfield_t] = None,
mu: fdfield_t = None, mu: Optional[fdfield_t] = None,
dxes: dx_lists_t = None, dxes: Optional[dx_lists_t] = None,
) -> fdfield_t: ) -> fdfield_t:
""" """
This is just from (h2 * h2 + e3 * e1) - (e1 * e1 + h0 * h2) This is just from (h2 * h2 + e3 * e1) - (e1 * e1 + h0 * h2)
@ -105,7 +108,10 @@ def delta_energy_e2h(dt: float,
return du return du
def delta_energy_j(j0: fdfield_t, e1: fdfield_t, dxes: dx_lists_t = None) -> fdfield_t: def delta_energy_j(j0: fdfield_t,
e1: fdfield_t,
dxes: Optional[dx_lists_t] = None,
) -> fdfield_t:
if dxes is None: if dxes is None:
dxes = tuple(tuple(numpy.ones(1) for _ in range(3)) for _ in range(2)) dxes = tuple(tuple(numpy.ones(1) for _ in range(3)) for _ in range(2))
@ -118,9 +124,9 @@ def delta_energy_j(j0: fdfield_t, e1: fdfield_t, dxes: dx_lists_t = None) -> fdf
def dxmul(ee: fdfield_t, def dxmul(ee: fdfield_t,
hh: fdfield_t, hh: fdfield_t,
epsilon: fdfield_t = None, epsilon: Optional[Union[fdfield_t, float]] = None,
mu: fdfield_t = None, mu: Optional[Union[fdfield_t, float]] = None,
dxes: dx_lists_t = None dxes: Optional[dx_lists_t] = None
) -> fdfield_t: ) -> fdfield_t:
if epsilon is None: if epsilon is None:
epsilon = 1 epsilon = 1

8
meanas/fdtd/pml.py
View File

@ -7,7 +7,7 @@ PML implementations
""" """
# TODO retest pmls! # TODO retest pmls!
from typing import List, Callable, Tuple, Dict from typing import List, Callable, Tuple, Dict, Any
import numpy import numpy
from ..fdmath import dx_lists_t, fdfield_t, fdfield_updater_t from ..fdmath import dx_lists_t, fdfield_t, fdfield_updater_t
@ -59,9 +59,9 @@ def cpml(direction: int,
else: else:
raise Exception('Bad polarity!') raise Exception('Bad polarity!')
expand_slice = [None] * 3 expand_slice_l: List[Any] = [None] * 3
expand_slice[direction] = slice(None) expand_slice_l[direction] = slice(None)
expand_slice = tuple(expand_slice) expand_slice = tuple(expand_slice_l)
def par(x): def par(x):
scaling = (x / thickness) ** m scaling = (x / thickness) ** m