Initial commit

fdtd/base.py

@@ -0,0 +1,75 @@
+"""
+Basic code snippets for opencl FDTD
+"""
+
+from typing import List
+import numpy
+
+
+def shape_source(shape: List[int] or numpy.ndarray) -> str:
+    """
+    Defines sx, sy, sz C constants specifying the shape of the grid in each of the 3 dimensions.
+
+    :param shape: [sx, sy, sz] values.
+    :return: String containing C source.
+    """
+    sxyz = """
+// Field sizes
+const int sx = {shape[0]};
+const int sy = {shape[1]};
+const int sz = {shape[2]};
+""".format(shape=shape)
+    return sxyz
+
+# Defines dix, diy, diz constants used for stepping in the x, y, z directions in a linear array
+#  (ie, given Ex[i] referring to position (x, y, z), Ex[i+diy] will refer to position (x, y+1, z))
+dixyz_source = """
+// Convert offset in field xyz to linear index offset
+const int dix = sz * sy;
+const int diy = sz;
+const int diz = 1;
+"""
+
+# Given a linear index i and shape sx, sy, sz, defines x, y, and z as the 3D indices of the current element (i).
+xyz_source = """
+// Convert linear index to field index (xyz)
+const int x = i / (sz * sy);
+const int y = (i - x * sz * sy) / sz;
+const int z = (i - y * sz - x * sz * sy);
+"""
+
+# Source code for updating the E field; maxes use of dixyz_source.
+maxwell_E_source = """
+// E update equations
+Ex[i] += dt / epsx[i] * ((Hz[i] - Hz[i-diy]) - (Hy[i] - Hy[i-diz]));
+Ey[i] += dt / epsy[i] * ((Hx[i] - Hx[i-diz]) - (Hz[i] - Hz[i-dix]));
+Ez[i] += dt / epsz[i] * ((Hy[i] - Hy[i-dix]) - (Hx[i] - Hx[i-diy]));
+"""
+
+# Source code for updating the H field; maxes use of dixyz_source and assumes mu=0
+maxwell_H_source = """
+// H update equations
+Hx[i] -= dt * ((Ez[i+diy] - Ez[i]) - (Ey[i+diz] - Ey[i]));
+Hy[i] -= dt * ((Ex[i+diz] - Ex[i]) - (Ez[i+dix] - Ez[i]));
+Hz[i] -= dt * ((Ey[i+dix] - Ey[i]) - (Ex[i+diy] - Ex[i]));
+"""
+
+
+def type_to_C(float_type: numpy.float32 or numpy.float64) -> str:
+    """
+    Returns a string corresponding to the C equivalent of a numpy type.
+    Only works for float32 and float64.
+
+    :param float_type: numpy.float32 or numpy.float64
+    :return: string containing the corresponding C type (eg. 'double')
+    """
+    if float_type == numpy.float32:
+        arg_type = 'float'
+    elif float_type == numpy.float64:
+        arg_type = 'double'
+    else:
+        raise Exception('Unsupported type')
+    return arg_type
+
+
+

fdtd/boundary.py

