fdtdx 0.1.0__py3-none-any.whl

fdtdx/__init__.py

@@ -0,0 +1,37 @@
+"""FDTDX: A JAX-based Finite-Difference Time-Domain (FDTD) Electromagnetic Simulation Package.
+
+A high-performance framework for electromagnetic simulations and inverse design optimization
+using the FDTD method, implemented in JAX for GPU acceleration and automatic differentiation.
+
+Key Features:
+    - Memory-efficient automatic differentiation using time-reversibility
+    - Multi-GPU support for large-scale simulations
+    - Intuitive object positioning with relative coordinates
+    - Built-in optimization capabilities for inverse design
+    - Comprehensive boundary condition implementations (e.g., PML)
+    - Field detectors and visualization tools
+    - JIT-compiled operations for maximum performance
+
+Example:
+    Basic simulation setup:
+    ```python
+    import fdtdx
+    from fdtdx.objects import Source, Detector
+
+    # Create simulation components
+    source = Source(...)
+    detector = Detector(...)
+
+    # Run simulation
+    fields = fdtdx.simulate(source, detector)
+    ```
+
+Notes:
+    The package is designed for both research and industrial applications in nanophotonics,
+    metamaterials, and photonic integrated circuits. It supports both 2D and 3D simulations
+    with automatic gradient computation for optimization tasks.
+
+References:
+    - Paper: https://github.com/ymahlau/fdtdx
+    - Documentation: https://github.com/ymahlau/fdtdx/docs
+"""

fdtdx/constraints/__init__.py

@@ -0,0 +1,14 @@
+"""Constraint modules for enforcing physical and fabrication requirements.
+
+This package provides constraint modules that can be applied during optimization
+to ensure designs meet physical realizability and fabrication requirements.
+Key constraints include:
+
+- Minimum feature size constraints
+- Binary material constraints (air/dielectric only)
+- Connectivity constraints (no floating material)
+- Fabrication constraints (no trapped air, etc.)
+
+The constraints are implemented as modules that can be chained together
+and applied during the optimization process.
+"""

fdtdx/constraints/binary_transform.py

