pynx.wavefront: Basic wavefront propagation (near and far field). Fresnel zone plate simulations

Description

This module allows to propagate 2D wavefront using either:

  • near field propagation

  • far field propagation

  • continuous propagation using the fractional Fourier Transform approach

  • Sub-module pynx.wavefront.fzp can be used to calculate the coherent illumination from a Fresnel Zone Plate

  • Sub-module pynx.wavefront.fresnel can be used to simulate

Calculations can be done using a GPU (OpenCL or CUDA), and use an ‘operator’ approach, where operations on a given wavefront can be simply written by multiplying that wavefront by the corresponding operator or by a series of operators.

API Reference

Note that the Python API uses an ‘operator’ approach, to enable writing complex operations in a more mathematical and natural way.

Wavefront

exception pynx.wavefront.wavefront.UserWarningWavefrontNearFieldPropagation
class pynx.wavefront.wavefront.Wavefront(d=None, z=0, pixel_size=5.5e-05, wavelength=1.5498e-10, copy_d=True)

A 2D Wavefront object

Create the Wavefront object.

Parameters:

  • z – current position of 2D Wavefront along direction propagation, in meters

  • pixel_size – current pixelsize at z, in meters

  • d –

    the original data array at z - will be converted to complex64 if needed. This can either be : - a 2D array - a 3D array which will be treated as a stack of 2D wavefront for multi-views/modes propagation. - None: will be initialized to a 512x512 data array filled with 1. - a string to use one image from either scipy or scikit-image data. These will be truncated to 512x512. Available string values are ‘ascent’, ‘face’, ‘camera’, ‘hubble’, ‘immunohistochemistry’. In the case of RGB images, all 3 components are loaded.

    The data should have its center at (0,0) to avoid requiring fftshifts during propagation.

    Any 2D data will be converted to 3D with a shape of (1, ny, nx)

  • wavelength – X-ray wavelength in meters.

copy(copy_d=True)

Creates a copy (without any reference passing) of this object, unless copy_d is False.

Parameters:

copy_d – if False, the new object will be a shallow copy, with d copied as a reference.

Returns:

a copy of the object.

get(shift=False)

Get the wavefront data array. This will automatically get the latest data, either from GPU or from the host memory, depending where the lat changes were made.

Parameters:

shift – if True, the data array will be fft-shifted so that the center of the data is in the center of the array, rather than in the corner (the default).

Returns:

the 2D or 3D (stack of wavefront) numpy data array

get_x_y()

Get 1D arrays of x and y coordinates, taking into account the pixel size. The arrays are centered at (0,0) - i.e. with the origin in the corner for FFT puroposes. x is an horizontal vector and y vertical.

Returns:

a tuple (x, y) of 2D numpy arrays

set(d, shift=False)

Set the wavefront data array.

Parameters:

  • d – the data array (complex64 numpy array)

  • shift – if True, the data array will be fft-shifted so that the center of the stored data is in the corner of the array (0,0). [default: the array is already shifted]

Returns:

nothing

Wavefront operators

This section lists the OpenCL operators, but identical operators are available for CUDA. The best set of operators (determined by querying available languages and devices) is imported automatically when performing from pynx.wavefront import * or from pynx.cdi.wavefront import *

class pynx.wavefront.cl_operator.BackPropagatePaganin(dz=1, delta=1e-06, beta=1e-09)

Back-propagation algorithm using the single-projection approach. Ref: Paganin et al., Journal of microscopy 206 (2002), 33–40. (DOI: 10.1046/j.1365-2818.2002.01010.x)

This operator is special since it will use only the intensity of the wavefront. Therefore it will first take the square modulus of the wavefront it is applied to, discarding any phase information. The result of the transformation is the calculated wavefront at the sample position, i.e. if T(r) is the estimated thickness of the sample, it is exp(-mu * T - 2*pi/lambda * T)

Parameters:

  • dz – distance between sample and detector (meter)

  • delta – real part of the refraction index, n = 1 - delta + i * beta

  • beta – imaginary part of the refraction index

