DiffASM

DiffASM is a differentiable angular spectrum method framework for 6G electromagnetic field simulation. It targets the regime between high-speed but approximate geometric methods and high-fidelity but expensive full-wave solvers, with an emphasis on GPU acceleration, gradient-friendly operators, and backend comparisons under one interface.

What Is ASM?

The angular spectrum method (ASM) propagates a wavefield by decomposing it into spatial-frequency components, applying a propagation phase factor in the spectral domain, and transforming the field back to space after each step. In practice, ASM is attractive because it captures diffraction and interference much more faithfully than geometric optics while remaining far cheaper than a full time-domain Maxwell solver.

For wireless sensing and 6G scene simulation, ASM is often a useful middle ground between ray tracing and full-wave solvers:

| Method | Physical Model | Diffraction / Interference | Cost | Typical Use | |---|---|---|---|---| | Ray tracing | Geometric optics / paths | Weak or approximate | Low to moderate | Large scenes, path analysis, coverage | | ASM | Wave propagation on a discretized field | Stronger than ray tracing, naturally wave-based | Moderate | Wave-aware propagation and inverse design | | FDTD | Full-wave time-domain Maxwell solver | Full electromagnetic wave physics | High to very high | High-fidelity EM simulation |

ASM is therefore attractive when ray tracing is too geometric but FDTD is too expensive.

Why Differentiable ASM?

One of the main reasons to use ASM in modern wireless research is that it can be made differentiable in a relatively clean way. Because propagation is expressed as structured linear field operators, gradients with respect to source parameters, field samples, or scene-dependent quantities are much more accessible than in many traditional high-frequency simulation pipelines.

This matters for:

• inverse design
• source and beam optimization
• learned scene parameters
• differentiable sensing and communication pipelines
• gradient-based calibration and reconstruction

Why DiffASM?

DiffASM is a GPU-accelerated, differentiable ASM project aimed at 6G research workflows where forward simulation and gradients both matter. The codebase provides:

• A unified simulator for multiple propagation backends.
• Differentiable field operations and gradient examples.
• Optional Slang/CUDA kernels for GPU acceleration.
• Side-by-side backend comparison plots for the same scene configuration.

The project currently focuses on 2D scalar field propagation in a vertical slice of the environment, where y is the transverse axis and z is the forward propagation axis.

Current Scope

The implementation includes:

• Scene construction with material obstacles.
• Source generation for Gaussian, planar, point, and beam-shaped launches.
• Multiple propagation backends under the same simulation interface.
• Optional GPU kernels through slangtorch.
• A reference comparison script in main.py that renders backend outputs and gradient fields.

The current comparison figure is shown below:

Backends

DiffASM currently exposes four propagation backends:

FFT: Fast Fourier Transform

The FFT backend is the standard ASM implementation. It propagates each step in the spectral domain using a transfer function and inverse FFT.

Advantages:

• Fastest and most direct ASM formulation on a uniform grid.
• Good numerical reference for the other backends.
• Clean fit for differentiable field propagation.

Tradeoffs:

• Assumes a uniform discretization of the transverse axis.
• Still requires discretizing the full field over the simulation grid.

DI: Direct Integration

The DI (Direct Integration) backend applies the same propagation model as ASM, but materializes the step operator in the spatial domain as a dense linear map.

Advantages:

• Useful as a dense operator reference under the same physical model.
• Easier to inspect as an explicit linear map.

Tradeoffs:

• More expensive than FFT for large grids.
• Less scalable because the propagation step is represented densely.

GBD: Gaussian Beam Decomposition

The GBD (Gaussian Beam Decomposition) backend projects the field into a Gaussian beamlet subspace and propagates inside that reduced representation.

Advantages:

• Compresses propagation through a structured beam basis.
• Can be a useful reduced-order approximation when beam-like structure is strong.

Tradeoffs:

• Approximate by construction.
• Accuracy depends on the beamlet basis and scene content.
• More sensitive than FFT to basis design and projection quality.

NUFFT: Non-Uniform Fast Fourier Transform

