mjo/symbol_sequence.py

   1 from sage.all import *
   2
   3 class SymbolSequence:
   4     r"""
   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 range(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', r'\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 range(2) for j in range(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             try:
 179                 item = iter(item)
 180                 result += self._flatten_list_(item)
 181             except TypeError:
 182                 result += [item]
 183
 184         return result
 185
 186
 187     def __getitem__(self, key):
 188         """
 189         This handles individual integer arguments, slices, and
 190         tuples. It just hands off the real work to
 191         self._subscript_foo_().
 192
 193         SETUP::
 194
 195             sage: from mjo.symbol_sequence import SymbolSequence
 196
 197         EXAMPLES:
 198
 199         An integer argument::
 200
 201             sage: a = SymbolSequence('a')
 202             sage: a.__getitem__(1)
 203             a_1
 204
 205         A tuple argument::
 206
 207             sage: a = SymbolSequence('a')
 208             sage: a.__getitem__((1,2))
 209             a_1_2
 210
 211         A slice argument::
 212
 213             sage: a = SymbolSequence('a')
 214             sage: a.__getitem__(slice(1,4))
 215             [a_1, a_2, a_3]
 216
 217         """
 218         if isinstance(key, tuple):
 219             return self._subscript_tuple_(key)
 220
 221         if isinstance(key, slice):
 222             return self._subscript_slice_(key)
 223
 224         # This is the most common case so it would make sense to have
 225         # this test first. But there are too many different "integer"
 226         # classes that you have to check for.
 227         return self._subscript_integer_(key)
 228
 229
 230     def _subscript_integer_(self, n):
 231         """
 232         The subscript is a single integer, or something that acts like
 233         one.
 234
 235         SETUP::
 236
 237             sage: from mjo.symbol_sequence import SymbolSequence
 238
 239         EXAMPLES::
 240
 241             sage: a = SymbolSequence('a')
 242             sage: a._subscript_integer_(123)
 243             a_123
 244
 245         """
 246         if n < 0:
 247             # Cowardly refuse to create a variable named "a-1".
 248             raise IndexError('Indices must be nonnegative')
 249
 250         try:
 251             return self._symbols[n]
 252         except KeyError:
 253             self._symbols[n] = self._create_symbol_(n)
 254             return self._symbols[n]
 255
 256
 257     def _subscript_slice_(self, s):
 258         """
 259         We were given a slice. Clean up some of its properties
 260         first. The start/step are default for lists. We make
 261         copies of these because they're read-only.
 262
 263         SETUP::
 264
 265             sage: from mjo.symbol_sequence import SymbolSequence
 266
 267         EXAMPLES::
 268
 269             sage: a = SymbolSequence('a')
 270             sage: a._subscript_slice_(slice(1,3))
 271             [a_1, a_2]
 272
 273         """
 274         (start, step) = (s.start, s.step)
 275         if start is None:
 276             start = 0
 277         if s.stop is None:
 278             # Would otherwise loop forever since our "length" is
 279             # undefined.
 280             raise ValueError('You must supply an terminal index')
 281         if step is None:
 282             step = 1
 283
 284         # If the user asks for a slice, we'll be returning a list
 285         # of symbols.
 286         return [ self._subscript_integer_(idx)
 287                  for idx in range(start, s.stop, step) ]
 288
 289
 290
 291     def _subscript_tuple_(self, args):
 292         """
 293         When we have more than one level of subscripts, we pick off
 294         the first one and generate the rest recursively.
 295
 296         SETUP::
 297
 298             sage: from mjo.symbol_sequence import SymbolSequence
 299
 300         EXAMPLES:
 301
 302         A simple two-tuple::
 303
 304             sage: a = SymbolSequence('a')
 305             sage: a._subscript_tuple_((1,8))
 306             a_1_8
 307
 308         Nested tuples::
 309
 310             sage: a._subscript_tuple_(( (1,2), (3,(4,5,6)) ))
 311             a_1_2_3_4_5_6
 312
 313         """
 314
 315         # We never call this method without an argument.
 316         key = args[0]
 317         args = args[1:] # Peel off the first arg, which we've called 'key'
 318
 319         # We don't know the type of 'key', but __getitem__ will figure
 320         # it out and dispatch properly.
 321         v = self[key]
 322         if len(args) == 0:
 323             # There was only one element left in the tuple.
 324             return v
 325
 326         # At this point, we know we were given at least a two-tuple.
 327         # The symbols corresponding to the first entry are already
 328         # computed, in 'v'. Here we recursively compute the symbols
 329         # corresponding to the second coordinate, with the first
 330         # coordinate(s) fixed.
 331         if isinstance(key, slice):
 332             ss = ( SymbolSequence(w._repr_(), w._latex_(), self._domain)
 333                    for w in v )
 334
 335             # This might be nested...
 336             maybe_nested_list = ( s._subscript_tuple_(args) for s in ss )
 337             return self._flatten_list_(maybe_nested_list)
 338
 339         else:
 340             # If it's not a slice, it's an integer.
 341             ss = SymbolSequence(v._repr_(), v._latex_(), self._domain)
 342             return ss._subscript_tuple_(args)