Files
fourierOptics/scripts/SPIE/Thesis/Adding a WFE in the focal plane.ipynb
T

3.4 MiB

Interpolating a provided WFE to a small subsection of the field

The combination of Zernike modes and the PSD description of a WFE combine to provide a fantastic WFE over a large area. However, the resolution is not suitable to apply this at or near the focal plane of a system. Therefore, we need to find a way to interpolate the WFE in a way that is suitable.

I see two primary ways of doing:

  1. Interpolate the original wavefront to a new area, assuming that the new area is a small part contained within the original large area. This will not add any data to the WFE. But then the WFE at the smaller area will primarily be a tip/tilt effect, without any smaller structure.
  2. Interpolate the phase information of the WFE, while retaining the PSD relationship. Because the phase is randomly chosen, we could assume that the phase is actually continous at smaller scales. Then the points in the original data are consistent with the zoomed in area, while adding new structure as well. At least, that is how I think this works...

Let's first try to recreate the original way of doing a WFE.

Method #1: Interpolate on the generated WFE amplitude

We start with the built in way, i.e. the way it's currently implemented in the fourierOptics package for the MICADO ADC.

In [1]:
import numpy as np
import sys
sys.path.insert(0, '/Users/joost/Dropbox/Werk/PhD/Documents/Published/FourierOptics GitHub release/')
from fourierOptics import *
In [14]:
l = 1.49e-3
mcd = ParaxialMICADO(n=3001, m=3001, dpix=50)
wfe = WavefrontError(mcd.n, mcd.m)

_, x3, y3 = mcd.propagate(l=l,surface=4, zoom=False, show=False, 
                            print_progress=False)
wfe_p1 = ADC_surfaces(x3,y3)
# wfe_p1.generate_ADC_wfe(zernRMS=20e-6, psdRMS=15e-6)
wfe_p1.generate_ADC_wfe(zernRMS=10e-6, psdRMS=5e-6)

mask = np.sqrt(wfe_p1.xz**2 + wfe_p1.yz**2) > 1

surf_p1 = wfe_p1.latest_surface.copy()
surf_p1[mask]=np.nan
figp1, axp1 = plot_real_valued_results(surf_p1[850:2150,850:2150]*1e6, x3[850:2150,850:2150], y3[850:2150,850:2150],
        colormap='jet', labeltext='Surface height (nm)', fig_size=None, show=False)
plt.show()
Your resolution is 3001x3001 with a pixel size of 50
This means your grid is 150050.0 mm by 150050.0 mm, with each pixel 50 mm.
Generating new ADC WFE...
In [15]:
from scipy import interpolate
# Get the coordinates of the Focal Plane that we want to add a WFE to.
_, xfp, yfp = mcd.propagate(l=l,surface=1, zoom=True, show=False, 
                            print_progress=False)

surf_og = wfe_p1.latest_surface.copy()
surf_og_ip = interpolate.interp2d(x3[0,:].real, y3[:,0].real, surf_og, 
                                  kind='cubic')

fp_wfe = surf_og_ip(xfp[0,:].real, yfp[:,0].real)

figp1, axp1 = plot_real_valued_results(surf_p1[850:2150,850:2150]*1e6, x3[850:2150,850:2150], y3[850:2150,850:2150],
        colormap='jet', labeltext='Surface height (nm)', fig_size=None, show=False)
plt.show()

figp2, axp2 = plot_real_valued_results(fp_wfe*1e6, xfp, yfp,
        colormap='jet', labeltext='Surface height (nm)', fig_size=None, show=False)
plt.show()

Unfortunately, this does not contain any information on the phase and PSD distribution anymore. So we need to return to the original functions to try approach #2.

Method #2: Rebuild the PSD function and use the phase information

If we assume that we want to obtain a small window of the original WFE, which covers the full plane. This small window has a diameter d and the full plane has a diameter D.

Also, the PSD curve should in principle be space invariant, where only the range of spatial frequencies that are sampled change. Based on this, we can say that we keep our original expression of the PSD, but change the spatial frequencies with the change in size. So,


u' = \left(\frac{d}{D}\right) u \\
v' = \left(\frac{d}{D}\right) v

and to keep things consistent, we should also scale the cut-off frequency \rho_c with the same factor,


\rho_c' = \left(\frac{d}{D}\right) \rho_c
In [3]:
import numpy as np
import pyfftw
import types
from math import ceil
from scipy.special import factorial, j1

