Coverage for src / chebpy / chebyshev.py: 93%

1"""Immutable representation of Chebyshev polynomials. 

2 

3This module provides a class for the immutable representation of Chebyshev 

4polynomials and various factory functions to construct such polynomials. 

5""" 

6 

7import warnings 

8from typing import Any 

9 

10import matplotlib.pyplot as plt 

11import numpy as np 

12import numpy.polynomial.chebyshev as cheb 

13from matplotlib.axes import Axes 

14 

15from .algorithms import coeffs2vals2, standard_chop, vals2coeffs2 

16from .settings import _preferences as prefs 

17 

18# Type aliases 

19ArrayLike = list[float] | tuple[float, ...] | np.ndarray 

20DomainLike = tuple[float, float] | list[float] | np.ndarray 

21ScalarLike = int | float | complex 

22 

23 

24class ChebyshevPolynomial(cheb.Chebyshev): 

25 """Immutable representation of a Chebyshev polynomial. 

26 

27 This class represents a Chebyshev polynomial using its coefficients in the 

28 Chebyshev basis. The polynomial is defined on a specific domain. 

29 

30 Attributes: 

31 coef (np.ndarray): The coefficients of the Chebyshev polynomial. 

32 domain (np.ndarray): The domain on which the polynomial is defined. 

33 window (np.ndarray): The window on which the polynomial is mapped. Please use [-1, +1] 

34 symbol (str): Symbol used to represent the independent variable. 

35 """ 

36 

37 def __init__(self, coef: ArrayLike, domain: DomainLike | None = None, window=None, symbol: str = "x") -> None: 

38 """Initialize a ChebyshevPolynomial object. 

39 

40 Args: 

41 coef: Chebyshev coefficients in order of increasing degree. 

42 domain: Domain to use. The interval [domain[0], domain[1]] is mapped 

43 to the interval [window[0], window[1]] by shifting and scaling. 

44 If None, the default domain [-1, 1] is used. 

45 window: Window to use. The interval [domain[0], domain[1]] is mapped 

46 to the interval [window[0], window[1]] by shifting and scaling. 

47 If None, the default window [-1, 1] is used. 

48 symbol: Symbol used to represent the independent variable in string 

49 representations of the polynomial expression. Default is 'x'. 

50 

51 Examples: 

52 >>> import numpy as np 

53 >>> p = ChebyshevPolynomial([1, 2, 3]) 

54 >>> p.coef.tolist() 

55 [1.0, 2.0, 3.0] 

56 >>> p.domain.tolist() 

57 [-1.0, 1.0] 

58 """ 

59 if window is None: 

60 window = np.array([-1, 1]) 

61 

62 super().__init__(coef, domain, window=window, symbol=symbol) 

63 

64 def copy(self) -> "ChebyshevPolynomial": 

65 """Create a copy of the ChebyshevPolynomial object. 

66 

67 Returns: 

68 A new ChebyshevPolynomial object with the same attributes. 

69 """ 

70 return ChebyshevPolynomial(coef=self.coef.copy(), domain=self.domain.copy(), symbol=self.symbol) 

71 

72 def real(self) -> "ChebyshevPolynomial": 

73 """Return the real part of the polynomial. 

74 

75 Returns: 

76 ChebyshevPolynomial: A new polynomial with the real part of the coefficients 

77 if the polynomial is complex, otherwise the original polynomial. 

78 """ 

79 if self.iscomplex: 

80 return ChebyshevPolynomial(coef=np.real(self.coef), domain=self.domain, symbol=f"{self.symbol}") 

81 else: 

82 return self 

83 

84 def imag(self) -> "ChebyshevPolynomial": 

85 """Return the imaginary part of the polynomial. 

86 

87 Returns: 

88 ChebyshevPolynomial: A new polynomial with the imaginary part of the coefficients 

89 if the polynomial is complex, otherwise the original polynomial. 

90 """ 

91 if self.iscomplex: 

92 return ChebyshevPolynomial(coef=np.imag(self.coef), domain=self.domain, symbol=f"{self.symbol}") 

93 else: 

94 return self 

95 

96 def __call__(self, arg: ScalarLike | ArrayLike) -> ScalarLike | np.ndarray: 