The NUFFT (Non-Uniform Fast Fourier Transform) backend is intended as a path toward nonuniform transverse sampling while still leveraging FFT-style propagation internally.

Advantages:

• Provides a bridge between uniform FFT propagation and nonuniform sampling strategies.
• Useful for experiments with resampling and generalized discretizations.

Tradeoffs:

• In the current codebase, the main use case still starts from a uniform simulation grid.
• Adds interpolation and operator construction overhead compared with plain FFT.

Differentiability

DiffASM is designed for field-level differentiation. The project already includes examples of:

• Forward propagation through linear wave operators.
• Source-shift sensitivity analysis.
• AD-style tangent propagation.
• Finite-difference gradient comparison.

This makes it suitable for inverse problems, optimization, learned scene parameters, and differentiable wireless simulation workflows where gradients with respect to source or scene variables are needed.

Limitations

Most of the important limitations come from ASM itself as a method, not just from this particular codebase.

Method-Level Limitations Of ASM

• ASM is fundamentally a field propagation method built around a discretized wavefield, so the full field has to be sampled on a grid. As the domain grows, memory and runtime grow with it.
• Standard ASM is most natural for predominantly one-way propagation along a chosen marching axis. It is less natural than full-wave methods for truly bidirectional, strongly multi-bounce, or fully volumetric electromagnetic interactions.
• The common scalar ASM formulation is not a full vector Maxwell solver. It does not natively capture the full electromagnetic vector structure of the field.
• Polarization handling is therefore limited or absent unless an additional vector-field formulation is built on top.
• ASM is much harder to scale to 3D than to 2D. In 2D, one propagation step typically uses a 1D FFT; in 3D, that becomes a 2D FFT per step, and the field tensor itself becomes much larger.
• Because the method is grid-based, obstacle geometry and field detail are only represented up to the chosen discretization. Resolution therefore directly affects fidelity.
• ASM is not a substitute for full-wave vector solvers such as FDTD when near-field coupling, broadband transients, polarization, or strongly resonant electromagnetic behavior are dominant.

Current Implementation Limits In DiffASM

• The present codebase focuses on 2D scalar propagation in a y-z slice.
• The main workflow is organized around edge-launched, forward propagation along the z axis.
• The current public examples and comparison scripts are built around that 2D scalar setting rather than a general 3D vector EM stack.

Project Structure

Key modules:

• main.py: comparison figure generation and backend timing display.
• engine/simulator.py: main simulation entry point.
• engine/propagation: propagation backends and transfer operators.
• engine/results.py: simulation result container.
• scene: scene, source, grid, and object definitions.
• ops: field, obstacle, transform, and window operations.
• kernels: optional Slang/CUDA kernels for accelerated execution.

Installation

This repository does not currently ship with a packaged pyproject.toml or requirements.txt, so installation is manual.

Core Requirements

You will need:

• Python 3.10 or newer
• torch
• numpy
• matplotlib

Optional but useful:

• scipy for Hankel-function evaluation in the Green's-function path
• slangtorch for GPU kernel compilation and execution
• A CUDA-capable GPU if you want accelerated Slang kernels

Example Setup

conda create -n diffasm python=3.11 -y
conda activate diffasm
pip install torch numpy matplotlib scipy

If you want GPU kernel support as well, install slangtorch in an environment where CUDA and PyTorch are already configured correctly.

Running The Demo

From the repository root:

python main.py

This generates the comparison figure at:

artifacts/figures/main_figure.png

When To Use DiffASM

DiffASM is a good fit when you want:

• Wave-based propagation beyond geometric ray tracing.
• Faster experimentation than FDTD for 6G scene studies.
• Field-level differentiation for optimization or learning.
• A backend comparison framework for propagation research.

It is a weaker fit when you need:

• Full vector electromagnetic modeling.
• Polarization-aware simulation.
• Strongly resonant or broadband transient full-wave physics.
• Large-scale 3D domains with strict physical completeness.

License

This project is released under the MIT License. See LICENSE.