diff --git a/.github/workflows/release.yml b/.github/workflows/release.yml index ade56e08c..bb3143169 100644 --- a/.github/workflows/release.yml +++ b/.github/workflows/release.yml @@ -9,6 +9,68 @@ on: - 'v*.*.*' jobs: + test-latest-submodules: + runs-on: ubuntu-latest + steps: + - name: Checkout repository with submodules + uses: actions/checkout@v4 + with: + submodules: 'recursive' + # This fetches only a single branch by default, so additional fetch is needed + fetch-depth: 0 # Optionally, set to 0 to fetch all history for all branches and tags + + - name: Initialize and update submodule + run: | + git submodule update --init --recursive + + - name: Check if submodules are up to date + shell: bash + run: | + NOTEBOOKS_PATH=docs/notebooks + FAQ_PATH=docs/faq + + # Checking out Notebooks submodule with the same branch as the main project/develop + echo "Checking $NOTEBOOKS_PATH for updates..." + cd $NOTEBOOKS_PATH + NOTEBOOKS_CURRENT_COMMIT=$(git rev-parse HEAD) + echo $(git fetch --all --verbose) + echo $(git remote get-url origin) + if git show-ref --verify refs/remotes/origin/develop; then + echo "Branch develop exists." + else + echo "::error::Branch develop does not exist on remote." + exit 1 + fi + NOTEBOOKS_LATEST_COMMIT=$(git rev-parse refs/remotes/origin/develop) + echo "NOTEBOOKS_LATEST_COMMIT: $NOTEBOOKS_LATEST_COMMIT" + echo "NOTEBOOKS_CURRENT_COMMIT: $NOTEBOOKS_CURRENT_COMMIT" + + + cd ../.. + if [ "$NOTEBOOKS_LATEST_COMMIT" != "$NOTEBOOKS_CURRENT_COMMIT" ]; then + echo "::error ::Submodule $NOTEBOOKS_PATH is not up to date with the develop branch. Please update it." + exit 1 + else + echo "Submodule $NOTEBOOKS_PATH is up to date with the develop branch." + fi + + # Checking FAQs only on the develop branch. + echo "Checking $FAQ_PATH for updates..." + cd $FAQ_PATH + FAQ_CURRENT_COMMIT=$(git rev-parse HEAD) + echo $(git fetch --all --verbose) + echo $(git remote get-url origin) + FAQ_LATEST_COMMIT=$(git rev-parse refs/remotes/origin/develop) + echo "FAQ_LATEST_COMMIT: $FAQ_LATEST_COMMIT" + echo "FAQ_CURRENT_COMMIT: $FAQ_CURRENT_COMMIT" + cd ../.. + if [ "$FAQ_LATEST_COMMIT" != "$FAQ_CURRENT_COMMIT" ]; then + echo "::error ::Submodule $FAQ_PATH is not up to date. Please update it." + exit 1 + else + echo "Submodule $FAQ_PATH is up to date." + fi + github-release: runs-on: ubuntu-latest steps: diff --git a/.github/workflows/run_tests.yml b/.github/workflows/run_tests.yml index 77a527eb4..46b399d10 100644 --- a/.github/workflows/run_tests.yml +++ b/.github/workflows/run_tests.yml @@ -27,86 +27,6 @@ jobs: pip install pre-commit pre-commit run # this should be really more agressive - test-latest-submodules: - runs-on: ubuntu-latest - steps: - - name: Checkout repository with submodules - uses: actions/checkout@v4 - with: - submodules: 'recursive' - # This fetches only a single branch by default, so additional fetch is needed - fetch-depth: 0 # Optionally, set to 0 to fetch all history for all branches and tags - - - name: Determine current branch or PR ref - id: get_branch - run: | - if [[ "${{ github.event_name }}" == "pull_request" ]]; then - # Extract the base branch of the PR - BRANCH_NAME="${{ github.event.pull_request.base.ref }}" - echo "BRANCH_NAME=$BRANCH_NAME" >> $GITHUB_ENV - else - # Assume it's a push event, extract the branch name from $GITHUB_REF - BRANCH_NAME=$(echo $GITHUB_REF | sed 's|refs/heads/||') - echo "BRANCH_NAME=$BRANCH_NAME" >> $GITHUB_ENV - fi - # Now echoing the BRANCH_NAME to verify - echo "BRANCH_NAME: $BRANCH_NAME" - shell: bash - env: - GITHUB_REF: ${{ github.ref }} - - - name: Initialize and update submodule - run: | - git submodule update --init --recursive - - - name: Check if submodules are up to date - shell: bash - run: | - NOTEBOOKS_PATH=docs/notebooks - FAQ_PATH=docs/faq - - # Checking out Notebooks submodule with the same branch as the main project - echo "Checking $NOTEBOOKS_PATH for updates..." - cd $NOTEBOOKS_PATH - NOTEBOOKS_CURRENT_COMMIT=$(git rev-parse HEAD) - echo $(git fetch --all --verbose) - echo $(git remote get-url origin) - if git show-ref --verify refs/remotes/origin/$BRANCH_NAME; then - echo "Branch $BRANCH_NAME exists." - else - echo "::error::Branch $BRANCH_NAME does not exist on remote." - exit 1 - fi - NOTEBOOKS_LATEST_COMMIT=$(git rev-parse refs/remotes/origin/${{ env.BRANCH_NAME }}) - echo "NOTEBOOKS_LATEST_COMMIT: $NOTEBOOKS_LATEST_COMMIT" - echo "NOTEBOOKS_CURRENT_COMMIT: $NOTEBOOKS_CURRENT_COMMIT" - - - cd ../.. - if [ "$NOTEBOOKS_LATEST_COMMIT" != "$NOTEBOOKS_CURRENT_COMMIT" ]; then - echo "::error ::Submodule $NOTEBOOKS_PATH is not up to date with the ${{ env.BRANCH_NAME }} branch. Please update it." - exit 1 - else - echo "Submodule $NOTEBOOKS_PATH is up to date with the ${{ env.BRANCH_NAME }} branch." - fi - - # Checking FAQs only on the develop branch. - echo "Checking $FAQ_PATH for updates..." - cd $FAQ_PATH - FAQ_CURRENT_COMMIT=$(git rev-parse HEAD) - echo $(git fetch --all --verbose) - echo $(git remote get-url origin) - FAQ_LATEST_COMMIT=$(git rev-parse refs/remotes/origin/develop) - echo "FAQ_LATEST_COMMIT: $FAQ_LATEST_COMMIT" - echo "FAQ_CURRENT_COMMIT: $FAQ_CURRENT_COMMIT" - cd ../.. - if [ "$FAQ_LATEST_COMMIT" != "$FAQ_CURRENT_COMMIT" ]; then - echo "::error ::Submodule $FAQ_PATH is not up to date. Please update it." - exit 1 - else - echo "Submodule $FAQ_PATH is up to date." - fi - build: name: Python ${{ matrix.python-version }} - ${{ matrix.platform }} runs-on: ${{ matrix.platform }} diff --git a/.github/workflows/sync-to-readthedocs-repo.yaml b/.github/workflows/sync-to-readthedocs-repo.yaml index 941ad4e17..251f1249e 100644 --- a/.github/workflows/sync-to-readthedocs-repo.yaml +++ b/.github/workflows/sync-to-readthedocs-repo.yaml @@ -42,6 +42,16 @@ jobs: token: ${{ secrets.GH_PAT }} ref: ${{ needs.extract_branch_or_tag.outputs.ref_name }} + - name: Push corresponding reference to mirror repo if a branch + if: contains(github.ref, 'refs/heads/') + run: | + git fetch --unshallow origin ${{ needs.extract_branch_or_tag.outputs.ref_name }} + git pull origin ${{ needs.extract_branch_or_tag.outputs.ref_name }} + 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 }} + # Conditional Checkout for Tag - name: Checkout Tag if tag-triggered-sync if: contains(github.ref, 'refs/tags/') @@ -53,11 +63,11 @@ jobs: ref: ${{ needs.extract_branch_or_tag.outputs.ref_name }} fetch-tags: true - - name: Push corresponding reference to mirror repo + + - name: Push corresponding reference to mirror repo if a tag + if: contains(github.ref, 'refs/tags/') run: | - git fetch --unshallow origin ${{ needs.extract_branch_or_tag.outputs.ref_name }} - git pull origin ${{ needs.extract_branch_or_tag.outputs.ref_name }} 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.GITHUB_TOKEN }} \ No newline at end of file diff --git a/.github/workflows/test_daily_latest_submodule.yaml b/.github/workflows/test_daily_latest_submodule.yaml new file mode 100644 index 000000000..e32ee3155 --- /dev/null +++ b/.github/workflows/test_daily_latest_submodule.yaml @@ -0,0 +1,107 @@ +name: "tidy3d-submodule-daily-tests" + +on: + workflow_dispatch: + schedule: + - cron: '0 9 * * *' # Runs at 9am UK-time every day + + +jobs: + test-latest-submodules: + runs-on: ubuntu-latest + steps: + - name: Checkout repository with submodules + uses: actions/checkout@v4 + with: + submodules: 'recursive' + # This fetches only a single branch by default, so additional fetch is needed + fetch-depth: 0 # Optionally, set to 0 to fetch all history for all branches and tags + + - name: Initialize and update submodule + run: | + git submodule update --init --recursive + + - name: Determine the latest pre/2.* branch + id: get_latest_pre_branch + run: | + # Fetch all branches + git fetch --all --quiet + + # List all remote branches for debugging purposes + echo "Available branches:" + git branch -r + + # List branches matching the pre/2.* pattern + BRANCHES=$(git branch -r | grep 'origin/pre/2\.' | sed 's|origin/||' | sort -V) + + # Debugging: Print out the matched branches + echo "Matched branches with pre/2.* pattern:" + echo "$BRANCHES" + + # Identify the latest branch + LATEST_BRANCH=$(echo "$BRANCHES" | tail -n 1) + + # Set the latest branch as an environment variable + echo "LATEST_BRANCH=$LATEST_BRANCH" >> $GITHUB_ENV + echo "Latest pre/2.* branch is: $LATEST_BRANCH" + + - name: Check submodules for multiple branches + shell: bash + run: | + BRANCHES=("develop" $LATEST_BRANCH) # Add your branches here + + for BRANCH in "${BRANCHES[@]}"; do + echo "Analyzing branch: $BRANCH" + + # Fetch all branches and tags + git fetch --all --verbose + + # Checkout the branch + git checkout $BRANCH + + NOTEBOOKS_PATH=docs/notebooks + FAQ_PATH=docs/faq + + # Checking Notebooks submodule + echo "Checking $NOTEBOOKS_PATH for updates..." + cd $NOTEBOOKS_PATH + NOTEBOOKS_CURRENT_COMMIT=$(git rev-parse HEAD) + echo $(git fetch --all --verbose) + echo $(git remote get-url origin) + if git show-ref --verify refs/remotes/origin/$BRANCH; then + echo "Branch $BRANCH exists." + else + echo "::error::Branch $BRANCH does not exist on remote." + exit 1 + fi + NOTEBOOKS_LATEST_COMMIT=$(git rev-parse refs/remotes/origin/${BRANCH}) + echo "NOTEBOOKS_LATEST_COMMIT: $NOTEBOOKS_LATEST_COMMIT" + echo "NOTEBOOKS_CURRENT_COMMIT: $NOTEBOOKS_CURRENT_COMMIT" + + cd ../.. + if [ "$NOTEBOOKS_LATEST_COMMIT" != "$NOTEBOOKS_CURRENT_COMMIT" ]; then + echo "::error::Submodule $NOTEBOOKS_PATH is not up to date with the $BRANCH branch. Please update it." + exit 1 + else + echo "Submodule $NOTEBOOKS_PATH is up to date with the $BRANCH branch." + fi + + # Checking FAQs only on the develop branch + if [[ "$BRANCH" == "develop" ]]; then + echo "Checking $FAQ_PATH for updates..." + cd $FAQ_PATH + FAQ_CURRENT_COMMIT=$(git rev-parse HEAD) + echo $(git fetch --all --verbose) + echo $(git remote get-url origin) + FAQ_LATEST_COMMIT=$(git rev-parse refs/remotes/origin/develop) + echo "FAQ_LATEST_COMMIT: $FAQ_LATEST_COMMIT" + echo "FAQ_CURRENT_COMMIT: $FAQ_CURRENT_COMMIT" + cd ../.. + if [ "$FAQ_LATEST_COMMIT" != "$FAQ_CURRENT_COMMIT" ]; then + echo "::error::Submodule $FAQ_PATH is not up to date. Please update it." + exit 1 + else + echo "Submodule $FAQ_PATH is up to date." + fi + fi + done \ No newline at end of file diff --git a/.github/workflows/test_pr_latest_submodule.yaml b/.github/workflows/test_pr_latest_submodule.yaml new file mode 100644 index 000000000..b40f3f4b1 --- /dev/null +++ b/.github/workflows/test_pr_latest_submodule.yaml @@ -0,0 +1,72 @@ +name: "tidy3d-submodule-PR-tests" + +on: + workflow_dispatch: + push: + branches: [ latest ] + pull_request: + branches: + - latest + +jobs: + test-latest-submodules: + runs-on: ubuntu-latest + steps: + - name: Checkout repository with submodules + uses: actions/checkout@v4 + with: + submodules: 'recursive' + # This fetches only a single branch by default, so additional fetch is needed + fetch-depth: 0 # Optionally, set to 0 to fetch all history for all branches and tags + + - name: Initialize and update submodule + run: | + git submodule update --init --recursive + + - name: Check if submodules are up to date + shell: bash + run: | + NOTEBOOKS_PATH=docs/notebooks + FAQ_PATH=docs/faq + + # Checking out Notebooks submodule with the same branch as the main project + echo "Checking $NOTEBOOKS_PATH for updates..." + cd $NOTEBOOKS_PATH + NOTEBOOKS_CURRENT_COMMIT=$(git rev-parse HEAD) + echo $(git fetch --all --verbose) + echo $(git remote get-url origin) + if git show-ref --verify refs/remotes/origin/develop; then + echo "Branch develop exists." + else + echo "::error::Branch develop does not exist on remote." + exit 1 + fi + NOTEBOOKS_LATEST_COMMIT=$(git rev-parse refs/remotes/origin/develop) + echo "NOTEBOOKS_LATEST_COMMIT: $NOTEBOOKS_LATEST_COMMIT" + echo "NOTEBOOKS_CURRENT_COMMIT: $NOTEBOOKS_CURRENT_COMMIT" + + + cd ../.. + if [ "$NOTEBOOKS_LATEST_COMMIT" != "$NOTEBOOKS_CURRENT_COMMIT" ]; then + echo "::error ::Submodule $NOTEBOOKS_PATH is not up to date with the develop branch. Please update it." + exit 1 + else + echo "Submodule $NOTEBOOKS_PATH is up to date with the develop branch." + fi + + # Checking FAQs only on the develop branch. + echo "Checking $FAQ_PATH for updates..." + cd $FAQ_PATH + FAQ_CURRENT_COMMIT=$(git rev-parse HEAD) + echo $(git fetch --all --verbose) + echo $(git remote get-url origin) + FAQ_LATEST_COMMIT=$(git rev-parse refs/remotes/origin/develop) + echo "FAQ_LATEST_COMMIT: $FAQ_LATEST_COMMIT" + echo "FAQ_CURRENT_COMMIT: $FAQ_CURRENT_COMMIT" + cd ../.. + if [ "$FAQ_LATEST_COMMIT" != "$FAQ_CURRENT_COMMIT" ]; then + echo "::error ::Submodule $FAQ_PATH is not up to date. Please update it." + exit 1 + else + echo "Submodule $FAQ_PATH is up to date." + fi diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 85a9a8615..8440b0412 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -1,6 +1,6 @@ repos: - repo: https://github.com/astral-sh/ruff-pre-commit - rev: "v0.4.8" + rev: "v0.5.5" hooks: - id: ruff args: [ --fix ] diff --git a/CHANGELOG.md b/CHANGELOG.md index 31da58387..94dfbeffd 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -3,8 +3,6 @@ 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/), and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html). -## [Unreleased] - ## [2.8.0rc1] ### Added @@ -22,27 +20,37 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ### Changed - Priority is given to `snapping_points` in `GridSpec` when close to structure boundaries, which reduces the chance of them being skipped. +- Gradients for autograd are computed server-side by default. They can be computed locally (requiring more data download) by passing `local_gradient=True` to the `web.run()` and related functions. ### Fixed - Significant speedup for field projection computations. - Fix numerical precision issue in `FieldProjectionCartesianMonitor`. -## [2.7.2] +## [2.7.2] - 2024-08-07 ### Added - Mode solver plugin now supports 'EMESimulation'. -- TriangleMesh class: automatic removal of zero-area faces, and functions fill_holes and fix_winding to attempt mesh repair. +- `TriangleMesh` class: automatic removal of zero-area faces, and functions `fill_holes` and `fix_winding` to attempt mesh repair. +- Added progress bar for EME simulations. +- Support for grid specifications from grid cell boundary coordinates via `CustomGridBoundaries` that subclasses from `GridSpec1d`. +- More convenient mesh importing from another simulation through `grid_spec = GridSpec.from_grid(sim.grid)`. +- `autograd` gradient calculations can be performed on the server by passing `local_gradient = False` into `web.run()` or `web.run_async()`. +- Automatic differentiation with `autograd` supports multiple frequencies through single, broadband adjoint simulation when the objective function can be formulated as depending on a single dataset in the output `SimulationData` with frequency dependence only. +- Convenience method `EMESimulation.subsection` to create a new EME simulation based on a subregion of an existing one. +- Added `flux` and `poynting` properties to `FieldProjectionCartesianData`. ### Changed -- Error if field projection monitors found in 2D simulations, except `FieldProjectionAngleMonitor` with `far_field_approx = True`. Support for other monitors and for exact field projection will be coming in a subsequent Tidy3D version. +- Mode solver now always operates on a reduced simulation copy. +- Moved `EMESimulation` size limit validators to preupload. +- Error if field projection monitors found in 2D simulations, except `FieldProjectionAngleMonitor` or `FieldProjectionCartesianMonitor` with `far_field_approx = True`. Support for other monitors and for exact field projection will be coming in a subsequent Tidy3D version. ### Fixed -- Error when loading a previously run Batch or ComponentModeler containing custom data. +- Error when loading a previously run `Batch` or `ComponentModeler` containing custom data. - 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. - +- 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. -## [2.7.1] +## [2.7.1] - 2024-07-10 ### Added - Support for differentiation with respect to `GeometryGroup.geometries` elements. diff --git a/docs/_static/img/apodization.png b/docs/_static/img/apodization.png new file mode 100755 index 000000000..c75df601a Binary files /dev/null and b/docs/_static/img/apodization.png differ diff --git a/docs/api/discretization.rst b/docs/api/discretization.rst index b9a3be7d9..a15172083 100644 --- a/docs/api/discretization.rst +++ b/docs/api/discretization.rst @@ -11,6 +11,7 @@ Discretization tidy3d.AutoGrid tidy3d.UniformGrid tidy3d.CustomGrid + tidy3d.CustomGridBoundaries tidy3d.Coords tidy3d.FieldGrid tidy3d.YeeGrid diff --git a/poetry.lock b/poetry.lock index 210dc3021..77c0ace7a 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.2 and should not be changed by hand. [[package]] name = "absl-py" @@ -2637,8 +2637,8 @@ files = [ numpy = [ {version = ">=1.26.0", markers = "python_version >= \"3.12\""}, {version = ">=1.23.3", markers = "python_version >= \"3.11\" and python_version < \"3.12\""}, - {version = ">1.20", markers = "python_version < \"3.10\""}, {version = ">=1.21.2", markers = "python_version >= \"3.10\" and python_version < \"3.11\""}, + {version = ">1.20", markers = "python_version < \"3.10\""}, ] [package.extras] @@ -4436,28 +4436,29 @@ files = [ [[package]] name = "ruff" -version = "0.4.8" +version = "0.5.5" description = "An extremely fast Python linter and code formatter, written in Rust." optional = true python-versions = ">=3.7" files = [ - {file = "ruff-0.4.8-py3-none-macosx_10_12_x86_64.whl", hash = "sha256:7663a6d78f6adb0eab270fa9cf1ff2d28618ca3a652b60f2a234d92b9ec89066"}, - {file = "ruff-0.4.8-py3-none-macosx_11_0_arm64.whl", hash = "sha256:eeceb78da8afb6de0ddada93112869852d04f1cd0f6b80fe464fd4e35c330913"}, - {file = "ruff-0.4.8-py3-none-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:aad360893e92486662ef3be0a339c5ca3c1b109e0134fcd37d534d4be9fb8de3"}, - {file = "ruff-0.4.8-py3-none-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:284c2e3f3396fb05f5f803c9fffb53ebbe09a3ebe7dda2929ed8d73ded736deb"}, - {file = "ruff-0.4.8-py3-none-manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:a7354f921e3fbe04d2a62d46707e569f9315e1a613307f7311a935743c51a764"}, - {file = "ruff-0.4.8-py3-none-manylinux_2_17_ppc64.manylinux2014_ppc64.whl", hash = "sha256:72584676164e15a68a15778fd1b17c28a519e7a0622161eb2debdcdabdc71883"}, - {file = "ruff-0.4.8-py3-none-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:9678d5c9b43315f323af2233a04d747409d1e3aa6789620083a82d1066a35199"}, - {file = "ruff-0.4.8-py3-none-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:704977a658131651a22b5ebeb28b717ef42ac6ee3b11e91dc87b633b5d83142b"}, - {file = "ruff-0.4.8-py3-none-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:d05f8d6f0c3cce5026cecd83b7a143dcad503045857bc49662f736437380ad45"}, - {file = "ruff-0.4.8-py3-none-musllinux_1_2_aarch64.whl", hash = "sha256:6ea874950daca5697309d976c9afba830d3bf0ed66887481d6bca1673fc5b66a"}, - {file = "ruff-0.4.8-py3-none-musllinux_1_2_armv7l.whl", hash = "sha256:fc95aac2943ddf360376be9aa3107c8cf9640083940a8c5bd824be692d2216dc"}, - {file = "ruff-0.4.8-py3-none-musllinux_1_2_i686.whl", hash = "sha256:384154a1c3f4bf537bac69f33720957ee49ac8d484bfc91720cc94172026ceed"}, - {file = "ruff-0.4.8-py3-none-musllinux_1_2_x86_64.whl", hash = "sha256:e9d5ce97cacc99878aa0d084c626a15cd21e6b3d53fd6f9112b7fc485918e1fa"}, - {file = "ruff-0.4.8-py3-none-win32.whl", hash = "sha256:6d795d7639212c2dfd01991259460101c22aabf420d9b943f153ab9d9706e6a9"}, - {file = "ruff-0.4.8-py3-none-win_amd64.whl", hash = "sha256:e14a3a095d07560a9d6769a72f781d73259655919d9b396c650fc98a8157555d"}, - {file = "ruff-0.4.8-py3-none-win_arm64.whl", hash = "sha256:14019a06dbe29b608f6b7cbcec300e3170a8d86efaddb7b23405cb7f7dcaf780"}, - {file = "ruff-0.4.8.tar.gz", hash = "sha256:16d717b1d57b2e2fd68bd0bf80fb43931b79d05a7131aa477d66fc40fbd86268"}, + {file = "ruff-0.5.5-py3-none-linux_armv6l.whl", hash = "sha256:605d589ec35d1da9213a9d4d7e7a9c761d90bba78fc8790d1c5e65026c1b9eaf"}, + {file = "ruff-0.5.5-py3-none-macosx_10_12_x86_64.whl", hash = "sha256:00817603822a3e42b80f7c3298c8269e09f889ee94640cd1fc7f9329788d7bf8"}, + {file = "ruff-0.5.5-py3-none-macosx_11_0_arm64.whl", hash = "sha256:187a60f555e9f865a2ff2c6984b9afeffa7158ba6e1eab56cb830404c942b0f3"}, + {file = "ruff-0.5.5-py3-none-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:fe26fc46fa8c6e0ae3f47ddccfbb136253c831c3289bba044befe68f467bfb16"}, + {file = "ruff-0.5.5-py3-none-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:4ad25dd9c5faac95c8e9efb13e15803cd8bbf7f4600645a60ffe17c73f60779b"}, + {file = "ruff-0.5.5-py3-none-manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:f70737c157d7edf749bcb952d13854e8f745cec695a01bdc6e29c29c288fc36e"}, + {file = "ruff-0.5.5-py3-none-manylinux_2_17_ppc64.manylinux2014_ppc64.whl", hash = "sha256:cfd7de17cef6ab559e9f5ab859f0d3296393bc78f69030967ca4d87a541b97a0"}, + {file = "ruff-0.5.5-py3-none-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:a09b43e02f76ac0145f86a08e045e2ea452066f7ba064fd6b0cdccb486f7c3e7"}, + {file = "ruff-0.5.5-py3-none-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:d0b856cb19c60cd40198be5d8d4b556228e3dcd545b4f423d1ad812bfdca5884"}, + {file = "ruff-0.5.5-py3-none-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:3687d002f911e8a5faf977e619a034d159a8373514a587249cc00f211c67a091"}, + {file = "ruff-0.5.5-py3-none-musllinux_1_2_aarch64.whl", hash = "sha256:ac9dc814e510436e30d0ba535f435a7f3dc97f895f844f5b3f347ec8c228a523"}, + {file = "ruff-0.5.5-py3-none-musllinux_1_2_armv7l.whl", hash = "sha256:af9bdf6c389b5add40d89b201425b531e0a5cceb3cfdcc69f04d3d531c6be74f"}, + {file = "ruff-0.5.5-py3-none-musllinux_1_2_i686.whl", hash = "sha256:d40a8533ed545390ef8315b8e25c4bb85739b90bd0f3fe1280a29ae364cc55d8"}, + {file = "ruff-0.5.5-py3-none-musllinux_1_2_x86_64.whl", hash = "sha256:cab904683bf9e2ecbbe9ff235bfe056f0eba754d0168ad5407832928d579e7ab"}, + {file = "ruff-0.5.5-py3-none-win32.whl", hash = "sha256:696f18463b47a94575db635ebb4c178188645636f05e934fdf361b74edf1bb2d"}, + {file = "ruff-0.5.5-py3-none-win_amd64.whl", hash = "sha256:50f36d77f52d4c9c2f1361ccbfbd09099a1b2ea5d2b2222c586ab08885cf3445"}, + {file = "ruff-0.5.5-py3-none-win_arm64.whl", hash = "sha256:3191317d967af701f1b73a31ed5788795936e423b7acce82a2b63e26eb3e89d6"}, + {file = "ruff-0.5.5.tar.gz", hash = "sha256:cc5516bdb4858d972fbc31d246bdb390eab8df1a26e2353be2dbc0c2d7f5421a"}, ] [[package]] @@ -5540,4 +5541,4 @@ vtk = ["vtk"] [metadata] lock-version = "2.0" python-versions = ">=3.9,<4.0.0" -content-hash = "9628ba0740a5ac4eb7f4d84a206e9911dd7a0ba3a9dba28b89df3bb0719dce55" +content-hash = "08216f09c8fa651fa543bb32d2305396afde84c069e4a6059d1761e752923394" diff --git a/pyproject.toml b/pyproject.toml index 2c0468472..dc4c1b89d 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -50,7 +50,7 @@ joblib = "*" ### Optional dependencies ### # development core bump-my-version = { version = "*", optional = true } -ruff = { version = "0.4.8", optional = true } +ruff = { version = "0.5.5", optional = true } coverage = { version = "*", optional = true } dill = { version = "*", optional = true } ipython = { version = "*", optional = true } diff --git a/tests/ruff.toml b/tests/ruff.toml index 2080302f5..56f5406cc 100644 --- a/tests/ruff.toml +++ b/tests/ruff.toml @@ -9,4 +9,5 @@ extend-ignore = [ "E402", # module-level import not at top of file "E731", # lambda assignment "F841", # unused local variable + "S101", # asserts allowed in tests ] diff --git a/tests/test_components/test_autograd.py b/tests/test_components/test_autograd.py index e7b0f0204..3fc342c56 100644 --- a/tests/test_components/test_autograd.py +++ b/tests/test_components/test_autograd.py @@ -5,6 +5,7 @@ import typing import warnings from importlib import reload +from os.path import join import autograd as ag import autograd.numpy as anp @@ -12,12 +13,15 @@ import numpy as np import pytest import tidy3d as td +import tidy3d.web as web +import xarray as xr from tidy3d.components.autograd.derivative_utils import DerivativeInfo +from tidy3d.components.data.sim_data import AdjointSourceInfo from tidy3d.plugins.polyslab import ComplexPolySlab -from tidy3d.web import run_async -from tidy3d.web.api.autograd.autograd import run +from tidy3d.web import run, run_async +from tidy3d.web.api.autograd.utils import FieldMap -from ..utils import SIM_FULL, AssertLogLevel, run_emulated +from ..utils import SIM_FULL, run_emulated """ Test configuration """ @@ -54,6 +58,7 @@ WVL = 1.0 FREQ0 = td.C_0 / WVL +FREQS = [FREQ0] # sim sizes LZ = 7 * WVL @@ -117,46 +122,120 @@ def use_emulated_run(monkeypatch): """If this fixture is used, the `tests.utils.run_emulated` function is used for simulation.""" + import tidy3d + if TEST_MODE in ("pipeline", "speed"): - import tidy3d.web.api.webapi as webapi + task_id_fwd = "task_fwd" + AUX_KEY_SIM_FIELDS_KEYS = "sim_fields_keys" - monkeypatch.setattr(webapi, "run", run_emulated) - _run_was_emulated[0] = True + cache = {} - # import here so it uses emulated run - from tidy3d.web.api.autograd import autograd + import tidy3d.web.api.webapi as webapi - reload(autograd) + # reload(tidy3d.web.api.autograd.autograd) + from tidy3d.web.api.autograd.autograd import ( + AUX_KEY_SIM_DATA_FWD, + AUX_KEY_SIM_DATA_ORIGINAL, + postprocess_adj, + postprocess_fwd, + ) + def emulated_run_fwd(simulation, task_name, **run_kwargs) -> td.SimulationData: + """What gets called instead of ``web/api/autograd/autograd.py::_run_tidy3d``.""" + task_id_fwd = task_name + if run_kwargs.get("simulation_type") == "autograd_fwd": + sim_original = simulation + sim_fields_keys = run_kwargs["sim_fields_keys"] + # add gradient monitors and make combined simulation + sim_combined = sim_original.with_adjoint_monitors(sim_fields_keys) + sim_data_combined = run_emulated(sim_combined, task_name=task_name) + + # store both original and fwd data aux_data + aux_data = {} + + _ = postprocess_fwd( + sim_data_combined=sim_data_combined, + sim_original=sim_original, + aux_data=aux_data, + ) + + # cache original and fwd data locally for test + cache[task_id_fwd] = copy.copy(aux_data) + cache[task_id_fwd][AUX_KEY_SIM_FIELDS_KEYS] = sim_fields_keys + # return original data only + return aux_data[AUX_KEY_SIM_DATA_ORIGINAL], task_id_fwd + else: + return run_emulated(simulation, task_name=task_name), task_id_fwd + + def emulated_run_bwd(simulation, task_name, **run_kwargs) -> td.SimulationData: + """What gets called instead of ``web/api/autograd/autograd.py::_run_tidy3d_bwd``.""" + + task_id_fwd = task_name[:-8] + + # run the adjoint sim + sim_data_adj = run_emulated(simulation, task_name="task_name") + + # grab the fwd and original data from the cache + aux_data_fwd = cache[task_id_fwd] + sim_data_orig = aux_data_fwd[AUX_KEY_SIM_DATA_ORIGINAL] + sim_data_fwd = aux_data_fwd[AUX_KEY_SIM_DATA_FWD] + + # 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, + ) -@pytest.fixture -def use_emulated_run_async(monkeypatch): - """If this fixture is used, the `tests.utils.run_emulated` function is used for simulation.""" + return traced_fields_vjp - if TEST_MODE in ("pipeline", "speed"): - import tidy3d.web.api.asynchronous as asynchronous + def emulated_run_async_fwd(simulations, **run_kwargs) -> td.SimulationData: + batch_data_orig, task_ids_fwd = {}, {} + sim_fields_keys_dict = run_kwargs.pop("sim_fields_keys_dict", None) + for task_name, simulation in simulations.items(): + if sim_fields_keys_dict is not None: + run_kwargs["sim_fields_keys"] = sim_fields_keys_dict[task_name] + sim_data_orig, task_id_fwd = emulated_run_fwd(simulation, task_name, **run_kwargs) + batch_data_orig[task_name] = sim_data_orig + task_ids_fwd[task_name] = task_id_fwd - def run_async_emulated(simulations, **kwargs): - """Mock version of ``run_async``.""" - return { - task_name: run_emulated(sim, task_name=task_name) - for task_name, sim in simulations.items() - } + return batch_data_orig, task_ids_fwd - monkeypatch.setattr(asynchronous, "run_async", run_async_emulated) - _run_was_emulated[0] = True + def emulated_run_async_bwd(simulations, **run_kwargs) -> td.SimulationData: + vjp_dict = {} + for task_name, simulation in simulations.items(): + task_id_fwd = task_name[:-8] + vjp_dict[task_name] = emulated_run_bwd(simulation, task_name, **run_kwargs) + return vjp_dict - # import here so it uses emulated run - from tidy3d.web.api.autograd import autograd + monkeypatch.setattr(webapi, "run", run_emulated) + monkeypatch.setattr(tidy3d.web.api.autograd.autograd, "_run_tidy3d", emulated_run_fwd) + monkeypatch.setattr(tidy3d.web.api.autograd.autograd, "_run_tidy3d_bwd", emulated_run_bwd) + monkeypatch.setattr( + tidy3d.web.api.autograd.autograd, "_run_async_tidy3d", emulated_run_async_fwd + ) + monkeypatch.setattr( + tidy3d.web.api.autograd.autograd, "_run_async_tidy3d_bwd", emulated_run_async_bwd + ) - reload(autograd) + _run_was_emulated[0] = True + return emulated_run_fwd, emulated_run_bwd def make_structures(params: anp.ndarray) -> dict[str, td.Structure]: """Make a dictionary of the structures given the parameters.""" + np.random.seed(0) + vector = np.random.random(N_PARAMS) - 0.5 - vector /= np.linalg.norm(vector) + vector = vector / np.linalg.norm(vector) # static components box = td.Box(center=(0, 0, 0), size=(1, 1, 1)) @@ -286,18 +365,6 @@ def make_structures(params: anp.ndarray) -> dict[str, td.Structure]: medium=td.Medium(permittivity=eps, conductivity=conductivity), ) - # geometry group - geo_group = td.Structure( - geometry=td.GeometryGroup( - geometries=[ - medium.geometry, - center_list.geometry, - size_element.geometry, - ], - ), - medium=td.Medium(permittivity=eps, conductivity=conductivity), - ) - # dispersive medium eps_inf = 1 + anp.abs(vector @ params) box = td.Box(center=(0, 0, 0), size=(1, 1, 1)) @@ -453,7 +520,7 @@ def plot_sim(sim: td.Simulation, plot_eps: bool = False) -> None: args = [("polyslab", "mode")] -# args = [("custom_med_vec", "mode")] +# args = [("custom_med", "mode")] def get_functions(structure_key: str, monitor_key: str) -> typing.Callable: @@ -569,7 +636,7 @@ def task_name_fn(i: int, sign: int) -> str: @pytest.mark.parametrize("structure_key, monitor_key", args) -def test_autograd_async(use_emulated_run_async, structure_key, monitor_key): +def test_autograd_async(use_emulated_run, structure_key, monitor_key): """Test an objective function through tidy3d autograd.""" fn_dict = get_functions(structure_key, monitor_key) @@ -601,7 +668,6 @@ def test_autograd_speed_num_structures(use_emulated_run): import time fn_dict = get_functions(ALL_KEY, ALL_KEY) - make_sim = fn_dict["sim"] monitor_key = "mode" structure_key = "size_element" @@ -629,104 +695,93 @@ def objective(*args): print(f"{num_structures_test} structures took {t2:.2e} seconds") -@pytest.mark.parametrize("structure_key", ("custom_med",)) -def test_sim_full_ops(structure_key): - """make sure the autograd operations don't error on a simulation containing everything.""" - - def objective(*params): - s = make_structures(*params)[structure_key] - s = s.updated_copy(geometry=s.geometry.updated_copy(center=(2, 2, 2), size=(0, 0, 0))) - sim_full_traced = SIM_FULL.updated_copy(structures=list(SIM_FULL.structures) + [s]) - - sim_full_static = sim_full_traced.to_static() - - 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 - - sim_traced = sim_full_static.insert_traced_fields(sim_fields) +@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.""" - assert sim_traced == sim_full_traced + fn_dict = get_functions(structure_key, monitor_key) + make_sim = fn_dict["sim"] + postprocess = fn_dict["postprocess"] - return anp.sum(sim_full_traced.structures[-1].medium.permittivity.values) + def objective(*args): + """Objective function.""" + sim = make_sim(*args) + data = run(sim, task_name="autograd_test", verbose=False, local_gradient=False) + value = postprocess(data) + return value - ag.grad(objective)(params0) + val, grad = ag.value_and_grad(objective)(params0) + print(val, grad) + assert anp.all(grad != 0.0), "some gradients are 0" + val, grad = ag.value_and_grad(objective)(params0) -def test_warning_no_adjoint_sources(log_capture, monkeypatch, use_emulated_run): - """Make sure we get the right warning with no adjoint sources, and no error.""" - monitor_key = "mode" - structure_key = "size_element" - monitor, postprocess = make_monitors()[monitor_key] +@pytest.mark.parametrize("structure_key, monitor_key", args) +def test_autograd_async_server(use_emulated_run, structure_key, monitor_key): + """Test an async objective function through tidy3d autograd.""" - def make_sim(*args): - structure = make_structures(*args)[structure_key] - return SIM_BASE.updated_copy(structures=[structure], monitors=[monitor]) + 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) - data = run(sim, task_name="autograd_test", verbose=False) - value = postprocess(data, data[monitor_key]) + sims = {"autograd_test1": sim, "autograd_test2": sim} + batch_data = run_async(sims, verbose=False, local_gradient=False) + value = 0.0 + for _, sim_data in batch_data.items(): + value = value + postprocess(sim_data) return value - monkeypatch.setattr(td.SimulationData, "make_adjoint_sources", lambda *args, **kwargs: []) - - with AssertLogLevel(log_capture, "WARNING", contains_str="No adjoint sources"): - ag.grad(objective)(params0) - + val, grad = ag.value_and_grad(objective)(params0) + print(val, grad) + assert anp.all(grad != 0.0), "some gradients are 0" -def test_web_failure_handling(log_capture, monkeypatch, use_emulated_run, use_emulated_run_async): - """Test what happens when autograd run pipeline fails.""" + val, grad = ag.value_and_grad(objective)(params0) - monitor_key = "mode" - structure_key = "size_element" - monitor, postprocess = make_monitors()[monitor_key] - def make_sim(*args): - structure = make_structures(*args)[structure_key] - return SIM_BASE.updated_copy(structures=[structure], monitors=[monitor]) +@pytest.mark.parametrize("structure_key", ("custom_med",)) +def test_sim_full_ops(structure_key): + """make sure the autograd operations don't error on a simulation containing everything.""" - def objective(*args): - """Objective function.""" - sim = make_sim(*args) - data = run(sim, task_name="autograd_test", verbose=False) - value = postprocess(data, data[monitor_key]) - return value + def objective(*params): + s = make_structures(*params)[structure_key] + s = s.updated_copy(geometry=s.geometry.updated_copy(center=(2, 2, 2), size=(0, 0, 0))) + sim_full_traced = SIM_FULL.updated_copy(structures=list(SIM_FULL.structures) + [s]) - def fail(*args, **kwargs): - """Just raise an exception.""" - raise ValueError("test") + sim_full_static = sim_full_traced.to_static() - """ if autograd run raises exception, raise a warning and continue with regular .""" + sim_fields = sim_full_traced.strip_traced_fields() - monkeypatch.setattr(td.web.api.autograd.autograd, "_run", fail) + # note: there is one traced structure in SIM_FULL already with 6 fields + 1 = 7 + assert len(sim_fields) == 7 - with AssertLogLevel( - log_capture, "WARNING", contains_str="If you received this warning, please file an issue" - ): - ag.grad(objective)(params0) + sim_traced = sim_full_static.insert_traced_fields(sim_fields) - def objective_async(*args): - sims = {"key": make_sim(*args)} - data = run_async(sims, verbose=False) - value = 0.0 - for _, val in data.items(): - value += postprocess(val, val[monitor_key]) - return value + assert sim_traced == sim_full_traced - """ if autograd run_async raises exception, raise a warning and continue with regular .""" + return anp.sum(sim_full_traced.structures[-1].medium.permittivity.values) - monkeypatch.setattr(td.web.api.autograd.autograd, "_run_async", fail) + ag.grad(objective)(params0) - with AssertLogLevel( - log_capture, "WARNING", contains_str="If you received this warning, please file an issue" - ): - ag.grad(objective_async)(params0) - """ if the regular web functions are called, raise custom exception and catch it in tests .""" +@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 + from file, and then converting back, returns the same object.""" + s = make_structures(params0)[structure_key] + s = s.updated_copy(geometry=s.geometry.updated_copy(center=(2, 2, 2), size=(0, 0, 0))) + sim_full_traced = SIM_FULL.updated_copy(structures=list(SIM_FULL.structures) + [s]) + sim_fields = sim_full_traced.strip_traced_fields() + + field_map = FieldMap.from_autograd_field_map(sim_fields) + field_map_file = join(tmp_path, "test_sim_fields.hdf5.gz") + field_map.to_file(field_map_file) + autograd_field_map = FieldMap.from_file(field_map_file).to_autograd_field_map + for path, data in sim_fields.items(): + assert np.all(data == autograd_field_map[path]) def test_web_incompatible_inputs(log_capture, monkeypatch): @@ -737,6 +792,7 @@ def catch(*args, **kwargs): raise AssertionError() monkeypatch.setattr(td.web.api.webapi, "run", catch) + monkeypatch.setattr(td.web.api.container.Job, "run", catch) monkeypatch.setattr(td.web.api.asynchronous, "run_async", catch) from tidy3d.web.api.autograd import autograd @@ -1050,3 +1106,219 @@ def f(x): * no copy : 16 sec * no to_static(): 13 sec """ + +FREQ1 = FREQ0 * 1.6 + +mnt_single = td.ModeMonitor( + size=(2, 2, 0), + center=(0, 0, LZ / 2 - WVL), + mode_spec=td.ModeSpec(num_modes=2), + freqs=[FREQ0], + name="single", +) + +mnt_multi = td.ModeMonitor( + size=(2, 2, 0), + center=(0, 0, LZ / 2 - WVL), + mode_spec=td.ModeSpec(num_modes=2), + freqs=[FREQ0, FREQ1], + name="multi", +) + + +def make_objective(postprocess_fn: typing.Callable, structure_key: str) -> typing.Callable: + def objective(params): + structure_traced = make_structures(params)[structure_key] + sim = SIM_BASE.updated_copy( + structures=[structure_traced], + monitors=list(SIM_BASE.monitors) + [mnt_single, mnt_multi], + ) + data = run(sim, task_name="multifreq_test") + return postprocess_fn(data) + + return objective + + +def get_amps(sim_data: td.SimulationData, mnt_name: str) -> xr.DataArray: + return sim_data[mnt_name].amps + + +def power(amps: xr.DataArray) -> float: + """Reduce a selected DataArray into just a float for objective function.""" + return anp.sum(anp.abs(amps.values) ** 2) + + +def postprocess_0_src(sim_data: td.SimulationData) -> float: + """Postprocess function that should return 0 adjoint sources.""" + return 0.0 + + +def compute_grad(postprocess_fn: typing.Callable, structure_key: str) -> typing.Callable: + objective = make_objective(postprocess_fn, structure_key=structure_key) + params = params0 + 1.0 # +1 is to avoid a warning in size_element with value 0 + return ag.grad(objective)(params) + + +def check_1_src_single(log_capture, structure_key): + def postprocess(sim_data: td.SimulationData) -> float: + """Postprocess function that should return 1 adjoint sources.""" + amps = get_amps(sim_data, "single").sel(mode_index=0, direction="+") + return power(amps) + + return postprocess + + +def check_2_src_single(log_capture, structure_key): + def postprocess(sim_data: td.SimulationData) -> float: + """Postprocess function that should return 2 different adjoint sources.""" + amps = get_amps(sim_data, "single").sel(mode_index=0) + return power(amps) + + return postprocess + + +def check_1_src_multi(log_capture, structure_key): + def postprocess(sim_data: td.SimulationData) -> float: + """Postprocess function that should return 1 adjoint sources.""" + amps = get_amps(sim_data, "multi").sel(mode_index=0, direction="+", f=FREQ0) + return power(amps) + + return postprocess + + +def check_2_src_multi(log_capture, structure_key): + def postprocess(sim_data: td.SimulationData) -> float: + """Postprocess function that should return 2 different adjoint sources.""" + amps = get_amps(sim_data, "multi").sel(mode_index=0, f=FREQ1) + return power(amps) + + return postprocess + + +def check_2_src_both(log_capture, structure_key): + def postprocess(sim_data: td.SimulationData) -> float: + """Postprocess function that should return 2 different adjoint sources.""" + amps_single = get_amps(sim_data, "single").sel(mode_index=0, direction="+") + amps_multi = get_amps(sim_data, "multi").sel(mode_index=0, direction="+", f=FREQ0) + return power(amps_single) + power(amps_multi) + + return postprocess + + +def check_1_multisrc(log_capture, structure_key): + def postprocess(sim_data: td.SimulationData) -> float: + """Postprocess function that should raise ValueError because diff sources, diff freqs.""" + amps_single = get_amps(sim_data, "single").sel(mode_index=0, direction="+") + amps_multi = get_amps(sim_data, "multi").sel(mode_index=0, direction="+", f=FREQ1) + return power(amps_single) + power(amps_multi) + + return postprocess + + +def check_2_multisrc(log_capture, structure_key): + def postprocess(sim_data: td.SimulationData) -> float: + """Postprocess function that should raise ValueError because diff sources, diff freqs.""" + amps_single = get_amps(sim_data, "single").sel(mode_index=0, direction="+") + amps_multi = get_amps(sim_data, "multi").sel(mode_index=0, direction="+") + return power(amps_single) + power(amps_multi) + + return postprocess + + +def check_1_src_broadband(log_capture, structure_key): + def postprocess(sim_data: td.SimulationData) -> float: + """Postprocess function that should return 1 broadband adjoint sources with many freqs.""" + amps = get_amps(sim_data, "multi").sel(mode_index=0, direction="+") + return power(amps) + + return postprocess + + +MULT_FREQ_TEST_CASES = dict( + src_1_freq_1=check_1_src_single, + src_2_freq_1=check_2_src_single, + src_1_freq_2=check_1_src_multi, + src_2_freq_1_mon_1=check_1_src_multi, + src_2_freq_1_mon_2=check_2_src_both, + src_2_freq_2_mon_1=check_1_multisrc, + src_2_freq_2_mon_2=check_2_multisrc, + src_1_freq_2_broadband=check_1_src_broadband, +) + +checks = list(MULT_FREQ_TEST_CASES.items()) + + +@pytest.mark.parametrize("label, check_fn", checks) +@pytest.mark.parametrize("structure_key", ("custom_med",)) +def test_multi_freq_edge_cases( + log_capture, use_emulated_run, structure_key, label, check_fn, monkeypatch +): + # test multi-frequency adjoint handling + + import tidy3d.components.data.sim_data as sd + + monkeypatch.setattr(sd, "RESIDUAL_CUTOFF_ADJOINT", 1) + reload(td) + + postprocess_fn = check_fn(structure_key=structure_key, log_capture=log_capture) + + def objective(params): + structure_traced = make_structures(params)[structure_key] + sim = SIM_BASE.updated_copy( + structures=[structure_traced], + monitors=list(SIM_BASE.monitors) + [mnt_single, mnt_multi], + ) + data = run(sim, task_name="multifreq_test") + return postprocess_fn(data) + + if label == "src_2_freq_2_mon_2": + with pytest.raises(NotImplementedError): + g = ag.grad(objective)(params0) + else: + g = ag.grad(objective)(params0) + print(g) + + +@pytest.mark.parametrize("structure_key", structure_keys_) +def test_multi_frequency_equivalence(use_emulated_run, structure_key): + """Test an objective function through tidy3d autograd.""" + + def objective_indi(params, structure_key) -> float: + power_sum = 0.0 + + for f in mnt_multi.freqs: + structure_traced = make_structures(params)[structure_key] + sim = SIM_BASE.updated_copy( + structures=[structure_traced], + monitors=list(SIM_BASE.monitors) + [mnt_multi], + ) + + sim_data = web.run(sim, task_name="multifreq_test") + amps_i = get_amps(sim_data, "multi").sel(mode_index=0, direction="+", f=f) + power_i = power(amps_i) + power_sum = power_sum + power_i + + return power_sum + + def objective_multi(params, structure_key) -> float: + structure_traced = make_structures(params)[structure_key] + sim = SIM_BASE.updated_copy( + structures=[structure_traced], + monitors=list(SIM_BASE.monitors) + [mnt_multi], + ) + sim_data = web.run(sim, task_name="multifreq_test") + amps = get_amps(sim_data, "multi").sel(mode_index=0, direction="+") + return power(amps) + + params0_ = params0 + 1.0 + + # J_indi = objective_indi(params0_, structure_key) + # J_multi = objective_multi(params0_, structure_key) + + # np.testing.assert_allclose(J_indi, J_multi) + + grad_indi = ag.grad(objective_indi)(params0_, structure_key=structure_key) + grad_multi = ag.grad(objective_multi)(params0_, structure_key=structure_key) + + assert not np.any(np.isclose(grad_indi, 0)) + assert not np.any(np.isclose(grad_multi, 0)) diff --git a/tests/test_components/test_custom.py b/tests/test_components/test_custom.py index 8d04d37c6..718aecaaa 100644 --- a/tests/test_components/test_custom.py +++ b/tests/test_components/test_custom.py @@ -604,6 +604,19 @@ def verify_custom_medium_methods(mat, reduced_fields): sim.subsection(subsection, remove_outside_custom_mediums=False) sim.subsection(subsection, remove_outside_custom_mediums=True) + # eme + eme_sim = td.EMESimulation( + axis=2, + freqs=[td.C_0], + size=(1, 1, 1), + grid_spec=td.GridSpec.auto(wavelength=1.0), + structures=(struct,), + eme_grid_spec=td.EMEUniformGrid(num_cells=1, mode_spec=td.EMEModeSpec()), + ) + _ = eme_sim.grid + eme_sim.subsection(subsection, remove_outside_custom_mediums=False) + eme_sim.subsection(subsection, remove_outside_custom_mediums=True) + def test_anisotropic_custom_medium(): """Anisotropic CustomMedium.""" diff --git a/tests/test_components/test_eme.py b/tests/test_components/test_eme.py index 4ca218965..1d2379566 100644 --- a/tests/test_components/test_eme.py +++ b/tests/test_components/test_eme.py @@ -296,6 +296,22 @@ def test_eme_simulation(log_capture): _ = sim2.plot(y=0, ax=AX) _ = sim2.plot(z=0, ax=AX) + # must be 3D + with pytest.raises(pd.ValidationError): + _ = td.EMESimulation( + size=(0, 2, 2), + freqs=[td.C_0], + axis=2, + eme_grid_spec=td.EMEUniformGrid(num_cells=2, mode_spec=td.EMEModeSpec()), + ) + with pytest.raises(pd.ValidationError): + _ = td.EMESimulation( + size=(2, 2, 0), + freqs=[td.C_0], + axis=2, + eme_grid_spec=td.EMEUniformGrid(num_cells=2, mode_spec=td.EMEModeSpec()), + ) + # need at least one freq with pytest.raises(pd.ValidationError): _ = sim.updated_copy(freqs=[]) @@ -324,7 +340,7 @@ def test_eme_simulation(log_capture): ) # test port offsets - with pytest.raises(SetupError): + with pytest.raises(ValidationError): _ = sim.updated_copy(port_offsets=[sim.size[sim.axis] * 2 / 3, sim.size[sim.axis] * 2 / 3]) # test duplicate freqs @@ -416,20 +432,25 @@ def test_eme_simulation(log_capture): _ = sim.updated_copy(boundary_spec=td.BoundarySpec.all_sides(td.Periodic())) # test max sim size and freqs + sim_bad = sim.updated_copy(size=(1000, 1000, 1000)) with pytest.raises(SetupError): - _ = sim.updated_copy(size=(1000, 1000, 1000)) + sim_bad.validate_pre_upload() + sim_bad = sim.updated_copy(size=(1000, 500, 3), monitors=[], store_port_modes=True) with pytest.raises(SetupError): - _ = sim.updated_copy(size=(1000, 500, 3), monitors=[], store_port_modes=True) + sim_bad.validate_pre_upload() + sim_bad = sim.updated_copy(size=(1000, 500, 3), monitors=[], store_port_modes=False) with pytest.raises(SetupError): - _ = sim.updated_copy(size=(1000, 500, 3), monitors=[], store_port_modes=False) + sim_bad.validate_pre_upload() + sim_bad = sim.updated_copy(size=(500, 500, 3), monitors=[]) with AssertLogLevel(log_capture, "WARNING", "slow-down"): - _ = sim.updated_copy(size=(500, 500, 3), monitors=[]) + sim_bad.validate_pre_upload() + sim_bad = sim.updated_copy( + freqs=list(sim.freqs) + list(1e14 * np.linspace(1, 2, 1000)), + grid_spec=sim.grid_spec.updated_copy(wavelength=1), + ) with pytest.raises(SetupError): - _ = sim.updated_copy( - freqs=list(sim.freqs) + list(1e14 * np.linspace(1, 2, 1000)), - grid_spec=sim.grid_spec.updated_copy(wavelength=1), - ) + sim_bad.validate_pre_upload() large_monitor = sim.monitors[2].updated_copy(size=(td.inf, td.inf, td.inf)) _ = sim.updated_copy( size=(10, 10, 10), @@ -437,27 +458,30 @@ def test_eme_simulation(log_capture): freqs=list(1e14 * np.linspace(1, 2, 1)), grid_spec=sim.grid_spec.updated_copy(wavelength=1), ) + sim_bad = sim.updated_copy( + size=(10, 10, 10), + monitors=[large_monitor], + freqs=list(1e14 * np.linspace(1, 2, 5)), + grid_spec=sim.grid_spec.updated_copy(wavelength=1), + ) with AssertLogLevel(log_capture, "WARNING", contains_str="estimated storage"): - _ = sim.updated_copy( - size=(10, 10, 10), - monitors=[large_monitor], - freqs=list(1e14 * np.linspace(1, 2, 5)), - grid_spec=sim.grid_spec.updated_copy(wavelength=1), - ) + sim_bad.validate_pre_upload() + sim_bad = sim.updated_copy( + size=(10, 10, 10), + monitors=[large_monitor], + freqs=list(1e14 * np.linspace(1, 2, 20)), + grid_spec=sim.grid_spec.updated_copy(wavelength=1), + ) with pytest.raises(SetupError): - _ = sim.updated_copy( - size=(10, 10, 10), - monitors=[large_monitor], - freqs=list(1e14 * np.linspace(1, 2, 20)), - grid_spec=sim.grid_spec.updated_copy(wavelength=1), - ) + sim_bad.validate_pre_upload() + sim_bad = sim.updated_copy( + size=(10, 10, 10), + monitors=[large_monitor, large_monitor.updated_copy(name="lmon2")], + freqs=list(1e14 * np.linspace(1, 2, 5)), + grid_spec=sim.grid_spec.updated_copy(wavelength=1), + ) with pytest.raises(SetupError): - _ = sim.updated_copy( - size=(10, 10, 10), - monitors=[large_monitor, large_monitor.updated_copy(name="lmon2")], - freqs=list(1e14 * np.linspace(1, 2, 5)), - grid_spec=sim.grid_spec.updated_copy(wavelength=1), - ) + sim_bad.validate_pre_upload() # test monitor that does not intersect any EME cells mode_monitor = td.EMEModeSolverMonitor( @@ -483,7 +507,14 @@ def test_eme_simulation(log_capture): assert sim_tmp._monitor_num_freqs(monitor=sim_tmp.monitors[0]) == 1 # test sweep - _ = sim.updated_copy(sweep_spec=td.EMELengthSweep(scale_factors=list(np.linspace(1, 2, 10)))) + sweep_sim = sim.updated_copy( + sweep_spec=td.EMELengthSweep(scale_factors=list(np.linspace(1, 2, 10))) + ) + assert sweep_sim._sweep_cells + assert not sweep_sim._sweep_interfaces + assert sweep_sim._num_sweep_cells == 10 + assert sweep_sim._num_sweep_interfaces == 1 + assert sweep_sim._num_sweep_modes == 1 _ = sim.updated_copy( sweep_spec=td.EMELengthSweep( scale_factors=np.stack((np.linspace(1, 2, 7), np.linspace(1, 2, 7))) @@ -491,6 +522,17 @@ def test_eme_simulation(log_capture): ) with pytest.raises(SetupError): _ = sim.updated_copy(sweep_spec=td.EMELengthSweep(scale_factors=[])) + with pytest.raises(SetupError): + _ = sim.updated_copy( + sweep_spec=td.EMELengthSweep( + scale_factors=np.stack( + ( + np.stack((np.linspace(1, 2, 7), np.linspace(1, 2, 7))), + np.stack((np.linspace(1, 2, 7), np.linspace(1, 2, 7))), + ) + ) + ) + ) # second shape of length sweep must equal number of cells with pytest.raises(SetupError): _ = sim.updated_copy(sweep_spec=td.EMELengthSweep(scale_factors=np.array([[1, 2], [3, 4]]))) @@ -498,32 +540,36 @@ def test_eme_simulation(log_capture): # test sweep size limit with pytest.raises(SetupError): _ = sim.updated_copy(sweep_spec=td.EMELengthSweep(scale_factors=[])) + sim_bad = sim.updated_copy( + sweep_spec=td.EMELengthSweep(scale_factors=list(np.linspace(1, 2, 200))) + ) with pytest.raises(SetupError): - _ = sim.updated_copy( - sweep_spec=td.EMELengthSweep(scale_factors=list(np.linspace(1, 2, 200))) - ) + sim_bad.validate_pre_upload() # can't exceed max num modes with pytest.raises(SetupError): _ = sim.updated_copy(sweep_spec=td.EMEModeSweep(num_modes=list(np.arange(150, 200)))) # don't warn in these two cases with AssertLogLevel(log_capture, None): - _ = sim.updated_copy( + sim_good = sim.updated_copy( constraint="passive", eme_grid_spec=td.EMEUniformGrid(num_cells=1, mode_spec=td.EMEModeSpec(num_modes=40)), grid_spec=sim.grid_spec.updated_copy(wavelength=1), ) - _ = sim.updated_copy( + sim_good.validate_pre_upload() + sim_good = sim.updated_copy( constraint=None, eme_grid_spec=td.EMEUniformGrid(num_cells=1, mode_spec=td.EMEModeSpec(num_modes=60)), grid_spec=sim.grid_spec.updated_copy(wavelength=1), ) + sim_good.validate_pre_upload() # warn about num modes with constraint + sim_bad = sim.updated_copy( + constraint="passive", + eme_grid_spec=td.EMEUniformGrid(num_cells=1, mode_spec=td.EMEModeSpec(num_modes=60)), + ) with AssertLogLevel(log_capture, "WARNING", contains_str="constraint"): - _ = sim.updated_copy( - constraint="passive", - eme_grid_spec=td.EMEUniformGrid(num_cells=1, mode_spec=td.EMEModeSpec(num_modes=60)), - ) + sim_bad.validate_pre_upload() _ = sim.port_modes_monitor @@ -1175,3 +1221,30 @@ def test_eme_sim_data(): assert "mode_index" not in field_in_basis.Ex.coords field_in_basis = sim_data.field_in_basis(field=sim_data["field"], modes=modes_in0, port_index=1) assert "mode_index" not in field_in_basis.Ex.coords + + +def test_eme_sim_subsection(): + eme_sim = td.EMESimulation( + axis=2, + size=(2, 2, 2), + freqs=[td.C_0], + grid_spec=td.GridSpec.auto(), + eme_grid_spec=td.EMEUniformGrid(num_cells=2, mode_spec=td.EMEModeSpec()), + ) + # check 3d subsection + region = td.Box(size=(2, 2, 1)) + subsection = eme_sim.subsection(region=region) + assert subsection.size[2] == 1 + + # check 3d subsection with identical eme grid + region = td.Box(size=(2, 2, 1)) + subsection = eme_sim.subsection(region=region, eme_grid_spec="identical") + assert subsection.size[2] == 2 + region = td.Box(size=(2, 2, 0.5), center=(0, 0, 0.5)) + subsection = eme_sim.subsection(region=region, eme_grid_spec="identical") + assert subsection.size[2] == 1 + + # 2d subsection errors + region = td.Box(size=(2, 2, 0)) + with pytest.raises(pd.ValidationError): + subsection = eme_sim.subsection(region=region) diff --git a/tests/test_components/test_field_projection.py b/tests/test_components/test_field_projection.py index 0e1523b87..23b37309b 100644 --- a/tests/test_components/test_field_projection.py +++ b/tests/test_components/test_field_projection.py @@ -344,6 +344,8 @@ def test_proj_clientside(): far_fields_cartesian.fields_cartesian far_fields_cartesian.radar_cross_section far_fields_cartesian.power + far_fields_cartesian.poynting + far_fields_cartesian.flux for val in far_fields_cartesian.field_components.values(): val.sel(f=f0) far_fields_cartesian.renormalize_fields(proj_distance=5e6) @@ -366,7 +368,198 @@ def test_proj_clientside(): exact_fields_cartesian.fields_cartesian exact_fields_cartesian.radar_cross_section exact_fields_cartesian.power + exact_fields_cartesian.poynting + exact_fields_cartesian.flux for val in exact_fields_cartesian.field_components.values(): val.sel(f=f0) with pytest.raises(DataError): exact_fields_cartesian.renormalize_fields(proj_distance=5e6) + + +def make_2d_proj_monitors(center, size, freqs, plane): + """Helper function to make near-to-far monitors in 2D simulations.""" + + if plane == "xy": + thetas = [np.pi / 2] + phis = np.linspace(0, 2 * np.pi, 100) + far_size = 10 * WAVELENGTH + Ns = 40 + xs = np.linspace(-far_size, far_size, Ns) + ys = [0] + projection_axis = 0 + elif plane == "yz": + thetas = np.linspace(0, np.pi, 1) + phis = [np.pi / 2] + far_size = 10 * WAVELENGTH + Ns = 40 + xs = [0] + ys = np.linspace(-far_size, far_size, Ns) + projection_axis = 1 + elif plane == "xz": + thetas = np.linspace(0, np.pi, 100) + phis = [0] + far_size = 10 * WAVELENGTH + Ns = 40 + xs = [0] + ys = np.linspace(-far_size, far_size, Ns) + projection_axis = 0 + else: + raise ValueError("Invalid plane. Use 'xy', 'yz', or 'xz'.") + + n2f_angle_monitor_2d = td.FieldProjectionAngleMonitor( + center=center, + size=size, + freqs=freqs, + name="far_field_angle", + phi=list(phis), + theta=list(thetas), + proj_distance=R_FAR, + far_field_approx=True, # Fields are far enough for geometric far field approximations + ) + + n2f_car_monitor_2d = td.FieldProjectionCartesianMonitor( + center=center, + size=size, + freqs=freqs, + name="far_field_cartesian", + x=list(xs), + y=list(ys), + 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) + + +def make_2d_proj(plane): + center = (0, 0, 0) + f0 = 1e13 + + if plane == "xy": + sim_size = (5, 5, 0) + monitor_size = (0, 2, td.inf) + # boundary conditions + boundary_conds = td.BoundarySpec( + x=td.Boundary.pml(), + y=td.Boundary.pml(), + z=td.Boundary.periodic(), + ) + # data coordinates + x = np.array([0.0]) + y = np.linspace(-1, 1, 10) + z = np.array([0.0]) + coords = dict(x=x, y=y, z=z, f=[f0]) + scalar_field = td.ScalarFieldDataArray( + (1 + 1j) * np.random.random((1, 10, 1, 1)), coords=coords + ) + elif plane == "yz": + sim_size = (0, 5, 5) + monitor_size = (td.inf, 0, 2) + # boundary conditions + boundary_conds = td.BoundarySpec( + x=td.Boundary.periodic(), + y=td.Boundary.pml(), + z=td.Boundary.pml(), + ) + # data coordinates + x = np.array([0.0]) + y = np.array([0.0]) + z = np.linspace(-1, 1, 10) + coords = dict(x=x, y=y, z=z, f=[f0]) + scalar_field = td.ScalarFieldDataArray( + (1 + 1j) * np.random.random((1, 1, 10, 1)), coords=coords + ) + elif plane == "xz": + sim_size = (5, 0, 5) + monitor_size = (0, td.inf, 2) + # boundary conditions + boundary_conds = td.BoundarySpec( + x=td.Boundary.pml(), + y=td.Boundary.periodic(), + z=td.Boundary.pml(), + ) + # data coordinates + x = np.array([0.0]) + y = np.array([0.0]) + z = np.linspace(-1, 1, 10) + coords = dict(x=x, y=y, z=z, f=[f0]) + scalar_field = td.ScalarFieldDataArray( + (1 + 1j) * np.random.random((1, 1, 10, 1)), coords=coords + ) + else: + raise ValueError("Invalid plane. Use 'xy', 'yz', or 'xz'.") + + monitor = td.FieldMonitor( + center=center, size=monitor_size, freqs=[f0], name="near_field", colocate=False + ) + + # Set up the simulation + sim = td.Simulation( + size=sim_size, + grid_spec=td.GridSpec.auto(wavelength=td.C_0 / f0), + boundary_spec=boundary_conds, + monitors=[monitor], + run_time=1e-12, + ) + + data = td.FieldData( + monitor=monitor, + Ex=scalar_field, + Ey=scalar_field, + Ez=scalar_field, + Hx=scalar_field, + Hy=scalar_field, + Hz=scalar_field, + symmetry=sim.symmetry, + symmetry_center=sim.center, + grid_expanded=sim.discretize_monitor(monitor), + ) + + sim_data = td.SimulationData(simulation=sim, data=(data,)) + + proj = td.FieldProjector.from_near_field_monitors( + sim_data=sim_data, + near_monitors=[monitor], + normal_dirs=["+"], + ) + + # make near-to-far monitors + ( + n2f_angle_monitor_2d, + n2f_cart_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) + + # compute far field quantities + far_fields_angular_2d.r + far_fields_angular_2d.theta + far_fields_angular_2d.phi + far_fields_angular_2d.fields_spherical + far_fields_angular_2d.fields_cartesian + far_fields_angular_2d.radar_cross_section + far_fields_angular_2d.power + for val in far_fields_angular_2d.field_components.values(): + val.sel(f=f0) + far_fields_angular_2d.renormalize_fields(proj_distance=5e6) + + far_fields_cartesian_2d.x + far_fields_cartesian_2d.y + far_fields_cartesian_2d.z + far_fields_cartesian_2d.fields_spherical + far_fields_cartesian_2d.fields_cartesian + far_fields_cartesian_2d.radar_cross_section + far_fields_cartesian_2d.power + for val in far_fields_cartesian_2d.field_components.values(): + val.sel(f=f0) + far_fields_cartesian_2d.renormalize_fields(proj_distance=5e6) + + +def test_2d_proj_clientside(): + # Run simulations and tests for all three planes + planes = ["xy", "yz", "xz"] + + for plane in planes: + make_2d_proj(plane) diff --git a/tests/test_components/test_grid_spec.py b/tests/test_components/test_grid_spec.py index d11226558..5000a40dc 100644 --- a/tests/test_components/test_grid_spec.py +++ b/tests/test_components/test_grid_spec.py @@ -344,3 +344,40 @@ def test_zerosize_dimensions(): ), run_time=1e-12, ) + + +def test_custom_grid_boundaries(): + custom = td.CustomGridBoundaries(coords=np.linspace(-1, 1, 11)) + grid_spec = td.GridSpec(grid_x=custom, grid_y=custom, grid_z=custom) + source = td.PointDipole( + source_time=td.GaussianPulse(freq0=3e14, fwidth=1e14), polarization="Ex" + ) + + # matches exactly + sim = td.Simulation( + size=(2, 2, 2), + sources=[source], + grid_spec=grid_spec, + run_time=1e-12, + medium=td.Medium(permittivity=4), + boundary_spec=td.BoundarySpec.all_sides(boundary=td.Periodic()), + ) + assert np.allclose(sim.grid.boundaries.x, custom.coords) + + # chop off + sim_chop = sim.updated_copy(size=(1, 1, 1)) + assert np.allclose(sim_chop.grid.boundaries.x, np.linspace(-0.4, 0.4, 5)) + + sim_chop = sim.updated_copy(size=(1.2, 1, 1)) + assert np.allclose(sim_chop.grid.boundaries.x, np.linspace(-0.6, 0.6, 7)) + + # expand + sim_expand = sim.updated_copy(size=(4, 4, 4)) + assert np.allclose(sim_expand.grid.boundaries.x, np.linspace(-2, 2, 21)) + + # pml + num_layers = 10 + sim_pml = sim.updated_copy( + boundary_spec=td.BoundarySpec.all_sides(boundary=td.PML(num_layers=num_layers)) + ) + assert np.allclose(sim_pml.grid.boundaries.x, np.linspace(-3, 3, 31)) diff --git a/tests/test_components/test_source.py b/tests/test_components/test_source.py index 3fbf84ddf..05ee39062 100644 --- a/tests/test_components/test_source.py +++ b/tests/test_components/test_source.py @@ -307,8 +307,20 @@ def test_custom_source_time(log_capture): atol=ATOL, ) + # all times out of range + _ = cst.amp_time([-1]) + _ = cst.amp_time(-1) + assert np.allclose(cst.amp_time([2]), np.exp(-1j * 2 * np.pi * 2 * freq0), rtol=0, atol=ATOL) + assert_log_level(log_capture, None) + vals = td.components.data.data_array.TimeDataArray([1, 2], coords=dict(t=[-1, -0.5])) + dataset = td.components.data.dataset.TimeDataset(values=vals) + cst = td.CustomSourceTime(source_time_dataset=dataset, freq0=freq0, fwidth=0.1e12) + source = td.PointDipole(center=(0, 0, 0), source_time=cst, polarization="Ex") + with AssertLogLevel(log_capture, "WARNING", contains_str="defined over a time range"): + sim = sim.updated_copy(sources=[source]) + # test normalization warning with AssertLogLevel(log_capture, "WARNING"): sim = sim.updated_copy(normalize_index=0) diff --git a/tests/test_data/test_monitor_data.py b/tests/test_data/test_monitor_data.py index 3ceb8fbae..772cacfd3 100644 --- a/tests/test_data/test_monitor_data.py +++ b/tests/test_data/test_monitor_data.py @@ -294,6 +294,20 @@ def test_mode_solver_data(): data_reversed = data.time_reversed_copy assert data_reversed.monitor.store_fields_direction == "-" + # check mode summary table with and without fields + modes_info = data.modes_info + assert all( + np.shape(modes_info[key]) != () + for key in ["TE (Ex) fraction", "wg TE fraction", "wg TM fraction", "mode area"] + ) + + data_no_fields = data.updated_copy(Ex=None) + modes_info = data_no_fields.modes_info + assert all( + np.shape(modes_info[key]) == () + for key in ["TE (Ex) fraction", "wg TE fraction", "wg TM fraction", "mode area"] + ) + def test_permittivity_data(): data = make_permittivity_data() diff --git a/tests/test_data/test_sim_data.py b/tests/test_data/test_sim_data.py index ac6f2de54..d31214419 100644 --- a/tests/test_data/test_sim_data.py +++ b/tests/test_data/test_sim_data.py @@ -303,7 +303,7 @@ def test_to_hdf5(tmp_path): sim_data.to_file(fname=FNAME) # Make sure data is loaded as Tidy3dDataArray and not xarray DataArray for data, data2 in zip(sim_data.data, sim_data2.data): - assert type(data) == type(data2) + assert type(data) is type(data2) assert sim_data == sim_data2 diff --git a/tests/test_plugins/autograd/test_differential_operators.py b/tests/test_plugins/autograd/test_differential_operators.py new file mode 100644 index 000000000..21fbe1b1a --- /dev/null +++ b/tests/test_plugins/autograd/test_differential_operators.py @@ -0,0 +1,23 @@ +import autograd.numpy as np +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 + + +def test_value_and_grad(rng): + """Test the custom value_and_grad function against autograd's implementation""" + x = 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)) + + (v, g), aux = vg_fun(x) + v_ag, g_ag = vg_fun_ag(x) + + # 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 diff --git a/tests/test_plugins/test_invdes.py b/tests/test_plugins/test_invdes.py index 59c20765b..07813555a 100644 --- a/tests/test_plugins/test_invdes.py +++ b/tests/test_plugins/test_invdes.py @@ -6,10 +6,9 @@ import tidy3d as td import tidy3d.plugins.invdes as tdi -from ..utils import AssertLogLevel, run_async_emulated, run_emulated - # use single threading pipeline -from .test_adjoint import use_emulated_run, use_emulated_run_async # noqa: F401 +from ..test_components.test_autograd import use_emulated_run # noqa: F401 +from ..utils import AssertLogLevel, run_emulated FREQ0 = 1e14 L_SIM = 1.0 @@ -46,19 +45,21 @@ @pytest.fixture -def use_emulated_run_autograd(monkeypatch): - """Emulate the tidy3d web function used in autograd web API.""" - import tidy3d.web.api.autograd.autograd as web_ag - - monkeypatch.setattr(web_ag, "_run_tidy3d", run_emulated) - - -@pytest.fixture -def use_emulated_run_autograd_async(monkeypatch): - """Emulate the tidy3d web function used in autograd web API.""" - import tidy3d.web.api.autograd.autograd as web_ag - - monkeypatch.setattr(web_ag, "_run_async_tidy3d", run_async_emulated) +def use_emulated_to_sim_data(monkeypatch): + """Emulate the InverseDesign.to_simulation_data to call emulated run.""" + monkeypatch.setattr( + tdi.InverseDesign, + "to_simulation_data", + lambda self, params, **kwargs: run_emulated(self.simulation, task_name="test"), + ) + monkeypatch.setattr( + tdi.InverseDesignMulti, + "to_simulation_data", + lambda self, params, **kwargs: { + task_name: run_emulated(sim, task_name=task_name) + for task_name, sim in zip(self.task_names, self.simulations) + }, + ) def make_design_region(): @@ -187,9 +188,11 @@ def __getitem__(self, name: str) -> MockDataArray: return MockDataArray() -def test_invdes_simulation_data(use_emulated_run): # noqa: F811 +def test_invdes_simulation_data(use_emulated_run, use_emulated_to_sim_data): # noqa: F811 """Test convenience function to convert ``InverseDesign`` to simulation and run it.""" + # monkeypatch.setattr(tdi.InverseDesign, "to_simulation_data", lambda self, params, **kwargs: run_emulated(self.simulation, task_name='test')) + invdes = make_invdes() params = invdes.design_region.params_random invdes.to_simulation_data(params=params, task_name="test") @@ -271,7 +274,7 @@ def make_optimizer(): ) -def make_result(use_emulated_run_autograd): +def make_result(use_emulated_run): # noqa: F811 """Test running the optimization defined in the ``InverseDesign`` object.""" optimizer = make_optimizer() @@ -281,7 +284,7 @@ def make_result(use_emulated_run_autograd): return optimizer.run(params0=PARAMS_0, post_process_fn=post_process_fn) -def test_default_params(use_emulated_run_autograd): # noqa: F811 +def test_default_params(use_emulated_run): # noqa: F811 """Test default paramns running the optimization defined in the ``InverseDesign`` object.""" optimizer = make_optimizer() @@ -291,7 +294,7 @@ def test_default_params(use_emulated_run_autograd): # noqa: F811 optimizer.run(post_process_fn=post_process_fn) -def test_warn_zero_grad(log_capture, use_emulated_run_autograd): +def test_warn_zero_grad(log_capture, use_emulated_run): # noqa: F811 """Test default paramns running the optimization defined in the ``InverseDesign`` object.""" optimizer = make_optimizer() @@ -301,7 +304,7 @@ def test_warn_zero_grad(log_capture, use_emulated_run_autograd): _ = optimizer.run(post_process_fn=post_process_fn_untraced) -def make_result_multi(use_emulated_run_autograd_async): +def make_result_multi(use_emulated_run): # noqa: F811 """Test running the optimization defined in the ``InverseDesignMulti`` object.""" optimizer = make_optimizer() @@ -314,7 +317,7 @@ def make_result_multi(use_emulated_run_autograd_async): return optimizer.run(params0=PARAMS_0, post_process_fn=post_process_fn_multi) -def test_result_store_full_results_is_false(use_emulated_run_autograd): # noqa: F811 +def test_result_store_full_results_is_false(use_emulated_run): # noqa: F811 """Test running the optimization defined in the ``InverseDesign`` object.""" optimizer = make_optimizer() @@ -337,9 +340,9 @@ def test_result_store_full_results_is_false(use_emulated_run_autograd): # noqa: _ = result.last["params"] -def test_continue_run_fns(use_emulated_run_autograd): +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_autograd) + result_orig = make_result(use_emulated_run) optimizer = make_optimizer() result_full = optimizer.continue_run(result=result_orig, post_process_fn=post_process_fn) @@ -350,9 +353,9 @@ def test_continue_run_fns(use_emulated_run_autograd): ), "wrong number of elements in the combined run history." -def test_continue_run_from_file(use_emulated_run_autograd): +def test_continue_run_from_file(use_emulated_run): # noqa: F811 """Test continuing an already run inverse design from file.""" - result_orig = make_result(use_emulated_run_autograd) + 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) @@ -368,13 +371,12 @@ def test_continue_run_from_file(use_emulated_run_autograd): def test_result( use_emulated_run, # noqa: F811 - use_emulated_run_autograd, - use_emulated_run_autograd_async, + use_emulated_to_sim_data, tmp_path, ): """Test methods of the ``InverseDesignResult`` object.""" - result = make_result(use_emulated_run_autograd) + result = make_result(use_emulated_run) with pytest.raises(KeyError): _ = result.get_last("blah") @@ -387,21 +389,20 @@ def test_result( _ = result.sim_data_last(task_name="last") -def test_result_data(use_emulated_run, use_emulated_run_autograd): # noqa: F811 +def test_result_data(use_emulated_run, use_emulated_to_sim_data): # noqa: F811 """Test methods of the ``InverseDesignResult`` object.""" - result = make_result(use_emulated_run_autograd) + result = make_result(use_emulated_run) _ = result.sim_last _ = result.sim_data_last(task_name="last") def test_result_data_multi( - use_emulated_run_async, # noqa: F811 - use_emulated_run_autograd_async, - use_emulated_run_autograd, + use_emulated_to_sim_data, # noqa: F811 + use_emulated_run, # noqa: F811 tmp_path, ): - result_multi = make_result_multi(use_emulated_run_autograd_async) + result_multi = make_result_multi(use_emulated_run) _ = result_multi.sim_last _ = result_multi.sim_data_last() @@ -413,10 +414,10 @@ def test_result_empty(): result_empty.get_last("params") -def test_invdes_io(tmp_path, log_capture, use_emulated_run_autograd): +def test_invdes_io(tmp_path, log_capture, use_emulated_run): # noqa: F811 """Test saving a loading ``invdes`` components to file.""" - result = make_result(use_emulated_run_autograd) + result = make_result(use_emulated_run) optimizer = make_optimizer() design = optimizer.design @@ -430,7 +431,7 @@ def test_invdes_io(tmp_path, log_capture, use_emulated_run_autograd): assert obj2.json() == obj.json() -def test_objective_utilities(use_emulated_run_autograd): +def test_objective_utilities(use_emulated_run): # noqa: F811 """Test objective function helpers.""" sim_data = run_emulated(simulation, task_name="test") diff --git a/tests/test_plugins/test_mode_solver.py b/tests/test_plugins/test_mode_solver.py index fa54e74ae..2481116f9 100644 --- a/tests/test_plugins/test_mode_solver.py +++ b/tests/test_plugins/test_mode_solver.py @@ -372,6 +372,39 @@ def test_mode_solver_simple(mock_remote_api, local): assert len(nS_add_mode_solver_monitor.monitors) == len(simulation.monitors) + 1 +@responses.activate +def test_mode_solver_remote_after_local(mock_remote_api): + """Test that running a remote solver after a local one modifies the stored data. This is to + catch a bug if ``_cached_properties["data"]`` is inadvertently used.""" + + simulation = td.Simulation( + size=SIM_SIZE, + grid_spec=td.GridSpec(wavelength=1.0), + structures=[WAVEGUIDE], + run_time=1e-12, + symmetry=(0, 0, 1), + boundary_spec=td.BoundarySpec.all_sides(boundary=td.Periodic()), + sources=[SRC], + ) + mode_spec = td.ModeSpec( + num_modes=3, + target_neff=2.0, + filter_pol="tm", + track_freq="lowest", + ) + + ms = ModeSolver( + simulation=simulation, + plane=PLANE, + mode_spec=mode_spec, + freqs=[td.C_0 / 1.0], + direction="-", + ) + data_local = ms.data + data_remote = msweb.run(ms) + assert np.all(data_local.n_eff != data_remote.n_eff) + + @pytest.mark.parametrize("local", [True, False]) @responses.activate def test_mode_solver_custom_medium(mock_remote_api, local, tmp_path): @@ -999,3 +1032,5 @@ def test_modes_eme_sim(mock_remote_api, local): with pytest.raises(SetupError): _ = msweb.run(solver) _ = msweb.run(solver.to_fdtd_mode_solver()) + + _ = solver.reduced_simulation_copy diff --git a/tests/test_web/test_webapi_eme.py b/tests/test_web/test_webapi_eme.py index 07c1a7d20..93fcf43dd 100644 --- a/tests/test_web/test_webapi_eme.py +++ b/tests/test_web/test_webapi_eme.py @@ -6,6 +6,7 @@ from botocore.exceptions import ClientError from responses import matchers from tidy3d import EMESimulation +from tidy3d.exceptions import SetupError from tidy3d.web.api.asynchronous import run_async from tidy3d.web.api.container import Batch, Job from tidy3d.web.api.webapi import ( @@ -16,6 +17,8 @@ get_info, get_run_info, load_simulation, + monitor, + real_cost, run, upload, ) @@ -238,6 +241,13 @@ def mock_webapi( """Mocks all webapi operation.""" +@responses.activate +def test_preupload_validation(mock_upload): + sim = make_eme_sim().updated_copy(size=(1000, 1000, 1000)) + with pytest.raises(SetupError): + upload(sim, TASK_NAME, PROJECT_NAME) + + @responses.activate def test_upload(mock_upload): sim = make_eme_sim() @@ -300,6 +310,17 @@ def test_run(mock_webapi, monkeypatch, tmp_path): ) +@responses.activate +def test_monitor(mock_get_info, mock_monitor): + monitor(TASK_ID, verbose=True) + monitor(TASK_ID, verbose=False) + + +@responses.activate +def test_real_cost(mock_get_info): + assert real_cost(TASK_ID) == FLEX_UNIT + + @responses.activate def test_abort_task(set_api_key, mock_get_info): responses.add( diff --git a/tests/utils.py b/tests/utils.py index 2fea27295..5090d6e81 100644 --- a/tests/utils.py +++ b/tests/utils.py @@ -16,6 +16,8 @@ """ utilities shared between all tests """ np.random.seed(4) +# function used to generate the data for emulated runs +DATA_GEN_FN = np.random.random FREQS = np.array([1.90, 2.01, 2.2]) * 1e12 SIM_MONITORS = td.Simulation( @@ -889,7 +891,7 @@ def make_data( """make a random DataArray out of supplied coordinates and data_type.""" data_shape = [len(coords[k]) for k in data_array_type._dims] np.random.seed(1) - data = np.random.random(data_shape) + data = DATA_GEN_FN(data_shape) 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 @@ -948,7 +950,7 @@ def make_mode_solver_data(monitor: td.ModeSolverMonitor) -> td.ModeSolverData: index_coords["mode_index"] = np.arange(monitor.mode_spec.num_modes) index_data_shape = (len(index_coords["f"]), len(index_coords["mode_index"])) index_data = ModeIndexDataArray( - (1 + 1j) * np.random.random(index_data_shape), coords=index_coords + (1 + 1j) * DATA_GEN_FN(index_data_shape), coords=index_coords ) for field_name in ["Ex", "Ey", "Ez", "Hx", "Hy", "Hz"]: coords = get_spatial_coords_dict(simulation, monitor, field_name) @@ -986,7 +988,7 @@ def make_diff_data(monitor: td.DiffractionMonitor) -> td.DiffractionData: orders_x = np.linspace(-1, 1, 3) orders_y = np.linspace(-2, 2, 5) coords = dict(orders_x=orders_x, orders_y=orders_y, f=f) - values = np.random.random((len(orders_x), len(orders_y), len(f))) + values = DATA_GEN_FN((len(orders_x), len(orders_y), len(f))) data = td.DiffractionDataArray(values, coords=coords) field_data = {field: data for field in ("Er", "Etheta", "Ephi", "Hr", "Htheta", "Hphi")} return td.DiffractionData(monitor=monitor, sim_size=(1, 1), bloch_vecs=(0, 0), **field_data) diff --git a/tidy3d/__init__.py b/tidy3d/__init__.py index bce015872..5d35d9e47 100644 --- a/tidy3d/__init__.py +++ b/tidy3d/__init__.py @@ -122,7 +122,13 @@ from .components.geometry.polyslab import PolySlab from .components.geometry.primitives import Cylinder, Sphere from .components.grid.grid import Coords, Coords1D, FieldGrid, Grid, YeeGrid -from .components.grid.grid_spec import AutoGrid, CustomGrid, GridSpec, UniformGrid +from .components.grid.grid_spec import ( + AutoGrid, + CustomGrid, + CustomGridBoundaries, + GridSpec, + UniformGrid, +) from .components.heat_charge.boundary import ( ConvectionBC, CurrentBC, @@ -301,6 +307,7 @@ def set_logging_level(level: str) -> None: "UniformGrid", "CustomGrid", "AutoGrid", + "CustomGridBoundaries", "Box", "Sphere", "Cylinder", diff --git a/tidy3d/components/apodization.py b/tidy3d/components/apodization.py index d4a848c27..9dd861409 100644 --- a/tidy3d/components/apodization.py +++ b/tidy3d/components/apodization.py @@ -15,7 +15,13 @@ class ApodizationSpec(Tidy3dBaseModel): Example ------- - >>> apod_spec = ApodizationSpec(start=1, end=2, width=0.5) + >>> apod_spec = ApodizationSpec(start=1, end=2, width=0.2) + + + .. image:: ../../_static/img/apodization.png + :width: 80% + :align: center + """ start: pd.NonNegativeFloat = pd.Field( @@ -35,7 +41,7 @@ class ApodizationSpec(Tidy3dBaseModel): width: pd.PositiveFloat = pd.Field( None, title="Apodization Width", - description="Characteristic decay length of the apodization function.", + description="Characteristic decay length of the apodization function, i.e., the width of the ramping up of the scaling function from 0 to 1.", units=SECOND, ) diff --git a/tidy3d/components/autograd/__init__.py b/tidy3d/components/autograd/__init__.py index 1e2da4aa4..98bcec266 100644 --- a/tidy3d/components/autograd/__init__.py +++ b/tidy3d/components/autograd/__init__.py @@ -11,7 +11,6 @@ from .utils import get_static __all__ = [ - "DerivativeInfo", "TracedFloat", "TracedSize1D", "TracedSize", @@ -20,6 +19,5 @@ "AutogradTraced", "AutogradFieldMap", "get_static", - "integrate_within_bounds", "interpn", ] diff --git a/tidy3d/components/autograd/utils.py b/tidy3d/components/autograd/utils.py index d2e0eb2e3..7b120b6fe 100644 --- a/tidy3d/components/autograd/utils.py +++ b/tidy3d/components/autograd/utils.py @@ -10,6 +10,13 @@ def get_static(x: typing.Any) -> typing.Any: return getval(x) +def split_list(x: list[typing.Any], index: int) -> (list[typing.Any], list[typing.Any]): + """Split a list at a given index.""" + x = list(x) + return x[:index], x[index:] + + __all__ = [ "get_static", + "split_list", ] diff --git a/tidy3d/components/base.py b/tidy3d/components/base.py index e08a247ee..86daf7f0d 100644 --- a/tidy3d/components/base.py +++ b/tidy3d/components/base.py @@ -1014,8 +1014,12 @@ def insert_value(x, path: tuple[str, ...], sub_dict: dict): sub_element = current_dict[final_key] if isinstance(sub_element, DataArray): current_dict[final_key] = sub_element.copy(deep=False, data=x) - if AUTOGRAD_KEY in sub_element.attrs: + 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/data/data_array.py b/tidy3d/components/data/data_array.py index 4c4a2b595..5fead6825 100644 --- a/tidy3d/components/data/data_array.py +++ b/tidy3d/components/data/data_array.py @@ -330,11 +330,13 @@ def interp( obj = self if assume_sorted else self.sortby(list(coords.keys())) + 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(coords.get(k, obj.coords[k]) for k in obj.dims) + xi = tuple(out_coords.values()) + vals = interpn(points, obj.tracers, xi, method=method) - da = DataArray(vals, dict(obj.coords) | coords) # tracers go into .attrs + 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 diff --git a/tidy3d/components/data/dataset.py b/tidy3d/components/data/dataset.py index 3864ce5bf..dcae12c7a 100644 --- a/tidy3d/components/data/dataset.py +++ b/tidy3d/components/data/dataset.py @@ -2898,7 +2898,7 @@ def _check_same_coordinates( # we can have xarray.DataArray's of different types but still same coordinates # we will deal with that case separately both_xarrays = isinstance(a, xr.DataArray) and isinstance(b, xr.DataArray) - if (not both_xarrays) and type(a) != type(b): + if (not both_xarrays) and type(a) is not type(b): return False if isinstance(a, UnstructuredGridDataset): diff --git a/tidy3d/components/data/monitor_data.py b/tidy3d/components/data/monitor_data.py index 54b523aeb..d707cc00e 100644 --- a/tidy3d/components/data/monitor_data.py +++ b/tidy3d/components/data/monitor_data.py @@ -129,7 +129,7 @@ def _updated(self, update: Dict) -> MonitorData: data_dict.update(update) return type(self).parse_obj(data_dict) - def make_adjoint_sources(self, dataset_names: list[str]) -> list[Source]: + def make_adjoint_sources(self, dataset_names: list[str], fwidth: float) -> list[Source]: """Generate adjoint sources for this ``MonitorData`` instance.""" # TODO: if there's data in the MonitorData, but no adjoint source, then @@ -1021,16 +1021,16 @@ def to_source( ) def make_adjoint_sources( - self, dataset_names: list[str] + self, dataset_names: list[str], fwidth: float ) -> List[Union[CustomCurrentSource, PointDipole]]: """Converts a :class:`.FieldData` to a list of adjoint current or point sources.""" if np.allclose(self.monitor.size, 0): - return self.to_adjoint_point_sources() + return self.to_adjoint_point_sources(fwidth=fwidth) - return self.to_adjoint_field_sources() + return self.to_adjoint_field_sources(fwidth=fwidth) - def to_adjoint_point_sources(self) -> List[PointDipole]: + def to_adjoint_point_sources(self, fwidth: float) -> List[PointDipole]: """Create adjoint point dipole source if this field data contains one item.""" sources = [] @@ -1053,7 +1053,7 @@ def to_adjoint_point_sources(self) -> List[PointDipole]: polarization=polarization, source_time=GaussianPulse( freq0=freq0, - fwidth=freq0 / 10, # TODO: how to set this properly? + fwidth=fwidth, amplitude=abs(adj_amp), phase=adj_phase, ), @@ -1064,7 +1064,7 @@ def to_adjoint_point_sources(self) -> List[PointDipole]: return sources - def to_adjoint_field_sources(self) -> List[CustomCurrentSource]: + def to_adjoint_field_sources(self, fwidth: float) -> List[CustomCurrentSource]: """Create adjoint custom field sources if this field data has some dimensionality.""" sources = [] @@ -1114,7 +1114,7 @@ def shift_value(coords) -> float: size=source_geo.size, source_time=GaussianPulse( freq0=freq0, - fwidth=freq0 / 10, # TODO: how to set this properly? + fwidth=fwidth, ), current_dataset=dataset, interpolate=True, @@ -1724,10 +1724,10 @@ def modes_info(self) -> xr.Dataset: "n eff": self.n_eff, "k eff": self.k_eff, "loss (dB/cm)": loss_db_cm, - f"TE (E{self._tangential_dims[0]}) fraction": self.pol_fraction["te"], - "wg TE fraction": self.pol_fraction_waveguide["te"], - "wg TM fraction": self.pol_fraction_waveguide["tm"], - "mode area": self.mode_area, + f"TE (E{self._tangential_dims[0]}) fraction": None, + "wg TE fraction": None, + "wg TM fraction": None, + "mode area": None, "group index": self.n_group_raw, # Use raw field to avoid issuing a warning "dispersion (ps/(nm km))": self.dispersion_raw, # Use raw field to avoid issuing a warning } @@ -1759,14 +1759,14 @@ def to_dataframe(self) -> DataFrame: return dataset.drop_vars(drop).to_dataframe() - def make_adjoint_sources(self, dataset_names: list[str]) -> list[ModeSource]: + def make_adjoint_sources(self, dataset_names: list[str], fwidth: float) -> list[ModeSource]: """Get all adjoint sources for the ``ModeMonitorData``.""" adjoint_sources = [] for name in dataset_names: if name == "amps": - adjoint_sources += self.make_adjoint_sources_amps() + adjoint_sources += self.make_adjoint_sources_amps(fwidth=fwidth) else: log.warning( f"Can't create adjoint source for 'ModeData.{type(self)}.{name}'. " @@ -1777,7 +1777,7 @@ def make_adjoint_sources(self, dataset_names: list[str]) -> list[ModeSource]: return adjoint_sources - def make_adjoint_sources_amps(self) -> list[ModeSource]: + def make_adjoint_sources_amps(self, fwidth: float) -> list[ModeSource]: """Generate adjoint sources for ``ModeMonitorData.amps``.""" coords = self.amps.coords @@ -1793,12 +1793,12 @@ def make_adjoint_sources_amps(self) -> list[ModeSource]: if self.get_amplitude(amp_single) == 0.0: continue - adjoint_source = self.adjoint_source_amp(amp=amp_single) + adjoint_source = self.adjoint_source_amp(amp=amp_single, fwidth=fwidth) adjoint_sources.append(adjoint_source) return adjoint_sources - def adjoint_source_amp(self, amp: DataArray) -> ModeSource: + def adjoint_source_amp(self, amp: DataArray, fwidth: float) -> ModeSource: """Generate an adjoint ``ModeSource`` for a single amplitude.""" monitor = self.monitor @@ -1821,7 +1821,7 @@ def adjoint_source_amp(self, amp: DataArray) -> ModeSource: amplitude=abs(src_amp), phase=np.angle(src_amp), freq0=freq0, - fwidth=freq0 / 10, # TODO: how to set this properly? + fwidth=fwidth, ), mode_spec=monitor.mode_spec, size=monitor.size, @@ -2202,7 +2202,7 @@ def propagation_factor(dist: Union[float, None], k: complex, is_2d_simulation: b return 1.0 if is_2d_simulation: - return -np.exp(1j * k * dist) * np.sqrt(1j * k / (8 * np.pi * dist)) + return np.exp(1j * k * dist) * np.sqrt(-1j * k / (8 * np.pi * dist)) return -1j * k * np.exp(1j * k * dist) / (4 * np.pi * dist) @@ -2506,6 +2506,35 @@ def z(self) -> np.ndarray: """Z positions.""" return self.Etheta.z.values + @property + def tangential_dims(self): + tangential_dims = ["x", "y", "z"] + tangential_dims.pop(self.monitor.proj_axis) + return tangential_dims + + @property + def poynting(self) -> xr.DataArray: + """Time-averaged Poynting vector for field data associated to a Cartesian field projection monitor.""" + fc = self.fields_cartesian + dim1, dim2 = self.tangential_dims + + e1 = fc["E" + dim1] + e2 = fc["E" + dim2] + h1 = fc["H" + dim1] + h2 = fc["H" + dim2] + + e1_h2 = e1 * h2.conj() + e2_h1 = e2 * h1.conj() + + e_x_h_star = e1_h2 - e2_h1 + return 0.5 * np.real(e_x_h_star) + + @cached_property + def flux(self) -> FluxDataArray: + """Flux for projecteded field data corresponding to a Cartesian field projection monitor.""" + flux = self.poynting.integrate(self.tangential_dims) + return FluxDataArray(flux) + def renormalize_fields(self, proj_distance: float) -> FieldProjectionCartesianData: """Return a :class:`.FieldProjectionCartesianData` with fields re-normalized to a new projection distance, by applying a phase factor based on ``proj_distance``. @@ -2920,13 +2949,13 @@ def _make_dataset(self, fields: Tuple[np.ndarray, ...], keys: Tuple[str, ...]) - """ Autograd code """ - def make_adjoint_sources(self, dataset_names: list[str]) -> list[PlaneWave]: + def make_adjoint_sources(self, dataset_names: list[str], fwidth: float) -> list[PlaneWave]: """Get all adjoint sources for the ``DiffractionMonitor.amps``.""" # NOTE: everything just goes through `.amps`, any post-processing is encoded in E-fields - return self.make_adjoint_sources_amps() + return self.make_adjoint_sources_amps(fwidth=fwidth) - def make_adjoint_sources_amps(self) -> list[PlaneWave]: + def make_adjoint_sources_amps(self, fwidth: float) -> list[PlaneWave]: """Make adjoint sources for outputs that depend on DiffractionData.`amps`.""" amps = self.amps @@ -2953,13 +2982,13 @@ def make_adjoint_sources_amps(self) -> list[PlaneWave]: continue # compute a plane wave for this amplitude (if propagating / not None) - adjoint_source = self.adjoint_source_amp(amp=amp_single) + adjoint_source = self.adjoint_source_amp(amp=amp_single, fwidth=fwidth) if adjoint_source is not None: adjoint_sources.append(adjoint_source) return adjoint_sources - def adjoint_source_amp(self, amp: DataArray) -> PlaneWave: + def adjoint_source_amp(self, amp: DataArray, fwidth: float) -> PlaneWave: """Generate an adjoint ``PlaneWave`` for a single amplitude.""" monitor = self.monitor @@ -3003,7 +3032,7 @@ def adjoint_source_amp(self, amp: DataArray) -> PlaneWave: amplitude=abs(src_amp), phase=np.angle(src_amp), freq0=freq0, - fwidth=freq0 / 10, # TODO: how to set this properly? + fwidth=fwidth, ), direction=self.flip_direction(monitor.normal_dir), angle_theta=angle_theta, diff --git a/tidy3d/components/data/sim_data.py b/tidy3d/components/data/sim_data.py index 7a4696d30..8e90f879e 100644 --- a/tidy3d/components/data/sim_data.py +++ b/tidy3d/components/data/sim_data.py @@ -13,17 +13,20 @@ import pydantic.v1 as pd import xarray as xr +from ...constants import C_0, inf from ...exceptions import DataError, FileError, Tidy3dKeyError from ...log import log -from ..base import JSON_TAG +from ..autograd.utils import split_list +from ..base import JSON_TAG, Tidy3dBaseModel from ..base_sim.data.sim_data import AbstractSimulationData from ..file_util import replace_values from ..monitor import Monitor from ..simulation import Simulation -from ..source import Source +from ..source import GaussianPulse, SourceType from ..structure import Structure from ..types import Ax, Axis, ColormapType, FieldVal, PlotScale, annotate_type from ..viz import add_ax_if_none, equal_aspect +from .data_array import FreqDataArray from .monitor_data import ( AbstractFieldData, FieldTimeData, @@ -36,6 +39,33 @@ # maps monitor type (string) to the class of the corresponding data DATA_TYPE_NAME_MAP = {val.__fields__["monitor"].type_.__name__: val for val in MonitorDataTypes} +# residuals below this are considered good fits for broadband adjoint source creation +RESIDUAL_CUTOFF_ADJOINT = 1e-6 + + +class AdjointSourceInfo(Tidy3dBaseModel): + """Stores information about the adjoint sources to pass to autograd pipeline.""" + + sources: tuple[SourceType, ...] = pd.Field( + ..., + title="Adjoint Sources", + description="Set of processed sources to include in the adjoint simulation.", + ) + + post_norm: Union[float, FreqDataArray] = pd.Field( + ..., + title="Post Normalization Values", + description="Factor to multiply the adjoint fields by after running " + "given the adjoint source pipeline used.", + ) + + normalize_sim: bool = pd.Field( + ..., + title="Normalize Adjoint Simulation", + description="Whether the adjoint simulation needs to be normalized " + "given the adjoint source pipeline used.", + ) + class AbstractYeeGridSimulationData(AbstractSimulationData, ABC): """Data from an :class:`.AbstractYeeGridSimulation` involving @@ -951,26 +981,73 @@ def source_spectrum_fn(freqs): return self.copy(update=dict(simulation=simulation, data=data_normalized)) + def split_adjoint_data(self: SimulationData, num_mnts_original: int) -> tuple[list, list]: + """Split data list into original, adjoint field, and adjoint permittivity.""" + + data_all = list(self.data) + num_mnts_adjoint = (len(data_all) - num_mnts_original) // 2 + + log.info( + f" -> {num_mnts_original} monitors, {num_mnts_adjoint} adjoint field monitors, {num_mnts_adjoint} adjoint eps monitors." + ) + + data_original, data_adjoint = split_list(data_all, index=num_mnts_original) + + return data_original, data_adjoint + + def split_original_fwd(self, num_mnts_original: int) -> Tuple[SimulationData, SimulationData]: + """Split this simulation data into original and fwd data from number of original mnts.""" + + # split the data and monitors into the original ones & adjoint gradient ones (for 'fwd') + data_original, data_fwd = self.split_adjoint_data(num_mnts_original=num_mnts_original) + monitors_orig, monitors_fwd = split_list(self.simulation.monitors, index=num_mnts_original) + + # reconstruct the simulation data for the user, using original sim, and data for original mnts + sim_original = self.simulation.updated_copy(monitors=monitors_orig) + sim_data_original = self.updated_copy( + simulation=sim_original, + data=data_original, + deep=False, + ) + + # construct the 'forward' simulation and its data, which is only used for for gradient calc. + sim_fwd = self.simulation.updated_copy(monitors=monitors_fwd) + sim_data_fwd = self.updated_copy( + simulation=sim_fwd, + data=data_fwd, + deep=False, + ) + + return sim_data_original, sim_data_fwd + def make_adjoint_sim( self, data_vjp_paths: set[tuple], adjoint_monitors: list[Monitor] - ) -> Simulation: + ) -> tuple[Simulation, AdjointSourceInfo]: """Make the adjoint simulation from the original simulation and the VJP-containing data.""" sim_original = self.simulation - # generate the adjoint sources - sources_adj = self.make_adjoint_sources(data_vjp_paths=data_vjp_paths) + # generate the adjoint sources {mnt_name : list[Source]} + sources_adj_dict = self.make_adjoint_sources(data_vjp_paths=data_vjp_paths) + adj_srcs = [] + for src_list in sources_adj_dict.values(): + adj_srcs += list(src_list) + + adjoint_source_info = self.process_adjoint_sources(adj_srcs=adj_srcs) # grab boundary conditions with flipped Bloch vectors (for adjoint) bc_adj = sim_original.boundary_spec.flipped_bloch_vecs # fields to update the 'fwd' simulation with to make it 'adj' sim_adj_update_dict = dict( - sources=sources_adj, + sources=adjoint_source_info.sources, boundary_spec=bc_adj, monitors=adjoint_monitors, ) + if not adjoint_source_info.normalize_sim: + sim_adj_update_dict["normalize_index"] = None + # set the ADJ grid spec wavelength to the original wavelength (for same meshing) grid_spec_original = sim_original.grid_spec if sim_original.sources and grid_spec_original.wavelength is None: @@ -978,27 +1055,222 @@ 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) + return sim_original.updated_copy(**sim_adj_update_dict), adjoint_source_info - def make_adjoint_sources(self, data_vjp_paths: set[tuple]) -> list[Source]: + 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.""" - # TODO: determine if we can do multi-frequency sources - # map of index into 'self.data' to the list of datasets we need adjoint sources for adj_src_map = defaultdict(list) for _, index, dataset_name in data_vjp_paths: adj_src_map[index].append(dataset_name) - # gather a list of adjoint sources for every monitor data in the VJP that needs one - sources_adj_all = [] + # gather a dict of adjoint sources for every monitor data in the VJP that needs one + sources_adj_all = defaultdict(list) for data_index, dataset_names in adj_src_map.items(): mnt_data = self.data[data_index] - sources_adj = mnt_data.make_adjoint_sources(dataset_names=dataset_names) - sources_adj_all += sources_adj + sources_adj = mnt_data.make_adjoint_sources( + dataset_names=dataset_names, fwidth=self.fwidth_adj + ) + sources_adj_all[mnt_data.monitor.name] = sources_adj return sources_adj_all + @property + def fwidth_adj(self) -> float: + # fwidth of forward pass, try as default for adjoint + normalize_index_fwd = self.simulation.normalize_index or 0 + return self.simulation.sources[normalize_index_fwd].source_time.fwidth + + def process_adjoint_sources(self, adj_srcs: list[SourceType]) -> AdjointSourceInfo: + """Compute list of final sources along with a post run normalization for adj fields.""" + + # dictionary mapping hash of sources with same freq dependence to list of time-dependencies + hashes_to_sources = defaultdict(None) + hashes_to_src_times = defaultdict(list) + + tmp_src_time = GaussianPulse(freq0=C_0, fwidth=inf) + for src in adj_srcs: + tmp_src = src.updated_copy(source_time=tmp_src_time) + tmp_src_hash = tmp_src._hash_self() + hashes_to_sources[tmp_src_hash] = src + hashes_to_src_times[tmp_src_hash].append(src.source_time) + + num_ports = len(hashes_to_src_times) + num_unique_freqs = len({src.source_time.freq0 for src in adj_srcs}) + + # next, figure out which treatment / normalization to apply + if num_unique_freqs == 1: + log.info("Adjoint source creation: one unique frequency, no normalization.") + return AdjointSourceInfo(sources=adj_srcs, post_norm=1.0, normalize_sim=True) + + if num_ports == 1 and len(adj_srcs) == num_unique_freqs: + log.info("Adjoint source creation: one spatial port detected.") + adj_srcs, post_norm = self.process_adjoint_sources_broadband(adj_srcs) + return AdjointSourceInfo(sources=adj_srcs, post_norm=post_norm, normalize_sim=True) + + # if several spatial ports and several frequencies, try to fit + log.info("Adjoint source creation: trying multifrequency fit.") + adj_srcs, post_norm = self.process_adjoint_sources_fit( + adj_srcs=adj_srcs, + hashes_to_src_times=hashes_to_src_times, + hashes_to_sources=hashes_to_sources, + ) + return AdjointSourceInfo(sources=adj_srcs, post_norm=post_norm, normalize_sim=False) + + """ SIMPLE APPROACH """ + + def process_adjoint_sources_broadband( + self, adj_srcs: list[SourceType] + ) -> tuple[list[SourceType], xr.DataArray]: + """Process adjoint sources for the case of several sources at the same freq.""" + + src_broadband = self._make_broadband_source(adj_srcs=adj_srcs) + post_norm_amps = self._make_post_norm_amps(adj_srcs=adj_srcs) + + log.info( + "Several adjoint sources, from one monitor. " + "Only difference between them is the source time. " + "Constructing broadband adjoint source and performing post-run normalization " + f"of fields with {len(post_norm_amps)} frequencies." + ) + + return [src_broadband], post_norm_amps + + def _make_broadband_source(self, adj_srcs: list[SourceType]) -> SourceType: + """Make a broadband source for a set of adjoint sources.""" + + source_index = self.simulation.normalize_index or 0 + src_time_base = self.simulation.sources[source_index].source_time.copy() + src_broadband = adj_srcs[0].updated_copy(source_time=src_time_base) + + return src_broadband + + @staticmethod + def _make_post_norm_amps(adj_srcs: list[SourceType]) -> xr.DataArray: + """Make a ``DataArray`` containing the complex amplitudes to multiply with adjoint field.""" + + freqs = [] + amps_complex = [] + for src in adj_srcs: + src_time = src.source_time + freqs.append(src_time.freq0) + amp_complex = src_time.amplitude * np.exp(1j * src_time.phase) + amps_complex.append(amp_complex) + + coords = dict(f=freqs) + amps_complex = np.array(amps_complex) + return xr.DataArray(amps_complex, coords=coords) + + """ FITTING APPROACH """ + + def process_adjoint_sources_fit( + self, + adj_srcs: list[SourceType], + hashes_to_src_times: dict[str, GaussianPulse], + hashes_to_sources: dict[str, list[SourceType]], + ) -> tuple[list[SourceType], float]: + """Process the adjoint sources using a least squared fit to the derivative data.""" + + raise NotImplementedError( + "Can't perform multi-frequency autograd with several adjoint sources yet. " + "In the meantime, please construct a single 'Simulation' per output data " + "(can be multi-frequency) and run in parallel using 'web.run_async'. For example, " + "if your problem has 'P' outuput ports, e.g. waveguides, please make a 'Simulation' " + "corresponding to the objective function contribution at each port." + ) + + # new adjoint sources + new_adj_srcs = [] + for src_hash, source_times in hashes_to_src_times.items(): + src = hashes_to_sources[src_hash] + new_sources = self.correct_adjoint_sources( + src=src, fwidth=self.fwidth_adj, source_times=source_times + ) + new_adj_srcs += new_sources + + # compute amplitudes of each adjoint source, and the norm + adj_src_amps = [] + for src in new_adj_srcs: + amp = src.source_time.amp_complex + adj_src_amps.append(amp) + norm_amps = np.linalg.norm(adj_src_amps) + + # normalize all of the adjoint sources by this and return the normalization term used + adj_srcs_norm = [] + for src in new_adj_srcs: + src_time = src.source_time + amp = src_time.amp_complex + src_time_norm = src_time.from_amp_complex(amp=amp / norm_amps) + src_nrm = src.updated_copy(source_time=src_time_norm) + adj_srcs_norm.append(src_nrm) + + return adj_srcs_norm, norm_amps + + def correct_adjoint_sources( + self, src: SourceType, fwidth: float, source_times: list[GaussianPulse] + ) -> [SourceType]: + """Corret a set of spectrally overlapping adjoint sources to give correct E_adj.""" + + freqs = [st.freq0 for st in source_times] + times = self.simulation.tmesh + dt = self.simulation.dt + + def get_spectrum(source_time: GaussianPulse, freqs: list[float]) -> complex: + """Get the spectrum of a source time at a given frequency.""" + return source_time.spectrum(times=times, freqs=freqs, dt=dt) + + # compute matrix coupling the spectra of Gaussian pulses centered at each adjoint freq + def get_coupling_matrix(fwidth: float) -> np.ndarray: + """Matrix coupling the spectra of Gaussian pulses centered at each adjoint freq.""" + + return np.array( + [ + get_spectrum( + source_time=GaussianPulse(freq0=source_time.freq0, fwidth=fwidth), + freqs=freqs, + ) + for source_time in source_times + ] + ).T + + amps_adj = np.array([src_time.amp_complex for src_time in source_times]) + + # compute the corrected set of amps to inject at each freq to take coupling into account + def get_amps_corrected(fwidth: float) -> tuple[np.ndarray, float]: + """New set of new adjoint source amps that generate the desired response at each f.""" + J_coupling = get_coupling_matrix(fwidth=fwidth) + + amps_adj_new, *info = np.linalg.lstsq(J_coupling, amps_adj, rcond=None) + # amps_adj_new = np.linalg.solve(J_coupling, amps_adj) + residual = J_coupling @ amps_adj_new - amps_adj + residual_norm = np.linalg.norm(residual) / np.linalg.norm(amps_adj) + return amps_adj_new, residual_norm + + # get the corrected amplitudes + amps_corrected, res_norm = get_amps_corrected(self.fwidth_adj) + + if res_norm > RESIDUAL_CUTOFF_ADJOINT: + raise ValueError( + f"Residual of {res_norm:.5e} found when trying to fit adjoint source spectrum. " + f"This is above our accuracy cutoff of {RESIDUAL_CUTOFF_ADJOINT:.5e} and therefore " + "we are not able to process this adjoint simulation in a broadband way. " + "To fix, split your simulation into a set of simulations, one for each port, and " + "run parallel, broadband simulations using 'web.run_async'. " + ) + + # construct the new adjoint sources with the corrected amplitudes + src_times_corrected = [ + src_time.from_amp_complex(amp=amp, fwidth=self.fwidth_adj) + for src_time, amp in zip(source_times, amps_corrected) + ] + srcs_corrected = [] + for src_time in src_times_corrected: + src_new = src.updated_copy(source_time=src_time) + srcs_corrected.append(src_new) + + return srcs_corrected + def get_adjoint_data(self, structure_index: int, data_type: str) -> MonitorDataType: """Grab the field or permittivity data for a given structure index.""" diff --git a/tidy3d/components/eme/grid.py b/tidy3d/components/eme/grid.py index 1eca48a42..7d63978dc 100644 --- a/tidy3d/components/eme/grid.py +++ b/tidy3d/components/eme/grid.py @@ -210,11 +210,12 @@ def _validate_boundaries(cls, val, values): raise ValidationError( "There must be exactly one fewer item in 'boundaries' than " "in 'mode_specs'." ) - rmin = boundaries[0] - for rmax in boundaries[1:]: - if rmax < rmin: - raise ValidationError("The 'boundaries' must be increasing.") - rmin = rmax + if len(boundaries) > 0: + rmin = boundaries[0] + for rmax in boundaries[1:]: + if rmax < rmin: + raise ValidationError("The 'boundaries' must be increasing.") + rmin = rmax return val def make_grid(self, center: Coordinate, size: Size, axis: Axis) -> EMEGrid: @@ -237,12 +238,15 @@ def make_grid(self, center: Coordinate, size: Size, axis: Axis) -> EMEGrid: """ sim_rmin = center[axis] - size[axis] / 2 sim_rmax = center[axis] + size[axis] / 2 - if self.boundaries[0] < sim_rmin - fp_eps: - raise ValidationError( - "The first item in 'boundaries' is outside the simulation domain." - ) - if self.boundaries[-1] > sim_rmax + fp_eps: - raise ValidationError("The last item in 'boundaries' is outside the simulation domain.") + if len(self.boundaries) > 0: + if self.boundaries[0] < sim_rmin - fp_eps: + raise ValidationError( + "The first item in 'boundaries' is outside the simulation domain." + ) + if self.boundaries[-1] > sim_rmax + fp_eps: + raise ValidationError( + "The last item in 'boundaries' is outside the simulation domain." + ) boundaries = [sim_rmin] + list(self.boundaries) + [sim_rmax] return EMEGrid( diff --git a/tidy3d/components/eme/monitor.py b/tidy3d/components/eme/monitor.py index af1b390db..8e8dc5176 100644 --- a/tidy3d/components/eme/monitor.py +++ b/tidy3d/components/eme/monitor.py @@ -105,6 +105,14 @@ class EMEModeSolverMonitor(EMEMonitor): """EME mode solver monitor. Records EME modes computed in planes intersecting the monitor geometry. + Note + ---- + + This is different than a :class:`.ModeSolverMonitor`, which computes modes within + its planar geometry. In contrast, this monitor does not compute new modes; instead, + it records the modes used for EME expansion and propagation, but only within the + monitor geometry. + Example ------- >>> monitor = EMEModeSolverMonitor( diff --git a/tidy3d/components/eme/simulation.py b/tidy3d/components/eme/simulation.py index 84590d7fa..8a3ca8d94 100644 --- a/tidy3d/components/eme/simulation.py +++ b/tidy3d/components/eme/simulation.py @@ -2,28 +2,29 @@ from __future__ import annotations -from typing import Dict, List, Literal, Optional, Tuple +from typing import Dict, List, Literal, Optional, Tuple, Union import matplotlib as mpl import numpy as np import pydantic.v1 as pd -from ...exceptions import SetupError +from ...exceptions import SetupError, ValidationError from ...log import log from ..base import cached_property from ..boundary import BoundarySpec, PECBoundary +from ..geometry.base import Box from ..grid.grid import Grid from ..grid.grid_spec import GridSpec from ..medium import FullyAnisotropicMedium -from ..monitor import AbstractModeMonitor, ModeSolverMonitor, Monitor +from ..monitor import AbstractModeMonitor, ModeSolverMonitor, Monitor, MonitorType from ..scene import Scene from ..simulation import AbstractYeeGridSimulation, Simulation -from ..source import GaussianPulse, ModeSource +from ..source import GaussianPulse, PointDipole from ..structure import Structure -from ..types import Ax, Axis, FreqArray, annotate_type +from ..types import Ax, Axis, FreqArray, Symmetry, annotate_type from ..validators import MIN_FREQUENCY, validate_freqs_min, validate_freqs_not_empty from ..viz import add_ax_if_none, equal_aspect -from .grid import EMECompositeGrid, EMEGrid, EMEGridSpec, EMEGridSpecType +from .grid import EMECompositeGrid, EMEExplicitGrid, EMEGrid, EMEGridSpec, EMEGridSpecType from .monitor import EMEFieldMonitor, EMEModeSolverMonitor, EMEMonitor, EMEMonitorType from .sweep import EMEFreqSweep, EMELengthSweep, EMEModeSweep, EMESweepSpecType @@ -231,6 +232,16 @@ class EMESimulation(AbstractYeeGridSimulation): _freqs_not_empty = validate_freqs_not_empty() _freqs_lower_bound = validate_freqs_min() + @pd.validator("size", always=True) + def _validate_fully_3d(cls, val): + """An EME simulation must be fully 3D.""" + if val.count(0.0) != 0: + raise ValidationError( + "'EMESimulation' cannot have any component of 'size' equal to " + f"zero, given 'size={val}'." + ) + return val + @pd.validator("grid_spec", always=True) def _validate_auto_grid_wavelength(cls, val, values): """Handle the case where grid_spec is auto and wavelength is not provided.""" @@ -547,15 +558,19 @@ def _post_init_validators(self) -> None: _ = self.grid _ = self.eme_grid _ = self.mode_solver_monitors - log.begin_capture() self._validate_too_close_to_edges() self._validate_sweep_spec() + self._validate_symmetry() self._validate_monitor_setup() self._validate_sources_and_boundary() + + def validate_pre_upload(self) -> None: + """Validate the fully initialized EME simulation is ok for upload to our servers.""" + log.begin_capture() + self._validate_sweep_spec_size() self._validate_size() self._validate_monitor_size() self._validate_modes_size() - self._validate_symmetry() self._validate_constraint() # self._warn_monitor_interval() log.end_capture(self) @@ -606,7 +621,7 @@ def _validate_port_offsets(self): size = self.size axis = self.axis if size[axis] < total_offset: - raise SetupError( + raise ValidationError( "The sum of the two 'port_offset' fields " "cannot exceed the simulation 'size' in the 'axis' direction." ) @@ -628,19 +643,24 @@ def _validate_symmetry(self): # "it always monitors every EME cell." # ) - def _validate_sweep_spec(self): - """Validate sweep spec.""" + def _validate_sweep_spec_size(self): + """Make sure sweep spec is not too large.""" if self.sweep_spec is None: return - # mode sweep can't exceed max num modes num_sweep = self.sweep_spec.num_sweep - if num_sweep == 0: - raise SetupError("Simulation 'sweep_spec' has 'num_sweep=0'.") if num_sweep > MAX_NUM_SWEEP: raise SetupError( f"Simulation 'sweep_spec' has 'num_sweep={num_sweep}, " f"which exceeds the maximum allowed '{MAX_NUM_SWEEP}'." ) + + def _validate_sweep_spec(self): + """Validate sweep spec.""" + if self.sweep_spec is None: + return + num_sweep = self.sweep_spec.num_sweep + if num_sweep == 0: + raise SetupError("Simulation 'sweep_spec' has 'num_sweep=0'.") if isinstance(self.sweep_spec, EMEModeSweep): if any(self.sweep_spec.num_modes > self.max_num_modes): raise SetupError( @@ -652,6 +672,11 @@ def _validate_sweep_spec(self): ) elif isinstance(self.sweep_spec, EMELengthSweep): scale_factors_shape = self.sweep_spec.scale_factors.shape + if len(scale_factors_shape) > 2: + raise SetupError( + "Simulation 'sweep_spec.scale_factors' must " + "have either one or two dimensions." + ) if len(scale_factors_shape) == 2: num_scale_factors = scale_factors_shape[1] if num_scale_factors != self.eme_grid.num_cells: @@ -674,6 +699,8 @@ def _validate_sweep_spec(self): def _validate_monitor_setup(self): """Check monitor setup.""" for monitor in self.monitors: + if isinstance(monitor, EMEMonitor): + _ = self._monitor_eme_cell_indices(monitor=monitor) if ( hasattr(monitor, "freqs") and monitor.freqs is not None @@ -871,6 +898,41 @@ def _sweep_modes(self) -> bool: """Whether the sweep changes the modes.""" return self.sweep_spec is not None and isinstance(self.sweep_spec, EMEFreqSweep) + @property + def _num_sweep_modes(self) -> pd.PositiveInt: + """Number of sweep indices for modes.""" + if self._sweep_modes: + return self._num_sweep + return 1 + + @property + def _sweep_interfaces(self) -> bool: + """Whether the sweep changes the cell interface scattering matrices.""" + return self.sweep_spec is not None and isinstance( + self.sweep_spec, (EMEFreqSweep, EMEModeSweep) + ) + + @property + def _num_sweep_interfaces(self) -> pd.PositiveInt: + """Number of sweep indices for interfaces.""" + if self._sweep_interfaces: + return self._num_sweep + return 1 + + @property + def _sweep_cells(self) -> bool: + """Whether the sweep changes the propagation within a cell.""" + return self.sweep_spec is not None and isinstance( + self.sweep_spec, (EMELengthSweep, EMEFreqSweep, EMEModeSweep) + ) + + @property + def _num_sweep_cells(self) -> pd.PositiveInt: + """Number of sweep indices for cells.""" + if self._sweep_cells: + return self._num_sweep + return 1 + def _monitor_num_sweep(self, monitor: EMEMonitor) -> pd.PositiveInt: """Number of sweep indices for a certain monitor.""" if self.sweep_spec is None: @@ -965,12 +1027,10 @@ def grid(self) -> Grid: ) plane = self.eme_grid.mode_planes[0] sources.append( - ModeSource( + PointDipole( center=plane.center, - size=plane.size, source_time=GaussianPulse(freq0=freqs[0], fwidth=0.1 * freqs[0]), - direction="+", - mode_spec=self.eme_grid.mode_specs[0], + polarization="Ez", ) ) @@ -1010,12 +1070,10 @@ def _to_fdtd_sim(self) -> Simulation: plane = self.eme_grid.mode_planes[0] freq0 = self.freqs[0] source_time = GaussianPulse(freq0=freq0, fwidth=0.1 * freq0) - source = ModeSource( + source = PointDipole( center=plane.center, - size=plane.size, source_time=source_time, - direction="+", - mode_spec=self.eme_grid.mode_specs[0]._to_mode_spec(), + polarization="Ez", ) # copy over all FDTD monitors too monitors = [monitor for monitor in self.monitors if not isinstance(monitor, EMEMonitor)] @@ -1033,3 +1091,83 @@ def _to_fdtd_sim(self) -> Simulation: sources=[source], monitors=monitors, ) + + def subsection( + self, + region: Box, + grid_spec: Union[GridSpec, Literal["identical"]] = None, + eme_grid_spec: Union[EMEGridSpec, Literal["identical"]] = None, + symmetry: Tuple[Symmetry, Symmetry, Symmetry] = None, + monitors: Tuple[MonitorType, ...] = None, + remove_outside_structures: bool = True, + remove_outside_custom_mediums: bool = False, + **kwargs, + ) -> EMESimulation: + """Generate a simulation instance containing only the ``region``. + Same as in :class:`.AbstractYeeGridSimulation`, except also restricting EME grid. + + Parameters + ---------- + region : :class:.`Box` + New simulation domain. + grid_spec : :class:.`GridSpec` = None + New grid specification. If ``None``, then it is inherited from the original + simulation. If ``identical``, then the original grid is transferred directly as a + :class:.`CustomGrid`. Note that in the latter case the region of the new simulation is + snapped to the original grid lines. + eme_grid_spec: :class:`.EMEGridSpec` = None + New EME grid specification. If ``None``, then it is inherited from the original + simulation. If ``identical``, then the original grid is transferred directly as a + :class:`.EMEExplicitGrid`. Noe that in the latter case the region of the new simulation + is expanded to contain full EME cells. + symmetry : Tuple[Literal[0, -1, 1], Literal[0, -1, 1], Literal[0, -1, 1]] = None + New simulation symmetry. If ``None``, then it is inherited from the original + simulation. Note that in this case the size and placement of new simulation domain + must be commensurate with the original symmetry. + monitors : Tuple[MonitorType, ...] = None + New list of monitors. If ``None``, then the monitors intersecting the new simulation + domain are inherited from the original simulation. + remove_outside_structures : bool = True + Remove structures outside of the new simulation domain. + remove_outside_custom_mediums : bool = True + Remove custom medium data outside of the new simulation domain. + **kwargs + Other arguments passed to new simulation instance. + """ + + new_region = region + if eme_grid_spec is None: + eme_grid_spec = self.eme_grid_spec + elif isinstance(eme_grid_spec, str) and eme_grid_spec == "identical": + axis = self.axis + mode_specs = self.eme_grid.mode_specs + boundaries = self.eme_grid.boundaries + indices = self.eme_grid.cell_indices_in_box(box=region) + + new_boundaries = boundaries[indices[0] : indices[-1] + 2] + new_mode_specs = mode_specs[indices[0] : indices[-1] + 1] + + rmin = list(region.bounds[0]) + rmax = list(region.bounds[1]) + rmin[axis] = min(rmin[axis], new_boundaries[0]) + rmax[axis] = max(rmax[axis], new_boundaries[-1]) + new_region = Box.from_bounds(rmin=rmin, rmax=rmax) + + # remove outer boundaries for explicit grid + new_boundaries = new_boundaries[1:-1] + + eme_grid_spec = EMEExplicitGrid(mode_specs=new_mode_specs, boundaries=new_boundaries) + + new_sim = super().subsection( + region=new_region, + grid_spec=grid_spec, + symmetry=symmetry, + monitors=monitors, + remove_outside_structures=remove_outside_structures, + remove_outside_custom_mediums=remove_outside_custom_mediums, + **kwargs, + ) + + new_sim = new_sim.updated_copy(eme_grid_spec=eme_grid_spec) + + return new_sim diff --git a/tidy3d/components/eme/sweep.py b/tidy3d/components/eme/sweep.py index 669668674..92e83556e 100644 --- a/tidy3d/components/eme/sweep.py +++ b/tidy3d/components/eme/sweep.py @@ -8,7 +8,7 @@ import pydantic.v1 as pd from ..base import Tidy3dBaseModel -from ..types import ArrayFloat1D, ArrayFloat2D, ArrayInt1D +from ..types import ArrayFloat1D, ArrayInt1D, ArrayLike class EMESweepSpec(Tidy3dBaseModel, ABC): @@ -23,7 +23,7 @@ def num_sweep(self) -> pd.PositiveInt: class EMELengthSweep(EMESweepSpec): """Spec for sweeping EME cell lengths.""" - scale_factors: Union[ArrayFloat1D, ArrayFloat2D] = pd.Field( + scale_factors: ArrayLike = pd.Field( ..., title="Length Scale Factor", description="Length scale factors to be used in the EME propagation step. " diff --git a/tidy3d/components/field_projection.py b/tidy3d/components/field_projection.py index 2f5cc1dbf..e9b2ae457 100644 --- a/tidy3d/components/field_projection.py +++ b/tidy3d/components/field_projection.py @@ -2,7 +2,7 @@ from __future__ import annotations -from typing import Dict, List, Tuple, Union +from typing import List, Tuple, Union import numpy as np import pydantic.v1 as pydantic @@ -91,13 +91,6 @@ class FieldProjector(Tidy3dBaseModel): units=MICROMETER, ) - currents: Dict[str, xr.Dataset] = pydantic.Field( - None, - title="Surface current densities", - description="Dictionary mapping monitor name to an ``xarray.Dataset`` storing the " - "surface current densities.", - ) - @cached_property def is_2d_simulation(self) -> bool: non_zero_dims = sum(1 for size in self.sim_data.simulation.size if size != 0) @@ -362,7 +355,7 @@ def integrate_1d( pts_u: np.ndarray, ): """Trapezoidal integration in two dimensions.""" - return np.trapz(np.squeeze(function) * np.squeeze(phase), pts_u, axis=0) + return np.trapz(np.squeeze(function) * np.squeeze(phase), pts_u, axis=0) # noqa: NPY201 def integrate_2d( self, @@ -372,7 +365,7 @@ def integrate_2d( 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) + return np.trapz(np.trapz(np.squeeze(function) * phase, pts_u, axis=0), pts_v, axis=0) # noqa: NPY201 def _far_fields_for_surface( self, @@ -420,10 +413,15 @@ def _far_fields_for_surface( _, source_names = surface.monitor.pop_axis(("x", "y", "z"), axis=surface.axis) # integration dimension for 2d far field projection - zero_dim = (dim for dim, size in enumerate(self.sim_data.simulation.size) if size == 0) + zero_dim = [dim for dim, size in enumerate(self.sim_data.simulation.size) if size == 0] if self.is_2d_simulation: + # Ensure zero_dim has a single element since {zero_dim} expects a value + if len(zero_dim) != 1: + raise ValueError("Expected exactly one dimension with size 0 for 2D simulation") + + zero_dim = zero_dim[0] integration_axis = {0, 1, 2} - {zero_dim, surface.axis} - idx_int_1d = integration_axis.pop() # Get the remaining axis as an integer + idx_int_1d = integration_axis.pop() idx_u, idx_v = idx_uv cmp_1, cmp_2 = source_names diff --git a/tidy3d/components/geometry/base.py b/tidy3d/components/geometry/base.py index 68d104946..d0f436383 100644 --- a/tidy3d/components/geometry/base.py +++ b/tidy3d/components/geometry/base.py @@ -729,7 +729,8 @@ def parse_xyz_kwargs(**xyz) -> Tuple[Axis, float]: Index into xyz axis (0,1,2) and position along that axis. """ xyz_filtered = {k: v for k, v in xyz.items() if v is not None} - assert len(xyz_filtered) == 1, "exactly one kwarg in [x,y,z] must be specified." + if len(xyz_filtered) != 1: + raise ValueError("exactly one kwarg in [x,y,z] must be specified.") axis_label, position = list(xyz_filtered.items())[0] axis = "xyz".index(axis_label) return axis, position @@ -3289,7 +3290,10 @@ def compute_derivatives(self, derivative_info: DerivativeInfo) -> AutogradFieldM ) vjp_dict_geo = geo.compute_derivatives(geo_info) grad_vjp_values = list(vjp_dict_geo.values()) - assert len(grad_vjp_values) == 1, "Got multiple gradients for single geometry field." + + if len(grad_vjp_values) != 1: + raise AssertionError("Got multiple gradients for single geometry field.") + grad_vjps[field_path] = grad_vjp_values[0] return grad_vjps diff --git a/tidy3d/components/geometry/polyslab.py b/tidy3d/components/geometry/polyslab.py index 814da9654..817d6e823 100644 --- a/tidy3d/components/geometry/polyslab.py +++ b/tidy3d/components/geometry/polyslab.py @@ -1373,9 +1373,8 @@ def _surface_area(self, bounds: Bound) -> float: def compute_derivatives(self, derivative_info: DerivativeInfo) -> AutogradFieldMap: """Compute the adjoint derivatives for this object.""" - assert derivative_info.paths == [ - ("vertices",) - ], "only support derivative wrt 'PolySlab.vertices'." + if derivative_info.paths != [("vertices",)]: + raise ValueError("only support derivative wrt 'PolySlab.vertices'.") vjp_vertices = self.compute_derivative_vertices(derivative_info=derivative_info) @@ -1397,7 +1396,8 @@ def compute_derivative_vertices(self, derivative_info: DerivativeInfo) -> Traced edge_centers_axis = self.center_axis * np.ones(num_vertices) edge_centers_xyz = self.unpop_axis_vect(edge_centers_axis, edge_centers_plane) - assert edge_centers_xyz.shape == (num_vertices, 3), "something bad happened" + 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( diff --git a/tidy3d/components/grid/grid.py b/tidy3d/components/grid/grid.py index 85371af9c..430cc7aa7 100644 --- a/tidy3d/components/grid/grid.py +++ b/tidy3d/components/grid/grid.py @@ -517,7 +517,8 @@ def discretize_inds(self, box: Box, extend: bool = False) -> List[Tuple[int, int # for each dimension for axis, (pt_min, pt_max) in enumerate(zip(pts_min, pts_max)): bound_coords = np.array(boundaries.to_list[axis]) - assert pt_min <= pt_max, "min point was greater than max point" + if pt_min > pt_max: + raise AssertionError("min point was greater than max point") # index of smallest coord greater than pt_max inds_gt_pt_max = np.where(bound_coords > pt_max)[0] diff --git a/tidy3d/components/grid/grid_spec.py b/tidy3d/components/grid/grid_spec.py index 5c06473fe..2c401705d 100644 --- a/tidy3d/components/grid/grid_spec.py +++ b/tidy3d/components/grid/grid_spec.py @@ -8,7 +8,7 @@ import numpy as np import pydantic.v1 as pd -from ...constants import C_0, MICROMETER, fp_eps +from ...constants import C_0, MICROMETER, inf from ...exceptions import SetupError from ...log import log from ..base import Tidy3dBaseModel @@ -136,6 +136,78 @@ def _add_pml_to_bounds(num_layers: Tuple[int, int], bounds: Coords1D) -> Coords1 add_right = bounds[-1] + last_step * np.arange(1, num_layers[1] + 1) return np.concatenate((add_left, bounds, add_right)) + @staticmethod + def _postprocess_unaligned_grid( + axis: Axis, + simulation_box: Box, + machine_error_relaxation: bool, + bound_coords: Coords1D, + ) -> Coords1D: + """Postprocess grids whose two ends might be aligned with simulation boundaries. + This is to be used in `_make_coords_initial`. + + Parameters + ---------- + axis : Axis + Axis of this direction. + structures : List[StructureType] + List of structures present in simulation, the first one being the simulation domain. + machine_error_relaxation : bool + When operations such as translation are applied to the 1d grids, fix the bounds + were numerically within the simulation bounds but were still chopped off. + bound_coords : Coord1D + 1D grids potentially unaligned with the simulation boundary + + Returns + ------- + :class:`.Coords1D`: + 1D coords to be used as grid boundaries. + + """ + center, size = simulation_box.center[axis], simulation_box.size[axis] + # chop off any coords outside of simulation bounds, beyond some buffer region + # to take numerical effects into account + bound_min = np.nextafter(center - size / 2, -inf, dtype=np.float32) + bound_max = np.nextafter(center + size / 2, inf, dtype=np.float32) + + if bound_max < bound_coords[0] or bound_min > bound_coords[-1]: + axis_name = "xyz"[axis] + raise SetupError( + f"Simulation domain does not overlap with the provided grid in '{axis_name}' direction." + ) + + if size == 0: + # in case of zero-size dimension return the boundaries between which simulation falls + ind = np.searchsorted(bound_coords, center, side="right") + + # in case when the center coincides with the right most boundary + if ind >= len(bound_coords): + ind = len(bound_coords) - 1 + + return bound_coords[ind - 1 : ind + 1] + + else: + bound_coords = bound_coords[bound_coords <= bound_max] + bound_coords = bound_coords[bound_coords >= bound_min] + + # if not extending to simulation bounds, repeat beginning and end + dl_min = bound_coords[1] - bound_coords[0] + dl_max = bound_coords[-1] - bound_coords[-2] + while bound_coords[0] - dl_min >= bound_min: + bound_coords = np.insert(bound_coords, 0, bound_coords[0] - dl_min) + while bound_coords[-1] + dl_max <= bound_max: + bound_coords = np.append(bound_coords, bound_coords[-1] + dl_max) + + # in case operations are applied to coords, it's possible the bounds were numerically within + # the simulation bounds but were still chopped off, which is fixed here + if machine_error_relaxation: + if np.isclose(bound_coords[0] - dl_min, bound_min): + bound_coords = np.insert(bound_coords, 0, bound_coords[0] - dl_min) + if np.isclose(bound_coords[-1] + dl_max, bound_max): + bound_coords = np.append(bound_coords, bound_coords[-1] + dl_max) + + return bound_coords + class UniformGrid(GridSpec1d): """Uniform 1D grid. The most standard way to define a simulation is to use a constant grid size in each of the three directions. @@ -176,8 +248,6 @@ def _make_coords_initial( Axis of this direction. structures : List[StructureType] List of structures present in simulation, the first one being the simulation domain. - **kwargs: - Other arguments all go here. Returns ------- @@ -199,6 +269,50 @@ def _make_coords_initial( return center - size / 2 + np.arange(num_cells + 1) * dl_snapped +class CustomGridBoundaries(GridSpec1d): + """Custom 1D grid supplied as a list of grid cell boundary coordinates. + + Example + ------- + >>> grid_1d = CustomGridBoundaries(coords=[-0.2, 0.0, 0.2, 0.4, 0.5, 0.6, 0.7]) + """ + + coords: Coords1D = pd.Field( + ..., + title="Grid Boundary Coordinates", + description="An array of grid boundary coordinates.", + units=MICROMETER, + ) + + def _make_coords_initial( + self, + axis: Axis, + structures: List[StructureType], + **kwargs, + ) -> Coords1D: + """Customized 1D coords to be used as grid boundaries. + + Parameters + ---------- + axis : Axis + Axis of this direction. + structures : List[StructureType] + List of structures present in simulation, the first one being the simulation domain. + + Returns + ------- + :class:`.Coords1D`: + 1D coords to be used as grid boundaries. + """ + + return self._postprocess_unaligned_grid( + axis=axis, + simulation_box=structures[0].geometry, + machine_error_relaxation=False, + bound_coords=self.coords, + ) + + class CustomGrid(GridSpec1d): """Custom 1D grid supplied as a list of grid cell sizes centered on the simulation center. @@ -241,8 +355,6 @@ def _make_coords_initial( Axis of this direction. structures : List[StructureType] List of structures present in simulation, the first one being the simulation domain. - *kwargs - Other arguments all go here. Returns ------- @@ -250,7 +362,7 @@ def _make_coords_initial( 1D coords to be used as grid boundaries. """ - center, size = structures[0].geometry.center[axis], structures[0].geometry.size[axis] + center = structures[0].geometry.center[axis] # get bounding coordinates dl = np.array(self.dl) @@ -263,49 +375,12 @@ def _make_coords_initial( else: bound_coords += self.custom_offset - # chop off any coords outside of simulation bounds, beyond some buffer region - # to take numerical effects into account - buffer = fp_eps * size - bound_min = center - size / 2 - buffer - bound_max = center + size / 2 + buffer - - if bound_max < bound_coords[0] or bound_min > bound_coords[-1]: - axis_name = "xyz"[axis] - raise SetupError( - f"Simulation domain does not overlap with the provided custom grid in '{axis_name}' direction." - ) - - if size == 0: - # in case of zero-size dimension return the boundaries between which simulation falls - ind = np.searchsorted(bound_coords, center, side="right") - - # in case when the center coincides with the right most boundary - if ind >= len(bound_coords): - ind = len(bound_coords) - 1 - - return bound_coords[ind - 1 : ind + 1] - - else: - bound_coords = bound_coords[bound_coords <= bound_max] - bound_coords = bound_coords[bound_coords >= bound_min] - - # if not extending to simulation bounds, repeat beginning and end - dl_min = dl[0] - dl_max = dl[-1] - while bound_coords[0] - dl_min >= bound_min: - bound_coords = np.insert(bound_coords, 0, bound_coords[0] - dl_min) - while bound_coords[-1] + dl_max <= bound_max: - bound_coords = np.append(bound_coords, bound_coords[-1] + dl_max) - - # in case a `custom_offset` is provided, it's possible the bounds were numerically within - # the simulation bounds but were still chopped off, which is fixed here - if self.custom_offset is not None: - if np.isclose(bound_coords[0] - dl_min, bound_min): - bound_coords = np.insert(bound_coords, 0, bound_coords[0] - dl_min) - if np.isclose(bound_coords[-1] + dl_max, bound_max): - bound_coords = np.append(bound_coords, bound_coords[-1] + dl_max) - - return bound_coords + return self._postprocess_unaligned_grid( + axis=axis, + simulation_box=structures[0].geometry, + machine_error_relaxation=self.custom_offset is not None, + bound_coords=bound_coords, + ) class AutoGrid(GridSpec1d): @@ -454,7 +529,7 @@ def _make_coords_initial( return np.array(bound_coords) -GridType = Union[UniformGrid, CustomGrid, AutoGrid] +GridType = Union[UniformGrid, CustomGrid, AutoGrid, CustomGridBoundaries] class GridSpec(Tidy3dBaseModel): @@ -657,6 +732,14 @@ def make_grid( coords = Coords(**coords_dict) return Grid(boundaries=coords) + @classmethod + def from_grid(cls, grid: Grid) -> GridSpec: + """Import grid directly from another simulation, e.g. ``grid_spec = GridSpec.from_grid(sim.grid)``.""" + grid_dict = {} + for dim in "xyz": + grid_dict["grid_" + dim] = CustomGridBoundaries(coords=grid.boundaries.to_dict[dim]) + return cls(**grid_dict) + @classmethod def auto( cls, diff --git a/tidy3d/components/medium.py b/tidy3d/components/medium.py index add2904b2..361b7153f 100644 --- a/tidy3d/components/medium.py +++ b/tidy3d/components/medium.py @@ -1157,7 +1157,7 @@ def derivative_eps_complex_volume( ) vjp_value += vjp_value_fld - return vjp_value + return vjp_value.sum("f") class AbstractCustomMedium(AbstractMedium, ABC): @@ -1402,7 +1402,7 @@ 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) + E_der_dim_interp = E_der_dim.interp(**coords_interp).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) @@ -2558,8 +2558,11 @@ def _derivative_field_cmp( eps_data: PermittivityDataset, dim: str, ) -> np.ndarray: - coords_interp = {key: val for key, val in eps_data.coords.items() if len(val) > 1} - dims_sum = {dim for dim in eps_data.coords.keys() if dim not in coords_interp} + """Compute derivative with respect to the ``dim`` components within the custom medium.""" + + coords_interp = {key: eps_data.coords[key] for key in "xyz"} + coords_interp = {key: val for key, val in coords_interp.items() if len(val) > 1} + dims_sum = [dim for dim in "xyz" if dim not in coords_interp] # compute sizes along each of the interpolation dimensions sizes_list = [] @@ -2588,8 +2591,10 @@ 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) - vjp_array = np.array(E_der_dim_interp.values).astype(complex) + 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.sum("f") + + vjp_array = np.array(E_der_dim_interp.values, dtype=float) vjp_array = vjp_array.reshape(eps_data.shape) # multiply by volume elements (if possible, being defensive here..) diff --git a/tidy3d/components/parameter_perturbation.py b/tidy3d/components/parameter_perturbation.py index e68c1f64d..afb5ba18c 100644 --- a/tidy3d/components/parameter_perturbation.py +++ b/tidy3d/components/parameter_perturbation.py @@ -13,7 +13,7 @@ from ..components.data.validators import validate_no_nans from ..components.types import TYPE_TAG_STR, ArrayLike, Ax, Complex, FieldVal, InterpMethod from ..components.viz import add_ax_if_none -from ..constants import CMCUBE, EPSILON_0, KELVIN, PERCMCUBE, inf +from ..constants import CMCUBE, EPSILON_0, HERTZ, KELVIN, PERCMCUBE, inf from ..exceptions import DataError from ..log import log from .base import Tidy3dBaseModel, cached_property @@ -1263,6 +1263,7 @@ class IndexPerturbation(Tidy3dBaseModel): ..., title="Frequency", description="Frequency to evaluate permittivity at (Hz).", + units=HERTZ, ) @pd.root_validator(skip_on_failure=True) diff --git a/tidy3d/components/simulation.py b/tidy3d/components/simulation.py index 330b80495..e03daa706 100644 --- a/tidy3d/components/simulation.py +++ b/tidy3d/components/simulation.py @@ -16,7 +16,6 @@ from ..exceptions import SetupError, Tidy3dError, Tidy3dImportError, ValidationError from ..log import log from ..updater import Updater -from .autograd import AutogradFieldMap from .base import cached_property, skip_if_fields_missing from .base_sim.simulation import AbstractSimulation from .boundary import ( @@ -37,7 +36,7 @@ from .geometry.utils import flatten_groups, traverse_geometries from .geometry.utils_2d import get_bounds, get_thickened_geom, snap_coordinate_to_grid, subdivide from .grid.grid import Coords, Coords1D, Grid -from .grid.grid_spec import AutoGrid, CustomGrid, GridSpec, UniformGrid +from .grid.grid_spec import AutoGrid, GridSpec, UniformGrid from .lumped_element import LumpedElementType from .medium import ( AbstractCustomMedium, @@ -203,6 +202,14 @@ class AbstractYeeGridSimulation(AbstractSimulation, ABC): "``True`` to apply the default subpixel averaging methods corresponding to ``SubpixelSpec()`` " ", 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.", + ) + """ Supply :class:`SubpixelSpec` to select subpixel averaging methods separately for dielectric, metal, and PEC material interfaces. Alternatively, supply ``True`` to use default subpixel averaging methods, @@ -1269,7 +1276,12 @@ def get_dls(geom: Geometry, axis: Axis, num_dls: int) -> List[float]: def snap_to_grid(geom: Geometry, axis: Axis) -> Geometry: """Snap a 2D material to the Yee grid.""" center = get_bounds(geom, axis)[0] - assert get_bounds(geom, axis)[0] == get_bounds(geom, axis)[1] + if get_bounds(geom, axis)[0] != get_bounds(geom, axis)[1]: + raise AssertionError( + "Unexpected error encountered while processing 2D material. " + "The upper and lower bounds of the geometry in the normal direction are not equal. " + "If you encounter this error, please create an issue in the Tidy3D github repository." + ) snapped_center = snap_coordinate_to_grid(self.grid, center, axis) return geom._update_from_bounds(bounds=(snapped_center, snapped_center), axis=axis) @@ -1344,6 +1356,213 @@ def suggest_mesh_overrides(self, **kwargs) -> List[MeshOverrideStructure]: return mesh_overrides + def subsection( + self, + region: Box, + boundary_spec: BoundarySpec = None, + grid_spec: Union[GridSpec, Literal["identical"]] = None, + symmetry: Tuple[Symmetry, Symmetry, Symmetry] = None, + sources: Tuple[SourceType, ...] = None, + monitors: Tuple[MonitorType, ...] = None, + remove_outside_structures: bool = True, + remove_outside_custom_mediums: bool = False, + include_pml_cells: bool = False, + **kwargs, + ) -> AbstractYeeGridSimulation: + """Generate a simulation instance containing only the ``region``. + + Parameters + ---------- + region : :class:.`Box` + New simulation domain. + boundary_spec : :class:.`BoundarySpec` = None + New boundary specification. If ``None``, then it is inherited from the original + simulation. + grid_spec : :class:.`GridSpec` = None + New grid specification. If ``None``, then it is inherited from the original + simulation. If ``identical``, then the original grid is transferred directly as a + :class:.`CustomGrid`. Note that in the latter case the region of the new simulation is + snapped to the original grid lines. + symmetry : Tuple[Literal[0, -1, 1], Literal[0, -1, 1], Literal[0, -1, 1]] = None + New simulation symmetry. If ``None``, then it is inherited from the original + simulation. Note that in this case the size and placement of new simulation domain + must be commensurate with the original symmetry. + sources : Tuple[SourceType, ...] = None + New list of sources. If ``None``, then the sources intersecting the new simulation + domain are inherited from the original simulation. + monitors : Tuple[MonitorType, ...] = None + New list of monitors. If ``None``, then the monitors intersecting the new simulation + domain are inherited from the original simulation. + remove_outside_structures : bool = True + Remove structures outside of the new simulation domain. + remove_outside_custom_mediums : bool = True + Remove custom medium data outside of the new simulation domain. + include_pml_cells : bool = False + Keep PML cells in simulation boundaries. Note that retained PML cells will be converted + to regular cells, and the simulation domain boundary will be moved accordingly. + **kwargs + Other arguments passed to new simulation instance. + """ + + # must intersect the original domain + if not self.intersects(region): + raise SetupError("Requested region does not intersect simulation domain") + + # restrict to the original simulation domain + if include_pml_cells: + new_bounds = Box.bounds_intersection(self.simulation_bounds, region.bounds) + else: + new_bounds = Box.bounds_intersection(self.bounds, region.bounds) + new_bounds = [list(new_bounds[0]), list(new_bounds[1])] + + # grid spec inheritace + if grid_spec is None: + grid_spec = self.grid_spec + elif isinstance(grid_spec, str) and grid_spec == "identical": + # create a custom grid from existing one + grids_1d = self.grid.boundaries.to_list + grid_spec = GridSpec.from_grid(self.grid) + + # adjust region bounds to perfectly coincide with the grid + # note, sometimes (when a box already seems to perfrecty align with the grid) + # this causes the new region to expand one more pixel because of numerical roundoffs + # To help to avoid that we shrink new region by a small amount. + center = [(bmin + bmax) / 2 for bmin, bmax in zip(*new_bounds)] + size = [max(0.0, bmax - bmin - 2 * fp_eps) for bmin, bmax in zip(*new_bounds)] + aux_box = Box(center=center, size=size) + grid_inds = self.grid.discretize_inds(box=aux_box) + + for dim in range(3): + # preserve zero size dimensions + if new_bounds[0][dim] != new_bounds[1][dim]: + new_bounds[0][dim] = grids_1d[dim][grid_inds[dim][0]] + new_bounds[1][dim] = grids_1d[dim][grid_inds[dim][1]] + + # if symmetry is not overriden we inherit it from the original simulation where is needed + if symmetry is None: + # start with no symmetry + symmetry = [0, 0, 0] + + # now check in each dimension whether we cross symmetry plane + for dim in range(3): + if self.symmetry[dim] != 0: + crosses_symmetry = ( + new_bounds[0][dim] < self.center[dim] + and new_bounds[1][dim] > self.center[dim] + ) + + # inherit symmetry only if we cross symmetry plane, otherwise we don't impose + # symmetry even if the original simulation had symmetry + if crosses_symmetry: + symmetry[dim] = self.symmetry[dim] + center = (new_bounds[0][dim] + new_bounds[1][dim]) / 2 + + if not math.isclose(center, self.center[dim]): + log.warning( + f"The original simulation is symmetric along {'xyz'[dim]} direction. " + "The requested new simulation region does cross the symmetry plane but is " + "not symmetric with respect to it. To preserve correct symmetry, " + "the requested simulation region is expanded symmetrically." + ) + new_bounds[0][dim] = 2 * self.center[dim] - new_bounds[1][dim] + + # symmetry and grid spec treatments could change new simulation bounds + # thus, recreate a box instance + new_box = Box.from_bounds(*new_bounds) + + # inheritance of structures, sources, monitors, and boundary specs + if remove_outside_structures: + new_structures = [strc for strc in self.structures if new_box.intersects(strc.geometry)] + else: + new_structures = self.structures + + if sources is None: + sources = [src for src in self.sources if new_box.intersects(src)] + + if monitors is None: + monitors = [mnt for mnt in self.monitors if new_box.intersects(mnt)] + + if boundary_spec is None: + boundary_spec = self.boundary_spec + + # set boundary conditions in zero-size dimension to periodic + for dim in range(3): + if new_bounds[0][dim] == new_bounds[1][dim] and not isinstance( + boundary_spec.to_list[dim][0], Periodic + ): + axis_name = "xyz"[dim] + log.warning( + f"The resulting simulation subsection has size zero along axis '{axis_name}'. " + "Periodic boundary conditions are automatically set along this dimension." + ) + boundary_spec = boundary_spec.updated_copy(**{"xyz"[dim]: Boundary.periodic()}) + + # reduction of custom medium data + new_sim_medium = self.medium + if remove_outside_custom_mediums: + # check for special treatment in case of PML + if any( + any(isinstance(edge, (PML, StablePML, Absorber)) for edge in boundary) + for boundary in boundary_spec.to_list + ): + # if we need to cut out outside custom medium we have to be careful about PML/Absorber + # we should include data in PML so that there is no artificial reflection at PML boundaries + + # to do this, we first create an auxiliary simulation + aux_sim = self.updated_copy( + center=new_box.center, + size=new_box.size, + grid_spec=grid_spec, + boundary_spec=boundary_spec, + monitors=[], + sources=sources, # need wavelength in case of auto grid + symmetry=symmetry, + structures=new_structures, + ) + + # then use its bounds as region for data cut off + new_bounds = aux_sim.simulation_bounds + + # Note that this is not a full proof strategy. For example, if grid_spec is AutoGrid + # then after outside custom medium data is removed the grid sizes and, thus, + # pml extents can change as well + + # now cut out custom medium data + new_structures_reduced_data = [] + + for structure in new_structures: + medium = structure.medium + if isinstance(medium, AbstractCustomMedium): + new_structure_bounds = Box.bounds_intersection( + new_bounds, structure.geometry.bounds + ) + new_medium = medium.sel_inside(bounds=new_structure_bounds) + new_structure = structure.updated_copy(medium=new_medium) + new_structures_reduced_data.append(new_structure) + else: + new_structures_reduced_data.append(structure) + + new_structures = new_structures_reduced_data + + if isinstance(self.medium, AbstractCustomMedium): + new_sim_medium = self.medium.sel_inside(bounds=new_bounds) + + # finally, create an updated copy with all modifications + new_sim = self.updated_copy( + center=new_box.center, + size=new_box.size, + medium=new_sim_medium, + grid_spec=grid_spec, + boundary_spec=boundary_spec, + monitors=monitors, + sources=sources, + symmetry=symmetry, + structures=new_structures, + **kwargs, + ) + + return new_sim + class Simulation(AbstractYeeGridSimulation): """ @@ -2673,14 +2892,14 @@ def _projection_monitors_distance(cls, val, values): def _projection_monitors_2d(cls, val, values): """ Validate if the field projection monitor is set up for a 2D simulation and - ensure the observation angle is configured correctly. + ensure the observation parameters are configured correctly. - - For a 2D simulation in the x-y plane, 'theta' should be set to 'pi/2'. - - For a 2D simulation in the y-z plane, 'phi' should be set to 'pi/2'. - - For a 2D simulation in the x-z plane, 'phi' should be set to '0'. + - For a 2D simulation in the x-y plane, ``theta`` should be set to ``pi/2``. + - For a 2D simulation in the y-z plane, ``phi`` should be set to ``pi/2`` or ``3*pi/2``. + - For a 2D simulation in the x-z plane, ``phi`` should be set to ``0`` or ``pi``. Note: Exact far field projection is not available yet. Currently, only - 'far_field_approx = True' is supported. + ``far_field_approx = True`` is supported. """ if val is None: @@ -2693,19 +2912,12 @@ def _projection_monitors_2d(cls, val, values): if non_zero_dims == 3: return val - plane = None - phi_value = None - theta_value = None - if sim_size[0] == 0: plane = "y-z" - phi_value = np.pi / 2 elif sim_size[1] == 0: plane = "x-z" - phi_value = 0 elif sim_size[2] == 0: plane = "x-y" - theta_value = np.pi / 2 for monitor in val: if isinstance(monitor, AbstractFieldProjectionMonitor): @@ -2714,37 +2926,63 @@ def _projection_monitors_2d(cls, val, values): f"Monitor '{monitor.name}' is not supported in 1D simulations." ) - if isinstance( - monitor, (FieldProjectionCartesianMonitor, FieldProjectionKSpaceMonitor) - ): + 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, FieldProjectionAngleMonitor): - if not monitor.far_field_approx: - raise SetupError( - "Exact far field projection in 2D simulations is coming soon." - "Please set 'far_field_approx = True'." - ) - if plane == "y-z" and (len(monitor.phi) != 1 or monitor.phi[0] != phi_value): + 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( - "For a 2D simulation in the y-z plane, the observation angle 'phi' " - f"of monitor '{monitor.name}' should be set to 'np.pi/2'." + 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}." ) - elif plane == "x-z" and (len(monitor.phi) != 1 or monitor.phi[0] != phi_value): - raise SetupError( - "For a 2D simulation in the x-z plane, the observation angle 'phi' " - f"of monitor '{monitor.name}' should be set to '0'." + + 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"}, + "x-z": {"valid_value": [0, np.pi], "coord": "phi"}, + "x-y": {"valid_value": [np.pi / 2], "coord": "theta"}, + }[plane] + + coord = getattr(monitor, config["coord"]) + if not all(value in config["valid_value"] for value in coord): + replacements = { + np.pi: "np.pi", + np.pi / 2: "np.pi/2", + 3 * np.pi / 2: "3*np.pi/2", + 0: "0", + } + valid_values_str = ", ".join( + replacements.get(val) for val in config["valid_value"] ) - elif plane == "x-y" and ( - len(monitor.theta) != 1 or monitor.theta[0] != theta_value - ): raise SetupError( - "For a 2D simulation in the x-y plane, the observation angle 'theta' " - f"of monitor '{monitor.name}' should be set to 'np.pi/2'." + f"For a 2D simulation in the {plane} plane, the observation " + f"angle '{config['coord']}' of monitor " + f"'{monitor.name}' should be set to " + f"'{valid_values_str}'" ) + return val @pydantic.validator("monitors", always=True) @@ -2931,6 +3169,24 @@ def _post_init_validators(self) -> None: self._validate_no_structures_pml() self._validate_tfsf_nonuniform_grid() self._validate_nonlinear_specs() + self._validate_custom_source_time() + + def _validate_custom_source_time(self): + """Warn if all simulation times are outside CustomSourceTime definition range.""" + run_time = self._run_time + for idx, source in enumerate(self.sources): + if isinstance(source.source_time, CustomSourceTime): + if source.source_time._all_outside_range(run_time=run_time): + data_times = source.source_time.data_times + mint = np.min(data_times) + maxt = np.max(data_times) + log.warning( + f"'CustomSourceTime' at 'sources[{idx}]' is defined over a time range " + f"'({mint}, {maxt})' which does not include any of the 'Simulation' " + f"times '({0, run_time})'. The envelope will be constant extrapolated " + "from the first or last value in the 'CustomSourceTime', which may not " + "be the desired outcome." + ) def _validate_no_structures_pml(self) -> None: """Ensure no structures terminate / have bounds inside of PML.""" @@ -3306,11 +3562,11 @@ def _warn_time_monitors_outside_run_time(self) -> None: """ Autograd adjoint support """ - def with_adjoint_monitors(self, sim_fields: AutogradFieldMap) -> Simulation: + def with_adjoint_monitors(self, sim_fields_keys: list) -> Simulation: """Copy of self with adjoint field and permittivity monitors for every traced structure.""" # set of indices in the structures needing adjoint monitors - structure_indices = {index for (_, index, *_), _ in sim_fields.items()} + structure_indices = {index for (_, index, *_) in sim_fields_keys} mnts_fld, mnts_eps = self.make_adjoint_monitors(structure_indices=structure_indices) monitors = list(self.monitors) + list(mnts_fld) + list(mnts_eps) @@ -3344,14 +3600,6 @@ def freqs_adjoint(self) -> list[float]: if isinstance(mnt, FreqMonitor): freqs.update(mnt.freqs) freqs = sorted(freqs) - - if len(freqs) > 1: - raise ValueError( - "Only the same, single frequency is supported in all monitors " - "when using autograd differentiation. " - f"Found {len(freqs)} distinct frequencies in the monitors." - ) - return freqs """ Accounting """ @@ -4291,207 +4539,3 @@ def from_scene(cls, scene: Scene, **kwargs) -> Simulation: medium=scene.medium, **kwargs, ) - - def subsection( - self, - region: Box, - boundary_spec: BoundarySpec = None, - grid_spec: Union[GridSpec, Literal["identical"]] = None, - symmetry: Tuple[Symmetry, Symmetry, Symmetry] = None, - sources: Tuple[SourceType, ...] = None, - monitors: Tuple[MonitorType, ...] = None, - remove_outside_structures: bool = True, - remove_outside_custom_mediums: bool = False, - **kwargs, - ) -> Simulation: - """Generate a simulation instance containing only the ``region``. - - Parameters - ---------- - region : :class:.`Box` - New simulation domain. - boundary_spec : :class:.`BoundarySpec` = None - New boundary specification. If ``None``, then it is inherited from the original - simulation. - grid_spec : :class:.`GridSpec` = None - New grid specification. If ``None``, then it is inherited from the original - simulation. If ``identical``, then the original grid is transferred directly as a - :class:.`CustomGrid`. Note that in the latter case the region of the new simulation is - snapped to the original grid lines. - symmetry : Tuple[Literal[0, -1, 1], Literal[0, -1, 1], Literal[0, -1, 1]] = None - New simulation symmetry. If ``None``, then it is inherited from the original - simulation. Note that in this case the size and placement of new simulation domain - must be commensurate with the original symmetry. - sources : Tuple[SourceType, ...] = None - New list of sources. If ``None``, then the sources intersecting the new simulation - domain are inherited from the original simulation. - monitors : Tuple[MonitorType, ...] = None - New list of monitors. If ``None``, then the monitors intersecting the new simulation - domain are inherited from the original simulation. - remove_outside_structures : bool = True - Remove structures outside of the new simulation domain. - remove_outside_custom_mediums : bool = True - Remove custom medium data outside of the new simulation domain. - **kwargs - Other arguments passed to new simulation instance. - """ - - # must intersect the original domain - if not self.intersects(region): - raise SetupError("Requested region does not intersect simulation domain") - - # restrict to the original simulation domain - new_bounds = Box.bounds_intersection(self.bounds, region.bounds) - new_bounds = [list(new_bounds[0]), list(new_bounds[1])] - - # grid spec inheritace - if grid_spec is None: - grid_spec = self.grid_spec - elif isinstance(grid_spec, str) and grid_spec == "identical": - # create a custom grid from existing one - grids_1d = self.grid.boundaries.to_list - new_grids = [ - CustomGrid(dl=tuple(np.diff(grids_1d[dim])), custom_offset=grids_1d[dim][0]) - for dim in range(3) - ] - grid_spec = GridSpec(grid_x=new_grids[0], grid_y=new_grids[1], grid_z=new_grids[2]) - - # adjust region bounds to perfectly coincide with the grid - # note, sometimes (when a box already seems to perfrecty align with the grid) - # this causes the new region to expand one more pixel because of numerical roundoffs - # To help to avoid that we shrink new region by a small amount. - center = [(bmin + bmax) / 2 for bmin, bmax in zip(*new_bounds)] - size = [max(0.0, bmax - bmin - 2 * fp_eps) for bmin, bmax in zip(*new_bounds)] - aux_box = Box(center=center, size=size) - grid_inds = self.grid.discretize_inds(box=aux_box) - - for dim in range(3): - # preserve zero size dimensions - if new_bounds[0][dim] != new_bounds[1][dim]: - new_bounds[0][dim] = grids_1d[dim][grid_inds[dim][0]] - new_bounds[1][dim] = grids_1d[dim][grid_inds[dim][1]] - - # if symmetry is not overriden we inherit it from the original simulation where is needed - if symmetry is None: - # start with no symmetry - symmetry = [0, 0, 0] - - # now check in each dimension whether we cross symmetry plane - for dim in range(3): - if self.symmetry[dim] != 0: - crosses_symmetry = ( - new_bounds[0][dim] < self.center[dim] - and new_bounds[1][dim] > self.center[dim] - ) - - # inherit symmetry only if we cross symmetry plane, otherwise we don't impose - # symmetry even if the original simulation had symmetry - if crosses_symmetry: - symmetry[dim] = self.symmetry[dim] - center = (new_bounds[0][dim] + new_bounds[1][dim]) / 2 - - if not math.isclose(center, self.center[dim]): - log.warning( - f"The original simulation is symmetric along {'xyz'[dim]} direction. " - "The requested new simulation region does cross the symmetry plane but is " - "not symmetric with respect to it. To preserve correct symmetry, " - "the requested simulation region is expanded symmetrically." - ) - new_bounds[0][dim] = 2 * self.center[dim] - new_bounds[1][dim] - - # symmetry and grid spec treatments could change new simulation bounds - # thus, recreate a box instance - new_box = Box.from_bounds(*new_bounds) - - # inheritance of structures, sources, monitors, and boundary specs - if remove_outside_structures: - new_structures = [strc for strc in self.structures if new_box.intersects(strc.geometry)] - else: - new_structures = self.structures - - if sources is None: - sources = [src for src in self.sources if new_box.intersects(src)] - - if monitors is None: - monitors = [mnt for mnt in self.monitors if new_box.intersects(mnt)] - - if boundary_spec is None: - boundary_spec = self.boundary_spec - - # set boundary conditions in zero-size dimension to periodic - for dim in range(3): - if new_bounds[0][dim] == new_bounds[1][dim] and not isinstance( - boundary_spec.to_list[dim][0], Periodic - ): - axis_name = "xyz"[dim] - log.warning( - f"The resulting simulation subsection has size zero along axis '{axis_name}'. " - "Periodic boundary conditions are automatically set along this dimension." - ) - boundary_spec = boundary_spec.updated_copy(**{"xyz"[dim]: Boundary.periodic()}) - - # reduction of custom medium data - new_sim_medium = self.medium - if remove_outside_custom_mediums: - # check for special treatment in case of PML - if any( - any(isinstance(edge, (PML, StablePML, Absorber)) for edge in boundary) - for boundary in boundary_spec.to_list - ): - # if we need to cut out outside custom medium we have to be careful about PML/Absorber - # we should include data in PML so that there is no artificial reflection at PML boundaries - - # to do this, we first create an auxiliary simulation - aux_sim = self.updated_copy( - center=new_box.center, - size=new_box.size, - grid_spec=grid_spec, - boundary_spec=boundary_spec, - monitors=[], - sources=sources, # need wavelength in case of auto grid - symmetry=symmetry, - structures=new_structures, - ) - - # then use its bounds as region for data cut off - new_bounds = aux_sim.simulation_bounds - - # Note that this is not a full proof strategy. For example, if grid_spec is AutoGrid - # then after outside custom medium data is removed the grid sizes and, thus, - # pml extents can change as well - - # now cut out custom medium data - new_structures_reduced_data = [] - - for structure in new_structures: - medium = structure.medium - if isinstance(medium, AbstractCustomMedium): - new_structure_bounds = Box.bounds_intersection( - new_bounds, structure.geometry.bounds - ) - new_medium = medium.sel_inside(bounds=new_structure_bounds) - new_structure = structure.updated_copy(medium=new_medium) - new_structures_reduced_data.append(new_structure) - else: - new_structures_reduced_data.append(structure) - - new_structures = new_structures_reduced_data - - if isinstance(self.medium, AbstractCustomMedium): - new_sim_medium = self.medium.sel_inside(bounds=new_bounds) - - # finally, create an updated copy with all modifications - new_sim = self.updated_copy( - center=new_box.center, - size=new_box.size, - medium=new_sim_medium, - grid_spec=grid_spec, - boundary_spec=boundary_spec, - monitors=monitors, - sources=sources, - symmetry=symmetry, - structures=new_structures, - **kwargs, - ) - - return new_sim diff --git a/tidy3d/components/source.py b/tidy3d/components/source.py index 82e4cfb81..663ad4a5d 100644 --- a/tidy3d/components/source.py +++ b/tidy3d/components/source.py @@ -201,6 +201,27 @@ def end_time(self) -> float | None: return self.offset * self.twidth + END_TIME_FACTOR_GAUSSIAN * self.twidth + @property + def amp_complex(self) -> complex: + """Grab the complex amplitude from a ``GaussianPulse``.""" + phase = np.exp(1j * self.phase) + return self.amplitude * phase + + @classmethod + def from_amp_complex(cls, amp: complex, **kwargs) -> GaussianPulse: + """Set the complex amplitude of a ``GaussianPulse``. + + Parameters + ---------- + amp : complex + Complex-valued amplitude to set in the returned ``GaussianPulse``. + kwargs : dict + Keyword arguments passed to ``GaussianPulse()``, excluding ``amplitude`` & ``phase``. + """ + amplitude = abs(amp) + phase = np.angle(amp) + return cls(amplitude=amplitude, phase=phase, **kwargs) + class ContinuousWave(Pulse): """Source time dependence that ramps up to continuous oscillation @@ -331,6 +352,31 @@ def from_values( source_time_dataset=source_time_dataset, ) + @property + def data_times(self) -> ArrayFloat1D: + """Times of envelope definition.""" + if self.source_time_dataset is None: + return [] + data_times = self.source_time_dataset.values.coords["t"].values.squeeze() + return data_times + + def _all_outside_range(self, run_time: float) -> bool: + """Whether all times are outside range of definition.""" + + # can't validate if data isn't loaded + if self.source_time_dataset is None: + return False + + # make time a numpy array for uniform handling + data_times = self.data_times + + # shift time + twidth = 1.0 / (2 * np.pi * self.fwidth) + max_time_shifted = run_time - self.offset * twidth + min_time_shifted = -self.offset * twidth + + return (max_time_shifted < min(data_times)) | (min_time_shifted > max(data_times)) + def amp_time(self, time: float) -> complex: """Complex-valued source amplitude as a function of time. @@ -349,8 +395,8 @@ def amp_time(self, time: float) -> complex: return None # make time a numpy array for uniform handling - times = np.array([time] if isinstance(time, float) else time) - data_times = self.source_time_dataset.values.coords["t"].values.squeeze() + times = np.array([time] if isinstance(time, (int, float)) else time) + data_times = self.data_times # shift time twidth = 1.0 / (2 * np.pi * self.fwidth) @@ -363,12 +409,13 @@ def amp_time(self, time: float) -> complex: envelope = np.zeros(len(time_shifted), dtype=complex) values = self.source_time_dataset.values envelope[mask] = values.sel(t=time_shifted[mask], method="nearest").to_numpy() - envelope[~mask] = values.interp(t=time_shifted[~mask]).to_numpy() + if not all(mask): + envelope[~mask] = values.interp(t=time_shifted[~mask]).to_numpy() # modulation, phase, amplitude omega0 = 2 * np.pi * self.freq0 offset = np.exp(1j * self.phase) - oscillation = np.exp(-1j * omega0 * time) + oscillation = np.exp(-1j * omega0 * times) amp = self.amplitude return offset * oscillation * amp * envelope @@ -812,7 +859,7 @@ class CustomFieldSource(FieldSource, PlanarSource): ... run_time = 1e-6, ... shutoff=1e-6, ... ) - >>> sim_empty = sim.updated_copy(monitors = [Flux_monitor], + >>> sim_empty = sim.updated_copy(monitors = [Flux_monitor], # doctest: +SKIP ... structures = [], ... grid_spec= sim.grid_spec.updated_copy(override_structures = sim.structures) ... ) diff --git a/tidy3d/plugins/adjoint/web.py b/tidy3d/plugins/adjoint/web.py index ce3beaeab..b3e1da163 100644 --- a/tidy3d/plugins/adjoint/web.py +++ b/tidy3d/plugins/adjoint/web.py @@ -286,6 +286,15 @@ class AdjointBatch(Batch): description="Containers of information needed to reconstruct JaxSimulation for each item.", ) + jobs_cached: Dict[str, AdjointJob] = pd.Field( + None, + title="Jobs (Cached)", + description="Optional field to specify ``jobs``. Only used as a workaround internally " + "so that ``jobs`` is written when ``Batch.to_file()`` and then the proper task is loaded " + "from ``Batch.from_file()``. We recommend leaving unset as setting this field along with " + "fields that were not used to create the task will cause errors.", + ) + _job_type = AdjointJob def start(self) -> None: diff --git a/tidy3d/plugins/autograd/README.md b/tidy3d/plugins/autograd/README.md index 577b2ef86..557c813f4 100644 --- a/tidy3d/plugins/autograd/README.md +++ b/tidy3d/plugins/autograd/README.md @@ -109,54 +109,60 @@ Finally, the regular `web.run()` and `web.run_async()` functions have their deri The following components are traceable as inputs to the `td.Simulation` +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` -The following components are traceable as outputs of the `td.SimulationData` +The following data types are traceable as outputs of the `td.SimulationData` - `ModeData.amps` - `DiffractionData.amps` - `FieldData.field_components` +- `FieldData` operations: + - `FieldData.flux` + - `SimulationData.get_intensity(fld_monitor_name)` + - `SimulationData.get_poynting(fld_monitor_name)` -We currently have the following restrictions: +We also support the following high level features -- All monitors in the `Simulation` must be single frequency only. -- Only 500 max structures containing tracers can be added to the `Simulation` to cut down on processing time. In the future, `GeometryGroup` support will allow us to relax this restriction. -- `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. -- Gradient calculations are done client-side, meaning the field data in the traced structure regions must be downloaded. This can be a large amount of data for large, 3D structures. +- 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. -### To be supported soon +We currently have the following restrictions: -Next on our roadmap (targeting 2.8 and 2.9, summer 2024) is to support: +- 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. -- support for multi-frequency monitors in certain situations (single adjoint source). -- server-side gradient processing. +### To be supported soon -- `FieldData` operations: - - `FieldData.flux` - - `SimulationData.get_intensity` - - `SimulationData.get_poynting` +Next on our roadmap (targeting 2.8 and 2.9, fall 2024) is to support: -- `PoleResidue` and other dispersive models. -- custom (spatially-dependent) dispersive models, allowing topology optimization with metals. +- Field projection support +- Automatic handling of all broadband objective functions using a fitting approach. Later this year (2024), we plan to support: @@ -177,5 +183,6 @@ To convert existing tidy3d front end code to be autograd compatible, will need t - `float()` is not supported as far as I can tell. - `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 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. +- 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. diff --git a/tidy3d/plugins/autograd/__init__.py b/tidy3d/plugins/autograd/__init__.py index e69de29bb..dcbf05eb0 100644 --- a/tidy3d/plugins/autograd/__init__.py +++ b/tidy3d/plugins/autograd/__init__.py @@ -0,0 +1,5 @@ +from .differential_operators import value_and_grad + +__all__ = [ + "value_and_grad", +] diff --git a/tidy3d/plugins/autograd/differential_operators.py b/tidy3d/plugins/autograd/differential_operators.py new file mode 100644 index 000000000..e83c439e1 --- /dev/null +++ b/tidy3d/plugins/autograd/differential_operators.py @@ -0,0 +1,65 @@ +from typing import Any, 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 + +__all__ = [ + "value_and_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. + + 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. + x : ArrayLike + The point at which to evaluate the function and its gradient. + has_aux : bool = False + If True, the function 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) + + Raises + ------ + TypeError + If the function does not return a scalar value. + + 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) + + 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." + ) + + return (ans, vjp((vspace(ans).ones(), None))), aux diff --git a/tidy3d/plugins/dispersion/fit.py b/tidy3d/plugins/dispersion/fit.py index 80b65bfb0..b9de1df42 100644 --- a/tidy3d/plugins/dispersion/fit.py +++ b/tidy3d/plugins/dispersion/fit.py @@ -167,7 +167,8 @@ def _unpack_coeffs(coeffs): Tuple[np.ndarray[complex], np.ndarray[complex]] "a" and "c" poles for the PoleResidue model. """ - assert len(coeffs) % 4 == 0, "len(coeffs) must be multiple of 4." + if len(coeffs) % 4 != 0: + raise ValueError(f"len(coeffs) must be multiple of 4, got {len(coeffs)=}.") a_real = coeffs[0::4] a_imag = coeffs[1::4] @@ -692,8 +693,10 @@ def from_file(cls, fname: str, **loadtxt_kwargs) -> DispersionFitter: A :class:`DispersionFitter` instance. """ data = np.loadtxt(fname, **loadtxt_kwargs) - assert len(data.shape) == 2, "data must contain [wavelength, ndata, kdata] in columns" - assert data.shape[-1] in (2, 3), "data must have either 2 or 3 rows (if k data)" + if len(data.shape) != 2: + raise ValueError("data must contain [wavelength, ndata, kdata] in columns") + if data.shape[-1] not in (2, 3): + raise ValueError("data must have either 2 or 3 rows (if k data)") if data.shape[-1] == 2: wvl_um, n_data = data.T k_data = None diff --git a/tidy3d/plugins/invdes/design.py b/tidy3d/plugins/invdes/design.py index ab5a4074b..06266ae93 100644 --- a/tidy3d/plugins/invdes/design.py +++ b/tidy3d/plugins/invdes/design.py @@ -129,8 +129,9 @@ 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) - sim_data = web.run(simulation, verbose=self.verbose, **kwargs) - return sim_data + return web.run(simulation, verbose=self.verbose, **kwargs) + # sim_data = job.run() + # return sim_data class InverseDesignMulti(AbstractInverseDesign): diff --git a/tidy3d/plugins/invdes/transformation.py b/tidy3d/plugins/invdes/transformation.py index df61ce370..e6719e384 100644 --- a/tidy3d/plugins/invdes/transformation.py +++ b/tidy3d/plugins/invdes/transformation.py @@ -27,7 +27,10 @@ def __call__(self, *args, **kwargs) -> float: class FilterProject(InvdesBaseModel): """Transformation involving convolution by a conic filter followed by a ``tanh`` projection. - .. image:: ../../_static/img/filter_project.png + Notes + ----- + + .. image:: ../../_static/img/filter_project.png """ diff --git a/tidy3d/plugins/microwave/custom_path_integrals.py b/tidy3d/plugins/microwave/custom_path_integrals.py index 1fca7a80a..251e97a57 100644 --- a/tidy3d/plugins/microwave/custom_path_integrals.py +++ b/tidy3d/plugins/microwave/custom_path_integrals.py @@ -138,7 +138,11 @@ def compute_integral( elif isinstance(em_field, FieldTimeData): return TimeDataArray(data=result.data, coords=result.coords) else: - assert isinstance(em_field, ModeSolverData) + if not isinstance(em_field, ModeSolverData): + raise TypeError( + f"Unsupported 'em_field' type: {type(em_field)}. " + "Expected one of 'FieldData', 'FieldTimeData', 'ModeSolverData'." + ) return FreqModeDataArray(data=result.data, coords=result.coords) @staticmethod diff --git a/tidy3d/plugins/microwave/path_integrals.py b/tidy3d/plugins/microwave/path_integrals.py index 76e769ed7..60101f1a1 100644 --- a/tidy3d/plugins/microwave/path_integrals.py +++ b/tidy3d/plugins/microwave/path_integrals.py @@ -142,7 +142,12 @@ def compute_integral(self, scalar_field: EMScalarFieldType) -> IntegralResultTyp elif isinstance(scalar_field, ScalarFieldTimeDataArray): return TimeDataArray(data=result.data, coords=result.coords) else: - assert isinstance(scalar_field, ScalarModeFieldDataArray) + if not isinstance(scalar_field, ScalarModeFieldDataArray): + raise TypeError( + f"Unsupported 'scalar_field' type: {type(scalar_field)}. " + "Expected one of 'ScalarFieldDataArray', 'ScalarFieldTimeDataArray', " + "'ScalarModeFieldDataArray'." + ) return FreqModeDataArray(data=result.data, coords=result.coords) def _get_field_along_path(self, scalar_field: EMScalarFieldType) -> EMScalarFieldType: diff --git a/tidy3d/plugins/mode/mode_solver.py b/tidy3d/plugins/mode/mode_solver.py index b5399a0ca..2e4310290 100644 --- a/tidy3d/plugins/mode/mode_solver.py +++ b/tidy3d/plugins/mode/mode_solver.py @@ -33,6 +33,7 @@ from ...components.simulation import Simulation from ...components.source import ModeSource, SourceTime from ...components.types import ( + TYPE_TAG_STR, ArrayComplex3D, ArrayComplex4D, ArrayFloat1D, @@ -74,6 +75,7 @@ MODE_SIMULATION_TYPE = Union[Simulation, EMESimulation] MODE_SIMULATION_DATA_TYPE = Union[SimulationData, EMESimulationData] +MODE_PLANE_TYPE = Union[Box, ModeSource, ModeMonitor, ModeSolverMonitor] def require_fdtd_simulation(fn): @@ -118,8 +120,11 @@ class ModeSolver(Tidy3dBaseModel): discriminator="type", ) - plane: Box = pydantic.Field( - ..., title="Plane", description="Cross-sectional plane in which the mode will be computed." + plane: MODE_PLANE_TYPE = pydantic.Field( + ..., + title="Plane", + description="Cross-sectional plane in which the mode will be computed.", + discriminator=TYPE_TAG_STR, ) mode_spec: ModeSpec = pydantic.Field( @@ -336,56 +341,58 @@ 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_coords = self.plane.pop_axis( - self._solver_grid.boundaries.to_list, axis=self.normal_axis + solver = self.reduced_simulation_copy + + _, _solver_coords = solver.plane.pop_axis( + solver._solver_grid.boundaries.to_list, axis=solver.normal_axis ) # Compute and store the modes at all frequencies - n_complex, fields, eps_spec = self._solve_all_freqs( - coords=_solver_coords, symmetry=self.solver_symmetry + n_complex, fields, eps_spec = solver._solve_all_freqs( + coords=_solver_coords, symmetry=solver.solver_symmetry ) # start a dictionary storing the data arrays for the ModeSolverData index_data = ModeIndexDataArray( np.stack(n_complex, axis=0), coords=dict( - f=list(self.freqs), - mode_index=np.arange(self.mode_spec.num_modes), + f=list(solver.freqs), + mode_index=np.arange(solver.mode_spec.num_modes), ), ) data_dict = {"n_complex": index_data} # Construct the field data on Yee grid for field_name in ("Ex", "Ey", "Ez", "Hx", "Hy", "Hz"): - xyz_coords = self.grid_snapped[field_name].to_list + xyz_coords = solver.grid_snapped[field_name].to_list scalar_field_data = ScalarModeFieldDataArray( np.stack([field_freq[field_name] for field_freq in fields], axis=-2), coords=dict( x=xyz_coords[0], y=xyz_coords[1], z=xyz_coords[2], - f=list(self.freqs), - mode_index=np.arange(self.mode_spec.num_modes), + f=list(solver.freqs), + mode_index=np.arange(solver.mode_spec.num_modes), ), ) data_dict[field_name] = scalar_field_data # finite grid corrections - grid_factors = self._grid_correction( - simulation=self.simulation, - plane=self.plane, - mode_spec=self.mode_spec, + grid_factors = solver._grid_correction( + simulation=solver.simulation, + plane=solver.plane, + mode_spec=solver.mode_spec, n_complex=index_data, - direction=self.direction, + direction=solver.direction, ) # make mode solver data on the Yee grid - mode_solver_monitor = self.to_mode_solver_monitor(name=MODE_MONITOR_NAME, colocate=False) - grid_expanded = self.simulation.discretize_monitor(mode_solver_monitor) + mode_solver_monitor = solver.to_mode_solver_monitor(name=MODE_MONITOR_NAME, colocate=False) + grid_expanded = solver.simulation.discretize_monitor(mode_solver_monitor) mode_solver_data = ModeSolverData( monitor=mode_solver_monitor, - symmetry=self.simulation.symmetry, - symmetry_center=self.simulation.center, + symmetry=solver.simulation.symmetry, + symmetry_center=solver.simulation.center, grid_expanded=grid_expanded, grid_primal_correction=grid_factors[0], grid_dual_correction=grid_factors[1], @@ -1500,6 +1507,12 @@ def reduced_simulation_copy(self): This might significantly reduce upload time in the presence of custom mediums. """ + # for now, we handle EME simulation by converting to FDTD simulation + # because we can't take planar subsection of an EME simulation. + # eventually, we will convert to ModeSimulation + if isinstance(self.simulation, EMESimulation): + return self.to_fdtd_mode_solver().reduced_simulation_copy + # we preserve extra cells along the normal direction to ensure there is enough data for # subpixel extended_grid = self._get_solver_grid(keep_additional_layers=True, truncate_symmetry=False) @@ -1536,6 +1549,7 @@ def reduced_simulation_copy(self): boundary_spec=new_bspec, remove_outside_custom_mediums=True, remove_outside_structures=True, + include_pml_cells=True, ) return self.updated_copy(simulation=new_sim) diff --git a/tidy3d/plugins/smatrix/ports/rectangular_lumped.py b/tidy3d/plugins/smatrix/ports/rectangular_lumped.py index fd7db9f22..bed7ada69 100644 --- a/tidy3d/plugins/smatrix/ports/rectangular_lumped.py +++ b/tidy3d/plugins/smatrix/ports/rectangular_lumped.py @@ -194,9 +194,16 @@ def compute_current(self, sim_data: SimulationData) -> FreqDataArray: inject_center = h_cap_coords_along_injection[orth_index] # Some sanity checks, tangent H field coordinates should be directly above # and below the coordinates of the resistive sheet - assert orth_index > 0 - assert inject_center < h_coords_along_injection[orth_index] - assert h_coords_along_injection[orth_index - 1] < inject_center + error_message = ( + "Unexpected error encountered when setting up the current computation for a 'LumpedPort'. " + "If you encounter this error, please create an issue in the Tidy3D github repository." + ) + if orth_index <= 0: + raise AssertionError(error_message) + if inject_center >= h_coords_along_injection[orth_index]: + raise AssertionError(error_message) + if h_coords_along_injection[orth_index - 1] >= inject_center: + raise AssertionError(error_message) # Distance between the h1_field and h2_field, a single cell size dcap = h_coords_along_injection[orth_index] - h_coords_along_injection[orth_index - 1] diff --git a/tidy3d/plugins/smatrix/ports/wave.py b/tidy3d/plugins/smatrix/ports/wave.py index 6ee439de0..28a5d759f 100644 --- a/tidy3d/plugins/smatrix/ports/wave.py +++ b/tidy3d/plugins/smatrix/ports/wave.py @@ -118,7 +118,7 @@ def to_mode_solver(self, simulation: Simulation, freqs: FreqArray) -> ModeSolver """Helper to create a :class:`.ModeSolver` instance.""" mode_solver = ModeSolver( simulation=simulation, - plane=self, + plane=self.geometry, mode_spec=self.mode_spec, freqs=freqs, direction=self.direction, diff --git a/tidy3d/web/api/autograd/autograd.py b/tidy3d/web/api/autograd/autograd.py index 40a8bd079..686f74832 100644 --- a/tidy3d/web/api/autograd/autograd.py +++ b/tidy3d/web/api/autograd/autograd.py @@ -1,7 +1,8 @@ # autograd wrapper for web functions -import traceback +import tempfile import typing +from collections import defaultdict import numpy as np from autograd.builtins import dict as dict_ag @@ -10,17 +11,26 @@ 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 BatchData +from ..container import Batch, BatchData, Job from ..tidy3d_stub import SimulationDataType, SimulationType from ..webapi import run as run_webapi -from .utils import get_derivative_maps, split_data_list, split_list +from .utils import FieldMap, TracerKeys, get_derivative_maps # keys for data into auxiliary dictionary AUX_KEY_SIM_DATA_ORIGINAL = "sim_data" AUX_KEY_SIM_DATA_FWD = "sim_data_fwd_adjoint" +AUX_KEY_FWD_TASK_ID = "task_id_fwd" +AUX_KEY_SIM_ORIGINAL = "sim_original" +# 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?" "assignees=tylerflex&labels=adjoint&projects=&template=autograd_bug.md" @@ -29,20 +39,8 @@ MAX_NUM_TRACED_STRUCTURES = 500 - -def warn_autograd(fn_name: str, exc: Exception) -> str: - """Get warning message.""" - - exc_str = exc.__repr__() - traceback_str = "".join(traceback.format_tb(exc.__traceback__)) - - td.log.warning( - f"Autograd compatible '{fn_name}' failed, running original '{fn_name}'. " - "If you received this warning, please file an issue at the tidy3d front end with this " - f"message pasted in and we will investigate. \n\n " - f"{URL_LINK}.\n\n" - f"{exc_str} {traceback_str}.\n\n" - ) +# default value for whether to do local gradient calculation (True) or server side (False) +LOCAL_GRADIENT = False def is_valid_for_autograd(simulation: td.Simulation) -> bool: @@ -93,6 +91,7 @@ def run( worker_group: str = None, simulation_type: str = "tidy3d", parent_tasks: list[str] = None, + local_gradient: bool = LOCAL_GRADIENT, ) -> SimulationDataType: """ Submits a :class:`.Simulation` to server, starts running, monitors progress, downloads, @@ -123,6 +122,9 @@ def run( target solver version. worker_group: str = None worker group + local_gradient: bool = False + Whether to perform gradient calculation locally, requiring more downloads but potentially + more stable with experimental features. Returns ------- @@ -170,23 +172,21 @@ def run( """ if is_valid_for_autograd(simulation): - try: - return _run( - simulation=simulation, - task_name=task_name, - folder_name=folder_name, - path=path, - callback_url=callback_url, - verbose=verbose, - progress_callback_upload=progress_callback_upload, - progress_callback_download=progress_callback_download, - solver_version=solver_version, - worker_group=worker_group, - simulation_type="tidy3d_autograd", - parent_tasks=parent_tasks, - ) - except Exception as exc: - warn_autograd("web.run()", exc=exc) + return _run( + simulation=simulation, + task_name=task_name, + folder_name=folder_name, + path=path, + callback_url=callback_url, + verbose=verbose, + progress_callback_upload=progress_callback_upload, + progress_callback_download=progress_callback_download, + solver_version=solver_version, + worker_group=worker_group, + simulation_type="tidy3d_autograd", + parent_tasks=parent_tasks, + local_gradient=local_gradient, + ) return run_webapi( simulation=simulation, @@ -213,6 +213,7 @@ def run_async( verbose: bool = True, simulation_type: str = "tidy3d", parent_tasks: dict[str, list[str]] = None, + local_gradient: bool = LOCAL_GRADIENT, ) -> BatchData: """Submits a set of Union[:class:`.Simulation`, :class:`.HeatSimulation`, :class:`.EMESimulation`] objects to server, starts running, monitors progress, downloads, and loads results as a :class:`.BatchData` object. @@ -234,6 +235,9 @@ def run_async( Number of tasks to submit at once in a batch, if None, will run all at the same time. verbose : bool = True If ``True``, will print progressbars and status, otherwise, will run silently. + local_gradient: bool = False + Whether to perform gradient calculations locally, requiring more downloads but potentially + more stable with experimental features. Returns ------ @@ -252,19 +256,17 @@ def run_async( """ if is_valid_for_autograd_async(simulations): - try: - return _run_async( - simulations=simulations, - folder_name=folder_name, - path_dir=path_dir, - callback_url=callback_url, - num_workers=num_workers, - verbose=verbose, - simulation_type="tidy3d_autograd_async", - parent_tasks=parent_tasks, - ) - except Exception as exc: - warn_autograd("web.run_async()", exc=exc) + return _run_async( + simulations=simulations, + folder_name=folder_name, + path_dir=path_dir, + callback_url=callback_url, + num_workers=num_workers, + verbose=verbose, + simulation_type="tidy3d_autograd_async", + parent_tasks=parent_tasks, + local_gradient=local_gradient, + ) return run_async_webapi( simulations=simulations, @@ -281,7 +283,9 @@ def run_async( """ User-facing ``run`` and `run_async`` functions, compatible with ``autograd`` """ -def _run(simulation: td.Simulation, task_name: str, **run_kwargs) -> td.SimulationData: +def _run( + simulation: td.Simulation, task_name: str, local_gradient: bool = LOCAL_GRADIENT, **run_kwargs +) -> td.SimulationData: """User-facing ``web.run`` function, compatible with ``autograd`` differentiation.""" traced_fields_sim = setup_run(simulation=simulation) @@ -294,7 +298,8 @@ def _run(simulation: td.Simulation, task_name: str, **run_kwargs) -> td.Simulati "to the 'Simulation'. If this is unexpected, double check your objective function " "pre-processing. Running regular tidy3d simulation." ) - return _run_tidy3d(simulation, task_name=task_name, **run_kwargs) + data, _ = _run_tidy3d(simulation, task_name=task_name, **run_kwargs) + return data # will store the SimulationData for original and forward so we can access them later aux_data = {} @@ -305,6 +310,7 @@ def _run(simulation: td.Simulation, task_name: str, **run_kwargs) -> td.Simulati sim_original=simulation.to_static(), task_name=task_name, aux_data=aux_data, + local_gradient=local_gradient, **run_kwargs, ) @@ -312,7 +318,7 @@ def _run(simulation: td.Simulation, task_name: str, **run_kwargs) -> td.Simulati def _run_async( - simulations: dict[str, td.Simulation], **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.""" @@ -333,6 +339,7 @@ def _run_async( traced_fields_sim_dict, # if you pass as a kwarg it will not trace :/ sims_original=sims_original, aux_data_dict=aux_data_dict, + local_gradient=local_gradient, **run_async_kwargs, ) @@ -374,16 +381,46 @@ def _run_primitive( sim_original: td.Simulation, task_name: str, aux_data: dict, + local_gradient: bool, **run_kwargs, ) -> AutogradFieldMap: """Autograd-traced 'run()' function: runs simulation, strips tracer data, caches fwd data.""" td.log.info("running primitive '_run_primitive()'") - sim_combined = setup_fwd(sim_fields=sim_fields, sim_original=sim_original) - sim_data_combined = _run_tidy3d(sim_combined, task_name=task_name, **run_kwargs) - return postprocess_fwd( - sim_data_combined=sim_data_combined, sim_original=sim_original, aux_data=aux_data - ) + + 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( + sim_data_combined=sim_data_combined, + sim_original=sim_original, + aux_data=aux_data, + ) + + else: + 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()) + + sim_data_orig, task_id_fwd = _run_tidy3d( + sim_original, + task_name=task_name, + **run_kwargs, + ) + + # TODO: put this in postprocess? + aux_data[AUX_KEY_FWD_TASK_ID] = task_id_fwd + aux_data[AUX_KEY_SIM_DATA_ORIGINAL] = sim_data_orig + field_map = sim_data_orig.strip_traced_fields( + include_untraced_data_arrays=True, starting_path=("data",) + ) + + return field_map @primitive @@ -391,66 +428,87 @@ def _run_async_primitive( sim_fields_dict: dict[str, AutogradFieldMap], sims_original: dict[str, td.Simulation], aux_data_dict: dict[dict[str, typing.Any]], + local_gradient: bool, **run_async_kwargs, ) -> dict[str, AutogradFieldMap]: task_names = sim_fields_dict.keys() - sims_combined = {} - for task_name in task_names: - sim_fields = sim_fields_dict[task_name] - sim_original = sims_original[task_name] - sims_combined[task_name] = setup_fwd(sim_fields=sim_fields, sim_original=sim_original) + if local_gradient: + sims_combined = {} + for task_name in task_names: + sim_fields = sim_fields_dict[task_name] + sim_original = sims_original[task_name] + sims_combined[task_name] = setup_fwd(sim_fields=sim_fields, sim_original=sim_original) + + batch_data_combined, _ = _run_async_tidy3d(sims_combined, **run_async_kwargs) + + field_map_fwd_dict = {} + for task_name in task_names: + sim_data_combined = batch_data_combined[task_name] + sim_original = sims_original[task_name] + aux_data = aux_data_dict[task_name] + field_map_fwd_dict[task_name] = postprocess_fwd( + sim_data_combined=sim_data_combined, + sim_original=sim_original, + aux_data=aux_data, + ) - batch_data_combined = _run_async_tidy3d(sims_combined, **run_async_kwargs) + else: + run_async_kwargs["simulation_type"] = "autograd_fwd" + run_async_kwargs["sim_fields_keys_dict"] = {} + for task_name, sim_fields in sim_fields_dict.items(): + run_async_kwargs["sim_fields_keys_dict"][task_name] = list(sim_fields.keys()) - field_map_fwd_dict = {} - for task_name in task_names: - sim_data_combined = batch_data_combined[task_name] - sim_original = sims_original[task_name] - aux_data = aux_data_dict[task_name] - field_map_fwd_dict[task_name] = postprocess_fwd( - sim_data_combined=sim_data_combined, sim_original=sim_original, aux_data=aux_data + sims_original = { + task_name: sim.updated_copy(simulation_type="autograd_fwd", deep=False) + for task_name, sim in sims_original.items() + } + + sim_data_orig_dict, task_ids_fwd_dict = _run_async_tidy3d( + sims_original, + **run_async_kwargs, ) + field_map_fwd_dict = {} + for task_name, task_id_fwd in task_ids_fwd_dict.items(): + sim_data_orig = sim_data_orig_dict[task_name] + aux_data_dict[task_name][AUX_KEY_FWD_TASK_ID] = task_id_fwd + aux_data_dict[task_name][AUX_KEY_SIM_DATA_ORIGINAL] = sim_data_orig + field_map = sim_data_orig.strip_traced_fields( + include_untraced_data_arrays=True, starting_path=("data",) + ) + field_map_fwd_dict[task_name] = field_map + return field_map_fwd_dict -def setup_fwd(sim_fields: AutogradFieldMap, sim_original: td.Simulation) -> td.Simulation: +def setup_fwd( + sim_fields: AutogradFieldMap, + sim_original: td.Simulation, + local_gradient: bool = LOCAL_GRADIENT, +) -> td.Simulation: """Set up the combined forward simulation.""" - # make and run a sim with combined original & adjoint monitors - return sim_original.with_adjoint_monitors(sim_fields) + # if local gradient, make and run a sim with combined original & adjoint monitors + if local_gradient: + return sim_original.with_adjoint_monitors(sim_fields) + + # if remote gradient, add them later + return sim_original def postprocess_fwd( - sim_data_combined: td.SimulationData, sim_original: td.Simulation, aux_data: dict + sim_data_combined: td.SimulationData, + sim_original: td.Simulation, + aux_data: dict, ) -> AutogradFieldMap: """Postprocess the combined simulation data into an Autograd field map.""" - sim_combined = sim_data_combined.simulation - num_mnts_original = len(sim_original.monitors) - - # split the data and monitors into the original ones & adjoint gradient ones (for 'fwd') - data_original, data_fwd = split_data_list( - sim_data=sim_data_combined, num_mnts_original=num_mnts_original - ) - _, monitors_fwd = split_list(sim_combined.monitors, index=num_mnts_original) - - # reconstruct the simulation data for the user, using original sim, and data for original mnts - sim_data_original = sim_data_combined.updated_copy( - simulation=sim_original, data=data_original, deep=False + sim_data_original, sim_data_fwd = sim_data_combined.split_original_fwd( + num_mnts_original=num_mnts_original ) - # construct the 'forward' simulation and its data, which is only used for for gradient calc. - sim_fwd = sim_combined.updated_copy(monitors=monitors_fwd) - sim_data_fwd = sim_data_combined.updated_copy( - simulation=sim_fwd, - data=data_fwd, - deep=False, - ) - - # cache these two SimulationData objects for later (note: the Simulations are already inside) aux_data[AUX_KEY_SIM_DATA_ORIGINAL] = sim_data_original aux_data[AUX_KEY_SIM_DATA_FWD] = sim_data_fwd @@ -463,60 +521,104 @@ def postprocess_fwd( return data_traced +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, + ) + + """ VJP maker for ADJ pass.""" +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) + return field_map.to_autograd_field_map + + def _run_bwd( data_fields_original: AutogradFieldMap, sim_fields_original: AutogradFieldMap, sim_original: td.Simulation, task_name: str, aux_data: dict, + local_gradient: bool, **run_kwargs, ) -> typing.Callable[[AutogradFieldMap], AutogradFieldMap]: """VJP-maker for ``_run_primitive()``. Constructs and runs adjoint simulation, computes grad.""" # get the fwd epsilon and field data from the cached aux_data sim_data_orig = aux_data[AUX_KEY_SIM_DATA_ORIGINAL] - sim_data_fwd = aux_data[AUX_KEY_SIM_DATA_FWD] + # strip the sim fields keys + sim_fields_keys = list(sim_fields_original.keys()) + + if local_gradient: + sim_data_fwd = aux_data[AUX_KEY_SIM_DATA_FWD] td.log.info("constructing custom vjp function for backwards pass.") 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 = setup_adj( + sim_adj, adjoint_source_info = setup_adj( data_fields_vjp=data_fields_vjp, sim_data_orig=sim_data_orig, - sim_data_fwd=sim_data_fwd, - sim_fields_original=sim_fields_original, + sim_fields_keys=sim_fields_keys, ) - # no adjoint sources, no gradient for you :( - if not len(sim_adj.sources): - td.log.warning( - "No adjoint sources generated. " - "There is likely zero output in the data, or you have no traceable monitors. " - "As a result, the 'SimulationData' returned has no contribution to the gradient. " - "Skipping the adjoint simulation. " - "If this is unexpected, please double check the post-processing function to ensure " - "there is a path from the 'SimulationData' to the objective function return value." - ) - - # TODO: add a test for this - # construct a VJP of all zeros for all tracers in the original simulation - return {path: 0 * value for path, value in sim_fields_original.items()} - # run adjoint simulation task_name_adj = str(task_name) + "_adjoint" - sim_data_adj = _run_tidy3d(sim_adj, task_name=task_name_adj, **run_kwargs) - return postprocess_adj( - sim_data_adj=sim_data_adj, - sim_data_orig=sim_data_orig, - sim_data_fwd=sim_data_fwd, - sim_fields_original=sim_fields_original, - ) + if local_gradient: + sim_data_adj, _ = _run_tidy3d(sim_adj, task_name=task_name_adj, **run_kwargs) + + vjp_traced_fields = 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, + ) + + else: + task_id_fwd = aux_data[AUX_KEY_FWD_TASK_ID] + run_kwargs["parent_tasks"] = [task_id_fwd] + run_kwargs["simulation_type"] = "autograd_bwd" + sim_adj = sim_adj.updated_copy(simulation_type="autograd_bwd", deep=False) + + vjp_traced_fields = _run_tidy3d_bwd( + sim_adj, + task_name=task_name_adj, + adjoint_source_info=adjoint_source_info, + **run_kwargs, + ) + + return vjp_traced_fields return vjp @@ -526,6 +628,7 @@ def _run_async_bwd( sim_fields_original_dict: dict[str, AutogradFieldMap], sims_original: dict[str, td.Simulation], aux_data_dict: dict[str, dict[str, typing.Any]], + local_gradient: bool, **run_async_kwargs, ) -> typing.Callable[[dict[str, AutogradFieldMap]], dict[str, AutogradFieldMap]]: """VJP-maker for ``_run_primitive()``. Constructs and runs adjoint simulation, computes grad.""" @@ -535,10 +638,15 @@ def _run_async_bwd( # get the fwd epsilon and field data from the cached aux_data sim_data_orig_dict = {} sim_data_fwd_dict = {} + sim_fields_keys_dict = {} for task_name in task_names: aux_data = aux_data_dict[task_name] sim_data_orig_dict[task_name] = aux_data[AUX_KEY_SIM_DATA_ORIGINAL] - sim_data_fwd_dict[task_name] = aux_data[AUX_KEY_SIM_DATA_FWD] + # strip the sim fields keys + sim_fields_keys_dict[task_name] = list(sim_fields_original_dict[task_name].keys()) + + if local_gradient: + sim_data_fwd_dict[task_name] = aux_data[AUX_KEY_SIM_DATA_FWD] td.log.info("constructing custom vjp function for backwards pass.") @@ -548,39 +656,64 @@ def vjp(data_fields_dict_vjp: dict[str, AutogradFieldMap]) -> dict[str, Autograd task_names_adj = {task_name + "_adjoint" for task_name in task_names} sims_adj = {} + adjoint_source_info_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_data_fwd = sim_data_fwd_dict[task_name] - sim_fields_original = sim_fields_original_dict[task_name] + sim_fields_keys = sim_fields_keys_dict[task_name] - sim_adj = setup_adj( + sim_adj, adjoint_source_info = setup_adj( data_fields_vjp=data_fields_vjp, sim_data_orig=sim_data_orig, - sim_data_fwd=sim_data_fwd, - sim_fields_original=sim_fields_original, + sim_fields_keys=sim_fields_keys, ) sims_adj[task_name_adj] = sim_adj - + adjoint_source_info_dict[task_name_adj] = adjoint_source_info # TODO: handle case where no adjoint sources? - # run adjoint simulation - batch_data_adj = _run_async_tidy3d(sims_adj, **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] - sim_data_orig = sim_data_orig_dict[task_name] - sim_data_fwd = sim_data_fwd_dict[task_name] - sim_fields_original = sim_fields_original_dict[task_name] + if local_gradient: + # run adjoint simulation + batch_data_adj, _ = _run_async_tidy3d(sims_adj, **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] + 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_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 = postprocess_adj( - sim_data_adj=sim_data_adj, - sim_data_orig=sim_data_orig, - sim_data_fwd=sim_data_fwd, - sim_fields_original=sim_fields_original, + else: + parent_tasks = {} + for task_name_fwd, task_name_adj in zip(task_names, 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 = { + task_name: sim.updated_copy(simulation_type="autograd_bwd", deep=False) + for task_name, sim in sims_adj.items() + } + sim_fields_vjp_dict_adj_keys = _run_async_tidy3d_bwd( + simulations=sims_adj, + adjoint_source_info_dict=adjoint_source_info_dict, + **run_async_kwargs, ) - sim_fields_vjp_dict[task_name] = sim_fields_vjp + + # 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): + sim_fields_vjp_dict[task_name_fwd] = sim_fields_vjp_dict_adj_keys[task_name_adj] return sim_fields_vjp_dict @@ -590,9 +723,8 @@ def vjp(data_fields_dict_vjp: dict[str, AutogradFieldMap]) -> dict[str, Autograd def setup_adj( data_fields_vjp: AutogradFieldMap, sim_data_orig: td.SimulationData, - sim_data_fwd: td.SimulationData, - sim_fields_original: AutogradFieldMap, -) -> td.Simulation: + sim_fields_keys: list[tuple], +) -> tuple[td.Simulation, AdjointSourceInfo]: """Construct an adjoint simulation from a set of data_fields for the VJP.""" td.log.info("Running custom vjp (adjoint) pipeline.") @@ -607,31 +739,35 @@ def setup_adj( # make adjoint simulation from that SimulationData data_vjp_paths = set(data_fields_vjp.keys()) - sim_adj = sim_data_vjp.make_adjoint_sim( - data_vjp_paths=data_vjp_paths, adjoint_monitors=sim_data_fwd.simulation.monitors + + num_monitors = len(sim_data_orig.simulation.monitors) + adjoint_monitors = sim_data_orig.simulation.with_adjoint_monitors(sim_fields_keys).monitors[ + num_monitors: + ] + + sim_adj, adjoint_source_info = sim_data_vjp.make_adjoint_sim( + data_vjp_paths=data_vjp_paths, adjoint_monitors=adjoint_monitors ) td.log.info(f"Adjoint simulation created with {len(sim_adj.sources)} sources.") - return sim_adj + return sim_adj, adjoint_source_info def postprocess_adj( sim_data_adj: td.SimulationData, sim_data_orig: td.SimulationData, sim_data_fwd: td.SimulationData, - sim_fields_original: AutogradFieldMap, + 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.""" # map of index into 'structures' to the list of paths we need vjps for - sim_vjp_map = {} - for _, structure_index, *structure_path in sim_fields_original.keys(): + sim_vjp_map = defaultdict(list) + for _, structure_index, *structure_path in sim_fields_keys: structure_path = tuple(structure_path) - if structure_index in sim_vjp_map: - sim_vjp_map[structure_index].append(structure_path) - else: - sim_vjp_map[structure_index] = [structure_path] + sim_vjp_map[structure_index].append(structure_path) # store the derivative values given the forward and adjoint data sim_fields_vjp = {} @@ -642,6 +778,14 @@ def postprocess_adj( fld_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 + + fld_adj = fld_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 @@ -655,8 +799,7 @@ def postprocess_adj( # todo: handle multi-frequency, move to a property? frequencies = {src.source_time.freq0 for src in sim_data_adj.simulation.sources} frequencies = list(frequencies) - assert len(frequencies) == 1, "Multiple adjoint freqs found" - freq_adj = frequencies[0] + freq_adj = frequencies[0] or None eps_in = np.mean(structure.medium.eps_model(freq_adj)) eps_out = np.mean(sim_data_orig.simulation.medium.eps_model(freq_adj)) @@ -691,17 +834,98 @@ def postprocess_adj( """ The fundamental Tidy3D run and run_async functions used above. """ -def _run_tidy3d(simulation: td.Simulation, task_name: str, **run_kwargs) -> td.SimulationData: +def parse_run_kwargs(**run_kwargs): + """Parse the ``run_kwargs`` to extract what should be passed to the ``Job`` initialization.""" + job_fields = list(Job._upload_fields) + ["solver_version"] + job_init_kwargs = {k: v for k, v in run_kwargs.items() if k in job_fields} + return job_init_kwargs + + +def _run_tidy3d( + simulation: td.Simulation, task_name: str, **run_kwargs +) -> (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) + td.log.info(f"running {job.simulation_type} simulation with '_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() + return data, job.task_id + + +def _run_tidy3d_bwd( + simulation: td.Simulation, task_name: str, adjoint_source_info: AdjointSourceInfo, **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() + return get_vjp_traced_fields(task_id_adj=job.task_id, verbose=job.verbose) + + +def _run_async_tidy3d( + simulations: dict[str, td.Simulation], **run_kwargs +) -> tuple[BatchData, dict[str, str]]: + """Run a simulation without any tracers using regular web.run().""" + batch_init_kwargs = parse_run_kwargs(**run_kwargs) + path_dir = run_kwargs.pop("path_dir", None) + batch = Batch(simulations=simulations, **batch_init_kwargs) + td.log.info(f"running {batch.simulation_type} batch with '_run_async_tidy3d()'") + + if batch.simulation_type == "autograd_fwd": + verbose = run_kwargs.get("verbose", False) + # Need to upload to get the task_ids + sims = { + task_name: sim.updated_copy(simulation_type="autograd_fwd", deep=False) + for task_name, sim in batch.simulations.items() + } + batch = batch.updated_copy(simulations=sims) + + batch.upload() + task_ids = {key: job.task_id for key, job in batch.jobs.items()} + for task_name, sim_fields_keys in run_kwargs["sim_fields_keys_dict"].items(): + task_id = task_ids[task_name] + upload_sim_fields_keys(sim_fields_keys, task_id=task_id, verbose=verbose) + + if path_dir: + batch_data = batch.run(path_dir) + else: + batch_data = batch.run() + + task_ids = {key: job.task_id for key, job in batch.jobs.items()} + return batch_data, task_ids + + +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().""" - td.log.info("running regular simulation with '_run_tidy3d()'") - # TODO: set task_type to "tidy3d adjoint autograd?" - data = run_webapi(simulation, task_name=task_name, **run_kwargs) - return data - - -def _run_async_tidy3d(simulations: dict[str, td.Simulation], **run_async_kwargs) -> BatchData: - """Run a simulation without any tracers using regular ``web.run_async``.""" - td.log.info("running batch of simulations with '_run_async_tidy3d()'") - # TODO: set task_type to "tidy3d adjoint autograd?" - batch_data = run_async_webapi(simulations, **run_async_kwargs) - return batch_data + + batch_init_kwargs = parse_run_kwargs(**run_kwargs) + _ = run_kwargs.pop("path_dir") + 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() + + vjp_traced_fields_dict = {} + for task_name, job in batch.jobs.items(): + task_id = job.task_id + vjp = get_vjp_traced_fields(task_id_adj=task_id, verbose=batch.verbose) + vjp_traced_fields_dict[task_name] = vjp + + return vjp_traced_fields_dict diff --git a/tidy3d/web/api/autograd/utils.py b/tidy3d/web/api/autograd/utils.py index 13367330d..1d5cd136c 100644 --- a/tidy3d/web/api/autograd/utils.py +++ b/tidy3d/web/api/autograd/utils.py @@ -1,32 +1,14 @@ # utility functions for autograd web API +from __future__ import annotations import typing -import tidy3d as td - -""" generic data manipulation """ - - -def split_data_list(sim_data: td.SimulationData, num_mnts_original: int) -> tuple[list, list]: - """Split data list into original, adjoint field, and adjoint permittivity.""" - - data_all = list(sim_data.data) - num_mnts_adjoint = (len(data_all) - num_mnts_original) // 2 - - td.log.info( - f" -> {num_mnts_original} monitors, {num_mnts_adjoint} adjoint field monitors, {num_mnts_adjoint} adjoint eps monitors." - ) - - data_original, data_adjoint = split_list(data_all, index=num_mnts_original) - - return data_original, data_adjoint - - -def split_list(x: list[typing.Any], index: int) -> (list[typing.Any], list[typing.Any]): - """Split a list at a given index.""" - x = list(x) - return x[:index], x[index:] +import pydantic as pd +import tidy3d as td +from tidy3d.components.autograd.types import AutogradFieldMap, dict_ag +from tidy3d.components.base import Tidy3dBaseModel +from tidy3d.components.types import ArrayLike, tidycomplex """ E and D field gradient map calculation helpers. """ @@ -83,3 +65,46 @@ def get_field_key(dim: str, fld_data: typing.Union[td.FieldData, td.Permittivity mult = cmp_1 * cmp_2 field_components[key_1] = mult return fld_1.updated_copy(**field_components) + + +class Tracer(Tidy3dBaseModel): + """Class to store a single traced field.""" + + path: tuple[typing.Any, ...] = pd.Field( + ..., + title="Path to the traced object in the model dictionary.", + ) + + data: typing.Union[float, tidycomplex, ArrayLike] = pd.Field(..., title="Tracing data") + + +class FieldMap(Tidy3dBaseModel): + """Class to store a collection of traced fields.""" + + tracers: tuple[Tracer, ...] = pd.Field( + ..., + title="Collection of Tracers.", + ) + + @property + def to_autograd_field_map(self) -> AutogradFieldMap: + """Convert to ``AutogradFieldMap`` autograd dictionary.""" + return dict_ag({tracer.path: tracer.data for tracer in self.tracers}) + + @classmethod + def from_autograd_field_map(cls, autograd_field_map) -> FieldMap: + """Initialize from an ``AutogradFieldMap`` autograd dictionary.""" + tracers = [] + for path, data in autograd_field_map.items(): + tracers.append(Tracer(path=path, data=data)) + + return cls(tracers=tuple(tracers)) + + +class TracerKeys(Tidy3dBaseModel): + """Class to store a collection of tracer keys.""" + + keys: tuple[tuple[typing.Any, ...], ...] = pd.Field( + ..., + title="Collection of tracer keys.", + ) diff --git a/tidy3d/web/api/material_fitter.py b/tidy3d/web/api/material_fitter.py index d73d82e47..0d91f825e 100644 --- a/tidy3d/web/api/material_fitter.py +++ b/tidy3d/web/api/material_fitter.py @@ -73,8 +73,10 @@ def submit(cls, fitter: DispersionFitter, options: FitterOptions) -> MaterialFit options: FitterOptions fitter options """ - assert fitter - assert options + if not isinstance(fitter, DispersionFitter): + raise TypeError(f"fitter must be an instance of 'DispersionFitter', got {type(fitter)}") + if not isinstance(options, FitterOptions): + raise TypeError(f"options must be an instance of 'FitterOptions', got {type(options)}") data = np.asarray(list(zip(fitter.wvl_um, fitter.n_data, fitter.k_data))) with tempfile.NamedTemporaryFile(suffix=".csv") as temp: np.savetxt(temp, data, delimiter=",", header="Wavelength,n,k") diff --git a/tidy3d/web/api/mode.py b/tidy3d/web/api/mode.py index 4e759fde8..f08303b52 100644 --- a/tidy3d/web/api/mode.py +++ b/tidy3d/web/api/mode.py @@ -626,6 +626,7 @@ def get_result( self.mode_solver._cached_properties["data_raw"] = data # Perform symmetry expansion + self.mode_solver._cached_properties.pop("data", None) return self.mode_solver.data def get_log( diff --git a/tidy3d/web/api/tidy3d_stub.py b/tidy3d/web/api/tidy3d_stub.py index acec41677..f36983ca8 100644 --- a/tidy3d/web/api/tidy3d_stub.py +++ b/tidy3d/web/api/tidy3d_stub.py @@ -139,6 +139,8 @@ def validate_pre_upload(self, source_required) -> None: """Perform some pre-checks on instances of component""" if isinstance(self.simulation, Simulation): self.simulation.validate_pre_upload(source_required) + elif isinstance(self.simulation, EMESimulation): + self.simulation.validate_pre_upload() class Tidy3dStubData(BaseModel, TaskStubData): diff --git a/tidy3d/web/api/webapi.py b/tidy3d/web/api/webapi.py index 45304cf08..b7dd5291b 100644 --- a/tidy3d/web/api/webapi.py +++ b/tidy3d/web/api/webapi.py @@ -437,32 +437,47 @@ def monitor_preprocess() -> None: time.sleep(REFRESH_TIME) # while running but percentage done is available - if verbose and task_type == "FDTD": + if verbose: # verbose case, update progressbar console.log("running solver") - with Progress(console=console) as progress: - pbar_pd = progress.add_task("% done", total=100) - perc_done, _ = get_run_info(task_id) + if task_type == "FDTD": + with Progress(console=console) as progress: + pbar_pd = progress.add_task("% done", total=100) + perc_done, _ = get_run_info(task_id) + + while ( + perc_done is not None and perc_done < 100 and get_status(task_id) == "running" + ): + perc_done, field_decay = get_run_info(task_id) + new_description = f"solver progress (field decay = {field_decay:.2e})" + progress.update(pbar_pd, completed=perc_done, description=new_description) + time.sleep(RUN_REFRESH_TIME) - while perc_done is not None and perc_done < 100 and get_status(task_id) == "running": perc_done, field_decay = get_run_info(task_id) - new_description = f"solver progress (field decay = {field_decay:.2e})" - progress.update(pbar_pd, completed=perc_done, description=new_description) - time.sleep(RUN_REFRESH_TIME) - - perc_done, field_decay = get_run_info(task_id) - if perc_done is not None and perc_done < 100 and field_decay > 0: - console.log(f"early shutoff detected at {perc_done:1.0f}%, exiting.") + if perc_done is not None and perc_done < 100 and field_decay > 0: + console.log(f"early shutoff detected at {perc_done:1.0f}%, exiting.") - new_description = f"solver progress (field decay = {field_decay:.2e})" - progress.update(pbar_pd, completed=100, refresh=True, description=new_description) - - # TODO: progressbar for EME + new_description = f"solver progress (field decay = {field_decay:.2e})" + progress.update(pbar_pd, completed=100, refresh=True, description=new_description) + elif task_type == "EME": + with Progress(console=console) as progress: + pbar_pd = progress.add_task("% done", total=100) + perc_done, _ = get_run_info(task_id) + + while ( + perc_done is not None and perc_done < 100 and get_status(task_id) == "running" + ): + perc_done, _ = get_run_info(task_id) + new_description = "solver progress" + progress.update(pbar_pd, completed=perc_done, description=new_description) + time.sleep(RUN_REFRESH_TIME) + + perc_done, _ = get_run_info(task_id) + new_description = "solver progress" + progress.update(pbar_pd, completed=100, refresh=True, description=new_description) else: # non-verbose case, just keep checking until status is not running or perc_done >= 100 - if verbose: - console.log("running solver") perc_done, _ = get_run_info(task_id) while perc_done is not None and perc_done < 100 and get_status(task_id) == "running": perc_done, field_decay = get_run_info(task_id)