All tests pass and PEP8 compliance

src/sage/combinat/key_polynomial.py

@@ -1,11 +1,12 @@
 r"""
-Key polynomials (a.k.a Demazure characters)
+Key polynomials
 
-Key polynomials are defined by applying an operator `\pi_\sigma` to a monomial corresponding
-to an integer partition `\mu \vdash n`. For a full definition, see 
-`SymmetricFunctions.com<https://www.symmetricfunctions.com/key.htm>`_. Key polynomials are
-indexed by weak compositions, and `\sigma` is the permutation of shortest length which sorts
-the indexing composition into a partition.
+Key polynomials (a.k.a. Demazure characters) are defined by applying an
+operator `\pi_\sigma` to a monomial corresponding to an integer partition `\mu
+\vdash n`. For a full definition, see
+`SymmetricFunctions.com<https://www.symmetricfunctions.com/key.htm>`_. Key
+polynomials are indexed by weak compositions, and `\sigma` is the permutation
+of shortest length which sorts the indexing composition into a partition.
 
 .. SEEALSO::
 
@@ -28,16 +29,16 @@
     We can expand them in the standard monomial basis::
 
         sage: k([3,0,1,2]).expand()
-        z_3^2*z_2*z_0^3 + z_3^2*z_1*z_0^3 + z_3*z_2^2*z_0^3 
+        z_3^2*z_2*z_0^3 + z_3^2*z_1*z_0^3 + z_3*z_2^2*z_0^3
          + 2*z_3*z_2*z_1*z_0^3 + z_3*z_1^2*z_0^3 + z_2^2*z_1*z_0^3
          + z_2*z_1^2*z_0^3
 
     If we have a polynomial, we can express it in the key-basis::
 
         sage: z, = R.gens()
-        sage: k.from_polynomial(z[2]^2*z[1]*z[0]) ## this one fails
-        k[2, 1, 1] - k[1, 2, 1]
-        
+        sage: k.from_polynomial(z[2]^2*z[1]*z[0])
+        k[1, 1, 2] - k[1, 2, 1]
+
         sage: f = z[3]^2*z[2]*z[0]^3 + z[3]^2*z[1]*z[0]^3 + z[3]*z[2]^2*z[0]^3 + \
         ....: 2*z[3]*z[2]*z[1]*z[0]^3 + z[3]*z[1]^2*z[0]^3 + z[2]^2*z[1]*z[0]^3 + \
         ....: z[2]*z[1]^2*z[0]^3
@@ -56,19 +57,19 @@
 
         sage: s = SymmetricFunctions(QQ).schur()
         sage: k([1,2,3]).expand()
-        z_2^3*z_1^2*z_0 + z_2^3*z_1*z_0^2 + z_2^2*z_1^3*z_0 
-         + 2*z_2^2*z_1^2*z_0^2 + z_2^2*z_1*z_0^3 + z_2*z_1^3*z_0^2 
+        z_2^3*z_1^2*z_0 + z_2^3*z_1*z_0^2 + z_2^2*z_1^3*z_0
+         + 2*z_2^2*z_1^2*z_0^2 + z_2^2*z_1*z_0^3 + z_2*z_1^3*z_0^2
          + z_2*z_1^2*z_0^3
         sage: s[3,2,1].expand(3)
         x0^3*x1^2*x2 + x0^2*x1^3*x2 + x0^3*x1*x2^2 + 2*x0^2*x1^2*x2^2 + x0*x1^3*x2^2 + x0^2*x1*x2^3 + x0*x1^2*x2^3
 
     The polynomial expansions can be computed using crystals and expressed in
-    terms of the key-basis
+    terms of the key-basis::
 
-        # sage: T = crystals.Tableaux(['A',3],shape=[2,1])
-        # sage: f = T.demazure_character([3,2,1])
-        # sage: k.from_polynomial(f)
-        # k[1, 0, 0, 2]
+        sage: T = crystals.Tableaux(['A',3],shape=[2,1])
+        sage: f = T.demazure_character([3,2,1])
+        sage: k.from_polynomial(f)
+        k[1, 0, 0, 2]
 
 AUTHORS:
 
@@ -97,6 +98,7 @@
 from sage.rings.polynomial.infinite_polynomial_element import InfinitePolynomial_sparse
 from sage.rings.rational_field import QQ
 
+
 class KeyPolynomial(CombinatorialFreeModule.Element):
 
     def expand(self):
@@ -114,16 +116,12 @@ def expand(self):
 
             sage: f = k([4,3,2,1])
             sage: f.expand()
-            z_3*z_2^2*z_1^3*z+0^4
+            z_3*z_2^2*z_1^3*z_0^4
 
             sage: f = k([1,2,3])
             sage: f.expand()
-            z_2^3*z_1^2*z_0 + z_2^3*z_1*z_0^2 + z_2^2*z_1^3*z_0
-             + 2*z_2^2*z_1^2*z_0^2 + z_2^2*z_1*z_0 + z_2*z_1^3*z_0^2
-             + z_2*z_1^2*z_0^3
+            z_2^3*z_1^2*z_0 + z_2^3*z_1*z_0^2 + z_2^2*z_1^3*z_0 + 2*z_2^2*z_1^2*z_0^2 + z_2^2*z_1*z_0^3 + z_2*z_1^3*z_0^2 + z_2*z_1^2*z_0^3
         """
