Compare commits

...

4 Commits

Author SHA1 Message Date
5b16295769 update deps 2024-07-30 23:19:31 -07:00
fe684d83d3 add ruff and mypy configs 2024-07-30 23:19:21 -07:00
7f620a9bb3 modern rng approach 2024-07-30 23:19:07 -07:00
50b30d31fb fixes driven by ruff & mypy 2024-07-30 23:18:54 -07:00
4 changed files with 131 additions and 72 deletions

View File

@ -1,4 +1,7 @@
from .simulation import Simulation, type_to_C
from .simulation import (
Simulation as Simulation,
type_to_C as type_to_C,
)
__author__ = 'Jan Petykiewicz'
__version__ = '0.4'

View File

@ -2,12 +2,13 @@
Class for constructing and holding the basic FDTD operations and fields
"""
from typing import Callable, Type, Sequence
from collections import OrderedDict
from collections.abc import Callable, Sequence
import numpy
from numpy.typing import NDArray
from numpy import floating, complexfloating
import jinja2
import warnings
import logging
import pyopencl
import pyopencl.array
@ -16,13 +17,18 @@ from pyopencl.elementwise import ElementwiseKernel
from meanas.fdmath import vec
__author__ = 'Jan Petykiewicz'
logger = logging.getLogger(__name__)
# Create jinja2 env on module load
jinja_env = jinja2.Environment(loader=jinja2.PackageLoader(__name__.split('.')[0], 'kernels'))
class FDTDError(Exception):
""" Custom exception for opencl_fdtd """
pass
class Simulation:
r"""
Constructs and holds the basic FDTD operations and related fields
@ -64,7 +70,7 @@ class Simulation:
dt: float
inv_dxes: list[pyopencl.array.Array]
arg_type: Type
arg_type: type
context: pyopencl.Context
queue: pyopencl.CommandQueue
@ -85,7 +91,7 @@ class Simulation:
initial_fields: dict[str, NDArray] | None = None,
context: pyopencl.Context | None = None,
queue: pyopencl.CommandQueue | None = None,
float_type: Type = numpy.float32,
float_type: type = numpy.float32,
do_poynting: bool = True,
do_fieldsrc: bool = False,
) -> None:
@ -134,7 +140,7 @@ class Simulation:
if dxes is None:
dxes = 1.0
if isinstance(dxes, (float, int)):
if isinstance(dxes, float | int):
uniform_dx = dxes
min_dx = dxes
else:
@ -143,13 +149,14 @@ class Simulation:
min_dx = min(min(dxn) for dxn in dxes[0] + dxes[1])
max_dt = min_dx * .99 / numpy.sqrt(3)
logger.info(f'{min_dx=}, {max_dt=}, {dt=}')
if dt is None:
self.dt = max_dt
elif dt > max_dt:
warnings.warn(f'Warning: unstable dt: {dt}')
warnings.warn(f'Warning: unstable dt: {dt}', stacklevel=2)
elif dt <= 0:
raise Exception(f'Invalid dt: {dt}')
raise FDTDError(f'Invalid dt: {dt}')
else:
self.dt = dt
@ -173,28 +180,31 @@ class Simulation:
def ptr(arg: str) -> str:
return ctype + ' *' + arg
base_fields = OrderedDict()
base_fields[ptr('E')] = self.E
base_fields[ptr('H')] = self.H
base_fields[ctype + ' dt'] = self.dt
base_fields = {
ptr('E'): self.E,
ptr('H'): self.H,
ctype + ' dt': self.dt,
}
if uniform_dx is False:
inv_dx_names = ['inv_d' + eh + r for eh in 'eh' for r in 'xyz']
for name, field in zip(inv_dx_names, self.inv_dxes):
for name, field in zip(inv_dx_names, self.inv_dxes, strict=True):
base_fields[ptr(name)] = field
eps_field = OrderedDict()
eps_field[ptr('eps')] = self.eps
eps_field = {ptr('eps'): self.eps}
if bloch_boundaries:
base_fields[ptr('F')] = self.F
base_fields[ptr('G')] = self.G
base_fields |= {
ptr('F'): self.F,
ptr('G'): self.G,
}
bloch_fields = OrderedDict()
bloch_fields[ptr('E')] = self.F
bloch_fields[ptr('H')] = self.G
bloch_fields[ctype + ' dt'] = self.dt
bloch_fields[ptr('F')] = self.E
bloch_fields[ptr('G')] = self.H
bloch_fields = {
ptr('E'): self.F,
ptr('H'): self.G,
ctype + ' dt': self.dt,
ptr('F'): self.E,
ptr('G'): self.H,
}
common_source = jinja_env.get_template('common.cl').render(
ftype=ctype,
@ -216,8 +226,8 @@ class Simulation:
if bloch_boundaries:
bloch_args = jinja_args.copy()
bloch_args['do_poynting'] = False
bloch_args['bloch'] = [
{'axis': b['axis'],
bloch_args['bloch'] = [{
'axis': b['axis'],
'real': b['imag'],
'imag': b['real'],
}
@ -227,7 +237,7 @@ class Simulation:
self.sources['F'] = F_source
self.sources['G'] = G_source
S_fields = OrderedDict()
S_fields = {}
if do_poynting:
self.S = pyopencl.array.zeros_like(self.E)
S_fields[ptr('S')] = self.S
@ -237,7 +247,7 @@ class Simulation:
S_fields[ptr('S0')] = self.S0
S_fields[ptr('S1')] = self.S1
J_fields = OrderedDict()
J_fields = {}
if do_fieldsrc:
J_source = jinja_env.get_template('update_j.cl').render(**jinja_args)
self.sources['J'] = J_source
@ -247,37 +257,36 @@ class Simulation:
J_fields[ptr('Jr')] = self.Jr
J_fields[ptr('Ji')] = self.Ji
'''
PML
'''
#
# PML
#
pml_e_fields, pml_h_fields = self._create_pmls(pmls)
if bloch_boundaries:
pml_f_fields, pml_g_fields = self._create_pmls(pmls)
'''
Create operations
'''
#
# Create operations
#
self.update_E = self._create_operation(E_source, (base_fields, eps_field, pml_e_fields))
self.update_H = self._create_operation(H_source, (base_fields, pml_h_fields, S_fields))
if bloch_boundaries:
self.update_F = self._create_operation(F_source, (bloch_fields, eps_field, pml_f_fields))
self.update_G = self._create_operation(G_source, (bloch_fields, pml_g_fields))
if do_fieldsrc:
args = OrderedDict()
[args.update(d) for d in (base_fields, J_fields)]
args = base_fields | J_fields
var_args = [ctype + ' ' + v for v in 'cs'] + ['uint ' + r + m for r in 'xyz' for m in ('min', 'max')]
update = ElementwiseKernel(self.context, operation=J_source,
arguments=', '.join(list(args.keys()) + var_args))
self.update_J = lambda e, *a: update(*args.values(), *a, wait_for=e)
def _create_pmls(self, pmls):
def _create_pmls(self, pmls: Sequence[dict[str, float]]) -> tuple[dict[str, pyopencl.array.Array], ...]:
ctype = type_to_C(self.arg_type)
def ptr(arg: str) -> str:
return ctype + ' *' + arg
pml_e_fields = OrderedDict()
pml_h_fields = OrderedDict()
pml_e_fields = {}
pml_h_fields = {}
for pml in pmls:
a = 'xyz'.find(pml['axis'])
@ -285,7 +294,9 @@ class Simulation:
kappa_max = numpy.sqrt(pml['mu_eff'] * pml['epsilon_eff'])
alpha_max = pml['cfs_alpha']
def par(x):
print(sigma_max, kappa_max, alpha_max, pml['thickness'], self.dt)
def par(x, pml=pml, sigma_max=sigma_max, kappa_max=kappa_max, alpha_max=alpha_max): # noqa: ANN001, ANN202
scaling = (x / pml['thickness']) ** pml['m']
sigma = scaling * sigma_max
kappa = 1 + scaling * (kappa_max - 1)
@ -301,8 +312,13 @@ class Simulation:
elif pml['polarity'] == 'n':
xh -= 0.5
logger.debug(f'{pml=}')
logger.debug(f'{xe=}')
logger.debug(f'{xh=}')
logger.debug(f'{par(xe)=}')
logger.debug(f'{par(xh)=}')
pml_p_names = [['p' + pml['axis'] + i + eh + pml['polarity'] for i in '012'] for eh in 'eh']
for name_e, name_h, pe, ph in zip(pml_p_names[0], pml_p_names[1], par(xe), par(xh)):
for name_e, name_h, pe, ph in zip(pml_p_names[0], pml_p_names[1], par(xe), par(xh), strict=True):
pml_e_fields[ptr(name_e)] = pyopencl.array.to_device(self.queue, pe)
pml_h_fields[ptr(name_h)] = pyopencl.array.to_device(self.queue, ph)
@ -313,13 +329,13 @@ class Simulation:
psi_shape = list(self.shape)
psi_shape[a] = pml['thickness']
for ne, nh in zip(*psi_names):
for ne, nh in zip(*psi_names, strict=True):
pml_e_fields[ptr(ne)] = pyopencl.array.zeros(self.queue, tuple(psi_shape), dtype=self.arg_type)
pml_h_fields[ptr(nh)] = pyopencl.array.zeros(self.queue, tuple(psi_shape), dtype=self.arg_type)
return pml_e_fields, pml_h_fields
def _create_operation(self, source, args_fields) -> Callable[..., pyopencl.Event]:
args = OrderedDict()
def _create_operation(self, source: str, args_fields: Sequence[dict[str, pyopencl.array.Array]]) -> Callable[..., pyopencl.Event]:
args = {}
for d in args_fields:
args.update(d)
update = ElementwiseKernel(
@ -334,38 +350,30 @@ class Simulation:
context: pyopencl.Context | None = None,
queue: pyopencl.CommandQueue | None = None,
) -> None:
if context is None:
self.context = pyopencl.create_some_context()
else:
self.context = context
if queue is None:
self.queue = pyopencl.CommandQueue(self.context)
else:
self.queue = queue
self.context = context or pyopencl.create_some_context()
self.queue = queue or pyopencl.CommandQueue(self.context)
def _create_eps(self, epsilon: NDArray) -> pyopencl.array.Array:
if len(epsilon) != 3:
raise Exception('Epsilon must be a list with length of 3')
if not all((e.shape == epsilon[0].shape for e in epsilon[1:])):
raise Exception('All epsilon grids must have the same shape. Shapes are {}', [e.shape for e in epsilon])
raise FDTDError('Epsilon must be a list with length of 3')
if not all(e.shape == epsilon[0].shape for e in epsilon[1:]):
raise FDTDError('All epsilon grids must have the same shape. Shapes are {}', [e.shape for e in epsilon])
if not epsilon[0].shape == self.shape:
raise Exception(f'Epsilon shape mismatch. Expected {self.shape}, got {epsilon[0].shape}')
raise FDTDError(f'Epsilon shape mismatch. Expected {self.shape}, got {epsilon[0].shape}')
self.eps = pyopencl.array.to_device(self.queue, vec(epsilon).astype(self.arg_type))
def _create_field(self, initial_value: NDArray | None = None) -> pyopencl.array.Array:
if initial_value is None:
return pyopencl.array.zeros_like(self.eps)
else:
if len(initial_value) != 3:
Exception('Initial field value must be a list of length 3')
if not all((f.shape == self.shape for f in initial_value)):
Exception('Initial field list elements must have same shape as epsilon elements')
raise FDTDError('Initial field value must be a list of length 3')
if not all(f.shape == self.shape for f in initial_value):
raise FDTDError('Initial field list elements must have same shape as epsilon elements')
return pyopencl.array.to_device(self.queue, vec(initial_value).astype(self.arg_type))
def type_to_C(
float_type: Type,
float_type: type,
) -> str:
"""
Returns a string corresponding to the C equivalent of a numpy type.
@ -385,7 +393,7 @@ def type_to_C(
numpy.complex128: 'cdouble_t',
}
if float_type not in types:
raise Exception(f'Unsupported type: {float_type}')
raise FDTDError(f'Unsupported type: {float_type}')
return types[float_type]

View File

@ -323,8 +323,9 @@ class JdotE_3DUniformDX(unittest.TestCase):
self.epsilon = numpy.full(shape, 4, dtype=float)
self.epsilon[self.src_mask] = 2
e = numpy.random.randint(-128, 128 + 1, size=shape).astype(float)
h = numpy.random.randint(-128, 128 + 1, size=shape).astype(float)
rng = numpy.random.default_rng()
e = rng.integers(-128, 128 + 1, size=shape, dtype=float)
h = rng.integers(-128, 128 + 1, size=shape, dtype=float)
self.es = [e]
self.hs = [h]
self.ss = []

View File

@ -32,10 +32,10 @@ classifiers = [
"License :: OSI Approved :: GNU Affero General Public License v3",
"Topic :: Scientific/Engineering",
]
requires-python = ">=3.8"
requires-python = ">=3.11"
dynamic = ["version"]
dependencies = [
"numpy~=1.21",
"numpy>=1.26",
"pyopencl",
"jinja2",
"meanas>=0.3",
@ -44,3 +44,50 @@ dependencies = [
[tool.hatch.version]
path = "opencl_fdtd/__init__.py"
[tool.ruff]
exclude = [
".git",
"dist",
]
line-length = 145
indent-width = 4
lint.dummy-variable-rgx = "^(_+|(_+[a-zA-Z0-9_]*[a-zA-Z0-9]+?))$"
lint.select = [
"NPY", "E", "F", "W", "B", "ANN", "UP", "SLOT", "SIM", "LOG",
"C4", "ISC", "PIE", "PT", "RET", "TCH", "PTH", "INT",
"ARG", "PL", "R", "TRY",
"G010", "G101", "G201", "G202",
"Q002", "Q003", "Q004",
]
lint.ignore = [
#"ANN001", # No annotation
"ANN002", # *args
"ANN003", # **kwargs
"ANN401", # Any
"ANN101", # self: Self
"SIM108", # single-line if / else assignment
"RET504", # x=y+z; return x
"PIE790", # unnecessary pass
"ISC003", # non-implicit string concatenation
"C408", # dict(x=y) instead of {'x': y}
"PLR09", # Too many xxx
"PLR2004", # magic number
"PLC0414", # import x as x
"TRY003", # Long exception message
]
[[tool.mypy.overrides]]
module = [
"scipy",
"scipy.optimize",
"scipy.linalg",
"scipy.sparse",
"scipy.sparse.linalg",
"pyopencl",
"pyopencl.array",
"pyopencl.elementwise",
"pyopencl.reduction",
]
ignore_missing_imports = true