Skip to content

rafael-fuente/diffractsim

Repository files navigation

Diffractsim: A diffraction simulator for exploring and visualizing physical optics

DOI

animation

Flexible, and easy-to-use Python diffraction simulator that focuses on visualizing physical optics phenomena. The simulator uses mainly scalar diffraction techniques for light propagation, provides a nice interface for simulation set up, and includes several plotting options, counting with CIE Color matching functions for accurate color reproduction. The implemented methods are differentiable when using the JAX-backend, enabling optimization and design of the optical setups.

The basic use of this simulator using the angular spectrum method is described in this article. Take a look to the videos to see the animated simulations!

Features

  • Arbitrary apertures and light spectrums
  • Full-path optical propagation and arbitrary zoom in the region of interest
  • Lenses
  • Phase holograms generation and reconstruction
  • GPU acceleration
  • Differentiable propagation
  • Incoherent Light (coming soon)

Installation

pip install diffractsim

Alternatively, to download the examples and the apertures as well, you can also build from source by cloning the repository and running from the main folder project on the command prompt:

python install .

Examples

To perform the simulations, just run from the examples subdirectory the corresponding Python scripts on the command prompt. To compute your own diffraction pattern, you'll need to specify in the script the aperture as an image and input its size.

python hexagon_monochromatic.py

N|Solid

python hexagon_polychromatic.py

N|Solid

python rectangular_grating_small.py

N|Solid

python rectangular_grating_big.py

N|Solid

python bahtinov_mask.py

N|Solid

python rings.py

N|Solid

python hexagonal_grating.py

N|Solid

python diffraction_text.py

N|Solid

The examples from the video about diffraction with lenses can be found here.

GPU acceleration requires having either CuPy or JAX installed along with CUDA on your computer. To use GPU acceleration with CuPy in your simulations, after import diffractsim add the line:

diffractsim.set_backend("CUDA")

Or when using JAX:

diffractsim.set_backend("JAX")

JAX, in addition to offering GPU acceleration, will allow differentiable programming. Cupy, JAX and CUDA aren't required to install and use this package, but they can offer a significant speed boost.

The first GPU accelerated run can be slow because Python is caching the required functions. The next time it can be about 10x and 100x faster than a CPU backend depending on your GPU. The speed boost raises as the grid gets larger.

Diffractsim can also be used to compute and visualize longitudinal profiles. Since the computation of each propagation distance is independent, it can be fully parallelized, and therefore GPU use is highly recommended.

python lens_longitudinal_profile.py

N|Solid

python axicon_longitudinal_profile.py

N|Solid

The problem of phase retrieval is a classic one in optics and arises when one is interested in retrieving the wavefront from two intensity measurements acquired in two different planes. Diffractsim provides a simple implementation of this problem.

In the following example, the GitHub logo is recovered at the Fourier plane from a coherently illuminated square-shaped aperture. The script generates a phase mask, which is stored as an image using an HSV colormap and then placed on the aperture to reconstruct the desired image.

python phase_hologram_github_logo_generation_and_reconstruction.py

animation