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 

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 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: 

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 if self.isempty: # pragma: no cover 

189 return 0.0 

190 

191 a, b = self.domain 

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

193 integral = ch.integ() 

194 return float(integral(b) - integral(a)) 

195 

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

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

198 

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

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

201 

202 Args: 

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

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

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

206 

207 Returns: 

208 The axes on which the plot was created. 

209 """ 

210 ax = ax or plt.gca() 

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

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

213 ff = self(xx) 

214 if self.iscomplex: 

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

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

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

218 else: 

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

220 return ax 

221 

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

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

224 

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

226 

227 Returns: 

228 ChebyshevPolynomial: The derivative of the polynomial. 

229 """ 

230 # Get the coefficients of the derivative 

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

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

233 

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

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

236 

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

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

239 and the integration constant set to 0. 

240 

241 Returns: 

242 ChebyshevPolynomial: The antiderivative of the polynomial. 

243 """ 

244 # Get the coefficients of the antiderivative 

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

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

247 

248 @property 

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

250 """Get function values at Chebyshev points. 

251 

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

253 using the coeffs2vals2 algorithm. 

254 

255 Returns: 

256 np.ndarray: Function values at Chebyshev points. 

257 """ 

258 return coeffs2vals2(self.coef) 

259 

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

261 """Return a ChebyshevPolynomial of length n. 

262 

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

264 

265 Args: 

266 n: The desired number of coefficients. 

267 

268 Returns: 

269 ChebyshevPolynomial: A new polynomial with n coefficients. 

270 

271 Note: 

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

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

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

275 In all cases, a deep copy is returned. 

276 """ 

277 m = self.size 

278 ak = self.coef 

279 

280 if n < m: 

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

282 elif n > m: 

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

284 else: 

285 return self.copy() 

286 

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

288 

289 

290def from_coefficients( 

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

292) -> ChebyshevPolynomial: 

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

294 

295 Args: 

296 coef: Chebyshev coefficients in order of increasing degree. 

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

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

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

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

301 [-1, 1] is used. 

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

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

304 

305 Returns: 

306 A new Chebyshev polynomial with the given coefficients. 

307 

308 Raises: 

309 ValueError: If the coefficient array is empty. 

310 

311 Examples: 

312 >>> import numpy as np 

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

314 >>> p.coef.tolist() 

315 [3.14] 

316 >>> float(p(0)) 

317 3.14 

318 """ 

319 if len(coef) == 0: 

320 raise ValueError(coef) 

321 

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

323 

324 

325def from_values( 

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

327) -> ChebyshevPolynomial: 

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

329 

330 Constructs a Chebyshev polynomial that interpolates the given values at 

331 Chebyshev points of the second kind. 

332 

333 Args: 

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

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

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

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

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

339 [-1, 1] is used. 

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

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

342 

343 Returns: 

344 A new Chebyshev polynomial that interpolates the given values. 

345 

346 Raises: 

347 ValueError: If the values array is empty. 

348 """ 

349 if len(values) == 0: 

350 raise ValueError(values) 

351 

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

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

354 

355 

356def from_roots( 

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

358) -> ChebyshevPolynomial: 

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

360 

361 Constructs a Chebyshev polynomial that has the specified roots. 

362 

363 Args: 

364 roots: Sequence of root values. 

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

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

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

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

369 [-1, 1] is used. 

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

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

372 

373 Returns: 

374 A new Chebyshev polynomial with the specified roots. 

375 

376 Raises: 

377 ValueError: If the roots array is empty. 

378 """ 

379 if len(roots) == 0: 

380 raise ValueError(roots) 

381 

382 coef = cheb.chebfromroots(roots) 

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

384 

385 

386def from_constant( 

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

388) -> ChebyshevPolynomial: 

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

390 

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

392 

393 Args: 

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

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

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

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

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

399 [-1, 1] is used. 

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

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

402 

403 Returns: 

404 A new Chebyshev polynomial representing the constant value. 

405 

406 Raises: 

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

408 

409 Examples: 

410 >>> import numpy as np 

411 >>> p = from_constant(3.14) 

412 >>> float(p(0)) 

413 3.14 

414 >>> float(p(1)) 

415 3.14 

416 """ 

417 if not np.isscalar(c): 

418 raise ValueError(c) 

419 

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

421 if isinstance(c, int): 

422 c = float(c) 

423 

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

425 

426 

427def from_function( 

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

429 domain: DomainLike | None = None, 

430 window: DomainLike | None = None, 

431 symbol: str = "x", 

432 n: int | None = None, 

433) -> ChebyshevPolynomial: 

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

435 

436 Constructs a Chebyshev polynomial that approximates the given function. 

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

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

439 number of degrees of freedom. 

440 

441 Args: 

442 fun: Callable function to approximate. 

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

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

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

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

447 [-1, 1] is used. 

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

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

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

451 

452 Returns: 

453 A new Chebyshev polynomial that approximates the given function. 

454 """ 

455 from .algorithms import chebpts2, vals2coeffs2 

456 

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

458 

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

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

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