def power_spectral_density(s0, p, rho_c, dx, dy, A, m=1001, n=1001):
    # Lorentzian equation used as PSD function
    UmList = [(M - ( m + 1) / 2) / (m * dx) for M in range(m)]
    VnList = [(N - ( n + 1) / 2) / (n * dy) for N in range(n)]
    VnList, UmList = np.meshgrid(UmList, VnList)

    _rho_mn_list = np.sqrt(UmList**2 + VnList**2)
    _h0list = 1 / ( 1 + np.power( ( _rho_mn_list / rho_c ), p) )
    _h0_norm = np.sum(_h0list)
    _rho_mn_list = np.asarray(_rho_mn_list).reshape(m,n)

    _PSD = (s0**2 * A) / (_h0_norm * (1 + np.power((_rho_mn_list/rho_c), p)))
    PSD2D = _PSD

    return _PSD, _rho_mn_list


def phase_map(phase, m=1001, n=1001):
    # Uniform distribution of the phase
    _phase_in	 		= pyfftw.empty_aligned((n,m), dtype='complex128')
    _phase_in[:] 		= phase
    _phase_out	 		= pyfftw.empty_aligned((n, m), dtype='complex128')
    _FFTW_Obj 	 		= pyfftw.FFTW( _phase_in, _phase_out, axes=(0,1), direction='FFTW_FORWARD', threads=6, flags=('FFTW_ESTIMATE',))
    _FFTW_Obj()
    _phaseAngle 		= np.angle(_phase_out)
    return _phaseAngle


def generate_height_map(s0, p, rho_c, m=1001, n=1001):
    '''
    This function generates a surface height distribution based
    on a Lorentzian PSD profile. Following the paper by E. Sidick (2009).

    2020-01-20: Removed "mean=1e-4" from function def. --> Did this do anything?
    '''

    # Calculation of the necessary constants
    _dx 				= 1 / m
    _dy 				= 1 / n
    _A 					= (m * _dx) * (n * _dy)
    _phase 				= np.random.uniform(0,1,(n,m))

    # Calculating the PSD and phase
    _PSD, _rho_mn_list 	= power_spectral_density(s0, p, rho_c, _dx, _dy, _A)
    _phiMap 			= phase_map(_phase)

    # Shifting the coordinates for the FFT
    _PSD 				= np.fft.ifftshift(_PSD)

    # Doing the fourier transform
    _in	 				= pyfftw.empty_aligned((n,m), dtype='complex128')
    _in[:] 				= np.exp( 1j * _phiMap) * np.sqrt(_A * np.abs(_PSD))
    _out	 			= pyfftw.empty_aligned((n,m), dtype='complex128')
    _FFTW_Obj 	 		= pyfftw.FFTW( _in, _out, axes=(0,1), direction='FFTW_BACKWARD', threads=6, flags=('FFTW_ESTIMATE',))
    _FFTW_Obj()	

    # The final height distribution
    _hMap 				= 1 / ( _dx * _dy ) * _out
    data = _hMap.real
    rms 			= np.sqrt(np.mean(np.asarray(data)**2))
    ptv 			= np.abs(np.max(data)) + np.abs(np.min(data))
    return data
In [38]:
hsurf = generate_height_map(15e-6, 2, 1)
print(hsurf)
plt.figure()
plt.imshow(hsurf, cmap='jet')
plt.show()
[[1.93633170e-05 2.10425308e-05 2.77771267e-05 ... 1.44274571e-05
  2.27325109e-05 2.23404535e-05]
 [1.91990992e-05 3.48035221e-05 3.07693450e-05 ... 2.45813303e-05
  2.74385316e-05 3.09157219e-05]
 [3.95044135e-05 3.58946453e-05 3.75796845e-05 ... 1.44757173e-05
  2.79688859e-05 3.20674296e-05]
 ...
 [1.26774660e-05 1.01140328e-05 2.02602288e-05 ... 1.63181162e-05
  7.44747354e-06 1.35122524e-05]
 [1.02753006e-05 1.66528385e-06 7.30292121e-06 ... 7.90939382e-08
  2.18146860e-06 5.33645131e-06]
 [2.51955094e-05 1.47003151e-05 2.25325249e-05 ... 1.68272061e-05
  5.81322513e-06 2.11535448e-05]]
<IPython.core.display.Javascript object>
In [99]:
def generate_height_map_ip(s0, p, rho_c, xy_og, xy_new, m=1001, n=1001, phase=None):
    '''
    This function generates a surface height distribution based
    on a Lorentzian PSD profile. Following the paper by E. Sidick (2009).

    2020-01-20: Removed "mean=1e-4" from function def. --> Did this do anything?
    '''
    # Normalise coordinates
    xog, yog = xy_og
    xn, yn = xy_new
    xnorm = xog/xog.max()
    ynorm = yog/yog.max()
    xwin = xn/xog.max()
    ywin = yn/yog.max()
    scale = (xwin.real.max() - xwin.real.min())/(xnorm.real.max() - xnorm.real.min())
    
    # Calculation of the necessary constants
    _dx 				= 1 / m
    _dy 				= 1 / n
    _A 					= (m * _dx) * (n * _dy)
    
    if phase is None:
        _phase 				= np.random.uniform(0,1,(n,m))
    else:
        _phase = phase
    
    # Interpolation of the phase info
