예제 #1
0
def test_indexed_idx_sum():
    i = symbols('i', cls=Idx)
    r = Indexed('r', i)
    assert Sum(r,
               (i, 0, 3)).doit() == sum([r.xreplace({i: j}) for j in range(4)])
    assert Product(r, (i, 0, 3)).doit() == prod(
        [r.xreplace({i: j}) for j in range(4)])

    j = symbols('j', integer=True)
    assert Sum(r, (i, j, j + 2)).doit() == sum(
        [r.xreplace({i: Idx(j + k)}) for k in range(3)])
    assert Product(r, (i, j, j + 2)).doit() == prod(
        [r.xreplace({i: Idx(j + k)}) for k in range(3)])

    k = Idx('k', range=(1, 3))
    A = IndexedBase('A')
    assert Sum(A[k], k).doit() == sum([A[Idx(j, (1, 3))] for j in range(1, 4)])
    assert Product(A[k],
                   k).doit() == prod([A[Idx(j, (1, 3))] for j in range(1, 4)])

    raises(ValueError, lambda: Sum(A[k], (k, 1, 4)))
    raises(ValueError, lambda: Sum(A[k], (k, 0, 3)))
    raises(ValueError, lambda: Sum(A[k], (k, 2, oo)))

    raises(ValueError, lambda: Product(A[k], (k, 1, 4)))
    raises(ValueError, lambda: Product(A[k], (k, 0, 3)))
    raises(ValueError, lambda: Product(A[k], (k, 2, oo)))
예제 #2
0
def test_issue_18604():
    m = symbols("m")
    assert Idx("i", m).name == 'i'
    assert Idx("i", m).lower == 0
    assert Idx("i", m).upper == m - 1
    m = symbols("m", real=False)
    raises(TypeError, lambda: Idx("i", m))
예제 #3
0
def test_Idx_subs():
    i, a, b = symbols('i a b', integer=True)
    assert Idx(i, a).subs(a, b) == Idx(i, b)
    assert Idx(i, a).subs(i, b) == Idx(b, a)

    assert Idx(i).subs(i, 2) == Idx(2)
    assert Idx(i, a).subs(a, 2) == Idx(i, 2)
    assert Idx(i, (a, b)).subs(i, 2) == Idx(2, (a, b))
예제 #4
0
def test_get_indices(c, u_i, tau_i_j):
    """ Check that all indices of an EinsteinTerm are parsed correctly. """

    assert c.get_indices() == []  # Scalar
    assert u_i.get_indices() == [Idx(Symbol("i", integer=True))]  # Vector
    assert tau_i_j.get_indices() == [
        Idx(Symbol("i", integer=True)),
        Idx(Symbol("j", integer=True))
    ]  # Tensor
예제 #5
0
def test_get_expanded(c, u_i, tau_i_j):
    """ Check that the indices of an EinsteinTerm are expanded correctly to produce a new EinsteinTerm object. """

    index_map = [(Idx(Symbol("i", integer=True)), 0),
                 (Idx(Symbol("j", integer=True)), 1)]
    assert c.get_expanded(
        index_map) == c  # Scalar. No indices should be expanded here.
    assert u_i.get_expanded(index_map) == EinsteinTerm("u0")  # Vector.
    assert tau_i_j.get_expanded(index_map) == EinsteinTerm("tau01")  # Tensor.
예제 #6
0
def test_Idx_func_args():
    i, a, b = symbols('i a b', integer=True)
    ii = Idx(i)
    assert ii.func(*ii.args) == ii
    ii = Idx(i, a)
    assert ii.func(*ii.args) == ii
    ii = Idx(i, (a, b))
    assert ii.func(*ii.args) == ii
예제 #7
0
def test_Indexed_properties():
    i, j = symbols('i j', integer=True)
    A = Indexed('A', i, j)
    assert A.rank == 2
    assert A.indices == (i, j)
    assert A.base == IndexedBase('A')
    assert A.ranges == [None, None]
    raises(IndexException, 'A.shape')

    n, m = symbols('n m', integer=True)
    assert Indexed('A', Idx(i, m),
                   Idx(j, n)).ranges == [Tuple(0, m - 1),
                                         Tuple(0, n - 1)]
    assert Indexed('A', Idx(i, m), Idx(j, n)).shape == Tuple(m, n)
    raises(IndexException, 'Indexed("A", Idx(i, m), Idx(j)).shape')
