├── beamz
├── visual
│ ├── core.py
│ ├── README.md
│ ├── __init__.py
│ └── helpers.py
├── devices
│ ├── __init__.py
│ ├── sources
│ │ ├── __init__.py
│ │ ├── signals.py
│ │ ├── gaussian.py
│ │ └── solve.py
│ ├── monitors
│ │ ├── __init__.py
│ │ └── README.md
│ ├── README.md
│ └── core.py
├── optimization
│ ├── __init__.py
│ ├── README.md
│ ├── topology.py
│ └── autodiff.py
├── design
│ ├── library.py
│ ├── __init__.py
│ ├── README.md
│ ├── io.py
│ └── materials.py
├── simulation
│ ├── __init__.py
│ ├── README.md
│ ├── boundaries.py
│ ├── core.py
│ └── fields.py
├── const.py
├── README.md
└── __init__.py
├── MANIFEST.in
├── docs
├── assets
│ ├── z.png
│ ├── head_icon.png
│ ├── flat_logo_2.png
│ └── icon_black_contrast.png
├── index.md
├── getting-started
│ └── installation.md
├── examples
│ ├── simple_demo.ipynb
│ ├── dipole.md
│ ├── ring-resonator.md
│ └── mmi-splitter.md
├── stylesheets
│ └── extra.css
├── jupyter-integration.md
├── unidirectional_mode_source.md
├── jax_performance_guide.md
└── adjoint_optimization_guide.md
├── docs-requirements.txt
├── pytest.ini
├── AGENTS.md
├── examples
├── secondary
│ ├── 3D_dipole.py
│ ├── 6_gds_import.py
│ ├── 5_design_export.py
│ ├── 3D_mode_test.py
│ ├── modesolver_1d_waveguide.py
│ ├── 3D_gaussian_test.py
│ ├── 3D_mmi.py
│ ├── modesolver_2d_waveguide.py
│ └── 4_topo_ideal.py
├── 0_dipole.py
├── 3_waveguide.py
├── 2_resring.py
└── 1_mmi.py
├── README.md
├── LICENSE
├── .gitignore
├── ascii.txt
├── DEVELOPMENT.md
├── mkdocs.yml
├── setup.py
├── pyproject.toml
├── patch_wheel.py
└── release_version.py
/beamz/visual/core.py:
--------------------------------------------------------------------------------
1 |
--------------------------------------------------------------------------------
/beamz/devices/__init__.py:
--------------------------------------------------------------------------------
1 |
--------------------------------------------------------------------------------
/MANIFEST.in:
--------------------------------------------------------------------------------
1 | include README.md
2 | include pyproject.toml
--------------------------------------------------------------------------------
/docs/assets/z.png:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/QuentinWach/beamz/HEAD/docs/assets/z.png
--------------------------------------------------------------------------------
/docs/assets/head_icon.png:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/QuentinWach/beamz/HEAD/docs/assets/head_icon.png
--------------------------------------------------------------------------------
/docs/assets/flat_logo_2.png:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/QuentinWach/beamz/HEAD/docs/assets/flat_logo_2.png
--------------------------------------------------------------------------------
/docs/assets/icon_black_contrast.png:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/QuentinWach/beamz/HEAD/docs/assets/icon_black_contrast.png
--------------------------------------------------------------------------------
/beamz/visual/README.md:
--------------------------------------------------------------------------------
1 | # Visual
2 |
3 | Handles user interactions through plotting, animation, and the command line interface.
--------------------------------------------------------------------------------
/beamz/optimization/__init__.py:
--------------------------------------------------------------------------------
1 | """Adjoint-based optimization helpers for BEAMZ."""
2 |
3 | from . import topology
4 |
5 | __all__ = ["topology"]
--------------------------------------------------------------------------------
/beamz/devices/sources/__init__.py:
--------------------------------------------------------------------------------
1 | from .mode import ModeSource
2 | from .gaussian import GaussianSource
3 |
4 | __all__ = ["ModeSource", "GaussianSource"]
5 |
--------------------------------------------------------------------------------
/docs-requirements.txt:
--------------------------------------------------------------------------------
1 | mkdocs>=1.4.0
2 | mkdocs-material>=9.0.0
3 | mkdocs-jupyter>=0.25.0
4 | ipywidgets>=8.0.0
5 | matplotlib>=3.5.0
6 | numpy>=1.20.0
7 | jupyter>=1.0.0
8 | nbconvert>=7.2.0
--------------------------------------------------------------------------------
/beamz/design/library.py:
--------------------------------------------------------------------------------
1 | # Material Library
2 |
3 | # Vacuum
4 |
5 | # Air
6 |
7 | # SiN
8 |
9 | # SiO2
10 |
11 | # Si3N4
12 |
13 | # Gold
14 |
15 | # Aluminum
16 |
17 | # Copper
18 |
19 |
--------------------------------------------------------------------------------
/beamz/devices/monitors/__init__.py:
--------------------------------------------------------------------------------
1 | """
2 | Monitors module for BEAMZ - Contains field and power monitors.
3 | """
4 |
5 | from beamz.devices.monitors.monitors import Monitor
6 |
7 | __all__ = ['Monitor']
8 |
9 |
--------------------------------------------------------------------------------
/beamz/simulation/__init__.py:
--------------------------------------------------------------------------------
1 | """
2 | Simulation module for BEAMZ - Contains FDTD simulation and field operations.
3 | """
4 |
5 | from beamz.design.meshing import RegularGrid
6 | from beamz.simulation.core import Simulation
7 |
8 | __all__ = ['RegularGrid', 'Simulation']
--------------------------------------------------------------------------------
/beamz/const.py:
--------------------------------------------------------------------------------
1 | LIGHT_SPEED = 299792458.0 # m/s
2 | VAC_PERMITTIVITY = 8.85e-12 # Vacuum permittivity
3 | VAC_PERMEABILITY = 1.256e-6 # Vacuum permeability
4 | EPS_0 = 8.85e-12
5 | MU_0 = 1.256e-6
6 |
7 | # SI units
8 | um = 1e-6
9 | µm = 1e-6
10 | nm = 1e-9
11 |
--------------------------------------------------------------------------------
/beamz/simulation/README.md:
--------------------------------------------------------------------------------
1 | # Simulation
2 |
3 | Modules that orchestrate how the EM fields evolve depending on the design and devices.
4 |
5 | + core.py / Main module to orchestrate the simulation.
6 | + fields.py / Contains the Field class which owns the field data and defines the field update.
7 | + ops.py / Contains the operations used by the field updates.
--------------------------------------------------------------------------------
/pytest.ini:
--------------------------------------------------------------------------------
1 | [pytest]
2 | testpaths = tests
3 | python_files = test_*.py
4 | python_classes = Test*
5 | python_functions = test_*
6 | addopts = -v --cov=beamz
7 | markers =
8 | slow: marks tests as slow (deselect with '-m "not slow"')
9 | design: marks tests related to design module
10 | simulation: marks tests related to simulation module
11 | optimization: marks tests related to optimization module
--------------------------------------------------------------------------------
/beamz/devices/README.md:
--------------------------------------------------------------------------------
1 | # Devices
2 |
3 | Modules that interact with the EM fields of the simulation. This includes injecting (sources), detecting (monitors), and other kinds of regional field manipulation (boundaries) that don't stem from the design or its materials themselves.
4 |
5 | + core.py / Main module that defines ...
6 | + boundaries/
7 | + pml.py
8 | + monitors/
9 | + ...
10 | + sources/
11 | + ...
--------------------------------------------------------------------------------
/AGENTS.md:
--------------------------------------------------------------------------------
1 | + Use the full 120 characters of a line if possible, both for the code but also for arguments of classes and defs.
2 | + Format if-statements into single lines if possible.
3 | + Avoid writing exceptions.
4 | + Don't specify types of variables arguments (except for classes!!) etc. and don't use modules for typing.
5 | + Only write single line descriptions in """ """ under functions rather than long argument, return explanations etc.
6 | + Don't use data classes.
7 | + Don't use decorators.
8 | + Don't create functions with just a single line or return.
--------------------------------------------------------------------------------
/examples/secondary/3D_dipole.py:
--------------------------------------------------------------------------------
1 | from beamz import *
2 | import numpy as np
3 |
4 | WL = 0.6*µm # wavelength of the source
5 | TIME = 40*WL/LIGHT_SPEED # total simulation duration
6 | N_CLAD = 1; N_CORE = 2 # refractive indices of the core and cladding
7 | DX, DT = calc_optimal_fdtd_params(WL, max(N_CORE, N_CLAD)) # optimal grid size and time step
8 |
9 | # Create the design
10 | design = Design(8*µm, 8*µm, 8*µm, material=Material(N_CLAD**2), pml_size=WL*1.5)
11 | design += Rectangle(width=4*µm, height=4*µm, depth=4*µm, material=Material(N_CORE**2))
12 | design.show()
--------------------------------------------------------------------------------
/README.md:
--------------------------------------------------------------------------------
1 |
2 |
3 |

