You cannot select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
144 lines
4.3 KiB
Python
144 lines
4.3 KiB
Python
"""
|
|
Math functions for finite difference simulations
|
|
|
|
Basic discrete calculus etc.
|
|
"""
|
|
from typing import Sequence, Callable
|
|
|
|
import numpy
|
|
from numpy.typing import NDArray
|
|
|
|
from .types import fdfield_t, fdfield_updater_t
|
|
|
|
|
|
def deriv_forward(
|
|
dx_e: Sequence[NDArray[numpy.float_]] | None = None,
|
|
) -> tuple[fdfield_updater_t, fdfield_updater_t, fdfield_updater_t]:
|
|
"""
|
|
Utility operators for taking discretized derivatives (backward variant).
|
|
|
|
Args:
|
|
dx_e: Lists of cell sizes for all axes
|
|
`[[dx_0, dx_1, ...], [dy_0, dy_1, ...], ...]`.
|
|
|
|
Returns:
|
|
List of functions for taking forward derivatives along each axis.
|
|
"""
|
|
if dx_e is not 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=2) - f) / dx_e[2][None, None, :])
|
|
else:
|
|
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=2) - f)
|
|
return derivs
|
|
|
|
|
|
def deriv_back(
|
|
dx_h: Sequence[NDArray[numpy.float_]] | None = None,
|
|
) -> tuple[fdfield_updater_t, fdfield_updater_t, fdfield_updater_t]:
|
|
"""
|
|
Utility operators for taking discretized derivatives (forward variant).
|
|
|
|
Args:
|
|
dx_h: Lists of cell sizes for all axes
|
|
`[[dx_0, dx_1, ...], [dy_0, dy_1, ...], ...]`.
|
|
|
|
Returns:
|
|
List of functions for taking forward derivatives along each axis.
|
|
"""
|
|
if dx_h is not 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=2)) / dx_h[2][None, None, :])
|
|
else:
|
|
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=2))
|
|
return derivs
|
|
|
|
|
|
def curl_forward(
|
|
dx_e: Sequence[NDArray[numpy.float_]] | None = None,
|
|
) -> fdfield_updater_t:
|
|
"""
|
|
Curl operator for use with the E field.
|
|
|
|
Args:
|
|
dx_e: Lists of cell sizes for all axes
|
|
`[[dx_0, dx_1, ...], [dy_0, dy_1, ...], ...]`.
|
|
|
|
Returns:
|
|
Function `f` for taking the discrete forward curl of a field,
|
|
`f(E)` -> curlE $= \\nabla_f \\times E$
|
|
"""
|
|
Dx, Dy, Dz = deriv_forward(dx_e)
|
|
|
|
def ce_fun(e: fdfield_t) -> fdfield_t:
|
|
output = numpy.empty_like(e)
|
|
output[0] = Dy(e[2])
|
|
output[1] = Dz(e[0])
|
|
output[2] = Dx(e[1])
|
|
output[0] -= Dz(e[1])
|
|
output[1] -= Dx(e[2])
|
|
output[2] -= Dy(e[0])
|
|
return output
|
|
|
|
return ce_fun
|
|
|
|
|
|
def curl_back(
|
|
dx_h: Sequence[NDArray[numpy.float_]] | None = None,
|
|
) -> fdfield_updater_t:
|
|
"""
|
|
Create a function which takes the backward curl of a field.
|
|
|
|
Args:
|
|
dx_h: Lists of cell sizes for all axes
|
|
`[[dx_0, dx_1, ...], [dy_0, dy_1, ...], ...]`.
|
|
|
|
Returns:
|
|
Function `f` for taking the discrete backward curl of a field,
|
|
`f(H)` -> curlH $= \\nabla_b \\times H$
|
|
"""
|
|
Dx, Dy, Dz = deriv_back(dx_h)
|
|
|
|
def ch_fun(h: fdfield_t) -> fdfield_t:
|
|
output = numpy.empty_like(h)
|
|
output[0] = Dy(h[2])
|
|
output[1] = Dz(h[0])
|
|
output[2] = Dx(h[1])
|
|
output[0] -= Dz(h[1])
|
|
output[1] -= Dx(h[2])
|
|
output[2] -= Dy(h[0])
|
|
return output
|
|
|
|
return ch_fun
|
|
|
|
|
|
def curl_forward_parts(
|
|
dx_e: Sequence[NDArray[numpy.float_]] | None = None,
|
|
) -> Callable:
|
|
Dx, Dy, Dz = deriv_forward(dx_e)
|
|
|
|
def mkparts_fwd(e: fdfield_t) -> tuple[tuple[fdfield_t, fdfield_t], ...]:
|
|
return ((-Dz(e[1]), Dy(e[2])),
|
|
( Dz(e[0]), -Dx(e[2])),
|
|
(-Dy(e[0]), Dx(e[1])))
|
|
|
|
return mkparts_fwd
|
|
|
|
|
|
def curl_back_parts(
|
|
dx_h: Sequence[NDArray[numpy.float_]] | None = None,
|
|
) -> Callable:
|
|
Dx, Dy, Dz = deriv_back(dx_h)
|
|
|
|
def mkparts_back(h: fdfield_t) -> tuple[tuple[fdfield_t, fdfield_t], ...]:
|
|
return ((-Dz(h[1]), Dy(h[2])),
|
|
( Dz(h[0]), -Dx(h[2])),
|
|
(-Dy(h[0]), Dx(h[1])))
|
|
|
|
return mkparts_back
|