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     SETUP::
  24
  25         sage: from mjo.symbol_sequence import SymbolSequence
  26
  27     EXAMPLES:
  28
  29     The simplest use case::
  30
  31         sage: a = SymbolSequence('a')
  32         sage: a[0]
  33         a_0
  34         sage: a[1]
  35         a_1
  36
  37     Create polynomials with symbolic coefficients of arbitrary
  38     degree::
  39
  40         sage: a = SymbolSequence('a')
  41         sage: p = sum( a[i]*x^i for i in xrange(5) )
  42         sage: p
  43         a_4*x^4 + a_3*x^3 + a_2*x^2 + a_1*x + a_0
  44
  45     Using a different latex name since 'lambda' is reserved::
  46
  47         sage: l = SymbolSequence('l', '\lambda')
  48         sage: l[0]
  49         l_0
  50         sage: latex(l[0])
  51         \lambda_{0}
  52
  53     Using multiple indices::
  54
  55         sage: a = SymbolSequence('a')
  56         sage: a[0,1,2]
  57         a_0_1_2
  58         sage: latex(a[0,1,2])
  59         a_{0}_{1}_{2}
  60         sage: [ a[i,j] for i in xrange(2) for j in xrange(2) ]
  61         [a_0_0, a_0_1, a_1_0, a_1_1]
  62
  63     You can pass slices instead of integers to obtain a list of
  64     symbols::
  65
  66         sage: a = SymbolSequence('a')
  67         sage: a[5:7]
  68         [a_5, a_6]
  69
  70     This even works for the second, third, etc. indices::
  71
  72         sage: a = SymbolSequence('a')
  73         sage: a[0:2, 0:2]
  74         [a_0_0, a_0_1, a_1_0, a_1_1]
  75
  76     TESTS:
  77
  78     We shouldn't overwrite variables in the global namespace::
  79
  80         sage: a = SymbolSequence('a')
  81         sage: a_0 = 4
  82         sage: a[0]
  83         a_0
  84         sage: a_0
  85         4
  86
  87     The symbol at a given index should always be the same, even when
  88     the symbols themselves are unnamed. We store the string
  89     representation and compare because the output is unpredictable::
  90
  91         sage: a = SymbolSequence()
  92         sage: a0str = str(a[0])
  93         sage: str(a[0]) == a0str
  94         True
  95
  96     Slices and single indices work when combined::
  97
  98         sage: a = SymbolSequence('a')
  99         sage: a[3, 0:2]
 100         [a_3_0, a_3_1]
 101         sage: a[0:2, 3]
 102         [a_0_3, a_1_3]
 103
 104     """
 105
 106     def __init__(self, name=None, latex_name=None, domain=None):
 107         # We store a dict of already-created symbols so that we don't
 108         # recreate a symbol which already exists. This is especially
 109         # helpful when using unnamed variables, if you want e.g. a[0]
 110         # to return the same variable each time.
 111         self._symbols = {}
 112
 113         self._name = name
 114         self._latex_name = latex_name
 115         self._domain = domain
 116
 117
 118     def _create_symbol_(self, subscript):
 119         """
 120         Return a symbol with the given subscript. Creates the
 121         appropriate name and latex_name before delegating to
 122         SR.symbol().
 123
 124         SETUP::
 125
 126             sage: from mjo.symbol_sequence import SymbolSequence
 127
 128         EXAMPLES::
 129
 130             sage: a = SymbolSequence('a', 'alpha', 'real')
 131             sage: a_1 = a._create_symbol_(1)
 132             sage: a_1
 133             a_1
 134             sage: latex(a_1)
 135             alpha_{1}
 136
 137         """
 138         # Allow creating unnamed symbols, for consistency with
 139         # SR.symbol().
 140         name = None
 141         if self._name is not None:
 142             name = '%s_%d' % (self._name, subscript)
 143
 144         latex_name = None
 145         if self._latex_name is not None:
 146             latex_name = r'%s_{%d}' % (self._latex_name, subscript)
 147
 148         return SR.symbol(name, latex_name, self._domain)
 149
 150
 151     def _flatten_list_(self, l):
 152         """
 153         Recursively flatten the given list, allowing for some elements
 154         to be non-iterable. This is slow, but also works, which is
 155         more than can be said about some of the snappier solutions of
 156         lore.
 157
 158         SETUP::
 159
 160             sage: from mjo.symbol_sequence import SymbolSequence
 161
 162         EXAMPLES::
 163
 164             sage: a = SymbolSequence('a')
 165             sage: a._flatten_list_([1,2,3])
 166             [1, 2, 3]
 167             sage: a._flatten_list_([1,[2,3]])
 168             [1, 2, 3]
 169             sage: a._flatten_list_([1,[2,[3]]])
 170             [1, 2, 3]
 171             sage: a._flatten_list_([[[[[1,[2,[3]]]]]]])
 172             [1, 2, 3]
 173
 174         """
 175         result = []
 176
 177         for item in l:
 178             if isinstance(item, list):
 179                 result += self._flatten_list_(item)
 180             else:
 181                 result += [item]
 182
 183         return result
 184
 185
 186     def __getitem__(self, key):
 187         """
 188         This handles individual integer arguments, slices, and
 189         tuples. It just hands off the real work to
 190         self._subscript_foo_().
 191
 192         SETUP::
 193
 194             sage: from mjo.symbol_sequence import SymbolSequence
 195
 196         EXAMPLES:
 197
 198         An integer argument::
 199
 200             sage: a = SymbolSequence('a')
 201             sage: a.__getitem__(1)
 202             a_1
 203
 204         A tuple argument::
 205
 206             sage: a = SymbolSequence('a')
 207             sage: a.__getitem__((1,2))
 208             a_1_2
 209
 210         A slice argument::
 211
 212             sage: a = SymbolSequence('a')
 213             sage: a.__getitem__(slice(1,4))
 214             [a_1, a_2, a_3]
 215
 216         """
 217         if isinstance(key, tuple):
 218             return self._subscript_tuple_(key)
 219
 220         if isinstance(key, slice):
 221             return self._subscript_slice_(key)
 222
 223         # This is the most common case so it would make sense to have
 224         # this test first. But there are too many different "integer"
 225         # classes that you have to check for.
 226         return self._subscript_integer_(key)
 227
 228
 229     def _subscript_integer_(self, n):
 230         """
 231         The subscript is a single integer, or something that acts like
 232         one.
 233
 234         SETUP::
 235
 236             sage: from mjo.symbol_sequence import SymbolSequence
 237
 238         EXAMPLES::
 239
 240             sage: a = SymbolSequence('a')
 241             sage: a._subscript_integer_(123)
 242             a_123
 243
 244         """
 245         if n < 0:
 246             # Cowardly refuse to create a variable named "a-1".
 247             raise IndexError('Indices must be nonnegative')
 248
 249         try:
 250             return self._symbols[n]
 251         except KeyError:
 252             self._symbols[n] = self._create_symbol_(n)
 253             return self._symbols[n]
 254
 255
 256     def _subscript_slice_(self, s):
 257         """
 258         We were given a slice. Clean up some of its properties
 259         first. The start/step are default for lists. We make
 260         copies of these because they're read-only.
 261
 262         SETUP::
 263
 264             sage: from mjo.symbol_sequence import SymbolSequence
 265
 266         EXAMPLES::
 267
 268             sage: a = SymbolSequence('a')
 269             sage: a._subscript_slice_(slice(1,3))
 270             [a_1, a_2]
 271
 272         """
 273         (start, step) = (s.start, s.step)
 274         if start is None:
 275             start = 0
 276         if s.stop is None:
 277             # Would otherwise loop forever since our "length" is
 278             # undefined.
 279             raise ValueError('You must supply an terminal index')
 280         if step is None:
 281             step = 1
 282
 283         # If the user asks for a slice, we'll be returning a list
 284         # of symbols.
 285         return [ self._subscript_integer_(idx)
 286                  for idx in xrange(start, s.stop, step) ]
 287
 288
 289
 290     def _subscript_tuple_(self, args):
 291         """
 292         When we have more than one level of subscripts, we pick off
 293         the first one and generate the rest recursively.
 294
 295         SETUP::
 296
 297             sage: from mjo.symbol_sequence import SymbolSequence
 298
 299         EXAMPLES:
 300
 301         A simple two-tuple::
 302
 303             sage: a = SymbolSequence('a')
 304             sage: a._subscript_tuple_((1,8))
 305             a_1_8
 306
 307         Nested tuples::
 308
 309             sage: a._subscript_tuple_(( (1,2), (3,(4,5,6)) ))
 310             a_1_2_3_4_5_6
 311
 312         """
 313
 314         # We never call this method without an argument.
 315         key = args[0]
 316         args = args[1:] # Peel off the first arg, which we've called 'key'
 317
 318         # We don't know the type of 'key', but __getitem__ will figure
 319         # it out and dispatch properly.
 320         v = self[key]
 321         if len(args) == 0:
 322             # There was only one element left in the tuple.
 323             return v
 324
 325         # At this point, we know we were given at least a two-tuple.
 326         # The symbols corresponding to the first entry are already
 327         # computed, in 'v'. Here we recursively compute the symbols
 328         # corresponding to the second coordinate, with the first
 329         # coordinate(s) fixed.
 330         if isinstance(key, slice):
 331             ss = ( SymbolSequence(w._repr_(), w._latex_(), self._domain)
 332                    for w in v )
 333
 334             # This might be nested...
 335             maybe_nested_list = [ s._subscript_tuple_(args) for s in ss ]
 336             return self._flatten_list_(maybe_nested_list)
 337
 338         else:
 339             # If it's not a slice, it's an integer.
 340             ss = SymbolSequence(v._repr_(), v._latex_(), self._domain)
 341             return ss._subscript_tuple_(args)