TST: Add units tests for direct comparison with matlab

azure-pipelines.yml

@@ -46,6 +46,7 @@ jobs:
       python=$(python.version)
       cudatoolkit=$(CUDA_VERSION)
       "openmpi=*=h*"
+      h5py
     displayName: Create build environment
 
   - script: conda list -n tike
@@ -101,6 +102,7 @@ jobs:
       pytest
       python=$(python.version)
       cudatoolkit=$(CUDA_VERSION)
+      h5py
     displayName: Create build environment
 
   - script: conda list -n tike

src/tike/ptycho/probe.py

@@ -475,7 +475,7 @@ def add_modes_cartesian_hermite(probe, nmodes: int):
     X, Y = np.meshgrid(
         np.arange(probe.shape[-2]) - (probe.shape[-2] // 2 - 1),
         np.arange(probe.shape[-1]) - (probe.shape[-2] // 2 - 1),
-        indexing='xy',
+        indexing='ij',
     )
 
     cenx = np.sum(

tests/matlab/fft_test.m

@@ -0,0 +1,47 @@
+pd = makedist('Uniform', -1, 1);
+
+x0 = gpuArray(cast(1 + random(pd, [64, 64, 10]) + 1i * random(pd, [64, 64, 10]), 'double'));
+x = x0;
+error = zeros(100);
+
+for i = 1:100
+    x = fft2(x);
+    x = ifft2(x);
+    if i == 1
+        y_double = gather(x);
+    end
+    error(i) = mean(reshape(abs((x-x0)./x0), 1, []));
+end
+
+figure();
+plot(1:100, error);
+title({['Mean Absolute Relative Error'], ['for Double Precision on GPU']});
+xlabel('Calls to FFT/iFFT');
+
+x0_double = gather(x0);
+
+x0 = cast(x0, 'single');
+x = x0;
+error = zeros(100);
+
+for i = 1:100
+    if i == 1
+        x0_single = gather(x);
+    end
+    x = fft2(x);
+    if i == 1
+        y_single = gather(x);
+    end
+    x = ifft2(x);
+    if i == 1
+        x_single = gather(x);
+    end
+    error(i) = mean(reshape(abs((x-x0)./x0), 1, []));
+end
+
+figure();
+plot(1:100, error);
+title({['Mean Absolute Relative Error'], ['for Single Precision on GPU']});
+xlabel('Calls to FFT/iFFT');
+
+save('fft.mat', 'x0_double', 'y_double', 'x_single', 'x0_single', 'y_single', '-v7.3');

tests/matlab/fft_test.py

@@ -0,0 +1,106 @@
+import os
+
+import cupy as cp
+import h5py
+import matplotlib.pyplot as plt
+import numpy as np
+import tike.random
+import tike.view
+
+_dir = os.path.dirname(__file__)
+_test = os.path.join(os.path.dirname(_dir), 'result', 'matlab')
+
+os.makedirs(_test, exist_ok=True)
+
+def test_consistent_fft():
+
+    with h5py.File(os.path.join(_dir, 'fft.mat'), 'r') as inputs:
+        x = inputs['x0_double'][...].view('complex128')
+        print(x.shape, x.dtype)
+
+        y = inputs['y_double'][...].view('complex128')
+        print(y.shape, y.dtype)
+
+    y1 = cp.fft.ifft2(cp.fft.fft2(cp.asarray(x, order='C'))).get()
+
+    plt.figure()
+    tike.view.plot_complex(np.fft.ifftshift(y1[0]))
+    plt.savefig(os.path.join(_test, 'fft-00.png'))
+
+    plt.figure()
+    tike.view.plot_complex(np.fft.ifftshift(y1[0] - y[0]))
+    plt.suptitle('FFT Error between MATLAB and CUPY for Double Precision.')
+    plt.savefig(os.path.join(_test, 'fft-01.svg'))
+
+    np.testing.assert_array_equal(y1, y)
+
+    with h5py.File(os.path.join(_dir, 'fft.mat'), 'r') as inputs:
+        x0 = inputs['x0_single'][...].view('complex64')
+        print(x0.shape, x0.dtype)
+
+        x = inputs['x_single'][...].view('complex64')
+        print(x.shape, x.dtype)
+
+        y = inputs['y_single'][...].view('complex64')
+        print(y.shape, y.dtype)
+
+    y1 = cp.fft.fft2(cp.asarray(x0, order='C'), norm='backward')
+    x1 = cp.fft.ifft2(y1, norm='backward')
+
+    x1 = x1.get()
+    y1 = y1.get()
+
+    plt.figure()
+    tike.view.plot_complex(np.fft.ifftshift(y1[0]))
+    plt.savefig(os.path.join(_test, 'fft-02.png'))
+
+    plt.figure()
+    plt.suptitle('FFT Error between MATLAB and CUPY for Single Precision.')
+    tike.view.plot_complex(np.fft.ifftshift(y1[0] - y[0]))
+    plt.savefig(os.path.join(_test, 'fft-03.svg'))
+
+    np.testing.assert_array_equal(y1, y)
+    np.testing.assert_array_equal(x1, x)
+
+def test_repeated_fft():
+
+    with h5py.File(os.path.join(_dir, 'fft.mat'), 'r') as inputs:
+        x = inputs['x0_double'][...].view('complex128')
+    x = x.astype('complex128')
+
+    x = cp.asarray(x, order='C')
+    x0 = x.copy()
+    error = list()
+
+    for _ in range(100):
+        x = cp.fft.fft2( x, norm='ortho')
+        x = cp.fft.ifft2(x, norm='ortho')
+        error.append(cp.mean(cp.abs((x-x0)/x0)).get())
+
+    plt.figure()
+    plt.plot(error)
+    plt.title('Mean Absolute Relative Error for Double Precision on GPU')
+    plt.xlabel('Calls to FFT/iFFT')
+    plt.tight_layout()
+    plt.savefig(os.path.join(_test, 'fft-04.svg'))
+
+    x = cp.asarray(x, order='C').astype('complex64')
+    x0 = x.copy()
+    error = list()
+
+    for _ in range(100):
+        x = cp.fft.fft2( x, norm='ortho')
+        x = cp.fft.ifft2(x, norm='ortho')
+        error.append(cp.mean(cp.abs((x-x0)/x0)).get())
+
+    plt.figure()
+    plt.plot(error)
+    plt.title('Mean Absolute Relative Error for Single Precision on GPU')
+    plt.xlabel('Calls to FFT/iFFT')
+    plt.tight_layout()
+    plt.savefig(os.path.join(_test, 'fft-05.svg'))
+
+    plt.figure()
+    tike.view.plot_complex(x[0].get() - x0[0].get())
+    plt.savefig(os.path.join(_test, 'fft-06.png'))
+    plt.close('all')

tests/matlab/forward_test.m

@@ -0,0 +1,127 @@
+load('forward-in.mat');
+pd = makedist('Normal', 0, 1);
+peps = cast(padarray(imread('peppers.png'), [64, 0, 0]), 'single');
+corn = cast(imread('ngc6543a.jpg'), 'single');
+coins = cast(imread('coins_2048.png'), 'single');
+satyre = cast(imread('satyre_2048.png'), 'single');
+
+Np_p = self.Np_p;
+self.probe{1} = 0 .* self.probe{1} + peps(:, :, 1) + 1i .* corn(1:512, 1:512, 2);
+probe0 = gather(self.probe{1});
+self.probe{2} = 0 .* self.probe{2} + peps(:, :, 2) + 1i .* corn(1:512, 1:512, 3);
+probe1 = gather(self.probe{2});
+probe_evolution = self.probe_evolution(g_ind, :);
+z_distance = self.z_distance;
+modes = self.modes;
+self.object{1} = 0 .* self.object{1} + satyre(1:1568, 1:1568, 2) + 1i .* coins(1:1568, 1:1568);
+object = gather(self.object{1});
+object_modes = par.object_modes;
+probe_modes = par.probe_modes;
+variable_intensity = par.variable_intensity;
+Nlayers = par.Nlayers;
+apply_subpix_shift = par.apply_subpix_shift;
+apodwin = cache.apodwin;
+cache.oROI_s{1}{1}(g_ind(1), :) = 0;
+cache.oROI_s{1}{1}(g_ind(2), :) = size(object, 1) - 512;
+cache.oROI_s{1}{1}(g_ind(3), :) = 0;
+cache.oROI_s{1}{1}(g_ind(4), :) = size(object, 1) - 512;
+cache.oROI_s{1}{1}(g_ind(5), :) = 512;
+cache.oROI_s{1}{2}(g_ind(1), :) = 0;
+cache.oROI_s{1}{2}(g_ind(2), :) = 0;
+cache.oROI_s{1}{2}(g_ind(3), :) = size(object, 2) - 512;
+cache.oROI_s{1}{2}(g_ind(4), :) = size(object, 2) - 512;
+cache.oROI_s{1}{2}(g_ind(5), :) = 512;
+positions0 = cache.oROI_s{1}{1}(g_ind, 1);
+positions1 = cache.oROI_s{1}{2}(g_ind, 1);
+
+save('forward-in1.mat', 'Np_p',...
+'probe0',...
+'probe1',...    
+'probe_evolution',...
+'z_distance',...
+'modes',...
+'object',...
+'object_modes',...
+'probe_modes',...
+'variable_intensity',...
+'Nlayers',...
+'apply_subpix_shift',...
+'positions0', 'positions1', ...
+'apodwin', '-v7.3');
+
+import engines.GPU.shared.*
+import engines.GPU.GPU_wrapper.*
+import engines.GPU.LSQML.*
+import math.*
+import utils.*
+import plotting.*
+
+if isempty(obj_proj{1})
+    for ll = 1:par.object_modes
+        obj_proj{ll} = Gzeros([self.Np_p, 0], true);
+    end
+end
+probe = self.probe; 
+
+% get illumination probe 
+for ll = 1:par.probe_modes
+    if (ll == 1 && (par.variable_probe || par.variable_intensity))
+        % add variable probe (OPRP) part into the constant illumination 
+        probe{ll,1} =  get_variable_probe(self.probe{ll}, self.probe_evolution(g_ind,:),p_ind{ll});
+    else
+        % store the normal (constant) probe(s)
+        probe{ll,1} = self.probe{min(ll,end)}(:,:,min(end,p_ind{ll}),1);
+    end
+
+%     if (ll == 1 && par.apply_subpix_shift && isinf(self.z_distance(end)))  || is_used(par,'fly_scan')
+%         % only in farfield mode 
+%         probe{ll} = apply_subpx_shift(probe{ll}, self.modes{min(end,ll)}.sub_px_shift(g_ind,:) );
+%     end
+%     if (ll == 1)
+%         probe{ll} = apply_subpx_shift_fft(probe{ll}, self.modes{1}.probe_fourier_shift(g_ind,:)); 
+%     end
+end
+
+% get projection of the object and probe 
+for layer = 1:par.Nlayers
+   for ll = 1:max(par.object_modes, par.probe_modes)
+       llo = min(ll, par.object_modes); 
+       llp = min(ll, par.probe_modes); 
+        % get objects projections 
+        obj_proj{llo} = get_views(self.object, obj_proj{llo},layer_ids(layer),llo, g_ind, cache, scan_ids,[]);
+        if (ll == 1 && par.apply_subpix_shift && ~isinf(self.z_distance(end)))
+            % only in nearfield mode , apply shift in the opposite direction 
+            obj_proj{ll} = apply_subpx_shift(obj_proj{ll} .* cache.apodwin, -self.modes{min(end,ll)}.sub_px_shift(g_ind,:) ) ./ cache.apodwin;
+        end
+
+        % get exitwave after each layer
+        psi{ll} =probe{llp,layer} .* obj_proj{llo};
+        % fourier propagation  
+%         [psi{ll}] = fwd_fourier_proj(psi{ll} , self.modes{layer}, g_ind);  
+        if par.Nlayers > 1
+             probe{llp,layer+1} = psi{llp};
+        end
+   end
+end
+
+% figure();
+% imshow(real(object), [0, 255]);
+% savefig('forward-00.png');
+
+% figure();
+% for i = 1:5
+%     subplot(5, 1, i);
+%     imshow(real(obj_proj{1}(:, :, i)), [0, 255]);
+% end
+% savefig('forward-01.png');
+
+psi1 = gather(psi{1});
+psi2 = gather(psi{2});
+obj_proj_ = gather(obj_proj{1});
+probe0 = gather(probe{1});
+probe1 = gather(probe{2});
+
+save('forward-out.mat', 'probe0', 'probe1', 'obj_proj_', '-v7.3');
+save('forward-out2.mat', 'psi2', '-v7.3');
+save('forward-out1.mat', 'psi1', '-v7.3');
+