| 20 | return np.polynomial.polynomial.Polynomial(phi_c)(x) * (1-mask) |
| 21 | |
| 22 | def get_phi_psi(k, base): |
| 23 | |
| 24 | x = Symbol('x') |
| 25 | phi_coeff = np.zeros((k,k)) |
| 26 | phi_2x_coeff = np.zeros((k,k)) |
| 27 | if base == 'legendre': |
| 28 | for ki in range(k): |
| 29 | coeff_ = Poly(legendre(ki, 2*x-1), x).all_coeffs() |
| 30 | phi_coeff[ki,:ki+1] = np.flip(np.sqrt(2*ki+1) * np.array(coeff_).astype(np.float64)) |
| 31 | coeff_ = Poly(legendre(ki, 4*x-1), x).all_coeffs() |
| 32 | phi_2x_coeff[ki,:ki+1] = np.flip(np.sqrt(2) * np.sqrt(2*ki+1) * np.array(coeff_).astype(np.float64)) |
| 33 | |
| 34 | psi1_coeff = np.zeros((k, k)) |
| 35 | psi2_coeff = np.zeros((k, k)) |
| 36 | for ki in range(k): |
| 37 | psi1_coeff[ki,:] = phi_2x_coeff[ki,:] |
| 38 | for i in range(k): |
| 39 | a = phi_2x_coeff[ki,:ki+1] |
| 40 | b = phi_coeff[i, :i+1] |
| 41 | prod_ = np.convolve(a, b) |
| 42 | prod_[np.abs(prod_)<1e-8] = 0 |
| 43 | proj_ = (prod_ * 1/(np.arange(len(prod_))+1) * np.power(0.5, 1+np.arange(len(prod_)))).sum() |
| 44 | psi1_coeff[ki,:] -= proj_ * phi_coeff[i,:] |
| 45 | psi2_coeff[ki,:] -= proj_ * phi_coeff[i,:] |
| 46 | for j in range(ki): |
| 47 | a = phi_2x_coeff[ki,:ki+1] |
| 48 | b = psi1_coeff[j, :] |
| 49 | prod_ = np.convolve(a, b) |
| 50 | prod_[np.abs(prod_)<1e-8] = 0 |
| 51 | proj_ = (prod_ * 1/(np.arange(len(prod_))+1) * np.power(0.5, 1+np.arange(len(prod_)))).sum() |
| 52 | psi1_coeff[ki,:] -= proj_ * psi1_coeff[j,:] |
| 53 | psi2_coeff[ki,:] -= proj_ * psi2_coeff[j,:] |
| 54 | |
| 55 | a = psi1_coeff[ki,:] |
| 56 | prod_ = np.convolve(a, a) |
| 57 | prod_[np.abs(prod_)<1e-8] = 0 |
| 58 | norm1 = (prod_ * 1/(np.arange(len(prod_))+1) * np.power(0.5, 1+np.arange(len(prod_)))).sum() |
| 59 | |
| 60 | a = psi2_coeff[ki,:] |
| 61 | prod_ = np.convolve(a, a) |
| 62 | prod_[np.abs(prod_)<1e-8] = 0 |
| 63 | norm2 = (prod_ * 1/(np.arange(len(prod_))+1) * (1-np.power(0.5, 1+np.arange(len(prod_))))).sum() |
| 64 | norm_ = np.sqrt(norm1 + norm2) |
| 65 | psi1_coeff[ki,:] /= norm_ |
| 66 | psi2_coeff[ki,:] /= norm_ |
| 67 | psi1_coeff[np.abs(psi1_coeff)<1e-8] = 0 |
| 68 | psi2_coeff[np.abs(psi2_coeff)<1e-8] = 0 |
| 69 | |
| 70 | phi = [np.poly1d(np.flip(phi_coeff[i,:])) for i in range(k)] |
| 71 | psi1 = [np.poly1d(np.flip(psi1_coeff[i,:])) for i in range(k)] |
| 72 | psi2 = [np.poly1d(np.flip(psi2_coeff[i,:])) for i in range(k)] |
| 73 | |
| 74 | elif base == 'chebyshev': |
| 75 | for ki in range(k): |
| 76 | if ki == 0: |
| 77 | phi_coeff[ki,:ki+1] = np.sqrt(2/np.pi) |
| 78 | phi_2x_coeff[ki,:ki+1] = np.sqrt(2/np.pi) * np.sqrt(2) |
| 79 | else: |