예제 #8
0
    def generate(self, sorder):
        """
        Generate the numerical code.

        Parameters
        ----------
        sorder : list
            the order of nv, nx, ny and nz
        """
        from .generator import For
        from .symbolic import nx, ny, nz, indexed, ix

        ns = int(self.stencil.nv_ptr[-1])
        dim = self.stencil.dim

        istore, iload, ncond = self._get_istore_iload_symb(dim)

        idx = Idx(ix, (0, ncond))
        fstore = indexed('f', [ns, nx, ny, nz],
                         index=[istore[idx, k] for k in range(dim + 1)],
                         priority=sorder)
        fload = indexed('f', [ns, nx, ny, nz],
                        index=[iload[0][idx, k] for k in range(dim + 1)],
                        priority=sorder)

        self.generator.add_routine((self.name, For(idx, Eq(fstore, fload))))
예제 #9
0
    def generate(self, sorder):
        """
        Generate the numerical code.

        Parameters
        ----------
        sorder : list
            the order of nv, nx, ny and nz
        """
        from .generator import For
        from .symbolic import nx, ny, nz, indexed, ix

        ns = int(self.stencil.nv_ptr[-1])
        dim = self.stencil.dim

        istore, iload, ncond = self._get_istore_iload_symb(dim)
        rhs, dist = self._get_rhs_dist_symb(ncond)

        idx = Idx(ix, (0, ncond))
        fstore = indexed('f', [ns, nx, ny, nz],
                         index=[istore[idx, k] for k in range(dim + 1)],
                         priority=sorder)
        fload0 = indexed('f', [ns, nx, ny, nz],
                         index=[iload[0][idx, k] for k in range(dim + 1)],
                         priority=sorder)
        fload1 = indexed('f', [ns, nx, ny, nz],
                         index=[iload[1][idx, k] for k in range(dim + 1)],
                         priority=sorder)

        self.generator.add_routine(
            ('Bouzidi_anti_bounce_back',
             For(
                 idx,
                 Eq(fstore, -dist[idx] * fload0 + (1 - dist[idx]) * fload1 +
                    rhs[idx]))))
예제 #10
0
def test_Idx_inequalities_current_fails():
    i14 = Idx("i14", (1, 4))

    assert S(5) >= i14
    assert S(5) > i14
    assert not (S(5) <= i14)
    assert not (S(5) < i14)
def test_Assignment():
    x, y = symbols("x, y")
    A = MatrixSymbol('A', 3, 1)
    mat = Matrix([1, 2, 3])
    B = IndexedBase('B')
    n = symbols("n", integer=True)
    i = Idx("i", n)
    # Here we just do things to show they don't error
    Assignment(x, y)
    Assignment(x, 0)
    Assignment(A, mat)
    Assignment(A[1, 0], 0)
    Assignment(A[1, 0], x)
    Assignment(B[i], x)
    Assignment(B[i], 0)
    a = Assignment(x, y)
    assert a.func(*a.args) == a
    # Here we test things to show that they error
    # Matrix to scalar
    raises(ValueError, lambda: Assignment(B[i], A))
    raises(ValueError, lambda: Assignment(B[i], mat))
    raises(ValueError, lambda: Assignment(x, mat))
    raises(ValueError, lambda: Assignment(x, A))
    raises(ValueError, lambda: Assignment(A[1, 0], mat))
    # Scalar to matrix
    raises(ValueError, lambda: Assignment(A, x))
    raises(ValueError, lambda: Assignment(A, 0))
    # Non-atomic lhs
    raises(TypeError, lambda: Assignment(mat, A))
    raises(TypeError, lambda: Assignment(0, x))
    raises(TypeError, lambda: Assignment(x * x, 1))
    raises(TypeError, lambda: Assignment(A + A, mat))
    raises(TypeError, lambda: Assignment(B, 0))

    assert Relational(x, y, ':=') == Assignment(x, y)
예제 #12
0
def test_Indexed_coeff():
    N = Symbol("N", integer=True)
    len_y = N
    i = Idx("i", len_y - 1)
    y = IndexedBase("y", shape=(len_y,))
    a = (1 / y[i + 1] * y[i]).coeff(y[i])
    b = (y[i] / y[i + 1]).coeff(y[i])
    assert a == b