#     _phase_ip = interpolate.interp2d(xnorm[0,:].real, ynorm[:,0].real, _phase, kind='cubic')
    _phase_ip = interpolate.RectBivariateSpline(ynorm[:,0].real, xnorm[0,:].real, _phase, kx=5, ky=5)
#     _phase_window = _phase_ip(xwin[0,:].real, ywin[:,0].real)
    _phase_window = _phase_ip(ywin[:,0].real, xwin[0,:].real)

    # Calculating the PSD and phase
    _PSD, _rho_mn_list 	= power_spectral_density(s0, p, rho_c, _dx, _dy, _A, m=m, n=n)
    _phiMap 			= phase_map(_phase, m=m, n=n)

    # Shifting the coordinates for the FFT
    _PSD 				= np.fft.ifftshift(_PSD)
    
    # Doing the fourier transform
    _in	 				= pyfftw.empty_aligned((n,m), dtype='complex128')
    _in[:] 				= np.exp( 1j * _phiMap) * np.sqrt(_A * np.abs(_PSD))
    _out	 			= pyfftw.empty_aligned((n,m), dtype='complex128')
    _FFTW_Obj 	 		= pyfftw.FFTW( _in, _out, axes=(0,1), direction='FFTW_BACKWARD', threads=6, flags=('FFTW_ESTIMATE',))
    _FFTW_Obj()	    
    
    # The final height distribution
    _hMap 				= 1 / ( _dx * _dy ) * _out
    data = _hMap.real
    rms 			= np.sqrt(np.mean(np.asarray(data)**2))
    ptv 			= np.abs(np.max(data)) + np.abs(np.min(data))
    
    # Now the smaller window
    # Calculating the PSD and phase for the window
    s0window = rms # The rms based on the window of the original data.
    print(s0window)
    _PSDw, _rho_mn_listw 	= power_spectral_density(s0window, p, rho_c, _dx, _dy, _A, m=m, n=n)
    # Calculating the scaled phase
    _phiMap_window  	= phase_map(_phase_window, m=m, n=n)

    # Shifting the coordinates for the FFT
    _PSDw = np.fft.ifftshift(_PSDw)
    
    # Doing the fourier transform of the window
    _inw	 				= pyfftw.empty_aligned((n,m), dtype='complex128')
    _inw[:] 				= np.exp( 1j * _phiMap_window) * np.sqrt(_A * np.abs(_PSD))
    _outw	 			= pyfftw.empty_aligned((n,m), dtype='complex128')
    _FFTW_Objw 	 		= pyfftw.FFTW( _inw, _outw, axes=(0,1), direction='FFTW_BACKWARD', threads=6, flags=('FFTW_ESTIMATE',))
    _FFTW_Objw()	

    _hMapw 				= 1 / ( _dx * _dy ) * _outw
    dataw = _hMapw.real
    rmsw 			= np.sqrt(np.mean(np.asarray(dataw)**2))
    ptvw 			= np.abs(np.max(dataw)) + np.abs(np.min(dataw))
    print(rmsw)
    _hMap_ip = interpolate.interp2d(xnorm[0,:].real, ynorm[:,0].real, data, kind='cubic')
    dataip = _hMap_ip(xwin[0,:].real, ywin[:,0].real)
    
    return data, dataw, dataip
In [100]:
plt.close('all')
xy3 = (x3, y3)
xyfp = (xfp, yfp)

xfp2 = xfp + 1e-2 * xfp.max()
yfp2 = yfp
xyfp2 = (xfp2, yfp2)

x3n = x3 / x3.max()
y3n = y3 / y3.max()

# phasegen2 = np.random.uniform(0,1,x3.shape)
# phasegen = np.random.normal(0,1,x3.shape)
d, dw, dip = generate_height_map_ip(25e-6, 2, 1, xy3, xyfp, m=3001, n=3001, phase=phasegen2)
d2,dw2, dip2 = generate_height_map_ip(25e-6, 2, 1, xy3, xyfp2, m=3001, n=3001, phase=phasegen2)

