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

135 statements  

« prev     ^ index     » next       coverage.py v7.10.2, created at 2025-08-07 10:30 +0000

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 if window is None: 

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

53 

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

55 

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

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

58 

59 Returns: 

60 A new ChebyshevPolynomial object with the same attributes. 

61 """ 

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

63 

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

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

66 

67 Returns: 

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

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

70 """ 

71 if self.iscomplex: 

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

73 else: 

74 return self 

75 

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

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

78 

79 Returns: 

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

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

82 """ 

83 if self.iscomplex: 

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

85 else: 

86 return self 

87 

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

89 """Evaluate the polynomial at points x. 

90 

91 Args: 

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

93 

94 Returns: 

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

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

97 """ 

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

99 if np.isscalar(arg): 

100 # Map the input to the window 

101 mapped_arg = np.asarray(arg) 

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

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

104 ) + self.window[0] 

105 

106 # Evaluate the polynomial using the chebval function 

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

108 

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

110 return super().__call__(arg) 

111 

112 @property 

113 def iscomplex(self) -> bool: 

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

115 

116 Returns: 

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

118 """ 

119 return np.iscomplexobj(self.coef) 

120 

121 @property 

122 def size(self) -> int: 

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

124 

125 Returns: 

126 int: The number of coefficients in the polynomial. 

127 """ 

128 return self.coef.size 

129 

130 @property 

131 def isempty(self) -> bool: 

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

133 

134 Returns: 

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

136 """ 

137 return self.size == 0 

138 

139 @property 

140 def isconst(self) -> bool: 

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

142 

143 Returns: 

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

145 """ 

146 return self.size == 1 

147 

148 @property 

149 def vscale(self) -> float: 

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

151 

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

153 evaluated at Chebyshev points. 

154 

155 Returns: 

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

157 """ 

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

159 

160 def sum(self) -> float: 

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

162 

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

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

165 

166 Returns: 

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

168 """ 

169 if self.isempty: 

170 return 0.0 

171 

172 a, b = self.domain 

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

174 integral = ch.integ() 

175 return integral(b) - integral(a) 

176 

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

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

179 

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

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

182 

183 Args: 

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

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

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

187 

188 Returns: 

189 The axes on which the plot was created. 

190 """ 

191 ax = ax or plt.gca() 

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

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

194 ff = self(xx) 

195 if self.iscomplex: 

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

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

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

199 else: 

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

201 return ax 

202 

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

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

205 

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

207 

208 Returns: 

209 ChebyshevPolynomial: The derivative of the polynomial. 

210 """ 

211 # Get the coefficients of the derivative 

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

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

214 

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

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

217 

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

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

220 and the integration constant set to 0. 

221 

222 Returns: 

223 ChebyshevPolynomial: The antiderivative of the polynomial. 

224 """ 

225 # Get the coefficients of the antiderivative 

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

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

228 

229 @property 

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

231 """Get function values at Chebyshev points. 

232 

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

234 using the coeffs2vals2 algorithm. 

235 

236 Returns: 

237 np.ndarray: Function values at Chebyshev points. 

238 """ 

239 return coeffs2vals2(self.coef) 

240 

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

242 """Return a ChebyshevPolynomial of length n. 

243 

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

245 

246 Args: 

247 n: The desired number of coefficients. 

248 

249 Returns: 

250 ChebyshevPolynomial: A new polynomial with n coefficients. 

251 

252 Note: 

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

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

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

256 In all cases, a deep copy is returned. 

257 """ 

258 m = self.size 

259 ak = self.coef 

260 

261 if n < m: 

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

263 elif n > m: 

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

265 else: 

266 return self.copy() 

267 

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

269 

270 

271def from_coefficients( 

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

273) -> ChebyshevPolynomial: 

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

275 

276 Args: 

277 coef: Chebyshev coefficients in order of increasing degree. 

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

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

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

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

282 [-1, 1] is used. 

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

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

