feectools 0.1.0__py3-none-any.whl

feectools/linalg/tests/test_fft.py

@@ -0,0 +1,106 @@
+import pytest
+import scipy.fft as scifft
+import numpy as np
+
+from feectools.ddm.mpi import mpi as MPI
+from feectools.linalg.fft import *
+from feectools.ddm.cart               import DomainDecomposition, CartDecomposition
+from feectools.linalg.stencil import StencilVector
+#===============================================================================
+def compute_global_starts_ends(domain_decomposition, npts):
+    ndims         = len(npts)
+    global_starts = [None]*ndims
+    global_ends   = [None]*ndims
+
+    for axis in range(ndims):
+        es = domain_decomposition.global_element_starts[axis]
+        ee = domain_decomposition.global_element_ends  [axis]
+
+        global_ends  [axis]     = ee.copy()
+        global_ends  [axis][-1] = npts[axis]-1
+        global_starts[axis]     = np.array([0] + (global_ends[axis][:-1]+1).tolist())
+
+    return global_starts, global_ends
+#===============================================================================
+def decode_fft_type(ffttype):
+    if ffttype == 'fft':
+        return (complex, DistributedFFT, scifft.fftn)
+    elif ffttype == 'ifft':
+        return (complex, DistributedIFFT, scifft.ifftn)
+    elif ffttype == 'dct':
+        return (complex, DistributedDCT, scifft.dctn)
+    elif ffttype == 'idct':
+        return (complex, DistributedIDCT, scifft.idctn)
+    elif ffttype == 'dst':
+        return (complex, DistributedDST, scifft.dstn)
+    elif ffttype == 'idst':
+        return (complex, DistributedIDST, scifft.idstn)
+    else:
+        raise NotImplementedError()
+
+def method_test(seed, comm, config, dtype, classtype, comparison, verbose=False):
+    np.random.seed(seed)
+
+    if comm is None:
+        rank = -1
+    else:
+        rank = comm.Get_rank()
+    
+    npts, pads, periods = config
+
+    if verbose:
+        print(f'[{rank}] Test start', flush=True)
+
+    comm = MPI.COMM_WORLD
+    D = DomainDecomposition(npts, periods=periods, comm=comm)
+
+    # Partition the points
+    global_starts, global_ends = compute_global_starts_ends(D, npts)
+
+    cart = CartDecomposition(D, npts, global_starts, global_ends, pads=pads, shifts=[1]*len(pads))
+
+    # vector spaces
+    V = StencilVectorSpace(cart, dtype=dtype)
+    localslice = tuple([slice(s, e+1) for s, e in zip(V.starts, V.ends)])
+
+    if verbose:
+        print(f'[{rank}] Vector spaces built', flush=True)
+
+    if np.dtype(dtype).kind == 'c':
+        Y_glob = np.random.random(V.npts) + np.random.random(V.npts) * 1j
+    else:
+        Y_glob = np.random.random(V.npts)
+
+    # vector to solve for (Y)
+    Y = StencilVector(V)
+    Y[localslice] = Y_glob[localslice]
+    Y.update_ghost_regions()
+
+    if verbose:
+        print(f'[{rank}] Vector built', flush=True)
+
+    X_glob = comparison(Y_glob)
+
+    compare = classtype(V)
+    X = compare.dot(Y)
+
+    if verbose:
+        print(f'[{rank}] Functions have been run', flush=True)
+
+    assert np.allclose(X_glob[localslice], X[localslice], 1e-10, 1e-10)
+
+@pytest.mark.parametrize( 'seed', [0, 2] )
+@pytest.mark.parametrize( 'params', [([8], [2], [False]), ([8,9], [2,3], [False,True]), ([8,9,17], [2,3,7], [False,True,False])] )
+@pytest.mark.parametrize( 'ffttype', ['fft', 'ifft', 'dct', 'idct', 'dst', 'idst'] )
+def test_kron_fft_ser(seed, params, ffttype):
+    method_test(seed, None, params, *decode_fft_type(ffttype), verbose=False)
+
+@pytest.mark.parametrize( 'seed', [0, 2] )
+@pytest.mark.parametrize( 'params', [([16], [2], [False]), ([16,18], [2,3], [False,True]), ([16,18,37], [2,3,7], [False,True,False])] )
+@pytest.mark.parametrize( 'ffttype', ['fft', 'ifft', 'dct', 'idct', 'dst', 'idst'] )
+@pytest.mark.parallel
+def test_kron_fft_par(seed, params, ffttype):
+    method_test(seed, MPI.COMM_WORLD, params, *decode_fft_type(ffttype), verbose=False)
+
+if __name__ == '__main__':
+    method_test(0, MPI.COMM_WORLD, ([8,9,5], [2,3,7], [False,True,False]), *decode_fft_type('fft'), verbose=True)