fig, axes = plt.subplots(ncols=3, nrows=2, sharex=True, sharey=True, figsize=(8,8))
axes[0,0].imshow(d, cmap='jet', extent=[x3.min(), x3.max(), y3.min(), y3.max()], vmin=-5e-5, vmax=5e-5)
axes[0,1].imshow(dw, cmap='jet', extent=[xfp.min(), xfp.max(), yfp.min(), yfp.max()], vmin=-5e-5, vmax=5e-5)
axes[0,2].imshow(dip, cmap='jet', extent=[xfp.min(), xfp.max(), yfp.min(), yfp.max()], vmin=-5e-5, vmax=5e-5)
axes[1,0].imshow(d2, cmap='jet', extent=[x3.min(), x3.max(), y3.min(), y3.max()], vmin=-5e-5, vmax=5e-5)
axes[1,1].imshow(dw2, cmap='jet', extent=[xfp2.min(), xfp2.max(), yfp2.min(), yfp2.max()], vmin=-5e-5, vmax=5e-5)
axes[1,2].imshow(dip2, cmap='jet', extent=[xfp2.min(), xfp2.max(), yfp2.min(), yfp2.max()], vmin=-5e-5, vmax=5e-5)
axes[0,0].set_xlim([xfp.min(),xfp.max()])
axes[0,0].set_ylim([yfp.min(),yfp.max()])
axes[1,0].set_xlim([xfp.min(),xfp.max()])
axes[1,0].set_ylim([yfp.min(),yfp.max()])
plt.tight_layout()
plt.show()
2.4733511996914156e-05
2.4733511996914153e-05
2.4733511996914156e-05
2.473351199691415e-05
<IPython.core.display.Javascript object>

In [72]:
fig, axes = plt.subplots(ncols=2, nrows=2, sharex=True, sharey=True, figsize=(10,8))
axes[0,0].imshow(d, cmap='jet', extent=[x3.min(), x3.max(), y3.min(), y3.max()], vmin=-5e-5, vmax=5e-5)
axes[0,1].imshow(dw, cmap='jet', extent=[xfp.min(), xfp.max(), yfp.min(), yfp.max()], vmin=-5e-5, vmax=5e-5)
axes[1,0].imshow(d2, cmap='jet', extent=[x3.min(), x3.max(), y3.min(), y3.max()], vmin=-5e-5, vmax=5e-5)
axes[1,1].imshow(dw2, cmap='jet', extent=[xfp2.min(), xfp2.max(), yfp2.min(), yfp2.max()], vmin=-5e-5, vmax=5e-5)

axes[0,0].set_xlim([-1,1])
axes[0,0].set_ylim([-1,1])
axes[1,0].set_xlim([-1,1])
axes[1,0].set_ylim([-1,1])
plt.show()

print(np.sqrt(np.mean(np.asarray(d)**2)))
print(np.sqrt(np.mean(np.asarray(dw)**2)))
print(np.sqrt(np.mean(np.asarray(dw2)**2)))
<IPython.core.display.Javascript object>
1.4840107198148489e-05
1.4840107198148484e-05
1.484010719814849e-05
In [62]:
phasegen = np.random.uniform(0,1,x3.shape)
phasegen_ip = interpolate.interp2d(x3[0,:].real, y3[:,0].real, phasegen, kind='cubic', fill_value=np.nan)
phase1 = phasegen_ip(xfp[0,:].real, yfp[:,0].real)
phasegen_ip2 = interpolate.interp2d(x3[0,:].real, y3[:,0].real, phasegen, kind='cubic', fill_value=np.nan)
phase2 = phasegen_ip2(xfp[0,:].real+0.5, yfp[:,0].real)


fig, axes = plt.subplots(ncols=3, nrows=1, sharex=True, sharey=True, figsize=(9,4.8))
axes[0].imshow(phasegen, cmap='jet', extent=[x3.min(), x3.max(), y3.min(), y3.max()], vmin=0, vmax=1)
axes[1].imshow(phase1, cmap='jet', extent=[xfp.min(), xfp.max(), yfp.min(), yfp.max()], vmin=0, vmax=1)
axes[2].imshow(phase2, cmap='jet', extent=[xfp.min()+0.5, xfp.max()+0.5, yfp.min(), yfp.max()], vmin=0, vmax=1)

axes[0].set_xlim([-1,1])
axes[0].set_ylim([-1,1])
plt.show()

fig = plt.figure()
ax = fig.add_subplot(111)
ax.imshow(phase1 - phase2, cmap='jet', extent=[xfp.min(), xfp.max(), yfp.min(), yfp.max()], vmin=0, vmax=1)
plt.show()
<IPython.core.display.Javascript object>
<IPython.core.display.Javascript object>
In [85]:
_, xfp, yfp = mcd.propagate(l=l,surface=1, zoom=True, show=True, 
								print_progress=False)
<IPython.core.display.Javascript object>
In [33]:
print(x3.min(), x3.max())
(-155.63238243799717+0j) (155.63238243799717+0j)
In [29]:
%matplotlib notebook
In [ ]:
// cash stands for chaos hash :D
int cash(int x, int y){   
    int h = seed + x*374761393 + y*668265263; //all constants are prime
    h = (h^(h >> 13))*1274126177;
    return h^(h >> 16);
}
In [11]:
import numpy as np
import matplotlib.pyplot as plt

