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     If for whatever reason you prefer square brackets, the following
  81     also works::
  82
  83         sage: a = SymbolSequence('a')
  84         sage: a[1,3]
  85         a13
  86
  87     TESTS:
  88
  89     We shouldn't overwrite variables in the global namespace::
  90
  91         sage: a = SymbolSequence('a')
  92         sage: a0 = 4
  93         sage: a(0)
  94         a0
  95         sage: a0
  96         4
  97
  98     The symbol at a given index should always be the same, even when
  99     the symbols themselves are unnamed. We store the string
 100     representation and compare because the output is unpredictable::
 101
 102         sage: a = SymbolSequence()
 103         sage: a0str = str(a(0))
 104         sage: str(a(0)) == a0str
 105         True
 106
 107     Slices and single indices work when combined::
 108
 109         sage: a = SymbolSequence('a')
 110         sage: a(3, slice(0,2))
 111         [a30, a31]
 112         sage: a(slice(0,2), 3)
 113         [a03, a13]
 114
 115     """
 116
 117     def __init__(self, name=None, latex_name=None, domain=None):
 118         # We store a dict of already-created symbols so that we don't
 119         # recreate a symbol which already exists. This is especially
 120         # helpful when using unnamed variables, if you want e.g. a(0)
 121         # to return the same variable each time.
 122         #
 123         # The entry corresponding to None is the un-subscripted symbol
 124         # with our name.
 125         unsubscripted = SR.symbol(name, latex_name, domain)
 126         self._symbols = { None: unsubscripted }
 127
 128         self._name = name
 129         self._latex_name = latex_name
 130         self._domain = domain
 131
 132
 133     def _create_symbol_(self, subscript):
 134         if self._name is None:
 135             # Allow creating unnamed symbols, for consistency with
 136             # SR.symbol().
 137             name = None
 138         else:
 139             name = '%s%d' % (self._name, subscript)
 140
 141         if self._latex_name is None:
 142             latex_name = None
 143         else:
 144             latex_name = r'%s_{%d}' % (self._latex_name, subscript)
 145
 146         return SR.symbol(name, latex_name, self._domain)
 147
 148
 149     def _flatten_list_(self, l):
 150         """
 151         Recursively flatten the given list, allowing for some elements
 152         to be non-iterable.
 153         """
 154         result = []
 155
 156         for item in l:
 157             if isinstance(item, list):
 158                 result += self._flatten_list_(item)
 159             else:
 160                 result += [item]
 161
 162         return result
 163
 164
 165     def __getitem__(self, key):
 166         if isinstance(key, tuple):
 167             return self(*key)
 168
 169         if isinstance(key, slice):
 170             # We were given a slice. Clean up some of its properties
 171             # first. The start/step are default for lists. We make
 172             # copies of these because they're read-only.
 173             (start, step) = (key.start, key.step)
 174             if start is None:
 175                 start = 0
 176             if key.stop is None:
 177                 # Would otherwise loop forever since our "length" is
 178                 # undefined.
 179                 raise ValueError('You must supply an terminal index')
 180             if step is None:
 181                step = 1
 182
 183             # If the user asks for a slice, we'll be returning a list
 184             # of symbols.
 185             return [ self(idx) for idx in range(start, key.stop, step) ]
 186
 187         return self(key)
 188
 189
 190     def __call__(self, *args):
 191         args = list(args)
 192
 193         if len(args) == 0:
 194             return self._symbols[None]
 195
 196         # This is safe after the len == 0 test.
 197         key = args[0]
 198         args = args[1:] # Peel off the first arg, which we've called 'key'
 199
 200         if isinstance(key, slice):
 201             if len(args) == 0:
 202                 return self[key]
 203             else:
 204                 v = self[key]
 205                 ss = [ SymbolSequence(w._repr_(), w._latex_(), self._domain)
 206                        for w in v ]
 207
 208                 # This might be nested...
 209                 maybe_nested_list = [ s(*args) for s in ss ]
 210                 return self._flatten_list_(maybe_nested_list)
 211
 212         if key < 0:
 213             # Cowardly refuse to create a variable named "a-1".
 214             raise IndexError('Indices must be nonnegative')
 215
 216         if len(args) == 0:
 217             # Base case, create a symbol and return it.
 218             try:
 219                 return self._symbols[key]
 220             except KeyError:
 221                 self._symbols[key] = self._create_symbol_(key)
 222                 return self._symbols[key]
 223         else:
 224             # If we're given more than one index, we want to create the
 225             # subsequences recursively. For example, if we're asked for
 226             # x(1,2), this is really SymbolSequence('x1')(2).
 227             v = self(key) # x(1) -> x1
 228             ss = SymbolSequence(v._repr_(), v._latex_(), self._domain)
 229             return ss(*args)