]> gitweb.michael.orlitzky.com - sage.d.git/blobdiff - mjo/symbol_sequence.py
COPYING,LICENSE: add (AGPL-3.0+)
[sage.d.git] / mjo / symbol_sequence.py
index 43844271876b2b935857dfd3f1d679f57c05d4ba..2a104aed958d877103e1962f1ad212e2ab07c981 100644 (file)
@@ -1,13 +1,14 @@
 from sage.all import *
 
 class SymbolSequence:
-    """
-    A callable object which imitates a function from ZZ^n to a
-    sequence with n subscripts.
+    r"""
+    An iterable object which acts like a sequence of symbolic
+    expressions (variables).
 
     INPUT:
 
-      - ``name`` -- The sequence name.
+      - ``name`` -- The sequence name. For example, if you name the
+        sequence `x`, the variables will be called `x_0`, `x_1`,...
 
       - ``latex_name`` -- An optional latex expression (string) to
         use instead of `name` when converting the symbols to latex.
@@ -17,75 +18,70 @@ class SymbolSequence:
 
     OUTPUT:
 
-    A callable object returning symbolic expressions.
+    An iterable object containing symbolic expressions.
+
+    SETUP::
+
+        sage: from mjo.symbol_sequence import SymbolSequence
 
     EXAMPLES:
 
-    Create coefficients for polynomials of arbitrary degree::
+    The simplest use case::
 
         sage: a = SymbolSequence('a')
-        sage: p = sum([ a(i)*x^i for i in range(0,5)])
+        sage: a[0]
+        a_0
+        sage: a[1]
+        a_1
+
+    Create polynomials with symbolic coefficients of arbitrary
+    degree::
+
+        sage: a = SymbolSequence('a')
+        sage: p = sum( a[i]*x^i for i in range(5) )
         sage: p
-        a4*x^4 + a3*x^3 + a2*x^2 + a1*x + a0
+        a_4*x^4 + a_3*x^3 + a_2*x^2 + a_1*x + a_0
 
     Using a different latex name since 'lambda' is reserved::
 
-        sage: l = SymbolSequence('l', '\lambda')
-        sage: l(0)
-        l0
-        sage: latex(l(0))
+        sage: l = SymbolSequence('l', r'\lambda')
+        sage: l[0]
+        l_0
+        sage: latex(l[0])
         \lambda_{0}
 
     Using multiple indices::
 
         sage: a = SymbolSequence('a')
-        sage: a(0,1,2)
-        a012
-        sage: latex(a(0,1,2))
+        sage: a[0,1,2]
+        a_0_1_2
+        sage: latex(a[0,1,2])
         a_{0}_{1}_{2}
-        sage: [ a(i,j) for i in range(0,2) for j in range(0,2) ]
-        [a00, a01, a10, a11]
-
-    If no index is given, an unsubscripted symbol is returned::
+        sage: [ a[i,j] for i in range(2) for j in range(2) ]
+        [a_0_0, a_0_1, a_1_0, a_1_1]
 
-        sage: a = SymbolSequence('a')
-        sage: a()
-        a
-
-    You can pass slice objects instead of integers to obtain a list of
+    You can pass slices instead of integers to obtain a list of
     symbols::
 
         sage: a = SymbolSequence('a')
-        sage: a(slice(5,7))
-        [a5, a6]
+        sage: a[5:7]
+        [a_5, a_6]
 
     This even works for the second, third, etc. indices::
 
         sage: a = SymbolSequence('a')
-        sage: a(slice(0,2), slice(0,2))
-        [a00, a01, a10, a11]
-
-    You can also index with the list index operator::
-
-        sage: a = SymbolSequence('a')
-        sage: a[1]
-        a1
-
-    This allows you to retrieve one-dimensional slices easily::
-
-        sage: a = SymbolSequence('a')
-        sage: a[0:5]
-        [a0, a1, a2, a3, a4]
+        sage: a[0:2, 0:2]
+        [a_0_0, a_0_1, a_1_0, a_1_1]
 
     TESTS:
 
     We shouldn't overwrite variables in the global namespace::
 
         sage: a = SymbolSequence('a')
-        sage: a0 = 4
-        sage: a(0)
-        a0
-        sage: a0
+        sage: a_0 = 4
+        sage: a[0]
+        a_0
+        sage: a_0
         4
 
     The symbol at a given index should always be the same, even when
@@ -93,30 +89,26 @@ class SymbolSequence:
     representation and compare because the output is unpredictable::
 
         sage: a = SymbolSequence()
-        sage: a0str = str(a(0))
-        sage: str(a(0)) == a0str
+        sage: a0str = str(a[0])
+        sage: str(a[0]) == a0str
         True
 
     Slices and single indices work when combined::
 
         sage: a = SymbolSequence('a')