97 """Evaluate the polynomial at points x. 

98 

99 Args: 

100 arg: Points at which to evaluate the polynomial. Can be a scalar or array-like. 

101 

102 Returns: 

103 If arg is a scalar, returns a scalar value. 

104 If arg is an array, returns an array of values. 

105 

106 Examples: 

107 >>> import numpy as np 

108 >>> p = ChebyshevPolynomial([1]) 

109 >>> float(p(0)) 

110 1.0 

111 >>> float(p(1)) 

112 1.0 

113 """ 

114 # If the input is a scalar, directly evaluate the polynomial 

115 if np.isscalar(arg): 

116 # Map the input to the window 

117 mapped_arg = np.asarray(arg) 

118 mapped_arg = (mapped_arg - self.domain[0]) / (self.domain[1] - self.domain[0]) * ( 

119 self.window[1] - self.window[0] 

120 ) + self.window[0] 

121 

122 # Evaluate the polynomial using the chebval function 

123 return cheb.chebval(mapped_arg, self.coef) 

124 

125 # For array inputs, call the parent class's __call__ method 

126 return super().__call__(arg) 

127 

128 @property 

129 def iscomplex(self) -> bool: 

130 """Determine whether the polynomial has complex coefficients. 

131 

132 Returns: 

133 bool: True if the polynomial has complex coefficients, False otherwise. 

134 """ 

135 return np.iscomplexobj(self.coef) 

136 

137 @property 

138 def size(self) -> int: 

139 """Return the size of the polynomial (number of coefficients). 

140 

141 Returns: 

142 int: The number of coefficients in the polynomial. 

143 """ 

144 return self.coef.size 

145 

146 @property 

147 def isempty(self) -> bool: 

148 """Return True if the polynomial is empty (has no coefficients). 

149 

150 Returns: 

151 bool: True if the polynomial has no coefficients, False otherwise. 

152 """ 

153 return self.size == 0 

154 

155 @property 

156 def isconst(self) -> bool: 

157 """Return True if the polynomial represents a constant (has only one coefficient). 

158 

159 Returns: 

160 bool: True if the polynomial has only one coefficient, False otherwise. 

161 """ 

162 return self.size == 1 

163 

164 @property 

165 def vscale(self) -> float: 

166 """Estimate the vertical scale of the polynomial. 

167 

168 The vertical scale is the maximum absolute value of the polynomial 

169 evaluated at Chebyshev points. 

170 

171 Returns: 

172 float: The maximum absolute value of the polynomial at Chebyshev points. 

173 """ 

174 return np.abs(self.values).max() 

175 

176 def sum(self) -> float: 

177 """Return the definite integral of the polynomial over its domain [a, b]. 

178 

179 Computes the definite integral of the polynomial over its domain using 

180 numpy.polynomial.chebyshev tools with correct domain/window logic. 

181 

182 Returns: 

183 float: The definite integral of the polynomial over its domain. 

184 """ 

185 if self.isempty: 

186 return 0.0 

187 

188 a, b = self.domain 

189 ch = ChebyshevPolynomial(self.coef, domain=self.domain) # window = [-1, 1] by default 

190 integral = ch.integ() 

191 return integral(b) - integral(a) 

192 

193 def plot(self, ax: Axes | None = None, n: int | None = None, **kwds: Any) -> Axes: 

194 """Plot the Chebyshev polynomial over its domain. 

195 

196 This method plots the Chebyshev polynomial over its domain using matplotlib. 

197 For complex-valued polynomials, it plots the real part against the imaginary part. 

198 

199 Args: 

200 ax: The axes on which to plot. If None, a new axes will be created. 

201 n: Number of points to use for plotting. If None, uses the value from preferences. 

202 **kwds: Additional keyword arguments to pass to matplotlib's plot function. 

203 

204 Returns: 

205 The axes on which the plot was created. 

206 """ 

207 ax = ax or plt.gca() 

208 n = n if n is not None else prefs.N_plot 

209 xx = np.linspace(self.domain[0], self.domain[1], n) 

210 ff = self(xx) 

211 if self.iscomplex: 