예제 #13
0
def test_Indexed_coeff():
    N = Symbol('N', integer=True)
    len_y = N
    i = Idx('i', len_y - 1)
    y = IndexedBase('y', shape=(len_y, ))
    a = (1 / y[i + 1] * y[i]).coeff(y[i])
    b = (y[i] / y[i + 1]).coeff(y[i])
    assert a == b
예제 #14
0
def test_Indexed_shape_precedence():
    i, j = symbols("i j", integer=True)
    o, p = symbols("o p", integer=True)
    n, m = symbols("n m", integer=True)
    a = IndexedBase("a", shape=(o, p))
    assert a.shape == Tuple(o, p)
    assert Indexed(a, Idx(i, m), Idx(j, n)).ranges == [Tuple(0, m - 1), Tuple(0, n - 1)]
    assert Indexed(a, Idx(i, m), Idx(j, n)).shape == Tuple(o, p)
    assert Indexed(a, Idx(i, m), Idx(j)).ranges == [Tuple(0, m - 1), Tuple(None, None)]
    assert Indexed(a, Idx(i, m), Idx(j)).shape == Tuple(o, p)
예제 #15
0
def test_indexed_by_grid(grid):
    """ Ensure that an Indexed object gets correctly indexed by the Grid indices. """

    idx = Idx(Symbol("i", integer=True))
    base = IndexedBase("test")
    i = base[idx]
    
    assert grid.indexed_by_grid(i) == base[grid.indices]

    return 
예제 #16
0
def test_cse_Indexed():
    len_y = 5
    y = IndexedBase('y', shape=(len_y, ))
    x = IndexedBase('x', shape=(len_y, ))
    i = Idx('i', len_y - 1)

    expr1 = (y[i + 1] - y[i]) / (x[i + 1] - x[i])
    expr2 = 1 / (x[i + 1] - x[i])
    replacements, reduced_exprs = cse([expr1, expr2])
    assert len(replacements) > 0
예제 #17
0
파일: test_cse.py 프로젝트: jdoherty7/sumpy
def test_cse_Indexed():  # noqa
    from sympy import IndexedBase, Idx
    len_y = 5
    y = IndexedBase('y', shape=(len_y,))
    x = IndexedBase('x', shape=(len_y,))
    Dy = IndexedBase('Dy', shape=(len_y-1,))  # noqa
    i = Idx('i', len_y-1)

    expr1 = (y[i+1]-y[i])/(x[i+1]-x[i])
    expr2 = 1/(x[i+1]-x[i])
    replacements, reduced_exprs = cse([expr1, expr2])
    assert len(replacements) > 0
예제 #18
0
    def copy(self, cdlt, start=None, stride=None, end=None, offset=None, **kwargs):
        obj = super(Loop, self).copy(cdlt, **kwargs)
        obj._start = start or copy(self.start)
        obj._end = end or copy(self.end)
        obj._stride = stride or copy(self.stride)
        obj._offset = offset or copy(self.offset)
        if obj.op_str not in obj.param_symbols:
            obj_idx = Idx(obj.op_str, (obj._start, obj._end))
            old_idx = obj.param_symbols.pop(self.op_str)
            new_idx = old_idx.subs(old_idx, obj_idx)
            obj.param_symbols[obj.op_str] = new_idx

        return obj
예제 #19
0
def test_Idx_func_args():
    i, a, b = symbols('i a b', integer=True)
    ii = Idx(i)
    assert ii.func(*ii.args) == ii
    ii = Idx(i, a)
    assert ii.func(*ii.args) == ii
    ii = Idx(i, (a, b))
    assert ii.func(*ii.args) == ii
예제 #20
0
def test_Idx_construction():
    i, a, b = symbols('i a b', integer=True)
    assert Idx(i) != Idx(i, 1)
    assert Idx(i, a) == Idx(i, (0, a - 1))
    assert Idx(i, oo) == Idx(i, (0, oo))

    x = symbols('x')
    raises(TypeError, "Idx(x)")
    raises(TypeError, "Idx(0.5)")
    raises(TypeError, "Idx(i, x)")
    raises(TypeError, "Idx(i, 0.5)")
    raises(TypeError, "Idx(i, (x, 5))")
    raises(TypeError, "Idx(i, (2, x))")
    raises(TypeError, "Idx(i, (2, 3.5))")
