Coverage for src/chebpy/chebtech.py: 100%
252 statements
« prev ^ index » next coverage.py v7.15.2, created at 2026-07-22 21:50 +0000
1"""Implementation of Chebyshev polynomial technology for function approximation.
3This module provides the Chebtech class, which is an abstract base class for
4representing functions using Chebyshev polynomial expansions. It serves as the
5foundation for the Chebtech class, which uses Chebyshev points of the second kind.
7The Chebtech classes implement core functionality for working with Chebyshev
8expansions, including:
9- Function evaluation using Clenshaw's algorithm or barycentric interpolation
10- Algebraic operations (addition, multiplication, etc.)
11- Calculus operations (differentiation, integration, etc.)
12- Rootfinding
13- Plotting
15These classes are primarily used internally by higher-level classes like Bndfun
16and Chebfun, rather than being used directly by end users.
17"""
19from abc import ABC
20from collections.abc import Callable
21from typing import Any
23import matplotlib.pyplot as plt
24import numpy as np
26from .algorithms import (
27 adaptive,
28 bary,
29 barywts2,
30 chebpts2,
31 clenshaw,
32 coeffmult,
33 coeffs2vals2,
34 newtonroots,
35 rootsunit,
36 standard_chop,
37 vals2coeffs2,
38)
39from .decorators import self_empty
40from .plotting import plotfun, plotfuncoeffs
41from .settings import _preferences as prefs
42from .smoothfun import Smoothfun
43from .utilities import Interval, coerce_list
46class Chebtech(Smoothfun, ABC):
47 """Abstract base class serving as the template for Chebtech1 and Chebtech subclasses.
49 Chebtech objects always work with first-kind coefficients, so much
50 of the core operational functionality is defined this level.
52 The user will rarely work with these classes directly so we make
53 several assumptions regarding input data types.
54 """
56 @classmethod
57 def initconst(cls, c: Any = None, *, interval: Any = None) -> Any:
58 """Initialise a Chebtech from a constant c."""
59 if not np.isscalar(c):
60 raise ValueError(c)
61 if isinstance(c, int):
62 c = float(c)
63 return cls(np.array([c]), interval=interval)
65 @classmethod
66 def initempty(cls, *, interval: Any = None) -> "Chebtech":
67 """Initialise an empty Chebtech."""
68 return cls(np.array([]), interval=interval)
70 @classmethod
71 def initidentity(cls, *, interval: Any = None) -> "Chebtech":
72 """Chebtech representation of f(x) = x on [-1,1]."""
73 return cls(np.array([0, 1]), interval=interval)
75 @classmethod
76 def initfun(cls, fun: Any = None, n: Any = None, *, interval: Any = None) -> Any:
77 """Convenience constructor to automatically select the adaptive or fixedlen constructor.
79 This constructor automatically selects between the adaptive or fixed-length
80 constructor based on the input arguments passed.
81 """
82 if n is None:
83 return cls.initfun_adaptive(fun, interval=interval)
84 else:
85 return cls.initfun_fixedlen(fun, n, interval=interval)
87 @classmethod
88 def initfun_fixedlen(cls, fun: Any = None, n: Any = None, *, interval: Any = None) -> Any:
89 """Initialise a Chebtech from the callable fun using n degrees of freedom.
91 This constructor creates a Chebtech representation of the function using
92 a fixed number of degrees of freedom specified by n.
93 """
94 if n is None:
95 raise ValueError("n must be specified for fixed-length initialization") # noqa: TRY003
96 points = cls._chebpts(int(n))
97 values = fun(points)
98 coeffs = vals2coeffs2(values)
99 return cls(coeffs, interval=interval)
101 @classmethod
102 def initfun_adaptive(cls, fun: Any = None, *, interval: Any = None) -> Any:
103 """Initialise a Chebtech from the callable fun utilising the adaptive constructor.
105 This constructor uses an adaptive algorithm to determine the appropriate
106 number of degrees of freedom needed to represent the function.
107 """
108 interval = interval if interval is not None else prefs.domain
109 interval = Interval(*interval)
110 coeffs = adaptive(cls, fun, hscale=interval.hscale)
111 return cls(coeffs, interval=interval)
113 @classmethod
114 def initvalues(cls, values: Any = None, *, interval: Any = None) -> Any:
115 """Initialise a Chebtech from an array of values at Chebyshev points."""
116 return cls(cls._vals2coeffs(values), interval=interval)
118 def __init__(self, coeffs: Any, interval: Any = None) -> None:
119 """Initialize a Chebtech object.
121 This method initializes a new Chebtech object with the given coefficients
122 and interval. If no interval is provided, the default interval from
123 preferences is used.
125 Args:
126 coeffs (array-like): The coefficients of the Chebyshev series.
127 interval (array-like, optional): The interval on which the function
128 is defined. Defaults to None, which uses the default interval
129 from preferences.
130 """
131 interval = interval if interval is not None else prefs.domain
132 self._coeffs = np.array(coeffs)
133 self._interval = Interval(*interval)
135 def __call__(self, x: Any, how: str = "clenshaw") -> Any:
136 """Evaluate the Chebtech at the given points.
138 Args:
139 x: Points at which to evaluate the Chebtech.
140 how (str, optional): Method to use for evaluation. Either "clenshaw" or "bary".
141 Defaults to "clenshaw".
143 Returns:
144 The values of the Chebtech at the given points.
146 Raises:
147 ValueError: If the specified method is not supported.
148 """
149 method: dict[str, Callable[[Any], Any]] = {
150 "clenshaw": self.__call__clenshaw,
151 "bary": self.__call__bary,
152 }
153 try:
154 return method[how](x)
155 except KeyError as err:
156 raise ValueError(how) from err
158 def __call__clenshaw(self, x: Any) -> Any:
159 """Evaluate at *x* using Clenshaw recurrence on the coefficients."""
160 return clenshaw(x, self.coeffs)
162 def __call__bary(self, x: Any) -> Any:
163 """Evaluate at *x* using the barycentric interpolation formula."""
164 fk = self.values()
165 xk = self._chebpts(fk.size)
166 vk = self._barywts(fk.size)
167 return bary(x, fk, xk, vk)
169 def __repr__(self) -> str: # pragma: no cover
170 """Return a string representation of the Chebtech.
172 Returns:
173 str: A string representation of the Chebtech.
174 """
175 out = f"<{self.__class__.__name__}{{{self.size}}}>"
176 return out
178 # ------------
179 # properties
180 # ------------
181 @property
182 def coeffs(self) -> np.ndarray:
183 """Chebyshev expansion coefficients in the T_k basis."""
184 return self._coeffs
186 @property
187 def interval(self) -> Interval:
188 """Interval that Chebtech is mapped to."""
189 return self._interval
191 @property
192 def size(self) -> int:
193 """Return the size of the object."""
194 return self.coeffs.size
196 @property
197 def isempty(self) -> bool:
198 """Return True if the Chebtech is empty."""
199 return self.size == 0
201 @property
202 def iscomplex(self) -> bool:
203 """Determine whether the underlying onefun is complex or real valued."""
204 return self._coeffs.dtype == complex
206 @property
207 def isconst(self) -> bool:
208 """Return True if the Chebtech represents a constant."""
209 return self.size == 1
211 @property
212 @self_empty(0.0)
213 def vscale(self) -> float:
214 """Estimate the vertical scale of a Chebtech."""
215 return float(np.abs(np.asarray(coerce_list(self.values()))).max())
217 # -----------
218 # utilities
219 # -----------
220 def copy(self) -> "Chebtech":
221 """Return a deep copy of the Chebtech."""
222 return self.__class__(self.coeffs.copy(), interval=self.interval.copy())
224 def imag(self) -> Any:
225 """Return the imaginary part of the Chebtech.
227 Returns:
228 Chebtech: A new Chebtech representing the imaginary part of this Chebtech.
229 If this Chebtech is real-valued, returns a zero Chebtech.
230 """
231 if self.iscomplex:
232 return self.__class__(np.imag(self.coeffs), self.interval)
233 else:
234 return self.initconst(0, interval=self.interval)
236 def prolong(self, n: int) -> "Chebtech":
237 """Return a Chebtech of length n.
239 Obtained either by truncating if n < self.size or zero-padding if n > self.size.
240 In all cases a deep copy is returned.
241 """
242 m = self.size
243 ak = self.coeffs
244 cls = self.__class__
245 if n - m < 0:
246 out = cls(ak[:n].copy(), interval=self.interval)
247 elif n - m > 0:
248 out = cls(np.append(ak, np.zeros(n - m)), interval=self.interval)
249 else:
250 out = self.copy()
251 return out
253 def real(self) -> "Chebtech":
254 """Return the real part of the Chebtech.
256 Returns:
257 Chebtech: A new Chebtech representing the real part of this Chebtech.
258 If this Chebtech is already real-valued, returns self.
259 """
260 if self.iscomplex:
261 return self.__class__(np.real(self.coeffs), self.interval)
262 else:
263 return self
265 def simplify(self) -> "Chebtech":
266 """Call standard_chop on the coefficients of self.
268 Returns a Chebtech comprised of a copy of the truncated coefficients.
269 """
270 # coefficients
271 oldlen = len(self.coeffs)
272 longself = self.prolong(max(17, oldlen))
273 cfs = longself.coeffs
274 # scale (decrease) tolerance by hscale
275 tol = prefs.eps * max(self.interval.hscale, 1)
276 # chop
277 npts = standard_chop(cfs, tol=tol)
278 npts = min(oldlen, npts)
279 # construct
280 return self.__class__(cfs[:npts].copy(), interval=self.interval)
282 def values(self) -> np.ndarray:
283 """Function values at Chebyshev points."""
284 return coeffs2vals2(self.coeffs)
286 # ---------
287 # algebra
288 # ---------
289 @self_empty()
290 def __add__(self, f: Any) -> Any:
291 """Add a scalar or another Chebtech to this Chebtech.
293 Args:
294 f: A scalar or another Chebtech to add to this Chebtech.
296 Returns:
297 Chebtech: A new Chebtech representing the sum.
298 """
299 cls = self.__class__
300 if np.isscalar(f):
301 if np.iscomplexobj(f):
302 dtype: Any = complex
303 else:
304 dtype = self.coeffs.dtype
305 cfs = np.array(self.coeffs, dtype=dtype)
306 cfs[0] += f
307 return cls(cfs, interval=self.interval)
308 else:
309 # TODO: is a more general decorator approach better here?
310 # TODO: for constant Chebtech, convert to constant and call __add__ again
311 if f.isempty:
312 return f.copy()
313 g = self
314 n, m = g.size, f.size
315 if n < m:
316 g = g.prolong(m)
317 elif m < n:
318 f = f.prolong(n)
319 cfs = f.coeffs + g.coeffs
321 # check for zero output
322 eps = prefs.eps
323 tol = 0.5 * eps * max([f.vscale, g.vscale])
324 if all(abs(cfs) < tol):
325 return cls.initconst(0.0, interval=self.interval)
326 else:
327 return cls(cfs, interval=self.interval)
329 @self_empty()
330 def __div__(self, f: Any) -> Any:
331 """Divide this Chebtech by a scalar or another Chebtech.
333 Args:
334 f: A scalar or another Chebtech to divide this Chebtech by.
336 Returns:
337 Chebtech: A new Chebtech representing the quotient.
338 """
339 cls = self.__class__
340 if np.isscalar(f):
341 cfs = 1.0 / np.asarray(f) * self.coeffs
342 return cls(cfs, interval=self.interval)
343 else:
344 # TODO: review with reference to __add__
345 if f.isempty:
346 return f.copy()
347 return cls.initfun_adaptive(lambda x: self(x) / f(x), interval=self.interval)
349 __truediv__ = __div__
351 @self_empty()
352 def __mul__(self, g: Any) -> Any:
353 """Multiply this Chebtech by a scalar or another Chebtech.
355 Args:
356 g: A scalar or another Chebtech to multiply this Chebtech by.
358 Returns:
359 Chebtech: A new Chebtech representing the product.
360 """
361 cls = self.__class__
362 if np.isscalar(g):
363 cfs = g * self.coeffs
364 return cls(cfs, interval=self.interval)
365 else:
366 # TODO: review with reference to __add__
367 if g.isempty:
368 return g.copy()
369 f = self
370 n = f.size + g.size - 1
371 f = f.prolong(n)
372 g = g.prolong(n)
373 cfs = coeffmult(f.coeffs, g.coeffs)
374 out = cls(cfs, interval=self.interval)
375 return out
377 def __neg__(self) -> "Chebtech":
378 """Return the negative of this Chebtech.
380 Returns:
381 Chebtech: A new Chebtech representing the negative of this Chebtech.
382 """
383 coeffs = -self.coeffs
384 return self.__class__(coeffs, interval=self.interval)
386 def __pos__(self) -> "Chebtech":
387 """Return this Chebtech (unary positive).
389 Returns:
390 Chebtech: This Chebtech (self).
391 """
392 return self
394 @self_empty()
395 def __pow__(self, f: Any) -> Any:
396 """Raise this Chebtech to a power.
398 Args:
399 f: The exponent, which can be a scalar or another Chebtech.
401 Returns:
402 Chebtech: A new Chebtech representing this Chebtech raised to the power f.
403 """
405 def powfun(fn: Any, x: Any) -> Any:
406 if np.isscalar(fn):
407 return fn
408 else:
409 return fn(x)
411 return self.__class__.initfun_adaptive(lambda x: np.power(self(x), powfun(f, x)), interval=self.interval)
413 def __rdiv__(self, f: Any) -> Any:
414 """Divide a scalar by this Chebtech.
416 This is called when f / self is executed and f is not a Chebtech.
418 Args:
419 f: A scalar to be divided by this Chebtech.
421 Returns:
422 Chebtech: A new Chebtech representing f divided by this Chebtech.
423 """
425 # Executed when __div__(f, self) fails, which is to say whenever f
426 # is not a Chebtech. We proceeed on the assumption f is a scalar.
427 def constfun(x: Any) -> Any:
428 return 0.0 * x + f
430 return self.__class__.initfun_adaptive(lambda x: constfun(x) / self(x), interval=self.interval)
432 __radd__ = __add__
434 def __rsub__(self, f: Any) -> Any:
435 """Subtract this Chebtech from a scalar.
437 This is called when f - self is executed and f is not a Chebtech.
439 Args:
440 f: A scalar from which to subtract this Chebtech.
442 Returns:
443 Chebtech: A new Chebtech representing f minus this Chebtech.
444 """
445 return -(self - f)
447 @self_empty()
448 def __rpow__(self, f: Any) -> Any:
449 """Raise a scalar to the power of this Chebtech.
451 This is called when f ** self is executed and f is not a Chebtech.
453 Args:
454 f: A scalar to be raised to the power of this Chebtech.
456 Returns:
457 Chebtech: A new Chebtech representing f raised to the power of this Chebtech.
458 """
459 return self.__class__.initfun_adaptive(lambda x: np.power(f, self(x)), interval=self.interval)
461 __rtruediv__ = __rdiv__
462 __rmul__ = __mul__
464 def __sub__(self, f: Any) -> Any:
465 """Subtract a scalar or another Chebtech from this Chebtech.
467 Args:
468 f: A scalar or another Chebtech to subtract from this Chebtech.
470 Returns:
471 Chebtech: A new Chebtech representing the difference.
472 """
473 return self + (-f)
475 # -------
476 # roots
477 # -------
478 def roots(self, sort: bool | None = None) -> np.ndarray:
479 """Compute the roots of the Chebtech on [-1,1].
481 Uses the coefficients in the associated Chebyshev series approximation.
482 """
483 sort = sort if sort is not None else prefs.sortroots
484 rts = rootsunit(self.coeffs)
485 rts = newtonroots(self, rts)
486 # fix problems with newton for roots that are numerically very close
487 rts = np.clip(rts, -1, 1) # if newton roots are just outside [-1,1]
488 rts = rts if not sort else np.sort(rts)
489 return rts
491 # ----------
492 # calculus
493 # ----------
494 # Note that function returns 0 for an empty Chebtech object; this is
495 # consistent with numpy, which returns zero for the sum of an empty array
496 @self_empty(resultif=0.0)
497 def sum(self) -> Any:
498 """Definite integral of a Chebtech on the interval [-1,1]."""
499 if self.isconst:
500 out = 2.0 * self(0.0)
501 else:
502 ak = self.coeffs.copy()
503 ak[1::2] = 0
504 kk = np.arange(2, ak.size)
505 ii = np.append([2, 0], 2 / (1 - kk**2))
506 out = (ak * ii).sum()
507 return out
509 @self_empty()
510 def cumsum(self) -> "Chebtech":
511 """Return a Chebtech object representing the indefinite integral.
513 Computes the indefinite integral of a Chebtech on the interval [-1,1].
514 The constant term is chosen such that F(-1) = 0.
515 """
516 n = self.size
517 ak = np.append(self.coeffs, [0, 0])
518 bk = np.zeros(n + 1, dtype=self.coeffs.dtype)
519 rk = np.arange(2, n + 1)
520 bk[2:] = 0.5 * (ak[1:n] - ak[3:]) / rk
521 bk[1] = ak[0] - 0.5 * ak[2]
522 vk = np.ones(n)
523 vk[1::2] = -1
524 bk[0] = (vk * bk[1:]).sum()
525 out = self.__class__(bk, interval=self.interval)
526 return out
528 @self_empty()
529 def diff(self) -> "Chebtech":
530 """Return a Chebtech object representing the derivative.
532 Computes the derivative of a Chebtech on the interval [-1,1].
533 """
534 if self.isconst:
535 out = self.__class__(np.array([0.0]), interval=self.interval)
536 else:
537 n = self.size
538 ak = self.coeffs
539 zk = np.zeros(n - 1, dtype=self.coeffs.dtype)
540 wk = 2 * np.arange(1, n)
541 vk = wk * ak[1:]
542 zk[-1::-2] = vk[-1::-2].cumsum()
543 zk[-2::-2] = vk[-2::-2].cumsum()
544 zk[0] = 0.5 * zk[0]
545 out = self.__class__(zk, interval=self.interval)
546 return out
548 @staticmethod
549 def _chebpts(n: int) -> np.ndarray:
550 """Return n Chebyshev points of the second-kind."""
551 return chebpts2(n)
553 @staticmethod
554 def _barywts(n: int) -> np.ndarray:
555 """Barycentric weights for Chebyshev points of 2nd kind."""
556 return barywts2(n)
558 @staticmethod
559 def _vals2coeffs(vals: Any) -> np.ndarray:
560 """Map function values at Chebyshev points of 2nd kind.
562 Converts values at Chebyshev points of 2nd kind to first-kind Chebyshev polynomial coefficients.
563 """
564 return vals2coeffs2(vals)
566 @staticmethod
567 def _coeffs2vals(coeffs: Any) -> np.ndarray:
568 """Map first-kind Chebyshev polynomial coefficients.
570 Converts first-kind Chebyshev polynomial coefficients to function values at Chebyshev points of 2nd kind.
571 """
572 return coeffs2vals2(coeffs)
574 # ----------
575 # plotting
576 # ----------
577 def plot(self, ax: Any = None, **kwargs: Any) -> Any:
578 """Plot the Chebtech on the interval [-1, 1].
580 Args:
581 ax (matplotlib.axes.Axes, optional): The axes on which to plot. Defaults to None.
582 **kwargs: Additional keyword arguments to pass to the plot function.
584 Returns:
585 matplotlib.lines.Line2D: The line object created by the plot.
586 """
587 return plotfun(self, (-1, 1), ax=ax, **kwargs)
589 def plotcoeffs(self, ax: Any = None, **kwargs: Any) -> Any:
590 """Plot the absolute values of the Chebyshev coefficients.
592 Args:
593 ax (matplotlib.axes.Axes, optional): The axes on which to plot. Defaults to None.
594 **kwargs: Additional keyword arguments to pass to the plot function.
596 Returns:
597 matplotlib.lines.Line2D: The line object created by the plot.
598 """
599 ax = ax or plt.gca()
600 return plotfuncoeffs(abs(self.coeffs), ax=ax, **kwargs)