ENH: Add support for symbol to polynomial package (#16154)

Adds a symbol attribute to the polynomials from the np.polynomial package to allow the user to control/modify the symbol used to represent the independent variable for a polynomial expression. This attribute corresponds to the variable attribute of the poly1d class from the old np.lib.polynomial module.

Marked as draft for now as it depends on #15666 - all _str* and _repr* methods of ABCPolyBase and derived classes would need to be modified (and tested) to support this change.

Co-authored-by: Warren Weckesser <warren.weckesser@gmail.com>

1 files changed, 79 insertions, 36 deletions

diff --git a/numpy/polynomial/_polybase.py b/numpy/polynomial/_polybase.py
index 155d72805..6382732dc 100644
--- a/numpy/polynomial/_polybase.py
+++ b/numpy/polynomial/_polybase.py
@@ -37,6 +37,12 @@ class ABCPolyBase(abc.ABC):
     window : (2,) array_like, optional
         Window, see domain for its use. The default value is the
         derived class window.
+    symbol : str, optional
+        Symbol used to represent the independent variable in string 
+        representations of the polynomial expression, e.g. for printing.
+        The symbol must be a valid Python identifier. Default value is 'x'.
+
+        .. versionadded:: 1.24
 
     Attributes
     ----------
@@ -46,6 +52,8 @@ class ABCPolyBase(abc.ABC):
         Domain that is mapped to window.
     window : (2,) ndarray
         Window that domain is mapped to.
+    symbol : str
+        Symbol representing the independent variable.
 
     Class Attributes
     ----------------
@@ -100,6 +108,10 @@ class ABCPolyBase(abc.ABC):
     _use_unicode = not os.name == 'nt'
 
     @property
+    def symbol(self):
+        return self._symbol
+
+    @property
     @abc.abstractmethod
     def domain(self):
         pass
@@ -284,10 +296,12 @@ class ABCPolyBase(abc.ABC):
                 raise TypeError("Domains differ")
             elif not np.all(self.window == other.window):
                 raise TypeError("Windows differ")
+            elif self.symbol != other.symbol:
+                raise ValueError("Polynomial symbols differ")
             return other.coef
         return other
 
-    def __init__(self, coef, domain=None, window=None):
+    def __init__(self, coef, domain=None, window=None, symbol='x'):
         [coef] = pu.as_series([coef], trim=False)
         self.coef = coef
 
@@ -303,12 +317,27 @@ class ABCPolyBase(abc.ABC):
                 raise ValueError("Window has wrong number of elements.")
             self.window = window
 
+        # Validation for symbol
+        try:
+            if not symbol.isidentifier():
+                raise ValueError(
+                    "Symbol string must be a valid Python identifier"
+                )
+        # If a user passes in something other than a string, the above
+        # results in an AttributeError. Catch this and raise a more
+        # informative exception
+        except AttributeError:
+            raise TypeError("Symbol must be a non-empty string")
+
+        self._symbol = symbol
+
     def __repr__(self):
         coef = repr(self.coef)[6:-1]
         domain = repr(self.domain)[6:-1]
         window = repr(self.window)[6:-1]
         name = self.__class__.__name__
-        return f"{name}({coef}, domain={domain}, window={window})"
+        return (f"{name}({coef}, domain={domain}, window={window}, "
+                f"symbol='{self.symbol}')")
 
     def __format__(self, fmt_str):
         if fmt_str == '':
@@ -353,7 +382,7 @@ class ABCPolyBase(abc.ABC):
             except TypeError:
                 next_term = f"+ {coef}"
             # Polynomial term
-            next_term += term_method(power, "x")
+            next_term += term_method(power, self.symbol)
             # Length of the current line with next term added
             line_len = len(out.split('\n')[-1]) + len(next_term)
             # If not the last term in the polynomial, it will be two
@@ -412,18 +441,18 @@ class ABCPolyBase(abc.ABC):
         # get the scaled argument string to the basis functions
         off, scale = self.mapparms()
         if off == 0 and scale == 1:
-            term = 'x'
+            term = self.symbol
             needs_parens = False
         elif scale == 1:
-            term = f"{self._repr_latex_scalar(off)} + x"
+            term = f"{self._repr_latex_scalar(off)} + {self.symbol}"
             needs_parens = True
         elif off == 0:
-            term = f"{self._repr_latex_scalar(scale)}x"
+            term = f"{self._repr_latex_scalar(scale)}{self.symbol}"
             needs_parens = True
         else:
             term = (
                 f"{self._repr_latex_scalar(off)} + "
-                f"{self._repr_latex_scalar(scale)}x"
+                f"{self._repr_latex_scalar(scale)}{self.symbol}"
             )
             needs_parens = True
 
@@ -459,7 +488,7 @@ class ABCPolyBase(abc.ABC):
             # in case somehow there are no coefficients at all
             body = '0'
 
-        return rf"$x \mapsto {body}$"
+        return rf"${self.symbol} \mapsto {body}$"
 
 
 
@@ -470,6 +499,7 @@ class ABCPolyBase(abc.ABC):
         ret['coef'] = self.coef.copy()
         ret['domain'] = self.domain.copy()
         ret['window'] = self.window.copy()
+        ret['symbol'] = self.symbol.copy()
         return ret
 
     def __setstate__(self, dict):
@@ -491,7 +521,9 @@ class ABCPolyBase(abc.ABC):
     # Numeric properties.
 
     def __neg__(self):
-        return self.__class__(-self.coef, self.domain, self.window)
+        return self.__class__(
+            -self.coef, self.domain, self.window, self.symbol
+        )
 
     def __pos__(self):
         return self
@@ -502,7 +534,7 @@ class ABCPolyBase(abc.ABC):
             coef = self._add(self.coef, othercoef)
         except Exception:
             return NotImplemented
-        return self.__class__(coef, self.domain, self.window)
+        return self.__class__(coef, self.domain, self.window, self.symbol)
 
     def __sub__(self, other):
         othercoef = self._get_coefficients(other)
@@ -510,7 +542,7 @@ class ABCPolyBase(abc.ABC):
             coef = self._sub(self.coef, othercoef)
         except Exception:
             return NotImplemented
-        return self.__class__(coef, self.domain, self.window)
+        return self.__class__(coef, self.domain, self.window, self.symbol)
 
     def __mul__(self, other):
         othercoef = self._get_coefficients(other)
@@ -518,7 +550,7 @@ class ABCPolyBase(abc.ABC):
             coef = self._mul(self.coef, othercoef)
         except Exception:
             return NotImplemented
-        return self.__class__(coef, self.domain, self.window)
+        return self.__class__(coef, self.domain, self.window, self.symbol)
 
     def __truediv__(self, other):
         # there is no true divide if the rhs is not a Number, although it
@@ -551,13 +583,13 @@ class ABCPolyBase(abc.ABC):
             raise
         except Exception:
             return NotImplemented
-        quo = self.__class__(quo, self.domain, self.window)
-        rem = self.__class__(rem, self.domain, self.window)
+        quo = self.__class__(quo, self.domain, self.window, self.symbol)
+        rem = self.__class__(rem, self.domain, self.window, self.symbol)
         return quo, rem
 
     def __pow__(self, other):
         coef = self._pow(self.coef, other, maxpower=self.maxpower)
-        res = self.__class__(coef, self.domain, self.window)
+        res = self.__class__(coef, self.domain, self.window, self.symbol)
         return res
 
     def __radd__(self, other):
@@ -565,21 +597,21 @@ class ABCPolyBase(abc.ABC):
             coef = self._add(other, self.coef)
         except Exception:
             return NotImplemented
-        return self.__class__(coef, self.domain, self.window)
+        return self.__class__(coef, self.domain, self.window, self.symbol)
 
     def __rsub__(self, other):
         try:
             coef = self._sub(other, self.coef)
         except Exception:
             return NotImplemented
-        return self.__class__(coef, self.domain, self.window)
+        return self.__class__(coef, self.domain, self.window, self.symbol)
 
     def __rmul__(self, other):
         try:
             coef = self._mul(other, self.coef)
         except Exception:
             return NotImplemented
-        return self.__class__(coef, self.domain, self.window)
+        return self.__class__(coef, self.domain, self.window, self.symbol)
 
     def __rdiv__(self, other):
         # set to __floordiv__ /.
@@ -609,8 +641,8 @@ class ABCPolyBase(abc.ABC):
             raise
         except Exception:
             return NotImplemented
-        quo = self.__class__(quo, self.domain, self.window)
-        rem = self.__class__(rem, self.domain, self.window)
+        quo = self.__class__(quo, self.domain, self.window, self.symbol)
+        rem = self.__class__(rem, self.domain, self.window, self.symbol)
         return quo, rem
 
     def __eq__(self, other):
@@ -618,7 +650,8 @@ class ABCPolyBase(abc.ABC):
                np.all(self.domain == other.domain) and
                np.all(self.window == other.window) and
                (self.coef.shape == other.coef.shape) and
-               np.all(self.coef == other.coef))
+               np.all(self.coef == other.coef) and
+               (self.symbol == other.symbol))
         return res
 
     def __ne__(self, other):