예제 #21
0
파일: hakaru.py 프로젝트: ccshan/pykaru
 def enter(self, avoid):
     # Rename self.variable in self.after to avoid free_symbols
     var = self._args[1]
     aft = self._args[2]
     if isinstance(self._args[0], Lebesgue):
         x = Dummy(var.name,
                   real = True,
                   positive = self._args[0].lower >= 0 and self._args[0].upper >= 0,
                   negative = self._args[0].lower <= 0 and self._args[0].upper <= 0)
         return (x, aft.subs(var, x))
     elif isinstance(self._args[0], Counting):
         x = Dummy(var.name,
                   integer = True,
                   positive = self._args[0].lower >= 0 and self._args[0].upper >= 0,
                   negative = self._args[0].lower <= 0 and self._args[0].upper <= 0)
         x = Idx(x, self._args[0].lower, self._args[0].upper)
         return (x, aft.subs(var, x))
     elif any(self.variable in object.free_symbols for object in avoid):
         x = Dummy(var.name)
         return (x, aft.subs(var, x))
     else:
         return (var, aft)
예제 #22
0
def test_Indexed_properties():
    i, j = symbols("i j", integer=True)
    A = Indexed("A", i, j)
    assert A.name == "A[i, j]"
    assert A.rank == 2
    assert A.indices == (i, j)
    assert A.base == IndexedBase("A")
    assert A.ranges == [None, None]
    raises(IndexException, lambda: A.shape)

    n, m = symbols("n m", integer=True)
    assert Indexed("A", Idx(i, m), Idx(j, n)).ranges == [
        Tuple(0, m - 1),
        Tuple(0, n - 1),
    ]
    assert Indexed("A", Idx(i, m), Idx(j, n)).shape == Tuple(m, n)
    raises(IndexException, lambda: Indexed("A", Idx(i, m), Idx(j)).shape)
예제 #23
0
def test_Idx_fixed_bounds():
    i, a, b, x = symbols('i a b x', integer=True)
    assert Idx(x).lower is None
    assert Idx(x).upper is None
    assert Idx(x, a).lower == 0
    assert Idx(x, a).upper == a - 1
    assert Idx(x, 5).lower == 0
    assert Idx(x, 5).upper == 4
    assert Idx(x, oo).lower == 0
    assert Idx(x, oo).upper is oo
    assert Idx(x, (a, b)).lower == a
    assert Idx(x, (a, b)).upper == b
    assert Idx(x, (1, 5)).lower == 1
    assert Idx(x, (1, 5)).upper == 5
    assert Idx(x, (-oo, oo)).lower is -oo
    assert Idx(x, (-oo, oo)).upper is oo
예제 #24
0
def test_issue_12780():
    n = symbols("n")
    i = Idx("i", (0, n))
    raises(TypeError, lambda: i.subs(n, 1.5))
예제 #25
0
def test_Idx_bounds():
    i, a, b = symbols('i a b', integer=True)
    assert Idx(i).lower is None
    assert Idx(i).upper is None
    assert Idx(i, a).lower == 0
    assert Idx(i, a).upper == a - 1
    assert Idx(i, 5).lower == 0
    assert Idx(i, 5).upper == 4
    assert Idx(i, oo).lower == 0
    assert Idx(i, oo).upper is oo
    assert Idx(i, (a, b)).lower == a
    assert Idx(i, (a, b)).upper == b
    assert Idx(i, (1, 5)).lower == 1
    assert Idx(i, (1, 5)).upper == 5
    assert Idx(i, (-oo, oo)).lower is -oo
    assert Idx(i, (-oo, oo)).upper is oo
예제 #26
0
def test_Idx_properties():
    i, a, b = symbols('i a b', integer=True)
    assert Idx(i).is_integer
    assert Idx(i).name == 'i'
    assert Idx(i + 2).name == 'i + 2'
    assert Idx('foo').name == 'foo'