def cash(x,y):
    try:
        seed = 12345
        h = seed + x*374761393 + y*668265263
        print((h >> 13))
        h = (h**(h >> 13))*1274126177
    except ValueError:
        print((h >> 13) < 0)
    return h**(h >> 16)

y, x = np.mgrid[0:3000:3000j, 0:3000:3000j].astype(int)

qrf = cash(x,y)

plt.imshow(qrf, origin='lower')
plt.colorbar()
plt.show()
[[        1     45748     91495 ... 137104480 137150228 137195975]
 [    81576    127324    173071 ... 137186056 137231803 137277550]
 [   163152    208899    254646 ... 137267631 137313378 137359126]
 ...
 [244481323 244527071 244572818 ... 381585803 381631550 381677297]
 [244562899 244608646 244654393 ... 381667378 381713125 381758873]
 [244644474 244690221 244735969 ... 381748953 381794701 381840448]]
---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
/var/folders/yz/y0zt1db902bfbj6p8rfsrkxh0000gn/T/ipykernel_37535/882201994.py in <module>
     14 y, x = np.mgrid[0:3000:3000j, 0:3000:3000j].astype(int)
     15 
---> 16 qrf = cash(x,y)
     17 
     18 plt.imshow(qrf, origin='lower')

/var/folders/yz/y0zt1db902bfbj6p8rfsrkxh0000gn/T/ipykernel_37535/882201994.py in cash(x, y)
     10     except ValueError:
     11         print((h >> 13) < 0)
---> 12     return h**(h >> 16)
     13 
     14 y, x = np.mgrid[0:3000:3000j, 0:3000:3000j].astype(int)

ValueError: Integers to negative integer powers are not allowed.
In [105]:
xr, yr = x3.real, y3.real
print((xr,yr))
(array([[-155.63238244, -155.52862752, -155.42487259, ...,  155.42487259,
         155.52862752,  155.63238244],
       [-155.63238244, -155.52862752, -155.42487259, ...,  155.42487259,
         155.52862752,  155.63238244],
       [-155.63238244, -155.52862752, -155.42487259, ...,  155.42487259,
         155.52862752,  155.63238244],
       ...,
       [-155.63238244, -155.52862752, -155.42487259, ...,  155.42487259,
         155.52862752,  155.63238244],
       [-155.63238244, -155.52862752, -155.42487259, ...,  155.42487259,
         155.52862752,  155.63238244],
       [-155.63238244, -155.52862752, -155.42487259, ...,  155.42487259,
         155.52862752,  155.63238244]]), array([[-155.63238244, -155.63238244, -155.63238244, ..., -155.63238244,
        -155.63238244, -155.63238244],
       [-155.52862752, -155.52862752, -155.52862752, ..., -155.52862752,
        -155.52862752, -155.52862752],
       [-155.42487259, -155.42487259, -155.42487259, ..., -155.42487259,
        -155.42487259, -155.42487259],
       ...,
       [ 155.42487259,  155.42487259,  155.42487259, ...,  155.42487259,
         155.42487259,  155.42487259],
       [ 155.52862752,  155.52862752,  155.52862752, ...,  155.52862752,
         155.52862752,  155.52862752],
       [ 155.63238244,  155.63238244,  155.63238244, ...,  155.63238244,
         155.63238244,  155.63238244]]))
In [1]:
%matplotlib notebook
In [2]:
import numpy as np
import matplotlib.pyplot as plt

def fourierSeries(x,y,c):
    f = np.zeros(np.shape(x)).astype('complex128')
    mm,nn = np.meshgrid(range(len(x[0,:])), range(len(y[:,0])))
    print(mm, nn)
    for j, m in enumerate(mm):
        for i, n in enumerate(nn):
            px = 1/(n+1)
            py = 1/(m+1)
            fij = c * np.exp(1j*(2*np.pi/px)*x + (2*np.pi/py)*y)
#             fij = c[i,j] * np.exp(1j*n*x+m*y)
            f += fij
    return f

def fourierSeries1D(x,c, P=50, xs=25):
    '''
    x - x coordinates
    c - Coefficients
    P - Period before repeating the pattern
    xs - shift of the full field x-coordinates. To remove symmetry around 0.
    '''
    xpos = x - xs
    norm_factor = np.sum(c)
    print(norm_factor)
    f = np.asarray([c[i]*np.exp(1j*i* 2*np.pi/P * xpos) for i in range(len(c))]).sum(axis=0)
    print(type(f[0]))
    return f/norm_factor

# 1D Fourier series test
np.random.seed(1)
coeffs = np.random.uniform(-1,1,100)
# coeffs[0] = 1
# coeffs[1] = 1
# coeffs[2] = 1

