mjo/symbol_sequence.py

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