Fea nn functional em

[loliverhennigh]

PhysicsNeMo Pull Request

Description

Getting in old Electromagnetic FDTD work we did in PhysicsNemo. This is a bit of a test for the new functional structure we have.

[greptile-apps]

Greptile Summary

This PR adds electromagnetic FDTD (Finite-Difference Time-Domain) functionality to PhysicsNeMo by implementing an electric_field_update functional for 3D Yee grid simulations with periodic boundary conditions.

Key Changes

• Implements electric field update for one FDTD timestep with both PyTorch and Warp backends
• Follows the new functional structure pattern with FunctionSpec for multi-backend dispatch
• Supports scalar and spatial material fields (permittivity eps and conductivity sigma_e)
• Includes optional impressed current sources with offset positioning
• Provides both in-place and out-of-place update modes
• Full backward pass support for differentiable electromagnetics simulations

Implementation Quality

The implementation demonstrates high quality:

• Comprehensive documentation: The main functional has excellent docstrings with LaTeX equations explaining the FDTD update equations and material coefficients
• Robust validation: Extensive input validation with clear error messages
• Thorough testing: Test suite covers forward/backward parity between implementations, all material mode combinations, in-place behavior, edge cases, and gradient handling
• Performance optimization: Warp implementation includes 8 specialized kernels for different scalar/field combinations to minimize overhead
• Proper abstraction: Clean separation between shared utilities, torch implementation, and warp implementation

Important Files Changed

Filename	Overview
physicsnemo/nn/functional/electromagnetism/electric_field_update/electric_field_update.py	Core FDTD electric field update functional with FunctionSpec pattern, comprehensive docstring with LaTeX equations, includes make_inputs for benchmarking
physicsnemo/nn/functional/electromagnetism/electric_field_update/_torch_impl.py	Pure PyTorch implementation of electric field update with periodic curl computation and impressed current handling
physicsnemo/nn/functional/electromagnetism/electric_field_update/_warp_impl.py	Warp-accelerated implementation with 8 specialized kernels (4 forward, 4 backward) for different scalar/field material combinations, includes custom op registration
physicsnemo/nn/functional/electromagnetism/electric_field_update/utils.py	Shared validation and normalization utilities for electric field update implementations, handles material field conversion and input validation
test/nn/functional/electromagnetism/test_electric_field_update.py	Comprehensive test suite verifying warp/torch parity for forward/backward passes across all material modes, current variants, in-place behavior, and edge cases

_{Last reviewed commit: a17aa5e}

[greptile-apps]

_{12 files reviewed, no comments}

_{Edit Code Review Agent Settings | Greptile}