예제 #27
0
def test_Idx_inequalities():
    i14 = Idx("i14", (1, 4))
    i79 = Idx("i79", (7, 9))
    i46 = Idx("i46", (4, 6))
    i35 = Idx("i35", (3, 5))

    assert i14 <= 5
    assert i14 < 5
    assert not (i14 >= 5)
    assert not (i14 > 5)

    assert 5 >= i14
    assert 5 > i14
    assert not (5 <= i14)
    assert not (5 < i14)

    assert LessThan(i14, 5)
    assert StrictLessThan(i14, 5)
    assert not GreaterThan(i14, 5)
    assert not StrictGreaterThan(i14, 5)

    assert i14 <= 4
    assert isinstance(i14 < 4, StrictLessThan)
    assert isinstance(i14 >= 4, GreaterThan)
    assert not (i14 > 4)

    assert isinstance(i14 <= 1, LessThan)
    assert not (i14 < 1)
    assert i14 >= 1
    assert isinstance(i14 > 1, StrictGreaterThan)

    assert not (i14 <= 0)
    assert not (i14 < 0)
    assert i14 >= 0
    assert i14 > 0

    from sympy.abc import x

    assert isinstance(i14 < x, StrictLessThan)
    assert isinstance(i14 > x, StrictGreaterThan)
    assert isinstance(i14 <= x, LessThan)
    assert isinstance(i14 >= x, GreaterThan)

    assert i14 < i79
    assert i14 <= i79
    assert not (i14 > i79)
    assert not (i14 >= i79)

    assert i14 <= i46
    assert isinstance(i14 < i46, StrictLessThan)
    assert isinstance(i14 >= i46, GreaterThan)
    assert not (i14 > i46)

    assert isinstance(i14 < i35, StrictLessThan)
    assert isinstance(i14 > i35, StrictGreaterThan)
    assert isinstance(i14 <= i35, LessThan)
    assert isinstance(i14 >= i35, GreaterThan)

    iNone1 = Idx("iNone1")
    iNone2 = Idx("iNone2")

    assert isinstance(iNone1 < iNone2, StrictLessThan)
    assert isinstance(iNone1 > iNone2, StrictGreaterThan)
    assert isinstance(iNone1 <= iNone2, LessThan)
    assert isinstance(iNone1 >= iNone2, GreaterThan)
예제 #28
0
from sympy.codegen.ast import (
    Assignment, Attribute, aug_assign, CodeBlock, For, Type, Variable, Pointer,
    Declaration, AddAugmentedAssignment, SubAugmentedAssignment,
    MulAugmentedAssignment, DivAugmentedAssignment, ModAugmentedAssignment,
    value_const, pointer_const, integer, real, complex_, int8, uint8, float16
    as f16, float32 as f32, float64 as f64, float80 as f80, float128 as f128,
    complex64 as c64, complex128 as c128, While, Scope, String, Print,
    QuotedString, FunctionPrototype, FunctionDefinition, Return, FunctionCall,
    untyped, IntBaseType, intc, Node, none, NoneToken, Token, Comment)

x, y, z, t, x0, x1, x2, a, b = symbols("x, y, z, t, x0, x1, x2, a, b")
n = symbols("n", integer=True)
A = MatrixSymbol('A', 3, 1)
mat = Matrix([1, 2, 3])
B = IndexedBase('B')
i = Idx("i", n)
A22 = MatrixSymbol('A22', 2, 2)
B22 = MatrixSymbol('B22', 2, 2)


def test_Assignment():
    # Here we just do things to show they don't error
    Assignment(x, y)
    Assignment(x, 0)
    Assignment(A, mat)
    Assignment(A[1, 0], 0)
    Assignment(A[1, 0], x)
    Assignment(B[i], x)
    Assignment(B[i], 0)
    a = Assignment(x, y)
    assert a.func(*a.args) == a
예제 #29
0
파일: opt.py 프로젝트: aleleo97/proxy_nav
#! /usr/bin/env python
import cvxpy as cp
import numpy as np
from sympy import symbols, IndexedBase, Idx
#let's define the variables of the class (u inputs and x states)
u = IndexedBase('u')
n_in = symbols('n_in ', integer=True)
u[n_in]
#you can change the number of input but not the name
n_in = Idx('n_in', 2)
x = IndexedBase('x')
n_states = symbols('n_states', integer=True)
x[n_states]
#You can change the number of states not the name
n_states = Idx('n_states', 3)


class ConvexOpt():
    def __init__(self, N, x_init, x_fin, u_in, A_list, B_list, C_list):
        #init the variables of the class
        self.N = N
        self.x_init = x_init
        self.x_fin = x_fin
        self.u_in = u_in
        self.Ad_list = A_list
        self.Bd_list = B_list
        self.Cd_list = C_list

    def CVXOPT(self, opt_power=False, opt_velocity=False):
        #save the number of states and inputs
        x_len = (int)(x[n_states].shape[0])
