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.
110 lines
3.3 KiB
Python
110 lines
3.3 KiB
Python
"""
|
|
Math functions for finite difference simulations
|
|
|
|
Basic discrete calculus etc.
|
|
"""
|
|
from typing import List, Callable, Tuple, Dict
|
|
import numpy
|
|
|
|
from .types import fdfield_t, fdfield_updater_t
|
|
|
|
|
|
def deriv_forward(dx_e: List[numpy.ndarray] = None) -> 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:
|
|
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: List[numpy.ndarray] = None) -> 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:
|
|
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: List[numpy.ndarray] = 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: List[numpy.ndarray] = 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
|
|
|
|
|