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 coefficients for polynomials of arbitrary degree::
  34
  35         sage: a = SymbolSequence('a')
  36         sage: p = sum([ a[i]*x^i for i in range(0,5)])
  37         sage: p
  38         a4*x^4 + a3*x^3 + a2*x^2 + a1*x + a0
  39
  40     Using a different latex name since 'lambda' is reserved::
  41
  42         sage: l = SymbolSequence('l', '\lambda')
  43         sage: l[0]
  44         l0
  45         sage: latex(l[0])
  46         \lambda_{0}
  47
  48     Using multiple indices::
  49
  50         sage: a = SymbolSequence('a')
  51         sage: a[0,1,2]
  52         a012
  53         sage: latex(a[0,1,2])
  54         a_{0}_{1}_{2}
  55         sage: [ a[i,j] for i in range(0,2) for j in range(0,2) ]
  56         [a00, a01, a10, a11]
  57
  58     You can pass slice objects instead of integers to obtain a list of
  59     symbols::
  60
  61         sage: a = SymbolSequence('a')
  62         sage: a[5:7]
  63         [a5, a6]
  64
  65     This even works for the second, third, etc. indices::
  66
  67         sage: a = SymbolSequence('a')
  68         sage: a[0:2, 0:2]
  69         [a00, a01, a10, a11]
  70
  71     TESTS:
  72
  73     We shouldn't overwrite variables in the global namespace::
  74
  75         sage: a = SymbolSequence('a')
  76         sage: a0 = 4
  77         sage: a[0]
  78         a0
  79         sage: a0
  80         4
  81
  82     The symbol at a given index should always be the same, even when
  83     the symbols themselves are unnamed. We store the string
  84     representation and compare because the output is unpredictable::
  85
  86         sage: a = SymbolSequence()
  87         sage: a0str = str(a[0])
  88         sage: str(a(0)) == a0str
  89         True
  90
  91     Slices and single indices work when combined::
  92
  93         sage: a = SymbolSequence('a')
  94         sage: a[3, 0:2]
  95         [a30, a31]
  96         sage: a[0:2, 3]
  97         [a03, a13]
  98
  99     """
 100
 101     def __init__(self, name=None, latex_name=None, domain=None):
 102         # We store a dict of already-created symbols so that we don't
 103         # recreate a symbol which already exists. This is especially
 104         # helpful when using unnamed variables, if you want e.g. a(0)
 105         # to return the same variable each time.
 106         #
 107         # The entry corresponding to None is the un-subscripted symbol
 108         # with our name.
 109         unsubscripted = SR.symbol(name, latex_name, domain)
 110         self._symbols = { None: unsubscripted }
 111
 112         self._name = name
 113         self._latex_name = latex_name
 114         self._domain = domain
 115
 116
 117     def _create_symbol_(self, subscript):
 118         if self._name is None:
 119             # Allow creating unnamed symbols, for consistency with
 120             # SR.symbol().
 121             name = None
 122         else:
 123             name = '%s%d' % (self._name, subscript)
 124
 125         if self._latex_name is None:
 126             latex_name = None
 127         else:
 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.
 137         """
 138         result = []
 139
 140         for item in l:
 141             if isinstance(item, list):
 142                 result += self._flatten_list_(item)
 143             else:
 144                 result += [item]
 145
 146         return result
 147
 148
 149     def __getitem__(self, key):
 150         if isinstance(key, tuple):
 151             return self(*key)
 152
 153         if isinstance(key, slice):
 154             # We were given a slice. Clean up some of its properties
 155             # first. The start/step are default for lists. We make
 156             # copies of these because they're read-only.
 157             (start, step) = (key.start, key.step)
 158             if start is None:
 159                 start = 0
 160             if key.stop is None:
 161                 # Would otherwise loop forever since our "length" is
 162                 # undefined.
 163                 raise ValueError('You must supply an terminal index')
 164             if step is None:
 165                step = 1
 166
 167             # If the user asks for a slice, we'll be returning a list
 168             # of symbols.
 169             return [ self(idx) for idx in range(start, key.stop, step) ]
 170
 171         return self(key)
 172
 173
 174     def __call__(self, *args):
 175         args = list(args)
 176
 177         if len(args) == 0:
 178             return self._symbols[None]
 179
 180         # This is safe after the len == 0 test.
 181         key = args[0]
 182         args = args[1:] # Peel off the first arg, which we've called 'key'
 183
 184         if isinstance(key, slice):
 185             if len(args) == 0:
 186                 return self[key]
 187             else:
 188                 v = self[key]
 189                 ss = [ SymbolSequence(w._repr_(), w._latex_(), self._domain)
 190                        for w in v ]
 191
 192                 # This might be nested...
 193                 maybe_nested_list = [ s(*args) for s in ss ]
 194                 return self._flatten_list_(maybe_nested_list)
 195
 196         if key < 0:
 197             # Cowardly refuse to create a variable named "a-1".
 198             raise IndexError('Indices must be nonnegative')
 199
 200         if len(args) == 0:
 201             # Base case, create a symbol and return it.
 202             try:
 203                 return self._symbols[key]
 204             except KeyError:
 205                 self._symbols[key] = self._create_symbol_(key)
 206                 return self._symbols[key]
 207         else:
 208             # If we're given more than one index, we want to create the
 209             # subsequences recursively. For example, if we're asked for
 210             # x(1,2), this is really SymbolSequence('x1')(2).
 211             v = self(key) # x(1) -> x1
 212             ss = SymbolSequence(v._repr_(), v._latex_(), self._domain)
 213             return ss(*args)