212 ax.plot(np.real(ff), np.imag(ff), **kwds) 

213 ax.set_xlabel(kwds.pop("ylabel", "real")) 

214 ax.set_ylabel(kwds.pop("xlabel", "imag")) 

215 else: 

216 ax.plot(xx, ff, **kwds) 

217 return ax 

218 

219 def diff(self) -> "ChebyshevPolynomial": 

220 """Return the derivative as a new ChebyshevPolynomial. 

221 

222 Computes the first derivative of the polynomial with respect to its variable. 

223 

224 Returns: 

225 ChebyshevPolynomial: The derivative of the polynomial. 

226 """ 

227 # Get the coefficients of the derivative 

228 deriv_coef = cheb.chebder(self.coef, m=1) 

229 return ChebyshevPolynomial(coef=deriv_coef, domain=self.domain, symbol=f"{self.symbol}") 

230 

231 def cumsum(self) -> "ChebyshevPolynomial": 

232 """Return the antiderivative as a new ChebyshevPolynomial. 

233 

234 Computes the first antiderivative of the polynomial with respect to its variable. 

235 The antiderivative is calculated with the lower bound set to the lower bound of the domain 

236 and the integration constant set to 0. 

237 

238 Returns: 

239 ChebyshevPolynomial: The antiderivative of the polynomial. 

240 """ 

241 # Get the coefficients of the antiderivative 

242 integ_coef = cheb.chebint(self.coef, m=1, lbnd=self.domain[0], k=0) 

243 return ChebyshevPolynomial(coef=integ_coef, domain=self.domain, symbol=f"{self.symbol}") 

244 

245 @property 

246 def values(self) -> np.ndarray: 

247 """Get function values at Chebyshev points. 

248 

249 Computes the values of the polynomial at Chebyshev points of the second kind 

250 using the coeffs2vals2 algorithm. 

251 

252 Returns: 

253 np.ndarray: Function values at Chebyshev points. 

254 """ 

255 return coeffs2vals2(self.coef) 

256 

257 def prolong(self, n: int) -> "ChebyshevPolynomial": 

258 """Return a ChebyshevPolynomial of length n. 

259 

260 Creates a new ChebyshevPolynomial with a specified number of coefficients. 

261 

262 Args: 

263 n: The desired number of coefficients. 

264 

265 Returns: 

266 ChebyshevPolynomial: A new polynomial with n coefficients. 

267 

268 Note: 

269 If n < self.size, the result is a truncated copy. 

270 If n > self.size, the result is zero-padded. 

271 If n == self.size, a copy of the original polynomial is returned. 

272 In all cases, a deep copy is returned. 

273 """ 

274 m = self.size 

275 ak = self.coef 

276 

277 if n < m: 

278 new_coeffs = ak[:n].copy() 

279 elif n > m: 

280 new_coeffs = np.concatenate([ak, np.zeros(n - m, dtype=ak.dtype)]) 

281 else: 

282 return self.copy() 

283 

284 return ChebyshevPolynomial(new_coeffs, domain=self.domain, symbol=self.symbol) 

285 

286 

287def from_coefficients( 

288 coef: ArrayLike, domain: DomainLike | None = None, window: DomainLike | None = None, symbol: str = "x" 

289) -> ChebyshevPolynomial: 

290 """Create a Chebyshev polynomial from its coefficients. 

291 

292 Args: 

293 coef: Chebyshev coefficients in order of increasing degree. 

294 domain: Domain to use. The interval [domain[0], domain[1]] is mapped 

295 to the interval [window[0], window[1]] by shifting and scaling. 

296 If None, the default domain [-1, 1] is used. 

297 window: Window, see domain for its use. If None, the default window 

298 [-1, 1] is used. 

299 symbol: Symbol used to represent the independent variable in string 

300 representations of the polynomial expression. Default is 'x'. 

301 

302 Returns: 

303 A new Chebyshev polynomial with the given coefficients. 

304 

305 Raises: 

306 ValueError: If the coefficient array is empty. 

307 

308 Examples: 

309 >>> import numpy as np 

310 >>> p = from_coefficients([3.14]) 

311 >>> p.coef.tolist() 

312 [3.14] 

313 >>> float(p(0)) 

314 3.14 

315 """ 