@@ -637,7 +670,7 @@ class ABCPolyBase(abc.ABC):
             Copy of self.
 
         """
-        return self.__class__(self.coef, self.domain, self.window)
+        return self.__class__(self.coef, self.domain, self.window, self.symbol)
 
     def degree(self):
         """The degree of the series.
@@ -698,7 +731,7 @@ class ABCPolyBase(abc.ABC):
 
         """
         coef = pu.trimcoef(self.coef, tol)
-        return self.__class__(coef, self.domain, self.window)
+        return self.__class__(coef, self.domain, self.window, self.symbol)
 
     def truncate(self, size):
         """Truncate series to length `size`.
@@ -727,7 +760,7 @@ class ABCPolyBase(abc.ABC):
             coef = self.coef
         else:
             coef = self.coef[:isize]
-        return self.__class__(coef, self.domain, self.window)
+        return self.__class__(coef, self.domain, self.window, self.symbol)
 
     def convert(self, domain=None, kind=None, window=None):
         """Convert series to a different kind and/or domain and/or window.
@@ -764,7 +797,7 @@ class ABCPolyBase(abc.ABC):
             domain = kind.domain
         if window is None:
             window = kind.window
-        return self(kind.identity(domain, window=window))
+        return self(kind.identity(domain, window=window, symbol=self.symbol))
 
     def mapparms(self):
         """Return the mapping parameters.