-        sage: a(3, slice(0,2))
-        [a30, a31]
-        sage: a(slice(0,2), 3)
-        [a03, a13]
+        sage: a[3, 0:2]
+        [a_3_0, a_3_1]
+        sage: a[0:2, 3]
+        [a_0_3, a_1_3]
 
     """
 
     def __init__(self, name=None, latex_name=None, domain=None):
         # We store a dict of already-created symbols so that we don't
         # recreate a symbol which already exists. This is especially
-        # helpful when using unnamed variables, if you want e.g. a(0)
+        # helpful when using unnamed variables, if you want e.g. a[0]
         # to return the same variable each time.
-        #
-        # The entry corresponding to None is the un-subscripted symbol
-        # with our name.
-        unsubscripted = SR.symbol(name, latex_name, domain)
-        self._symbols = { None: unsubscripted }
+        self._symbols = {}
 
         self._name = name
         self._latex_name = latex_name
@@ -124,16 +116,33 @@ class SymbolSequence:
 
 
     def _create_symbol_(self, subscript):
-        if self._name is None:
-            # Allow creating unnamed symbols, for consistency with
-            # SR.symbol().
-            name = None
-        else:
-            name = '%s%d' % (self._name, subscript)
+        """
+        Return a symbol with the given subscript. Creates the
+        appropriate name and latex_name before delegating to
+        SR.symbol().
 
-        if self._latex_name is None:
-            latex_name = None
-        else:
+        SETUP::
+
+            sage: from mjo.symbol_sequence import SymbolSequence
+
+        EXAMPLES::
+
+            sage: a = SymbolSequence('a', 'alpha', 'real')
+            sage: a_1 = a._create_symbol_(1)
+            sage: a_1
+            a_1
+            sage: latex(a_1)
+            alpha_{1}
+
+        """
+        # Allow creating unnamed symbols, for consistency with
+        # SR.symbol().
+        name = None
+        if self._name is not None:
+            name = '%s_%d' % (self._name, subscript)
+
+        latex_name = None
+        if self._latex_name is not None:
             latex_name = r'%s_{%d}' % (self._latex_name, subscript)
 
         return SR.symbol(name, latex_name, self._domain)
@@ -142,78 +151,192 @@ class SymbolSequence:
     def _flatten_list_(self, l):
         """
         Recursively flatten the given list, allowing for some elements
-        to be non-iterable.
+        to be non-iterable. This is slow, but also works, which is
+        more than can be said about some of the snappier solutions of
+        lore.
+
+        SETUP::
+
+            sage: from mjo.symbol_sequence import SymbolSequence
+
+        EXAMPLES::
+
+            sage: a = SymbolSequence('a')
+            sage: a._flatten_list_([1,2,3])
+            [1, 2, 3]
+            sage: a._flatten_list_([1,[2,3]])
+            [1, 2, 3]
+            sage: a._flatten_list_([1,[2,[3]]])
+            [1, 2, 3]
+            sage: a._flatten_list_([[[[[1,[2,[3]]]]]]])
+            [1, 2, 3]
+
         """
         result = []
 
         for item in l:
-            if isinstance(item, list):
+            try:
+                item = iter(item)
                 result += self._flatten_list_(item)
-            else:
+            except TypeError:
                 result += [item]
 
         return result
 
 
     def __getitem__(self, key):
-        if isinstance(key, slice):
-            # We were given a slice. Clean up some of its properties
-            # first. The start/step are default for lists. We make
-            # copies of these because they're read-only.
-            (start, step) = (key.start, key.step)
-            if start is None:
-                start = 0
-            if key.stop is None:
-                # Would otherwise loop forever since our "length" is
-                # undefined.
-                raise ValueError('You must supply an terminal index')
-            if step is None:
-               step = 1
-
-            # If the user asks for a slice, we'll be returning a list
-            # of symbols.
-            return [ self(idx) for idx in range(start, key.stop, step) ]
-        else:
-            return self(key)
+        """
+        This handles individual integer arguments, slices, and
+        tuples. It just hands off the real work to
+        self._subscript_foo_().
 
+        SETUP::
 
-    def __call__(self, *args):
-        args = list(args)
+            sage: from mjo.symbol_sequence import SymbolSequence
 
-        if len(args) == 0:
-            return self._symbols[None]
+        EXAMPLES:
 
