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

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 collections.abc import Callable 

9from typing import Any, TypeAlias, cast 

10 

11import matplotlib.pyplot as plt 

12import numpy as np 

13import numpy.polynomial.chebyshev as cheb 

14from matplotlib.axes import Axes 

15 

16from .algorithms import chebpts2, coeffs2vals2, standard_chop, vals2coeffs2 

17from .settings import _preferences as prefs 

18 

19# Type aliases 

20ArrayLike: TypeAlias = list[float] | tuple[float, ...] | np.ndarray 

21DomainLike: TypeAlias = tuple[float, float] | list[float] | np.ndarray 

22ScalarLike: TypeAlias = int | float | complex 

23 

24 

25class ChebyshevPolynomial(cheb.Chebyshev): 

26 """Immutable representation of a Chebyshev polynomial. 

27 

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

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

30 

31 Attributes: 

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

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

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

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

36 """ 

37 

38 def __init__( 

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

40 ) -> None: 

41 """Initialize a ChebyshevPolynomial object. 

42 

43 Args: 

44 coef: Chebyshev coefficients in order of increasing degree. 

45 domain: Domain 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 domain [-1, 1] is used. 

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

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

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

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

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

53 

54 Examples: 

55 >>> import numpy as np 

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

57 >>> p.coef.tolist() 

58 [1.0, 2.0, 3.0] 

59 >>> p.domain.tolist() 

60 [-1.0, 1.0] 

61 """ 

62 if window is None: 

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

64 

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

66 

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

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

69 

70 Returns: 

71 A new ChebyshevPolynomial object with the same attributes. 

72 """ 

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

74 

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

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

77 

78 Returns: 

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

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

81 """ 

82 if self.iscomplex: 

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

84 else: 

85 return self 

86 

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

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

89 

90 Returns: 

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

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

93 """ 

94 if self.iscomplex: 

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

96 else: 

97 return self 

98 

99 def __call__(self, arg: ScalarLike | ArrayLike) -> ScalarLike | np.ndarray: # type: ignore[override] # numpy ABCPolyBase.__call__ uses a looser arg/return type 

100 """Evaluate the polynomial at points x. 

101 

102 Args: 

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

104 

105 Returns: 

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

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

108 

109 Examples: 

110 >>> import numpy as np 

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

112 >>> float(p(0)) 

113 1.0 

114 >>> float(p(1)) 

115 1.0 

116 """ 

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

118 if np.isscalar(arg): 

119 # Map the input to the window 

120 mapped_arg = np.asarray(arg) 

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

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

123 ) + self.window[0] 

124 

125 # Evaluate the polynomial using the chebval function 

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

127 

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

129 return super().__call__(arg) 

130 

131 @property 

132 def iscomplex(self) -> bool: 

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

134 

135 Returns: 

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

137 """ 

138 return np.iscomplexobj(self.coef) 

139 

140 @property 

141 def size(self) -> int: 

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

143 

144 Returns: 

145 int: The number of coefficients in the polynomial. 

146 """ 

147 return self.coef.size 

148 

149 @property 

150 def isempty(self) -> bool: 

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

152 

153 Returns: 

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

155 """ 

156 return self.size == 0 

157 

158 @property 

159 def isconst(self) -> bool: 

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

161 

162 Returns: 

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

164 """ 

165 return self.size == 1 

166 

167 @property 

168 def vscale(self) -> float: 

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

170 

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

172 evaluated at Chebyshev points. 

173 

174 Returns: 

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

176 """ 

177 return float(np.abs(self.values).max()) 

178 

179 def sum(self) -> float: 

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

181 

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

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

184 

185 Returns: 

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

187 """ 

188 # Defensive: numpy rejects an empty coefficient array at construction. 

189 if self.isempty: # pragma: no cover - defensive, see above 

190 return 0.0 

191 

192 a, b = self.domain 

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

194 integral = ch.integ() 

195 return float(cast(Any, integral(b) - integral(a))) 

196 

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

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

199 

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

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

202 

203 Args: 

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

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

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

207 

208 Returns: 

209 The axes on which the plot was created. 

210 """ 

211 ax = ax or plt.gca() 

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

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

214 ff = self(xx) 

215 if self.iscomplex: 

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

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

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

219 else: 

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

221 return ax 

222 

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

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

225 

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

227 

228 Returns: 

229 ChebyshevPolynomial: The derivative of the polynomial. 

230 """ 

231 # Get the coefficients of the derivative 

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

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

234 

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

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

237 

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

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

240 and the integration constant set to 0. 

241 

242 Returns: 

243 ChebyshevPolynomial: The antiderivative of the polynomial. 

244 """ 

245 # Get the coefficients of the antiderivative 

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

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

248 

249 @property 

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

251 """Get function values at Chebyshev points. 

252 

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

254 using the coeffs2vals2 algorithm. 

255 

256 Returns: 

257 np.ndarray: Function values at Chebyshev points. 

258 """ 

259 return coeffs2vals2(self.coef) 

260 

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

262 """Return a ChebyshevPolynomial of length n. 

263 

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

265 

266 Args: 

267 n: The desired number of coefficients. 

268 

269 Returns: 

270 ChebyshevPolynomial: A new polynomial with n coefficients. 

271 

272 Note: 

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

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

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

276 In all cases, a deep copy is returned. 

277 """ 

