Architecture Extensions for Adding (Topology) Optimization

ascii.txt

@@ -0,0 +1,27 @@
+
+        @@@@@@@: @@@@@@@   =@@@    @@@#   @@@@        
+        @@@ :@@% @@@       @@@@@   @@@@. %@@@@        
+        @@@@@@@- @@@@@@*  @@@-@@.  @@=@@+@@%@@        
+        @@@  @@@ @@@     @@@@@@@@  @@ %@@@:%@@        
+        @@@@@@@: @@@@@@@+@@    @@@.@@  @@@ #@@        
+
+        *@@@@@@@@@%*=::::---:::---::---+%@*=*+        
+        -%#-=@@@@@@#+=------==-=%%+::::..-@@#         
+        *@+-*@+=#++*#***+===++%@@@=    .*@=-%-        
+        #@@@+.=--#@:         -+-+.   .=%:  -@@        
+        #@@---:=@-          #=#=  ..-%-   @@@         
+        #@@*=+@=          #+**:..-=*-   *%*           
+        *@@%@+          +*##-::-+#+   +#+             
+         @@+          *%#%*+==+**   =+=               
+                   -%@%@@%#***#.  :+=.                
+                 .@@@@@@%%%#%=   #=-                  
+                 %@@@@@@@@@#  .*==                    
+               +%@@@@@@@@@- .**+.          @*-        
+             =*%@@@@@@@@*. +%#-          *@@@#        
+           =%%@#%@@@@@#..+#%=          =**@@--        
+         *@@@%+%@@@@@:.:*%*          -#=@+ =::        
+        @@@@:-*@@@@=.:*#*          -%=%%   +::        
+        =#@--+%@@=..@%%%**+=--==*##%@@.   .**-        
+        -#@#%@@@#*+=#@*=++==-==+**#%=:.:=+%@@=        
+        @@*@@%+*#%@@@%+-:.......::-+=:=#@#@@@@        
+

beamz/__init__.py

@@ -15,16 +15,19 @@
     Rectangle, Circle, Ring, 
     CircularBend, Polygon, Taper
 )
-from beamz.design.sources import ModeSource, GaussianSource
-from beamz.design.monitors import Monitor
-from beamz.design.signals import ramped_cosine, plot_signal
-from beamz.design.mode import solve_modes, slab_mode_source
+from beamz.devices.sources import ModeSource, GaussianSource
+from beamz.devices.monitors import Monitor
+from beamz.devices.signals import ramped_cosine, plot_signal
+from beamz.devices.mode import solve_modes, slab_mode_source
 
 # Import simulation-related classes and functions
 from beamz.simulation.meshing import RegularGrid
 from beamz.simulation.fdtd import FDTD
 from beamz.simulation.backends import get_backend
 
+from beamz.optimization.optimizers import Optimizer
+from beamz.optimization.topology import compute_overlap_gradient
+
 # Import optimization-related classes
 # (Currently empty, to be filled as the module grows)
 
@@ -81,6 +84,10 @@
     'RegularGrid': RegularGrid,
     'FDTD': FDTD,
     'get_backend': get_backend,
+
+    # Optimization
+    'Optimizer': Optimizer,
+    'compute_overlap_gradient': compute_overlap_gradient,
 
     # UI helpers
     'display_header': display_header,

beamz/design/__init__.py

@@ -8,10 +8,10 @@
     Rectangle, Circle, Ring, CircularBend, Polygon, Taper
 )
 from beamz.design.pml import PML