285 

286 Returns: 

287 A new Chebyshev polynomial with the given coefficients. 

288 

289 Raises: 

290 ValueError: If the coefficient array is empty. 

291 """ 

292 if len(coef) == 0: 

293 raise ValueError("Empty coefficients") 

294 

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

296 

297 

298def from_values( 

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

300) -> ChebyshevPolynomial: 

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

302 

303 Constructs a Chebyshev polynomial that interpolates the given values at 

304 Chebyshev points of the second kind. 

305 

306 Args: 

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

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

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

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

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

312 [-1, 1] is used. 

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

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

315 

316 Returns: 

317 A new Chebyshev polynomial that interpolates the given values. 

318 

319 Raises: 

320 ValueError: If the values array is empty. 

321 """ 

322 if len(values) == 0: 

323 raise ValueError("Empty values") 

324 

325 coef = vals2coeffs2(values) 

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

327 

328 

329def from_roots( 

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

331) -> ChebyshevPolynomial: 

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

333 

334 Constructs a Chebyshev polynomial that has the specified roots. 

335 

336 Args: 

337 roots: Sequence of root values. 

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

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

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

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

342 [-1, 1] is used. 

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

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

345 

346 Returns: 

347 A new Chebyshev polynomial with the specified roots. 

348 

349 Raises: 

350 ValueError: If the roots array is empty. 

351 """ 

352 if len(roots) == 0: 

353 raise ValueError("Empty roots") 

354 

355 coef = cheb.chebfromroots(roots) 

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

357 

358 

359def from_constant( 

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

361) -> ChebyshevPolynomial: 

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

363 

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

365 

366 Args: 

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

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

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

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

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

372 [-1, 1] is used. 

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

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

375 

376 Returns: 

377 A new Chebyshev polynomial representing the constant value. 

378 

379 Raises: 

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

381 """ 

382 if not np.isscalar(c): 

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

384 

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

386 if isinstance(c, int): 

387 c = float(c) 

388 

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

390 

391 

392def from_function( 

393 fun: callable, 

394 domain: DomainLike | None = None, 

395 window: DomainLike | None = None, 

396 symbol: str = "x", 

397 n: int | None = None, 

398) -> ChebyshevPolynomial: 

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

400 

401 Constructs a Chebyshev polynomial that approximates the given function. 

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

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

404 number of degrees of freedom. 

405 

406 Args: 

407 fun: Callable function to approximate. 

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

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

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

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

412 [-1, 1] is used. 

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

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

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

416 

417 Returns: 

418 A new Chebyshev polynomial that approximates the given function. 

419 """ 

420 from .algorithms import chebpts2, vals2coeffs2 

421 

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

423 

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

425 def mapped_fun(x): 

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

427 a, b = domain_arr 

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

429 return fun(mapped_x) 

430 

431 if n is None: 

432 # Use adaptive algorithm 

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

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

435 else: 

436 # Use fixed number of degrees of freedom 

437 points = chebpts2(n) 

438 values = mapped_fun(points) 

439 coeffs = vals2coeffs2(values) 

440 

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

442 

443 

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

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

446 

447 This function implements an adaptive algorithm to determine the appropriate 

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

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

450 and checking if the resulting coefficients can be truncated. 

451 

452 Args: 

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

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

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

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

457 value from preferences. 

458 

459 Returns: 

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

461 

462 Warns: 

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

464 number of iterations. 

465 """ 

466 minpow2 = 4 # 17 points 

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

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

469 n = 2**k + 1 

470 points = cheb.chebpts2(n) 

471 values = fun(points) 

472 coeffs = vals2coeffs2(values) 

473 eps = prefs.eps 

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

475 chplen = standard_chop(coeffs, tol=tol) 

476 if chplen < coeffs.size: 

477 coeffs = coeffs[:chplen] 

478 break 

479 if k == maxpow2: 

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

481 break 

482 return coeffs