def some_basic_tests(): import pyfftw from mpi4py import MPI comm = MPI.COMM_WORLD N = 8 K0 = shenfun.fourier.bases.C2CBasis(N) K1 = shenfun.fourier.bases.C2CBasis(N) K2 = shenfun.fourier.bases.C2CBasis(N) K3 = shenfun.fourier.bases.R2CBasis(N) T = TensorProductSpace(comm, (K0, K1, K2, K3)) # Create data on rank 0 for testing if comm.Get_rank() == 0: f_g = np.random.random(T.shape()) f_g_hat = pyfftw.interfaces.numpy_fft.rfftn(f_g, axes=(0, 1, 2, 3)) else: f_g = np.zeros(T.shape()) f_g_hat = np.zeros(T.spectral_shape(), dtype=np.complex) # Distribute test data to all ranks comm.Bcast(f_g, root=0) comm.Bcast(f_g_hat, root=0) # Create a function in real space to hold the test data fj = shenfun.Array(T) fj[:] = f_g[T.local_slice(False)] # Perform forward transformation f_hat = T.forward(fj) assert np.allclose(f_g_hat[T.local_slice(True)], f_hat * N**4) # Perform backward transformation fj2 = shenfun.Array(T) fj2 = T.backward(f_hat) assert np.allclose(fj, fj2) f_hat = T.scalar_product(fj) # Padding # Needs new instances of bases because arrays have new sizes Kp0 = shenfun.fourier.bases.C2CBasis(N, padding_factor=1.5) Kp1 = shenfun.fourier.bases.C2CBasis(N, padding_factor=1.5) Kp2 = shenfun.fourier.bases.C2CBasis(N, padding_factor=1.5) Kp3 = shenfun.fourier.bases.R2CBasis(N, padding_factor=1.5) Tp = TensorProductSpace(comm, (Kp0, Kp1, Kp2, Kp3)) f_g_pad = Tp.backward(f_hat) f_hat2 = Tp.forward(f_g_pad) assert np.allclose(f_hat2, f_hat)
def test_to_ortho(basis, quad): N = 10 if basis.family() == 'legendre': B1 = lBasis[0](N, quad) #B3 = lBasis[0](N, quad) elif basis.family() == 'chebyshev': if basis.short_name() == 'DU': B1 = cBasisGC[0](N, quad) else: B1 = cBasis[0](N, quad) #B3 = cBasis[0](N, quad) elif basis.family() == 'laguerre': B1 = laBasis[0](N, quad) #B3 = laBasis[0](N, quad) B0 = basis(N, quad=quad) a = shenfun.Array(B0) a_hat = shenfun.Function(B0) b0_hat = shenfun.Function(B1) b1_hat = shenfun.Function(B1) a[:] = np.random.random(a.shape) a_hat = a.forward(a_hat) b0_hat = shenfun.project(a_hat, B1, output_array=b0_hat, fill=False, use_to_ortho=True) b1_hat = shenfun.project(a_hat, B1, output_array=b1_hat, fill=False, use_to_ortho=False) assert np.linalg.norm(b0_hat-b1_hat) < 1e-10 #B2 = basis(N, quad=quad) TD = shenfun.TensorProductSpace(shenfun.comm, (B0, B0)) TC = shenfun.TensorProductSpace(shenfun.comm, (B1, B1)) a = shenfun.Array(TD) a_hat = shenfun.Function(TD) b0_hat = shenfun.Function(TC) b1_hat = shenfun.Function(TC) a[:] = np.random.random(a.shape) a_hat = a.forward(a_hat) b0_hat = shenfun.project(a_hat, TC, output_array=b0_hat, fill=False, use_to_ortho=True) b1_hat = shenfun.project(a_hat, TC, output_array=b1_hat, fill=False, use_to_ortho=False) assert np.linalg.norm(b0_hat-b1_hat) < 1e-10 F0 = shenfun.FunctionSpace(N, 'F') TD = shenfun.TensorProductSpace(shenfun.comm, (B0, F0)) TC = shenfun.TensorProductSpace(shenfun.comm, (B1, F0)) a = shenfun.Array(TD) a_hat = shenfun.Function(TD) b0_hat = shenfun.Function(TC) b1_hat = shenfun.Function(TC) a[:] = np.random.random(a.shape) a_hat = a.forward(a_hat) b0_hat = shenfun.project(a_hat, TC, output_array=b0_hat, fill=False, use_to_ortho=True) b1_hat = shenfun.project(a_hat, TC, output_array=b1_hat, fill=False, use_to_ortho=False) assert np.linalg.norm(b0_hat-b1_hat) < 1e-10
def test_transforms(ST, quad, dim): kwargs = {} if not ST.family() == 'fourier': kwargs['quad'] = quad ST0 = ST(N, **kwargs) fj = shenfun.Array(ST0) fj[:] = np.random.random(fj.shape[0]) # Project function to space first f_hat = shenfun.Function(ST0) f_hat = ST0.forward(fj, f_hat) fj = ST0.backward(f_hat, fj) # Then check if transformations work as they should u0 = shenfun.Function(ST0) u1 = shenfun.Array(ST0) u0 = ST0.forward(fj, u0) u1 = ST0.backward(u0, u1) assert np.allclose(fj, u1, rtol=1e-5, atol=1e-6) u0 = ST0.forward(fj, u0) u1 = ST0.backward(u0, u1) assert np.allclose(fj, u1, rtol=1e-5, atol=1e-6) u0 = ST0.forward(fj, u0, fast_transform=False) u1 = ST0.backward(u0, u1, fast_transform=False) assert np.allclose(fj, u1, rtol=1e-5, atol=1e-6) # Multidimensional version for axis in range(dim): bc = [ np.newaxis, ] * dim bc[axis] = slice(None) fij = np.broadcast_to(fj[tuple(bc)], (N, ) * dim).copy() ST1 = ST(N, **kwargs) ST1.tensorproductspace = ABC(dim, ST0.coors) ST1.plan((N, ) * dim, axis, fij.dtype, {}) u00 = shenfun.Function(ST1) u11 = shenfun.Array(ST1) u00 = ST1.forward(fij, u00) u11 = ST1.backward(u00, u11) cc = [ 0, ] * dim cc[axis] = slice(None) cc = tuple(cc) assert np.allclose(fij[cc], u11[cc], rtol=1e-5, atol=1e-6) del ST1
def test_axis(ST, quad, axis): kwargs = {'plan': True} if not ST.family() == 'fourier': kwargs['quad'] = quad ST = ST(N, **kwargs) points, weights = ST.points_and_weights(N) f_hat = shenfun.Function(ST) f_hat[:] = np.random.random(f_hat.shape[0]) B = inner_product((ST, 0), (ST, 0)) c = shenfun.Function(ST) c = B.solve(f_hat, c) # Multidimensional version f0 = shenfun.Array(ST) bc = [ np.newaxis, ] * 3 bc[axis] = slice(None) ST.plan((N, ) * 3, axis, f0.dtype, {}) if hasattr(ST, 'bc'): ST.bc.set_tensor_bcs( ST ) # To set Dirichlet boundary conditions on multidimensional array ck = shenfun.Function(ST) fk = np.broadcast_to(f_hat[bc], ck.shape).copy() ck = B.solve(fk, ck, axis=axis) cc = [ 0, ] * 3 cc[axis] = slice(None) assert np.allclose(ck[cc], c)
def _setup_variational_problem(self): self._setup_function_space() u = sf.TrialFunction(self._V) v = sf.TestFunction(self._V) self._elastic_law.set_material_parameters(self._material_parameters) self._dw_int = self._elastic_law.dw_int(u, v) self._dw_ext = inner( v, sf.Array(self._V.get_orthogonal(), buffer=(0, ) * self._dim)) if self._body_forces is not None: V_body_forces = self._V.get_orthogonal() body_forces_quad = sf.Array(V_body_forces, buffer=self._body_forces) self._dw_ext = inner(v, body_forces_quad)
def test_axis(ST, quad, axis): kwargs = {} if not ST.family() == 'fourier': kwargs['quad'] = quad ST = ST(N, **kwargs) points, weights = ST.points_and_weights(N) f_hat = shenfun.Function(ST) f_hat[:] = np.random.random(f_hat.shape[0]) B = inner_product((ST, 0), (ST, 0)) c = shenfun.Function(ST) c = B.solve(f_hat, c) # Multidimensional version f0 = shenfun.Array(ST) bc = [np.newaxis,]*3 bc[axis] = slice(None) ST.tensorproductspace = ABC(3, ST.coors) ST.plan((N,)*3, axis, f0.dtype, {}) if ST.has_nonhomogeneous_bcs: ST.bc.set_tensor_bcs(ST, ST) # To set Dirichlet boundary conditions on multidimensional array ck = shenfun.Function(ST) fk = np.broadcast_to(f_hat[tuple(bc)], ck.shape).copy() ck = B.solve(fk, ck, axis=axis) cc = [1,]*3 cc[axis] = slice(None) assert np.allclose(ck[tuple(cc)], c, rtol=1e-5, atol=1e-6)
def test_shear_test(): h = 0.1 length_ratio = 20 ell = h * length_ratio N = (round(length_ratio / 5) * 30, 30) domain = ((0., ell), (0., h)) print('Starting shear test ...') for elastic_law in (LinearCauchyElasticity(), LinearGradientElasticity()): for name_suffix in ('DisplacementControlled', 'TractionControlled'): Shear = ShearTest(N, domain, elastic_law, name_suffix) Shear.solve() u_ana_dl = get_dimensionless_displacement(Shear.u_ana, Shear._l_ref, Shear._u_ref) error_center = sf.Array(Shear.solution.function_space(), buffer=u_ana_dl)[0, round(N[0] / 2), :] - \ Shear.solution.backward()[0, round(N[0] / 2), :] error = np.linalg.norm(error_center) assert error < 1e-5, 'Error tolerance not achieved' Shear.postprocess() print( f'Error {elastic_law._name} ({name_suffix}):\t {error}\t N = {N}' ) print('Finished tensile test (clamped)!')
def postprocess(self): dirs = ['results', str(self.name), str(self.elastic_law.name)] for d in dirs: if not os.path.exists(d): os.mkdir(d) os.chdir(d) output = [] # displacement u = self.get_dimensional_solution() V = u.function_space() fl_disp_name = 'displacement' fl_disp = sf.ShenfunFile(fl_disp_name, V, backend='hdf5', mode='w', uniform=True) output.append(fl_disp_name + '.h5') for i in range(self.dim): fl_disp.write(i, {'u': [u.backward(kind='uniform')]}, as_scalar=True) # stress stress, space = self.elastic_law.compute_cauchy_stresses(u) fl_stress_name = 'cauchy_stress' fl_stress = sf.ShenfunFile(fl_stress_name, space, backend='hdf5', mode='w', uniform=True) for i in range(self.dim): for j in range(self.dim): s = sf.Array(space, buffer=stress[i, j]) fl_stress.write(0, {f'S{i}{j}': [s]}, as_scalar=True) output.append(fl_stress_name + '.h5') # hyper stress if self.elastic_law.name == 'LinearGradientElasticity': stress, space = self.elastic_law.compute_hyper_stresses(u) fl_stress_name = 'hyper_stress' fl_stress = sf.ShenfunFile(fl_stress_name, space, backend='hdf5', mode='w', uniform=True) for i in range(self.dim): for j in range(self.dim): for k in range(self.dim): s = sf.Array(space, buffer=stress[i, j]) fl_stress.write(0, {f'S{i}{j}{k}': [s]}, as_scalar=True) output.append(fl_stress_name + '.h5') self.write_xdmf_file(output) os.chdir('../../..')
def test_tensor2(): B0 = shenfun.FunctionSpace(8, 'C') T = shenfun.TensorProductSpace(comm, (B0, B0)) x, y = sp.symbols('x,y') ue = x**2 + y**2 ua = shenfun.Array(T, buffer=ue) uh = ua.forward() M = shenfun.VectorSpace(T) gradu = shenfun.project(grad(uh), M)
def test_transforms(ST, quad): N = 10 kwargs = {} if not ST.family() == 'fourier': kwargs['quad'] = quad ST0 = ST(N, **kwargs) fj = shenfun.Array(ST0) fj[:] = np.random.random(N) fj = fj.forward().backward().copy() assert np.allclose(fj, fj.forward().backward()) u0 = shenfun.Function(ST0) u1 = shenfun.Array(ST0) u0 = ST0.forward(fj, u0, fast_transform=False) u1 = ST0.backward(u0, u1, fast_transform=False) assert np.allclose(fj, u1, rtol=1e-5, atol=1e-6) # Multidimensional version ST0 = ST(N, **kwargs) if ST0.short_name() in ('R2C', 'C2C'): F0 = shenfun.FunctionSpace(N, 'F', dtype='D') T0 = shenfun.TensorProductSpace(shenfun.comm, (F0, ST0)) else: F0 = shenfun.FunctionSpace(N, 'F', dtype='d') T0 = shenfun.TensorProductSpace(shenfun.comm, (F0, ST0)) fij = shenfun.Array(T0) fij[:] = np.random.random(T0.shape(False)) fij = fij.forward().backward().copy() assert np.allclose(fij, fij.forward().backward()) if ST0.short_name() in ('R2C', 'C2C'): F0 = shenfun.FunctionSpace(N, 'F', dtype='D') F1 = shenfun.FunctionSpace(N, 'F', dtype='D') T = shenfun.TensorProductSpace(shenfun.comm, (F0, F1, ST0), dtype=ST0.dtype.char) else: F0 = shenfun.FunctionSpace(N, 'F', dtype='d') F1 = shenfun.FunctionSpace(N, ST.family()) T = shenfun.TensorProductSpace(shenfun.comm, (F0, ST0, F1)) fij = shenfun.Array(T) fij[:] = np.random.random(T.shape(False)) fij = fij.forward().backward().copy() assert np.allclose(fij, fij.forward().backward())
def test_transforms(ST, quad, axis): kwargs = {'plan': True} if not ST.family() == 'fourier': kwargs['quad'] = quad ST = ST(N, **kwargs) points, weights = ST.points_and_weights(N) fj = shenfun.Array(ST) fj[:] = np.random.random(fj.shape[0]) # Project function to space first f_hat = shenfun.Function(ST) f_hat = ST.forward(fj, f_hat) fj = ST.backward(f_hat, fj) # Then check if transformations work as they should u0 = shenfun.Function(ST) u1 = shenfun.Array(ST) u0 = ST.forward(fj, u0) u1 = ST.backward(u0, u1) assert np.allclose(fj, u1) u0 = ST.forward(fj, u0) u1 = ST.backward(u0, u1) assert np.allclose(fj, u1) # Multidimensional version bc = [ np.newaxis, ] * 3 bc[axis] = slice(None) fj = np.broadcast_to(fj[bc], (N, ) * 3).copy() ST.plan((N, ) * 3, axis, fj.dtype, {}) if hasattr(ST, 'bc'): ST.bc.set_slices(ST) # To set Dirichlet boundary conditions u00 = shenfun.Function(ST) u11 = shenfun.Array(ST) u00 = ST.forward(fj, u00) u11 = ST.backward(u00, u11) cc = [ 0, ] * 3 cc[axis] = slice(None) assert np.allclose(fj[cc], u11[cc])
def __call__(self, a_hat, b_hat, ab_hat=None): """Compute convolution of a_hat and b_hat without truncation Parameters ---------- a_hat : Function b_hat : Function ab_hat : Function """ Tp = self.padding_space T = self.newspace if ab_hat is None: ab_hat = shenfun.Function(T) a = shenfun.Array(Tp) b = shenfun.Array(Tp) a = Tp.backward(a_hat, a) b = Tp.backward(b_hat, b) ab_hat = T.forward(a * b, ab_hat) return ab_hat
def compute_numerical_error(u_ana, u_hat): assert isinstance(u_hat, sf.Function) V = u_hat.function_space() # evaluate u_ana at quadrature points error_array = sf.Array(V, buffer=u_ana) # subtract numerical solution error_array -= u_hat.backward() # compute integral error error = np.sqrt(sf.inner((1, 1), error_array**2)) return error
def test_scalarproduct(ST, quad): """Test fast scalar product against Vandermonde computed version""" kwargs = {} if not ST.family() == 'fourier': kwargs['quad'] = quad ST = ST(N, **kwargs) f = x*x+cos(pi*x) fj = shenfun.Array(ST, buffer=f) u0 = shenfun.Function(ST) u1 = shenfun.Function(ST) u0 = ST.scalar_product(fj, u0, fast_transform=True) u1 = ST.scalar_product(fj, u1, fast_transform=False) assert np.allclose(u1, u0) assert not np.all(u1 == u0) # Check that fast is not the same as slow
def test_tensor2(): B0 = shenfun.FunctionSpace(8, 'C') T = shenfun.TensorProductSpace(comm, (B0, B0)) x, y = sp.symbols('x,y') ue = x**2 + y**2 ua = shenfun.Array(T, buffer=ue) uh = ua.forward() M = shenfun.VectorSpace(T) gradu = shenfun.project(grad(uh), M) V = shenfun.TensorSpace(T) gradgradu = shenfun.project(grad(grad(uh)), V) g = gradgradu.backward() assert np.allclose(g.v[0], 2) assert np.allclose(g.v[1], 0) assert np.allclose(g.v[2], 0) assert np.allclose(g.v[3], 2)
def test_project_1D(basis): ue = sin(2*np.pi*x)*(1-x**2) T = basis(12) u = shenfun.TrialFunction(T) v = shenfun.TestFunction(T) u_tilde = shenfun.Function(T) X = T.mesh() ua = shenfun.Array(T, buffer=ue) u_tilde = shenfun.inner(v, ua, output_array=u_tilde) M = shenfun.inner(u, v) u_p = shenfun.Function(T) u_p = M.solve(u_tilde, u=u_p) u_0 = shenfun.Function(T) u_0 = shenfun.project(ua, T) assert np.allclose(u_0, u_p) u_1 = shenfun.project(ue, T) assert np.allclose(u_1, u_p)
def test_scalarproduct(ST, quad): """Test fast scalar product against Vandermonde computed version""" kwargs = {} if not ST.family() == 'fourier': kwargs['quad'] = quad ST = ST(N, **kwargs) points, weights = ST.points_and_weights(N) f = x * x + cos(pi * x) fl = lambdify(x, f, 'numpy') fj = shenfun.Array(ST) fj[:] = fl(points) u0 = shenfun.Function(ST) u1 = shenfun.Function(ST) u0 = ST.scalar_product(fj, u0, fast_transform=True) u1 = ST.scalar_product(fj, u1, fast_transform=False) assert np.allclose(u1, u0) assert not np.all(u1 == u0) # Check that fast is not the same as slow
def test_eval(ST, quad): """Test eval against fast inverse""" kwargs = {} if not ST.family() == 'fourier': kwargs['quad'] = quad ST = ST(N, **kwargs) points, weights = ST.mpmath_points_and_weights(N) fk = shenfun.Function(ST) fj = shenfun.Array(ST) fj[:] = np.random.random(fj.shape[0]) fk = ST.forward(fj, fk) fj = ST.backward(fk, fj) fk = ST.forward(fj, fk) f = ST.eval(points, fk) assert np.allclose(fj, f, rtol=1e-5, atol=1e-6), np.linalg.norm(fj-f) fj = ST.backward(fk, fj, fast_transform=False) fk = ST.forward(fj, fk, fast_transform=False) f = ST.eval(points, fk) assert np.allclose(fj, f, rtol=1e-5, atol=1e-6)
def check_with_pde(self, u_hat, material_parameters, body_forces): assert isinstance(u_hat, sf.Function) assert len(material_parameters) == self._n_material_parameters for comp in body_forces: assert isinstance(comp, (sp.Expr, float, int)) lambd, mu = material_parameters V = u_hat.function_space().get_orthogonal() # left hand side of pde lhs = (lambd + mu) * grad(div(u_hat)) + mu * div(grad(u_hat)) error_array = sf.Array(V, buffer=body_forces) error_array += sf.project(lhs, V).backward() error = np.sqrt(inner((1, 1), error_array ** 2)) # scale by magnitude of solution scale = np.sqrt(inner((1, 1), u_hat.backward() ** 2)) return error / scale
def test_eval(ST, quad): """Test eval agains fast inverse""" kwargs = {} if not ST.family() == 'fourier': kwargs['quad'] = quad ST = ST(N, **kwargs) points, weights = ST.points_and_weights(N) fk = shenfun.Function(ST) fj = shenfun.Array(ST) fj[:] = np.random.random(fj.shape[0]) fk = ST.forward(fj, fk) fj = ST.backward(fk, fj) fk = ST.forward(fj, fk) f = ST.eval(points, fk) #from IPython import embed; embed() assert np.allclose(fj, f) fj = ST.backward(fk, fj, fast_transform=False) fk = ST.forward(fj, fk, fast_transform=False) f = ST.eval(points, fk) assert np.allclose(fj, f)
def test_CXXmat(test, trial): test = test(N) trial = trial(N) CT = cBasis[0](N) Cm = inner_product((test, 0), (trial, 1)) S2 = Cm.trialfunction[0] S1 = Cm.testfunction[0] fj = shenfun.Array(S2, buffer=np.random.randn(N)) # project to S2 f_hat = fj.forward() fj = f_hat.backward(fj) # Check S1.scalar_product(f) equals Cm*S2.forward(f) f_hat = S2.forward(fj, f_hat) cs = np.zeros_like(f_hat) cs = Cm.matvec(f_hat, cs) df = shenfun.project(shenfun.grad(f_hat), CT).backward() cs2 = np.zeros(N) cs2 = S1.scalar_product(df, cs2) s = S1.slice() assert np.allclose(cs[s], cs2[s], rtol=1e-5, atol=1e-6) # Multidimensional version f_hat = f_hat.repeat(4 * 4).reshape( (N, 4, 4)) + 1j * f_hat.repeat(4 * 4).reshape((N, 4, 4)) df = df.repeat(4 * 4).reshape((N, 4, 4)) + 1j * df.repeat(4 * 4).reshape( (N, 4, 4)) cs = np.zeros_like(f_hat) cs = Cm.matvec(f_hat, cs) cs2 = np.zeros((N, 4, 4), dtype=np.complex) S1.tensorproductspace = ABC(3, S1.coors) S1.plan((N, 4, 4), 0, np.complex, {}) cs2 = S1.scalar_product(df, cs2) assert np.allclose(cs[s], cs2[s], rtol=1e-5, atol=1e-6)
def test_CXXmat(test, trial): test = test(N) trial = trial(N) CT = cBasis[0](N) Cm = inner_product((test, 0), (trial, 1)) S2 = Cm.trialfunction[0] S1 = Cm.testfunction[0] fj = shenfun.Array(S2, buffer=np.random.randn(N)) # project to S2 f_hat = fj.forward() fj = f_hat.backward(fj) # Check S1.scalar_product(f) equals Cm*S2.forward(f) f_hat = S2.forward(fj, f_hat) cs = np.zeros_like(f_hat) cs = Cm.matvec(f_hat, cs) df = shenfun.project(shenfun.grad(f_hat), CT).backward() cs2 = np.zeros(N) cs2 = S1.scalar_product(df, cs2) s = S1.slice() assert np.allclose(cs[s], cs2[s], rtol=1e-5, atol=1e-6)
import shenfun as sf import numpy as np import sympy as sp import matplotlib.pyplot as plt N = 10 L = sf.FunctionSpace(N, family='legendre', bc=(0, 0)) print(L.__class__) # plot quadrature points quad_points = L.mpmath_points_and_weights()[0] plt.scatter(quad_points, np.zeros_like(quad_points)) # a projection x = sp.symbols('x', real=True) f = 1 - 1 / 2 * (3 * x**2 - 1) # the first Shen-Dirichlet basis function P = sf.Function(L) func = sf.project(f, L) func_quad = sf.Array(L, buffer=f) plt.scatter(quad_points, func_quad) # should be orthogonal to all others except 0th and 2nd scalprod = L.scalar_product(func_quad) assert np.allclose(scalprod[np.r_[1, 3:N - 1]], 0) assert np.logical_not(np.allclose(scalprod[np.r_[0, 2]], 0))
e = shenfun.Expr(basis) e2 = -e assert np.allclose(np.array(e.scales()).astype(np.int), (-np.array(e2.scales())).astype(np.int)) K0 = shenfun.FunctionSpace(N, 'F', dtype='D') K1 = shenfun.FunctionSpace(N, 'F', dtype='D') K2 = shenfun.FunctionSpace(N, 'F', dtype='d') K3 = shenfun.FunctionSpace(N, 'C', dtype='d') T = shenfun.TensorProductSpace(comm, (K0, K1, K2)) C = shenfun.TensorProductSpace(comm, (K1, K2, K3)) TT = shenfun.VectorTensorProductSpace(T) CC = shenfun.VectorTensorProductSpace(C) VT = shenfun.MixedTensorProductSpace([TT, T]) KK = shenfun.MixedTensorProductSpace([T, T, C]) vf = shenfun.Function(TT) va = shenfun.Array(TT) cf = shenfun.Function(CC) ca = shenfun.Array(CC) df = shenfun.Function(KK) da = shenfun.Array(KK) @pytest.mark.parametrize('u', (va, vf, cf, ca, df, da)) def test_index(u): va0 = u[0] va1 = u[1] va2 = u[2] assert (va0.index(), va1.index(), va2.index()) == (0, 1, 2) assert va0.function_space() is u.function_space()[0] assert va1.function_space() is u.function_space()[1] assert va2.function_space() is u.function_space()[2]