xx0 = np.linspace(-10,40,11)
yy0 = fourierSeries1D(xx0, coeffs, P=1000)
print(yy0)
xx1 = np.linspace(-50, 50, 201)
yy1 = fourierSeries1D(xx1,coeffs, P=1000)
xx2 = np.linspace(-30, 45, 21)
yy2 = fourierSeries1D(xx2,coeffs, P=1000)
plt.figure()
plt.plot(xx0,yy0.real, 'g', marker='x', linestyle='--', zorder=2)
plt.plot(xx1,yy1.real, 'b-', zorder=2)
plt.plot(xx2,yy2.real, 'r-', zorder=3, linewidth=1)
plt.grid(alpha=0.3, zorder=0)
plt.show()

'''
# coeffs = np.random.uniform(0,1,(101,101))
# coeffs = np.zeros((101,101))
# coeffs[0,0] = 1
x = np.linspace(-np.pi,np.pi,101)
y = np.linspace(-np.pi,np.pi,101)
xx, yy = np.meshgrid(x, y)

data = fourierSeries(xx,yy,coeffs)
print(data)

plt.figure()
plt.imshow(np.abs(data), cmap='jet')
plt.colorbar()
plt.show()
'''
Out [2]:
-2.8244144799708755
<class 'numpy.complex128'>
[ 0.92804617-0.14694j     0.25085172+0.14699884j -1.14854409-2.04597838j
  3.27431443-2.71612062j  1.90618351+1.83303848j -0.20603149-0.96766014j
  2.12881562-0.41981614j  1.        -0.j          2.12881562+0.41981614j
 -0.20603149+0.96766014j  1.90618351-1.83303848j]
-2.8244144799708755
<class 'numpy.complex128'>
-2.8244144799708755
<class 'numpy.complex128'>
<IPython.core.display.Javascript object>
"\n# coeffs = np.random.uniform(0,1,(101,101))\n# coeffs = np.zeros((101,101))\n# coeffs[0,0] = 1\nx = np.linspace(-np.pi,np.pi,101)\ny = np.linspace(-np.pi,np.pi,101)\nxx, yy = np.meshgrid(x, y)\n\ndata = fourierSeries(xx,yy,coeffs)\nprint(data)\n\nplt.figure()\nplt.imshow(np.abs(data), cmap='jet')\nplt.colorbar()\nplt.show()\n"
In [3]:
# 1D Fourier series test
np.random.seed(1)
coeffs = np.random.uniform(-1,1,100)
# coeffs[0] = 1
# coeffs[1] = 1
# coeffs[2] = 1

xx0 = np.linspace(-10,10,11)
yy0 = fourierSeries1D(xx0, coeffs, P=1000)
print(yy0)
xx1 = np.linspace(-5, 15, 11)
yy1 = fourierSeries1D(xx1,coeffs, P=1000)
# xx2 = np.linspace(-30, 45, 21)
# yy2 = fourierSeries1D(xx2,coeffs, P=1000)
plt.figure()
plt.plot(xx0,yy0.real, 'g', marker='x', linestyle='--', zorder=2)
plt.plot(xx1,yy1.real, 'b-', zorder=2)

# plt.plot(xx2,yy2.real, 'r-', zorder=3, linewidth=1)
plt.grid(alpha=0.3, zorder=0)
plt.show()
-2.8244144799708755
<class 'numpy.complex128'>
[ 0.92804617-0.14694j     0.73482068-0.1828004j   0.52660171+0.0308276j
 -0.15380869+0.14587545j -1.03250078-0.52772232j -1.14854409-2.04597838j
  0.15220504-3.4237516j   2.32344441-3.40161665j  3.87639295-1.69179571j
  3.65598886+0.5920894j   1.90618351+1.83303848j]
-2.8244144799708755
<class 'numpy.complex128'>
<IPython.core.display.Javascript object>
In [6]:
import time

def fourierSeries2D(x, y, c, Px=500, Py=500, xs=250, ys=250):
    '''
    x - x coordinates
    c - Coefficients
    P - Period before repeating the pattern
    xs - shift of the full field x-coordinates. To remove symmetry around 0.
    '''
    xpos = x #- np.sqrt(2)*xs
    ypos = y #- np.sqrt(2)*ys
    fout = []#np.zeros(c.shape).astype('complex128')
    for i in range(c.shape[0]):
        for j in range(c.shape[1]):
            f = c[i,j] * np.exp(1j * i * 2*np.pi*xpos/Px) * np.exp(1j * j * 2*np.pi*ypos/Py)
            fout.append(f)
    fout = np.asarray(fout).sum(axis=0)
    return fout


