mjo/symbol_sequence.py

   1 from sage.all import *
   2
   3 class SymbolSequence:
   4     """
   5     A callable object which imitates a function from ZZ^n to a
   6     sequence with n subscripts.
   7
   8     INPUT:
   9
  10       - ``name`` -- The sequence name.
  11
  12       - ``latex_name`` -- An optional latex expression (string) to
  13         use instead of `name` when converting the symbols to latex.
  14
  15       - ``domain`` -- A string representing the domain of the symbol,
  16         either 'real', 'complex', or 'positive'.
  17
  18     OUTPUT:
  19
  20     A callable object returning symbolic expressions.
  21
  22     EXAMPLES:
  23
  24     Create coefficients for polynomials of arbitrary degree::
  25
  26         sage: a = SymbolSequence('a')
  27         sage: p = sum([ a(i)*x^i for i in range(0,5)])
  28         sage: p
  29         a4*x^4 + a3*x^3 + a2*x^2 + a1*x + a0
  30
  31     Using a different latex name since 'lambda' is reserved::
  32
  33         sage: l = SymbolSequence('l', '\lambda')
  34         sage: l(0)
  35         l0
  36         sage: latex(l(0))
  37         \lambda_{0}
  38
  39     Using multiple indices::
  40
  41         sage: a = SymbolSequence('a')
  42         sage: a(0,1,2)
  43         a012
  44         sage: latex(a(0,1,2))
  45         a_{0}_{1}_{2}
  46         sage: [ a(i,j) for i in range(0,2) for j in range(0,2) ]
  47         [a00, a01, a10, a11]
  48
  49     If no index is given, an unsubscripted symbol is returned::
  50
  51         sage: a = SymbolSequence('a')
  52         sage: a()
  53         a
  54
  55     You can pass slice objects instead of integers to obtain a list of
  56     symbols::
  57
  58         sage: a = SymbolSequence('a')
  59         sage: a(slice(5,7))
  60         [a5, a6]
  61
  62     This even works for the second, third, etc. indices::
  63
  64         sage: a = SymbolSequence('a')
  65         sage: a(slice(0,2), slice(0,2))
  66         [a00, a01, a10, a11]
  67
  68     You can also index with the list index operator::
  69
  70         sage: a = SymbolSequence('a')
  71         sage: a[1]
  72         a1
  73
  74     This allows you to retrieve one-dimensional slices easily::
  75
  76         sage: a = SymbolSequence('a')
  77         sage: a[0:5]
  78         [a0, a1, a2, a3, a4]
  79
  80     TESTS:
  81
  82     We shouldn't overwrite variables in the global namespace::
  83
  84         sage: a = SymbolSequence('a')
  85         sage: a0 = 4
  86         sage: a(0)
  87         a0
  88         sage: a0
  89         4
  90
  91     The symbol at a given index should always be the same, even when
  92     the symbols themselves are unnamed. We store the string
  93     representation and compare because the output is unpredictable::
  94
  95         sage: a = SymbolSequence()
  96         sage: a0str = str(a(0))
  97         sage: str(a(0)) == a0str
  98         True
  99
 100     Slices and single indices work when combined::
 101
 102         sage: a = SymbolSequence('a')
 103         sage: a(3, slice(0,2))
 104         [a30, a31]
 105         sage: a(slice(0,2), 3)
 106         [a03, a13]
 107
 108     """
 109
 110     def __init__(self, name=None, latex_name=None, domain=None):
 111         # We store a dict of already-created symbols so that we don't
 112         # recreate a symbol which already exists. This is especially
 113         # helpful when using unnamed variables, if you want e.g. a(0)
 114         # to return the same variable each time.
 115         #
 116         # The entry corresponding to None is the un-subscripted symbol
 117         # with our name.
 118         unsubscripted = SR.symbol(name, latex_name, domain)
 119         self._symbols = { None: unsubscripted }
 120
 121         self._name = name
 122         self._latex_name = latex_name
 123         self._domain = domain
 124
 125
 126     def _create_symbol_(self, subscript):
 127         if self._name is None:
 128             # Allow creating unnamed symbols, for consistency with
 129             # SR.symbol().
 130             name = None
 131         else:
 132             name = '%s%d' % (self._name, subscript)
 133
 134         if self._latex_name is None:
 135             latex_name = None
 136         else:
 137             latex_name = r'%s_{%d}' % (self._latex_name, subscript)
 138
 139         return SR.symbol(name, latex_name, self._domain)
 140
 141
 142     def _flatten_list_(self, l):
 143         """
 144         Recursively flatten the given list, allowing for some elements
 145         to be non-iterable.
 146         """
 147         result = []
 148
 149         for item in l:
 150             if isinstance(item, list):
 151                 result += self._flatten_list_(item)
 152             else:
 153                 result += [item]
 154
 155         return result
 156
 157
 158     def __getitem__(self, key):
 159         if isinstance(key, slice):
 160             # We were given a slice. Clean up some of its properties
 161             # first. The start/step are default for lists. We make
 162             # copies of these because they're read-only.
 163             (start, step) = (key.start, key.step)
 164             if start is None:
 165                 start = 0
 166             if key.stop is None:
 167                 # Would otherwise loop forever since our "length" is
 168                 # undefined.
 169                 raise ValueError('You must supply an terminal index')
 170             if step is None:
 171                step = 1
 172
 173             # If the user asks for a slice, we'll be returning a list
 174             # of symbols.
 175             return [ self(idx) for idx in range(start, key.stop, step) ]
 176         else:
 177             return self(key)
 178
 179
 180     def __call__(self, *args):
 181         args = list(args)
 182
 183         if len(args) == 0:
 184             return self._symbols[None]
 185
 186         # This is safe after the len == 0 test.
 187         key = args[0]
 188         args = args[1:] # Peel off the first arg, which we've called 'key'
 189
 190         if isinstance(key, slice):
 191             if len(args) == 0:
 192                 return self[key]
 193             else:
 194                 v = self[key]
 195                 ss = [ SymbolSequence(w._repr_(), w._latex_(), self._domain)
 196                        for w in v ]
 197
 198                 # This might be nested...
 199                 maybe_nested_list = [ s(*args) for s in ss ]
 200                 return self._flatten_list_(maybe_nested_list)
 201
 202         if key < 0:
 203             # Cowardly refuse to create a variable named "a-1".
 204             raise IndexError('Indices must be nonnegative')
 205
 206         if len(args) == 0:
 207             # Base case, create a symbol and return it.
 208             try:
 209                 return self._symbols[key]
 210             except KeyError:
 211                 self._symbols[key] = self._create_symbol_(key)
 212                 return self._symbols[key]
 213         else:
 214             # If we're given more than one index, we want to create the
 215             # subsequences recursively. For example, if we're asked for
 216             # x(1,2), this is really SymbolSequence('x1')(2).
 217             v = self(key) # x(1) -> x1
 218             ss = SymbolSequence(v._repr_(), v._latex_(), self._domain)
 219             return ss(*args)