feectools/linalg/tests/test_kron_stencil_matrix.py

@@ -0,0 +1,114 @@
+from functools import reduce
+
+import pytest
+import numpy as np
+from scipy.sparse import kron
+
+from feectools.ddm.cart       import DomainDecomposition, CartDecomposition
+from feectools.linalg.stencil import StencilVectorSpace
+from feectools.linalg.stencil import StencilVector
+from feectools.linalg.stencil import StencilMatrix
+from feectools.linalg.kron    import KroneckerStencilMatrix
+#===============================================================================
+def compute_global_starts_ends(domain_decomposition, npts):
+    ndims         = len(npts)
+    global_starts = [None]*ndims
+    global_ends   = [None]*ndims
+
+    for axis in range(ndims):
+        ee = domain_decomposition.global_element_ends  [axis]
+
+        global_ends  [axis]     = ee.copy()
+        global_ends  [axis][-1] = npts[axis]-1
+        global_starts[axis]     = np.array([0] + (global_ends[axis][:-1]+1).tolist())
+
+    return tuple(global_starts), tuple(global_ends)
+
+#==============================================================================
+@pytest.mark.parametrize('dtype', [float])
+@pytest.mark.parametrize('npts', [(5, 7, 8)])
+@pytest.mark.parametrize('pads', [(2, 3, 5)])
+@pytest.mark.parametrize('periodic', [(True, False, False)])
+
+def test_KroneckerStencilMatrix(dtype, npts, pads, periodic):
+
+    # Extract input parameters
+    n1, n2, n3 = npts
+    p1, p2, p3 = pads
+    P1, P2, P3 = periodic
+
+    # Define data type with a factor
+    if dtype==complex:
+        factor=1j
+    else:
+        factor=1
+
+    # Create domain decomposition
+    D = DomainDecomposition([n1-1,n2-1, n3-1], periods=[P1,P2,P3])
+
+    # Partition the points
+    global_starts, global_ends = compute_global_starts_ends(D, npts)
+
+    cart = CartDecomposition(D, npts, global_starts, global_ends, pads=[p1,p2,p3], shifts=[1,1,1])
+
+    # 3D vector space and element
+    W = StencilVectorSpace( cart, dtype=dtype)
+    w = StencilVector(W)
+
+    # 1D vector space
+
+    D1 = DomainDecomposition([n1-1], periods=[P1])
+    D2 = DomainDecomposition([n2-1], periods=[P2])
+    D3 = DomainDecomposition([n3-1], periods=[P3])
+
+    # Partition the points
+    global_starts1, global_ends1 = compute_global_starts_ends(D1, [n1])
+    global_starts2, global_ends2 = compute_global_starts_ends(D2, [n2])
+    global_starts3, global_ends3 = compute_global_starts_ends(D3, [n3])
+
+    cart1 = CartDecomposition(D1, [n1], global_starts1, global_ends1, pads=[p1], shifts=[1])
+    cart2 = CartDecomposition(D2, [n2], global_starts2, global_ends2, pads=[p2], shifts=[1])
+    cart3 = CartDecomposition(D3, [n3], global_starts3, global_ends3, pads=[p3], shifts=[1])
+
+    V1 = StencilVectorSpace( cart1, dtype=dtype )
+    V2 = StencilVectorSpace( cart2, dtype=dtype )
+    V3 = StencilVectorSpace( cart3, dtype=dtype )
+
+    # 1D stencil matrices
+    M1 = StencilMatrix(V1, V1)
+    M2 = StencilMatrix(V2, V2)
+    M3 = StencilMatrix(V3, V3)
+
+    # ...
+    # Fill in stencil matrix values
+    for k1 in range(-p1, p1+1):
+        M1[:, k1] = 10 + k1*factor
+
+    for k2 in range(-p2, p2+1):
+        M2[:, k2] = 20 + k2*factor
+
+    for k3 in range(-p3, p3+1):
+        M3[:, k3] = 40 + k3*factor
+
+    M1.remove_spurious_entries()
+    M2.remove_spurious_entries()
+    M3.remove_spurious_entries()
+    # ...
+
+    # Fill in vector values
+    w[:, :, :] = factor
+
+    # Create Kronecker matrix 
+    M = KroneckerStencilMatrix(W, W, M1, M2, M3)
+
+    # Scipy sparse matrices used for comparison
+    M1_sp = M1.tosparse().tocsr()
+    M2_sp = M2.tosparse().tocsr()
+    M3_sp = M3.tosparse().tocsr()
+    M_sp  = reduce(kron, (M1_sp, M2_sp, M3_sp)).tocsr()
+
+    # Test transpose
+    assert (M_sp.T - M.T.tosparse().tocsr()).count_nonzero() == 0
+
+    # Test dot product
+    assert np.array_equal(M_sp.dot(w.toarray()), M.dot(w).toarray())