-        ## TODO: check if there is a better way to expand. Assaf/Quijada 2019?
-
         R = self.parent().polynomial_ring()
         out = R.zero()
         z, = R.gens()
@@ -140,12 +138,13 @@ def expand(self):
 
             # P.reduced_word() acts on indices 1, 2, ...
             # w adjusts for the zero-index.
-            w = [i-1 for i in P.reduced_word()] 
+            w = [i-1 for i in reversed(P.reduced_word())]
 
             out += c * self.parent().pi(w, monom)
 
         return out
 
+
 class KeyPolynomialRing(CombinatorialFreeModule, metaclass=ClasscallMetaclass):
 
     Element = KeyPolynomial
@@ -158,7 +157,7 @@ def __init__(self, R):
             sage: k = KeyPolynomialRing(QQ)
             sage: k == loads(dumps(k))
             True
-            sage: #TestSuite(k).run()
+            sage: TestSuite(k).run()
         """
         self._name = "Ring of key polynomials"
         self._repr_option_bracket = False
@@ -181,11 +180,12 @@ def _element_constructor_(self, alpha):
             k[9, 1, 2, 3]
 
             sage: R.<z> = InfinitePolynomialRing(QQ)
-            sage: k(z[4]^3*z[3]^2*z[2]*z[1]*z[0])
+            sage: k(z[4]*z[3]*z[2]*z[1]^2*z[0]^3)
             k[3, 2, 1, 1, 1]
 
-            sage: z0, z1, z2 = var('z0, z1, z2')
-            sage: k(z2^9*z1^5*z0^4)
+            sage: z0, z1, z2 = var('z_0, z_1, z_2')
+            sage: f = z2^4*z1^5*z0^9
+            sage: k(R(f))
             k[9, 5, 4]
         """
         if isinstance(alpha, (list, IntegerVector, Composition)):
@@ -195,6 +195,32 @@ def _element_constructor_(self, alpha):
             return self.from_polynomial(alpha)
         raise ValueError("alpha can not be interpreted as a weak composition or polynomial")
 
