from sage.all import * class SymbolSequence: """ An iterable object which acts like a sequence of symbolic expressions (variables). INPUT: - ``name`` -- The sequence name. For example, if you name the sequence `x`, the variables will be called `x0`, `x1`,... - ``latex_name`` -- An optional latex expression (string) to use instead of `name` when converting the symbols to latex. - ``domain`` -- A string representing the domain of the symbol, either 'real', 'complex', or 'positive'. OUTPUT: An iterable object containing symbolic expressions. EXAMPLES: The simplest use case:: sage: a = SymbolSequence('a') sage: a[0] a0 sage: a[1] a1 Create polynomials with symbolic coefficients of arbitrary degree:: sage: a = SymbolSequence('a') sage: p = sum([ a[i]*x^i for i in range(0,5)]) sage: p a4*x^4 + a3*x^3 + a2*x^2 + a1*x + a0 Using a different latex name since 'lambda' is reserved:: sage: l = SymbolSequence('l', '\lambda') sage: l[0] l0 sage: latex(l[0]) \lambda_{0} Using multiple indices:: sage: a = SymbolSequence('a') sage: a[0,1,2] a012 sage: latex(a[0,1,2]) a_{0}_{1}_{2} sage: [ a[i,j] for i in range(0,2) for j in range(0,2) ] [a00, a01, a10, a11] You can pass slices instead of integers to obtain a list of symbols:: sage: a = SymbolSequence('a') sage: a[5:7] [a5, a6] This even works for the second, third, etc. indices:: sage: a = SymbolSequence('a') sage: a[0:2, 0:2] [a00, a01, a10, a11] TESTS: We shouldn't overwrite variables in the global namespace:: sage: a = SymbolSequence('a') sage: a0 = 4 sage: a[0] a0 sage: a0 4 The symbol at a given index should always be the same, even when the symbols themselves are unnamed. We store the string representation and compare because the output is unpredictable:: sage: a = SymbolSequence() sage: a0str = str(a[0]) sage: str(a[0]) == a0str True Slices and single indices work when combined:: sage: a = SymbolSequence('a') sage: a[3, 0:2] [a30, a31] sage: a[0:2, 3] [a03, a13] """ def __init__(self, name=None, latex_name=None, domain=None): # We store a dict of already-created symbols so that we don't # recreate a symbol which already exists. This is especially # helpful when using unnamed variables, if you want e.g. a[0] # to return the same variable each time. self._symbols = {} self._name = name self._latex_name = latex_name self._domain = domain def _create_symbol_(self, subscript): """ Return a symbol with the given subscript. Creates the appropriate name and latex_name before delegating to SR.symbol(). """ # Allow creating unnamed symbols, for consistency with # SR.symbol(). name = None if self._name is not None: name = '%s%d' % (self._name, subscript) latex_name = None if self._latex_name is not None: latex_name = r'%s_{%d}' % (self._latex_name, subscript) return SR.symbol(name, latex_name, self._domain) def _flatten_list_(self, l): """ Recursively flatten the given list, allowing for some elements to be non-iterable. This is slow, but also works, which is more than can be said about some of the snappier solutions of lore. """ result = [] for item in l: if isinstance(item, list): result += self._flatten_list_(item) else: result += [item] return result def __getitem__(self, key): """ This handles individual integer arguments, slices, and tuples. It just hands off the real work to self._subscript_foo_(). """ if isinstance(key, tuple): return self._subscript_tuple_(key) if isinstance(key, slice): return self._subscript_slice_(key) # This is the most common case so it would make sense to have # this test first. But there are too many different "integer" # classes that you have to check for. return self._subscript_integer_(key) def _subscript_integer_(self, n): """ The subscript is a single integer, or something that acts like one. """ if n < 0: # Cowardly refuse to create a variable named "a-1". raise IndexError('Indices must be nonnegative') try: return self._symbols[n] except KeyError: self._symbols[n] = self._create_symbol_(n) return self._symbols[n] def _subscript_slice_(self, s): """ We were given a slice. Clean up some of its properties first. The start/step are default for lists. We make copies of these because they're read-only. """ (start, step) = (s.start, s.step) if start is None: start = 0 if s.stop is None: # Would otherwise loop forever since our "length" is # undefined. raise ValueError('You must supply an terminal index') if step is None: step = 1 # If the user asks for a slice, we'll be returning a list # of symbols. return [ self._subscript_integer_(idx) for idx in range(start, s.stop, step) ] def _subscript_tuple_(self, args): """ When we have more than one level of subscripts, we pick off the first one and generate the rest recursively. """ # We never call this method without an argument. key = args[0] args = args[1:] # Peel off the first arg, which we've called 'key' # We don't know the type of 'key', but __getitem__ will figure # it out and dispatch properly. v = self[key] if len(args) == 0: # There was only one element left in the tuple. return v # At this point, we know we were given at least a two-tuple. # The symbols corresponding to the first entry are already # computed, in 'v'. Here we recursively compute the symbols # corresponding to the second coordinate, with the first # coordinate(s) fixed. if isinstance(key, slice): ss = [ SymbolSequence(w._repr_(), w._latex_(), self._domain) for w in v ] # This might be nested... maybe_nested_list = [ s._subscript_tuple_(args) for s in ss ] return self._flatten_list_(maybe_nested_list) else: # If it's not a slice, it's an integer. ss = SymbolSequence(v._repr_(), v._latex_(), self._domain) return ss._subscript_tuple_(args)