4 |
5 | BEAMZ is an experimental **electromagnetic simulation** package using the FDTD method. It features a **high-level API** that enables fast prototyping and procedural design with just a few lines of code, made for (but not limited to) photonic integrated circuits.
6 |
7 |
8 | ```bash
9 | pip3 install beamz
10 | ```
11 | 
12 | 
13 |
--------------------------------------------------------------------------------
/beamz/design/__init__.py:
--------------------------------------------------------------------------------
1 | """
2 | Design module for BEAMZ - Contains components for designing photonic structures.
3 | """
4 |
5 | from beamz.design.materials import Material, CustomMaterial
6 | from beamz.design.core import Design
7 | from beamz.design.structures import Rectangle, Circle, Ring, CircularBend, Polygon, Taper
8 | from beamz.design.meshing import RegularGrid, RegularGrid3D, create_mesh
9 |
10 | __all__ = ['Material', 'CustomMaterial', 'Design', 'Rectangle', 'Circle', 'Ring', 'CircularBend', 'Polygon', 'Taper',
11 | 'RegularGrid', 'RegularGrid3D', 'create_mesh']
12 |
--------------------------------------------------------------------------------
/beamz/design/README.md:
--------------------------------------------------------------------------------
1 | # Design
2 |
3 | Module to define the design of complex structures parametrically as well as mesh it into grids of material values.
4 |
5 | + core.py / Main module to define and organiz complex design geometries, materials, ...
6 | + structures.py / Polygon objects to define geometry within the design
7 | + materials.py / Material response implemenations (Sellmeier, Drude, etc.)
8 | + library.py / Instances of materials with exp. measurements for various materials (Si, SiO, InP, ...)
9 | + meshing.py / Turns parametric design into rasterized grid.
10 | + io.py / Import and export of the design as .gds, .gltf, etc.
--------------------------------------------------------------------------------
/docs/index.md:
--------------------------------------------------------------------------------
1 | BEAMZ is an **electromagnetic simulation and inverse / generative design** package with support for multiple backends. It features a **high-level API** that enables fast prototyping and design with just a few lines of code, ideal for engineers. For researchers, BEAMZ also exposes low-level functionality, making it a flexible playground for developing and testing novel simulation and optimization methods.
2 |
3 |
4 | ```bash
5 | pip install beamz
6 | ```
7 |
8 | > **Note**: _BEAMZ is currently in a premature state and not publically advertised. If you're a friend of Quentin Wach, feel free to get a first peek, otherwise, consider this package yet to be released._
--------------------------------------------------------------------------------
/docs/getting-started/installation.md:
--------------------------------------------------------------------------------
1 | # Installation
2 |
3 | ```bash
4 | pip install beamz
5 | ```
6 |
7 | is the easiest way to install BEAMZ since you only need [Python 3.8 or higher](https://www.python.org/downloads/) and [pip (Python package installer)](https://pip.pypa.io/en/stable/installation/).
8 |
9 | Alternatively, you can install BEAMZ directly from the source code. To do so, clone the repository:
10 | ```bash
11 | git clone https://github.com/QuentinWach/beamz.git
12 | cd beamz
13 | ```
14 | and then install it using:
15 | ```bash
16 | pip install -e .
17 | ```
18 |
19 | You can verify that Beamz is installed correctly, you can run:
20 | ```python
21 | import beamz
22 | print(beamz.__version__)
23 | ```
24 |
25 | For more help, please visit our [GitHub issues page](https://github.com/QuentinWach/beamz/issues).
--------------------------------------------------------------------------------
/beamz/devices/core.py:
--------------------------------------------------------------------------------
1 | class Device:
2 | """Base class for all simulation devices (sources, monitors, etc.)."""
3 |
4 | def inject(self, fields, t, dt, current_step, resolution, design):
5 | """Inject source fields directly into the simulation grid.
6 |
7 | This method is called before the field update step to add source contributions directly
8 | to the field arrays.
9 | """
10 | pass
11 |
12 | def get_source_terms(self, fields, t, dt, current_step, resolution, design):
13 | """Return source current terms for FDTD update.
14 |
15 | Returns:
16 | source_j: dict mapping field components to (current_array, indices) tuples
17 | source_m: dict mapping field components to (current_array, indices) tuples
18 | """
19 | return {}, {} # Override in subclasses
--------------------------------------------------------------------------------
/examples/0_dipole.py:
--------------------------------------------------------------------------------
1 | from beamz import *
2 | import numpy as np
3 |
4 | WL = 0.6*µm # wavelength of the source
5 | TIME = 25*WL/LIGHT_SPEED # total simulation duration
6 | N_CLAD = 1; N_CORE = 2 # refractive indices of the core and cladding
7 | DX, DT = calc_optimal_fdtd_params(WL, max(N_CORE, N_CLAD), dims=2, safety_factor=0.999, points_per_wavelength=8)
8 |
9 | # Create the design
10 | design = Design(8*µm, 8*µm, material=Material(N_CLAD**2))
11 | design += Rectangle(width=4*µm, height=4*µm, material=Material(N_CORE**2))
12 |
13 | time_steps = np.arange(0, TIME, DT)
14 | signal = ramped_cosine(time_steps, amplitude=1.0, frequency=LIGHT_SPEED/WL, ramp_duration=3*WL/LIGHT_SPEED, t_max=TIME/2)
15 | source = GaussianSource(position=(4*µm, 5*µm), width=WL/6, signal=signal)
16 |
17 | # Add PML boundaries to simulation (not design)
18 | sim = Simulation(design=design, devices=[source], boundaries=[PML(edges='all', thickness=2*WL)], time=time_steps, resolution=DX)
19 | sim.run(animate_live="Ez", animation_interval=1, clean_visualization=True)
--------------------------------------------------------------------------------
/LICENSE:
--------------------------------------------------------------------------------
1 | Copyright © 2025 Quentin Wach
2 |
3 | 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:
4 |
5 | The above copyright notice and this permission notice shall be included in all copies or substantial portions of the Software.
6 |
7 | 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.
8 |
--------------------------------------------------------------------------------
/.gitignore:
--------------------------------------------------------------------------------
1 | # Python
2 | __pycache__/
3 | *.py[cod]
4 | *$py.class
5 | *.so
6 | .Python
7 | build/
8 | develop-eggs/
9 | dist/
10 | downloads/
11 | eggs/
12 | .eggs/
13 | lib/
14 | lib64/
15 | parts/
16 | sdist/
17 | var/
18 | wheels/
19 | *.egg-info/
20 | .installed.cfg
21 | *.egg
22 | *.excalidraw
23 | wip/
24 |
25 | # Virtual Environment
26 | venv/
27 | env/
28 | ENV/
29 |
30 | # IDE
31 | .idea/
32 | .vscode/
33 | *.swp
34 | *.swo
35 |
36 | # Testing
37 | .coverage
38 | htmlcov/
39 | .pytest_cache/
40 | .tox/
41 | .cursor
42 |
43 | # Distribution
44 | dist/
45 | build/
46 | *.egg-info/
47 |
48 | # Jupyter Notebook
49 | .ipynb_checkpoints
50 |
51 | # Local development settings
52 | .env
53 | .env.local
54 |
55 | archive
56 | NOTES.md
57 | *.pdf
58 | simulation_results/
59 | *.h5
60 | *.gds
61 | mat_lib.txt
62 | tape_coupler.py
63 | BENCHMARKS.md
64 | ARTICLES.md
65 | analyze_project.py
66 | *.mp4
67 | pre_README.md
68 | data.txt
69 | plot_data.py
70 | site/
71 | project_ideas/
72 | MARKETING_2.md
73 | TODO.md
74 | *.png
75 | *.npy
76 |
77 | beamz/devices/reference/
78 | notes/
79 | .mplconfig/
80 | .pytest_cache/
81 | .python_cache/
82 | .venv/
--------------------------------------------------------------------------------
/beamz/visual/__init__.py:
--------------------------------------------------------------------------------
1 | """
2 | Visual module for BEAMZ - Contains visualization and UI helpers.
3 | """
4 |
5 | from beamz.visual.viz import (
6 | draw_polygon,
7 | show_design,
8 | show_design_2d,
9 | show_design_3d,
10 | plot_fdtd_field,
11 | animate_fdtd_live,
12 | save_fdtd_animation,
13 | plot_fdtd_power,
14 | close_fdtd_figure,
15 | animate_manual_field
16 | )
17 |
18 | from beamz.visual.helpers import (
19 | display_status,
20 | display_header,
21 | display_parameters,
22 | display_results,
23 | create_rich_progress,
24 | get_si_scale_and_label,
25 | check_fdtd_stability,
26 | calc_optimal_fdtd_params
27 | )
28 |
29 | __all__ = [
30 | 'draw_polygon',
31 | 'show_design',
32 | 'show_design_2d',
33 | 'show_design_3d',
34 | 'plot_fdtd_field',
35 | 'animate_fdtd_live',
36 | 'save_fdtd_animation',
37 | 'plot_fdtd_power',
38 | 'close_fdtd_figure',
39 | 'animate_manual_field',
40 | 'display_status',
41 | 'display_header',
42 | 'display_parameters',
43 | 'display_results',
44 | 'create_rich_progress',
45 | 'get_si_scale_and_label',
46 | 'check_fdtd_stability',
47 | 'calc_optimal_fdtd_params'
48 | ]
49 |
50 |
--------------------------------------------------------------------------------
/ascii.txt:
--------------------------------------------------------------------------------
1 | @@@@@@@: @@@@@@@ =@@@ @@@# @@@@
2 | @@@ :@@% @@@ @@@@@ @@@@. %@@@@
3 | @@@@@@@- @@@@@@* @@@-@@. @@=@@+@@%@@
4 | @@@ @@@ @@@ @@@@@@@@ @@ %@@@:%@@
5 | @@@@@@@: @@@@@@@+@@ @@@.@@ @@@ #@@
6 |
7 | *@@@@@@@@@%*=::::---:::---::---+%@*=*+
8 | -%#-=@@@@@@#+=------==-=%%+::::..-@@#
9 | *@+-*@+=#++*#***+===++%@@@= .*@=-%-
10 | #@@@+.=--#@: -+-+. .=%: -@@
11 | #@@---:=@- #=#= ..-%- @@@
12 | #@@*=+@= #+**:..-=*- *%*
13 | *@@%@+ +*##-::-+#+ +#+
14 | @@+ *%#%*+==+** =+=
15 | -%@%@@%#***#. :+=.
16 | .@@@@@@%%%#%= #=-
17 | %@@@@@@@@@# .*==
18 | +%@@@@@@@@@- .**+. @*-
19 | =*%@@@@@@@@*. +%#- *@@@#
20 | =%%@#%@@@@@#..+#%= =**@@--
21 | *@@@%+%@@@@@:.:*%* -#=@+ =::
22 | @@@@:-*@@@@=.:*#* -%=%% +::
23 | =#@--+%@@=..@%%%**+=--==*##%@@. .**-
24 | -#@#%@@@#*+=#@*=++==-==+**#%=:.:=+%@@=
25 | @@*@@%+*#%@@@%+-:.......::-+=:=#@#@@@@
26 |
--------------------------------------------------------------------------------
/beamz/devices/monitors/README.md:
--------------------------------------------------------------------------------
1 | A Monitor is a functional mapping
2 |
3 | 𝑀
4 | :
5 | (
6 | 𝐸
7 | ,
8 | 𝐻
9 | ,
10 | 𝑡
11 | )
12 | →
13 | observable
14 | (
15 | 𝑡
16 | )
17 | M:(E,H,t)→observable(t)
18 |
19 | Implementation:
20 |
21 | class Monitor:
22 | def __init__(self, region, op, temporal_reduction=None, spatial_reduction=None):
23 | self.region = region # bounding box or surface
24 | self.op = op # function on E,H per cell
25 | self.temporal_reduction = temporal_reduction # "instant", "DFT", "accumulate"
26 | self.spatial_reduction = spatial_reduction # "none", "integrate", "average"
27 |
28 |
29 | op could be a lambda or GPU shader snippet like dot(E, H) or abs(E)**2.
30 |
31 | temporal_reduction defines whether to store time-series, accumulate FFTs, or just instantaneous values.
32 |
33 | spatial_reduction defines whether to integrate or store field maps.
34 |
35 | Then:
36 |
37 | “flux monitor” = op = cross(E,H).n, temporal=DFT, spatial=integrate
38 |
39 | “field snapshot” = op = E, temporal=instant, spatial=none
40 |
41 | “mode monitor” = op = dot(E, E_mode*), temporal=DFT, spatial=integrate
42 |
43 | “energy monitor” = op = |E|^2 + |H|^2, temporal=accumulate, spatial=integrate
44 |
45 | All the same code path, just different parameters.
--------------------------------------------------------------------------------
/examples/secondary/6_gds_import.py:
--------------------------------------------------------------------------------
1 | from beamz import *
2 | import numpy as np
3 |
4 | # Define materials for different layers
5 | LAYER_MATERIALS = {
6 | 0: Material(permittivity=11.67),
7 | 1: Material(permittivity=2.085)
8 | }
9 | LAYER_PROPERTIES = {
10 | 1: {"depth": 1*µm, "z": 0.0}, # Silicon layer at bottom
11 | 0: {"depth": 0.6*µm, "z": 1*µm} # Oxide layer in middle
12 | }
13 |
14 | # Import a GDS file
15 | gds_design = Design.import_gds("mmi.gds")
16 |
17 | # Initialize design with appropriate size
18 | max_width = max(max(v[0] for v in s.vertices) for l in gds_design.layers.values() for s in l)
19 | max_height = max(max(v[1] for v in s.vertices) for l in gds_design.layers.values() for s in l)
20 | total_depth = sum(prop["depth"] for prop in LAYER_PROPERTIES.values())
21 | design = Design(width=max_width, height=max_height, depth=total_depth*1.2, material=Material(1))
22 |
23 | # Add structures from each layer with proper materials and z-positions
24 | for layer_num, structures in gds_design.layers.items():
25 | material = LAYER_MATERIALS.get(layer_num, Material(1))
26 | properties = LAYER_PROPERTIES.get(layer_num, {"depth": 0.5*µm, "z": 0.0})
27 | for structure in structures:
28 | design += Polygon(vertices=structure.vertices, material=material, depth=properties["depth"], z=properties["z"])
29 |
30 | # This will show the 3D visualization
31 | design.show()
--------------------------------------------------------------------------------
/examples/secondary/5_design_export.py:
--------------------------------------------------------------------------------
1 | from beamz import *
2 | import numpy as np
3 |
4 | # Parameters
5 | X, Y = 20*µm, 10*µm # domain size
6 | WL = 1.55*µm # wavelength
7 | TIME = 40*WL/LIGHT_SPEED # total simulation duration
8 | N_CORE = 2.04 # Si3N4
9 | N_CLAD = 1.444 # SiO2
10 | WG_W = 0.565*µm # width of the waveguide
11 | H = 3.5*µm # height of the MMI
12 | W = 9*µm # length of the MMI (in propagation direction)
13 | OFFSET = 1.05*µm # offset of the output waveguides from center of the MMI
14 | DX, DT = calc_optimal_fdtd_params(WL, max(N_CORE, N_CLAD), dims=2, safety_factor=0.60)
15 |
16 | # Design the MMI with input and output waveguides
17 | design = Design(width=X, height=Y, material=Material(N_CLAD**2), pml_size=WL)
18 | design += Rectangle(position=(0, Y/2-WG_W/2), width=X/2, height=WG_W, material=Material(N_CORE**2))
19 | design += Rectangle(position=(X/2, Y/2 + OFFSET - WG_W/2), width=X/2, height=WG_W, material=Material(N_CORE**2))
20 | design += Rectangle(position=(X/2, Y/2 - OFFSET - WG_W/2), width=X/2, height=WG_W, material=Material(N_CORE**2))
21 | design += Rectangle(position=(X/2-W/2, Y/2-H/2), width=W, height=H, material=Material(N_CORE**2))
22 | design.show()
23 |
24 | # TODO: Make this work:
25 | # When no layers are specified, we merge all shapes of the same material that are touching and then seperate the merged polygons by layer,
26 | # i.e. each polygon is a new layer. Note though that we still make polygons early in the list of structures of the Design lower layers.
27 | design.export_gds("mmi.gds")
--------------------------------------------------------------------------------
/beamz/README.md:
--------------------------------------------------------------------------------
1 | ## Module Structure
2 |
3 | ### `design/` - Parametric Design and Geometry
4 | Defines the physical structure of the device through parametric geometry and materials.
5 |
6 | ### `devices/` - Field Sources and Monitors
7 | Handles electromagnetic field injection (sources) and detection (monitors) that interact with the simulation fields.
8 |
9 | ### `simulation/` - FDTD Engine
10 | Orchestrates the finite-difference time-domain (FDTD) simulation and field evolution.
11 |
12 | ### `optimization/` - Gradient-Based Optimization
13 | Provides topology optimization tools using JAX for automatic differentiation and the adjoint method.
14 |
15 | ### `visual/` - Visualization and UI
16 | Handles plotting, animation, and command-line interface interactions.
17 |
18 | ### `const.py` - Physical Constants
19 | Defines fundamental physical constants (light speed, vacuum permittivity/permeability) and unit conversions (µm, nm).
20 |
21 | ## Code Architecture
22 |
23 | The codebase follows a **modular high-level design** with **object-oriented patterns within modules**:
24 |
25 | - **Design-centric**: The `Design` class serves as the central container for structures, materials, sources, and monitors.
26 | - **Simulation orchestration**: The `Simulation` class references a `Design` and manages the FDTD time-stepping, delegating field updates to the `Fields` class.
27 | - **Device abstraction**: Sources and monitors inherit from the `Device` base class, providing a unified interface for field manipulation.
28 | - **Separation of concerns**: Design geometry is separate from simulation execution, which is separate from optimization and visualization.
29 |
--------------------------------------------------------------------------------
/DEVELOPMENT.md:
--------------------------------------------------------------------------------
1 | ## Development Guide / Notes
2 |
3 |
4 |
5 | ## Documentation
6 | To test out the documentation locally, do:
7 | ```bash
8 | mkdocs serve
9 | ```
10 | To deploy it, type:
11 | ```bash
12 | mkdocs gh-deploy
13 | ```
14 | which will then create all the needed files on the gh-deploy branch and, well, deploy it there as a github-page.
15 |
16 |
17 | ## Version Release
18 |
19 | To create a new version release and GitHub tag:
20 |
21 | ```bash
22 | python release_version.py 0.1.6
23 | ```
24 |
25 | This will:
26 | 1. Update version in `setup.py`, `pyproject.toml`, and `beamz/__init__.py`
27 | 2. Create a git tag `v0.1.6`
28 | 3. Push the tag to the remote repository
29 |
30 | To also create a GitHub release (requires GitHub token):
31 |
32 | ```bash
33 | export GITHUB_TOKEN=your_token_here
34 | python scripts/release_version.py 0.1.6 --message "Release notes here"
35 | ```
36 |
37 | Or pass the token directly:
38 |
39 | ```bash
40 | python scripts/release_version.py 0.1.6 --github-token your_token_here
41 | ```
42 |
43 | Additional options:
44 | - `--tag-only`: Only create git tag, don't push or create GitHub release
45 | - `--no-push`: Don't push tag to remote
46 | - `--draft`: Create draft GitHub release
47 | - `--force`: Force overwrite existing tag
48 | - `--skip-version-update`: Skip updating version files
49 |
50 | ## Package Publishing
51 | First update the version numbers in the `setup.py` file and others! Then
52 | ```bash
53 | python -m build
54 | ```
55 | then
56 | ```bash
57 | python patch_wheel.py
58 | ```
59 | then
60 | ```bash
61 | python -m twine upload dist/beamz-0.1.0-py3-none-any.whl
62 | ```
63 | (though with the correct version) in order to publis the newest version of the package to pypi.
--------------------------------------------------------------------------------
/examples/3_waveguide.py:
--------------------------------------------------------------------------------
1 | from beamz import *
2 | import numpy as np
3 | from beamz import calc_optimal_fdtd_params
4 |
5 | WL = 1.55*µm
6 | TIME = 90*WL/LIGHT_SPEED
7 | N_CORE, N_CLAD = 2.04, 1.444 # Si3N4, SiO2
8 | WG_WIDTH = 0.565*µm
9 | DX, DT = calc_optimal_fdtd_params(WL, max(N_CORE, N_CLAD), safety_factor=0.999, points_per_wavelength=20)
10 |
11 | # Create the design
12 | design = Design(width=18*µm, height=7*µm, material=Material(N_CLAD**2))
13 | design += Rectangle(position=(0,3.5*µm-WG_WIDTH/2), width=18*µm, height=WG_WIDTH, material=Material(N_CORE**2))
14 | #design += Rectangle(position=(9*µm-WG_WIDTH/2,0), width=WG_WIDTH, height=7*µm, material=Material(N_CORE**2))
15 | design.show()
16 |
17 | # Rasterize the design
18 | grid = design.rasterize(resolution=DX)
19 | #grid.show(field="permittivity")
20 |
21 | # Create the signal & source
22 | time_steps = np.arange(0, TIME, DT)
23 | signal = ramped_cosine(
24 | time_steps,
25 | amplitude=1.0,
26 | frequency=LIGHT_SPEED / WL,
27 | phase=0,
28 | ramp_duration=WL * 20 / LIGHT_SPEED,
29 | t_max=TIME / 2,
30 | )
31 | source = ModeSource(
32 | grid=grid,
33 | center=(design.width/2, design.height/2),
34 | width=WG_WIDTH * 1.2, # Slightly wider than waveguide to capture mode tails, but not so wide to hit PML/boundaries
35 | wavelength=WL,
36 | pol="tm",
37 | signal=signal,
38 | direction="+x",
39 | )
40 |
41 | # Run the simulation
42 | sim = Simulation(
43 | design=design,
44 | devices=[source],
45 | boundaries=[PML(edges='all', thickness=1.2*WL)],
46 | time=time_steps,
47 | resolution=DX
48 | )
49 | sim.run(animate_live="Ez",
50 | animation_interval=20,
51 | axis_scale=[-9e-5, 9e-5],
52 | cmap="twilight_zero",
53 | clean_visualization=True)
54 |
--------------------------------------------------------------------------------
/examples/2_resring.py:
--------------------------------------------------------------------------------
1 | from beamz import *
2 | import numpy as np
3 |
4 | # Parameters
5 | WL = 1.55*µm
6 | TIME = 120*WL/LIGHT_SPEED
7 | X, Y = 20*µm, 19*µm
8 | N_CORE, N_CLAD = 2.04, 1.444 # Si3N4, SiO2
9 | DX, DT = calc_optimal_fdtd_params(WL, max(N_CORE, N_CLAD), dims=2, safety_factor=0.999, points_per_wavelength=8)
10 | RING_RADIUS, WG_WIDTH = 6*µm, 0.5*µm #0.565*µm
11 |
12 | # Create the design
13 | design = Design(width=X, height=Y, material=Material(N_CLAD**2))
14 | design += Rectangle(position=(0,WL*2), width=X, height=WG_WIDTH, material=Material(N_CORE**2))
15 | design += Ring(position=(X/2, WL*2+WG_WIDTH+RING_RADIUS+WG_WIDTH/2+0.2*WG_WIDTH),
16 | inner_radius=RING_RADIUS-WG_WIDTH/2, outer_radius=RING_RADIUS+WG_WIDTH/2,
17 | material=Material(N_CORE**2))
18 | #design.show()
19 |
20 | # Rasterize the design
21 | grid = design.rasterize(resolution=DX)
22 | #grid.show(field="permittivity")
23 |
24 | # Define the signal & source
25 | time_steps = np.arange(0, TIME, DT)
26 | signal = ramped_cosine(time_steps, amplitude=1.0, frequency=LIGHT_SPEED/WL, phase=0,
27 | ramp_duration=WL*6/LIGHT_SPEED, t_max=TIME/2.5)
28 | source = ModeSource(
29 | grid=grid,
30 | center=(WL*2, WL*2+WG_WIDTH/2),
31 | width=WG_WIDTH * 1.2, # Slightly wider than waveguide to capture mode tails
32 | wavelength=WL,
33 | pol="tm",
34 | signal=signal,
35 | direction="+x",
36 | )
37 |
38 | # Run the simulation
39 | sim = Simulation(
40 | design=design,
41 | devices=[source],
42 | boundaries=[PML(edges='all', thickness=1.2*WL)],
43 | time=time_steps,
44 | resolution=DX
45 | )
46 | sim.run(animate_live="Ez",
47 | animation_interval=2,
48 | axis_scale=[-1.5e-4, 1.5e-4],
49 | cmap="twilight_zero",
50 | clean_visualization=False)
--------------------------------------------------------------------------------
/examples/secondary/3D_mode_test.py:
--------------------------------------------------------------------------------
1 | import os, sys
2 | repo_root = os.path.abspath(os.path.join(os.path.dirname(__file__), '..'))
3 | if repo_root not in sys.path:
4 | sys.path.insert(0, repo_root)
5 |
6 | from beamz import *
7 | import numpy as np
8 |
9 | # Small 3D waveguide box with a fundamental mode source
10 | WL = 1.55*µm
11 | X, Y, Z = 6*µm, 3*µm, 2*µm
12 | DX, DT = calc_optimal_fdtd_params(WL, 2.1, dims=3, safety_factor=0.45)
13 | TIME = 20*WL/LIGHT_SPEED
14 | time_steps = np.arange(0, TIME, DT)
15 |
16 | # Materials
17 | n_clad = 1.444; n_core = 2.04
18 |
19 | design = Design(width=X, height=Y, depth=Z, material=Material(n_clad**2), pml_size=WL/2)
20 |
21 | # Core strip in the center along x, finite thickness in z
22 | core_w = 0.7*µm
23 | core_t = 0.4*µm
24 | design += Rectangle(position=(0, Y/2-core_w/2, Z/2-core_t/2), width=X*0.6, height=core_w, depth=core_t, material=Material(n_core**2))
25 |
26 | # Monitor at half depth
27 | monitor = Monitor(start=(0, 0, Z/2), end=(X, Y, Z/2), record_fields=True, accumulate_power=True, live_update=True, record_interval=2)
28 | design += monitor
29 |
30 | # Mode source at x ~ 1 µm, oriented to +x, with 3D cross-section
31 | signal = ramped_cosine(time_steps, amplitude=1.0, frequency=LIGHT_SPEED/WL, phase=0, ramp_duration=3*WL/LIGHT_SPEED, t_max=TIME/2)
32 | src = ModeSource(design=design, position=(1.0*µm, Y/2, Z/2), width=1.2*µm, height=0.8*µm, direction="+x", signal=signal, grid_resolution=1200, num_modes=1)
33 | design += src
34 |
35 | design.show()
36 | monitor.start_live_visualization('Ez')
37 |
38 | sim = FDTD(design=design, time=time_steps, mesh="regular", resolution=DX, backend="numpy")
39 | sim.run(live=True, axis_scale=[-0.05, 0.05], save=False, save_memory_mode=True, accumulate_power=True)
40 | print("Power records:", len(monitor.power_history))
41 |
42 |
--------------------------------------------------------------------------------
/examples/secondary/modesolver_1d_waveguide.py:
--------------------------------------------------------------------------------
1 | import numpy as np
2 | import matplotlib.pyplot as plt
3 | from beamz.const import LIGHT_SPEED, µm
4 | from beamz.devices.mode import solve_modes
5 |
6 | plt.switch_backend("Agg")
7 |
8 | # Wavelength and grid
9 | wavelength = 1.55 * µm
10 | omega = 2 * np.pi * LIGHT_SPEED / wavelength
11 | dL = 0.005 * µm # 20 nm transverse sampling
12 |
13 | # 1D slab waveguide profile (y-direction)
14 | height = 4.0 * µm
15 | ny = int(np.round(height / dL))
16 | y = np.linspace(-height / 2, height / 2, ny)
17 |
18 | n_core = 3.45
19 | n_clad = 1.44
20 | core_thickness = 0.5 * µm
21 |
22 | eps = np.full(ny, n_clad**2, dtype=float)
23 | core_mask = np.abs(y) <= (core_thickness / 2)
24 | eps[core_mask] = n_core**2
25 |
26 | # Solve for the first two modes propagating along +x
27 | neff, e_fields, h_fields, prop_axis = solve_modes(
28 | eps=eps,
29 | omega=omega,
30 | dL=dL,
31 | npml=20,
32 | m=2,
33 | direction="+x",
34 | return_fields=True,
35 | )
36 |
37 | print("Effective indices:", [float(np.real(n)) for n in neff])
38 |
39 | # Plot permittivity and fundamental |Ez| profile
40 | Ez0 = np.squeeze(e_fields[0][2]) # component index 2 is Ez for TE in 2D
41 | intensity0 = np.abs(Ez0) ** 2
42 |
43 | fig, ax1 = plt.subplots(figsize=(7, 4))
44 | ax1.plot(y / µm, eps, color="black", label="εr")
45 | ax1.set_xlabel("y (µm)")
46 | ax1.set_ylabel("εr", color="black")
47 | ax1.tick_params(axis='y', labelcolor='black')
48 | ax1.grid(True, alpha=0.3)
49 |
50 | ax2 = ax1.twinx()
51 | ax2.plot(y / µm, intensity0 / np.max(intensity0 + 1e-18), color="crimson", label="|Ez|² (norm)")
52 | ax2.set_ylabel("|Ez|² (norm)")
53 |
54 | ax1.set_xlim(-0.5, 0.5)
55 |
56 | plt.title(f"1D Slab Modes, λ = {wavelength/µm:.2f} µm, neff₀ = {float(np.real(neff[0])):.3f}")
57 | fig.tight_layout()
58 | fig.savefig("modesolver_1d_mode0.png", dpi=200)
59 |
60 | plt.close(fig)
--------------------------------------------------------------------------------
/examples/1_mmi.py:
--------------------------------------------------------------------------------
1 | from beamz import *
2 | import numpy as np
3 |
4 | # Parameters
5 | X, Y = 20*µm, 10*µm # domain width, height
6 | WL = 1.55*µm # wavelength
7 | TIME = 40*WL/LIGHT_SPEED # total simulation duration
8 | N_CORE, N_CLAD = 2.04, 1.444 # Si3N4, SiO2
9 | WG_W = 0.565*µm # width of the waveguide
10 | H, W, OFFSET = 3.5*µm, 9*µm, 1.05*µm # height, length, offset of the MMI
11 | DX, DT = calc_optimal_fdtd_params(WL, max(N_CORE, N_CLAD), dims=2, safety_factor=0.999, points_per_wavelength=20)
12 |
13 | # Design the MMI with input and output waveguides
14 | design = Design(width=X, height=Y, material=Material(N_CLAD**2))
15 | design += Rectangle(position=(0, Y/2-WG_W/2), width=X/2, height=WG_W, material=Material(N_CORE**2))
16 | design += Rectangle(position=(X/2, Y/2 + OFFSET - WG_W/2), width=X/2, height=WG_W, material=Material(N_CORE**2))
17 | design += Rectangle(position=(X/2, Y/2 - OFFSET - WG_W/2), width=X/2, height=WG_W, material=Material(N_CORE**2))
18 | design += Rectangle(position=(X/2-W/2, Y/2-H/2), width=W, height=H, material=Material(N_CORE**2))
19 | #design.show()
20 |
21 | # Rasterize the design
22 | grid = design.rasterize(resolution=DX)
23 | #grid.show(field="permittivity")
24 |
25 | # Define the source
26 | time_steps = np.arange(0, TIME, DT)
27 | signal = ramped_cosine(time_steps, amplitude=1.0, frequency=LIGHT_SPEED/WL, ramp_duration=WL*6/LIGHT_SPEED, t_max=TIME/2)
28 | # Prefer TE polarization and restrict transverse width to single-mode core to avoid exciting higher-order lobes
29 | source = ModeSource(grid=grid, center=(3*µm, Y/2), width=WG_W, wavelength=WL, pol="tm", signal=signal, direction="+x")
30 |
31 | # Run the simulation and show results
32 | sim = Simulation(design=design, devices=[source], boundaries=[PML(edges='all', thickness=1.2*WL)], time=time_steps, resolution=DX)
33 | sim.run(animate_live="Ez",
34 | animation_interval=12,
35 | axis_scale=[-7e-5, 7e-5],
36 | clean_visualization=True,
37 | line_color="gray")
--------------------------------------------------------------------------------
/mkdocs.yml:
--------------------------------------------------------------------------------
1 | site_name: BEAMZ Docs
2 | site_description: Documentation for the BEAMZ project
3 | site_author: BEAMZ Team
4 | site_url: https://quentinwach.com/beamz/
5 | repo_url: https://github.com/quentinwach/beamz
6 | repo_name: beamz
7 |
8 | theme:
9 | name: material
10 | default_theme: slate
11 | custom_dir: docs/stylesheets
12 | logo: assets/z.png
13 | favicon: assets/z.png
14 | features:
15 | - navigation.sections
16 | - navigation.expand
17 | - search.highlight
18 | - search.share
19 | - content.code.copy
20 | palette:
21 | - scheme: slate
22 | primary: black
23 | accent: black
24 | toggle:
25 | icon: material/brightness-4
26 | name: Switch to light mode
27 | - scheme: default
28 | primary: black
29 | accent: black
30 | toggle:
31 | icon: material/brightness-7
32 | name: Switch to dark mode
33 |
34 | extra_css:
35 | - stylesheets/extra.css
36 |
37 | plugins:
38 | - search
39 | - mkdocs-jupyter:
40 | execute: true
41 | allow_errors: false
42 | theme: dark
43 | include_source: true
44 | ignore_h1_titles: true
45 | show_input: true
46 | remove_tag_config:
47 | remove_cell_tags:
48 | - hide_code
49 | remove_all_outputs_tags:
50 | - hide_output
51 | remove_input_tags:
52 | - hide_input
53 |
54 | markdown_extensions:
55 | - pymdownx.highlight
56 | - pymdownx.superfences
57 | - pymdownx.inlinehilite
58 | - pymdownx.tabbed
59 | - admonition
60 | - footnotes
61 | - toc:
62 | permalink: true
63 |
64 | nav:
65 | - Home: index.md
66 | - Getting Started:
67 | - Installation: getting-started/installation.md
68 | - Interactive Notebooks: jupyter-integration.md
69 | - Examples:
70 | - MMI Splitter: examples/mmi-splitter.md
71 | - Interactive Demo: examples/simple_demo.ipynb
72 | - Interactive MMI Tutorial: examples/MMI_powersplitter.ipynb
73 | - Basic Dipole: examples/dipole.md
74 | - Ring Resonator: examples/ring-resonator.md
--------------------------------------------------------------------------------
/docs/examples/simple_demo.ipynb:
--------------------------------------------------------------------------------
1 | {
2 | "cells": [
3 | {
4 | "cell_type": "markdown",
5 | "id": "a2e436f2",
6 | "metadata": {},
7 | "source": [
8 | "# Interactive BEAMZ Demo\n",
9 | "\n",
10 | "This is a simple interactive demonstration of BEAMZ functionality."
11 | ]
12 | },
13 | {
14 | "cell_type": "code",
15 | "execution_count": null,
16 | "id": "6cea25f1",
17 | "metadata": {},
18 | "outputs": [],
19 | "source": [
20 | "import numpy as np\n",
21 | "import matplotlib.pyplot as plt\n",
22 | "\n",
23 | "# Create some sample data\n",
24 | "x = np.linspace(0, 10, 100)\n",
25 | "y = np.sin(x)\n",
26 | "\n",
27 | "# Plot the data\n",
28 | "plt.figure(figsize=(10, 6))\n",
29 | "plt.plot(x, y, \"b-\", linewidth=2)\n",
30 | "plt.title(\"Simple Sine Wave\")\n",
31 | "plt.xlabel(\"x\")\n",
32 | "plt.ylabel(\"sin(x)\")\n",
33 | "plt.grid(True, alpha=0.3)\n",
34 | "plt.show()"
35 | ]
36 | },
37 | {
38 | "cell_type": "markdown",
39 | "id": "60f5058e",
40 | "metadata": {},
41 | "source": [
42 | "## Interactive Parameters\n",
43 | "\n",
44 | "The following example shows how to create interactive parameters:"
45 | ]
46 | },
47 | {
48 | "cell_type": "code",
49 | "execution_count": null,
50 | "id": "a8f1d82b",
51 | "metadata": {},
52 | "outputs": [],
53 | "source": [
54 | "# Simple parameter example\n",
55 | "from ipywidgets import interact\n",
56 | "import ipywidgets as widgets\n",
57 | "\n",
58 | "@interact(frequency=(0.1, 5.0, 0.1), amplitude=(0.1, 3.0, 0.1))\n",
59 | "def interactive_plot(frequency=1.0, amplitude=1.0):\n",
60 | " x = np.linspace(0, 10, 100)\n",
61 | " y = amplitude * np.sin(frequency * x)\n",
62 | " \n",
63 | " plt.figure(figsize=(10, 6))\n",
64 | " plt.plot(x, y, \"r-\", linewidth=2)\n",
65 | " plt.title(f\"sin({frequency:.1f}x) * {amplitude:.1f}\")\n",
66 | " plt.xlabel(\"x\")\n",
67 | " plt.ylabel(\"y\")\n",
68 | " plt.grid(True, alpha=0.3)\n",
69 | " plt.ylim(-3, 3)\n",
70 | " plt.show()"
71 | ]
72 | }
73 | ],
74 | "metadata": {},
75 | "nbformat": 4,
76 | "nbformat_minor": 5
77 | }
78 |
--------------------------------------------------------------------------------
/setup.py:
--------------------------------------------------------------------------------
1 | from setuptools import setup, find_packages
2 | import os
3 |
4 | # Remove LICENSE file if it exists to prevent setuptools from adding it
5 | if os.path.exists('LICENSE'):
6 | os.rename('LICENSE', 'LICENSE.bak')
7 |
8 | try:
9 | with open("README.md", "r", encoding="utf-8") as fh:
10 | long_description = fh.read()
11 |
12 | setup(
13 | name="beamz",
14 | version="0.1.7",
15 | author="Quentin Wach",
16 | author_email="quentin.wach+beamz@gmail.com",
17 | description="EM package to create inverse / generative designs for your photonic devices with ease and efficiency.",
18 | long_description=long_description,
19 | long_description_content_type="text/markdown",
20 | url="https://github.com/QuentinWach/beamz",
21 | packages=find_packages(),
22 | classifiers=[
23 | "Development Status :: 3 - Alpha",
24 | "Intended Audience :: Science/Research",
25 | "Operating System :: OS Independent",
26 | "Programming Language :: Python :: 3",
27 | "Programming Language :: Python :: 3.8",
28 | "Programming Language :: Python :: 3.9",
29 | "Programming Language :: Python :: 3.10",
30 | "Programming Language :: Python :: 3.11",
31 | "Topic :: Scientific/Engineering :: Physics",
32 | ],
33 | python_requires=">=3.8",
34 | install_requires=[
35 | "numpy>=1.24.4",
36 | "matplotlib>=3.7.5",
37 | "gdspy>=1.6.0",
38 | "scipy>=1.13.0",
39 | "rich>=13.9.4",
40 | "shapely>=2.0.6",
41 | "jax>=0.4.0",
42 | "jaxlib>=0.4.0",
43 | "optax>=0.1.0"
44 | ],
45 | extras_require={
46 | "dev": [
47 | "pytest>=7.0.0",
48 | "pytest-cov>=4.0.0",
49 | "black>=22.0.0",
50 | "isort>=5.0.0",
51 | "flake8>=4.0.0",
52 | "myst-parser>=2.0.0"
53 | ],
54 | "gpu": [
55 | "torch>=2.6.0",
56 | ],
57 | },
58 | include_package_data=True,
59 | )
60 | finally:
61 | # Restore the LICENSE file
62 | if os.path.exists('LICENSE.bak'):
63 | os.rename('LICENSE.bak', 'LICENSE')
--------------------------------------------------------------------------------
/examples/secondary/3D_gaussian_test.py:
--------------------------------------------------------------------------------
1 | import os, sys
2 | repo_root = os.path.abspath(os.path.join(os.path.dirname(__file__), '..'))
3 | if repo_root not in sys.path:
4 | sys.path.insert(0, repo_root)
5 |
6 | from beamz import *
7 | import numpy as np
8 |
9 | # Small 3D domain for a fast smoke test
10 | WL = 0.6*µm
11 | X, Y, Z = 2.0*µm, 2.0*µm, 1.0*µm
12 | N_MAX = 1.0
13 |
14 | # Stable grid/time step for 3D
15 | DX, DT = calc_optimal_fdtd_params(WL, N_MAX, dims=3, safety_factor=0.45)
16 | TIME = 10*WL/LIGHT_SPEED
17 | time_steps = np.arange(0, TIME, DT)
18 |
19 | # Homogeneous medium with thinner PML to leave a larger propagation region
20 | design = Design(width=X, height=Y, depth=Z, material=Material(permittivity=1.0), pml_size=WL/4)
21 |
22 | # 3D Gaussian source at center
23 | src_pos = (X/2, Y/2, Z/2)
24 | src_width = WL/8
25 | signal = ramped_cosine(time_steps, amplitude=2.0, frequency=LIGHT_SPEED/WL, phase=0, ramp_duration=2*WL/LIGHT_SPEED, t_max=TIME/2)
26 | design += GaussianSource(position=src_pos, width=src_width, signal=signal)
27 |
28 | # Add a 3D plane monitor at mid-Z for interactive field visualization
29 | monitor = Monitor(start=(0, 0, Z/3), end=(X, Y, Z/3),
30 | record_fields=True, accumulate_power=True,
31 | live_update=True, record_interval=2)
32 | design += monitor
33 |
34 | # Show the 3D design domain (uses Plotly if available, otherwise falls back to 2D)
35 | design.show()
36 | monitor.start_live_visualization(field_component='Ez')
37 |
38 | # Run 3D FDTD (lightweight z-slice evolution)
39 | sim = FDTD(design=design, time=time_steps, mesh="regular", resolution=DX, backend="numpy")
40 | # Run with live visualization of Ez (center z-slice for 3D). Use tight color scale based on expected amplitude.
41 | results = sim.run(live=True, axis_scale=[-0.1,0.1], save=False, save_memory_mode=True, accumulate_power=False)
42 |
43 | # Report Ez magnitude statistics at the end
44 | Ez = sim.backend.to_numpy(sim.Ez)
45 | Ez_abs_max = float(np.max(np.abs(Ez)))
46 | Ez_abs_mean = float(np.mean(np.abs(Ez)))
47 | print(f"3D Gaussian test: Ez |max| = {Ez_abs_max:.3e}, |mean| = {Ez_abs_mean:.3e}, shape = {Ez.shape}")
48 |
49 | # Plot a mid-z slice of Ez for visual confirmation
50 | try:
51 | mid_z = Ez.shape[0]//2
52 | sim.plot_field(field="Ez", z_slice=mid_z)
53 | except Exception as e:
54 | print(f"Plot skipped: {e}")
55 |
56 |
57 |
--------------------------------------------------------------------------------
/examples/secondary/3D_mmi.py:
--------------------------------------------------------------------------------
1 | import os, sys
2 | # Ensure local package is used when running from the repo
3 | repo_root = os.path.abspath(os.path.join(os.path.dirname(__file__), '..'))
4 | if repo_root not in sys.path:
5 | sys.path.insert(0, repo_root)
6 | from beamz import *
7 | import numpy as np
8 |
9 | # Parameters
10 | X, Y, Z = 20*µm, 10*µm, 5*µm # domain width, height, depth
11 | WL = 1.55*µm # wavelength
12 | TIME = 40*WL/LIGHT_SPEED # total simulation duration
13 | N_CORE, N_CLAD = 2.04, 1.444 # Si3N4, SiO2
14 | WG_W = 0.565*µm # width of the waveguide
15 | H, W, OFFSET = 3.5*µm, 9*µm, 1.05*µm # height, length, offset of the MMI
16 | DX, DT = calc_optimal_fdtd_params(WL, max(N_CORE, N_CLAD), dims=3, safety_factor=0.40)
17 |
18 | # Design the MMI with input and output waveguides
19 | design = Design(width=X, height=Y, depth=Z, material=Material(permittivity=1.0, permeability=1.0, conductivity=0.0), pml_size=WL)
20 | design += Rectangle(position=(0, 0, 0), width=X, height=Y, depth=Z/2, material=Material(N_CLAD**2))
21 | design += Rectangle(position=(0, Y/2-WG_W/2, Z/2), width=X/2, height=WG_W, depth=Z/8, material=Material(N_CORE**2))
22 | design += Rectangle(position=(X/2, Y/2 + OFFSET - WG_W/2, Z/2), width=X/2, height=WG_W, depth=Z/8, material=Material(N_CORE**2))
23 | design += Rectangle(position=(X/2, Y/2 - OFFSET - WG_W/2, Z/2), width=X/2, height=WG_W, depth=Z/8, material=Material(N_CORE**2))
24 | design += Rectangle(position=(X/2-W/2, Y/2-H/2, Z/2), width=W, height=H, depth=Z/8, material=Material(N_CORE**2))
25 |
26 | # Add monitors
27 | monitor = Monitor(start=(0, 0, Z/2+Z/16), end=(X, Y, Z/2+Z/16), record_fields=True, accumulate_power=True, live_update=True, record_interval=2)
28 | design += monitor
29 |
30 | # Define the source
31 | time_steps = np.arange(0, TIME, DT)
32 | signal = ramped_cosine(time_steps, amplitude=1.0, frequency=LIGHT_SPEED/WL, phase=0, ramp_duration=WL*5/LIGHT_SPEED, t_max=TIME/2)
33 | source = ModeSource(design=design, position=(2.5*µm, Y/2, Z/2),
34 | width=2*µm, height=2*µm, direction="+x", signal=signal)
35 | design += source
36 | design.show()
37 |
38 | source.show()
39 |
40 |
41 | # Kick off live visualization for the monitor (updates will occur during run)
42 | monitor.start_live_visualization(field_component='Ez')
43 |
44 | # Run the simulation
45 | sim = FDTD(design=design, time=time_steps, mesh="regular", resolution=DX)
46 | sim.run(live=True, save_memory_mode=True, accumulate_power=True)
47 | sim.plot_power(db_colorbar=True)
48 |
49 |
50 |
51 |
--------------------------------------------------------------------------------
/pyproject.toml:
--------------------------------------------------------------------------------
1 | [build-system]
2 | requires = ["setuptools>=45", "wheel"]
3 | build-backend = "setuptools.build_meta"
4 |
5 | [project]
6 | name = "beamz"
7 | version = "0.1.7" # Match setup.py
8 | authors = [
9 | { name = "Quentin Wach", email = "quentin.wach+beamz@gmail.com" },
10 | ]
11 | description = "EM package to create inverse / generative designs for your photonic devices with ease and efficiency."
12 | readme = "README.md"
13 | requires-python = ">=3.8"
14 | license = { file = "LICENSE" } # Assuming LICENSE file exists at the root
15 | keywords = ["electromagnetics", "photonics", "inverse design", "generative design", "simulation"]
16 | classifiers = [
17 | "Development Status :: 3 - Alpha",
18 | "Intended Audience :: Science/Research",
19 | "Operating System :: OS Independent",
20 | "Programming Language :: Python :: 3",
21 | "Programming Language :: Python :: 3.8",
22 | "Programming Language :: Python :: 3.9",
23 | "Programming Language :: Python :: 3.10",
24 | "Programming Language :: Python :: 3.11",
25 | "Topic :: Scientific/Engineering :: Physics",
26 | ]
27 | dependencies = [
28 | "numpy>=1.24.4",
29 | "matplotlib>=3.7.5",
30 | "gdspy>=1.6.0",
31 | "scipy>=1.10.1",
32 | "rich>=13.9.4",
33 | "shapely>=2.0.6",
34 | "jax>=0.4.0",
35 | "jaxlib>=0.4.0",
36 | "optax>=0.1.0",
37 | # "mkdocs>=1.5.0", # These seem like dev/docs dependencies, not core
38 | # "mkdocs-material>=9.0.0", # Moved to dev optional-dependencies
39 | ]
40 |
41 | [project.urls]
42 | Homepage = "https://github.com/QuentinWach/beamz"
43 | Repository = "https://github.com/QuentinWach/beamz"
44 | # Bug Tracker, Documentation etc. can be added here
45 |
46 | [project.optional-dependencies]
47 | dev = [
48 | "pytest>=7.0.0",
49 | "pytest-cov>=4.0.0",
50 | "black>=22.0.0",
51 | "isort>=5.0.0",
52 | "flake8>=4.0.0",
53 | "myst-parser>=2.0.0",
54 | "mkdocs>=1.5.0", # Moved from core dependencies
55 | "mkdocs-material>=9.0.0", # Moved from core dependencies
56 | ]
57 | gpu = [
58 | "torch>=2.6.0",
59 | ]
60 | test = [ # Adding a 'test' extra that includes pytest and pytest-cov for convenience
61 | "pytest>=7.0.0",
62 | "pytest-cov>=4.0.0",
63 | ]
64 |
65 |
66 | [tool.black]
67 | line-length = 88
68 | target-version = ['py38']
69 | include = '\.pyi?$'
70 |
71 | [tool.isort]
72 | profile = "black"
73 | multi_line_output = 3
74 |
75 | [tool.pytest.ini_options]
76 | testpaths = ["tests"]
77 | python_files = ["test_*.py"]
78 | addopts = "-v --cov=beamz --cov-report=term-missing"
--------------------------------------------------------------------------------
/beamz/devices/sources/signals.py:
--------------------------------------------------------------------------------
1 | import math as m
2 | import numpy as np
3 | import matplotlib.pyplot as plt
4 |
5 | def cosine(t,amplitude, frequency, phase):
6 | return amplitude * np.cos(2 * np.pi * frequency * t + phase)
7 |
8 | def sigmoid(t, duration=1, min=0, max=1, t0=0):
9 | return min + (max - min) * (1 / (1 + np.exp(-10 * (t - duration/2 - t0) / duration)))
10 |
11 | def ramped_cosine(t, amplitude, frequency, phase=None, ramp_duration=None, t_max=None):
12 | if phase is None: phase = 0
13 | if ramp_duration is None: ramp_duration = t_max / 10
14 | signal = sigmoid(t, min=0, max=1, duration=ramp_duration, t0=0)
15 | signal *= cosine(t, amplitude=amplitude, frequency=frequency, phase=phase)
16 | signal *= sigmoid(t, min=1, max=0, duration=ramp_duration, t0=t_max - ramp_duration)
17 | return signal
18 |
19 | def gaussian(t, amplitude, center, width):
20 | return amplitude * np.exp(-(t - center)**2 / (2 * width**2))
21 |
22 | def gaussian_pulse(t, amplitude, center, width, frequency, phase):
23 | return gaussian(t, amplitude, center, width) * cosine(t, amplitude, frequency, phase)
24 |
25 | def plot_signal(signals, t):
26 | """Create a single signal or a list of signals on the same plot."""
27 | # Convert time to seconds
28 | t_seconds = t
29 | # Determine appropriate time unit and scaling factor
30 | if t_seconds[-1] < 1e-12:
31 | t_scaled = t_seconds * 1e15 # Convert to fs
32 | unit = 'fs'
33 | elif t_seconds[-1] < 1e-9: # Less than 1 ns
34 | t_scaled = t_seconds * 1e12 # Convert to ps
35 | unit = 'ps'
36 | elif t_seconds[-1] < 1e-6: # Less than 1 µs
37 | t_scaled = t_seconds * 1e9 # Convert to ns
38 | unit = 'ns'
39 | elif t_seconds[-1] < 1e-3: # Less than 1 ms
40 | t_scaled = t_seconds * 1e6 # Convert to µs
41 | unit = 'µs'
42 | elif t_seconds[-1] < 1: # Less than 1 s
43 | t_scaled = t_seconds * 1e3 # Convert to ms
44 | unit = 'ms'
45 | else:
46 | t_scaled = t_seconds
47 | unit = 's'
48 | # Create the figure and axis
49 | fig, ax = plt.subplots(figsize=(9, 4))
50 | if isinstance(signals, list):
51 | i = 0
52 | for signal in signals:
53 | ax.plot(t_scaled, signal, label=f'Signal {i}')
54 | i += 1
55 | ax.legend()
56 | else:
57 | ax.plot(t_scaled, signals, color='black')
58 | ax.set_xlim(t_scaled[0], t_scaled[-1])
59 | ax.set_xlabel(f'Time ({unit})')
60 | ax.set_ylabel('Amplitude')
61 | ax.set_title('Signal')
62 | plt.tight_layout()
63 | plt.show()
--------------------------------------------------------------------------------
/docs/examples/dipole.md:
--------------------------------------------------------------------------------
1 | # Basic Dipole Simulation
2 |
3 | This tutorial demonstrates how to create a simple dipole simulation using BEAMZ. We'll simulate light propagation from a point source in a dielectric medium.
4 |
5 | ## Full Code Example
6 |
7 | You heard that right. This is it! This is all it takes!
8 |
9 | ```python
10 | from beamz import *
11 | import numpy as np
12 |
13 | # Define simulation parameters
14 | WL = 0.6*µm # wavelength of the source
15 | TIME = 40*WL/LIGHT_SPEED # total simulation duration
16 | N_CLAD = 1 # refractive index of cladding
17 | N_CORE = 2 # refractive index of core
18 | DX, DT = calc_optimal_fdtd_params(WL, max(N_CORE, N_CLAD))
19 |
20 | # Create the design
21 | design = Design(8*µm, 8*µm, material=Material(N_CLAD**2), pml_size=WL*1.5)
22 | design += Rectangle(width=4*µm, height=4*µm, material=Material(N_CORE**2))
23 |
24 | # Define the signal
25 | time_steps = np.arange(0, TIME, DT)
26 | signal = ramped_cosine(time_steps, amplitude=1.0, frequency=LIGHT_SPEED/WL,
27 | phase=0, ramp_duration=3*WL/LIGHT_SPEED, t_max=TIME/2)
28 | design += GaussianSource(position=(4*µm, 5*µm), width=WL/6, signal=signal)
29 |
30 | # Run the simulation
31 | sim = FDTD(design=design, time=time_steps, mesh="regular", resolution=DX)
32 | sim.run(live=True, save_memory_mode=True, accumulate_power=True)
33 | sim.plot_power(db_colorbar=True)
34 | ```
35 |
36 | ## Step-by-Step Explanation
37 |
38 | But just in case, let's walk through it with some additional details.
39 |
40 | ### 1. Import Required Libraries
41 | ```python
42 | from beamz import *
43 | import numpy as np
44 | ```
45 | We import the BEAMZ library and NumPy for numerical operations.
46 |
47 | ### 2. Define Simulation Parameters
48 | ```python
49 | WL = 0.6*µm # wavelength
50 | TIME = 40*WL/LIGHT_SPEED # simulation duration
51 | N_CLAD = 1 # cladding refractive index
52 | N_CORE = 2 # core refractive index
53 | ```
54 | These parameters define the basic properties of our simulation.
55 |
56 | ### 3. Create the Design
57 | ```python
58 | design = Design(8*µm, 8*µm, material=Material(N_CLAD**2), pml_size=WL*1.5)
59 | design += Rectangle(width=4*µm, height=4*µm, material=Material(N_CORE**2))
60 | ```
61 | We create a simulation domain with a rectangular dielectric region.
62 |
63 | ### 4. Define the Source
64 | ```python
65 | time_steps = np.arange(0, TIME, DT)
66 | signal = ramped_cosine(time_steps, amplitude=1.0, frequency=LIGHT_SPEED/WL,
67 | phase=0, ramp_duration=3*WL/LIGHT_SPEED, t_max=TIME/2)
68 | design += GaussianSource(position=(4*µm, 5*µm), width=WL/6, signal=signal)
69 | ```
70 | We create a Gaussian source with a ramped cosine signal.
71 |
72 | ### 5. Run the Simulation
73 | ```python
74 | sim = FDTD(design=design, time=time_steps, mesh="regular", resolution=DX)
75 | sim.run(live=True, axis_scale=[-1, 1], save_memory_mode=True, accumulate_power=True)
76 | sim.plot_power(db_colorbar=True)
77 | ```
78 | We run the FDTD simulation and visualize the results.
--------------------------------------------------------------------------------
/patch_wheel.py:
--------------------------------------------------------------------------------
1 | #!/usr/bin/env python3
2 | import os
3 | import zipfile
4 | import tempfile
5 | import shutil
6 | import re
7 |
8 | def patch_wheel(wheel_path):
9 | """Patch a wheel file to remove the license-file entry from METADATA."""
10 | temp_dir = tempfile.mkdtemp()
11 | try:
12 | # Unzip the wheel to the temp directory
13 | with zipfile.ZipFile(wheel_path, 'r') as wheel_zip:
14 | wheel_zip.extractall(temp_dir)
15 |
16 | # Find the METADATA file
17 | for root, dirs, files in os.walk(temp_dir):
18 | for dir_name in dirs:
19 | if dir_name.endswith('.dist-info'):
20 | metadata_path = os.path.join(root, dir_name, 'METADATA')
21 | if os.path.exists(metadata_path):
22 | # Read and modify the METADATA file
23 | with open(metadata_path, 'r') as f:
24 | metadata_content = f.read()
25 |
26 | # Remove the License-File line
27 | new_metadata = []
28 | for line in metadata_content.split('\n'):
29 | if not line.startswith('License-File:'):
30 | new_metadata.append(line)
31 |
32 | # Clean up Dynamic fields
33 | final_metadata = []
34 | skip_next = False
35 | for i, line in enumerate(new_metadata):
36 | if skip_next:
37 | skip_next = False
38 | continue
39 |
40 | if line.startswith('Dynamic:'):
41 | # Remove all Dynamic fields
42 | skip_next = True
43 | else:
44 | final_metadata.append(line)
45 |
46 | # Write the modified content back
47 | with open(metadata_path, 'w') as f:
48 | f.write('\n'.join(final_metadata))
49 |
50 | # Create a new zip file
51 | new_wheel_path = wheel_path.replace('.whl', '_patched.whl')
52 | with zipfile.ZipFile(new_wheel_path, 'w', zipfile.ZIP_DEFLATED) as new_wheel:
53 | for root, dirs, files in os.walk(temp_dir):
54 | for file in files:
55 | file_path = os.path.join(root, file)
56 | arcname = os.path.relpath(file_path, temp_dir)
57 | new_wheel.write(file_path, arcname)
58 |
59 | # Replace the original wheel
60 | shutil.move(new_wheel_path, wheel_path)
61 | print(f"Successfully patched: {wheel_path}")
62 | finally:
63 | # Clean up
64 | shutil.rmtree(temp_dir)
65 |
66 | if __name__ == "__main__":
67 | # Find wheel file in dist directory
68 | for file in os.listdir('dist'):
69 | if file.endswith('.whl'):
70 | wheel_path = os.path.join('dist', file)
71 | patch_wheel(wheel_path)
--------------------------------------------------------------------------------
/examples/secondary/modesolver_2d_waveguide.py:
--------------------------------------------------------------------------------
1 | import numpy as np
2 | import matplotlib.pyplot as plt
3 | from beamz.const import LIGHT_SPEED, µm
4 | from beamz.devices.mode import tidy3d_mode_computation_wrapper
5 |
6 | plt.switch_backend("Agg")
7 | wavelength = 1.55 * µm
8 | frequency = LIGHT_SPEED / wavelength
9 | # 2D cross-section grid (y, z)
10 | dy = dz = 0.0125 * µm
11 | lateral_span = 2.0 * µm
12 | substrate_thickness = 1.0 * µm
13 | air_thickness = 1.0 * µm
14 | total_height = substrate_thickness + air_thickness
15 | ny = int(np.round(lateral_span / dy))
16 | nz = int(np.round(total_height / dz))
17 | n_core = 3.45
18 | n_sub = 1.44
19 | n_air = 1.00
20 | core_wy = 0.6 * µm
21 | core_wz = 0.22 * µm
22 |
23 | # Build permittivity grid directly
24 | core_offset = -0.1 * µm
25 | y_samples = np.linspace(-lateral_span / 2 + dy/2, lateral_span / 2 - dy/2, ny)
26 | z_samples = np.linspace(-substrate_thickness + dz/2 - core_offset, air_thickness - dz/2 - core_offset, nz)
27 | eps = np.full((ny, nz), n_air**2, dtype=float)
28 | eps[:, z_samples < 0] = n_sub**2
29 | core_y_mask = np.abs(y_samples) <= core_wy / 2
30 | core_z_mask = (z_samples >= core_offset) & (z_samples <= core_offset + core_wz)
31 | eps[np.ix_(core_y_mask, core_z_mask)] = n_core**2
32 | y_edges = np.linspace(-lateral_span / 2, lateral_span / 2, ny + 1)
33 | z_edges = np.linspace(-substrate_thickness, air_thickness, nz + 1)
34 | coords = [y_edges / µm, z_edges / µm]
35 | y_centered = y_samples
36 | z_centered = z_samples
37 |
38 | modes = tidy3d_mode_computation_wrapper(
39 | frequency=frequency,
40 | permittivity_cross_section=eps,
41 | coords=coords,
42 | direction="+",
43 | num_modes=8,
44 | precision="double",
45 | )
46 |
47 | # Sort by descending neff (real part)
48 | modes = sorted(modes, key=lambda m: float(np.real(m.neff)), reverse=True)
49 | n_cols = min(3, len(modes))
50 |
51 | # Define window size
52 | window_um = (core_wy / 2 + 0.4 * core_wy) / µm
53 |
54 | # Detailed field components (Ey, Ez, Hy, Hz) for first mode
55 | if modes:
56 | mode0 = modes[0]
57 | Ex = np.array(mode0.Ex)
58 | Ey = np.array(mode0.Ey)
59 | Ez = np.array(mode0.Ez)
60 | Hx = np.array(mode0.Hx)
61 | Hy = np.array(mode0.Hy)
62 | Hz = np.array(mode0.Hz)
63 |
64 | fields = {
65 | "Ex": Ex,
66 | "Ey": Ey,
67 | "Ez": Ez,
68 | "Hx": Hx,
69 | "Hy": Hy,
70 | "Hz": Hz,
71 | }
72 |
73 | fig_comp, ax_comp = plt.subplots(2, 3, figsize=(12, 8), constrained_layout=True)
74 | i = 0
75 | for (name, field), ax in zip(fields.items(), ax_comp.ravel()):
76 | ax.text(-0.05, 1.05, f"{chr(ord('a') + i)}", transform=ax.transAxes, fontsize=16, fontweight="bold")
77 | i += 1
78 |
79 | real_field = np.real(field)
80 | vmax = np.max(np.abs(real_field)) or 1.0
81 | im = ax.imshow(real_field.T / vmax, origin="lower",
82 | extent=(y_centered[0] / µm, y_centered[-1] / µm, z_centered[0] / µm, z_centered[-1] / µm),
83 | cmap="viridis", aspect="equal", vmin=-1, vmax=1)
84 | ax.set_xlim(-window_um, window_um)
85 | ax.set_ylim(-window_um, window_um)
86 | ax.set_aspect('equal')
87 | ax.set_title(f"Mode 0: Re({name}) (norm)")
88 | ax.set_xlabel("y (µm)")
89 | ax.set_ylabel("z (µm)")
90 |
91 | fig_comp.savefig("2D_modes.png", dpi=200, bbox_inches="tight")
92 | plt.close(fig_comp)
--------------------------------------------------------------------------------
/docs/stylesheets/extra.css:
--------------------------------------------------------------------------------
1 | /* Custom styles for Beamz documentation */
2 |
3 | /* Improve readability of main content */
4 | .md-content {
5 | font-size: 16px;
6 | line-height: 1.25;
7 | }
8 |
9 | .md-typeset {
10 | -webkit-print-color-adjust: exact;
11 | color-adjust: exact;
12 | font-size: .8rem;
13 | line-height: 1.3;
14 | overflow-wrap: break-word;
15 | }
16 |
17 | /* Style code blocks */
18 | .md-typeset pre {
19 | border-radius: 16px;
20 | box-shadow: 0 2px 4px rgba(0,0,0,0.1);
21 | }
22 |
23 | .md-typeset code {
24 | font-size: 0.9em;
25 | padding: 0.2em 0.4em;
26 | border-radius: 8px;
27 | }
28 |
29 | /* Style headers */
30 | .md-typeset h1 {
31 | font-weight: 700;
32 | margin-bottom: 1.2em;
33 | }
34 |
35 | .md-typeset h2 {
36 | }
37 |
38 | /* Style navigation */
39 | .md-nav {
40 | font-size: .7rem;
41 | line-height: 1.1;
42 | }
43 |
44 | .md-nav__title {
45 | font-weight: 600;
46 | }
47 |
48 | .md-nav__link {
49 | transition: color 0.1s ease;
50 | }
51 |
52 | .md-nav__link:hover {
53 | color: var(--md-primary-fg-color);
54 | }
55 |
56 | /* Style admonitions */
57 | .md-typeset .admonition {
58 | border-radius: 8px;
59 | box-shadow: 0 2px 4px rgba(0,0,0,0.1);
60 | }
61 |
62 | .md-typeset .admonition-title {
63 | font-weight: 600;
64 | }
65 |
66 | /* Style tables */
67 | .md-typeset table {
68 | border-radius: 8px;
69 | overflow: hidden;
70 | box-shadow: 0 2px 4px rgba(0,0,0,0.1);
71 | }
72 |
73 | .md-typeset table th {
74 | background-color: var(--md-primary-fg-color);
75 | color: var(--md-primary-bg-color);
76 | font-weight: 600;
77 | }
78 |
79 | .md-nav__item--section {
80 | display: block;
81 | margin: 1.5em 0;
82 | }
83 |
84 | /* Style buttons and links */
85 | .md-typeset a {
86 | text-decoration: none;
87 | transition: color 0.1s ease;
88 | }
89 |
90 | .md-typeset a:hover {
91 | color: var(--md-accent-fg-color);
92 | text-decoration: underline;
93 | }
94 |
95 | /* Style search */
96 | .md-search__input {
97 | border-radius: 8px;
98 | background-color: var(--md-default-bg-color);
99 | }
100 |
101 | /* Style footer */
102 | .md-footer {
103 | margin-top: 3em;
104 | }
105 |
106 | /* Style tabs */
107 | .md-tabs__link {
108 | font-weight: 500;
109 | }
110 |
111 | /* Style copy button for code blocks */
112 | .md-clipboard {
113 | color: var(--md-default-fg-color--light);
114 | }
115 |
116 | .md-clipboard:hover {
117 | color: var(--md-primary-fg-color);
118 | }
119 |
120 | .md-typeset ol li, .md-typeset ul li {
121 | margin-bottom: .1em;
122 | }
123 |
124 |
125 | /* Style scrollbar */
126 | ::-webkit-scrollbar {
127 | width: 8px;
128 | height: 8px;
129 | }
130 |
131 | ::-webkit-scrollbar-track {
132 | background: var(--md-default-bg-color);
133 | }
134 |
135 | ::-webkit-scrollbar-thumb {
136 | background: var(--md-default-fg-color--light);
137 | border-radius: 4px;
138 | }
139 |
140 | ::-webkit-scrollbar-thumb:hover {
141 | background: var(--md-primary-fg-color);
142 | }
143 |
144 | /* Invert logo and favicon to white */
145 | .md-header__button.md-logo img, .md-header__button.md-logo svg {
146 | filter: invert(0) brightness(2);
147 | }
148 |
149 | /* Invert favicon (browser tab icon) */
150 | link[rel="icon"], link[rel="shortcut icon"] {
151 | filter: invert(0) brightness(2) !important;
152 | }
153 |
154 | /* Always show vertical scrollbar to prevent layout shift */
155 | html {
156 | overflow-y: scroll;
157 | }
--------------------------------------------------------------------------------
/docs/jupyter-integration.md:
--------------------------------------------------------------------------------
1 | # Interactive Jupyter Notebooks in MKDocs
2 |
3 | This documentation site supports **interactive Jupyter notebooks** that are automatically converted and executed when the site is built. This allows for rich, interactive documentation with live code examples, plots, and widgets.
4 |
5 | ## Features
6 |
7 | ### ✅ What Works
8 | - **Live Code Execution**: Notebooks are executed during build time
9 | - **Interactive Plots**: Matplotlib, Plotly, and other plotting libraries
10 | - **Interactive Widgets**: ipywidgets for interactive parameters
11 | - **Rich Output**: Images, HTML, LaTeX, and more
12 | - **Code Highlighting**: Syntax highlighting for all code cells
13 | - **Dark/Light Theme**: Automatically matches the site theme
14 |
15 | ### 📝 How to Add Notebooks
16 |
17 | 1. **Create or copy** your `.ipynb` file to the `docs/` directory structure
18 | 2. **Add to navigation** in `mkdocs.yml`:
19 | ```yaml
20 | nav:
21 | - Examples:
22 | - My Interactive Tutorial: examples/my_notebook.ipynb
23 | ```
24 | 3. **Notebooks are automatically processed** when you run `mkdocs serve` or `mkdocs build`
25 |
26 | ### 🎛️ Interactive Features
27 |
28 | #### Basic Plots
29 | ```python
30 | import matplotlib.pyplot as plt
31 | import numpy as np
32 |
33 | x = np.linspace(0, 10, 100)
34 | y = np.sin(x)
35 | plt.plot(x, y)
36 | plt.show()
37 | ```
38 |
39 | #### Interactive Widgets
40 | ```python
41 | import ipywidgets as widgets
42 | from IPython.display import display
43 |
44 | @widgets.interact(x=(0, 10, 0.1))
45 | def plot_sine(x=1.0):
46 | plt.plot(np.sin(x * np.linspace(0, 10, 100)))
47 | plt.show()
48 | ```
49 |
50 | ### 🏷️ Cell Tags for Control
51 |
52 | You can use cell tags to control how cells are displayed:
53 |
54 | - `hide_code` - Hides the input code, shows only output
55 | - `hide_output` - Hides the output, shows only code
56 | - `hide_input` - Hides the entire input cell
57 |
58 | To add tags in Jupyter:
59 | 1. View → Cell Toolbar → Tags
60 | 2. Add tags to individual cells
61 |
62 | ### ⚙️ Configuration
63 |
64 | The Jupyter integration is configured in `mkdocs.yml`:
65 |
66 | ```yaml
67 | plugins:
68 | - mkdocs-jupyter:
69 | execute: true # Execute notebooks during build
70 | allow_errors: false # Stop build on errors
71 | theme: dark # Match site theme
72 | include_source: true # Show source code
73 | ignore_h1_titles: true # Don't duplicate H1 titles
74 | ```
75 |
76 | ### 📚 Examples
77 |
78 | Check out these interactive examples:
79 | - [Interactive Demo](examples/simple_demo.ipynb) - Basic plots and widgets
80 | - [MMI Tutorial](examples/MMI_powersplitter.ipynb) - Advanced BEAMZ simulation
81 |
82 | ### 🔧 Development Tips
83 |
84 | 1. **Test locally**: Always run `mkdocs serve` to test notebooks locally
85 | 2. **Keep notebooks clean**: Remove unnecessary outputs before committing
86 | 3. **Use relative imports**: Import your package modules relatively
87 | 4. **Handle dependencies**: Ensure all required packages are installed
88 | 5. **Error handling**: Set `allow_errors: false` to catch issues early
89 |
90 | ### 🚀 Advanced Features
91 |
92 | #### Custom CSS for Notebooks
93 | Add custom styling in `docs/stylesheets/extra.css`:
94 |
95 | ```css
96 | /* Notebook-specific styling */
97 | .jp-Notebook {
98 | /* Custom notebook styles */
99 | }
100 | ```
101 |
102 | #### Execution Control
103 | You can control execution per notebook by adding metadata:
104 |
105 | ```json
106 | {
107 | "metadata": {
108 | "mkdocs": {
109 | "execute": false
110 | }
111 | }
112 | }
113 | ```
114 |
115 | This setup gives you powerful, interactive documentation that combines the best of Jupyter notebooks with the convenience of MKDocs!
--------------------------------------------------------------------------------
/beamz/__init__.py:
--------------------------------------------------------------------------------
1 | """
2 | BeamZ - A Python package for electromagnetic simulations.
3 | """
4 |
5 | # Import constants from the const module
6 | from beamz.const import (
7 | LIGHT_SPEED, VAC_PERMITTIVITY, VAC_PERMEABILITY,
8 | EPS_0, MU_0, um, nm, µm, μm
9 | )
10 |
11 | # Import design-related classes and functions
12 | from beamz.design.core import Design
13 | from beamz.design.materials import Material, CustomMaterial
14 | from beamz.design.structures import (
15 | Rectangle, Circle, Ring,
16 | CircularBend, Polygon, Taper
17 | )
18 | from beamz.devices.sources import ModeSource, GaussianSource
19 | from beamz.devices.monitors import Monitor
20 | from beamz.devices.sources.signals import ramped_cosine, plot_signal
21 | from beamz.devices.sources.mode import solve_modes
22 |
23 | # Import simulation-related classes and functions
24 | from beamz.design.meshing import RegularGrid
25 | from beamz.simulation.core import Simulation
26 | from beamz.simulation.boundaries import Boundary, PML, ABC, PeriodicBoundary
27 |
28 | # from beamz.optimization.optimizers import Optimizer # TODO: Re-enable when optimizers module is created
29 | from beamz.optimization.topology import compute_overlap_gradient
30 |
31 | # Import optimization-related classes
32 | # (Currently empty, to be filled as the module grows)
33 |
34 | # Import UI helpers
35 | from beamz.visual.helpers import (
36 | display_header, display_status, create_rich_progress,
37 | display_parameters, display_results,
38 | display_simulation_status, display_optimization_progress,
39 | display_time_elapsed, tree_view, code_preview, get_si_scale_and_label, calc_optimal_fdtd_params
40 | )
41 |
42 | # Prepare a dictionary of all our exports
43 | _exports = {
44 | # Constants
45 | 'LIGHT_SPEED': LIGHT_SPEED,
46 | 'VAC_PERMITTIVITY': VAC_PERMITTIVITY,
47 | 'VAC_PERMEABILITY': VAC_PERMEABILITY,
48 | 'EPS_0': EPS_0,
49 | 'MU_0': MU_0,
50 | 'um': um,
51 | 'nm': nm,
52 | 'µm': µm,
53 | 'μm': μm,
54 |
55 | # Materials
56 | 'Material': Material,
57 | 'CustomMaterial': CustomMaterial,
58 |
59 | # Structures
60 | 'Design': Design,
61 | 'Rectangle': Rectangle,
62 | 'Circle': Circle,
63 | 'Ring': Ring,
64 | 'CircularBend': CircularBend,
65 | 'Polygon': Polygon,
66 | 'Taper': Taper,
67 |
68 | # Sources
69 | 'ModeSource': ModeSource,
70 | 'GaussianSource': GaussianSource,
71 |
72 | # Monitors
73 | 'Monitor': Monitor,
74 |
75 | # Signals
76 | 'ramped_cosine': ramped_cosine,
77 | 'plot_signal': plot_signal,
78 |
79 | # Mode calculations
80 | 'solve_modes': solve_modes,
81 |
82 | # Simulation
83 | 'RegularGrid': RegularGrid,
84 | 'Simulation': Simulation,
85 |
86 | # Boundaries
87 | 'Boundary': Boundary,
88 | 'PML': PML,
89 | 'ABC': ABC,
90 | 'PeriodicBoundary': PeriodicBoundary,
91 |
92 | # Optimization
93 | # 'Optimizer': Optimizer, # TODO: Re-enable when optimizers module is created
94 | 'compute_overlap_gradient': compute_overlap_gradient,
95 |
96 | # UI helpers
97 | 'display_header': display_header,
98 | 'display_status': display_status,
99 | 'create_rich_progress': create_rich_progress,
100 | 'display_parameters': display_parameters,
101 | 'display_results': display_results,
102 | 'display_simulation_status': display_simulation_status,
103 | 'display_optimization_progress': display_optimization_progress,
104 | 'display_time_elapsed': display_time_elapsed,
105 | 'tree_view': tree_view,
106 | 'code_preview': code_preview,
107 | 'get_si_scale_and_label': get_si_scale_and_label,
108 | 'calc_optimal_fdtd_params': calc_optimal_fdtd_params,
109 | }
110 |
111 | # Update module's dictionary with our exports
112 | globals().update(_exports)
113 |
114 | # Define what should be available with "from beamz import *"
115 | __all__ = list(_exports.keys())
116 |
117 | # Version information
118 | __version__ = "0.1.7"
--------------------------------------------------------------------------------
/examples/secondary/4_topo_ideal.py:
--------------------------------------------------------------------------------
1 | import numpy as np
2 | from beamz import *
3 | from beamz.optimization import topology as topo
4 |
5 | # General parameters for the simulation
6 | W, H = 15*µm, 15*µm
7 | WG_W = 0.5*µm
8 | N_CORE, N_CLAD = 2.25, 1.444
9 | WL = 1.55*µm
10 | OPT_STEPS, LR = 50, 0.5
11 | TIME = 15*WL/LIGHT_SPEED
12 | DX, DT = calc_optimal_fdtd_params(WL, max(N_CORE, N_CLAD), dims=2, safety_factor=0.95, points_per_wavelength=10)
13 |
14 | # Design with a custom inhomogeneous material which we will update during optimization
15 | design = Design(width=W, height=H, pml_size=2*µm)
16 | design += Rectangle(position=(0*µm, H/2-WG_W/2), width=3.5*µm, height=WG_W, material=Material(permittivity=N_CORE**2))
17 | design += Rectangle(position=(W/2-WG_W/2, H), width=WG_W, height=-3.5*µm, material=Material(permittivity=N_CORE**2))
18 | design_material = CustomMaterial(permittivity_func= lambda x, y, z=None: np.random.uniform(N_CLAD**2, N_CORE**2))
19 | design += Rectangle(position=(W/2-4*µm, H/2-4*µm), width=8*µm, height=8*µm, material=design_material)
20 |
21 | # Define the signal
22 | t = np.linspace(0, TIME, int(TIME/DT))
23 | signal = ramped_cosine(t=t, amplitude=1.0, frequency=LIGHT_SPEED/WL, t_max=TIME, ramp_duration=WL*5/LIGHT_SPEED, phase=0)
24 |
25 | # Rasterize the initial design
26 | grid = RegularGrid(design=design, resolution=DX)
27 |
28 | # Mask and density initialization for topology updates
29 | design_region = design.structures[-1]
30 | mask = design_region.make_mask()
31 | density = topo.initialize_density_from_region(design_region, DX)
32 |
33 | # Define the objective function
34 | def objective_function(monitor: Monitor, normalize: bool = True) -> float:
35 | total_power = 0.0
36 | if monitor.power_history: total_power = sum(monitor.power_history)
37 | else:
38 | if monitor.accumulate_power and monitor.power_accumulated is not None:
39 | total_power = float(np.sum(monitor.power_accumulated))
40 | if normalize and monitor.power_history: total_power /= len(monitor.power_history)
41 | return -float(total_power)
42 |
43 | # Define the optimizer
44 | optimizer = Optimizer(method="adam", learning_rate=LR)
45 | objective_history = []
46 | for opt_step in range(OPT_STEPS):
47 |
48 | # Add filters and contraints to the design region permittivity density
49 | blurred = topo.blur_density(density, radius=WL / DX)
50 | projected = topo.project_density(blurred, beta=2.0, eta=0.5)
51 | density = np.where(mask, projected, density)
52 | topo.update_design_region_material(design_region, density)
53 |
54 | # Run the forward FDTD simulation
55 | forward = FDTD(design=grid, devices=[
56 | ModeSource(grid=grid, center=(2.5*µm, H/2), width=WG_W*4, wavelength=WL, pol="tm", signal=signal, direction="+x", mode=0),
57 | Monitor(design=design, start=(W/2-WG_W*2, H-2.5*µm), end=(W/2+WG_W*2, H-2.5*µm), objective_function=objective_function)
58 | ], time=t) # TODO: Integrate the objective function into the FDTD simulation
59 | forward_fields, objective_value = forward.run(live=True, axis_scale=[-1, 1], save_memory_mode=True) # TODO: only save the Ez field!!!
60 | objective_history.append(objective_value)
61 |
62 | # Run the adjoint FDTD simulation step-by-step and accumulate the overlap field
63 | adjoint = FDTD(design=grid, devices=[ModeSource(grid=grid, center=(W/2, H-2.5*µm),
64 | width=WG_W*4, wavelength=WL, pol="tm", signal=signal, direction="-y", mode=0)], time=t)
65 | overlap_gradient = np.zeros_like(forward_fields["Ez"]) # TODO: Initialize with the correct shape!!!
66 | for step in t:
67 | adjoint_field = adjoint.step() # Simulate one step of the adjoint FDTD simulation
68 | overlap_gradient += compute_overlap_gradient(forward_fields, adjoint_field) / len(t) # Accumulate overlap gradient
69 | forward_fields.pop() # Delete the forward field that was just used to free up memory
70 |
71 | # Update the grid permittivity with the overlap gradient & clip values to the permittivity range
72 | update = optimizer.step(overlap_gradient)
73 | density = np.clip(density + update, N_CLAD**2, N_CORE**2)
74 |
75 | # Run final simulation
76 |
77 | # Show the final design
78 |
79 | # Convert and save the design as GDS
--------------------------------------------------------------------------------
/docs/examples/ring-resonator.md:
--------------------------------------------------------------------------------
1 | # Ring Resonator Simulation
2 |
3 | This tutorial demonstrates how to simulate a ring resonator using BEAMZ. Ring resonators are important components in integrated photonics for filtering and sensing applications.
4 |
5 | ## Overview
6 |
7 | In this tutorial, you will learn:
8 |
9 | - How to create a ring resonator structure
10 | - How to simulate light coupling between a waveguide and a ring
11 | - How to analyze resonance effects
12 | - How to visualize the results
13 |
14 | ## Code Example
15 |
16 | ```python
17 | from beamz import *
18 | import numpy as np
19 |
20 | # Parameters
21 | WL = 1.55*µm # wavelength
22 | TIME = 120*WL/LIGHT_SPEED # simulation duration
23 | X = 20*µm # domain width
24 | Y = 19*µm # domain height
25 | N_CORE = 2.04 # Si3N4 refractive index
26 | N_CLAD = 1.444 # SiO2 refractive index
27 | WG_WIDTH = 0.565*µm # waveguide width
28 | RING_RADIUS = 6*µm # ring radius
29 | DX, DT = calc_optimal_fdtd_params(WL, max(N_CORE, N_CLAD))
30 |
31 | # Create the design
32 | design = Design(width=X, height=Y, material=Material(N_CLAD**2), pml_size=WL)
33 | # Create the bus waveguide
34 | design += Rectangle(position=(0,WL*2), width=X, height=WG_WIDTH,
35 | material=Material(N_CORE**2))
36 | # Create the ring resonator
37 | design += Ring(position=(X/2, WL*2+WG_WIDTH+RING_RADIUS+WG_WIDTH/2+0.2*WG_WIDTH),
38 | inner_radius=RING_RADIUS-WG_WIDTH/2,
39 | outer_radius=RING_RADIUS+WG_WIDTH/2,
40 | material=Material(N_CORE**2))
41 |
42 | # Define the signal & source
43 | time_steps = np.arange(0, TIME, DT)
44 | signal = ramped_cosine(time_steps, amplitude=1.0, frequency=LIGHT_SPEED/WL,
45 | phase=0, ramp_duration=WL*20/LIGHT_SPEED, t_max=TIME/3)
46 | design += ModeSource(design=design,
47 | start=(WL*2, WL*2+WG_WIDTH/2-1.5*µm),
48 | end=(WL*2, WL*2+WG_WIDTH/2+1.5*µm),
49 | wavelength=WL, signal=signal)
50 | design.show()
51 |
52 | # Run the simulation
53 | sim = FDTD(design=design, time=time_steps, mesh="regular", resolution=DX, backend="numpy")
54 | sim.run(live=True, save_memory_mode=True, accumulate_power=True)
55 | sim.plot_power(db_colorbar=True)
56 | ```
57 |
58 | ## Step-by-Step Explanation
59 |
60 | ### 1. Import Required Libraries
61 | ```python
62 | from beamz import *
63 | import numpy as np
64 | ```
65 | We import the necessary libraries for the simulation.
66 |
67 | ### 2. Define Simulation Parameters
68 | ```python
69 | WL = 1.55*µm # wavelength
70 | TIME = 120*WL/LIGHT_SPEED # simulation duration
71 | X = 20*µm # domain width
72 | Y = 19*µm # domain height
73 | N_CORE = 2.04 # Si3N4 refractive index
74 | N_CLAD = 1.444 # SiO2 refractive index
75 | WG_WIDTH = 0.565*µm # waveguide width
76 | RING_RADIUS = 6*µm # ring radius
77 | ```
78 | These parameters define the ring resonator properties and simulation settings.
79 |
80 | ### 3. Create the Design
81 | ```python
82 | design = Design(width=X, height=Y, material=Material(N_CLAD**2), pml_size=WL)
83 | # Create the bus waveguide
84 | design += Rectangle(position=(0,WL*2), width=X, height=WG_WIDTH,
85 | material=Material(N_CORE**2))
86 | # Create the ring resonator
87 | design += Ring(position=(X/2, WL*2+WG_WIDTH+RING_RADIUS+WG_WIDTH/2+0.2*WG_WIDTH),
88 | inner_radius=RING_RADIUS-WG_WIDTH/2,
89 | outer_radius=RING_RADIUS+WG_WIDTH/2,
90 | material=Material(N_CORE**2))
91 | ```
92 | We create a bus waveguide and a ring resonator with specific dimensions.
93 |
94 | ### 4. Define the Source
95 | ```python
96 | time_steps = np.arange(0, TIME, DT)
97 | signal = ramped_cosine(time_steps, amplitude=1.0, frequency=LIGHT_SPEED/WL,
98 | phase=0, ramp_duration=WL*20/LIGHT_SPEED, t_max=TIME/3)
99 | design += ModeSource(design=design,
100 | start=(WL*2, WL*2+WG_WIDTH/2-1.5*µm),
101 | end=(WL*2, WL*2+WG_WIDTH/2+1.5*µm),
102 | wavelength=WL, signal=signal)
103 | ```
104 | We create a mode source to excite the bus waveguide.
105 |
106 | ### 5. Run the Simulation
107 | ```python
108 | sim = FDTD(design=design, time=time_steps, mesh="regular", resolution=DX, backend="numpy")
109 | sim.run(live=True, save_memory_mode=True, accumulate_power=True)
110 | sim.plot_power(db_colorbar=True)
111 | ```
112 | We run the FDTD simulation and visualize the results.
--------------------------------------------------------------------------------
/beamz/optimization/README.md:
--------------------------------------------------------------------------------
1 | # Beamz Optimization Module
2 |
3 | This module provides tools for gradient-based optimization of electromagnetic devices, with a primary focus on topology optimization using the adjoint method and JAX for autodifferentiation.
4 |
5 | ## Structure
6 |
7 | ### `topology.py`
8 | The high-level interface for topology optimization.
9 | - **`TopologyManager` Class**: Manages the optimization lifecycle.
10 | - **Initialization**: Sets up the design, mask, optimizer (using `optax`), and filter parameters.
11 | - **`update_design(step)`**: Updates the physical material distribution based on the current optimization step (handling beta-continuation).
12 | - **`apply_gradient(grad_eps, beta)`**: Backpropagates the gradient from the epsilon (permittivity) domain to the design parameter domain and performs an optimizer step using `optax`.
13 | - **`get_physical_density(beta)`**: Transforms the latent design parameters into a physical density (0 to 1) using filtering and projection.
14 | - **Helper Functions**:
15 | - `compute_overlap_gradient`: Calculates the gradient of the overlap integral (mode matching) using forward and adjoint fields.
16 | - `create_optimization_mask`: Generates a boolean mask defining the design region.
17 | - `get_fixed_structure_mask`: Identifies fixed structures (e.g., waveguides) outside the design region to ensure proper connectivity.
18 |
19 | ### `autodiff.py`
20 | A library of JAX-based differentiable operations used for density filtering and projection.
21 | - **Morphological Filters**:
22 | - `grayscale_erosion`, `grayscale_dilation`: Differentiable grayscale morphology using smooth min/max approximations (LogSumExp).
23 | - `grayscale_opening`, `grayscale_closing`: Compound operations for noise removal and feature size control.
24 | - `masked_morphological_filter`: Applies filters with support for a "fixed structure mask" to prevent erosion at waveguide connections.
25 | - **Conic Filters**:
26 | - `masked_conic_filter`: A filter with a linear decay kernel (cone), used for enforcing geometric constraints like minimum linewidth and spacing.
27 | - **Blurring**:
28 | - `masked_box_blur`: Standard box blur implementation.
29 | - **Projection**:
30 | - `smoothed_heaviside`: A differentiable step function (using `tanh`) to binarize the density field.
31 | - **Backpropagation**:
32 | - `compute_parameter_gradient_vjp`: Uses JAX's vector-Jacobian product (VJP) to automatically compute gradients through the entire filter-project pipeline.
33 |
34 | ## Key Features
35 |
36 | 1. **Differentiable Morphology**: Unlike standard blurring, this module supports differentiable morphological operations (erosion, dilation, opening, closing). This allows for strict control over minimum feature sizes and avoids "gray" boundaries often seen with Gaussian blurs.
37 | 2. **Geometric Constraints**: The **conic filter** option provides a method to enforce minimum length scales (linewidth and spacing) by using a cone-shaped kernel, as described in topology optimization literature.
38 | 3. **Connectivity Preservation**: The filtering pipeline includes a mechanism to "pad" the design region with information from fixed structures (like input/output waveguides). This prevents the optimization from creating gaps or disconnecting the device from the external circuit.
39 | 4. **Beta-Continuation**: Supports a beta-schedule for the Heaviside projection, gradually increasing the sharpness of the binarization to avoid getting stuck in local minima while ensuring a final binary design.
40 | 5. **JAX Integration**: All heavy lifting for density transformation and gradient chain-rule calculation is handled efficiently by JAX. Uses `optax` for JAX-native optimizer implementations (Adam, SGD).
41 |
42 | ## Usage Example
43 |
44 | ```python
45 | from beamz.optimization.topology import TopologyManager, create_optimization_mask
46 |
47 | # 1. Setup Design and Mask
48 | mask = create_optimization_mask(grid, opt_region)
49 |
50 | # 2. Initialize Manager
51 | opt = TopologyManager(
52 | design=design,
53 | region_mask=mask,
54 | resolution=DX,
55 | filter_type='conic', # Options: 'morphological', 'conic', 'blur'
56 | filter_radius=0.15*µm, # Physical units (e.g. microns)
57 | simple_smooth_radius=0.03*µm # Optional smoothing (physical units)
58 | )
59 |
60 | # 3. Optimization Loop
61 | for step in range(STEPS):
62 | # Get current physical density
63 | beta, phys_density = opt.update_design(step, STEPS)
64 |
65 | # Update grid permittivity
66 | grid.permittivity[mask] = EPS_MIN + phys_density[mask] * (EPS_MAX - EPS_MIN)
67 |
68 | # ... Run FDTD & Compute Gradient (grad_eps) ...
69 |
70 | # Update Parameters
71 | opt.apply_gradient(grad_eps, beta)
72 | ```
73 |
--------------------------------------------------------------------------------
/docs/examples/mmi-splitter.md:
--------------------------------------------------------------------------------
1 | # MMI Coupler Simulation
2 |
3 | This tutorial demonstrates how to simulate a Multi-Mode Interference (MMI) coupler using BEAMZ. MMI couplers are important components in integrated photonics for power splitting and combining.
4 |
5 | ## Overview
6 |
7 | In this tutorial, you will learn:
8 |
9 | - How to create an MMI coupler structure
10 | - How to simulate light propagation in multi-mode regions
11 | - How to analyze power splitting
12 | - How to visualize the results
13 |
14 | ## Code Example
15 |
16 | ```python
17 | from beamz import *
18 | import numpy as np
19 |
20 | # Parameters
21 | X = 20*µm # domain width
22 | Y = 10*µm # domain height
23 | WL = 1.55*µm # wavelength
24 | TIME = 40*WL/LIGHT_SPEED # simulation duration
25 | N_CORE = 2.04 # Si3N4 refractive index
26 | N_CLAD = 1.444 # SiO2 refractive index
27 | WG_W = 0.565*µm # width of the waveguide
28 | H = 3.5*µm # height of the MMI
29 | W = 9*µm # length of the MMI
30 | OFFSET = 1.05*µm # offset of the output waveguides
31 | DX, DT = calc_optimal_fdtd_params(WL, max(N_CORE, N_CLAD), dims=2, safety_factor=0.60)
32 |
33 | # Design the MMI with input and output waveguides
34 | design = Design(width=X, height=Y, material=Material(N_CLAD**2), pml_size=WL)
35 | # Input waveguide
36 | design += Rectangle(position=(0, Y/2-WG_W/2), width=X/2, height=WG_W,
37 | material=Material(N_CORE**2))
38 | # Output waveguides
39 | design += Rectangle(position=(X/2, Y/2 + OFFSET - WG_W/2), width=X/2, height=WG_W,
40 | material=Material(N_CORE**2))
41 | design += Rectangle(position=(X/2, Y/2 - OFFSET - WG_W/2), width=X/2, height=WG_W,
42 | material=Material(N_CORE**2))
43 | # MMI section
44 | design += Rectangle(position=(X/2-W/2, Y/2-H/2), width=W, height=H,
45 | material=Material(N_CORE**2))
46 |
47 | # Define the source
48 | time_steps = np.arange(0, TIME, DT)
49 | signal = ramped_cosine(time_steps, amplitude=1.0, frequency=LIGHT_SPEED/WL,
50 | phase=0, ramp_duration=WL*5/LIGHT_SPEED, t_max=TIME/2)
51 | source = ModeSource(design=design, start=(2*µm, Y/2-1.2*µm),
52 | end=(2*µm, Y/2+1.2*µm), wavelength=WL, signal=signal)
53 | design += source
54 | design.show()
55 |
56 | # Run the simulation
57 | sim = FDTD(design=design, time=time_steps, mesh="regular", resolution=DX)
58 | sim.run(live=True, save_memory_mode=True, accumulate_power=True)
59 | sim.plot_power(db_colorbar=True)
60 | ```
61 |
62 | ## Step-by-Step Explanation
63 |
64 | ### 1. Import Required Libraries
65 | ```python
66 | from beamz import *
67 | import numpy as np
68 | ```
69 | We import the necessary libraries for the simulation.
70 |
71 | ### 2. Define Simulation Parameters
72 | ```python
73 | X = 20*µm # domain width
74 | Y = 10*µm # domain height
75 | WL = 1.55*µm # wavelength
76 | TIME = 40*WL/LIGHT_SPEED # simulation duration
77 | N_CORE = 2.04 # Si3N4 refractive index
78 | N_CLAD = 1.444 # SiO2 refractive index
79 | WG_W = 0.565*µm # waveguide width
80 | H = 3.5*µm # MMI height
81 | W = 9*µm # MMI length
82 | OFFSET = 1.05*µm # output waveguide offset
83 | ```
84 | These parameters define the MMI coupler properties and simulation settings.
85 |
86 | ### 3. Create the Design
87 | ```python
88 | design = Design(width=X, height=Y, material=Material(N_CLAD**2), pml_size=WL)
89 | # Input waveguide
90 | design += Rectangle(position=(0, Y/2-WG_W/2), width=X/2, height=WG_W,
91 | material=Material(N_CORE**2))
92 | # Output waveguides
93 | design += Rectangle(position=(X/2, Y/2 + OFFSET - WG_W/2), width=X/2, height=WG_W,
94 | material=Material(N_CORE**2))
95 | design += Rectangle(position=(X/2, Y/2 - OFFSET - WG_W/2), width=X/2, height=WG_W,
96 | material=Material(N_CORE**2))
97 | # MMI section
98 | design += Rectangle(position=(X/2-W/2, Y/2-H/2), width=W, height=H,
99 | material=Material(N_CORE**2))
100 | ```
101 | We create the MMI coupler structure with input and output waveguides.
102 |
103 | ### 4. Define the Source
104 | ```python
105 | time_steps = np.arange(0, TIME, DT)
106 | signal = ramped_cosine(time_steps, amplitude=1.0, frequency=LIGHT_SPEED/WL,
107 | phase=0, ramp_duration=WL*5/LIGHT_SPEED, t_max=TIME/2)
108 | source = ModeSource(design=design, start=(2*µm, Y/2-1.2*µm),
109 | end=(2*µm, Y/2+1.2*µm), wavelength=WL, signal=signal)
110 | design += source
111 | ```
112 | We create a mode source to excite the input waveguide.
113 |
114 | ### 5. Run the Simulation
115 | ```python
116 | sim = FDTD(design=design, time=time_steps, mesh="regular", resolution=DX)
117 | sim.run(live=True, save_memory_mode=True, accumulate_power=True)
118 | sim.plot_power(db_colorbar=True)
119 | ```
120 | We run the FDTD simulation and visualize the results.
--------------------------------------------------------------------------------
/beamz/design/io.py:
--------------------------------------------------------------------------------
1 | import gdspy
2 | from beamz.visual.helpers import display_status
3 |
4 |
5 | def import_gds(gds_file: str, default_depth=1e-6):
6 | """Import a GDS file and return polygon and layer data.
7 |
8 | Args:
9 | gds_file (str): Path to the GDS file
10 | default_depth (float): Default depth/thickness for imported structures in meters
11 | """
12 | from beamz.design.core import Design
13 | from beamz.design.structures import Polygon
14 |
15 | gds_lib = gdspy.GdsLibrary(infile=gds_file)
16 | design = Design() # Create Design instance
17 | cells = gds_lib.cells # Get all cells from the library
18 | total_polygons_imported = 0
19 |
20 | for _cell_name, cell in cells.items():
21 | # Get polygons by spec, which returns a dict: {(layer, datatype): [poly1_points, poly2_points,...]}
22 | gdspy_polygons_by_spec = cell.get_polygons(by_spec=True)
23 | for (layer_num, _datatype), list_of_polygon_points in gdspy_polygons_by_spec.items():
24 | if layer_num not in design.layers: design.layers[layer_num] = []
25 | for polygon_points in list_of_polygon_points:
26 | # Convert points from microns to meters and ensure CCW ordering
27 | vertices_2d = [(point[0] * 1e-6, point[1] * 1e-6) for point in polygon_points]
28 | # Create polygon with appropriate depth
29 | beamz_polygon = Polygon(vertices=vertices_2d, depth=default_depth)
30 | design.layers[layer_num].append(beamz_polygon)
31 | design.structures.append(beamz_polygon)
32 | total_polygons_imported += 1
33 |
34 | # Set 3D flag if we have depth
35 | if default_depth > 0:
36 | design.is_3d = True
37 | design.depth = default_depth
38 |
39 | print(f"Imported {total_polygons_imported} polygons from '{gds_file}' into Design object.")
40 | if design.is_3d: print(f"3D design with depth: {design.depth:.2e} m")
41 | return design
42 |
43 | def export_gds(self, output_file):
44 | """Export a BEAMZ design (including only the structures, not sources or monitors) to a GDS file.
45 |
46 | For 3D designs, structures with the same material that touch (in 3D) will be placed in the same layer.
47 | """
48 | from beamz.design.structures import Polygon, Rectangle, Circle, Ring, CircularBend, Taper
49 | from beamz.devices.sources import ModeSource, GaussianSource
50 | from beamz.devices.monitors import Monitor
51 |
52 | # Create library with micron units (1e-6) and nanometer precision (1e-9)
53 | lib = gdspy.GdsLibrary(unit=1e-6, precision=1e-9)
54 | cell = lib.new_cell("main")
55 | # First, we unify the polygons given their material and if they touch
56 | self.unify_polygons()
57 | # Scale factor to convert from meters to microns
58 | scale = 1e6 # 1 meter = 1e6 microns
59 |
60 | # Group structures by material properties
61 | material_groups = {}
62 | for structure in self.structures:
63 | # Skip PML visualizations, sources, monitors
64 | if hasattr(structure, 'is_pml') and structure.is_pml: continue
65 | if isinstance(structure, (ModeSource, GaussianSource, Monitor)): continue
66 | # Create material key based on material properties
67 | material = getattr(structure, 'material', None)
68 | if material is None: continue
69 | material_key = (
70 | getattr(material, 'permittivity', 1.0),
71 | getattr(material, 'permeability', 1.0),
72 | getattr(material, 'conductivity', 0.0)
73 | )
74 | if material_key not in material_groups: material_groups[material_key] = []
75 | material_groups[material_key].append(structure)
76 |
77 | # Export each material group as a separate layer
78 | for layer_num, (material_key, structures) in enumerate(material_groups.items()):
79 | for structure in structures:
80 | # Get vertices based on structure type
81 | if isinstance(structure, Polygon):
82 | vertices = structure.vertices
83 | interiors = structure.interiors if hasattr(structure, 'interiors') else []
84 | elif isinstance(structure, Rectangle):
85 | x, y = structure.position[0:2] # Take only x,y from position
86 | w, h = structure.width, structure.height
87 | vertices = [(x, y, 0), (x + w, y, 0), (x + w, y + h, 0), (x, y + h, 0)]
88 | interiors = []
89 | elif isinstance(structure, (Circle, Ring, CircularBend, Taper)):
90 | if hasattr(structure, 'to_polygon'):
91 | poly = structure.to_polygon()
92 | vertices = poly.vertices
93 | interiors = getattr(poly, 'interiors', [])
94 | else: continue
95 | else: continue
96 |
97 | # Project vertices to 2D and scale to microns
98 | vertices_2d = [(x * scale, y * scale) for x, y, _ in vertices]
99 | if not vertices_2d: continue
100 | # Scale and project interiors if they exist
101 | interior_2d = []
102 | if interiors:
103 | for interior in interiors:
104 | interior_2d.append([(x * scale, y * scale) for x, y, _ in interior])
105 | try:
106 | # Create gdspy polygon for this layer
107 | if interior_2d: gdspy_poly = gdspy.Polygon(vertices_2d, layer=layer_num, holes=interior_2d)
108 | else: gdspy_poly = gdspy.Polygon(vertices_2d, layer=layer_num)
109 | cell.add(gdspy_poly)
110 | except Exception as e:
111 | print(f"Warning: Failed to create GDS polygon: {e}")
112 | continue
113 |
114 | # Write the GDS file
115 | lib.write_gds(output_file)
116 | print(f"GDS file saved as '{output_file}' with {len(material_groups)} material-based layers")
117 | # Print material information for each layer
118 | for layer_num, (material_key, structures) in enumerate(material_groups.items()):
119 | print(f"Layer {layer_num}: εᵣ={material_key[0]:.1f}, μᵣ={material_key[1]:.1f}, σ={material_key[2]:.2e} S/m")
120 | display_status(f"Created design with size: {self.width:.2e} x {self.height:.2e} x {self.depth:.2e} m")
121 |
--------------------------------------------------------------------------------
/docs/unidirectional_mode_source.md:
--------------------------------------------------------------------------------
1 | # Unidirectional Mode Source Implementation
2 |
3 | ## Overview
4 |
5 | The `ModeSource` class implements a unidirectional electromagnetic mode source using Huygens surface equivalent currents. This source injects fields that propagate in only one direction (±x), preventing unwanted reflections and backpropagation.
6 |
7 | ## Physics Background
8 |
9 | ### Maxwell's Equations with Sources
10 |
11 | For 2D simulations with fields varying in the xy-plane, the relevant Maxwell equations with electric current **J** and magnetic current **M** are:
12 |
13 | ```
14 | ∂_t H_x = -(1/μ)(∂_y E_z + M_x)
15 | ∂_t H_y = (1/μ)(∂_x E_z - M_y)
16 | ∂_t E_z = (1/ε)(∂_x H_y - ∂_y H_x - J_z)
17 | ```
18 |
19 | ### Huygens Surface Equivalent Currents
20 |
21 | To inject a mode unidirectionally, we place electric and magnetic surface currents on a virtual surface (Huygens surface) with normal vector **n̂**:
22 |
23 | ```
24 | J_s = n̂ × H_mode
25 | M_s = -n̂ × E_mode
26 | ```
27 |
28 | For propagation in the +x direction (n̂ = x̂), with field components E_z and H_y:
29 |
30 | ```
31 | J_z = H_y^mode
32 | M_y = E_z^mode
33 | ```
34 |
35 | These currents create a wave that propagates forward (+x) while canceling the backward wave (-x), achieving unidirectional injection.
36 |
37 | ### Sign Conventions for FDTD Integration
38 |
39 | In the FDTD update equations:
40 |
41 | - **J_z** is subtracted from the curl term in the E_z update:
42 | ```
43 | E_z^(n+1) = ... + (dt/ε)[curlH - J_z]
44 | ```
45 | Therefore, we pass J_z with a negative sign.
46 |
47 | - **M_y** is added to the curl_E_y term, which is then subtracted in the H_y update:
48 | ```
49 | H_y^(n+1) = ... - (dt/μ)[curlE_y + M_y]
50 | ```
51 | Therefore, we pass M_y with a positive sign.
52 |
53 | ## Yee Grid Implementation
54 |
55 | ### Field Component Placement
56 |
57 | In the staggered Yee grid for 2D:
58 |
59 | - **E_z**: Cell centers at (i+1/2, j+1/2)
60 | - **H_x**: Cell edges at (i+1/2, j)
61 | - **H_y**: Cell edges at (i, j+1/2)
62 |
63 | ### Source Placement
64 |
65 | For a vertical source plane at x = x_s implementing the Total-Field/Scattered-Field (TFSF) boundary:
66 |
67 | - **J_z** is injected at E_z positions: column index `x_ez_idx` where `(x_ez_idx + 0.5) * dx ≈ x_s`
68 | - **M_y** is injected at H_y positions with directional offset:
69 | - For **+x propagation**: `x_hy_idx = x_ez_idx - 1` (one column to the LEFT)
70 | - For **-x propagation**: `x_hy_idx = x_ez_idx + 1` (one column to the RIGHT)
71 |
72 | This spatial offset is critical for unidirectional behavior. The staggering creates the proper TFSF boundary where:
73 | - E_z at column `i` is at physical position x = `(i + 0.5) * dx`
74 | - H_y at column `i-1` is at physical position x = `(i - 1) * dx`
75 | - The offset enables the forward and backward wave components to interfere destructively in one direction while reinforcing in the desired propagation direction.
76 |
77 | ## Usage Example
78 |
79 | ```python
80 | from beamz import *
81 | import numpy as np
82 |
83 | # Setup
84 | X, Y = 20*µm, 10*µm
85 | WL = 1.55*µm
86 | TIME = 40*WL/LIGHT_SPEED
87 | N_CORE, N_CLAD = 2.04, 1.444
88 | WG_W = 0.6*µm
89 | DX, DT = calc_optimal_fdtd_params(WL, max(N_CORE, N_CLAD), dims=2, points_per_wavelength=10)
90 |
91 | # Create waveguide design
92 | design = Design(width=X, height=Y, material=Material(N_CLAD**2))
93 | design += Rectangle(position=(0, Y/2-WG_W/2), width=X, height=WG_W, material=Material(N_CORE**2))
94 |
95 | # Create time signal
96 | time_steps = np.arange(0, TIME, DT)
97 | signal = ramped_cosine(time_steps, amplitude=0.1, frequency=LIGHT_SPEED/WL,
98 | ramp_duration=WL*6/LIGHT_SPEED, t_max=TIME/2)
99 |
100 | # Create unidirectional mode source
101 | source = ModeSource(
102 | grid=design.rasterize(resolution=DX),
103 | center=(3*µm, Y/2), # Source plane position
104 | width=2.4*µm, # Vertical extent
105 | wavelength=WL,
106 | pol="tm", # Polarization (te or tm)
107 | signal=signal,
108 | direction="+x" # Propagation direction (+x or -x)
109 | )
110 |
111 | # Run simulation
112 | sim = Simulation(design=design, devices=[source],
113 | boundaries=[PML(edges='all', thickness=1.2*WL)],
114 | time=time_steps, resolution=DX)
115 | sim.run(animate_live="Ez")
116 | ```
117 |
118 | ## Implementation Details
119 |
120 | ### Polarization Handling
121 |
122 | The mode solver returns fields in a coordinate system where:
123 |
124 | - **TE** mode: E transverse (E_x, E_y), H has z-component (H_z)
125 | - **TM** mode: H transverse (H_x, H_y), E has z-component (E_z)
126 |
127 | For 2D FDTD with x-propagation:
128 | - `pol="te"` → Uses E_mode[2] (E_z) and H_mode[1] (H_y)
129 | - `pol="tm"` → Uses E_mode[1] (E_y) and H_mode[2] (H_z) mapped as equivalent fields
130 |
131 | ### Direction Handling
132 |
133 | For **-x propagation**:
134 | - Both J_z and M_y are negated relative to +x propagation
135 | - This is automatically handled by checking the Poynting vector direction
136 |
137 | ### Key Methods
138 |
139 | - `initialize(permittivity, resolution)`: Computes the mode and sets up source currents
140 | - `get_source_terms(...)`: Returns (source_j, source_m) dictionaries with current arrays and indices
141 | - `_phase_align(field)`: Aligns field phase to be mostly real at peak amplitude
142 | - `_enforce_propagation_direction(E, H, axis)`: Ensures correct propagation direction via Poynting vector
143 |
144 | ## Verification
145 |
146 | The implementation has been verified to:
147 |
148 | 1. ✓ Correctly apply J_z = H_y^mode at E_z Yee grid positions
149 | 2. ✓ Correctly apply M_y = E_z^mode at H_y Yee grid positions with proper spatial offset
150 | 3. ✓ Use proper sign conventions matching Maxwell's equations (both currents subtracted)
151 | 4. ✓ Place H_y source one column offset from E_z source for TFSF boundary
152 | 5. ✓ Handle both +x and -x propagation directions with correct spatial offsets
153 | 6. ✓ Generate non-zero field amplitudes from mode solver with correct impedance relation
154 |
155 | ## References
156 |
157 | - Taflove & Hagness, "Computational Electrodynamics: The Finite-Difference Time-Domain Method"
158 | - Oskooi et al., "MEEP: A flexible free-software package for electromagnetic simulations"
159 | - Schneider, "Understanding the Finite-Difference Time-Domain Method"
160 |
161 | ## Notes
162 |
163 | - The source automatically initializes on first call to `get_source_terms()`
164 | - Mode effective index (neff) is printed during initialization
165 | - Both TE and TM polarizations are supported for 2D simulations
166 | - For proper unidirectional behavior, ensure PML boundaries are used to absorb outgoing waves
167 |
168 |
--------------------------------------------------------------------------------
/docs/jax_performance_guide.md:
--------------------------------------------------------------------------------
1 | # JAX Backend Performance Guide
2 |
3 | ## Overview
4 |
5 | The JAX backend provides **significant performance improvements** for FDTD simulations through:
6 |
7 | - **JIT Compilation**: Automatic optimization and compilation to GPU/CPU native code
8 | - **GPU Acceleration**: Seamless GPU acceleration without code changes
9 | - **Vectorized Operations**: Highly optimized array operations
10 | - **Multi-device Support**: Automatic parallelization across multiple GPUs
11 | - **Memory Efficiency**: Optimized memory usage and reduced overhead
12 |
13 | **Important**: This backend focuses on **pure performance** - no automatic differentiation overhead!
14 |
15 | ## Installation
16 |
17 | ### Basic JAX Installation
18 |
19 | ```bash
20 | # CPU-only version
21 | pip install jax
22 |
23 | # For GPU support (NVIDIA)
24 | pip install jax[cuda12_pip]
25 |
26 | # For TPU support (Google Cloud)
27 | pip install jax[tpu]
28 | ```
29 |
30 | ### Verify Installation
31 |
32 | ```python
33 | import jax
34 | print(f"JAX version: {jax.__version__}")
35 | print(f"Available devices: {jax.devices()}")
36 | print(f"Default backend: {jax.default_backend()}")
37 | ```
38 |
39 | ## Usage
40 |
41 | ### Basic Usage
42 |
43 | ```python
44 | from beamz.simulation import FDTD
45 | from beamz.simulation.backends import get_backend
46 |
47 | # Use JAX backend for maximum performance
48 | sim = FDTD(
49 | design=design,
50 | time=time_steps,
51 | backend="jax", # This is all you need!
52 | resolution=resolution
53 | )
54 |
55 | # Run simulation - automatically optimized!
56 | sim.run()
57 | ```
58 |
59 | ### Advanced Configuration
60 |
61 | ```python
62 | # Configure JAX backend for specific needs
63 | backend_options = {
64 | "use_jit": True, # Enable JIT compilation (recommended)
65 | "use_64bit": False, # Use 64-bit precision if needed
66 | "device": "auto", # Auto-select best device (gpu > cpu)
67 | "use_pmap": True, # Enable multi-device parallelization
68 | }
69 |
70 | sim = FDTD(
71 | design=design,
72 | time=time_steps,
73 | backend="jax",
74 | backend_options=backend_options,
75 | resolution=resolution
76 | )
77 | ```
78 |
79 | ## Performance Optimization Tips
80 |
81 | ### 1. Grid Size Optimization
82 |
83 | JAX performs best with larger grids due to parallelization overhead:
84 |
85 | ```python
86 | # ✅ Good: Large enough for GPU efficiency
87 | resolution = 20*nm # Creates ~500x500 grid for 10μm design
88 |
89 | # ⚠️ Suboptimal: Too small for GPU efficiency
90 | resolution = 5*nm # Creates ~2000x2000 grid (might be too large)
91 | ```
92 |
93 | ### 2. Memory Management
94 |
95 | ```python
96 | # Monitor memory usage
97 | backend = get_backend("jax")
98 | memory_info = backend.memory_usage()
99 | print(f"GPU memory: {memory_info}")
100 | ```
101 |
102 | ### 3. JIT Compilation Warmup
103 |
104 | ```python
105 | # Pre-compile functions to avoid first-call overhead
106 | if hasattr(sim.backend, 'warmup_compilation'):
107 | field_shapes = {
108 | "Hx": sim.Hx.shape,
109 | "Hy": sim.Hy.shape,
110 | "Ez": sim.Ez.shape
111 | }
112 | sim.backend.warmup_compilation(field_shapes)
113 | ```
114 |
115 | ### 4. Fused Operations
116 |
117 | For maximum performance, use fused field updates when available:
118 |
119 | ```python
120 | # Use fused updates for better performance
121 | if hasattr(sim.backend, 'update_fields_fused'):
122 | Hx_new, Hy_new, Ez_new = sim.backend.update_fields_fused(
123 | sim.Hx, sim.Hy, sim.Ez, sim.sigma, sim.epsilon_r,
124 | sim.dx, sim.dy, sim.dt, mu0, eps0
125 | )
126 | ```
127 |
128 | ## Performance Comparison
129 |
130 | ### Expected Speedups
131 |
132 | | Hardware | Typical Speedup vs NumPy | Best Case |
133 | |----------|---------------------------|-----------|
134 | | Modern CPU (8+ cores) | 2-5x | 10x |
135 | | NVIDIA RTX 3080/4080 | 10-50x | 100x |
136 | | NVIDIA A100/H100 | 50-200x | 500x |
137 | | Google TPU v3/v4 | 20-100x | 300x |
138 |
139 | ### Benchmark Example
140 |
141 | ```python
142 | from examples.jax_performance_benchmark import run_jax_performance_benchmark
143 |
144 | # Run comprehensive performance comparison
145 | run_jax_performance_benchmark()
146 | ```
147 |
148 | This will show you actual speedups on your hardware!
149 |
150 | ## Troubleshooting
151 |
152 | ### Common Issues
153 |
154 | 1. **JAX not using GPU**:
155 | ```python
156 | import jax
157 | print(jax.devices()) # Should show GPU devices
158 |
159 | # If only CPU shown, reinstall with GPU support:
160 | # pip uninstall jax jaxlib
161 | # pip install jax[cuda12_pip]
162 | ```
163 |
164 | 2. **Out of memory errors**:
165 | ```python
166 | # Reduce grid size or enable 32-bit precision
167 | backend_options = {"use_64bit": False}
168 | ```
169 |
170 | 3. **Slow first iteration**:
171 | - This is normal! JIT compilation happens on first call
172 | - Use `warmup_compilation()` to pre-compile
173 | - Subsequent iterations will be much faster
174 |
175 | 4. **Poor multi-GPU performance**:
176 | ```python
177 | # Ensure arrays are large enough for multi-device
178 | # Minimum ~1M points per device recommended
179 | backend_options = {"use_pmap": False} # Disable if needed
180 | ```
181 |
182 | ### Performance Debugging
183 |
184 | ```python
185 | # Benchmark specific operations
186 | backend = get_backend("jax")
187 |
188 | # Test array operations
189 | shapes = {"large": (1000, 1000)}
190 | numpy_backend = get_backend("numpy")
191 | results = backend.compare_with_numpy(numpy_backend, shapes)
192 | ```
193 |
194 | ## System Requirements
195 |
196 | ### Minimum
197 | - Python 3.8+
198 | - 8GB RAM
199 | - Modern CPU (4+ cores)
200 |
201 | ### Recommended for GPU
202 | - NVIDIA GPU with CUDA Compute Capability 3.5+
203 | - 16GB+ RAM
204 | - CUDA 11.2+ or 12.0+
205 |
206 | ### Optimal for Large Simulations
207 | - NVIDIA RTX 4090 or A100
208 | - 32GB+ RAM
209 | - Multiple GPUs for multi-device parallelization
210 |
211 | ## When to Use JAX Backend
212 |
213 | ### ✅ Perfect for:
214 | - Large FDTD simulations (>100k grid points)
215 | - GPU-accelerated computing
216 | - Production simulations requiring maximum speed
217 | - Parameter sweeps and batch simulations
218 | - Real-time applications
219 |
220 | ### ⚠️ Consider alternatives for:
221 | - Very small simulations (<10k grid points)
222 | - Debugging and development (use NumPy)
223 | - Systems without GPU acceleration
224 | - When you need automatic differentiation (wait for our inverse design features!)
225 |
226 | ## Example Performance Results
227 |
228 | Typical results for a 1000×1000 grid, 500 time steps:
229 |
230 | ```
231 | Backend Time/Step Speedup Throughput
232 | numpy 0.012000s 1.00x 41M pts/s
233 | jax 0.000800s 15.0x 625M pts/s
234 | torch 0.001200s 10.0x 416M pts/s
235 | ```
236 |
237 | Your mileage may vary based on hardware! 🚀
--------------------------------------------------------------------------------
/beamz/devices/sources/gaussian.py:
--------------------------------------------------------------------------------
1 | import numpy as np
2 | from beamz.const import EPS_0
3 |
4 | class GaussianSource:
5 | """Gaussian spatial source for FDTD simulations.
6 |
7 | Injects a Gaussian spatial profile into the Ez field (and other E components in 3D).
8 | Useful for dipole-like excitations.
9 | """
10 |
11 | def __init__(self, position, width, signal):
12 | """Initialize the Gaussian source.
13 |
14 | Args:
15 | position: (x, y) for 2D or (x, y, z) for 3D - center of Gaussian
16 | width: Standard deviation of Gaussian profile
17 | signal: Time-dependent signal function s(t) or array
18 | """
19 | self.position = position
20 | self.width = width
21 | self.signal = signal
22 | self._spatial_profile_ez = None
23 | self._grid_indices = None
24 |
25 | def _get_signal_value(self, time, dt):
26 | """Interpolate signal value at arbitrary time."""
27 | # Handle array signal
28 | if isinstance(self.signal, (list, np.ndarray)):
29 | idx_float = float(time / dt)
30 | idx_low = int(np.floor(idx_float))
31 | idx_high = idx_low + 1
32 | frac = idx_float - idx_low
33 |
34 | if 0 <= idx_low < len(self.signal) - 1:
35 | return (1.0 - frac) * self.signal[idx_low] + frac * self.signal[idx_high]
36 | elif idx_low == len(self.signal) - 1:
37 | return self.signal[idx_low]
38 | else:
39 | return 0.0
40 | # Handle callable signal
41 | elif callable(self.signal):
42 | return self.signal(time)
43 | else:
44 | return 0.0
45 |
46 | def inject(self, fields, t, dt, current_step, resolution, design):
47 | """Inject source fields directly into the simulation grid before the FDTD update step.
48 |
49 | Args:
50 | fields: The Fields object containing E and H arrays
51 | t: Current simulation time
52 | dt: Time step size
53 | current_step: Current time step index
54 | resolution: Spatial resolution (dx, dy)
55 | design: The simulation Design object
56 | """
57 | dx = dy = resolution
58 |
59 | # Check dimensionality from fields
60 | if fields.permittivity.ndim == 3:
61 | self._inject_3d(fields, t, dt, resolution)
62 | else:
63 | self._inject_2d(fields, t, dt, resolution)
64 |
65 | def _inject_2d(self, fields, t, dt, resolution):
66 | """Inject into 2D grid (Ez component)."""
67 | ny, nx = fields.Ez.shape
68 |
69 | # Initialize spatial profile if needed (do this once)
70 | if self._spatial_profile_ez is None:
71 | x0, y0 = self.position
72 |
73 | # Determine bounding box for Gaussian (e.g., +/- 4 sigma) to save computation
74 | # Convert position to grid indices
75 | sigma_grid = self.width / resolution
76 | radius_grid = int(np.ceil(4 * sigma_grid))
77 |
78 | cx = int(round(x0 / resolution))
79 | cy = int(round(y0 / resolution))
80 |
81 | # Define ROI limits
82 | x_start = max(0, cx - radius_grid)
83 | x_end = min(nx, cx + radius_grid + 1)
84 | y_start = max(0, cy - radius_grid)
85 | y_end = min(ny, cy + radius_grid + 1)
86 |
87 | self._grid_indices = (slice(y_start, y_end), slice(x_start, x_end))
88 |
89 | # Generate coordinate grids for the ROI
90 | # Ez is at (i+0.5, j+0.5) * resolution? Or integer?
91 | # Assuming standard integer/half-integer staggering for Ez
92 | # Beamz convention: Ez at (i, j) corresponds to integer steps usually?
93 | # Let's use meshgrid relative to center
94 |
95 | # Coordinate arrays for ROI
96 | x_coords = (np.arange(x_start, x_end) + 0.5) * resolution
97 | y_coords = (np.arange(y_start, y_end) + 0.5) * resolution
98 |
99 | X, Y = np.meshgrid(x_coords, y_coords)
100 |
101 | # Compute Gaussian
102 | dist_sq = (X - x0)**2 + (Y - y0)**2
103 | profile = np.exp(-dist_sq / (2 * self.width**2))
104 | self._spatial_profile_ez = profile
105 |
106 | # Get signal value
107 | # Inject at t + 0.5 dt because Ez is updated at n+1 from n
108 | # Soft source: J_z is added.
109 | # Ez_new = Ez_old + ... - dt/eps * J_z
110 | # We want to add to Ez.
111 | # Typically soft source adds to E directly or J term.
112 | # ModeSource adds: fields.Ez += injection
113 | # injection = -jz * dt / (eps * eps0)
114 | # Here we treat GaussianSource as a J_z source.
115 |
116 | signal_val = self._get_signal_value(t + 0.5 * dt, dt)
117 |
118 | # Get permittivity in the region
119 | eps_region = fields.permittivity[self._grid_indices]
120 |
121 | # Calculate injection term
122 | # J(x, t) = Profile(x) * s(t)
123 | # Update: E += -J * dt / (eps * eps0)
124 | # We can absorb the negative sign into the signal definition if we want E to follow signal
125 | # But strictly J is current.
126 | # Let's just add it as a "forcing function" to E.
127 | # If we want E ~ Signal, we might just add Profile * Signal * Scaling
128 | # Let's follow ModeSource physics: inject current J.
129 |
130 | term = self._spatial_profile_ez * signal_val
131 | injection = -term * dt / (EPS_0 * eps_region)
132 |
133 | # Inject
134 | fields.Ez[self._grid_indices] += injection
135 |
136 | def _inject_3d(self, fields, t, dt, resolution):
137 | """Inject into 3D grid (Ez component, could be expanded)."""
138 | # Similar to 2D but with Z coordinate
139 | # Only implementing Ez injection for dipole-like behavior along Z
140 | nz, ny, nx = fields.Ez.shape
141 |
142 | if self._spatial_profile_ez is None:
143 | x0, y0, z0 = self.position
144 |
145 | sigma_grid = self.width / resolution
146 | radius_grid = int(np.ceil(4 * sigma_grid))
147 |
148 | cx, cy, cz = int(round(x0/resolution)), int(round(y0/resolution)), int(round(z0/resolution))
149 |
150 | x_start, x_end = max(0, cx - radius_grid), min(nx, cx + radius_grid + 1)
151 | y_start, y_end = max(0, cy - radius_grid), min(ny, cy + radius_grid + 1)
152 | z_start, z_end = max(0, cz - radius_grid), min(nz, cz + radius_grid + 1)
153 |
154 | self._grid_indices = (slice(z_start, z_end), slice(y_start, y_end), slice(x_start, x_end))
155 |
156 | x_coords = (np.arange(x_start, x_end) + 0.5) * resolution
157 | y_coords = (np.arange(y_start, y_end) + 0.5) * resolution
158 | z_coords = (np.arange(z_start, z_end) + 0.5) * resolution # Ez staggered in Z?
159 |
160 | # Ez in 3D is at (i, j, k+0.5)?
161 | # Assuming cell centers for simplicity for Gaussian "blob"
162 |
163 | Z, Y, X = np.meshgrid(z_coords, y_coords, x_coords, indexing='ij')
164 |
165 | dist_sq = (X - x0)**2 + (Y - y0)**2 + (Z - z0)**2
166 | self._spatial_profile_ez = np.exp(-dist_sq / (2 * self.width**2))
167 |
168 | signal_val = self._get_signal_value(t + 0.5 * dt, dt)
169 | eps_region = fields.permittivity[self._grid_indices]
170 |
171 | term = self._spatial_profile_ez * signal_val
172 | injection = -term * dt / (EPS_0 * eps_region)
173 |
174 | fields.Ez[self._grid_indices] += injection
175 |
--------------------------------------------------------------------------------
/release_version.py:
--------------------------------------------------------------------------------
1 | #!/usr/bin/env python3
2 | """Script to update version and create GitHub tag/release for beamz."""
3 |
4 | import re
5 | import sys
6 | import os
7 | import subprocess
8 | import argparse
9 | from pathlib import Path
10 |
11 | def update_version_in_file(filepath, version, pattern, replacement_template):
12 | """Update version in a file using regex pattern."""
13 | filepath = Path(filepath)
14 | if not filepath.exists():
15 | print(f"Warning: {filepath} does not exist, skipping")
16 | return False
17 |
18 | content = filepath.read_text()
19 | new_content = re.sub(pattern, replacement_template.format(version=version), content)
20 |
21 | if content != new_content:
22 | filepath.write_text(new_content)
23 | print(f"Updated version in {filepath}")
24 | return True
25 | else:
26 | print(f"No change needed in {filepath}")
27 | return False
28 |
29 | def update_version(version):
30 | """Update version in all relevant files."""
31 | changes = []
32 |
33 | # Update setup.py
34 | changes.append(update_version_in_file(
35 | "setup.py",
36 | version,
37 | r'version="[^"]+"',
38 | 'version="{version}"'
39 | ))
40 |
41 | # Update pyproject.toml
42 | changes.append(update_version_in_file(
43 | "pyproject.toml",
44 | version,
45 | r'version = "[^"]+"',
46 | 'version = "{version}"'
47 | ))
48 |
49 | # Update beamz/__init__.py
50 | changes.append(update_version_in_file(
51 | "beamz/__init__.py",
52 | version,
53 | r'__version__ = "[^"]+"',
54 | '__version__ = "{version}"'
55 | ))
56 |
57 | return any(changes)
58 |
59 | def validate_version(version):
60 | """Validate version string format (semantic versioning)."""
61 | pattern = r'^\d+\.\d+\.\d+(-[a-zA-Z0-9]+)?$'
62 | if not re.match(pattern, version):
63 | raise ValueError(f"Invalid version format: {version}. Expected format: X.Y.Z or X.Y.Z-suffix")
64 | return version
65 |
66 | def get_current_branch():
67 | """Get current git branch."""
68 | result = subprocess.run(
69 | ["git", "rev-parse", "--abbrev-ref", "HEAD"],
70 | capture_output=True,
71 | text=True,
72 | check=True
73 | )
74 | return result.stdout.strip()
75 |
76 | def check_git_status():
77 | """Check if git working directory is clean."""
78 | result = subprocess.run(
79 | ["git", "status", "--porcelain"],
80 | capture_output=True,
81 | text=True
82 | )
83 | return result.stdout.strip() == ""
84 |
85 | def create_git_tag(version, message=None):
86 | """Create a git tag for the version."""
87 | if message is None:
88 | message = f"Release version {version}"
89 |
90 | # Check if tag already exists
91 | result = subprocess.run(
92 | ["git", "tag", "-l", f"v{version}"],
93 | capture_output=True,
94 | text=True
95 | )
96 | if result.stdout.strip():
97 | print(f"Tag v{version} already exists. Use --force to overwrite.")
98 | return False
99 |
100 | subprocess.run(
101 | ["git", "tag", "-a", f"v{version}", "-m", message],
102 | check=True
103 | )
104 | print(f"Created git tag: v{version}")
105 | return True
106 |
107 | def push_tag(version, remote="origin"):
108 | """Push tag to remote repository."""
109 | subprocess.run(
110 | ["git", "push", remote, f"v{version}"],
111 | check=True
112 | )
113 | print(f"Pushed tag v{version} to {remote}")
114 |
115 | def create_github_release(version, token=None, message=None, draft=False):
116 | """Create a GitHub release using GitHub API."""
117 | if token is None:
118 | print("GitHub token not provided. Skipping GitHub release creation.")
119 | print("Set GITHUB_TOKEN environment variable or use --github-token flag.")
120 | return False
121 |
122 | if message is None:
123 | message = f"Release version {version}"
124 |
125 | import json
126 | import urllib.request
127 | import urllib.error
128 |
129 | repo = "QuentinWach/beamz" # From setup.py
130 | url = f"https://api.github.com/repos/{repo}/releases"
131 |
132 | data = {
133 | "tag_name": f"v{version}",
134 | "name": f"v{version}",
135 | "body": message,
136 | "draft": draft
137 | }
138 |
139 | req = urllib.request.Request(url, data=json.dumps(data).encode())
140 | req.add_header("Authorization", f"token {token}")
141 | req.add_header("Content-Type", "application/json")
142 |
143 | try:
144 | with urllib.request.urlopen(req) as response:
145 | result = json.loads(response.read())
146 | print(f"Created GitHub release: {result['html_url']}")
147 | return True
148 | except urllib.error.HTTPError as e:
149 | error_body = e.read().decode()
150 | print(f"Error creating GitHub release: {e.code} - {error_body}")
151 | return False
152 |
153 | def main():
154 | parser = argparse.ArgumentParser(
155 | description="Update version and create GitHub tag/release for beamz"
156 | )
157 | parser.add_argument(
158 | "version",
159 | type=validate_version,
160 | help="Version string (e.g., 0.1.6)"
161 | )
162 | parser.add_argument(
163 | "--message", "-m",
164 | help="Release message (default: 'Release version X.Y.Z')"
165 | )
166 | parser.add_argument(
167 | "--tag-only",
168 | action="store_true",
169 | help="Only create git tag, don't push or create GitHub release"
170 | )
171 | parser.add_argument(
172 | "--no-push",
173 | action="store_true",
174 | help="Don't push tag to remote"
175 | )
176 | parser.add_argument(
177 | "--github-token",
178 | help="GitHub personal access token for creating releases"
179 | )
180 | parser.add_argument(
181 | "--draft",
182 | action="store_true",
183 | help="Create draft GitHub release"
184 | )
185 | parser.add_argument(
186 | "--force",
187 | action="store_true",
188 | help="Force overwrite existing tag"
189 | )
190 | parser.add_argument(
191 | "--skip-version-update",
192 | action="store_true",
193 | help="Skip updating version files (use existing version)"
194 | )
195 |
196 | args = parser.parse_args()
197 |
198 | # Check we're on main branch
199 | branch = get_current_branch()
200 | if branch != "main":
201 | print(f"Warning: Not on main branch (currently on {branch})")
202 | response = input("Continue anyway? (y/N): ")
203 | if response.lower() != 'y':
204 | print("Aborted.")
205 | sys.exit(1)
206 |
207 | # Check git status
208 | if not check_git_status():
209 | print("Warning: Working directory is not clean. Uncommitted changes detected.")
210 | response = input("Continue anyway? (y/N): ")
211 | if response.lower() != 'y':
212 | print("Aborted.")
213 | sys.exit(1)
214 |
215 | # Update version files
216 | if not args.skip_version_update:
217 | if not update_version(args.version):
218 | print("No version files were updated.")
219 | response = input("Continue with tag creation? (y/N): ")
220 | if response.lower() != 'y':
221 | print("Aborted.")
222 | sys.exit(1)
223 |
224 | # Create git tag
225 | if args.force:
226 | # Delete existing tag if it exists
227 | subprocess.run(
228 | ["git", "tag", "-d", f"v{args.version}"],
229 | capture_output=True
230 | )
231 | subprocess.run(
232 | ["git", "push", "origin", "--delete", f"v{args.version}"],
233 | capture_output=True
234 | )
235 |
236 | if not create_git_tag(args.version, args.message):
237 | if not args.force:
238 | print("Tag creation failed. Use --force to overwrite existing tag.")
239 | sys.exit(1)
240 | create_git_tag(args.version, args.message)
241 |
242 | if args.tag_only:
243 | print("Tag-only mode: skipping push and GitHub release")
244 | return
245 |
246 | # Push tag
247 | if not args.no_push:
248 | try:
249 | push_tag(args.version)
250 | except subprocess.CalledProcessError as e:
251 | print(f"Error pushing tag: {e}")
252 | sys.exit(1)
253 |
254 | # Create GitHub release
255 | token = args.github_token or os.environ.get("GITHUB_TOKEN")
256 | if token:
257 | create_github_release(args.version, token, args.message, args.draft)
258 | else:
259 | print("\nTo create a GitHub release, set GITHUB_TOKEN environment variable or use --github-token")
260 |
261 | print(f"\n✓ Version {args.version} released successfully!")
262 | print(f" - Version files updated")
263 | print(f" - Git tag v{args.version} created")
264 | if not args.no_push:
265 | print(f" - Tag pushed to remote")
266 |
267 | if __name__ == "__main__":
268 | import os
269 | main()
270 |
271 |
--------------------------------------------------------------------------------
/beamz/visual/helpers.py:
--------------------------------------------------------------------------------
1 | import sys
2 | import datetime
3 | from typing import Dict, Any, Optional
4 | from rich.console import Console
5 | from rich.panel import Panel
6 | from rich.progress import Progress, SpinnerColumn, TextColumn, BarColumn, TaskProgressColumn, TimeRemainingColumn
7 | from rich.table import Table
8 | from rich.text import Text
9 | from rich.tree import Tree
10 | from rich.syntax import Syntax
11 | import numpy as np
12 | from beamz.const import LIGHT_SPEED
13 |
14 | def get_si_scale_and_label(value):
15 | """Convert a value to appropriate SI unit and return scale factor and label."""
16 | if value >= 1e-3: return 1e3, 'mm'
17 | elif value >= 1e-6: return 1e6, 'µm'
18 | elif value >= 1e-9: return 1e9, 'nm'
19 | else: return 1e12, 'pm'
20 |
21 | def check_fdtd_stability(dt, dx, dy=None, dz=None, n_max=1.0, safety_factor=1.0):
22 | """
23 | Check FDTD stability with the Courant-Friedrichs-Lewy (CFL) condition.
24 |
25 | Args:
26 | dt: Time step
27 | dx: Grid spacing in x direction
28 | dy: Grid spacing in y direction (None for 1D)
29 | dz: Grid spacing in z direction (None for 1D/2D)
30 | n_max: Maximum refractive index in the simulation
31 | safety_factor: Factor to apply to the theoretical Courant limit (0-1).
32 | Use 1.0 to evaluate against the theoretical limit 1/sqrt(dims).
33 |
34 | Returns:
35 | tuple: (is_stable, courant_number, max_allowed)
36 | """
37 | # Determine dimensionality
38 | dims = 1
39 | min_spacing = dx
40 | if dy is not None:
41 | dims = 2
42 | min_spacing = min(dx, dy)
43 | if dz is not None:
44 | dims = 3
45 | min_spacing = min(dx, dy, dz)
46 | # Courant number defined with vacuum speed (conservative and standard for Yee grid)
47 | c0 = LIGHT_SPEED
48 | courant = c0 * dt / min_spacing
49 | # Theoretical stability limit
50 | max_allowed = 1.0 / np.sqrt(dims)
51 | # Apply safety factor
52 | safe_limit = safety_factor * max_allowed
53 | return courant <= safe_limit, courant, safe_limit
54 |
55 | def calc_optimal_fdtd_params(wavelength, n_max, dims=2, safety_factor=0.999, points_per_wavelength=10):
56 | """
57 | Calculate optimal FDTD grid resolution and time step based on wavelength and material properties.
58 |
59 | Args:
60 | wavelength: Light wavelength in vacuum
61 | n_max: Maximum refractive index in the simulation
62 | dims: Dimensionality of simulation (1, 2, or 3)
63 | safety_factor: Fraction of the theoretical Courant limit to target (0-1).
64 | 0.95 operates close to the limit; reduce for additional margin.
65 | points_per_wavelength: Number of grid points per wavelength in the highest index material
66 |
67 | Returns:
68 | tuple: (resolution, dt) - optimal spatial resolution and time step
69 | """
70 | # Calculate wavelength in the highest index material
71 | lambda_material = wavelength / n_max
72 | # Calculate optimal grid resolution based on desired points per wavelength
73 | resolution = lambda_material / points_per_wavelength
74 | # Calculate theoretical Courant limit (dt_max = dx / (c * sqrt(dims)))
75 | dt_max = resolution / (LIGHT_SPEED * np.sqrt(dims))
76 | # Apply safety factor (vacuum-based Courant condition)
77 | dt = safety_factor * dt_max
78 | # Verify stability
79 | _, courant, limit = check_fdtd_stability(dt, resolution,
80 | dy=resolution if dims >= 2 else None,
81 | dz=resolution if dims >= 3 else None,
82 | n_max=n_max,
83 | safety_factor=1.0) # Use 1.0 here to get theoretical limit
84 | # Double-check our calculation
85 | assert courant <= limit + 1e-15, "Internal error: calculated time step exceeds stability limit"
86 |
87 | return resolution, dt
88 |
89 | # Initialize rich console
90 | console = Console()
91 |
92 | def progress_bar(progress:int, total:int, length:int=50):
93 | """Print a progress bar to the console."""
94 | percent = 100 * (progress / float(total))
95 | filled_length = int(length * progress // total)
96 | bar = '█' * filled_length + '-' * (length - filled_length - 1)
97 | sys.stdout.write(f'\r|{bar}| {percent:.2f}%')
98 | sys.stdout.flush()
99 |
100 | def display_header(title: str, subtitle: Optional[str] = None) -> None:
101 | """Display a formatted header with optional subtitle."""
102 | console.print(Panel(f"[bold blue]{title}[/]", subtitle=subtitle, expand=False))
103 |
104 | def display_status(status: str, status_type: str = "info") -> None:
105 | """Display a status message with appropriate styling."""
106 | style_map = {
107 | "info": "blue",
108 | "success": "green",
109 | "warning": "yellow",
110 | "error": "red",
111 | }
112 | style = style_map.get(status_type, "white")
113 | console.print(f"[{style}]● {status}[/]")
114 |
115 | def create_rich_progress() -> Progress:
116 | """Create and return a rich progress bar for tracking processes."""
117 | return Progress(
118 | SpinnerColumn(),
119 | TextColumn("[bold blue]{task.description}"),
120 | BarColumn(bar_width=None),
121 | "[progress.percentage]{task.percentage:>3.0f}%",
122 | TimeRemainingColumn(),
123 | )
124 |
125 | def display_parameters(params: Dict[str, Any], title: str = "Parameters") -> None:
126 | """Display a dictionary of parameters in a clean, formatted table."""
127 | table = Table(title=title)
128 | table.add_column("Parameter", style="cyan")
129 | table.add_column("Value", style="green")
130 | for key, value in params.items(): table.add_row(str(key), str(value))
131 | console.print(table)
132 |
133 | def display_results(results: Dict[str, Any], title: str = "Results") -> None:
134 | """Display simulation or optimization results in a formatted table."""
135 | table = Table(title=title)
136 | table.add_column("Metric", style="cyan")
137 | table.add_column("Value", style="green")
138 | for key, value in results.items():
139 | if isinstance(value, (int, float)): value_str = f"{value:.6g}"
140 | else: value_str = str(value)
141 | table.add_row(str(key), value_str)
142 | console.print(table)
143 |
144 | def display_simulation_status(progress: float, metrics: Dict[str, Any] = None) -> None:
145 | """
146 | Display current simulation status with progress and metrics.
147 |
148 | Args:
149 | progress: Progress percentage (0-100)
150 | metrics: Current simulation metrics
151 | """
152 | progress_text = f"Simulation Progress: [bold cyan]{progress:.1f}%[/]"
153 | console.print(progress_text)
154 |
155 | if metrics:
156 | metrics_panel = Panel(
157 | "\n".join([f"[blue]{k}:[/] {v}" for k, v in metrics.items()]),
158 | title="Current Metrics",
159 | expand=False
160 | )
161 | console.print(metrics_panel)
162 |
163 | def display_optimization_progress(iteration: int, total: int, best_value: float, parameters: Dict[str, Any] = None) -> None:
164 | """Display optimization progress information."""
165 | console.rule(f"[bold magenta]Optimization - Iteration {iteration}/{total}[/]")
166 | console.print(f"Best objective value: [bold green]{best_value:.6g}[/]")
167 | if parameters:
168 | table = Table(title="Best Parameters")
169 | table.add_column("Parameter", style="cyan")
170 | table.add_column("Value", style="green")
171 | for key, value in parameters.items():
172 | if isinstance(value, float): table.add_row(str(key), f"{value:.6g}")
173 | else: table.add_row(str(key), str(value))
174 | console.print(table)
175 |
176 | def display_time_elapsed(start_time: datetime.datetime) -> None:
177 | """Display the time elapsed since the start time."""
178 | elapsed = datetime.datetime.now() - start_time
179 | hours, remainder = divmod(elapsed.total_seconds(), 3600)
180 | minutes, seconds = divmod(remainder, 60)
181 | time_str = f"[bold]Time elapsed:[/] "
182 | if hours > 0: time_str += f"{int(hours)}h "
183 | if minutes > 0 or hours > 0: time_str += f"{int(minutes)}m "
184 | time_str += f"{seconds:.1f}s"
185 | console.print(time_str)
186 |
187 | def code_preview(code: str, language: str = "python") -> None:
188 | """Display formatted code with syntax highlighting."""
189 | syntax = Syntax(code, language, theme="monokai", line_numbers=True)
190 | console.print(syntax)
191 |
192 | def tree_view(data: Dict[str, Any], title: str = "Structure") -> None:
193 | """
194 | Display nested data in a tree view."""
195 | tree = Tree(f"[bold]{title}[/]")
196 | def _add_to_tree(tree_node, data_node):
197 | if isinstance(data_node, dict):
198 | for key, value in data_node.items():
199 | if isinstance(value, (dict, list)):
200 | branch = tree_node.add(f"[blue]{key}[/]")
201 | _add_to_tree(branch, value)
202 | else: tree_node.add(f"[blue]{key}:[/] {value}")
203 | elif isinstance(data_node, list):
204 | for i, item in enumerate(data_node):
205 | if isinstance(item, (dict, list)):
206 | branch = tree_node.add(f"[green]{i}[/]")
207 | _add_to_tree(branch, item)
208 | else: tree_node.add(f"[green]{i}:[/] {item}")
209 | _add_to_tree(tree, data)
210 | console.print(tree)
--------------------------------------------------------------------------------
/beamz/simulation/boundaries.py:
--------------------------------------------------------------------------------
1 | import numpy as np
2 | from beamz.const import µm, EPS_0, MU_0
3 |
4 | class Boundary:
5 | """Abstract base class for all boundary conditions."""
6 | def __init__(self, edges, thickness):
7 | """
8 | Args:
9 | edges: list of edge names or 'all'
10 | 2D: ['left', 'right', 'top', 'bottom']
11 | 3D: ['left', 'right', 'top', 'bottom', 'front', 'back']
12 | thickness: physical thickness of boundary region
13 | """
14 | if edges == 'all':
15 | # Will be determined based on dimensionality in apply()
16 | self.edges = 'all'
17 | else:
18 | self.edges = edges if isinstance(edges, list) else [edges]
19 | self.thickness = thickness
20 |
21 | def apply(self, fields, design, resolution, dt):
22 | """Apply boundary condition to fields. Must be implemented by subclasses."""
23 | raise NotImplementedError
24 |
25 | def _get_edges_for_dimensionality(self, is_3d):
26 | """Resolve 'all' edges based on dimensionality."""
27 | if self.edges == 'all':
28 | return ['left', 'right', 'top', 'bottom', 'front', 'back'] if is_3d else ['left', 'right', 'top', 'bottom']
29 | return self.edges
30 |
31 | class PML(Boundary):
32 | """Perfectly Matched Layer boundary condition for FDTD simulations."""
33 |
34 | def __init__(self, edges='all', thickness=1*µm, sigma_max=None, m=3, kappa_max=1, alpha_max=0):
35 | """
36 | Initialize UPML with stretched-coordinate parameters.
37 |
38 | Args:
39 | edges: edges to apply PML
40 | thickness: PML thickness
41 | sigma_max: maximum conductivity (auto-calculated if None)
42 | m: conductivity grading order
43 | kappa_max: maximum real coordinate stretching
44 | alpha_max: maximum CFS alpha parameter (for better absorption at low frequencies)
45 | """
46 | super().__init__(edges, thickness)
47 | self.sigma_max = sigma_max
48 | self.m = m
49 | self.kappa_max = kappa_max
50 | self.alpha_max = alpha_max
51 |
52 | def apply(self, fields, design, resolution, dt):
53 | """Apply PML by modifying field update equations with PML conductivity."""
54 | # This method is now deprecated - use modify_conductivity instead
55 | pass
56 |
57 | def create_pml_regions(self, fields, design, resolution, dt, plane_2d='xy'):
58 | """Create permanent PML region masks and stretched-coordinate parameters.
59 |
60 | Returns dict with:
61 | - mask: boolean arrays indicating PML cells
62 | - sigma_x, sigma_y, sigma_z: conductivity profiles
63 | - kappa_x, kappa_y, kappa_z: real stretching factors
64 | - alpha_x, alpha_y, alpha_z: CFS alpha parameters
65 | """
66 | # Calculate optimal sigma_max if not provided
67 | if self.sigma_max is None:
68 | eta = np.sqrt(MU_0 / (EPS_0 * 1.0))
69 | self.sigma_max = 0.8 * (self.m + 1) / (eta * resolution)
70 |
71 | # Create graded profiles for each direction based on plane
72 | pml_data = self._create_pml_profiles_2d(fields, design, resolution, dt, plane_2d)
73 | return pml_data
74 |
75 | def _create_pml_profiles_2d(self, fields, design, resolution, dt, plane_2d):
76 | """Create UPML stretched-coordinate profiles for 2D plane."""
77 | # Grid shape from material grid (same as field shape for collocated/main grid)
78 | # Assuming design.width/height/depth match the grid dimensions
79 |
80 | if plane_2d == 'xy':
81 | shape = fields.permittivity.shape # (ny, nx)
82 | dim1, dim2 = shape
83 | len1, len2 = design.height, design.width
84 | labels = ['y', 'x']
85 | elif plane_2d == 'yz':
86 | shape = fields.permittivity.shape # (nz, ny)
87 | dim1, dim2 = shape
88 | len1, len2 = (design.depth if design.depth else 0), design.height # Assuming depth is defined for yz slice?
89 | # If 2D sim in yz, design.depth might be relevant or height/width mapping changes.
90 | # Assuming standard: design.width (x), design.height (y), design.depth (z)
91 | labels = ['z', 'y']
92 | elif plane_2d == 'xz':
93 | shape = fields.permittivity.shape # (nz, nx)
94 | dim1, dim2 = shape
95 | len1, len2 = (design.depth if design.depth else 0), design.width
96 | labels = ['z', 'x']
97 |
98 | # Initialize profile arrays
99 | profiles = {}
100 | for axis in ['x', 'y', 'z']:
101 | profiles[f'sigma_{axis}'] = np.zeros(shape)
102 | profiles[f'kappa_{axis}'] = np.ones(shape)
103 | profiles[f'alpha_{axis}'] = np.zeros(shape)
104 |
105 | # Create coordinate arrays
106 | coords1 = np.linspace(0, len1, dim1) # axis 0 (y or z)
107 | coords2 = np.linspace(0, len2, dim2) # axis 1 (x or y)
108 |
109 | edges = self._get_edges_for_dimensionality(False)
110 |
111 | # Map edges to axes based on plane
112 | # xy: Left/Right -> x (axis 1). Bottom/Top -> y (axis 0).
113 | # yz: Left/Right -> y (axis 1)? Or z? Usually Left/Right is horizontal on screen.
114 | # yz plane: horizontal=y, vertical=z? Or horizontal=y, vertical=z.
115 | # Let's assume consistent mapping:
116 | # dim2 is "horizontal" (second index), dim1 is "vertical" (first index).
117 | # xy: x is horizontal (dim2), y is vertical (dim1).
118 | # yz: y is horizontal (dim2), z is vertical (dim1).
119 | # xz: x is horizontal (dim2), z is vertical (dim1).
120 |
121 | # Edges mapping:
122 | # Left/Right -> dim2 (coords2)
123 | # Bottom/Top -> dim1 (coords1)
124 |
125 | # Determine which sigma/kappa/alpha component to set
126 | # xy: dim2->x, dim1->y
127 | # yz: dim2->y, dim1->z
128 | # xz: dim2->x, dim1->z
129 |
130 | axis1 = labels[0] # vertical axis name
131 | axis2 = labels[1] # horizontal axis name
132 |
133 | for edge in edges:
134 | if edge == 'left': # Start of horizontal axis (dim2)
135 | for i in range(dim2):
136 | if coords2[i] < self.thickness:
137 | dist = self.thickness - coords2[i]
138 | profiles[f'sigma_{axis2}'][:, i] = self._sigma_profile(dist, self.thickness)
139 | profiles[f'kappa_{axis2}'][:, i] = self._kappa_profile(dist, self.thickness)
140 | profiles[f'alpha_{axis2}'][:, i] = self._alpha_profile(dist, self.thickness)
141 |
142 | elif edge == 'right': # End of horizontal axis (dim2)
143 | for i in range(dim2):
144 | if coords2[i] > (len2 - self.thickness):
145 | dist = coords2[i] - (len2 - self.thickness)
146 | profiles[f'sigma_{axis2}'][:, i] = self._sigma_profile(dist, self.thickness)
147 | profiles[f'kappa_{axis2}'][:, i] = self._kappa_profile(dist, self.thickness)
148 | profiles[f'alpha_{axis2}'][:, i] = self._alpha_profile(dist, self.thickness)
149 |
150 | elif edge == 'bottom': # Start of vertical axis (dim1)
151 | for j in range(dim1):
152 | if coords1[j] < self.thickness:
153 | dist = self.thickness - coords1[j]
154 | profiles[f'sigma_{axis1}'][j, :] = self._sigma_profile(dist, self.thickness)
155 | profiles[f'kappa_{axis1}'][j, :] = self._kappa_profile(dist, self.thickness)
156 | profiles[f'alpha_{axis1}'][j, :] = self._alpha_profile(dist, self.thickness)
157 |
158 | elif edge == 'top': # End of vertical axis (dim1)
159 | for j in range(dim1):
160 | if coords1[j] > (len1 - self.thickness):
161 | dist = coords1[j] - (len1 - self.thickness)
162 | profiles[f'sigma_{axis1}'][j, :] = self._sigma_profile(dist, self.thickness)
163 | profiles[f'kappa_{axis1}'][j, :] = self._kappa_profile(dist, self.thickness)
164 | profiles[f'alpha_{axis1}'][j, :] = self._alpha_profile(dist, self.thickness)
165 |
166 | # Create PML mask (True where any PML sigma is active)
167 | pml_mask = (profiles[f'sigma_{axis1}'] > 0) | (profiles[f'sigma_{axis2}'] > 0)
168 |
169 | result = {'mask': pml_mask}
170 | result.update(profiles)
171 | return result
172 |
173 | def _modify_conductivity_3d(self, fields, design, resolution, dt, edges):
174 | """Modify conductivity arrays to include PML absorption in 3D."""
175 | # TODO: Implement 3D PML conductivity modification
176 | # For now, just pass to avoid breaking 3D simulations
177 | pass
178 |
179 | def _sigma_profile(self, dist, thickness):
180 | """Graded conductivity profile."""
181 | return self.sigma_max * (dist / thickness) ** self.m
182 |
183 | def _kappa_profile(self, dist, thickness):
184 | """Real coordinate stretching profile."""
185 | return 1 + (self.kappa_max - 1) * (dist / thickness) ** self.m
186 |
187 | def _alpha_profile(self, dist, thickness):
188 | """CFS alpha profile for low-frequency absorption."""
189 | return self.alpha_max * ((thickness - dist) / thickness) ** self.m
190 |
191 | class ABC(Boundary):
192 | """Absorbing Boundary Condition (Mur, Liao, etc.) - placeholder."""
193 | def apply(self, fields, design, resolution, dt):
194 | pass # TODO: implement
195 |
196 | class PeriodicBoundary(Boundary):
197 | """Periodic boundary condition - placeholder."""
198 | def apply(self, fields, design, resolution, dt):
199 | pass # TODO: implement
--------------------------------------------------------------------------------
/beamz/simulation/core.py:
--------------------------------------------------------------------------------
1 | import numpy as np
2 | from beamz.const import *
3 | from beamz.design.core import Design
4 | from beamz.devices.core import Device
5 | from beamz.simulation.fields import Fields
6 | from beamz.simulation.boundaries import Boundary, PML
7 | from beamz.visual.viz import animate_manual_field, close_fdtd_figure
8 |
9 | class Simulation:
10 | """FDTD simulation class supporting both 2D and 3D electromagnetic simulations."""
11 | def __init__(self,
12 | design:Design=None,
13 | devices:list[Device]=[],
14 | boundaries:list[Boundary]=[],
15 | resolution:float=0.02*µm,
16 | time:np.ndarray=None,
17 | plane_2d:str='xy'):
18 | self.design = design
19 | self.resolution = resolution
20 | self.is_3d = design.is_3d and design.depth > 0
21 | self.plane_2d = plane_2d.lower()
22 | if self.plane_2d not in ['xy', 'yz', 'xz']: self.plane_2d = 'xy'
23 |
24 | # Get material grids from design (design owns the material grids, we reference them)
25 | permittivity, conductivity, permeability = design.get_material_grids(resolution)
26 |
27 | # Initialize time stepping first
28 | if time is None or len(time) < 2: raise ValueError("FDTD requires a time array with at least two entries")
29 | self.time, self.dt, self.num_steps = time, float(time[1] - time[0]), len(time)
30 | self.t, self.current_step = 0, 0
31 |
32 | # Create field storage (fields owns the E/H field arrays, references material grids)
33 | self.fields = Fields(permittivity, conductivity, permeability, resolution, plane_2d=self.plane_2d)
34 |
35 | # Initialize PML regions if present
36 | pml_boundaries = [b for b in boundaries if isinstance(b, PML)]
37 | if pml_boundaries:
38 | # Create PML regions (do this once, not every timestep)
39 | pml_data = {}
40 | for pml in pml_boundaries:
41 | pml_data.update(pml.create_pml_regions(self.fields, design, resolution, self.dt, plane_2d=self.plane_2d))
42 | self.pml_data = pml_data
43 |
44 | # Initialize split fields in Fields object - DEPRECATED/REMOVED in favor of effective conductivity
45 | # self.fields._init_upml_fields(pml_data)
46 |
47 | # Set effective conductivity for PML
48 | self.fields.set_pml_conductivity(pml_data)
49 | else:
50 | self.pml_data = None
51 |
52 | # Store device references (no duplication)
53 | self.devices = devices
54 |
55 | # Store boundary references (no duplication)
56 | self.boundaries = boundaries
57 |
58 | def step(self):
59 | """Perform one FDTD time step."""
60 | if self.current_step >= self.num_steps: return False
61 |
62 | # Inject source fields (if any) directly into the grid before update
63 | self._inject_sources()
64 |
65 | # Collect source terms from legacy devices (if any)
66 | source_j, source_m = self._collect_source_terms()
67 |
68 | # Update fields (legacy sources passed, new sources already injected)
69 | self.fields.update(self.dt, source_j=source_j, source_m=source_m)
70 |
71 | # Record monitor data (if monitors are in devices)
72 | self._record_monitors()
73 |
74 | # Update time and step counter
75 | self.t += self.dt
76 | self.current_step += 1
77 | return True
78 |
79 | def _record_monitors(self):
80 | """Record data from Monitor devices during simulation."""
81 | for device in self.devices:
82 | if hasattr(device, 'should_record') and hasattr(device, 'record_fields'):
83 | if device.should_record(self.current_step):
84 | if not self.is_3d:
85 | device.record_fields(self.fields.Ez, self.fields.Hx, self.fields.Hy,
86 | self.t, self.resolution, self.resolution, self.current_step)
87 | else:
88 | device.record_fields(self.fields.Ex, self.fields.Ey, self.fields.Ez,
89 | self.fields.Hx, self.fields.Hy, self.fields.Hz,
90 | self.t, self.resolution, self.resolution, self.resolution, self.current_step)
91 |
92 | def _inject_sources(self):
93 | """Inject source fields directly into the simulation grid."""
94 | for device in self.devices:
95 | if hasattr(device, 'inject'):
96 | device.inject(self.fields, self.t, self.dt, self.current_step, self.resolution, self.design)
97 |
98 | def _collect_source_terms(self):
99 | """Collect electric and magnetic current sources from all devices."""
100 | source_j = {} # Electric currents for E-field update
101 | source_m = {} # Magnetic currents for H-field update
102 |
103 | for device in self.devices:
104 | if hasattr(device, 'get_source_terms'):
105 | j, m = device.get_source_terms(self.fields, self.t, self.dt, self.current_step, self.resolution, self.design)
106 | if j: source_j.update(j)
107 | if m: source_m.update(m)
108 |
109 | return source_j, source_m
110 |
111 |
112 | def run(self, animate_live=None, animation_interval=10, axis_scale=None, cmap='twilight_zero', clean_visualization=False, wavelength=None, line_color='gray', line_opacity=0.5, save_fields=None, field_subsample=1):
113 | """Run complete FDTD simulation with optional live field visualization.
114 |
115 | Args:
116 | animate_live: Field component to animate ('Ez', 'Hx', 'Hy', 'Ex', 'Ey', etc.) or None to disable
117 | animation_interval: Update visualization every N steps (higher = faster but less smooth)
118 | axis_scale: Tuple (min, max) for fixed color scale during animation, or None for auto-scaling
119 | cmap: Matplotlib colormap name (default: 'twilight_zero')
120 | clean_visualization: If True, hide axes, title, and colorbar (only show field and structures)
121 | wavelength: Wavelength for scale bar calculation (if None, tries to extract from devices)
122 | line_color: Color for structure and PML boundary outlines (default: 'gray')
123 | line_opacity: Opacity/transparency of structure and PML boundary outlines (0.0 to 1.0, default: 0.5)
124 | save_fields: List of field components to save ('Ez', 'Hx', etc.) or None to disable
125 | field_subsample: Save fields every N steps (default: 1, save all steps)
126 |
127 | Returns:
128 | dict with keys:
129 | - 'fields': dict of field histories if save_fields was provided
130 | - 'monitors': list of Monitor objects with recorded data
131 | """
132 | # Handle 3D simulations - require monitor for now (not implemented yet)
133 | if animate_live and self.is_3d:
134 | print("Live animation for 3D simulations requires a monitor (not yet implemented)")
135 | animate_live = None
136 |
137 | # Initialize animation context if requested
138 | viz_context = None
139 | if animate_live:
140 | # Validate field component exists
141 | if animate_live not in self.fields.available_components():
142 | print(f"Warning: Field '{animate_live}' not found. Available: {self.fields.available_components()}")
143 | animate_live = None
144 |
145 | # Extract wavelength from devices if not provided
146 | if wavelength is None:
147 | for device in self.devices:
148 | if hasattr(device, 'wavelength'):
149 | wavelength = device.wavelength
150 | break
151 |
152 | # Initialize field storage if requested
153 | field_history = {}
154 | if save_fields:
155 | for field_name in save_fields:
156 | field_history[field_name] = []
157 |
158 | try:
159 | # Main simulation loop
160 | while self.step():
161 | # Save field history if requested
162 | # current_step is incremented in step(), so we check after increment
163 | if save_fields and (self.current_step % field_subsample == 0):
164 | for field_name in save_fields:
165 | if hasattr(self.fields, field_name):
166 | field_history[field_name].append(getattr(self.fields, field_name).copy())
167 |
168 | # Update live animation if enabled
169 | if animate_live and self.current_step % animation_interval == 0:
170 | field_data = getattr(self.fields, animate_live)
171 | # Convert to V/µm for display
172 | field_display = field_data * 1e-6 if 'E' in animate_live else field_data
173 | extent = (0, self.design.width, 0, self.design.height)
174 | title = f'{animate_live} at t = {self.t:.2e} s (step {self.current_step}/{self.num_steps})'
175 | viz_context = animate_manual_field(field_display, context=viz_context, extent=extent,
176 | title=title, units='V/µm' if 'E' in animate_live else 'A/m',
177 | design=self.design, boundaries=self.boundaries, pause=0.001,
178 | axis_scale=axis_scale, cmap=cmap, clean_visualization=clean_visualization,
179 | wavelength=wavelength, line_color=line_color, line_opacity=line_opacity,
180 | plane_2d=self.plane_2d)
181 | finally:
182 | # Cleanup: keep the final frame visible
183 | if viz_context and viz_context.get('fig'):
184 | import matplotlib.pyplot as plt
185 | plt.show(block=False)
186 | print("Simulation complete. Close the plot window to continue.")
187 |
188 | # Collect monitor data
189 | monitors = [device for device in self.devices if hasattr(device, 'power_history')]
190 |
191 | # Return results if requested
192 | result = {}
193 | if save_fields:
194 | result['fields'] = field_history
195 | if monitors:
196 | result['monitors'] = monitors
197 |
198 | return result if result else None
--------------------------------------------------------------------------------
/beamz/optimization/topology.py:
--------------------------------------------------------------------------------
1 | """Topology optimization manager and helpers."""
2 |
3 | from __future__ import annotations
4 |
5 | import numpy as np
6 | from typing import Optional, Tuple
7 | from beamz.const import LIGHT_SPEED
8 |
9 | from .autodiff import transform_density, compute_parameter_gradient_vjp
10 |
11 | # Defer imports to avoid circular dependencies if any,
12 | # or import at top level if safe. design shouldn't depend on optimization.
13 | from beamz.design.core import Design
14 | from beamz.design.materials import Material
15 | from beamz.design.meshing import RegularGrid
16 |
17 | class TopologyManager:
18 | """
19 | High-level manager for topology optimization.
20 |
21 | Handles:
22 | - Density parameter storage
23 | - Physical density transformation (JAX-based)
24 | - Gradient backpropagation (JAX-based)
25 | - Optimizer stepping
26 | - Material grid updates
27 | """
28 |
29 | def __init__(
30 | self,
31 | design,
32 | region_mask: np.ndarray,
33 | optimizer: str = "Adam",
34 | learning_rate: float = 0.1,
35 | filter_radius: float = 0.0,
36 | simple_smooth_radius: float = 0.0,
37 | projection_eta: float = 0.5,
38 | beta_schedule: tuple[float, float] = (1.0, 20.0),
39 | eps_min: float = 1.0,
40 | eps_max: float = 12.0,
41 | resolution: float = None,
42 | filter_type: str = 'conic', # 'blur', 'morphological', or 'conic'
43 | morphology_operation: str = 'openclose', # 'opening', 'closing', 'openclose'
44 | **kwargs
45 | ):
46 | """
47 | Args:
48 | filter_radius: Primary filter radius in physical units (e.g. microns).
49 | For 'conic', this enforces minimum feature size.
50 | simple_smooth_radius: Optional post-filter smoothing radius in physical units.
51 | Use to remove grid artifacts (e.g. 0.02 * um).
52 | filter_type: 'conic' (geometric constraints), 'morphological', or 'blur'.
53 | morphology_operation: 'opening', 'closing', or 'openclose' (for morphological filter).
54 | """
55 | self.design = design
56 | self.mask = region_mask.astype(bool)
57 |
58 | # Setup optimizer using optax (JAX-native)
59 | try:
60 | import optax
61 | except ImportError:
62 | raise ImportError("optax is required for optimization. Install with: pip install optax")
63 |
64 | if optimizer.lower() == "adam":
65 | self.optax_optimizer = optax.adam(learning_rate=learning_rate)
66 | elif optimizer.lower() == "sgd":
67 | self.optax_optimizer = optax.sgd(learning_rate=learning_rate)
68 | else:
69 | raise ValueError(f"Unknown optimizer '{optimizer}'. Supported: 'adam', 'sgd'")
70 |
71 | # Initialize optimizer state (will be created on first use)
72 | self._opt_state = None
73 |
74 | # Parameters
75 | self.filter_radius = filter_radius
76 | self.simple_smooth_radius = simple_smooth_radius
77 | self.projection_eta = projection_eta
78 | self.beta_start, self.beta_end = beta_schedule
79 | self.eps_min = eps_min
80 | self.eps_max = eps_max
81 | self.resolution = resolution or getattr(design.rasterize(resolution=0.1), "dx") # Fallback resolution check?
82 |
83 | # Filter settings
84 | self.filter_type = filter_type
85 | self.morphology_operation = morphology_operation
86 | self.morphology_smooth_tau = kwargs.get('morphology_smooth_tau', 0.01) # Default good value
87 |
88 | # Convert filter radius to cells
89 | self.filter_radius_cells = int(round(filter_radius / self.resolution)) if self.resolution else 0
90 | self.smooth_radius_cells = int(round(simple_smooth_radius / self.resolution)) if self.resolution else 0
91 |
92 | # Initialize density parameters (0.5 inside mask)
93 | self.design_density = np.zeros_like(self.mask, dtype=float)
94 | self.design_density[self.mask] = 0.5
95 |
96 | # Store base grid for fixed structure detection
97 | self.base_grid = design.rasterize(self.resolution)
98 | self.fixed_structure_mask = get_fixed_structure_mask(
99 | self.base_grid, self.eps_min, self.eps_max, self.mask
100 | )
101 |
102 | # History
103 | self.objective_history = []
104 |
105 | def get_current_beta(self, step: int, total_steps: int) -> float:
106 | """Calculate projection beta for current step."""
107 | if total_steps <= 1: return self.beta_end
108 | frac = step / (total_steps - 1)
109 | return self.beta_start + frac * (self.beta_end - self.beta_start)
110 |
111 | def get_physical_density(self, beta: float) -> np.ndarray:
112 | """Compute physical density from design parameters using JAX transform."""
113 | import jax.numpy as jnp
114 | d_jax = jnp.array(self.design_density)
115 | m_jax = jnp.array(self.mask)
116 | fixed_jax = jnp.array(self.fixed_structure_mask) if self.fixed_structure_mask is not None else None
117 |
118 | p_jax = transform_density(
119 | d_jax, m_jax,
120 | beta, self.projection_eta, self.filter_radius_cells,
121 | filter_type=self.filter_type,
122 | morphology_operation=self.morphology_operation,
123 | morphology_tau=self.morphology_smooth_tau,
124 | fixed_structure_mask=fixed_jax,
125 | post_smooth_radius=self.smooth_radius_cells
126 | )
127 | return np.array(p_jax)
128 |
129 | def update_design(self, step: int, total_steps: int) -> tuple[float, np.ndarray]:
130 | """
131 | Update the design's material grid based on current parameters.
132 | Returns (current_beta, physical_density).
133 | """
134 | beta = self.get_current_beta(step, total_steps)
135 | physical_density = self.get_physical_density(beta)
136 |
137 |
138 | return beta, physical_density
139 |
140 | def apply_gradient(self, grad_eps: np.ndarray, beta: float):
141 | """
142 | Apply gradient update:
143 | 1. Convert dJ/dEps -> dJ/dPhysical
144 | 2. Backprop dJ/dPhysical -> dJ/dParams (using JAX)
145 | 3. Optimizer step
146 | """
147 | import jax.numpy as jnp
148 |
149 | # dJ/dPhysical = dJ/dEps * (eps_max - eps_min)
150 | grad_physical = grad_eps * (self.eps_max - self.eps_min)
151 |
152 | fixed_jax = jnp.array(self.fixed_structure_mask) if self.fixed_structure_mask is not None else None
153 |
154 | # JAX Backprop
155 | grad_param_jax = compute_parameter_gradient_vjp(
156 | jnp.array(self.design_density),
157 | jnp.array(grad_physical),
158 | jnp.array(self.mask),
159 | beta,
160 | self.projection_eta,
161 | self.filter_radius_cells,
162 | filter_type=self.filter_type,
163 | morphology_operation=self.morphology_operation,
164 | morphology_tau=self.morphology_smooth_tau,
165 | fixed_structure_mask=fixed_jax,
166 | post_smooth_radius=self.smooth_radius_cells
167 | )
168 | grad_param = np.array(grad_param_jax)
169 |
170 | # Optimizer step (maximize objective -> ascent -> negative grad for minimizer)
171 | # Convert to JAX array for optax
172 | import jax.numpy as jnp
173 | grad_jax = jnp.array(-grad_param) # Negative because we want to maximize
174 |
175 | # Initialize optimizer state on first call
176 | if self._opt_state is None:
177 | params_init = jnp.array(self.design_density)
178 | self._opt_state = self.optax_optimizer.init(params_init)
179 |
180 | # Compute updates
181 | updates, self._opt_state = self.optax_optimizer.update(grad_jax, self._opt_state)
182 | update = np.array(updates)
183 |
184 | # Apply update
185 | self.design_density[self.mask] += update[self.mask]
186 | self.design_density = np.clip(self.design_density, 0.0, 1.0)
187 |
188 | return np.max(np.abs(update))
189 |
190 |
191 | def compute_overlap_gradient(forward_fields_history, adjoint_fields_history, field_key="Ez"):
192 | """
193 | Compute the gradient of the overlap integral with respect to epsilon.
194 | Gradient = Re(E_fwd * E_adj) integrated over time.
195 | """
196 | grad = np.zeros_like(forward_fields_history[0], dtype=float)
197 |
198 |
199 | n_steps = min(len(forward_fields_history), len(adjoint_fields_history))
200 |
201 | for i in range(n_steps):
202 |
203 | grad += forward_fields_history[i] * adjoint_fields_history[n_steps - 1 - i]
204 |
205 | return grad
206 |
207 | def create_optimization_mask(grid, region_structure):
208 | """
209 | Helper to create a boolean mask from a structure on a grid.
210 | Uses rasterization to ensure exact alignment with how structures are mapped to the grid.
211 | """
212 | # Create temp design to rasterize mask exactly as grid does
213 | temp_design = Design(width=grid.width, height=grid.height,
214 | material=Material(permittivity=1.0))
215 |
216 | # Copy structure to avoid modifying original
217 | if hasattr(region_structure, 'copy'):
218 | struct_copy = region_structure.copy()
219 | else:
220 | # Fallback if no copy method, reuse (risky if material modified, but we set it)
221 | struct_copy = region_structure
222 |
223 | # Set to a distinct permittivity to detect it
224 | struct_copy.material = Material(permittivity=2.0)
225 | temp_design.add(struct_copy)
226 |
227 | # Rasterize
228 | temp_grid = RegularGrid(temp_design, resolution=grid.dx)
229 |
230 | # Mask is where permittivity > background
231 | # Use a safe threshold to include any partial fill
232 | mask = temp_grid.permittivity > 1.001
233 |
234 | return mask
235 |
236 | def get_fixed_structure_mask(grid, eps_min, eps_max, design_mask):
237 | """
238 | Identify fixed structures (waveguides) from the base permittivity grid.
239 | Returns boolean mask where fixed solid material exists outside design region.
240 | """
241 | # Threshold for "solid" material. E.g. > 90% of core-clad difference
242 | threshold = eps_min + 0.9 * (eps_max - eps_min)
243 |
244 | # High permittivity regions
245 | high_eps = grid.permittivity >= threshold
246 |
247 | # Exclude design region (we only care about fixed structures outside)
248 | fixed_structures = high_eps & (~design_mask)
249 |
250 | return fixed_structures
251 |
--------------------------------------------------------------------------------
/beamz/design/materials.py:
--------------------------------------------------------------------------------
1 | # Medium: Dispersionless medium.
2 | class Material:
3 | def __init__(self, permittivity=1.0, permeability=1.0, conductivity=0.0):
4 | self.permittivity = permittivity
5 | self.permeability = permeability
6 | self.conductivity = conductivity
7 |
8 | def get_sample(self):
9 | return self.permittivity, self.permeability, self.conductivity
10 |
11 |
12 | # CustomMaterial: Function-based material for inverse design
13 | class CustomMaterial:
14 | def __init__(self, permittivity_func=None, permeability_func=None, conductivity_func=None,
15 | permittivity_grid=None, permeability_grid=None, conductivity_grid=None,
16 | bounds=None, interpolation='linear'):
17 | """
18 | Custom material with spatially-varying properties for inverse design.
19 |
20 | Args:
21 | permittivity_func: Function that takes (x, y) or (x, y, z) and returns permittivity
22 | permeability_func: Function that takes (x, y) or (x, y, z) and returns permeability
23 | conductivity_func: Function that takes (x, y) or (x, y, z) and returns conductivity
24 | permittivity_grid: 2D numpy array of permittivity values for grid-based interpolation
25 | permeability_grid: 2D numpy array of permeability values for grid-based interpolation
26 | conductivity_grid: 2D numpy array of conductivity values for grid-based interpolation
27 | bounds: Tuple ((x_min, x_max), (y_min, y_max)) defining the spatial bounds for grid interpolation
28 | interpolation: 'linear', 'cubic', or 'nearest' for grid interpolation
29 |
30 | Examples:
31 | # Function-based material
32 | def perm_func(x, y):
33 | return 2.0 + 0.5 * np.sin(x) * np.cos(y)
34 | material = CustomMaterial(permittivity_func=perm_func)
35 |
36 | # Grid-based material for inverse design
37 | perm_grid = np.ones((50, 50)) * 2.0
38 | perm_grid[20:30, 20:30] = 4.0 # High index region
39 | material = CustomMaterial(
40 | permittivity_grid=perm_grid,
41 | bounds=((0, 10e-6), (0, 10e-6)) # 10 micron x 10 micron
42 | )
43 | """
44 | import numpy as np
45 |
46 | # Store function-based definitions
47 | self.permittivity_func = permittivity_func
48 | self.permeability_func = permeability_func
49 | self.conductivity_func = conductivity_func
50 |
51 | # Store grid-based definitions
52 | self.permittivity_grid = permittivity_grid
53 | self.permeability_grid = permeability_grid
54 | self.conductivity_grid = conductivity_grid
55 |
56 | # Spatial bounds for grid interpolation
57 | self.bounds = bounds
58 | self.interpolation = interpolation
59 |
60 | # Default values
61 | self.default_permittivity = 1.0
62 | self.default_permeability = 1.0
63 | self.default_conductivity = 0.0
64 |
65 | # Create interpolation functions for grids
66 | if permittivity_grid is not None and bounds is not None:
67 | self._create_grid_interpolator('permittivity')
68 | if permeability_grid is not None and bounds is not None:
69 | self._create_grid_interpolator('permeability')
70 | if conductivity_grid is not None and bounds is not None:
71 | self._create_grid_interpolator('conductivity')
72 |
73 | @property
74 | def permittivity(self):
75 | """Return representative permittivity for display purposes."""
76 | if self.permittivity_grid is not None:
77 | import numpy as np
78 | return f"grid({np.min(self.permittivity_grid):.3f}-{np.max(self.permittivity_grid):.3f})"
79 | elif self.permittivity_func is not None:
80 | return "function"
81 | else:
82 | return self.default_permittivity
83 |
84 | @property
85 | def permeability(self):
86 | """Return representative permeability for display purposes."""
87 | if self.permeability_grid is not None:
88 | import numpy as np
89 | return f"grid({np.min(self.permeability_grid):.3f}-{np.max(self.permeability_grid):.3f})"
90 | elif self.permeability_func is not None:
91 | return "function"
92 | else:
93 | return self.default_permeability
94 |
95 | @property
96 | def conductivity(self):
97 | """Return representative conductivity for display purposes."""
98 | if self.conductivity_grid is not None:
99 | import numpy as np
100 | return f"grid({np.min(self.conductivity_grid):.3f}-{np.max(self.conductivity_grid):.3f})"
101 | elif self.conductivity_func is not None:
102 | return "function"
103 | else:
104 | return self.default_conductivity
105 |
106 | def _create_grid_interpolator(self, property_name):
107 | """Create scipy interpolator for grid-based material property."""
108 | try:
109 | from scipy.interpolate import RegularGridInterpolator
110 | import numpy as np
111 |
112 | grid = getattr(self, f'{property_name}_grid')
113 | if grid is None:
114 | return
115 |
116 | # Create coordinate arrays
117 | x_coords = np.linspace(self.bounds[0][0], self.bounds[0][1], grid.shape[1])
118 | y_coords = np.linspace(self.bounds[1][0], self.bounds[1][1], grid.shape[0])
119 |
120 | # Create interpolator
121 | interpolator = RegularGridInterpolator(
122 | (y_coords, x_coords), grid,
123 | method=self.interpolation,
124 | bounds_error=False,
125 | fill_value=getattr(self, f'default_{property_name}')
126 | )
127 |
128 | # Store interpolator
129 | setattr(self, f'_{property_name}_interpolator', interpolator)
130 |
131 | except ImportError:
132 | print("Warning: scipy not available, using nearest neighbor interpolation")
133 | setattr(self, f'_{property_name}_interpolator', None)
134 |
135 | def get_permittivity(self, x, y, z=None):
136 | """Get permittivity at spatial coordinates (x, y, z)."""
137 | if self.permittivity_func is not None:
138 | if z is not None:
139 | return self.permittivity_func(x, y, z)
140 | else:
141 | return self.permittivity_func(x, y)
142 | elif hasattr(self, '_permittivity_interpolator') and self._permittivity_interpolator is not None:
143 | import numpy as np
144 | points = np.column_stack([np.atleast_1d(y), np.atleast_1d(x)])
145 | return self._permittivity_interpolator(points)
146 | else:
147 | return self.default_permittivity
148 |
149 | def get_permeability(self, x, y, z=None):
150 | """Get permeability at spatial coordinates (x, y, z)."""
151 | if self.permeability_func is not None:
152 | if z is not None:
153 | return self.permeability_func(x, y, z)
154 | else:
155 | return self.permeability_func(x, y)
156 | elif hasattr(self, '_permeability_interpolator') and self._permeability_interpolator is not None:
157 | import numpy as np
158 | points = np.column_stack([np.atleast_1d(y), np.atleast_1d(x)])
159 | return self._permeability_interpolator(points)
160 | else:
161 | return self.default_permeability
162 |
163 | def get_conductivity(self, x, y, z=None):
164 | """Get conductivity at spatial coordinates (x, y, z)."""
165 | if self.conductivity_func is not None:
166 | if z is not None:
167 | return self.conductivity_func(x, y, z)
168 | else:
169 | return self.conductivity_func(x, y)
170 | elif hasattr(self, '_conductivity_interpolator') and self._conductivity_interpolator is not None:
171 | import numpy as np
172 | points = np.column_stack([np.atleast_1d(y), np.atleast_1d(x)])
173 | return self._conductivity_interpolator(points)
174 | else:
175 | return self.default_conductivity
176 |
177 | def get_sample(self, x=0, y=0, z=None):
178 | """Get material properties at spatial coordinates for backward compatibility."""
179 | return (self.get_permittivity(x, y, z),
180 | self.get_permeability(x, y, z),
181 | self.get_conductivity(x, y, z))
182 |
183 | def update_grid(self, property_name, new_grid):
184 | """Update material property grid (for optimization)."""
185 | if property_name == 'permittivity':
186 | self.permittivity_grid = new_grid
187 | self._create_grid_interpolator('permittivity')
188 | elif property_name == 'permeability':
189 | self.permeability_grid = new_grid
190 | self._create_grid_interpolator('permeability')
191 | elif property_name == 'conductivity':
192 | self.conductivity_grid = new_grid
193 | self._create_grid_interpolator('conductivity')
194 | else:
195 | raise ValueError(f"Unknown property: {property_name}")
196 |
197 | def copy(self):
198 | """Create a deep copy of the CustomMaterial."""
199 | import numpy as np
200 |
201 | # Deep copy grids if they exist
202 | perm_grid = self.permittivity_grid.copy() if self.permittivity_grid is not None else None
203 | permeability_grid = self.permeability_grid.copy() if self.permeability_grid is not None else None
204 | cond_grid = self.conductivity_grid.copy() if self.conductivity_grid is not None else None
205 |
206 | # Create new CustomMaterial with copied data
207 | return CustomMaterial(
208 | permittivity_func=self.permittivity_func, # Functions can be shared
209 | permeability_func=self.permeability_func,
210 | conductivity_func=self.conductivity_func,
211 | permittivity_grid=perm_grid, # Deep copied grids
212 | permeability_grid=permeability_grid,
213 | conductivity_grid=cond_grid,
214 | bounds=self.bounds, # Bounds can be shared (tuples are immutable)
215 | interpolation=self.interpolation
216 | )
217 |
218 |
219 | # ================================
220 |
221 | # PoleResidue: A dispersive medium described by the pole-residue pair model.
222 |
223 | # Lorentz: A dispersive medium described by the Lorentz model.
224 |
225 | # Sellmeier: A dispersive medium described by the Sellmeier model.
226 |
227 | # Drude: A dispersive medium described by the Drude model.
228 |
229 | # Debye: A dispersive medium described by the Debye model.
--------------------------------------------------------------------------------
/beamz/optimization/autodiff.py:
--------------------------------------------------------------------------------
1 | """JAX-based autodifferentiation helpers for topology optimization."""
2 |
3 | import jax
4 | import jax.numpy as jnp
5 | from jax.scipy.signal import convolve2d
6 | from functools import partial
7 |
8 | @partial(jax.jit, static_argnames=['radius'])
9 | def generate_conic_kernel(radius: int):
10 | """
11 | Generate a 2D conic kernel (linear decay).
12 | w(r) = max(0, 1 - r/R)
13 | """
14 | radius = int(max(1, radius))
15 | kernel_size = 2 * radius + 1
16 | center = radius
17 |
18 | # Create coordinate grids
19 | y, x = jnp.ogrid[-radius:radius+1, -radius:radius+1]
20 |
21 | # Calculate distance from center
22 | dist = jnp.sqrt(x**2 + y**2)
23 |
24 | # Conic weights: linear decay, 0 outside radius
25 | weights = jnp.maximum(0.0, 1.0 - dist / radius)
26 |
27 | # Normalize
28 | weights = weights / jnp.sum(weights)
29 |
30 | return weights
31 |
32 | @partial(jax.jit, static_argnames=['radius'])
33 | def masked_conic_filter(values, mask, radius: int):
34 | """
35 | Apply a masked conic filter (linear decay).
36 | Used for geometric constraints (minimum feature size).
37 | """
38 | radius = int(max(0, radius))
39 | if radius <= 0:
40 | return jnp.where(mask, values, 0.0), jnp.where(mask, 1.0, 1.0)
41 |
42 | masked_values = jnp.where(mask, values, 0.0)
43 | float_mask = mask.astype(float)
44 |
45 | # Create conic kernel
46 | kernel = generate_conic_kernel(radius)
47 |
48 | # Pad input
49 | padded_values = jnp.pad(masked_values, radius, mode='edge')
50 | padded_mask = jnp.pad(float_mask, radius, mode='constant', constant_values=0.0)
51 |
52 | # Convolve
53 | weighted_sum = convolve2d(padded_values, kernel, mode='valid')
54 | weights = convolve2d(padded_mask, kernel, mode='valid')
55 |
56 | # Avoid division by zero
57 | weights = jnp.where(weights == 0.0, 1.0, weights)
58 |
59 | filtered = weighted_sum / weights
60 | filtered = jnp.where(mask, filtered, 0.0)
61 |
62 | return filtered, weights
63 |
64 | @partial(jax.jit, static_argnames=['radius'])
65 | def masked_box_blur(values, mask, radius: int):
66 | """
67 | Apply a masked box blur using JAX convolutions.
68 | """
69 | radius = int(max(0, radius))
70 | if radius <= 0:
71 | return jnp.where(mask, values, 0.0), jnp.where(mask, 1.0, 1.0)
72 |
73 | masked_values = jnp.where(mask, values, 0.0)
74 | float_mask = mask.astype(float)
75 |
76 | # Create box kernel
77 | kernel_size = 2 * radius + 1
78 | kernel = jnp.ones((kernel_size, kernel_size))
79 |
80 | # Pad input to handle edges manually to match numpy 'edge' padding behavior roughly
81 | # For simplicity and efficiency in JAX, we use standard padding
82 | padded_values = jnp.pad(masked_values, radius, mode='edge')
83 | padded_mask = jnp.pad(float_mask, radius, mode='constant', constant_values=0.0)
84 |
85 | # Convolve
86 | weighted_sum = convolve2d(padded_values, kernel, mode='valid')
87 | weights = convolve2d(padded_mask, kernel, mode='valid')
88 |
89 | # Avoid division by zero
90 | weights = jnp.where(weights == 0.0, 1.0, weights)
91 |
92 | blurred = weighted_sum / weights
93 | blurred = jnp.where(mask, blurred, 0.0)
94 | weights = jnp.where(mask, weights, 1.0)
95 |
96 | return blurred, weights
97 |
98 | @jax.jit
99 | def smoothed_heaviside(value, beta, eta):
100 | """
101 | Smoothed Heaviside projection using tanh.
102 | """
103 | beta = jnp.maximum(beta, 1e-6)
104 | # Use tanh projection: (tanh(beta*eta) + tanh(beta*(x-eta))) / (tanh(beta*eta) + tanh(beta*(1-eta)))
105 | num = jnp.tanh(beta * eta) + jnp.tanh(beta * (value - eta))
106 | den = jnp.tanh(beta * eta) + jnp.tanh(beta * (1.0 - eta))
107 | return num / den
108 |
109 | @partial(jax.jit, static_argnames=['axis'])
110 | def smooth_max(x, axis=None, tau=0.1):
111 | """
112 | Smooth maximum approximation: tau * log(sum(exp(x/tau)))
113 | Also known as LogSumExp.
114 | """
115 | return tau * jax.scipy.special.logsumexp(x / tau, axis=axis)
116 |
117 | @partial(jax.jit, static_argnames=['axis'])
118 | def smooth_min(x, axis=None, tau=0.1):
119 | """
120 | Smooth minimum approximation: -smooth_max(-x)
121 | """
122 | return -smooth_max(-x, axis=axis, tau=tau)
123 |
124 | @partial(jax.jit, static_argnames=['radius'])
125 | def grayscale_erosion(values, radius, tau=0.05):
126 | """
127 | Grayscale erosion using smooth minimum filter with a disk structuring element.
128 | Uses 2D shifts to implement isotropic erosion.
129 | """
130 | radius = int(max(0, radius))
131 | if radius <= 0: return values
132 |
133 | # 2D shift helper
134 | def shift_2d(arr, dy, dx):
135 | if dy == 0 and dx == 0: return arr
136 |
137 | # Handle y shift
138 | if dy > 0: arr = jnp.pad(arr[:-dy, :], ((dy,0), (0,0)), mode='edge')
139 | elif dy < 0: arr = jnp.pad(arr[-dy:, :], ((0,-dy), (0,0)), mode='edge')
140 |
141 | # Handle x shift
142 | if dx > 0: arr = jnp.pad(arr[:, :-dx], ((0,0), (dx,0)), mode='edge')
143 | elif dx < 0: arr = jnp.pad(arr[:, -dx:], ((0,0), (0,-dx)), mode='edge')
144 |
145 | return arr
146 |
147 | # Generate disk offsets
148 | # This loop runs at trace time since radius is static
149 | shifts = []
150 | for dy in range(-radius, radius + 1):
151 | for dx in range(-radius, radius + 1):
152 | if dy*dy + dx*dx <= radius*radius:
153 | shifts.append((dy, dx))
154 |
155 | # Create stack of shifted images
156 | stack = jnp.stack([shift_2d(values, dy, dx) for dy, dx in shifts], axis=0)
157 |
158 | # Compute smooth min over the stack
159 | eroded = smooth_min(stack, axis=0, tau=tau)
160 |
161 | return eroded
162 |
163 | @partial(jax.jit, static_argnames=['radius'])
164 | def grayscale_dilation(values, radius, tau=0.05):
165 | """
166 | Grayscale dilation using smooth maximum filter.
167 | Separable implementation.
168 | """
169 | radius = int(max(0, radius))
170 | if radius <= 0: return values
171 |
172 | # Use relationship: Dilation(f) = -Erosion(-f)
173 | return -grayscale_erosion(-values, radius, tau)
174 |
175 | @partial(jax.jit, static_argnames=['radius'])
176 | def grayscale_opening(values, radius, tau=0.05):
177 | """Opening: Erosion followed by Dilation."""
178 | return grayscale_dilation(grayscale_erosion(values, radius, tau), radius, tau)
179 |
180 | @partial(jax.jit, static_argnames=['radius'])
181 | def grayscale_closing(values, radius, tau=0.05):
182 | """Closing: Dilation followed by Erosion."""
183 | return grayscale_erosion(grayscale_dilation(values, radius, tau), radius, tau)
184 |
185 | @partial(jax.jit, static_argnames=['radius', 'operation'])
186 | def masked_morphological_filter(values, mask, radius, operation='openclose', tau=0.05, fixed_structure_mask=None):
187 | """
188 | Apply masked morphological filtering.
189 |
190 | Args:
191 | values: Density field
192 | mask: Design region mask
193 | radius: Filter radius in cells
194 | operation: 'erosion', 'dilation', 'opening', 'closing', 'openclose' (opening then closing)
195 | tau: Smoothness temperature for differentiable min/max
196 | fixed_structure_mask: Optional boolean mask of fixed solid structures (e.g. waveguides)
197 | used to pad the filter input to prevent boundary erosion.
198 | """
199 | # Isolate design region values.
200 | # For morphology, boundaries are important.
201 |
202 | # Pad with fixed structures if provided
203 | # Treat fixed structures as solid (1.0) to provide context for erosion/dilation
204 | filter_input = values
205 | if fixed_structure_mask is not None:
206 | # We assume values is already density [0,1].
207 | # We override fixed structure locations with 1.0
208 | # NOTE: fixed_structure_mask should be a JAX array (tracer or concrete)
209 | filter_input = jnp.where(fixed_structure_mask, 1.0, values)
210 |
211 | # Apply filter to the padded/context-aware field
212 | filtered = filter_input
213 |
214 | if operation == 'erosion':
215 | filtered = grayscale_erosion(filtered, radius, tau)
216 | elif operation == 'dilation':
217 | filtered = grayscale_dilation(filtered, radius, tau)
218 | elif operation == 'opening':
219 | filtered = grayscale_opening(filtered, radius, tau)
220 | elif operation == 'closing':
221 | filtered = grayscale_closing(filtered, radius, tau)
222 | elif operation == 'openclose':
223 | # Opening then Closing is a standard noise removal filter
224 | filtered = grayscale_closing(grayscale_opening(filtered, radius, tau), radius, tau)
225 |
226 | # Re-apply mask constraints
227 | # We only care about the result inside the design region
228 | return jnp.where(mask, filtered, 0.0)
229 |
230 | @partial(jax.jit, static_argnames=['radius', 'filter_type', 'morphology_operation', 'post_smooth_radius'])
231 | def transform_density(density, mask, beta, eta, radius, filter_type='blur', morphology_operation='openclose', morphology_tau=0.05, fixed_structure_mask=None, post_smooth_radius=0):
232 | """
233 | Full density transform: Filter -> Project.
234 | Returns the physical density [0, 1].
235 |
236 | Args:
237 | filter_type: 'blur', 'morphological', or 'conic'
238 | morphology_operation: 'opening', 'closing', 'openclose'
239 | fixed_structure_mask: Optional mask for fixed structures (morphological filter only)
240 | post_smooth_radius: Optional blur after morphology or filter
241 | """
242 | if filter_type == 'morphological':
243 | # Morphological filter
244 | filtered = masked_morphological_filter(density, mask, radius, morphology_operation, morphology_tau, fixed_structure_mask)
245 | elif filter_type == 'conic':
246 | # Conic filter (for geometric constraints)
247 | filtered, _ = masked_conic_filter(density, mask, radius)
248 | else:
249 | # Standard box blur
250 | filtered, _ = masked_box_blur(density, mask, radius)
251 |
252 | # Universal Post-Smoothing: Apply a small blur to smooth edges/artifacts
253 | if post_smooth_radius > 0:
254 | smoothed, _ = masked_box_blur(filtered, mask, post_smooth_radius)
255 | filtered = smoothed
256 |
257 | projected = smoothed_heaviside(filtered, beta, eta)
258 | return jnp.where(mask, projected, 0.0)
259 |
260 | @partial(jax.jit, static_argnames=['radius', 'filter_type', 'morphology_operation', 'post_smooth_radius'])
261 | def compute_parameter_gradient_vjp(density, grad_physical, mask, beta, eta, radius, filter_type='blur', morphology_operation='openclose', morphology_tau=0.05, fixed_structure_mask=None, post_smooth_radius=0):
262 | """
263 | Compute gradient w.r.t. design density using VJP.
264 | Supports both blur and morphological filters.
265 | """
266 | # Define a wrapper for the transform to differentiate
267 | def transform_wrapper(d):
268 | return transform_density(d, mask, beta, eta, radius, filter_type, morphology_operation, morphology_tau, fixed_structure_mask, post_smooth_radius)
269 |
270 | # Compute VJP
271 | _, vjp_fun = jax.vjp(transform_wrapper, density)
272 | grad_density = vjp_fun(grad_physical)[0]
273 |
274 | return grad_density
275 |
--------------------------------------------------------------------------------
/beamz/devices/sources/solve.py:
--------------------------------------------------------------------------------
1 | # Adapted from FDTDx by Yannik Mahlau
2 | from collections import namedtuple
3 | from types import SimpleNamespace
4 | from typing import List, Literal, Tuple, Union
5 |
6 | import numpy as np
7 | import tidy3d
8 | from tidy3d.components.mode.solver import compute_modes as _compute_modes
9 |
10 | ModeTupleType = namedtuple("Mode", ["neff", "Ex", "Ey", "Ez", "Hx", "Hy", "Hz"])
11 | """A named tuple containing the mode fields and effective index."""
12 |
13 | def compute_mode_polarization_fraction(
14 | mode: ModeTupleType,
15 | tangential_axes: tuple[int, int],
16 | pol: Literal["te", "tm"],
17 | ) -> float:
18 | E_fields = [mode.Ex, mode.Ey, mode.Ez]
19 | E1 = E_fields[tangential_axes[0]]
20 | E2 = E_fields[tangential_axes[1]]
21 |
22 | if pol == "te": numerator = np.sum(np.abs(E1) ** 2)
23 | elif pol == "tm": numerator = np.sum(np.abs(E2) ** 2)
24 | else: raise ValueError(f"pol must be 'te' or 'tm', but got {pol}")
25 |
26 | denominator = np.sum(np.abs(E1) ** 2 + np.abs(E2) ** 2) + 1e-18
27 | return numerator / denominator
28 |
29 | def sort_modes(
30 | modes: list[ModeTupleType],
31 | filter_pol: Union[Literal["te", "tm"], None],
32 | tangential_axes: tuple[int, int],
33 | ) -> list[ModeTupleType]:
34 | if filter_pol is None:
35 | return sorted(modes, key=lambda m: float(np.real(m.neff)), reverse=True)
36 |
37 | def is_matching(mode: ModeTupleType) -> bool:
38 | frac = compute_mode_polarization_fraction(mode, tangential_axes, filter_pol)
39 | return frac >= 0.5
40 |
41 | matching = [m for m in modes if is_matching(m)]
42 | non_matching = [m for m in modes if not is_matching(m)]
43 |
44 | matching_sorted = sorted(matching, key=lambda m: float(np.real(m.neff)), reverse=True)
45 | non_matching_sorted = sorted(non_matching, key=lambda m: float(np.real(m.neff)), reverse=True)
46 |
47 | return matching_sorted + non_matching_sorted
48 |
49 | def compute_mode(
50 | frequency: float,
51 | inv_permittivities: np.ndarray,
52 | inv_permeabilities: Union[np.ndarray, float],
53 | resolution: float,
54 | direction: Literal["+", "-"],
55 | mode_index: int = 0,
56 | filter_pol: Union[Literal["te", "tm"], None] = None,
57 | target_neff: Union[float, None] = None,
58 | ) -> tuple[np.ndarray, np.ndarray, complex, int]:
59 | inv_permittivities = np.asarray(inv_permittivities, dtype=np.complex128)
60 | if inv_permittivities.ndim == 1: inv_permittivities = inv_permittivities[np.newaxis, :, np.newaxis]
61 | elif inv_permittivities.ndim == 2: inv_permittivities = inv_permittivities[np.newaxis, :, :]
62 | elif inv_permittivities.ndim > 3: raise ValueError(f"Invalid shape of inv_permittivities: {inv_permittivities.shape}")
63 |
64 | if isinstance(inv_permeabilities, np.ndarray):
65 | inv_permeabilities = np.asarray(inv_permeabilities, dtype=np.complex128)
66 | if inv_permeabilities.ndim == 1: inv_permeabilities = inv_permeabilities[np.newaxis, :, np.newaxis]
67 | elif inv_permeabilities.ndim == 2: inv_permeabilities = inv_permeabilities[np.newaxis, :, :]
68 | elif inv_permeabilities.ndim > 3: raise ValueError(f"Invalid shape of inv_permeabilities: {inv_permeabilities.shape}")
69 | else:
70 | inv_permeabilities = np.asarray(inv_permeabilities, dtype=np.complex128)
71 |
72 | singleton_axes = [idx for idx, size in enumerate(inv_permittivities.shape) if size == 1]
73 | if not singleton_axes: raise ValueError("At least one singleton dimension is required to denote the propagation axis")
74 | propagation_axis = singleton_axes[0]
75 |
76 | cross_axes = [ax for ax in range(inv_permittivities.ndim) if ax != propagation_axis]
77 | if not cross_axes: raise ValueError("Need at least one transverse axis for mode computation")
78 |
79 | permittivities = 1 / inv_permittivities
80 | coords = [np.arange(permittivities.shape[dim] + 1) * resolution / 1e-6 for dim in cross_axes]
81 | permittivity_squeezed = np.take(permittivities, indices=0, axis=propagation_axis)
82 | if permittivity_squeezed.ndim == 1: permittivity_squeezed = permittivity_squeezed[:, np.newaxis]
83 |
84 | if inv_permeabilities.ndim == inv_permittivities.ndim:
85 | permeability = 1 / inv_permeabilities
86 | permeability_squeezed = np.take(permeability, indices=0, axis=propagation_axis)
87 | if permeability_squeezed.ndim == 1: permeability_squeezed = permeability_squeezed[:, np.newaxis]
88 | else:
89 | permeability_squeezed = 1 / inv_permeabilities.item()
90 |
91 | tangential_axes_map = {0: (1, 2), 1: (0, 2), 2: (0, 1)}
92 |
93 | modes = tidy3d_mode_computation_wrapper(
94 | frequency=frequency,
95 | permittivity_cross_section=permittivity_squeezed,
96 | permeability_cross_section=permeability_squeezed,
97 | coords=coords,
98 | direction=direction,
99 | num_modes=2 * (mode_index + 1) + 5,
100 | target_neff=target_neff,
101 | )
102 | tangential_axes = tangential_axes_map.get(propagation_axis, (0, 1))
103 | modes = sort_modes(modes, filter_pol, tangential_axes)
104 | if mode_index >= len(modes): raise ValueError(f"Requested mode index {mode_index}, but only {len(modes)} modes available")
105 |
106 | mode = modes[mode_index]
107 |
108 | if propagation_axis == 0:
109 | E = np.stack([mode.Ez, mode.Ex, mode.Ey], axis=0).astype(np.complex128)
110 | H = np.stack([mode.Hz, mode.Hx, mode.Hy], axis=0).astype(np.complex128)
111 | elif propagation_axis == 1:
112 | E = np.stack([mode.Ex, mode.Ez, mode.Ey], axis=0).astype(np.complex128)
113 | H = -np.stack([mode.Hx, mode.Hz, mode.Hy], axis=0).astype(np.complex128)
114 | else:
115 | E = np.stack([mode.Ex, mode.Ey, mode.Ez], axis=0).astype(np.complex128)
116 | H = np.stack([mode.Hx, mode.Hy, mode.Hz], axis=0).astype(np.complex128)
117 |
118 | H *= tidy3d.constants.ETA_0
119 |
120 | E_norm, H_norm = _normalize_by_poynting_flux(E, H, axis=propagation_axis)
121 | return E_norm, H_norm, np.asarray(mode.neff, dtype=np.complex128), propagation_axis
122 |
123 |
124 | def solve_modes(
125 | eps: np.ndarray,
126 | omega: float,
127 | dL: float,
128 | npml: int = 0,
129 | m: int = 2,
130 | direction: Literal["+x", "-x", "+y", "-y", "+z", "-z"] = "+x",
131 | filter_pol: Union[Literal["te", "tm"], None] = None,
132 | return_fields: bool = False,
133 | propagation_axis: Union[Literal["+x", "-x", "+y", "-y", "+z", "-z"], None] = None,
134 | target_neff: Union[float, None] = None,
135 | ) -> Union[Tuple[np.ndarray, np.ndarray], Tuple[np.ndarray, np.ndarray, np.ndarray, int]]:
136 | if eps.ndim != 1: raise ValueError("solve_modes expects a 1D permittivity array")
137 |
138 | freq = omega / (2 * np.pi)
139 | inv_eps = (1.0 / np.asarray(eps, dtype=np.complex128)).reshape(1, eps.size, 1)
140 | direction_flag = "+" if direction.startswith("+") else "-"
141 | axis_hint = propagation_axis if propagation_axis is not None else direction
142 |
143 | neffs: list[complex] = []
144 | e_fields: list[np.ndarray] = []
145 | h_fields: list[np.ndarray] = []
146 | mode_vectors: list[np.ndarray] = []
147 |
148 | for mode_index in range(m):
149 | E_full, H_full, neff, prop_axis = compute_mode(
150 | frequency=freq,
151 | inv_permittivities=inv_eps,
152 | inv_permeabilities=1.0,
153 | resolution=dL,
154 | direction=direction_flag,
155 | mode_index=mode_index,
156 | filter_pol=filter_pol,
157 | target_neff=target_neff,
158 | )
159 |
160 | neffs.append(neff)
161 | if return_fields:
162 | e_fields.append(E_full)
163 | h_fields.append(H_full)
164 | else:
165 | component_norms = [np.linalg.norm(np.squeeze(E_full[i])) for i in range(3)]
166 | component_idx = int(np.argmax(component_norms))
167 | field_line = np.squeeze(E_full[component_idx])
168 | if field_line.ndim > 1: field_line = field_line[:, 0]
169 | max_amp = np.max(np.abs(field_line)) or 1.0
170 | mode_vectors.append(field_line / max_amp)
171 |
172 | neff_array = np.asarray(neffs, dtype=np.complex128)
173 |
174 | if return_fields:
175 | return (
176 | neff_array,
177 | np.stack(e_fields) if e_fields else np.empty((0, 3, 0, 0)),
178 | np.stack(h_fields) if h_fields else np.empty((0, 3, 0, 0)),
179 | prop_axis,
180 | )
181 |
182 | if not mode_vectors: return neff_array, np.zeros((eps.size, 0), dtype=np.complex128)
183 |
184 | return neff_array, np.column_stack(mode_vectors)
185 |
186 | def tidy3d_mode_computation_wrapper(
187 | frequency: float,
188 | permittivity_cross_section: np.ndarray,
189 | coords: List[np.ndarray],
190 | direction: Literal["+", "-"],
191 | permeability_cross_section: Union[np.ndarray, None] = None,
192 | target_neff: Union[float, None] = None,
193 | angle_theta: float = 0.0,
194 | angle_phi: float = 0.0,
195 | num_modes: int = 10,
196 | precision: Literal["single", "double"] = "double",
197 | ) -> List[ModeTupleType]:
198 | mode_spec = SimpleNamespace(
199 | num_modes=num_modes,
200 | target_neff=target_neff,
201 | num_pml=(0, 0),
202 | angle_theta=angle_theta,
203 | angle_phi=angle_phi,
204 | bend_radius=None,
205 | bend_axis=None,
206 | precision=precision,
207 | track_freq="central",
208 | group_index_step=False,
209 | )
210 | od = np.zeros_like(permittivity_cross_section)
211 | eps_cross = [permittivity_cross_section if i in {0, 4, 8} else od for i in range(9)]
212 | mu_cross = None
213 | if permeability_cross_section is not None:
214 | mu_cross = [permeability_cross_section if i in {0, 4, 8} else od for i in range(9)]
215 |
216 | EH, neffs, _ = _compute_modes(
217 | eps_cross=eps_cross,
218 | coords=coords,
219 | freq=frequency,
220 | precision=precision,
221 | mode_spec=mode_spec,
222 | direction=direction,
223 | mu_cross=mu_cross,
224 | )
225 | (Ex, Ey, Ez), (Hx, Hy, Hz) = EH.squeeze()
226 |
227 | if num_modes == 1: return [ModeTupleType(Ex=Ex, Ey=Ey, Ez=Ez, Hx=Hx, Hy=Hy, Hz=Hz, neff=complex(neffs))]
228 |
229 | return [
230 | ModeTupleType(
231 | Ex=Ex[..., i],
232 | Ey=Ey[..., i],
233 | Ez=Ez[..., i],
234 | Hx=Hx[..., i],
235 | Hy=Hy[..., i],
236 | Hz=Hz[..., i],
237 | neff=neffs[i],
238 | )
239 | for i in range(min(num_modes, Ex.shape[-1]))
240 | ]
241 |
242 | def _normalize_by_poynting_flux(E: np.ndarray, H: np.ndarray, axis: int) -> tuple[np.ndarray, np.ndarray]:
243 | S = np.cross(E, np.conjugate(H), axis=0)
244 | power = float(np.real(np.sum(S[axis])))
245 |
246 | # Debug: check which normalization path is taken
247 | # print(f"[DEBUG normalize] axis={axis}, power={power:.3e}, E_norm={np.linalg.norm(E):.3e}, H_norm={np.linalg.norm(H):.3e}")
248 |
249 | # Guard against tiny/negative/NaN power from numerical noise
250 | if not np.isfinite(power) or abs(power) < 1e-18:
251 | # Fallback: normalize by field amplitude
252 | e_norm = float(np.linalg.norm(E))
253 | if e_norm > 1e-18 and np.isfinite(e_norm):
254 | # print(f"[DEBUG normalize] Using E-norm fallback: {e_norm:.3e}")
255 | return E / e_norm, H / e_norm
256 | return E, H
257 | # Normalize by magnitude of power to avoid sqrt of negative
258 | scale = np.sqrt(abs(power))
259 | if scale == 0.0 or not np.isfinite(scale):
260 | # Fallback: normalize by field amplitude
261 | e_norm = float(np.linalg.norm(E))
262 | if e_norm > 1e-18 and np.isfinite(e_norm):
263 | return E / e_norm, H / e_norm
264 | return E, H
265 | E_norm = E / scale
266 | H_norm = H / scale
267 | # Final NaN check
268 | if not np.all(np.isfinite(E_norm)) or not np.all(np.isfinite(H_norm)):
269 | return E, H
270 | return E_norm, H_norm
271 |
--------------------------------------------------------------------------------
/docs/adjoint_optimization_guide.md:
--------------------------------------------------------------------------------
1 | # Adjoint Optimization Guide
2 |
3 | ## Overview
4 |
5 | BEAMZ now includes adjoint optimization capabilities for inverse photonic design. The `examples/bend_opt.py` script demonstrates how to optimize a photonic bend structure using the adjoint method with gradient-based optimization.
6 |
7 | ## Key Features
8 |
9 | ### ✅ Complete Adjoint Implementation
10 | - **Forward simulation**: Computes electromagnetic fields propagating from source to target
11 | - **Backward (adjoint) simulation**: Computes adjoint fields propagating backwards from target
12 | - **Gradient computation**: Uses field overlap between forward and adjoint fields
13 | - **Material optimization**: Optimizes permittivity distribution in design region
14 |
15 | ### ✅ Robust Optimization Framework
16 | - **External optimizer**: Uses `scipy.optimize.minimize` with L-BFGS-B algorithm
17 | - **Bounded optimization**: Constrains permittivity values between cladding and core
18 | - **Error handling**: Comprehensive exception handling for robust operation
19 | - **Component testing**: Built-in testing to validate individual components
20 |
21 | ### ✅ Practical Photonic Design
22 | - **Realistic parameters**: Silicon nitride (Si3N4) photonic structures
23 | - **Design regions**: Pixelated optimization regions for flexible design
24 | - **Performance monitoring**: Real-time objective and gradient tracking
25 | - **Visualization**: Automatic plotting of optimization results
26 |
27 | ## Usage
28 |
29 | ### Basic Usage
30 |
31 | ```python
32 | from examples.bend_opt import run_optimization
33 |
34 | # Run the complete optimization
35 | result, optimized_design = run_optimization()
36 | ```
37 |
38 | ### Custom Parameters
39 |
40 | ```python
41 | # Modify parameters at the top of bend_opt.py:
42 | WL = 1.55*µm # Wavelength
43 | N_CORE = 2.04 # Core index (Si3N4)
44 | N_CLAD = 1.444 # Cladding index (SiO2)
45 | initial_permittivity = np.ones((20, 20)) * N_CLAD**2 # Grid size
46 | ```
47 |
48 | ### Advanced Customization
49 |
50 | ```python
51 | # Custom objective function
52 | def compute_objective(permittivity_dist):
53 | design = create_design_with_permittivity(permittivity_dist)
54 | forward_data = forward_sim(design)
55 |
56 | # Custom power calculation
57 | ez_final = forward_data['Ez'][-1]
58 | power = np.sum(np.abs(ez_final)**2)
59 |
60 | return -power # Minimize negative = maximize positive
61 |
62 | # Custom optimizer settings
63 | result = minimize(
64 | fun=objective_function,
65 | x0=x0,
66 | method='L-BFGS-B',
67 | jac=gradient_function,
68 | bounds=bounds,
69 | options={
70 | 'maxiter': 100, # More iterations
71 | 'ftol': 1e-8, # Tighter tolerance
72 | 'gtol': 1e-8,
73 | 'disp': True
74 | }
75 | )
76 | ```
77 |
78 | ## How It Works
79 |
80 | ### 1. Forward Simulation
81 | ```python
82 | def forward_sim(design):
83 | # Add source at input
84 | source = ModeSource(design=design, start=..., end=..., wavelength=WL, signal=signal)
85 |
86 | # Add monitor at output
87 | monitor = Monitor(design=design, start=..., end=..., record_fields=True)
88 |
89 | # Run FDTD simulation
90 | sim = FDTD(design=design, time=time_steps, resolution=DX)
91 | field_history = sim.run()
92 |
93 | return monitor.fields
94 | ```
95 |
96 | ### 2. Backward (Adjoint) Simulation
97 | ```python
98 | def backward_sim(design, target_fields):
99 | # Add adjoint source at output (backwards direction)
100 | adjoint_source = ModeSource(design=design, start=..., end=...,
101 | wavelength=WL, signal=signal, direction="-y")
102 |
103 | # Add monitor at input for field overlap
104 | monitor = Monitor(design=design, start=..., end=..., record_fields=True)
105 |
106 | # Run adjoint FDTD simulation
107 | sim = FDTD(design=design, time=time_steps, resolution=DX)
108 | adjoint_fields = sim.run()
109 |
110 | return monitor.fields
111 | ```
112 |
113 | ### 3. Gradient Computation
114 | ```python
115 | def compute_gradient(permittivity_dist):
116 | # Run forward and backward simulations
117 | forward_data = forward_sim(design)
118 | backward_data = backward_sim(design, forward_data)
119 |
120 | # Compute field overlap for gradient
121 | for i, j in design_region:
122 | gradient[i, j] = -Re(forward_field * conj(backward_field))
123 |
124 | return gradient
125 | ```
126 |
127 | ### 4. Optimization Loop
128 | ```python
129 | # Scipy optimization with gradient
130 | result = minimize(
131 | fun=objective_function, # Minimize -power
132 | x0=initial_permittivity, # Starting design
133 | method='L-BFGS-B', # Gradient-based optimizer
134 | jac=gradient_function, # Provide gradients
135 | bounds=material_bounds # Physical constraints
136 | )
137 | ```
138 |
139 | ## Understanding the Results
140 |
141 | ### Optimization Output
142 | ```
143 | Starting adjoint optimization for photonic bend...
144 | Testing individual components...
145 | 1. Testing design creation...
146 | Created design with 6 structures
147 | 2. Testing forward simulation...
148 | Forward sim: recorded 267 time steps
149 | Forward sim keys: ['Ez', 'Hx', 'Hy', 't']
150 | 3. Testing backward simulation...
151 | Backward sim: recorded 267 time steps
152 | Backward sim keys: ['Ez', 'Hx', 'Hy', 't']
153 | 4. Testing objective computation...
154 | Computed power: 2.345e-12
155 | Objective value: -2.345e-12
156 | 5. Testing gradient computation...
157 | Gradient computed, norm: 1.234e-15
158 | Gradient shape: (400,), norm: 1.234e-15
159 |
160 | Objective value: -2.345e-12
161 | Gradient norm: 1.234e-15
162 | ... optimization iterations ...
163 |
164 | Optimization complete!
165 | Success: True
166 | Final objective: -5.678e-12
167 | Iterations: 15
168 | Function evaluations: 18
169 | ```
170 |
171 | ### Visualization
172 | The script automatically generates a 4-panel plot:
173 | 1. **Initial Permittivity**: Starting material distribution
174 | 2. **Optimized Permittivity**: Final optimized distribution
175 | 3. **Initial Design**: Geometric layout before optimization
176 | 4. **Optimized Design**: Final optimized device structure
177 |
178 | ## Performance Tips
179 |
180 | ### 1. Grid Resolution
181 | ```python
182 | # Coarse grid for testing (faster)
183 | initial_permittivity = np.ones((10, 10)) * N_CLAD**2
184 |
185 | # Fine grid for production (slower but more accurate)
186 | initial_permittivity = np.ones((50, 50)) * N_CLAD**2
187 | ```
188 |
189 | ### 2. Simulation Parameters
190 | ```python
191 | # Faster simulation (less accurate)
192 | DX, DT = calc_optimal_fdtd_params(WL, max(N_CORE, N_CLAD),
193 | dims=2, safety_factor=0.8,
194 | points_per_wavelength=10)
195 |
196 | # Slower simulation (more accurate)
197 | DX, DT = calc_optimal_fdtd_params(WL, max(N_CORE, N_CLAD),
198 | dims=2, safety_factor=0.4,
199 | points_per_wavelength=20)
200 | ```
201 |
202 | ### 3. Optimization Settings
203 | ```python
204 | # Quick convergence (fewer iterations)
205 | options = {'maxiter': 10, 'ftol': 1e-4, 'gtol': 1e-4}
206 |
207 | # Thorough optimization (more iterations)
208 | options = {'maxiter': 100, 'ftol': 1e-8, 'gtol': 1e-8}
209 | ```
210 |
211 | ## Common Issues & Solutions
212 |
213 | ### 1. Low Gradient Magnitude
214 | **Problem**: Gradient norm is very small (~1e-15)
215 | **Solution**:
216 | - Check field overlap calculation
217 | - Increase simulation time
218 | - Adjust monitor positions
219 | - Use different objective function
220 |
221 | ### 2. Optimization Not Converging
222 | **Problem**: L-BFGS-B terminates without improvement
223 | **Solution**:
224 | - Increase `maxiter`
225 | - Relax tolerances (`ftol`, `gtol`)
226 | - Try different optimizer (`SLSQP`, `trust-constr`)
227 | - Check gradient accuracy
228 |
229 | ### 3. Memory Issues
230 | **Problem**: Out of memory during simulation
231 | **Solution**:
232 | - Reduce grid size (`initial_permittivity` shape)
233 | - Reduce simulation time
234 | - Use `save_memory_mode=True`
235 | - Enable field decimation
236 |
237 | ### 4. Slow Performance
238 | **Problem**: Each iteration takes too long
239 | **Solution**:
240 | - Use JAX backend for GPU acceleration
241 | - Reduce grid resolution
242 | - Parallelize if multiple devices available
243 | - Use coarser time stepping
244 |
245 | ## Extensions
246 |
247 | ### Multi-Objective Optimization
248 | ```python
249 | def compute_multi_objective(permittivity_dist):
250 | design = create_design_with_permittivity(permittivity_dist)
251 | forward_data = forward_sim(design)
252 |
253 | # Transmission efficiency
254 | power = compute_power(forward_data)
255 |
256 | # Device footprint penalty
257 | footprint = np.sum(permittivity_dist > N_CLAD**2 + 0.1)
258 |
259 | # Combined objective
260 | return -(power - 0.01 * footprint)
261 | ```
262 |
263 | ### Wavelength Sweep Optimization
264 | ```python
265 | wavelengths = [1.50*µm, 1.55*µm, 1.60*µm]
266 |
267 | def broadband_objective(permittivity_dist):
268 | total_obj = 0
269 | for wl in wavelengths:
270 | # Update wavelength and recompute
271 | global WL
272 | WL = wl
273 | obj = compute_objective(permittivity_dist)
274 | total_obj += obj
275 | return total_obj / len(wavelengths)
276 | ```
277 |
278 | ### Fabrication Constraints
279 | ```python
280 | def apply_fabrication_constraints(permittivity_dist):
281 | # Minimum feature size constraint
282 | from scipy import ndimage
283 |
284 | # Erosion/dilation for minimum linewidth
285 | binary_design = permittivity_dist > (N_CLAD**2 + N_CORE**2) / 2
286 | min_feature = ndimage.binary_erosion(binary_design, iterations=2)
287 | min_feature = ndimage.binary_dilation(min_feature, iterations=2)
288 |
289 | # Apply constraint
290 | constrained = np.where(min_feature, N_CORE**2, N_CLAD**2)
291 | return constrained
292 | ```
293 |
294 | ## Future Enhancements
295 |
296 | ### Planned Features
297 | - [ ] **Automatic differentiation**: Native AD support with JAX
298 | - [ ] **Multi-physics optimization**: Thermal and mechanical constraints
299 | - [ ] **Fabrication-aware design**: Built-in manufacturing constraints
300 | - [ ] **Topology optimization**: Level-set and SIMP methods
301 | - [ ] **Multi-port devices**: Complex routing and switching structures
302 |
303 | ### Optimization Algorithms
304 | - [ ] **Population-based**: Genetic algorithms, particle swarm
305 | - [ ] **Bayesian optimization**: For expensive objective functions
306 | - [ ] **Multi-objective**: Pareto-optimal trade-offs
307 | - [ ] **Robust optimization**: Manufacturing tolerance aware
308 |
309 | ## Examples
310 |
311 | ### Waveguide Bend
312 | The default example optimizes a 90° bend for maximum transmission:
313 | - **Input**: Straight waveguide at input
314 | - **Output**: Straight waveguide at 90° angle
315 | - **Objective**: Maximize power transmission
316 | - **Constraints**: Fixed input/output positions
317 |
318 | ### Beam Splitter
319 | Modify for 1×2 beam splitter:
320 | ```python
321 | # Add two output monitors
322 | monitor1 = Monitor(design=design, start=..., end=..., record_fields=True)
323 | monitor2 = Monitor(design=design, start=..., end=..., record_fields=True)
324 |
325 | # Objective: equal power splitting
326 | def splitter_objective(permittivity_dist):
327 | power1 = compute_power_at_monitor(monitor1)
328 | power2 = compute_power_at_monitor(monitor2)
329 |
330 | # Maximize total power and balance
331 | total_power = power1 + power2
332 | balance = 1 - abs(power1 - power2) / (power1 + power2)
333 |
334 | return -(total_power * balance)
335 | ```
336 |
337 | ### Wavelength Demultiplexer
338 | Multi-wavelength optimization:
339 | ```python
340 | def demux_objective(permittivity_dist):
341 | total_obj = 0
342 |
343 | for i, wl in enumerate([1.50*µm, 1.55*µm, 1.60*µm]):
344 | # Run simulation at wavelength wl
345 | power_correct_port = get_power_at_port(i, wl)
346 | power_other_ports = get_power_at_other_ports(i, wl)
347 |
348 | # Reward power in correct port, penalize crosstalk
349 | obj = power_correct_port - 0.1 * power_other_ports
350 | total_obj += obj
351 |
352 | return -total_obj
353 | ```
354 |
355 | This adjoint optimization framework provides a solid foundation for inverse photonic design with BEAMZ! 🎯
--------------------------------------------------------------------------------
/beamz/simulation/fields.py:
--------------------------------------------------------------------------------
1 | """Field storage and update logic for FDTD simulations."""
2 |
3 | from __future__ import annotations
4 | import numpy as np
5 | from beamz.simulation import ops
6 |
7 |
8 | class Fields:
9 | """Container for E/H field arrays on staggered Yee grid with FDTD update logic."""
10 |
11 | def __init__(self, permittivity, conductivity, permeability, resolution, pml_regions=None, plane_2d='xy'):
12 | """Initialize field arrays on a Yee grid for 2D (all 6 components) or 3D (Ex, Ey, Ez, Hx, Hy, Hz) simulations."""
13 | self.resolution = resolution
14 | self.plane_2d = plane_2d
15 | # Store references to material grids owned by Design (no copying)
16 | self.permittivity = permittivity
17 | self.conductivity = conductivity
18 | self.permeability = permeability
19 |
20 | # Initialize PML regions if present
21 | if pml_regions:
22 | # We don't use split-field anymore, but we might store regions for visualization
23 | self.has_pml = True
24 | self.pml_regions = pml_regions
25 | else:
26 | self.has_pml = False
27 |
28 | # Infer dimensionality and shape from material arrays
29 | is_3d = self.permittivity.ndim == 3
30 | grid_shape = self.permittivity.shape
31 |
32 | if is_3d:
33 | nz, ny, nx = grid_shape
34 | self._init_fields_3d(nx, ny, nz)
35 | self.update = self._update_3d
36 | self._curl_e_to_h = ops.curl_e_to_h_3d
37 | self._curl_h_to_e = ops.curl_h_to_e_3d
38 | self._material_slice = ops.material_slice_for_e_3d
39 | self.sigma_m_hx, self.sigma_m_hy, self.sigma_m_hz = ops.magnetic_conductivity_terms_3d(
40 | self.conductivity, self.permeability, self.Hx.shape, self.Hy.shape, self.Hz.shape)
41 | self.eps_x, self.sig_x, self.region_x = self._material_slice(self.permittivity, self.conductivity, orientation="x")
42 | self.eps_y, self.sig_y, self.region_y = self._material_slice(self.permittivity, self.conductivity, orientation="y")
43 | self.eps_z, self.sig_z, self.region_z = self._material_slice(self.permittivity, self.conductivity, orientation="z")
44 | else:
45 | dim1, dim2 = grid_shape
46 | self._init_fields_2d(dim1, dim2)
47 | self.update = self._update_2d
48 | self._curl_e_to_h = ops.curl_e_to_h_2d
49 | self._curl_h_to_e = ops.curl_h_to_e_2d
50 | # Material slicing will be handled dynamically or initialized here based on plane
51 | # Initial dummy initialization, will be properly set in set_pml_conductivity or init
52 | self._init_material_parameters_2d()
53 |
54 | def set_pml_conductivity(self, pml_data):
55 | """Set effective conductivity for PML regions (replaces split-field initialization)."""
56 | self.has_pml = True
57 | self.pml_data = pml_data
58 |
59 | # Recalculate material parameters with PML conductivity added
60 | if self.permittivity.ndim == 3: # 3D
61 | # TODO: Implement 3D PML effective conductivity
62 | pass
63 | else: # 2D
64 | self._init_material_parameters_2d()
65 |
66 | def _init_material_parameters_2d(self):
67 | """Initialize 2D material parameters including PML conductivity if present."""
68 | # Base conductivity from design
69 | base_sigma = self.conductivity
70 |
71 | # Add PML conductivity if present
72 | if self.has_pml and hasattr(self, 'pml_data'):
73 | # Combine profiles based on plane
74 | sigma_pml = np.zeros_like(base_sigma)
75 | if self.plane_2d == 'xy':
76 | if 'sigma_x' in self.pml_data: sigma_pml += self.pml_data['sigma_x']
77 | if 'sigma_y' in self.pml_data: sigma_pml += self.pml_data['sigma_y']
78 | elif self.plane_2d == 'yz':
79 | if 'sigma_y' in self.pml_data: sigma_pml += self.pml_data['sigma_y']
80 | if 'sigma_z' in self.pml_data: sigma_pml += self.pml_data['sigma_z']
81 | elif self.plane_2d == 'xz':
82 | if 'sigma_x' in self.pml_data: sigma_pml += self.pml_data['sigma_x']
83 | if 'sigma_z' in self.pml_data: sigma_pml += self.pml_data['sigma_z']
84 |
85 | total_sigma = base_sigma + sigma_pml
86 | else:
87 | total_sigma = base_sigma
88 |
89 | # Initialize slicing and magnetic terms based on plane
90 | # Note: We use the same total_sigma for all components as a simplification
91 | # Ideally, we should stagger sigma for each component, but grid-colocated sigma is a standard approx
92 |
93 | # Setup slices for all 3 E-components
94 | self.eps_x, self.sig_x, self.region_x = ops.material_slice_for_e_2d_component(
95 | self.permittivity, total_sigma, "x", self.plane_2d)
96 | self.eps_y, self.sig_y, self.region_y = ops.material_slice_for_e_2d_component(
97 | self.permittivity, total_sigma, "y", self.plane_2d)
98 | self.eps_z, self.sig_z, self.region_z = ops.material_slice_for_e_2d_component(
99 | self.permittivity, total_sigma, "z", self.plane_2d)
100 |
101 | # Setup magnetic conductivity for H-field updates
102 | self.sigma_m_hx, self.sigma_m_hy, self.sigma_m_hz = ops.magnetic_conductivity_terms_2d_full(
103 | total_sigma, self.permeability, self.Hx.shape, self.Hy.shape, self.Hz.shape, self.plane_2d)
104 |
105 | def _init_upml_fields(self, pml_regions):
106 | """Initialize auxiliary fields for split-field UPML. (Deprecated/No-op)"""
107 | pass
108 |
109 | def _init_split_fields_2d(self):
110 | """Initialize split-field components for 2D UPML. (Deprecated/No-op)"""
111 | pass
112 |
113 | def _init_split_fields_3d(self):
114 | """Initialize split-field components for 3D UPML."""
115 | if self.has_pml:
116 | # 3D: all components split
117 | self.Ex_y = np.zeros_like(self.Ex)
118 | self.Ex_z = np.zeros_like(self.Ex)
119 | self.Ey_x = np.zeros_like(self.Ey)
120 | self.Ey_z = np.zeros_like(self.Ey)
121 | self.Ez_x = np.zeros_like(self.Ez)
122 | self.Ez_y = np.zeros_like(self.Ez)
123 | # Similar for H fields...
124 |
125 | def _init_fields_3d(self, nx, ny, nz):
126 | """Initialize 3D field arrays (Ex, Ey, Ez, Hx, Hy, Hz) with proper Yee grid staggering."""
127 | self.Ex = np.zeros((nz, ny, nx - 1))
128 | self.Ey = np.zeros((nz, ny - 1, nx))
129 | self.Ez = np.zeros((nz - 1, ny, nx))
130 | self.Hx = np.zeros((nz - 1, ny - 1, nx))
131 | self.Hy = np.zeros((nz - 1, ny, nx - 1))
132 | self.Hz = np.zeros((nz, ny - 1, nx - 1))
133 |
134 | def _init_fields_2d(self, dim1, dim2):
135 | """Initialize 2D field arrays (Ex, Ey, Ez, Hx, Hy, Hz) on staggered Yee grid for the selected plane."""
136 | # dim1, dim2 correspond to the two active dimensions
137 | # xy: (y, x), yz: (z, y), xz: (z, x)
138 |
139 | if self.plane_2d == 'xy':
140 | ny, nx = dim1, dim2
141 | # TM set (Ez, Hx, Hy)
142 | self.Ez = np.zeros((ny, nx))
143 | self.Hx = np.zeros((ny, nx - 1))
144 | self.Hy = np.zeros((ny - 1, nx))
145 | # TE set (Hz, Ex, Ey)
146 | self.Hz = np.zeros((ny - 1, nx - 1))
147 | self.Ex = np.zeros((ny, nx - 1))
148 | self.Ey = np.zeros((ny - 1, nx))
149 |
150 | elif self.plane_2d == 'yz':
151 | nz, ny = dim1, dim2
152 | # Invariant in x. We map standard 3D staggering to 2D slice.
153 | # 3D: Ex(z,y,x-1/2), Ey(z,y-1/2,x), Ez(z-1/2,y,x)
154 | # 2D yz (x-invariant):
155 | # Ex is normal to plane -> (nz, ny) [like Ez in xy]
156 | # Ey is in plane -> (nz, ny-1)
157 | # Ez is in plane -> (nz-1, ny)
158 |
159 | # TE-like set (Ex, Hy, Hz)
160 | self.Ex = np.zeros((nz, ny))
161 | self.Hy = np.zeros((nz, ny - 1))
162 | self.Hz = np.zeros((nz - 1, ny))
163 |
164 | # TM-like set (Hx, Ey, Ez)
165 | self.Hx = np.zeros((nz - 1, ny - 1))
166 | self.Ey = np.zeros((nz, ny - 1))
167 | self.Ez = np.zeros((nz - 1, ny))
168 |
169 | elif self.plane_2d == 'xz':
170 | nz, nx = dim1, dim2
171 | # Invariant in y.
172 | # Ey is normal to plane -> (nz, nx)
173 | # Ex is in plane -> (nz, nx-1)
174 | # Ez is in plane -> (nz-1, nx)
175 |
176 | # TE-like set (Ey, Hx, Hz)
177 | self.Ey = np.zeros((nz, nx))
178 | self.Hx = np.zeros((nz, nx - 1))
179 | self.Hz = np.zeros((nz - 1, nx))
180 |
181 | # TM-like set (Hy, Ex, Ez)
182 | self.Hy = np.zeros((nz - 1, nx - 1))
183 | self.Ex = np.zeros((nz, nx - 1))
184 | self.Ez = np.zeros((nz - 1, nx))
185 |
186 | def available_components(self):
187 | """Return list of available field components."""
188 | if self.permittivity.ndim == 3:
189 | return ['Ex', 'Ey', 'Ez', 'Hx', 'Hy', 'Hz']
190 | else:
191 | # All 6 are available in 2D full mode
192 | return ['Ex', 'Ey', 'Ez', 'Hx', 'Hy', 'Hz']
193 |
194 | def _update_2d(self, dt, source_j=None, source_m=None):
195 | """Execute one 2D FDTD time step (all 6 components) for the selected plane."""
196 | # 1. Update H fields from E fields
197 | curlE_x, curlE_y, curlE_z = self._curl_e_to_h((self.Ex, self.Ey, self.Ez), self.resolution, plane=self.plane_2d)
198 |
199 | # Inject magnetic sources (if any)
200 | if source_m:
201 | if 'Hx' in source_m:
202 | val, indices = source_m['Hx']
203 | curlE_x[indices] += val # Note: Physical sign of source injection depends on equation form, using += for standard source current
204 | if 'Hy' in source_m:
205 | val, indices = source_m['Hy']
206 | curlE_y[indices] += val
207 | if 'Hz' in source_m:
208 | val, indices = source_m['Hz']
209 | curlE_z[indices] += val
210 |
211 | self.Hx = ops.advance_h_field(self.Hx, curlE_x, self.sigma_m_hx, dt)
212 | self.Hy = ops.advance_h_field(self.Hy, curlE_y, self.sigma_m_hy, dt)
213 | self.Hz = ops.advance_h_field(self.Hz, curlE_z, self.sigma_m_hz, dt)
214 |
215 | # 2. Update E fields from H fields
216 | curlH_x, curlH_y, curlH_z = self._curl_h_to_e((self.Hx, self.Hy, self.Hz), self.resolution,
217 | (self.Ex.shape, self.Ey.shape, self.Ez.shape), plane=self.plane_2d)
218 |
219 | # Inject electric sources (if any)
220 | if source_j:
221 | if 'Ex' in source_j:
222 | val, indices = source_j['Ex']
223 | curlH_x[indices] += val # Standard convention: dE/dt = ... - J/eps -> usually subtracted in update or added to curl H?
224 | # Beamz convention: dE/dt = ... + source * curlH_with_source.
225 | # Usually Ampere's law: curl H = J + dD/dt -> dE/dt = (curl H - J)/eps.
226 | # ops.advance_e_field adds source*curl_region.
227 | # If J is added to curlH, it means dE/dt ~ (curlH + J).
228 | # This corresponds to a source J that OPPOSES conduction current?
229 | # Legacy code: curlH_z_with_source += j_z. So we add J.
230 |
231 | if 'Ey' in source_j:
232 | val, indices = source_j['Ey']
233 | curlH_y[indices] += val
234 | if 'Ez' in source_j:
235 | val, indices = source_j['Ez']
236 | curlH_z[indices] += val
237 |
238 | self.Ex = ops.advance_e_field(self.Ex, curlH_x, self.sig_x, self.eps_x, dt, self.region_x)
239 | self.Ey = ops.advance_e_field(self.Ey, curlH_y, self.sig_y, self.eps_y, dt, self.region_y)
240 | self.Ez = ops.advance_e_field(self.Ez, curlH_z, self.sig_z, self.eps_z, dt, self.region_z)
241 |
242 | def _update_3d(self, dt):
243 | """Execute one 3D FDTD time step: H from curl(E) via Faraday's law, then E from curl(H) via Ampere's law."""
244 | curlE_x, curlE_y, curlE_z = self._curl_e_to_h(self.Ex, self.Ey, self.Ez, self.resolution)
245 | self.Hx = ops.advance_h_field(self.Hx, curlE_x, self.sigma_m_hx, dt)
246 | self.Hy = ops.advance_h_field(self.Hy, curlE_y, self.sigma_m_hy, dt)
247 | self.Hz = ops.advance_h_field(self.Hz, curlE_z, self.sigma_m_hz, dt)
248 | curlH_x, curlH_y, curlH_z = self._curl_h_to_e(self.Hx, self.Hy, self.Hz, self.resolution)
249 | self.Ex = ops.advance_e_field(self.Ex, curlH_x, self.sig_x, self.eps_x, dt, self.region_x)
250 | self.Ey = ops.advance_e_field(self.Ey, curlH_y, self.sig_y, self.eps_y, dt, self.region_y)
251 | self.Ez = ops.advance_e_field(self.Ez, curlH_z, self.sig_z, self.eps_z, dt, self.region_z)
252 |
253 | def update(self, dt, source_j=None, source_m=None):
254 | """Execute one FDTD time step with optional source injection."""
255 | if not self.permittivity.ndim == 3: # 2D
256 | self._update_2d(dt, source_j=source_j, source_m=source_m)
257 | else: # 3D
258 | self._update_3d(dt)
--------------------------------------------------------------------------------