@@ -0,0 +1,185 @@
+from typing import List, Dict
+import numpy
+
+
+def conductor(direction: int,
+              polarity: int,
+              ) -> List[str]:
+    """
+    Create source code for conducting boundary conditions.
+
+    :param direction: integer in range(3), corresponding to x,y,z.
+    :param polarity: -1 or 1, specifying eg. a -x or +x boundary.
+    :return: [E_source, H_source] source code for E and H boundary update steps.
+    """
+    if direction not in range(3):
+        raise Exception('Invalid direction: {}'.format(direction))
+
+    if polarity not in (-1, 1):
+        raise Exception('Invalid polarity: {}'.format(polarity))
+
+    r = 'xyz'[direction]
+    uv = 'xyz'.replace(r, '')
+
+    if polarity < 0:
+        bc_E = """
+if ({r} == 0) {{
+    E{r}[i] = 0;
+    E{u}[i] = E{u}[i+di{r}];
+    E{v}[i] = E{v}[i+di{r}];
+}}
+"""
+        bc_H = """
+if ({r} == 0) {{
+    H{r}[i] = H{r}[i+di{r}];
+    H{u}[i] = 0;
+    H{v}[i] = 0;
+}}
+"""
+
+    elif polarity > 0:
+        bc_E = """
+if ({r} == s{r} - 1) {{
+    E{r}[i] = -E{r}[i-2*di{r}];
+    E{u}[i] = +E{u}[i-di{r}];
+    E{v}[i] = +E{v}[i-di{r}];
+}} else if ({r} == s{r} - 2) {{
+    E{r}[i] = 0;
+}}
+"""
+        bc_H = """
+if ({r} == s{r} - 1) {{
+    H{r}[i] = +H{r}[i-di{r}];
+    H{u}[i] = -H{u}[i-2*di{r}];
+    H{v}[i] = -H{v}[i-2*di{r}];
+}} else if ({r} == s{r} - 2) {{
+    H{u}[i] = 0;
+    H{v}[i] = 0;
+}}
+"""
+    else:
+        raise Exception()
+
+    replacements = {'r': r, 'u': uv[0], 'v': uv[1]}
+    return [s.format(**replacements) for s in (bc_E, bc_H)]
+
+
+def cpml(direction: int,
+         polarity: int,
+         dt: float,
+         thickness: int=8,
+         epsilon_eff: float=1,
+         ) -> Dict:
+    """
+    Generate source code for complex phase matched layer (cpml) absorbing boundaries.
+    These are not full boundary conditions and require a conducting boundary to be added
+     in the same direction as the pml.
+
+    :param direction: integer in range(3), corresponding to x, y, z directions.
+    :param polarity: -1 or 1, corresponding to eg. -x or +x direction.
+    :param dt: timestep used by the simulation
+    :param thickness: Number of cells used by the pml (the grid is NOT expanded to add these cells). Default 8.
+    :param epsilon_eff: Effective epsilon_r of the pml layer. Default 1.
+    :return: Dict with entries 'E', 'H' (update equations for E and H) and 'psi_E', 'psi_H' (lists of str,
+            specifying the field names of the cpml fields used in the E and H update steps. Eg.,
+            Psi_xn_Ex for the complex Ex component for the x- pml.)
+    """
+    if direction not in range(3):
+        raise Exception('Invalid direction: {}'.format(direction))
+
+    if polarity not in (-1, 1):
+        raise Exception('Invalid polarity: {}'.format(polarity))
+
+    if thickness <= 2:
+        raise Exception('It would be wise to have a pml with 4+ cells of thickness')
+
+    if epsilon_eff <= 0:
+        raise Exception('epsilon_eff must be positive')
+
+    m = (3.5, 1)
+    sigma_max = 0.8 * (m[0] + 1) / numpy.sqrt(epsilon_eff)
+    alpha_max = 0  # TODO: Decide what to do about non-zero alpha
+    transverse = numpy.delete(range(3), direction)
+
+    r = 'xyz'[direction]
+    np = 'nVp'[numpy.sign(polarity)+1]
+    uv = ['xyz'[i] for i in transverse]
+
+    xE = numpy.arange(1, thickness+1, dtype=float)[::-1]
+    xH = numpy.arange(1, thickness+1, dtype=float)[::-1]
+    if polarity > 0:
+        xE -= 0.5
+    elif polarity < 0:
+        xH -= 0.5
+
+    def par(x):
+        sigma = ((x / thickness) ** m[0]) * sigma_max
+        alpha = ((1 - x / thickness) ** m[1]) * alpha_max
+        p0 = numpy.exp(-(sigma + alpha) * dt)
+        p1 = sigma / (sigma + alpha) * (p0 - 1)
+        return p0, p1
+    p0e, p1e = par(xE)
+    p0h, p1h = par(xH)
+
+    vals = {'r': r,
+            'u': uv[0],
+            'v': uv[1],
+            'np':   np,
+            'th':   thickness,
+            'p0e': ', '.join((str(x) for x in p0e)),
+            'p1e': ', '.join((str(x) for x in p1e)),
+            'p0h': ', '.join((str(x) for x in p0h)),
+            'p1h': ', '.join((str(x) for x in p1h)),
+            'se': '-+'[direction % 2],
+            'sh': '+-'[direction % 2]}
+
+    if polarity < 0:
+        bounds_if = """
+if ( 0 < {r} && {r} < {th} + 1 ) {{
+    const int ir = {r} - 1;                             // index into pml parameters
+    const int ip = {v} + {u} * s{v} + ir * s{v} * s{u};   // linear index into Psi
+"""
+    elif polarity > 0:
+        bounds_if = """
+if ( (s{r} - 1) > {r} && {r} > (s{r} - 1) - ({th} + 1) ) {{
+    const int ir = (s{r} - 1) - ({r} + 1);              // index into pml parameters
+    const int ip = {v} + {u} * s{v} + ir * s{v} * s{u};   // linear index into Psi
+"""
+    else:
+        raise Exception('Bad polarity (=0)')
+
+    code_E = """
+    // pml parameters:
+    const float p0[{th}] = {{ {p0e} }};
+    const float p1[{th}] = {{ {p1e} }};
+
+    Psi_{r}{np}_E{u}[ip] = p0[ir] * Psi_{r}{np}_E{u}[ip] + p1[ir] * (H{v}[i] - H{v}[i-di{r}]);
+    Psi_{r}{np}_E{v}[ip] = p0[ir] * Psi_{r}{np}_E{v}[ip] + p1[ir] * (H{u}[i] - H{u}[i-di{r}]);
+
+    E{u}[i] {se}= dt / eps{u}[i] * Psi_{r}{np}_E{u}[ip];
+    E{v}[i] {sh}= dt / eps{v}[i] * Psi_{r}{np}_E{v}[ip];
+}}
+"""
+    code_H = """
+    // pml parameters:
+    const float p0[{th}] = {{ {p0h} }};
+    const float p1[{th}] = {{ {p1h} }};
+
+    Psi_{r}{np}_H{u}[ip] = p0[ir] * Psi_{r}{np}_H{u}[ip] + p1[ir] * (E{v}[i+di{r}] - E{v}[i]);
+    Psi_{r}{np}_H{v}[ip] = p0[ir] * Psi_{r}{np}_H{v}[ip] + p1[ir] * (E{u}[i+di{r}] - E{u}[i]);
+
+    H{u}[i] {sh}= dt * Psi_{r}{np}_H{u}[ip];
+    H{v}[i] {se}= dt * Psi_{r}{np}_H{v}[ip];
+}}
+"""
+
+    pml_data = {
+        'E': (bounds_if + code_E).format(**vals),
+        'H': (bounds_if + code_H).format(**vals),
+        'psi_E': ['Psi_{r}{np}_E{u}'.format(**vals),
+                  'Psi_{r}{np}_E{v}'.format(**vals)],
+        'psi_H': ['Psi_{r}{np}_H{u}'.format(**vals),
+                  'Psi_{r}{np}_H{v}'.format(**vals)],
+    }
+
+    return pml_data