# 1D Fourier series test
np.random.seed(2)
coeffs = np.random.uniform(-1,1,(11,11))
x0 = np.linspace(-10,10,11)
y0 = np.linspace(-10,10,11)
xx, yy = np.meshgrid(x0,y0)
start = time.time()
zz = fourierSeries2D(xx, yy, coeffs)
end = time.time()
print('Elapsed time: {:.4f} s'.format(end - start))

plt.figure()
plt.imshow(zz.real, extent=[-150,150,-150,150], cmap='jet', origin='lower')
plt.grid()
plt.colorbar()
plt.show()
Elapsed time: 0.0068 s
<IPython.core.display.Javascript object>
In [7]:
from numba import njit, jit, complex64, float32,float64,complex128
import numba
import concurrent 

# def fourierSeries2Dopt(x, y, c, Px=500, Py=500, xs=0, ys=0):
#     '''
#     x - x coordinates
#     c - Coefficients
#     P - Period before repeating the pattern
#     xs - shift of the full field x-coordinates. To remove symmetry around 0.
#     '''
#     xpos = x - xs# * np.sqrt(2)
#     ypos = y - ys# * np.sqrt(2)
#     j,i = np.mgrid[0:(c.shape[0]-1):(c.shape[0])*1j,0:(c.shape[1]-1):c.shape[1]*1j]
#     print(i)
#     print(xpos)
#     ixt = np.asarray([i*2*np.pi*xpos/Px for i in range(c.shape[0])])
#     jyt = np.asarray([j*2*np.pi*ypos/Py for j in range(c.shape[1])])
#     f = c * np.exp(1j * (ixt + jyt)) * np.exp(1j * j * 2*np.pi*ypos/Py).sum(axis=0)
# #     f = c * np.exp(1j * (ixt + jyt)).sum(axis=0)
#     return f
def fourierSeries2Dopt2(x, y, c, Px=500, Py=500, xs=250, ys=250):
    '''
    x - x coordinates
    c - Coefficients
    P - Period before repeating the pattern
    xs - shift of the full field x-coordinates. To remove symmetry around 0.
    '''
    xpos = x + np.sqrt(2)*xs
    ypos = y + np.sqrt(2)*ys
    cs0, cs1 = c.shape
    xt = 2*np.pi*xpos/Px
    yt = 2*np.pi*ypos/Py
    fout = [c[i,j] * np.exp(1j * (i * xt + j * yt)) for i  in range(cs0) for j in range(cs1)]
    fout = np.array(fout).sum(axis=0)
    return fout


def meshgrid(x, y):
    xx = np.empty(shape=(len(x), len(y)), dtype=np.int64)
    yy = np.empty(shape=(len(x), len(y)), dtype=np.int64)
    for j in y:
        for k in x:
            xx[j,k] = k  # change to x[k] if indexing xy
            yy[j,k] = j  # change to y[j] if indexing xy
    return xx, yy


def fourierSeries2Dopt(x, y, c, Px=500, Py=500, xs=250, ys=250):
    '''
    x - x coordinates
    c - Coefficients
    P - Period before repeating the pattern
    xs - shift of the full field x-coordinates. To remove symmetry around 0.
    '''
    xpos = x + np.sqrt(2)*xs
    ypos = y + np.sqrt(2)*ys
    cs0, cs1 = c.shape
    normf = np.sum(c)
    xt = 2*np.pi*xpos/Px
    yt = 2*np.pi*ypos/Py
    c_ij = c.flatten()
    i,j = meshgrid(range(cs0), range(cs1))
    I = i.flatten()
    J = j.flatten()
    fout = f(c, i, j, xt, yt)
    return fout/normf

@njit
def f(c, i, j, xt, yt):
    cs0, cs1 = c.shape
    fout = np.zeros(xt.shape, dtype=np.float64)
    for j in range(cs1):
        for i in range(cs0):
            v = c[i,j] * np.exp(1j * (i * xt + j * yt)) 
            fout += v.real
    return fout

# def expo_helper(var):
#     i, j, xt, yt, c = var
#     return c[i,j] * np.exp(1j * (i * xt + j * yt)) 
    

# def fourierSeries2Dopt3(x, y, c, Px=500, Py=500, xs=250, ys=250):
#     '''
#     x - x coordinates
#     c - Coefficients
#     P - Period before repeating the pattern
#     xs - shift of the full field x-coordinates. To remove symmetry around 0.
#     '''
#     xpos = x + np.sqrt(2)*xs
#     ypos = y + np.sqrt(2)*ys
#     cs0, cs1 = c.shape
#     xt = 2*np.pi*xpos/Px
#     yt = 2*np.pi*ypos/Py
#     c_ij = c.flatten()
#     i,j = meshgrid(range(cs0), range(cs1))
#     I = i.flatten()
#     J = j.flatten()
#     var_list = [(c_ij[k], I[k], J[k], xt, yt) for k in range(len(c_ij))]
#     fout = np.zeros(np.shape(var_list))
#     executor = concurrent.futures.ProcessPoolExecutor(10)
#     futures = [executor.submit(expo_helper, var) for var in var_list]
#     concurrent.futures.wait(futures)

