fdtdx 0.1.0__tar.gz

fdtdx-0.1.0/.editorconfig

@@ -0,0 +1,6 @@
+[*]
+max_line_length = 120
+
+[*.json]
+indent_style = space
+indent_size = 4

fdtdx-0.1.0/.github/actions/setup-python-env/action.yml

@@ -0,0 +1,30 @@
+name: "Setup Python Environment"
+description: "Set up Python environment for the given Python version"
+
+inputs:
+  python-version:
+    description: "Python version to use"
+    required: true
+    default: "3.12"
+  uv-version:
+    description: "uv version to use"
+    required: true
+    default: "0.5.8"
+
+runs:
+  using: "composite"
+  steps:
+    - uses: actions/setup-python@v5
+      with:
+        python-version: ${{ inputs.python-version }}
+
+    - name: Install uv
+      uses: astral-sh/setup-uv@v2
+      with:
+        version: ${{ inputs.uv-version }}
+        enable-cache: "true"
+        cache-suffix: ${{ matrix.python-version }}
+
+    - name: Install Python dependencies
+      run: uv sync --frozen
+      shell: bash

fdtdx-0.1.0/.github/workflows/main.yml

@@ -0,0 +1,71 @@
+name: Main
+
+on:
+  pull_request:
+    types: [opened, synchronize, reopened]
+  push:
+    branches: [main]
+
+jobs:
+  checks:
+    runs-on: ubuntu-latest
+    steps:
+      - name: Check out
+        uses: actions/checkout@v4
+        with:
+          persist-credentials: false
+
+      - uses: actions/cache@v4
+        with:
+          path: ~/.cache/pre-commit
+          key: pre-commit-${{ hashFiles('.pre-commit-config.yaml') }}
+
+      - name: Set up the environment
+        uses: ./.github/actions/setup-python-env
+
+      - name: Run pre-commit
+        run: uv run pre-commit run -a --show-diff-on-failure
+
+      - name: Check lock file consistency
+        run: uv sync --locked
+
+  tests:
+    runs-on: ubuntu-latest
+    strategy:
+      matrix:
+        python-version: ["3.11", "3.12"]
+      fail-fast: false
+    defaults:
+      run:
+        shell: bash
+    steps:
+      - name: Check out
+        uses: actions/checkout@v4
+        with:
+          persist-credentials: false
+
+      - name: Set up the environment
+        uses: ./.github/actions/setup-python-env
+        with:
+          python-version: ${{ matrix.python-version }}
+
+      - name: Run tests
+        run: uv run python -m pytest tests --cov --cov-config=pyproject.toml --cov-report=xml
+
+      - name: Upload coverage reports to Codecov with GitHub Action on Python 3.12 for Ubuntu
+        uses: codecov/codecov-action@v3
+        if: ${{ matrix.python-version == '3.12' }}
+
+  deploy-docs:
+    runs-on: ubuntu-latest
+    steps:
+      - name: Check out
+        uses: actions/checkout@v4
+        with:
+          persist-credentials: false
+
+      - name: Set up the environment
+        uses: ./.github/actions/setup-python-env
+
+      - name: Deploy documentation
+        run: uv run mkdocs gh-deploy --force

fdtdx-0.1.0/.gitignore

@@ -0,0 +1,107 @@
+multirun
+outputs
+*.png
+tmp
+slurm-*
+.jax_cache
+.vscode
+.pytest_cache
+.ruff_cache
+planos
+scripts/configs/run_config.yaml
+.ragatouille
+*.jsonl
+.aider*
+.DS_Store
+site
+
+# Byte-compiled / optimized / DLL files
+__pycache__/
+*.py[cod]
+*$py.class
+
+# C extensions
+*.so
+
+# Distribution / packaging
+.Python
+env/
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+*.egg-info/
+.installed.cfg
+*.egg
+
+# PyInstaller
+#  Usually these files are written by a python script from a template
+#  before PyInstaller builds the exe, so as to inject date/other infos into it.
+*.manifest
+*.spec
+
+# Installer logs
+pip-log.txt
+pip-delete-this-directory.txt
+
+# Unit test / coverage reports
+htmlcov/
+.tox/
+.coverage
+.coverage.*
+.cache
+nosetests.xml
+coverage.xml
+*,cover
+.hypothesis/
+
+# Translations
+*.mo
+*.pot
+
+# Django stuff:
+*.log
+local_settings.py
+
+# Flask stuff:
+instance/
+.webassets-cache
+
+# Scrapy stuff:
+.scrapy
+
+# Sphinx documentation
+docs/_build/
+
+# PyBuilder
+target/
+
+# IPython Notebook
+.ipynb_checkpoints
+
+# pyenv
+.python-version
+
+# celery beat schedule file
+celerybeat-schedule
+
+# dotenv
+.env
+
+# virtualenv
+.venv/
+venv/
+ENV/
+
+# Spyder project settings
+.spyderproject
+
+# Rope project settings
+.ropeproject

