PyTorch Backend Basics
[1]:
import optiland.backend as be
from optiland.samples.objectives import CookeTriplet
from optiland.psf import FFTPSF
[2]:
be.set_backend("torch") # all calculations will be done using torch
be.set_precision("float32") # or "float64" for double precision
be.grad_mode.enable() # enable gradient calculation. Can also be set to "disable" to save memory if gradients are not needed.
be.set_device("cuda") # run on GPU if available, otherwise use "cpu" option
From this point on, all calculations are identical regardless of the backend used (NumPy or PyTorch).
[3]:
lens = CookeTriplet()
[4]:
psf = FFTPSF(lens, field=(0, 1), wavelength=0.55)
psf.view(projection="2d", num_points=256)
C:\Users\kdani\AppData\Local\Temp\ipykernel_17828\2951478777.py:2: UserWarning: The PSF view has a high oversampling factor (3.46). Results may be inaccurate.
psf.view(projection="2d", num_points=256)
[4]:
(<Figure size 700x550 with 2 Axes>,
<Axes: title={'center': 'ScalarFFT PSF'}, xlabel='X (µm)', ylabel='Y (µm)'>)