+    def one(self):
+        r"""
+        Return the multiplicative identity
+
+        EXAMPLES::
+
+            sage: from sage.combinat.key_polynomial import KeyPolynomialRing
+            sage: k = KeyPolynomialRing(QQ)
+            sage: k([4,2,5,3]) * k.one()
+            k[4, 2, 5, 3]
+        """
+        return self([])
+
+    def one_basis(self):
+        r"""
+        Return the basis element indexing the identity.
+
+        EXAMPLES::
+
+            sage: from sage.combinat.key_polynomial import KeyPolynomialRing
+            sage: k = KeyPolynomialRing(QQ)
+            sage: k.one_basis()
+            []
+        """
+        return _IntegerVectors([])
+
     def polynomial_ring(self):
         r"""
         Return the polynomial ring associated to ``self``.
@@ -204,7 +230,7 @@ def polynomial_ring(self):
             sage: from sage.combinat.key_polynomial import KeyPolynomialRing
             sage: k = KeyPolynomialRing(QQ)
             sage: k.polynomial_ring()
-            Infinite polynomial ring in z over Rational field
+            Infinite polynomial ring in z over Rational Field
         """
         return InfinitePolynomialRing(self.base_ring(), 'z')
 
@@ -221,23 +247,23 @@ def from_polynomial(self, f):
             sage: K.from_polynomial(p)
             k[4, 1, 2, 1]
 
-            sage: all(K(c) == K.from_polynomial(c).expand())
-            ....:     for c in Compositions(5))
+            sage: all(K(c) == K.from_polynomial(K(c).expand()) for c in Compositions(5))
             True
 
             sage: T = crystals.Tableaux(['A', 4], shape=[4,2,1,1])
             sage: K.from_polynomial(T.demazure_character([2]))
             k[4, 1, 2, 1]
 
         """
-        from sage.symbolic.ring import SymbolicRing
-        if f.parent() is SymbolicRing:  # to accept Demazure characters from crystals
-            f.substitute(list(d == var(f'z_{i}') for i,d in enumerate(f.variables())))
-
-        if not f in self.polynomial_ring():  #todo: change this to accept coerce
-            raise ValueError(f"f must be an element of {self.polynomial_ring()}")
+        if f not in self.polynomial_ring():
+            try:  # to accept elements of SymbolicRing
+                from sage.calculus.var import var
+                f = f.substitute(list(d == var(f'z_{i}')
+                                 for i, d in enumerate(f.variables())))
+                f = self.polynomial_ring()(f)
+            except AttributeError:
+                raise ValueError(f"f must be an element of {self.polynomial_ring()}")
 
-
         out = self.zero()
         counter = 0
 
@@ -255,18 +281,28 @@ def from_polynomial(self, f):
 
     def pi(self, w, f):
         r"""
+        Apply the operator `\pi_w` to the polynomial `f`.
+        ``w`` may be either a single index or a list of
+        indices of simple transpositions.
+
         EXAMPLES::
+
             sage: from sage.combinat.key_polynomial import KeyPolynomialRing
             sage: K = KeyPolynomialRing(QQ)
             sage: R = K.polynomial_ring()
             sage: z, = R.gens()
             sage: w = [2, 3]
             sage: K.pi(w, z[1]^4*z[2]^2*z[3]*z[4])
+            z_4*z_3^2*z_2*z_1^4 + z_4*z_3*z_2^2*z_1^4
+
             sage: w = [3, 2]
             sage: K.pi(w, z[1]^4*z[2]^2*z[3]*z[4])
+            z_4^2*z_3*z_2*z_1^4 + z_4*z_3^2*z_2*z_1^4 + z_4*z_3*z_2^2*z_1^4
 
-            sage: w = [2, 3, 2, 1]
+            sage: w = [3, 2, 1]
             sage: K.pi(w, z[1]^2*z[2])
+            z_4^2*z_1 + z_4*z_3*z_1 + z_4*z_2*z_1 + z_4*z_1^2 + z_3^2*z_1
+             + z_3*z_2*z_1 + z_3*z_1^2 + z_2^2*z_1 + z_2*z_1^2
         """
         if not hasattr(w, "__iter__"):  # this allows us to pass i instead of a word
             return self._pi_i(w, f)
@@ -276,6 +312,8 @@ def pi(self, w, f):
 
     def _pi_i(self, i, f):
         r"""