fdtd/simulation.py

@@ -0,0 +1,207 @@
+"""
+Class for constructing and holding the basic FDTD operations and fields
+"""
+
+from typing import List, Dict, Callable
+import numpy
+import warnings
+
+import pyopencl
+import pyopencl.array
+from pyopencl.elementwise import ElementwiseKernel
+
+from . import boundary, base
+from .base import type_to_C
+
+
+class Simulation(object):
+    """
+    Constructs and holds the basic FDTD operations and related fields
+    """
+    E = None    # type: List[pyopencl.array.Array]
+    H = None    # type: List[pyopencl.array.Array]
+    eps = None  # type: List[pyopencl.array.Array]
+    dt = None   # type: float
+
+    arg_type = None     # type: numpy.float32 or numpy.float64
+
+    context = None      # type: pyopencl.Context
+    queue = None        # type: pyopencl.CommandQueue
+
+    update_E = None     # type: Callable[[],pyopencl.Event]
+    update_H = None     # type: Callable[[],pyopencl.Event]
+
+    conductor_E = None  # type: Callable[[],pyopencl.Event]
+    conductor_H = None  # type: Callable[[],pyopencl.Event]
+
+    cpml_E = None       # type: Callable[[],pyopencl.Event]
+    cpml_H = None       # type: Callable[[],pyopencl.Event]
+
+    cpml_psi_E = None   # type: Dict[str, pyopencl.array.Array]
+    cpml_psi_H = None   # type: Dict[str, pyopencl.array.Array]
+
+    def __init__(self,
+                 epsilon: List[numpy.ndarray],
+                 dt: float=.99/numpy.sqrt(3),
+                 initial_E: List[numpy.ndarray]=None,
+                 initial_H: List[numpy.ndarray]=None,
+                 context: pyopencl.Context=None,
+                 queue: pyopencl.CommandQueue=None,
+                 float_type: numpy.float32 or numpy.float64=numpy.float32):
+        """
+        Initialize the simulation.
+
+        :param epsilon: List containing [eps_r,xx, eps_r,yy, eps_r,zz], where each element is a Yee-shifted ndarray
+                spanning the simulation domain. Relative epsilon is used.
+        :param dt: Time step. Default is the Courant factor.
+        :param initial_E: Initial E-field (default is 0 everywhere). Same format as epsilon.
+        :param initial_H: Initial H-field (default is 0 everywhere). Same format as epsilon.
+        :param context: pyOpenCL context. If not given, pyopencl.create_some_context(False) is called.
+        :param queue: pyOpenCL command queue. If not given, pyopencl.CommandQueue(context) is called.
+        :param float_type: numpy.float32 or numpy.float64. Default numpy.float32.
+        """
+
+        if len(epsilon) != 3:
+            Exception('Epsilon must be a list with length of 3')
+        if not all((e.shape == epsilon[0].shape for e in epsilon[1:])):
+            Exception('All epsilon grids must have the same shape. Shapes are {}', [e.shape for e in epsilon])
+
+        if context is None:
+            self.context = pyopencl.create_some_context(False)
+        else:
+            self.context = context
+
+        if queue is None:
+            self.queue = pyopencl.CommandQueue(self.context)
+        else:
+            self.queue = queue
+
+        if dt > .99/numpy.sqrt(3):
+            warnings.warn('Warning: unstable dt: {}'.format(dt))
+        elif dt <= 0:
+            raise Exception('Invalid dt: {}'.format(dt))
+        else:
+            self.dt = dt
+
+        self.arg_type = float_type
+
+        self.eps = [pyopencl.array.to_device(self.queue, e.astype(float_type)) for e in epsilon]
+
+        if initial_E is None:
+            self.E = [pyopencl.array.zeros_like(self.eps[0]) for _ in range(3)]
+        else:
+            if len(initial_E) != 3:
+                Exception('Initial_E must be a list of length 3')
+            if not all((E.shape == epsilon[0].shape for E in initial_E)):
+                Exception('Initial_E list elements must have same shape as epsilon elements')
+            self.E = [pyopencl.array.to_device(self.queue, E.astype(float_type)) for E in initial_E]
+
+        if initial_H is None:
+            self.H = [pyopencl.array.zeros_like(self.eps[0]) for _ in range(3)]
+        else:
+            if len(initial_H) != 3:
+                Exception('Initial_H must be a list of length 3')
+            if not all((H.shape == epsilon[0].shape for H in initial_H)):
+                Exception('Initial_H list elements must have same shape as epsilon elements')
+            self.H = [pyopencl.array.to_device(self.queue, H.astype(float_type)) for H in initial_H]
+
+        E_args = [type_to_C(self.arg_type) + ' *E' + c for c in 'xyz']
+        H_args = [type_to_C(self.arg_type) + ' *H' + c for c in 'xyz']
+        eps_args = [type_to_C(self.arg_type) + ' *eps' + c for c in 'xyz']
+        dt_arg = [type_to_C(self.arg_type) + ' dt']
+
+        sxyz = base.shape_source(epsilon[0].shape)
+        E_source = sxyz + base.dixyz_source + base.maxwell_E_source
+        H_source = sxyz + base.dixyz_source + base.maxwell_H_source
+
+        E_update = ElementwiseKernel(self.context, operation=E_source,
+                                     arguments=', '.join(E_args + H_args + dt_arg + eps_args))
+
+        H_update = ElementwiseKernel(self.context, operation=H_source,
+                                     arguments=', '.join(E_args + H_args + dt_arg))
+
+        self.update_E = lambda e: E_update(*self.E, *self.H, self.dt, *self.eps, wait_for=e)
+        self.update_H = lambda e: H_update(*self.E, *self.H, self.dt, wait_for=e)
+
+    def init_cpml(self, pml_args: List[Dict]):
+        """
+        Initialize absorbing layers (cpml: complex phase matched layer). PMLs are not actual
+         boundary conditions, so you should add a conducting boundary (.init_conductors()) for
+         all directions in which you add PMLs.
+        Allows use of self.cpml_E(events) and self.cpml_H(events).
+        All necessary additional fields are created on the opencl device.
+
+        :param pml_args: A list containing dictionaries which are passed to .boundary.cpml(...).
+            The dt argument is set automatically, but the others must be passed in each entry
+             of pml_args.
+        """
+        sxyz = base.shape_source(self.eps[0].shape)
+
+        # Prepare per-iteration constants for later use
+        pml_E_source = sxyz + base.dixyz_source + base.xyz_source
+        pml_H_source = sxyz + base.dixyz_source + base.xyz_source
+
+        psi_E = []
+        psi_H = []
+        psi_E_names = []
+        psi_H_names = []
+        for arg_set in pml_args:
+            pml_data = boundary.cpml(dt=self.dt, **arg_set)
+
+            pml_E_source += pml_data['E']
+            pml_H_source += pml_data['H']
+
+            ti = numpy.delete(range(3), arg_set['direction'])
+            trans = [self.eps[0].shape[i] for i in ti]
+            psi_shape = (8, trans[0], trans[1])
+
+            psi_E += [pyopencl.array.zeros(self.queue, psi_shape, dtype=self.arg_type)
+                      for _ in pml_data['psi_E']]
+            psi_H += [pyopencl.array.zeros(self.queue, psi_shape, dtype=self.arg_type)
+                      for _ in pml_data['psi_H']]
+
+            psi_E_names += pml_data['psi_E']
+            psi_H_names += pml_data['psi_H']
+
+        E_args = [type_to_C(self.arg_type) + ' *E' + c for c in 'xyz']
+        H_args = [type_to_C(self.arg_type) + ' *H' + c for c in 'xyz']
+        eps_args = [type_to_C(self.arg_type) + ' *eps' + c for c in 'xyz']
+        dt_arg = [type_to_C(self.arg_type) + ' dt']
+        arglist_E = [type_to_C(self.arg_type) + ' *' + psi for psi in psi_E_names]
+        arglist_H = [type_to_C(self.arg_type) + ' *' + psi for psi in psi_H_names]
+        pml_E_args = ', '.join(E_args + H_args + dt_arg + eps_args + arglist_E)
+        pml_H_args = ', '.join(E_args + H_args + dt_arg + arglist_H)
+
+        pml_E = ElementwiseKernel(self.context, arguments=pml_E_args, operation=pml_E_source)
+        pml_H = ElementwiseKernel(self.context, arguments=pml_H_args, operation=pml_H_source)
+
+        self.cpml_E = lambda e: pml_E(*self.E, *self.H, self.dt, *self.eps, *psi_E, wait_for=e)
+        self.cpml_H = lambda e: pml_H(*self.E, *self.H, self.dt, *psi_H, wait_for=e)
+        self.cmpl_psi_E = {k: v for k, v in zip(psi_E_names, psi_E)}
+        self.cmpl_psi_H = {k: v for k, v in zip(psi_H_names, psi_H)}
+
+    def init_conductors(self, conductor_args: List[Dict]):
+        """
+        Initialize reflecting boundary conditions.
+        Allows use of self.conductor_E(events) and self.conductor_H(events).
+
+        :param conductor_args: List of dictionaries with which to call .boundary.conductor(...).
+        """
+
+        sxyz = base.shape_source(self.eps[0].shape)
+
+        # Prepare per-iteration constants
+        bc_E_source = sxyz + base.dixyz_source + base.xyz_source
+        bc_H_source = sxyz + base.dixyz_source + base.xyz_source
+        for arg_set in conductor_args:
+            [e, h] = boundary.conductor(**arg_set)
+            bc_E_source += e
+            bc_H_source += h
+
+        E_args = [type_to_C(self.arg_type) + ' *E' + c for c in 'xyz']
+        H_args = [type_to_C(self.arg_type) + ' *H' + c for c in 'xyz']
+        bc_E = ElementwiseKernel(self.context, arguments=E_args, operation=bc_E_source)
+        bc_H = ElementwiseKernel(self.context, arguments=H_args, operation=bc_H_source)
+
+        self.conductor_E = lambda e: bc_E(*self.E, wait_for=e)
+        self.conductor_H = lambda e: bc_H(*self.H, wait_for=e)