from sage.all import * class SymbolSequence: r""" 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 `x_0`, `x_1`,... - ``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. SETUP:: sage: from mjo.symbol_sequence import SymbolSequence EXAMPLES: The simplest use case:: sage: a = SymbolSequence('a') sage: a[0] a_0 sage: a[1] a_1 Create polynomials with symbolic coefficients of arbitrary degree:: sage: a = SymbolSequence('a') sage: p = sum( a[i]*x^i for i in range(5) ) sage: p a_4*x^4 + a_3*x^3 + a_2*x^2 + a_1*x + a_0 Using a different latex name since 'lambda' is reserved:: sage: l = SymbolSequence('l', r'\lambda') sage: l[0] l_0 sage: latex(l[0]) \lambda_{0} Using multiple indices:: sage: a = SymbolSequence('a') sage: a[0,1,2] a_0_1_2 sage: latex(a[0,1,2]) a_{0}_{1}_{2} sage: [ a[i,j] for i in range(2) for j in range(2) ] [a_0_0, a_0_1, a_1_0, a_1_1] You can pass slices instead of integers to obtain a list of symbols:: sage: a = SymbolSequence('a') sage: a[5:7] [a_5, a_6] This even works for the second, third, etc. indices:: sage: a = SymbolSequence('a') sage: a[0:2, 0:2] [a_0_0, a_0_1, a_1_0, a_1_1] TESTS: We shouldn't overwrite variables in the global namespace:: sage: a = SymbolSequence('a') sage: a_0 = 4 sage: a[0] a_0 sage: a_0 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] [a_3_0, a_3_1] sage: a[0:2, 3] [a_0_3, a_1_3] """ 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(). SETUP:: sage: from mjo.symbol_sequence import SymbolSequence EXAMPLES:: sage: a = SymbolSequence('a', 'alpha', 'real') sage: a_1 = a._create_symbol_(1) sage: a_1 a_1 sage: latex(a_1) alpha_{1} """ # 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. SETUP:: sage: from mjo.symbol_sequence import SymbolSequence EXAMPLES:: sage: a = SymbolSequence('a') sage: a._flatten_list_([1,2,3]) [1, 2, 3] sage: a._flatten_list_([1,[2,3]]) [1, 2, 3] sage: a._flatten_list_([1,[2,[3]]]) [1, 2, 3] sage: a._flatten_list_([[[[[1,[2,[3]]]]]]]) [1, 2, 3] """ result = [] for item in l: try: item = iter(item) result += self._flatten_list_(item) except TypeError: 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_(). SETUP:: sage: from mjo.symbol_sequence import SymbolSequence EXAMPLES: An integer argument:: sage: a = SymbolSequence('a') sage: a.__getitem__(1) a_1 A tuple argument:: sage: a = SymbolSequence('a') sage: a.__getitem__((1,2)) a_1_2 A slice argument:: sage: a = SymbolSequence('a') sage: a.__getitem__(slice(1,4)) [a_1, a_2, a_3] """ 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. SETUP:: sage: from mjo.symbol_sequence import SymbolSequence EXAMPLES:: sage: a = SymbolSequence('a') sage: a._subscript_integer_(123) a_123 """ 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. SETUP:: sage: from mjo.symbol_sequence import SymbolSequence EXAMPLES:: sage: a = SymbolSequence('a') sage: a._subscript_slice_(slice(1,3)) [a_1, a_2] """ (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. SETUP:: sage: from mjo.symbol_sequence import SymbolSequence EXAMPLES: A simple two-tuple:: sage: a = SymbolSequence('a') sage: a._subscript_tuple_((1,8)) a_1_8 Nested tuples:: sage: a._subscript_tuple_(( (1,2), (3,(4,5,6)) )) a_1_2_3_4_5_6 """ # 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)