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 `x_0`, `x_1`,...
  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         a_0
  30         sage: a[1]
  31         a_1
  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         a_4*x^4 + a_3*x^3 + a_2*x^2 + a_1*x + a_0
  40
  41     Using a different latex name since 'lambda' is reserved::
  42
  43         sage: l = SymbolSequence('l', '\lambda')
  44         sage: l[0]
  45         l_0
  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         a_0_1_2
  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         [a_0_0, a_0_1, a_1_0, a_1_1]
  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         [a_5, a_6]
  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         [a_0_0, a_0_1, a_1_0, a_1_1]
  71
  72     TESTS:
  73
  74     We shouldn't overwrite variables in the global namespace::
  75
  76         sage: a = SymbolSequence('a')
  77         sage: a_0 = 4
  78         sage: a[0]
  79         a_0
  80         sage: a_0
  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         [a_3_0, a_3_1]
  97         sage: a[0:2, 3]
  98         [a_0_3, a_1_3]
  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         EXAMPLES::
 121
 122             sage: a = SymbolSequence('a', 'alpha', 'real')
 123             sage: a_1 = a._create_symbol_(1)
 124             sage: a_1
 125             a_1
 126             sage: latex(a_1)
 127             alpha_{1}
 128
 129         """
 130         # Allow creating unnamed symbols, for consistency with
 131         # SR.symbol().
 132         name = None
 133         if self._name is not None:
 134             name = '%s_%d' % (self._name, subscript)
 135
 136         latex_name = None
 137         if self._latex_name is not None:
 138             latex_name = r'%s_{%d}' % (self._latex_name, subscript)
 139
 140         return SR.symbol(name, latex_name, self._domain)
 141
 142
 143     def _flatten_list_(self, l):
 144         """
 145         Recursively flatten the given list, allowing for some elements
 146         to be non-iterable. This is slow, but also works, which is
 147         more than can be said about some of the snappier solutions of
 148         lore.
 149
 150         EXAMPLES::
 151
 152             sage: a = SymbolSequence('a')
 153             sage: a._flatten_list_([1,2,3])
 154             [1, 2, 3]
 155             sage: a._flatten_list_([1,[2,3]])
 156             [1, 2, 3]
 157             sage: a._flatten_list_([1,[2,[3]]])
 158             [1, 2, 3]
 159             sage: a._flatten_list_([[[[[1,[2,[3]]]]]]])
 160             [1, 2, 3]
 161
 162         """
 163         result = []
 164
 165         for item in l:
 166             if isinstance(item, list):
 167                 result += self._flatten_list_(item)
 168             else:
 169                 result += [item]
 170
 171         return result
 172
 173
 174     def __getitem__(self, key):
 175         """
 176         This handles individual integer arguments, slices, and
 177         tuples. It just hands off the real work to
 178         self._subscript_foo_().
 179
 180         EXAMPLES:
 181
 182         An integer argument::
 183
 184             sage: a = SymbolSequence('a')
 185             sage: a.__getitem__(1)
 186             a_1
 187
 188         A tuple argument::
 189
 190             sage: a = SymbolSequence('a')
 191             sage: a.__getitem__((1,2))
 192             a_1_2
 193
 194         A slice argument::
 195
 196             sage: a = SymbolSequence('a')
 197             sage: a.__getitem__(slice(1,4))
 198             [a_1, a_2, a_3]
 199
 200         """
 201         if isinstance(key, tuple):
 202             return self._subscript_tuple_(key)
 203
 204         if isinstance(key, slice):
 205             return self._subscript_slice_(key)
 206
 207         # This is the most common case so it would make sense to have
 208         # this test first. But there are too many different "integer"
 209         # classes that you have to check for.
 210         return self._subscript_integer_(key)
 211
 212
 213     def _subscript_integer_(self, n):
 214         """
 215         The subscript is a single integer, or something that acts like
 216         one.
 217
 218         EXAMPLES::
 219
 220             sage: a = SymbolSequence('a')
 221             sage: a._subscript_integer_(123)
 222             a_123
 223
 224         """
 225         if n < 0:
 226             # Cowardly refuse to create a variable named "a-1".
 227             raise IndexError('Indices must be nonnegative')
 228
 229         try:
 230             return self._symbols[n]
 231         except KeyError:
 232             self._symbols[n] = self._create_symbol_(n)
 233             return self._symbols[n]
 234
 235
 236     def _subscript_slice_(self, s):
 237         """
 238         We were given a slice. Clean up some of its properties
 239         first. The start/step are default for lists. We make
 240         copies of these because they're read-only.
 241
 242         EXAMPLES::
 243
 244             sage: a = SymbolSequence('a')
 245             sage: a._subscript_slice_(slice(1,3))
 246             [a_1, a_2]
 247
 248         """
 249         (start, step) = (s.start, s.step)
 250         if start is None:
 251             start = 0
 252         if s.stop is None:
 253             # Would otherwise loop forever since our "length" is
 254             # undefined.
 255             raise ValueError('You must supply an terminal index')
 256         if step is None:
 257             step = 1
 258
 259         # If the user asks for a slice, we'll be returning a list
 260         # of symbols.
 261         return [ self._subscript_integer_(idx)
 262                  for idx in range(start, s.stop, step) ]
 263
 264
 265
 266     def _subscript_tuple_(self, args):
 267         """
 268         When we have more than one level of subscripts, we pick off
 269         the first one and generate the rest recursively.
 270
 271         EXAMPLES:
 272
 273         A simple two-tuple::
 274
 275             sage: a = SymbolSequence('a')
 276             sage: a._subscript_tuple_((1,8))
 277             a_1_8
 278
 279         Nested tuples::
 280
 281             sage: a._subscript_tuple_(( (1,2), (3,(4,5,6)) ))
 282             a_1_2_3_4_5_6
 283
 284         """
 285
 286         # We never call this method without an argument.
 287         key = args[0]
 288         args = args[1:] # Peel off the first arg, which we've called 'key'
 289
 290         # We don't know the type of 'key', but __getitem__ will figure
 291         # it out and dispatch properly.
 292         v = self[key]
 293         if len(args) == 0:
 294             # There was only one element left in the tuple.
 295             return v
 296
 297         # At this point, we know we were given at least a two-tuple.
 298         # The symbols corresponding to the first entry are already
 299         # computed, in 'v'. Here we recursively compute the symbols
 300         # corresponding to the second coordinate, with the first
 301         # coordinate(s) fixed.
 302         if isinstance(key, slice):
 303             ss = [ SymbolSequence(w._repr_(), w._latex_(), self._domain)
 304                    for w in v ]
 305
 306             # This might be nested...
 307             maybe_nested_list = [ s._subscript_tuple_(args) for s in ss ]
 308             return self._flatten_list_(maybe_nested_list)
 309
 310         else:
 311             # If it's not a slice, it's an integer.
 312             ss = SymbolSequence(v._repr_(), v._latex_(), self._domain)
 313             return ss._subscript_tuple_(args)