@@ -0,0 +1,407 @@
+import jax
+import jax.numpy as jnp
+
+from fdtdx.core.misc import PaddingConfig, advanced_padding
+
+
+def remove_floating_polymer(
+    matrix: jax.Array,  # 1 is polymer, zero is air, shape (x, y, z)
+) -> jax.Array:
+    """Removes polymer regions that are not connected to the substrate.
+
+    Uses flood-fill algorithm to identify polymer regions connected to the bottom layer
+    and removes any floating polymer regions that are not connected.
+
+    Args:
+        matrix: Binary array where 1 represents polymer and 0 represents air.
+               Shape is (x, y, z) representing the 3D grid.
+
+    Returns:
+        Modified binary array with floating polymer regions removed.
+    """
+    connected = compute_polymer_connection(matrix)
+    non_connected_polymer = jnp.invert(connected) & matrix
+    matrix = matrix & jnp.invert(non_connected_polymer)
+    return matrix
+
+
+def remove_polymer_non_connected_to_x_max_middle(
+    matrix: jax.Array,  # 1 is polymer, zero is air, shape (x, y, z)
+) -> jax.Array:
+    """Removes polymer regions not connected to the middle of the x-max boundary.
+
+    Uses flood-fill algorithm starting from the middle point of the x-max boundary
+    to identify connected polymer regions and removes unconnected regions.
+
+    Args:
+        matrix: Binary array where 1 represents polymer and 0 represents air.
+               Shape is (x, y, z) representing the 3D grid.
+
+    Returns:
+        Modified binary array with unconnected polymer regions removed.
+    """
+    y_middle = round(matrix.shape[1] / 2)
+    x_middle = round(matrix.shape[0] / 2)
+    connected = compute_polymer_connection(
+        matrix,
+        connected_slice=(x_middle, y_middle, None),
+    )
+    non_connected_polymer = jnp.invert(connected) & matrix
+    matrix = matrix & jnp.invert(non_connected_polymer)
+    return matrix
+
+
+def connect_holes_and_structures(
+    matrix: jax.Array,  # 1 is polymer, zero is air, shape (x, y, z)
+) -> jax.Array:
+    """Connects disconnected polymer structures and air holes in the matrix.
+
+    Performs a two-pass algorithm:
+    1. Connects disconnected polymer structures by adding polymer material
+    2. Connects disconnected air regions by removing polymer material
+
+    This ensures both structural integrity and proper air ventilation.
+
+    Args:
+        matrix: Binary array where 1 represents polymer and 0 represents air.
+               Shape is (x, y, z) representing the 3D grid.
+
+    Returns:
+        Modified binary array with connected structures and holes.
+    """
+    matrix = matrix.astype(bool)
+    # first connect all polymer structures  # TODO: convert to fori-loops
+    for i in range(matrix.shape[-1] - 1):
+        connected = compute_polymer_connection(matrix)
+        if i == 0:
+            lower_slice = jnp.ones_like(matrix[..., 0])
+        else:
+            lower_slice = matrix[..., i - 1]
+        new_middle, new_upper = connect_slice(
+            lower_slice=lower_slice,
+            middle_slice=matrix[..., i],
+            upper_slice=matrix[..., i + 1],
+            upper_save_points=connected[..., i + 1],
+        )
+        matrix = matrix.at[..., i].set(new_middle)
+        matrix = matrix.at[..., i + 1].set(new_upper)
+    # then connect all air
+    for i in range(matrix.shape[-1], 0, -1):
+        air_connected = compute_air_connection(matrix)
+        if i == matrix.shape[-1]:
+            lower_slice = jnp.ones_like(matrix[..., 0])
+        else:
+            lower_slice = jnp.invert(matrix[..., i + 1])
+        new_middle, new_upper = connect_slice(
+            lower_slice=lower_slice,
+            middle_slice=jnp.invert(matrix[..., i]),
+            upper_slice=jnp.invert(matrix[..., i - 1]),
+            upper_save_points=air_connected[..., i - 1],
+        )
+        matrix = matrix.at[..., i].set(jnp.invert(new_middle))
+        matrix = matrix.at[..., i - 1].set(jnp.invert(new_upper))
+    # throw away all polymer that got free floating due to air connection in last step
+    connected = compute_polymer_connection(matrix)
+    non_connected_polymer = jnp.invert(connected) & matrix
+    matrix = matrix & jnp.invert(non_connected_polymer)
+    return matrix
+
+
+def compute_air_connection(matrix: jax.Array) -> jax.Array:
+    """Computes a mask of air regions connected to the boundaries.
+
+    Uses iterative dilation to identify air regions (zeros in the matrix)
+    that are connected to the simulation boundaries. This is used to ensure
+    proper ventilation in the structure.
+
+    Args:
+        matrix: Binary array where 1 represents polymer and 0 represents air.
+
+    Returns:
+        Boolean array marking air regions connected to boundaries.
+    """
+    inv_matrix = jnp.invert(matrix)
+    n = max([matrix.shape[0], matrix.shape[1], matrix.shape[2]])
+    n4_kernel = jnp.asarray(
+        [
+            [0, 1, 0],
+            [1, 1, 1],
+            [0, 1, 0],
+        ],
+        dtype=bool,
+    )
+    connected = jnp.zeros_like(inv_matrix, dtype=bool)
+    connected = connected.at[:, :, -1].set(True)
+    connected = connected.at[0, :, :].set(True)
+    connected = connected.at[-1, :, :].set(True)
+    connected = connected.at[:, 0, :].set(True)
+    connected = connected.at[:, -1, :].set(True)
+    connected = connected & inv_matrix
+
+    def _body_fn(_, arr):
+        arr = seperated_3d_dilation(
+            arr_3d=arr,
+            kernel_xy=n4_kernel,
+            kernel_xz=n4_kernel,
+            kernel_yz=n4_kernel,
+            reduction_arr=inv_matrix,
+        )
+        return arr
+
+    connected = jax.lax.fori_loop(0, n, _body_fn, connected)
+
+    return connected
+
+
+def compute_polymer_connection(
+    matrix: jax.Array,
+    connected_slice: tuple[int | None, int | None, int | None] | None = None,
+) -> jax.Array:
+    """Computes a mask of polymer regions connected to specified points.
+
+    Uses iterative dilation to identify polymer regions (ones in the matrix)
+    that are connected either to the bottom layer or to specified points
+    in connected_slice.
+
+    Args:
+        matrix: Binary array where 1 represents polymer and 0 represents air.
+        connected_slice: Optional tuple of indices specifying starting points
+                       for the connection computation. If None, uses bottom layer.
+
+    Returns:
+        Boolean array marking connected polymer regions.
+    """
+    n = max([matrix.shape[0], matrix.shape[1], matrix.shape[2]])
+    padded = False
+    if matrix.shape[2] == 1:
+        padded = True
+        matrix = jnp.pad(matrix, pad_width=((0, 0), (0, 0), (1, 1)))
+    n4_kernel = jnp.asarray(
+        [
+            [0, 1, 0],
+            [1, 1, 1],
+            [0, 1, 0],
+        ],
+        dtype=bool,
+    )
+    connected = jnp.zeros_like(matrix, dtype=bool)
+    if connected_slice is None:
+        connected = connected.at[..., 0].set(True)
+    else:
+        connected = connected.at[connected_slice].set(True)
+
+    def _body_fn(_, arr):
+        arr = seperated_3d_dilation(
+            arr_3d=arr,
+            kernel_xy=n4_kernel,
+            kernel_xz=n4_kernel,
+            kernel_yz=n4_kernel,
+            reduction_arr=matrix,
+        )
+        return arr
+
+    connected = jax.lax.fori_loop(0, n, _body_fn, connected)
+
+    if padded:
+        connected = connected[..., 1:2]
+    return connected
+
+
+def connect_slice(
+    lower_slice: jax.Array,
+    middle_slice: jax.Array,  # this layer is optimized
+    upper_slice: jax.Array,  # this layer is optimized,
+    upper_save_points: jax.Array,  # points which are surely connected
+) -> tuple[jax.Array, jax.Array]:
+    """Connects polymer regions between adjacent slices.
+
+    Attempts to connect disconnected polymer regions between three adjacent
+    z-slices by modifying the middle and upper slices while preserving
+    specified connection points.
+
+    Args:
+        lower_slice: Binary array representing the lower z-slice.
+        middle_slice: Binary array representing the middle z-slice to be optimized.
+        upper_slice: Binary array representing the upper z-slice to be optimized.
+        upper_save_points: Boolean mask of points in upper slice that must remain connected.
+
+    Returns:
+        Tuple of (modified_middle_slice, modified_upper_slice) with connected regions.
+    """
+    n = max(lower_slice.shape[0], lower_slice.shape[1])
+
+    # define kernels
+    n4_kernel = jnp.asarray(
+        [
+            [0, 1, 0],
+            [1, 1, 1],
+            [0, 1, 0],
+        ],
+        dtype=bool,
+    )
+    n8_kernel = jnp.asarray(
+        [
+            [1, 1, 1],
+            [1, 1, 1],
+            [1, 1, 1],
+        ],
+        dtype=bool,
+    )
+    direction_kernels = jnp.asarray(
+        [
+            [[0, 1, 0], [0, 0, 0], [0, 0, 0]],
+            [[0, 0, 0], [0, 0, 1], [0, 0, 0]],
+            [[0, 0, 0], [0, 0, 0], [0, 1, 0]],
+            [[0, 0, 0], [1, 0, 0], [0, 0, 0]],
+        ],
+        dtype=bool,
+    )
+
+    # first determine all well-supported points
+    connected_points = upper_slice & middle_slice
+    connected_points = connected_points | upper_save_points
+
+    for _ in range(n):
+        connected_points = dilate_jax(connected_points, n4_kernel)
+        connected_points = connected_points & upper_slice
+    upper_air = jnp.invert(upper_slice)
+    non_connected_points = jnp.invert(upper_air | connected_points)
+
+    # first try to connect by setting additional polymers in lower array
+    connectable_region_lower = dilate_jax(middle_slice, n4_kernel)
+    connectable_region_lower = connectable_region_lower | lower_slice  # connectable either vertical or horiz.
+    connectable_by_lower = non_connected_points & connectable_region_lower
+    # update matrix, non-connected and connected points
+    middle_slice = middle_slice | connectable_by_lower
+    connected_points = connected_points | connectable_by_lower
+    for _ in range(n):
+        connected_points = dilate_jax(connected_points, n4_kernel)
+        connected_points = connected_points & upper_slice
+    non_connected_points = jnp.invert(upper_air | connected_points)
+
+    # then try to connect by adding polymer in upper array
+    connectable_region_upper = dilate_jax(connected_points, n8_kernel)
+    connectable_by_upper = non_connected_points & connectable_region_upper
+    connection_points = jax.vmap(dilate_jax, in_axes=(None, 0))(connectable_by_upper, direction_kernels)
+    valid_connection_points = connectable_region_upper & connection_points
+    valid_connection_points = jnp.sum(valid_connection_points, axis=0).astype(bool)
+    # update matrix, non-connected and connected points
+    upper_slice = upper_slice | valid_connection_points
+    for _ in range(n):
+        connected_points = dilate_jax(connected_points, n4_kernel)
+        connected_points = connected_points & upper_slice
+    non_connected_points = jnp.invert(upper_air | connected_points)
+
+    # delete all non-connected
+    upper_slice = upper_slice & jnp.invert(non_connected_points)
+
+    return middle_slice, upper_slice
+
+
+def dilate_jax(image: jax.Array, kernel: jax.Array) -> jax.Array:
+    """Performs morphological dilation on a binary image using JAX.
+
+    Args:
+        image: Binary input array to be dilated.
+        kernel: Binary kernel defining the dilation shape.
+
+    Returns:
+        Dilated binary array.
+    """
+    conv = jax.scipy.signal.convolve2d(image, kernel, mode="same", boundary="fill")
+    binary_arr = jnp.asarray(conv, dtype=bool)
+    return binary_arr
+
+
+def erode_jax(image: jax.Array, kernel: jax.Array) -> jax.Array:
+    """Performs morphological erosion on a binary image using JAX.
+
+    Args:
+        image: Binary input array to be eroded.
+        kernel: Binary kernel defining the erosion shape.
+
+    Returns:
+        Eroded binary array.
+    """
+    conv = jax.scipy.signal.convolve2d(~image, kernel, mode="same", boundary="fill")
+    binary_arr = jnp.asarray(conv, dtype=bool)
+    return ~binary_arr
+
+
+def seperated_3d_dilation(
+    arr_3d: jax.Array,
+    kernel_xy: jax.Array,
+    kernel_yz: jax.Array,
+    kernel_xz: jax.Array,
+    reduction_arr: jax.Array,
+) -> jax.Array:
+    """Performs separated 3D dilation along each axis plane.
+
+    Applies 2D dilation kernels separately in the XY, YZ, and XZ planes
+    to approximate 3D dilation while being more computationally efficient.
+    The result is masked by the reduction array after each operation.
+
+    Args:
+        arr_3d: 3D binary array to be dilated.
+        kernel_xy: 2D kernel for XY plane dilation.
+        kernel_yz: 2D kernel for YZ plane dilation.
+        kernel_xz: 2D kernel for XZ plane dilation.
+        reduction_arr: Binary mask to constrain dilation results.
+
+    Returns:
+        Dilated 3D binary array.
+    """
+
+    def convolve_partial(image: jax.Array, kernel: jax.Array):
+        return jax.scipy.signal.convolve2d(image, kernel, mode="same", boundary="fill")
+
+    arr_3d = jax.vmap(convolve_partial, in_axes=(2, None), out_axes=(2))(arr_3d, kernel_xy)
+    arr_3d = jnp.asarray(arr_3d, dtype=bool)
+    arr_3d = arr_3d & reduction_arr
+
+    arr_3d = jax.vmap(convolve_partial, in_axes=(1, None), out_axes=(1))(arr_3d, kernel_xz)
+    arr_3d = jnp.asarray(arr_3d, dtype=bool)
+    arr_3d = arr_3d & reduction_arr
+
+    arr_3d = jax.vmap(convolve_partial, in_axes=(0, None), out_axes=(0))(arr_3d, kernel_yz)
+    arr_3d = jnp.asarray(arr_3d, dtype=bool)
+    arr_3d = arr_3d & reduction_arr
+
+    return arr_3d
+
+
+def binary_median_filter(
+    arr_3d: jax.Array,
+    kernel_sizes: tuple[int, int, int],
+    padding_cfg: PaddingConfig,
+) -> jax.Array:
+    """Applies a binary median filter to a 3D array.
+
+    Implements a median filter for binary data by using convolution
+    and thresholding. The filter is applied separately along each axis
+    using the specified kernel sizes.
+
+    Args:
+        arr_3d: 3D binary input array.
+        kernel_sizes: Tuple of (kx, ky, kz) specifying filter size in each dimension.
+        padding_cfg: Configuration for padding the input array.
+
+    Returns:
+        Filtered binary array.
+    """
+    # padding
+    padded_arr, orig_slice = advanced_padding(arr_3d, padding_cfg)
+    padded_arr = padded_arr.astype(jnp.float32)
+    # create filter kernel and convolve (avg filter)
+    for axis, k_size in enumerate(kernel_sizes):
+        shape_list = [1, 1, 1]
+        shape_list[axis] = k_size
+        shape_tpl = tuple(shape_list)
+        kernel = jnp.ones(shape=shape_tpl, dtype=jnp.float32)
+        padded_arr = jax.scipy.signal.convolve(padded_arr, kernel, mode="same", method="direct")
+    # discretize again to get median
+    kernel_sum = jnp.prod(jnp.asarray(kernel_sizes))
+    result = padded_arr / kernel_sum
+    result = result[orig_slice]
+    result = jnp.round(result).astype(arr_3d.dtype)
+    return result