op(w: Wavefront)
Parameters:

w – the Wavefront object this operator applies to

Returns:

the updated Wavefront object

class pynx.wavefront.cl_operator.CircularMask(radius, invert=False)

Multiplies the wavefront by a binary circular mask with a given radius

Parameters:

  • radius – radius of the mask (in meters)

  • invert – if True, the inside of the circle will be masked rather than the outside

op(w: Wavefront)
Parameters:

w – the Wavefront object this operator applies to

Returns:

the updated Wavefront object

class pynx.wavefront.cl_operator.FT(processing_unit=None)

Forward Fourier transform.

op(w: Wavefront)
Parameters:

w – the Wavefront object this operator applies to

Returns:

the updated Wavefront object

class pynx.wavefront.cl_operator.FreePU(processing_unit=None)

Operator freeing OpenCL memory. The gpyfft data reference in self.processing_unit is removed, as well as any OpenCL pyopencl.array.Array attribute in the supplied wavefront.

op(w: Wavefront)
Parameters:

w – the Wavefront object this operator applies to

Returns:

the updated Wavefront object

timestamp_increment(cdi)

Increment the timestamp counter corresponding to the processing language used (OpenCL or CUDA) Virtual, must be derived.

Parameters:

w – the object (e.g. wavefront) the operator will be applied to.

Returns:

class pynx.wavefront.cl_operator.IFT(processing_unit=None)

Inverse Fourier transform

op(w: Wavefront)
Parameters:

w – the Wavefront object this operator applies to

Returns:

the updated Wavefront object

class pynx.wavefront.cl_operator.PropagateFRT(dz, forward=True)
Wavefront propagator using a fractional Fourier transform

References:

    1. Mas, J. Garcia, C. Ferreira, L.M. Bernardo, and F. Marinho, Optics Comm 164, 233 (1999)

    1. García, D. Mas, and R.G. Dorsch, Applied Optics 35, 7013 (1996)

Notes:

  • the computed phase is only ‘valid’ for dz<N*pixel_size**2/wavelength, i.e near-to-medium field (I am not sure this is strictly true - the phase will vary quickly from pixel to pixel as for a spherical wave propagation but the calculation is still correct)

  • the amplitude remains correct even in the far field

  • only forward propagation works correctly (z>0)

Parameters:

  • dz – the propagation distance

  • forward – if True (default), forward propagation. Note that backward propagation is purely experimental and not fully correct.

op(w: Wavefront)
Parameters:

w – the Wavefront object this operator applies to

Returns:

the updated Wavefront object

class pynx.wavefront.cl_operator.PropagateFarField(dz, forward=True, no_far_field_quadphase=True)

Far field propagator

Parameters:

  • dz – propagation distance in meters

  • forward – if True, forward propagation, otherwise backward

  • no_far_field_quadphase – if True (default), no quadratic phase is applied in the far field

op(w: Wavefront)
Parameters:

w – the Wavefront object this operator applies to

Returns:

the updated Wavefront object

class pynx.wavefront.cl_operator.PropagateNearField(dz, magnification=None, verbose=False)

Near field propagator

Parameters:

  • dz – propagation distance (in meters)

  • magnification – if not None, the destination pixel size will will be multiplied by this factor. Note that it creates important restrictions on the validity domain of the calculation, both near and far.

  • verbose – if True, prints the propagation limit for a valid phase.

op(w: Wavefront)
Parameters:

w – the Wavefront object this operator applies to

Returns:

the updated Wavefront object

class pynx.wavefront.cl_operator.QuadraticPhase(factor, scale=1)

Operator applying a quadratic phase factor

Application of a quadratic phase factor, and optionally a scale factor.

The actual factor is: \(scale * e^{i * factor * (ix^2 + iy^2)}\) where ix and iy are the integer indices of the pixels

Parameters:

  • factor – the factor for the phase calculation.

  • scale – the data will be scaled by this factor. Useful to normalize after a Fourier transform, without accessing twice the array data.

op(w: Wavefront)
Parameters:

w – the Wavefront object this operator applies to