278 m = self.size 

279 ak = self.coef 

280 

281 if n < m: 

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

283 elif n > m: 

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

285 else: 

286 return self.copy() 

287 

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

289 

290 

291def from_coefficients( 

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

293) -> ChebyshevPolynomial: 

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

295 

296 Args: 

297 coef: Chebyshev coefficients in order of increasing degree. 

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

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

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

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

302 [-1, 1] is used. 

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

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

305 

306 Returns: 

307 A new Chebyshev polynomial with the given coefficients. 

308 

309 Raises: 

310 ValueError: If the coefficient array is empty. 

311 

312 Examples: 

313 >>> import numpy as np 

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

315 >>> p.coef.tolist() 

316 [3.14] 

317 >>> float(p(0)) 

318 3.14 

319 """ 

320 if len(coef) == 0: 

321 raise ValueError(coef) 

322 

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

324 

325 

326def from_values( 

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

328) -> ChebyshevPolynomial: 

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

330 

331 Constructs a Chebyshev polynomial that interpolates the given values at 

332 Chebyshev points of the second kind. 

333 

334 Args: 

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

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

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

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

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

340 [-1, 1] is used. 

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

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

343 

344 Returns: 

345 A new Chebyshev polynomial that interpolates the given values. 

346 

347 Raises: 

348 ValueError: If the values array is empty. 

349 """ 

350 if len(values) == 0: 

351 raise ValueError(values) 

352 

353 coef = vals2coeffs2(np.asarray(values)) 

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

355 

356 

357def from_roots( 

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

359) -> ChebyshevPolynomial: 

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

361 

362 Constructs a Chebyshev polynomial that has the specified roots. 

363 

364 Args: 

365 roots: Sequence of root values. 

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

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

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

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

370 [-1, 1] is used. 

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

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

373 

374 Returns: 

375 A new Chebyshev polynomial with the specified roots. 

376 

377 Raises: 

378 ValueError: If the roots array is empty. 

379 """ 

380 if len(roots) == 0: 

381 raise ValueError(roots) 

382 

383 coef = cheb.chebfromroots(roots) 

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

385 

386 

387def from_constant( 

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

389) -> ChebyshevPolynomial: 

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

391 

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

393 

394 Args: 

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

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

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

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

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

400 [-1, 1] is used. 

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

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

403 

404 Returns: 

405 A new Chebyshev polynomial representing the constant value. 

406 

407 Raises: 

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

409 

410 Examples: 

411 >>> import numpy as np 

412 >>> p = from_constant(3.14) 

413 >>> float(p(0)) 

414 3.14 

415 >>> float(p(1)) 

416 3.14 

417 """ 

418 if not np.isscalar(c): 

419 raise ValueError(c) 

420 

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

422 if isinstance(c, int): 

423 c = float(c) 

424 

425 return ChebyshevPolynomial(cast(ArrayLike, [c]), domain, window, symbol) 

426 

427 

428def from_function( 

429 fun: Callable[..., Any], 

430 domain: DomainLike | None = None, 

431 window: DomainLike | None = None, 

432 symbol: str = "x", 

433 n: int | None = None, 

434) -> ChebyshevPolynomial: 

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

436 

437 Constructs a Chebyshev polynomial that approximates the given function. 

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

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

440 number of degrees of freedom. 

441 

442 Args: 

443 fun: Callable function to approximate. 

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

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

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

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

448 [-1, 1] is used. 

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

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

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

452 

453 Returns: 

454 A new Chebyshev polynomial that approximates the given function. 

455 """ 

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

457 

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

459 def mapped_fun(x: np.ndarray) -> np.ndarray: 

460 """Evaluate the user function after mapping *x* from [-1, 1] to the custom domain.""" 

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

462 a, b = domain_arr 

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

464 return np.asarray(fun(mapped_x)) 

465 

466 if n is None: 

467 # Use adaptive algorithm 

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

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

470 else: 

471 # Use fixed number of degrees of freedom 

472 points = chebpts2(n) 

473 values = mapped_fun(points) 

474 coeffs = vals2coeffs2(values) 

475 

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

477 

478 

479def __adaptive(cls: type, fun: Callable[..., Any], hscale: float = 1, maxpow2: int | None = None) -> np.ndarray: 

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

481 

482 This function implements an adaptive algorithm to determine the appropriate 

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

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

485 and checking if the resulting coefficients can be truncated. 

486 

487 Args: 

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

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

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

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

492 value from preferences. 

493 

494 Returns: 

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

496 

497 Warns: 

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

499 number of iterations. 

500 """ 

501 minpow2 = 4 # 17 points 

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

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

504 n = 2**k + 1 

505 points = cheb.chebpts2(n) 

506 values = fun(points) 

507 coeffs = vals2coeffs2(values) 

508 eps = prefs.eps 

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

510 chplen = standard_chop(coeffs, tol=tol) 

511 if chplen < coeffs.size: 

512 coeffs = coeffs[:chplen] 

513 break 

514 if k == maxpow2: 

515 warnings.warn(f"The {cls.__name__} constructor did not converge: using {n} points", stacklevel=2) 

516 break 

517 return coeffs