316 if len(coef) == 0: 

317 raise ValueError("Empty coefficients") 

318 

319 return ChebyshevPolynomial(coef, domain, window, symbol) 

320 

321 

322def from_values( 

323 values: ArrayLike, domain: DomainLike | None = None, window: DomainLike | None = None, symbol: str = "x" 

324) -> ChebyshevPolynomial: 

325 """Create a Chebyshev polynomial from values at Chebyshev points. 

326 

327 Constructs a Chebyshev polynomial that interpolates the given values at 

328 Chebyshev points of the second kind. 

329 

330 Args: 

331 values: Values at Chebyshev points of the second kind. 

332 domain: Domain to use. The interval [domain[0], domain[1]] is mapped 

333 to the interval [window[0], window[1]] by shifting and scaling. 

334 If None, the default domain [-1, 1] is used. 

335 window: Window, see domain for its use. If None, the default window 

336 [-1, 1] is used. 

337 symbol: Symbol used to represent the independent variable in string 

338 representations of the polynomial expression. Default is 'x'. 

339 

340 Returns: 

341 A new Chebyshev polynomial that interpolates the given values. 

342 

343 Raises: 

344 ValueError: If the values array is empty. 

345 """ 

346 if len(values) == 0: 

347 raise ValueError("Empty values") 

348 

349 coef = vals2coeffs2(values) 

350 return ChebyshevPolynomial(coef, domain, window, symbol) 

351 

352 

353def from_roots( 

354 roots: ArrayLike, domain: DomainLike | None = None, window: DomainLike | None = None, symbol: str = "x" 

355) -> ChebyshevPolynomial: 

356 """Create a Chebyshev polynomial from its roots. 

357 

358 Constructs a Chebyshev polynomial that has the specified roots. 

359 

360 Args: 

361 roots: Sequence of root values. 

362 domain: Domain to use. The interval [domain[0], domain[1]] is mapped 

363 to the interval [window[0], window[1]] by shifting and scaling. 

364 If None, the default domain [-1, 1] is used. 

365 window: Window, see domain for its use. If None, the default window 

366 [-1, 1] is used. 

367 symbol: Symbol used to represent the independent variable in string 

368 representations of the polynomial expression. Default is 'x'. 

369 

370 Returns: 

371 A new Chebyshev polynomial with the specified roots. 

372 

373 Raises: 

374 ValueError: If the roots array is empty. 

375 """ 

376 if len(roots) == 0: 

377 raise ValueError("Empty roots") 

378 

379 coef = cheb.chebfromroots(roots) 

380 return ChebyshevPolynomial(coef, domain, window, symbol) 

381 

382 

383def from_constant( 

384 c: ScalarLike, domain: DomainLike | None = None, window: DomainLike | None = None, symbol: str = "x" 

385) -> ChebyshevPolynomial: 

386 """Create a Chebyshev polynomial representing a constant value. 

387 

388 Constructs a Chebyshev polynomial of degree 0 that represents a constant value. 

389 

390 Args: 

391 c: The constant value (must be a scalar). 

392 domain: Domain to use. The interval [domain[0], domain[1]] is mapped 

393 to the interval [window[0], window[1]] by shifting and scaling. 

394 If None, the default domain [-1, 1] is used. 

395 window: Window, see domain for its use. If None, the default window 

396 [-1, 1] is used. 

397 symbol: Symbol used to represent the independent variable in string 

398 representations of the polynomial expression. Default is 'x'. 

399 

400 Returns: 

401 A new Chebyshev polynomial representing the constant value. 

402 

403 Raises: 

404 ValueError: If the input is not a scalar value. 

405 

406 Examples: 

407 >>> import numpy as np 

408 >>> p = from_constant(3.14) 

409 >>> float(p(0)) 

410 3.14 

411 >>> float(p(1)) 

412 3.14 

413 """ 

414 if not np.isscalar(c): 

415 raise ValueError("Input must be a scalar value") 

416 

417 # Convert integer to float to match behavior in other parts of the codebase 

418 if isinstance(c, int): 

419 c = float(c) 