@@ -826,7 +859,7 @@ class ABCPolyBase(abc.ABC):
         else:
             lbnd = off + scl*lbnd
         coef = self._int(self.coef, m, k, lbnd, 1./scl)
-        return self.__class__(coef, self.domain, self.window)
+        return self.__class__(coef, self.domain, self.window, self.symbol)
 
     def deriv(self, m=1):
         """Differentiate.
@@ -848,7 +881,7 @@ class ABCPolyBase(abc.ABC):
         """
         off, scl = self.mapparms()
         coef = self._der(self.coef, m, scl)
-        return self.__class__(coef, self.domain, self.window)
+        return self.__class__(coef, self.domain, self.window, self.symbol)
 
     def roots(self):
         """Return the roots of the series polynomial.
@@ -899,7 +932,7 @@ class ABCPolyBase(abc.ABC):
 
     @classmethod
     def fit(cls, x, y, deg, domain=None, rcond=None, full=False, w=None,
-        window=None):
+        window=None, symbol='x'):
         """Least squares fit to data.
 
         Return a series instance that is the least squares fit to the data
@@ -948,6 +981,8 @@ class ABCPolyBase(abc.ABC):
             value is the default class domain
 
             .. versionadded:: 1.6.0
+        symbol : str, optional
+            Symbol representing the independent variable. Default is 'x'.
 
         Returns
         -------
@@ -980,13 +1015,15 @@ class ABCPolyBase(abc.ABC):
         res = cls._fit(xnew, y, deg, w=w, rcond=rcond, full=full)
         if full:
             [coef, status] = res
-            return cls(coef, domain=domain, window=window), status
+            return (
+                cls(coef, domain=domain, window=window, symbol=symbol), status
+            )
         else:
             coef = res
-            return cls(coef, domain=domain, window=window)
+            return cls(coef, domain=domain, window=window, symbol=symbol)
 
     @classmethod
-    def fromroots(cls, roots, domain=[], window=None):
+    def fromroots(cls, roots, domain=[], window=None, symbol='x'):
         """Return series instance that has the specified roots.
 
         Returns a series representing the product
@@ -1004,6 +1041,8 @@ class ABCPolyBase(abc.ABC):
         window : {None, array_like}, optional
             Window for the returned series. If None the class window is
             used. The default is None.
+        symbol : str, optional
+            Symbol representing the independent variable. Default is 'x'.
 
         Returns
         -------
@@ -1024,10 +1063,10 @@ class ABCPolyBase(abc.ABC):
         off, scl = pu.mapparms(domain, window)
         rnew = off + scl*roots
         coef = cls._fromroots(rnew) / scl**deg
-        return cls(coef, domain=domain, window=window)
+        return cls(coef, domain=domain, window=window, symbol=symbol)
 
     @classmethod
-    def identity(cls, domain=None, window=None):
+    def identity(cls, domain=None, window=None, symbol='x'):
         """Identity function.
 
         If ``p`` is the returned series, then ``p(x) == x`` for all
@@ -1044,6 +1083,8 @@ class ABCPolyBase(abc.ABC):
             ``[beg, end]``, where ``beg`` and ``end`` are the endpoints of
             the window. If None is given then the class window is used. The
             default is None.
+        symbol : str, optional
+            Symbol representing the independent variable. Default is 'x'.
 
         Returns
         -------
@@ -1057,10 +1098,10 @@ class ABCPolyBase(abc.ABC):
             window = cls.window
         off, scl = pu.mapparms(window, domain)
         coef = cls._line(off, scl)
-        return cls(coef, domain, window)
+        return cls(coef, domain, window, symbol)
 
     @classmethod
-    def basis(cls, deg, domain=None, window=None):
+    def basis(cls, deg, domain=None, window=None, symbol='x'):
         """Series basis polynomial of degree `deg`.
 
         Returns the series representing the basis polynomial of degree `deg`.
@@ -1080,6 +1121,8 @@ class ABCPolyBase(abc.ABC):
             ``[beg, end]``, where ``beg`` and ``end`` are the endpoints of
             the window. If None is given then the class window is used. The
             default is None.
+        symbol : str, optional
+            Symbol representing the independent variable. Default is 'x'.
 
         Returns
         -------
@@ -1096,7 +1139,7 @@ class ABCPolyBase(abc.ABC):
 
         if ideg != deg or ideg < 0:
             raise ValueError("deg must be non-negative integer")
-        return cls([0]*ideg + [1], domain, window)
+        return cls([0]*ideg + [1], domain, window, symbol)
 
     @classmethod
     def cast(cls, series, domain=None, window=None):