Returns:

the updated Wavefront object

class pynx.wavefront.cl_operator.RectangularMask(width, height, invert=False)

Multiplies the wavefront by a rectangular mask with a given width and height

Parameters:

  • width – width of the mask (in meters)

  • height – height of the mask (in meters)

  • invert – if True, the inside of the rectangle will be masked rather than the outside

op(w: Wavefront)
Parameters:

w – the Wavefront object this operator applies to

Returns:

the updated Wavefront object

class pynx.wavefront.cl_operator.Scale(x)

Multiply the wavefront by a scalar (real or complex).

Parameters:

x – the scaling factor

op(w: Wavefront)
Parameters:

w – the Wavefront object this operator applies to

Returns:

the updated Wavefront object

class pynx.wavefront.cl_operator.ThinLens(focal_length)

Multiplies the wavefront by a quadratic phase factor corresponding to a thin lens with a given focal length. The phase factor is: \(e^{-\frac{i * \pi * (x^2 + y^2)}{\lambda * focal\_length}}\)

Note that a too short focal_length can lead to aliasing, which will occur when the phase varies from more than pi from one pixel to the next. A warning will be written if this occurs at half the distance from the center (i.e. a quarter of the array size).

Parameters:

focal_length – focal length (in meters)

op(w: Wavefront)
Parameters:

w – the Wavefront object this operator applies to

Returns:

the updated Wavefront object

Fresnel propagation

Warning: this code is old and was mostly written as a proof-of-concept, but is no longer tested/developed. It is recommended to use the Wavefront operators.

pynx.wavefront.fresnel.Fresnel_thread(v1, vx1, vy1, vx2, vy2, vdz2, wavelength=1e-10, verbose=False, gpu_name='GTX', cl_platform='')

Compute the Fresnel propagation between an origin and a destination plane, using direct calculation. Uses OpenCL.

NOTE: this code is experimental, mostly used for tests !

Parameters:

  • v1 – complex 2D array of the field to propagate, size=nx1*ny1

  • vx1 – 1D vectors of x and y coordinates of v1 (nx1, ny1)

  • vy1 – 1D vectors of x and y coordinates of v1 (nx1, ny1)

  • vx2 – 1D vectors of x and y coordinates of v2 (nx1, ny1)

  • vy2 – 1D vectors of x and y coordinates of v2 (nx1, ny1)

  • vdz2 – distance (m) between the origin and destination plane

  • wavelength – wavelength

  • verbose – if True, print calculcation messages

  • gpu_name – name of the GPU to use (string)

  • cl_platform – OpenCL platform name to use (string)

Returns:

a complex array of the propagated wavefront with the same shape as vx2, vy2.

Illumination from a Fresnel Zone Place

Warning: this code is old and was mostly written as a proof-of-concept, but is no longer tested/developed. It is recommended to use the Wavefront operators.

pynx.wavefront.fzp.FZP_thread(x, y, z, sourcex=0, sourcey=0, sourcez=-50, wavelength=1, focal_length=0.129, rmax=0.0001, r_cs=0, osa_z=0, osa_r=1000000.0, nr=512, ntheta=256, fzp_xmin=None, fzp_xmax=None, fzp_ymin=None, fzp_ymax=None, fzp_nx=None, fzp_ny=None, gpu_name='GTX', cl_platform='', verbose=False)

Compute illumination from a FZP, itself illuminated from a single monochromatic point source. Uses OpenCL.

The integration is either made on the fill circular shaphe of the FZP, or a rectangular area.

All units are SI.

NOTE: this code is experimental, mostly used for tests !