-        # This is safe after the len == 0 test.
-        key = args[0]
-        args = args[1:] # Peel off the first arg, which we've called 'key'
+        An integer argument::
+
+            sage: a = SymbolSequence('a')
+            sage: a.__getitem__(1)
+            a_1
+
+        A tuple argument::
+
+            sage: a = SymbolSequence('a')
+            sage: a.__getitem__((1,2))
+            a_1_2
+
+        A slice argument::
+
+            sage: a = SymbolSequence('a')
+            sage: a.__getitem__(slice(1,4))
+            [a_1, a_2, a_3]
+
+        """
+        if isinstance(key, tuple):
+            return self._subscript_tuple_(key)
 
         if isinstance(key, slice):
-            if len(args) == 0:
-                return self[key]
-            else:
-                v = self[key]
-                ss = [ SymbolSequence(w._repr_(), w._latex_(), self._domain)
-                       for w in v ]
-
-                # This might be nested...
-                maybe_nested_list = [ s(*args) for s in ss ]
-                return self._flatten_list_(maybe_nested_list)
-
-        if key < 0:
+            return self._subscript_slice_(key)
+
+        # This is the most common case so it would make sense to have
+        # this test first. But there are too many different "integer"
+        # classes that you have to check for.
+        return self._subscript_integer_(key)
+
+
+    def _subscript_integer_(self, n):
+        """
+        The subscript is a single integer, or something that acts like
+        one.
+
+        SETUP::
+
+            sage: from mjo.symbol_sequence import SymbolSequence
+
+        EXAMPLES::
+
+            sage: a = SymbolSequence('a')
+            sage: a._subscript_integer_(123)
+            a_123
+
+        """
+        if n < 0:
             # Cowardly refuse to create a variable named "a-1".
             raise IndexError('Indices must be nonnegative')
 
+        try:
+            return self._symbols[n]
+        except KeyError:
+            self._symbols[n] = self._create_symbol_(n)
+            return self._symbols[n]
+
+
+    def _subscript_slice_(self, s):
+        """
+        We were given a slice. Clean up some of its properties
+        first. The start/step are default for lists. We make
+        copies of these because they're read-only.
+
+        SETUP::
+
+            sage: from mjo.symbol_sequence import SymbolSequence
+
+        EXAMPLES::
+
+            sage: a = SymbolSequence('a')
+            sage: a._subscript_slice_(slice(1,3))
+            [a_1, a_2]
+
+        """
+        (start, step) = (s.start, s.step)
+        if start is None:
+            start = 0
+        if s.stop is None:
+            # Would otherwise loop forever since our "length" is
+            # undefined.
+            raise ValueError('You must supply an terminal index')
+        if step is None:
+            step = 1
+
+        # If the user asks for a slice, we'll be returning a list
+        # of symbols.
+        return [ self._subscript_integer_(idx)
+                 for idx in range(start, s.stop, step) ]
+
+
+
+    def _subscript_tuple_(self, args):
+        """
+        When we have more than one level of subscripts, we pick off
+        the first one and generate the rest recursively.
+
+        SETUP::
+
+            sage: from mjo.symbol_sequence import SymbolSequence
+
+        EXAMPLES:
+
+        A simple two-tuple::
+
+            sage: a = SymbolSequence('a')
+            sage: a._subscript_tuple_((1,8))
+            a_1_8
+
+        Nested tuples::
+
+            sage: a._subscript_tuple_(( (1,2), (3,(4,5,6)) ))
+            a_1_2_3_4_5_6
+
+        """
+
+        # We never call this method without an argument.
+        key = args[0]
+        args = args[1:] # Peel off the first arg, which we've called 'key'
+
+        # We don't know the type of 'key', but __getitem__ will figure
+        # it out and dispatch properly.
+        v = self[key]
         if len(args) == 0:
-            # Base case, create a symbol and return it.
-            try:
-                return self._symbols[key]
-            except KeyError:
-                self._symbols[key] = self._create_symbol_(key)
-                return self._symbols[key]
+            # There was only one element left in the tuple.
+            return v
+
+        # At this point, we know we were given at least a two-tuple.
+        # The symbols corresponding to the first entry are already
+        # computed, in 'v'. Here we recursively compute the symbols
+        # corresponding to the second coordinate, with the first
+        # coordinate(s) fixed.
+        if isinstance(key, slice):
+            ss = ( SymbolSequence(w._repr_(), w._latex_(), self._domain)
+                   for w in v )
+
+            # This might be nested...
+            maybe_nested_list = ( s._subscript_tuple_(args) for s in ss )
+            return self._flatten_list_(maybe_nested_list)
+
         else:
-            # If we're given more than one index, we want to create the
-            # subsequences recursively. For example, if we're asked for
-            # x(1,2), this is really SymbolSequence('x1')(2).
-            v = self(key) # x(1) -> x1
+            # If it's not a slice, it's an integer.
             ss = SymbolSequence(v._repr_(), v._latex_(), self._domain)
-            return ss(*args)
+            return ss._subscript_tuple_(args)