예제 #30
0
def split_transfer(cdlt: 'Codelet', outer_xfer: 'Transfer',
                   inner_xfer: 'Transfer'):
    full_path = outer_xfer.path.copy()
    all_transfers = outer_xfer.transfers.copy()

    outer_xfer.path = full_path[:2]
    inner_xfer.path = full_path[1:]

    outer_xfer_key = tuple(full_path[:2])

    outer_xfer.transfers = {outer_xfer_key: all_transfers[outer_xfer_key]}
    inner_xfer.transfers.pop(outer_xfer_key)

    # Update dependencies
    new_inner_deps = []
    dep_map = {}
    dep_symbols = {}
    if inner_xfer.loop_level > outer_xfer.loop_level:

        for d in inner_xfer.dependencies:
            dep_op = cdlt.op_map[d]
            for level, name in dep_op.split_map.items():
                dep_map[name] = d
                dep_symbols[d] = Idx(name, (dep_op.start, dep_op.end))
                new_inner_deps.append(name)
        inner_xfer.dependencies = new_inner_deps

        new_offset = []

        for o in outer_xfer.transfers[outer_xfer_key]._src_offset:
            if isinstance(o, Basic):
                sym_map = {i: dep_symbols[str(i)] for i in list(o.atoms(Idx))}
                new_offset.append(o.subs(sym_map))

        inner_xfer.transfers[tuple(full_path[1:3])]._src_offset = new_offset
        outer_xfer.transfers[outer_xfer_key].compute_src_size(cdlt)
        for _, v in inner_xfer.transfers.items():
            v.compute_src_size(cdlt)
    else:
        for d in outer_xfer.dependencies:
            dep_op = cdlt.op_map[d]
            if dep_op.op_type == "compute":
                new_inner_deps.append(d)
                inner_xfer.dependencies.remove(d)
            else:
                for level, name in dep_op.split_map.items():
                    dep_map[name] = d
                    new_inner_deps.append(name)
                    if dep_op.op_type == "loop":
                        dep_symbols[d] = Idx(name, (dep_op.start, dep_op.end))

        outer_xfer.dependencies = new_inner_deps

        for path, xfer in inner_xfer.transfers.items():

            new_offset = []
            for o in xfer._dst_offset:
                if isinstance(o, Basic):
                    sym_map = {
                        i: dep_symbols[str(i)]
                        for i in list(o.atoms(Idx))
                    }
                    new_offset.append(o.subs(sym_map))
                outer_xfer.transfers[outer_xfer_key]._dst_offset = new_offset
            xfer.compute_dst_size(cdlt)

        outer_xfer.transfers[outer_xfer_key].compute_dst_size(cdlt)

    return inner_xfer
예제 #31
0
def test_issue_12780():
    n = symbols("n")
    i = Idx("i", (0, n))
    raises(TypeError, lambda: i.subs(n, 1.5))