fdtdx-0.1.0/.pre-commit-config.yaml

@@ -0,0 +1,35 @@
+# See https://pre-commit.com for more information
+# See https://pre-commit.com/hooks.html for more hooks
+repos:
+  - repo: https://github.com/pre-commit/pre-commit-hooks
+    rev: v5.0.0
+    hooks:
+      - id: check-yaml
+      - id: check-added-large-files
+      - id: check-shebang-scripts-are-executable
+      - id: check-toml
+      - id: check-merge-conflict
+        args: [--assume-in-merge]
+  - repo: https://github.com/astral-sh/ruff-pre-commit
+    # Ruff version.
+    rev: v0.8.2
+    hooks:
+      # Run the linter.
+      - id: ruff
+        args:
+          - --fix
+      # Run the formatter.
+      - id: ruff-format
+  - repo: local
+    hooks:
+      - id: mkdocs-build
+        name: mkdocs-build
+        entry: mkdocs build
+        language: system
+        pass_filenames: false # Or true if you want to pass files as arguments
+        always_run: true # Usually you want to build every time
+        stages: [pre-commit]
+  - repo: https://github.com/woodruffw/zizmor-pre-commit
+    rev: v0.9.2
+    hooks:
+      - id: zizmor

fdtdx-0.1.0/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2024 Florian Maas
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

fdtdx-0.1.0/PKG-INFO