-from beamz.design.sources import ModeSource, GaussianSource
-from beamz.design.monitors import Monitor
-from beamz.design.signals import ramped_cosine, plot_signal
-from beamz.design.mode import solve_modes, slab_mode_source
+from beamz.devices.sources import ModeSource, GaussianSource
+from beamz.devices.monitors import Monitor
+from beamz.devices.signals import ramped_cosine, plot_signal
+from beamz.devices.mode import solve_modes, slab_mode_source
 
 __all__ = [
     'Material', 'CustomMaterial',

beamz/design/core.py

@@ -7,8 +7,8 @@
 from beamz.helpers import display_status
 
 from beamz.design.materials import Material
-from beamz.design.sources import ModeSource, GaussianSource
-from beamz.design.monitors import Monitor
+from beamz.devices.sources import ModeSource, GaussianSource
+from beamz.devices.monitors import Monitor
 from beamz.design.structures import Polygon, Rectangle, Circle, Ring, CircularBend, Taper
 from beamz.design.pml import PML
 
@@ -178,8 +178,8 @@ def unify_polygons(self):
                 non_polygon_structures.append(structure)
                 continue
             # Import at function level to avoid circular imports
-            from beamz.design.sources import ModeSource, GaussianSource
-            from beamz.design.monitors import Monitor
+            from beamz.devices.sources import ModeSource, GaussianSource
+            from beamz.devices.monitors import Monitor
             if isinstance(structure, (ModeSource, GaussianSource, Monitor)):
                 non_polygon_structures.append(structure)
                 continue

beamz/design/structures.py

@@ -23,7 +23,7 @@ def _process_vertices(self, vertices, z=0):
             vertices_3d = [(x, y, vertices_3d[i][2] if len(vertices_3d[i]) > 2 else z) 
                           for i, (x, y) in enumerate(vertices_2d)]
         return vertices_3d
-    
+
     def _process_vertices_preserve_orientation(self, vertices, z=0):
         if not vertices: return []
         vertices_3d = self._ensure_3d_vertices(vertices)
@@ -345,7 +345,7 @@ def add_to_plot(self, ax, facecolor=None, edgecolor="black", alpha=None, linesty
         if alpha is None: alpha = 1
         if linestyle is None: linestyle = '-'
         return super().add_to_plot(ax, facecolor=facecolor, edgecolor=edgecolor, alpha=alpha, linestyle=linestyle)
-    
+
     def copy(self):
         return Ring(position=self.position, 
                     inner_radius=self.inner_radius, 
@@ -412,7 +412,7 @@ def add_to_plot(self, ax, facecolor=None, edgecolor="black", alpha=None, linesty
         if linestyle is None: linestyle = '-'
         # Use parent polygon drawing
         return super().add_to_plot(ax, facecolor=facecolor, edgecolor=edgecolor, alpha=alpha, linestyle=linestyle)
-
+        
     def copy(self):
         return CircularBend(self.position, self.inner_radius, self.outer_radius, 
                             self.angle, self.rotation, self.material, self.color, self.optimize, 

beamz/design/sources.py → beamz/devices/sources.py

@@ -2,7 +2,7 @@
 from beamz.const import LIGHT_SPEED, µm
 import matplotlib.pyplot as plt
 import matplotlib.patches as patches
-from beamz.design.mode import solve_modes
+from beamz.devices.mode import solve_modes
 
 class GaussianSource():
     """A Gaussian current distribution in space.
@@ -126,7 +126,7 @@ def __init__(self, design, position=None, width=None, height=None, wavelength=1.
             # This assumes a simple rectangular waveguide structure
             # Extract core/cladding indices for analytical solver
             try:
-                from beamz.design.mode import slab_mode_source
+                from beamz.devices.mode import slab_mode_source
 
                 # Sample x coordinates along the cross-section
                 num_points = eps_1d.size

beamz/optimization/__init__.py

@@ -1,8 +1,6 @@
-"""
-Optimization module for BEAMZ - Contains adjoint optimization functionality.
-"""
+"""Adjoint-based optimization helpers for BEAMZ."""
 
-# Currently empty as the module is under development
-# Will be populated as optimization features are added
+from . import topology
+from .optimizers import Optimizer
 
-__all__ = [] 
+__all__ = ["topology", "Optimizer"]

beamz/optimization/optimizers.py

@@ -0,0 +1,112 @@
+"""Optimization utilities for topology updates.
+
+This module provides a small wrapper that keeps track of optimizer state
+across iterations.  The intent is to hide method-specific bookkeeping from the
+application layer so that examples can simply call ``step(gradient)`` and feed
+the returned update back into their design variables.
+
+Only the Adam optimizer is implemented at the moment, but the structure allows
+additional optimizers to be registered transparently.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from typing import Callable, Dict, Mapping, MutableMapping, Optional
+
+import numpy as np
+
+
+UpdateFn = Callable[[np.ndarray, np.ndarray, MutableMapping[str, np.ndarray]], np.ndarray]
+
+
+def _adam_update(
+    grad: np.ndarray,
+    learning_rate: float,
+    state: MutableMapping[str, np.ndarray],
+    *,
+    beta1: float = 0.9,
+    beta2: float = 0.999,
+    epsilon: float = 1e-8,
+) -> np.ndarray:
+    """Single Adam update step.
+
+    Args:
+        grad: Gradient array (same shape as parameters).
+        learning_rate: Scalar step size.
+        state: Mutable mapping storing ``m`` (first moment), ``v`` (second
+            moment) and ``t`` (step counter).
+
+    Returns:
+        Update array to be added to the parameters.
+    """
+
+    if "m" not in state:
+        state["m"] = np.zeros_like(grad)
+    if "v" not in state:
+        state["v"] = np.zeros_like(grad)
+    state["t"] = state.get("t", 0) + 1
+
+    m = state["m"]
+    v = state["v"]
+    t = state["t"]
+
+    m *= beta1
+    m += (1.0 - beta1) * grad
+
+    v *= beta2
+    v += (1.0 - beta2) * (grad ** 2)
+
+    m_hat = m / (1.0 - beta1 ** t)
+    v_hat = v / (1.0 - beta2 ** t)
+
+    return -learning_rate * m_hat / (np.sqrt(v_hat) + epsilon)
+
+
+_OPTIMIZER_REGISTRY: Dict[str, UpdateFn] = {
+    "adam": _adam_update,
+}
+
+
+@dataclass
+class Optimizer:
+    """General optimizer wrapper supporting multiple methods.
+
+    Example
+    -------
+    >>> opt = Optimizer(method="adam", learning_rate=1e-2)
+    >>> update = opt.step(gradient)
+    >>> density += update
+    """
+
+    method: str = "adam"
+    learning_rate: float = 1e-2
+    options: Mapping[str, float] = field(default_factory=dict)
+    _state: MutableMapping[str, np.ndarray] = field(default_factory=dict, init=False)
+
+    def __post_init__(self) -> None:
+        method_key = self.method.lower()
+        if method_key not in _OPTIMIZER_REGISTRY:
+            available = ", ".join(sorted(_OPTIMIZER_REGISTRY))
+            raise ValueError(f"Unknown optimizer method '{self.method}'. Available: {available}.")
+        self._update_fn = _OPTIMIZER_REGISTRY[method_key]
+
+    def step(self, gradient: np.ndarray) -> np.ndarray:
+        """Compute the parameter update for a given gradient."""
+
+        if not isinstance(gradient, np.ndarray):
+            raise TypeError("Gradient must be a numpy.ndarray")
+        if gradient.size == 0:
+            return np.zeros_like(gradient)
+
+        return self._update_fn(gradient, self.learning_rate, self._state, **self.options)
+
+    def reset(self) -> None:
+        """Reset internal optimizer state."""
+
+        self._state.clear()
+
+
+__all__ = ["Optimizer"]
+
+