diff --git a/.github/workflows/release.yml b/.github/workflows/release.yml index bb3143169..15ff62060 100644 --- a/.github/workflows/release.yml +++ b/.github/workflows/release.yml @@ -81,7 +81,7 @@ jobs: if: contains(github.ref, 'rc') == false uses: everlytic/branch-merge@1.1.2 with: - github_token: ${{ secrets.GITHUB_TOKEN }} + github_token: ${{ secrets.GH_PAT }} source_ref: ${{ github.ref }} target_branch: "latest" commit_message_template: ':tada: RELEASE: Merged {source_ref} into target {target_branch}' @@ -90,7 +90,7 @@ jobs: with: generate_release_notes: true env: - GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }} + GITHUB_TOKEN: ${{ secrets.GH_PAT }} pypi-release: runs-on: ubuntu-latest steps: @@ -119,7 +119,7 @@ jobs: if: contains(github.ref, 'rc') == false uses: everlytic/branch-merge@1.1.2 with: - github_token: ${{ secrets.GITHUB_TOKEN }} + github_token: ${{ secrets.GH_PAT }} source_ref: "latest" target_branch: "develop" commit_message_template: ':tada: RELEASE: Synced latest into develop' diff --git a/.github/workflows/sync-to-readthedocs-repo.yaml b/.github/workflows/sync-to-readthedocs-repo.yaml index 251f1249e..b0915c80b 100644 --- a/.github/workflows/sync-to-readthedocs-repo.yaml +++ b/.github/workflows/sync-to-readthedocs-repo.yaml @@ -50,7 +50,7 @@ jobs: git remote add mirror https://github.com/flexcompute-readthedocs/tidy3d-docs.git git push mirror ${{ needs.extract_branch_or_tag.outputs.ref_name }} --force # overwrites always env: - GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }} + GITHUB_TOKEN: ${{ secrets.GH_PAT }} # Conditional Checkout for Tag - name: Checkout Tag if tag-triggered-sync @@ -70,4 +70,4 @@ jobs: git remote add mirror https://github.com/flexcompute-readthedocs/tidy3d-docs.git git push mirror ${{ needs.extract_branch_or_tag.outputs.ref_name }} --force # overwrites always env: - GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }} \ No newline at end of file + GITHUB_TOKEN: ${{ secrets.GH_PAT }} diff --git a/CHANGELOG.md b/CHANGELOG.md index d8076c8b0..e6d13b23d 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,3 +1,5 @@ +# Changelog + All notable changes to this project will be documented in this file. The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/), @@ -5,6 +7,94 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ## [2.8.0rc1] +## [2.7.7] - 2024-11-15 + +### Added +- Autograd support for local field projections using `FieldProjectionKSpaceMonitor`. +- Function `components.geometry.utils.flatten_groups` now also flattens transformed groups when requested. +- Differentiable `smooth_min`, `smooth_max`, and `least_squares` functions in `tidy3d.plugins.autograd`. +- Differential operators `grad` and `value_and_grad` in `tidy3d.plugins.autograd` that behave similarly to the autograd operators but support auxiliary data via `aux_data=True` as well as differentiation w.r.t. `DataArray`. +- `@scalar_objective` decorator in `tidy3d.plugins.autograd` that wraps objective functions to ensure they return a scalar value and performs additional checks to ensure compatibility of objective functions with autograd. Used by default in `tidy3d.plugins.autograd.value_and_grad` as well as `tidy3d.plugins.autograd.grad`. +- Autograd support for simulations without adjoint sources in `run` as well as `run_async`, which will not attempt to run the simulation but instead return zero gradients. This can sometimes occur if the objective function gradient does not depend on some simulations, for example when using `min` or `max` in the objective. + +### Changed +- `CustomMedium` design regions require far less data when performing inverse design by reducing adjoint field monitor size for dims with one pixel. +- Calling `.values` on `DataArray` no longer raises a `DeprecationWarning` during automatic differentiation. +- Minimum number of PML layers set to 6. +- `Structure.background_permittivity : float` for specifying background medium for shape differentiation deprecated in favor of `Structure.background_medium : Medium` for more generality. +- Validate mode objects so that if a bend radius is defined, it is not smaller than half the size of the modal plane along the radial direction (i.e. the bend center is not inside the mode plane). + +### Fixed +- Regression in local field projection leading to incorrect projection results. +- Bug when differentiating with respect to `Cylinder.center`. +- `xarray` 2024.10.0 compatibility for autograd. +- Some failing examples in the expressions plugin documentation. +- Inaccuracy in transforming gradients from edge to `PolySlab.vertices`. +- Bug in `run_async` where an adjoint simulation would sometimes be assigned to the wrong forward simulation. +- Validate against nonlinearity or modulation in `FullyAnisotropicMedium.from_diagonal`. +- Add warning to complex-field nonlinearities, which may require more careful physical interpretation. + + +## [2.7.6] - 2024-10-30 + +### Added +- Users can pass the `shading` argument in the `SimulationData.plot_field` to `Xarray.plot` method. When `shading='gouraud'`, the image is interpolated, producing a smoother visualization. +- Users can manually specify the background medium for a structure to be used for geometry gradient calculations by supplying `Structure.background_permittivity`. This is useful when there are overlapping structures or structures embedded in other mediums. +- Autograd functions can now be called directly on `DataArray` (e.g., `np.sum(data_array)`) in objective functions. +- Automatic differentiation support for local field projections with `FieldProjectionAngleMonitor` and `FieldProjectionCartesianMonitor` using `FieldProjector.project_fields(far_field_monitor)`. + +### Changed +- Improved autograd tracer handling in `DataArray`, resulting in significant speedups for differentiation involving large monitors. +- Triangulation of `PolySlab` polygons now supports polygons with collinear vertices. +- Frequency and wavelength utilities under `tidy3d.frequencies` and `tidy3d.wavelengths`. + +### Fixed +- Minor gradient direction and normalization fixes for polyslab, field monitors, and diffraction monitors in autograd. +- Resolved an issue where temporary files for adjoint simulations were not being deleted properly. +- Resolve several edge cases where autograd boxes were incorrectly converted to numpy arrays. +- Resolve issue where scalar frequencies in metric definitions (`ModeAmp(f=freq)` instead of `ModeAmp(f=[freq])`) would erroneously fail validation. + +## [2.7.5] - 2024-10-16 + +### Added +- `TopologyDesignRegion` is now invariant in `z` by default and supports assigning dimensions along which a design should be uniform via `TopologyDesignRegion(uniform=(bool, bool, bool))`. +- Support for arbitrary padding sizes for all padding modes in `tidy3d.plugins.autograd.functions.pad`. +- `Expression.filter(target_type, target_field)` method for extracting object instances and fields from nested expressions. +- Additional constraints and validation logic to ensure correct setup of optimization problems in `invdes` plugin. +- `tidy3d.plugins.pytorch` to wrap autograd functions for interoperability with PyTorch via the `to_torch` wrapper. + +### Changed +- Renamed `Metric.freqs` --> `Metric.f` and made frequency argument optional, in which case all frequencies from the relevant monitor will be extracted. Metrics can still be initialized with both `f` or `freqs`. + +### Fixed +- Some validation fixes for design region. +- Bug in adjoint source creation that included empty sources for extraneous `FieldMonitor` objects, triggering unnecessary errors. +- Correct sign in objective function history depending on `Optimizer.maximize`. +- Fix to batch mode solver run that could create multiple copies of the same folder. +- Fixed ``ModeSolver.plot`` method when the simulation is not at the origin. +- Gradient calculation is orders of magnitude faster for large datasets and many structures by applying more efficient handling of field interpolation and passing to structures. +- Bug with infinite coordinates in `ClipOperation` not working with shapely. + +## [2.7.4] - 2024-09-25 + +### Added +- New `tidy3d.plugins.expressions` module for constructing and serializing mathematical expressions and simulation metrics like `ModeAmp` and `ModePower`. +- Support for serializable expressions in the `invdes` plugin (`InverseDesign(metric=ModePower(...))`). +- Added `InitializationSpec` as the default way to initialize design region parameters in the `invdes` plugin (`DesignRegion(initialization_spec=RandomInitializationSpec(...))`). +- Callback support in `invdes.Optimizer` and support for running the optimizer for a fixed number of steps via the `num_steps` argument in `Optimizer.continue_run()`. +- Convenience method `Structure.from_permittivity_array(geometry, eps_data)`, which creates structure containing `CustomMedium` with `eps_data` array sampled within `geometry.bounds`. + +### Changed +- All filter functions in `plugins/autograd` now accept either an absolute size in pixels or a `radius` and `dl` argument. +- Reverted fix for TFSF close to simulation boundaries that was introduced in 2.7.3 as it could cause different results in some cases with nonuniform mesh along the propagation direction. + +### Fixed +- Ensure `path` argument in `run()` function is respected when running under autograd or the adjoint plugin. +- Bug in `Simulation.subsection` (used in the mode solver) when nonlinear materials rely on information about sources outside of the region. + + +## [2.7.3] - 2024-09-12 + ### Added - Support for differentiation with respect to `ComplexPolySlab.vertices`. - Introduce RF material library. Users can now import `rf_material_library` from `tidy3d.plugins.microwave`. @@ -47,6 +137,29 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 - Fields stored in `ModeMonitor` objects were computed colocated to the grid boundaries, but the built-in `ModeMonitor.colocate=False` was resulting in wrong results in some cases, most notably if symmetries are also involved. - Small inaccuracy when applying a mode solver `bend_radius` when the simulation grid is not symmetric w.r.t. the mode plane center. Previously, the radius was defined w.r.t. the middle grid coordinate, while now it is correctly applied w.r.t. the plane center. - Silence warning in graphene from checking fit quality at large frequencies. +- Added value_and_grad function to the autograd plugin, importable via `from tidy3d.plugins.autograd import value_and_grad`. Supports differentiating functions with auxiliary data (`value_and_grad(f, has_aux=True)`). +- `Simulation.num_computational_grid_points` property to examine the number of grid cells that compose the computational domain corresponding to the simulation. This can differ from `Simulation.num_cells` based on boundary conditions and symmetries. +- Support for `dilation` argument in `JaxPolySlab`. +- Support for autograd differentiation with respect to `Cylinder.radius` and `Cylinder.center` (for elements not along axis dimension). +- `Cylinder.to_polyslab(num_pts_circumference, **kwargs)` to convert a cylinder into a discretized version represented by a `PolySlab`. +- `tidy3d.plugins.invdes.Optimizer` now accepts an optional optimization direction via the `maximize` argument (defaults to `maximize=True`). + +### Changed +- `PolySlab` now raises error when differentiating and dilation causes damage to the polygon. +- Validator `boundaries_for_zero_dims` to raise error when Bloch boundaries are used along 0-sized dims. +- `FieldProjectionKSpaceMonitor` support for 2D simulations with `far_field_approx = True`. + +### Fixed +- `DataArray` interpolation failure due to incorrect ordering of coordinates when interpolating with autograd tracers. +- Error in `CustomSourceTime` when evaluating at a list of times entirely outside of the range of the envelope definition times. +- Improved passivity enforcement near high-Q poles in `FastDispersionFitter`. Failed passivity enforcement could lead to simulation divergences. +- More helpful error and suggestion if users try to differentiate w.r.t. unsupported `FluxMonitor` output. +- Removed positive warnings in Simulation validators for Bloch boundary conditions. +- Improve accuracy in `Box` shifting boundary gradients. +- Improve accuracy in `FieldData` operations involving H fields (like `.flux`). +- Better error and warning handling in autograd pipeline. +- Added the option to specify the `num_freqs` argument and `kwargs` to the `.to_source` method for both `ModeSolver` and `ComponentModeler`. +- Fixes to TFSF source in some 2D simulations, and in some cases when the injection plane is close to the simulation domain boundaries. ## [2.7.2] - 2024-08-07 @@ -71,6 +184,11 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 - Error when plotting mode plane PML and the simulation has symmetry. - Validators using `TriangleMesh.intersections_plane` will fall back on bounding box in case the method fails for a non-watertight mesh. - Bug when running the same `ModeSolver` first locally then remotely, or vice versa, in which case the cached data from the first run is always returned. +- Gradient monitors for `PolySlab` only store fields at the center location along axis, reducing data usage. +- Validate the forward simulation on the client side even when using `local_gradient=False` for server-side gradient processing. +- Gradient inaccuracies in `PolySlab.vertices`, `Medium.conductivity`, and `DiffractionData` s-polarization. +- Adjoint field monitors no longer store H fields, which aren't needed for gradient calculation. +- `MeshOverrideStructures` in a `Simulation.GridSpec` are properly handled to remove any derivative tracers. ## [2.7.1] - 2024-07-10 @@ -1312,6 +1430,13 @@ which fields are to be projected is now determined automatically based on the me [Unreleased]: https://github.com/flexcompute/tidy3d/compare/v2.8.0rc1...pre/2.8 [2.8.0rc1]: https://github.com/flexcompute/tidy3d/compare/v2.7.1...v2.8.0rc1 +[Unreleased]: https://github.com/flexcompute/tidy3d/compare/v2.7.7...develop +[2.7.7]: https://github.com/flexcompute/tidy3d/compare/v2.7.6...v2.7.7 +[2.7.6]: https://github.com/flexcompute/tidy3d/compare/v2.7.5...v2.7.6 +[2.7.5]: https://github.com/flexcompute/tidy3d/compare/v2.7.4...v2.7.5 +[2.7.4]: https://github.com/flexcompute/tidy3d/compare/v2.7.3...v2.7.4 +[2.7.3]: https://github.com/flexcompute/tidy3d/compare/v2.7.2...v2.7.3 +[2.7.2]: https://github.com/flexcompute/tidy3d/compare/v2.7.1...v2.7.2 [2.7.1]: https://github.com/flexcompute/tidy3d/compare/v2.7.0...v2.7.1 [2.7.0]: https://github.com/flexcompute/tidy3d/compare/v2.6.4...v2.7.0 [2.6.4]: https://github.com/flexcompute/tidy3d/compare/v2.6.3...v2.6.4 diff --git a/README.md b/README.md index 545d8adae..b6eedaeff 100644 --- a/README.md +++ b/README.md @@ -62,7 +62,7 @@ python -c "import tidy3d; tidy3d.web.test()" It should pass without any errors if the API key is set up correctly. -To get started, our documentation has a lot of [examples](https://docs.flexcompute.com/projects/tidy3d/en/latest/notebooks/index.html) for inspiration. +To get started, our documentation has a lot of [examples](https://docs.flexcompute.com/projects/tidy3d/en/latest/notebooks/docs/index.html) for inspiration. ## Common Documentation References diff --git a/docs/_static/css/custom.css b/docs/_static/css/custom.css index dfac73147..34f82bdaa 100644 --- a/docs/_static/css/custom.css +++ b/docs/_static/css/custom.css @@ -29,6 +29,14 @@ iframe {width: 100%;border: unset} .toctree-wrapper li[class^=toctree-l3]>a {font-size: 0.9em; text-decoration: none} .toctree-wrapper li[class^=toctree-l4]>a {font-size: 0.9em; text-decoration: none} +.toctree-wrapper.example-notebook-toc li[class^=toctree-l1] > a::before { + content: "📘 "; /* Emoji before the link */ +} + +.toctree-wrapper.example-notebook-toc li[class^=toctree-l1] > a:hover::before { + content: "📖 "; /* Change emoji on hover */ +} + html[data-theme="light"] { --pst-color-primary: #b5445b; --pst-color-secondary: #4772ae; diff --git a/docs/_static/img/beam_waist.png b/docs/_static/img/beam_waist.png new file mode 100644 index 000000000..af0f0c284 Binary files /dev/null and b/docs/_static/img/beam_waist.png differ diff --git a/docs/api/_custom_autosummary/tidy3d.SimulationData.rst b/docs/api/_custom_autosummary/tidy3d.SimulationData.rst deleted file mode 100644 index 9b9830697..000000000 --- a/docs/api/_custom_autosummary/tidy3d.SimulationData.rst +++ /dev/null @@ -1,5 +0,0 @@ -ï»ż.. autosummary:: - :recursive: - :toctree: ../_autosummary/ - - tidy3d.components.base.Tidy3dBaseModel.from_file \ No newline at end of file diff --git a/docs/api/abstract_base.rst b/docs/api/abstract_base.rst index 2eee39fe8..9934fa8e2 100644 --- a/docs/api/abstract_base.rst +++ b/docs/api/abstract_base.rst @@ -3,15 +3,24 @@ Abstract Base Models ===================== -Base classes that represent abstractions of the core elements of a Simulation. +Base classes that represent abstractions of the core elements of a common components. Provide inherited functionality. .. autosummary:: :toctree: _autosummary/ :template: module.rst - tidy3d.components.base_sim.simulation.AbstractSimulation + tidy3d.components.base_sim.data.sim_data.AbstractSimulationData tidy3d.components.base_sim.monitor.AbstractMonitor + tidy3d.components.base_sim.simulation.AbstractSimulation tidy3d.components.base_sim.source.AbstractSource - tidy3d.components.source.Source - tidy3d.components.monitor.Monitor + tidy3d.components.data.dataset.AbstractFieldDataset + tidy3d.components.data.monitor_data.AbstractFieldProjectionData + tidy3d.components.parameter_perturbation.AbstractPerturbation + tidy3d.components.parameter_perturbation.AbstractPerturbation + tidy3d.components.medium.AbstractCustomMedium + tidy3d.components.medium.AbstractMedium + tidy3d.components.simulation.AbstractYeeGridSimulation + tidy3d.components.structure.AbstractStructure + tidy3d.components.time.AbstractTimeDependence + diff --git a/docs/api/abstract_models.rst b/docs/api/abstract_models.rst index dc230424f..47cddc373 100644 --- a/docs/api/abstract_models.rst +++ b/docs/api/abstract_models.rst @@ -1,19 +1,15 @@ .. currentmodule:: tidy3d -Abstract Models +Base Models =============== -These are some classes that are used to organize the tidy3d components, but aren't to be used directly in the code. Documented here mainly for reference. +These are some classes that are used to organize the tidy3d components, but aren't to be used directly in the code. Documented here mainly for reference of inherited components. .. autosummary:: :toctree: _autosummary/ :template: module.rst - tidy3d.components.base.Tidy3dBaseModel - tidy3d.components.base_sim.simulation.AbstractSimulation - tidy3d.components.simulation.AbstractYeeGridSimulation - tidy3d.components.boundary.AbsorberSpec tidy3d.Geometry tidy3d.components.geometry.base.Centered tidy3d.components.geometry.base.Planar @@ -47,12 +43,37 @@ These are some classes that are used to organize the tidy3d components, but aren tidy3d.components.base_sim.data.sim_data.SimulationData tidy3d.components.data.sim_data.AbstractYeeGridSimulationData tidy3d.components.data.sim_data.SimulationData + tidy3d.components.base.Tidy3dBaseModel + tidy3d.components.boundary.AbsorberSpec tidy3d.components.data.data_array.DataArray - tidy3d.components.data.monitor_data.MonitorData - tidy3d.components.data.monitor_data.AbstractFieldProjectionData - tidy3d.components.data.monitor_data.ElectromagneticFieldData - tidy3d.components.data.monitor_data.AbstractMonitorData - tidy3d.components.data.dataset.AbstractFieldDataset tidy3d.components.data.dataset.FieldDataset tidy3d.components.data.dataset.FieldTimeDataset tidy3d.components.data.dataset.ModeSolverDataset + tidy3d.components.data.monitor_data.ElectromagneticFieldData + tidy3d.components.data.monitor_data.MonitorData + tidy3d.components.data.sim_data.SimulationData + tidy3d.components.geometry.base.Centered + tidy3d.components.geometry.base.Circular + tidy3d.components.geometry.base.Planar + tidy3d.components.geometry.base.SimplePlaneIntersection + tidy3d.components.grid.grid_spec.GridSpec1d + tidy3d.components.lumped_element.LumpedElement + tidy3d.components.medium.CustomDispersiveMedium + tidy3d.components.medium.DispersiveMedium + tidy3d.components.monitor.FreqMonitor + tidy3d.components.monitor.Monitor + tidy3d.components.monitor.PlanarMonitor + tidy3d.components.monitor.TimeMonitor + tidy3d.components.source.AngledFieldSource + tidy3d.components.source.BroadbandSource + tidy3d.components.source.CurrentSource + tidy3d.components.source.DirectionalSource + tidy3d.components.source.FieldSource + tidy3d.components.source.PlanarSource + tidy3d.components.source.Pulse + tidy3d.components.source.ReverseInterpolatedSource + tidy3d.components.source.Source + tidy3d.components.source.SourceTime + tidy3d.components.source.VolumeSource + + diff --git a/docs/api/constants.rst b/docs/api/constants.rst index c109b8340..6a2c0ebca 100644 --- a/docs/api/constants.rst +++ b/docs/api/constants.rst @@ -4,7 +4,7 @@ Constants ========= Physical Constants ------------------- +--------------------- .. autosummary:: :toctree: _autosummary/ @@ -27,10 +27,9 @@ Tidy3D Special Constants :template: module.rst tidy3d.constants.inf - tidy3d.constants.PEC Tidy3D Configuration --------------------- +--------------------- .. autosummary:: :toctree: _autosummary/ @@ -39,7 +38,7 @@ Tidy3D Configuration tidy3d.config.Tidy3dConfig Default Absorber Parameters ---------------------------- +---------------------------- .. autosummary:: :toctree: _autosummary/ @@ -83,7 +82,7 @@ Units Precision & Comparator Values ---------------------------- +------------------------------ .. autosummary:: :toctree: _autosummary/ diff --git a/docs/api/eme/index.rst b/docs/api/eme/index.rst index 353a2e472..48b58bc2b 100644 --- a/docs/api/eme/index.rst +++ b/docs/api/eme/index.rst @@ -1,5 +1,5 @@ -EME -============= +EME |:rainbow:| +=============== .. toctree:: :hidden: diff --git a/docs/api/heat_charge/boundary_conditions.rst b/docs/api/heat_charge/boundary_conditions.rst index 32f1adf3f..6d20f8eea 100644 --- a/docs/api/heat_charge/boundary_conditions.rst +++ b/docs/api/heat_charge/boundary_conditions.rst @@ -4,7 +4,7 @@ Thermal/Charge Boundary Conditions ----------------------------- Specifications -'''''''''''''''''' +^^^^^^^^^^^^^^^^^ .. autosummary:: :toctree: ../_autosummary/ @@ -15,7 +15,7 @@ Specifications Types -'''''''''''''''''' +^^^^^^^^^^^^^^^^^ .. autosummary:: :toctree: ../_autosummary/ @@ -30,7 +30,7 @@ Types Placement -'''''''''''''''''' +^^^^^^^^^^^^^^^^^ .. autosummary:: :toctree: ../_autosummary/ diff --git a/docs/api/index.rst b/docs/api/index.rst index 985f65ace..6196f37de 100644 --- a/docs/api/index.rst +++ b/docs/api/index.rst @@ -15,6 +15,7 @@ API |:computer:| sources monitors mode + field_projector lumped_elements discretization subpixel_averaging @@ -38,6 +39,7 @@ API |:computer:| .. include:: /api/sources.rst .. include:: /api/monitors.rst .. include:: /api/mode.rst +.. include:: /api/field_projector.rst .. include:: /api/lumped_elements.rst .. include:: /api/discretization.rst .. include:: /api/subpixel_averaging.rst diff --git a/docs/api/mediums.rst b/docs/api/mediums.rst index 46db62eb7..82f5331a5 100644 --- a/docs/api/mediums.rst +++ b/docs/api/mediums.rst @@ -127,3 +127,13 @@ Material Library .. toctree:: material_library rf_material_library + + +Abstract Classes +----------------- + +.. autosummary:: + :toctree: _autosummary/ + + tidy3d.components.medium.AbstractPerturbationMedium + tidy3d.components.medium.NonlinearModel diff --git a/docs/api/plugins/adjoint.rst b/docs/api/plugins/adjoint.rst index 6ce31d2d9..97b8252e3 100644 --- a/docs/api/plugins/adjoint.rst +++ b/docs/api/plugins/adjoint.rst @@ -5,7 +5,7 @@ Adjoint .. toctree:: - ../../../tidy3d/plugins/adjoint/README + ./../../../tidy3d/plugins/adjoint/README.md .. autosummary:: :toctree: ../_autosummary/ diff --git a/docs/api/plugins/autograd.rst b/docs/api/plugins/autograd.rst index 09f9bca6a..1fdb5bc2b 100644 --- a/docs/api/plugins/autograd.rst +++ b/docs/api/plugins/autograd.rst @@ -5,7 +5,7 @@ Automatic Differentiation with Autograd .. toctree:: - ../../../tidy3d/plugins/autograd/README + ./../../../tidy3d/plugins/autograd/README .. autosummary:: :toctree: ../_autosummary/ @@ -23,12 +23,12 @@ Automatic Differentiation with Autograd tidy3d.plugins.autograd.functions.pad tidy3d.plugins.autograd.functions.convolve - tidy3d.plugins.autograd.primitives.gaussian_filter + tidy3d.plugins.autograd.utilities.chain + tidy3d.plugins.autograd.utilities.make_kernel + tidy3d.plugins.autograd.utilities.get_kernel_size_px - tidy3d.plugins.autograd.types.PaddingType - tidy3d.plugins.autograd.types.KernelType + tidy3d.plugins.autograd.primitives.gaussian_filter - tidy3d.plugins.autograd.invdes.get_kernel_size_px tidy3d.plugins.autograd.invdes.grey_indicator tidy3d.plugins.autograd.invdes.make_circular_filter tidy3d.plugins.autograd.invdes.make_conic_filter @@ -38,3 +38,7 @@ Automatic Differentiation with Autograd tidy3d.plugins.autograd.invdes.make_filter_and_project tidy3d.plugins.autograd.invdes.ramp_projection tidy3d.plugins.autograd.invdes.tanh_projection + + tidy3d.plugins.autograd.types.PaddingType + tidy3d.plugins.autograd.types.KernelType + diff --git a/docs/api/plugins/design.rst b/docs/api/plugins/design.rst index fea173080..9010e6842 100644 --- a/docs/api/plugins/design.rst +++ b/docs/api/plugins/design.rst @@ -11,9 +11,11 @@ Design Space Exploration :toctree: ../_autosummary/ :template: module.rst + tidy3d.plugins.design.parameter.Parameter tidy3d.plugins.design.ParameterFloat tidy3d.plugins.design.ParameterInt tidy3d.plugins.design.ParameterAny + tidy3d.plugins.design.method.Method tidy3d.plugins.design.MethodGrid tidy3d.plugins.design.MethodMonteCarlo tidy3d.plugins.design.MethodBayOpt diff --git a/docs/api/plugins/index.rst b/docs/api/plugins/index.rst index d60fc3691..b21f411a9 100644 --- a/docs/api/plugins/index.rst +++ b/docs/api/plugins/index.rst @@ -4,17 +4,17 @@ Plugins .. toctree:: :hidden: - mode_solver - dispersion - polyslab - smatrix - resonance - autograd - adjoint - design - invdes - waveguide - microwave + ./mode_solver + ./dispersion + ./polyslab + ./smatrix + ./resonance + ./autograd + ./adjoint + ./invdes + ./design + ./waveguide + ./microwave .. include:: /api/plugins/mode_solver.rst diff --git a/docs/api/plugins/invdes.rst b/docs/api/plugins/invdes.rst index cbd85ec10..ff2bfe27f 100644 --- a/docs/api/plugins/invdes.rst +++ b/docs/api/plugins/invdes.rst @@ -5,7 +5,7 @@ Inverse Design Plugin .. toctree:: - ../../../tidy3d/plugins/invdes/README + ./../../../tidy3d/plugins/invdes/README.md .. autosummary:: :toctree: ../_autosummary/ @@ -18,9 +18,9 @@ Inverse Design Plugin tidy3d.plugins.invdes.InverseDesignMulti tidy3d.plugins.invdes.InverseDesignResult tidy3d.plugins.invdes.AdamOptimizer - tidy3d.plugins.invdes.get_amps - tidy3d.plugins.invdes.get_field_component - tidy3d.plugins.invdes.get_intensity - tidy3d.plugins.invdes.sum_array - tidy3d.plugins.invdes.sum_abs_squared - tidy3d.plugins.invdes.get_phase + tidy3d.plugins.invdes.utils.get_amps + tidy3d.plugins.invdes.utils.get_field_component + tidy3d.plugins.invdes.utils.get_intensity + tidy3d.plugins.invdes.utils.sum_array + tidy3d.plugins.invdes.utils.sum_abs_squared + tidy3d.plugins.invdes.utils.get_phase diff --git a/docs/api/submit_simulations.rst b/docs/api/submit_simulations.rst index d4707003d..7f247947c 100644 --- a/docs/api/submit_simulations.rst +++ b/docs/api/submit_simulations.rst @@ -4,8 +4,8 @@ Submitting Simulations ====================== -Through python API ------------------- +Generic Web API +---------------- .. autosummary:: :toctree: _autosummary/ @@ -26,8 +26,8 @@ Through python API tidy3d.web.api.webapi.load_simulation tidy3d.web.api.asynchronous.run_async -Convenience for Single and Batch --------------------------------- +Job and Batch Containers +------------------------- .. autosummary:: :toctree: _autosummary/ @@ -46,3 +46,15 @@ Information Containers tidy3d.web.core.task_info.TaskInfo tidy3d.web.core.task_info.TaskStatus + + +Mode Solver Web API +-------------------- + +.. autosummary:: + :toctree: _autosummary/ + :template: module.rst + + tidy3d.web.api.mode.run + tidy3d.web.api.mode.run_batch + tidy3d.web.api.mode.ModeSolverTask \ No newline at end of file diff --git a/docs/conf.py b/docs/conf.py index d020264b7..ff8bde333 100644 --- a/docs/conf.py +++ b/docs/conf.py @@ -18,6 +18,7 @@ # absolute, like shown here. # import datetime +import logging import os import re import subprocess @@ -32,14 +33,14 @@ # TODO sort this out here = os.path.abspath(os.path.dirname(__file__)) sys.path.insert(0, os.path.abspath("_ext")) -sys.path.insert(0, os.path.abspath("source")) -sys.path.insert(0, os.path.abspath("notebooks")) -sys.path.insert(0, os.path.abspath("")) -sys.path.insert(0, os.path.abspath("../tidy3d")) -sys.path.insert(0, os.path.abspath("../tidy3d/components")) -sys.path.insert(0, os.path.abspath("../tidy3d/components/base_sim")) -sys.path.insert(0, os.path.abspath("../tidy3d/web")) -sys.path.insert(0, os.path.abspath("../tidy3d/plugins")) +# sys.path.insert(0, os.path.abspath("source")) +# sys.path.insert(0, os.path.abspath("notebooks")) +# # sys.path.insert(0, os.path.abspath("")) +# sys.path.insert(0, os.path.abspath("../tidy3d")) +# sys.path.insert(0, os.path.abspath("../tidy3d/components")) +# sys.path.insert(0, os.path.abspath("../tidy3d/components/base_sim")) +# sys.path.insert(0, os.path.abspath("../tidy3d/web")) +# sys.path.insert(0, os.path.abspath("../tidy3d/plugins")) # -- Project information ----------------------------------------------------- @@ -75,7 +76,18 @@ copybutton_prompt_is_regexp = True custom_sitemap_excludes = [r"/notebooks/"] # divparams_enable_postprocessing = True # TODO FIX -exclude_patterns = ["_build", "Thumbs.db", ".DS_Store", "**.ipynb_checkpoints", "faq/_faqs/*"] +exclude_patterns = [ + "_docs/", + "_templates/", + "_ext/", + "**.ipynb_checkpoints", + ".DS_Store", + "Thumbs.db", + "faq/_faqs/*", + "scripts/*", + "tests/*", + ".github/*", +] extensions = [ "IPython.sphinxext.ipython_directive", "IPython.sphinxext.ipython_console_highlighting", @@ -152,6 +164,17 @@ } latex_engine = "xelatex" language = "en" +include_patterns = [ + "tidy3d/*", + "faq/docs/**", + "notebooks/*.ipynb", + "notebooks/docs/*", + "**.rst", + "**.png", + "**.svg", + "**.txt", + "**/sitemap.xml", +] napoleon_google_docstring = False napoleon_numpy_docstring = True napoleon_include_init_with_doc = False @@ -239,3 +262,49 @@ # # # # 'figure_align': 'htbp', # } + + +class ImportWarningFilter(logging.Filter): + def filter(self, record): + # Suppress specific autosummary import warnings + message = record.getMessage() + if "autosummary: failed to import" in message and any( + phrase in message + for phrase in ["ModuleNotFoundError", "ValueError", "KeyError", "AttributeError"] + ): + return False + return True + + +class AutosummaryFilter(logging.Filter): + """ + This is basically a hack until I finally get round to writing our own custom sphinx extension which will customise + the way we represent our documentation properly. The goal of adding these filters is that at least we'll get useful + information on errors, rather than those related to the docs memory - stub page generation tradeoff. + """ + + def filter(self, record): + # Suppress "autosummary: stub file not found" warnings + if "autosummary" in record.getMessage() and "stub file not found" in record.getMessage(): + return False + return True + + +def add_import_warning_filter(app): + # Get the Sphinx logger + logger = logging.getLogger("sphinx") + # Add the custom filter to the logger + logger.addFilter(ImportWarningFilter()) + + +def add_autosummary_filter(app): + # Get the Sphinx logger + logger = logging.getLogger("sphinx") + # Add the custom filter to the logger + logger.addFilter(AutosummaryFilter()) + + +def setup(app): + # Apply the custom filter early in the build process + app.connect("builder-inited", add_autosummary_filter) + app.connect("builder-inited", add_import_warning_filter) diff --git a/docs/development/documentation.rst b/docs/development/documentation.rst index 1773ac7ed..09efe18e5 100644 --- a/docs/development/documentation.rst +++ b/docs/development/documentation.rst @@ -21,7 +21,7 @@ Common Updates -------------- Adding a new notebook -'''''''''''''''''''''' +^^^^^^^^^^^^^^^^^^^^^^ This process is self-contained in ``tidy3d-notebooks``. @@ -34,7 +34,7 @@ This submodule commit process can be done by running ``git add docs/notebooks`` If you want to locally develop notebooks in ``tidy3d/docs/notebooks`` then just use that submodule as your main development repository and commit to your local branch. Then when you are ready to publish, just make sure to commit the submodule to the latest ``develop`` branch. You can then build the documentation locally easily using this approach before it is published. Updating Docstrings -''''''''''''''''''''' +^^^^^^^^^^^^^^^^^^^^^^ The ``tidy3d develop`` suite includes a utility command ``replace-in-files``, which is designed to recursively find and replace strings in files within a specified directory. This functionality is particularly useful for updating docstrings across the codebase when there are changes in the API, ensuring that the documentation remains consistent with multiple version updates. This is useful when updating the API and you want to update the docstrings to reflect the changes from multiple versions. @@ -88,4 +88,19 @@ Further Guidance - The sphinx warnings are OK as long as the build occurs, errors will cause the crash the build. - Make sure all your internal API references start with ``tidy3d.`` -- In notebooks, always have absolute links, otherwise the links will break when the user downloads them. \ No newline at end of file +- In notebooks, always have absolute links, otherwise the links will break when the user downloads them. + + +Writing Documentation +^^^^^^^^^^^^^^^^^^^^^^^^ + +... raw:: + + Normally, there are no heading levels assigned to certain characters as the structure is determined from the succession of headings. However, this convention is used in Python Developer’s Guide for documenting which you may follow: + # with overline, for parts + * with overline, for chapters + = for sections + - for subsections + ^ for subsubsections + " for paragraphs + diff --git a/docs/development/index.rst b/docs/development/index.rst index ac79cf547..3f4a36306 100644 --- a/docs/development/index.rst +++ b/docs/development/index.rst @@ -2,22 +2,20 @@ Development Guide |:hammer_and_wrench:| ******************************************************************** -Welecome to the ``tidy3d`` developers guide! These are just some recommendations I've compiled, but we can change anything as we think might help the development cycle more. +Welecome to the ``tidy3d`` developers guide! + .. toctree:: - :maxdepth: 1 - :hidden: + :maxdepth: 2 - project_structure + introduction/index installation usage - documentation recommendations release/index -.. include:: /development/project_structure.rst +.. include:: /development/introduction/index.rst .. include:: /development/installation.rst .. include:: /development/usage.rst -.. include:: /development/documentation.rst .. include:: /development/recommendations.rst -.. include:: /development/release/index.rst \ No newline at end of file +.. include:: /development/release/index.rst diff --git a/docs/development/installation.rst b/docs/development/installation.rst index d485b788e..d45602444 100644 --- a/docs/development/installation.rst +++ b/docs/development/installation.rst @@ -1,20 +1,15 @@ -Installation -============== - -Beginners Guide -^^^^^^^^^^^^^^^ - - +Development Environment Installation +===================================== The Fast Lane -^^^^^^^^^^^^^ +-------------- Maybe you already have ``tidy3d`` installed in some form. After installing version ``tidy3d>=2.6``, you can use a few terminal commands to set you up on the correct environment and perform common development tasks. Just run in your terminal, :code:`tidy3d develop` to get the latest list of commands. It does not matter how you have installed ``tidy3d`` before as long as you have any form of ``tidy3d>=2.6`` in your environment. This can help you transition from a standard user installation to a development environment installation. Quick Start -'''''''''''' +^^^^^^^^^^^^^ Instructions for anyone who wants to migrate to the development flow from a version before 2.6: @@ -36,7 +31,7 @@ Now you can run the following ``tidy3d`` cli commands to test them. Automatic Environment Installation *Beta* -'''''''''''''''''''''''''''''''''''''''''' +"""""""""""""""""""""""""""""""""""""""""""""" If you are transitioning from the old development flow, to this new one, there are a list of commands you can run to make your life easier and set you up well: @@ -52,7 +47,7 @@ The way this command works is dependent on the operating system you are running. This command will first check if you already have installed the development requirements, and if not, it will run the installation scripts for ``pipx``, ``poetry``, and ask you to install the required version of ``pandoc``. It will also install the development requirements and ``tidy3d`` package in a specific ``poetry`` environment. Environment Verification -'''''''''''''''''''''''' +"""""""""""""""""""""""""""" If you rather install ``poetry``, ``pipx`` and ``pandoc`` yourself, you can run the following command to verify that your environment conforms to the reproducible development environment which would be equivalent to the one installed automatically above and described in :ref:`The Detailed Lane`. @@ -61,16 +56,14 @@ If you rather install ``poetry``, ``pipx`` and ``pandoc`` yourself, you can run tidy3d develop verify-dev-environment -.. _The Detailed Lane:: - The Detailed Lane -^^^^^^^^^^^^^^^^^ +------------------ If you do not have any of the above tools already installed and want to install them manually, let's go through the process of setting things up from scratch: Environment Requirements -'''''''''''''''''''''''''' +^^^^^^^^^^^^^^^^^^^^^^^^ Make sure you have installed ``pipx``. We provide common installation flows below: @@ -182,7 +175,7 @@ If you want to contribute to the project, read the following section: More Contribution Requirements -'''''''''''''''''''''''''''''' +^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ If you want to contribute to the development of ``tidy3d``, you can follow the instructions below to set up your development environment. This will allow you to run the tests, build the documentation, and run the examples. Another thing you need to do before committing to the project is to install the pre-commit hooks. This will ensure that your code is formatted correctly and passes the tests before you commit it. To do this, run the following command: @@ -204,7 +197,7 @@ You can also run the checks manually on all files by running the following comma Packaging Equivalent Functionality -''''''''''''''''''''''''''''''''''' +^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ This package installation process should be approximately equivalent to the previous ``setup.py`` installation flow. Independent of the ``poetry`` development flow, it is possible to run any of the following commands in any particular virtual environment you have configured: diff --git a/docs/development/introduction/code_quality_principles.rst b/docs/development/introduction/code_quality_principles.rst new file mode 100644 index 000000000..bf8890168 --- /dev/null +++ b/docs/development/introduction/code_quality_principles.rst @@ -0,0 +1,104 @@ +Code Quality Principles +------------------------ + +When writing a code snippet, remember the saying: "code is read more than written". We want to maintain our code maintainable, readable and high quality. + +Linting & Formatting +^^^^^^^^^^^^^^^^^^^^^^^^^^ + +To maintain code quality, we use `Ruff `_ as a linter and code formatter. A linter analyzes code to identify and flag potential errors, stylistic issues, and code that doesn't adhere to defined standards (such as `PEP8 `_). A code formatter automatically restructures the code to ensure it is consistently styled and properly formatted, making it consistent across the code base. + +Run ``ruff format`` to format all Python files: + +.. code-block:: bash + + poetry run ruff format . + +Run ``ruff check`` to check for style and other issues. Many common warnings can be automatically fixed with the ``--fix`` flag: + +.. code-block:: bash + + poetry run ruff check tidy3d --fix + +The configuration defining what ``ruff`` will correct lives in ``pyproject.toml`` under the ``[tool.ruff]`` section. + +When submitting code, for tests to pass, ``ruff`` should give no warnings. + +Documentation +^^^^^^^^^^^^^^^ + +Document all code you write using `NumPy-style docstrings `_. + +Testing +------- + +Here we will discuss how tests are defined and run in Tidy3d. + +Unit Testing +^^^^^^^^^^^^^^ + +The tests live in ``tests/`` directory. + +We use `pytest `_ package for our testing. + +To run all of the tests, call: + +.. code-block:: bash + + poetry run pytest -rA tests + +This command will trigger ``pytest`` to go through each file in ``tests/`` called ``test*.py`` and run each function in that file with a name starting with ``test``. + +If all of these functions run without any exceptions being raised, the tests pass! + +The specific configuration we use for ``pytest`` lives in the ``[tool.pytest.ini_options]`` section of ``pyproject.toml``. + +These tests are automatically run when code is submitted using GitHub Actions, which tests on Python 3.9 through 3.12 running on Ubuntu, MacOS, and Windows operating systems, as well as Flexcompute's servers. + +Note: The ``-rA`` flag is optional but produces output that is easily readable. + +Note: You may notice warnings and errors in the ``pytest`` output, this is because many of the tests intentionally trigger these warnings and errors to ensure they occur in certain situations. The important information about the success of the test is printed out at the bottom of the ``pytest`` output for each file. + +To get a code coverage report: + +.. code-block:: bash + + pip install pytest-cov + +if not already installed + +To run coverage tests with results printed to STDOUT: + +.. code-block:: bash + + pytest tests --cov-report term-missing --cov=tidy3d + +To run coverage tests and get output as .html (more intuitive): + +.. code-block:: bash + + pytest tests --cov-report=html --cov=tidy3d + open htmlcov/index.html + +Automated Testing +^^^^^^^^^^^^^^^^^^^^^^^^^^ + +We use GitHub Actions to perform these tests automatically and across different operating systems. + +On commits, each of the ``pytest`` tests are run using Python 3.9 - 3.12 installed on Ubuntu, MacOS, and Windows operating systems. + +See the "actions" tab for details on previous tests and ``.github/workflows/run_tests.yml`` for the configuration and to see the specific tests run. + +See `this `_ for more explanation. + +Other Tests +^^^^^^^^^^^^^^^^^^^^^^^^^^ + +There are additional tests in both the `documentation `_ and our private backend code. The same practices outlined here apply to those tests. + +More Resources on Testing +^^^^^^^^^^^^^^^^^^^^^^^^^^ + +A useful explanation for those curious to learn more about the reasoning behind these decisions: + +`https://www.youtube.com/watch?v=DhUpxWjOhME ` \ No newline at end of file diff --git a/docs/development/introduction/index.rst b/docs/development/introduction/index.rst new file mode 100644 index 000000000..e9301961e --- /dev/null +++ b/docs/development/introduction/index.rst @@ -0,0 +1,26 @@ +An Introduction to the Development Flow +======================================== + +This page hopefully will get you started to develop Tidy3D. + +**TLDR:** + +- Branch off of the target branch (usually ``develop`` or ``pre/x.x``), work on your branch, and submit a PR when ready. +- Use isolated development environments with ``poetry``. +- Use ``ruff`` to lint and format code, and install the pre-commit hook via ``pre-commit install`` to automate this. +- Document code using NumPy-style docstrings. +- Write unit tests for new features and try to maintain high test coverage. + +.. toctree:: + :maxdepth: 1 + + understanding_virtual_environments + understanding_poetry + code_quality_principles + project_structure + + +.. include:: /development/introduction/understanding_virtual_environments.rst +.. include:: /development/introduction/understanding_poetry.rst +.. include:: /development/introduction/code_quality_principles.rst +.. include:: /development/introduction/project_structure.rst \ No newline at end of file diff --git a/docs/development/project_structure.rst b/docs/development/introduction/project_structure.rst similarity index 97% rename from docs/development/project_structure.rst rename to docs/development/introduction/project_structure.rst index 736ef5512..522700ea4 100644 --- a/docs/development/project_structure.rst +++ b/docs/development/introduction/project_structure.rst @@ -1,5 +1,5 @@ -Project Structure -================= +``tidy3d`` Project Structure +----------------------------- As of ``tidy3d>=2.6``, the frontend has been restructured to improve the development cycle. The project directories follow the following structure, which is derived from some recommended `Python project architecture guides `_. This is a handy structure because many tools, such as ``sphinx``, integrate quite well with this type of project layout. @@ -31,7 +31,7 @@ It is important to note the new tools we are using to manage our development env - ``pipx`` Important Branches ------------------- +^^^^^^^^^^^^^^^^^^^ We currently have *three* main branches that have to be kept track of when creating a release, each with different functionality. diff --git a/docs/development/introduction/understanding_poetry.rst b/docs/development/introduction/understanding_poetry.rst new file mode 100644 index 000000000..a415066a8 --- /dev/null +++ b/docs/development/introduction/understanding_poetry.rst @@ -0,0 +1,96 @@ +Using `poetry` for package management +-------------------------------------- + +What is Poetry +^^^^^^^^^^^^^^^^^^^^^^^^^^ + +`Poetry `_ is a package management tool for Python. + +Among other things, it provides a nice way to: + +- Manage dependencies +- Publish packages +- Set up and use virtual environments + +Effectively, it is a command line utility (similar to ``pip``) that is a bit more convenient and allows more customization. + +Why do we want to use it +^^^^^^^^^^^^^^^^^^^^^^^^^^ + +1. To improve our dependency management, which is used to be all over the place. We have several ``requirements.txt`` files that get imported into ``setup.py`` and parsed depending on the extra arguments passed to ``pip install``. ``Poetry`` handles this much more elegantly through a ``pyproject.toml`` file that defines the dependency configuration very explicitly in a simple data format. +2. Reproducible development virtual environments means that everyone is using the exact same dependencies, without conflicts. This also improves our packaging and release flow. + +How to install it? +^^^^^^^^^^^^^^^^^^^^^^^^^^ + +We provide custom installation instructions and an installation script on TODO ADD LINK SECTION. However, you can read more information here: see the poetry documentation for a guide to `installation `_ and `basic use `_. + + +Usage Examples +^^^^^^^^^^^^^^^^^^^^^^^^^^ + +To add poetry to a project +"""""""""""""""""""""""""""" + +To initialize a new basic project with poetry configured, run: + +.. code-block:: bash + + poetry new poetry-demo + +To add poetry to an existing project, ``cd`` to the project directory and run: + +.. code-block:: bash + + poetry init + +Configuring dependencies +^^^^^^^^^^^^^^^^^^^^^^^^^^ + +The dependency configuration is in the editable file called ``pyproject.toml``. Here you can specify whatever dependencies you want in your project, their versions, and even different levels of dependencies (e.g., ``dev``). + +To add a dependency to the project (e.g., ``numpy``), run: + +.. code-block:: bash + + poetry add numpy + +You can then verify that it was added to the ``tool.poetry.dependencies`` section of ``pyproject.toml``. + +For many more options on defining dependencies, see `here `_. + +Virtual environments +^^^^^^^^^^^^^^^^^^^^^^^^^^ + +Now that the project has had poetry configured and the correct dependencies are specified, we can use poetry to run our scripts/shell commands from a virtual environment without much effort. There are a few ways to do this: + +**Poetry run**: One way is to precede any shell command you’d normally run with ``poetry run``. For example, if you want to run ``python tidy_script.py`` from the virtual environment set up by poetry, you’d do: + +.. code-block:: bash + + poetry run python tidy3d_script.py + +**Poetry shell**: + +If you want to open up a shell session with the environment activated, you can run: + +.. code-block:: bash + + poetry shell + +And then run your commands. To return to the original shell, run ``exit``. + +There are many more advanced options explained `here `_. + +Publishing Package +^^^^^^^^^^^^^^^^^^^^^^^^^^ + +To upload the package to PyPI: + +.. code-block:: bash + + poetry build + + poetry publish + +Note that some `configuration `_ must be set up before this would work properly. diff --git a/docs/development/introduction/understanding_virtual_environments.rst b/docs/development/introduction/understanding_virtual_environments.rst new file mode 100644 index 000000000..165344cf5 --- /dev/null +++ b/docs/development/introduction/understanding_virtual_environments.rst @@ -0,0 +1,64 @@ +Understanding Virtual Environments +---------------------------------- + +Introduction +^^^^^^^^^^^^^ + +In larger projects, it's crucial to have a *separate* Python environment for each feature or branch you work on. This practice ensures isolation and reproducibility, simplifying testing and debugging by allowing issues to be traced back to specific environments. It also facilitates smoother integration and deployment processes, ensuring controlled and consistent development. +Managing multiple environments might seem daunting, but it's straightforward with the right tools. Follow the steps below to set up and manage your environments efficiently. + +Benefits +^^^^^^^^^^^^^ + +- **Isolation**: Avoids conflicts between dependencies of different features. +- **Reproducibility**: Each environment can be easily replicated. +- **Simplified Testing**: Issues are contained within their respective environments. +- **Smooth Integration**: Ensures features are developed in a consistent setting. + +Prerequisites +^^^^^^^^^^^^^^ + +Make sure that you have ``poetry`` installed. This can be done system-wide with ``pipx`` or within a ``conda`` environment. Note that we use ``conda`` only for setting up the interpreter (Python version) and ``poetry``, not for managing dependencies. +Refer to the official development guide for detailed instructions: + +`https://docs.flexcompute.com/projects/tidy3d/en/stable/development/index.html#installation `_ + +Setting Up a New Environment +^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + +1. Check out the branch: + + .. code-block:: bash + + git checkout branch + +2. Set up the environment with ``conda`` (skip this step if you don’t use ``conda``): + + .. code-block:: bash + + conda create -n branch_env python=3.11 poetry + conda activate branch_env + poetry env use system + poetry env info # verify you're running the right environment now + +3. Install dependencies with ``poetry``: + + .. code-block:: bash + + poetry install -E dev + poetry run pre-commit install + +4. Update the environment when switching to a different branch: + + .. code-block:: bash + + poetry install -E dev + + + +Multiple Folders or Worktrees +^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + +If you have multiple folders (e.g., multiple clones or git worktrees), you will need to repeat the environment setup for each folder. Ensure that each folder has its own isolated environment. + +By following these steps, you can maintain isolated and reproducible environments for each branch and feature, leading to a more efficient and error-free development process. \ No newline at end of file diff --git a/docs/development/recommendations.rst b/docs/development/recommendations.rst index e2964df19..d6d3d042e 100644 --- a/docs/development/recommendations.rst +++ b/docs/development/recommendations.rst @@ -50,7 +50,7 @@ However, if we do decide to commit with emojis, I believe it would be worth havi Package Speedup Best Practices ----------------------------- +-------------------------------- ``tidy3d`` is a pretty big project already, and will get bigger. We want to optimise the performance of the codebase throughout the multiple operations that we perform. @@ -77,7 +77,7 @@ This is because the latter will import the entire package, which is not necessar Managing Optional Dependencies On-The-Fly ----------------------------- +------------------------------------------ If you look within ``pyproject.toml``, it is possible to see that we have different packages relating to different functionalities that are optional. diff --git a/docs/development/release/documentation.rst b/docs/development/release/documentation.rst index 8b8050ce4..1f73fb01f 100644 --- a/docs/development/release/documentation.rst +++ b/docs/development/release/documentation.rst @@ -19,7 +19,7 @@ The `latest` branch holds the state of the docs that we want to host in `latest` The `stable` version of the docs on our website is built based on the last version tag which is not a pre-release tag (no `rc` ending). Hot Fix & Submodule Updates -''''''''''''''''''''''''''' +^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ To make a “hot fix” (eg fix a typo, add a notebook, update the release FAQ), just update the ``latest`` branch in ``tidy3d`` repo. This should automatically sync to `tidy3d-docs`, and trigger a docs rebuild. **However, we should avoid this as this will cause the ``develop`` and ``latest branches`` to diverge.** Ideally, these hot fixes could wait until the next pre/post-release to be propagated through. diff --git a/docs/development/release/flow.rst b/docs/development/release/flow.rst index accac7adf..55df801e7 100644 --- a/docs/development/release/flow.rst +++ b/docs/development/release/flow.rst @@ -1,15 +1,37 @@ +Feature Contribution +----------------------- + + + Feature Development Workflow ------------------------------- +^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + + +1. Create a branch off of ``develop`` +"""""""""""""""""""""""""""""""""""""""" + +Our ``pre/x.x`` branches are where new features, bug fixes, and other things get added before a major release. To switch to the ``pre/x.x`` branch: + +.. code-block:: bash + + git checkout pre/x.x + +And then you can create a new branch from here with your GitHub username pre-pended: + +.. code-block:: bash + + git checkout -b myusername/cool-feature Currently most of our release development flow is made under the latest ``pre/*`` branch under the main frontend -tidy3d repository. You want to fork from this latest branch to develop your feature in order for it to be included -under that release. +tidy3d repository. You want to fork from this latest branch to develop your feature in order for it to be included under that release. We are using a variation of the `gitflow workflow `__ -- so this is the first thing to familiarize yourselves with. The -splitting of branches into ``main``, ``develop`` and separate feature -branches is as explained there. Most importantly, **all contributions +- so this is the first thing to familiarize yourselves with. The link provided explains it very well, but to summarize: features get added to a pre-release branch (``pre/x.x``), and once all the features for a particular release have been implemented, the pre-release branch gets merged into ``develop``. The ``latest`` branch holds the code we want most users to be using. When we wish to release a new version, we simply merge the state of ``develop`` into ``latest``, propagating all of the changes at once. We will describe this process in more detail below. + +When developing new features, we ask that you create a branch off of whichever branch you aim to contribute to. This is typically either the current pre-release branch named ``pre/x.x`` or ``develop`` depending on what stage of development we are currently in. You then work on your branch, and when the feature is ready to merge in, we prefer ``git rebase`` to ``git merge``. This creates a cleaner, linear history. You can read about why we do it and what a rebase is at `this link `_. And see Momchil’s more specific notes `here `_. + +Most importantly, **all contributions should happen through a PR from a feature branch into the develop branch.** @@ -29,23 +51,36 @@ a clean branch on your machine. .. code-block:: bash # from the main tidy3d frontend repo - git checkout develop - git pull origin develop + git checkout pre/x.x + git pull origin pre/x.x git checkout -b my_name/new_feature +2. Writing code +"""""""""""""""""""""""""""""""""""""""" + +Develop your code in this new branch, committing your changes when it seems like a natural time to “save your progress”. + +If you are working on a new feature, make sure you add a line in the `CHANGELOG.md `_ file (if it exists in that repository) to summarize your changes. + + +3. Create a pull request on GitHub +"""""""""""""""""""""""""""""""""""""""" + +First, push your changes to your branch on GitHub. + +In the GitHub website, create a pull request to merge your branch into ``pre/x.x``. -Create your feature rebase -'''''''''''''''''''''''''''''' +Write some comments or a summary of changes in your pull request to be clear about what is being added/changed and why. Before rebasing, you should make sure you have the latest version of ``develop``, in case other work has been merged meanwhile. .. code-block:: bash - git checkout develop - git pull origin develop + git checkout pre/x.x + git pull origin pre/x.x git checkout my_name/new_feature - git rebase -i develop + git rebase -i pre/x.x This will now open an editor that will allow you to edit the commits in the feature branch. There is plenty of explanations of the various @@ -76,8 +111,22 @@ rebasing has changed its history. git push -f origin my_name/new_feature -Submitting to PR -''''''''''''''''' + +4. Submit for review +""""""""""""""""""""" + +Every PR must have the following before it can be merged: + +- At least one review. +- A description in the CHANGELOG of what has been done. + +Every new major feature must also pass all of the following before it can be merged: + +- Frontend and backend tests by the developer (unless no code has changed on one or the other), as well as a new example notebook or a modification to an existing example notebook that utilizes the new feature. Intermediate reviews can happen, but these conditions must be met for the feature to begin to be considered for a merge. +- Ensure any known limitations are listed at the top message in the PR conversation (e.g., does the feature work with the mode solver? The auto grid? Does it work, but not as well as it should?). The feature can be merged given the limitations if we make a decision to do that, but only if an error or warning is issued whenever a user could encounter them, and after the list has been moved to another PR or an issue to keep track. +- If backend changes are present, review by one of the people well-versed with the solver (Momchil, Weiliang, Shashwat, Daniil). +- If frontend changes are present, review by any member of the team and additional approval by Momchil or Tyler. +- QA from any member of the team: playing around with the new feature and trying to find limitations. The goal is not to construct one successful example but to figure out if there is any allowed usage that may be problematic. An extra example notebook may or may not come out of this. After this, you can notify Momchil that the branch is ready to to be merged. In the comment you can optionally also say things like “Fixes @@ -94,4 +143,3 @@ rebasing allows you to clean everything up. feature, i.e. all PR comments have been addressed, etc. This is not critical, but is nicer to only rebase in the end so as not to muddle up the PR discussion when you force push the new branch (see below). - diff --git a/docs/development/release/version.rst b/docs/development/release/version.rst index d13ffa44f..25f48ee05 100644 --- a/docs/development/release/version.rst +++ b/docs/development/release/version.rst @@ -4,9 +4,8 @@ Releasing a new ``tidy3d`` version This document contains the relevant information to create and publish a new tidy3d version. Version Information Management -^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ +^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ The ``pyproject.toml`` is declarative (ie static) and provides information to the packaging tools like PyPi on what version is ``tidy3d``. However, we also have a ``version.py`` file so that we can dynamically query ``tidy3d.__version__`` within our python version. These two files need to be kept with the same version. This is achieved by using the ``bump-my-version`` utility as described in the following section. **These files should not be manually updated.** -The configuration of the way the version bumping occurs is described in the ``pyproject.toml``. - +The configuration of the way the version bumping occurs is described in the ``pyproject.toml``. \ No newline at end of file diff --git a/docs/development/usage.rst b/docs/development/usage.rst index 1bacada72..37a706fc9 100644 --- a/docs/development/usage.rst +++ b/docs/development/usage.rst @@ -2,7 +2,7 @@ Using the Development Flow ========================== Developing ``tidy3d`` with ``poetry`` -^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ +-------------------------------------------------- `poetry `_ is an incredibly powerful tool for reproducible package development environments and dependency management. @@ -19,7 +19,7 @@ It is important to note the function above is equivalent to ``pip install tidy3d ``poetry`` with an external virtual environment --------------------------------------------------- +^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ It is recommended to use ``poetry`` for package development. However, there are some cases where you might need to use an external virtual environment for some operations. There are a few workarounds where you can leverage the reproducibility of the ``poetry`` managed environment with the freedom of a standard virtual environment. There are a few more instructions and explanations in `the poetry env docs `_ . F See the following example: @@ -35,7 +35,7 @@ It is recommended to use ``poetry`` for package development. However, there are There are also other methodologies of implementing common dependencies management. Common Utilities -^^^^^^^^^^^^^^^^ +"""""""""""""""""""" There are a range of handy development functions that you might want to use to streamline your development experience. diff --git a/docs/index.rst b/docs/index.rst index 803f50124..d99a5d16e 100644 --- a/docs/index.rst +++ b/docs/index.rst @@ -187,7 +187,7 @@ Further Information About our Solver Github Repositories -~~~~~~~~~~~~~~~~~ +-------------------- .. list-table:: :header-rows: 1 diff --git a/docs/install.rst b/docs/install.rst index 44cea860b..64039c931 100644 --- a/docs/install.rst +++ b/docs/install.rst @@ -135,7 +135,7 @@ Some users or systems may require a more specialized installation, which we will If you already have Python installed on your computer, it is possible that some packages in your current environment could have version conflicts with Tidy3D. To avoid this, we strongly recommend that you create a clean Python virtual environment to install Tidy3D. - We recommend using the Mamba package management system to manage your Python virtual environment as well as installing Tidy3D. You can install Mamba conveniently following `these instructions `__. + We recommend using the Mamba package management system to manage your Python virtual environment as well as installing Tidy3D. You can install Mamba conveniently following `these instructions `__. After you install Anaconda, open the Anaconda Prompt and enter @@ -181,4 +181,4 @@ To see some other examples of Tidy3D being used in large scale photonics simulat To learn more about the many features of Tidy3D, check out our `Feature Walkthrough <./notebooks/Simulation.html>`_. -Or, if you're interested in the API documentation, see `API Reference <./api/index.html>`_. \ No newline at end of file +Or, if you're interested in the API documentation, see `API Reference <./api/index.html>`_. diff --git a/docs/lectures/fdtd_workshop.rst b/docs/lectures/fdtd_workshop.rst new file mode 100644 index 000000000..36640fac7 --- /dev/null +++ b/docs/lectures/fdtd_workshop.rst @@ -0,0 +1,18 @@ +*************************** +Future-Ready FDTD Workshop +*************************** + +The Future-Ready FDTD Workshop Series is a collaborative effort between Tidy3D, iOptics, and Optica. The series aims to present a range of workshops focused on FDTD simulations for photonic devices. The workshop sessions took place from February 9th to March 8th, 2024, and covered theoretical foundations of concepts such as the FDTD method, Bragg filters and reflectors, multimode photonic design, photonic crystals, and inverse design. +Each session included a live demonstration of setting up and running an FDTD simulation. Additionally, a challenge was proposed to enhance participants' skills. Challenge solutions are available upon request through our technical support. + +`Lecture 1: The FDTD Method Demystified `_ + +`Lecture 2: Multimode Photonics Design Made Easy `_ + +`Lecture 3: Designing Filters and Reflectors with Bragg Gratings `_ + +`Lecture 4: Photonic Crystal Slabs Controlling and Confining Light on the Nanoscale `_ + +`Lecture 5: Inverse Design in Photonics An Introduction `_ + +More lectures coming soon! diff --git a/docs/lectures/index.rst b/docs/lectures/index.rst index 0c1438764..2ee700759 100644 --- a/docs/lectures/index.rst +++ b/docs/lectures/index.rst @@ -8,7 +8,10 @@ Welcome to our lecture series! :hidden: fdtd101 + fdtd_workshop inversedesign + .. include:: /lectures/fdtd101.rst +.. include:: /lectures/fdtd_workshop.rst .. include:: /lectures/inversedesign.rst \ No newline at end of file diff --git a/docs/notebooks b/docs/notebooks index 3ad84c09b..a2f1abe29 160000 --- a/docs/notebooks +++ b/docs/notebooks @@ -1 +1 @@ -Subproject commit 3ad84c09bfa80ceb3cb4cdf894472430acf2a1fb +Subproject commit a2f1abe2980bfae56b8bdeb969a96e6bed5f4a72 diff --git a/poetry.lock b/poetry.lock index f22a85d8a..279a9c4fc 100644 --- a/poetry.lock +++ b/poetry.lock @@ -1,4 +1,4 @@ -# This file is automatically @generated by Poetry 1.8.3 and should not be changed by hand. +# This file is automatically @generated by Poetry 1.8.4 and should not be changed by hand. [[package]] name = "absl-py" @@ -227,18 +227,20 @@ tests-mypy = ["mypy (>=1.11.1)", "pytest-mypy-plugins"] [[package]] name = "autograd" -version = "1.6.2" -description = "Efficiently computes derivatives of numpy code." +version = "1.7.0" +description = "Efficiently computes derivatives of NumPy code." optional = false -python-versions = "*" +python-versions = ">=3.8" files = [ - {file = "autograd-1.6.2-py3-none-any.whl", hash = "sha256:208dde2a938e63b4f8f5049b1985505139e529068b0d26f8cd7771fd3eb145d5"}, - {file = "autograd-1.6.2.tar.gz", hash = "sha256:8731e08a0c4e389d8695a40072ada4512641c113b6cace8f4cfbe8eb7e9aedeb"}, + {file = "autograd-1.7.0-py3-none-any.whl", hash = "sha256:49680300f842f3a8722b060ac0d3ed7aca071d1ad4d3d38c9fdadafdcc73c30b"}, + {file = "autograd-1.7.0.tar.gz", hash = "sha256:de743fd368d6df523cd37305dcd171861a9752a144493677d2c9f5a56983ff2f"}, ] [package.dependencies] -future = ">=0.15.2" -numpy = ">=1.12" +numpy = "*" + +[package.extras] +scipy = ["scipy"] [[package]] name = "babel" @@ -1323,7 +1325,7 @@ tqdm = ["tqdm"] name = "future" version = "1.0.0" description = "Clean single-source support for Python 3 and 2" -optional = false +optional = true python-versions = ">=2.6, !=3.0.*, !=3.1.*, !=3.2.*" files = [ {file = "future-1.0.0-py3-none-any.whl", hash = "sha256:929292d34f5872e70396626ef385ec22355a1fae8ad29e1a734c3e43f9fbc216"}, @@ -2770,13 +2772,30 @@ files = [ numpy = [ {version = ">=1.26.0", markers = "python_version >= \"3.12\" and python_version < \"3.13\""}, {version = ">=1.23.3", markers = "python_version >= \"3.11\" and python_version < \"3.12\""}, - {version = ">=1.21", markers = "python_version < \"3.10\""}, {version = ">=1.21.2", markers = "python_version >= \"3.10\" and python_version < \"3.11\""}, + {version = ">=1.21", markers = "python_version < \"3.10\""}, ] [package.extras] dev = ["absl-py", "pyink", "pylint (>=2.6.0)", "pytest", "pytest-xdist"] +[[package]] +name = "mpmath" +version = "1.3.0" +description = "Python library for arbitrary-precision floating-point arithmetic" +optional = true +python-versions = "*" +files = [ + {file = "mpmath-1.3.0-py3-none-any.whl", hash = "sha256:a0b2b9fe80bbcd81a6647ff13108738cfb482d481d826cc0e02f5b35e5c88d2c"}, + {file = "mpmath-1.3.0.tar.gz", hash = "sha256:7a28eb2a9774d00c7bc92411c19a89209d5da7c4c9a9e227be8330a23a25b91f"}, +] + +[package.extras] +develop = ["codecov", "pycodestyle", "pytest (>=4.6)", "pytest-cov", "wheel"] +docs = ["sphinx"] +gmpy = ["gmpy2 (>=2.1.0a4)"] +tests = ["pytest (>=4.6)"] + [[package]] name = "msgpack" version = "1.1.0" @@ -3153,6 +3172,161 @@ files = [ {file = "numpy-1.26.4.tar.gz", hash = "sha256:2a02aba9ed12e4ac4eb3ea9421c420301a0c6460d9830d74a9df87efa4912010"}, ] +[[package]] +name = "nvidia-cublas-cu12" +version = "12.4.5.8" +description = "CUBLAS native runtime libraries" +optional = true +python-versions = ">=3" +files = [ + {file = "nvidia_cublas_cu12-12.4.5.8-py3-none-manylinux2014_aarch64.whl", hash = "sha256:0f8aa1706812e00b9f19dfe0cdb3999b092ccb8ca168c0db5b8ea712456fd9b3"}, + {file = "nvidia_cublas_cu12-12.4.5.8-py3-none-manylinux2014_x86_64.whl", hash = "sha256:2fc8da60df463fdefa81e323eef2e36489e1c94335b5358bcb38360adf75ac9b"}, + {file = "nvidia_cublas_cu12-12.4.5.8-py3-none-win_amd64.whl", hash = "sha256:5a796786da89203a0657eda402bcdcec6180254a8ac22d72213abc42069522dc"}, +] + +[[package]] +name = "nvidia-cuda-cupti-cu12" +version = "12.4.127" +description = "CUDA profiling tools runtime libs." +optional = true +python-versions = ">=3" +files = [ + {file = "nvidia_cuda_cupti_cu12-12.4.127-py3-none-manylinux2014_aarch64.whl", hash = "sha256:79279b35cf6f91da114182a5ce1864997fd52294a87a16179ce275773799458a"}, + {file = "nvidia_cuda_cupti_cu12-12.4.127-py3-none-manylinux2014_x86_64.whl", hash = "sha256:9dec60f5ac126f7bb551c055072b69d85392b13311fcc1bcda2202d172df30fb"}, + {file = "nvidia_cuda_cupti_cu12-12.4.127-py3-none-win_amd64.whl", hash = "sha256:5688d203301ab051449a2b1cb6690fbe90d2b372f411521c86018b950f3d7922"}, +] + +[[package]] +name = "nvidia-cuda-nvrtc-cu12" +version = "12.4.127" +description = "NVRTC native runtime libraries" +optional = true +python-versions = ">=3" +files = [ + {file = "nvidia_cuda_nvrtc_cu12-12.4.127-py3-none-manylinux2014_aarch64.whl", hash = "sha256:0eedf14185e04b76aa05b1fea04133e59f465b6f960c0cbf4e37c3cb6b0ea198"}, + {file = "nvidia_cuda_nvrtc_cu12-12.4.127-py3-none-manylinux2014_x86_64.whl", hash = "sha256:a178759ebb095827bd30ef56598ec182b85547f1508941a3d560eb7ea1fbf338"}, + {file = "nvidia_cuda_nvrtc_cu12-12.4.127-py3-none-win_amd64.whl", hash = "sha256:a961b2f1d5f17b14867c619ceb99ef6fcec12e46612711bcec78eb05068a60ec"}, +] + +[[package]] +name = "nvidia-cuda-runtime-cu12" +version = "12.4.127" +description = "CUDA Runtime native Libraries" +optional = true +python-versions = ">=3" +files = [ + {file = "nvidia_cuda_runtime_cu12-12.4.127-py3-none-manylinux2014_aarch64.whl", hash = "sha256:961fe0e2e716a2a1d967aab7caee97512f71767f852f67432d572e36cb3a11f3"}, + {file = "nvidia_cuda_runtime_cu12-12.4.127-py3-none-manylinux2014_x86_64.whl", hash = "sha256:64403288fa2136ee8e467cdc9c9427e0434110899d07c779f25b5c068934faa5"}, + {file = "nvidia_cuda_runtime_cu12-12.4.127-py3-none-win_amd64.whl", hash = "sha256:09c2e35f48359752dfa822c09918211844a3d93c100a715d79b59591130c5e1e"}, +] + +[[package]] +name = "nvidia-cudnn-cu12" +version = "9.1.0.70" +description = "cuDNN runtime libraries" +optional = true +python-versions = ">=3" +files = [ + {file = "nvidia_cudnn_cu12-9.1.0.70-py3-none-manylinux2014_x86_64.whl", hash = "sha256:165764f44ef8c61fcdfdfdbe769d687e06374059fbb388b6c89ecb0e28793a6f"}, + {file = "nvidia_cudnn_cu12-9.1.0.70-py3-none-win_amd64.whl", hash = "sha256:6278562929433d68365a07a4a1546c237ba2849852c0d4b2262a486e805b977a"}, +] + +[package.dependencies] +nvidia-cublas-cu12 = "*" + +[[package]] +name = "nvidia-cufft-cu12" +version = "11.2.1.3" +description = "CUFFT native runtime libraries" +optional = true +python-versions = ">=3" +files = [ + {file = "nvidia_cufft_cu12-11.2.1.3-py3-none-manylinux2014_aarch64.whl", hash = "sha256:5dad8008fc7f92f5ddfa2101430917ce2ffacd86824914c82e28990ad7f00399"}, + {file = "nvidia_cufft_cu12-11.2.1.3-py3-none-manylinux2014_x86_64.whl", hash = "sha256:f083fc24912aa410be21fa16d157fed2055dab1cc4b6934a0e03cba69eb242b9"}, + {file = "nvidia_cufft_cu12-11.2.1.3-py3-none-win_amd64.whl", hash = "sha256:d802f4954291101186078ccbe22fc285a902136f974d369540fd4a5333d1440b"}, +] + +[package.dependencies] +nvidia-nvjitlink-cu12 = "*" + +[[package]] +name = "nvidia-curand-cu12" +version = "10.3.5.147" +description = "CURAND native runtime libraries" +optional = true +python-versions = ">=3" +files = [ + {file = "nvidia_curand_cu12-10.3.5.147-py3-none-manylinux2014_aarch64.whl", hash = "sha256:1f173f09e3e3c76ab084aba0de819c49e56614feae5c12f69883f4ae9bb5fad9"}, + {file = "nvidia_curand_cu12-10.3.5.147-py3-none-manylinux2014_x86_64.whl", hash = "sha256:a88f583d4e0bb643c49743469964103aa59f7f708d862c3ddb0fc07f851e3b8b"}, + {file = "nvidia_curand_cu12-10.3.5.147-py3-none-win_amd64.whl", hash = "sha256:f307cc191f96efe9e8f05a87096abc20d08845a841889ef78cb06924437f6771"}, +] + +[[package]] +name = "nvidia-cusolver-cu12" +version = "11.6.1.9" +description = "CUDA solver native runtime libraries" +optional = true +python-versions = ">=3" +files = [ + {file = "nvidia_cusolver_cu12-11.6.1.9-py3-none-manylinux2014_aarch64.whl", hash = "sha256:d338f155f174f90724bbde3758b7ac375a70ce8e706d70b018dd3375545fc84e"}, + {file = "nvidia_cusolver_cu12-11.6.1.9-py3-none-manylinux2014_x86_64.whl", hash = "sha256:19e33fa442bcfd085b3086c4ebf7e8debc07cfe01e11513cc6d332fd918ac260"}, + {file = "nvidia_cusolver_cu12-11.6.1.9-py3-none-win_amd64.whl", hash = "sha256:e77314c9d7b694fcebc84f58989f3aa4fb4cb442f12ca1a9bde50f5e8f6d1b9c"}, +] + +[package.dependencies] +nvidia-cublas-cu12 = "*" +nvidia-cusparse-cu12 = "*" +nvidia-nvjitlink-cu12 = "*" + +[[package]] +name = "nvidia-cusparse-cu12" +version = "12.3.1.170" +description = "CUSPARSE native runtime libraries" +optional = true +python-versions = ">=3" +files = [ + {file = "nvidia_cusparse_cu12-12.3.1.170-py3-none-manylinux2014_aarch64.whl", hash = "sha256:9d32f62896231ebe0480efd8a7f702e143c98cfaa0e8a76df3386c1ba2b54df3"}, + {file = "nvidia_cusparse_cu12-12.3.1.170-py3-none-manylinux2014_x86_64.whl", hash = "sha256:ea4f11a2904e2a8dc4b1833cc1b5181cde564edd0d5cd33e3c168eff2d1863f1"}, + {file = "nvidia_cusparse_cu12-12.3.1.170-py3-none-win_amd64.whl", hash = "sha256:9bc90fb087bc7b4c15641521f31c0371e9a612fc2ba12c338d3ae032e6b6797f"}, +] + +[package.dependencies] +nvidia-nvjitlink-cu12 = "*" + +[[package]] +name = "nvidia-nccl-cu12" +version = "2.21.5" +description = "NVIDIA Collective Communication Library (NCCL) Runtime" +optional = true +python-versions = ">=3" +files = [ + {file = "nvidia_nccl_cu12-2.21.5-py3-none-manylinux2014_x86_64.whl", hash = "sha256:8579076d30a8c24988834445f8d633c697d42397e92ffc3f63fa26766d25e0a0"}, +] + +[[package]] +name = "nvidia-nvjitlink-cu12" +version = "12.4.127" +description = "Nvidia JIT LTO Library" +optional = true +python-versions = ">=3" +files = [ + {file = "nvidia_nvjitlink_cu12-12.4.127-py3-none-manylinux2014_aarch64.whl", hash = "sha256:4abe7fef64914ccfa909bc2ba39739670ecc9e820c83ccc7a6ed414122599b83"}, + {file = "nvidia_nvjitlink_cu12-12.4.127-py3-none-manylinux2014_x86_64.whl", hash = "sha256:06b3b9b25bf3f8af351d664978ca26a16d2c5127dbd53c0497e28d1fb9611d57"}, + {file = "nvidia_nvjitlink_cu12-12.4.127-py3-none-win_amd64.whl", hash = "sha256:fd9020c501d27d135f983c6d3e244b197a7ccad769e34df53a42e276b0e25fa1"}, +] + +[[package]] +name = "nvidia-nvtx-cu12" +version = "12.4.127" +description = "NVIDIA Tools Extension" +optional = true +python-versions = ">=3" +files = [ + {file = "nvidia_nvtx_cu12-12.4.127-py3-none-manylinux2014_aarch64.whl", hash = "sha256:7959ad635db13edf4fc65c06a6e9f9e55fc2f92596db928d169c0bb031e88ef3"}, + {file = "nvidia_nvtx_cu12-12.4.127-py3-none-manylinux2014_x86_64.whl", hash = "sha256:781e950d9b9f60d8241ccea575b32f5105a5baf4c2351cab5256a24869f12a1a"}, + {file = "nvidia_nvtx_cu12-12.4.127-py3-none-win_amd64.whl", hash = "sha256:641dccaaa1139f3ffb0d3164b4b84f9d253397e38246a4f2f36728b48566d485"}, +] + [[package]] name = "opt-einsum" version = "3.3.0" @@ -4711,6 +4885,11 @@ files = [ {file = "scikit_learn-1.5.2-cp312-cp312-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:f60021ec1574e56632be2a36b946f8143bf4e5e6af4a06d85281adc22938e0dd"}, {file = "scikit_learn-1.5.2-cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:394397841449853c2290a32050382edaec3da89e35b3e03d6cc966aebc6a8ae6"}, {file = "scikit_learn-1.5.2-cp312-cp312-win_amd64.whl", hash = "sha256:57cc1786cfd6bd118220a92ede80270132aa353647684efa385a74244a41e3b1"}, + {file = "scikit_learn-1.5.2-cp313-cp313-macosx_10_13_x86_64.whl", hash = "sha256:e9a702e2de732bbb20d3bad29ebd77fc05a6b427dc49964300340e4c9328b3f5"}, + {file = "scikit_learn-1.5.2-cp313-cp313-macosx_12_0_arm64.whl", hash = "sha256:b0768ad641981f5d3a198430a1d31c3e044ed2e8a6f22166b4d546a5116d7908"}, + {file = "scikit_learn-1.5.2-cp313-cp313-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:178ddd0a5cb0044464fc1bfc4cca5b1833bfc7bb022d70b05db8530da4bb3dd3"}, + {file = "scikit_learn-1.5.2-cp313-cp313-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:f7284ade780084d94505632241bf78c44ab3b6f1e8ccab3d2af58e0e950f9c12"}, + {file = "scikit_learn-1.5.2-cp313-cp313-win_amd64.whl", hash = "sha256:b7b0f9a0b1040830d38c39b91b3a44e1b643f4b36e36567b80b7c6bd2202a27f"}, {file = "scikit_learn-1.5.2-cp39-cp39-macosx_10_9_x86_64.whl", hash = "sha256:757c7d514ddb00ae249832fe87100d9c73c6ea91423802872d9e74970a0e40b9"}, {file = "scikit_learn-1.5.2-cp39-cp39-macosx_12_0_arm64.whl", hash = "sha256:52788f48b5d8bca5c0736c175fa6bdaab2ef00a8f536cda698db61bd89c551c1"}, {file = "scikit_learn-1.5.2-cp39-cp39-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:643964678f4b5fbdc95cbf8aec638acc7aa70f5f79ee2cdad1eec3df4ba6ead8"}, @@ -5245,6 +5424,23 @@ pure-eval = "*" [package.extras] tests = ["cython", "littleutils", "pygments", "pytest", "typeguard"] +[[package]] +name = "sympy" +version = "1.13.1" +description = "Computer algebra system (CAS) in Python" +optional = true +python-versions = ">=3.8" +files = [ + {file = "sympy-1.13.1-py3-none-any.whl", hash = "sha256:db36cdc64bf61b9b24578b6f7bab1ecdd2452cf008f34faa33776680c26d66f8"}, + {file = "sympy-1.13.1.tar.gz", hash = "sha256:9cebf7e04ff162015ce31c9c6c9144daa34a93bd082f54fd8f12deca4f47515f"}, +] + +[package.dependencies] +mpmath = ">=1.1.0,<1.4" + +[package.extras] +dev = ["hypothesis (>=6.70.0)", "pytest (>=7.1.0)"] + [[package]] name = "synced-collections" version = "1.0.0" @@ -5399,6 +5595,94 @@ files = [ {file = "toolz-0.12.1.tar.gz", hash = "sha256:ecca342664893f177a13dac0e6b41cbd8ac25a358e5f215316d43e2100224f4d"}, ] +[[package]] +name = "torch" +version = "2.5.1" +description = "Tensors and Dynamic neural networks in Python with strong GPU acceleration" +optional = true +python-versions = ">=3.8.0" +files = [ + {file = "torch-2.5.1-cp310-cp310-manylinux1_x86_64.whl", hash = "sha256:71328e1bbe39d213b8721678f9dcac30dfc452a46d586f1d514a6aa0a99d4744"}, + {file = "torch-2.5.1-cp310-cp310-manylinux2014_aarch64.whl", hash = "sha256:34bfa1a852e5714cbfa17f27c49d8ce35e1b7af5608c4bc6e81392c352dbc601"}, + {file = "torch-2.5.1-cp310-cp310-win_amd64.whl", hash = "sha256:32a037bd98a241df6c93e4c789b683335da76a2ac142c0973675b715102dc5fa"}, + {file = "torch-2.5.1-cp310-none-macosx_11_0_arm64.whl", hash = "sha256:23d062bf70776a3d04dbe74db950db2a5245e1ba4f27208a87f0d743b0d06e86"}, + {file = "torch-2.5.1-cp311-cp311-manylinux1_x86_64.whl", hash = "sha256:de5b7d6740c4b636ef4db92be922f0edc425b65ed78c5076c43c42d362a45457"}, + {file = "torch-2.5.1-cp311-cp311-manylinux2014_aarch64.whl", hash = "sha256:340ce0432cad0d37f5a31be666896e16788f1adf8ad7be481196b503dad675b9"}, + {file = "torch-2.5.1-cp311-cp311-win_amd64.whl", hash = "sha256:603c52d2fe06433c18b747d25f5c333f9c1d58615620578c326d66f258686f9a"}, + {file = "torch-2.5.1-cp311-none-macosx_11_0_arm64.whl", hash = "sha256:31f8c39660962f9ae4eeec995e3049b5492eb7360dd4f07377658ef4d728fa4c"}, + {file = "torch-2.5.1-cp312-cp312-manylinux1_x86_64.whl", hash = "sha256:ed231a4b3a5952177fafb661213d690a72caaad97d5824dd4fc17ab9e15cec03"}, + {file = "torch-2.5.1-cp312-cp312-manylinux2014_aarch64.whl", hash = "sha256:3f4b7f10a247e0dcd7ea97dc2d3bfbfc90302ed36d7f3952b0008d0df264e697"}, + {file = "torch-2.5.1-cp312-cp312-win_amd64.whl", hash = "sha256:73e58e78f7d220917c5dbfad1a40e09df9929d3b95d25e57d9f8558f84c9a11c"}, + {file = "torch-2.5.1-cp312-none-macosx_11_0_arm64.whl", hash = "sha256:8c712df61101964eb11910a846514011f0b6f5920c55dbf567bff8a34163d5b1"}, + {file = "torch-2.5.1-cp313-cp313-manylinux1_x86_64.whl", hash = "sha256:9b61edf3b4f6e3b0e0adda8b3960266b9009d02b37555971f4d1c8f7a05afed7"}, + {file = "torch-2.5.1-cp39-cp39-manylinux1_x86_64.whl", hash = "sha256:1f3b7fb3cf7ab97fae52161423f81be8c6b8afac8d9760823fd623994581e1a3"}, + {file = "torch-2.5.1-cp39-cp39-manylinux2014_aarch64.whl", hash = "sha256:7974e3dce28b5a21fb554b73e1bc9072c25dde873fa00d54280861e7a009d7dc"}, + {file = "torch-2.5.1-cp39-cp39-win_amd64.whl", hash = "sha256:46c817d3ea33696ad3b9df5e774dba2257e9a4cd3c4a3afbf92f6bb13ac5ce2d"}, + {file = "torch-2.5.1-cp39-none-macosx_11_0_arm64.whl", hash = "sha256:8046768b7f6d35b85d101b4b38cba8aa2f3cd51952bc4c06a49580f2ce682291"}, +] + +[package.dependencies] +filelock = "*" +fsspec = "*" +jinja2 = "*" +networkx = "*" +nvidia-cublas-cu12 = {version = "12.4.5.8", markers = "platform_system == \"Linux\" and platform_machine == \"x86_64\""} +nvidia-cuda-cupti-cu12 = {version = "12.4.127", markers = "platform_system == \"Linux\" and platform_machine == \"x86_64\""} +nvidia-cuda-nvrtc-cu12 = {version = "12.4.127", markers = "platform_system == \"Linux\" and platform_machine == \"x86_64\""} +nvidia-cuda-runtime-cu12 = {version = "12.4.127", markers = "platform_system == \"Linux\" and platform_machine == \"x86_64\""} +nvidia-cudnn-cu12 = {version = "9.1.0.70", markers = "platform_system == \"Linux\" and platform_machine == \"x86_64\""} +nvidia-cufft-cu12 = {version = "11.2.1.3", markers = "platform_system == \"Linux\" and platform_machine == \"x86_64\""} +nvidia-curand-cu12 = {version = "10.3.5.147", markers = "platform_system == \"Linux\" and platform_machine == \"x86_64\""} +nvidia-cusolver-cu12 = {version = "11.6.1.9", markers = "platform_system == \"Linux\" and platform_machine == \"x86_64\""} +nvidia-cusparse-cu12 = {version = "12.3.1.170", markers = "platform_system == \"Linux\" and platform_machine == \"x86_64\""} +nvidia-nccl-cu12 = {version = "2.21.5", markers = "platform_system == \"Linux\" and platform_machine == \"x86_64\""} +nvidia-nvjitlink-cu12 = {version = "12.4.127", markers = "platform_system == \"Linux\" and platform_machine == \"x86_64\""} +nvidia-nvtx-cu12 = {version = "12.4.127", markers = "platform_system == \"Linux\" and platform_machine == \"x86_64\""} +setuptools = {version = "*", markers = "python_version >= \"3.12\""} +sympy = {version = "1.13.1", markers = "python_version >= \"3.9\""} +triton = {version = "3.1.0", markers = "platform_system == \"Linux\" and platform_machine == \"x86_64\" and python_version < \"3.13\""} +typing-extensions = ">=4.8.0" + +[package.extras] +opt-einsum = ["opt-einsum (>=3.3)"] +optree = ["optree (>=0.12.0)"] + +[[package]] +name = "torch" +version = "2.5.1+cpu" +description = "Tensors and Dynamic neural networks in Python with strong GPU acceleration" +optional = true +python-versions = ">=3.8.0" +files = [ + {file = "torch-2.5.1+cpu-cp310-cp310-linux_x86_64.whl", hash = "sha256:7f91a2200e352745d70e22396bd501448e28350fbdbd8d8b1c83037e25451150"}, + {file = "torch-2.5.1+cpu-cp310-cp310-win_amd64.whl", hash = "sha256:df93157482b672892d29134d3fae9d38ba3219702faedd79f407eb36774c56ce"}, + {file = "torch-2.5.1+cpu-cp311-cp311-linux_x86_64.whl", hash = "sha256:07d7c9e069123d5af08b0cf0013d74f680b2d8be7d9e2cf561a52c90c55d9409"}, + {file = "torch-2.5.1+cpu-cp311-cp311-win_amd64.whl", hash = "sha256:81531d4d5ca74163dc9574b87396531e546a60cceb6253303c7db6a21e867fdf"}, + {file = "torch-2.5.1+cpu-cp312-cp312-linux_x86_64.whl", hash = "sha256:4856f9d6925121d13c2df07aa7580b767f449dfe71ae5acde9c27535d5da4840"}, + {file = "torch-2.5.1+cpu-cp312-cp312-win_amd64.whl", hash = "sha256:a6b720410350765d3d77c01a5ce098a6c45af446284e45e87a98b8a16e7d564d"}, + {file = "torch-2.5.1+cpu-cp313-cp313-linux_x86_64.whl", hash = "sha256:5dbbdf83caa90d0bcaa50e4933ca424889133b35226db79000877d4ec5d9ea37"}, + {file = "torch-2.5.1+cpu-cp39-cp39-linux_x86_64.whl", hash = "sha256:a3ad26468abc5ee601aba49ff02f72387ae734b0900aa589b890c80d72b7b26b"}, + {file = "torch-2.5.1+cpu-cp39-cp39-win_amd64.whl", hash = "sha256:2ebd0b6135dc60b96ce51349c92c9757b2b9634a6b90045dfab3eb4921a4d62f"}, +] + +[package.dependencies] +filelock = "*" +fsspec = "*" +jinja2 = "*" +networkx = "*" +setuptools = {version = "*", markers = "python_version >= \"3.12\""} +sympy = {version = "1.13.1", markers = "python_version >= \"3.9\""} +typing-extensions = ">=4.8.0" + +[package.extras] +opt-einsum = ["opt-einsum (>=3.3)"] +optree = ["optree (>=0.12.0)"] + +[package.source] +type = "legacy" +url = "https://download.pytorch.org/whl/cpu" +reference = "torch-cpu" + [[package]] name = "tornado" version = "6.4.1" @@ -5502,6 +5786,28 @@ easy = ["chardet", "colorlog", "embreex", "httpx", "jsonschema", "lxml", "manifo recommend = ["cascadio", "fast-simplification", "glooey", "meshio", "openctm", "psutil", "pyglet (<2)", "scikit-image", "sympy"] test = ["coveralls", "ezdxf", "matplotlib", "pyinstrument", "pymeshlab", "pyright", "pytest", "pytest-beartype", "pytest-cov", "ruff"] +[[package]] +name = "triton" +version = "3.1.0" +description = "A language and compiler for custom Deep Learning operations" +optional = true +python-versions = "*" +files = [ + {file = "triton-3.1.0-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:6b0dd10a925263abbe9fa37dcde67a5e9b2383fc269fdf59f5657cac38c5d1d8"}, + {file = "triton-3.1.0-cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:0f34f6e7885d1bf0eaaf7ba875a5f0ce6f3c13ba98f9503651c1e6dc6757ed5c"}, + {file = "triton-3.1.0-cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:c8182f42fd8080a7d39d666814fa36c5e30cc00ea7eeeb1a2983dbb4c99a0fdc"}, + {file = "triton-3.1.0-cp38-cp38-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:6dadaca7fc24de34e180271b5cf864c16755702e9f63a16f62df714a8099126a"}, + {file = "triton-3.1.0-cp39-cp39-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:aafa9a20cd0d9fee523cd4504aa7131807a864cd77dcf6efe7e981f18b8c6c11"}, +] + +[package.dependencies] +filelock = "*" + +[package.extras] +build = ["cmake (>=3.20)", "lit"] +tests = ["autopep8", "flake8", "isort", "llnl-hatchet", "numpy", "pytest", "scipy (>=1.7.1)"] +tutorials = ["matplotlib", "pandas", "tabulate"] + [[package]] name = "typeguard" version = "2.13.3" @@ -5780,11 +6086,12 @@ test = ["big-O", "importlib-resources", "jaraco.functools", "jaraco.itertools", type = ["pytest-mypy"] [extras] -dev = ["bayesian-optimization", "bump-my-version", "cma", "coverage", "devsim", "dill", "gdspy", "gdstk", "grcwa", "ipython", "ipython", "jax", "jaxlib", "jinja2", "jupyter", "memory_profiler", "myst-parser", "nbconvert", "nbdime", "nbsphinx", "networkx", "optax", "pre-commit", "pydata-sphinx-theme", "pygad", "pylint", "pyswarms", "pytest", "pytest-timeout", "rtree", "ruff", "sax", "scikit-rf", "signac", "sphinx", "sphinx-book-theme", "sphinx-copybutton", "sphinx-favicon", "sphinx-notfound-page", "sphinx-sitemap", "sphinx-tabs", "sphinxemoji", "tmm", "tox", "trimesh", "vtk"] +dev = ["bayesian-optimization", "bump-my-version", "cma", "coverage", "devsim", "dill", "gdspy", "gdstk", "grcwa", "ipython", "ipython", "jax", "jaxlib", "jinja2", "jupyter", "memory_profiler", "myst-parser", "nbconvert", "nbdime", "nbsphinx", "networkx", "optax", "pre-commit", "pydata-sphinx-theme", "pygad", "pylint", "pyswarms", "pytest", "pytest-timeout", "rtree", "ruff", "sax", "scikit-rf", "signac", "sphinx", "sphinx-book-theme", "sphinx-copybutton", "sphinx-favicon", "sphinx-notfound-page", "sphinx-sitemap", "sphinx-tabs", "sphinxemoji", "tmm", "torch", "torch", "tox", "trimesh", "vtk"] docs = ["cma", "devsim", "gdstk", "grcwa", "ipython", "jinja2", "jupyter", "myst-parser", "nbconvert", "nbdime", "nbsphinx", "optax", "pydata-sphinx-theme", "pylint", "sax", "signac", "sphinx", "sphinx-book-theme", "sphinx-copybutton", "sphinx-favicon", "sphinx-notfound-page", "sphinx-sitemap", "sphinx-tabs", "sphinxemoji", "tmm"] gdspy = ["gdspy"] gdstk = ["gdstk"] jax = ["jax", "jaxlib"] +ruff = ["ruff"] scikit-rf = ["scikit-rf"] trimesh = ["networkx", "rtree", "trimesh"] vtk = ["vtk"] @@ -5792,4 +6099,4 @@ vtk = ["vtk"] [metadata] lock-version = "2.0" python-versions = ">=3.9,<4.0.0" -content-hash = "b4f66c2dea65df1ff21e0373cc0b7f2a7385a368384b972412ec064bf84db3b5" +content-hash = "ce0b77b2f40357c755bb7e6df7d404d0d30a93c6a64cce08e04ba7ee7646861a" diff --git a/pyproject.toml b/pyproject.toml index e8d99efd3..548c3158a 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -23,20 +23,20 @@ documentation = "https://docs.flexcompute.com/projects/tidy3d/en/latest/" [tool.poetry.dependencies] python = ">=3.9,<4.0.0" pyroots = ">=0.5.0" -xarray = ">=0.16.2" +xarray = ">=2023.08.0" importlib-metadata = ">=6.0.0" h5netcdf = "1.0.2" h5py = "^3.0.0" rich = "^13.0" -numpy = "<2" +numpy = "*" matplotlib = "*" shapely = "^2.0" -pandas = "<=2.2.1" +pandas = "*" pydantic = "^2.0" PyYAML = "*" dask = "*" toml = "*" -autograd = "1.6.2" +autograd = ">=1.7.0" scipy = "*" ### NOT CORE boto3 = "^1.28.0" @@ -57,7 +57,7 @@ ipython = { version = "*", optional = true } memory_profiler = { version = "*", optional = true } pre-commit = { version = "*", optional = true } pylint = { version = "*", optional = true } -pytest = { version = "*", optional = true } +pytest = { version = ">=8.1", optional = true } pytest-timeout = { version = "*", optional = true } tox = { version = "*", optional = true } @@ -68,8 +68,8 @@ gdspy = { version = "*", optional = true } gdstk = { version = ">=0.9.49", optional = true } # jax -jaxlib = { version = "0.4.25", source = "jaxsource", optional = true } -jax = { version = "0.4.25", extras = [ +jaxlib = { version = ">=0.4.25", source = "jaxsource", optional = true } +jax = { version = ">=0.4.25", extras = [ "cpu", ], source = "jaxsource", optional = true } @@ -78,6 +78,12 @@ bayesian-optimization = { version = "*", optional = true } pygad = { version = "3.3.1", optional = true } pyswarms = { version = "*", optional = true } +# pytorch +torch = [ + { version = "^2.1.0", source = "PyPI", platform = "darwin", optional = true }, + { version = "^2.1.0", source = "torch-cpu", platform = "!=darwin", optional = true }, +] + # scikit-rf scikit-rf = { version = "*", optional = true } @@ -127,6 +133,7 @@ dev = [ 'ipython', 'jax', 'jaxlib', + 'torch', 'jinja2', 'jupyter', 'myst-parser', @@ -197,6 +204,7 @@ jax = ["jaxlib", "jax"] scikit-rf = ["scikit-rf"] trimesh = ["trimesh", "networkx", "rtree"] vtk = ["vtk"] +ruff = ["ruff"] [tool.poetry.scripts] tidy3d = "tidy3d.web.cli:tidy3d_cli" @@ -210,6 +218,11 @@ name = "jaxsource" url = "https://storage.googleapis.com/jax-releases/jax_releases.html" priority = "supplemental" +[[tool.poetry.source]] +name = "torch-cpu" +url = "https://download.pytorch.org/whl/cpu" +priority = "explicit" + [build-system] requires = ["poetry-core>=1.0.0"] build-backend = "poetry.core.masonry.api" diff --git a/tests/_test_notebooks/full_test_notebooks.py b/tests/_test_notebooks/full_test_notebooks.py index da37b6ffb..077bb70ab 100644 --- a/tests/_test_notebooks/full_test_notebooks.py +++ b/tests/_test_notebooks/full_test_notebooks.py @@ -30,6 +30,9 @@ run_only = [] skip = [ + # WIP + "Autograd10YBranchLevelSet", + "Autograd13Metasurface", # long time (excluding most adjoint) "8ChannelDemultiplexer", "90BendPolarizationSplitterRotator", @@ -53,7 +56,7 @@ notebook_filenames.append(fname) """ -as of June 08 2024 +as of Sept 04 2024 '8ChannelDemultiplexer', '90BendPolarizationSplitterRotator', '90OpticalHybrid', @@ -78,15 +81,21 @@ 'AnimationTutorial', 'AntiResonantHollowCoreFiber', 'AutoGrid', -'Autograd10YBranchLevelSet', +'Autograd0Quickstart', +'Autograd15Antenna', +'Autograd16BilayerCoupler', +'Autograd17BandPassFilter', 'Autograd1Intro', 'Autograd2GradientChecking', 'Autograd3InverseDesign', 'Autograd4MultiObjective', 'Autograd5BoundaryGradients', 'Autograd6GratingCoupler', +'Autograd7Metalens', 'Autograd8WaveguideBend', +'Autograd9WDM', 'Bandstructure', +'BatchModeSolver', 'BilayerSiNEdgeCoupler', 'BilevelPSR', 'BiosensorGrating', @@ -98,6 +107,7 @@ 'CMOSRGBSensor', 'CavityFOM', 'CharacteristicImpedanceCalculator', +'CircularlyPolarizedPatchAntenna', 'CoupledLineBandpassFilter', 'CreatingGeometryUsingTrimesh', 'CustomFieldSource', @@ -129,11 +139,14 @@ 'HeatSolver', 'HighQGe', 'HighQSi', +'IntegratedVivaldiAntenna', 'InverseDesign', +'LNOIPolarizationSplitterRotator', 'MIMResonator', 'MMI1x4', 'MachZehnderModulator', 'MetalHeaterPhaseShifter', +'MetalOxideSunscreen', 'Metalens', 'MicrowaveFrequencySelectiveSurface', 'MidIRMetalens', @@ -141,6 +154,8 @@ 'ModalSourcesMonitors', 'ModeSolver', 'ModesBentAngled', +'MultiplexingMMI', +'NanobeamCavity', 'NanostructuredBoronNitride', 'Near2FarSphereRCS', 'NonHermitianMetagratings', @@ -173,6 +188,7 @@ 'Symmetry', 'TFSF', 'THzDemultiplexerFilter', +'TaperedWgDispersion', 'ThermallyTunedRingResonator', 'ThermoOpticDopedModulator', 'TimeModulationTutorial', @@ -180,6 +196,7 @@ 'UnstructuredData', 'VizData', 'VizSimulation', +'VortexMetasurface', 'WaveguideBendSimulator', 'WaveguideCrossing', 'WaveguideGratingAntenna', diff --git a/tests/conftest.py b/tests/conftest.py index 26c2c9d9c..36da32800 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -1,3 +1,4 @@ +import numpy as np import pytest import tidy3d as td @@ -21,3 +22,9 @@ def log_capture(monkeypatch): log_capture = CaptureHandler() monkeypatch.setitem(td.log.handlers, "pytest_capture", log_capture) return log_capture.records + + +@pytest.fixture +def rng(): + seed = 36523525 + return np.random.default_rng(seed) diff --git a/tests/ruff.toml b/tests/ruff.toml index 56f5406cc..816bcd68e 100644 --- a/tests/ruff.toml +++ b/tests/ruff.toml @@ -10,4 +10,5 @@ extend-ignore = [ "E731", # lambda assignment "F841", # unused local variable "S101", # asserts allowed in tests + "NPY201", # numpy 2.* compatibility check ] diff --git a/tests/test_components/test_autograd.py b/tests/test_components/test_autograd.py index c3e2761e9..403f9b788 100644 --- a/tests/test_components/test_autograd.py +++ b/tests/test_components/test_autograd.py @@ -11,17 +11,20 @@ import autograd.numpy as anp import matplotlib.pylab as plt import numpy as np +import numpy.testing as npt import pytest import tidy3d as td import tidy3d.web as web import xarray as xr +from autograd.test_util import check_grads from tidy3d.components.autograd.derivative_utils import DerivativeInfo -from tidy3d.components.data.sim_data import AdjointSourceInfo +from tidy3d.components.autograd.utils import is_tidy_box +from tidy3d.components.data.data_array import DataArray from tidy3d.plugins.polyslab import ComplexPolySlab from tidy3d.web import run, run_async from tidy3d.web.api.autograd.utils import FieldMap -from ..utils import SIM_FULL, run_emulated +from ..utils import SIM_FULL, AssertLogLevel, run_emulated, tracer_arr """ Test configuration """ @@ -36,8 +39,11 @@ TEST_CUSTOM_MEDIUM_SPEED = False TEST_POLYSLAB_SPEED = False +# whether to run numerical gradient tests, off by default because it runs real simulations +RUN_NUMERICAL = False +_NUMERICAL_COMBINATION = ("size_element", "mode") -TEST_MODES = ("pipeline", "adjoint", "numerical", "speed") +TEST_MODES = ("pipeline", "adjoint", "speed") TEST_MODE = "speed" if TEST_POLYSLAB_SPEED else "pipeline" # number of elements in the parameters / input to the objective function @@ -59,16 +65,16 @@ WVL = 1.0 FREQ0 = td.C_0 / WVL FREQS = [FREQ0] +FWIDTH = FREQ0 / 10 # sim sizes -LZ = 7 * WVL +LZ = 7.0 * WVL -# NOTE: regular stuff is broken in 2D need to change volume and face integration to handle this -IS_3D = True +IS_3D = False # TODO: test 2D and 3D parameterized -LX = 4 * WVL if IS_3D else 0.0 +LX = 0.5 * WVL if IS_3D else 0.0 PML_X = True if IS_3D else False @@ -77,23 +83,34 @@ DA_SHAPE = (DA_SHAPE_X, 1_000, 1_000) if TEST_CUSTOM_MEDIUM_SPEED else (DA_SHAPE_X, 12, 12) # number of vertices in the polyslab -NUM_VERTICES = 100_000 if TEST_POLYSLAB_SPEED else 12 +NUM_VERTICES = 100_000 if TEST_POLYSLAB_SPEED else 110 + +PNT_DIPOLE = td.PointDipole( + center=(0, 0, -LZ / 2 + WVL), + polarization="Ey", + source_time=td.GaussianPulse( + freq0=FREQ0, + fwidth=FWIDTH, + amplitude=1.0, + ), +) + +PLANE_WAVE = td.PlaneWave( + center=(0, 0, -LZ / 2 + WVL), + size=(td.inf, td.inf, 0), + direction="+", + source_time=td.GaussianPulse( + freq0=FREQ0, + fwidth=FWIDTH, + amplitude=1.0, + ), +) # sim that we add traced structures and monitors to SIM_BASE = td.Simulation( - size=(LX, 3, LZ), - run_time=1e-12, - sources=[ - td.PointDipole( - center=(0, 0, -LZ / 2 + WVL), - polarization="Ey", - source_time=td.GaussianPulse( - freq0=FREQ0, - fwidth=FREQ0 / 10.0, - amplitude=1.0, - ), - ) - ], + size=(LX, 3.15, LZ), + run_time=200 / FWIDTH, + sources=[PLANE_WAVE], structures=[ td.Structure( geometry=td.Box( @@ -111,7 +128,8 @@ name="extraneous", ) ], - boundary_spec=td.BoundarySpec.pml(x=False, y=False, z=True), + boundary_spec=td.BoundarySpec.pml(x=False, y=True, z=True), + grid_spec=td.GridSpec.uniform(dl=0.01 * td.C_0 / FREQ0), ) # variable to store whether the emulated run as used @@ -183,15 +201,12 @@ def emulated_run_bwd(simulation, task_name, **run_kwargs) -> td.SimulationData: # get the original traced fields sim_fields_keys = cache[task_id_fwd][AUX_KEY_SIM_FIELDS_KEYS] - adjoint_source_info = AdjointSourceInfo(sources=[], post_norm=1.0, normalize_sim=True) - # postprocess (compute adjoint gradients) traced_fields_vjp = postprocess_adj( sim_data_adj=sim_data_adj, sim_data_orig=sim_data_orig, sim_data_fwd=sim_data_fwd, sim_fields_keys=sim_fields_keys, - adjoint_source_info=adjoint_source_info, ) return traced_fields_vjp @@ -239,14 +254,17 @@ def make_structures(params: anp.ndarray) -> dict[str, td.Structure]: # static components box = td.Box(center=(0, 0, 0), size=(1, 1, 1)) - med = td.Medium(permittivity=2.0) + med = td.Medium(permittivity=3.0) # Structure with variable .medium eps = 1 + anp.abs(vector @ params) - conductivity = eps / 10.0 + sigma = 0.1 * (anp.tanh(vector @ params) + 1) + + permittivity, conductivity = eps, sigma + medium = td.Structure( geometry=box, - medium=td.Medium(permittivity=eps, conductivity=conductivity), + medium=td.Medium(permittivity=permittivity, conductivity=conductivity), ) # Structure with variable Box.center @@ -264,13 +282,13 @@ def make_structures(params: anp.ndarray) -> dict[str, td.Structure]: size_element = td.Structure( geometry=td.Box(center=(0, 0, 0), size=(1, size_y, 1)), medium=med, - background_permittivity=5.0, + background_medium=td.Medium(permittivity=5.0), ) # custom medium with variable permittivity data len_arr = np.prod(DA_SHAPE) matrix = np.random.random((len_arr, N_PARAMS)) - matrix /= np.linalg.norm(matrix) + # matrix /= np.linalg.norm(matrix) eps_arr = 1.01 + 0.5 * (anp.tanh(matrix @ params).reshape(DA_SHAPE) + 1) @@ -309,9 +327,12 @@ def make_structures(params: anp.ndarray) -> dict[str, td.Structure]: ) # Polyslab with variable radius about origin - matrix = np.random.random((NUM_VERTICES, N_PARAMS)) - params_01 = 0.5 * (anp.tanh(matrix @ params) + 1) - radii = 1.0 + 0.1 * params_01 + # matrix = np.random.random((NUM_VERTICES, N_PARAMS)) - 0.5 + # params_01 = 0.5 * (anp.tanh(matrix @ params / 3) + 1) + matrix = np.random.random((N_PARAMS,)) - 0.5 + params_01 = 0.5 * (anp.tanh(matrix @ params / 3) + 1) + + radii = 1.0 + 0.5 * params_01 phis = 2 * anp.pi * anp.linspace(0, 1, NUM_VERTICES + 1)[:NUM_VERTICES] xs = radii * anp.cos(phis) @@ -321,7 +342,7 @@ def make_structures(params: anp.ndarray) -> dict[str, td.Structure]: geometry=td.PolySlab( vertices=vertices, slab_bounds=(-0.5, 0.5), - axis=1, + axis=0, sidewall_angle=0.01, dilation=0.01, ), @@ -399,6 +420,17 @@ def make_structures(params: anp.ndarray) -> dict[str, td.Structure]: custom_med_pole_res = td.CustomPoleResidue(eps_inf=eps_inf, poles=[(a1, c1), (a2, c2)]) custom_pole_res = td.Structure(geometry=box, medium=custom_med_pole_res) + radius = 0.4 * (1 + anp.abs(vector @ params)) + cyl_center_y = vector @ params + cyl_center_z = -vector @ params + cylinder_geo = td.Cylinder( + radius=anp.mean(radii) * 0.5, + center=(0, cyl_center_y, cyl_center_z), + axis=0, + length=LX / 2 if IS_3D else td.inf, + ) + cylinder = td.Structure(geometry=cylinder_geo, medium=polyslab.medium) + return dict( medium=medium, center_list=center_list, @@ -410,15 +442,18 @@ def make_structures(params: anp.ndarray) -> dict[str, td.Structure]: complex_polyslab=complex_polyslab_geo_group, pole_res=pole_res, custom_pole_res=custom_pole_res, + cylinder=cylinder, ) def make_monitors() -> dict[str, tuple[td.Monitor, typing.Callable[[td.SimulationData], float]]]: """Make a dictionary of all the possible monitors in the simulation.""" + X = 0.75 + mode_mnt = td.ModeMonitor( size=(2, 2, 0), - center=(0, 0, LZ / 2 - WVL), + center=(0, 0, +LZ / 2 - X * WVL), mode_spec=td.ModeSpec(), freqs=[FREQ0], name="mode", @@ -429,18 +464,18 @@ def mode_postprocess_fn(sim_data, mnt_data): diff_mnt = td.DiffractionMonitor( size=(td.inf, td.inf, 0), - center=(0, 0, -LZ / 2 + WVL), + center=(0, 0, +LZ / 2 - 2 * WVL), freqs=[FREQ0], normal_dir="+", name="diff", ) def diff_postprocess_fn(sim_data, mnt_data): - return anp.sum(abs(mnt_data.amps.values) ** 2) + return anp.sum(abs(mnt_data.amps.sel(polarization=["s", "p"]).values) ** 2) field_vol = td.FieldMonitor( size=(1, 1, 0), - center=(0, 0, -LZ / 2 + WVL), + center=(0, 0, +LZ / 2 - X * WVL), freqs=[FREQ0], name="field_vol", ) @@ -448,7 +483,7 @@ def diff_postprocess_fn(sim_data, mnt_data): def field_vol_postprocess_fn(sim_data, mnt_data): value = 0.0 for _, val in mnt_data.field_components.items(): - value += abs(anp.sum(val.values)) + value = value + abs(anp.sum(val.values)) intensity = anp.nan_to_num(anp.sum(sim_data.get_intensity(mnt_data.monitor.name).values)) value += intensity value += anp.sum(mnt_data.flux.values) @@ -456,7 +491,7 @@ def field_vol_postprocess_fn(sim_data, mnt_data): field_point = td.FieldMonitor( size=(0, 0, 0), - center=(0, 0, -LZ / 2 + WVL), + center=(0, 0, LZ / 2 - WVL), freqs=[FREQ0], name="field_point", ) @@ -464,7 +499,7 @@ def field_vol_postprocess_fn(sim_data, mnt_data): def field_point_postprocess_fn(sim_data, mnt_data): value = 0.0 for _, val in mnt_data.field_components.items(): - value += abs(anp.sum(val.values)) + value += abs(anp.sum(abs(val.values))) value += anp.sum(sim_data.get_intensity(mnt_data.monitor.name).values) return value @@ -476,9 +511,11 @@ def field_point_postprocess_fn(sim_data, mnt_data): ) -def plot_sim(sim: td.Simulation, plot_eps: bool = False) -> None: +def plot_sim(sim: td.Simulation, plot_eps: bool = True) -> None: """Plot the simulation.""" + sim = sim.to_static() + plot_fn = sim.plot_eps if plot_eps else sim.plot f, (ax1, ax2, ax3) = plt.subplots(1, 3, tight_layout=True) @@ -500,6 +537,7 @@ def plot_sim(sim: td.Simulation, plot_eps: bool = False) -> None: "geo_group", "pole_res", "custom_pole_res", + "cylinder", ) monitor_keys_ = ("mode", "diff", "field_vol", "field_point") @@ -520,7 +558,7 @@ def plot_sim(sim: td.Simulation, plot_eps: bool = False) -> None: args = [("polyslab", "mode")] -# args = [("custom_med", "mode")] +# args = [("size_element", "mode")] def get_functions(structure_key: str, monitor_key: str) -> typing.Callable: @@ -552,7 +590,12 @@ def make_sim(*args) -> td.Simulation: for structure_key in structure_keys: structures.append(structures_traced_dict[structure_key]) - return SIM_BASE.updated_copy(structures=structures, monitors=monitors) + sim = SIM_BASE + if "diff" in monitor_dict: + sim = sim.updated_copy(boundary_spec=td.BoundarySpec.pml(x=False, y=False, z=True)) + sim = sim.updated_copy(structures=structures, monitors=monitors) + + return sim def postprocess(data: td.SimulationData) -> float: """Postprocess the dataset.""" @@ -562,6 +605,125 @@ def postprocess(data: td.SimulationData) -> float: return dict(sim=make_sim, postprocess=postprocess) +@pytest.mark.parametrize("axis", (0, 1, 2)) +def test_polyslab_axis_ops(axis): + vertices = ((0, 0), (0, 1), (1, 1), (1, 0)) + p = td.PolySlab(vertices=vertices, axis=axis, slab_bounds=(0, 1)) + + ax_coords = np.array([0, 1, 2, 3]) + plane_coords = np.array([[4, 5], [6, 7], [8, 9], [10, 11]]) + coord = p.unpop_axis_vect(ax_coords=ax_coords, plane_coords=plane_coords) + + assert np.all(coord[:, axis] == ax_coords) + + _ax_coords, _plane_coords = p.pop_axis_vect(coord=coord) + + assert np.all(_ax_coords == ax_coords) + assert np.all(_plane_coords == plane_coords) + + vertices_next = np.roll(vertices, axis=0, shift=-1) + edges = vertices_next - vertices + + basis_vecs = p.edge_basis_vectors(edges=edges) + + +@pytest.mark.skipif(not RUN_NUMERICAL, reason="Numerical gradient tests runs through web API.") +@pytest.mark.parametrize("structure_key, monitor_key", (_NUMERICAL_COMBINATION,)) +def test_autograd_numerical(structure_key, monitor_key): + """Test an objective function through tidy3d autograd.""" + + import tidy3d.web as web + + fn_dict = get_functions(structure_key, monitor_key) + make_sim = fn_dict["sim"] + postprocess = fn_dict["postprocess"] + + def objective(*args): + """Objective function.""" + sim = make_sim(*args) + if PLOT_SIM: + plot_sim(sim, plot_eps=True) + data = web.run(sim, task_name="autograd_test_numerical", verbose=False) + value = postprocess(data) + return value + + val, grad = ag.value_and_grad(objective)(params0) + print(val, grad) + assert anp.all(grad != 0.0), "some gradients are 0" + + # numerical gradients + delta = 1e-3 + sims_numerical = {} + + params_num = np.zeros((N_PARAMS, N_PARAMS)) + + def task_name_fn(i: int, sign: int) -> str: + """Task name for a given index into grad num and sign.""" + pm_string = "+" if sign > 0 else "-" + return f"{i}_{pm_string}" + + for i in range(N_PARAMS): + for j, sign in enumerate((-1, 1)): + task_name = task_name_fn(i, sign) + params_i = np.copy(params0) + params_i[i] += sign * delta + params_num[:, j] = params_i.copy() + sim_i = make_sim(params_i) + sims_numerical[task_name] = sim_i + + datas = web.Batch(simulations=sims_numerical).run(path_dir="data") + + grad_num = np.zeros_like(grad) + objectives_num = np.zeros((len(params0), 2)) + for i in range(N_PARAMS): + for j, sign in enumerate((-1, 1)): + task_name = task_name_fn(i, sign) + sim_data_i = datas[task_name] + obj_i = postprocess(sim_data_i) + objectives_num[i, j] = obj_i + grad_num[i] += sign * obj_i / 2 / delta + + print("adjoint: ", grad) + print("numerical: ", grad_num) + + print(objectives_num) + + grad_normalized = grad / np.linalg.norm(grad) + grad_num_normalized = grad_num / np.linalg.norm(grad_num) + + rms_error = np.linalg.norm(grad_normalized - grad_num_normalized) + norm_factor = np.linalg.norm(grad) / np.linalg.norm(grad_num) + + diff_objectives_num = np.mean(abs(np.diff(objectives_num, axis=-1))) + + print(f"rms_error = {rms_error:.4f}") + print(f"|grad| / |grad_num| = {norm_factor:.4f}") + print(f"avg(diff(objectives)) = {diff_objectives_num:.4f}") + + +def test_run_zero_grad(use_emulated_run, log_capture): + """Test warning if no adjoint sim is run (no adjoint sources). + + This checks the case where a simulation is still part of the computational + graph (i.e. the output technically depends on the simulation), + but no adjoint sources are placed because their amplitudes are zero and thus + no adjoint simulation is run. + """ + + # only needs to be checked for one monitor + fn_dict = get_functions(args[0][0], args[0][1]) + make_sim = fn_dict["sim"] + postprocess = fn_dict["postprocess"] + + def objective(*args): + sim = make_sim(*args) + sim_data = run(sim, task_name="adjoint_test", verbose=False) + return 0 * postprocess(sim_data) + + with AssertLogLevel(log_capture, "WARNING", contains_str="no sources"): + grad = ag.grad(objective)(params0) + + @pytest.mark.parametrize("structure_key, monitor_key", args) def test_autograd_objective(use_emulated_run, structure_key, monitor_key): """Test an objective function through tidy3d autograd.""" @@ -594,46 +756,6 @@ def objective(*args): print(val, grad) assert anp.all(grad != 0.0), "some gradients are 0" - # if 'numerical', we do a numerical gradient check - if TEST_MODE == "numerical": - import tidy3d.web as web - - delta = 1e-8 - sims_numerical = {} - - params_num = np.zeros((N_PARAMS, N_PARAMS)) - - def task_name_fn(i: int, sign: int) -> str: - """Task name for a given index into grad num and sign.""" - pm_string = "+" if sign > 0 else "-" - return f"{i}_{pm_string}" - - for i in range(N_PARAMS): - for j, sign in enumerate((-1, 1)): - task_name = task_name_fn(i, sign) - params_i = np.copy(params0) - params_i[i] += sign * delta - params_num[:, j] = params_i.copy() - sim_i = make_sim(params_i) - sims_numerical[task_name] = sim_i - - datas = web.Batch(simulations=sims_numerical).run(path_dir="data") - - grad_num = np.zeros_like(grad) - objectives_num = np.zeros((len(params0), 2)) - for i in range(N_PARAMS): - for j, sign in enumerate((-1, 1)): - task_name = task_name_fn(i, sign) - sim_data_i = datas[task_name] - obj_i = postprocess(sim_data_i) - objectives_num[i, j] = obj_i - grad_num[i] += sign * obj_i / 2 / delta - - print("adjoint: ", grad) - print("numerical: ", grad_num) - - assert np.allclose(grad, grad_num), "gradients dont match" - @pytest.mark.parametrize("structure_key, monitor_key", args) def test_autograd_async(use_emulated_run, structure_key, monitor_key): @@ -646,8 +768,6 @@ def test_autograd_async(use_emulated_run, structure_key, monitor_key): task_names = {"1", "2", "3", "4"} def objective(*args): - """Objective function.""" - sims = {task_name: make_sim(*args) for task_name in task_names} batch_data = run_async(sims, verbose=False) value = 0.0 @@ -660,6 +780,51 @@ def objective(*args): assert anp.all(grad != 0.0), "some gradients are 0" +@pytest.mark.parametrize("structure_key, monitor_key", args) +def test_autograd_async_some_zero_grad(use_emulated_run, log_capture, structure_key, monitor_key): + """Test objective where only some simulations in batch have adjoint sources.""" + + fn_dict = get_functions(structure_key, monitor_key) + make_sim = fn_dict["sim"] + postprocess = fn_dict["postprocess"] + + task_names = {"1", "2", "3", "4"} + + def objective(*args): + sims = {task_name: make_sim(*args) for task_name in task_names} + batch_data = run_async(sims, verbose=False) + values = [] + for _, sim_data in batch_data.items(): + values.append(postprocess(sim_data)) + return min(values) + + # with AssertLogLevel(log_capture, "DEBUG", contains_str="no sources"): + val, grad = ag.value_and_grad(objective)(params0) + + assert anp.all(grad != 0.0), "some gradients are 0" + + +def test_autograd_async_all_zero_grad(use_emulated_run, log_capture): + """Test objective where no simulation in batch has adjoint sources.""" + + fn_dict = get_functions(args[0][0], args[0][1]) + make_sim = fn_dict["sim"] + postprocess = fn_dict["postprocess"] + + task_names = {"1", "2", "3", "4"} + + def objective(*args): + sims = {task_name: make_sim(*args) for task_name in task_names} + batch_data = run_async(sims, verbose=False) + values = [] + for _, sim_data in batch_data.items(): + values.append(postprocess(sim_data)) + return 0 * sum(values) + + with AssertLogLevel(log_capture, "WARNING", contains_str="contains adjoint sources"): + grad = ag.grad(objective)(params0) + + def test_autograd_speed_num_structures(use_emulated_run): """Test an objective function through tidy3d autograd.""" @@ -695,6 +860,72 @@ def objective(*args): print(f"{num_structures_test} structures took {t2:.2e} seconds") +@pytest.mark.parametrize("monitor_key", ("mode",)) +def test_autograd_polyslab_cylinder(use_emulated_run, monitor_key): + """Test an objective function through tidy3d autograd.""" + + t = 1.0 + axis = 0 + + num_pts = 89 + + monitor, postprocess = make_monitors()[monitor_key] + + def make_cylinder(radius, x0, y0): + return td.Cylinder( + center=td.Cylinder.unpop_axis(0.0, (x0, y0), axis=axis), + radius=radius, + length=t, + axis=axis, + ) # .to_polyslab(num_pts) + + def make_polyslab(radius, x0, y0): + phis = anp.linspace(0, 2 * np.pi, num_pts + 1)[:-1] + + xs = radius * anp.cos(phis) + x0 + ys = radius * anp.sin(phis) + y0 + + vertices = anp.stack((xs, ys), axis=-1) + + return td.PolySlab( + vertices=vertices, + axis=axis, + slab_bounds=(-t / 2, t / 2), + ) + + def make_sim(params, geo_maker): + geo = geo_maker(*params) + structure = td.Structure(geometry=geo, medium=td.Medium(permittivity=2)) + + return SIM_BASE.updated_copy(structures=[structure], monitors=[monitor]) + + p0 = [1.0, 0.0, 0.0] + + def objective_polyslab(params): + """Objective function.""" + sim = make_sim(params, geo_maker=make_polyslab) + if PLOT_SIM: + plot_sim(sim, plot_eps=True) + data = run(sim, task_name="autograd_test", verbose=False) + return anp.sum(anp.abs(data[monitor.name].amps)).item() + + val_polyslab, grad_polyslab = ag.value_and_grad(objective_polyslab)(p0) + print(val_polyslab, grad_polyslab) + assert anp.all(grad_polyslab != 0.0), "some gradients are 0" + + def objective_cylinder(params): + """Objective function.""" + sim = make_sim(params, geo_maker=make_cylinder) + if PLOT_SIM: + plot_sim(sim, plot_eps=True) + data = run(sim, task_name="autograd_test", verbose=False) + return anp.sum(anp.abs(data[monitor.name].amps)).item() + + val_cylinder, grad_cylinder = ag.value_and_grad(objective_cylinder)(p0) + print(val_cylinder, grad_cylinder) + assert anp.all(grad_cylinder != 0.0), "some gradients are 0" + + @pytest.mark.parametrize("structure_key, monitor_key", args) def test_autograd_server(use_emulated_run, structure_key, monitor_key): """Test an objective function through tidy3d autograd.""" @@ -756,7 +987,7 @@ def objective(*params): sim_fields = sim_full_traced.strip_traced_fields() # note: there is one traced structure in SIM_FULL already with 6 fields + 1 = 7 - assert len(sim_fields) == 7 + assert len(sim_fields) == 10 sim_traced = sim_full_static.insert_traced_fields(sim_fields) @@ -767,6 +998,21 @@ def objective(*params): ag.grad(objective)(params0) +def test_sim_traced_override_structures(log_capture): + """Make sure that sims with traced override structures are handled properly.""" + + def f(x): + override_structure = td.MeshOverrideStructure( + geometry=td.Box(center=(0, 0, 0), size=(1, 1, x)), + dl=[1, 1, 1], + ) + sim = SIM_FULL.updated_copy(override_structures=[override_structure], path="grid_spec") + return sim.grid_spec.override_structures[0].geometry.size[2] + + with AssertLogLevel(log_capture, "WARNING", contains_str="override structures"): + ag.grad(f)(1.0) + + @pytest.mark.parametrize("structure_key", ("custom_med",)) def test_sim_fields_io(structure_key, tmp_path): """Test that converging and AutogradFieldMap dictionary to a FieldMap object, saving and loading @@ -866,7 +1112,7 @@ def objective(args): data = run(sim, task_name="autograd_test", verbose=False) if objtype == "flux": - return anp.sum(data[monitor.name].flux.values) + return data[monitor.name].flux.item() elif objtype == "intensity": return anp.sum(data.get_intensity(monitor.name).values) @@ -874,6 +1120,123 @@ def objective(args): assert np.any(grads > 0) +@pytest.mark.parametrize("far_field_approx", [True, False]) +@pytest.mark.parametrize("projection_type", ["angular", "cartesian", "kspace"]) +@pytest.mark.parametrize("sim_2d", [True, False]) +class TestFieldProjection: + @staticmethod + def setup(far_field_approx, projection_type, sim_2d): + if sim_2d and not far_field_approx: + pytest.skip("Exact field projection not implemented for 2d simulations") + + r_proj = 50 * WVL + monitor = td.FieldMonitor( + center=(0, SIM_BASE.size[1] / 2 - 0.1, 0), + size=(td.inf, 0, td.inf), + freqs=[FREQ0], + name="near_field", + colocate=False, + ) + + if projection_type == "angular": + theta_proj = np.linspace(np.pi / 10, np.pi - np.pi / 10, 2) + phi_proj = np.linspace(np.pi / 10, np.pi - np.pi / 10, 3) + monitor_far = td.FieldProjectionAngleMonitor( + center=monitor.center, + size=monitor.size, + freqs=monitor.freqs, + phi=tuple(phi_proj), + theta=tuple(theta_proj), + proj_distance=r_proj, + far_field_approx=far_field_approx, + name="far_field", + ) + elif projection_type == "cartesian": + x_proj = np.linspace(-10, 10, 2) + y_proj = np.linspace(-10, 10, 3) + monitor_far = td.FieldProjectionCartesianMonitor( + center=monitor.center, + size=monitor.size, + freqs=monitor.freqs, + x=x_proj, + y=y_proj, + proj_axis=1, + proj_distance=r_proj, + far_field_approx=far_field_approx, + name="far_field", + ) + elif projection_type == "kspace": + ux = np.linspace(-0.7, 0.7, 2) + uy = np.linspace(-0.7, 0.7, 3) + monitor_far = td.FieldProjectionKSpaceMonitor( + center=monitor.center, + size=monitor.size, + freqs=monitor.freqs, + ux=ux, + uy=uy, + proj_axis=1, + proj_distance=r_proj, + far_field_approx=far_field_approx, + name="far_field", + ) + + sim = SIM_BASE.updated_copy(monitors=[monitor]) + + if sim_2d and IS_3D: + sim = sim.updated_copy(size=(0, *sim.size[1:])) + + return sim, monitor_far + + @staticmethod + def objective(sim_data, monitor_far): + projector = td.FieldProjector.from_near_field_monitors( + sim_data=sim_data, + near_monitors=[sim_data.simulation.monitors[0]], + normal_dirs=["+"], + ) + + projected_fields = projector.project_fields(monitor_far) + + return projected_fields.power.sum().item() + + def test_field_projection_grad_prop( + self, use_emulated_run, far_field_approx, projection_type, sim_2d + ): + """Tests whether field projection gradients are propagated through simulation. + x0 <-> structures <-> sim <-> run <-> fields <-> projection <-> objval + Does _not_ test gradient accuracy! + """ + sim_base, monitor_far = self.setup(far_field_approx, projection_type, sim_2d) + + def objective(args): + structures_traced_dict = make_structures(args) + structures = list(SIM_BASE.structures) + for structure_key in structure_keys_: + structures.append(structures_traced_dict[structure_key]) + + sim = sim_base.updated_copy(structures=structures) + sim_data = run(sim, task_name="field_projection_test") + + return self.objective(sim_data, monitor_far) + + grads = ag.grad(objective)(params0) + assert np.linalg.norm(grads) > 0 + + def test_field_projection_grads( + self, use_emulated_run, far_field_approx, projection_type, sim_2d + ): + """Tests projection gradient accuracy w.r.t. fields. + fields <-> projection <-> objval + """ + sim_base, monitor_far = self.setup(far_field_approx, projection_type, sim_2d) + + def objective(x0): + sim_data = run_emulated(sim_base, task_name="field_projection_test", x0=x0) + return self.objective(sim_data, monitor_far) + + check_grads(objective, modes=["rev"], order=1)(1.0) + + def test_autograd_deepcopy(): """make sure deepcopy works as expected in autograd.""" @@ -937,6 +1300,10 @@ def J(eps): paths=field_paths, E_der_map={}, D_der_map={}, + E_fwd={}, + D_fwd={}, + E_adj={}, + D_adj={}, eps_data={}, eps_in=2.0, eps_out=1.0, @@ -1012,6 +1379,10 @@ def J(eps): paths=field_paths, E_der_map={}, D_der_map={}, + E_fwd={}, + D_fwd={}, + E_adj={}, + D_adj={}, eps_data={}, eps_in=2.0, eps_out=1.0, @@ -1322,3 +1693,142 @@ def objective_multi(params, structure_key) -> float: assert not np.any(np.isclose(grad_indi, 0)) assert not np.any(np.isclose(grad_multi, 0)) + + +def test_error_flux(use_emulated_run, log_capture): + """Make sure proper error raised if differentiating w.r.t. FluxData.""" + + def objective(params): + structure_traced = make_structures(params)["medium"] + sim = SIM_BASE.updated_copy( + structures=[structure_traced], + monitors=[td.FluxMonitor(size=(1, 1, 0), center=(0, 0, 0), freqs=[FREQ0], name="flux")], + ) + data = run(sim, task_name="flux_error") + return anp.sum(data["flux"].flux.values) + + with pytest.raises( + NotImplementedError, match="Could not formulate adjoint source for 'FluxMonitor' output" + ): + g = ag.grad(objective)(params0) + + +def test_extraneous_field(use_emulated_run, log_capture): + """Make sure this doesnt fail.""" + + def objective(params): + structure_traced = make_structures(params)["medium"] + sim = SIM_BASE.updated_copy( + structures=[structure_traced], + monitors=[ + SIM_BASE.monitors[0], + td.ModeMonitor( + size=(1, 1, 0), + center=(0, 0, 0), + mode_spec=td.ModeSpec(), + freqs=[FREQ0 * 0.9, FREQ0 * 1.1], + name="mode", + ), + ], + ) + data = run(sim, task_name="extra_field") + amp = data["mode"].amps.sel(direction="+", f=FREQ0 * 0.9, mode_index=0).values + return abs(amp.item()) ** 2 + + g = ag.grad(objective)(params0) + + +def test_background_medium(log_capture): + geo = td.Box(size=(1, 1, 1), center=(0, 0, 0)) + med = td.Medium(permittivity=2.0) + + background_permittivity = 5.0 + background_medium = td.Medium(permittivity=background_permittivity) + + # nothing + s = td.Structure( + geometry=geo, + medium=med, + ) + + # both supplied, consistent + td.Structure( + geometry=geo, + medium=med, + background_permittivity=background_permittivity, + background_medium=background_medium, + ) + + # both supplied, inconsistent + with pytest.raises(ValueError): + td.Structure( + geometry=geo, + medium=med, + background_permittivity=background_permittivity + 1, + background_medium=background_medium, + ) + + # background medium (preferred) + s = td.Structure( + geometry=geo, + medium=med, + background_medium=background_medium, + ) + + # background permittivity (deprecated) + with AssertLogLevel(log_capture, "WARNING", contains_str="deprecated"): + s_warn = td.Structure( + geometry=geo, + medium=med, + background_permittivity=background_permittivity, + ) + + assert s_warn.background_medium is not None + assert s_warn.background_medium.permittivity == background_permittivity + + +class TestTidyArrayBox: + def test_is_tidy_box(self): + da = DataArray(tracer_arr, dims=map(str, range(tracer_arr.ndim))) + assert is_tidy_box(da.data) + + def test_real(self): + npt.assert_allclose(tracer_arr.real._value, tracer_arr._value.real) + + def test_imag(self): + npt.assert_allclose(tracer_arr.imag._value, tracer_arr._value.imag) + + def test_conj(self): + npt.assert_allclose(tracer_arr.conj()._value, tracer_arr._value.conj()) + + def test_item(self): + assert tracer_arr.item() == tracer_arr._value.item() + + +class TestDataArrayGrads: + @pytest.mark.parametrize("attr", ["real", "imag", "conj"]) + def test_custom_methods_grads(self, attr): + """Test grads of TidyArrayBox methods implemented in autograd/boxes.py""" + + def objective(x, attr): + da = DataArray(x) + attr_value = getattr(da, attr) + val = attr_value() if callable(attr_value) else attr_value + return val.item() + + x = np.array([1.0]) + check_grads(objective, modes=["fwd", "rev"], order=2)(x, attr) + + def test_multiply_at_grads(self, rng): + """Test grads of DataArray.multiply_at method""" + + def objective(a, b): + coords = {str(i): np.arange(a.shape[i]) for i in range(a.ndim)} + da = DataArray(a, coords=coords) + da_mult = da.multiply_at(b, "0", [0, 1]) ** 2 + return np.sum(da_mult).item() + + a = rng.uniform(-1, 1, (3, 3)) + b = 1.0 + check_grads(lambda x: objective(x, b), modes=["fwd", "rev"], order=2)(a) + check_grads(lambda x: objective(a, x), modes=["fwd", "rev"], order=2)(b) diff --git a/tests/test_components/test_field_projection.py b/tests/test_components/test_field_projection.py index 23b37309b..5b03b53c4 100644 --- a/tests/test_components/test_field_projection.py +++ b/tests/test_components/test_field_projection.py @@ -3,6 +3,7 @@ import numpy as np import pytest import tidy3d as td +from tidy3d.components.field_projection import FieldProjector from tidy3d.exceptions import DataError MEDIUM = td.Medium(permittivity=3) @@ -386,6 +387,8 @@ def make_2d_proj_monitors(center, size, freqs, plane): Ns = 40 xs = np.linspace(-far_size, far_size, Ns) ys = [0] + kx = np.linspace(-0.7, 0.7, Ns) + ky = [0] projection_axis = 0 elif plane == "yz": thetas = np.linspace(0, np.pi, 1) @@ -394,6 +397,8 @@ def make_2d_proj_monitors(center, size, freqs, plane): Ns = 40 xs = [0] ys = np.linspace(-far_size, far_size, Ns) + kx = [0] + ky = np.linspace(-0.7, 0.7, Ns) projection_axis = 1 elif plane == "xz": thetas = np.linspace(0, np.pi, 100) @@ -402,6 +407,8 @@ def make_2d_proj_monitors(center, size, freqs, plane): Ns = 40 xs = [0] ys = np.linspace(-far_size, far_size, Ns) + kx = [0] + ky = np.linspace(-0.7, 0.7, Ns) projection_axis = 0 else: raise ValueError("Invalid plane. Use 'xy', 'yz', or 'xz'.") @@ -429,7 +436,19 @@ def make_2d_proj_monitors(center, size, freqs, plane): far_field_approx=True, # Fields are far enough for geometric far field approximations ) - return (n2f_angle_monitor_2d, n2f_car_monitor_2d) + n2f_k_monitor_2d = td.FieldProjectionKSpaceMonitor( + center=center, + size=size, + freqs=freqs, + name="far_field_kspace", + ux=list(kx), + uy=list(ky), + proj_axis=projection_axis, + proj_distance=R_FAR, + far_field_approx=True, # Fields are far enough for geometric far field approximations + ) + + return (n2f_angle_monitor_2d, n2f_car_monitor_2d, n2f_k_monitor_2d) def make_2d_proj(plane): @@ -528,10 +547,12 @@ def make_2d_proj(plane): ( n2f_angle_monitor_2d, n2f_cart_monitor_2d, + n2f_kspace_monitor_2d, ) = make_2d_proj_monitors(center, monitor_size, [f0], plane) far_fields_angular_2d = proj.project_fields(n2f_angle_monitor_2d) far_fields_cartesian_2d = proj.project_fields(n2f_cart_monitor_2d) + far_fields_kspace_2d = proj.project_fields(n2f_kspace_monitor_2d) # compute far field quantities far_fields_angular_2d.r @@ -556,6 +577,17 @@ def make_2d_proj(plane): val.sel(f=f0) far_fields_cartesian_2d.renormalize_fields(proj_distance=5e6) + far_fields_kspace_2d.ux + far_fields_kspace_2d.uy + far_fields_kspace_2d.r + far_fields_kspace_2d.fields_spherical + far_fields_kspace_2d.fields_cartesian + far_fields_kspace_2d.radar_cross_section + far_fields_kspace_2d.power + for val in far_fields_kspace_2d.field_components.values(): + val.sel(f=f0) + far_fields_kspace_2d.renormalize_fields(proj_distance=5e6) + def test_2d_proj_clientside(): # Run simulations and tests for all three planes @@ -563,3 +595,25 @@ def test_2d_proj_clientside(): for plane in planes: make_2d_proj(plane) + + +@pytest.mark.parametrize( + "array, pts, axes, expected", + [ + # 1D array, integrate over axis 0 + (np.array([1, 2, 3]), np.array([0, 1, 2]), 0, 4.0), + # 2D array, integrate over axis 0 + (np.array([[1, 2, 3], [4, 5, 6]]), np.array([0, 1]), 0, np.array([2.5, 3.5, 4.5])), + # 2D array, integrate over axis 1 + (np.array([[1, 2], [3, 4], [5, 6]]), np.array([0, 1]), 1, np.array([1.5, 3.5, 5.5])), + # 3D array, integrate over axes 0 and 1 + (np.ones((2, 2, 2)), [np.array([0, 1]), np.array([0, 1])], [0, 1], np.array([1.0, 1.0])), + # one element along integration axis but two points in pts + (np.array([[1, 1], [2, 2], [3, 3]]), np.array([0, 1]), 1, np.array([1.0, 2.0, 3.0])), + # 2D array of shape (1, 3), integrate over both axes + (np.array([[1, 2, 3]]), [np.array([0]), np.array([0, 1, 2])], [0, 1], 4.0), + ], +) +def test_trapezoid(array, pts, axes, expected): + result = FieldProjector.trapezoid(array, pts, axes) + assert np.allclose(result, expected) diff --git a/tests/test_components/test_frequencies.py b/tests/test_components/test_frequencies.py new file mode 100644 index 000000000..30c8f2571 --- /dev/null +++ b/tests/test_components/test_frequencies.py @@ -0,0 +1,17 @@ +import numpy as np +import pytest +import tidy3d as td + + +def test_classification(): + assert td.frequencies.classification(1) == ("near static",) + assert td.wavelengths.classification(td.C_0) == ("near static",) + assert td.frequencies.classification(td.C_0 / 1.55) == ("infrared", "NIR") + assert td.wavelengths.classification(1.55) == ("infrared", "NIR") + + +@pytest.mark.parametrize("band", ["O", "E", "S", "C", "L", "U"]) +def test_bands(band): + freqs = getattr(td.frequencies, band.lower() + "_band")() + ldas = getattr(td.wavelengths, band.lower() + "_band")() + assert np.allclose(freqs, td.C_0 / np.array(ldas)) diff --git a/tests/test_components/test_geometry.py b/tests/test_components/test_geometry.py index 77119d91b..1c5279a4d 100644 --- a/tests/test_components/test_geometry.py +++ b/tests/test_components/test_geometry.py @@ -203,6 +203,17 @@ def test_intersections_plane(component): assert len(component.intersections_plane(x=10000)) == 0 +def test_intersections_plane_inf(): + a = ( + td.Cylinder(radius=3.2, center=(0.45, 9, 0), length=td.inf) + + td.Box(center=(0, 0, 0), size=(0.9, 24, td.inf)) + + td.Box(center=(0, 0, 0), size=(7.3, 18, td.inf)) + ) + b = td.Cylinder(radius=2.9, center=(-0.45, 9, 0), length=td.inf) + c = a - b + assert len(c.intersections_plane(y=0)) == 1 + + def test_bounds_base(): assert all(a == b for a, b in zip(Planar.bounds.fget(POLYSLAB), POLYSLAB.bounds)) @@ -254,6 +265,10 @@ def test_slanted_cylinder_infinite_length_validate(): ) +def test_cylinder_to_polyslab(): + ps = CYLINDER.to_polyslab(num_pts_circumference=10, dilation=0.02) + + def test_box_from_bounds(): b = td.Box.from_bounds(rmin=(-td.inf, 0, 0), rmax=(td.inf, 0, 0)) assert b.center[0] == 0.0 @@ -523,6 +538,33 @@ def test_flattening(): for g in flat ) + t0 = np.array([[2, 0, 0, 0], [3, 2, 0, 0], [1, 0, 2, 0], [0, 0, 0, 1.0]]) + g0 = td.Sphere(radius=1) + t1 = np.array([[2, 0, 5, 0], [0, 1, 0, 0], [-1, 0, 1, 0], [0, 0, 0, 1.0]]) + g1 = td.Box(size=(1, 2, 3)) + flat = list( + flatten_groups( + td.Transformed( + transform=t0, + geometry=td.ClipOperation( + operation="union", + geometry_a=g0, + geometry_b=td.Transformed(transform=t1, geometry=g1), + ), + ), + flatten_transformed=True, + ) + ) + assert len(flat) == 2 + + assert isinstance(flat[0], td.Transformed) + assert flat[0].geometry == g0 + assert np.allclose(flat[0].transform, t0) + + assert isinstance(flat[1], td.Transformed) + assert flat[1].geometry == g1 + assert np.allclose(flat[1].transform, t0 @ t1) + def test_geometry_traversal(): geometries = list(traverse_geometries(td.Box(size=(1, 1, 1)))) @@ -1019,3 +1061,9 @@ def test_snap_box_to_grid(snap_location, snap_behavior): assert math.isclose(new_box.bounds[1][1], xyz[3]) # Check that the box boundary outside the grid was snapped to the smallest grid coordinate assert math.isclose(new_box.bounds[0][2], xyz[0]) + + +def test_triangulation_with_collinear_vertices(): + xr = np.linspace(0, 1, 6) + a = np.array([[x, -0.5] for x in xr] + [[x, 0.5] for x in xr[::-1]]) + assert len(td.components.geometry.triangulation.triangulate(a)) == 10 diff --git a/tests/test_components/test_grid.py b/tests/test_components/test_grid.py index 2ca7e613e..587a09ddb 100644 --- a/tests/test_components/test_grid.py +++ b/tests/test_components/test_grid.py @@ -108,7 +108,7 @@ def test_sim_nonuniform_small(): # tests when the nonuniform grid does not cover the simulation size size_x = 18 - num_layers_pml_x = 2 + num_layers_pml_x = 6 grid_size_x = [2, 1, 3] sim = td.Simulation( center=(1, 0, 0), @@ -153,15 +153,12 @@ def test_sim_nonuniform_small(): for dl in dls: assert dl in grid_size_x - # tests that it gives exactly what we expect - assert np.all(bound_coords == np.array([-12, -10, -8, -6, -4, -2, 0, 1, 4, 7, 10, 13, 16])) - def test_sim_nonuniform_large(): # tests when the nonuniform grid extends beyond the simulation size size_x = 18 - num_layers_pml_x = 2 + num_layers_pml_x = 6 grid_size_x = [2, 3, 4, 1, 2, 1, 3, 1, 2, 3, 4] sim = td.Simulation( center=(1, 0, 0), @@ -230,9 +227,9 @@ def test_sim_symmetry_grid(): size=(11, 11, 11), grid_spec=td.GridSpec(grid_x=grid_1d, grid_y=grid_1d, grid_z=grid_1d), boundary_spec=td.BoundarySpec( - x=td.Boundary.pml(num_layers=2), - y=td.Boundary.pml(num_layers=2), - z=td.Boundary.pml(num_layers=2), + x=td.Boundary.pml(num_layers=6), + y=td.Boundary.pml(num_layers=6), + z=td.Boundary.pml(num_layers=6), ), symmetry=(0, 1, -1), run_time=1e-12, @@ -257,20 +254,20 @@ def test_sim_pml_grid(): size=(4, 4, 4), grid_spec=td.GridSpec.uniform(1.0), boundary_spec=td.BoundarySpec( - x=td.Boundary.pml(num_layers=2), - y=td.Boundary.absorber(num_layers=2), - z=td.Boundary.stable_pml(num_layers=2), + x=td.Boundary.pml(num_layers=6), + y=td.Boundary.absorber(num_layers=6), + z=td.Boundary.stable_pml(num_layers=6), ), run_time=1e-12, ) for dim in "xyz": c = sim.grid.centers.dict()[dim] - assert np.all(c == np.array([-3.5, -2.5, -1.5, -0.5, 0.5, 1.5, 2.5, 3.5])) + assert np.all(c == np.arange(-7.5, 8, 1)) for dim in "xyz": b = sim.grid.boundaries.dict()[dim] - assert np.all(b == np.array([-4, -3, -2, -1, 0, 1, 2, 3, 4])) + assert np.all(b == np.arange(-8, 8.5, 1)) def test_sim_discretize_vol(): diff --git a/tests/test_components/test_grid_spec.py b/tests/test_components/test_grid_spec.py index 5000a40dc..0c611c4d8 100644 --- a/tests/test_components/test_grid_spec.py +++ b/tests/test_components/test_grid_spec.py @@ -210,9 +210,9 @@ def test_autogrid_2dmaterials(): structures=[box], sources=[src], boundary_spec=td.BoundarySpec( - x=td.Boundary.pml(num_layers=5), - y=td.Boundary.pml(num_layers=5), - z=td.Boundary.pml(num_layers=5), + x=td.Boundary.pml(num_layers=6), + y=td.Boundary.pml(num_layers=6), + z=td.Boundary.pml(num_layers=6), ), grid_spec=td.GridSpec.auto(), run_time=1e-12, @@ -231,9 +231,9 @@ def test_autogrid_2dmaterials(): structures=[box2], sources=[src], boundary_spec=td.BoundarySpec( - x=td.Boundary.pml(num_layers=5), - y=td.Boundary.pml(num_layers=5), - z=td.Boundary.pml(num_layers=5), + x=td.Boundary.pml(num_layers=6), + y=td.Boundary.pml(num_layers=6), + z=td.Boundary.pml(num_layers=6), ), grid_spec=td.GridSpec.auto(), run_time=1e-12, @@ -250,9 +250,9 @@ def test_autogrid_2dmaterials(): structures=[box, box2], sources=[src], boundary_spec=td.BoundarySpec( - x=td.Boundary.pml(num_layers=5), - y=td.Boundary.pml(num_layers=5), - z=td.Boundary.pml(num_layers=5), + x=td.Boundary.pml(num_layers=6), + y=td.Boundary.pml(num_layers=6), + z=td.Boundary.pml(num_layers=6), ), grid_spec=td.GridSpec.auto(), run_time=1e-12, diff --git a/tests/test_components/test_medium.py b/tests/test_components/test_medium.py index 6be6c009e..e316bd4e2 100644 --- a/tests/test_components/test_medium.py +++ b/tests/test_components/test_medium.py @@ -537,6 +537,38 @@ def test_fully_anisotropic_media(): assert all(np.isin(np.round(perm_d), np.round(np.diag(perm_diag)))) assert all(np.isin(np.round(cond_d), np.round(np.diag(cond_diag)))) + with pytest.raises(ValidationError): + _ = td.FullyAnisotropicMedium.from_diagonal( + xx=td.Medium( + permittivity=2, + nonlinear_spec=td.NonlinearSpec( + models=[ + td.NonlinearSusceptibility(chi3=2), + td.TwoPhotonAbsorption(beta=1.3), + td.KerrNonlinearity(n2=1.3), + ] + ), + ), + yy=td.Medium(permittivity=4), + zz=td.Medium(permittivity=1), + rotation=td.RotationAroundAxis(axis=2, angle=np.pi / 4), + ) + + with pytest.raises(ValidationError): + _ = td.FullyAnisotropicMedium.from_diagonal( + xx=td.Medium(permittivity=2), + yy=td.Medium( + permittivity=4, + modulation_spec=td.ModulationSpec( + permittivity=td.SpaceTimeModulation( + time_modulation=td.ContinuousWaveTimeModulation(freq0=1e12, amplitude=0.02) + ) + ), + ), + zz=td.Medium(permittivity=1), + rotation=td.RotationAroundAxis(axis=2, angle=np.pi / 4), + ) + def test_nonlinear_medium(log_capture): med = td.Medium( @@ -551,15 +583,15 @@ def test_nonlinear_medium(log_capture): ) # complex parameters - med = td.Medium( - nonlinear_spec=td.NonlinearSpec( - models=[ - td.KerrNonlinearity(n2=-1 + 1j, n0=1), - ], - num_iters=20, + with AssertLogLevel(log_capture, "WARNING", contains_str="preferred"): + med = td.Medium( + nonlinear_spec=td.NonlinearSpec( + models=[ + td.KerrNonlinearity(n2=-1 + 1j, n0=1), + ], + num_iters=20, + ) ) - ) - assert_log_level(log_capture, None) # warn about deprecated api med = td.Medium(nonlinear_spec=td.NonlinearSusceptibility(chi3=1.5)) @@ -610,10 +642,11 @@ def test_nonlinear_medium(log_capture): with pytest.raises(ValidationError): med = td.Medium(nonlinear_spec=td.NonlinearSpec(models=[td.KerrNonlinearity(n2=-1j, n0=1)])) - med = td.Medium( - nonlinear_spec=td.NonlinearSpec(models=[td.TwoPhotonAbsorption(beta=-1, n0=1)]), - allow_gain=True, - ) + with AssertLogLevel(log_capture, "WARNING", contains_str="phenomenological"): + med = td.Medium( + nonlinear_spec=td.NonlinearSpec(models=[td.TwoPhotonAbsorption(beta=-1, n0=1)]), + allow_gain=True, + ) # automatic detection of n0 and freq0 n0 = 2 @@ -637,6 +670,15 @@ def test_nonlinear_medium(log_capture): assert n0 == nonlinear_spec.models[0]._get_n0(n0=None, medium=medium, freqs=[freq0]) assert freq0 == nonlinear_spec.models[0]._get_freq0(freq0=None, freqs=[freq0]) + # subsection with nonlinear materials needs to hardcode source info + sim2 = sim.updated_copy(center=(-4, -4, -4), path="sources/0") + sim2 = sim2.updated_copy( + models=[td.TwoPhotonAbsorption(beta=1)], path="structures/0/medium/nonlinear_spec" + ) + sim2 = sim2.subsection(region=td.Box(center=(0, 0, 0), size=(1, 1, 0))) + assert sim2.structures[0].medium.nonlinear_spec.models[0].n0 == n0 + assert sim2.structures[0].medium.nonlinear_spec.models[0].freq0 == freq0 + # can't detect n0 with different source freqs source_time2 = source_time.updated_copy(freq0=2 * freq0) source2 = source.updated_copy(source_time=source_time2) diff --git a/tests/test_components/test_monitor.py b/tests/test_components/test_monitor.py index 8daccb0a7..115e49fde 100644 --- a/tests/test_components/test_monitor.py +++ b/tests/test_components/test_monitor.py @@ -296,6 +296,24 @@ def test_monitor_num_modes(log_capture, num_modes, log_level): assert_log_level(log_capture, log_level) +def test_mode_bend_radius(): + """Test that small bend radius fails.""" + + with pytest.raises(ValueError): + mnt = td.ModeMonitor( + size=(5, 0, 1), + freqs=np.linspace(1e14, 2e14, 100), + name="test", + mode_spec=td.ModeSpec(num_modes=1, bend_radius=1, bend_axis=1), + ) + _ = td.Simulation( + size=(2, 2, 2), + run_time=1e-12, + monitors=[mnt], + grid_spec=td.GridSpec.uniform(dl=0.1), + ) + + def test_diffraction_validators(): # ensure error if boundaries are not periodic boundary_spec = td.BoundarySpec( diff --git a/tests/test_components/test_simulation.py b/tests/test_components/test_simulation.py index 543bb320a..da12fdfe8 100644 --- a/tests/test_components/test_simulation.py +++ b/tests/test_components/test_simulation.py @@ -108,6 +108,27 @@ def test_sim_init(): sim.epsilon(m) +def test_num_cells(): + """Test num_cells and num_computational_grid_points.""" + + sim = td.Simulation( + size=(1, 1, 1), + run_time=1e-12, + grid_spec=td.GridSpec.uniform(dl=0.1), + sources=[ + td.PointDipole( + center=(0, 0, 0), + polarization="Ex", + source_time=td.GaussianPulse(freq0=2e14, fwidth=1e14), + ) + ], + ) + assert sim.num_computational_grid_points > sim.num_cells # due to extra pixels at boundaries + + sim = sim.updated_copy(symmetry=(1, 0, 0)) + assert sim.num_computational_grid_points < sim.num_cells # due to symmetry + + def test_monitors_data_size(): """make sure a simulation can be initialized""" @@ -402,6 +423,14 @@ def test_validate_plane_wave_boundaries(log_capture): angle_theta=np.pi / 4, ) + mnt = td.DiffractionMonitor( + center=(0, 0, 0), + size=(td.inf, td.inf, 0), + freqs=[250e12, 300e12], + name="monitor_diffraction", + normal_dir="+", + ) + bspec1 = td.BoundarySpec( x=td.Boundary.pml(), y=td.Boundary.absorber(), @@ -459,6 +488,7 @@ def test_validate_plane_wave_boundaries(log_capture): run_time=1e-12, sources=[src2], boundary_spec=bspec3, + monitors=[mnt], ) # angled incidence plane wave with wrong Bloch vector should warn @@ -807,9 +837,9 @@ def test_sim_structure_gap(log_capture, box_size, log_level): structures=[box], sources=[src], boundary_spec=td.BoundarySpec( - x=td.Boundary.pml(num_layers=5), - y=td.Boundary.pml(num_layers=5), - z=td.Boundary.pml(num_layers=5), + x=td.Boundary.pml(num_layers=6), + y=td.Boundary.pml(num_layers=6), + z=td.Boundary.pml(num_layers=6), ), run_time=1e-12, ) @@ -2022,9 +2052,9 @@ def test_sim_volumetric_structures(log_capture, tmp_path): structures=[struct], sources=[src], boundary_spec=td.BoundarySpec( - x=td.Boundary.pml(num_layers=5), - y=td.Boundary.pml(num_layers=5), - z=td.Boundary.pml(num_layers=5), + x=td.Boundary.pml(num_layers=6), + y=td.Boundary.pml(num_layers=6), + z=td.Boundary.pml(num_layers=6), ), grid_spec=td.GridSpec.uniform(dl=grid_dl), run_time=1e-12, @@ -2064,9 +2094,9 @@ def test_sim_volumetric_structures(log_capture, tmp_path): sources=[src], monitors=[monitor], boundary_spec=td.BoundarySpec( - x=td.Boundary.pml(num_layers=5), - y=td.Boundary.pml(num_layers=5), - z=td.Boundary.pml(num_layers=5), + x=td.Boundary.pml(num_layers=6), + y=td.Boundary.pml(num_layers=6), + z=td.Boundary.pml(num_layers=6), ), grid_spec=td.GridSpec.uniform(dl=grid_dl), run_time=1e-12, @@ -2089,9 +2119,9 @@ def test_sim_volumetric_structures(log_capture, tmp_path): sources=[src], monitors=[monitor], boundary_spec=td.BoundarySpec( - x=td.Boundary.pml(num_layers=5), - y=td.Boundary.pml(num_layers=5), - z=td.Boundary.pml(num_layers=5), + x=td.Boundary.pml(num_layers=6), + y=td.Boundary.pml(num_layers=6), + z=td.Boundary.pml(num_layers=6), ), grid_spec=td.GridSpec.uniform(dl=grid_dl), run_time=1e-12, @@ -2120,9 +2150,9 @@ def test_sim_volumetric_structures(log_capture, tmp_path): structures=[below_half, box], sources=[src], boundary_spec=td.BoundarySpec( - x=td.Boundary.pml(num_layers=5), - y=td.Boundary.pml(num_layers=5), - z=td.Boundary.pml(num_layers=5), + x=td.Boundary.pml(num_layers=6), + y=td.Boundary.pml(num_layers=6), + z=td.Boundary.pml(num_layers=6), ), grid_spec=td.GridSpec.uniform(dl=grid_dl), run_time=1e-12, @@ -2136,9 +2166,9 @@ def test_sim_volumetric_structures(log_capture, tmp_path): structures=[box, below], sources=[src], boundary_spec=td.BoundarySpec( - x=td.Boundary.pml(num_layers=5), - y=td.Boundary.pml(num_layers=5), - z=td.Boundary.pml(num_layers=5), + x=td.Boundary.pml(num_layers=6), + y=td.Boundary.pml(num_layers=6), + z=td.Boundary.pml(num_layers=6), ), grid_spec=td.GridSpec.uniform(dl=grid_dl), run_time=1e-12, @@ -2154,9 +2184,9 @@ def test_sim_volumetric_structures(log_capture, tmp_path): sources=[src], medium=box.medium, boundary_spec=td.BoundarySpec( - x=td.Boundary.pml(num_layers=5), - y=td.Boundary.pml(num_layers=5), - z=td.Boundary.pml(num_layers=5), + x=td.Boundary.pml(num_layers=6), + y=td.Boundary.pml(num_layers=6), + z=td.Boundary.pml(num_layers=6), ), grid_spec=td.GridSpec.uniform(dl=grid_dl), run_time=1e-12, @@ -2183,9 +2213,9 @@ def test_sim_volumetric_structures(log_capture, tmp_path): structures=[struct], sources=[src], boundary_spec=td.BoundarySpec( - x=td.Boundary.pml(num_layers=5), - y=td.Boundary.pml(num_layers=5), - z=td.Boundary.pml(num_layers=5), + x=td.Boundary.pml(num_layers=6), + y=td.Boundary.pml(num_layers=6), + z=td.Boundary.pml(num_layers=6), ), grid_spec=td.GridSpec.uniform(dl=grid_dl), run_time=1e-12, @@ -2430,7 +2460,10 @@ def test_sim_subsection(unstructured, nz, log_capture): sim_red = SIM_FULL.subsection(region=region, monitors=[]) assert len(sim_red.monitors) == 0 sim_red = SIM_FULL.subsection(region=region, remove_outside_structures=False) - assert sim_red.structures == SIM_FULL.structures + assert len(sim_red.structures) == len(SIM_FULL.structures) + for strc_red, strc in zip(sim_red.structures, SIM_FULL.structures): + if strc.medium.nonlinear_spec is None: + assert strc == strc_red sim_red = SIM_FULL.subsection(region=region, remove_outside_custom_mediums=True) perm = td.SpatialDataArray( @@ -2903,9 +2936,9 @@ def test_validate_low_num_cells_in_mode_objects(): grid_spec=td.GridSpec(wavelength=1.0), sources=[mode_source], boundary_spec=td.BoundarySpec( - x=td.Boundary.pml(num_layers=5), + x=td.Boundary.pml(num_layers=6), y=td.Boundary.pec(), - z=td.Boundary.pml(num_layers=5), + z=td.Boundary.pml(num_layers=6), ), ) sim2d._validate_num_cells_in_mode_objects() diff --git a/tests/test_components/test_source.py b/tests/test_components/test_source.py index 05ee39062..560559e33 100644 --- a/tests/test_components/test_source.py +++ b/tests/test_components/test_source.py @@ -46,6 +46,22 @@ def test_dir_vector(): assert DirectionalSource._dir_vector.fget(S) is None +def test_mode_bend_radius(): + """Test that small bend radius fails.""" + + with pytest.raises(ValueError): + src = td.ModeSource( + size=(1, 0, 5), + source_time=ST, + mode_spec=td.ModeSpec(num_modes=1, bend_radius=1, bend_axis=0), + ) + _ = td.Simulation( + size=(2, 2, 2), + run_time=1e-12, + sources=[src], + ) + + def test_UniformCurrentSource(): g = td.GaussianPulse(freq0=1e12, fwidth=0.1e12) diff --git a/tests/test_components/test_structure.py b/tests/test_components/test_structure.py index 54fc32231..151a8634b 100644 --- a/tests/test_components/test_structure.py +++ b/tests/test_components/test_structure.py @@ -1,3 +1,5 @@ +import autograd as ag +import autograd.numpy as anp import gdstk import numpy as np import pydantic.v1 as pd @@ -211,3 +213,26 @@ def test_validation_of_structures_with_2d_materials(): for geom in not_allowed_geometries: with pytest.raises(pd.ValidationError): _ = td.Structure(geometry=geom, medium=med2d) + + +def test_from_permittivity_array(): + nx, ny, nz = 10, 1, 12 + box = td.Box(size=(2, td.inf, 4), center=(0, 0, 0)) + + eps_data = 1.0 + np.random.random((nx, ny, nz)) + + structure = td.Structure.from_permittivity_array(geometry=box, eps_data=eps_data, name="test") + + assert structure.name == "test" + + assert np.all(structure.medium.permittivity.values == eps_data) + + def f(x): + eps_data = (1 + x) * (1 + np.random.random((nx, ny, nz))) + structure = td.Structure.from_permittivity_array( + geometry=box, eps_data=eps_data, name="test" + ) + return anp.sum(structure.medium.permittivity).item() + + grad = ag.grad(f)(1.0) + assert not np.isclose(grad, 0.0) diff --git a/tests/test_components/test_viz.py b/tests/test_components/test_viz.py index 93174b0ac..e88279e6d 100644 --- a/tests/test_components/test_viz.py +++ b/tests/test_components/test_viz.py @@ -57,7 +57,7 @@ def test_2d_boundary_plot(): # Simulation details per_boundary = td.Boundary.periodic() - pml_boundary = td.Boundary.pml(num_layers=2) + pml_boundary = td.Boundary.pml(num_layers=6) sim = td.Simulation( size=(0, 1, 1), diff --git a/tests/test_data/test_data_arrays.py b/tests/test_data/test_data_arrays.py index 152844b2b..fa975741d 100644 --- a/tests/test_data/test_data_arrays.py +++ b/tests/test_data/test_data_arrays.py @@ -5,6 +5,7 @@ import numpy as np import pytest import tidy3d as td +import xarray.testing as xrt from tidy3d.exceptions import DataError np.random.seed(4) @@ -436,3 +437,17 @@ def test_uniform_check(): coords=dict(x=[0, 1], y=[1, 2], z=[2, 3]), ) assert not arr.is_uniform + + +@pytest.mark.parametrize("method", ["nearest", "linear"]) +@pytest.mark.parametrize("scalar_index", [True, False]) +def test_interp(method, scalar_index): + data = make_scalar_field_data_array("Ex") + + f = 1.5e14 + if not scalar_index: + f = [f] + + xr_interp = data.interp(f=f) + ag_interp = data._ag_interp(f=f) + xrt.assert_allclose(xr_interp, ag_interp) diff --git a/tests/test_data/test_sim_data.py b/tests/test_data/test_sim_data.py index d31214419..a4a597c38 100644 --- a/tests/test_data/test_sim_data.py +++ b/tests/test_data/test_sim_data.py @@ -153,6 +153,10 @@ def test_plot(phase): "field", field_cmp, val="imag", f=1e14, phase=phase, **xyz_kwargs ) plt.close() + for shading in ["gouraud", "nearest", "auto"]: + _ = sim_data.plot_field( + "field", field_cmp, val="imag", f=1e14, phase=phase, shading=shading, **xyz_kwargs + ) for axis_name in "xyz": xyz_kwargs = {axis_name: 0} _ = sim_data.plot_field("field", "int", f=1e14, phase=phase, **xyz_kwargs) diff --git a/tests/test_package/test_log.py b/tests/test_package/test_log.py index 335482cbe..1a352df2e 100644 --- a/tests/test_package/test_log.py +++ b/tests/test_package/test_log.py @@ -129,7 +129,7 @@ def test_logging_warning_capture(): medium=td.Medium(permittivity=2, frequency_range=[0.5, 1]), ) - # 2 warnings: inside pml + # 1 warning: inside pml box_in_pml = td.Structure( geometry=td.Box(center=(0, 0, 0), size=(domain_size * 1.001, 5, 5)), medium=td.Medium(permittivity=10), @@ -216,7 +216,7 @@ def test_logging_warning_capture(): sim.validate_pre_upload() warning_list = td.log.captured_warnings() print(json.dumps(warning_list, indent=4)) - assert len(warning_list) == 32 + assert len(warning_list) == 31 td.log.set_capture(False) # check that capture doesn't change validation errors diff --git a/tests/test_plugins/autograd/__init__.py b/tests/test_plugins/autograd/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/tests/test_plugins/autograd/invdes/test_filters.py b/tests/test_plugins/autograd/invdes/test_filters.py new file mode 100644 index 000000000..2be94fedd --- /dev/null +++ b/tests/test_plugins/autograd/invdes/test_filters.py @@ -0,0 +1,80 @@ +import pytest +from tidy3d.plugins.autograd.invdes.filters import ( + _get_kernel_size, + make_circular_filter, + make_conic_filter, + make_filter, +) +from tidy3d.plugins.autograd.types import PaddingType + + +@pytest.mark.parametrize( + "radius, dl, size_px, expected", + [ + (1, 0.1, None, (21,)), + (1, [0.1, 0.2], None, (21, 11)), + ([1, 2], 0.1, None, (21, 41)), + ([1, 1], [0.1, 0.2], None, (21, 11)), + ([1, 2], [0.1, 0.1], None, (21, 41)), + (None, None, 5, (5,)), + (None, None, (5, 7), (5, 7)), + ], +) +def test_get_kernel_size(radius, dl, size_px, expected): + result = _get_kernel_size(radius, dl, size_px) + assert result == expected + + +def test_get_kernel_size_invalid_arguments(): + with pytest.raises( + ValueError, match="Either 'size_px' or both 'radius' and 'dl' must be provided." + ): + _get_kernel_size(None, None, None) + + +@pytest.mark.parametrize("radius", [1, 2, (1, 2)]) +@pytest.mark.parametrize("dl", [0.1, 0.2, (0.1, 0.2)]) +@pytest.mark.parametrize("size_px", [None, 5, (5, 7)]) +@pytest.mark.parametrize("normalize", [True, False]) +@pytest.mark.parametrize("padding", PaddingType.__args__) +class TestMakeFilter: + @pytest.mark.parametrize("filter_type", ["circular", "conic"]) + def test_make_filter(self, rng, filter_type, radius, dl, size_px, normalize, padding): + """Test make_filter function for various parameters.""" + filter_func = make_filter( + radius=radius, + dl=dl, + size_px=size_px, + normalize=normalize, + padding=padding, + filter_type=filter_type, + ) + array = rng.random((51, 51)) + result = filter_func(array) + assert result.shape == array.shape + + def test_make_circular_filter(self, rng, radius, dl, size_px, normalize, padding): + """Test make_circular_filter function for various parameters.""" + filter_func = make_circular_filter( + radius=radius, + dl=dl, + size_px=size_px, + normalize=normalize, + padding=padding, + ) + array = rng.random((51, 51)) + result = filter_func(array) + assert result.shape == array.shape + + def test_make_conic_filter(self, rng, radius, dl, size_px, normalize, padding): + """Test make_conic_filter function for various parameters.""" + filter_func = make_conic_filter( + radius=radius, + dl=dl, + size_px=size_px, + normalize=normalize, + padding=padding, + ) + array = rng.random((51, 51)) + result = filter_func(array) + assert result.shape == array.shape diff --git a/tests/test_plugins/autograd/invdes/test_parametrizations.py b/tests/test_plugins/autograd/invdes/test_parametrizations.py new file mode 100644 index 000000000..aa31dec92 --- /dev/null +++ b/tests/test_plugins/autograd/invdes/test_parametrizations.py @@ -0,0 +1,26 @@ +import autograd.numpy as np +import pytest +from tidy3d.plugins.autograd.invdes.parametrizations import make_filter_and_project +from tidy3d.plugins.autograd.types import PaddingType + + +@pytest.mark.parametrize("radius", [1, 2, (1, 2)]) +@pytest.mark.parametrize("dl", [0.1, 0.2, (0.1, 0.2)]) +@pytest.mark.parametrize("size_px", [None, 5, (5, 7)]) +@pytest.mark.parametrize("filter_type", ["circular", "conic"]) +@pytest.mark.parametrize("padding", PaddingType.__args__) +def test_make_filter_and_project(rng, radius, dl, size_px, filter_type, padding): + """Test make_filter_and_project function for various parameters.""" + filter_and_project_func = make_filter_and_project( + radius=radius, + dl=dl, + size_px=size_px, + beta=10, + eta=0.5, + filter_type=filter_type, + padding=padding, + ) + array = rng.random((51, 51)) + result = filter_and_project_func(array) + assert result.shape == array.shape + assert np.all(result >= 0) and np.all(result <= 1) diff --git a/tests/test_plugins/autograd/invdes/test_penalties.py b/tests/test_plugins/autograd/invdes/test_penalties.py new file mode 100644 index 000000000..8ddad5267 --- /dev/null +++ b/tests/test_plugins/autograd/invdes/test_penalties.py @@ -0,0 +1,24 @@ +import pytest +from tidy3d.plugins.autograd.invdes.penalties import make_erosion_dilation_penalty +from tidy3d.plugins.autograd.types import PaddingType + + +@pytest.mark.parametrize("radius", [1, 2, (1, 2)]) +@pytest.mark.parametrize("dl", [0.1, 0.2, (0.1, 0.2)]) +@pytest.mark.parametrize("size_px", [None, 5, (5, 7)]) +@pytest.mark.parametrize("padding", PaddingType.__args__) +def test_make_erosion_dilation_penalty(rng, radius, dl, size_px, padding): + """Test make_erosion_dilation_penalty function for various parameters.""" + erosion_dilation_penalty_func = make_erosion_dilation_penalty( + radius=radius, + dl=dl, + size_px=size_px, + beta=10, + eta=0.5, + delta_eta=0.01, + padding=padding, + ) + array = rng.random((51, 51)) + result = erosion_dilation_penalty_func(array) + assert isinstance(result, float) + assert result >= 0 diff --git a/tests/test_plugins/autograd/test_differential_operators.py b/tests/test_plugins/autograd/test_differential_operators.py index 21fbe1b1a..75879ea98 100644 --- a/tests/test_plugins/autograd/test_differential_operators.py +++ b/tests/test_plugins/autograd/test_differential_operators.py @@ -1,23 +1,66 @@ import autograd.numpy as np +import pytest +from autograd import grad as grad_ag from autograd import value_and_grad as value_and_grad_ag from numpy.testing import assert_allclose -from tidy3d.plugins.autograd.differential_operators import value_and_grad +from tidy3d.components.data.data_array import DataArray +from tidy3d.plugins.autograd import grad, value_and_grad -def test_value_and_grad(rng): +@pytest.mark.parametrize("argnum", [0, 1]) +@pytest.mark.parametrize("has_aux", [True, False]) +def test_grad(rng, argnum, has_aux): """Test the custom value_and_grad function against autograd's implementation""" x = rng.random(10) + y = rng.random(10) aux_val = "aux" - vg_fun = value_and_grad(lambda x: (np.linalg.norm(x), aux_val), has_aux=True) - vg_fun_ag = value_and_grad_ag(lambda x: np.linalg.norm(x)) + def f(x, y): + ret = DataArray(x * y).sum() # still DataArray + if has_aux: + return ret, aux_val + return ret - (v, g), aux = vg_fun(x) - v_ag, g_ag = vg_fun_ag(x) + grad_fun = grad(f, argnum=argnum, has_aux=has_aux) + grad_fun_ag = grad_ag( + lambda x, y: f(x, y)[0].item() if has_aux else f(x, y).item(), argnum=argnum + ) + + if has_aux: + g, aux = grad_fun(x, y) + assert aux == aux_val + else: + g = grad_fun(x, y) + g_ag = grad_fun_ag(x, y) - # assert that values and gradients match - assert_allclose(v, v_ag) assert_allclose(g, g_ag) - # check that auxiliary output is correctly returned - assert aux == aux_val + +@pytest.mark.parametrize("argnum", [0, 1]) +@pytest.mark.parametrize("has_aux", [True, False]) +def test_value_and_grad(rng, argnum, has_aux): + """Test the custom value_and_grad function against autograd's implementation""" + x = rng.random(10) + y = rng.random(10) + aux_val = "aux" + + def f(x, y): + ret = DataArray(np.linalg.norm(x * y)).sum() # still DataArray + if has_aux: + return ret, aux_val + return ret + + vg_fun = value_and_grad(f, argnum=argnum, has_aux=has_aux) + vg_fun_ag = value_and_grad_ag( + lambda x, y: f(x, y)[0].item() if has_aux else f(x, y).item(), argnum=argnum + ) + + if has_aux: + (v, g), aux = vg_fun(x, y) + assert aux == aux_val + else: + v, g = vg_fun(x, y) + v_ag, g_ag = vg_fun_ag(x, y) + + assert_allclose(v, v_ag) + assert_allclose(g, g_ag) diff --git a/tests/test_plugins/autograd/test_functions.py b/tests/test_plugins/autograd/test_functions.py index eeb3c2bf6..a54723ecb 100644 --- a/tests/test_plugins/autograd/test_functions.py +++ b/tests/test_plugins/autograd/test_functions.py @@ -5,19 +5,24 @@ import scipy.ndimage from autograd.test_util import check_grads from scipy.signal import convolve as convolve_sp -from tidy3d.plugins.autograd.functions import ( +from tidy3d.plugins.autograd import ( + add_at, convolve, grey_closing, grey_dilation, grey_erosion, grey_opening, interpn, + least_squares, morphological_gradient, morphological_gradient_external, morphological_gradient_internal, pad, rescale, + smooth_max, + smooth_min, threshold, + trapz, ) from tidy3d.plugins.autograd.types import PaddingType @@ -32,7 +37,7 @@ @pytest.mark.parametrize("mode", PaddingType.__args__) @pytest.mark.parametrize("size", [3, 4, (3, 3), (4, 4), (3, 4), (3, 3, 3), (4, 4, 4), (3, 4, 5)]) -@pytest.mark.parametrize("pad_width", [0, 1, 2, (0, 0), (0, 1), (1, 0), (1, 2)]) +@pytest.mark.parametrize("pad_width", [0, 1, 2, 4, 5, (0, 0), (0, 1), (1, 0), (1, 2)]) @pytest.mark.parametrize("axis", [None, 0, -1]) class TestPad: def test_pad_val(self, rng, mode, size, pad_width, axis): @@ -53,7 +58,7 @@ def test_pad_val(self, rng, mode, size, pad_width, axis): def test_pad_grad(self, rng, mode, size, pad_width, axis): """Test gradients of padding function for various modes, sizes, pad widths, and axes.""" x = rng.random(size) - check_grads(pad, modes=["fwd", "rev"], order=1)(x, pad_width, mode=mode, axis=axis) + check_grads(pad, modes=["fwd", "rev"], order=2)(x, pad_width, mode=mode, axis=axis) class TestPadExceptions: @@ -64,31 +69,24 @@ def test_invalid_pad_width_size(self): with pytest.raises(ValueError, match="Padding width must have one or two elements"): pad(self.array, (1, 2, 3)) - def test_padding_larger_than_input_size(self): - """Test that an exception is raised when padding is larger than the input size.""" - with pytest.raises( - NotImplementedError, match="Padding larger than the input size is not supported" - ): - pad(self.array, (3, 3)) - def test_negative_padding(self): """Test that an exception is raised when padding is negative.""" - with pytest.raises(ValueError, match="Padding must be positive"): + with pytest.raises(ValueError, match="Padding must be non-negative"): pad(self.array, (-1, 1)) def test_unsupported_padding_mode(self): """Test that an exception is raised when an unsupported padding mode is used.""" - with pytest.raises(KeyError, match="Unsupported padding mode"): + with pytest.raises(ValueError, match="Unsupported padding mode"): pad(self.array, (1, 1), mode="unsupported_mode") def test_axis_out_of_range(self): """Test that an exception is raised when the axis is out of range.""" - with pytest.raises(IndexError, match="Axis out of range"): + with pytest.raises(IndexError, match="out of range"): pad(self.array, (1, 1), axis=2) def test_negative_axis_out_of_range(self): """Test that an exception is raised when a negative axis is out of range.""" - with pytest.raises(IndexError, match="Axis out of range"): + with pytest.raises(IndexError, match="out of range"): pad(self.array, (1, 1), axis=-3) @@ -141,7 +139,7 @@ def test_convolve_grad(self, rng, mode, padding, ary_size, kernel_size, square_k ) x, k = self._ary_and_kernel(rng, ary_size, kernel_size, square_kernel) - check_grads(convolve, modes=["rev"], order=1)(x, k, padding=padding, mode=mode) + check_grads(convolve, modes=["rev"], order=2)(x, k, padding=padding, mode=mode) class TestConvolveExceptions: @@ -199,7 +197,7 @@ def test_morphology_val_size(self, rng, op, sp_op, mode, ary_size, kernel_size): def test_morphology_val_grad(self, rng, op, sp_op, mode, ary_size, kernel_size): """Test gradients of morphological operations for various modes, array sizes, and kernel sizes.""" x = rng.random(ary_size) - check_grads(op, modes=["rev"], order=1)(x, size=kernel_size, mode=mode) + check_grads(op, modes=["rev"], order=2)(x, size=kernel_size, mode=mode) @pytest.mark.parametrize( "full", @@ -243,7 +241,7 @@ def test_morphology_val_structure_grad( ): """Test gradients of morphological operations for various kernel structures.""" x, k = self._ary_and_kernel(rng, ary_size, kernel_size, full, square, flat) - check_grads(op, modes=["rev"], order=1)(x, size=kernel_size, mode=mode) + check_grads(op, modes=["rev"], order=2)(x, size=kernel_size, mode=mode) @pytest.mark.parametrize( @@ -318,15 +316,13 @@ def test_interpn_val(self, rng, dim, method): points, values, xi = self.generate_points_values_xi(rng, dim) xi_grid = np.meshgrid(*xi, indexing="ij") - result_custom = interpn(points, values, xi, method=method) + result_custom = interpn(points, values, tuple(xi_grid), method=method) result_scipy = scipy.interpolate.interpn(points, values, tuple(xi_grid), method=method) npt.assert_allclose(result_custom, result_scipy) - @pytest.mark.parametrize("order", [1, 2]) - @pytest.mark.parametrize("mode", ["fwd", "rev"]) - def test_interpn_values_grad(self, rng, dim, method, order, mode): + def test_interpn_values_grad(self, rng, dim, method): points, values, xi = self.generate_points_values_xi(rng, dim) - check_grads(lambda v: interpn(points, v, xi, method=method), modes=[mode], order=order)( + check_grads(lambda v: interpn(points, v, xi, method=method), modes=["fwd", "rev"], order=2)( values ) @@ -337,3 +333,184 @@ def test_invalid_method(self, rng): points, values, xi = TestInterpn.generate_points_values_xi(rng, 2) with pytest.raises(ValueError, match="interpolation method"): interpn(points, values, xi, method="invalid_method") + + +@pytest.mark.parametrize("axis", [0, -1]) +@pytest.mark.parametrize("shape", [(10,), (10, 10)]) +@pytest.mark.parametrize("use_x", [True, False]) +class TestTrapz: + @staticmethod + def generate_y_x_dx(rng, shape, use_x): + y = rng.uniform(-1, 1, shape) + if use_x: + x = rng.random(shape) + dx = 1.0 # dx is not used when x is provided + else: + x = None + dx = rng.random() + 0.1 # ensure dx is not zero + return y, x, dx + + def test_trapz_val(self, rng, shape, axis, use_x): + """Test trapz values against NumPy for different array dimensions and integration axes.""" + y, x, dx = self.generate_y_x_dx(rng, shape, use_x) + result_custom = trapz(y, x=x, dx=dx, axis=axis) + result_numpy = np.trapz(y, x=x, dx=dx, axis=axis) + npt.assert_allclose(result_custom, result_numpy) + + def test_trapz_grad(self, rng, shape, axis, use_x): + """Test gradients of trapz function for different array dimensions and integration axes.""" + y, x, dx = self.generate_y_x_dx(rng, shape, use_x) + check_grads(lambda y: trapz(y, x=x, dx=dx, axis=axis), modes=["fwd", "rev"], order=2)(y) + + +@pytest.mark.parametrize("shape", [(10,), (10, 10)]) +@pytest.mark.parametrize("indices", [(0,), (slice(3, 8),)]) +class TestAddAt: + @staticmethod + def generate_x_y(rng, shape, indices): + x = rng.uniform(-1, 1, shape) + y = rng.uniform(-1, 1, x[tuple(indices)].shape) + return x, y + + def test_add_at_val(self, rng, shape, indices): + """Test add_at values against NumPy for different array dimensions and indices.""" + x, y = self.generate_x_y(rng, shape, indices) + result_custom = add_at(x, indices, y) + result_numpy = np.array(x) + result_numpy[indices] += y + npt.assert_allclose(result_custom, result_numpy) + + def test_add_at_grad(self, rng, shape, indices): + """Test gradients of add_at function for different array dimensions and indices.""" + x, y = self.generate_x_y(rng, shape, indices) + check_grads(lambda x: add_at(x, indices, y), modes=["fwd", "rev"], order=2)(x) + check_grads(lambda y: add_at(x, indices, y), modes=["fwd", "rev"], order=2)(y) + + +@pytest.mark.parametrize("shape", [(5,), (5, 5), (5, 5, 5)]) +@pytest.mark.parametrize("tau", [1e-3, 1.0]) +@pytest.mark.parametrize("axis", [None, 0, 1, -1]) +class TestSmoothMax: + def test_smooth_max_values(self, rng, shape, tau, axis): + """Test `smooth_max` values for various shapes, tau, and axes.""" + + if axis == 1 and len(shape) == 1: + pytest.skip() + + x = rng.uniform(-10, 10, size=shape) + result = smooth_max(x, tau=tau, axis=axis) + + expected = np.max(x, axis=axis) + npt.assert_allclose(result, expected, atol=10 * tau) + + def test_smooth_max_grad(self, rng, shape, tau, axis): + """Test gradients of `smooth_max` for various parameters.""" + + if axis == 1 and len(shape) == 1: + pytest.skip() + + x = rng.uniform(-1, 1, size=shape) + func = lambda x: smooth_max(x, tau=tau, axis=axis) + check_grads(func, modes=["fwd", "rev"], order=2)(x) + + +@pytest.mark.parametrize("shape", [(5,), (5, 5), (5, 5, 5)]) +@pytest.mark.parametrize("tau", [1e-3, 1.0]) +@pytest.mark.parametrize("axis", [None, 0, 1, -1]) +class TestSmoothMin: + def test_smooth_min_values(self, rng, shape, tau, axis): + """Test `smooth_min` values for various shapes, tau, and axes.""" + + if axis == 1 and len(shape) == 1: + pytest.skip() + + x = rng.uniform(-10, 10, size=shape) + result = smooth_min(x, tau=tau, axis=axis) + + expected = np.min(x, axis=axis) + npt.assert_allclose(result, expected, atol=10 * tau) + + def test_smooth_min_grad(self, rng, shape, tau, axis): + """Test gradients of `smooth_min` for various parameters.""" + + if axis == 1 and len(shape) == 1: + pytest.skip() + + x = rng.uniform(-1, 1, size=shape) + func = lambda x: smooth_min(x, tau=tau, axis=axis) + check_grads(func, modes=["fwd", "rev"], order=2)(x) + + +class TestLeastSquares: + @pytest.mark.parametrize( + "model, params_true, initial_guess, x, y", + [ + ( + lambda x, a, b: a * x + b, + np.array([2.0, -3.0]), + (0.0, 0.0), + np.linspace(0, 10, 50), + 2.0 * np.linspace(0, 10, 50) - 3.0, + ), + ( + lambda x, a, b, c: a * x**2 + b * x + c, + np.array([1.0, -2.0, 1.0]), + (0.0, 0.0, 0.0), + np.linspace(-5, 5, 100), + 1.0 * np.linspace(-5, 5, 100) ** 2 - 2.0 * np.linspace(-5, 5, 100) + 1.0, + ), + ( + lambda x, a, b: a * np.exp(b * x), + np.array([1.5, 0.5]), + (1.0, 0.0), + np.linspace(0, 2, 50), + 1.5 * np.exp(0.5 * np.linspace(0, 2, 50)), + ), + ], + ) + def test_least_squares(self, model, params_true, initial_guess, x, y): + """Test least_squares function with different models.""" + params_estimated = least_squares(model, x, y, initial_guess) + npt.assert_allclose(params_estimated, params_true, rtol=1e-5) + + def test_least_squares_with_noise(self, rng): + """Test least_squares function with noisy data.""" + + model = lambda x, a, b: a * x + b + a_true, b_true = -1.0, 4.0 + params_true = np.array([a_true, b_true]) + x = np.linspace(0, 10, 100) + noise = rng.normal(scale=0.1, size=x.shape) + y = a_true * x + b_true + noise + initial_guess = (0.0, 0.0) + + params_estimated = least_squares(model, x, y, initial_guess) + + npt.assert_allclose(params_estimated, params_true, rtol=1e-1) + + def test_least_squares_no_convergence(self): + """Test that least_squares function raises an error when not converging.""" + + def constant_model(x, a): + return a + + x = np.linspace(0, 10, 50) + y = 2.0 * x - 3.0 # Linear data + initial_guess = (0.0,) + + with pytest.raises(np.linalg.LinAlgError): + least_squares(constant_model, x, y, initial_guess, max_iterations=10, tol=1e-12) + + def test_least_squares_gradient(self): + """Test gradients of least_squares function with respect to parameters.""" + + def linear_model(x, a, b): + return a * x + b + + x = np.linspace(0, 10, 50) + y = 2.0 * x - 3.0 + initial_guess = (1.0, 0.0) + + check_grads( + lambda params: least_squares(linear_model, x, y, params), modes=["fwd", "rev"], order=2 + )(initial_guess) diff --git a/tests/test_plugins/autograd/test_utilities.py b/tests/test_plugins/autograd/test_utilities.py index bf5e1cdea..68cf7a970 100644 --- a/tests/test_plugins/autograd/test_utilities.py +++ b/tests/test_plugins/autograd/test_utilities.py @@ -1,7 +1,15 @@ import numpy as np import numpy.testing as npt import pytest -from tidy3d.plugins.autograd.utilities import chain, make_kernel +import xarray as xr +from tidy3d.exceptions import Tidy3dError +from tidy3d.plugins.autograd import ( + chain, + get_kernel_size_px, + make_kernel, + scalar_objective, + value_and_grad, +) @pytest.mark.parametrize("size", [(3, 3), (4, 4), (5, 5)]) @@ -42,6 +50,27 @@ def test_make_kernel_invalid_size(self): make_kernel("circular", size) +@pytest.mark.parametrize( + "radius, dl, expected", + [ + (1, 0.1, 21), + (1, [0.1, 0.2], [21, 11]), + ([1, 2], 0.1, [21, 41]), + ([1, 1], [0.1, 0.2], [21, 11]), + ([1, 2], [0.1, 0.1], [21, 41]), + ], +) +def test_get_kernel_size_px_with_radius_and_dl(radius, dl, expected): + result = get_kernel_size_px(radius, dl) + assert result == expected + + +def test_get_kernel_size_px_invalid_arguments(): + """Test get_kernel_size_px function with invalid arguments.""" + with pytest.raises(ValueError, match="must be provided"): + get_kernel_size_px() + + class TestChain: def test_chain_functions(self): """Test chain function with multiple functions.""" @@ -83,3 +112,55 @@ def add_one(x): funcs = [add_one, "not_a_function"] with pytest.raises(TypeError, match="All elements in funcs must be callable"): chain(funcs) + + +class TestScalarObjective: + def test_scalar_objective_no_aux(self): + """Test scalar_objective decorator without auxiliary data.""" + + @scalar_objective + def objective(x): + da = xr.DataArray(x) + return da.sum() + + x = np.array([1.0, 2.0, 3.0]) + result, grad = value_and_grad(objective)(x) + assert np.allclose(grad, np.ones_like(grad)) + assert np.isclose(result, 6.0) + + def test_scalar_objective_with_aux(self): + """Test scalar_objective decorator with auxiliary data.""" + + @scalar_objective(has_aux=True) + def objective(x): + da = xr.DataArray(x) + return da.sum(), "auxiliary_data" + + x = np.array([1.0, 2.0, 3.0]) + (result, grad), aux_data = value_and_grad(objective, has_aux=True)(x) + assert np.allclose(grad, np.ones_like(grad)) + assert np.isclose(result, 6.0) + assert aux_data == "auxiliary_data" + + def test_scalar_objective_invalid_return(self): + """Test scalar_objective decorator with invalid return value.""" + + @scalar_objective + def objective(x): + da = xr.DataArray(x) + return da # Returning the array directly, not a scalar + + x = np.array([1, 2, 3]) + with pytest.raises(Tidy3dError, match="must be a scalar"): + objective(x) + + def test_scalar_objective_float(self): + """Test scalar_objective decorator with a Python float return value.""" + + @scalar_objective + def objective(x): + return x**2 + + result, grad = value_and_grad(objective)(3.0) + assert np.isclose(grad, 6.0) + assert np.isclose(result, 9.0) diff --git a/tests/test_plugins/expressions/__init__.py b/tests/test_plugins/expressions/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/tests/test_plugins/expressions/test_functions.py b/tests/test_plugins/expressions/test_functions.py new file mode 100644 index 000000000..d7942b545 --- /dev/null +++ b/tests/test_plugins/expressions/test_functions.py @@ -0,0 +1,35 @@ +import numpy as np +import pytest +from tidy3d.plugins.expressions.functions import Cos, Exp, Log, Log10, Sin, Sqrt, Tan +from tidy3d.plugins.expressions.variables import Constant + +FUNCTIONS = [ + (Sin, np.sin), + (Cos, np.cos), + (Tan, np.tan), + (Exp, np.exp), + (Log, np.log), + (Log10, np.log10), + (Sqrt, np.sqrt), +] + + +@pytest.fixture(params=[1, 2.5, 1 + 2j, np.array([1, 2, 3]), np.array([1 + 2j, 3 - 4j, 5 + 6j])]) +def value(request): + return request.param + + +@pytest.mark.parametrize("tidy3d_func, numpy_func", FUNCTIONS) +def test_functions_evaluate(tidy3d_func, numpy_func, value): + constant = Constant(value) + func = tidy3d_func(constant) + result = func.evaluate(constant.evaluate()) + np.testing.assert_allclose(result, numpy_func(value)) + + +@pytest.mark.parametrize("tidy3d_func, numpy_func", FUNCTIONS) +def test_functions_type(tidy3d_func, numpy_func, value): + constant = Constant(value) + func = tidy3d_func(constant) + result = func.evaluate(constant.evaluate()) + assert isinstance(result, type(numpy_func(value))) diff --git a/tests/test_plugins/expressions/test_operators.py b/tests/test_plugins/expressions/test_operators.py new file mode 100644 index 000000000..d794a928f --- /dev/null +++ b/tests/test_plugins/expressions/test_operators.py @@ -0,0 +1,59 @@ +import operator + +import numpy as np +import pytest +from tidy3d.plugins.expressions.operators import ( + Abs, + Add, + Divide, + FloorDivide, + MatMul, + Modulus, + Multiply, + Negate, + Power, + Subtract, +) + +UNARY_OPS = [ + (Abs, operator.abs), + (Negate, operator.neg), +] + +BINARY_OPS = [ + (Add, operator.add), + (Subtract, operator.sub), + (Multiply, operator.mul), + (Divide, operator.truediv), + (FloorDivide, operator.floordiv), + (Modulus, operator.mod), + (Power, operator.pow), + (MatMul, operator.matmul), +] + + +@pytest.mark.parametrize("tidy3d_op, python_op", BINARY_OPS) +@pytest.mark.parametrize( + "x, y", [(1, 2), (2.5, 3.7), (1 + 2j, 3 - 4j), (np.array([1, 2, 3]), np.array([4, 5, 6]))] +) +def test_binary_operators(tidy3d_op, python_op, x, y): + if any(isinstance(p, complex) for p in (x, y)) and any( + tidy3d_op == op for op in (FloorDivide, Modulus) + ): + pytest.skip("operation undefined for complex inputs") + if tidy3d_op == MatMul and (np.isscalar(x) or np.isscalar(y)): + pytest.skip("matmul operation undefined for scalar inputs") + + tidy3d_result = tidy3d_op(left=x, right=y).evaluate(x) + python_result = python_op(x, y) + + np.testing.assert_allclose(tidy3d_result, python_result) + + +@pytest.mark.parametrize("tidy3d_op, python_op", UNARY_OPS) +@pytest.mark.parametrize("x", [1, 2.5, 1 + 2j, np.array([1 + 2j, 3 - 4j, 5 + 6j])]) +def test_unary_operators(tidy3d_op, python_op, x): + tidy3d_result = tidy3d_op(operand=x).evaluate(x) + python_result = python_op(x) + + np.testing.assert_allclose(tidy3d_result, python_result) diff --git a/tests/test_plugins/expressions/test_variables.py b/tests/test_plugins/expressions/test_variables.py new file mode 100644 index 000000000..0d4818c9a --- /dev/null +++ b/tests/test_plugins/expressions/test_variables.py @@ -0,0 +1,114 @@ +import numpy as np +import pytest +from tidy3d.plugins.expressions.variables import Constant, Variable + + +@pytest.fixture(params=[1, 2.5, 1 + 2j, np.array([1, 2, 3])]) +def value(request): + return request.param + + +def test_constant_evaluate(value): + constant = Constant(value) + result = constant.evaluate() + np.testing.assert_allclose(result, value) + + +def test_constant_type(value): + constant = Constant(value) + result = constant.evaluate() + assert isinstance(result, type(value)) + + +def test_variable_evaluate_positional(value): + variable = Variable() + result = variable.evaluate(value) + np.testing.assert_allclose(result, value) + + +def test_variable_evaluate_named(value): + variable = Variable(name="x") + result = variable.evaluate(x=value) + np.testing.assert_allclose(result, value) + + +def test_variable_missing_positional(): + variable = Variable() + with pytest.raises(ValueError, match="No positional argument provided for unnamed variable."): + variable.evaluate() + + +def test_variable_missing_named(): + variable = Variable(name="x") + with pytest.raises(ValueError, match="Variable 'x' not provided."): + variable.evaluate() + + +def test_variable_wrong_named(): + variable = Variable(name="x") + with pytest.raises(ValueError, match="Variable 'x' not provided."): + variable.evaluate(y=5) + + +def test_variable_repr(): + variable_unnamed = Variable() + variable_named = Variable(name="x") + assert repr(variable_unnamed) == "Variable()" + assert repr(variable_named) == "x" + + +def test_variable_in_expression_positional(value): + variable = Variable() + expr = variable + 2 + result = expr(value) + expected = value + 2 + np.testing.assert_allclose(result, expected) + + +def test_variable_in_expression_named(value): + variable = Variable(name="x") + expr = variable + 2 + result = expr(x=value) + expected = value + 2 + np.testing.assert_allclose(result, expected) + + +def test_variable_mixed_args(): + variable_unnamed = Variable() + variable_named = Variable(name="x") + expr = variable_unnamed + variable_named + result = expr(5, x=3) + expected = 5 + 3 + np.testing.assert_allclose(result, expected) + + +def test_variable_missing_args(): + variable_unnamed = Variable() + variable_named = Variable(name="x") + expr = variable_unnamed + variable_named + with pytest.raises(ValueError, match="No positional argument provided for unnamed variable."): + expr(x=3) + with pytest.raises(ValueError, match="Variable 'x' not provided."): + expr(5) + + +def test_variable_multiple_positional_args(): + variable1 = Variable() + variable2 = Variable() + expr = variable1 + variable2 + with pytest.raises(ValueError, match="Multiple positional arguments"): + expr(5, 3) + + +def test_single_unnamed_variable_multiple_args(): + variable = Variable() + expr = variable * 2 + with pytest.raises(ValueError, match="Multiple positional arguments"): + expr(5, 3) + + +def test_multiple_unnamed_variables(): + variable1 = Variable() + variable2 = Variable() + expr = variable1 + variable2 + assert expr(5) == 10 diff --git a/tests/test_plugins/pytorch/test_wrapper.py b/tests/test_plugins/pytorch/test_wrapper.py new file mode 100644 index 000000000..2e80996c0 --- /dev/null +++ b/tests/test_plugins/pytorch/test_wrapper.py @@ -0,0 +1,43 @@ +import autograd.numpy as anp +import torch +from autograd import elementwise_grad +from numpy.testing import assert_allclose +from tidy3d.plugins.pytorch.wrapper import to_torch + + +def test_to_torch_no_kwargs(rng): + x_np = rng.uniform(-1, 1, 10).astype("f4") + x_torch = torch.tensor(x_np, requires_grad=True) + + def f_np(x): + return x * anp.sin(x) ** 2 + + f_torch = to_torch(f_np) + + val = f_torch(x_torch) + val.backward(torch.ones(x_torch.shape)) + + grad = x_torch.grad.numpy() + expected_grad = elementwise_grad(f_np)(x_np) + + assert_allclose(grad, expected_grad) + + +def test_to_torch_with_kwargs(rng): + x_np = rng.uniform(-1, 1, 10).astype("f4") + y_np = rng.uniform(-1, 1, 10).astype("f4") + x_torch = torch.tensor(x_np, requires_grad=True) + y_torch = torch.tensor(y_np, requires_grad=True) + + def f_np(x, y): + return y * anp.sin(x) ** 2 + + f_torch = to_torch(f_np) + + val = f_torch(y=y_torch, x=x_torch) + val.backward(torch.ones(x_torch.shape)) + + grad = x_torch.grad.numpy(), y_torch.grad.numpy() + expected_grad = elementwise_grad(f_np, argnum=[0, 1])(x_np, y_np) + + assert_allclose(grad, expected_grad) diff --git a/tests/test_plugins/test_adjoint.py b/tests/test_plugins/test_adjoint.py index 9c8807f04..795412293 100644 --- a/tests/test_plugins/test_adjoint.py +++ b/tests/test_plugins/test_adjoint.py @@ -1858,7 +1858,7 @@ def test_nonlinear_warn(log_capture): ) # make the nonlinear objects to add to the JaxSimulation one by one - nl_model = td.KerrNonlinearity(n2=1) + nl_model = td.NonlinearSusceptibility(chi3=1) nl_medium = td.Medium(nonlinear_spec=td.NonlinearSpec(models=[nl_model])) struct_static_nl = struct_static.updated_copy(medium=nl_medium) input_struct_nl = JaxStructureStaticMedium(geometry=struct.geometry, medium=nl_medium) @@ -2034,17 +2034,21 @@ def test_to_gds(tmp_path): [(0, 0), (1, 0), (1, 1), (0, 1), (0, 0.9), (0, 0.11)], # notched rectangle ], ) -@pytest.mark.parametrize("subdivide", [0, 1, 5]) +@pytest.mark.parametrize("subdivide", [0, 1, 3]) @pytest.mark.parametrize("sidewall_angle_deg", [0, 10]) +@pytest.mark.parametrize("dilation", [-0.02, 0.0, 0.02]) class TestJaxComplexPolySlab: slab_bounds = (-0.25, 0.25) EPS = 1e-12 RTOL = 1e-2 @staticmethod - def objfun(vertices, slab_bounds, sidewall_angle): + def objfun(vertices, slab_bounds, sidewall_angle, dilation): p = JaxComplexPolySlab( - vertices=vertices, slab_bounds=slab_bounds, sidewall_angle=sidewall_angle + vertices=vertices, + slab_bounds=slab_bounds, + sidewall_angle=sidewall_angle, + dilation=dilation, ) obj = 0.0 for s in p.sub_polyslabs: @@ -2072,11 +2076,12 @@ def vertices(self, base_vertices, subdivide): def sidewall_angle(self, sidewall_angle_deg): return np.deg2rad(sidewall_angle_deg) - def test_matches_complexpolyslab(self, vertices, sidewall_angle): + def test_matches_complexpolyslab(self, vertices, sidewall_angle, dilation): kwargs = dict( vertices=vertices, sidewall_angle=sidewall_angle, slab_bounds=self.slab_bounds, + dilation=dilation, axis=POLYSLAB_AXIS, ) cp = ComplexPolySlab(**kwargs) @@ -2087,23 +2092,44 @@ def test_matches_complexpolyslab(self, vertices, sidewall_angle): for cps, jcps in zip(cp.sub_polyslabs, jcp.sub_polyslabs): assert_allclose(cps.vertices, jcps.vertices) - def test_vertices_grads(self, vertices, sidewall_angle): + def test_vertices_grads(self, vertices, sidewall_angle, dilation): check_grads( - lambda x: self.objfun(x, self.slab_bounds, sidewall_angle), + lambda x: self.objfun(x, self.slab_bounds, sidewall_angle, dilation), (vertices,), order=1, rtol=self.RTOL, eps=self.EPS, ) - @pytest.mark.skip(reason="No VJP implemented yet") - def test_slab_bounds_grads(self, vertices, sidewall_angle): + def test_dilation_grads(self, vertices, sidewall_angle, dilation): + if sidewall_angle != 0: + pytest.xfail("Dilation gradients only work if no sidewall angle.") + check_grads( + lambda x: self.objfun(vertices, self.slab_bounds, sidewall_angle, x), + (dilation,), + order=1, + rtol=self.RTOL, + eps=self.EPS, + ) + + def test_slab_bounds_grads(self, vertices, sidewall_angle, dilation): + if sidewall_angle != 0: + pytest.xfail("Slab bound gradients only work if no sidewall angle.") check_grads( - lambda x: self.objfun(vertices, x, sidewall_angle), (self.slab_bounds,), order=1 + lambda x: self.objfun(vertices, x, sidewall_angle, dilation), + (self.slab_bounds,), + order=1, + rtol=self.RTOL, + eps=self.EPS, ) - @pytest.mark.skip(reason="No VJP implemented yet") - def test_sidewall_angle_grads(self, vertices, sidewall_angle): + def test_sidewall_angle_grads(self, vertices, sidewall_angle, dilation): + if sidewall_angle != 0: + pytest.xfail("Sidewall gradients only work for small angles.") check_grads( - lambda x: self.objfun(vertices, self.slab_bounds, x), (sidewall_angle,), order=1 + lambda x: self.objfun(vertices, self.slab_bounds, x, dilation), + (sidewall_angle,), + order=1, + rtol=self.RTOL, + eps=self.EPS, ) diff --git a/tests/test_plugins/test_dispersion_fitter.py b/tests/test_plugins/test_dispersion_fitter.py index 518a6d150..0330637bc 100644 --- a/tests/test_plugins/test_dispersion_fitter.py +++ b/tests/test_plugins/test_dispersion_fitter.py @@ -272,3 +272,16 @@ def test_dispersion_guess(random_data): medium, rms = fitter.fit(num_tries=2) medium_new, rms_new = fitter.fit(num_tries=1, guess=medium) + + +def test_dispersion_loss_samples(): + wvls = np.array([275e-3, 260e-3, 255e-3]) + n_nAlGaN = np.array([2.72, 2.68, 2.53]) + + nAlGaN_fitter = FastDispersionFitter(wvl_um=wvls, n_data=n_nAlGaN) + nAlGaN_mat, _ = nAlGaN_fitter.fit() + + freq_list = nAlGaN_mat.angular_freq_to_Hz(nAlGaN_mat._imag_ep_extrema_with_samples()) + ep = nAlGaN_mat.eps_model(freq_list) + for e in ep: + assert e.imag >= 0 diff --git a/tests/test_plugins/test_invdes.py b/tests/test_plugins/test_invdes.py index 07813555a..b96d7e9ab 100644 --- a/tests/test_plugins/test_invdes.py +++ b/tests/test_plugins/test_invdes.py @@ -2,9 +2,16 @@ import autograd.numpy as anp import numpy as np +import numpy.testing as npt import pytest import tidy3d as td import tidy3d.plugins.invdes as tdi +from tidy3d.plugins.expressions import ModeAmp, ModePower +from tidy3d.plugins.invdes.initialization import ( + CustomInitializationSpec, + RandomInitializationSpec, + UniformInitializationSpec, +) # use single threading pipeline from ..test_components.test_autograd import use_emulated_run # noqa: F401 @@ -88,6 +95,18 @@ def test_region_params(): _ = design_region.params_zeros +def test_region_uniform(): + """Test parameter shape for uniform dimensions""" + region = make_design_region() + shape = region.params_shape + + test_region = region.updated_copy(uniform=(1, 1, 1)) + assert test_region.params_shape == (1, 1, 1) + + test_region = region.updated_copy(uniform=(1, 0, 1)) + assert test_region.params_shape == (1, *shape[1:]) + + def test_region_penalties(): """Test evaluation of penalties of a ``TopologyDesignRegion``.""" @@ -278,10 +297,9 @@ def make_result(use_emulated_run): # noqa: F811 """Test running the optimization defined in the ``InverseDesign`` object.""" optimizer = make_optimizer() + optimizer.validate_pre_upload() - PARAMS_0 = np.random.random(optimizer.design.design_region.params_shape) - - return optimizer.run(params0=PARAMS_0, post_process_fn=post_process_fn) + return optimizer.run(post_process_fn=post_process_fn) def test_default_params(use_emulated_run): # noqa: F811 @@ -289,8 +307,6 @@ def test_default_params(use_emulated_run): # noqa: F811 optimizer = make_optimizer() - _ = np.random.random(optimizer.design.design_region.params_shape) - optimizer.run(post_process_fn=post_process_fn) @@ -312,9 +328,7 @@ def make_result_multi(use_emulated_run): # noqa: F811 optimizer = optimizer.updated_copy(design=design) - PARAMS_0 = np.random.random(optimizer.design.design_region.params_shape) - - return optimizer.run(params0=PARAMS_0, post_process_fn=post_process_fn_multi) + return optimizer.run(post_process_fn=post_process_fn_multi) def test_result_store_full_results_is_false(use_emulated_run): # noqa: F811 @@ -323,9 +337,7 @@ def test_result_store_full_results_is_false(use_emulated_run): # noqa: F811 optimizer = make_optimizer() optimizer = optimizer.updated_copy(store_full_results=False, num_steps=3) - PARAMS_0 = np.random.random(optimizer.design.design_region.params_shape) - - result = optimizer.run(params0=PARAMS_0, post_process_fn=post_process_fn) + result = optimizer.run(post_process_fn=post_process_fn) # these store at the very beginning and at the end of every iteration # but when ``store_full_results == False``, they only store the last one @@ -344,12 +356,15 @@ def test_continue_run_fns(use_emulated_run): # noqa: F811 """Test continuing an already run inverse design from result.""" result_orig = make_result(use_emulated_run) optimizer = make_optimizer() - result_full = optimizer.continue_run(result=result_orig, post_process_fn=post_process_fn) + num_steps_continue = 2 + result_full = optimizer.continue_run( + result=result_orig, num_steps=num_steps_continue, post_process_fn=post_process_fn + ) num_steps_orig = len(result_orig.history["params"]) num_steps_full = len(result_full.history["params"]) assert ( - num_steps_full == num_steps_orig + optimizer.num_steps + num_steps_full == num_steps_orig + num_steps_continue ), "wrong number of elements in the combined run history." @@ -358,15 +373,23 @@ def test_continue_run_from_file(use_emulated_run): # noqa: F811 result_orig = make_result(use_emulated_run) optimizer_orig = make_optimizer() optimizer = optimizer_orig.updated_copy(num_steps=optimizer_orig.num_steps + 1) - result_full = optimizer.continue_run_from_file(HISTORY_FNAME, post_process_fn=post_process_fn) + num_steps_continue = 2 + result_full = optimizer.continue_run_from_file( + HISTORY_FNAME, num_steps=2, post_process_fn=post_process_fn + ) num_steps_orig = len(result_orig.history["params"]) - num_steps_full = len(result_full.history["params"]) + num_steps_new = len(result_full.history["params"]) assert ( - num_steps_full == num_steps_orig + optimizer.num_steps + num_steps_new == num_steps_orig + num_steps_continue ), "wrong number of elements in the combined run history." # test the convenience function to load it from file - result_full = optimizer.continue_run_from_history(post_process_fn=post_process_fn) + result_full = optimizer.continue_run_from_history(num_steps=2, post_process_fn=post_process_fn) + num_steps_orig = num_steps_new + num_steps_new = len(result_full.history["params"]) + assert ( + num_steps_new == num_steps_orig + num_steps_continue + ), "wrong number of elements in the combined run history." def test_result( @@ -383,7 +406,7 @@ def test_result( val_last1 = result.last["params"] val_last2 = result.get_last("params") - assert np.allclose(val_last1, val_last2) + npt.assert_allclose(val_last1, val_last2) result.plot_optimization() _ = result.sim_data_last(task_name="last") @@ -478,3 +501,127 @@ def test_pixel_size_warn_validator(log_capture): with AssertLogLevel(log_capture, "WARNING", contains_str="pixel_size"): invdes_multi = invdes_multi.updated_copy(design_region=region_too_coarse) + + +def test_invdes_with_metric_objective(use_emulated_run, use_emulated_to_sim_data): # noqa: F811 + """Test using a metric as an objective function in InverseDesign.""" + + # Create a metric as the objective function + metric = 2 * ModePower(monitor_name=MNT_NAME2, f=[FREQ0]) ** 2 + + invdes = tdi.InverseDesign( + simulation=simulation, + design_region=make_design_region(), + task_name="test_metric", + metric=metric, + ) + + optimizer = tdi.AdamOptimizer( + design=invdes, + learning_rate=0.2, + num_steps=1, + ) + + optimizer.run() + + +@pytest.mark.parametrize( + "spec_class, spec_kwargs, expected_shape", + [ + (RandomInitializationSpec, {"min_value": 0.0, "max_value": 1.0}, (3, 3)), + (UniformInitializationSpec, {"value": 0.5}, (2, 2)), + (CustomInitializationSpec, {"params": np.zeros((3, 3, 3))}, (3, 3, 3)), + ], +) +def test_parameter_spec(spec_class, spec_kwargs, expected_shape): + """Test the creation of parameter arrays from different InitializationSpec classes.""" + spec = spec_class(**spec_kwargs) + params = spec.create_parameters(expected_shape) + assert params.shape == expected_shape + + +def test_parameter_spec_with_inverse_design(use_emulated_run, use_emulated_to_sim_data): # noqa: F811 + """Test InitializationSpec with InverseDesign class.""" + + metric = 2 * ModePower(monitor_name=MNT_NAME2, f=[FREQ0]) ** 2 + + initialization_spec = RandomInitializationSpec() + design_region = make_design_region() + design_region = design_region.updated_copy(initialization_spec=initialization_spec) + + invdes = tdi.InverseDesign( + simulation=simulation, + design_region=design_region, + task_name="test_metric", + metric=metric, + ) + + optimizer = tdi.AdamOptimizer( + design=invdes, + learning_rate=0.2, + num_steps=1, + ) + + optimizer.run() + + +def test_initial_simulation(): + """Test the initial_simulation property for InverseDesign.""" + invdes = make_invdes() + initial_sim = invdes.initial_simulation + assert isinstance(initial_sim, td.Simulation) + assert initial_sim.structures[-1] == invdes.design_region.to_structure( + invdes.design_region.initial_parameters + ) + + +def test_initial_simulation_multi(): + """Test the initial_simulation property for InverseDesignMulti.""" + invdes_multi = make_invdes_multi() + initial_sims = invdes_multi.initial_simulation + assert isinstance(initial_sims, dict) + for sim in initial_sims.values(): + assert isinstance(sim, td.Simulation) + assert sim.structures[-1] == invdes_multi.design_region.to_structure( + invdes_multi.design_region.initial_parameters + ) + + +def test_metric_scalar_freq(): + invdes = make_invdes() + metric = ModePower(monitor_name=MNT_NAME2, mode_index=0, f=FREQ0) + monitor = mnt2.updated_copy(freqs=[FREQ0, FREQ0 / 2]) + invdes = invdes.updated_copy( + metric=metric, + simulation=simulation.updated_copy(monitors=[monitor]), + ) + + +def test_validate_invdes_metric(): + """Test the _validate_metric_monitor_name validator.""" + invdes = make_invdes() + metric = ModePower(monitor_name="invalid_monitor", f=[FREQ0]) + with pytest.raises(ValueError, match="monitors"): + invdes.updated_copy(metric=metric) + + metric = ModePower(monitor_name=MNT_NAME2, mode_index=10, f=[FREQ0]) + with pytest.raises(ValueError, match="mode index"): + invdes.updated_copy(metric=metric) + + metric = ModePower(monitor_name=MNT_NAME2, mode_index=0, f=[FREQ0 / 2]) + with pytest.raises(ValueError, match="frequencies"): + invdes.updated_copy(metric=metric) + + metric = ModePower(monitor_name=MNT_NAME2, mode_index=0) + monitor = mnt2.updated_copy(freqs=[FREQ0, FREQ0 / 2]) + invdes = invdes.updated_copy(simulation=simulation.updated_copy(monitors=[monitor])) + with pytest.raises(ValueError, match="single frequency"): + invdes.updated_copy(metric=metric) + + metric = ModeAmp(monitor_name=MNT_NAME2, mode_index=0) + ModePower( + monitor_name=MNT_NAME2, mode_index=0 + ) + monitor = mnt2.updated_copy(freqs=[FREQ0]) + invdes = invdes.updated_copy(simulation=simulation.updated_copy(monitors=[monitor])) + with pytest.raises(ValueError, match="must return a real"): + invdes.updated_copy(metric=metric) diff --git a/tests/test_plugins/test_mode_solver.py b/tests/test_plugins/test_mode_solver.py index 851b0ea1d..47357cec1 100644 --- a/tests/test_plugins/test_mode_solver.py +++ b/tests/test_plugins/test_mode_solver.py @@ -1087,3 +1087,14 @@ def test_modes_eme_sim(mock_remote_api, local): _ = msweb.run(solver.to_fdtd_mode_solver()) _ = solver.reduced_simulation_copy + + +def test_mode_small_bend_radius_fail(): + """Test that small bend radius fails.""" + + with pytest.raises(ValueError): + ms = ModeSolver( + plane=PLANE, + freqs=np.linspace(1e14, 2e14, 100), + mode_spec=td.ModeSpec(num_modes=1, bend_radius=1, bend_axis=0), + ) diff --git a/tests/utils.py b/tests/utils.py index 1449ec818..210f2a932 100644 --- a/tests/utils.py +++ b/tests/utils.py @@ -96,7 +96,7 @@ def cartesian_to_unstructured( shape = np.shape(XYZ[0]) - XYZp = XYZ.copy() + XYZp = np.array(XYZ).copy() rng = np.random.default_rng(seed=seed) x_pert = (1 - 2 * rng.random(shape)) * pert @@ -401,6 +401,11 @@ def make_custom_data(lims, unstructured): size=(8.0, 8.0, 8.0), run_time=1e-12, structures=[ + td.Structure( + geometry=td.Cylinder(length=1, center=(-1 * tracer, 0, 0), radius=tracer, axis=2), + medium=td.Medium(permittivity=1 + tracer, name="dieletric"), + name="traced_dieletric_cylinder", + ), td.Structure( geometry=td.Box(size=(1, tracer, tracer), center=(-1 * tracer, 0, 0)), medium=td.Medium(permittivity=1 + tracer, name="dieletric"), @@ -890,6 +895,8 @@ def run_emulated(simulation: td.Simulation, path=None, **kwargs) -> td.Simulatio """Emulates a simulation run.""" from scipy.ndimage.filters import gaussian_filter + x = kwargs.get("x0", 1.0) + def make_data( coords: dict, data_array_type: type, is_complex: bool = False ) -> td.components.data.data_array.DataArray: @@ -900,7 +907,7 @@ def make_data( data = (1 + 0.5j) * data if is_complex else data data = gaussian_filter(data, sigma=1.0) # smooth out the data a little so it isnt random - data_array = data_array_type(data, coords=coords) + data_array = data_array_type(x * data, coords=coords) return data_array def make_field_data(monitor: td.FieldMonitor) -> td.FieldData: @@ -1018,6 +1025,13 @@ def make_mode_data(monitor: td.ModeMonitor) -> td.ModeData: grid_expanded=simulation.discretize_monitor(monitor), ) + def make_flux_data(monitor: td.FluxMonitor) -> td.FluxData: + """make a random ModeData from a ModeMonitor.""" + + coords = dict(f=list(monitor.freqs)) + flux = make_data(coords=coords, data_array_type=td.FluxDataArray, is_complex=False) + return td.FluxData(monitor=monitor, flux=flux) + MONITOR_MAKER_MAP = { td.FieldMonitor: make_field_data, td.FieldTimeMonitor: make_field_time_data, @@ -1025,6 +1039,7 @@ def make_mode_data(monitor: td.ModeMonitor) -> td.ModeData: td.ModeMonitor: make_mode_data, td.PermittivityMonitor: make_eps_data, td.DiffractionMonitor: make_diff_data, + td.FluxMonitor: make_flux_data, } data = [MONITOR_MAKER_MAP[type(mnt)](mnt) for mnt in simulation.monitors] diff --git a/tidy3d/__init__.py b/tidy3d/__init__.py index 73fe5b04c..f428c4993 100644 --- a/tidy3d/__init__.py +++ b/tidy3d/__init__.py @@ -108,7 +108,6 @@ EMEMonitor, ) -# EME # EME from .components.eme.simulation import EMESimulation from .components.eme.sweep import EMEFreqSweep, EMELengthSweep, EMEModeSweep @@ -116,6 +115,9 @@ # field projection from .components.field_projection import FieldProjector +# frequency conversion utilities +from .components.frequencies import frequencies, wavelengths + # geometry from .components.geometry.base import Box, ClipOperation, Geometry, GeometryGroup, Transformed from .components.geometry.mesh import TriangleMesh @@ -455,6 +457,8 @@ def set_logging_level(level: str) -> None: "Q_e", "K_B", "inf", + "frequencies", + "wavelengths", "material_library", "Graphene", "AbstractMedium", diff --git a/tidy3d/components/autograd/__init__.py b/tidy3d/components/autograd/__init__.py index 98bcec266..e80b43fd4 100644 --- a/tidy3d/components/autograd/__init__.py +++ b/tidy3d/components/autograd/__init__.py @@ -1,3 +1,4 @@ +from .boxes import TidyArrayBox from .functions import interpn from .types import ( AutogradFieldMap, @@ -8,9 +9,10 @@ TracedSize1D, TracedVertices, ) -from .utils import get_static +from .utils import get_static, is_tidy_box, split_list __all__ = [ + "TidyArrayBox", "TracedFloat", "TracedSize1D", "TracedSize", @@ -20,4 +22,8 @@ "AutogradFieldMap", "get_static", "interpn", + "split_list", + "is_tidy_box", + "trapz", + "add_at", ] diff --git a/tidy3d/components/autograd/boxes.py b/tidy3d/components/autograd/boxes.py new file mode 100644 index 000000000..437005d9d --- /dev/null +++ b/tidy3d/components/autograd/boxes.py @@ -0,0 +1,159 @@ +# Adds some functionality to the autograd arraybox and related autograd patches +# NOTE: we do not subclass ArrayBox since that would break autograd's internal checks + +import importlib +from typing import Any, Callable, Dict, List, Tuple + +import autograd.numpy as anp +from autograd.extend import VJPNode, defjvp, register_notrace +from autograd.numpy.numpy_boxes import ArrayBox +from autograd.numpy.numpy_wrapper import _astype + +TidyArrayBox = ArrayBox # NOT a subclass + +_autograd_module_cache = {} # cache for imported autograd modules + +register_notrace(VJPNode, anp.full_like) + +defjvp( + _astype, + lambda g, ans, A, dtype, order="K", casting="unsafe", subok=True, copy=True: _astype(g, dtype), +) + +anp.astype = _astype +anp.permute_dims = anp.transpose + + +@classmethod +def from_arraybox(cls, box: ArrayBox) -> TidyArrayBox: + """Construct a TidyArrayBox from an ArrayBox.""" + return cls(box._value, box._trace, box._node) + + +def __array_function__( + self: Any, + func: Callable, + types: List[Any], + args: Tuple[Any, ...], + kwargs: Dict[str, Any], +) -> Any: + """ + Handle the dispatch of NumPy functions to autograd's numpy implementation. + + Parameters + ---------- + self : Any + The instance of the class. + func : Callable + The NumPy function being called. + types : List[Any] + The types of the arguments that implement __array_function__. + args : Tuple[Any, ...] + The positional arguments to the function. + kwargs : Dict[str, Any] + The keyword arguments to the function. + + Returns + ------- + Any + The result of the function call, or NotImplemented. + + Raises + ------ + NotImplementedError + If the function is not implemented in autograd.numpy. + + See Also + -------- + https://numpy.org/doc/stable/reference/arrays.classes.html#numpy.class.__array_function__ + """ + if not all(t in TidyArrayBox.type_mappings for t in types): + return NotImplemented + + module_name = func.__module__ + + if module_name.startswith("numpy"): + anp_module_name = "autograd." + module_name + else: + return NotImplemented + + # Use the cached module if available + anp_module = _autograd_module_cache.get(anp_module_name) + if anp_module is None: + try: + anp_module = importlib.import_module(anp_module_name) + _autograd_module_cache[anp_module_name] = anp_module + except ImportError: + return NotImplemented + + f = getattr(anp_module, func.__name__, None) + if f is None: + return NotImplemented + + if f.__name__ == "nanmean": # somehow xarray always dispatches to nanmean + f = anp.mean + kwargs.pop("dtype", None) # autograd mean vjp doesn't support dtype + + return f(*args, **kwargs) + + +def __array_ufunc__( + self: Any, + ufunc: Callable, + method: str, + *inputs: Any, + **kwargs: Dict[str, Any], +) -> Any: + """ + Handle the dispatch of NumPy ufuncs to autograd's numpy implementation. + + Parameters + ---------- + self : Any + The instance of the class. + ufunc : Callable + The universal function being called. + method : str + The method of the ufunc being called. + inputs : Any + The input arguments to the ufunc. + kwargs : Dict[str, Any] + The keyword arguments to the ufunc. + + Returns + ------- + Any + The result of the ufunc call, or NotImplemented. + + See Also + -------- + https://numpy.org/doc/stable/reference/arrays.classes.html#numpy.class.__array_ufunc__ + """ + if method != "__call__": + return NotImplemented + + ufunc_name = ufunc.__name__ + + anp_ufunc = getattr(anp, ufunc_name, None) + if anp_ufunc is not None: + return anp_ufunc(*inputs, **kwargs) + + return NotImplemented + + +def item(self): + if self.size != 1: + raise ValueError("Can only convert an array of size 1 to a scalar") + return anp.ravel(self)[0] + + +TidyArrayBox._tidy = True +TidyArrayBox.from_arraybox = from_arraybox +TidyArrayBox.__array_namespace__ = lambda self, *, api_version=None: anp +TidyArrayBox.__array_ufunc__ = __array_ufunc__ +TidyArrayBox.__array_function__ = __array_function__ +TidyArrayBox.__repr__ = str +TidyArrayBox.real = property(anp.real) +TidyArrayBox.imag = property(anp.imag) +TidyArrayBox.conj = anp.conj +TidyArrayBox.item = item diff --git a/tidy3d/components/autograd/derivative_utils.py b/tidy3d/components/autograd/derivative_utils.py index ee7efbc79..ef3f038a6 100644 --- a/tidy3d/components/autograd/derivative_utils.py +++ b/tidy3d/components/autograd/derivative_utils.py @@ -42,6 +42,34 @@ class DerivativeInfo(Tidy3dBaseModel): "of this dataset is used when computing adjoint gradients for shifting boundaries.", ) + E_fwd: FieldData = pd.Field( + ..., + title="Forward Electric Fields", + description='Dataset where the field components ``("Ex", "Ey", "Ez")`` represent the ' + "forward electric fields used for computing gradients for a given structure.", + ) + + E_adj: FieldData = pd.Field( + ..., + title="Adjoint Electric Fields", + description='Dataset where the field components ``("Ex", "Ey", "Ez")`` represent the ' + "adjoint electric fields used for computing gradients for a given structure.", + ) + + D_fwd: FieldData = pd.Field( + ..., + title="Forward Displacement Fields", + description='Dataset where the field components ``("Ex", "Ey", "Ez")`` represent the ' + "forward displacement fields used for computing gradients for a given structure.", + ) + + D_adj: FieldData = pd.Field( + ..., + title="Adjoint Displacement Fields", + description='Dataset where the field components ``("Ex", "Ey", "Ez")`` represent the ' + "adjoint displacement fields used for computing gradients for a given structure.", + ) + eps_data: PermittivityData = pd.Field( ..., title="Permittivity Dataset", @@ -89,16 +117,79 @@ def updated_paths(self, paths: list[PathType]) -> DerivativeInfo: """Update this ``DerivativeInfo`` with new set of paths.""" return self.updated_copy(paths=paths) + def grad_in_bases(self, spatial_coords: np.ndarray, basis_vectors: dict) -> dict: + """Get the ``D_norm``, ``E_edge`` ``E_slab`` components of the gradient contributions.""" + + # unpack electric and displacement fields + E_fwd = self.E_fwd + E_adj = self.E_adj + D_fwd = self.D_fwd + D_adj = self.D_adj + + # compute the E and D fields at the edge centers + E_fwd_at_coords = self.evaluate_flds_at(fld_dataset=E_fwd, spatial_coords=spatial_coords) + E_adj_at_coords = self.evaluate_flds_at(fld_dataset=E_adj, spatial_coords=spatial_coords) + D_fwd_at_coords = self.evaluate_flds_at(fld_dataset=D_fwd, spatial_coords=spatial_coords) + D_adj_at_coords = self.evaluate_flds_at(fld_dataset=D_adj, spatial_coords=spatial_coords) + + # project the relevant field quantities into their respective basis for gradient calculation + D_fwd_norm = self.project_in_basis(D_fwd_at_coords, basis_vector=basis_vectors["norm"]) + D_adj_norm = self.project_in_basis(D_adj_at_coords, basis_vector=basis_vectors["norm"]) + + E_fwd_perp1 = self.project_in_basis(E_fwd_at_coords, basis_vector=basis_vectors["perp1"]) + E_adj_perp1 = self.project_in_basis(E_adj_at_coords, basis_vector=basis_vectors["perp1"]) + + E_fwd_perp2 = self.project_in_basis(E_fwd_at_coords, basis_vector=basis_vectors["perp2"]) + E_adj_perp2 = self.project_in_basis(E_adj_at_coords, basis_vector=basis_vectors["perp2"]) + + # multiply forward and adjoint + D_der_norm = D_fwd_norm * D_adj_norm + E_der_perp1 = E_fwd_perp1 * E_adj_perp1 + E_der_perp2 = E_fwd_perp2 * E_adj_perp2 + + return dict(D_norm=D_der_norm, E_perp1=E_der_perp1, E_perp2=E_der_perp2) + + @staticmethod + def evaluate_flds_at( + fld_dataset: dict[str, ScalarFieldDataArray], + spatial_coords: np.ndarray, # (N, 3) + ) -> dict[str, ScalarFieldDataArray]: + """Compute the value of an dict with keys Ex, Ey, Ez at a set of spatial locations.""" + + xs, ys, zs = spatial_coords.T + edge_index_dim = "edge_index" + + interp_kwargs = {} + for dim, locations_dim in zip("xyz", (xs, ys, zs)): + # only include dims where the data has more than 1 coord, to avoid warnings and errors + if True or all(np.array(fld.coords).size > 1 for fld in fld_dataset.values()): + interp_kwargs[dim] = xr.DataArray(locations_dim, dims=edge_index_dim) + + components = {} + for fld_name, arr in fld_dataset.items(): + components[fld_name] = arr.interp(**interp_kwargs, assume_sorted=True).sum("f") + + return components + + @staticmethod + def project_in_basis( + der_dataset: xr.Dataset, + basis_vector: np.ndarray, + ) -> xr.DataArray: + """Project a derivative dataset along a supplied basis vector.""" + + value = 0.0 + for coeffs, dim in zip(basis_vector.T, "xyz"): + value += coeffs * der_dataset[f"E{dim}"] + return value + # TODO: could we move this into a DataArray method? def integrate_within_bounds(arr: xr.DataArray, dims: list[str], bounds: Bound) -> xr.DataArray: """integrate a data array within bounds, assumes bounds are [2, N] for N dims.""" - _arr = arr.copy() - # order bounds with dimension first (N, 2) - bounds = np.array(bounds).T - + bounds = np.asarray(bounds).T all_coords = {} # loop over all dimensions @@ -106,15 +197,14 @@ def integrate_within_bounds(arr: xr.DataArray, dims: list[str], bounds: Bound) - bmin = get_static(bmin) bmax = get_static(bmax) - coord_values = _arr.coords[dim].values + coord_values = np.copy(arr.coords[dim].data) # reset all coordinates outside of bounds to the bounds, so that dL = 0 in integral - coord_values[coord_values < bmin] = bmin - coord_values[coord_values > bmax] = bmax + np.clip(coord_values, bmin, bmax, out=coord_values) all_coords[dim] = coord_values - _arr = _arr.assign_coords(**all_coords) + _arr = arr.assign_coords(**all_coords) # uses trapezoidal rule # https://docs.xarray.dev/en/stable/generated/xarray.DataArray.integrate.html diff --git a/tidy3d/components/autograd/functions.py b/tidy3d/components/autograd/functions.py index cee028159..2bff330a8 100644 --- a/tidy3d/components/autograd/functions.py +++ b/tidy3d/components/autograd/functions.py @@ -2,6 +2,9 @@ import autograd.numpy as anp import numpy as np +from autograd.extend import defjvp, defvjp, primitive +from autograd.numpy.numpy_jvps import broadcast +from autograd.numpy.numpy_vjps import unbroadcast_f from numpy.typing import NDArray from scipy.interpolate import RegularGridInterpolator @@ -18,17 +21,17 @@ def _evaluate_nearest( Parameters ---------- - indices : NDArray[np.int64] + indices : np.ndarray[np.int64] Indices of the lower bounds of the grid cell containing the interpolation point. - norm_distances : NDArray[np.float64] + norm_distances : np.ndarray[np.float64] Normalized distances from the lower bounds of the grid cell to the interpolation point, for each dimension. - values : NDArray[np.float64] + values : np.ndarray[np.float64] The n-dimensional array of values to interpolate from. Returns ------- - NDArray[np.float64] + np.ndarray[np.float64] The value of the nearest neighbor to the interpolation point. """ idx_res = tuple(anp.where(yi <= 0.5, i, i + 1) for i, yi in zip(indices, norm_distances)) @@ -49,17 +52,17 @@ def _evaluate_linear( Parameters ---------- - indices : NDArray[np.int64] + indices : np.ndarray[np.int64] Indices of the lower bounds of the grid cell containing the interpolation point. - norm_distances : NDArray[np.float64] + norm_distances : np.ndarray[np.float64] Normalized distances from the lower bounds of the grid cell to the interpolation point, for each dimension. - values : NDArray[np.float64] + values : np.ndarray[np.float64] The n-dimensional array of values to interpolate from. Returns ------- - NDArray[np.float64] + np.ndarray[np.float64] The interpolated value at the desired point. """ # Create a slice object for broadcasting over trailing dimensions @@ -101,18 +104,18 @@ def interpn( Parameters ---------- - points : tuple[NDArray[np.float64], ...] + points : tuple[np.ndarray[np.float64], ...] The points defining the rectilinear grid in n dimensions. - values : NDArray[np.float64] + values : np.ndarray[np.float64] The data values on the rectilinear grid. - xi : tuple[NDArray[np.float64], ...] + xi : tuple[np.ndarray[np.float64], ...] The coordinates to sample the gridded data at. method : InterpolationType = "linear" The method of interpolation to perform. Supported are "linear" and "nearest". Returns ------- - NDArray[np.float64] + np.ndarray[np.float64] The interpolated values. Raises @@ -132,7 +135,6 @@ def interpn( raise ValueError(f"Unsupported interpolation method: {method}") itrp = RegularGridInterpolator(points, values, method=method) - grid = anp.meshgrid(*xi, indexing="ij") # Prepare the grid for interpolation # This step reshapes the grid, checks for NaNs and out-of-bounds values @@ -142,12 +144,12 @@ def interpn( # - number of dimensions # - boolean array indicating NaN positions # - (discarded) boolean array for out-of-bounds values - grid, shape, ndim, nans, _ = itrp._prepare_xi(tuple(grid)) + xi, shape, ndim, nans, _ = itrp._prepare_xi(xi) # Find the indices of the grid cells containing the interpolation points # and calculate the normalized distances (ranging from 0 at lower grid point to 1 # at upper grid point) within these cells - indices, norm_distances = itrp._find_indices(grid.T) + indices, norm_distances = itrp._find_indices(xi.T) result = interp_fn(indices, norm_distances, values) nans = anp.reshape(nans, (-1,) + (1,) * (result.ndim - 1)) @@ -155,6 +157,89 @@ def interpn( return anp.reshape(result, shape[:-1] + values.shape[ndim:]) +def trapz(y: NDArray, x: NDArray = None, dx: float = 1.0, axis: int = -1) -> float: + """ + Integrate along the given axis using the composite trapezoidal rule. + + Parameters + ---------- + y : np.ndarray + Input array to integrate. + x : np.ndarray = None + The sample points corresponding to the y values. If None, the sample points are assumed to be evenly spaced + with spacing `dx`. + dx : float = 1.0 + The spacing between sample points when `x` is None. Default is 1.0. + axis : int = -1 + The axis along which to integrate. Default is the last axis. + + Returns + ------- + float + Definite integral as approximated by the trapezoidal rule. + """ + if x is None: + d = dx + elif x.ndim == 1: + d = np.diff(x) + shape = [1] * y.ndim + shape[axis] = d.shape[0] + d = np.reshape(d, shape) + else: + d = np.diff(x, axis=axis) + + slice1 = [slice(None)] * y.ndim + slice2 = [slice(None)] * y.ndim + slice1[axis] = slice(1, None) + slice2[axis] = slice(None, -1) + + return anp.sum((y[tuple(slice1)] + y[tuple(slice2)]) * d / 2, axis=axis) + + +@primitive +def add_at(x: NDArray, indices_x: tuple, y: NDArray) -> NDArray: + """ + Add values to specified indices of an array. + + This function creates a copy of the input array `x`, adds the values from `y` to the specified + indices `indices_x`, and returns the modified array. + + Parameters + ---------- + x : np.ndarray + Input array to which values will be added. + indices_x : tuple + Indices of `x` where values from `y` will be added. + y : np.ndarray + Values to add to the specified indices of `x`. + + Returns + ------- + np.ndarray + The modified array with values added at the specified indices. + """ + out = np.copy(x) # Copy to preserve 'x' for gradient computation + out[tuple(indices_x)] += y + return out + + +defvjp( + add_at, + lambda ans, x, indices_x, y: unbroadcast_f(x, lambda g: g), + lambda ans, x, indices_x, y: lambda g: g[tuple(indices_x)], + argnums=(0, 2), +) + +defjvp( + add_at, + lambda g, ans, x, indices_x, y: broadcast(g, ans), + lambda g, ans, x, indices_x, y: add_at(anp.zeros_like(ans), indices_x, g), + argnums=(0, 2), +) + + __all__ = [ "interpn", + "trapz", + "add_at", ] diff --git a/tidy3d/components/autograd/utils.py b/tidy3d/components/autograd/utils.py index 7b120b6fe..0a1fbfd43 100644 --- a/tidy3d/components/autograd/utils.py +++ b/tidy3d/components/autograd/utils.py @@ -16,7 +16,13 @@ def split_list(x: list[typing.Any], index: int) -> (list[typing.Any], list[typin return x[:index], x[index:] +def is_tidy_box(x: typing.Any) -> bool: + """Check if a value is a tidy box.""" + return getattr(x, "_tidy", False) + + __all__ = [ "get_static", "split_list", + "is_tidy_box", ] diff --git a/tidy3d/components/base.py b/tidy3d/components/base.py index d2f547443..8c14325d2 100644 --- a/tidy3d/components/base.py +++ b/tidy3d/components/base.py @@ -26,7 +26,7 @@ from ..log import log from .autograd.types import AutogradFieldMap, Box from .autograd.utils import get_static -from .data.data_array import AUTOGRAD_KEY, DATA_ARRAY_MAP, DataArray +from .data.data_array import DATA_ARRAY_MAP, DataArray from .file_util import compress_file_to_gzip, extract_gzip_file from .types import TYPE_TAG_STR, ComplexNumber, Literal @@ -91,22 +91,28 @@ def _get_valid_extension(fname: str) -> str: ) -def skip_if_fields_missing(fields: List[str]): +def skip_if_fields_missing(fields: List[str], root=False): """Decorate ``validator`` to check that other fields have passed validation.""" def actual_decorator(validator): @wraps(validator) - def _validator(cls, val, values): + def _validator(cls, *args, **kwargs): """New validator function.""" + values = kwargs.get("values") + if values is None: + values = args[0] if root else args[1] for field in fields: if field not in values: log.warning( f"Could not execute validator '{validator.__name__}' because field " f"'{field}' failed validation." ) - return val + if root: + return values + else: + return kwargs.get("val") if "val" in kwargs.keys() else args[0] - return validator(cls, val, values) + return validator(cls, *args, **kwargs) return _validator @@ -965,27 +971,19 @@ def handle_value(x: Any, path: tuple[str, ...]) -> None: if isbox(x): field_mapping[path] = x - # for data arrays, need to be more careful as their tracers are stored in attrs - elif isinstance(x, DataArray): - # try to grab the traced values from the `attrs` (if traced) - if AUTOGRAD_KEY in x.attrs: - field_mapping[path] = x.attrs[AUTOGRAD_KEY] - - # or just grab the static value out of the values - elif include_untraced_data_arrays: - field_mapping[path] = get_static(x.values) + # for data arrays, need to be more careful as their tracers are stored in .data + elif isinstance(x, xr.DataArray) and (isbox(x.data) or include_untraced_data_arrays): + field_mapping[path] = x.data # for sequences, add (i,) to the path and handle each value individually elif isinstance(x, (list, tuple)): for i, val in enumerate(x): - sub_paths = path + (i,) - handle_value(val, path=sub_paths) + handle_value(val, path=path + (i,)) # for dictionaries, add the (key,) to the path and handle each value individually elif isinstance(x, dict): for key, val in x.items(): - sub_paths = path + (key,) - handle_value(val, path=sub_paths) + handle_value(val, path=path + (key,)) # recursively parse the dictionary of this object self_dict = self.dict() @@ -1018,13 +1016,8 @@ def insert_value(x, path: tuple[str, ...], sub_dict: dict): current_dict[final_key] = list(current_dict[final_key]) sub_element = current_dict[final_key] - if isinstance(sub_element, DataArray): + if isinstance(sub_element, xr.DataArray): current_dict[final_key] = sub_element.copy(deep=False, data=x) - if isbox(x): - current_dict[final_key].attrs[AUTOGRAD_KEY] = x - - elif AUTOGRAD_KEY in current_dict[final_key].attrs: - current_dict[final_key].attrs.pop(AUTOGRAD_KEY) else: current_dict[final_key] = x diff --git a/tidy3d/components/base_sim/simulation.py b/tidy3d/components/base_sim/simulation.py index b6c782f57..ea2fe6c9e 100644 --- a/tidy3d/components/base_sim/simulation.py +++ b/tidy3d/components/base_sim/simulation.py @@ -163,6 +163,7 @@ def _structures_not_at_edges(cls, val, values): "use td.inf as a size variable instead to make this explicit.", custom_loc=["structures", istruct], ) + continue return val diff --git a/tidy3d/components/boundary.py b/tidy3d/components/boundary.py index 76441f9f3..923da4ea8 100644 --- a/tidy3d/components/boundary.py +++ b/tidy3d/components/boundary.py @@ -16,6 +16,8 @@ from .source import TFSF, GaussianBeam, ModeSource, PlaneWave from .types import TYPE_TAG_STR, Axis, Complex +MIN_NUM_PML_LAYERS = 6 + class BoundaryEdge(ABC, Tidy3dBaseModel): """Electromagnetic boundary condition at a domain edge.""" @@ -260,10 +262,11 @@ class PMLParams(AbsorberParams): class AbsorberSpec(BoundaryEdge): """Specifies the generic absorber properties along a single dimension.""" - num_layers: pd.NonNegativeInt = pd.Field( + num_layers: int = pd.Field( ..., title="Number of Layers", description="Number of layers of standard PML.", + ge=MIN_NUM_PML_LAYERS, ) parameters: AbsorberParams = pd.Field( ..., @@ -376,10 +379,11 @@ class PML(AbsorberSpec): """ - num_layers: pd.NonNegativeInt = pd.Field( + num_layers: int = pd.Field( 12, title="Number of Layers", description="Number of layers of standard PML.", + ge=MIN_NUM_PML_LAYERS, ) parameters: PMLParams = pd.Field( @@ -413,8 +417,11 @@ class StablePML(AbsorberSpec): * `Introduction to perfectly matched layer (PML) tutorial `__ """ - num_layers: pd.NonNegativeInt = pd.Field( - 40, title="Number of Layers", description="Number of layers of 'stable' PML." + num_layers: int = pd.Field( + 40, + title="Number of Layers", + description="Number of layers of 'stable' PML.", + ge=MIN_NUM_PML_LAYERS, ) parameters: PMLParams = pd.Field( @@ -463,10 +470,11 @@ class Absorber(AbsorberSpec): * `How to troubleshoot a diverged FDTD simulation <../../notebooks/DivergedFDTDSimulation.html>`_ """ - num_layers: pd.NonNegativeInt = pd.Field( + num_layers: int = pd.Field( 40, title="Number of Layers", description="Number of layers of absorber to add to + and - boundaries.", + ge=MIN_NUM_PML_LAYERS, ) parameters: AbsorberParams = pd.Field( diff --git a/tidy3d/components/data/data_array.py b/tidy3d/components/data/data_array.py index 5fead6825..6a90ab7d0 100644 --- a/tidy3d/components/data/data_array.py +++ b/tidy3d/components/data/data_array.py @@ -11,9 +11,13 @@ import numpy as np import pandas import xarray as xr -from autograd.tracer import Box, getval, isbox -from xarray.core.types import InterpOptions -from xarray.core.utils import either_dict_or_kwargs +from autograd.tracer import isbox +from xarray.core import alignment, missing +from xarray.core.indexes import PandasIndex +from xarray.core.indexing import _outer_to_numpy_indexer +from xarray.core.types import InterpOptions, Self +from xarray.core.utils import OrderedSet, either_dict_or_kwargs +from xarray.core.variable import as_variable from ...constants import ( HERTZ, @@ -24,7 +28,7 @@ WATT, ) from ...exceptions import DataError, FileError -from ..autograd.functions import interpn +from ..autograd import TidyArrayBox, get_static, interpn, is_tidy_box from ..types import Axis, Bound # maps the dimension names to their attributes @@ -56,9 +60,6 @@ # name of the DataArray.values in the hdf5 file (xarray's default name too) DATA_ARRAY_VALUE_NAME = "__xarray_dataarray_variable__" -# name for the autograd-traced part of the DataArray -AUTOGRAD_KEY = "AUTOGRAD" - class DataArray(xr.DataArray): """Subclass of ``xr.DataArray`` that requires _dims to match the keys of the coords.""" @@ -71,25 +72,17 @@ class DataArray(xr.DataArray): _data_attrs: Dict[str, str] = {} def __init__(self, data, *args, **kwargs): - """Initialize ``DataArray``.""" - - # initialize with untraced data - super().__init__(getval(data), *args, **kwargs) - # and put tracers in .attrs - if isbox(data): - self.attrs[AUTOGRAD_KEY] = data - - @property - def tracers(self) -> Box: - if self.data.size == 0: - return None - elif AUTOGRAD_KEY not in self.attrs and not isbox(self.data.flat[0]): - # no tracers - return None - elif isbox(self.data.flat[0]): # traced values take precedence over traced attrs - return anp.array(self.values.tolist()) - else: - return self.attrs[AUTOGRAD_KEY] + # if data is a vanilla autograd box, convert to our box + if isbox(data) and not is_tidy_box(data): + data = TidyArrayBox.from_arraybox(data) + # do the same for xr.Variable or xr.DataArray type + elif ( + isinstance(data, (xr.Variable, xr.DataArray)) + and isbox(data.data) + and not is_tidy_box(data.data) + ): + data.data = TidyArrayBox.from_arraybox(data.data) + super().__init__(data, *args, **kwargs) @classmethod def __get_validators__(cls): @@ -133,11 +126,6 @@ def _interp_validator(self, field_name: str = None) -> None: This does not check every 'DataArray' by default. Instead, when required, this check can be called from a validator, as is the case with 'CustomMedium' and 'CustomFieldSource'. """ - # skip this validator if currently tracing for autograd because - # self.values will be dtype('object') and not interpolatable - if self.tracers is not None: - return - if field_name is None: field_name = "DataArray" @@ -221,6 +209,22 @@ def __eq__(self, other) -> bool: return False return True + @property + def values(self): + """ + The array's data converted to a numpy.ndarray. + + Returns + ------- + np.ndarray + The values of the DataArray. + """ + return self.data if isbox(self.data) else super().values + + @values.setter + def values(self, value: Any) -> None: + self.variable.values = value + @property def abs(self): """Absolute value of data array.""" @@ -248,7 +252,7 @@ def to_hdf5_handle(self, f_handle: h5py.File, group_path: str) -> None: """Save an xr.DataArray to the hdf5 file handle with a given path to the group.""" sub_group = f_handle.create_group(group_path) - sub_group[DATA_ARRAY_VALUE_NAME] = self.values + sub_group[DATA_ARRAY_VALUE_NAME] = get_static(self.data) for key, val in self.coords.items(): if val.dtype == " None: sub_group[key] = val @classmethod - def from_hdf5(cls, fname: str, group_path: str) -> DataArray: + def from_hdf5(cls, fname: str, group_path: str) -> Self: """Load an DataArray from an hdf5 file with a given path to the group.""" with h5py.File(fname, "r") as f: sub_group = f[group_path] @@ -268,7 +272,7 @@ def from_hdf5(cls, fname: str, group_path: str) -> DataArray: return cls(values, coords=coords, dims=cls._dims) @classmethod - def from_file(cls, fname: str, group_path: str) -> DataArray: + def from_file(cls, fname: str, group_path: str) -> Self: """Load an DataArray from an hdf5 file with a given path to the group.""" if ".hdf5" not in fname: raise FileError( @@ -282,20 +286,32 @@ def __hash__(self) -> int: token_str = dask.base.tokenize(self) return hash(token_str) - def multiply_at(self, value: complex, coord_name: str, indices: List[int]) -> DataArray: - """Multiply self by value at indices into .""" + def multiply_at(self, value: complex, coord_name: str, indices: List[int]) -> Self: + """Multiply self by value at indices.""" + if isbox(self.data) or isbox(value): + return self._ag_multiply_at(value, coord_name, indices) + self_mult = self.copy() self_mult[{coord_name: indices}] *= value return self_mult + def _ag_multiply_at(self, value: complex, coord_name: str, indices: List[int]) -> Self: + """Autograd multiply_at override when tracing.""" + key = {coord_name: indices} + _, index_tuple, _ = self.variable._broadcast_indexes(key) + idx = _outer_to_numpy_indexer(index_tuple, self.data.shape) + mask = np.zeros(self.data.shape, dtype="?") + mask[idx] = True + return self.copy(deep=False, data=anp.where(mask, self.data * value, self.data)) + def interp( self, - coords: Union[Mapping[Any, Any], None] = None, + coords: Mapping[Any, Any] | None = None, method: InterpOptions = "linear", assume_sorted: bool = False, - kwargs: Union[Mapping[str, Any], None] = None, + kwargs: Mapping[str, Any] | None = None, **coords_kwargs: Any, - ): + ) -> Self: """Interpolate this DataArray to new coordinate values. Parameters @@ -321,47 +337,181 @@ def interp( KeyError If any of the specified coordinates are not in the DataArray. """ - if self.tracers is not None: # use custom interp if using traced data - coords = either_dict_or_kwargs(coords, coords_kwargs, "interp") + if isbox(self.data): + return self._ag_interp(coords, method, assume_sorted, kwargs, **coords_kwargs) - missing_keys = set(coords) - set(self.coords) - if missing_keys: - raise KeyError(f"Cannot interpolate: {missing_keys} not in coords.") + return super().interp(coords, method, assume_sorted, kwargs, **coords_kwargs) - obj = self if assume_sorted else self.sortby(list(coords.keys())) + def _ag_interp( + self, + coords: Union[Mapping[Any, Any], None] = None, + method: InterpOptions = "linear", + assume_sorted: bool = False, + kwargs: Union[Mapping[str, Any], None] = None, + **coords_kwargs: Any, + ) -> Self: + """Autograd interp override when tracing over self.data. - out_coords = {k: coords.get(k, obj.coords[k]) for k in obj.dims} - points = tuple(obj.coords[k] for k in obj.dims) - xi = tuple(out_coords.values()) + This implementation closely follows the interp implementation of xarray + to match its behavior as closely as possible while supporting autograd. - vals = interpn(points, obj.tracers, xi, method=method) + See: + - https://docs.xarray.dev/en/latest/generated/xarray.DataArray.interp.html + - https://docs.xarray.dev/en/latest/generated/xarray.Dataset.interp.html + """ + if kwargs is None: + kwargs = {} - da = DataArray(vals, out_coords) # tracers go into .attrs - if isbox(self.values.flat[0]): # if tracing .values instead of .attrs - da = da.copy(deep=False, data=vals) # copy over tracers + ds = self._to_temp_dataset() - return da + coords = either_dict_or_kwargs(coords, coords_kwargs, "interp") + indexers = dict(ds._validate_interp_indexers(coords)) - return super().interp( - coords=coords, - method=method, - assume_sorted=assume_sorted, - kwargs=kwargs, - **coords_kwargs, - ) + if coords: + # Find shared dimensions between the dataset and the indexers + sdims = ( + set(ds.dims) + .intersection(*[set(nx.dims) for nx in indexers.values()]) + .difference(coords.keys()) + ) + indexers.update({d: ds.variables[d] for d in sdims}) + + obj = ds if assume_sorted else ds.sortby(list(coords)) + + # workaround to get a variable for a dimension without a coordinate + validated_indexers = { + k: (obj._variables.get(k, as_variable((k, range(obj.sizes[k])))), v) + for k, v in indexers.items() + } + + for k, v in validated_indexers.items(): + obj, newidx = missing._localize(obj, {k: v}) + validated_indexers[k] = newidx[k] + + variables = {} + reindex = False + for name, var in obj._variables.items(): + if name in indexers: + continue + dtype_kind = var.dtype.kind + if dtype_kind in "uifc": + # Interpolation for numeric types + var_indexers = {k: v for k, v in validated_indexers.items() if k in var.dims} + variables[name] = self._ag_interp_func(var, var_indexers, method, **kwargs) + elif dtype_kind in "ObU" and (validated_indexers.keys() & var.dims): + # Stepwise interpolation for non-numeric types + reindex = True + elif all(d not in indexers for d in var.dims): + # Keep variables not dependent on interpolated coords + variables[name] = var + + if reindex: + # Reindex for non-numeric types + reindex_indexers = {k: v for k, (_, v) in validated_indexers.items() if v.dims == (k,)} + reindexed = alignment.reindex( + obj, + indexers=reindex_indexers, + method="nearest", + exclude_vars=variables.keys(), + ) + indexes = dict(reindexed._indexes) + variables.update(reindexed.variables) + else: + # Get the indexes that are not being interpolated along + indexes = {k: v for k, v in obj._indexes.items() if k not in indexers} + + # Get the coords that also exist in the variables + coord_names = obj._coord_names & variables.keys() + selected = ds._replace_with_new_dims(variables.copy(), coord_names, indexes=indexes) + + # Attach indexer as coordinate + for k, v in indexers.items(): + if v.dims == (k,): + index = PandasIndex(v, k, coord_dtype=v.dtype) + index_vars = index.create_variables({k: v}) + indexes[k] = index + variables.update(index_vars) + else: + variables[k] = v - def conj(self, *args: Any, **kwargs: Any): - """Return the complex conjugate of this DataArray.""" - if self.tracers is not None: - return self.__array_wrap__(anp.conj(self.tracers)) - return super().conj(*args, **kwargs) + # Extract coordinates from indexers + coord_vars, new_indexes = selected._get_indexers_coords_and_indexes(coords) + variables.update(coord_vars) + indexes.update(new_indexes) - @property - def real(self): - """Return the real part of this DataArray.""" - if self.tracers is not None: - return self.__array_wrap__(anp.real(self.tracers)) - return super().real + coord_names = obj._coord_names & variables.keys() | coord_vars.keys() + ds = ds._replace_with_new_dims(variables, coord_names, indexes=indexes) + return self._from_temp_dataset(ds) + + @staticmethod + def _ag_interp_func(var, indexes_coords, method, **kwargs): + """ + Interpolate the variable `var` along the coordinates specified in `indexes_coords` using the given `method`. + + The implementation follows xarray's interp implementation in xarray.core.missing, + but replaces some of the pre-processing as well as the actual interpolation + function with an autograd-compatible approach. + + + Parameters + ---------- + var : xr.Variable + The variable to be interpolated. + indexes_coords : dict + A dictionary mapping dimension names to coordinate values for interpolation. + method : str + The interpolation method to use. + **kwargs : dict + Additional keyword arguments to pass to the interpolation function. + + Returns + ------- + xr.Variable + The interpolated variable. + """ + if not indexes_coords: + return var.copy() + result = var + for indep_indexes_coords in missing.decompose_interp(indexes_coords): + var = result + + # target dimensions + dims = list(indep_indexes_coords) + x, new_x = zip(*[indep_indexes_coords[d] for d in dims]) + destination = missing.broadcast_variables(*new_x) + + broadcast_dims = [d for d in var.dims if d not in dims] + original_dims = broadcast_dims + dims + new_dims = broadcast_dims + list(destination[0].dims) + + x, new_x = missing._floatize_x(x, new_x) + + permutation = [var.dims.index(dim) for dim in original_dims] + combined_permutation = permutation[-len(x) :] + permutation[: -len(x)] + data = anp.transpose(var.data, combined_permutation) + xi = anp.stack([anp.ravel(new_xi.data) for new_xi in new_x], axis=-1) + + result = interpn( + [xn.data for xn in x], + data, + xi, + method=method, + ) + + result = anp.moveaxis(result, 0, -1) + result = anp.reshape(result, result.shape[:-1] + new_x[0].shape) + + result = xr.Variable(new_dims, result, attrs=var.attrs, fastpath=True) + + out_dims: OrderedSet = OrderedSet() + for d in var.dims: + if d in dims: + out_dims.update(indep_indexes_coords[d][1].dims) + else: + out_dims.add(d) + if len(out_dims) > 1: + result = result.transpose(*out_dims) + return result class FreqDataArray(DataArray): diff --git a/tidy3d/components/data/monitor_data.py b/tidy3d/components/data/monitor_data.py index d707cc00e..0885c49df 100644 --- a/tidy3d/components/data/monitor_data.py +++ b/tidy3d/components/data/monitor_data.py @@ -10,13 +10,13 @@ import pydantic.v1 as pd import xarray as xr from pandas import DataFrame +from xarray.core.types import Self from ...constants import C_0, ETA_0, MICROMETER, MU_0 from ...exceptions import DataError, SetupError, Tidy3dNotImplementedError, ValidationError from ...log import log from ..base import TYPE_TAG_STR, cached_property, skip_if_fields_missing from ..base_sim.data.monitor_data import AbstractMonitorData -from ..geometry.base import Box from ..grid.grid import Coords, Grid from ..medium import Medium, MediumType from ..monitor import ( @@ -426,7 +426,7 @@ def _plane_grid_centers(self) -> Tuple[Coords1D, Coords1D]: return [(bs[1:] + bs[:-1]) / 2 for bs in self._plane_grid_boundaries] @property - def _diff_area(self) -> xr.DataArray: + def _diff_area(self) -> DataArray: """For a 2D monitor data, return the area of each cell in the plane, for use in numerical integrations. This assumes that data is colocated to grid boundaries, and uses the difference in the surrounding grid centers to compute the area. @@ -463,7 +463,7 @@ def _diff_area(self) -> xr.DataArray: sizes_dim0 = coords[0][1:] - coords[0][:-1] if bounds[0].size > 1 else [1.0] sizes_dim1 = coords[1][1:] - coords[1][:-1] if bounds[1].size > 1 else [1.0] - return xr.DataArray(np.outer(sizes_dim0, sizes_dim1), dims=self._tangential_dims) + return DataArray(np.outer(sizes_dim0, sizes_dim1), dims=self._tangential_dims) def _tangential_corrected(self, fields: Dict[str, DataArray]) -> Dict[str, DataArray]: """For a 2D monitor data, extract the tangential components from fields and orient them @@ -603,7 +603,7 @@ def poynting(self) -> ScalarFieldDataArray: return poynting - def package_flux_results(self, flux_values: xr.DataArray) -> Any: + def package_flux_results(self, flux_values: DataArray) -> Any: """How to package flux""" return FluxDataArray(flux_values) @@ -855,10 +855,10 @@ def _outer_fn_summation( outer_dim_2: str, sum_dims: List[str], fn: Callable, - ) -> xr.DataArray: + ) -> DataArray: """ Loop over ``outer_dim_1`` and ``outer_dim_2``, apply ``fn`` to ``fields_1`` and ``fields_2``, and sum over ``sum_dims``. - The resulting ``xr.DataArray`` has has dimensions any dimensions in the fields which are not contained in sum_dims. + The resulting ``DataArray`` has has dimensions any dimensions in the fields which are not contained in sum_dims. This can be more memory efficient than vectorizing over the ``outer_dims``, which can involve broadcasting and reshaping data. It also converts to numpy arrays outside the loops to minimize xarray overhead. """ @@ -910,7 +910,7 @@ def _outer_fn_summation( data_curr = np.sum(summand_curr, axis=tuple(sum_axes)) data[tuple(idx_data)] = data_curr - return xr.DataArray(data, coords=coords) + return DataArray(data, coords=coords) @property def time_reversed_copy(self) -> FieldData: @@ -1041,10 +1041,13 @@ def to_adjoint_point_sources(self, fwidth: float) -> List[PointDipole]: for freq0 in field_component.coords["f"]: omega0 = 2 * np.pi * freq0 - scaling_factor = 1 / (MU_0 * omega0) + scaling_factor = 33 / (MU_0 * omega0) forward_amp = self.get_amplitude(field_component.sel(f=freq0)) + if forward_amp == 0.0: + continue + adj_phase = np.pi + np.angle(forward_amp) adj_amp = scaling_factor * forward_amp @@ -1068,47 +1071,37 @@ def to_adjoint_field_sources(self, fwidth: float) -> List[CustomCurrentSource]: """Create adjoint custom field sources if this field data has some dimensionality.""" sources = [] + source_geo = self.monitor.geometry + freqs = self.monitor.freqs - # Define source geometry based on coordinates in the data - data_mins = [] - data_maxs = [] - - def shift_value(coords) -> float: - """How much to shift the geometry by along a dimension (only if > 1D).""" - return SHIFT_VALUE_ADJ_FLD_SRC if len(coords) > 1 else 0 - - for _, field_component in self.field_components.items(): - coords = field_component.coords - data_mins.append({key: min(val) + shift_value(val) for key, val in coords.items()}) - data_maxs.append({key: max(val) + shift_value(val) for key, val in coords.items()}) - - rmin = [] - rmax = [] - for dim in "xyz": - rmin.append(max(val[dim] for val in data_mins)) - rmax.append(min(val[dim] for val in data_maxs)) - - source_geo = Box.from_bounds(rmin=rmin, rmax=rmax) - - # Define source dataset - # Offset coordinates by source center since local coords are assumed in CustomCurrentSource - - for freq0 in tuple(self.field_components.values())[0].coords["f"]: + for freq0 in freqs: src_field_components = {} for name, field_component in self.field_components.items(): + # get the VJP values at frequency and apply adjoint phase field_component = field_component.sel(f=freq0) - forward_amps = field_component.values - values = -1j * forward_amps + values = 2 * -1j * field_component.values + + # make source go backwards + if "H" in name: + values *= -1 + + # make coords that are shifted relative to geometry (0,0,0) = geometry.center coords = dict(field_component.coords.copy()) for dim, key in enumerate("xyz"): coords[key] = np.array(coords[key]) - source_geo.center[dim] coords["f"] = np.array([freq0]) values = np.expand_dims(values, axis=-1) + + # ignore zero components if not np.all(values == 0): src_field_components[name] = ScalarFieldDataArray(values, coords=coords) - dataset = FieldDataset(**src_field_components) + # dont include this source if no data + if all(fld_cmp is None for fld_cmp in src_field_components.values()): + continue + # construct custom Current source + dataset = FieldDataset(**src_field_components) custom_source = CustomCurrentSource( center=source_geo.center, size=source_geo.size, @@ -1607,7 +1600,7 @@ def time_reversed_copy(self) -> FieldData: new_data["monitor"] = mnt.updated_copy(store_fields_direction=new_dir) return self.copy(update=new_data) - def _colocated_propagation_axes_field(self, field_name: Literal["E", "H"]) -> xr.DataArray: + def _colocated_propagation_axes_field(self, field_name: Literal["E", "H"]) -> DataArray: """Collect a field DataArray containing all 3 field components and rotate from frame with normal axis along z to frame with propagation axis along z. """ @@ -1633,7 +1626,7 @@ def _colocated_propagation_axes_field(self, field_name: Literal["E", "H"]) -> xr for dim in fields["Ex"].dims: coords.update({dim: fields["Ex"].coords[dim]}) - return xr.DataArray(data=field, coords=coords) + return DataArray(data=field, coords=coords) @cached_property def pol_fraction(self) -> xr.Dataset: @@ -1767,7 +1760,7 @@ def make_adjoint_sources(self, dataset_names: list[str], fwidth: float) -> list[ for name in dataset_names: if name == "amps": adjoint_sources += self.make_adjoint_sources_amps(fwidth=fwidth) - else: + elif not np.all(self.n_complex.values == 0.0): log.warning( f"Can't create adjoint source for 'ModeData.{type(self)}.{name}'. " f"for monitor '{self.monitor.name}'. " @@ -1947,6 +1940,24 @@ class FluxData(MonitorData): ..., title="Flux", description="Flux values in the frequency-domain." ) + def make_adjoint_sources( + self, dataset_names: list[str], fwidth: float + ) -> List[Union[CustomCurrentSource, PointDipole]]: + """Converts a :class:`.FieldData` to a list of adjoint current or point sources.""" + + # avoids error in edge case where there are extraneous flux monitors not used in objective + if np.all(self.flux.values == 0.0): + return [] + + raise NotImplementedError( + "Could not formulate adjoint source for 'FluxMonitor' output. To compute derivatives " + "with respect to flux data, please use a 'FieldMonitor' and call '.flux' on the " + "resulting 'FieldData' object. Using 'FluxMonitor' directly is not supported as " + "the full field information is required to construct the adjoint source for this " + "problem. The 'FluxData' does not contain the information necessary for gradient " + "computation." + ) + def normalize(self, source_spectrum_fn) -> FluxData: """Return copy of self after normalization is applied using source spectrum function.""" source_freq_amps = source_spectrum_fn(self.flux.f) @@ -2142,9 +2153,9 @@ def dims(self) -> Tuple[str, ...]: """Dimensions of the radiation vectors contained.""" return self.Etheta.dims - def make_data_array(self, data: np.ndarray) -> xr.DataArray: - """Make an xr.DataArray with data and same coords and dims as fields of self.""" - return xr.DataArray(data=data, coords=self.coords, dims=self.dims) + def make_data_array(self, data: np.ndarray) -> DataArray: + """Make an DataArray with data and same coords and dims as fields of self.""" + return DataArray(data=data, coords=self.coords, dims=self.dims) def make_dataset(self, keys: Tuple[str, ...], vals: Tuple[np.ndarray, ...]) -> xr.Dataset: """Make an xr.Dataset with keys and data with same coords and dims as fields.""" @@ -2258,7 +2269,7 @@ def fields_cartesian(self) -> xr.Dataset: return self.make_dataset(keys=keys, vals=field_components) @property - def power(self) -> xr.DataArray: + def power(self) -> DataArray: """Get power measured on the projection grid relative to the monitor's local origin. Returns @@ -2266,14 +2277,14 @@ def power(self) -> xr.DataArray: ``xarray.DataArray`` Power at points relative to the local origin. """ - power_theta = 0.5 * np.real(self.Etheta.values * np.conj(self.Hphi.values)) - power_phi = 0.5 * np.real(-self.Ephi.values * np.conj(self.Htheta.values)) + power_theta = 0.5 * np.real(self.Etheta * self.Hphi.conj()) + power_phi = 0.5 * np.real(-self.Ephi * self.Htheta.conj()) power = power_theta + power_phi return self.make_data_array(data=power) @property - def radar_cross_section(self) -> xr.DataArray: + def radar_cross_section(self) -> DataArray: """Radar cross section in units of incident power.""" _, index_k = self.nk @@ -2513,7 +2524,7 @@ def tangential_dims(self): return tangential_dims @property - def poynting(self) -> xr.DataArray: + def poynting(self) -> DataArray: """Time-averaged Poynting vector for field data associated to a Cartesian field projection monitor.""" fc = self.fields_cartesian dim1, dim2 = self.tangential_dims @@ -2857,17 +2868,17 @@ def uy(self) -> np.ndarray: ) @property - def angles(self) -> Tuple[xr.DataArray]: + def angles(self) -> Tuple[DataArray]: """The (theta, phi) angles corresponding to each allowed pair of diffraction orders storeds as data arrays. Disallowed angles are set to ``np.nan``. """ thetas, phis = self.compute_angles(self.reciprocal_vectors) - theta_data = xr.DataArray(thetas, coords=self.coords) - phi_data = xr.DataArray(phis, coords=self.coords) + theta_data = DataArray(thetas, coords=self.coords) + phi_data = DataArray(phis, coords=self.coords) return theta_data, phi_data @property - def amps(self) -> xr.DataArray: + def amps(self) -> DataArray: """Complex power amplitude in each order for 's' and 'p' polarizations, normalized so that the power carried by the wave of that order and polarization equals ``abs(amps)^2``. """ @@ -2886,10 +2897,10 @@ def amps(self) -> xr.DataArray: coords["orders_y"] = np.atleast_1d(self.orders_y) coords["f"] = np.atleast_1d(self.f) coords["polarization"] = ["s", "p"] - return xr.DataArray(np.stack([amp_phi, amp_theta], axis=3), coords=coords) + return DataArray(np.stack([amp_phi, amp_theta], axis=3), coords=coords) @property - def power(self) -> xr.DataArray: + def power(self) -> Self: """Total power in each order, summed over both polarizations.""" return (np.abs(self.amps) ** 2).sum(dim="polarization") @@ -2944,7 +2955,7 @@ def _make_dataset(self, fields: Tuple[np.ndarray, ...], keys: Tuple[str, ...]) - """Make an xr.Dataset for fields with given field names.""" data_arrays = [] for field in fields: - data_arrays.append(xr.DataArray(data=field, coords=self.coords, dims=self.dims)) + data_arrays.append(DataArray(data=field, coords=self.coords, dims=self.dims)) return xr.Dataset(dict(zip(keys, data_arrays))) """ Autograd code """ diff --git a/tidy3d/components/data/sim_data.py b/tidy3d/components/data/sim_data.py index 8e90f879e..e3a07ae7e 100644 --- a/tidy3d/components/data/sim_data.py +++ b/tidy3d/components/data/sim_data.py @@ -46,7 +46,7 @@ class AdjointSourceInfo(Tidy3dBaseModel): """Stores information about the adjoint sources to pass to autograd pipeline.""" - sources: tuple[SourceType, ...] = pd.Field( + sources: Tuple[annotate_type(SourceType), ...] = pd.Field( ..., title="Adjoint Sources", description="Set of processed sources to include in the adjoint simulation.", @@ -447,6 +447,7 @@ def plot_field_monitor_data( vmin: float = None, vmax: float = None, ax: Ax = None, + shading: str = "flat", **sel_kwargs, ) -> Ax: """Plot the field data for a monitor with simulation plot overlaid. @@ -481,6 +482,8 @@ def plot_field_monitor_data( inferred from the data and other keyword arguments. ax : matplotlib.axes._subplots.Axes = None matplotlib axes to plot on, if not specified, one is created. + shading: str = 'flat' + Shading argument for Xarray plot method ('flat','nearest','goraud') sel_kwargs : keyword arguments used to perform ``.sel()`` selection in the monitor data. These kwargs can select over the spatial dimensions (``x``, ``y``, ``z``), frequency or time dimensions (``f``, ``t``) or ``mode_index``, if applicable. @@ -493,7 +496,6 @@ def plot_field_monitor_data( matplotlib.axes._subplots.Axes The supplied or created matplotlib axes. """ - # get the DataArray corresponding to the monitor_name and field_name # deprecated intensity if field_name == "int": @@ -640,6 +642,8 @@ def plot_field_monitor_data( vmax=vmax, cmap_type=cmap_type, ax=ax, + shading=shading, + infer_intervals=True if shading == "flat" else False, ) def plot_field( @@ -654,6 +658,7 @@ def plot_field( vmin: float = None, vmax: float = None, ax: Ax = None, + shading: str = "flat", **sel_kwargs, ) -> Ax: """Plot the field data for a monitor with simulation plot overlaid. @@ -689,6 +694,8 @@ def plot_field( inferred from the data and other keyword arguments. ax : matplotlib.axes._subplots.Axes = None matplotlib axes to plot on, if not specified, one is created. + shading: str = 'flat' + Shading argument for Xarray plot method ('flat','nearest','goraud') sel_kwargs : keyword arguments used to perform ``.sel()`` selection in the monitor data. These kwargs can select over the spatial dimensions (``x``, ``y``, ``z``), frequency or time dimensions (``f``, ``t``) or ``mode_index``, if applicable. @@ -703,6 +710,7 @@ def plot_field( """ field_monitor_data = self.load_field_monitor(field_monitor_name) + return self.plot_field_monitor_data( field_monitor_data=field_monitor_data, field_name=field_name, @@ -714,6 +722,7 @@ def plot_field( vmin=vmin, vmax=vmax, ax=ax, + shading=shading, **sel_kwargs, ) @@ -731,6 +740,7 @@ def plot_scalar_array( vmax: float = None, cmap_type: ColormapType = "divergent", ax: Ax = None, + **kwargs, ) -> Ax: """Plot the field data for a monitor with simulation plot overlaid. @@ -763,6 +773,7 @@ def plot_scalar_array( Type of color map to use for plotting. ax : matplotlib.axes._subplots.Axes = None matplotlib axes to plot on, if not specified, one is created. + **kwargs : Extra arguments to ``DataArray.plot``. Returns ------- @@ -802,6 +813,7 @@ def plot_scalar_array( robust=robust, center=center, cbar_kwargs={"label": field_data.name}, + **kwargs, ) # plot the simulation epsilon @@ -1021,8 +1033,10 @@ def split_original_fwd(self, num_mnts_original: int) -> Tuple[SimulationData, Si return sim_data_original, sim_data_fwd def make_adjoint_sim( - self, data_vjp_paths: set[tuple], adjoint_monitors: list[Monitor] - ) -> tuple[Simulation, AdjointSourceInfo]: + self, + data_vjp_paths: set[tuple], + adjoint_monitors: list[Monitor], + ) -> Simulation | None: """Make the adjoint simulation from the original simulation and the VJP-containing data.""" sim_original = self.simulation @@ -1033,6 +1047,9 @@ def make_adjoint_sim( for src_list in sources_adj_dict.values(): adj_srcs += list(src_list) + if not any(adj_srcs): + return None + adjoint_source_info = self.process_adjoint_sources(adj_srcs=adj_srcs) # grab boundary conditions with flipped Bloch vectors (for adjoint) @@ -1043,6 +1060,7 @@ def make_adjoint_sim( sources=adjoint_source_info.sources, boundary_spec=bc_adj, monitors=adjoint_monitors, + post_norm=adjoint_source_info.post_norm, ) if not adjoint_source_info.normalize_sim: @@ -1055,7 +1073,7 @@ def make_adjoint_sim( grid_spec_adj = grid_spec_original.updated_copy(wavelength=wavelength_original) sim_adj_update_dict["grid_spec"] = grid_spec_adj - return sim_original.updated_copy(**sim_adj_update_dict), adjoint_source_info + return sim_original.updated_copy(**sim_adj_update_dict) def make_adjoint_sources(self, data_vjp_paths: set[tuple]) -> dict[str, SourceType]: """Generate all of the non-zero sources for the adjoint simulation given the VJP data.""" diff --git a/tidy3d/components/dispersion_fitter.py b/tidy3d/components/dispersion_fitter.py index ee015cbfd..07fecf866 100644 --- a/tidy3d/components/dispersion_fitter.py +++ b/tidy3d/components/dispersion_fitter.py @@ -363,7 +363,7 @@ def get_default_weights(cls, eps: ArrayComplex1D) -> Tuple[float, float]: def pole_residue(self) -> Tuple[float, ArrayComplex1D, ArrayComplex1D]: """Parameters for pole-residue model in original units.""" if self.eps_inf is None or self.poles is None: - return None + return 1, [], [] poles = self.poles / self.scale_factor residues = self.residues / self.scale_factor eps_inf = self.eps_inf @@ -381,13 +381,47 @@ def values(self) -> ArrayComplex1D: """Evaluate model at sample frequencies.""" return self.evaluate(self.omega) + @cached_property + def loss_omega_pole_samples(self) -> ArrayFloat1D: + """Samples to check around each pole for passivity.""" + ranges_omega = [] + for pole, res in zip(self.poles, self.residues): + cr = np.real(res) + ci = np.imag(res) + ar = np.real(pole) + ai = np.imag(pole) + # no extra checking needed for marginally stable poles; these are handled later + if ar == 0: + continue + if cr == 0: + pole_extrema = [-ai] + else: + disc = ci**2 + cr**2 + pole_extrema = [ + -ai + ar * (ci + np.sqrt(disc)) / cr, + -ai + ar * (ci - np.sqrt(disc)) / cr, + ] + + ranges_omega.append(np.abs(pole_extrema)) + if len(ranges_omega) == 0: + return [] + return np.concatenate(ranges_omega) + + @cached_property + def loss_omega_samples(self) -> ArrayFloat1D: + """Frequencies to sample loss to ensure it is within bounds.""" + # let's check a big range in addition to the imag_extrema + range_omega = np.logspace(LOSS_CHECK_MIN, LOSS_CHECK_MAX, LOSS_CHECK_NUM) + range_omega_poles = self.loss_omega_pole_samples + return np.concatenate((range_omega, range_omega_poles)) + @cached_property def loss_in_bounds_violations(self) -> ArrayFloat1D: """Return list of frequencies where model violates loss bounds.""" extrema_list = imag_resp_extrema_locs(poles=self.poles, residues=self.residues) - # let's check a big range in addition to the imag_extrema - range_omega = np.logspace(LOSS_CHECK_MIN, LOSS_CHECK_MAX, LOSS_CHECK_NUM) + range_omega = self.loss_omega_samples + omega = np.concatenate((range_omega, extrema_list)) loss = self.evaluate(omega).imag bmin, bmax = self.loss_bounds @@ -673,7 +707,7 @@ def enforce_passivity( model = self.updated_copy(passivity_optimized=True) violations = model.loss_in_bounds_violations - range_omega = np.logspace(LOSS_CHECK_MIN, LOSS_CHECK_MAX, LOSS_CHECK_NUM) + range_omega = model.loss_omega_samples violations = np.unique(np.concatenate((violations, range_omega))) # only need one iteration since poles are fixed @@ -906,7 +940,6 @@ def make_configs(): "Unweighted RMS error %.3g", best_model.unweighted_rms_error, ) - return ( best_model.pole_residue, best_model.rms_error, diff --git a/tidy3d/components/eme/grid.py b/tidy3d/components/eme/grid.py index 7d63978dc..21ed3317e 100644 --- a/tidy3d/components/eme/grid.py +++ b/tidy3d/components/eme/grid.py @@ -14,7 +14,7 @@ from ..geometry.base import Box from ..grid.grid import Coords1D from ..mode import ModeSpec -from ..types import ArrayFloat1D, Axis, Coordinate, Size, TrackFreq, annotate_type +from ..types import ArrayFloat1D, Axis, Coordinate, Size, TrackFreq # grid limits MAX_NUM_MODES = 100 @@ -277,7 +277,7 @@ class EMECompositeGrid(EMEGridSpec): ... ) """ - subgrids: List[annotate_type(EMESubgridType)] = pd.Field( + subgrids: List[EMESubgridType] = pd.Field( ..., title="Subgrids", description="Subgrids in the composite grid." ) diff --git a/tidy3d/components/field_projection.py b/tidy3d/components/field_projection.py index 8e6a9d0b2..7135e5520 100644 --- a/tidy3d/components/field_projection.py +++ b/tidy3d/components/field_projection.py @@ -2,8 +2,9 @@ from __future__ import annotations -from typing import List, Tuple, Union +from typing import Iterable, List, Tuple, Union +import autograd.numpy as anp import numpy as np import pydantic.v1 as pydantic import xarray as xr @@ -12,6 +13,7 @@ from ..constants import C_0, EPSILON_0, ETA_0, MICROMETER, MU_0 from ..exceptions import SetupError from ..log import get_logging_console +from .autograd.functions import add_at, trapz from .base import Tidy3dBaseModel, cached_property, skip_if_fields_missing from .data.data_array import ( FieldProjectionAngleDataArray, @@ -278,7 +280,7 @@ def _fields_to_currents(field_data: FieldData, surface: FieldProjectionSurface) @staticmethod def _resample_surface_currents( - currents: xr.Dataset, + currents: FieldData, sim_data: SimulationData, surface: FieldProjectionSurface, medium: MediumType, @@ -288,7 +290,7 @@ def _resample_surface_currents( Parameters ---------- - currents : xarray.Dataset + currents : :class:`.FieldData` Surface currents defined on the original Yee grid. sim_data : :class:`.SimulationData` Container for simulation data containing the near field monitors. @@ -348,24 +350,38 @@ def _resample_surface_currents( currents = currents.colocate(*colocation_points) return currents - def integrate_1d( - self, - function: np.ndarray, - phase: np.ndarray, - pts_u: np.ndarray, + @staticmethod + def trapezoid( + ary: np.ndarray, + pts: Union[Iterable[np.ndarray], np.ndarray], + axes: Union[Iterable[int], int] = 0, ): - """Trapezoidal integration in two dimensions.""" - return np.trapz(np.squeeze(function) * np.squeeze(phase), pts_u, axis=0) # noqa: NPY201 + """Trapezoidal integration in n dimensions. - def integrate_2d( - self, - function: np.ndarray, - phase: np.ndarray, - pts_u: np.ndarray, - pts_v: np.ndarray, - ): - """Trapezoidal integration in two dimensions.""" - return np.trapz(np.trapz(np.squeeze(function) * phase, pts_u, axis=0), pts_v, axis=0) # noqa: NPY201 + Parameters + ---------- + ary : np.ndarray + Array to integrate. + pts : Iterable[np.ndarray] + Iterable of points for each dimension. + axes : Union[Iterable[int], int] + Iterable of axes along which to integrate. If not an iterable, assume 1D integration. + + Returns + ------- + np.ndarray + Integrated array. + """ + if not isinstance(axes, Iterable): + axes = [axes] + pts = [pts] + + for idx, (axis, pt) in enumerate(zip(axes, pts)): + if ary.shape[axis - idx] > 1: + ary = trapz(ary, pt, axis=axis - idx) + else: # array has only one element along axis + ary = ary[(slice(None),) * (axis - idx) + (0,)] + return ary def _far_fields_for_surface( self, @@ -375,7 +391,7 @@ def _far_fields_for_surface( surface: FieldProjectionSurface, currents: xr.Dataset, medium: MediumType, - ): + ) -> np.ndarray: """Compute far fields at an angle in spherical coordinates for a given set of surface currents and observation angles. @@ -397,11 +413,10 @@ def _far_fields_for_surface( Returns ------- - tuple(numpy.ndarray[float], ...) - ``Er``, ``Etheta``, ``Ephi``, ``Hr``, ``Htheta``, ``Hphi`` for the given surface. + np.ndarray + With leading dimension containing ``Er``, ``Etheta``, ``Ephi``, ``Hr``, ``Htheta``, ``Hphi`` + projected fields for each frequency. """ - pts = [currents[name].values for name in ["x", "y", "z"]] - try: currents_f = currents.sel(f=frequency) except Exception as e: @@ -429,83 +444,61 @@ def _far_fields_for_surface( theta = np.atleast_1d(theta) phi = np.atleast_1d(phi) - sin_theta = np.sin(theta) - cos_theta = np.cos(theta) - sin_phi = np.sin(phi) - cos_phi = np.cos(phi) - - J = np.zeros((3, len(theta), len(phi)), dtype=complex) - M = np.zeros_like(J) - - phase = [None] * 3 propagation_factor = -1j * AbstractFieldProjectionData.wavenumber( medium=medium, frequency=frequency ) - def integrate_for_one_theta(i_th: int): - """Perform integration for a given theta angle index""" - - for j_ph in np.arange(len(phi)): - phase[0] = np.exp(propagation_factor * pts[0] * sin_theta[i_th] * cos_phi[j_ph]) - phase[1] = np.exp(propagation_factor * pts[1] * sin_theta[i_th] * sin_phi[j_ph]) - phase[2] = np.exp(propagation_factor * pts[2] * cos_theta[i_th]) + sin_theta = np.sin(theta) + cos_theta = np.cos(theta) + sin_phi = np.sin(phi) + cos_phi = np.cos(phi) - phase_ij = phase[idx_u][:, None] * phase[idx_v][None, :] * phase[idx_w] + pts = [currents[name].values for name in ["x", "y", "z"]] - if self.is_2d_simulation: - J[idx_u, i_th, j_ph] = self.integrate_1d( - currents_f[f"E{cmp_1}"].values, phase_ij, pts[idx_int_1d] - ) + phase_0 = np.exp(np.einsum("i,j,k->ijk", propagation_factor * pts[0], sin_theta, cos_phi)) + phase_1 = np.exp(np.einsum("i,j,k->ijk", propagation_factor * pts[1], sin_theta, sin_phi)) + phase_2 = np.exp(np.einsum("i,j->ij", propagation_factor * pts[2], cos_theta)) - J[idx_v, i_th, j_ph] = self.integrate_1d( - currents_f[f"E{cmp_2}"].values, phase_ij, pts[idx_int_1d] - ) + E1 = "E" + cmp_1 + E2 = "E" + cmp_2 + H1 = "H" + cmp_1 + H2 = "H" + cmp_2 - M[idx_u, i_th, j_ph] = self.integrate_1d( - currents_f[f"H{cmp_1}"].values, phase_ij, pts[idx_int_1d] - ) + def contract(currents): + return anp.einsum("xtp,ytp,zt,xyz->xyztp", phase_0, phase_1, phase_2, currents) - M[idx_v, i_th, j_ph] = self.integrate_1d( - currents_f[f"H{cmp_2}"].values, phase_ij, pts[idx_int_1d] - ) - else: - J[idx_u, i_th, j_ph] = self.integrate_2d( - currents_f[f"E{cmp_1}"].values, phase_ij, pts[idx_u], pts[idx_v] - ) + jm = [] + for field_component in (E1, E2, H1, H2): + currents = currents_f[field_component].data + currents = anp.reshape(currents, currents_f[field_component].shape) + currents_phase = contract(currents) - J[idx_v, i_th, j_ph] = self.integrate_2d( - currents_f[f"E{cmp_2}"].values, phase_ij, pts[idx_u], pts[idx_v] - ) + if self.is_2d_simulation: + jm_i = self.trapezoid(currents_phase, pts[idx_int_1d], idx_int_1d) + else: + jm_i = self.trapezoid(currents_phase, (pts[idx_u], pts[idx_v]), (idx_u, idx_v)) - M[idx_u, i_th, j_ph] = self.integrate_2d( - currents_f[f"H{cmp_1}"].values, phase_ij, pts[idx_u], pts[idx_v] - ) + jm.append(anp.reshape(jm_i, (len(theta), len(phi)))) - M[idx_v, i_th, j_ph] = self.integrate_2d( - currents_f[f"H{cmp_2}"].values, phase_ij, pts[idx_u], pts[idx_v] - ) + order = [idx_u, idx_v, idx_w] + zeros = np.zeros(jm[0].shape) - if len(theta) < 2: - integrate_for_one_theta(0) - else: - for i_th in track( - np.arange(len(theta)), - description=f"Processing surface monitor '{surface.monitor.name}'...", - console=get_logging_console(), - ): - integrate_for_one_theta(i_th) + # for each index (0, 1, 2), if it’s in the first two elements of order, + # select the corresponding jm element for J or the offset element (+2) for M + J = anp.array([jm[order.index(i)] if i in order[:2] else zeros for i in range(3)]) + M = anp.array([jm[order.index(i) + 2] if i in order[:2] else zeros for i in range(3)]) - cos_th_cos_phi = cos_theta[:, None] * cos_phi[None, :] - cos_th_sin_phi = cos_theta[:, None] * sin_phi[None, :] + cos_theta_cos_phi = cos_theta[:, None] * cos_phi[None, :] + cos_theta_sin_phi = cos_theta[:, None] * sin_phi[None, :] # Ntheta (8.33a) - Ntheta = J[0] * cos_th_cos_phi + J[1] * cos_th_sin_phi - J[2] * sin_theta[:, None] + Ntheta = J[0] * cos_theta_cos_phi + J[1] * cos_theta_sin_phi - J[2] * sin_theta[:, None] # Nphi (8.33b) Nphi = -J[0] * sin_phi[None, :] + J[1] * cos_phi[None, :] # Ltheta (8.34a) - Ltheta = M[0] * cos_th_cos_phi + M[1] * cos_th_sin_phi - M[2] * sin_theta[:, None] + Ltheta = M[0] * cos_theta_cos_phi + M[1] * cos_theta_sin_phi - M[2] * sin_theta[:, None] # Lphi (8.34b) Lphi = -M[0] * sin_phi[None, :] + M[1] * cos_phi[None, :] @@ -514,12 +507,12 @@ def integrate_for_one_theta(i_th: int): Etheta = -(Lphi + eta * Ntheta) Ephi = Ltheta - eta * Nphi - Er = np.zeros_like(Ephi) + Er = anp.zeros_like(Ephi) Htheta = -Ephi / eta Hphi = Etheta / eta - Hr = np.zeros_like(Hphi) + Hr = anp.zeros_like(Hphi) - return Er, Etheta, Ephi, Hr, Htheta, Hphi + return anp.array([Er, Etheta, Ephi, Hr, Htheta, Hphi]) @staticmethod def apply_window_to_currents( @@ -604,9 +597,7 @@ def _project_fields_angular( # compute projected fields for the dataset associated with each monitor field_names = ("Er", "Etheta", "Ephi", "Hr", "Htheta", "Hphi") - fields = [ - np.zeros((1, len(theta), len(phi), len(freqs)), dtype=complex) for _ in field_names - ] + fields = np.zeros((len(field_names), 1, len(theta), len(phi), len(freqs)), dtype=complex) medium = monitor.medium if monitor.medium else self.medium k = AbstractFieldProjectionData.wavenumber(medium=medium, frequency=freqs) @@ -630,8 +621,7 @@ def _project_fields_angular( currents=currents, medium=medium, ) - for field, _field in zip(fields, _fields): - field[..., idx_f] += _field * phase[idx_f] + fields = add_at(fields, [..., idx_f], _fields[:, None] * phase[idx_f]) else: iter_coords = [ ([_theta, _phi], [i, j]) @@ -647,8 +637,9 @@ def _project_fields_angular( _fields = self._fields_for_surface_exact( x=_x, y=_y, z=_z, surface=surface, currents=currents, medium=medium ) - for field, _field in zip(fields, _fields): - field[0, i, j, :] += _field + where = (slice(None), 0, i, j) + _fields = anp.reshape(_fields, fields[where].shape) + fields = add_at(fields, where, _fields) coords = {"r": np.atleast_1d(monitor.proj_distance), "theta": theta, "phi": phi, "f": freqs} fields = { @@ -687,9 +678,7 @@ def _project_fields_cartesian( # compute projected fields for the dataset associated with each monitor field_names = ("Er", "Etheta", "Ephi", "Hr", "Htheta", "Hphi") - fields = [ - np.zeros((len(x), len(y), len(z), len(freqs)), dtype=complex) for _ in field_names - ] + fields = np.zeros((len(field_names), len(x), len(y), len(z), len(freqs)), dtype=complex) medium = monitor.medium if monitor.medium else self.medium wavenumber = AbstractFieldProjectionData.wavenumber(medium=medium, frequency=freqs) @@ -728,14 +717,16 @@ def _project_fields_cartesian( currents=currents, medium=medium, ) - for field, _field in zip(fields, _fields): - field[i, j, k, idx_f] += _field * phase[idx_f] + where = (slice(None), i, j, k, idx_f) + _fields = anp.reshape(_fields, fields[where].shape) + fields = add_at(fields, where, _fields * phase[idx_f]) else: _fields = self._fields_for_surface_exact( x=_x, y=_y, z=_z, surface=surface, currents=currents, medium=medium ) - for field, _field in zip(fields, _fields): - field[i, j, k, :] += _field + where = (slice(None), i, j, k) + _fields = anp.reshape(_fields, fields[where].shape) + fields = add_at(fields, where, _fields) coords = {"x": x, "y": y, "z": z, "f": freqs} fields = { @@ -768,7 +759,7 @@ def _project_fields_kspace( # compute projected fields for the dataset associated with each monitor field_names = ("Er", "Etheta", "Ephi", "Hr", "Htheta", "Hphi") - fields = [np.zeros((len(ux), len(uy), 1, len(freqs)), dtype=complex) for _ in field_names] + fields = np.zeros((len(field_names), len(ux), len(uy), 1, len(freqs)), dtype=complex) medium = monitor.medium if monitor.medium else self.medium k = AbstractFieldProjectionData.wavenumber(medium=medium, frequency=freqs) @@ -802,16 +793,17 @@ def _project_fields_kspace( currents=currents, medium=medium, ) - for field, _field in zip(fields, _fields): - field[i, j, 0, idx_f] += _field * phase[idx_f] - + where = (slice(None), i, j, 0, idx_f) + _fields = anp.reshape(_fields, fields[where].shape) + fields = add_at(fields, where, _fields * phase[idx_f]) else: _x, _y, _z = monitor.sph_2_car(monitor.proj_distance, theta, phi) _fields = self._fields_for_surface_exact( x=_x, y=_y, z=_z, surface=surface, currents=currents, medium=medium ) - for field, _field in zip(fields, _fields): - field[i, j, 0, :] += _field + where = (slice(None), i, j, 0) + _fields = anp.reshape(_fields, fields[where].shape) + fields = add_at(fields, where, _fields) coords = { "ux": np.array(monitor.ux), @@ -837,7 +829,7 @@ def _fields_for_surface_exact( surface: FieldProjectionSurface, currents: xr.Dataset, medium: MediumType, - ): + ) -> np.ndarray: """Compute projected fields in spherical coordinates at a given projection point on a Cartesian grid for a given set of surface currents using the exact homogeneous medium Green's function without geometric approximations. @@ -859,11 +851,11 @@ def _fields_for_surface_exact( Returns ------- - tuple(np.ndarray, np.ndarray, np.ndarray, np.ndarray, np.ndarray, np.ndarray) - ``Er``, ``Etheta``, ``Ephi``, ``Hr``, ``Htheta``, ``Hphi`` projected fields for - each frequency. + np.ndarray + With leading dimension containing ``Er``, ``Etheta``, ``Ephi``, ``Hr``, ``Htheta``, ``Hphi`` + projected fields for each frequency. """ - freqs = np.array(self.frequencies) + freqs = anp.array(self.frequencies) i_omega = 1j * 2.0 * np.pi * freqs[None, None, None, :] wavenumber = AbstractFieldProjectionData.wavenumber(frequency=freqs, medium=medium) wavenumber = wavenumber[None, None, None, :] # add space dimensions @@ -886,15 +878,14 @@ def _fields_for_surface_exact( cmp_1, cmp_2 = source_names # set the surface current density Cartesian components - J = [np.atleast_1d(0)] * 3 - M = [np.atleast_1d(0)] * 3 - - J[idx_u] = currents[f"E{cmp_1}"].values - J[idx_v] = currents[f"E{cmp_2}"].values - J[idx_w] = np.zeros_like(J[idx_u]) - M[idx_u] = currents[f"H{cmp_1}"].values - M[idx_v] = currents[f"H{cmp_2}"].values - M[idx_w] = np.zeros_like(M[idx_u]) + J = [None] * 3 + M = [None] * 3 + J[idx_u] = currents[f"E{cmp_1}"].data + J[idx_v] = currents[f"E{cmp_2}"].data + J[idx_w] = anp.zeros(J[idx_u].shape) + M[idx_u] = currents[f"H{cmp_1}"].data + M[idx_v] = currents[f"H{cmp_2}"].data + M[idx_w] = anp.zeros(M[idx_u].shape) # observation point in the new spherical system r, theta_obs, phi_obs = surface.monitor.car_2_sph( @@ -902,14 +893,14 @@ def _fields_for_surface_exact( ) # angle terms - sin_theta = np.sin(theta_obs) - cos_theta = np.cos(theta_obs) - sin_phi = np.sin(phi_obs) - cos_phi = np.cos(phi_obs) + sin_theta = anp.sin(theta_obs) + cos_theta = anp.cos(theta_obs) + sin_phi = anp.sin(phi_obs) + cos_phi = anp.cos(phi_obs) # Green's function and terms related to its derivatives ikr = 1j * wavenumber * r - G = np.exp(ikr) / (4.0 * np.pi * r) + G = anp.exp(ikr) / (4.0 * np.pi * r) dG_dr = G * (ikr - 1.0) / r d2G_dr2 = dG_dr * (ikr - 1.0) / r + G / (r**2) @@ -982,12 +973,12 @@ def potential_terms(current: Tuple[np.ndarray, ...], const: complex): ) # integrate over the surface - e_x = self.integrate_2d(e_x_integrand, 1.0, pts[idx_u], pts[idx_v]) - e_y = self.integrate_2d(e_y_integrand, 1.0, pts[idx_u], pts[idx_v]) - e_z = self.integrate_2d(e_z_integrand, 1.0, pts[idx_u], pts[idx_v]) - h_x = self.integrate_2d(h_x_integrand, 1.0, pts[idx_u], pts[idx_v]) - h_y = self.integrate_2d(h_y_integrand, 1.0, pts[idx_u], pts[idx_v]) - h_z = self.integrate_2d(h_z_integrand, 1.0, pts[idx_u], pts[idx_v]) + e_x = self.trapezoid(e_x_integrand, (pts[idx_u], pts[idx_v]), (idx_u, idx_v)) + e_y = self.trapezoid(e_y_integrand, (pts[idx_u], pts[idx_v]), (idx_u, idx_v)) + e_z = self.trapezoid(e_z_integrand, (pts[idx_u], pts[idx_v]), (idx_u, idx_v)) + h_x = self.trapezoid(h_x_integrand, (pts[idx_u], pts[idx_v]), (idx_u, idx_v)) + h_y = self.trapezoid(h_y_integrand, (pts[idx_u], pts[idx_v]), (idx_u, idx_v)) + h_z = self.trapezoid(h_z_integrand, (pts[idx_u], pts[idx_v]), (idx_u, idx_v)) # observation point in the original spherical system _, theta_obs, phi_obs = surface.monitor.car_2_sph(x, y, z) @@ -996,4 +987,4 @@ def potential_terms(current: Tuple[np.ndarray, ...], const: complex): e_r, e_theta, e_phi = surface.monitor.car_2_sph_field(e_x, e_y, e_z, theta_obs, phi_obs) h_r, h_theta, h_phi = surface.monitor.car_2_sph_field(h_x, h_y, h_z, theta_obs, phi_obs) - return [e_r, e_theta, e_phi, h_r, h_theta, h_phi] + return anp.array([e_r, e_theta, e_phi, h_r, h_theta, h_phi]) diff --git a/tidy3d/components/frequencies.py b/tidy3d/components/frequencies.py new file mode 100644 index 000000000..6a14fb5b6 --- /dev/null +++ b/tidy3d/components/frequencies.py @@ -0,0 +1,231 @@ +"""Frequency utilities.""" + +import numpy as np +import pydantic as pd + +from ..constants import C_0 +from .base import Tidy3dBaseModel + +O_BAND = (1.260, 1.360) +E_BAND = (1.360, 1.460) +S_BAND = (1.460, 1.530) +C_BAND = (1.530, 1.565) +L_BAND = (1.565, 1.625) +U_BAND = (1.625, 1.675) + + +class FrequencyUtils(Tidy3dBaseModel): + """Class for general frequency/wavelength utilities.""" + + use_wavelength: bool = pd.Field( + False, + title="Use wavelength", + description="Indicate whether to use wavelengths instead of frequencies for the return " + "values of functions and parameters.", + ) + + def classification(self, value: float) -> tuple[str]: + """Band classification for a given frequency/wavelength. + + Frequency values must be given in hertz (Hz). Wavelengths must be + given in micrometers (ÎŒm). + + Parameters + ---------- + value : float + Value to classify. + + Returns + ------- + tuple[str] + String tuple with classification. + """ + if self.use_wavelength: + value = C_0 / value + if value < 3: + return ("near static",) + elif value < 300e6: + if value < 30: + return ("radio wave", "ELF") + elif value < 300: + return ("radio wave", "SLF") + elif value < 3e3: + return ("radio wave", "ULF") + elif value < 30e3: + return ("radio wave", "VLF") + elif value < 300e3: + return ("radio wave", "LF") + elif value < 3e6: + return ("radio wave", "MF") + elif value < 30e6: + return ("radio wave", "HF") + return ("radio wave", "VHF") + elif value < 300e9: + if value < 3e9: + return ("microwave", "UHF") + elif value < 30e9: + return ("microwave", "SHF") + return ("microwave", "EHF") + elif value < 400e12: + if value < 6e12: + return ("infrared", "FIR") + elif value < 100e12: + return ("infrared", "MIR") + return ("infrared", "NIR") + elif value < 790e12: + if value < 480e12: + return ("visible", "red") + elif value < 510e12: + return ("visible", "orange") + elif value < 530e12: + return ("visible", "yellow") + elif value < 600e12: + return ("visible", "green") + elif value < 620e12: + return ("visible", "cyan") + elif value < 670e12: + return ("visible", "blue") + return ("visible", "violet") + elif value < 30e15: + if value < 1e15: + return ("ultraviolet", "NUV") + elif value < 1.5e15: + return ("ultraviolet", "MUV") + elif value < 2.47e15: + return ("ultraviolet", "FUV") + return ("ultraviolet", "EUV") + if value < 30e18: + if value < 3e18: + return ("X-ray", "soft X-ray") + return ("X-ray", "hard X-ray") + return ("Îł-ray",) + + def o_band(self, n: int = 11) -> list[float]: + """ + Optical O band frequencies/wavelengths sorted by wavelength. + + The returned samples are equally spaced in wavelength. + + Parameters + ---------- + n : int + Desired number of samples. + + Returns + ------- + list[float] + Samples list. + """ + values = np.linspace(*O_BAND, n) + if not self.use_wavelength: + values = C_0 / values + return values.tolist() + + def e_band(self, n: int = 11) -> list[float]: + """ + Optical E band frequencies/wavelengths sorted by wavelength. + + The returned samples are equally spaced in wavelength. + + Parameters + ---------- + n : int + Desired number of samples. + + Returns + ------- + list[float] + Samples list. + """ + values = np.linspace(*E_BAND, n) + if not self.use_wavelength: + values = C_0 / values + return values.tolist() + + def s_band(self, n: int = 15) -> list[float]: + """ + Optical S band frequencies/wavelengths sorted by wavelength. + + The returned samples are equally spaced in wavelength. + + Parameters + ---------- + n : int + Desired number of samples. + + Returns + ------- + list[float] + Samples list. + """ + values = np.linspace(*S_BAND, n) + if not self.use_wavelength: + values = C_0 / values + return values.tolist() + + def c_band(self, n: int = 8) -> list[float]: + """ + Optical C band frequencies/wavelengths sorted by wavelength. + + The returned samples are equally spaced in wavelength. + + Parameters + ---------- + n : int + Desired number of samples. + + Returns + ------- + list[float] + Samples list. + """ + values = np.linspace(*C_BAND, n) + if not self.use_wavelength: + values = C_0 / values + return values.tolist() + + def l_band(self, n: int = 13) -> list[float]: + """ + Optical L band frequencies/wavelengths sorted by wavelength. + + The returned samples are equally spaced in wavelength. + + Parameters + ---------- + n : int + Desired number of samples. + + Returns + ------- + list[float] + Samples list. + """ + values = np.linspace(*L_BAND, n) + if not self.use_wavelength: + values = C_0 / values + return values.tolist() + + def u_band(self, n: int = 11) -> list[float]: + """ + Optical U band frequencies/wavelengths sorted by wavelength. + + The returned samples are equally spaced in wavelength. + + Parameters + ---------- + n : int + Desired number of samples. + + Returns + ------- + list[float] + Samples list. + """ + values = np.linspace(*U_BAND, n) + if not self.use_wavelength: + values = C_0 / values + return values.tolist() + + +frequencies = FrequencyUtils(use_wavelength=False) +wavelengths = FrequencyUtils(use_wavelength=True) diff --git a/tidy3d/components/geometry/base.py b/tidy3d/components/geometry/base.py index d0f436383..2f19a3928 100644 --- a/tidy3d/components/geometry/base.py +++ b/tidy3d/components/geometry/base.py @@ -1059,11 +1059,11 @@ def kspace_2_sph(ux: float, uy: float, axis: Axis) -> Tuple[float, float]: return theta_local, phi_local x = np.cos(theta_local) - y = np.sin(theta_local) * np.sin(phi_local) - z = -np.sin(theta_local) * np.cos(phi_local) + y = np.sin(theta_local) * np.cos(phi_local) + z = np.sin(theta_local) * np.sin(phi_local) if axis == 1: - x, y, z = -z, x, -y + x, y, z = y, x, z theta = np.arccos(z) phi = np.arctan2(y, x) @@ -2510,6 +2510,7 @@ def derivative_face( # if not enough data, just use best guess using eps in medium and simulation needs_eps_approx = any(len(eps.coords[dim_normal]) <= num_cells_in for eps in eps_xyz) + if derivative_info.eps_approx or needs_eps_approx: eps_xyz_inside = 3 * [derivative_info.eps_in] eps_xyz_outside = 3 * [derivative_info.eps_out] @@ -2520,7 +2521,7 @@ def derivative_face( if min_max_index == 0: index_out, index_in = (0, num_cells_in - 1) else: - index_out, index_in = (-1, -num_cells_in - 1) + index_out, index_in = (-1, -num_cells_in) eps_xyz_inside = [eps.isel(**{dim_normal: index_in}) for eps in eps_xyz] eps_xyz_outside = [eps.isel(**{dim_normal: index_out}) for eps in eps_xyz] @@ -2535,7 +2536,7 @@ def derivative_face( def integrate_face(arr: xr.DataArray) -> complex: """Interpolate and integrate a scalar field data over the face using bounds.""" - arr_at_face = arr.interp(**{dim_normal: coord_normal_face}) + arr_at_face = arr.interp(**{dim_normal: coord_normal_face}, assume_sorted=True) integral_result = integrate_within_bounds( arr=arr_at_face, @@ -2929,12 +2930,10 @@ def intersections_tilted_plane( For more details refer to `Shapely's Documentation `_. """ - geom_a = Geometry.evaluate_inf_shape( - shapely.unary_union(self.geometry_a.intersections_tilted_plane(normal, origin, to_2D)) - ) - geom_b = Geometry.evaluate_inf_shape( - shapely.unary_union(self.geometry_b.intersections_tilted_plane(normal, origin, to_2D)) - ) + a = self.geometry_a.intersections_tilted_plane(normal, origin, to_2D) + b = self.geometry_b.intersections_tilted_plane(normal, origin, to_2D) + geom_a = shapely.unary_union([Geometry.evaluate_inf_shape(g) for g in a]) + geom_b = shapely.unary_union([Geometry.evaluate_inf_shape(g) for g in b]) return ClipOperation.to_polygon_list(self._shapely_operation(geom_a, geom_b)) def intersections_plane( @@ -2958,12 +2957,10 @@ def intersections_plane( For more details refer to `Shapely's Documentaton `_. """ - geom_a = Geometry.evaluate_inf_shape( - shapely.unary_union(self.geometry_a.intersections_plane(x, y, z)) - ) - geom_b = Geometry.evaluate_inf_shape( - shapely.unary_union(self.geometry_b.intersections_plane(x, y, z)) - ) + a = self.geometry_a.intersections_plane(x, y, z) + b = self.geometry_b.intersections_plane(x, y, z) + geom_a = shapely.unary_union([Geometry.evaluate_inf_shape(g) for g in a]) + geom_b = shapely.unary_union([Geometry.evaluate_inf_shape(g) for g in b]) return ClipOperation.to_polygon_list(self._shapely_operation(geom_a, geom_b)) @cached_property @@ -3286,7 +3283,7 @@ def compute_derivatives(self, derivative_info: DerivativeInfo) -> AutogradFieldM _, index, *geo_path = field_path geo = self.geometries[index] geo_info = derivative_info.updated_copy( - paths=[geo_path], bounds=geo.bounds, eps_approx=True + paths=[geo_path], bounds=geo.bounds, eps_approx=True, deep=False ) vjp_dict_geo = geo.compute_derivatives(geo_info) grad_vjp_values = list(vjp_dict_geo.values()) diff --git a/tidy3d/components/geometry/polyslab.py b/tidy3d/components/geometry/polyslab.py index 3930e96d1..eafdd668b 100644 --- a/tidy3d/components/geometry/polyslab.py +++ b/tidy3d/components/geometry/polyslab.py @@ -9,7 +9,7 @@ import autograd.numpy as np import pydantic.v1 as pydantic import shapely -import xarray as xr +from autograd.tracer import isbox from matplotlib import path from ...constants import LARGE_NUMBER, MICROMETER, fp_eps @@ -19,7 +19,6 @@ from ..autograd import AutogradFieldMap, TracedVertices, get_static from ..autograd.derivative_utils import DerivativeInfo from ..base import cached_property, skip_if_fields_missing -from ..data.dataset import ElectromagneticFieldDataset from ..types import ( ArrayFloat2D, ArrayLike, @@ -1318,6 +1317,13 @@ def _heal_polygon(vertices: np.ndarray) -> np.ndarray: shapely_poly = PolySlab.make_shapely_polygon(vertices) if shapely_poly.is_valid: return vertices + elif isbox(vertices): + raise NotImplementedError( + "The dilation caused damage to the polygon. " + "Automatically healing this is currently not supported when " + "differentiating w.r.t. the vertices. Try increasing the spacing " + "between vertices or reduce the amount of dilation." + ) # perform healing poly_heal = shapely.make_valid(shapely_poly) return PolySlab._proper_vertices(list(poly_heal.exterior.coords)) @@ -1390,7 +1396,7 @@ def compute_derivative_vertices(self, derivative_info: DerivativeInfo) -> Traced # compute edges between vertices - vertices_next = np.roll(self.vertices, axis=0, shift=1) + vertices_next = np.roll(self.vertices, axis=0, shift=-1) edges = vertices_next - vertices # compute center positions between each edge @@ -1401,21 +1407,17 @@ def compute_derivative_vertices(self, derivative_info: DerivativeInfo) -> Traced if edge_centers_xyz.shape != (num_vertices, 3): raise AssertionError("something bad happened") - # compute the E and D fields at the edge centers - E_der_at_edges = self.der_at_centers( - der_map=derivative_info.E_der_map, edge_centers=edge_centers_xyz - ) - D_der_at_edges = self.der_at_centers( - der_map=derivative_info.D_der_map, edge_centers=edge_centers_xyz - ) - - # compute the basis vectors along each edge + # get basis vectors for every edge segment basis_vectors = self.edge_basis_vectors(edges=edges) - # project the D and E fields into the basis vectors - D_der_norm = self.project_in_basis(D_der_at_edges, basis_vector=basis_vectors["norm"]) - E_der_edge = self.project_in_basis(E_der_at_edges, basis_vector=basis_vectors["edge"]) - E_der_slab = self.project_in_basis(E_der_at_edges, basis_vector=basis_vectors["slab"]) + grad_bases = derivative_info.grad_in_bases( + spatial_coords=edge_centers_xyz, basis_vectors=basis_vectors + ) + + # unpack gradient contributions from different bases + D_der_norm = grad_bases["D_norm"] + E_der_edge = grad_bases["E_perp1"] + E_der_slab = grad_bases["E_perp2"] # approximate permittivity in and out delta_eps_inv = 1.0 / derivative_info.eps_in - 1.0 / derivative_info.eps_out @@ -1438,12 +1440,9 @@ def compute_derivative_vertices(self, derivative_info: DerivativeInfo) -> Traced edge_areas = edge_lengths # correction to edge area based on sidewall distance along slab axis - dim_axis = "xyz"[self.axis] - field_coords_axis = derivative_info.E_der_map[f"E{dim_axis}"].coords[dim_axis] - if len(field_coords_axis) > 1: - slab_height = abs(float(np.squeeze(np.diff(self.slab_bounds)))) - if not np.isinf(slab_height): - edge_areas *= slab_height + slab_height = abs(float(np.squeeze(np.diff(self.slab_bounds)))) + if not np.isinf(slab_height): + edge_areas *= slab_height vjps_edges *= edge_areas @@ -1451,54 +1450,48 @@ def compute_derivative_vertices(self, derivative_info: DerivativeInfo) -> Traced vjps_edges_in_plane = vjps_edges.values.reshape((num_vertices, 1)) * normal_vectors_in_plane - vjps_vertices = vjps_edges_in_plane + np.roll(vjps_edges_in_plane, axis=0, shift=-1) + vjps_vertices = vjps_edges_in_plane + np.roll(vjps_edges_in_plane, axis=0, shift=1) + vjps_vertices /= 2.0 # each vertex is effected only 1/2 by each edge # sign change if counter clockwise, because normal direction is flipped if self.is_ccw: vjps_vertices *= -1 - # TODO: verify sign, or if this is rather when `not self.is_ccw` return vjps_vertices.real - def der_at_centers( + def edge_basis_vectors( self, - der_map: ElectromagneticFieldDataset, - edge_centers: np.ndarray, # (N, 3) - ) -> xr.Dataset: - """Compute the value of an ``ElectromagneticFieldDataset`` at a set of edge centers.""" + edges: np.ndarray, # (N, 2) + ) -> dict[str, np.ndarray]: # (N, 3) + """Normalized basis vectors for 'normal' direction, 'slab' tangent direction and 'edge'.""" - xs, ys, zs = edge_centers.T - edge_index_dim = "edge_index" + num_vertices, _ = edges.shape + zeros = np.zeros(num_vertices) + ones = np.ones(num_vertices) - interp_kwargs = {} - for dim, centers_dim in zip("xyz", edge_centers.T): - # only include dims where the data has more than 1 coord, to avoid warnings and errors - coords_data = der_map[f"E{dim}"].coords - if np.array(coords_data).size > 1: - interp_kwargs[dim] = xr.DataArray(centers_dim, dims=edge_index_dim) + # normalized vectors along edges + edges_norm_in_plane = self.normalize_vect(edges) + edges_norm_xyz = self.unpop_axis_vect(zeros, edges_norm_in_plane) - components = {} - for fld_name, arr in der_map.items(): - components[fld_name] = arr.interp(**interp_kwargs).sum("f") + # normalized vectors from base of edges to tops of edges + slabs_axis_components = np.cos(self.sidewall_angle) * ones + axis_norm = self.unpop_axis(1.0, (0.0, 0.0), axis=self.axis) + slab_normal_xyz = -np.sin(self.sidewall_angle) * np.cross(edges_norm_xyz, axis_norm) + _, slab_normal_in_plane = self.pop_axis_vect(slab_normal_xyz) + slabs_norm_xyz = self.unpop_axis_vect(slabs_axis_components, slab_normal_in_plane) - return xr.Dataset(components) + # normalized vectors pointing in normal direction of edge + normals_norm_xyz = np.cross(edges_norm_xyz, slabs_norm_xyz) - def project_in_basis( - self, - der_dataset: xr.Dataset, - basis_vector: np.ndarray, - ) -> xr.DataArray: - """Project a derivative dataset along a supplied basis vector.""" + if self.axis != 1: + normals_norm_xyz *= -1 - value = 0.0 - for coeffs, dim in zip(basis_vector.T, "xyz"): - value += coeffs * der_dataset.data_vars[f"E{dim}"] - return value + return dict(norm=normals_norm_xyz, perp1=edges_norm_xyz, perp2=slabs_norm_xyz) def unpop_axis_vect(self, ax_coords: np.ndarray, plane_coords: np.ndarray) -> np.ndarray: """Combine coordinate along axis with coordinates on the plane tangent to the axis. - ax_coords.shape == [N] or [N, 1] + ax_coords.shape == [N] plane_coords.shape == [N, 2] return shape == [N, 3] @@ -1523,33 +1516,9 @@ def pop_axis_vect(self, coord: np.ndarray) -> Tuple[np.ndarray, Tuple[np.ndarray @staticmethod def normalize_vect(arr: np.ndarray) -> np.ndarray: """normalize an array shaped (N, d) along the `d` axis and return (N, 1).""" - return arr / np.linalg.norm(arr, axis=-1)[..., None] - - def edge_basis_vectors( - self, - edges: np.ndarray, # (N, 2) - ) -> dict[str, np.ndarray]: # (N, 3) - """Normalized basis vectors for 'normal' direction, 'slab' tangent direction and 'edge'.""" - - num_vertices, _ = edges.shape - zeros = np.zeros(num_vertices) - ones = np.ones(num_vertices) - - # normalized vectors along edges - edges_norm_in_plane = self.normalize_vect(edges) - edges_norm_xyz = self.unpop_axis_vect(zeros, edges_norm_in_plane) - - # normalized vectors from base of edges to tops of edges - slabs_axis_components = np.cos(self.sidewall_angle) * ones - axis_norm = self.unpop_axis(1.0, (0.0, 0.0), axis=self.axis) - slab_normal_xyz = -np.sin(self.sidewall_angle) * np.cross(edges_norm_xyz, axis_norm) - _, slab_normal_in_plane = self.pop_axis_vect(slab_normal_xyz) - slabs_norm_xyz = self.unpop_axis_vect(slabs_axis_components, slab_normal_in_plane) - - # normalized vectors pointing in normal direction of edge - normals_norm_xyz = np.cross(edges_norm_xyz, slabs_norm_xyz) - - return dict(edge=edges_norm_xyz, norm=normals_norm_xyz, slab=slabs_norm_xyz) + norm = np.linalg.norm(arr, axis=-1, keepdims=True) + norm = np.where(norm == 0, 1, norm) + return arr / norm class ComplexPolySlabBase(PolySlab): @@ -1776,7 +1745,7 @@ def intersections_tilted_plane( return [ shapely.unary_union( [ - shape + base.Geometry.evaluate_inf_shape(shape) for polyslab in self.sub_polyslabs for shape in polyslab.intersections_tilted_plane(normal, origin, to_2D) ] diff --git a/tidy3d/components/geometry/primitives.py b/tidy3d/components/geometry/primitives.py index 2c88c5f62..2edba94d9 100644 --- a/tidy3d/components/geometry/primitives.py +++ b/tidy3d/components/geometry/primitives.py @@ -5,16 +5,20 @@ from math import isclose from typing import List +import autograd.numpy as anp import numpy as np import pydantic.v1 as pydantic import shapely -from ...constants import LARGE_NUMBER, MICROMETER +from ...constants import C_0, LARGE_NUMBER, MICROMETER from ...exceptions import SetupError, ValidationError from ...packaging import verify_packages_import +from ..autograd import AutogradFieldMap, TracedSize1D +from ..autograd.derivative_utils import DerivativeInfo from ..base import cached_property, skip_if_fields_missing from ..types import Axis, Bound, Coordinate, MatrixReal4x4, Shapely, Tuple from . import base +from .polyslab import PolySlab # for sampling conical frustum in visualization _N_SAMPLE_CURVE_SHAPELY = 40 @@ -22,6 +26,12 @@ # for shapely circular shapes discretization in visualization _N_SHAPELY_QUAD_SEGS = 200 +# Default number of points to discretize polyslab in `Cylinder.to_polyslab()` +_N_PTS_CYLINDER_POLYSLAB = 51 + +# Default number of points per wvl in material for discretizing cylinder in autograd derivative +_PTS_PER_WVL_MAT_CYLINDER_DISCRETIZE = 10 + class Sphere(base.Centered, base.Circular): """Spherical geometry. @@ -185,7 +195,7 @@ class Cylinder(base.Centered, base.Circular, base.Planar): """ # Provide more explanations on where radius is defined - radius: pydantic.NonNegativeFloat = pydantic.Field( + radius: TracedSize1D = pydantic.Field( ..., title="Radius", description="Radius of geometry at the ``reference_plane``.", @@ -215,6 +225,124 @@ def _only_middle_for_infinite_length_slanted_cylinder(cls, val, values): ) return val + def to_polyslab( + self, num_pts_circumference: int = _N_PTS_CYLINDER_POLYSLAB, **kwargs + ) -> PolySlab: + """Convert instance of ``Cylinder`` into a discretized version using ``PolySlab``. + + Parameters + ---------- + num_pts_circumference : int = 51 + Number of points in the circumference of the discretized polyslab. + **kwargs: + Extra keyword arguments passed to ``PolySlab()``, such as ``dilation``. + + Returns + ------- + PolySlab + Extruded polygon representing a discretized version of the cylinder. + """ + + center_axis = self.center_axis + length_axis = self.length_axis + slab_bounds = (center_axis - length_axis / 2.0, center_axis + length_axis / 2.0) + + if num_pts_circumference < 3: + raise ValueError("'PolySlab' from 'Cylinder' must have 3 or more radius points.") + + _, (x0, y0) = self.pop_axis(self.center, axis=self.axis) + + xs_, ys_ = self._points_unit_circle(num_pts_circumference=num_pts_circumference) + + xs = x0 + self.radius * xs_ + ys = y0 + self.radius * ys_ + + vertices = anp.stack((xs, ys), axis=-1) + + return PolySlab( + vertices=vertices, + axis=self.axis, + slab_bounds=slab_bounds, + sidewall_angle=self.sidewall_angle, + reference_plane=self.reference_plane, + **kwargs, + ) + + def _points_unit_circle( + self, num_pts_circumference: int = _N_PTS_CYLINDER_POLYSLAB + ) -> np.ndarray: + """Set of x and y points for the unit circle when discretizing cylinder as a polyslab.""" + angles = np.linspace(0, 2 * np.pi, num_pts_circumference, endpoint=False) + xs = np.cos(angles) + ys = np.sin(angles) + return np.stack((xs, ys), axis=0) + + def compute_derivatives(self, derivative_info: DerivativeInfo) -> AutogradFieldMap: + """Compute the adjoint derivatives for this object.""" + + # compute number of points in the circumference of the polyslab using resolution info + wvl0 = C_0 / derivative_info.frequency + wvl_mat = wvl0 / max(1.0, np.sqrt(abs(derivative_info.eps_in))) + + circumference = 2 * np.pi * self.radius + wvls_in_circumference = circumference / wvl_mat + + num_pts_circumference = int( + np.ceil(_PTS_PER_WVL_MAT_CYLINDER_DISCRETIZE * wvls_in_circumference) + ) + num_pts_circumference = max(3, num_pts_circumference) + + # construct equivalent polyslab and compute the derivatives + polyslab = self.to_polyslab(num_pts_circumference=num_pts_circumference) + + derivative_info_polyslab = derivative_info.updated_copy(paths=[("vertices",)], deep=False) + vjps_polyslab = polyslab.compute_derivatives(derivative_info_polyslab) + + vjps_vertices_xs, vjps_vertices_ys = vjps_polyslab[("vertices",)].T + + # transform polyslab vertices derivatives into Cylinder parameter derivatives + xs_, ys_ = self._points_unit_circle(num_pts_circumference=num_pts_circumference) + vjp_xs = np.sum(xs_ * vjps_vertices_xs) + vjp_ys = np.sum(ys_ * vjps_vertices_ys) + + vjps = {} + for path in derivative_info.paths: + if path == ("radius",): + vjps[path] = vjp_xs + vjp_ys + + elif "center" in path: + _, center_index = path + if center_index == self.axis: + raise NotImplementedError( + "Currently cannot differentiate Cylinder with respect to its 'center' along" + " the axis. If you would like this feature added, please feel free to raise" + " an issue on the tidy3d front end repository." + ) + + _, (index_x, index_y) = self.pop_axis((0, 1, 2), axis=self.axis) + if center_index == index_x: + vjps[path] = np.sum(vjps_vertices_xs) + elif center_index == index_y: + vjps[path] = np.sum(vjps_vertices_ys) + else: + raise ValueError( + "Something unexpected happened. Was asked to differentiate " + f"with respect to 'Cylinder.center[{center_index}]', but this was not " + "detected as being one of the parallel axis with " + f"'Cylinder.axis' of '{self.axis}'. If you received this error, please raise " + "an issue on the tidy3d front end repository with details about how you " + "defined your 'Cylinder' in the objective function." + ) + + else: + raise NotImplementedError( + f"Differentiation with respect to 'Cylinder' '{path}' field not supported. " + "If you would like this feature added, please feel free to raise " + "an issue on the tidy3d front end repository." + ) + + return vjps + @property def center_axis(self): """Gets the position of the center of the geometry in the out of plane dimension.""" @@ -359,13 +487,15 @@ def _intersections_normal(self, z: float): `Shapely's Documentation `_. """ + static_self = self.to_static() + # radius at z - radius_offset = self._radius_z(z) + radius_offset = static_self._radius_z(z) if radius_offset <= 0: return [] - _, (x0, y0) = self.pop_axis(self.center, axis=self.axis) + _, (x0, y0) = self.pop_axis(static_self.center, axis=self.axis) return [shapely.Point(x0, y0).buffer(radius_offset, quad_segs=_N_SHAPELY_QUAD_SEGS)] def _intersections_side(self, position, axis): diff --git a/tidy3d/components/geometry/triangulation.py b/tidy3d/components/geometry/triangulation.py index e581c987f..a6a6e00d2 100644 --- a/tidy3d/components/geometry/triangulation.py +++ b/tidy3d/components/geometry/triangulation.py @@ -4,6 +4,7 @@ import numpy as np import shapely +from ...exceptions import Tidy3dError from ..types import ArrayFloat1D, ArrayFloat2D @@ -17,8 +18,8 @@ class Vertex: Vertex coordinate. index : int Vertex index in the original polygon. - is_convex : bool = False - Flag indicating whether this is a convex vertex in the polygon. + convexity : float = 0.0 + Value representing the convexity (> 0) or concavity (< 0) of the vertex in the polygon. is_ear : bool = False Flag indicating whether this is an ear of the polygon. """ @@ -27,12 +28,12 @@ class Vertex: index: int - is_convex: bool = False + convexity: float - is_ear: bool = False + is_ear: bool -def update_convexity(vertices: List[Vertex], i: int) -> None: +def update_convexity(vertices: List[Vertex], i: int) -> int: """Update the convexity of a vertex in a polygon. Parameters @@ -41,15 +42,30 @@ def update_convexity(vertices: List[Vertex], i: int) -> None: Vertices of the polygon. i : int Index of the vertex to be updated. + + Returns + ------- + int + Value indicating vertex convexity change w.r.t. 0. See note below. + + Note + ---- + Besides updating the vertex, this function returns a value indicating whether the updated vertex + convexity changed to or from 0 (0 convexity means the vertex is collinear with its neighbors). + If the convexity changes from zero to non-zero, return -1. If it changes from non-zero to zero, + return +1. Return 0 in any other case. This allows the main triangulation loop to keep track of + the total number of collinear vertices in the polygon. + """ + result = -1 if vertices[i].convexity == 0.0 else 0 j = (i + 1) % len(vertices) - vertices[i].is_convex = ( - np.cross( - vertices[i].coordinate - vertices[i - 1].coordinate, - vertices[j].coordinate - vertices[i].coordinate, - ) - > 0 + vertices[i].convexity = np.cross( + vertices[i].coordinate - vertices[i - 1].coordinate, + vertices[j].coordinate - vertices[i].coordinate, ) + if vertices[i].convexity == 0.0: + result += 1 + return result def is_inside( @@ -87,10 +103,10 @@ def update_ear_flag(vertices: List[Vertex], i: int) -> None: h = (i - 1) % len(vertices) j = (i + 1) % len(vertices) triangle = (vertices[h].coordinate, vertices[i].coordinate, vertices[j].coordinate) - vertices[i].is_ear = vertices[i].is_convex and not any( + vertices[i].is_ear = vertices[i].convexity > 0 and not any( is_inside(v.coordinate, triangle) for k, v in enumerate(vertices) - if not (v.is_convex or k == h or k == i or k == j) + if not (v.convexity > 0 or k == h or k == i or k == j) ) @@ -110,35 +126,50 @@ def triangulate(vertices: ArrayFloat2D) -> List[Tuple[int, int, int]]: List of indices of the vertices of the triangles. """ is_ccw = shapely.LinearRing(vertices).is_ccw - vertices = [Vertex(v, i) for i, v in enumerate(vertices)] + + # Initialize vertices as non-collinear because we will update the actual value below and count + # the number of collinear vertices. + vertices = [Vertex(v, i, -1.0, False) for i, v in enumerate(vertices)] if not is_ccw: vertices.reverse() + collinears = 0 for i in range(len(vertices)): - update_convexity(vertices, i) + collinears += update_convexity(vertices, i) for i in range(len(vertices)): update_ear_flag(vertices, i) triangles = [] + ear_found = True while len(vertices) > 3: + if not ear_found: + raise Tidy3dError( + "Impossible to triangulate polygon. Verify that the polygon is valid." + ) + ear_found = False i = 0 while i < len(vertices): if vertices[i].is_ear: - j = (i + 1) % len(vertices) - triangles.append((vertices[i - 1].index, vertices[i].index, vertices[j].index)) - vertices.pop(i) - if len(vertices) == 3: - break + removed = vertices.pop(i) h = (i - 1) % len(vertices) j = i % len(vertices) - if not vertices[h].is_convex: - update_convexity(vertices, h) - if not vertices[j].is_convex: - update_convexity(vertices, j) - update_ear_flag(vertices, h) - update_ear_flag(vertices, j) + collinears += update_convexity(vertices, h) + collinears += update_convexity(vertices, j) + if collinears == len(vertices): + # Undo removal because only collinear vertices remain + vertices.insert(i, removed) + collinears += update_convexity(vertices, (i - 1) % len(vertices)) + collinears += update_convexity(vertices, (i + 1) % len(vertices)) + i += 1 + else: + ear_found = True + triangles.append((vertices[h].index, removed.index, vertices[j].index)) + update_ear_flag(vertices, h) + update_ear_flag(vertices, j) + if len(vertices) == 3: + break else: i += 1 diff --git a/tidy3d/components/geometry/utils.py b/tidy3d/components/geometry/utils.py index 043d2f932..72415f5f3 100644 --- a/tidy3d/components/geometry/utils.py +++ b/tidy3d/components/geometry/utils.py @@ -4,7 +4,7 @@ from enum import Enum from math import isclose -from typing import Tuple, Union +from typing import Optional, Tuple, Union import numpy as np import pydantic as pydantic @@ -30,17 +30,25 @@ ] -def flatten_groups(*geometries: GeometryType, flatten_nonunion_type: bool = False) -> GeometryType: +def flatten_groups( + *geometries: GeometryType, + flatten_nonunion_type: bool = False, + flatten_transformed: bool = False, + transform: Optional[MatrixReal4x4] = None, +) -> GeometryType: """Iterates over all geometries, flattening groups and unions. Parameters ---------- *geometries : GeometryType Geometries to flatten. - flatten_nonunion_type : bool = False If ``False``, only flatten geometry unions (and ``GeometryGroup``). If ``True``, flatten all clip operations. + flatten_transformed : bool = False + If ``True``, ``Transformed`` groups are flattened into individual transformed geometries. + transform : Optional[MatrixReal4x4] + Accumulated transform from parents. Only used when ``flatten_transformed`` is ``True``. Yields ------ @@ -50,7 +58,10 @@ def flatten_groups(*geometries: GeometryType, flatten_nonunion_type: bool = Fals for geometry in geometries: if isinstance(geometry, base.GeometryGroup): yield from flatten_groups( - *geometry.geometries, flatten_nonunion_type=flatten_nonunion_type + *geometry.geometries, + flatten_nonunion_type=flatten_nonunion_type, + flatten_transformed=flatten_transformed, + transform=transform, ) elif isinstance(geometry, base.ClipOperation) and ( flatten_nonunion_type or geometry.operation == "union" @@ -59,7 +70,21 @@ def flatten_groups(*geometries: GeometryType, flatten_nonunion_type: bool = Fals geometry.geometry_a, geometry.geometry_b, flatten_nonunion_type=flatten_nonunion_type, + flatten_transformed=flatten_transformed, + transform=transform, + ) + elif flatten_transformed and isinstance(geometry, base.Transformed): + new_transform = geometry.transform + if transform is not None: + new_transform = np.matmul(transform, new_transform) + yield from flatten_groups( + geometry.geometry, + flatten_nonunion_type=flatten_nonunion_type, + flatten_transformed=flatten_transformed, + transform=new_transform, ) + elif flatten_transformed and transform is not None: + yield base.Transformed(geometry=geometry, transform=transform) else: yield geometry diff --git a/tidy3d/components/grid/grid_spec.py b/tidy3d/components/grid/grid_spec.py index 2c401705d..872f594c2 100644 --- a/tidy3d/components/grid/grid_spec.py +++ b/tidy3d/components/grid/grid_spec.py @@ -717,11 +717,24 @@ def make_grid( ) grids_1d = [self.grid_x, self.grid_y, self.grid_z] + all_structures = list(structures) + [s.to_static() for s in self.override_structures] + + if any(s.strip_traced_fields() for s in self.override_structures): + log.warning( + "The override structures were detected as having a dependence on the objective " + "function parameters. This is not supported by our automatic differentiation " + "framework. The derivative will be un-traced through the override structures. " + "To make this explicit and remove this warning, use 'y = autograd.tracer.getval(x)'" + " to remove any derivative information from values being passed to create " + "override structures. Alternatively, 'obj = obj.to_static()' will create a copy of " + "an instance without any autograd tracers." + ) + coords_dict = {} for idim, (dim, grid_1d) in enumerate(zip("xyz", grids_1d)): coords_dict[dim] = grid_1d.make_coords( axis=idim, - structures=list(structures) + list(self.override_structures), + structures=all_structures, symmetry=symmetry, periodic=periodic[idim], wavelength=wavelength, diff --git a/tidy3d/components/medium.py b/tidy3d/components/medium.py index c6426a932..020078b78 100644 --- a/tidy3d/components/medium.py +++ b/tidy3d/components/medium.py @@ -168,6 +168,12 @@ def _validate_medium_freqs(self, medium: AbstractMedium, freqs: List[pd.Positive """Any additional validation that depends on the central frequencies of the sources.""" pass + def _hardcode_medium_freqs( + self, medium: AbstractMedium, freqs: List[pd.PositiveFloat] + ) -> NonlinearSpec: + """Update the nonlinear model to hardcode information on medium and freqs.""" + return self + def _get_freq0(self, freq0, freqs: List[pd.PositiveFloat]) -> float: """Get a single value for freq0.""" @@ -435,8 +441,24 @@ def _validate_medium_freqs(self, medium: AbstractMedium, freqs: List[pd.Positive "gain medium are unstable, and are likely to diverge." ) + def _hardcode_medium_freqs( + self, medium: AbstractMedium, freqs: List[pd.PositiveFloat] + ) -> TwoPhotonAbsorption: + """Update the nonlinear model to hardcode information on medium and freqs.""" + n0 = self._get_n0(n0=self.n0, medium=medium, freqs=freqs) + freq0 = self._get_freq0(freq0=self.freq0, freqs=freqs) + return self.updated_copy(n0=n0, freq0=freq0) + def _validate_medium(self, medium: AbstractMedium): """Check that the model is compatible with the medium.""" + log.warning( + "Found a medium with a 'TwoPhotonAbsorption' nonlinearity. " + "This uses a phenomenological model based on complex fields, " + "so care should be taken in interpreting the results. For more " + "information on the model, see the documentation at " + "'https://docs.flexcompute.com/projects/tidy3d/en/latest/api/_autosummary/tidy3d.TwoPhotonAbsorption.html' or the following reference: " + "'N. Suzuki, \"FDTD Analysis of Two-Photon Absorption and Free-Carrier Absorption in Si High-Index-Contrast Waveguides,\" J. Light. Technol. 25, 9 (2007).'." + ) # if n0 is specified, we can go ahead and validate passivity if self.n0 is not None: self._validate_medium_freqs(medium, []) @@ -476,7 +498,7 @@ class KerrNonlinearity(NonlinearModel): The fields in this equation are complex-valued, allowing a direct implementation of the Kerr nonlinearity. In contrast, the model :class:`.NonlinearSusceptibility` implements a chi3 nonlinear susceptibility using real-valued fields, giving rise to Kerr nonlinearity - as well as third-harmonic generation. The relationship between the parameters is given by + as well as third-harmonic generation and other effects. The relationship between the parameters is given by :math:`n_2 = \\frac{3}{4} \\frac{1}{\\varepsilon_0 c_0 n_0 \\operatorname{Re}(n_0)} \\chi_3`. The additional factor of :math:`\\frac{3}{4}` comes from the usage of complex-valued fields for the Kerr nonlinearity and real-valued fields for the nonlinear susceptibility. @@ -524,8 +546,22 @@ def _validate_medium_freqs(self, medium: AbstractMedium, freqs: List[pd.Positive "gain medium are unstable, and are likely to diverge." ) + def _hardcode_medium_freqs( + self, medium: AbstractMedium, freqs: List[pd.PositiveFloat] + ) -> KerrNonlinearity: + """Update the nonlinear model to hardcode information on medium and freqs.""" + n0 = self._get_n0(n0=self.n0, medium=medium, freqs=freqs) + return self.updated_copy(n0=n0) + def _validate_medium(self, medium: AbstractMedium): """Check that the model is compatible with the medium.""" + log.warning( + "Found a medium with a 'KerrNonlinearity'. Usually, " + "'NonlinearSusceptibility' is preferred, as it captures " + "additional physical effects by acting on the underlying real fields. " + "The relation between the parameters is " + "'chi3 = (4/3) * eps_0 * c_0 * n0 * Re(n0) * n2'." + ) # if n0 is specified, we can go ahead and validate passivity if self.n0 is not None: self._validate_medium_freqs(medium, []) @@ -593,6 +629,16 @@ def _validate_num_iters(cls, val, values): ) return val + def _hardcode_medium_freqs( + self, medium: AbstractMedium, freqs: List[pd.PositiveFloat] + ) -> NonlinearSpec: + """Update the nonlinear spec to hardcode information on medium and freqs.""" + new_models = [] + for model in self.models: + new_model = model._hardcode_medium_freqs(medium=medium, freqs=freqs) + new_models.append(new_model) + return self.updated_copy(models=new_models) + class AbstractMedium(ABC, Tidy3dBaseModel): """A medium within which electromagnetic waves propagate.""" @@ -1455,7 +1501,9 @@ def _derivative_field_cmp( # TODO: probably this could be more robust. eg if the DataArray has weird edge cases E_der_dim = E_der_map[f"E{dim}"] - E_der_dim_interp = E_der_dim.interp(**coords_interp).fillna(0.0).sum(dims_sum).sum("f") + E_der_dim_interp = ( + E_der_dim.interp(**coords_interp, assume_sorted=True).fillna(0.0).sum(dims_sum).sum("f") + ) vjp_array = np.array(E_der_dim_interp.values).astype(complex) vjp_array = vjp_array.reshape(eps_data.shape) @@ -1652,6 +1700,8 @@ def from_nk(cls, n: float, k: float, freq: float, **kwargs): Imaginary part of refrative index. freq : float Frequency to evaluate permittivity at (Hz). + kwargs: dict + Keyword arguments passed to the medium construction. Returns ------- @@ -1696,7 +1746,10 @@ def derivative_eps_sigma_volume( freqs = vjp_eps_complex.coords["f"].values values = vjp_eps_complex.values - eps_vjp, sigma_vjp = self.eps_complex_to_eps_sigma(eps_complex=values, freq=freqs) + # vjp of eps_complex_to_eps_sigma + omegas = 2 * np.pi * freqs + eps_vjp = np.real(values) + sigma_vjp = -np.imag(values) / omegas / EPSILON_0 eps_vjp = np.sum(eps_vjp) sigma_vjp = np.sum(sigma_vjp) @@ -2423,6 +2476,8 @@ def from_nk( interp_method : :class:`.InterpMethod`, optional Interpolation method to obtain permittivity values that are not supplied at the Yee grids. + kwargs: dict + Keyword arguments passed to the medium construction. Note ---- @@ -2644,7 +2699,9 @@ def _derivative_field_cmp( # TODO: probably this could be more robust. eg if the DataArray has weird edge cases E_der_dim = E_der_map[f"E{dim}"] - E_der_dim_interp = E_der_dim.interp(**coords_interp).fillna(0.0).sum(dims_sum).real + E_der_dim_interp = ( + E_der_dim.interp(**coords_interp, assume_sorted=True).fillna(0.0).sum(dims_sum).real + ) E_der_dim_interp = E_der_dim_interp.sum("f") vjp_array = np.array(E_der_dim_interp.values, dtype=float) @@ -4181,6 +4238,8 @@ def from_nk(cls, n: float, k: float, freq: float, **kwargs): Imaginary part of refrative index. freq : float Frequency to evaluate permittivity at (Hz). + kwargs: dict + Keyword arguments passed to the medium construction. Returns ------- @@ -4219,6 +4278,7 @@ def from_nk(cls, n: float, k: float, freq: float, **kwargs): coeffs=[ (eps_i, fp, delta_p), ], + **kwargs, ) @@ -5476,6 +5536,18 @@ def from_diagonal(cls, xx: Medium, yy: Medium, zz: Medium, rotation: RotationTyp Resulting fully anisotropic medium. """ + if any(comp.nonlinear_spec is not None for comp in [xx, yy, zz]): + raise ValidationError( + "Nonlinearities are not currently supported for the components " + "of a fully anisotropic medium." + ) + + if any(comp.modulation_spec is not None for comp in [xx, yy, zz]): + raise ValidationError( + "Modulation is not currently supported for the components " + "of a fully anisotropic medium." + ) + permittivity_diag = np.diag([comp.permittivity for comp in [xx, yy, zz]]).tolist() conductivity_diag = np.diag([comp.conductivity for comp in [xx, yy, zz]]).tolist() @@ -6777,6 +6849,8 @@ def medium_from_nk(n: float, k: float, freq: float, **kwargs) -> Union[Medium, L Imaginary part of refrative index. freq : float Frequency to evaluate permittivity at (Hz). + kwargs: dict + Keyword arguments passed to the medium construction. Returns ------- diff --git a/tidy3d/components/scene.py b/tidy3d/components/scene.py index 11ca0a3ee..13ee33b0f 100644 --- a/tidy3d/components/scene.py +++ b/tidy3d/components/scene.py @@ -133,7 +133,7 @@ def _validate_num_geometries(cls, val): return val for i, structure in enumerate(val): - for geometry in flatten_groups(structure.geometry): + for geometry in flatten_groups(structure.geometry, flatten_transformed=True): count = sum( 1 for g in traverse_geometries(geometry) @@ -257,6 +257,7 @@ def intersecting_media( List[:class:`.AbstractMedium`] Set of distinct mediums that intersect with the given planar object. """ + structures = [s.to_static() for s in structures] if test_object.size.count(0.0) == 1: # get all merged structures on the test_object, which is already planar structures_merged = Scene._filter_structures_plane_medium(structures, test_object) @@ -417,7 +418,7 @@ def plot_structures( """ medium_shapes = self._get_structures_2dbox( - structures=self.structures, x=x, y=y, z=z, hlim=hlim, vlim=vlim + structures=self.to_static().structures, x=x, y=y, z=z, hlim=hlim, vlim=vlim ) medium_map = self.medium_map for medium, shape in medium_shapes: @@ -654,7 +655,7 @@ def _filter_structures_plane( if _shape.is_empty or not shape.intersects(_shape): continue - diff_shape = _shape - shape + diff_shape = (_shape - shape).buffer(0) # different prop, remove intersection from background shape if prop != _prop and len(diff_shape.bounds) > 0: diff --git a/tidy3d/components/simulation.py b/tidy3d/components/simulation.py index 0e4c509d5..7b45310a6 100644 --- a/tidy3d/components/simulation.py +++ b/tidy3d/components/simulation.py @@ -5,7 +5,7 @@ import math import pathlib from abc import ABC, abstractmethod -from typing import Dict, List, Set, Tuple, Union +from typing import Dict, List, Optional, Set, Tuple, Union import autograd.numpy as np import matplotlib as mpl @@ -30,6 +30,7 @@ PMCBoundary, StablePML, ) +from .data.data_array import FreqDataArray from .data.dataset import CustomSpatialDataType, Dataset from .geometry.base import Box, Geometry from .geometry.mesh import TriangleMesh @@ -85,7 +86,11 @@ from .structure import MeshOverrideStructure, Structure from .subpixel_spec import SubpixelSpec from .types import TYPE_TAG_STR, Ax, Axis, FreqBound, InterpMethod, Literal, Symmetry, annotate_type -from .validators import assert_objects_in_sim_bounds, validate_mode_objects_symmetry +from .validators import ( + assert_objects_in_sim_bounds, + validate_mode_objects_symmetry, + validate_mode_plane_radius, +) from .viz import ( PlotParams, add_ax_if_none, @@ -204,11 +209,20 @@ class AbstractYeeGridSimulation(AbstractSimulation, ABC): ", or ``False`` to apply staircasing.", ) - simulation_type: Literal["autograd_fwd", "autograd_bwd", None] = pydantic.Field( - None, - title="Simulation Type", - description="Tag used internally to distinguish types of simulations for " - "``autograd`` gradient processing.", + simulation_type: Optional[Literal["autograd_fwd", "autograd_bwd", "tidy3d", None]] = ( + pydantic.Field( + "tidy3d", + title="Simulation Type", + description="Tag used internally to distinguish types of simulations for " + "``autograd`` gradient processing.", + ) + ) + + post_norm: Union[float, FreqDataArray] = pydantic.Field( + 1.0, + title="Post Normalization Values", + description="Factor to multiply the fields by after running, " + "given the adjoint source pipeline used. Note: this is used internally only.", ) """ @@ -255,6 +269,13 @@ class AbstractYeeGridSimulation(AbstractSimulation, ABC): * `Dielectric constant assignment on Yee grids `_ """ + @pydantic.validator("simulation_type", always=True) + def _validate_simulation_type_tidy3d(cls, val): + """Enforce the simulation_type is 'tidy3d' if passed as None for bkwrds compatibility.""" + if val is None: + return "tidy3d" + return val + @pydantic.validator("lumped_elements", always=True) @skip_if_fields_missing(["structures"]) def _validate_num_lumped_elements(cls, val, values): @@ -1292,7 +1313,7 @@ def snap_to_grid(geom: Geometry, axis: Axis) -> Geometry: lumped_structures.append(lumped_element.to_structure(self.grid)) # Begin volumetric structures grid - all_structures = list(self.structures) + lumped_structures + all_structures = list(self.static_structures) + lumped_structures # For 1D and 2D simulations, a nonzero size is needed for the polygon operations in subdivide placeholder_size = tuple(i if i > 0 else inf for i in self.geometry.size) @@ -1482,7 +1503,7 @@ def subsection( if remove_outside_structures: new_structures = [strc for strc in self.structures if new_box.intersects(strc.geometry)] else: - new_structures = self.structures + new_structures = list(self.structures) # If ``validate_geometries=False``, use aux structures whose geometry is replaced by its bounding box # so that other validations are still performed. @@ -1500,6 +1521,21 @@ def subsection( if sources is None: sources = [src for src in self.sources if new_box.intersects(src)] + # some nonlinear materials depend on the central frequency + # we update them with hardcoded freq0 + freqs = np.array([source.source_time.freq0 for source in self.sources]) + for i, structure in enumerate(new_structures): + medium = structure.medium + nonlinear_spec = medium.nonlinear_spec + if nonlinear_spec is not None: + new_nonlinear_spec = nonlinear_spec._hardcode_medium_freqs( + medium=medium, freqs=freqs + ) + new_structure = structure.updated_copy( + nonlinear_spec=new_nonlinear_spec, path="medium" + ) + new_structures[i] = new_structure + if monitors is None: monitors = [mnt for mnt in self.monitors if new_box.intersects(mnt)] @@ -2381,6 +2417,46 @@ def plane_wave_boundaries(cls, val, values): ) return val + @pydantic.validator("monitors", always=True) + @skip_if_fields_missing(["boundary_spec", "medium", "size", "structures", "sources"]) + def bloch_boundaries_diff_mnt(cls, val, values): + """Error if there are diffraction monitors incompatible with boundary conditions.""" + + monitors = val + + if not val or not any(isinstance(mnt, DiffractionMonitor) for mnt in monitors): + return val + + boundaries = values.get("boundary_spec").to_list + sources = values.get("sources") + size = values.get("size") + sim_medium = values.get("medium") + structures = values.get("structures") + for source_ind, source in enumerate(sources): + if not isinstance(source, PlaneWave): + continue + + _, tan_dirs = cls.pop_axis([0, 1, 2], axis=source.injection_axis) + medium_set = Scene.intersecting_media(source, structures) + medium = medium_set.pop() if medium_set else sim_medium + + for tan_dir in tan_dirs: + boundary = boundaries[tan_dir] + + # check the Bloch boundary + angled plane wave case + num_bloch = sum(isinstance(bnd, (Periodic, BlochBoundary)) for bnd in boundary) + if num_bloch > 0: + cls._check_bloch_vec( + source=source, + source_ind=source_ind, + bloch_vec=boundary[0].bloch_vec, + dim=tan_dir, + medium=medium, + domain_size=size[tan_dir], + has_diff_mnt=True, + ) + return val + @pydantic.validator("boundary_spec", always=True) @skip_if_fields_missing(["medium", "center", "size", "structures", "sources"]) def tfsf_boundaries(cls, val, values): @@ -2468,7 +2544,7 @@ def tfsf_with_symmetry(cls, val, values): @pydantic.validator("boundary_spec", always=True) @skip_if_fields_missing(["size", "symmetry"]) def boundaries_for_zero_dims(cls, val, values): - """Error if absorbing boundaries, unmatching pec/pmc, or symmetry is used along a zero dimension.""" + """Error if absorbing boundaries, bloch boundaries, unmatching pec/pmc, or symmetry is used along a zero dimension.""" boundaries = val.to_list size = values.get("size") symmetry = values.get("symmetry") @@ -2478,6 +2554,7 @@ def boundaries_for_zero_dims(cls, val, values): if size_dim == 0: axis = axis_names[dim] num_absorbing_bdries = sum(isinstance(bnd, AbsorberSpec) for bnd in boundary) + num_bloch_bdries = sum(isinstance(bnd, BlochBoundary) for bnd in boundary) if num_absorbing_bdries > 0: raise SetupError( @@ -2486,6 +2563,14 @@ def boundaries_for_zero_dims(cls, val, values): f"Use either 'Periodic' or 'BlochBoundary' along {axis}." ) + if num_bloch_bdries > 0: + raise SetupError( + f"The simulation has zero size along the {axis} axis, " + "using a Bloch boundary along such an axis is not supported because of " + "the Bloch vector definition in units of '2 * pi / (size along dimension)'. Use a small " + "but nonzero size along the dimension instead." + ) + if symmetry_dim != 0: raise SetupError( f"The simulation has zero size along the {axis} axis, so " @@ -2526,19 +2611,20 @@ def _validate_2d_geometry_has_2d_medium(cls, val, values): if val is None: return val - for i, structure in enumerate(val): - if isinstance(structure.medium, Medium2D): - continue - for geom in flatten_groups(structure.geometry): - zero_dims = geom.zero_dims - if len(zero_dims) > 0: - log.warning( - f"Structure at 'structures[{i}]' has geometry with zero size along " - f"dimensions {zero_dims}, and with a medium that is not a 'Medium2D'. " - "This is probably not correct, since the resulting simulation will " - "depend on the details of the numerical grid. Consider either " - "giving the geometry a nonzero thickness or using a 'Medium2D'." - ) + with log as consolidated_logger: + for i, structure in enumerate(val): + if isinstance(structure.medium, Medium2D): + continue + for geom in flatten_groups(structure.geometry): + zero_dims = geom.zero_dims + if len(zero_dims) > 0: + consolidated_logger.warning( + f"Structure at 'structures[{i}]' has geometry with zero size along " + f"dimensions {zero_dims}, and with a medium that is not a 'Medium2D'. " + "This is probably not correct, since the resulting simulation will " + "depend on the details of the numerical grid. Consider either " + "giving the geometry a nonzero thickness or using a 'Medium2D'." + ) return val @@ -2937,7 +3023,7 @@ def _projection_monitors_distance(cls, val, values): @pydantic.validator("monitors", always=True) @skip_if_fields_missing(["size"]) - def _projection_monitors_2d(cls, val, values): + def _projection_mnts_2d(cls, val, values): """ Validate if the field projection monitor is set up for a 2D simulation and ensure the observation parameters are configured correctly. @@ -2974,38 +3060,6 @@ def _projection_monitors_2d(cls, val, values): f"Monitor '{monitor.name}' is not supported in 1D simulations." ) - if isinstance(monitor, (FieldProjectionKSpaceMonitor)): - raise SetupError( - f"Monitor '{monitor.name}' in 2D simulations is coming soon. " - "Please use 'FieldProjectionAngleMonitor' instead." - "Please use 'FieldProjectionAngleMonitor' or 'FieldProjectionCartesianMonitor' instead." - ) - - if isinstance(monitor, (FieldProjectionCartesianMonitor)): - config = { - "y-z": {"valid_proj_axes": [1, 2], "coord": ["x", "x"]}, - "x-z": {"valid_proj_axes": [0, 2], "coord": ["x", "y"]}, - "x-y": {"valid_proj_axes": [0, 1], "coord": ["y", "y"]}, - }[plane] - - valid_proj_axes = config["valid_proj_axes"] - invalid_proj_axis = [i for i in range(3) if i not in valid_proj_axes] - - if monitor.proj_axis in invalid_proj_axis: - raise SetupError( - f"For a 2D simulation in the {plane} plane, the 'proj_axis' of " - f"monitor '{monitor.name}' should be set to one of {valid_proj_axes}." - ) - - for idx, axis in enumerate(valid_proj_axes): - coord = getattr(monitor, config["coord"][idx]) - if monitor.proj_axis == axis and not all(value in [0] for value in coord): - raise SetupError( - f"For a 2D simulation in the {plane} plane with " - f"'proj_axis = {monitor.proj_axis}', '{config['coord'][idx]}' of monitor " - f"'{monitor.name}' should be set to '[0]'." - ) - if isinstance(monitor, FieldProjectionAngleMonitor): config = { "y-z": {"valid_value": [np.pi / 2, 3 * np.pi / 2], "coord": "phi"}, @@ -3031,6 +3085,39 @@ def _projection_monitors_2d(cls, val, values): f"'{valid_values_str}'" ) + continue + + if isinstance(monitor, (FieldProjectionCartesianMonitor)): + config = { + "y-z": {"valid_proj_axes": [1, 2], "coord": ["x", "x"]}, + "x-z": {"valid_proj_axes": [0, 2], "coord": ["x", "y"]}, + "x-y": {"valid_proj_axes": [0, 1], "coord": ["y", "y"]}, + }[plane] + elif isinstance(monitor, (FieldProjectionKSpaceMonitor)): + config = { + "y-z": {"valid_proj_axes": [1, 2], "coord": ["ux", "ux"]}, + "x-z": {"valid_proj_axes": [0, 2], "coord": ["ux", "uy"]}, + "x-y": {"valid_proj_axes": [0, 1], "coord": ["uy", "uy"]}, + }[plane] + + valid_proj_axes = config["valid_proj_axes"] + invalid_proj_axis = [i for i in range(3) if i not in valid_proj_axes] + + if monitor.proj_axis in invalid_proj_axis: + raise SetupError( + f"For a 2D simulation in the {plane} plane, the 'proj_axis' of " + f"monitor '{monitor.name}' should be set to one of {valid_proj_axes}." + ) + + for idx, axis in enumerate(valid_proj_axes): + coord = getattr(monitor, config["coord"][idx]) + if monitor.proj_axis == axis and not all(value in [0] for value in coord): + raise SetupError( + f"For a 2D simulation in the {plane} plane with " + f"'proj_axis = {monitor.proj_axis}', '{config['coord'][idx]}' of monitor " + f"'{monitor.name}' should be set to '[0]'." + ) + return val @pydantic.validator("monitors", always=True) @@ -3218,6 +3305,24 @@ def _post_init_validators(self) -> None: self._validate_tfsf_nonuniform_grid() self._validate_nonlinear_specs() self._validate_custom_source_time() + self._validate_mode_object_bends() + + def _validate_mode_object_bends(self) -> None: + """Error if any mode sources or monitors with bends have a radius that is too small.""" + for imnt, monitor in enumerate(self.monitors): + if isinstance(monitor, AbstractModeMonitor): + validate_mode_plane_radius( + mode_spec=monitor.mode_spec, + plane=monitor.geometry, + msg_prefix=f"Monitor at 'monitors[{imnt}]' ", + ) + for isrc, source in enumerate(self.sources): + if isinstance(source, ModeSource): + validate_mode_plane_radius( + mode_spec=source.mode_spec, + plane=source.geometry, + msg_prefix=f"Source at 'sources[{isrc}]' ", + ) def _validate_custom_source_time(self): """Warn if all simulation times are outside CustomSourceTime definition range.""" @@ -3246,6 +3351,7 @@ def _validate_no_structures_pml(self) -> None: with log as consolidated_logger: for i, structure in enumerate(self.structures): geo_bounds = structure.geometry.bounds + warn = False # will only warn once per structure for sim_bound, geo_bound, pml_thick, bound_dim, pm_val in zip( sim_bounds, geo_bounds, pml_thicks, bound_spec, (-1, 1) ): @@ -3256,14 +3362,16 @@ def _validate_no_structures_pml(self) -> None: in_pml_plus = (pm_val > 0) and (sim_pos < geo_pos <= sim_pos_pml) in_pml_mnus = (pm_val < 0) and (sim_pos > geo_pos >= sim_pos_pml) if not isinstance(bound_edge, Absorber) and (in_pml_plus or in_pml_mnus): - consolidated_logger.warning( - f"A bound of Simulation.structures[{i}] was detected as being " - "within the simulation PML. We recommend extending structures to " - "infinity or completely outside of the simulation PML to avoid " - "unexpected effects when the structures are not translationally " - "invariant within the PML.", - custom_loc=["structures", i], - ) + warn = True + if warn: + consolidated_logger.warning( + f"A bound of Simulation.structures[{i}] was detected as being " + "within the simulation PML. We recommend extending structures to " + "infinity or completely outside of the simulation PML to avoid " + "unexpected effects when the structures are not translationally " + "invariant within the PML.", + custom_loc=["structures", i], + ) def _validate_tfsf_nonuniform_grid(self) -> None: """Warn if the grid is nonuniform along the directions tangential to the injection plane, @@ -3818,6 +3926,7 @@ def _check_bloch_vec( dim: Axis, medium: MediumType, domain_size: float, + has_diff_mnt: bool = False, ): """Helper to check if a given Bloch vector is consistent with a given source.""" @@ -3830,10 +3939,13 @@ def _check_bloch_vec( if bloch_vec != expected_bloch_vec: test_val = np.real(expected_bloch_vec - bloch_vec) - if np.isclose(test_val % 1, 0) and not np.isclose(test_val, 0): + test_val_is_int = np.isclose(test_val, np.round(test_val)) + src_name = f" '{source.name}'" if source.name else "" + + if has_diff_mnt and test_val_is_int and not np.isclose(test_val, 0): # the given Bloch vector is offset by an integer log.warning( - f"The wave vector of source '{source.name}' along dimension " + f"The wave vector of source{src_name} along dimension " f"'{dim}' is equal to the Bloch vector of the simulation " "boundaries along that dimension plus an integer reciprocal " "lattice vector. If using a 'DiffractionMonitor', diffraction " @@ -3841,12 +3953,13 @@ def _check_bloch_vec( "of the source. Consider using 'BlochBoundary.from_source()'.", custom_loc=["boundary_spec", "xyz"[dim]], ) - elif not np.isclose(test_val % 1, 0): + + if not test_val_is_int: # the given Bloch vector is neither equal to the expected value, nor # off by an integer log.warning( f"The Bloch vector along dimension '{dim}' may be incorrectly " - f"set with respect to the source '{source.name}'. The absolute " + f"set with respect to the source{src_name}. The absolute " "difference between the expected and provided values in " "bandstructure units, up to an integer offset, is greater than " "1e-6. Consider using ``BlochBoundary.from_source()``, or " @@ -4297,7 +4410,7 @@ def grid(self) -> Grid: @cached_property def num_cells(self) -> int: - """Number of cells in the simulation. + """Number of cells in the simulation grid. Returns ------- @@ -4307,10 +4420,40 @@ def num_cells(self) -> int: return np.prod(self.grid.num_cells, dtype=np.int64) + @property + def _num_computational_grid_points_dim(self): + """Number of cells in the computational domain for this simulation along each dimension.""" + num_cells = self.grid.num_cells + num_cells_comp_domain = [] + # symmetry overrides other boundaries so should be checked first + for sym, npts, boundary in zip(self.symmetry, num_cells, self.boundary_spec.to_list): + if sym != 0: + num_cells_comp_domain.append(npts // 2 + 2) + elif isinstance(boundary[0], Periodic): + num_cells_comp_domain.append(npts) + else: + num_cells_comp_domain.append(npts + 2) + return num_cells_comp_domain + + @property + def num_computational_grid_points(self): + """Number of cells in the computational domain for this simulation. This is usually + different from ``num_cells`` due to the boundary conditions. Specifically, all boundary + conditions apart from ``Periodic`` require an extra pixel at the end of the simulation + domain. On the other hand, if a symmetry is present along a given dimension, only half of + the grid cells along that dimension will be in the computational domain. + + Returns + ------- + int + Number of yee cells in the computational domain corresponding to the simulation. + """ + return np.prod(self._num_computational_grid_points_dim, dtype=np.int64) + def get_refractive_indices(self, freq: float) -> list[float]: """List of refractive indices in the simulation at a given frequency.""" - eps_values = [structure.medium.eps_model(freq) for structure in self.structures] + eps_values = [structure.medium.eps_model(freq) for structure in self.static_structures] eps_values.append(self.medium.eps_model(freq)) return [AbstractMedium.eps_complex_to_nk(eps)[0] for eps in eps_values] diff --git a/tidy3d/components/source.py b/tidy3d/components/source.py index 663ad4a5d..da2f1f2d6 100644 --- a/tidy3d/components/source.py +++ b/tidy3d/components/source.py @@ -101,7 +101,7 @@ def frequency_range(self, num_fwidth: float = 4.0) -> FreqBound: """Frequency range within plus/minus ``num_fwidth * fwidth`` of the central frequency.""" @abstractmethod - def end_time(self) -> float | None: + def end_time(self) -> Optional[float]: """Time after which the source is effectively turned off / close to zero amplitude.""" @@ -192,7 +192,7 @@ def amp_time(self, time: float) -> complex: return pulse_amp - def end_time(self) -> float | None: + def end_time(self) -> Optional[float]: """Time after which the source is effectively turned off / close to zero amplitude.""" # TODO: decide if we should continue to return an end_time if the DC component remains @@ -251,7 +251,7 @@ def amp_time(self, time: float) -> complex: return const * offset * oscillation * amp - def end_time(self) -> float | None: + def end_time(self) -> Optional[float]: """Time after which the source is effectively turned off / close to zero amplitude.""" return None @@ -420,7 +420,7 @@ def amp_time(self, time: float) -> complex: return offset * oscillation * amp * envelope - def end_time(self) -> float | None: + def end_time(self) -> Optional[float]: """Time after which the source is effectively turned off / close to zero amplitude.""" if self.source_time_dataset is None: @@ -1134,6 +1134,14 @@ class GaussianBeam(AngledFieldSource, PlanarSource, BroadbandSource): ... direction='+', ... waist_radius=1.0) + Notes + -------- + If one wants the focus 'in front' of the source, a negative value of ``beam_distance`` is needed. + + .. image:: ../../_static/img/beam_waist.png + :width: 30% + :align: center + See Also -------- @@ -1152,10 +1160,11 @@ class GaussianBeam(AngledFieldSource, PlanarSource, BroadbandSource): 0.0, title="Waist Distance", description="Distance from the beam waist along the propagation direction. " - "When ``direction`` is ``+`` and ``waist_distance`` is positive, the waist " - "is on the ``-`` side (behind) the source plane. When ``direction`` is ``+`` and " - " ``waist_distance``is negative, the waist is on the ``+`` side (in front) of " - "the source plane.", + "A positive value means the waist is positioned behind the source, considering the propagation direction. " + "For example, for a beam propagating in the ``+`` direction, a positive value of ``beam_distance`` " + "means the beam waist is positioned in the ``-`` direction (behind the source). " + "A negative value means the beam waist is in the ``+`` direction (in front of the source). " + "For an angled source, the distance is defined along the rotated propagation direction.", units=MICROMETER, ) @@ -1213,10 +1222,11 @@ class TFSF(AngledFieldSource, VolumeSource): Notes ----- - The TFSF source injects :math:`\\frac{1 W}{\\mu m^2}` of power along the :attr:`injection_axis`. Note that in the - case of angled incidence, :math:`\\frac{1 W}{\\mu m^2}` is still injected along the source's :attr:`injection_axis`, - and not the propagation direction, unlike a :class:`PlaneWave` source. This allows computing - scattering and absorption cross-sections without the need for additional normalization. + The TFSF source injects :math:`1 W` of power per :math:`\\mu m^2` of source area along the :attr:`injection_axis`. + Hence, the normalization for the incident field is :math:`|E_0|^2 = \\frac{2}{c\\epsilon_0}`, for any source size. + Note that in the case of angled incidence, the same power is injected along the source's :attr:`injection_axis`, + and not the propagation direction. This allows computing scattering and absorption cross-sections + without the need for additional normalization. The TFSF source allows specifying a box region into which a plane wave is injected. Fields inside this region can be interpreted as the superposition of the incident field and the scattered field due to any scatterers diff --git a/tidy3d/components/structure.py b/tidy3d/components/structure.py index 497df175b..c8fc67b0c 100644 --- a/tidy3d/components/structure.py +++ b/tidy3d/components/structure.py @@ -4,20 +4,24 @@ import pathlib from collections import defaultdict -from typing import Callable, Optional, Tuple, Union +from typing import Optional, Tuple, Union +import autograd.numpy as anp import numpy as np import pydantic.v1 as pydantic -from ..constants import MICROMETER, PERMITTIVITY +from ..constants import MICROMETER from ..exceptions import SetupError, Tidy3dError, Tidy3dImportError +from ..log import log from .autograd.derivative_utils import DerivativeInfo -from .autograd.types import AutogradFieldMap +from .autograd.types import AutogradFieldMap, Box from .autograd.utils import get_static from .base import Tidy3dBaseModel, skip_if_fields_missing +from .data.data_array import ScalarFieldDataArray +from .geometry.polyslab import PolySlab from .geometry.utils import GeometryType, validate_no_transformed_polyslabs from .grid.grid import Coords -from .medium import AbstractCustomMedium, Medium2D, MediumType +from .medium import AbstractCustomMedium, CustomMedium, Medium, Medium2D, MediumType from .monitor import FieldMonitor, PermittivityMonitor from .types import TYPE_TAG_STR, Ax, Axis from .validators import validate_name_str @@ -54,11 +58,47 @@ class AbstractStructure(Tidy3dBaseModel): None, ge=1.0, title="Background Permittivity", - description="Relative permittivity used for the background of this structure " + description="DEPRECATED: Use ``Structure.background_medium``. " + "Relative permittivity used for the background of this structure " "when performing shape optimization with autograd.", - units=PERMITTIVITY, ) + background_medium: MediumType = pydantic.Field( + None, + title="Background Medium", + description="Medium used for the background of this structure " + "when performing shape optimization with autograd. This is required when the " + "structure is embedded in another structure as autograd will use the permittivity of the " + "``Simulation`` by default to compute the shape derivatives.", + ) + + @pydantic.root_validator(skip_on_failure=True) + def _handle_background_mediums(cls, values): + """Handle background medium combinations, including deprecation.""" + + background_permittivity = values.get("background_permittivity") + background_medium = values.get("background_medium") + + # old case, only permittivity supplied, warn and set the Medium automatically + if background_medium is None and background_permittivity is not None: + log.warning( + "'Structure.background_permittivity' is deprecated, " + "set the 'Structure.background_medium' directly using a 'Medium'. " + "Handling automatically using the supplied relative permittivity." + ) + values["background_medium"] = Medium(permittivity=background_permittivity) + + # both present, just make sure they are consistent, error if not + if background_medium is not None and background_permittivity is not None: + is_medium = isinstance(background_medium, Medium) + if not (is_medium and background_medium.permittivity == background_permittivity): + raise ValueError( + "Inconsistent 'background_permittivity' and 'background_medium'. " + "Use 'background_medium' only as 'background_permittivity' is deprecated." + ) + + return values + _name_validator = validate_name_str() @pydantic.validator("geometry") @@ -212,13 +252,33 @@ def make_adjoint_monitors( box = self.geometry.bounding_box # we dont want these fields getting traced by autograd, otherwise it messes stuff up - size = tuple(get_static(x) for x in box.size) # TODO: expand slightly? - center = tuple(get_static(x) for x in box.center) + + size = [get_static(x) for x in box.size] # TODO: expand slightly? + center = [get_static(x) for x in box.center] + + # polyslab only needs fields at the midpoint along axis + if isinstance(self.geometry, PolySlab): + size[self.geometry.axis] = 0 + + # custom medium only needs fields at center locations of unit cells. + if isinstance(self.medium, CustomMedium): + for axis, dim in enumerate("xyz"): + if self.medium.permittivity is not None: + if len(self.medium.permittivity.coords[dim]) == 1: + size[axis] = 0 + if self.medium.eps_dataset is not None: + zero_size = True + for _, fld in self.medium.eps_dataset.field_components.items(): + if len(fld.coords[dim]) != 1: + zero_size = False + if zero_size: + size[axis] = 0 mnt_fld = FieldMonitor( size=size, center=center, freqs=freqs, + fields=("Ex", "Ey", "Ez"), name=self.get_monitor_name(index=index, data_type="fld"), colocate=False, ) @@ -233,24 +293,6 @@ def make_adjoint_monitors( return mnt_fld, mnt_eps - @property - def derivative_function_map(self) -> dict[tuple[str, str], Callable]: - """Map path to the right derivative function function.""" - return { - ("medium", "permittivity"): self.derivative_medium_permittivity, - ("medium", "conductivity"): self.derivative_medium_conductivity, - ("geometry", "size"): self.derivative_geometry_size, - ("geometry", "center"): self.derivative_geometry_center, - } - - def get_derivative_function(self, path: tuple[str, ...]) -> Callable: - """Get the derivative function function.""" - - derivative_map = self.derivative_function_map - if path not in derivative_map: - raise NotImplementedError(f"Can't compute derivative for structure field path: {path}.") - return derivative_map[path] - def compute_derivatives(self, derivative_info: DerivativeInfo) -> AutogradFieldMap: """Compute adjoint gradients given the forward and adjoint fields""" @@ -274,7 +316,7 @@ def compute_derivatives(self, derivative_info: DerivativeInfo) -> AutogradFieldM for med_or_geo, field_paths in structure_fields_map.items(): # grab derivative values {field_name -> vjp_value} med_or_geo_field = self.medium if med_or_geo == "medium" else self.geometry - info = derivative_info.updated_copy(paths=field_paths) + info = derivative_info.updated_copy(paths=field_paths, deep=False) derivative_values_map = med_or_geo_field.compute_derivatives(derivative_info=info) # construct map of {field path -> derivative value} @@ -537,6 +579,66 @@ def to_gds_file( pathlib.Path(fname).parent.mkdir(parents=True, exist_ok=True) library.write_gds(fname) + @classmethod + def from_permittivity_array( + cls, geometry: GeometryType, eps_data: np.ndarray, **kwargs + ) -> Structure: + """Create ``Structure`` with ``geometry`` and ``CustomMedium`` containing ``eps_data`` for + The ``permittivity`` field. Extra keyword arguments are passed to ``td.Structure()``. + """ + + rmin, rmax = geometry.bounds + + if not isinstance(eps_data, (np.ndarray, Box, list, tuple)): + raise ValueError("Must supply array-like object for 'eps_data'.") + + eps_data = anp.array(eps_data) + shape = eps_data.shape + + if len(shape) != 3: + raise ValueError( + "'Structure.from_permittivity_array' method only accepts 'eps_data' with 3 dimensions, " + f"corresponding to (x,y,z). Got array with {len(shape)} dimensions." + ) + + coords = {} + for key, pt_min, pt_max, num_pts in zip("xyz", rmin, rmax, shape): + if np.isinf(pt_min) and np.isinf(pt_max): + pt_min = 0.0 + pt_max = 0.0 + + coords_2x = np.linspace(pt_min, pt_max, 2 * num_pts + 1) + coords_centers = coords_2x[1:-1:2] + + if len(coords_centers) != num_pts: + raise ValueError( + "something went wrong, different number of coordinate values and data values. " + "Check your 'geometry', 'eps_data', and file a bug report." + ) + + # handle infinite size dimension edge case + coords_centers = np.nan_to_num(coords_centers, 0.0) + + _, count = np.unique(coords_centers, return_counts=True) + if np.any(count > 1): + raise ValueError( + "Found duplicates in the coordinates constructed from the supplied " + "'geometry' and 'eps_data'. This is likely due to having a geometry with an " + "infinite size in one dimension and a 'eps_data' with a 'shape' > 1 in that " + "dimension. " + ) + + coords[key] = coords_centers + + eps_data_array = ScalarFieldDataArray(eps_data, coords=coords) + custom_med = CustomMedium(permittivity=eps_data_array) + + return Structure( + geometry=geometry, + medium=custom_med, + **kwargs, + ) + class MeshOverrideStructure(AbstractStructure): """Defines an object that is only used in the process of generating the mesh. diff --git a/tidy3d/components/validators.py b/tidy3d/components/validators.py index b4c566498..c4c912e45 100644 --- a/tidy3d/components/validators.py +++ b/tidy3d/components/validators.py @@ -8,6 +8,7 @@ from .base import DATA_ARRAY_MAP, skip_if_fields_missing from .data.dataset import Dataset, FieldDataset from .geometry.base import Box +from .mode import ModeSpec from .types import Tuple """ Explanation of pydantic validators: @@ -136,8 +137,7 @@ def check_symmetry(cls, val, values): and geometric_object.center[dim] != sim_center[dim] ): raise SetupError( - f"Mode object '{geometric_object}' " - f"(at 'simulation.{field_name}[{position_index}]') " + f"{obj_type} at 'simulation.{field_name}[{position_index}]' " "in presence of symmetries must be in the main quadrant, " "or centered on the symmetry axis." ) @@ -178,17 +178,18 @@ def objects_in_sim_bounds(cls, val, values): # Do a strict check, unless simulation is 0D along a dimension strict_ineq = [size != 0 and strict_inequality for size in sim_size] - for position_index, geometric_object in enumerate(val): - if not sim_box.intersects(geometric_object.geometry, strict_inequality=strict_ineq): - message = ( - f"'simulation.{field_name}[{position_index}]' " - "is outside of the simulation domain." - ) - custom_loc = [field_name, position_index] + with log as consolidated_logger: + for position_index, geometric_object in enumerate(val): + if not sim_box.intersects(geometric_object.geometry, strict_inequality=strict_ineq): + message = ( + f"'simulation.{field_name}[{position_index}]' " + "is outside of the simulation domain." + ) + custom_loc = [field_name, position_index] - if error: - raise SetupError(message) - log.warning(message, custom_loc=custom_loc) + if error: + raise SetupError(message) + consolidated_logger.warning(message, custom_loc=custom_loc) return val @@ -396,3 +397,21 @@ def freqs_not_empty(cls, val): return val return freqs_not_empty + + +def validate_mode_plane_radius(mode_spec: ModeSpec, plane: Box, msg_prefix: str = ""): + """Validate that the radius of a mode spec with a bend is not smaller than half the size of + the plane along the radial direction.""" + + if not mode_spec.bend_radius: + return + + # radial axis is the plane axis that is not the bend axis + _, plane_axs = plane.pop_axis([0, 1, 2], plane.size.index(0.0)) + radial_ax = plane_axs[(mode_spec.bend_axis + 1) % 2] + + if np.abs(mode_spec.bend_radius) < plane.size[radial_ax] / 2: + raise ValueError( + f"{msg_prefix} bend radius is smaller than half the mode plane size " + "along the radial axis, which can produce wrong results." + ) diff --git a/tidy3d/packaging.py b/tidy3d/packaging.py index fd4d17b6e..43aeb8d4d 100644 --- a/tidy3d/packaging.py +++ b/tidy3d/packaging.py @@ -150,3 +150,26 @@ def _fn(*args, **kwargs): return fn(*args, **kwargs) return _fn + + +def get_numpy_major_version(module=np): + """ + Extracts the major version of the installed numpy accordingly. + + Parameters + ---------- + module : module + The module to extract the version from. Default is numpy. + + Returns + ------- + int + The major version of the module. + """ + # Get the version of the module + module_version = module.__version__ + + # Extract the major version number + major_version = int(module_version.split(".")[0]) + + return major_version diff --git a/tidy3d/plugins/adjoint/components/geometry.py b/tidy3d/plugins/adjoint/components/geometry.py index d4aa0351f..578d8f53a 100644 --- a/tidy3d/plugins/adjoint/components/geometry.py +++ b/tidy3d/plugins/adjoint/components/geometry.py @@ -304,6 +304,14 @@ class JaxPolySlab(JaxGeometry, PolySlab, JaxObject): stores_jax_for="sidewall_angle", ) + dilation_jax: JaxFloat = pd.Field( + default=0.0, + title="Dilation (Jax)", + description="Jax-traced float defining the dilation.", + units=MICROMETER, + stores_jax_for="dilation", + ) + @pd.validator("sidewall_angle", always=True) def no_sidewall(cls, val): """Warn if sidewall angle present.""" @@ -320,13 +328,6 @@ def no_sidewall(cls, val): ) return val - @pd.validator("dilation", always=True) - def no_dilation(cls, val): - """Don't allow dilation.""" - if not np.isclose(val, 0.0): - raise AdjointError("'JaxPolySlab' does not support dilation.") - return val - def _validate_web_adjoint(self) -> None: """Run validators for this component, only if using ``tda.web.run()``.""" self._limit_number_of_vertices() @@ -526,6 +527,19 @@ def _proper_vertices(vertices: ArrayFloat2D) -> jnp.ndarray: vertices_np = JaxPolySlab.vertices_to_array(vertices) return JaxPolySlab._orient(JaxPolySlab._remove_duplicate_vertices(vertices_np)) + @staticmethod + def _heal_polygon(vertices: jnp.ndarray) -> jnp.ndarray: + """heal a self-intersecting polygon.""" + shapely_poly = PolySlab.make_shapely_polygon(jax.lax.stop_gradient(vertices)) + if shapely_poly.is_valid: + return vertices + + raise NotImplementedError( + "The dilation caused damage to the polygon. Automatically healing this is " + "currently not supported for 'JaxPolySlab' objects. Try increasing the spacing " + "between vertices or reduce the amount of dilation." + ) + @staticmethod def vertices_to_array(vertices_tuple: ArrayFloat2D) -> jnp.ndarray: """Converts a list of tuples (vertices) to a jax array.""" @@ -543,7 +557,8 @@ def reference_polygon(self) -> jnp.ndarray: vertices = JaxPolySlab._proper_vertices(self.vertices_jax) if jnp.isclose(self.dilation, 0): return vertices - raise NotImplementedError("JaxPolySlab does not support dilation!") + offset_vertices = self._shift_vertices(vertices, self.dilation)[0] + return self._heal_polygon(offset_vertices) def edge_contrib( self, diff --git a/tidy3d/plugins/adjoint/web.py b/tidy3d/plugins/adjoint/web.py index b3e1da163..f1031ed21 100644 --- a/tidy3d/plugins/adjoint/web.py +++ b/tidy3d/plugins/adjoint/web.py @@ -1,5 +1,6 @@ """Adjoint-specific webapi.""" +import os import tempfile from functools import partial from typing import Dict, List, Tuple @@ -8,6 +9,7 @@ from jax import custom_vjp from jax.tree_util import register_pytree_node_class +import tidy3d as td from tidy3d.web.api.asynchronous import run_async as web_run_async from tidy3d.web.api.webapi import run as web_run from tidy3d.web.api.webapi import wait_for_connection @@ -219,26 +221,36 @@ def run_bwd( @wait_for_connection def upload_jax_info(jax_info: JaxInfo, task_id: str, verbose: bool) -> None: """Upload jax_info for a task with a given task_id.""" - - data_file = tempfile.NamedTemporaryFile(suffix=".json") - data_file.close() - jax_info.to_file(data_file.name) - upload_file( - task_id, - data_file.name, - JAX_INFO_FILE, - verbose=verbose, - ) + handle, fname = tempfile.mkstemp(suffix=".json") + os.close(handle) + try: + jax_info.to_file(fname) + upload_file( + task_id, + fname, + JAX_INFO_FILE, + verbose=verbose, + ) + except Exception as e: + td.log.error(f"Error occurred while uploading 'jax_info': {e}") + raise e + finally: + os.unlink(fname) @wait_for_connection def download_sim_vjp(task_id: str, verbose: bool) -> JaxSimulation: """Download the vjp loaded simulation from the server to return to jax.""" - - data_file = tempfile.NamedTemporaryFile(suffix=".hdf5") - data_file.close() - download_file(task_id, SIM_VJP_FILE, to_file=data_file.name, verbose=verbose) - return JaxSimulation.from_file(data_file.name) + handle, fname = tempfile.mkstemp(suffix=".hdf5") + os.close(handle) + try: + download_file(task_id, SIM_VJP_FILE, to_file=fname, verbose=verbose) + return JaxSimulation.from_file(fname) + except Exception as e: + td.log.error(f"Error occurred while downloading 'sim_vjp': {e}") + raise e + finally: + os.unlink(fname) AdjointSimulationType = Literal["tidy3d", "adjoint_fwd", "adjoint_bwd"] @@ -248,7 +260,7 @@ class AdjointJob(Job): """Job that uploads a jax_info object and also includes new fields for adjoint tasks.""" simulation_type: AdjointSimulationType = pd.Field( - None, + "tidy3d", title="Simulation Type", description="Type of simulation, used internally only.", ) @@ -331,7 +343,7 @@ def webapi_run_adjoint_fwd( jax_info=jax_info, ) - sim_data = job.run() + sim_data = job.run(path=path) return sim_data, job.task_id diff --git a/tidy3d/plugins/autograd/README.md b/tidy3d/plugins/autograd/README.md index 557c813f4..eac0364ce 100644 --- a/tidy3d/plugins/autograd/README.md +++ b/tidy3d/plugins/autograd/README.md @@ -1,48 +1,56 @@ # Automatic Differentiation in Tidy3D -### Context - -As of version 2.7.0, `tidy3d` supports the ability to differentiate functions involving a `web.run` of a `tidy3d` simulation. This allows users to optimize objective functions involving `tidy3d` simulations using gradient-descent. This gradient calculation is done under the hood using the adjoint method, which requires just 1 additional simulation, no matter how many design parameters are involved. +As of version 2.7.0, `tidy3d` supports the ability to differentiate functions involving a `web.run` of a `tidy3d` simulation. +This allows users to optimize objective functions involving `tidy3d` simulations using gradient descent. +This gradient calculation is done under the hood using the adjoint method, which requires just one additional simulation, no matter how many design parameters are involved. This functionality was previously available using the `adjoint` plugin, which used `jax`. There were a few issues with this approach: -1. `jax` was often quite difficult to install on many systems and often conflicted with other packages. +1. `jax` can be quite difficult to install on many systems and often conflicted with other packages. 2. Because we wanted `jax` to be an optional dependency, the `adjoint` plugin was separated from the regular `tidy3d` components, requiring a new set of `Jax_` classes. -3. Because we inherited these classes from their `tidy3d` components, for technical reasons, we needed to separate the `jax`-traced fields from the regular fields. For example, `JaxSimulation.input_structures` and `.output_monitors` were needed. +3. Because we inherited these classes from their `tidy3d` components, for technical reasons, we needed to separate the `jax`-traced fields from the regular fields. + For example, `JaxSimulation.input_structures` and `.output_monitors` were needed. -All of these limitations (among others) motivated us to come up with a new approach to automatic differentiation, which will be introduced as an experimental feature in `2.7`. The `adjoint` plugin will still be supported in the indefinite future, but will not be developed with new features. We also believe the new approach offers a far better user experience, so we encourage users to switch whenever is convenient. This guide will give some instructions on how to do so. +All of these limitations (among others) motivated us to come up with a new approach to automatic differentiation, which was introduced as an experimental feature in `2.7`. +The `adjoint` plugin will continue to be supported indefinitely, but no new features will be developed for it. +We also believe the new approach offers a far better user experience, so we encourage users to switch whenever is convenient. +This guide will give some instructions on how to do so. ## New implementation using `autograd` -Automatic differentiation in 2.7 is built directly into `tidy3d`. One can perform objective function differentiation similarly to what was possible in the `adjoint` plugin. However, this can be done using regular `td.` components, such as `td.Simulation`, `td.Structure`, and `td.Medium`. Also, the regular `web.run()` function is now differentiable so there is no need to import a wrapper. In short, users can take existing functional code and differentiate it without changing much: +Automatic differentiation in `2.7` is built directly into `tidy3d`. +One can perform objective function differentiation similarly to what was possible in the `adjoint` plugin. +However, this can be done using regular `td.` components, such as `td.Simulation`, `td.Structure`, and `td.Medium`. +Also, the regular `web.run()` function is now differentiable, so there is no need to import a wrapper. +In short, users can take existing functional code and differentiate it without changing much: ```py def objective(eps: float) -> float: + structure = td.Structure( + medium=td.Medium(permittivity=eps), + geometry=td.Box(...), + ) - structure = td.Structure( - medium=td.Medium(permittivity=eps), - geometry=td.Box(...), - ) - - sim = td.Simulation( - structures=[structure], - ... - ) + sim = td.Simulation( + structures=[structure], + ... + ) - data = td.web.run(sim) + data = td.web.run(sim) - return np.sum(abs(data['mode'].amps.sel(mode_index=0).values)) + return np.sum(np.abs(data["mode"].amps.sel(mode_index=0))).item() # compute derivative of objective(1.0) with respect to input autograd.grad(objective)(1.0) - ``` -Instead of using `jax`, we now use the [autograd](https://github.com/HIPS/autograd) package for our "core" automatic differentiation. Many `tidy3d` components now accept and are now compatible with `autograd` arrays. Because `autograd` is far lighter and has very few requirements, it was made a core dependency of `tidy3d`. +Instead of using `jax`, we now use the [autograd](https://github.com/HIPS/autograd) package for our "core" automatic differentiation. +Many `tidy3d` components now accept and are compatible with `autograd` arrays. +Due to its lightweight nature and minimal dependencies, `autograd` has been made a core dependency of `tidy3d`. -While we use `autograd` internally, in the future, we will include wrappers so you can use automatic differentiation frameworks of your choice (e.g. `jax`, `pytorch`) without much of a change to the syntax. +Although `autograd` is used internally, we provide wrappers for other automatic differentiation frameworks, allowing you to use your preferred AD framework (e.g., `jax`, `pytorch`) with minimal syntax changes. For instance, you can refer to our PyTorch wrapper [here](../pytorch/). -The usability of `autograd` is extremely similar to `jax` but with a couple modifications, which we'll outline below. +The usability of `autograd` is extremely similar to `jax` but with a couple of modifications, which we'll outline below. ### Migrating from jax to autograd @@ -65,17 +73,19 @@ There is also a `numpy` wrapper that can be similarly imported from `autograd.nu ```py import jax.numpy as jnp jnp.sum(...) - ``` becomes + ```py import autograd.numpy as anp anp.sum(...) - ``` -`Autograd` supports fewer features than `jax`. So, for example, the `has_aux` option is not supported in the `jax.grad()`, but one can write their own utilities to implement these features, as we show in the notebook examples. +`Autograd` supports fewer features than `jax`. +For example, the `has_aux` option is not supported in the default `autograd.grad()` function, but one can write their own utilities to implement these features, as we show in the notebook examples. +We also have a `value_and_grad` function in `tidy3d.plugins.autograd.differential_operators` that is similar to `jax.value_and_grad` and supports `has_aux`. +Additionally, `autograd` has a `grad_with_aux` function that can be used to compute gradients while returning auxiliary values, similar to `jax.grad` with `has_aux`. Otherwise, `jax` and `autograd` are very similar to each other in practice. @@ -97,92 +107,146 @@ import tidy3d as td td.Structure(...) ``` -These `td.` classes can be used directly in the differentiable objective functions. Like before, only some fields are traceable for differentiation, and we outline the full list of supported fields in the feature roadmap below. +These `td.` classes can be used directly in the differentiable objective functions. +Like before, only some fields are traceable for differentiation, and we outline the full list of supported fields in the feature roadmap below. -Furthermore, there is no need for separated fields in the `JaxSimulation`, so one can eliminate `output_monitors` and `input_structures` and put everything in `monitors` and `structures`, respectively. `tidy3d` will automatically determine which structure and monitor is traced for differentiation. +Furthermore, there is no need for separated fields in the `JaxSimulation`, so one can eliminate `output_monitors` and `input_structures` and put everything in `monitors` and `structures`, respectively. +`tidy3d` will automatically determine which structure and monitor is traced for differentiation. -Finally, the regular `web.run()` and `web.run_async()` functions have their derivatives registered with `autograd`, so there is no need to use special web API functions. If there are no tracers found in `web.run()` or `web.run_async()` simulations, the original (non-`autograd`) code will be called. +Finally, the regular `web.run()` and `web.run_async()` functions have their derivatives registered with `autograd`, so there is no need to use special web API functions. +If there are no tracers found in `web.run()` or `web.run_async()` simulations, the original (non-`autograd`) code will be called. -## Feature Roadmap +## Common Gotchas -### Currently Supported +Autograd has some limitations and quirks. +A good starting point to get familiar with them is the [autograd tutorial](https://github.com/HIPS/autograd/blob/master/docs/tutorial.md). -The following components are traceable as inputs to the `td.Simulation` +Some of the most important autograd "Don'ts" are: -rectangular prisms -- `Box.center` -- `Box.size` +- Do not use in-place assignment on numpy arrays, e.g., `x[i] = something`. + Often, you can formulate the assignment in terms of `np.where()`. +- Similarly, do not use in-place operators such as `+=`, `*=`, etc. +- Prefer numpy functions over array methods, e.g., use `np.sum(x)` over `x.sum()`. -polyslab (including those with dilation or slanted sidewalls) -- `PolySlab.vertices` +It is important to note that any function you use with autograd differential operators like `grad`, `value_and_grad`, `elementwise_grad`, etc., must return real values in the form of a float, a tuple of floats, or a numpy array. +Specifically, for `grad` and `value_and_grad`, the output must be either a scalar or a one-element array. -regular mediums -- `Medium.permittivity` -- `Medium.conductivity` +When extracting values from `SimulationData`, ensure that any output value is converted to a float or numpy array before returning. +This is because numpy operations on `DataArray` objects will yield other `DataArray` objects, which are not compatible with autograd's automatic differentiation when returned from the function. -spatially varying mediums (for topology optimization mainly) -- `CustomMedium.permittivity` -- `CustomMedium.eps_dataset` +For example: -groups of geometries with the same medium (for faster processing) -- `GeometryGroup.geometries` +```py +def objective(params: np.ndarray) -> float: + sim = make_simulation(params) + sim_data = td.web.run(sim) -complex and self-intersecting polyslabs -- `ComplexPolySlab.vertices` + amps = sim_data["mode_monitor"].amps + mode_power = np.abs(amps)**2 # mode_power is still a DataArray! -dispersive materials -- `PoleResidue.eps_inf` -- `PoleResidue.poles` + # either select out a specific value + objective_value = mode_power.sel(mode_index=0, f=freq0) + # or, for example, sum over all frequencies + objective_value = mode_power.sel(mode_index=0).sum() -spatially dependent dispersive materials: -- `CustomPoleResidue.eps_inf` -- `CustomPoleResidue.poles` + # just make sure that whatever you return is scalar and a numeric type by extracting the scalar value with item() + return objective_value.item() # alternatively, for single-element arrays: flux.data or flux.values (deprecated) +``` -The following data types are traceable as outputs of the `td.SimulationData` +For more complex objective functions, it is advisable to extract the `.data` attribute from the `DataArray` _before_ performing any numpy operations. +Although most autograd numpy functions are compatible with `DataArray` objects, there can be instances of unexpected behavior. +Therefore, working directly with the underlying data of the `DataArray` is generally a more robust approach. + +For example: + +```py +def objective(params: np.ndarray) -> float: + sim = make_simulation(params) + sim_data = td.web.run(sim) -- `ModeData.amps` -- `DiffractionData.amps` -- `FieldData.field_components` -- `FieldData` operations: - - `FieldData.flux` - - `SimulationData.get_intensity(fld_monitor_name)` - - `SimulationData.get_poynting(fld_monitor_name)` + fields = sim_data["field_monitor"] -We also support the following high level features + # extract the data from the DataArray + Ex = fields.Ex.data + Ey = fields.Ey.data + Ez = fields.Ez.data -- gradients for objective functions depending on multi-frequency data with single broadband adjoint source. -- server-side gradient processing by default, which saves transfer. This can be turned off by passing `local_gradient=True` to the `web.run()` and related functions. + # we can now use these just like regular numpy arrays + intensity = anp.abs(Ex) ** 2 + anp.abs(Ey) ** 2 + anp.abs(Ez) ** 2 # sim_data.get_intensity("field_monitor") would also work of course + norm_intensity = anp.linalg.norm(intensity) + + return norm_intensity # no .item() needed +``` + +## Feature Roadmap + +Please check out our [Adjoint Master Plan](https://github.com/flexcompute/tidy3d/issues/1548) on GitHub if you want to stay updated on the progress of planned features and contribute to the discussion. + +### Currently Supported + +The following components are traceable as inputs to the `td.Simulation` + +| Component Type | Traceable Attributes | +| ----------------------------------------------------------------- | ------------------------------------------------------- | +| rectangular prisms | `Box.center`, `Box.size` | +| polyslab (including those with dilation or slanted sidewalls) | `PolySlab.vertices` | +| regular mediums | `Medium.permittivity`, `Medium.conductivity` | +| spatially varying mediums (for topology optimization mainly) | `CustomMedium.permittivity`, `CustomMedium.eps_dataset` | +| groups of geometries with the same medium (for faster processing) | `GeometryGroup.geometries` | +| complex and self-intersecting polyslabs | `ComplexPolySlab.vertices` | +| dispersive materials | `PoleResidue.eps_inf`, `PoleResidue.poles` | +| spatially dependent dispersive materials | `CustomPoleResidue.eps_inf`, `CustomPoleResidue.poles` | +| cylinders | `Cylinder.radius`, `Cylinder.center` | + +The following components are traceable as outputs of the `td.SimulationData` + +| Data Type | Traceable Attributes & Methods | +| ----------------- | ------------------------------------------------------------- | +| `ModeData` | `amps` | +| `DiffractionData` | `amps` | +| `FieldData` | `field_components`, `flux` | +| `SimulationData` | `get_intensity(field_monitor)`, `get_poynting(field_monitor)` | + +We also support the following high-level features: + +- To manually set the background permittivity of a structure for purposes of shape optimization, one can set `Structure.background_medium`. + This is useful when there is a substrate or multiple overlapping structures as some geometries, such as `PolySlab`, do not automatically detect background permittivity and instead use the `Simulation.medium` by default. +- Compute gradients for objective functions that rely on multi-frequency data using a single broadband adjoint source. +- Enable server-side gradient processing by setting `local_gradient=False` in the web functions. + This can significantly reduce data storage time. + However, exercise caution when using this feature with multi-frequency monitors and large design regions, as it may result in substantial data storage on our servers. We currently have the following restrictions: -- Only 500 max structures containing tracers can be added to the `Simulation` to cut down on processing time. If you hit this, try using `GeometryGroup` to group any structures with the same `.medium` to relax the requirement. +- Only 500 max structures containing tracers can be added to the `Simulation` to cut down on processing time. + To bypass this restriction, use `GeometryGroup` to group structures with the same medium. +- `web.run_async` for simulations with tracers does not return a `BatchData` but rather a `dict` mapping task name to `SimulationData`. + There may be high memory usage with many simulations or a lot of data for each. +- Tidy3D can handle objective functions over a single simulation under any of the following conditions for the monitors that the objective function output depends on: + - Several monitors, all with the same frequency. + - One monitor with many frequencies where the data is being extracted out of a single coordinate (e.g., single mode_index or direction will work, multiple will not). + If your optimization does not fall into one of these categories, you must split it into separate simulations and run them with `web.run_async`. In all cases, the adjoint simulation bandwidth will be the same as the forward simulation. These limitations allow us to avoid the need to use methods that combine all adjoint sources into one simulation, which have the potential to degrade accuracy and increase the run time and cost significantly. That being said, we plan to offer support for more flexible and general broadband adjoint in the future. ### To be supported soon Next on our roadmap (targeting 2.8 and 2.9, fall 2024) is to support: -- Field projection support -- Automatic handling of all broadband objective functions using a fitting approach. - -Later this year (2024), we plan to support: - - `TriangleMesh`. - `GUI` integration of invdes plugin. ### Finally -If you have feature requests or questions, please feel free to file an issue or discussion on the `tidy3d` front end repository. +If you have feature requests or questions, please feel free to file an issue or discussion on this `tidy3d` front-end repository. Happy autogradding! ## Developer Notes To convert existing tidy3d front end code to be autograd compatible, will need to be aware of + - `numpy` -> `autograd.numpy` -- `x.real` -> `np.real(x)`` -- `float()` is not supported as far as I can tell. +- Casting to `float()` is not supported for autograd `ArrayBox` objects. - `isclose()` -> `np.isclose()` - `array[i] = something` needs a different approach (happens in mesher a lot) -- be careful with `+=` in autograd, it can fail silently.. -- whenever we pass things to other modules, like `shapely` especially, we need to be careful that they are untraced. -- I just made structures static (`obj = obj.to_static()`) before any meshing, as a cutoff point. So if we add a new `make_grid()` call somewhere, eg in a validator, just need to be aware. +- Whenever we pass things to other modules, like `shapely` especially, we need to be careful that they are untraced. +- I just made structures static before any meshing, as a cutoff point. So if we add a new `make_grid()` call somewhere, e.g. in a validator, just need to be aware. diff --git a/tidy3d/plugins/autograd/__init__.py b/tidy3d/plugins/autograd/__init__.py index dcbf05eb0..2e98df574 100644 --- a/tidy3d/plugins/autograd/__init__.py +++ b/tidy3d/plugins/autograd/__init__.py @@ -1,5 +1,77 @@ -from .differential_operators import value_and_grad +from .differential_operators import grad, value_and_grad +from .functions import ( + add_at, + convolve, + grey_closing, + grey_dilation, + grey_erosion, + grey_opening, + interpn, + least_squares, + morphological_gradient, + morphological_gradient_external, + morphological_gradient_internal, + pad, + rescale, + smooth_max, + smooth_min, + threshold, + trapz, +) +from .invdes import ( + CircularFilter, + ConicFilter, + ErosionDilationPenalty, + FilterAndProject, + grey_indicator, + make_circular_filter, + make_conic_filter, + make_curvature_penalty, + make_erosion_dilation_penalty, + make_filter, + make_filter_and_project, + ramp_projection, + tanh_projection, +) +from .primitives import gaussian_filter +from .utilities import chain, get_kernel_size_px, make_kernel, scalar_objective __all__ = [ + "CircularFilter", + "ConicFilter", + "ErosionDilationPenalty", + "FilterAndProject", + "make_filter", + "make_conic_filter", + "make_circular_filter", + "grey_indicator", + "convolve", + "pad", + "ramp_projection", + "tanh_projection", + "make_erosion_dilation_penalty", + "make_curvature_penalty", + "make_filter_and_project", + "gaussian_filter", + "make_kernel", + "get_kernel_size_px", + "chain", + "grey_closing", + "grey_dilation", + "grey_erosion", + "grey_opening", + "morphological_gradient", + "morphological_gradient_external", + "morphological_gradient_internal", + "rescale", + "threshold", "value_and_grad", + "smooth_min", + "smooth_max", + "add_at", + "interpn", + "least_squares", + "grad", + "scalar_objective", + "trapz", ] diff --git a/tidy3d/plugins/autograd/constants.py b/tidy3d/plugins/autograd/constants.py new file mode 100644 index 000000000..062be8e5c --- /dev/null +++ b/tidy3d/plugins/autograd/constants.py @@ -0,0 +1,2 @@ +BETA_DEFAULT = 1.0 +ETA_DEFAULT = 0.5 diff --git a/tidy3d/plugins/autograd/differential_operators.py b/tidy3d/plugins/autograd/differential_operators.py index e83c439e1..f61ea5894 100644 --- a/tidy3d/plugins/autograd/differential_operators.py +++ b/tidy3d/plugins/autograd/differential_operators.py @@ -1,65 +1,72 @@ -from typing import Any, Callable +from typing import Callable -from autograd import value_and_grad as value_and_grad_ag from autograd.builtins import tuple as atuple from autograd.core import make_vjp from autograd.extend import vspace from autograd.wrap_util import unary_to_nary from numpy.typing import ArrayLike +from .utilities import scalar_objective + __all__ = [ "value_and_grad", + "grad", ] @unary_to_nary -def value_and_grad( - fun: Callable, x: ArrayLike, *, has_aux: bool = False -) -> tuple[tuple[float, ArrayLike], Any]: - """Returns a function that returns both value and gradient. - - This function wraps and extends autograd's 'value_and_grad' function by adding - support for auxiliary data. +def grad(fun: Callable, x: ArrayLike, *, has_aux: bool = False) -> Callable: + """Returns a function that computes the gradient of `fun` with respect to `x`. Parameters ---------- fun : Callable - The function to differentiate. Should take a single argument and return - a scalar value, or a tuple where the first element is a scalar value if has_aux is True. + The function to differentiate. Should return a scalar value, or a tuple of + (scalar_value, auxiliary_data) if `has_aux` is True. x : ArrayLike - The point at which to evaluate the function and its gradient. + The point at which to evaluate the gradient. has_aux : bool = False - If True, the function returns auxiliary data as the second element of a tuple. + If True, `fun` returns auxiliary data as the second element of a tuple. Returns ------- - tuple[tuple[float, ArrayLike], Any] - A tuple containing: - - A tuple with the function value (float) and its gradient (ArrayLike) - - The auxiliary data returned by the function (if has_aux is True) + Callable + A function that takes the same arguments as `fun` and returns its gradient at `x`. + """ + wrapped_fun = scalar_objective(fun, has_aux=has_aux) + vjp, result = make_vjp(lambda x: atuple(wrapped_fun(x)) if has_aux else wrapped_fun(x), x) - Raises - ------ - TypeError - If the function does not return a scalar value. + if has_aux: + ans, aux = result + return vjp((vspace(ans).ones(), None)), aux + ans = result + return vjp(vspace(ans).ones()) - Notes - ----- - This function uses autograd for automatic differentiation. If the function - does not return auxiliary data (has_aux is False), it delegates to autograd's - value_and_grad function. The main extension is the support for auxiliary data - when has_aux is True. - """ - if not has_aux: - return value_and_grad_ag(fun)(x) - vjp, (ans, aux) = make_vjp(lambda x: atuple(fun(x)), x) +@unary_to_nary +def value_and_grad(fun: Callable, x: ArrayLike, *, has_aux: bool = False) -> Callable: + """Returns a function that computes both the value and gradient of `fun` with respect to `x`. + + Parameters + ---------- + fun : Callable + The function to differentiate. Should return a scalar value, or a tuple of + (scalar_value, auxiliary_data) if `has_aux` is True. + x : ArrayLike + The point at which to evaluate the function and its gradient. + has_aux : bool = False + If True, `fun` returns auxiliary data as the second element of a tuple. - if not vspace(ans).size == 1: - raise TypeError( - "value_and_grad only applies to real scalar-output " - "functions. Try jacobian, elementwise_grad or " - "holomorphic_grad." - ) + Returns + ------- + Callable + A function that takes the same arguments as `fun` and returns its value and gradient at `x`. + """ + wrapped_fun = scalar_objective(fun, has_aux=has_aux) + vjp, result = make_vjp(lambda x: atuple(wrapped_fun(x)) if has_aux else wrapped_fun(x), x) - return (ans, vjp((vspace(ans).ones(), None))), aux + if has_aux: + ans, aux = result + return (ans, vjp((vspace(ans).ones(), None))), aux + ans = result + return ans, vjp(vspace(ans).ones()) diff --git a/tidy3d/plugins/autograd/functions.py b/tidy3d/plugins/autograd/functions.py index 479d53160..0cdb3c952 100644 --- a/tidy3d/plugins/autograd/functions.py +++ b/tidy3d/plugins/autograd/functions.py @@ -1,15 +1,19 @@ -from typing import Iterable, List, Literal, Tuple, Union +from typing import Callable, Iterable, List, Literal, Tuple, Union import autograd.numpy as np +from autograd import jacobian from autograd.scipy.signal import convolve as convolve_ag +from autograd.scipy.special import logsumexp from numpy.typing import NDArray -from tidy3d.components.autograd.functions import interpn +from tidy3d.components.autograd.functions import add_at, interpn, trapz from .types import PaddingType __all__ = [ "interpn", + "trapz", + "add_at", "pad", "convolve", "grey_dilation", @@ -21,159 +25,94 @@ "morphological_gradient_external", "rescale", "threshold", + "smooth_min", + "smooth_max", ] -def _make_slices(rule: Union[int, slice], ndim: int, axis: int) -> Tuple[slice, ...]: - """Create a tuple of slices for indexing an array. +def _pad_indices(n: int, pad_width: Tuple[int, int], *, mode: PaddingType) -> NDArray: + """Compute the indices to pad an array along a single axis based on the padding mode. Parameters ---------- - rule : Union[int, slice] - The rule to apply on the specified axis. - ndim : int - The number of dimensions of the array. - axis : int - The axis to which the rule should be applied. - - Returns - ------- - Tuple[slice, ...] - A tuple of slices for indexing. - """ - return tuple(slice(None) if d != axis else rule for d in range(ndim)) - - -def _constant_pad( - array: NDArray, pad_width: Tuple[int, int], axis: int, *, constant_value: float = 0.0 -) -> NDArray: - """Pad an array with a constant value along a specified axis. - - Parameters - ---------- - array : NDArray - The input array to pad. + n : int + The size of the axis to pad. pad_width : Tuple[int, int] - The number of values padded to the edges of each axis. - axis : int - The axis along which to pad. - constant_value : float, optional - The constant value to pad with. Default is 0.0. + The number of values padded to the edges of the axis. + mode : PaddingType + The padding mode to use. Returns ------- - NDArray - The padded array. + np.ndarray + The indices for padding along the axis. """ - p = [(pad_width[0], pad_width[1]) if ax == axis else (0, 0) for ax in range(array.ndim)] - return np.pad(array, p, mode="constant", constant_values=constant_value) + total_pad = sum(pad_width) + if n == 0: + return np.zeros(total_pad, dtype=int) + idx = np.arange(-pad_width[0], n + pad_width[1]) -def _edge_pad(array: NDArray, pad_width: Tuple[int, int], axis: int) -> NDArray: - """Pad an array using the `edge` mode along a specified axis. + # Handle each padding mode + if mode == "constant": + return idx - Parameters - ---------- - array : NDArray - The input array to pad. - pad_width : Tuple[int, int] - The number of values padded to the edges of each axis. - axis : int - The axis along which to pad. + if mode == "edge": + return np.clip(idx, 0, n - 1) - Returns - ------- - NDArray - The padded array. - """ - left, right = (_make_slices(rule, array.ndim, axis) for rule in (0, -1)) + if mode == "reflect": + period = 2 * n - 2 if n > 1 else 1 + idx = np.mod(idx, period) + return np.where(idx >= n, period - idx, idx) - cat_arys = [] - if pad_width[0] > 0: - cat_arys.append(np.stack(pad_width[0] * [array[left]], axis=axis)) - cat_arys.append(array) - if pad_width[1] > 0: - cat_arys.append(np.stack(pad_width[1] * [array[right]], axis=axis)) + if mode == "symmetric": + period = 2 * n if n > 1 else 1 + idx = np.mod(idx, period) + return np.where(idx >= n, period - idx - 1, idx) - return np.concatenate(cat_arys, axis=axis) + if mode == "wrap": + return np.mod(idx, n) + raise ValueError(f"Unsupported padding mode: {mode}") -def _reflect_pad(array: NDArray, pad_width: Tuple[int, int], axis: int) -> NDArray: - """Pad an array using the `reflect` mode along a specified axis. - Parameters - ---------- - array : NDArray - The input array to pad. - pad_width : Tuple[int, int] - The number of values padded to the edges of each axis. - axis : int - The axis along which to pad. - - Returns - ------- - NDArray - The padded array. - """ - left, right = ( - _make_slices(rule, array.ndim, axis) - for rule in (slice(pad_width[0], 0, -1), slice(-2, -pad_width[1] - 2, -1)) - ) - return np.concatenate([array[left], array, array[right]], axis=axis) - - -def _symmetric_pad(array: NDArray, pad_width: Tuple[int, int], axis: int) -> NDArray: - """Pad an array using the `symmetric` mode along a specified axis. +def _pad_axis( + array: NDArray, + pad_width: Tuple[int, int], + axis: int, + *, + mode: PaddingType = "constant", + constant_value: float = 0.0, +) -> NDArray: + """Pad an array along a specified axis. Parameters ---------- - array : NDArray + array : np.ndarray The input array to pad. pad_width : Tuple[int, int] - The number of values padded to the edges of each axis. + The number of values padded to the edges of the axis. axis : int The axis along which to pad. + mode : PaddingType = "constant" + The padding mode to use. + constant_value : float = 0.0 + The constant value to pad with when mode is 'constant'. Returns ------- - NDArray + np.ndarray The padded array. """ - left, right = ( - _make_slices(rule, array.ndim, axis) - for rule in ( - slice(pad_width[0] - 1, None, -1) if pad_width[0] > 0 else slice(0, 0), - slice(-1, -pad_width[1] - 1, -1), - ) - ) - return np.concatenate([array[left], array, array[right]], axis=axis) - - -def _wrap_pad(array: NDArray, pad_width: Tuple[int, int], axis: int) -> NDArray: - """Pad an array using the `wrap` mode along a specified axis. - - Parameters - ---------- - array : NDArray - The input array to pad. - pad_width : Tuple[int, int] - The number of values padded to the edges of each axis. - axis : int - The axis along which to pad. + if mode == "constant": + padding = [(0, 0)] * array.ndim + padding[axis] = pad_width + return np.pad(array, padding, mode="constant", constant_values=constant_value) - Returns - ------- - NDArray - The padded array. - """ - left, right = ( - _make_slices(rule, array.ndim, axis) - for rule in ( - slice(-pad_width[0], None) if pad_width[0] > 0 else slice(0, 0), - slice(0, pad_width[1]) if pad_width[1] > 0 else slice(0, 0), - ) - ) - return np.concatenate([array[left], array, array[right]], axis=axis) + idx = _pad_indices(array.shape[axis], pad_width, mode=mode) + indexer = [slice(None)] * array.ndim + indexer[axis] = idx + return array[tuple(indexer)] def pad( @@ -184,94 +123,64 @@ def pad( axis: Union[int, Iterable[int], None] = None, constant_value: float = 0.0, ) -> NDArray: - """Pad an array along a specified axis with a given mode and padding width. + """Pad an array along specified axes with a given mode and padding width. Parameters ---------- - array : NDArray + array : np.ndarray The input array to pad. pad_width : Union[int, Tuple[int, int]] The number of values padded to the edges of each axis. If an integer is provided, it is used for both the left and right sides. If a tuple is provided, it specifies the padding for the left and right sides respectively. - mode : _pad_modes, optional - The padding mode to use. Default is "constant". - axis : Union[int, Iterable[int], None], optional - The axis or axes along which to pad. If None, padding is applied to all axes. Default is None. - constant_value : float, optional - The value to set the padded values for "constant" mode. Default is 0.0. + mode : PaddingType = "constant" + The padding mode to use. + axis : Union[int, Iterable[int], None] = None + The axis or axes along which to pad. If None, padding is applied to all axes. + constant_value : float = 0.0 + The value to set the padded values for "constant" mode. Returns ------- - NDArray + np.ndarray The padded array. Raises ------ ValueError If the padding width has more than two elements or if padding is negative. - NotImplementedError - If padding larger than the input size is requested. - KeyError - If an unsupported padding mode is specified. IndexError If an axis is out of range for the array dimensions. - - Examples - -------- - >>> import numpy as np - >>> from tidy3d.plugins.autograd.functions import pad - >>> array = np.array([[1, 2], [3, 4]]) - >>> pad(array, (1, 1), mode="constant", constant_value=0) - array([[0, 0, 0, 0], - [0, 1, 2, 0], - [0, 3, 4, 0], - [0, 0, 0, 0]]) - >>> pad(array, 1, mode="reflect") - array([[4, 3, 4, 3], - [2, 1, 2, 1], - [4, 3, 4, 3], - [2, 1, 2, 1]]) """ + # Normalize pad_width to a tuple of two elements pad_width = np.atleast_1d(pad_width) - - if pad_width.size == 1: - pad_width = np.array([pad_width[0], pad_width[0]]) - elif pad_width.size != 2: - raise ValueError("Padding width must have one or two elements, got {pad_width.size}.") - - if any(any(p >= s for p in pad_width) for s in array.shape): - raise NotImplementedError("Padding larger than the input size is not supported.") - if any(p < 0 for p in pad_width): - raise ValueError("Padding must be positive.") - if all(p == 0 for p in pad_width): + if pad_width.size > 2: + raise ValueError(f"Padding width must have one or two elements, got {pad_width.size}.") + pad_tuple = (pad_width[0], pad_width[0]) if pad_width.size == 1 else tuple(pad_width) + + # Validate padding values + if any(p < 0 for p in pad_tuple): + raise ValueError("Padding must be non-negative.") + if all(p == 0 for p in pad_tuple): return array - if axis is None: - axis = range(array.ndim) - - _mode_map = { - "constant": _constant_pad, - "edge": _edge_pad, - "reflect": _reflect_pad, - "symmetric": _symmetric_pad, - "wrap": _wrap_pad, - } - - try: - pad_fun = _mode_map[mode] - except KeyError as e: - raise KeyError(f"Unsupported padding mode: {mode}.") from e - - for ax in np.atleast_1d(axis): - if ax < 0: - ax += array.ndim - if ax < 0 or ax >= array.ndim: - raise IndexError(f"Axis out of range for array with {array.ndim} dimensions.") - - array = pad_fun(array, pad_width, axis=ax) - - return array + # Normalize and validate axes + axes = range(array.ndim) if axis is None else [axis] if isinstance(axis, int) else axis + axes = [ax + array.ndim if ax < 0 else ax for ax in axes] + if any(ax < 0 or ax >= array.ndim for ax in axes): + raise IndexError(f"Axis out of range for array with {array.ndim} dimensions.") + + # Apply padding to each axis + result = array + for ax in axes: + result = _pad_axis( + result, + pad_tuple, + axis=ax, + mode=mode, + constant_value=constant_value, + ) + return result def convolve( @@ -286,20 +195,20 @@ def convolve( Parameters ---------- - array : NDArray + array : np.ndarray The input array to be convolved. - kernel : NDArray + kernel : np.ndarray The kernel to convolve with the input array. All dimensions of the kernel must be odd. - padding : _pad_modes, optional - The padding mode to use. Default is "constant". - axes : Union[Tuple[List[int], List[int]], None], optional - The axes along which to perform the convolution. Default is None (all axes). - mode : Literal["full", "valid", "same"], optional - The convolution mode. Default is "same". + padding : PaddingType = "constant" + The padding mode to use. + axes : Union[Tuple[List[int], List[int]], None] = None + The axes along which to perform the convolution. + mode : Literal["full", "valid", "same"] = "same" + The convolution mode. Returns ------- - NDArray + np.ndarray The result of the convolution. Raises @@ -338,21 +247,20 @@ def grey_dilation( Parameters ---------- - array : NDArray + array : np.ndarray The input array to perform grey dilation on. - size : Union[Union[int, Tuple[int, int]], None], optional + size : Union[Union[int, Tuple[int, int]], None] = None The size of the structuring element. If None, `structure` must be provided. - Default is None. - structure : Union[NDArray, None], optional - The structuring element. If None, `size` must be provided. Default is None. - mode : _pad_modes, optional - The padding mode to use. Default is "reflect". - maxval : float, optional - Value to assume for infinite elements in the kernel. Default is 1e4. + structure : Union[np.ndarray, None] = None + The structuring element. If None, `size` must be provided. + mode : PaddingType = "reflect" + The padding mode to use. + maxval : float = 1e4 + Value to assume for infinite elements in the kernel. Returns ------- - NDArray + np.ndarray The result of the grey dilation operation. Raises @@ -393,21 +301,20 @@ def grey_erosion( Parameters ---------- - array : NDArray + array : np.ndarray The input array to perform grey dilation on. - size : Union[Union[int, Tuple[int, int]], None], optional + size : Union[Union[int, Tuple[int, int]], None] = None The size of the structuring element. If None, `structure` must be provided. - Default is None. - structure : Union[NDArray, None], optional - The structuring element. If None, `size` must be provided. Default is None. - mode : _pad_modes, optional - The padding mode to use. Default is "reflect". - maxval : float, optional - Value to assume for infinite elements in the kernel. Default is 1e4. + structure : Union[np.ndarray, None] = None + The structuring element. If None, `size` must be provided. + mode : PaddingType = "reflect" + The padding mode to use. + maxval : float = 1e4 + Value to assume for infinite elements in the kernel. Returns ------- - NDArray + np.ndarray The result of the grey dilation operation. Raises @@ -448,21 +355,20 @@ def grey_opening( Parameters ---------- - array : NDArray + array : np.ndarray The input array to perform grey opening on. - size : Union[Union[int, Tuple[int, int]], None], optional + size : Union[Union[int, Tuple[int, int]], None] = None The size of the structuring element. If None, `structure` must be provided. - Default is None. - structure : Union[NDArray, None], optional - The structuring element. If None, `size` must be provided. Default is None. - mode : _pad_modes, optional - The padding mode to use. Default is "reflect". - maxval : float, optional - Value to assume for infinite elements in the kernel. Default is 1e4. + structure : Union[np.ndarray, None] = None + The structuring element. If None, `size` must be provided. + mode : PaddingType = "reflect" + The padding mode to use. + maxval : float = 1e4 + Value to assume for infinite elements in the kernel. Returns ------- - NDArray + np.ndarray The result of the grey opening operation. """ array = grey_erosion(array, size, structure, mode=mode, maxval=maxval) @@ -482,21 +388,20 @@ def grey_closing( Parameters ---------- - array : NDArray + array : np.ndarray The input array to perform grey closing on. - size : Union[Union[int, Tuple[int, int]], None], optional + size : Union[Union[int, Tuple[int, int]], None] = None The size of the structuring element. If None, `structure` must be provided. - Default is None. - structure : Union[NDArray, None], optional - The structuring element. If None, `size` must be provided. Default is None. - mode : _pad_modes, optional - The padding mode to use. Default is "reflect". - maxval : float, optional - Value to assume for infinite elements in the kernel. Default is 1e4. + structure : Union[np.ndarray, None] = None + The structuring element. If None, `size` must be provided. + mode : PaddingType = "reflect" + The padding mode to use. + maxval : float = 1e4 + Value to assume for infinite elements in the kernel. Returns ------- - NDArray + np.ndarray The result of the grey closing operation. """ array = grey_dilation(array, size, structure, mode=mode, maxval=maxval) @@ -516,21 +421,20 @@ def morphological_gradient( Parameters ---------- - array : NDArray + array : np.ndarray The input array to compute the morphological gradient of. - size : Union[Union[int, Tuple[int, int]], None], optional + size : Union[Union[int, Tuple[int, int]], None] = None The size of the structuring element. If None, `structure` must be provided. - Default is None. - structure : Union[NDArray, None], optional - The structuring element. If None, `size` must be provided. Default is None. - mode : _pad_modes, optional - The padding mode to use. Default is "reflect". - maxval : float, optional - Value to assume for infinite elements in the kernel. Default is 1e4. + structure : Union[np.ndarray, None] = None + The structuring element. If None, `size` must be provided. + mode : PaddingType = "reflect" + The padding mode to use. + maxval : float = 1e4 + Value to assume for infinite elements in the kernel. Returns ------- - NDArray + np.ndarray The morphological gradient of the input array. """ return grey_dilation(array, size, structure, mode=mode, maxval=maxval) - grey_erosion( @@ -550,21 +454,20 @@ def morphological_gradient_internal( Parameters ---------- - array : NDArray + array : np.ndarray The input array to compute the internal morphological gradient of. - size : Union[Union[int, Tuple[int, int]], None], optional + size : Union[Union[int, Tuple[int, int]], None] = None The size of the structuring element. If None, `structure` must be provided. - Default is None. - structure : Union[NDArray, None], optional - The structuring element. If None, `size` must be provided. Default is None. - mode : _pad_modes, optional - The padding mode to use. Default is "reflect". - maxval : float, optional - Value to assume for infinite elements in the kernel. Default is 1e4. + structure : Union[np.ndarray, None] = None + The structuring element. If None, `size` must be provided. + mode : PaddingType = "reflect" + The padding mode to use. + maxval : float = 1e4 + Value to assume for infinite elements in the kernel. Returns ------- - NDArray + np.ndarray The internal morphological gradient of the input array. """ return array - grey_erosion(array, size, structure, mode=mode, maxval=maxval) @@ -582,21 +485,20 @@ def morphological_gradient_external( Parameters ---------- - array : NDArray + array : np.ndarray The input array to compute the external morphological gradient of. - size : Union[Union[int, Tuple[int, int]], None], optional + size : Union[Union[int, Tuple[int, int]], None] = None The size of the structuring element. If None, `structure` must be provided. - Default is None. - structure : Union[NDArray, None], optional - The structuring element. If None, `size` must be provided. Default is None. - mode : _pad_modes, optional - The padding mode to use. Default is "reflect". - maxval : float, optional - Value to assume for infinite elements in the kernel. Default is 1e4. + structure : Union[np.ndarray, None] = None + The structuring element. If None, `size` must be provided. + mode : PaddingType = "reflect" + The padding mode to use. + maxval : float = 1e4 + Value to assume for infinite elements in the kernel. Returns ------- - NDArray + np.ndarray The external morphological gradient of the input array. """ return grey_dilation(array, size, structure, mode=mode, maxval=maxval) - array @@ -610,20 +512,20 @@ def rescale( Parameters ---------- - array : NDArray + array : np.ndarray The input array to be rescaled. out_min : float The minimum value of the output range. out_max : float The maximum value of the output range. - in_min : float, optional - The minimum value of the input range. Default is 0.0. - in_max : float, optional - The maximum value of the input range. Default is 1.0. + in_min : float = 0.0 + The minimum value of the input range. + in_max : float = 1.0 + The maximum value of the input range. Returns ------- - NDArray + np.ndarray The rescaled array. """ @@ -648,18 +550,18 @@ def threshold( Parameters ---------- - array : NDArray + array : np.ndarray The input array to be thresholded. - vmin : float, optional - The value to assign to elements below the threshold. Default is 0.0. - vmax : float, optional - The value to assign to elements above the threshold. Default is 1.0. - level : Union[float, None], optional - The threshold level. If None, the threshold is set to the midpoint between `vmin` and `vmax`. Default is None. + vmin : float = 0.0 + The value to assign to elements below the threshold. + vmax : float = 1.0 + The value to assign to elements above the threshold. + level : Union[float, None] = None + The threshold level. If None, the threshold is set to the midpoint between `vmin` and `vmax`. Returns ------- - NDArray + np.ndarray The thresholded array. """ if vmin >= vmax: @@ -676,3 +578,110 @@ def threshold( ) return np.where(array < level, vmin, vmax) + + +def smooth_max( + x: NDArray, tau: float = 1.0, axis: Union[int, Tuple[int, ...], None] = None +) -> float: + """Compute the smooth maximum of an array using temperature parameter tau. + + Parameters + ---------- + x : np.ndarray + Input array. + tau : float = 1.0 + Temperature parameter controlling smoothness. Larger values make the maximum smoother. + axis : Union[int, Tuple[int, ...], None] = None + Axis or axes over which the smooth maximum is computed. By default, the smooth maximum is computed over the entire array. + + Returns + ------- + np.ndarray + The smooth maximum of the input array. + """ + return tau * logsumexp(x / tau, axis=axis) + + +def smooth_min( + x: NDArray, tau: float = 1.0, axis: Union[int, Tuple[int, ...], None] = None +) -> float: + """Compute the smooth minimum of an array using temperature parameter tau. + + Parameters + ---------- + x : np.ndarray + Input array. + tau : float = 1.0 + Temperature parameter controlling smoothness. Larger values make the minimum smoother. + axis : Union[int, Tuple[int, ...], None] = None + Axis or axes over which the smooth minimum is computed. By default, the smooth minimum is computed over the entire array. + + Returns + ------- + np.ndarray + The smooth minimum of the input array. + """ + return -smooth_max(-x, tau, axis=axis) + + +def least_squares( + func: Callable[[NDArray, float], NDArray], + x: NDArray, + y: NDArray, + initial_guess: Tuple[float, ...], + max_iterations: int = 100, + tol: float = 1e-6, +) -> NDArray: + """Perform least squares fitting to find the best-fit parameters for a model function. + + Parameters + ---------- + func : Callable[[np.ndarray, float], np.ndarray] + The model function to fit. It should accept the independent variable `x` and a tuple of parameters, + and return the predicted dependent variable values. + x : np.ndarray + Independent variable data. + y : np.ndarray + Dependent variable data. + initial_guess : Tuple[float, ...] + Initial guess for the parameters to be optimized. + max_iterations : int = 100 + Maximum number of iterations for the optimization process. + tol : float = 1e-6 + Tolerance for convergence. The optimization stops when the change in parameters is below this threshold. + + Returns + ------- + np.ndarray + The optimized parameters that best fit the model to the data. + + Raises + ------ + np.linalg.LinAlgError + If the optimization does not converge within the specified number of iterations. + + Example + ------- + >>> import numpy as np + >>> def linear_model(x, a, b): + ... return a * x + b + >>> x_data = np.linspace(0, 10, 50) + >>> y_data = 2.0 * x_data - 3.0 + >>> initial_guess = (0.0, 0.0) + >>> params = least_squares(linear_model, x_data, y_data, initial_guess) + >>> print(params) + [ 2. -3.] + """ + params = np.array(initial_guess, dtype="f8") + jac = jacobian(lambda params: func(x, *params)) + + for _ in range(max_iterations): + residuals = y - func(x, *params) + jacvec = jac(params) + pseudo_inv = np.linalg.pinv(jacvec) + delta = np.dot(pseudo_inv, residuals) + params = params + delta + if np.linalg.norm(delta) < tol: + break + + return params diff --git a/tidy3d/plugins/autograd/invdes/__init__.py b/tidy3d/plugins/autograd/invdes/__init__.py index 865e3f0a7..080f7dba8 100644 --- a/tidy3d/plugins/autograd/invdes/__init__.py +++ b/tidy3d/plugins/autograd/invdes/__init__.py @@ -1,18 +1,27 @@ -from .filters import make_circular_filter, make_conic_filter, make_filter -from .misc import get_kernel_size_px, grey_indicator -from .parametrizations import make_filter_and_project -from .penalties import make_curvature_penalty, make_erosion_dilation_penalty +from .filters import ( + CircularFilter, + ConicFilter, + make_circular_filter, + make_conic_filter, + make_filter, +) +from .misc import grey_indicator +from .parametrizations import FilterAndProject, make_filter_and_project +from .penalties import ErosionDilationPenalty, make_curvature_penalty, make_erosion_dilation_penalty from .projections import ramp_projection, tanh_projection __all__ = [ - "get_kernel_size_px", "grey_indicator", + "CircularFilter", + "ConicFilter", "make_circular_filter", "make_conic_filter", "make_curvature_penalty", "make_erosion_dilation_penalty", + "ErosionDilationPenalty", "make_filter", "make_filter_and_project", + "FilterAndProject", "ramp_projection", "tanh_projection", ] diff --git a/tidy3d/plugins/autograd/invdes/filters.py b/tidy3d/plugins/autograd/invdes/filters.py index c07fca06f..2bdf45e2b 100644 --- a/tidy3d/plugins/autograd/invdes/filters.py +++ b/tidy3d/plugins/autograd/invdes/filters.py @@ -1,66 +1,221 @@ -from functools import partial -from typing import Callable, Tuple, Union +from __future__ import annotations + +import abc +from functools import lru_cache, partial +from typing import Annotated, Callable, Iterable, Tuple, Union import numpy as np +import pydantic.v1 as pd from numpy.typing import NDArray +import tidy3d as td +from tidy3d.components.base import Tidy3dBaseModel +from tidy3d.components.types import TYPE_TAG_STR + from ..functions import convolve from ..types import KernelType, PaddingType -from ..utilities import make_kernel +from ..utilities import get_kernel_size_px, make_kernel + + +class AbstractFilter(Tidy3dBaseModel, abc.ABC): + """An abstract class for creating and applying convolution filters.""" + + kernel_size: Union[pd.PositiveInt, Tuple[pd.PositiveInt, ...]] = pd.Field( + ..., title="Kernel Size", description="Size of the kernel in pixels for each dimension." + ) + normalize: bool = pd.Field( + True, title="Normalize", description="Whether to normalize the kernel so that it sums to 1." + ) + padding: PaddingType = pd.Field( + "reflect", title="Padding", description="The padding mode to use." + ) + + @classmethod + def from_radius_dl( + cls, radius: Union[float, Tuple[float, ...]], dl: Union[float, Tuple[float, ...]], **kwargs + ) -> AbstractFilter: + """Create a filter from radius and grid spacing. + + Parameters + ---------- + radius : Union[float, Tuple[float, ...]] + The radius of the kernel. Can be a scalar or a tuple. + dl : Union[float, Tuple[float, ...]] + The grid spacing. Can be a scalar or a tuple. + **kwargs + Additional keyword arguments to pass to the filter constructor. + + Returns + ------- + AbstractFilter + An instance of the filter. + """ + kernel_size = get_kernel_size_px(radius=radius, dl=dl) + return cls(kernel_size=kernel_size, **kwargs) + + @staticmethod + @abc.abstractmethod + def get_kernel(size_px: Iterable[int], normalize: bool) -> NDArray: + """Get the kernel for the filter. + + Parameters + ---------- + size_px : Iterable[int] + Size of the kernel in pixels for each dimension. + normalize : bool + Whether to normalize the kernel so that it sums to 1. + + Returns + ------- + np.ndarray + The kernel. + """ + + def __call__(self, array: NDArray) -> NDArray: + """Apply the filter to an input array. + + Parameters + ---------- + array : np.ndarray + The input array to filter. + + Returns + ------- + np.ndarray + The filtered array. + """ + original_shape = array.shape + squeezed_array = np.squeeze(array) + size_px = tuple(np.atleast_1d(self.kernel_size)) + if len(size_px) != squeezed_array.ndim: + size_px *= squeezed_array.ndim + kernel = self.get_kernel(size_px, self.normalize) + convolved_array = convolve(squeezed_array, kernel, padding=self.padding) + return np.reshape(convolved_array, original_shape) + + +class ConicFilter(AbstractFilter): + """A conic filter for creating and applying convolution filters.""" + + @staticmethod + @lru_cache(maxsize=1) + def get_kernel(size_px: Iterable[int], normalize: bool) -> NDArray: + """Get the conic kernel. + + See Also + -------- + :func:`~filters.AbstractFilter.get_kernel` for full method documentation. + """ + return make_kernel(kernel_type="conic", size=size_px, normalize=normalize) + + +class CircularFilter(AbstractFilter): + """A circular filter for creating and applying convolution filters.""" + + @staticmethod + @lru_cache(maxsize=1) + def get_kernel(size_px: Iterable[int], normalize: bool) -> NDArray: + """Get the circular kernel. + + See Also + -------- + :func:`~filters.AbstractFilter.get_kernel` for full method documentation. + """ + return make_kernel(kernel_type="circular", size=size_px, normalize=normalize) + + +def _get_kernel_size( + radius: Union[float, Tuple[float, ...]], + dl: Union[float, Tuple[float, ...]], + size_px: Union[int, Tuple[int, ...]], +) -> Tuple[int, ...]: + """Determine the kernel size based on the provided radius, grid spacing, or size in pixels. + + Parameters + ---------- + radius : Union[float, Tuple[float, ...]] + The radius of the kernel. Can be a scalar or a tuple. + dl : Union[float, Tuple[float, ...]] + The grid spacing. Can be a scalar or a tuple. + size_px : Union[int, Tuple[int, ...]] + The size of the kernel in pixels for each dimension. Can be a scalar or a tuple. + + Returns + ------- + Tuple[int, ...] + The size of the kernel in pixels for each dimension. + + Raises + ------ + ValueError + If neither ``size_px`` nor both ``radius`` and ``dl`` are provided. + """ + if size_px is not None: + if radius is not None and dl is not None: + td.log.warning( + "Both 'size_px' and 'radius' and 'dl' are provided. 'size_px' will take precedence." + ) + return (size_px,) if np.isscalar(size_px) else tuple(size_px) + elif radius is not None and dl is not None: + kernel_size = get_kernel_size_px(radius=radius, dl=dl) + return (kernel_size,) if np.isscalar(kernel_size) else tuple(kernel_size) + else: + raise ValueError("Either 'size_px' or both 'radius' and 'dl' must be provided.") def make_filter( - filter_type: KernelType, - size: Union[int, Tuple[int, ...]], + radius: Union[float, Tuple[float, ...]] = None, + dl: Union[float, Tuple[float, ...]] = None, *, + size_px: Union[int, Tuple[int, ...]] = None, normalize: bool = True, padding: PaddingType = "reflect", -) -> Callable: + filter_type: KernelType, +) -> Callable[[NDArray], NDArray]: """Create a filter function based on the specified kernel type and size. Parameters ---------- - filter_type : KernelType - The type of kernel to create (`circular` or `conic`). - size : Union[int, Tuple[int, ...]] + radius : Union[float, Tuple[float, ...]] = None + The radius of the kernel. Can be a scalar or a tuple. + dl : Union[float, Tuple[float, ...]] = None + The grid spacing. Can be a scalar or a tuple. + size_px : Union[int, Tuple[int, ...]] = None The size of the kernel in pixels for each dimension. Can be a scalar or a tuple. - normalize : bool, optional - Whether to normalize the kernel so that it sums to 1. Default is True. - padding : PadMode, optional - The padding mode to use. Default is "reflect". + normalize : bool = True + Whether to normalize the kernel so that it sums to 1. + padding : PaddingType = "reflect" + The padding mode to use. + filter_type : KernelType + The type of kernel to create (``circular`` or ``conic``). Returns ------- - function + Callable[[np.ndarray], np.ndarray] A function that applies the created filter to an input array. """ - _kernel = {} - - def _filter(array: NDArray) -> NDArray: - original_shape = array.shape - squeezed_array = np.squeeze(array) - - if squeezed_array.ndim not in _kernel: - if np.isscalar(size): - kernel_size = (size,) * squeezed_array.ndim - else: - kernel_size = size - _kernel[squeezed_array.ndim] = make_kernel( - kernel_type=filter_type, size=kernel_size, normalize=normalize - ) + kernel_size = _get_kernel_size(radius, dl, size_px) - convolved_array = convolve(squeezed_array, _kernel[squeezed_array.ndim], padding=padding) - return np.reshape(convolved_array, original_shape) + if filter_type == "conic": + filter_class = ConicFilter + elif filter_type == "circular": + filter_class = CircularFilter + else: + raise ValueError( + f"Unsupported filter_type: {filter_type}. " + "Must be one of `CircularFilter` or `ConicFilter`." + ) - return _filter + filter_instance = filter_class(kernel_size=kernel_size, normalize=normalize, padding=padding) + return filter_instance make_conic_filter = partial(make_filter, filter_type="conic") -make_conic_filter.__doc__ = """make_filter() with a default filter_type value of `conic`. +make_conic_filter.__doc__ = """make_filter() with a default filter_type value of ``conic``. See Also -------- -make_filter : Function to create a filter based on the specified kernel type and size. +:func:`~filters.make_filter` : Function to create a filter based on the specified kernel type and size. """ make_circular_filter = partial(make_filter, filter_type="circular") @@ -68,5 +223,7 @@ def _filter(array: NDArray) -> NDArray: See Also -------- -make_filter : Function to create a filter based on the specified kernel type and size. +:func:`~filters.make_filter` : Function to create a filter based on the specified kernel type and size. """ + +FilterType = Annotated[Union[ConicFilter, CircularFilter], pd.Field(discriminator=TYPE_TAG_STR)] diff --git a/tidy3d/plugins/autograd/invdes/misc.py b/tidy3d/plugins/autograd/invdes/misc.py index d61561601..ea92c7828 100644 --- a/tidy3d/plugins/autograd/invdes/misc.py +++ b/tidy3d/plugins/autograd/invdes/misc.py @@ -12,7 +12,7 @@ def grey_indicator(array: NDArray) -> float: Parameters ---------- - array : NDArray + array : np.ndarray The input array for which the grey indicator is to be calculated. Returns @@ -21,22 +21,3 @@ def grey_indicator(array: NDArray) -> float: The calculated grey indicator. """ return np.mean(4 * array * (1 - array)) - - -def get_kernel_size_px(radius: float, dl: float) -> int: - """Calculate the size of the kernel in pixels based on the given radius and pixel size. - - Parameters - ---------- - radius : float - The radius of the kernel in micrometers. - dl : float - The size of each pixel in micrometers. - - Returns - ------- - float - The size of the kernel in pixels. - """ - radius_px = np.ceil(radius / dl) - return int(2 * radius_px + 1) diff --git a/tidy3d/plugins/autograd/invdes/parametrizations.py b/tidy3d/plugins/autograd/invdes/parametrizations.py index 23f3e7e16..22ef6d6f6 100644 --- a/tidy3d/plugins/autograd/invdes/parametrizations.py +++ b/tidy3d/plugins/autograd/invdes/parametrizations.py @@ -1,46 +1,96 @@ -from typing import Callable, Tuple +from __future__ import annotations +from typing import Callable, Tuple, Union + +import pydantic.v1 as pd from numpy.typing import NDArray +from tidy3d.components.base import Tidy3dBaseModel + +from ..constants import BETA_DEFAULT, ETA_DEFAULT from ..types import KernelType, PaddingType from .filters import make_filter from .projections import tanh_projection +class FilterAndProject(Tidy3dBaseModel): + """A class that combines filtering and projection operations.""" + + radius: Union[float, Tuple[float, ...]] = pd.Field( + ..., title="Radius", description="The radius of the kernel." + ) + dl: Union[float, Tuple[float, ...]] = pd.Field( + ..., title="Grid Spacing", description="The grid spacing." + ) + size_px: Union[int, Tuple[int, ...]] = pd.Field( + None, title="Size in Pixels", description="The size of the kernel in pixels." + ) + beta: pd.NonNegativeFloat = pd.Field( + BETA_DEFAULT, title="Beta", description="The beta parameter for the tanh projection." + ) + eta: pd.NonNegativeFloat = pd.Field( + ETA_DEFAULT, title="Eta", description="The eta parameter for the tanh projection." + ) + filter_type: KernelType = pd.Field( + "conic", title="Filter Type", description="The type of filter to create." + ) + padding: PaddingType = pd.Field( + "reflect", title="Padding", description="The padding mode to use." + ) + + def __call__(self, array: NDArray, beta: float = None, eta: float = None) -> NDArray: + """Apply the filter and projection to an input array. + + Parameters + ---------- + array : np.ndarray + The input array to filter and project. + beta : float = None + The beta parameter for the tanh projection. If None, uses the instance's beta. + eta : float = None + The eta parameter for the tanh projection. If None, uses the instance's eta. + + Returns + ------- + np.ndarray + The filtered and projected array. + """ + filter_instance = make_filter( + radius=self.radius, + dl=self.dl, + size_px=self.size_px, + filter_type=self.filter_type, + padding=self.padding, + ) + filtered = filter_instance(array) + beta = beta if beta is not None else self.beta + eta = eta if eta is not None else self.eta + projected = tanh_projection(filtered, beta, eta) + return projected + + def make_filter_and_project( - filter_size: Tuple[int, ...], - beta: float = 1.0, - eta: float = 0.5, + radius: Union[float, Tuple[float, ...]] = None, + dl: Union[float, Tuple[float, ...]] = None, + *, + size_px: Union[int, Tuple[int, ...]] = None, + beta: float = BETA_DEFAULT, + eta: float = ETA_DEFAULT, filter_type: KernelType = "conic", padding: PaddingType = "reflect", ) -> Callable: """Create a function that filters and projects an array. - This is the standard filter-and-project scheme used in topology optimization. - - Parameters - ---------- - filter_size : Tuple[int, ...] - The size of the filter kernel in pixels. - beta : float, optional - The beta parameter for the tanh projection, by default 1.0. - eta : float, optional - The eta parameter for the tanh projection, by default 0.5. - filter_type : KernelType, optional - The type of filter kernel to use, by default "conic". - padding : PaddingType, optional - The padding type to use for the filter, by default "reflect". - - Returns - ------- - function - A function that takes an array and applies the filter and projection. + See Also + -------- + :func:`~parametrizations.FilterAndProject`. """ - _filter = make_filter(filter_type, filter_size, padding=padding) - - def _filter_and_project(array: NDArray, beta: float = beta, eta: float = eta) -> NDArray: - array = _filter(array) - array = tanh_projection(array, beta, eta) - return array - - return _filter_and_project + return FilterAndProject( + radius=radius, + dl=dl, + size_px=size_px, + beta=beta, + eta=eta, + filter_type=filter_type, + padding=padding, + ) diff --git a/tidy3d/plugins/autograd/invdes/penalties.py b/tidy3d/plugins/autograd/invdes/penalties.py index cfa5e1039..9fc27d40a 100644 --- a/tidy3d/plugins/autograd/invdes/penalties.py +++ b/tidy3d/plugins/autograd/invdes/penalties.py @@ -1,72 +1,117 @@ from typing import Callable, Tuple, Union import autograd.numpy as np +import pydantic.v1 as pd from numpy.typing import NDArray +from tidy3d.components.base import Tidy3dBaseModel from tidy3d.components.types import ArrayFloat2D from ..types import PaddingType -from .parametrizations import make_filter_and_project +from .parametrizations import FilterAndProject + + +class ErosionDilationPenalty(Tidy3dBaseModel): + """A class that computes a penalty for erosion/dilation of a parameter map not being unity.""" + + radius: Union[float, Tuple[float, ...]] = pd.Field( + ..., title="Radius", description="The radius of the kernel." + ) + dl: Union[float, Tuple[float, ...]] = pd.Field( + ..., title="Grid Spacing", description="The grid spacing." + ) + size_px: Union[int, Tuple[int, ...]] = pd.Field( + None, title="Size in Pixels", description="The size of the kernel in pixels." + ) + beta: pd.NonNegativeFloat = pd.Field( + 20.0, title="Beta", description="The beta parameter for the tanh projection." + ) + eta: pd.NonNegativeFloat = pd.Field( + 0.5, title="Eta", description="The eta parameter for the tanh projection." + ) + filter_type: str = pd.Field( + "conic", title="Filter Type", description="The type of filter to create." + ) + padding: PaddingType = pd.Field( + "reflect", title="Padding", description="The padding mode to use." + ) + delta_eta: float = pd.Field( + 0.01, + title="Delta Eta", + description="The binarization threshold for erosion and dilation operations.", + ) + + def __call__(self, array: NDArray) -> float: + """Compute the erosion/dilation penalty for a given array. + + Parameters + ---------- + array : np.ndarray + The input array to compute the penalty for. + + Returns + ------- + float + The computed erosion/dilation penalty. + """ + filtproj = FilterAndProject( + radius=self.radius, + dl=self.dl, + size_px=self.size_px, + beta=self.beta, + eta=self.eta, + filter_type=self.filter_type, + padding=self.padding, + ) + + eta_dilate = 0.0 + self.delta_eta + eta_eroded = 1.0 - self.delta_eta + + def _dilate(arr: NDArray): + return filtproj(arr, eta=eta_dilate) + + def _erode(arr: NDArray): + return filtproj(arr, eta=eta_eroded) + + def _open(arr: NDArray): + return _dilate(_erode(arr)) + + def _close(arr: NDArray): + return _erode(_dilate(arr)) + + diff = _close(array) - _open(array) + + if not np.any(diff): + return 0.0 + + return np.linalg.norm(diff) / np.sqrt(diff.size) def make_erosion_dilation_penalty( - filter_size: Tuple[int, ...], - beta: float = 100.0, + radius: Union[float, Tuple[float, ...]], + dl: Union[float, Tuple[float, ...]], + *, + size_px: Union[int, Tuple[int, ...]] = None, + beta: float = 20.0, eta: float = 0.5, delta_eta: float = 0.01, padding: PaddingType = "reflect", ) -> Callable: """Computes a penalty for erosion/dilation of a parameter map not being unity. - Accepts a parameter array normalized between 0 and 1. Uses filtering and projection methods - to erode and dilate the features within this array. Measures the change in the array after - eroding and dilating (and also dilating and eroding). Returns a penalty proportional to the - magnitude of this change. The amount of change under dilation and erosion is minimized if - the structure has large feature sizes and large radius of curvature relative to the length scale. - - Parameters - ---------- - filter_size : Tuple[int, ...] - The size of the filter to be used for erosion and dilation. - beta : float, optional - Strength of the tanh projection. Default is 100.0. - eta : float, optional - Midpoint of the tanh projection. Default is 0.5. - delta_eta : float, optional - The binarization threshold for erosion and dilation operations. Default is 0.01. - padding : PaddingType, optional - The padding type to use for the filter. Default is "reflect". - - Returns - ------- - Callable - A function that computes the erosion/dilation penalty for a given array. + See Also + -------- + :func:`~penalties.ErosionDilationPenalty`. """ - filtproj = make_filter_and_project(filter_size, beta, eta, padding=padding) - eta_dilate = 0.0 + delta_eta - eta_eroded = 1.0 - delta_eta - - def _dilate(array: NDArray, beta: float) -> NDArray: - return filtproj(array, beta=beta, eta=eta_dilate) - - def _erode(array: NDArray, beta: float) -> NDArray: - return filtproj(array, beta=beta, eta=eta_eroded) - - def _open(array: NDArray, beta: float) -> NDArray: - return _dilate(_erode(array, beta=beta), beta=beta) - - def _close(array: NDArray, beta: float) -> NDArray: - return _erode(_dilate(array, beta=beta), beta=beta) - - def _erosion_dilation_penalty(array: NDArray, beta: float = beta) -> float: - diff = _close(array, beta) - _open(array, beta) - - if not np.any(diff): - return 0.0 - - return np.linalg.norm(diff) / np.sqrt(diff.size) - - return _erosion_dilation_penalty + return ErosionDilationPenalty( + radius=radius, + dl=dl, + size_px=size_px, + beta=beta, + eta=eta, + delta_eta=delta_eta, + padding=padding, + ) def curvature(dp: NDArray, ddp: NDArray) -> NDArray: @@ -74,14 +119,14 @@ def curvature(dp: NDArray, ddp: NDArray) -> NDArray: Parameters ---------- - dp : NDArray + dp : np.ndarray The first derivative at the point, with (x, y) entries in the first dimension. - ddp : NDArray + ddp : np.ndarray The second derivative at the point, with (x, y) entries in the first dimension. Returns ------- - NDArray + np.ndarray The curvature at the given point. Notes @@ -103,16 +148,16 @@ def bezier_with_grads( ---------- t : float The parameter at which to evaluate the Bezier curve. - p0 : NDArray + p0 : np.ndarray The first control point of the Bezier curve. - pc : NDArray + pc : np.ndarray The central control point of the Bezier curve. - p2 : NDArray + p2 : np.ndarray The last control point of the Bezier curve. Returns ------- - tuple[NDArray, NDArray, NDArray] + tuple[np.ndarray, np.ndarray, np.ndarray] A tuple containing the Bezier curve value, its first derivative, and its second derivative at the given point. """ p1 = 2 * pc - p0 / 2 - p2 / 2 @@ -128,16 +173,16 @@ def bezier_curvature(x: NDArray, y: NDArray, t: Union[NDArray, float] = 0.5) -> Parameters ---------- - x : NDArray + x : np.ndarray The x-coordinates of the control points. - y : NDArray + y : np.ndarray The y-coordinates of the control points. - t : Union[NDArray, float], optional - The parameter at which to evaluate the curvature, by default 0.5. + t : Union[np.ndarray, float] = 0.5 + The parameter at which to evaluate the curvature. Returns ------- - NDArray + np.ndarray The curvature of the Bezier curve at the given parameter t. """ p = np.stack((x, y), axis=1) @@ -154,12 +199,12 @@ def make_curvature_penalty( ---------- min_radius : float The minimum radius of curvature. - alpha : float, optional - Scaling factor for the penalty, by default 1.0. - kappa : float, optional - Exponential factor for the penalty, by default 10.0. - eps : float, optional - A small value to avoid division by zero, by default 1e-6. + alpha : float = 1.0 + Scaling factor for the penalty. + kappa : float = 10.0 + Exponential factor for the penalty. + eps : float = 1e-6 + A small value to avoid division by zero. Returns ------- diff --git a/tidy3d/plugins/autograd/invdes/projections.py b/tidy3d/plugins/autograd/invdes/projections.py index c95d8b0f7..51ba088db 100644 --- a/tidy3d/plugins/autograd/invdes/projections.py +++ b/tidy3d/plugins/autograd/invdes/projections.py @@ -1,6 +1,8 @@ import autograd.numpy as np from numpy.typing import NDArray +from ..constants import BETA_DEFAULT, ETA_DEFAULT + def ramp_projection(array: NDArray, width: float = 0.1, center: float = 0.5) -> NDArray: """Apply a piecewise linear ramp projection to an array. @@ -13,16 +15,16 @@ def ramp_projection(array: NDArray, width: float = 0.1, center: float = 0.5) -> Parameters ---------- - array : NDArray + array : np.ndarray The input array to be projected. - width : float, optional - The width of the ramp. Default is 0.1. - center : float, optional - The center of the ramp. Default is 0.5. + width : float = 0.1 + The width of the ramp. + center : float 0.5 + The center of the ramp. Returns ------- - NDArray + np.ndarray The array after applying the ramp projection. """ ll = array <= (center - width / 2) @@ -38,7 +40,9 @@ def ramp_projection(array: NDArray, width: float = 0.1, center: float = 0.5) -> ) -def tanh_projection(array: NDArray, beta: float = 1.0, eta: float = 0.5) -> NDArray: +def tanh_projection( + array: NDArray, beta: float = BETA_DEFAULT, eta: float = ETA_DEFAULT +) -> NDArray: """Apply a tanh-based soft-thresholding projection to an array. This function performs a tanh projection on the input array, which is a common @@ -47,17 +51,16 @@ def tanh_projection(array: NDArray, beta: float = 1.0, eta: float = 0.5) -> NDAr Parameters ---------- - array : NDArray + array : np.ndarray The input array to be projected. - beta : float, optional + beta : float = BETA_DEFAULT The steepness of the projection. Higher values result in a sharper transition. - Default is 1.0. - eta : float, optional - The midpoint of the projection. Default is 0.5. + eta : float = ETA_DEFAULT + The midpoint of the projection. Returns ------- - NDArray + np.ndarray The array after applying the tanh projection. """ if beta == 0: diff --git a/tidy3d/plugins/autograd/utilities.py b/tidy3d/plugins/autograd/utilities.py index aaef3759f..1303f9e1b 100644 --- a/tidy3d/plugins/autograd/utilities.py +++ b/tidy3d/plugins/autograd/utilities.py @@ -1,9 +1,13 @@ -from functools import reduce -from typing import Callable, Iterable, Union +from functools import reduce, wraps +from typing import Any, Callable, Iterable, List, Union +import autograd.numpy as anp import numpy as np +import xarray as xr from numpy.typing import NDArray +from tidy3d.exceptions import Tidy3dError + from .types import KernelType @@ -54,15 +58,15 @@ def make_kernel(kernel_type: KernelType, size: Iterable[int], normalize: bool = The type of kernel to create ('circular' or 'conic'). size : Iterable[int] The size of the kernel in pixels for each dimension. - normalize : bool, optional - Whether to normalize the kernel so that it sums to 1. Default is True. + normalize : bool = True + Whether to normalize the kernel so that it sums to 1. Returns ------- NDArray An n-dimensional array representing the specified type of kernel. """ - if not all(isinstance(dim, int) and dim > 0 for dim in size): + if not all(np.issubdtype(type(dim), int) and dim > 0 for dim in size): raise ValueError("'size' must be an iterable of positive integers.") if kernel_type == "circular": @@ -78,7 +82,45 @@ def make_kernel(kernel_type: KernelType, size: Iterable[int], normalize: bool = return kernel -def chain(*funcs: Union[Callable, Iterable[Callable]]) -> Callable: +def get_kernel_size_px( + radius: Union[float, Iterable[float]] = None, dl: Union[float, Iterable[float]] = None +) -> Union[int, List[int]]: + """Calculate the kernel size in pixels based on the provided radius and grid spacing. + + Parameters + ---------- + radius : Union[float, Iterable[float]] = None + The radius of the kernel. Can be a scalar or an iterable of floats. + dl : Union[float, Iterable[float]] = None + The grid spacing. Can be a scalar or an iterable of floats. + + Returns + ------- + Union[int, List[int]] + The size of the kernel in pixels for each dimension. Returns an integer if the radius is scalar, otherwise a list of integers. + + Raises + ------ + ValueError + If either 'radius' or 'dl' is not provided. + """ + if radius is None or dl is None: + raise ValueError("Either 'size_px' or both 'radius' and 'dl' must be provided.") + + if np.isscalar(radius): + radius = [radius] * len(dl) if isinstance(dl, Iterable) else [radius] + if np.isscalar(dl): + dl = [dl] * len(radius) + + radius_px = [np.ceil(r / g) for r, g in zip(radius, dl)] + return ( + [int(2 * r_px + 1) for r_px in radius_px] + if len(radius_px) > 1 + else int(2 * radius_px[0] + 1) + ) + + +def chain(*funcs: Union[Callable, Iterable[Callable]]): """Chain multiple functions together to apply them sequentially to an array. Parameters @@ -120,3 +162,80 @@ def chained(array: NDArray): return reduce(lambda x, y: y(x), funcs, array) return chained + + +def scalar_objective(func: Callable = None, *, has_aux: bool = False) -> Callable: + """Decorator to ensure the objective function returns a real scalar value. + + This decorator wraps an objective function to ensure that its return value is a real scalar. + If the function returns auxiliary data, it expects the return value to be a tuple of the form + (result, aux_data). + + Parameters + ---------- + func : Callable, optional + The objective function to be decorated. If not provided, the decorator should be used with + arguments. + has_aux : bool = False + If True, expects the function to return a tuple (result, aux_data). + + Returns + ------- + Callable + The wrapped function that ensures a real scalar return value. If `has_aux` is True, the + wrapped function returns a tuple (result, aux_data). + + Raises + ------ + Tidy3dError + If the return value is not a real scalar, or if `has_aux` is True and the function does not return a tuple of length 2. + """ + if func is None: + return lambda f: scalar_objective(f, has_aux=has_aux) + + @wraps(func) + def wrapper(*args: Any, **kwargs: Any) -> Any: + result = func(*args, **kwargs) + aux_data = None + + # Unpack auxiliary data if present + if has_aux: + if not isinstance(result, tuple) or len(result) != 2: + raise Tidy3dError( + "If 'has_aux' is True, the objective function must return " + "a tuple of length 2." + ) + result, aux_data = result + + # Extract data from xarray.DataArray + if isinstance(result, xr.DataArray): + result = result.data + + # Squeeze to remove singleton dimensions + result = anp.squeeze(result) + + # Attempt to extract scalar value + try: + result = result.item() + except AttributeError: + # If result is already a scalar, pass + if not isinstance(result, (float, int)): + raise Tidy3dError( + "An objective function's return value must be a scalar, " + "a Python float/int, or an array containing a single element." + ) + except ValueError as e: + # Result contains more than one element + raise Tidy3dError( + "An objective function's return value must be a scalar " + "but got an array with shape " + f"{getattr(result, 'shape', 'N/A')}." + ) from e + + # Ensure the result is real + if not anp.isreal(result): + raise Tidy3dError("An objective function's return value must be real.") + + return (result, aux_data) if aux_data is not None else result + + return wrapper diff --git a/tidy3d/plugins/expressions/README.md b/tidy3d/plugins/expressions/README.md new file mode 100644 index 000000000..83c2af9b9 --- /dev/null +++ b/tidy3d/plugins/expressions/README.md @@ -0,0 +1,303 @@ +# Expressions module + +The `expressions` module provides a way to construct and serialize complex mathematical expressions involving simulation metrics in Tidy3D. +It allows users to define objective functions for optimization tasks, for example in conjunction with the `invdes` plugin. +This module is essential for creating expressions that can be easily saved, loaded, and passed between different components of a simulation workflow. + +## Introduction + +The `expressions` module is designed to facilitate the construction of serializable mathematical expressions involving simulation metrics like mode power, mode coefficients, field intensity, and more. +These expressions can be combined using standard arithmetic operators. +This functionality is important when defining objective functions for optimization routines, especially in inverse design tasks where the objective functions may need to be transmitted or stored. + +### Creating metrics + +To start using the `expressions` module, you can create metrics using predefined metric classes such as `ModeAmp` and `ModePower`. + +Metrics are a special instance of a `Variable`, which itself is a subclass of `Expression`, that take a `SimulationData` as input and are intended to return a single scalar. + +```python +from tidy3d.plugins.expressions import ModeAmp, ModePower + +# Create a ModeAmp metric +mode_coeff = ModeAmp(monitor_name="monitor1", freqs=[1.0]) + +# Create a ModePower metric +mode_power = ModePower(monitor_name="monitor2", freqs=[1.0]) +``` + +### Combining metrics with operators + +Metrics can be combined using standard arithmetic operators like `+`, `-`, `*`, `/`, `**`, and functions like `abs()`. +Together, they will form an `Expression`. + +```python +# Define an objective function using metrics +f = abs(mode_coeff) - mode_power / 2 +``` + +### Functions + +The `expressions` module also provides a set of mathematical functions that can be used to create more complex expressions. + +```python +from tidy3d.plugins.expressions import Sin, Cos + +f = Sin(mode_coeff) + Cos(mode_power) +``` + +### Evaluating expressions + +Once you have a metric, you can evaluate it using simulation data. + +```python +# Assume "data" is a SimulationData object obtained from a simulation +result = f.evaluate(data) + +# ...or just +result = f(data) +``` + +### Serializing and deserializing expressions + +Expressions can be serialized to a file using the `to_file` method and deserialized using the `from_file` class method. + +```python +# Serialize the metric to a file +f.to_file("metric_expression.hdf5") + +# Deserialize the metric from a file +from tidy3d.plugins.expressions import Expression +loaded_expr = Expression.from_file("metric_expression.hdf5") +``` + +## Supported features + +### Operators and functions + +The `expressions` module supports various operators and functions to build complex expressions: + +| Type | Name | Description | Class Name | +|------------|------------|--------------------------------------------------|------------------| +| Operator | `+` | Addition | `Add` | +| Operator | `-` | Subtraction | `Subtract` | +| Operator | `*` | Multiplication | `Multiply` | +| Operator | `/` | Division | `Divide` | +| Operator | `**` | Power | `Power` | +| Operator | `//` | Floor Division | `FloorDivide` | +| Operator | `%` | Modulus | `Modulus` | +| Operator | `@` | Matrix Multiplication | `MatMul` | +| Operator | `abs` | Absolute Value | `Abs` | +| Function | `Sin` | Sine function | `Sin` | +| Function | `Cos` | Cosine function | `Cos` | +| Function | `Tan` | Tangent function | `Tan` | +| Function | `Exp` | Exponential function | `Exp` | +| Function | `Log` | Natural logarithm function | `Log` | +| Function | `Log10` | Base-10 logarithm function | `Log10` | +| Function | `Sqrt` | Square root function | `Sqrt` | + +### Variables and metrics + +The module provides predefined metrics and variables for constructing expressions. +Metrics represent specific simulation data, while variables act as placeholders. + +| Type | Name | Description | +|------------|------------|--------------------------------------------------| +| Metric | `ModeAmp` | Metric for calculating the mode coefficient | +| Metric | `ModePower`| Metric for calculating the mode power | +| Variable | `Constant` | Fixed value in expressions | +| Variable | `Variable` | Placeholder for values in expressions | + +## Examples + +### `ModeAmp` and `ModePower` + +In this example, we create metrics using only `ModeAmp` and `ModePower` and combine them to define an objective function. + +```python +from tidy3d.plugins.expressions import ModeAmp, ModePower + +# Create metrics +mode_coeff = ModeAmp(monitor_name="monitor1", freqs=[1.0]) +mode_power = ModePower(monitor_name="monitor2", freqs=[1.0]) + +# Define an objective function using metrics +f = abs(mode_coeff) - mode_power / 2 + +# Display the expression +print(f) +``` + +**Expected output:** + +```text +(abs(ModeAmp("monitor1")) - (ModePower("monitor2") / 2)) +``` + +**Evaluating the expression:** + +```python +# Assume "data" is a SimulationData object obtained from a simulation +result = f.evaluate(data) + +# Display the result +print(result) +``` + +### Using variables + +The `expressions` module provides `Variable` and `Constant` classes to represent placeholders and fixed values in your expressions. +Variables can be used to parameterize expressions, allowing you to supply values at evaluation time. +Note that a `Metric` such as `ModeAmp` is a subclass of `Variable`, so it can be used in the same way. + +#### Unnamed variables (positional arguments) + +If you create a `Variable` without a name, it expects its value to be provided as a single positional argument during evaluation. + +```python +from tidy3d.plugins.expressions import Variable + +# Create an unnamed variable +x = Variable() + +# Use the variable in an expression +expr = x**2 + 2 * x + 1 + +# Evaluate the expression with a positional argument +result = expr(3) + +print(result) # Outputs: 126 +``` + +#### Named variables (keyword arguments) + +If you create a `Variable` with a name, it expects its value to be provided as a keyword argument during evaluation. + +```python +from tidy3d.plugins.expressions import Variable + +# Create named variables +x = Variable(name="x") +y = Variable(name="y") +expr = x**2 + y**2 + +# Evaluate the expression with keyword arguments +result = expr(x=3, y=4) + +print(result) # Outputs: 25 +``` + +#### Mixing Positional and Keyword Arguments + +You can mix positional and keyword arguments when evaluating expressions that contain both unnamed and named variables. + +```python +# Create a mix of named and unnamed variables +x = Variable(name="x") +y = Variable() +expr = x**2 + y**2 + +result = expr(3, y=4) # x = 3, y = 4 + +print(result) # Outputs: 25 +``` + +#### A `Metric` is a `Variable` + +A `Metric` such as `ModeAmp` is a subclass of `Variable`, so all the rules for `Variable` apply to `Metric` as well. + +from tidy3d.plugins.expressions import ModeAmp, ModePower + +```python +# Define two named metrics +mode_coeff1 = ModeAmp(name="mode_coeff1", monitor_name="monitor1", freqs=[1.0]) +mode_power1 = ModePower(name="mode_power1", monitor_name="monitor2", freqs=[1.0]) + +# Create an expression using the metrics +expr = mode_coeff1 + mode_power1 + +# Assume "data1" and "data2" are SimulationData objects obtained from different simulations +result = expr(mode_coeff1=data1, mode_power1=data2) +``` + +#### Important notes on variable evaluation + +- **Unnamed Variables**: If you have an unnamed variable (`Variable()`), you must provide exactly one positional argument during evaluation. +Providing multiple positional arguments will raise a `ValueError`. + +- **Named Variables**: For each named variable, you must provide a corresponding keyword argument. +If a named variable is missing during evaluation, a `ValueError` will be raised. + +- **Multiple Unnamed Variables**: If your expression contains multiple unnamed variables, you cannot provide multiple positional arguments. +You should assign names to the variables and provide their values as keyword arguments instead. + +## Developer notes + +### Extending expressions + +To implement new metrics, follow these steps: + +1. **Subclass the `Metric` base class:** + + Create a new class that inherits from `Metric` and implement the required methods. + + ```python + from tidy3d.plugins.expressions.metrics import Metric + + class CustomMetric(Metric): + monitor_name: str + + def evaluate(self, data: SimulationData) -> NumberType: + # Implement custom evaluation logic + pass + ``` + +2. **Define attributes and methods:** + + - Add any necessary attributes with type annotations. + - Implement the `evaluate` method, which defines how the metric is calculated from simulation data. + +3. **Update forward references:** + + If your metrics reference other metrics or expressions, ensure you handle forward references appropriately. + +### Extending operators + +To extend the `expressions` module with additional operators: + +1. **Create a new operator class:** + + Subclass `UnaryOperator` or `BinaryOperator` depending on the operator's arity. + + ```python + from tidy3d.plugins.expressions.operators import BinaryOperator + + class CustomOperator(BinaryOperator): + _symbol = "??" # Replace with the operator symbol + _format = "({left} {symbol} {right})" + + def evaluate(self, x, y): + # Implement the operator logic + pass + ``` + +2. **Implement required methods:** + + - Define the `_symbol` class variable that represents the operator symbol. + - Implement the `evaluate` method with the operator's logic. + +3. **Update operator overloads:** + + If you want to enable the use of the operator with standard syntax (e.g., using `@` for matrix multiplication), override the corresponding magic methods in the `Expression` base class. + + ```python + class Expression: + # Existing methods... + + def __matmul__(self, other): + return MatMul(left=self, right=other) + ``` + +4. **Register the operator (if necessary):** + + Ensure the new operator is recognized by the serialization mechanism. If your operator introduces a new type, update any type maps or registries used in the `expressions` module. diff --git a/tidy3d/plugins/expressions/__init__.py b/tidy3d/plugins/expressions/__init__.py new file mode 100644 index 000000000..1f8f733f3 --- /dev/null +++ b/tidy3d/plugins/expressions/__init__.py @@ -0,0 +1,44 @@ +from .base import Expression +from .functions import Cos, Exp, Log, Log10, Sin, Sqrt, Tan +from .metrics import ModeAmp, ModePower, generate_validation_data +from .variables import Constant, Variable + +__all__ = [ + "Expression", + "Constant", + "Variable", + "ModeAmp", + "ModePower", + "generate_validation_data", + "Sin", + "Cos", + "Tan", + "Exp", + "Log", + "Log10", + "Sqrt", +] + +# The following code dynamically collects all classes that are subclasses of Expression +# from the specified modules and updates their forward references. This is necessary to handle +# cases where classes reference each other before they are fully defined. The local_vars dictionary +# is used to store these classes and any other necessary types for the forward reference updates. + +import importlib +import inspect + +from .types import ExpressionType + +_module_names = ["base", "variables", "functions", "metrics", "operators"] +_model_classes = set() +_local_vars = {"ExpressionType": ExpressionType} + +for module_name in _module_names: + module = importlib.import_module(f".{module_name}", package=__name__) + for name, obj in inspect.getmembers(module): + if inspect.isclass(obj) and issubclass(obj, Expression): + _model_classes.add(obj) + _local_vars[name] = obj + +for cls in _model_classes: + cls.update_forward_refs(**_local_vars) diff --git a/tidy3d/plugins/expressions/base.py b/tidy3d/plugins/expressions/base.py new file mode 100644 index 000000000..f24cd57cc --- /dev/null +++ b/tidy3d/plugins/expressions/base.py @@ -0,0 +1,228 @@ +from __future__ import annotations + +from abc import ABC, abstractmethod +from typing import TYPE_CHECKING, Any, Generator, Optional, Type + +from tidy3d.components.base import Tidy3dBaseModel +from tidy3d.components.types import TYPE_TAG_STR + +from .types import ExpressionType, NumberOrExpression, NumberType + +if TYPE_CHECKING: + from .operators import ( + Abs, + Add, + Divide, + FloorDivide, + MatMul, + Modulus, + Multiply, + Negate, + Power, + Subtract, + ) + +TYPE_TO_CLASS_MAP: dict[str, Any] = {} + + +class Expression(Tidy3dBaseModel, ABC): + """ + Base class for all expressions in the metrics module. + + This class serves as the foundation for all other components in the metrics module. + It provides common functionality and operator overloading for derived classes. + """ + + class Config: + smart_union = True + + @abstractmethod + def evaluate(self, *args: Any, **kwargs: Any) -> NumberType: + pass + + def __call__(self, *args: Any, **kwargs: Any) -> NumberType: + return self.evaluate(*args, **kwargs) + + def __init_subclass__(cls, **kwargs: dict[str, Any]) -> None: + super().__init_subclass__(**kwargs) + type_value = cls.__fields__.get(TYPE_TAG_STR) + if type_value and type_value.default: + TYPE_TO_CLASS_MAP[type_value.default] = cls + + @classmethod + def parse_obj(cls, obj: dict[str, Any]) -> ExpressionType: + if not isinstance(obj, dict): + raise TypeError("Input must be a dict") + type_value = obj.get(TYPE_TAG_STR) + if type_value is None: + raise ValueError('Missing "type" in data') + subclass = TYPE_TO_CLASS_MAP.get(type_value) + if subclass is None: + raise ValueError(f"Unknown type: {type_value}") + return subclass(**obj) + + def filter( + self, target_type: Type[Expression], target_field: Optional[str] = None + ) -> Generator[Expression, None, None]: + """ + Find all instances of a given type or field in the expression. + + Parameters + ---------- + target_type : Type[Expression] + The type of instances to find. + target_field : Optional[str] = None + The field to aggregate instead of the type. + + Yields + ------ + Expression + Instances of the specified type or field found in the expression. + """ + + def _find_instances(expr: Expression): + if isinstance(expr, target_type): + if target_field: + value = getattr(expr, target_field, None) + if value is not None: + yield value + else: + yield expr + for field in expr.__fields__.values(): + value = getattr(expr, field.name) + if isinstance(value, Expression): + yield from _find_instances(value) + elif isinstance(value, list): + for item in value: + if isinstance(item, Expression): + yield from _find_instances(item) + elif isinstance(value, dict): + for item in value.values(): + if isinstance(item, Expression): + yield from _find_instances(item) + + yield from _find_instances(self) + + @staticmethod + def _to_expression(other: NumberOrExpression | dict[str, Any]) -> ExpressionType: + if isinstance(other, Expression): + return other + elif isinstance(other, dict): + return Expression.parse_obj(other) + else: + from .variables import Constant + + return Constant(other) + + def __neg__(self) -> Negate: + from .operators import Negate + + return Negate(operand=self) + + def __add__(self, other: NumberOrExpression) -> Add: + from .operators import Add + + return Add(left=self, right=other) + + def __radd__(self, other: NumberOrExpression) -> Add: + return self.__add__(other) + + def __sub__(self, other: NumberOrExpression) -> Subtract: + from .operators import Subtract + + return Subtract(left=self, right=other) + + def __rsub__(self, other: NumberOrExpression) -> Subtract: + from .operators import Subtract + + return Subtract(left=other, right=self) + + def __mul__(self, other: NumberOrExpression) -> Multiply: + from .operators import Multiply + + return Multiply(left=self, right=other) + + def __rmul__(self, other: NumberOrExpression) -> Multiply: + return self.__mul__(other) + + def __abs__(self) -> Abs: + from .operators import Abs + + return Abs(operand=self) + + def __truediv__(self, other: NumberOrExpression) -> Divide: + from .operators import Divide + + return Divide(left=self, right=other) + + def __rtruediv__(self, other: NumberOrExpression) -> Divide: + from .operators import Divide + + return Divide(left=other, right=self) + + def __pow__(self, other: NumberOrExpression) -> Power: + from .operators import Power + + return Power(left=self, right=other) + + def __rpow__(self, other: NumberOrExpression) -> Power: + from .operators import Power + + return Power(left=other, right=self) + + def __mod__(self, other: NumberOrExpression) -> Modulus: + from .operators import Modulus + + return Modulus(left=self, right=other) + + def __rmod__(self, other: NumberOrExpression) -> Modulus: + from .operators import Modulus + + return Modulus(left=other, right=self) + + def __floordiv__(self, other: NumberOrExpression) -> FloorDivide: + from .operators import FloorDivide + + return FloorDivide(left=self, right=other) + + def __rfloordiv__(self, other: NumberOrExpression) -> FloorDivide: + from .operators import FloorDivide + + return FloorDivide(left=other, right=self) + + def __matmul__(self, other: NumberOrExpression) -> MatMul: + from .operators import MatMul + + return MatMul(left=self, right=other) + + def __rmatmul__(self, other: NumberOrExpression) -> MatMul: + from .operators import MatMul + + return MatMul(left=other, right=self) + + def __iadd__(self, other: NumberOrExpression) -> Add: + return self + other + + def __isub__(self, other: NumberOrExpression) -> Subtract: + return self - other + + def __imul__(self, other: NumberOrExpression) -> Multiply: + return self * other + + def __itruediv__(self, other: NumberOrExpression) -> Divide: + return self / other + + def __ifloordiv__(self, other: NumberOrExpression) -> FloorDivide: + return self // other + + def __imod__(self, other: NumberOrExpression) -> Modulus: + return self % other + + def __ipow__(self, other: NumberOrExpression) -> Power: + return self**other + + def __str__(self) -> str: + return self.__repr__() + + def __repr__(self) -> str: + return f"{self.__class__.__name__}()" diff --git a/tidy3d/plugins/expressions/functions.py b/tidy3d/plugins/expressions/functions.py new file mode 100644 index 000000000..5bcfcf5c1 --- /dev/null +++ b/tidy3d/plugins/expressions/functions.py @@ -0,0 +1,166 @@ +from typing import Any + +import autograd.numpy as anp +import pydantic.v1 as pd + +from .base import Expression +from .types import NumberOrExpression, NumberType + + +class Function(Expression): + """ + Base class for mathematical functions in expressions. + """ + + operand: NumberOrExpression = pd.Field( + ..., + title="Operand", + description="The operand for the function.", + ) + + _format: str = "{func}({operand})" + + @pd.validator("operand", pre=True, always=True) + def validate_operand(cls, v): + """ + Validate and convert operand to an expression. + """ + return cls._to_expression(v) + + def __init__(self, operand: NumberOrExpression, **kwargs: dict[str, Any]) -> None: + """ + Initialize the function with an operand. + + Parameters + ---------- + operand : NumberOrExpression + The operand for the function. + kwargs : dict[str, Any] + Additional keyword arguments. + """ + super().__init__(operand=operand, **kwargs) + + def __repr__(self): + """ + Return a string representation of the function. + """ + return self._format.format(func=self.type, operand=self.operand) + + +class Sin(Function): + """ + Sine function expression. + """ + + def evaluate(self, *args: Any, **kwargs: Any) -> NumberType: + """ + Evaluate the sine function. + + Returns + ------- + NumberType + Sine of the input value. + """ + return anp.sin(self.operand(*args, **kwargs)) + + +class Cos(Function): + """ + Cosine function expression. + """ + + def evaluate(self, *args: Any, **kwargs: Any) -> NumberType: + """ + Evaluate the cosine function. + + Returns + ------- + NumberType + Cosine of the input value. + """ + return anp.cos(self.operand(*args, **kwargs)) + + +class Tan(Function): + """ + Tangent function expression. + """ + + def evaluate(self, *args: Any, **kwargs: Any) -> NumberType: + """ + Evaluate the tangent function. + + Returns + ------- + NumberType + Tangent of the input value. + """ + return anp.tan(self.operand(*args, **kwargs)) + + +class Exp(Function): + """ + Exponential function expression. + """ + + def evaluate(self, *args: Any, **kwargs: Any) -> NumberType: + """ + Evaluate the exponential function. + + Returns + ------- + NumberType + Exponential of the input value. + """ + return anp.exp(self.operand(*args, **kwargs)) + + +class Log(Function): + """ + Natural logarithm function expression. + """ + + def evaluate(self, *args: Any, **kwargs: Any) -> NumberType: + """ + Evaluate the natural logarithm function. + + Returns + ------- + NumberType + Natural logarithm of the input value. + """ + return anp.log(self.operand(*args, **kwargs)) + + +class Log10(Function): + """ + Base-10 logarithm function expression. + """ + + def evaluate(self, *args: Any, **kwargs: Any) -> NumberType: + """ + Evaluate the base-10 logarithm function. + + Returns + ------- + NumberType + Base-10 logarithm of the input value. + """ + return anp.log10(self.operand(*args, **kwargs)) + + +class Sqrt(Function): + """ + Square root function expression. + """ + + def evaluate(self, *args: Any, **kwargs: Any) -> NumberType: + """ + Evaluate the square root function. + + Returns + ------- + NumberType + Square root of the input value. + """ + return anp.sqrt(self.operand(*args, **kwargs)) diff --git a/tidy3d/plugins/expressions/metrics.py b/tidy3d/plugins/expressions/metrics.py new file mode 100644 index 000000000..04aa79c40 --- /dev/null +++ b/tidy3d/plugins/expressions/metrics.py @@ -0,0 +1,134 @@ +from abc import ABC, abstractmethod +from typing import Any, Optional, Union + +import autograd.numpy as np +import pydantic.v1 as pd +import xarray as xr + +from tidy3d.components.monitor import ModeMonitor +from tidy3d.components.types import Direction, FreqArray + +from .base import Expression +from .types import NumberType +from .variables import Variable + + +def generate_validation_data(expr: Expression) -> dict[str, xr.Dataset]: + """Generate combined dummy simulation data for all metrics in the expression. + + Parameters + ---------- + expr : Expression + The expression containing metrics. + + Returns + ------- + dict[str, xr.Dataset] + The combined validation data. + """ + metrics = set(expr.filter(target_type=Metric)) + combined_data = {k: v for metric in metrics for k, v in metric._validation_data.items()} + return combined_data + + +class Metric(Variable, ABC): + """ + Base class for all metrics. + + To subclass Metric, you must implement an evaluate() method that takes a SimulationData + object and returns a scalar value. + """ + + @property + @abstractmethod + def _validation_data(self) -> Any: + """Return dummy data for this metric.""" + + def __repr__(self) -> str: + return f'{self.type}("{self.monitor_name}")' + + +class ModeAmp(Metric): + """ + Metric for calculating the mode coefficient from a ModeMonitor. + + Examples + -------- + >>> import tidy3d as td + >>> monitor = td.ModeMonitor(size=(1, 1, 0), freqs=[2e14], mode_spec=td.ModeSpec(), name="monitor1") + >>> mode_coeff = ModeAmp.from_mode_monitor(monitor) + >>> expr = abs(mode_coeff) ** 2 + >>> print(expr) + (abs(ModeAmp("monitor1")) ** 2) + """ + + monitor_name: str = pd.Field( + ..., + title="Monitor Name", + description="The name of the mode monitor. This needs to match the name of the monitor in the simulation.", + ) + f: Optional[Union[float, FreqArray]] = pd.Field( # type: ignore + None, + title="Frequency Array", + description="The frequency array. If None, all frequencies in the monitor will be used.", + alias="freqs", + ) + direction: Direction = pd.Field( + "+", + title="Direction", + description="The direction of propagation of the mode.", + ) + mode_index: pd.NonNegativeInt = pd.Field( + 0, + title="Mode Index", + description="The index of the mode.", + ) + + @classmethod + def from_mode_monitor( + cls, monitor: ModeMonitor, mode_index: int = 0, direction: Direction = "+" + ): + return cls( + monitor_name=monitor.name, f=monitor.freqs, mode_index=mode_index, direction=direction + ) + + @property + def _validation_data(self) -> Any: + """Return dummy data for this metric (complex array of mode amplitudes).""" + f = np.atleast_1d(self.f).tolist() if self.f is not None else [1.0] + amps_data = np.random.rand(len(f)) + 1j * np.random.rand(len(f)) + amps = xr.DataArray( + amps_data.reshape(1, 1, -1), + coords={ + "direction": [self.direction], + "mode_index": [self.mode_index], + "f": f, + }, + ) + return {self.monitor_name: xr.Dataset({"amps": amps})} + + def evaluate(self, *args: Any, **kwargs: Any) -> NumberType: + data = super().evaluate(*args, **kwargs) + amps = data[self.monitor_name].amps.sel( + direction=self.direction, mode_index=self.mode_index + ) + if self.f is not None: + f = list(self.f) if isinstance(self.f, tuple) else self.f + amps = amps.sel(f=f, method="nearest") + return np.squeeze(amps.data) + + +class ModePower(ModeAmp): + """ + Metric for calculating the mode power from a ModeMonitor. + + Examples + -------- + >>> import tidy3d as td + >>> monitor = td.ModeMonitor(size=(1, 1, 0), freqs=[2e14], mode_spec=td.ModeSpec(), name="monitor1") + >>> mode_power = ModePower.from_mode_monitor(monitor) + """ + + def evaluate(self, *args: Any, **kwargs: Any) -> NumberType: + amps = super().evaluate(*args, **kwargs) + return np.abs(amps) ** 2 diff --git a/tidy3d/plugins/expressions/operators.py b/tidy3d/plugins/expressions/operators.py new file mode 100644 index 000000000..e25004180 --- /dev/null +++ b/tidy3d/plugins/expressions/operators.py @@ -0,0 +1,174 @@ +from __future__ import annotations + +from typing import Any + +import pydantic.v1 as pd + +from .base import Expression +from .types import NumberOrExpression, NumberType + + +class UnaryOperator(Expression): + """ + Base class for unary operators in the metrics module. + + This class represents an operation with a single operand. + Subclasses should implement the evaluate method to define the specific operation. + """ + + operand: NumberOrExpression = pd.Field( + ..., + title="Operand", + description="The operand for the unary operator.", + ) + + _symbol: str + _format: str = "({symbol}{operand})" + + @pd.validator("operand", pre=True, always=True) + def validate_operand(cls, v): + return cls._to_expression(v) + + def __repr__(self) -> str: + return self._format.format(symbol=self._symbol, operand=self.operand) + + +class BinaryOperator(Expression): + """ + Base class for binary operators in the metrics module. + + This class represents an operation with two operands. + Subclasses should implement the evaluate method to define the specific operation. + """ + + left: NumberOrExpression = pd.Field( + ..., + title="Left", + description="The left operand for the binary operator.", + ) + right: NumberOrExpression = pd.Field( + ..., + title="Right", + description="The right operand for the binary operator.", + ) + + _symbol: str + _format: str = "({left} {symbol} {right})" + + @pd.validator("left", "right", pre=True, always=True) + def validate_operands(cls, v): + return cls._to_expression(v) + + def __repr__(self) -> str: + return self._format.format(left=self.left, symbol=self._symbol, right=self.right) + + +class Add(BinaryOperator): + """ + Represents the addition operation. + """ + + _symbol: str = "+" + + def evaluate(self, *args: Any, **kwargs: Any) -> NumberType: + return self.left(*args, **kwargs) + self.right(*args, **kwargs) + + +class Subtract(BinaryOperator): + """ + Represents the subtraction operation. + """ + + _symbol: str = "-" + + def evaluate(self, *args: Any, **kwargs: Any) -> NumberType: + return self.left(*args, **kwargs) - self.right(*args, **kwargs) + + +class Multiply(BinaryOperator): + """ + Represents the multiplication operation. + """ + + _symbol: str = "*" + + def evaluate(self, *args: Any, **kwargs: Any) -> NumberType: + return self.left(*args, **kwargs) * self.right(*args, **kwargs) + + +class Negate(UnaryOperator): + """ + Represents the negation operation. + """ + + _symbol: str = "-" + + def evaluate(self, *args: Any, **kwargs: Any) -> NumberType: + return -self.operand(*args, **kwargs) + + +class Abs(UnaryOperator): + """ + Represents the absolute value operation. + """ + + _symbol: str = "abs" + _format = "{symbol}({operand})" + + def evaluate(self, *args: Any, **kwargs: Any) -> NumberType: + return abs(self.operand(*args, **kwargs)) + + +class Divide(BinaryOperator): + """ + Represents the division operation. + """ + + _symbol: str = "/" + + def evaluate(self, *args: Any, **kwargs: Any) -> NumberType: + return self.left(*args, **kwargs) / self.right(*args, **kwargs) + + +class Power(BinaryOperator): + """ + Represents the exponentiation operation. + """ + + _symbol: str = "**" + + def evaluate(self, *args: Any, **kwargs: Any) -> NumberType: + return self.left(*args, **kwargs) ** self.right(*args, **kwargs) + + +class Modulus(BinaryOperator): + """ + Represents the modulus operation. + """ + + _symbol: str = "%" + + def evaluate(self, *args: Any, **kwargs: Any) -> NumberType: + return self.left(*args, **kwargs) % self.right(*args, **kwargs) + + +class FloorDivide(BinaryOperator): + """ + Represents the floor division operation. + """ + + _symbol: str = "//" + + def evaluate(self, *args: Any, **kwargs: Any) -> NumberType: + return self.left(*args, **kwargs) // self.right(*args, **kwargs) + + +class MatMul(BinaryOperator): + """ + Represents the matrix multiplication operation. + """ + + _symbol: str = "@" + + def evaluate(self, *args: Any, **kwargs: Any) -> NumberType: + return self.left(*args, **kwargs) @ self.right(*args, **kwargs) diff --git a/tidy3d/plugins/expressions/types.py b/tidy3d/plugins/expressions/types.py new file mode 100644 index 000000000..861e86353 --- /dev/null +++ b/tidy3d/plugins/expressions/types.py @@ -0,0 +1,71 @@ +from typing import TYPE_CHECKING, Annotated, Union + +from pydantic.v1 import Field + +from tidy3d.components.types import TYPE_TAG_STR, ArrayLike, Complex + +if TYPE_CHECKING: + from .functions import Cos, Exp, Log, Log10, Sin, Sqrt, Tan + from .metrics import ModeAmp, ModePower + from .operators import ( + Abs, + Add, + Divide, + FloorDivide, + MatMul, + Modulus, + Multiply, + Negate, + Power, + Subtract, + ) + from .variables import Constant, Variable + +NumberType = Union[int, float, Complex, ArrayLike] + +OperatorType = Annotated[ + Union[ + "Add", + "Subtract", + "Multiply", + "Divide", + "Power", + "Modulus", + "FloorDivide", + "MatMul", + "Negate", + "Abs", + ], + Field(discriminator=TYPE_TAG_STR), +] + +FunctionType = Annotated[ + Union[ + "Sin", + "Cos", + "Tan", + "Exp", + "Log", + "Log10", + "Sqrt", + ], + Field(discriminator=TYPE_TAG_STR), +] + +MetricType = Annotated[ + Union[ + "Constant", + "Variable", + "ModeAmp", + "ModePower", + ], + Field(discriminator=TYPE_TAG_STR), +] + +ExpressionType = Union[ + OperatorType, + FunctionType, + MetricType, +] + +NumberOrExpression = Union[NumberType, ExpressionType] diff --git a/tidy3d/plugins/expressions/variables.py b/tidy3d/plugins/expressions/variables.py new file mode 100644 index 000000000..ae32591c1 --- /dev/null +++ b/tidy3d/plugins/expressions/variables.py @@ -0,0 +1,96 @@ +from typing import Any, Optional + +import pydantic.v1 as pd + +from .base import Expression +from .types import NumberType + + +class Variable(Expression): + """ + Variable class represents a placeholder for a value provided at evaluation time. + + Attributes + ---------- + name : Optional[str] = None + The name of the variable used for lookup during evaluation. + + Methods + ------- + evaluate(*args, **kwargs) + Evaluates the variable by retrieving its value from provided arguments. + + Notes + ----- + - If `name` is `None`, the variable expects a single positional argument during evaluation. + - If `name` is provided, the variable expects a corresponding keyword argument during evaluation. + - Mixing positional and keyword arguments is allowed. + - Multiple positional arguments are disallowed and will raise a `ValueError`. + + Examples + -------- + >>> x = Variable() # unnamed + >>> y = Variable(name='y') + >>> z = Variable(name='z') + >>> expr = x + y + z + >>> expr(5, y=3, z=2) + 10 + """ + + name: Optional[str] = pd.Field( + None, + title="Name", + description="The name of the variable used for lookup during evaluation.", + ) + + def evaluate(self, *args: Any, **kwargs: Any) -> NumberType: + if self.name: + if self.name not in kwargs: + raise ValueError(f"Variable '{self.name}' not provided.") + return kwargs[self.name] + else: + if not args: + raise ValueError("No positional argument provided for unnamed variable.") + if len(args) > 1: + raise ValueError("Multiple positional arguments provided for unnamed variable.") + return args[0] + + def __repr__(self) -> str: + return self.name if self.name else "Variable()" + + +class Constant(Variable): + """ + Constant class represents a fixed value in an expression. + + Attributes + ---------- + value : NumberType + The fixed value of the constant. + + Methods + ------- + evaluate(*args, **kwargs) + Returns the value of the constant. + + Examples + -------- + >>> c = Constant(5) + >>> c.evaluate() + 5 + """ + + value: NumberType = pd.Field( + ..., + title="Value", + description="The fixed value of the constant.", + ) + + def __init__(self, value: NumberType, **kwargs: dict[str, Any]) -> None: + super().__init__(value=value, **kwargs) + + def evaluate(self, *args: Any, **kwargs: Any) -> NumberType: + return self.value + + def __repr__(self) -> str: + return f"{self.value}" diff --git a/tidy3d/plugins/invdes/__init__.py b/tidy3d/plugins/invdes/__init__.py index cb20ffef0..8bb3d5e96 100644 --- a/tidy3d/plugins/invdes/__init__.py +++ b/tidy3d/plugins/invdes/__init__.py @@ -2,6 +2,11 @@ from . import utils from .design import InverseDesign, InverseDesignMulti +from .initialization import ( + CustomInitializationSpec, + RandomInitializationSpec, + UniformInitializationSpec, +) from .optimizer import AdamOptimizer from .penalty import ErosionDilationPenalty from .region import TopologyDesignRegion @@ -16,5 +21,8 @@ "TopologyDesignRegion", "AdamOptimizer", "InverseDesignResult", + "RandomInitializationSpec", + "UniformInitializationSpec", + "CustomInitializationSpec", "utils", ) diff --git a/tidy3d/plugins/invdes/design.py b/tidy3d/plugins/invdes/design.py index 06266ae93..f8a4487b0 100644 --- a/tidy3d/plugins/invdes/design.py +++ b/tidy3d/plugins/invdes/design.py @@ -1,14 +1,19 @@ # container for everything defining the inverse design +from __future__ import annotations + import abc import typing import autograd.numpy as anp +import numpy as np import pydantic.v1 as pd import tidy3d as td -import tidy3d.web as web from tidy3d.components.autograd import get_static +from tidy3d.exceptions import ValidationError +from tidy3d.plugins.expressions.metrics import Metric, generate_validation_data +from tidy3d.plugins.expressions.types import ExpressionType from .base import InvdesBaseModel from .region import DesignRegionType @@ -38,31 +43,77 @@ class AbstractInverseDesign(InvdesBaseModel, abc.ABC): description="If ``True``, will print the regular output from ``web`` functions.", ) + metric: typing.Optional[ExpressionType] = pd.Field( + None, + title="Objective Metric", + description="Serializable expression defining the objective function.", + ) + def make_objective_fn( - self, post_process_fn: typing.Callable + self, post_process_fn: typing.Optional[typing.Callable] = None, maximize: bool = True ) -> typing.Callable[[anp.ndarray], tuple[float, dict]]: - """construct the objective function for this ``InverseDesignMulti`` object.""" + """Construct the objective function for this InverseDesign object.""" + + if (post_process_fn is None) and (self.metric is None): + raise ValueError("Either 'post_process_fn' or 'metric' must be provided.") + + if (post_process_fn is not None) and (self.metric is not None): + raise ValueError("Provide only one of 'post_process_fn' or 'metric', not both.") + + direction_multiplier = 1 if maximize else -1 def objective_fn(params: anp.ndarray, aux_data: dict = None) -> float: """Full objective function.""" - data = self.to_simulation_data(params=params) - # construct objective function values - post_process_val = post_process_fn(data) + if self.metric is None: + post_process_val = post_process_fn(data) + elif isinstance(data, td.SimulationData): + post_process_val = self.metric.evaluate(data) + elif getattr(data, "type", None) == "BatchData": + raise NotImplementedError("Metrics currently do not support 'BatchData'") + else: + raise ValueError(f"Invalid data type: {type(data)}") penalty_value = self.design_region.penalty_value(params) - objective_fn_val = post_process_val - penalty_value + objective_fn_val = direction_multiplier * post_process_val - penalty_value - # store things in ``aux_data`` passed by reference + # Store auxiliary data if provided if aux_data is not None: aux_data["penalty"] = get_static(penalty_value) aux_data["post_process_val"] = get_static(post_process_val) + aux_data["objective_fn_val"] = get_static(objective_fn_val) * direction_multiplier + if isinstance(data, td.SimulationData): + aux_data["sim_data"] = data.to_static() + else: + aux_data["sim_data"] = {k: v.to_static() for k, v in data.items()} + aux_data["params"] = params return objective_fn_val return objective_fn + @property + def initial_simulation(self) -> td.Simulation: + """Return a simulation with the initial design region parameters.""" + initial_params = self.design_region.initial_parameters + return self.to_simulation(initial_params) + + def run(self, simulation, **kwargs) -> td.SimulationData: + """Run a single tidy3d simulation.""" + from tidy3d.web import run + + kwargs.setdefault("verbose", self.verbose) + kwargs.setdefault("task_name", self.task_name) + return run(simulation, **kwargs) + + def run_async(self, simulations, **kwargs) -> web.BatchData: # noqa: F821 + """Run a batch of tidy3d simulations.""" + from tidy3d.web import run_async + + kwargs.setdefault("verbose", self.verbose) + return run_async(simulations, **kwargs) + class InverseDesign(AbstractInverseDesign): """Container for an inverse design problem.""" @@ -85,6 +136,81 @@ class InverseDesign(AbstractInverseDesign): _check_sim_pixel_size = check_pixel_size("simulation") + @pd.root_validator(pre=False) + def _validate_model(cls, values: dict) -> dict: + cls._validate_metric(values) + return values + + @staticmethod + def _validate_metric(values: dict) -> dict: + metric_expr = values.get("metric") + if not metric_expr: + return values + simulation = values.get("simulation") + for metric in metric_expr.filter(Metric): + InverseDesign._validate_metric_monitor_name(metric, simulation) + InverseDesign._validate_metric_mode_index(metric, simulation) + InverseDesign._validate_metric_f(metric, simulation) + InverseDesign._validate_metric_data(metric_expr, simulation) + return values + + @staticmethod + def _validate_metric_monitor_name(metric: Metric, simulation: td.Simulation) -> None: + """Validate that the monitor name of the metric exists in the simulation.""" + monitor = next((m for m in simulation.monitors if m.name == metric.monitor_name), None) + if monitor is None: + raise ValidationError( + f"Monitor named '{metric.monitor_name}' associated with the metric not found in the simulation monitors." + ) + + @staticmethod + def _validate_metric_mode_index(metric: Metric, simulation: td.Simulation) -> None: + """Validate that the mode index of the metric is within the bounds of the monitor's ``ModeSpec.num_modes``.""" + monitor = next((m for m in simulation.monitors if m.name == metric.monitor_name), None) + if metric.mode_index >= monitor.mode_spec.num_modes: + raise ValidationError( + f"Mode index '{metric.mode_index}' for metric associated with monitor " + f"'{metric.monitor_name}' is out of bounds. " + f"Maximum allowed mode index is '{monitor.mode_spec.num_modes - 1}'." + ) + + @staticmethod + def _validate_metric_f(metric: Metric, simulation: td.Simulation) -> None: + """Validate that the frequencies of the metric are present in the monitor.""" + monitor = next((m for m in simulation.monitors if m.name == metric.monitor_name), None) + if metric.f is not None: + metric_f_list = [metric.f] if isinstance(metric.f, float) else metric.f + if len(metric_f_list) != 1: + raise ValidationError("Only a single frequency is supported for the metric.") + for freq in metric_f_list: + if not any(np.isclose(freq, monitor.freqs, atol=1.0)): + raise ValidationError( + f"Frequency '{freq}' for metric associated with monitor " + f"'{metric.monitor_name}' not found in monitor frequencies." + ) + else: + if len(monitor.freqs) != 1: + raise ValidationError( + f"Monitor '{metric.monitor_name}' must contain only a single frequency when metric.f is None." + ) + + @staticmethod + def _validate_metric_data(expr: ExpressionType, simulation: td.Simulation) -> None: + """Validate that expression can be evaluated and returns a real scalar.""" + data = generate_validation_data(expr) + try: + result = expr(data) + except Exception as e: + raise ValidationError(f"Failed to evaluate the metric expression: {str(e)}") from e + if len(np.ravel(result)) > 1: + raise ValidationError( + f"The expression must return a scalar value or an array of length 1 (got {result})." + ) + if not np.all(np.isreal(result)): + raise ValidationError( + f"The expression must return a real (not complex) value (got {result})." + ) + def is_output_monitor(self, monitor: td.Monitor) -> bool: """Whether a monitor is added to the ``JaxSimulation`` as an ``output_monitor``.""" @@ -128,10 +254,7 @@ def to_simulation(self, params: anp.ndarray) -> td.Simulation: def to_simulation_data(self, params: anp.ndarray, **kwargs) -> td.SimulationData: """Convert the ``InverseDesign`` to a ``td.Simulation`` and run it.""" simulation = self.to_simulation(params=params) - kwargs.setdefault("task_name", self.task_name) - return web.run(simulation, verbose=self.verbose, **kwargs) - # sim_data = job.run() - # return sim_data + return self.run(simulation, **kwargs) class InverseDesignMulti(AbstractInverseDesign): @@ -201,11 +324,10 @@ def to_simulation(self, params: anp.ndarray) -> dict[str, td.Simulation]: simulation_list = [design.to_simulation(params) for design in self.designs] return dict(zip(self.task_names, simulation_list)) - def to_simulation_data(self, params: anp.ndarray, **kwargs) -> web.BatchData: + def to_simulation_data(self, params: anp.ndarray, **kwargs) -> web.BatchData: # noqa: F821 """Convert the ``InverseDesignMulti`` to a set of ``td.Simulation``s and run async.""" simulations = self.to_simulation(params) - kwargs.setdefault("verbose", self.verbose) - return web.run_async(simulations, **kwargs) + return self.run_async(simulations, **kwargs) InverseDesignType = typing.Union[InverseDesign, InverseDesignMulti] diff --git a/tidy3d/plugins/invdes/initialization.py b/tidy3d/plugins/invdes/initialization.py new file mode 100644 index 000000000..eb75a940c --- /dev/null +++ b/tidy3d/plugins/invdes/initialization.py @@ -0,0 +1,135 @@ +# module providing classes for initializing the parameters in an inverse design problem + +from __future__ import annotations + +from abc import ABC, abstractmethod +from typing import Optional, Union + +import numpy as np +import pydantic.v1 as pd +from numpy.typing import NDArray + +import tidy3d as td +from tidy3d.components.base import Tidy3dBaseModel +from tidy3d.components.types import ArrayLike +from tidy3d.exceptions import ValidationError + + +class AbstractInitializationSpec(Tidy3dBaseModel, ABC): + """Abstract base class for initialization specifications.""" + + @abstractmethod + def create_parameters(self, shape: tuple[int, ...]) -> NDArray: + """Generate the parameter array based on the specification.""" + pass + + +class RandomInitializationSpec(AbstractInitializationSpec): + """Specification for random initial parameters. + + When a seed is provided, a call to `create_parameters` will always return the same array. + """ + + min_value: float = pd.Field( + 0.0, + ge=0.0, + le=1.0, + title="Minimum Value", + description="Minimum value for the random parameters (inclusive).", + ) + max_value: float = pd.Field( + 1.0, + ge=0.0, + le=1.0, + title="Maximum Value", + description="Maximum value for the random parameters (exclusive).", + ) + seed: Optional[pd.NonNegativeInt] = pd.Field( + None, description="Seed for the random number generator." + ) + + @pd.root_validator(pre=False) + def _validate_max_ge_min(cls, values): + """Ensure that max_value is greater than or equal to min_value.""" + minval = values.get("min_value") + maxval = values.get("max_value") + if minval > maxval: + raise ValidationError( + f"'max_value' ({maxval}) must be greater or equal than 'min_value' ({minval})" + ) + return values + + def create_parameters(self, shape: tuple[int, ...]) -> NDArray: + """Generate the parameter array based on the specification.""" + rng = np.random.default_rng(self.seed) + return rng.uniform(self.min_value, self.max_value, shape) + + +class UniformInitializationSpec(AbstractInitializationSpec): + """Specification for uniform initial parameters.""" + + value: float = pd.Field( + 0.5, + ge=0.0, + le=1.0, + title="Value", + description="Value to use for all elements in the parameter array.", + ) + + def create_parameters(self, shape: tuple[int, ...]) -> NDArray: + """Generate the parameter array based on the specification.""" + return np.full(shape, self.value) + + +class CustomInitializationSpec(AbstractInitializationSpec): + """Specification for custom initial parameters provided by the user.""" + + params: ArrayLike = pd.Field( + ..., + title="Parameters", + description="Custom parameters provided by the user.", + ) + + @pd.validator("params") + def _validate_params_range(cls, value, values): + """Ensure that all parameter values are between 0 and 1.""" + if np.any((value < 0) | (value > 1)): + raise ValidationError("'params' need to be between 0 and 1.") + return value + + @pd.validator("params") + def _validate_params_dtype(cls, value, values): + """Ensure that params is real-valued.""" + if np.issubdtype(value.dtype, np.bool_): + td.log.warning( + "Got a boolean array for 'params'. " + "This will be treated as a floating point array." + ) + value = value.astype(float) + elif not np.issubdtype(value.dtype, np.floating): + raise ValidationError(f"'params' need to be real-valued, but got '{value.dtype}'.") + return value + + @pd.validator("params") + def _validate_params_3d(cls, value, values): + """Ensure that params is a 3D array.""" + if value.ndim != 3: + raise ValidationError(f"'params' must be 3D, but got {value.ndim}D.") + return value + + def create_parameters(self, shape: tuple[int, ...]) -> NDArray: + """Return the custom parameters provided by the user.""" + params = np.asarray(self.params) + if params.shape != shape: + raise ValueError( + f"Provided 'params.shape' ('{params.shape}') does not match " + f"the shape of the custom parameters ('{shape}')." + ) + return params + + +InitializationSpecType = Union[ + RandomInitializationSpec, + UniformInitializationSpec, + CustomInitializationSpec, +] diff --git a/tidy3d/plugins/invdes/optimizer.py b/tidy3d/plugins/invdes/optimizer.py index 27fd994e6..08758f12b 100644 --- a/tidy3d/plugins/invdes/optimizer.py +++ b/tidy3d/plugins/invdes/optimizer.py @@ -10,6 +10,7 @@ import pydantic.v1 as pd import tidy3d as td +from tidy3d.components.types import TYPE_TAG_STR from .base import InvdesBaseModel from .design import InverseDesignType @@ -23,14 +24,21 @@ class AbstractOptimizer(InvdesBaseModel, abc.ABC): ..., title="Inverse Design Specification", description="Specification describing the inverse design problem we wish to optimize.", + discriminator=TYPE_TAG_STR, ) - learning_rate: pd.NonNegativeFloat = pd.Field( + learning_rate: pd.PositiveFloat = pd.Field( ..., title="Learning Rate", description="Step size for the gradient descent optimizer.", ) + maximize: bool = pd.Field( + True, + title="Direction of Optimization", + description="If ``True``, the optimizer will maximize the objective function. If ``False``, the optimizer will minimize the objective function.", + ) + num_steps: pd.PositiveInt = pd.Field( ..., title="Number of Steps", @@ -63,6 +71,10 @@ class AbstractOptimizer(InvdesBaseModel, abc.ABC): def initial_state(self, parameters: np.ndarray) -> dict: """The initial state of the optimizer.""" + def validate_pre_upload(self) -> None: + """Validate the fully initialized optimizer is ok for upload to our servers.""" + self.design.simulation.validate_pre_upload() + def display_fn(self, result: InverseDesignResult, step_index: int) -> None: """Default display function while optimizing.""" print(f"step ({step_index + 1}/{self.num_steps})") @@ -71,43 +83,99 @@ def display_fn(self, result: InverseDesignResult, step_index: int) -> None: print(f"\tpost_process_val = {result.post_process_val[-1]:.3e}") print(f"\tpenalty = {result.penalty[-1]:.3e}") - def _initialize_result(self, params0: anp.ndarray = None) -> InverseDesignResult: - """Create an initially empty ``InverseDesignResult`` from the starting parameters.""" - - # initialize optimizer - if params0 is None: - params0 = self.design.design_region.params_half - params0 = anp.array(params0) - + def initialize_result( + self, params0: typing.Optional[anp.ndarray] = None + ) -> InverseDesignResult: + """ + Create an initially empty `InverseDesignResult` from the starting parameters. + + Returns + ------- + InverseDesignResult + An instance of `InverseDesignResult` initialized with the starting parameters and state. + """ + if params0 is not None: + td.log.warning( + "The 'params0' argument is deprecated and will be removed in the future. " + "Please use a 'DesignRegion.initialization_spec' in the design region " + "to specify initial parameters instead. For now, 'params0' will take precedence " + "over 'initialization_spec'." + ) + else: + params0 = self.design.design_region.initial_parameters state = self.initial_state(params0) # initialize empty result return InverseDesignResult(design=self.design, opt_state=[state], params=[params0]) def run( - self, post_process_fn: typing.Callable, params0: anp.ndarray = None + self, + post_process_fn: typing.Optional[typing.Callable] = None, + callback: typing.Optional[typing.Callable] = None, + params0: anp.ndarray = None, ) -> InverseDesignResult: - """Run this inverse design problem from an optional initial set of parameters.""" - self.design.design_region._check_params(params0) - starting_result = self._initialize_result(params0) - return self.continue_run(result=starting_result, post_process_fn=post_process_fn) + """Run this inverse design problem from an optional initial set of parameters. + + Parameters + ---------- + post_process_fn : Optional[Callable] = None + Function to apply on the simulation data results to produce the final objective function + value. If not provided, then ``Optimizer.design.metric`` must be defined. + callback : Optional[Callable] = None + Callback function to apply at every iteration step for extra functionality. Does not + need to be differentiable. This takes the optimizer ``result`` as a positional argument + and the ``step_index`` and ``aux_data`` as optional arguments. + params0 : anp.ndarray = None + Deprecated. Initial set of parameters. Use ``TopologyDesignRegion.intialization_spec`` instead. + """ + starting_result = self.initialize_result(params0) + return self.continue_run( + result=starting_result, + num_steps=self.num_steps, + post_process_fn=post_process_fn, + callback=callback, + ) def continue_run( - self, result: InverseDesignResult, post_process_fn: typing.Callable + self, + result: InverseDesignResult, + num_steps: int = None, + post_process_fn: typing.Optional[typing.Callable] = None, + callback: typing.Optional[typing.Callable] = None, ) -> InverseDesignResult: - """Run optimizer for a series of steps with an initialized state.""" + """Run optimizer for a series of steps with an initialized state. + + Parameters + ---------- + result : InverseDesignResult + Optimization result from previous run, or a starting optimization result data structure. + num_steps : int = None + Number of steps to continue the run for. If not provided, runs for the remainder of the + steps up to ``self.num_steps``. + post_process_fn : Optional[Callable] = None + Function to apply on the simulation data results to produce the final objective function + value. If not provided, then ``Optimizer.design.metric`` must be defined. + callback : Optional[Callable] = None + Callback function to apply at every iteration step for extra functionality. Does not + need to be differentiable. This takes the optimizer ``result`` as a positional argument + and the ``step_index`` and ``aux_data`` as optional arguments. + """ # get the last state of the optimizer and the last parameters opt_state = result.get_last("opt_state") params = result.get_last("params") history = deepcopy(result.history) + done_steps = len(history["objective_fn_val"]) # use autograd to take gradient the objective function - objective_fn = self.design.make_objective_fn(post_process_fn) + objective_fn = self.design.make_objective_fn(post_process_fn, maximize=self.maximize) val_and_grad_fn = ag.value_and_grad(objective_fn) + if num_steps is None: + num_steps = self.num_steps - done_steps + # main optimization loop - for step_index in range(self.num_steps): + for step_index in range(done_steps, done_steps + num_steps): aux_data = {} val, grad = val_and_grad_fn(params, aux_data=aux_data) @@ -116,9 +184,8 @@ def continue_run( "All elements of the gradient are almost zero. This likely indicates " "a problem with the optimization set up. This can occur if the symmetry of the " "simulation and design region are preventing any data to be recorded in the " - "'output_monitors'. In this case, we recommend creating the initial parameters " - " as 'params0 = DesignRegion.params_random' and passing this to " - "'Optimizer.run()' to break the symmetry in the design region. " + "'output_monitors'. In this case, we recommend initializing with a " + "'RandomInitializationSpec' to break the symmetry in the design region. " "This zero gradient can also occur if the objective function return value does " "not have a contribution from the input arguments. We recommend carefully " "inspecting your objective function to ensure that the variables passed to the " @@ -136,7 +203,7 @@ def continue_run( params = anp.clip(params, a_min=0.0, a_max=1.0) # save the history of scalar values - history["objective_fn_val"].append(val) + history["objective_fn_val"].append(aux_data["objective_fn_val"]) history["penalty"].append(penalty) history["post_process_val"].append(post_process_val) @@ -150,6 +217,8 @@ def continue_run( # display information result = InverseDesignResult(design=result.design, **history) self.display_fn(result, step_index=step_index) + if callback: + callback(result, step_index=step_index, aux_data=aux_data) # save current results to file if self.results_cache_fname: @@ -158,16 +227,33 @@ def continue_run( return InverseDesignResult(design=result.design, **history) def continue_run_from_file( - self, fname: str, post_process_fn: typing.Callable + self, + fname: str, + num_steps: int = None, + post_process_fn: typing.Optional[typing.Callable] = None, + callback: typing.Optional[typing.Callable] = None, ) -> InverseDesignResult: """Continue the optimization run from a ``.pkl`` file with an ``InverseDesignResult``.""" result = InverseDesignResult.from_file(fname) - return self.continue_run(result=result, post_process_fn=post_process_fn) + return self.continue_run( + result=result, + num_steps=num_steps, + post_process_fn=post_process_fn, + callback=callback, + ) - def continue_run_from_history(self, post_process_fn: typing.Callable) -> InverseDesignResult: + def continue_run_from_history( + self, + num_steps: int = None, + post_process_fn: typing.Optional[typing.Callable] = None, + callback: typing.Optional[typing.Callable] = None, + ) -> InverseDesignResult: """Continue the optimization run from a ``.pkl`` file with an ``InverseDesignResult``.""" return self.continue_run_from_file( - fname=self.results_cache_fname, post_process_fn=post_process_fn + fname=self.results_cache_fname, + num_steps=num_steps, + post_process_fn=post_process_fn, + callback=callback, ) @@ -176,17 +262,21 @@ class AdamOptimizer(AbstractOptimizer): beta1: float = pd.Field( 0.9, + ge=0.0, + le=1.0, title="Beta 1", description="Beta 1 parameter in the Adam optimization method.", ) beta2: float = pd.Field( 0.999, + ge=0.0, + le=1.0, title="Beta 2", description="Beta 2 parameter in the Adam optimization method.", ) - eps: float = pd.Field( + eps: pd.PositiveFloat = pd.Field( 1e-8, title="Epsilon", description="Epsilon parameter in the Adam optimization method.", diff --git a/tidy3d/plugins/invdes/penalty.py b/tidy3d/plugins/invdes/penalty.py index 73a20a668..621cbfb9e 100644 --- a/tidy3d/plugins/invdes/penalty.py +++ b/tidy3d/plugins/invdes/penalty.py @@ -7,7 +7,7 @@ import pydantic.v1 as pd from tidy3d.constants import MICROMETER -from tidy3d.plugins.autograd.invdes import get_kernel_size_px, make_erosion_dilation_penalty +from tidy3d.plugins.autograd.invdes import make_erosion_dilation_penalty from .base import InvdesBaseModel @@ -48,7 +48,7 @@ class ErosionDilationPenalty(AbstractPenalty): """ - length_scale: pd.NonNegativeFloat = pd.Field( + length_scale: pd.PositiveFloat = pd.Field( ..., title="Length Scale", description="Length scale of erosion and dilation. " @@ -58,24 +58,29 @@ class ErosionDilationPenalty(AbstractPenalty): units=MICROMETER, ) - beta: pd.PositiveFloat = pd.Field( + beta: float = pd.Field( 100.0, + ge=1.0, title="Projection Beta", description="Strength of the ``tanh`` projection. " "Corresponds to ``beta`` in the :class:`BinaryProjector. " "Higher values correspond to stronger discretization.", ) - eta0: pd.PositiveFloat = pd.Field( + eta0: float = pd.Field( 0.5, + ge=0.0, + le=1.0, title="Projection Midpoint", description="Value between 0 and 1 that sets the projection midpoint. In other words, " "for values of ``eta0``, the projected values are halfway between minimum and maximum. " "Corresponds to ``eta`` in the :class:`BinaryProjector`.", ) - delta_eta: pd.PositiveFloat = pd.Field( + delta_eta: float = pd.Field( 0.01, + ge=0.0, + le=1.0, title="Delta Eta Cutoff", description="The binarization threshold for erosion and dilation operations " "The thresholds are ``0 + delta_eta`` on the low end and ``1 - delta_eta`` on the high end. " @@ -85,9 +90,8 @@ class ErosionDilationPenalty(AbstractPenalty): def evaluate(self, x: anp.ndarray, pixel_size: float) -> float: """Evaluate this penalty.""" - filter_size = get_kernel_size_px(self.length_scale, pixel_size) penalty_fn = make_erosion_dilation_penalty( - filter_size, self.beta, self.eta0, self.delta_eta + self.length_scale, pixel_size, beta=self.beta, eta=self.eta0, delta_eta=self.delta_eta ) penalty_unweighted = penalty_fn(x) return self.weight * penalty_unweighted diff --git a/tidy3d/plugins/invdes/region.py b/tidy3d/plugins/invdes/region.py index 680c9d7c1..b94eaf868 100644 --- a/tidy3d/plugins/invdes/region.py +++ b/tidy3d/plugins/invdes/region.py @@ -2,15 +2,19 @@ import abc import typing +import warnings import autograd.numpy as anp import numpy as np import pydantic.v1 as pd +from autograd import elementwise_grad, grad import tidy3d as td -from tidy3d.components.types import Coordinate, Size +from tidy3d.components.types import TYPE_TAG_STR, Coordinate, Size +from tidy3d.exceptions import ValidationError from .base import InvdesBaseModel +from .initialization import InitializationSpecType, UniformInitializationSpec from .penalty import PenaltyType from .transformation import TransformationType @@ -36,6 +40,7 @@ class DesignRegion(InvdesBaseModel, abc.ABC): eps_bounds: typing.Tuple[float, float] = pd.Field( ..., + ge=1.0, title="Relative Permittivity Bounds", description="Minimum and maximum relative permittivity expressed to the design region.", ) @@ -58,6 +63,26 @@ class DesignRegion(InvdesBaseModel, abc.ABC): "inside of the penalties directly through the ``.weight`` field.", ) + initialization_spec: InitializationSpecType = pd.Field( + UniformInitializationSpec(value=0.5), + title="Initialization Specification", + description="Specification of how to initialize the parameters in the design region.", + discriminator=TYPE_TAG_STR, + ) + + def _post_init_validators(self): + """Automatically call any `_validate_XXX` method.""" + for attr_name in dir(self): + if attr_name.startswith("_validate") and callable(getattr(self, attr_name)): + getattr(self, attr_name)() + + def _validate_eps_bounds(self): + if self.eps_bounds[1] < self.eps_bounds[0]: + raise ValidationError( + f"Maximum relative permittivity ({self.eps_bounds[1]}) must be " + f"greater than minimum relative permittivity ({self.eps_bounds[0]})." + ) + @property def geometry(self) -> td.Box: """``Box`` corresponding to this design region.""" @@ -95,6 +120,13 @@ def evaluate_penalty(self, penalty: None) -> float: def to_structure(self, *args, **kwargs) -> td.Structure: """Convert this ``DesignRegion`` into a ``Structure`` with tracers. Implement in subs.""" + @property + def initial_parameters(self) -> np.ndarray: + """Generate initial parameters based on the initialization specification.""" + params0 = self.initialization_spec.create_parameters(self.params_shape) + self._check_params(params0) + return params0 + class TopologyDesignRegion(DesignRegion): """Design region as a pixellated permittivity grid.""" @@ -110,6 +142,13 @@ class TopologyDesignRegion(DesignRegion): "a value on the same order as the grid size.", ) + uniform: tuple[bool, bool, bool] = pd.Field( + (False, False, True), + title="Uniform", + description="Axes along which the design should be uniform. By default, the structure " + "is assumed to be uniform, i.e. invariant, in the z direction.", + ) + transformations: typing.Tuple[TransformationType, ...] = pd.Field( (), title="Transformations", @@ -138,12 +177,67 @@ class TopologyDesignRegion(DesignRegion): "Supplying ``False`` will completely leave out the override structure.", ) + def _validate_eps_values(self): + """Validate the epsilon values by evaluating the transformations.""" + try: + x = self.initial_parameters + self.eps_values(x) + except Exception as e: + raise ValidationError(f"Could not evaluate transformations: {str(e)}") from e + + def _validate_penalty_value(self): + """Validate the penalty values by evaluating the penalties.""" + try: + x = self.initial_parameters + self.penalty_value(x) + except Exception as e: + raise ValidationError(f"Could not evaluate penalties: {str(e)}") from e + + def _validate_gradients(self): + """Validate the gradients of the penalties and transformations.""" + x = self.initial_parameters + + penalty_independent = False + if self.penalties: + with warnings.catch_warnings(record=True) as w: + penalty_grad = grad(self.penalty_value)(x) + penalty_independent = any("independent" in str(warn.message).lower() for warn in w) + if np.any(np.isnan(penalty_grad) | np.isinf(penalty_grad)): + raise ValidationError("Penalty gradients contain 'NaN' or 'Inf' values.") + + eps_independent = False + if self.transformations: + with warnings.catch_warnings(record=True) as w: + eps_grad = elementwise_grad(self.eps_values)(x) + eps_independent = any("independent" in str(warn.message).lower() for warn in w) + if np.any(np.isnan(eps_grad) | np.isinf(eps_grad)): + raise ValidationError("Transformation gradients contain 'NaN' or 'Inf' values.") + + if penalty_independent and eps_independent: + raise ValidationError( + "Both penalty and transformation gradients appear to be independent of the input parameters. " + "This indicates that the optimization will not function correctly. " + "Please double-check the definitions of both the penalties and transformations." + ) + elif penalty_independent: + td.log.warning( + "Penalty gradient seems independent of input, meaning that it " + "will not contribute to the objective gradient during optimization. " + "This is likely not correct - double-check the penalties." + ) + elif eps_independent: + td.log.warning( + "Transformation gradient seems independent of input, meaning that it " + "will not contribute to the objective gradient during optimization. " + "This is likely not correct - double-check the transformations." + ) + @staticmethod def _check_params(params: anp.ndarray = None): """Ensure ``params`` are between 0 and 1.""" if params is None: return - if np.any(params < 0) or np.any(params > 1): + if np.any((params < 0) | (params > 1)): raise ValueError( "Parameters in the 'invdes' plugin's topology optimization feature " "are restricted to be between 0 and 1." @@ -152,36 +246,46 @@ def _check_params(params: anp.ndarray = None): @property def params_shape(self) -> typing.Tuple[int, int, int]: """Shape of the parameters array in (x, y, z), given the ``pixel_size`` and bounds.""" - # rmin, rmax = np.array(self.geometry.bounds) - # lengths = rmax - rmin side_lengths = np.array(self.size) num_pixels = np.ceil(side_lengths / self.pixel_size) # TODO: if the structure is infinite but the simulation is finite, need reduced bounds - num_pixels[np.isinf(num_pixels)] = 1 + num_pixels[np.logical_or(np.isinf(num_pixels), self.uniform)] = 1 return tuple(int(n) for n in num_pixels) + def _warn_deprecate_params(self): + td.log.warning( + "Parameter initialization via design region methods is deprecated and will be " + "removed in the future. Please specify this through the design region's " + "'initialization_spec' instead." + ) + def params_uniform(self, value: float) -> np.ndarray: """Make an array of parameters with all the same value.""" + self._warn_deprecate_params() return value * np.ones(self.params_shape) @property def params_random(self) -> np.ndarray: """Convenience for generating random parameters between (0,1) with correct shape.""" + self._warn_deprecate_params() return np.random.random(self.params_shape) @property def params_zeros(self): """Convenience for generating random parameters of all 0 values with correct shape.""" + self._warn_deprecate_params() return self.params_uniform(0.0) @property def params_half(self): """Convenience for generating random parameters of all 0.5 values with correct shape.""" + self._warn_deprecate_params() return self.params_uniform(0.5) @property def params_ones(self): """Convenience for generating random parameters of all 1 values with correct shape.""" + self._warn_deprecate_params() return self.params_uniform(1.0) @property @@ -216,7 +320,7 @@ def eps_values(self, params: anp.ndarray) -> anp.ndarray: return eps_arr.reshape(params.shape) def to_structure(self, params: anp.ndarray) -> td.Structure: - """Convert this ``DesignRegion`` into a custom ``JaxStructure``.""" + """Convert this ``DesignRegion`` into a custom ``Structure``.""" self._check_params(params) coords = self.coords diff --git a/tidy3d/plugins/invdes/transformation.py b/tidy3d/plugins/invdes/transformation.py index e6719e384..cb9e343ce 100644 --- a/tidy3d/plugins/invdes/transformation.py +++ b/tidy3d/plugins/invdes/transformation.py @@ -8,7 +8,7 @@ import tidy3d as td from tidy3d.plugins.autograd.functions import threshold -from tidy3d.plugins.autograd.invdes import get_kernel_size_px, make_filter_and_project +from tidy3d.plugins.autograd.invdes import make_filter_and_project from .base import InvdesBaseModel @@ -34,7 +34,7 @@ class FilterProject(InvdesBaseModel): """ - radius: float = pd.Field( + radius: pd.PositiveFloat = pd.Field( ..., title="Filter Radius", description="Radius of the filter to convolve with supplied spatial data. " @@ -50,13 +50,16 @@ class FilterProject(InvdesBaseModel): beta: float = pd.Field( 1.0, + ge=1.0, title="Beta", description="Steepness of the binarization, " "higher means more sharp transition " "at the expense of gradient accuracy and ease of optimization. ", ) - eta: float = pd.Field(0.5, title="Eta", description="Halfway point in projection function.") + eta: float = pd.Field( + 0.5, ge=0.0, le=1.0, title="Eta", description="Halfway point in projection function." + ) strict_binarize: bool = pd.Field( False, @@ -67,8 +70,9 @@ class FilterProject(InvdesBaseModel): def evaluate(self, spatial_data: anp.ndarray, design_region_dl: float) -> anp.ndarray: """Evaluate this transformation on spatial data, given some grid size in the region.""" - filter_size = get_kernel_size_px(self.radius, design_region_dl) - filt_proj = make_filter_and_project(filter_size, beta=self.beta, eta=self.eta) + filt_proj = make_filter_and_project( + self.radius, design_region_dl, beta=self.beta, eta=self.eta + ) data_projected = filt_proj(spatial_data) if self.strict_binarize: diff --git a/tidy3d/plugins/invdes/validators.py b/tidy3d/plugins/invdes/validators.py index 957b37d05..2e4ff499e 100644 --- a/tidy3d/plugins/invdes/validators.py +++ b/tidy3d/plugins/invdes/validators.py @@ -5,6 +5,7 @@ import pydantic.v1 as pd import tidy3d as td +from tidy3d.components.base import skip_if_fields_missing # warn if pixel size is > PIXEL_SIZE_WARNING_THRESHOLD * (minimum wavelength in material) PIXEL_SIZE_WARNING_THRESHOLD = 0.1 @@ -45,9 +46,9 @@ def check_pixel_size_sim(sim: td.Simulation, pixel_size: float, index: int = Non ) @pd.root_validator(allow_reuse=True) + @skip_if_fields_missing(["design_region"], root=True) def _check_pixel_size(cls, values): """Make sure region pixel_size isn't too large compared to sim's wavelength in material.""" - sim = values.get(sim_field_name) region = values.get("design_region") pixel_size = region.pixel_size diff --git a/tidy3d/plugins/mode/mode_solver.py b/tidy3d/plugins/mode/mode_solver.py index f3bc2fa2e..26c4fdea5 100644 --- a/tidy3d/plugins/mode/mode_solver.py +++ b/tidy3d/plugins/mode/mode_solver.py @@ -47,8 +47,12 @@ PlotScale, Symmetry, ) -from ...components.validators import validate_freqs_min, validate_freqs_not_empty -from ...components.viz import plot_params_pml +from ...components.validators import ( + validate_freqs_min, + validate_freqs_not_empty, + validate_mode_plane_radius, +) +from ...components.viz import make_ax, plot_params_pml from ...constants import C_0 from ...exceptions import SetupError, ValidationError from ...log import log @@ -174,6 +178,11 @@ def plane_in_sim_bounds(cls, val, values): raise SetupError("'ModeSolver.plane' must intersect 'ModeSolver.simulation'.") return val + def _post_init_validators(self) -> None: + validate_mode_plane_radius( + mode_spec=self.mode_spec, plane=self.plane, msg_prefix="Mode solver" + ) + @cached_property def normal_axis(self) -> Axis: """Axis normal to the mode plane.""" @@ -348,7 +357,20 @@ def data_raw(self) -> ModeSolverData: def _data_on_yee_grid(self) -> ModeSolverData: """Solve for all modes, and construct data with fields on the Yee grid.""" - solver = self.reduced_simulation_copy + + # we try to do reduced simulation copy for efficiency + # it should never fail -- if it does, this is likely due to an oversight + # in the Simulation.subsection method. but falling back to non-reduced + # simulation prevents unneeded errors in this case + try: + solver = self.reduced_simulation_copy + except Exception as e: + solver = self + log.warning( + "Mode solver reduced_simulation_copy failed. " + "Falling back to non-reduced simulation, which may be slower. " + f"Exception: {str(e)}" + ) _, _solver_coords = solver.plane.pop_axis( solver._solver_grid.boundaries.to_list, axis=solver.normal_axis @@ -985,6 +1007,8 @@ def to_source( source_time: SourceTime, direction: Direction = None, mode_index: pydantic.NonNegativeInt = 0, + num_freqs: pydantic.PositiveInt = 1, + **kwargs, ) -> ModeSource: """Creates :class:`.ModeSource` from a :class:`ModeSolver` instance plus additional specifications. @@ -1016,6 +1040,8 @@ def to_source( mode_spec=self.mode_spec, mode_index=mode_index, direction=direction, + num_freqs=num_freqs, + **kwargs, ) def to_monitor(self, freqs: List[float] = None, name: str = None) -> ModeMonitor: @@ -1264,7 +1290,7 @@ def plot( # Get the mode plane normal axis, center, and limits. a_center, h_lim, v_lim, _ = self._center_and_lims() - return self.simulation.plot( + ax = self.simulation.plot( x=a_center[0], y=a_center[1], z=a_center[2], @@ -1277,6 +1303,8 @@ def plot( **patch_kwargs, ) + return self.plot_pml(ax=ax) + def plot_eps( self, freq: float = None, @@ -1436,18 +1464,11 @@ def plot_pml( # Get the mode plane normal axis, center, and limits. a_center, h_lim, v_lim, t_axes = self._center_and_lims() - # Plot the mode plane is ax=None. + # Create ax if ax=None. if not ax: - ax = self.simulation.plot( - x=a_center[0], - y=a_center[1], - z=a_center[2], - hlim=h_lim, - vlim=v_lim, - source_alpha=0, - monitor_alpha=0, - ax=ax, - ) + ax = make_ax() + ax.set_xlim(h_lim) + ax.set_ylim(v_lim) # Mode plane grid. plane_grid = self.grid_snapped.centers.to_list @@ -1517,10 +1538,10 @@ def _center_and_lims(self) -> Tuple[List, List, List, List]: _, (h_min_s, v_min_s) = Box.pop_axis(self.simulation.bounds[0], axis=n_axis) _, (h_max_s, v_max_s) = Box.pop_axis(self.simulation.bounds[1], axis=n_axis) - h_min = a_center[n_axis] - self.plane.size[t_axes[0]] / 2 - h_max = a_center[n_axis] + self.plane.size[t_axes[0]] / 2 - v_min = a_center[n_axis] - self.plane.size[t_axes[1]] / 2 - v_max = a_center[n_axis] + self.plane.size[t_axes[1]] / 2 + h_min = self.plane.center[t_axes[0]] - self.plane.size[t_axes[0]] / 2 + h_max = self.plane.center[t_axes[0]] + self.plane.size[t_axes[0]] / 2 + v_min = self.plane.center[t_axes[1]] - self.plane.size[t_axes[1]] / 2 + v_max = self.plane.center[t_axes[1]] + self.plane.size[t_axes[1]] / 2 h_lim = [ h_min if abs(h_min) < abs(h_min_s) else h_min_s, @@ -1548,6 +1569,7 @@ def _validate_modes_size(self): ) def validate_pre_upload(self, source_required: bool = True): + """Validate the fully initialized mode solver is ok for upload to our servers.""" self._validate_modes_size() @cached_property diff --git a/tidy3d/plugins/pytorch/README.md b/tidy3d/plugins/pytorch/README.md new file mode 100644 index 000000000..74c6d9db2 --- /dev/null +++ b/tidy3d/plugins/pytorch/README.md @@ -0,0 +1,51 @@ +# Autograd to PyTorch Wrapper for Tidy3D + +This wrapper allows you to seamlessly convert autograd functions to PyTorch functions, enabling the use of Tidy3D simulations within PyTorch. + +## Examples + +### Basic Usage + +This module can be used to convert any autograd function to a PyTorch function: + +```python +import torch +import autograd.numpy as anp + +from tidy3d.plugins.pytorch.wrapper import to_torch + +@to_torch +def my_function(x): + return anp.sum(anp.sin(x)**2) + +x = torch.rand(10, requires_grad=True) +y = my_function(x) +y.backward() # backward works as expected, even though the function is defined in terms of autograd.numpy +print(x.grad) # gradients are available in the input tensor +``` + +### Usage with Tidy3D + +The `to_torch` wrapper can be used to convert an objective function that depends on Tidy3D simulations to a PyTorch function: + +```python +import torch +import autograd.numpy as anp + +import tidy3d as td +import tidy3d.web as web + +from tidy3d.plugins.pytorch.wrapper import to_torch + +@to_torch +def tidy3d_objective(params): + sim = make_sim(params) + sim_data = web.run(sim, task_name="pytorch_example") + flux = sim_data["flux"].flux.values + return anp.sum(flux) + +params = torch.rand(10, requires_grad=True) +y = tidy3d_objective(params) +y.backward() +print(params.grad) +``` diff --git a/tidy3d/plugins/pytorch/__init__.py b/tidy3d/plugins/pytorch/__init__.py new file mode 100644 index 000000000..e3bf690e2 --- /dev/null +++ b/tidy3d/plugins/pytorch/__init__.py @@ -0,0 +1,3 @@ +from .wrapper import to_torch + +__all__ = ["to_torch"] diff --git a/tidy3d/plugins/pytorch/wrapper.py b/tidy3d/plugins/pytorch/wrapper.py new file mode 100644 index 000000000..ffdabf0bd --- /dev/null +++ b/tidy3d/plugins/pytorch/wrapper.py @@ -0,0 +1,94 @@ +import inspect + +import torch +from autograd import make_vjp +from autograd.extend import vspace + + +def to_torch(fun): + """ + Converts an autograd function to a PyTorch function. + + Parameters + ---------- + fun : callable + The autograd function to be converted. + + Returns + ------- + callable + A PyTorch function that can be used with PyTorch tensors and supports + autograd differentiation. + + Examples + -------- + >>> import autograd.numpy as anp + >>> import torch + >>> from tidy3d.plugins.pytorch.wrapper import to_torch + >>> + >>> @to_torch + ... def f(x): + ... return anp.sum(anp.sin(x)) + >>> + >>> x = torch.tensor([0.0, anp.pi / 2, anp.pi], requires_grad=True) + >>> val = f(x) + >>> val.backward() + >>> torch.allclose(x.grad, torch.cos(x)) + True + """ + sig = inspect.signature(fun) + + class _Wrapper(torch.autograd.Function): + """A `torch.autograd.Function` wrapper for the autograd function `fun`. + + See Also + -------- + `PyTorch Autograd Function Documentation `_ + """ + + @staticmethod + def forward(ctx, *args): + numpy_args = [] + grad_argnums = [] + + # assume that all tensors are on the same device (cpu by default) + device = torch.device("cpu") + + for idx, arg in enumerate(args): + if torch.is_tensor(arg): + numpy_args.append(arg.detach().cpu().numpy()) + device = arg.device + if arg.requires_grad: + grad_argnums.append(idx) + else: + numpy_args.append(arg) + + # note: can't support differentiating w.r.t. keyword-only arguments because + # autograd's unary_to_nary decorator passes all function arguments as positional + _vjp = make_vjp(fun, argnum=grad_argnums) + vjp, ans = _vjp(*numpy_args) + + ctx.vjp = vjp + ctx.device = device + ctx.num_args = len(args) + ctx.grad_argnums = grad_argnums + + return torch.as_tensor(ans, device=device) + + @staticmethod + def backward(ctx, grad_output): + _grads = ctx.vjp(vspace(grad_output.detach().cpu().numpy()).ones()) + grads = [None] * ctx.num_args + for idx, grad in zip(ctx.grad_argnums, _grads): + grads[idx] = torch.as_tensor(grad, device=ctx.device) * grad_output + return tuple(grads) + + def apply(*args, **kwargs): + # we bind the full function signature including defaults so that we can pass + # all values as positional since torch.autograd.Function.apply only accepts + # positional arguments + bound_args = sig.bind(*args, **kwargs) + bound_args.apply_defaults() + return _Wrapper.apply(*bound_args.arguments.values()) + + return apply diff --git a/tidy3d/plugins/smatrix/component_modelers/modal.py b/tidy3d/plugins/smatrix/component_modelers/modal.py index 6894148fc..fe9a32f82 100644 --- a/tidy3d/plugins/smatrix/component_modelers/modal.py +++ b/tidy3d/plugins/smatrix/component_modelers/modal.py @@ -179,7 +179,9 @@ def to_monitor(self, port: Port) -> ModeMonitor: name=port.name, ) - def to_source(self, port: Port, mode_index: int) -> List[ModeSource]: + def to_source( + self, port: Port, mode_index: int, num_freqs: int = 1, **kwargs + ) -> List[ModeSource]: """Creates a list of mode sources from a given port.""" freq0 = np.mean(self.freqs) fdiff = max(self.freqs) - min(self.freqs) @@ -192,6 +194,8 @@ def to_source(self, port: Port, mode_index: int) -> List[ModeSource]: mode_index=mode_index, direction=port.direction, name=port.name, + num_freqs=num_freqs, + **kwargs, ) def shift_port(self, port: Port) -> Port: diff --git a/tidy3d/web/api/autograd/autograd.py b/tidy3d/web/api/autograd/autograd.py index 6a271306d..10a833c48 100644 --- a/tidy3d/web/api/autograd/autograd.py +++ b/tidy3d/web/api/autograd/autograd.py @@ -1,8 +1,10 @@ # autograd wrapper for web functions +import os import tempfile import typing from collections import defaultdict +from os.path import basename, dirname, join import numpy as np from autograd.builtins import dict as dict_ag @@ -11,15 +13,14 @@ import tidy3d as td from tidy3d.components.autograd import AutogradFieldMap, get_static from tidy3d.components.autograd.derivative_utils import DerivativeInfo -from tidy3d.components.data.sim_data import AdjointSourceInfo from ...core.s3utils import download_file, upload_file from ..asynchronous import DEFAULT_DATA_DIR from ..asynchronous import run_async as run_async_webapi -from ..container import Batch, BatchData, Job +from ..container import DEFAULT_DATA_PATH, Batch, BatchData, Job from ..tidy3d_stub import SimulationDataType, SimulationType from ..webapi import run as run_webapi -from .utils import FieldMap, TracerKeys, get_derivative_maps +from .utils import E_to_D, FieldMap, TracerKeys, get_derivative_maps # keys for data into auxiliary dictionary AUX_KEY_SIM_DATA_ORIGINAL = "sim_data" @@ -29,7 +30,6 @@ # server-side auxiliary files to upload/download SIM_VJP_FILE = "output/autograd_sim_vjp.hdf5" SIM_FIELDS_KEYS_FILE = "autograd_sim_fields_keys.hdf5" -ADJOINT_SOURCE_INFO_FILE = "autograd_adjoint_source_info_file.hdf5" ISSUE_URL = ( "https://github.com/flexcompute/tidy3d/issues/new?" @@ -42,6 +42,10 @@ # default value for whether to do local gradient calculation (True) or server side (False) LOCAL_GRADIENT = False +# if True, will plot the adjoint fields on the plane provided. used for debugging only +_INSPECT_ADJOINT_FIELDS = False +_INSPECT_ADJOINT_PLANE = td.Box(center=(0, 0, 0), size=(td.inf, td.inf, 0)) + def is_valid_for_autograd(simulation: td.Simulation) -> bool: """Check whether a supplied simulation can use autograd run.""" @@ -284,7 +288,10 @@ def run_async( def _run( - simulation: td.Simulation, task_name: str, local_gradient: bool = LOCAL_GRADIENT, **run_kwargs + simulation: td.Simulation, + task_name: str, + local_gradient: bool = LOCAL_GRADIENT, + **run_kwargs, ) -> td.SimulationData: """User-facing ``web.run`` function, compatible with ``autograd`` differentiation.""" @@ -318,7 +325,9 @@ def _run( def _run_async( - simulations: dict[str, td.Simulation], local_gradient: bool = LOCAL_GRADIENT, **run_async_kwargs + simulations: dict[str, td.Simulation], + local_gradient: bool = LOCAL_GRADIENT, + **run_async_kwargs, ) -> dict[str, td.SimulationData]: """User-facing ``web.run_async`` function, compatible with ``autograd`` differentiation.""" @@ -388,12 +397,14 @@ def _run_primitive( td.log.info("running primitive '_run_primitive()'") + # compute the combined simulation for both local and remote, so we can validate it + sim_combined = setup_fwd( + sim_fields=sim_fields, + sim_original=sim_original, + local_gradient=local_gradient, + ) + if local_gradient: - sim_combined = setup_fwd( - sim_fields=sim_fields, - sim_original=sim_original, - local_gradient=local_gradient, - ) sim_data_combined, _ = _run_tidy3d(sim_combined, task_name=task_name, **run_kwargs) field_map = postprocess_fwd( @@ -403,6 +414,7 @@ def _run_primitive( ) else: + sim_combined.validate_pre_upload() sim_original = sim_original.updated_copy(simulation_type="autograd_fwd", deep=False) run_kwargs["simulation_type"] = "autograd_fwd" run_kwargs["sim_fields_keys"] = list(sim_fields.keys()) @@ -523,30 +535,21 @@ def postprocess_fwd( def upload_sim_fields_keys(sim_fields_keys: list[tuple], task_id: str, verbose: bool = False): """Function to grab the VJP result for the simulation fields from the adjoint task ID.""" - data_file = tempfile.NamedTemporaryFile(suffix=".hdf5") - data_file.close() - TracerKeys(keys=sim_fields_keys).to_file(data_file.name) - upload_file( - task_id, - data_file.name, - SIM_FIELDS_KEYS_FILE, - verbose=verbose, - ) - - -def upload_adjoint_source_info( - adjoint_source_info: AdjointSourceInfo, task_id: str, verbose: bool = False -) -> None: - """Upload the adjoint source information for the adjoint run.""" - data_file = tempfile.NamedTemporaryFile(suffix=".hdf5") - data_file.close() - adjoint_source_info.to_file(data_file.name) - upload_file( - task_id, - data_file.name, - ADJOINT_SOURCE_INFO_FILE, - verbose=verbose, - ) + handle, fname = tempfile.mkstemp(suffix=".hdf5") + os.close(handle) + try: + TracerKeys(keys=sim_fields_keys).to_file(fname) + upload_file( + task_id, + fname, + SIM_FIELDS_KEYS_FILE, + verbose=verbose, + ) + except Exception as e: + td.log.error(f"Error occurred while uploading simulation fields keys: {e}") + raise e + finally: + os.unlink(fname) """ VJP maker for ADJ pass.""" @@ -554,10 +557,16 @@ def upload_adjoint_source_info( def get_vjp_traced_fields(task_id_adj: str, verbose: bool) -> AutogradFieldMap: """Function to grab the VJP result for the simulation fields from the adjoint task ID.""" - data_file = tempfile.NamedTemporaryFile(suffix=".hdf5") - data_file.close() - download_file(task_id_adj, SIM_VJP_FILE, to_file=data_file.name, verbose=verbose) - field_map = FieldMap.from_file(data_file.name) + handle, fname = tempfile.mkstemp(suffix=".hdf5") + os.close(handle) + try: + download_file(task_id_adj, SIM_VJP_FILE, to_file=fname, verbose=verbose) + field_map = FieldMap.from_file(fname) + except Exception as e: + td.log.error(f"Error occurred while getting VJP traced fields: {e}") + raise e + finally: + os.unlink(fname) return field_map.to_autograd_field_map @@ -585,12 +594,21 @@ def _run_bwd( def vjp(data_fields_vjp: AutogradFieldMap) -> AutogradFieldMap: """dJ/d{sim.traced_fields()} as a function of Function of dJ/d{data.traced_fields()}""" - sim_adj, adjoint_source_info = setup_adj( + sim_adj = setup_adj( data_fields_vjp=data_fields_vjp, sim_data_orig=sim_data_orig, sim_fields_keys=sim_fields_keys, ) + if sim_adj is None: + td.log.warning( + f"Adjoint simulation for task '{task_name}' contains no sources. " + "This can occur if the objective function does not depend on the " + "simulation's output. If this is unexpected, please review your " + "setup or contact customer support for assistance." + ) + return {k: 0 * v for k, v in sim_fields_original.items()} + # run adjoint simulation task_name_adj = str(task_name) + "_adjoint" @@ -602,7 +620,6 @@ def vjp(data_fields_vjp: AutogradFieldMap) -> AutogradFieldMap: sim_data_orig=sim_data_orig, sim_data_fwd=sim_data_fwd, sim_fields_keys=sim_fields_keys, - adjoint_source_info=adjoint_source_info, ) else: @@ -614,7 +631,6 @@ def vjp(data_fields_vjp: AutogradFieldMap) -> AutogradFieldMap: vjp_traced_fields = _run_tidy3d_bwd( sim_adj, task_name=task_name_adj, - adjoint_source_info=adjoint_source_info, **run_kwargs, ) @@ -653,66 +669,78 @@ def _run_async_bwd( def vjp(data_fields_dict_vjp: dict[str, AutogradFieldMap]) -> dict[str, AutogradFieldMap]: """dJ/d{sim.traced_fields()} as a function of Function of dJ/d{data.traced_fields()}""" - task_names_adj = {task_name + "_adjoint" for task_name in task_names} + task_names_adj = [task_name + "_adjoint" for task_name in task_names] sims_adj = {} - adjoint_source_info_dict = {} + sim_fields_vjp_dict = {} for task_name, task_name_adj in zip(task_names, task_names_adj): data_fields_vjp = data_fields_dict_vjp[task_name] sim_data_orig = sim_data_orig_dict[task_name] sim_fields_keys = sim_fields_keys_dict[task_name] - sim_adj, adjoint_source_info = setup_adj( + sim_adj = setup_adj( data_fields_vjp=data_fields_vjp, sim_data_orig=sim_data_orig, sim_fields_keys=sim_fields_keys, ) + + if sim_adj is None: + td.log.debug(f"Adjoint simulation for task '{task_name}' contains no sources. ") + sim_fields_vjp_dict[task_name] = { + k: 0 * v for k, v in sim_fields_original_dict[task_name].items() + } sims_adj[task_name_adj] = sim_adj - adjoint_source_info_dict[task_name_adj] = adjoint_source_info - # TODO: handle case where no adjoint sources? + + sims_to_run = {k: v for k, v in sims_adj.items() if v is not None} + + if not sims_to_run: + td.log.warning( + "No simulation in batch contains adjoint sources and thus all gradients are zero. " + "This likely indicates an issue with your setup, consider double-checking or contact support." + ) + return sim_fields_vjp_dict + + task_names_adj = list(sims_to_run.keys()) + task_names_fwd = [name.rstrip("_adjoint") for name in task_names_adj] if local_gradient: # run adjoint simulation - batch_data_adj, _ = _run_async_tidy3d(sims_adj, **run_async_kwargs) + batch_data_adj, _ = _run_async_tidy3d(sims_to_run, **run_async_kwargs) - sim_fields_vjp_dict = {} - for task_name, task_name_adj in zip(task_names, task_names_adj): - sim_data_adj = batch_data_adj[task_name_adj] + for task_name, task_name_adj in zip(task_names_fwd, task_names_adj): sim_data_orig = sim_data_orig_dict[task_name] sim_data_fwd = sim_data_fwd_dict[task_name] sim_fields_keys = sim_fields_keys_dict[task_name] - adjoint_source_info = adjoint_source_info_dict[task_name_adj] + + sim_data_adj = batch_data_adj.get(task_name_adj) sim_fields_vjp = postprocess_adj( sim_data_adj=sim_data_adj, sim_data_orig=sim_data_orig, sim_data_fwd=sim_data_fwd, sim_fields_keys=sim_fields_keys, - adjoint_source_info=adjoint_source_info, ) - sim_fields_vjp_dict[task_name] = sim_fields_vjp + sim_fields_vjp_dict[task_name] = sim_fields_vjp else: parent_tasks = {} - for task_name_fwd, task_name_adj in zip(task_names, task_names_adj): + for task_name_fwd, task_name_adj in zip(task_names_fwd, task_names_adj): task_id_fwd = aux_data_dict[task_name_fwd][AUX_KEY_FWD_TASK_ID] parent_tasks[task_name_adj] = [task_id_fwd] run_async_kwargs["parent_tasks"] = parent_tasks run_async_kwargs["simulation_type"] = "autograd_bwd" - sims_adj = { + simulations = { task_name: sim.updated_copy(simulation_type="autograd_bwd", deep=False) - for task_name, sim in sims_adj.items() + for task_name, sim in sims_to_run.items() } sim_fields_vjp_dict_adj_keys = _run_async_tidy3d_bwd( - simulations=sims_adj, - adjoint_source_info_dict=adjoint_source_info_dict, + simulations=simulations, **run_async_kwargs, ) # swap adjoint task_names for original task_names - sim_fields_vjp_dict = {} - for task_name_fwd, task_name_adj in zip(task_names, task_names_adj): + for task_name_fwd, task_name_adj in zip(task_names_fwd, task_names_adj): sim_fields_vjp_dict[task_name_fwd] = sim_fields_vjp_dict_adj_keys[task_name_adj] return sim_fields_vjp_dict @@ -724,15 +752,13 @@ def setup_adj( data_fields_vjp: AutogradFieldMap, sim_data_orig: td.SimulationData, sim_fields_keys: list[tuple], -) -> tuple[td.Simulation, AdjointSourceInfo]: +) -> typing.Optional[td.Simulation]: """Construct an adjoint simulation from a set of data_fields for the VJP.""" td.log.info("Running custom vjp (adjoint) pipeline.") # immediately filter out any data_vjps with all 0's in the data - data_fields_vjp = { - key: get_static(value) for key, value in data_fields_vjp.items() if not np.all(value == 0.0) - } + data_fields_vjp = {key: get_static(value) for key, value in data_fields_vjp.items()} # insert the raw VJP data into the .data of the original SimulationData sim_data_vjp = sim_data_orig.insert_traced_fields(field_mapping=data_fields_vjp) @@ -745,13 +771,39 @@ def setup_adj( num_monitors: ] - sim_adj, adjoint_source_info = sim_data_vjp.make_adjoint_sim( - data_vjp_paths=data_vjp_paths, adjoint_monitors=adjoint_monitors + sim_adj = sim_data_vjp.make_adjoint_sim( + data_vjp_paths=data_vjp_paths, + adjoint_monitors=adjoint_monitors, ) + if sim_adj is None: + return sim_adj + + if _INSPECT_ADJOINT_FIELDS: + adj_fld_mnt = td.FieldMonitor( + center=_INSPECT_ADJOINT_PLANE.center, + size=_INSPECT_ADJOINT_PLANE.size, + freqs=adjoint_monitors[0].freqs, + name="adjoint_fields", + ) + + import matplotlib.pylab as plt + + import tidy3d.web as web + + sim_data_new = web.run( + sim_adj.updated_copy(monitors=[adj_fld_mnt]), + task_name="adjoint_field_viz", + verbose=False, + ) + _, (ax1, ax2, ax3) = plt.subplots(1, 3, tight_layout=True, figsize=(10, 4)) + sim_data_new.plot_field("adjoint_fields", "Ex", "re", ax=ax1) + sim_data_new.plot_field("adjoint_fields", "Ey", "re", ax=ax2) + sim_data_new.plot_field("adjoint_fields", "Ez", "re", ax=ax3) + plt.show() td.log.info(f"Adjoint simulation created with {len(sim_adj.sources)} sources.") - return sim_adj, adjoint_source_info + return sim_adj def postprocess_adj( @@ -759,7 +811,6 @@ def postprocess_adj( sim_data_orig: td.SimulationData, sim_data_fwd: td.SimulationData, sim_fields_keys: list[tuple], - adjoint_source_info: AdjointSourceInfo, ) -> AutogradFieldMap: """Postprocess some data from the adjoint simulation into the VJP for the original sim flds.""" @@ -773,26 +824,29 @@ def postprocess_adj( sim_fields_vjp = {} for structure_index, structure_paths in sim_vjp_map.items(): # grab the forward and adjoint data - fld_fwd = sim_data_fwd.get_adjoint_data(structure_index, data_type="fld") + E_fwd = sim_data_fwd.get_adjoint_data(structure_index, data_type="fld") eps_fwd = sim_data_fwd.get_adjoint_data(structure_index, data_type="eps") - fld_adj = sim_data_adj.get_adjoint_data(structure_index, data_type="fld") + E_adj = sim_data_adj.get_adjoint_data(structure_index, data_type="fld") eps_adj = sim_data_adj.get_adjoint_data(structure_index, data_type="eps") # post normalize the adjoint fields if a single, broadband source fwd_flds_normed = {} - for key, val in fld_adj.field_components.items(): - fwd_flds_normed[key] = val * adjoint_source_info.post_norm + for key, val in E_adj.field_components.items(): + fwd_flds_normed[key] = val * sim_data_adj.simulation.post_norm - fld_adj = fld_adj.updated_copy(**fwd_flds_normed) + E_adj = E_adj.updated_copy(**fwd_flds_normed) # maps of the E_fwd * E_adj and D_fwd * D_adj, each as as td.FieldData & 'Ex', 'Ey', 'Ez' der_maps = get_derivative_maps( - fld_fwd=fld_fwd, eps_fwd=eps_fwd, fld_adj=fld_adj, eps_adj=eps_adj + fld_fwd=E_fwd, eps_fwd=eps_fwd, fld_adj=E_adj, eps_adj=eps_adj ) E_der_map = der_maps["E"] D_der_map = der_maps["D"] + D_fwd = E_to_D(E_fwd, eps_fwd) + D_adj = E_to_D(E_adj, eps_fwd) + # compute the derivatives for this structure structure = sim_data_fwd.simulation.structures[structure_index] @@ -805,13 +859,17 @@ def postprocess_adj( eps_out = np.mean(sim_data_orig.simulation.medium.eps_model(freq_adj)) # manually override simulation medium as the background structure - if structure.background_permittivity is not None: - eps_out = structure.background_permittivity + if structure.background_medium is not None: + eps_out = structure.background_medium.eps_model(freq_adj) derivative_info = DerivativeInfo( paths=structure_paths, E_der_map=E_der_map.field_components, D_der_map=D_der_map.field_components, + E_fwd=E_fwd.field_components, + E_adj=E_adj.field_components, + D_fwd=D_fwd.field_components, + D_adj=D_adj.field_components, eps_data=eps_fwd.field_components, eps_in=eps_in, eps_out=eps_out, @@ -847,7 +905,7 @@ def parse_run_kwargs(**run_kwargs): def _run_tidy3d( simulation: td.Simulation, task_name: str, **run_kwargs -) -> (td.SimulationData, str): +) -> tuple[td.SimulationData, str]: """Run a simulation without any tracers using regular web.run().""" job_init_kwargs = parse_run_kwargs(**run_kwargs) job = Job(simulation=simulation, task_name=task_name, **job_init_kwargs) @@ -855,18 +913,18 @@ def _run_tidy3d( if job.simulation_type == "autograd_fwd": verbose = run_kwargs.get("verbose", False) upload_sim_fields_keys(run_kwargs["sim_fields_keys"], task_id=job.task_id, verbose=verbose) - data = job.run() + path = run_kwargs.get("path", DEFAULT_DATA_PATH) + if task_name.endswith("_adjoint"): + path_parts = basename(path).split(".") + path = join(dirname(path), path_parts[0] + "_adjoint." + ".".join(path_parts[1:])) + data = job.run(path) return data, job.task_id -def _run_tidy3d_bwd( - simulation: td.Simulation, task_name: str, adjoint_source_info: AdjointSourceInfo, **run_kwargs -) -> AutogradFieldMap: +def _run_tidy3d_bwd(simulation: td.Simulation, task_name: str, **run_kwargs) -> AutogradFieldMap: """Run a simulation without any tracers using regular web.run().""" job_init_kwargs = parse_run_kwargs(**run_kwargs) job = Job(simulation=simulation, task_name=task_name, **job_init_kwargs) - verbose = run_kwargs.get("verbose", False) - upload_adjoint_source_info(adjoint_source_info, task_id=job.task_id, verbose=verbose) td.log.info(f"running {job.simulation_type} simulation with '_run_tidy3d_bwd()'") job.start() job.monitor() @@ -908,7 +966,6 @@ def _run_async_tidy3d( def _run_async_tidy3d_bwd( simulations: dict[str, td.Simulation], - adjoint_source_info_dict: dict[str, AdjointSourceInfo], **run_kwargs, ) -> dict[str, AutogradFieldMap]: """Run a simulation without any tracers using regular web.run().""" @@ -918,11 +975,6 @@ def _run_async_tidy3d_bwd( batch = Batch(simulations=simulations, **batch_init_kwargs) td.log.info(f"running {batch.simulation_type} simulation with '_run_tidy3d_bwd()'") - task_ids = {key: job.task_id for key, job in batch.jobs.items()} - for task_name, adjoint_source_info in adjoint_source_info_dict.items(): - task_id = task_ids[task_name] - upload_adjoint_source_info(adjoint_source_info, task_id=task_id, verbose=batch.verbose) - batch.start() batch.monitor() diff --git a/tidy3d/web/api/mode.py b/tidy3d/web/api/mode.py index f08303b52..ba1a9f9d8 100644 --- a/tidy3d/web/api/mode.py +++ b/tidy3d/web/api/mode.py @@ -231,6 +231,10 @@ def handle_mode_solver(index, progress, pbar): if verbose: console.log(f"[cyan]Running a batch of [deep_pink4]{num_mode_solvers} mode solvers.\n") + + # Create the common folder before running the parallel computation + _ = Folder.create(folder_name=folder_name) + with Progress(console=console) as progress: pbar = progress.add_task("Status:", total=num_mode_solvers) results = Parallel(n_jobs=max_workers, backend="threading")( diff --git a/tidy3d/web/api/webapi.py b/tidy3d/web/api/webapi.py index b7dd5291b..4c00f559d 100644 --- a/tidy3d/web/api/webapi.py +++ b/tidy3d/web/api/webapi.py @@ -329,7 +329,7 @@ def get_status(task_id) -> str: return "success" if status == "error": raise WebError( - f"Error running task {task_id}! Use 'web.download_log(task_id)' to " + f"Error running task {task_id}! Use 'web.download_log('{task_id}')' to " "download and examine the solver log, and/or contact customer support for help." ) return status @@ -645,7 +645,7 @@ def load( After the simulation is complete, you can load the results into a :class:`.SimulationData` object by its ``task_id`` using: - .. code-block:: python py + .. code-block:: python sim_data = web.load(task_id, path="outt/sim.hdf5", verbose=verbose) diff --git a/tidy3d/web/cli/develop/documentation.py b/tidy3d/web/cli/develop/documentation.py index 69fb1d971..fc3fc9d5e 100644 --- a/tidy3d/web/cli/develop/documentation.py +++ b/tidy3d/web/cli/develop/documentation.py @@ -174,7 +174,9 @@ def build_documentation(args=None): Additional arguments for the documentation build process. """ # Runs the documentation build from the poetry environment - echo_and_check_subprocess(["poetry", "run", "python", "-m", "sphinx", "docs/", "_docs/"]) + echo_and_check_subprocess( + ["poetry", "run", "python", "-m", "sphinx", "-j", "auto", "docs/", "_docs/"] + ) return 0 diff --git a/tidy3d/web/core/s3utils.py b/tidy3d/web/core/s3utils.py index 2c2f75a40..9e04c058a 100644 --- a/tidy3d/web/core/s3utils.py +++ b/tidy3d/web/core/s3utils.py @@ -335,6 +335,7 @@ def _download(_callback: Callable) -> None: Filename=str(to_file), Key=token.get_s3_key(), Callback=_callback, + Config=_s3_config, ) if progress_callback is not None: diff --git a/tidy3d/web/core/task_core.py b/tidy3d/web/core/task_core.py index f7d28508d..7d53f04f2 100644 --- a/tidy3d/web/core/task_core.py +++ b/tidy3d/web/core/task_core.py @@ -225,6 +225,11 @@ def create( :class:`SimulationTask` object containing info about status, size, credits of task and others. """ + + # handle backwards compatibility, "tidy3d" is the default simulation_type + if simulation_type is None: + simulation_type = "tidy3d" + folder = Folder.get(folder_name, create=True) resp = http.post( f"tidy3d/projects/{folder.folder_id}/tasks", diff --git a/tidy3d/web/core/task_info.py b/tidy3d/web/core/task_info.py index d3196be76..f33093b39 100644 --- a/tidy3d/web/core/task_info.py +++ b/tidy3d/web/core/task_info.py @@ -1,4 +1,4 @@ -"""Defnes information about a task""" +"""Defines information about a task""" from abc import ABC from datetime import datetime @@ -12,77 +12,156 @@ class TaskStatus(Enum): """The statuses that the task can be in.""" INIT = "initialized" + """The task has been initialized.""" + QUEUE = "queued" + """The task is in the queue.""" + PRE = "preprocessing" + """The task is in the preprocessing stage.""" + RUN = "running" + """The task is running.""" + POST = "postprocessing" + """The task is in the postprocessing stage.""" + SUCCESS = "success" + """The task has completed successfully.""" + ERROR = "error" + """The task has completed with an error.""" class TaskBase(pydantic.BaseModel, ABC): - """Base config for all task objects.""" + """Base configuration for all task objects.""" class Config: - """configure class""" + """Configuration for TaskBase""" arbitrary_types_allowed = True + """Allow arbitrary types to be used within the model.""" class ChargeType(str, Enum): - """The payment method of task.""" + """The payment method of the task.""" FREE = "free" + """No payment required.""" + PAID = "paid" + """Payment required.""" class TaskBlockInfo(TaskBase): - """The block info that task will be blocked by all three features of DE, - User limit and Insufficient balance""" + """Information about the task's block status. + + This includes details about how the task can be blocked by various features + such as user limits and insufficient balance. + """ chargeType: ChargeType = None + """The type of charge applicable to the task (free or paid).""" + maxFreeCount: int = None + """The maximum number of free tasks allowed.""" + maxGridPoints: int = None + """The maximum number of grid points permitted.""" + maxTimeSteps: int = None + """The maximum number of time steps allowed.""" class TaskInfo(TaskBase): - """General information about task.""" + """General information about a task.""" taskId: str + """Unique identifier for the task.""" + taskName: str = None + """Name of the task.""" + nodeSize: int = None + """Size of the node allocated for the task.""" + completedAt: Optional[datetime] = None + """Timestamp when the task was completed.""" + status: str = None + """Current status of the task.""" + realCost: float = None + """Actual cost incurred by the task.""" + timeSteps: int = None + """Number of time steps involved in the task.""" + solverVersion: str = None + """Version of the solver used for the task.""" + createAt: Optional[datetime] = None + """Timestamp when the task was created.""" + estCostMin: float = None + """Estimated minimum cost for the task.""" + estCostMax: float = None + """Estimated maximum cost for the task.""" + realFlexUnit: float = None + """Actual flexible units used by the task.""" + oriRealFlexUnit: float = None + """Original real flexible units.""" + estFlexUnit: float = None + """Estimated flexible units for the task.""" + estFlexCreditTimeStepping: float = None + """Estimated flexible credits for time stepping.""" + estFlexCreditPostProcess: float = None + """Estimated flexible credits for post-processing.""" + estFlexCreditMode: float = None + """Estimated flexible credits based on the mode.""" + s3Storage: float = None + """Amount of S3 storage used by the task.""" + startSolverTime: Optional[datetime] = None + """Timestamp when the solver started.""" + finishSolverTime: Optional[datetime] = None + """Timestamp when the solver finished.""" + totalSolverTime: int = None + """Total time taken by the solver.""" + callbackUrl: str = None + """Callback URL for task notifications.""" + taskType: str = None + """Type of the task.""" + metadataStatus: str = None + """Status of the metadata for the task.""" + taskBlockInfo: TaskBlockInfo = None + """Blocking information for the task.""" class RunInfo(TaskBase): - """Information about the run.""" + """Information about the run of a task.""" perc_done: pydantic.confloat(ge=0.0, le=100.0) + """Percentage of the task that is completed (0 to 100).""" + field_decay: pydantic.confloat(ge=0.0, le=1.0) + """Field decay from the maximum value (0 to 1).""" def display(self): - """Print some info.""" + """Print some info about the task's progress.""" print(f" - {self.perc_done:.2f} (%) done") print(f" - {self.field_decay:.2e} field decay from max")