예제 #32
0
def main():

    print(__doc__)

    # arrays are represented with IndexedBase, indices with Idx
    m = Symbol("m", integer=True)
    i = Idx("i", m)
    A = IndexedBase("A")
    B = IndexedBase("B")
    x = Symbol("x")

    print("Compiling ufuncs for radial harmonic oscillator solutions")

    # setup a basis of ho-solutions  (for l=0)
    basis_ho = {}
    for n in range(basis_dimension):

        # Setup the radial ho solution for this n
        expr = R_nl(n, orbital_momentum_l, omega2, x)

        # Reduce the number of operations in the expression by eval to float
        expr = expr.evalf(15)

        print("The h.o. wave function with l = %i and n = %i is" %
              (orbital_momentum_l, n))
        pprint(expr)

        # implement, compile and wrap it as a ufunc
        basis_ho[n] = ufuncify(x, expr)

    # now let's see if we can express a hydrogen radial wave in terms of
    # the ho basis.  Here's the solution we will approximate:
    H_ufunc = ufuncify(x, hydro_nl(hydrogen_n, orbital_momentum_l, 1, x))

    # The transformation to a different basis can be written like this,
    #
    #   psi(r) = sum_i c(i) phi_i(r)
    #
    # where psi(r) is the hydrogen solution, phi_i(r) are the H.O. solutions
    # and c(i) are scalar coefficients.
    #
    # So in order to express a hydrogen solution in terms of the H.O. basis, we
    # need to determine the coefficients c(i).  In position space, it means
    # that we need to evaluate an integral:
    #
    #  psi(r) = sum_i Integral(R**2*conj(phi(R))*psi(R), (R, 0, oo)) phi_i(r)
    #
    # To calculate the integral with autowrap, we notice that it contains an
    # element-wise sum over all vectors.  Using the Indexed class, it is
    # possible to generate autowrapped functions that perform summations in
    # the low-level code.  (In fact, summations are very easy to create, and as
    # we will see it is often necessary to take extra steps in order to avoid
    # them.)
    # we need one integration ufunc for each wave function in the h.o. basis
    binary_integrator = {}
    for n in range(basis_dimension):

        #
        # setup basis wave functions
        #
        # To get inline expressions in the low level code, we attach the
        # wave function expressions to a regular SymPy function using the
        # implemented_function utility.  This is an extra step needed to avoid
        # erroneous summations in the wave function expressions.
        #
        # Such function objects carry around the expression they represent,
        # but the expression is not exposed unless explicit measures are taken.
        # The benefit is that the routines that searches for repeated indices
        # in order to make contractions will not search through the wave
        # function expression.
        psi_ho = implemented_function(
            "psi_ho", Lambda(x, R_nl(n, orbital_momentum_l, omega2, x)))

        # We represent the hydrogen function by an array which will be an input
        # argument to the binary routine.  This will let the integrators find
        # h.o. basis coefficients for any wave function we throw at them.
        psi = IndexedBase("psi")

        #
        # setup expression for the integration
        #

        step = Symbol("step")  # use symbolic stepsize for flexibility

        # let i represent an index of the grid array, and let A represent the
        # grid array.  Then we can approximate the integral by a sum over the
        # following expression (simplified rectangular rule, ignoring end point
        # corrections):

        expr = A[i]**2 * psi_ho(A[i]) * psi[i] * step

        if n == 0:
            print("Setting up binary integrators for the integral:")
            pprint(Integral(x**2 * psi_ho(x) * Function("psi")(x), (x, 0, oo)))

        # Autowrap it.  For functions that take more than one argument, it is
        # a good idea to use the 'args' keyword so that you know the signature
        # of the wrapped function.  (The dimension m will be an optional
        # argument, but it must be present in the args list.)
        binary_integrator[n] = autowrap(expr,
                                        args=[A.label, psi.label, step, m])

        # Lets see how it converges with the grid dimension
        print("Checking convergence of integrator for n = %i" % n)
        for g in range(3, 8):
            grid, step = np.linspace(0, rmax, 2**g, retstep=True)
            print("grid dimension %5i, integral = %e" %
                  (2**g, binary_integrator[n](grid, H_ufunc(grid), step)))

    print("A binary integrator has been set up for each basis state")
    print("We will now use them to reconstruct a hydrogen solution.")

    # Note: We didn't need to specify grid or use gridsize before now
    grid, stepsize = np.linspace(0, rmax, gridsize, retstep=True)

    print("Calculating coefficients with gridsize = %i and stepsize %f" %
          (len(grid), stepsize))

    coeffs = {}
    for n in range(basis_dimension):
        coeffs[n] = binary_integrator[n](grid, H_ufunc(grid), stepsize)
        print("c(%i) = %e" % (n, coeffs[n]))

    print("Constructing the approximate hydrogen wave")
    hydro_approx = 0
    all_steps = {}
    for n in range(basis_dimension):
        hydro_approx += basis_ho[n](grid) * coeffs[n]
        all_steps[n] = hydro_approx.copy()
        if pylab:
            line = pylab.plot(grid, all_steps[n], ":", label="max n = %i" % n)

    # check error numerically
    diff = np.max(np.abs(hydro_approx - H_ufunc(grid)))
    print("Error estimate: the element with largest deviation misses by %f" %
          diff)
    if diff > 0.01:
        print("This is much, try to increase the basis size or adjust omega")
    else:
        print("Ah, that's a pretty good approximation!")

    # Check visually
    if pylab:
        print("Here's a plot showing the contribution for each n")
        line[0].set_linestyle("-")
        pylab.plot(grid, H_ufunc(grid), "r-", label="exact")
        pylab.legend()
        pylab.show()

    print("""Note:
    These binary integrators were specialized to find coefficients for a
    harmonic oscillator basis, but they can process any wave function as long
    as it is available as a vector and defined on a grid with equidistant
    points. That is, on any grid you get from numpy.linspace.

    To make the integrators even more flexible, you can setup the harmonic
    oscillator solutions with symbolic parameters omega and l.  Then the
    autowrapped binary routine will take these scalar variables as arguments,
    so that the integrators can find coefficients for *any* isotropic harmonic
    oscillator basis.

    """)