#     print(futures[0].result() )
    
#     print(np.array(futures).shape)
#     fout = np.array(futures).sum(axis=0)
#     return fout


# 2D Fourier series test
np.random.seed(2183)
coeffs = np.random.uniform(-1,1,(11,11))
# coeffs = np.ones((101,101))
x0 = np.linspace(-50,50,101)
y0 = np.linspace(-50,50,101)
xx, yy = np.meshgrid(x0,y0)
s0 = time.time()
# zz = fourierSeries2Dopt2(xx, yy, coeffs, )#Px=100, Py=100, xs=0, ys=0)
s1 = time.time()
zz = fourierSeries2Dopt(xx, yy, coeffs, )#Px=100, Py=100, xs=0, ys=0)
s2 = time.time()
print('Old one: {:.4f} s'.format(s1-s0))
print('New one: {:.4f} s'.format(s2-s1))


# Second set of coordinates
x1 = np.linspace( -10,10,101)
y1 = np.linspace(-10,40,101)
xx1, yy1 = np.meshgrid(x1,y1)
zz1 = fourierSeries2Dopt(xx1, yy1, coeffs, )#Px=100, Py=100, xs=0, ys=0)

ext = [x0.min(), x0.max(), y0.min(), y0.max()]
fig, axes = plt.subplots(ncols=3, nrows=1, figsize=(10,4), sharex=True)
im = axes[0].imshow(zz.real, extent=ext, cmap='jet', origin='lower')
# fig.colorbar(im, ax=axes[0])

axes[1].plot( xx[int(xx.shape[0]/2),:],  zz.real[int(xx.shape[0]/2),:],   'b-', zorder=2)
axes[1].plot(xx1[int(xx1.shape[0]/2),:], zz1.real[int(xx1.shape[0]/2),:], 'g-', marker='x', zorder=1)

axes[2].plot( yy[:,int(xx.shape[0]/2)],   zz.real[:,int(xx.shape[0]/2)],  'r-', zorder=2)
axes[2].plot(yy1[:,int(zz1.shape[0]/2)], zz1.real[:,int(xx1.shape[0]/2)], 'g-', marker='x', zorder=1)

axes[1].grid(alpha=0.3, zorder=1)
axes[2].grid(alpha=0.3, zorder=1)
plt.tight_layout()
plt.show()
Old one: 0.0000 s
New one: 0.4503 s
<IPython.core.display.Javascript object>
In [8]:
ext = [x0.min(), x0.max(), y0.min(), y0.max()]
fig, axes = plt.subplots(ncols=3, nrows=1, figsize=(10,4), sharex=True)
im = axes[0].imshow(zz.real, extent=ext, cmap='jet', origin='lower')
# fig.colorbar(im, ax=axes[0])

axes[1].plot( xx[int(xx.shape[0]/2),:],  zz.real[int(xx.shape[0]/2),:],   'b-', zorder=2)
axes[1].plot(xx1[int(xx1.shape[0]/2),:], zz1.real[int(xx1.shape[0]/2),:], 'g-', marker='x', zorder=1)

axes[2].plot( yy[:,int(xx.shape[0]/2)],   zz.real[:,int(xx.shape[0]/2)],  'r-', zorder=2)
axes[2].plot(yy1[:,int(zz1.shape[0]/2)], zz1.real[:,int(xx1.shape[0]/2)], 'g-', marker='x', zorder=1)

axes[1].grid(alpha=0.3, zorder=1)
axes[2].grid(alpha=0.3, zorder=1)
plt.tight_layout()
plt.show()
<IPython.core.display.Javascript object>
In [22]:
from matplotlib import cm
fig, axes = plt.subplots(ncols=3, nrows=1, figsize=(10,4), subplot_kw={"projection": "3d"}, 
                         sharex=True, sharey=True)
# fig.colorbar(im, ax=axes[0])
# surf = ax.plot_surface(X, Y, Z, cmap=cm.coolwarm, linewidth=0, antialiased=False)

surf1 = axes[0].plot_surface(xx, yy, zz.real, cmap=cm.coolwarm, linewidth=0, antialiased=False)
surf2 = axes[1].plot_surface(xx1, yy1, zz1.real, cmap=cm.coolwarm, linewidth=0, antialiased=False)

axes[2].plot_wireframe(xx, yy, zz.real, cmap=cm.coolwarm, linewidth=.8, antialiased=False)
axes[2].plot_surface(xx1, yy1, zz1.real, cmap=cm.coolwarm, linewidth=1, antialiased=False)


plt.show()
<IPython.core.display.Javascript object>
In [ ]: