Autograd0Overview.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "4995430a-4e1a-4d9f-abc7-9eacc36d6eda",
   "metadata": {},
   "source": [
    "# Inverse design overview\n",
    "\n",
    "*Effortlessly optimize complex photonic devices. Tidy3D integrates automatic differentiation with the efficient adjoint method, allowing you to compute gradients for thousands of parameters with just one extra simulation.*\n",
    "\n",
    "## **How Automatic Differentiation Works in Tidy3D**\n",
    "\n",
    "Tidy3D empowers users to perform gradient-based optimization and sensitivity analysis of photonic devices directly within their simulation workflow. This is achieved by making Tidy3D simulations **differentiable** – meaning we can efficiently compute the derivative (gradient) of a figure of merit (like transmission efficiency) with respect to any number of design parameters (like geometry dimensions or material properties).\n",
    "\n",
    "This capability relies on two core technologies working together:\n",
    "\n",
    "1. **Automatic Differentiation (AD) Framework (`autograd`):** Handles the differentiation of the overall Python code defining your objective function.\n",
    "2. **The Adjoint Method:** Provides an efficient way to calculate the specific derivatives related to the FDTD simulation step itself.\n",
    "\n",
    "Let's break down how these pieces fit together.\n",
    "\n",
    "### The Challenge: Differentiating Complex Simulations\n",
    "\n",
    "Optimizing photonic devices often involves finding the best set of design parameters (e.g., taper shape) that maximize or minimize a specific outcome (e.g., coupling efficiency). Gradient-based optimization methods are highly effective for this, but they require knowing how sensitive the outcome is to small changes in each parameter.\n",
    "\n",
    "Calculating this gradient for a complex FDTD simulation is challenging. A naive approach (like finite differences) would require running many simulations – one for each parameter – which quickly becomes computationally infeasible for designs with many variables.\n",
    "\n",
    "### Our Solution: Combining `autograd` and the Adjoint Method\n",
    "\n",
    "Tidy3D leverages the `autograd` library for automatic differentiation. When you write a Python function that defines your simulation setup, runs `tidy3d.web.run`, and calculates a final figure of merit, `autograd` can automatically track all the mathematical operations *outside* the simulation.\n",
    "\n",
    "<img src=\"img/autograd_overview_1.png\" width=\"700\" alt=\"Schematic of the inverse design\">\n",
    "\n",
    "However, `autograd` doesn't inherently know how to differentiate *through* the FDTD simulation node. This is where Tidy3D's custom integration comes in, using the **adjoint method**.\n",
    "\n",
    "We've essentially \"taught\" `autograd` the derivative rule for the `tidy3d.web.run` operation. This rule is implemented using the adjoint method, a powerful technique derived from electromagnetic theory.\n",
    "\n",
    "**Key Benefit of the Adjoint Method:** It allows us to compute the gradient of the figure of merit with respect to *all* design parameters using just **one additional simulation** (the adjoint simulation), regardless of how many parameters there are (hundreds, thousands, or even millions).\n",
    "\n",
    "### The Differentiation Pipeline: Forward and Backward Passes\n",
    "\n",
    "When you ask `autograd` to compute the gradient (e.g., using `autograd.grad(objective_function)`), here's a simplified view of what happens under the hood:\n",
    "\n",
    "1. **Forward Pass:**\n",
    "    - Your Python function executes normally.\n",
    "    - Tidy3D components (`td.Structure`, `td.Box`, `td.Medium`, etc.) are created, potentially using `autograd` tracked numbers (from `autograd.numpy`).\n",
    "    - `tidy3d.web.run(simulation)` is called. The standard **forward FDTD simulation** runs.\n",
    "    - Crucially, during this forward run, Tidy3D automatically stores necessary field information in the regions relevant to the differentiable parameters (e.g., near the boundaries of a shape-optimized geometry).\n",
    "    - The simulation results (`SimulationData`) are returned.\n",
    "    - Your function calculates the final scalar objective value using `autograd.numpy` operations. `autograd` builds the computational graph along the way.\n",
    "2. **Backward Pass (Vector-Jacobian Product - VJP):**\n",
    "    - `autograd` starts propagating gradient information backward through the computational graph using the chain rule.\n",
    "    - When it reaches the `td.web.run` node, Tidy3D's custom VJP function takes over.\n",
    "    - This function uses the gradient information flowing *into* it (representing the sensitivity of the final objective to the simulation's outputs, like monitor data) to set up **adjoint sources**.\n",
    "    - An **adjoint FDTD simulation** is automatically configured and run.\n",
    "    - Tidy3D combines the stored fields from the **forward simulation** with the results of the **adjoint simulation**.\n",
    "    - Using custom gradient rules, it efficiently calculates the partial derivative of the objective function with respect to *every single* tracked design parameter.\n",
    "    - These parameter gradients are packaged and returned to `autograd`.\n",
    "    - `autograd` continues the backward pass until it reaches the original input parameters, yielding the final overall gradient.\n",
    "\n",
    "<img src=\"img/autograd_overview_2.png\" width=\"700\" alt=\"Schematic of the adjoint method\">\n",
    "\n",
    "### How Tidy3D Components Compute Gradients\n",
    "\n",
    "Different types of parameters require different specific calculations within the adjoint method framework. Tidy3D handles this internally:\n",
    "\n",
    "- **Shape Optimization:** For parameters defining geometry (e.g., `Box.center`, `Box.size`, `PolySlab.vertices`), gradients are typically computed using *surface integrals* involving the forward and adjoint fields on the boundaries of the shape. Moving a boundary slightly changes the permittivity locally, and the adjoint method quantifies the impact of this change on the objective.\n",
    "- **Material Optimization (Topology Optimization):** For parameters defining material properties (e.g., the permittivity in each voxel of a `CustomMedium`), gradients are computed using the forward and adjoint fields within that material region.\n",
    "\n",
    "You don't need to worry about these formulas – Tidy3D components like `td.Box`, `td.Cylinder`, `td.Medium`, `td.CustomMedium`, etc., have their differentiation logic built-in.\n",
    "\n",
    "### The User Experience: Seamless Integration\n",
    "\n",
    "The beauty of this approach is its simplicity from the user's perspective:\n",
    "\n",
    "1. Define your objective function in Python using standard Tidy3D components and `autograd.numpy` for numerical operations.\n",
    "2. Call `td.web.run()` within your function as usual.\n",
    "3. Use `autograd.grad()` (or related functions like `value_and_grad`) to get the gradient.\n",
    "\n",
    "Tidy3D and `autograd` handle the complex forward simulation, field storage, adjoint simulation setup, adjoint run, and final gradient calculation automatically behind the scenes. You get the efficiency of the adjoint method without needing to implement it yourself.\n",
    "\n",
    "### Conclusion\n",
    "\n",
    "Tidy3D's integration with `autograd` via the adjoint method provides a powerful, flexible, and efficient platform for inverse design and sensitivity analysis. By defining a custom derivative rule for FDTD simulations, we unlock the ability to use gradient-based optimization on complex photonic design problems, requiring only one forward and one adjoint simulation per optimization step, regardless of the number of design parameters. This dramatically accelerates the design cycle for cutting-edge photonic devices.\n",
    "\n",
    "**Next Steps / Further Reading:**\n",
    "\n",
    "1. [Topology optimization of a mode converter](https://www.flexcompute.com/tidy3d/examples/notebooks/Autograd3InverseDesign/)\n",
    "2. [Shape optimization of a waveguide bend](https://www.flexcompute.com/tidy3d/examples/notebooks/Autograd8WaveguideBend/)\n",
    "3. [Level set optimization of a y-branch](https://www.flexcompute.com/tidy3d/examples/notebooks/Autograd10YBranchLevelSet/)\n",
    "4. [Fabrication-aware inverse design](https://www.flexcompute.com/tidy3d/examples/notebooks/Autograd23FabricationAwareInvdes/)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}