+        Apply `\pi_i` for a single simple transposition `s_i = (i, i+1)`.
+
         EXAMPLES::
 
             sage: from sage.combinat.key_polynomial import KeyPolynomialRing
@@ -294,6 +332,7 @@ def _pi_i(self, i, f):
 
     def _divided_difference(self, i, f):
         r"""
+        Apply the `i`th divided difference operator to the polynomial `f`.
 
         EXAMPLES::
 
@@ -304,7 +343,7 @@ def _divided_difference(self, i, f):
             sage: z, = R.gens()
             sage: f = z[1]*z[2]^3 + z[1]*z[2]*z[3]
             sage: K._divided_difference(2, f)
-            z_1*z_2^2 - z_1*z_3^2
+            z_3^2*z_1 + z_3*z_2*z_1 + z_2^2*z_1
 
         """
         out = self.polynomial_ring().zero()
@@ -317,8 +356,6 @@ def _divided_difference_on_monomial(self, i, m):
         r"""
         Apply the `i`th divided difference operator to `m`.
 
-        NOTE:: `i` can be 0. 
-
         EXAMPLES::
 
             sage: from sage.combinat.key_polynomial import KeyPolynomialRing
@@ -334,10 +371,6 @@ def _divided_difference_on_monomial(self, i, m):
             z_2*z_1
             sage: K._divided_difference_on_monomial(3, z[1]*z[2]*z[4])
             -z_2*z_1
-
-        TODO::
-            add examples where there are multiple factors of xi-x(i+1)
-
         """
         # The polynomial ring and generators
         R = self.polynomial_ring()
@@ -358,15 +391,15 @@ def _divided_difference_on_monomial(self, i, m):
         si_exp = exp[:i] + [exp[i+1], exp[i]] + exp[i+2:]
 
         # Create the corresponding list of variables in the monomial
-        terms = list(x[i]**j for i,j in enumerate(si_exp) if j)
+        terms = list(x[i]**j for i, j in enumerate(si_exp) if j)
 
         # Create the monomial from the list of variables
         si_m = R.prod(terms)
 
         # Create the numerator of the divided difference operator
         f = si_m - m
 
-        # Division using the / operator is not implemented for 
+        # Division using the / operator is not implemented for
         # InfinitePolynomialRing, so we want to remove the factor of
         # x[i+1] - x[i]. If x[i+1] - x[i] is not a factor, it is because
         # the numerator is zero.
@@ -377,7 +410,7 @@ def _divided_difference_on_monomial(self, i, m):
             return R.zero()
 
         factors_dict[x[i + 1] - x[i]] = factors_dict[x[i + 1] - x[i]] - 1
-        return R.prod(k**v for k,v in factors_dict.items()) * factors.unit()
+        return R.prod(k**v for k, v in factors_dict.items()) * factors.unit()
 
     def product_on_basis(self, a, b):
         r"""
@@ -395,6 +428,7 @@ def product_on_basis(self, a, b):
 
         return self.from_polynomial(p)
 
+
 def _sorting_permutation(alpha):
     r"""
     Get the sorting permutation for a composition ``alpha``.
@@ -422,8 +456,9 @@ def _sorting_permutation(alpha):
     # this yields the *shortest* permutation because
     # it does not introduce any unnecessary inversions,
     # and the number of inversions is the length of the permutation.
-    key = lambda x: (-x[1], x[0])
-    p = [i for i,j in sorted(enumerate(alpha,1), key=key)]
+    def key(x): return (-x[1], x[0])
+    p = [i for i, j in sorted(enumerate(alpha, 1), key=key)]
     return Permutation(p).inverse()
 
-_IntegerVectors = IntegerVectors()
+
+_IntegerVectors = IntegerVectors()