420 

421 return ChebyshevPolynomial([c], domain, window, symbol) 

422 

423 

424def from_function( 

425 fun: callable, 

426 domain: DomainLike | None = None, 

427 window: DomainLike | None = None, 

428 symbol: str = "x", 

429 n: int | None = None, 

430) -> ChebyshevPolynomial: 

431 """Create a Chebyshev polynomial from a callable function. 

432 

433 Constructs a Chebyshev polynomial that approximates the given function. 

434 If n is provided, uses a fixed number of degrees of freedom. 

435 If n is None, uses an adaptive algorithm to determine the appropriate 

436 number of degrees of freedom. 

437 

438 Args: 

439 fun: Callable function to approximate. 

440 domain: Domain to use. The interval [domain[0], domain[1]] is mapped 

441 to the interval [window[0], window[1]] by shifting and scaling. 

442 If None, the default domain [-1, 1] is used. 

443 window: Window, see domain for its use. If None, the default window 

444 [-1, 1] is used. 

445 symbol: Symbol used to represent the independent variable in string 

446 representations of the polynomial expression. Default is 'x'. 

447 n: Number of degrees of freedom to use. If None, uses an adaptive algorithm. 

448 

449 Returns: 

450 A new Chebyshev polynomial that approximates the given function. 

451 """ 

452 from .algorithms import chebpts2, vals2coeffs2 

453 

454 domain_arr = np.array([-1, 1]) if domain is None else np.array(domain) 

455 

456 # Create a wrapper function that maps points from [-1, 1] to the custom domain 

457 def mapped_fun(x): 

458 # Map x from [-1, 1] to the custom domain 

459 a, b = domain_arr 

460 mapped_x = 0.5 * (b - a) * (x + 1) + a 

461 return fun(mapped_x) 

462 

463 if n is None: 

464 # Use adaptive algorithm 

465 hscale = (domain_arr[1] - domain_arr[0]) / 2 

466 coeffs = __adaptive(ChebyshevPolynomial, mapped_fun, hscale=hscale) 

467 else: 

468 # Use fixed number of degrees of freedom 

469 points = chebpts2(n) 

470 values = mapped_fun(points) 

471 coeffs = vals2coeffs2(values) 

472 

473 return ChebyshevPolynomial(coeffs, domain, window, symbol) 

474 

475 

476def __adaptive(cls: type, fun: callable, hscale: float = 1, maxpow2: int = None) -> np.ndarray: 

477 """Adaptively determine the number of points needed to represent a function. 

478 

479 This function implements an adaptive algorithm to determine the appropriate 

480 number of points needed to represent a function to a specified tolerance. 

481 It cycles over powers of two, evaluating the function at Chebyshev points 

482 and checking if the resulting coefficients can be truncated. 

483 

484 Args: 

485 cls: The class that provides the _chebpts and _vals2coeffs methods. 

486 fun (callable): The function to be approximated. 

487 hscale (float, optional): Scale factor for the tolerance. Defaults to 1. 

488 maxpow2 (int, optional): Maximum power of 2 to try. If None, uses the 

489 value from preferences. 

490 

491 Returns: 

492 numpy.ndarray: Coefficients of the Chebyshev series representing the function. 

493 

494 Warns: 

495 UserWarning: If the constructor does not converge within the maximum 

496 number of iterations. 

497 """ 

498 minpow2 = 4 # 17 points 

499 maxpow2 = maxpow2 if maxpow2 is not None else prefs.maxpow2 

500 for k in range(minpow2, max(minpow2, maxpow2) + 1): 

501 n = 2**k + 1 

502 points = cheb.chebpts2(n) 

503 values = fun(points) 

504 coeffs = vals2coeffs2(values) 

505 eps = prefs.eps 

506 tol = eps * max(hscale, 1) # scale (decrease) tolerance by hscale 

507 chplen = standard_chop(coeffs, tol=tol) 

508 if chplen < coeffs.size: 

509 coeffs = coeffs[:chplen] 

510 break 

511 if k == maxpow2: 

512 warnings.warn(f"The {cls.__name__} constructor did not converge: using {n} points") 

513 break 

514 return coeffs