def fst(self, fj, fk): """Fast Shen transform for general BC. """ fk = self.fastShenScalar(fj, fk) N = fj.shape[0] k = self.wavenumbers(N) k1 = self.wavenumbers(N + 1) ak, bk = self.shenCoefficients(k, self.BC) ak1, bk1 = self.shenCoefficients(k1, self.BC) if self.quad == "GC": ck = ones(N - 2) ck[0] = 2 elif self.quad == "GL": ck = ones(N - 2) ck[0] = 2 ck[-1] = 2 a = (pi / 2) * (ck + ak**2 + bk**2) b = ones(N - 3) * (pi / 2) * (ak[:-1] + ak1[1:-1] * bk[:-1]) c = ones(N - 4) * (pi / 2) * bk[:-2] if len(fk.shape) == 3: if self.BC[0] == 0 and self.BC[1] == 1 and self.BC[ 2] == 0 and self.BC[3] == 0 and self.BC[ 4] == 1 and self.BC[5] == 0: fk[1:-2] = SFTc.PDMA_3D_complex(a[1:], b[1:], c[1:], fk[1:-2]) else: fk[:-2] = SFTc.PDMA_3D_complex(a, b, c, fk[:-2]) elif len(fk.shape) == 1: if self.BC[0] == 0 and self.BC[1] == 1 and self.BC[ 2] == 0 and self.BC[3] == 0 and self.BC[ 4] == 1 and self.BC[5] == 0: fk[1:-2] = SFTc.PDMA_1D(a[1:], b[1:], c[1:], fk[1:-2]) else: fk[:-2] = SFTc.PDMA_1D(a, b, c, fk[:-2]) return fk
def __call__(self, u, b): if self.dim=="1": SFTc.PDMA_1D(self.a0, self.b0, self.c0, self.d0, self.e0, u, b) elif self.dim=="3": SFTc.PDMA(self.a0, self.b0, self.c0, self.d0, self.e0, u, b) return b