Parameters:

  • x – numpy array of coordinates where the illumination will be calculated

  • y – numpy array of coordinates where the illumination will be calculated

  • z – numpy array of coordinates where the illumination will be calculated

  • sourcex – position of the point source illuminating the FZP (float)

  • sourcey – position of the point source illuminating the FZP (float)

  • sourcez – position of the point source illuminating the FZP (float)

  • wavelength – the wavelength of the incident beam

  • focal_length – the focal length of the FZP

  • rmax – max radius of the FZP

  • r_cs – radius of the central stop

  • osa_z – z position of the Order Sorting Aperture, relative to the FZP

  • osa_r – radius of the OSA

  • nr – number of radial steps for the integration

  • ntheta – number of angular (polar) steps for the integration

  • fzp_xmin – x min coordinate for a rectangular illumination

  • fzp_xmax – x max coordinate for a rectangular illumination

  • fzp_ymin – y min coordinate for a rectangular illumination

  • fzp_ymax – y max coordinate for a rectangular illumination

  • fzp_nx – number of x steps for the integration, for a rectangular illumination

  • fzp_ny – number of y steps for the integration, for a rectangular illumination

  • gpu_name – name (sub-string) of the gpu to be used

  • cl_platform – OpenCL platform to use (optional)

  • verbose – if true, will report on the progress of threaded calculations.

Returns:

a complex numpy array of the calculated illumination, with the same shape as x,y,z.

Examples

Propagation operators

# -*- coding: utf-8 -*-

# PyNX - Python tools for Nano-structures Crystallography
#   (c) 2016-present : ESRF-European Synchrotron Radiation Facility
#       authors:
#         Vincent Favre-Nicolin, favre@esrf.fr

import timeit
from pylab import *
# The following import will use CUDA if available, otherwise OpenCL
from pynx.wavefront import *


if True:
    # Near field propagation of a simple 20x200 microns slit
    w = Wavefront(d=np.zeros((512, 512), dtype=np.complex64), pixel_size=1e-6, wavelength=1.5e-10)
    a = 20e-6 / 2
    x, y = w.get_x_y()
    w.set((abs(y) < a) * (abs(x) < 100e-6))
    w = PropagateNearField(0.5) * w
    # w = PropagateFRT(3) * w
    w = ImshowRGBA(fig_num=1, title="Near field propagation (0.5m) of a 20x200 microns aperture") * w

if True:
    # Near field propagation of a simple 40x200 microns slit, displaying the propagated wavefront by steps
    # of 0.2 m propagation
    w = Wavefront(d=np.zeros((512, 512), dtype=np.complex64), pixel_size=1e-6, wavelength=1.5e-10)
    a = 40e-6 / 2
    x, y = w.get_x_y()
    w.set((abs(y) < a) * (abs(x) < 100e-6))
    # Perform 15 near field propagation of 0.2m steps, displaying the complex wavefront each time
    # the **15 expression allows to repeat the series of operators 15 times.
    w = (ImshowRGBA(fig_num=2) * PropagateNearField(0.2))**15 * w

if True:
    w = Wavefront(d=np.zeros((1024, 1024), dtype=np.complex64), pixel_size=1.3e-6, wavelength=1.e-10)
    a = 43e-6 / 2
    x, y = w.get_x_y()
    w.set((abs(y) < a) * (abs(x) < 100e-6))
    # w = PropagateNearField(3) * w
    w = PropagateFRT(3) * w
    # w = PropagateFarField(20) * w
    # w = PropagateNearField(-3) * PropagateNearField(3) * w
    # w = PropagateFRT(3, forward=False) * PropagateFRT(3) * w
    # w = PropagateFarField(100, forward=False) * PropagateFarField(100) * w
    w = ImshowRGBA(fig_num=3, title="Fractional Fourier transform propagation (3m) of a 43x200 microns slit") * w