@@ -0,0 +1,86 @@
+Metadata-Version: 2.4
+Name: fdtdx
+Version: 0.1.0
+Summary: A JAX framework for the finite-difference time-domain (FDTD) method for solving Maxwell's equations with a focus on inverse design of photonic devices.
+Project-URL: Repository, https://github.com/ymahlau/fdtdx
+Project-URL: Documentation, https://ymahlau.github.io/fdtdx/
+Author-email: Yannik Mahlau <mahlau@tnt.uni-hannover.de>, Frederik Schubert <schubert@tnt.uni-hannover.de>
+License-File: LICENSE
+Requires-Python: <3.13,>=3.11
+Requires-Dist: equinox>=0.11.10
+Requires-Dist: jax<0.4.36; platform_system != 'Linux'
+Requires-Dist: jax[cuda12]<0.4.36; platform_system == 'Linux'
+Requires-Dist: loguru>=0.7.3
+Requires-Dist: matfree>=0.1.1
+Requires-Dist: matplotlib>=3.9.4
+Requires-Dist: moviepy>=2.1.1
+Requires-Dist: optax>=0.2.4
+Requires-Dist: pytreeclass>=0.9.2
+Requires-Dist: rich>=13.9.4
+Requires-Dist: seaborn>=0.13.2
+Requires-Dist: sharedarray>=3.2.4
+Requires-Dist: tidy3d>=2.7.8
+Requires-Dist: trimesh>=4.5.3
+Requires-Dist: zstandard>=0.23.0
+Description-Content-Type: text/markdown
+
+# FDTDX
+
+FDTDX is an efficient open-source Python package for the inverse design of three-dimensional photonic nanostructures using the Finite-Difference Time-Domain (FDTD) method. Built on JAX, it provides native GPU support and automatic differentiation capabilities, making it ideal for large-scale 3D design in nanophotonics.
+
+## Key Features
+
+- **High Performance**: GPU-accelerated FDTD simulations with multi-GPU scaling capabilities
+- **Memory Efficient**: Leverages time-reversibility in Maxwell's equations for efficient gradient computation
+- **Automatic Differentiation**: Built-in gradient-based optimization for complex 3D structures
+- **User-Friendly API**: Intuitive positioning and sizing of objects in absolute or relative coordinates
+- **Large-Scale Design**: Capable of handling simulations with billions of grid cells
+- **Open Source**: Freely available for research and development
+
+## Documentation
+
+Visit our [documentation](https://ymahlau.github.io/fdtdx) for:
+- Detailed API reference
+- Tutorial guides
+- Example simulations
+- Best practices
+
+
+## Installation
+
+Install FDTDX using pip:
+
+```bash
+pip install fdtdx
+```
+
+For development installation, clone the repository and install in editable mode:
+
+```bash
+git clone https://github.com/ymahlau/fdtdx
+cd fdtdx
+pip install -e .
+```
+
+## Multi-GPU
+
+```bash
+# The following lines often lead to better memory usage in JAX
+# when using multiple GPU.
+export XLA_PYTHON_CLIENT_ALLOCATOR=platform
+export XLA_PYTHON_CLIENT_PREALLOCATE=false
+export NCCL_LL128_BUFFSIZE="-2"
+export NCCL_LL_BUFFSIZE="-2"
+export NCCL_PROTO="SIMPLE,LL,LL128"
+```
+
+## Citation
+If you find this repository helpful for you work, please consider citing:
+```
+@article{schubert2024quantized,
+  title={Quantized Inverse Design for Photonic Integrated Circuits},
+  author={Schubert, Frederik and Mahlau, Yannik  and Bethmann, Konrad and Hartmann, Fabian and Caspary, Reinhard and Munderloh, Marco and Ostermann, J{\"o}rn and Rosenhahn, Bodo},
+  journal={arXiv preprint arXiv:2407.10273},
+  year={2024}
+}
+```

fdtdx-0.1.0/README.md

@@ -0,0 +1,60 @@
+# FDTDX
+
+FDTDX is an efficient open-source Python package for the inverse design of three-dimensional photonic nanostructures using the Finite-Difference Time-Domain (FDTD) method. Built on JAX, it provides native GPU support and automatic differentiation capabilities, making it ideal for large-scale 3D design in nanophotonics.
+
+## Key Features
+
+- **High Performance**: GPU-accelerated FDTD simulations with multi-GPU scaling capabilities
+- **Memory Efficient**: Leverages time-reversibility in Maxwell's equations for efficient gradient computation
+- **Automatic Differentiation**: Built-in gradient-based optimization for complex 3D structures
+- **User-Friendly API**: Intuitive positioning and sizing of objects in absolute or relative coordinates
+- **Large-Scale Design**: Capable of handling simulations with billions of grid cells
+- **Open Source**: Freely available for research and development
+
+## Documentation
+
+Visit our [documentation](https://ymahlau.github.io/fdtdx) for:
+- Detailed API reference
+- Tutorial guides
+- Example simulations
+- Best practices
+
+
+## Installation
+
+Install FDTDX using pip:
+
+```bash
+pip install fdtdx
+```
+
+For development installation, clone the repository and install in editable mode:
+
+```bash
+git clone https://github.com/ymahlau/fdtdx
+cd fdtdx
+pip install -e .
+```
+
+## Multi-GPU
+
+```bash
+# The following lines often lead to better memory usage in JAX
+# when using multiple GPU.
+export XLA_PYTHON_CLIENT_ALLOCATOR=platform
+export XLA_PYTHON_CLIENT_PREALLOCATE=false
+export NCCL_LL128_BUFFSIZE="-2"
+export NCCL_LL_BUFFSIZE="-2"
+export NCCL_PROTO="SIMPLE,LL,LL128"
+```
+
+## Citation
+If you find this repository helpful for you work, please consider citing:
+```
+@article{schubert2024quantized,
+  title={Quantized Inverse Design for Photonic Integrated Circuits},
+  author={Schubert, Frederik and Mahlau, Yannik  and Bethmann, Konrad and Hartmann, Fabian and Caspary, Reinhard and Munderloh, Marco and Ostermann, J{\"o}rn and Rosenhahn, Bodo},
+  journal={arXiv preprint arXiv:2407.10273},
+  year={2024}
+}
+```

fdtdx-0.1.0/docs/api/constraints/binary.md

@@ -0,0 +1,17 @@
+# Binary Operations
+
+## Morphological Operations
+::: fdtdx.constraints.binary_transform.dilate_jax
+::: fdtdx.constraints.binary_transform.erode_jax
+::: fdtdx.constraints.binary_transform.seperated_3d_dilation
+::: fdtdx.constraints.binary_transform.binary_median_filter
+
+## Connectivity Analysis
+::: fdtdx.constraints.binary_transform.compute_polymer_connection
+::: fdtdx.constraints.binary_transform.compute_air_connection
+::: fdtdx.constraints.binary_transform.connect_slice
+
+## Structure Optimization
+::: fdtdx.constraints.binary_transform.remove_floating_polymer
+::: fdtdx.constraints.binary_transform.remove_polymer_non_connected_to_x_max_middle
+::: fdtdx.constraints.binary_transform.connect_holes_and_structures

fdtdx-0.1.0/docs/api/constraints/discrete.md

@@ -0,0 +1,17 @@
+# Discrete Constraints
+
+## Material Connectivity
+::: fdtdx.constraints.discrete.RemoveFloatingMaterial
+::: fdtdx.constraints.discrete.ConnectHolesAndStructures
+
+## Feature Size Control
+::: fdtdx.constraints.discrete.circular_brush
+::: fdtdx.constraints.discrete.BrushConstraint2D
+
+## Smoothing and Filtering
+::: fdtdx.constraints.discrete.BinaryMedianFilterModule
+::: fdtdx.constraints.discrete.BOTTOM_Z_PADDING_CONFIG
+::: fdtdx.constraints.discrete.BOTTOM_Z_PADDING_CONFIG_REPEAT
+
+## Pillar Constraints
+::: fdtdx.constraints.pillars.PillarMapping

fdtdx-0.1.0/docs/api/constraints/index.md

@@ -0,0 +1,11 @@
+# Constraints
+
+The constraints package provides tools for applying physical and geometric constraints during optimization.
+
+## Components
+
+- [Mapping](mapping.md) - Constraint mapping functionality
+- [Modules](modules.md) - Base classes and transformation modules
+- [Utilities](utils.md) - Index generation and value matching utilities  
+- [Binary Operations](binary.md) - Morphological operations and connectivity analysis
+- [Discrete Constraints](discrete.md) - Material connectivity and feature size control

fdtdx-0.1.0/docs/api/constraints/mapping.md

@@ -0,0 +1,3 @@
+# Constraint Mapping
+
+::: fdtdx.constraints.mapping.ConstraintMapping

fdtdx-0.1.0/docs/api/constraints/modules.md

@@ -0,0 +1,14 @@
+# Constraint Modules
+
+## Base Classes
+::: fdtdx.constraints.module.ConstraintInterface
+::: fdtdx.constraints.module.ConstraintModule
+
+## Value Range Transformations
+::: fdtdx.constraints.module.StandardToInversePermittivityRange
+::: fdtdx.constraints.module.StandardToCustomRange
+::: fdtdx.constraints.module.StandardToPlusOneMinusOneRange
+
+## Material Index Transformations
+::: fdtdx.constraints.module.ClosestIndex
+::: fdtdx.constraints.module.IndicesToInversePermittivities

fdtdx-0.1.0/docs/api/constraints/utils.md

@@ -0,0 +1,9 @@
+# Constraint Utilities
+
+## Index Generation
+::: fdtdx.constraints.utils.compute_allowed_indices
+::: fdtdx.constraints.utils.compute_allowed_indices_without_holes
+::: fdtdx.constraints.utils.compute_allowed_indices_without_holes_single_polymer_columns
+
+## Value Matching
+::: fdtdx.constraints.utils.nearest_index

fdtdx-0.1.0/docs/api/core/arrays.md

@@ -0,0 +1,12 @@
+# Array Operations
+
+::: fdtdx.core.misc.expand_matrix
+::: fdtdx.core.misc.ensure_slice_tuple
+::: fdtdx.core.misc.index_1d_array
+::: fdtdx.core.misc.index_by_slice
+::: fdtdx.core.misc.index_by_slice_take_1d
+::: fdtdx.core.misc.index_by_slice_take
+::: fdtdx.core.misc.mask_1d_from_slice
+::: fdtdx.core.misc.assimilate_shape
+::: fdtdx.core.misc.linear_interpolated_indexing
+::: fdtdx.core.misc.advanced_padding

fdtdx-0.1.0/docs/api/core/config.md

@@ -0,0 +1,4 @@
+# Configuration
+
+::: fdtdx.core.config.GradientConfig
+::: fdtdx.core.config.SimulationConfig

fdtdx-0.1.0/docs/api/core/distributed.md

@@ -0,0 +1,6 @@
+# Distributed Computing
+
+::: fdtdx.core.jax.sharding.get_named_sharding_from_shape
+::: fdtdx.core.jax.sharding.get_dtype_bytes
+::: fdtdx.core.jax.sharding.pretty_print_sharding
+::: fdtdx.core.jax.sharding.create_named_sharded_matrix

fdtdx-0.1.0/docs/api/core/index.md

@@ -0,0 +1,12 @@
+# Core Package
+
+The core package provides fundamental functionality for FDTD simulations.
+
+## Components
+
+- [Types](types.md) - Shape and slice type definitions
+- [Configuration](config.md) - Simulation and gradient configuration
+- [PyTree Utilities](pytree.md) - Tree class utilities and fields
+- [JAX Utilities](jax.md) - JAX-specific utilities
+- [Array Operations](arrays.md) - Array manipulation functions
+- [Distributed Computing](distributed.md) - Sharding and distributed computation

fdtdx-0.1.0/docs/api/core/jax.md

@@ -0,0 +1,10 @@
+# JAX Utilities
+
+## Shape and Type Checking
+::: fdtdx.core.jax.utils.check_shape_dtype
+
+## Utility Classes
+::: fdtdx.core.misc.PaddingConfig
+
+## Gradient Estimators
+::: fdtdx.core.jax.ste.straight_through_estimator

fdtdx-0.1.0/docs/api/core/pytree.md

@@ -0,0 +1,11 @@
+# PyTree Utilities
+
+## Core Utilities
+::: fdtdx.core.jax.pytrees.extended_autoinit
+::: fdtdx.core.jax.pytrees.field
+::: fdtdx.core.jax.pytrees.frozen_field
+::: fdtdx.core.jax.pytrees.frozen_private_field
+
+## Tree Classes and Fields
+::: fdtdx.core.jax.pytrees.ExtendedTreeClassIndexer
+::: fdtdx.core.jax.pytrees.ExtendedTreeClass

fdtdx-0.1.0/docs/api/core/types.md

@@ -0,0 +1,13 @@
+# Types
+
+## Shape Types
+::: fdtdx.core.jax.typing.RealShape3D
+::: fdtdx.core.jax.typing.PartialRealShape3D
+::: fdtdx.core.jax.typing.GridShape3D
+::: fdtdx.core.jax.typing.PartialGridShape3D
+
+## Slice Types
+::: fdtdx.core.jax.typing.SliceTuple3D
+::: fdtdx.core.jax.typing.PartialSliceTuple3D
+::: fdtdx.core.jax.typing.Slice3D
+::: fdtdx.core.jax.typing.PartialSlice3D

fdtdx-0.1.0/docs/api/fdtd/core.md

@@ -0,0 +1,29 @@
+# Core FDTD Algorithms
+
+## Memory-Efficient Implementations
+
+### Reversible FDTD
+::: fdtdx.fdtd.reversible_fdtd
+Time-reversal symmetric FDTD implementation with O(1) memory usage.
+
+### Checkpointed FDTD
+::: fdtdx.fdtd.checkpointed_fdtd
+Gradient checkpointing FDTD implementation for memory-performance tradeoff.
+
+## Custom Time Evolution
+::: fdtdx.fdtd.custom_fdtd_forward
+Customizable FDTD implementation for partial time evolution and analysis.
+
+## Forward Propagation
+::: fdtdx.fdtd.forward.forward
+Standard forward FDTD time stepping implementation.
+
+::: fdtdx.fdtd.forward.forward_single_args_wrapper
+JAX-compatible wrapper for forward propagation.
+
+## Backward Propagation
+::: fdtdx.fdtd.backward.full_backward
+Complete backward FDTD propagation from current state to start time.
+
+::: fdtdx.fdtd.backward.backward
+Single step backward FDTD propagation.

fdtdx-0.1.0/docs/api/fdtd/fields.md

@@ -0,0 +1,8 @@
+# Field Operations
+
+## Field Interpolation
+::: fdtdx.fdtd.curl.interpolate_fields
+
+## Curl Operations
+::: fdtdx.fdtd.curl.curl_E
+::: fdtdx.fdtd.curl.curl_H

fdtdx-0.1.0/docs/api/fdtd/index.md

@@ -0,0 +1,10 @@
+# FDTD Implementations
+
+This section documents the FDTD (Finite-Difference Time-Domain) implementations available in the package.
+
+## Contents
+
+- [Core FDTD Algorithms](core.md) - Memory-efficient and customizable FDTD implementations
+- [Field Operations](fields.md) - Field interpolation and curl operations
+- [Field Updates](updates.md) - Electric and magnetic field update functions
+- [Interface Management](interfaces.md) - Detector and interface handling

fdtdx-0.1.0/docs/api/fdtd/interfaces.md

@@ -0,0 +1,8 @@
+# Detector and Interface Management
+
+## Detector Updates
+::: fdtdx.fdtd.update.update_detector_states
+
+## Interface Handling
+::: fdtdx.fdtd.update.collect_interfaces
+::: fdtdx.fdtd.update.add_interfaces

fdtdx-0.1.0/docs/api/fdtd/updates.md

@@ -0,0 +1,9 @@
+# Field Updates
+
+## Electric Field Updates
+::: fdtdx.fdtd.update.update_E
+::: fdtdx.fdtd.update.update_E_reverse
+
+## Magnetic Field Updates  
+::: fdtdx.fdtd.update.update_H
+::: fdtdx.fdtd.update.update_H_reverse