if True:
    # Jacques et al 2012 single slit setup - here with simulated 1 micron pixel
    # Compare with figure 7 for a=43,88,142,82 microns
    figure(figsize=(15, 10))
    for a in np.array([22, 43, 88, 142, 182]) * 1e-6 / 2:
        w = Wavefront(d=np.zeros((1024, 1024), dtype=np.complex64), wavelength=1e-10, pixel_size=1.3e-6)
        x, y = w.get_x_y()
        w.set(abs(y) < (a) + x * 0)  # +x*0 needed to make sure the array is 2D
        # w = PropagateNearField(3) * w
        w = PropagateFRT(3) * w
        # w = PropagateFarField(3) * w
        icalc = fftshift(abs(w.get())).mean(axis=1) ** 2
        x, y = w.get_x_y()
        plot(fftshift(y) * 1e6, icalc, label=u'a=%5.1f µm' % (a * 2e6))
        text(0, icalc[len(icalc) // 2], u'a=%5.1f µm' % (a * 2e6))
        print('a=%5.1fum, dark spot at a^2/(2pi*lambda)=%5.2fm, I[0]=%5.2f' % (
            2 * a * 1e6, (2 * a) ** 2 / (2 * pi * w.wavelength), icalc[len(icalc) // 2]))
    title("Propagation of a slit at 3 meters, wavelength= 0.1nm")
    legend()
    xlim(-100, 100)
    xlabel(u'X (µm)')

if True:
    # propagation of a stack of A x 200 microns apertures, varying A
    w = Wavefront(d=np.zeros((16, 512, 512), dtype=np.complex64), pixel_size=1e-6, wavelength=1.5e-10)
    x, y = w.get_x_y()
    d = w.get()
    for i in range(16):
        a = 5e-6 / 2 * (i + 1)
        d[i] = ((abs(y) < a) * (abs(x) < 100e-6))
    w.set(d)
    w = PropagateFRT(1.2) * w
    figure(figsize=(15, 10))
    x, y = w.get_x_y()
    x *= 1e6
    y *= 1e6
    for i in range(16):
        subplot(4, 4, i + 1)
        imshow(abs(fftshift(w.get()[i])), extent=(x.min(), x.max(), y.min(), y.max()), origin='lower')
        title(u"A=%dµm" % (10 * (i + 1)))
        if i >= 12:
            xlabel(u'X (µm)')
        if i % 4 == 0:
            ylabel(u'Y (µm)')
        xlim(-150, 150)
        ylim(-100, 100)
    suptitle(u"Fractional Fourier propagation (0.5m) of a A x 200 µm aperture")

if True:
    # Using a lens plus near field propagation to focus a 40x80 microns^ wavefront
    w = Wavefront(d=np.zeros((512, 512), dtype=np.complex64), pixel_size=1e-6, wavelength=1.5e-10)
    x, y = w.get_x_y()
    w.set((abs(y) < 20e-6) * (abs(x) < 40e-6))
    w = PropagateNearField(1.5) * ThinLens(focal_length=2) * w
    w = ImshowAbs(fig_num=6, title="1.5m propagation after a f=2m lens") * w

if True:
    # Time propagation of stacks of 1024x1024 wavefronts
    for nz in [1, 1, 10, 50, 100, 200]:  # First size is repeated to avoid counting initializations
        t0 = timeit.default_timer()
        d = np.zeros((nz, 512, 512), dtype=np.complex64)
        w = Wavefront(d=d, pixel_size=1e-6, wavelength=1.5e-10)
        print("####################### Stack size: %4d x %4d *%4d ################" % w.get().shape)
        x, y = w.get_x_y()
        a = 20e-6 / 2
        d[:] = (abs(y) < a) * (abs(x) < 100e-6)
        w.set(d)
        t1 = timeit.default_timer()
        print("%30s: dt=%6.2fms" % ("Wavefront creation (CPU)", (t1 - t0) * 1000))
        w = PropagateFRT(1.2) * w
        t2 = timeit.default_timer()
        print("%30s: dt=%6.2fms" % ("Copy to GPU and propagation", (t2 - t1) * 1000))
        w = PropagateFRT(1.2) * w
        t3 = timeit.default_timer()
        print("%30s: dt=%6.2fms" % ("Propagation", (t3 - t2) * 1000))
        w = FreePU() * w  # We use FreePU() to make sure we release GPU memory as fast as possible
        t4 = timeit.default_timer()
        print("%30s: dt=%6.2fms" % ("Copy from GPU", (t4 - t3) * 1000))

Paganin operator

# -*- coding: utf-8 -*-

# PyNX - Python tools for Nano-structures Crystallography
#   (c) 2017-present : ESRF-European Synchrotron Radiation Facility
#       authors:
#         Vincent Favre-Nicolin, favre@esrf.fr

import numpy as np
from scipy import misc
from pylab import *
# The following import will use CUDA if available, otherwise OpenCL
from pynx.wavefront import *

################################################################################################################
# Create a wavefront as a simple transmission through a rectangular object
w = Wavefront(d=np.zeros((512, 512), dtype=np.complex64), pixel_size=1e-6, wavelength=1.5e-10)
a, b = 100e-6 / 2, 200e-6 / 2
x, y = w.get_x_y()
d = ((abs(y) < a) * (abs(x) < b))
delta = 1e-6
beta = 1e-9
thickness = 1e-6
mu = 4 * np.pi * beta / w.wavelength
k = 2 * np.pi / w.wavelength
print("       mu * t = %f\nk * delta * t = %f" % (mu * thickness, k * delta * thickness))
w.set(exp(1j * k * (-delta + 1j * beta) * thickness * d))

# Display the amplitude
w = ImshowAbs(fig_num=1) * w

# Propagate and display amplitude
w = ImshowAbs(fig_num=2) * PropagateNearField(0.5) * w

# Reconstruct original wavefield using Paganin's equation (only using propagated intensity, phase is discarded)
w = BackPropagatePaganin(dz=0.5) * w

# Display reconstructed Amplitude, phase
w = ImshowAngle(fig_num=3) * ImshowAbs(fig_num=4) * w

# Display the calculated thickness from the reconstructed phase
figure(5)
imshow(fftshift(-np.angle(w.get())/(k*delta)), extent=(x.min(), x.max(), y.min(), y.max()), origin='lower')
title("Back-propagation using Paganin's approach")
xlabel(u'X (µm)')
ylabel(u'Y (µm)')
c=colorbar()

################################################################################################################
# Slightly more complicated example
w = Wavefront(d=np.zeros((512, 512), dtype=np.complex64), pixel_size=1e-6, wavelength=1.5e-10)
delta = 1e-6
beta = 1e-9
mu = 4 * np.pi * beta / w.wavelength
k = 2 * np.pi / w.wavelength
print("       mu * t = %f\nk * delta * t = %f" % (mu * thickness, k * delta * thickness))
w.set(exp(1j * k * (-delta + 1j * beta) * fftshift(misc.ascent()) * 1e-7))
d1 = w.get(shift=True).copy()
# Propagate and reconstruct
dz = 0.5
w = PropagateNearField(dz) * w
d2 = w.get(shift=True).copy()
w = BackPropagatePaganin(dz=dz) * w
d3 = w.get(shift=True).copy()
# Display
figure(6, figsize=(10, 8))
subplot(221)
imshow(abs(d1))
title("Amplitude (z=0)")
colorbar()
subplot(222)
imshow(abs(d2))
title("Amplitude (z=%5.2fm)" % dz)
colorbar()
subplot(223)
imshow(abs(d3))
title("Reconstructed Amplitude (Paganin)")
colorbar()
subplot(224)
imshow(np.angle(d3))
title("Reconstructed Phase (Paganin)")
colorbar()

Illumination from a Fresnel Zone Place

# -*- coding: utf-8 -*-
from pynx.wavefront import fzp

from pynx.utils.plot_utils import complex2rgbalin, colorwheel
from pylab import *
import numpy as np

delta_min = 70e-9
wavelength = np.float32(12398.4 / 7000 * 1e-10)
rmax = np.float32(150e-6)  # radius of FZP
focal_length = 2 * rmax * delta_min / wavelength
print("FZP: diameter= %6.2fum, focal length=%6.2fcm" % (rmax * 2 * 1e6, focal_length * 100))
nr, ntheta = np.int32(1024), np.int32(512)  # number of points for integration on FZP
r_cs = np.float32(25e-6)  # Central stop radius
osa_z, osa_r = np.float32(focal_length - .021), np.float32(30e-6)  # OSA position and radius
sourcex = np.float32(0e-6)  # Source position
sourcey = np.float32(0e-6)
sourcez = np.float32(-90)
focal_point = 1 / (1 / focal_length - 1 / abs(sourcez))

gpu_name = "gpu" # will combine available gpu (experimental - it is safer to put part of the name of the gpu you want)

if True:
    # rectangular illumination ?
    xmin, xmax, ymin, ymax = 50e-6, 110e-6, -100e-6, 100e-6
    fzp_nx, fzp_ny = 512, 512
else:
    xmin, xmax, ymin, ymax = False, False, False, False
    fzp_nx, fzp_ny = None, None

if True:
    x = linspace(-.8e-6, .8e-6, 256)
    y = linspace(-.8e-6, .8e-6, 256)[:, newaxis]
    z = focal_point
else:
    x = linspace(-.8e-6, .8e-6, 256)[:, newaxis]
    y = 0
    z = focal_point + linspace(-2e-3, 2e-3, 256)

x = (x + (y + z) * 0).astype(float32)
y = (y + (x + z) * 0).astype(float32)
z = (z + (x + y) * 0).astype(float32)
nxyz = len(x.flat)

a_real = (x * 0).astype(np.float32)
a_imag = (x * 0).astype(np.float32)

a, dt, flop = fzp.FZP_thread(x, y, z, sourcex=sourcex, sourcey=sourcey, sourcez=sourcez, wavelength=wavelength,
                             focal_length=focal_length, rmax=rmax,
                             r_cs=r_cs, osa_z=osa_z, osa_r=osa_r, nr=nr, ntheta=ntheta,
                             fzp_xmin=xmin, fzp_xmax=xmax, fzp_ymin=ymin, fzp_ymax=ymax, fzp_nx=fzp_nx, fzp_ny=fzp_ny,
                             gpu_name=gpu_name, verbose=True)
print("clFZP dt=%9.5fs, %8.2f Gflop/s" % (dt, flop / 1e9 / dt))

if xmin is not None and xmax is not None:
  print("Correct phase for off-axis illumination")
  a *= exp(2j * pi * x * (xmin + xmax) / 2 / focal_length / wavelength)

clf()
zz = z - focal_point
if (x.max() - x.min()) <= 1e-8:
  # pylab.imshow(complex2rgba(a,amin=0.0,dlogs=2),extent=(z.min()*1e6,z.max()*1e6,y.min()*1e9,y.max()*1e9),aspect='equal',origin='lower')
  imshow(complex2rgbalin(a), extent=(zz.min() * 1e6, zz.max() * 1e6, y.min() * 1e9, y.max() * 1e9), aspect='equal', origin='lower')
  xlabel(r"$z\ (\mu m)$", fontsize=16)
  ylabel(r"$y\ (nm)$", fontsize=16)
elif (z.max() - z.min()) <= 1e-8:
  # pylab.imshow(complex2rgba(a,amin=0.0,dlogs=2),extent=(x.min()*1e6,x.max()*1e6,y.min()*1e6,y.max()*1e6),aspect='equal',origin='lower')
  imshow(complex2rgbalin(a), extent=(x.min() * 1e9, x.max() * 1e9, y.min() * 1e9, y.max() * 1e9), aspect='equal', origin='lower')
  xlabel(r"$x\ (nm)$", fontsize=16)
  ylabel(r"$y\ (nm)$", fontsize=16)
else:
  # pylab.imshow(complex2rgba(a,amin=0.0,dlogs=2),extent=(z.min()*1e3,z.max()*1e3,x.min()*1e6,x.max()*1e6),aspect='equal',origin='lower')
  imshow(complex2rgbalin(a), extent=(zz.min() * 1e6, zz.max() * 1e6, x.min() * 1e9, x.max() * 1e9), aspect='equal', origin='lower')
  xlabel(r"$z\ (\mu m)$", fontsize=16)
  ylabel(r"$x\ (nm)$", fontsize=16)

ax = axes((0.22, 0.76, 0.12, .12), facecolor='w')  # [left, bottom, width, height]
colorwheel()