text:
from __future__ import absolute_import

from collections import OrderedDict

from conftest import EVAL, dims, time, x, y, z, skipif_yask

import numpy as np
import pytest

from devito import (clear_cache, Grid, Eq, Operator, Constant, Function,
                    TimeFunction, SparseFunction, Dimension, configuration)
from devito.foreign import Operator as OperatorForeign
from devito.dle import retrieve_iteration_tree
from devito.ir.iet import IsPerfectIteration

def dimify(dimensions):
    assert isinstance(dimensions, str)
    return tuple(dims()[i] for i in dimensions.split())

def symbol(name, dimensions, value=0., shape=(3, 5), mode='function'):
    """Short-cut for symbol creation to test "function"
    and "indexed" API."""
    assert(mode in ['function', 'indexed'])
    s = Function(name=name, dimensions=dimensions, shape=shape)
    s.data[:] = value
    return s.indexify() if mode == 'indexed' else s

@skipif_yask
class TestAPI(object):

    @classmethod
    def setup_class(cls):
        clear_cache()

    def test_code_generation(self, const, a_dense):
        """
        Tests that we can actually generate code for a trivial operator
        using constant and array data objects.
        """
        eqn = Eq(a_dense, a_dense + 2.*const)
        op = Operator(eqn)
        assert len(op.parameters) == 6
        assert op.parameters[0].name == 'a_dense'
        assert op.parameters[0].is_TensorArgument
        assert op.parameters[1].name == 'constant'
        assert op.parameters[1].is_ScalarArgument
        assert op.parameters[2].name == 'i_size'
        assert op.parameters[2].is_ScalarArgument
        assert op.parameters[3].name == 'i_s'
        assert op.parameters[3].is_ScalarArgument
        assert op.parameters[4].name == 'i_e'
        assert op.parameters[4].is_ScalarArgument
        assert op.parameters[5].name == 'timings'
        assert op.parameters[5].is_PtrArgument
        assert 'a_dense[i] = 2.0F*constant + a_dense[i]' in str(op.ccode)

@skipif_yask
class TestArithmetic(object):

    @classmethod
    def setup_class(cls):
        clear_cache()

    @pytest.mark.parametrize('expr, result', [
        ('Eq(a, a + b + 5.)', 10.),
        ('Eq(a, b - a)', 1.),
        ('Eq(a, 4 * (b * a))', 24.),
        ('Eq(a, (6. / b) + (8. * a))', 18.),
    ])
    @pytest.mark.parametrize('mode', ['function'])
    def test_flat(self, expr, result, mode):
        """Tests basic point-wise arithmetic on two-dimensional data"""
        i, j = dimify('i j')
        a = symbol(name='a', dimensions=(i, j), value=2., mode=mode)
        b = symbol(name='b', dimensions=(i, j), value=3., mode=mode)
        fa = a.base.function if mode == 'indexed' else a
        fb = b.base.function if mode == 'indexed' else b

        eqn = eval(expr)
        Operator(eqn)(a=fa, b=fb)
        assert np.allclose(fa.data, result, rtol=1e-12)

    @pytest.mark.parametrize('expr, result', [
        ('Eq(a, a + b + 5.)', 10.),
        ('Eq(a, b - a)', 1.),
        ('Eq(a, 4 * (b * a))', 24.),
        ('Eq(a, (6. / b) + (8. * a))', 18.),
    ])
    @pytest.mark.parametrize('mode', ['function', 'indexed'])
    def test_deep(self, expr, result, mode):
        """Tests basic point-wise arithmetic on multi-dimensional data"""
        i, j, k, l = dimify('i j k l')
        a = symbol(name='a', dimensions=(i, j, k, l), shape=(3, 5, 7, 6),
                   value=2., mode=mode)
        b = symbol(name='b', dimensions=(j, k), shape=(5, 7),
                   value=3., mode=mode)
        fa = a.base.function if mode == 'indexed' else a
        fb = b.base.function if mode == 'indexed' else b

        eqn = eval(expr)
        Operator(eqn)(a=fa, b=fb)
        assert np.allclose(fa.data, result, rtol=1e-12)

    @pytest.mark.parametrize('expr, result', [
        ('Eq(a[j, l], a[j - 1 , l] + 1.)',
         np.meshgrid(np.arange(2., 8.), np.arange(2., 7.))[1]),
        ('Eq(a[j, l], a[j, l - 1] + 1.)',
         np.meshgrid(np.arange(2., 8.), np.arange(2., 7.))[0]),
    ])
    def test_indexed_increment(self, expr, result):
        """Tests point-wise increments with stencil offsets in one dimension"""
        j, l = dimify('j l')
        a = symbol(name='a', dimensions=(j, l), value=2., shape=(5, 6),
                   mode='indexed').base
        fa = a.function
        fa.data[1:, 1:] = 0

        eqn = eval(expr)
        Operator(eqn)(a=fa)
        assert np.allclose(fa.data, result, rtol=1e-12)

    @pytest.mark.parametrize('expr, result', [
        ('Eq(a[j, l], b[j - 1 , l] + 1.)', np.zeros((5, 6)) + 3.),
        ('Eq(a[j, l], b[j , l - 1] + 1.)', np.zeros((5, 6)) + 3.),
        ('Eq(a[j, l], b[j - 1, l - 1] + 1.)', np.zeros((5, 6)) + 3.),
        ('Eq(a[j, l], b[j + 1, l + 1] + 1.)', np.zeros((5, 6)) + 3.),
    ])
    def test_indexed_stencil(self, expr, result):
        """Test point-wise arithmetic with stencil offsets across two
        functions in indexed expression format"""
        j, l = dimify('j l')
        a = symbol(name='a', dimensions=(j, l), value=0., shape=(5, 6),
                   mode='indexed').base
        fa = a.function
        b = symbol(name='b', dimensions=(j, l), value=2., shape=(5, 6),
                   mode='indexed').base
        fb = b.function

        eqn = eval(expr)
        Operator(eqn)(a=fa, b=fb)
        assert np.allclose(fa.data[1:-1, 1:-1], result[1:-1, 1:-1], rtol=1e-12)

    @pytest.mark.parametrize('expr, result', [
        ('Eq(a[1, j, l], a[0, j - 1 , l] + 1.)', np.zeros((5, 6)) + 3.),
        ('Eq(a[1, j, l], a[0, j , l - 1] + 1.)', np.zeros((5, 6)) + 3.),
        ('Eq(a[1, j, l], a[0, j - 1, l - 1] + 1.)', np.zeros((5, 6)) + 3.),
        ('Eq(a[1, j, l], a[0, j + 1, l + 1] + 1.)', np.zeros((5, 6)) + 3.),
    ])
    def test_indexed_buffered(self, expr, result):
        """Test point-wise arithmetic with stencil offsets across a single
        functions with buffering dimension in indexed expression format"""
        i, j, l = dimify('i j l')
        a = symbol(name='a', dimensions=(i, j, l), value=2., shape=(3, 5, 6),
                   mode='indexed').base
        fa = a.function

        eqn = eval(expr)
        Operator(eqn)(a=fa)
        assert np.allclose(fa.data[1, 1:-1, 1:-1], result[1:-1, 1:-1], rtol=1e-12)

    @pytest.mark.parametrize('expr, result', [
        ('Eq(a[1, j, l], a[0, j - 1 , l] + 1.)', np.zeros((5, 6)) + 3.),
    ])
    def test_indexed_open_loops(self, expr, result):
        """Test point-wise arithmetic with stencil offsets and open loop
        boundaries in indexed expression format"""
        i, j, l = dimify('i j l')
        a = Function(name='a', dimensions=(i, j, l), shape=(3, 5, 6)).indexed
        fa = a.function
        fa.data[0, :, :] = 2.

        eqn = eval(expr)
        Operator(eqn)(a=fa)
        assert np.allclose(fa.data[1, 1:-1, 1:-1], result[1:-1, 1:-1], rtol=1e-12)

    def test_constant_time_dense(self):
        """Test arithmetic between different data objects, namely Constant
        and Function."""
        i, j = dimify('i j')
        const = Constant(name='truc', value=2.)
        a = Function(name='a', shape=(20, 20), dimensions=(i, j))
        a.data[:] = 2.
        eqn = Eq(a, a + 2.*const)
        op = Operator(eqn)
        op.apply(a=a, truc=const)
        assert(np.allclose(a.data, 6.))

        # Applying a different constant still works
        op.apply(a=a, truc=Constant(name='truc2', value=3.))
        assert(np.allclose(a.data, 12.))

@skipif_yask
class TestAllocation(object):

    @classmethod
    def setup_class(cls):
        clear_cache()

    @pytest.mark.parametrize('shape', [(20, 20),
                                       (20, 20, 20),
                                       (20, 20, 20, 20)])
    def test_first_touch(self, shape):
        dimensions = dimify('i j k l')[:len(shape)]
        grid = Grid(shape=shape, dimensions=dimensions)
        m = Function(name='m', grid=grid, first_touch=True)
        assert(np.allclose(m.data, 0))
        m2 = Function(name='m2', grid=grid, first_touch=False)
        assert(np.allclose(m2.data, 0))
        assert(np.array_equal(m.data, m2.data))

    @pytest.mark.parametrize('staggered', [
        (0, 0), (0, 1), (1, 0), (1, 1),
        (0, 0, 0), (1, 0, 0), (0, 1, 0), (0, 0, 1),
        (1, 1, 0), (1, 0, 1), (0, 1, 1), (1, 1, 1),
    ])
    def test_staggered(self, staggered):
        """
        Test the "deformed" allocation for staggered functions
        """
        grid = Grid(shape=tuple(11 for _ in staggered))
        f = Function(name='f', grid=grid, staggered=staggered)
        assert f.data.shape == tuple(11-i for i in staggered)
        # Add a non-staggered field to ensure that the auto-derived
        # dimension size arguments are at maximum
        g = Function(name='g', grid=grid)
        # Test insertion into a central point
        index = tuple(5 for _ in staggered)
        set_f = Eq(f.indexed[index], 2.)
        set_g = Eq(g.indexed[index], 3.)
        Operator([set_f, set_g])()
        assert f.data[index] == 2.

    @pytest.mark.parametrize('staggered', [
        (0, 0, 0), (0, 0, 1), (0, 1, 0), (0, 1, 1),
        (0, 0, 0), (0, 1, 0, 0), (0, 0, 1, 0), (0, 0, 0, 1),
        (0, 1, 1, 0), (0, 1, 0, 1), (0, 0, 1, 1), (0, 1, 1, 1),
    ])
    def test_staggered_time(self, staggered):
        """
        Test the "deformed" allocation for staggered functions
        """
        grid = Grid(shape=tuple(11 for _ in staggered[1:]))
        f = TimeFunction(name='f', grid=grid, staggered=staggered)
        # from IPython import embed; embed()
        assert f.data.shape[1:] == tuple(11-i for i in staggered[1:])
        # Add a non-staggered field to ensure that the auto-derived
        # dimension size arguments are at maximum
        g = TimeFunction(name='g', grid=grid)
        # Test insertion into a central point
        index = tuple([0] + [5 for _ in staggered[1:]])
        set_f = Eq(f.indexed[index], 2.)
        set_g = Eq(g.indexed[index], 3.)
        Operator([set_f, set_g])()
        assert f.data[index] == 2.

@skipif_yask
class TestArguments(object):

    @classmethod
    def setup_class(cls):
        clear_cache()

    def test_override_cache_aliasing(self):
        """Test that the call-time overriding of Operator arguments works"""
        i, j, k, l = dimify('i j k l')
        shape = (3, 5, 7, 6)
        a = symbol(name='a', dimensions=(i, j, k, l), value=2.,
                   shape=shape, mode='indexed').base.function
        a1 = symbol(name='a', dimensions=(i, j, k, l), value=3.,
                    shape=shape, mode='indexed').base.function
        a2 = symbol(name='a', dimensions=(i, j, k, l), value=4.,
                    shape=shape, mode='indexed').base.function
        eqn = Eq(a, a+3)
        op = Operator(eqn)
        op()
        op(a=a1)
        op(a=a2)

        assert(np.allclose(a.data, np.zeros(shape) + 5))
        assert(np.allclose(a1.data, np.zeros(shape) + 6))
        assert(np.allclose(a2.data, np.zeros(shape) + 7))

    def test_override_symbol(self):
        """Test call-time symbols overrides with other symbols"""
        i, j, k, l = dimify('i j k l')
        shape = (3, 5, 7, 6)
        a = symbol(name='a', dimensions=(i, j, k, l), shape=shape, value=2.)
        a1 = symbol(name='a1', dimensions=(i, j, k, l), shape=shape, value=3.)
        a2 = symbol(name='a2', dimensions=(i, j, k, l), shape=shape, value=4.)
        op = Operator(Eq(a, a + 3))
        op()
        op(a=a1)
        op(a=a2)

        assert(np.allclose(a.data, np.zeros(shape) + 5))
        assert(np.allclose(a1.data, np.zeros(shape) + 6))
        assert(np.allclose(a2.data, np.zeros(shape) + 7))

    def test_override_array(self):
        """Test call-time symbols overrides with numpy arrays"""
        i, j, k, l = dimify('i j k l')
        shape = (3, 5, 7, 6)
        a = symbol(name='a', shape=shape, dimensions=(i, j, k, l), value=2.)
        a1 = np.zeros(shape=shape, dtype=np.float32) + 3.
        a2 = np.zeros(shape=shape, dtype=np.float32) + 4.
        op = Operator(Eq(a, a + 3))
        op()
        op(a=a1)
        op(a=a2)

        assert(np.allclose(a.data, np.zeros(shape) + 5))
        assert(np.allclose(a1, np.zeros(shape) + 6))
        assert(np.allclose(a2, np.zeros(shape) + 7))

    def test_dimension_size_infer(self, nt=100):
        """Test that the dimension sizes are being inferred correctly"""
        grid = Grid(shape=(3, 5, 7))
        a = Function(name='a', grid=grid)
        b = TimeFunction(name='b', grid=grid, save=True, time_dim=nt)
        op = Operator(Eq(b, a))

        time = b.indices[0]
        op_arguments, _ = op.arguments()
        assert(op_arguments[time.start_name] == 0)
        assert(op_arguments[time.end_name] == nt)

    def test_arg_offset_adjust(self, nt=100):
        """Test that the dimension sizes are being inferred correctly"""
        i, j, k = dimify('i j k')
        shape = (10, 10, 10)
        grid = Grid(shape=shape, dimensions=(i, j, k))
        a = Function(name='a', grid=grid).indexed
        b = TimeFunction(name='b', grid=grid, save=True, time_dim=nt)
        time = b.indices[0]
        eqn = Eq(b.indexed[time + 1, i, j, k], b.indexed[time - 1, i, j, k]
                 + b.indexed[time, i, j, k] + a[i, j, k])
        op = Operator(eqn)
        args = {time.end_name: nt-10}
        op_arguments, _ = op.arguments(**args)
        assert(op_arguments[time.start_name] == 0)
        assert(op_arguments[time.end_name] == nt - 8)

    def test_dimension_offset_adjust(self, nt=100):
        """Test that the dimension sizes are being inferred correctly"""
        i, j, k = dimify('i j k')
        shape = (10, 10, 10)
        grid = Grid(shape=shape, dimensions=(i, j, k))
        a = Function(name='a', grid=grid).indexed
        b = TimeFunction(name='b', grid=grid, save=True, time_dim=nt)
        time = b.indices[0]
        eqn = Eq(b.indexed[time + 1, i, j, k], b.indexed[time - 1, i, j, k]
                 + b.indexed[time, i, j, k] + a[i, j, k])
        op = Operator(eqn)
        op_arguments, _ = op.arguments(time=nt-10)
        assert(op_arguments[time.start_name] == 0)
        assert(op_arguments[time.end_name] == nt - 8)

    def test_dimension_size_override(self):
        """Test explicit overrides for the leading time dimension"""
        grid = Grid(shape=(3, 5, 7))
        a = TimeFunction(name='a', grid=grid)
        one = Function(name='one', grid=grid)
        one.data[:] = 1.
        op = Operator(Eq(a.forward, a + one))

        # Test dimension override via the buffered dimenions
        a.data[0] = 0.
        op(a=a, t=6)
        assert(np.allclose(a.data[1], 5.))

        # Test dimension override via the parent dimenions
        a.data[0] = 0.
        op(a=a, time=5)
        assert(np.allclose(a.data[0], 4.))

    def test_override_composite_data(self):
        i, j = dimify('i j')
        grid = Grid(shape=(10, 10), dimensions=(i, j))
        original_coords = (1., 1.)
        new_coords = (2., 2.)
        p_dim = Dimension('p_src')
        u = TimeFunction(name='u', grid=grid, time_order=2, space_order=2)
        src1 = SparseFunction(name='src1', grid=grid, dimensions=[time, p_dim],
                              npoint=1, nt=10, coordinates=original_coords)
        src2 = SparseFunction(name='src1', grid=grid, dimensions=[time, p_dim],
                              npoint=1, nt=10, coordinates=new_coords)
        op = Operator(src1.inject(u, src1))

        # Move the source from the location where the setup put it so we can test
        # whether the override picks up the original coordinates or the changed ones

        # Operator.arguments() returns a tuple of (data, dimension_sizes)
        args = op.arguments(src1=src2)[0]
        arg_name = src1.name + "_coords"
        assert(np.array_equal(args[arg_name], np.asarray((new_coords,))))

    def test_start_end_2d(self):
        """ In a 2D square domain, ask the operator to operate over a smaller square
        and then ensure that it only operated on the requested region
        """
        i, j = dimify('i j')
        shape = (10, 10)
        start = 3
        end = 6
        a = Function(name='a', shape=shape, dimensions=(i, j))
        eqn = Eq(a, a + 1)
        op = Operator(eqn)
        op(i_s=start, i_e=end, j_s=start, j_e=end)
        mask = np.ones(shape, np.bool)
        mask[start:end, start:end] = 0
        assert(np.allclose(a.data[start:end, start:end], 1))
        assert(np.allclose(a.data[mask], 0))

    def test_start_end_3d(self):
        """ In a 3D cubical domain, ask the operator to operate over a smaller cube
        and then ensure that it only operated on the requested region
        """
        i, j, k = dimify('i j k')
        shape = (10, 10, 10)
        start = 3
        end = 6
        a = Function(name='a', shape=shape, dimensions=(i, j, k))
        eqn = Eq(a, a + 1)
        op = Operator(eqn)
        op(i_s=start, i_e=end, j_s=start, j_e=end, k_s=start, k_e=end)
        mask = np.ones(shape, np.bool)
        mask[start:end, start:end, start:end] = 0
        assert(np.allclose(a.data[start:end, start:end, start:end], 1))
        assert(np.allclose(a.data[mask], 0))

    def test_argument_derivation_order(self, nt=100):
        """ Ensure the precedence order of arguments is respected
        Defaults < (overriden by) Tensor Arguments < Dimensions < Scalar Arguments
        """
        i, j, k = dimify('i j k')
        shape = (10, 10, 10)
        grid = Grid(shape=shape, dimensions=(i, j, k))
        a = Function(name='a', grid=grid).indexed
        b_function = TimeFunction(name='b', grid=grid, save=True, time_dim=nt)
        b = b_function.indexed
        time = b_function.indices[0]
        b1 = TimeFunction(name='b1', grid=grid, save=True, time_dim=nt+1).indexed
        eqn = Eq(b[time, i, j, k], a[i, j, k])
        op = Operator(eqn)

        # Simple case, same as that tested above.
        # Repeated here for clarity of further tests.
        op_arguments, _ = op.arguments()
        assert(op_arguments[time.start_name] == 0)
        assert(op_arguments[time.end_name] == nt)

        # Providing a tensor argument should infer the dimension size from its shape
        op_arguments, _ = op.arguments(b=b1)
        assert(op_arguments[time.start_name] == 0)
        assert(op_arguments[time.end_name] == nt + 1)

        # Providing a dimension size explicitly should override the automatically inferred
        op_arguments, _ = op.arguments(b=b1, time=nt - 1)
        assert(op_arguments[time.start_name] == 0)
        assert(op_arguments[time.end_name] == nt - 1)

        # Providing a scalar argument explicitly should override the automatically\
        # inferred
        op_arguments, _ = op.arguments(b=b1, time=nt - 1, time_e=nt - 2)
        assert(op_arguments[time.start_name] == 0)
        assert(op_arguments[time.end_name] == nt - 2)

@skipif_yask
class TestDeclarator(object):

    @classmethod
    def setup_class(cls):
        clear_cache()

    def test_heap_1D_stencil(self, a, b):
        operator = Operator(Eq(a, a + b + 5.), dse='noop', dle=None)
        assert """\
  float (*a);
  posix_memalign((void**)&a, 64, sizeof(float[i_size]));
  struct timeval start_section_0, end_section_0;
  gettimeofday(&start_section_0, NULL);
  for (int i = i_s; i < i_e; i += 1)
  {
    a[i] = a[i] + b[i] + 5.0F;
  }
  gettimeofday(&end_section_0, NULL);
  timings->section_0 += (double)(end_section_0.tv_sec-start_section_0.tv_sec)\
+(double)(end_section_0.tv_usec-start_section_0.tv_usec)/1000000;
  free(a);
  return 0;""" in str(operator.ccode)

    def test_heap_perfect_2D_stencil(self, a, c):
        operator = Operator([Eq(a, c), Eq(c, c*a)], dse='noop', dle=None)
        assert """\
  float (*a);
  float (*c)[j_size];
  posix_memalign((void**)&a, 64, sizeof(float[i_size]));
  posix_memalign((void**)&c, 64, sizeof(float[i_size][j_size]));
  struct timeval start_section_0, end_section_0;
  gettimeofday(&start_section_0, NULL);
  for (int i = i_s; i < i_e; i += 1)
  {
    for (int j = j_s; j < j_e; j += 1)
    {
      a[i] = c[i][j];
      c[i][j] = a[i]*c[i][j];
    }
  }
  gettimeofday(&end_section_0, NULL);
  timings->section_0 += (double)(end_section_0.tv_sec-start_section_0.tv_sec)\
+(double)(end_section_0.tv_usec-start_section_0.tv_usec)/1000000;
  free(a);
  free(c);
  return 0;""" in str(operator.ccode)

    def test_heap_imperfect_2D_stencil(self, a, c):
        operator = Operator([Eq(a, 0.), Eq(c, c*a)], dse='noop', dle=None)
        assert """\
  float (*a);
  float (*c)[j_size];
  posix_memalign((void**)&a, 64, sizeof(float[i_size]));
  posix_memalign((void**)&c, 64, sizeof(float[i_size][j_size]));
  struct timeval start_section_1, end_section_1;
  gettimeofday(&start_section_1, NULL);
  for (int i = i_s; i < i_e; i += 1)
  {
    a[i] = 0.0F;
    for (int j = j_s; j < j_e; j += 1)
    {
      c[i][j] = a[i]*c[i][j];
    }
  }
  gettimeofday(&end_section_1, NULL);
  timings->section_1 += (double)(end_section_1.tv_sec-start_section_1.tv_sec)\
+(double)(end_section_1.tv_usec-start_section_1.tv_usec)/1000000;
  free(a);
  free(c);
  return 0;""" in str(operator.ccode)

    def test_stack_scalar_temporaries(self, a, t0, t1):
        operator = Operator([Eq(t0, 1.), Eq(t1, 2.), Eq(a, t0*t1*3.)],
                            dse='noop', dle=None)
        assert """\
  float (*a);
  posix_memalign((void**)&a, 64, sizeof(float[i_size]));
  struct timeval start_section_0, end_section_0;
  gettimeofday(&start_section_0, NULL);
  for (int i = i_s; i < i_e; i += 1)
  {
    float t0 = 1.00000000000000F;
    float t1 = 2.00000000000000F;
    a[i] = 3.0F*t0*t1;
  }
  gettimeofday(&end_section_0, NULL);
  timings->section_0 += (double)(end_section_0.tv_sec-start_section_0.tv_sec)\
+(double)(end_section_0.tv_usec-start_section_0.tv_usec)/1000000;
  free(a);
  return 0;""" in str(operator.ccode)

    def test_stack_vector_temporaries(self, c_stack, e):
        operator = Operator([Eq(c_stack, e*1.)], dse='noop', dle=None)
        assert """\
  float c_stack[i_size][j_size] __attribute__((aligned(64)));
  struct timeval start_section_0, end_section_0;
  gettimeofday(&start_section_0, NULL);
  for (int k = k_s; k < k_e; k += 1)
  {
    for (int s = s_s; s < s_e; s += 1)
    {
      for (int q = q_s; q < q_e; q += 1)
      {
        for (int i = i_s; i < i_e; i += 1)
        {
          for (int j = j_s; j < j_e; j += 1)
          {
            c_stack[i][j] = 1.0F*e[k][s][q][i][j];
          }
        }
      }
    }
  }
  gettimeofday(&end_section_0, NULL);
  timings->section_0 += (double)(end_section_0.tv_sec-start_section_0.tv_sec)\
+(double)(end_section_0.tv_usec-start_section_0.tv_usec)/1000000;
  return 0;""" in str(operator.ccode)

@skipif_yask
class TestLoopScheduler(object):

    def test_consistency_coupled_wo_ofs(self, tu, tv, ti0, t0, t1):
        """
        Test that no matter what is the order in which the equations are
        provided to an Operator, the resulting loop nest is the same.
        None of the array accesses in the equations use offsets.
        """
        eq1 = Eq(tu, tv*ti0*t0 + ti0*t1)
        eq2 = Eq(ti0, tu + t0*3.)
        eq3 = Eq(tv, ti0*tu)
        op1 = Operator([eq1, eq2, eq3], dse='noop', dle='noop')
        op2 = Operator([eq2, eq1, eq3], dse='noop', dle='noop')
        op3 = Operator([eq3, eq2, eq1], dse='noop', dle='noop')

        trees = [retrieve_iteration_tree(i) for i in [op1, op2, op3]]
        assert all(len(i) == 1 for i in trees)
        trees = [i[0] for i in trees]
        for tree in trees:
            assert IsPerfectIteration().visit(tree[0])
            assert len(tree[-1].nodes) == 3
        pivot = set([j.expr for j in trees[0][-1].nodes])
        assert all(set([j.expr for j in i[-1].nodes]) == pivot for i in trees)

    @pytest.mark.parametrize('exprs', [
        ('Eq(ti0[x,y,z], ti0[x,y,z] + ti1[x,y,z])', 'Eq(ti1[x,y,z], ti3[x,y,z])',
         'Eq(ti3[x,y,z], ti1[x,y,z] + 1.)'),
        ('Eq(ti0[x,y,z], ti0[x,y,z-1] + ti1[x,y,z+1])', 'Eq(ti1[x,y,z], ti3[x,y,z+1])',
         'Eq(ti3[x,y,z], ti3[x,y,z-1] + ti0[x,y,z])'),
        ('Eq(ti0[x,y,z+2], ti0[x,y,z-1] + ti1[x,y,z+1])',
         'Eq(ti1[x,y,z+3], ti3[x,y,z+1])',
         'Eq(ti3[x,y,z+2], ti0[x,y,z+1]*ti3[x,y,z-1])'),
        ('Eq(ti0[x,y,z], ti0[x-2,y-1,z-1] + ti1[x+2,y+3,z+1])',
         'Eq(ti1[x,y,z], ti3[x+1,y-4,z+1])',
         'Eq(ti3[x,y,z], ti3[x+5,y,z-1] - ti0[x+3,y-2,z+4])')
    ])
    def test_consistency_coupled_w_ofs(self, exprs, ti0, ti1, ti3):
        """
        Test that no matter what is the order in which the equations are
        provided to an Operator, the resulting loop nest is the same.
        The array accesses in the equations may or may not use offsets;
        these impact the loop bounds, but not the resulting tree
        structure.
        """
        eq1, eq2, eq3 = EVAL(exprs, ti0.base, ti1.base, ti3.base)
        op1 = Operator([eq1, eq2, eq3], dse='noop', dle='noop')
        op2 = Operator([eq2, eq1, eq3], dse='noop', dle='noop')
        op3 = Operator([eq3, eq2, eq1], dse='noop', dle='noop')

        trees = [retrieve_iteration_tree(i) for i in [op1, op2, op3]]
        assert all(len(i) == 1 for i in trees)
        trees = [i[0] for i in trees]
        for tree in trees:
            assert IsPerfectIteration().visit(tree[0])
            assert len(tree[-1].nodes) == 3
        pivot = set([j.expr for j in trees[0][-1].nodes])
        assert all(set([j.expr for j in i[-1].nodes]) == pivot for i in trees)

    def test_expressions_imperfect_loops(self, ti0, ti1, ti2, t0):
        """
        Test that equations depending only on a subset of all indices
        appearing across all equations are placed within earlier loops
        in the loop nest tree.
        """
        eq1 = Eq(ti2, t0*3.)
        eq2 = Eq(ti0, ti1 + 4. + ti2*5.)
        op = Operator([eq1, eq2], dse='noop', dle='noop')
        trees = retrieve_iteration_tree(op)
        assert len(trees) == 2
        outer, inner = trees
        assert len(outer) == 2 and len(inner) == 3
        assert all(i == j for i, j in zip(outer, inner[:-1]))
        assert outer[-1].nodes[0].expr.rhs == eq1.rhs
        assert inner[-1].nodes[0].expr.rhs == eq2.rhs

    def test_equations_emulate_bc(self, t0):
        """
        Test that bc-like equations get inserted into the same loop nest
        as the "main" equations.
        """
        grid = Grid(shape=(3, 3, 3))
        x, y, z = grid.dimensions
        time = grid.time_dim
        a = Function(name='a', grid=grid).indexed
        b = TimeFunction(name='b', grid=grid, save=True, time_dim=6).indexed
        main = Eq(b[time + 1, x, y, z], b[time - 1, x, y, z] + a[x, y, z] + 3.*t0)
        bcs = [Eq(b[time, 0, y, z], 0.),
               Eq(b[time, x, 0, z], 0.),
               Eq(b[time, x, y, 0], 0.)]
        op = Operator([main] + bcs, dse='noop', dle='noop')
        trees = retrieve_iteration_tree(op)
        assert len(trees) == 4
        assert all(id(trees[0][0]) == id(i[0]) for i in trees)

    def test_different_section_nests(self, tu, ti0, t0, t1):
        eq1 = Eq(ti0, t0*3.)
        eq2 = Eq(tu, ti0 + t1*3.)
        op = Operator([eq1, eq2], dse='noop', dle='noop')
        trees = retrieve_iteration_tree(op)
        assert len(trees) == 2
        assert trees[0][-1].nodes[0].expr.rhs == eq1.rhs
        assert trees[1][-1].nodes[0].expr.rhs == eq2.rhs

    @pytest.mark.parametrize('exprs', [
        ['Eq(ti0[x,y,z], ti0[x,y,z] + t0*2.)', 'Eq(ti0[0,0,z], 0.)'],
        ['Eq(ti0[x,y,z], ti0[x,y,z-1] + t0*2.)', 'Eq(ti0[0,0,z], 0.)'],
        ['Eq(ti0[x,y,z], ti0[x,y,z] + t0*2.)', 'Eq(ti0[0,y,0], 0.)'],
        ['Eq(ti0[x,y,z], ti0[x,y,z] + t0*2.)', 'Eq(ti0[0,y,z], 0.)'],
    ])
    def test_directly_indexed_expression(self, fa, ti0, t0, exprs):
        """
        Test that equations using integer indices are inserted in the right
        loop nest, at the right loop nest depth.
        """
        eqs = EVAL(exprs, ti0.base, t0)
        op = Operator(eqs, dse='noop', dle='noop')
        trees = retrieve_iteration_tree(op)
        assert len(trees) == 2
        assert trees[0][-1].nodes[0].expr.rhs == eqs[0].rhs
        assert trees[1][-1].nodes[0].expr.rhs == eqs[1].rhs

    @pytest.mark.parametrize('shape, dimensions', [((11, 11), (x, y)),
                                                   ((11, 11), (y, x)),
                                                   ((11, 11, 11), (x, y, z)),
                                                   ((11, 11, 11), (x, z, y)),
                                                   ((11, 11, 11), (y, x, z)),
                                                   ((11, 11, 11), (y, z, x)),
                                                   ((11, 11, 11), (z, x, y)),
                                                   ((11, 11, 11), (z, y, x))])
    def test_equations_mixed_densedata_timedata(self, shape, dimensions):
        """
        Test that equations using a mixture of Function and TimeFunction objects
        are embedded within the same time loop.
        """
        grid = Grid(shape=shape, dimensions=dimensions, time_dimension=time)
        a = TimeFunction(name='a', grid=grid, time_order=2, space_order=2)
        p_aux = Dimension(name='p_aux', size=10)
        b = Function(name='b', shape=shape + (10,), dimensions=dimensions + (p_aux,),
                     space_order=2)
        b.data[:] = 1.0
        b2 = Function(name='b2', shape=(10,) + shape, dimensions=(p_aux,) + dimensions,
                      space_order=2)
        b2.data[:] = 1.0
        eqns = [Eq(a.forward, a.laplace + 1.),
                Eq(b, time*b*a + b)]
        eqns2 = [Eq(a.forward, a.laplace + 1.),
                 Eq(b2, time*b2*a + b2)]
        subs = {x.spacing: 2.5, y.spacing: 1.5, z.spacing: 2.0}
        op = Operator(eqns, subs=subs, dle='noop')
        trees = retrieve_iteration_tree(op)
        assert len(trees) == 2
        assert all(trees[0][i] is trees[1][i] for i in range(3))

        op2 = Operator(eqns2, subs=subs, dle='noop')
        trees = retrieve_iteration_tree(op2)
        assert len(trees) == 2

        # Verify both operators produce the same result
        op(time=10)
        a.data[:] = 0.
        op2(time=10)

        for i in range(10):
            assert(np.allclose(b2.data[i, ...].reshape(-1) -
                               b.data[..., i].reshape(-1), 0.))

@skipif_yask
@pytest.mark.skipif(configuration['backend'] != 'foreign',
                    reason="'foreign' wasn't selected as backend on startup")
class TestForeign(object):

    def test_explicit_run(self):
        time_dim = 6
        grid = Grid(shape=(11, 11))
        a = TimeFunction(name='a', grid=grid, time_order=1,
                         time_dim=time_dim, save=True)
        eqn = Eq(a.forward, a + 1.)
        op = Operator(eqn)
        assert isinstance(op, OperatorForeign)
        args = OrderedDict(op.arguments())
        assert args['a'] is None
        # Emulate data feeding from outside
        array = np.ndarray(shape=a.shape, dtype=np.float32)
        array.fill(0.0)
        args['a'] = array
        op.cfunction(*list(args.values()))
        assert all(np.allclose(args['a'][i], i) for i in range(time_dim))


text:
# CS 315 Final Project by Brandon Cook and Matt Balint

import numpy as np
import pandas as pd

from draw_graph import graph_regression
from draw_graph import graph
from regression import calculate_regression

def get_states():
    return np.array(state_populations['State'].tolist())

# Remove any US territories or prisons from the data set
def remove_non_states():
    x = vaccine_data['location'] # get locations column
    i = len(x) - 1 # number of rows

    while i > 0:
        if x[i] not in states:
            vaccine_data.drop(vaccine_data.index[i], inplace=True)
        i -= 1

    vaccine_data.reset_index(drop=True, inplace=True) # update the indexes so they match with the length of the DataFrame

# Group data together for each state and put it in a dictionary
def order_by_state():
    num_rows = vaccine_data.shape[0]
    state = vaccine_data.at[0,'location']
    start_index = 0
    temp_df = pd.DataFrame()

    for i in range(num_rows):
        current_state = vaccine_data.at[i, 'location']
        if (current_state != state):
            temp_df = vaccine_data.loc[start_index:i-1]

            state_vaccinations.update({vaccine_data['location'][i - 1] : temp_df})
            start_index = i
        state = current_state

    # Need to add the last state in the list since the for loop won't cover it.
    state_vaccinations.update({vaccine_data['location'][start_index] : temp_df})

def get_percent_vaccinated(date, state): # returns the % of people vaccinated in that state on that particular date
    state_data = state_vaccinations.get(state)
    date_data = state_data.loc[state_data['date'] == date] # Gets row of date from dataframe

    num_vaccinated = date_data.iloc[0]['people_fully_vaccinated'] # Gets value of vaccinated from row

    state_data = state_populations[state_populations['State'] == state]
    population = state_data.iloc[0]['Population_Estimate']
    return num_vaccinated/population

def get_state_input():
    state = ""
    while True:
        print("\nPlease enter a US state: ")
        state = input()
        state = state.title()
        if state == "New York":
            print("\Invalid state name. Please enter \'New York State\'")
        elif state not in states:
            print("\nInvalid state name. Please enter a valid US state like Washington.")
        else:
            break

    return state

def menu():
    choice = 0
    while choice <= 0 or choice > 5:
        print("--------US Covid-19 Vaccine Tracker --------\n")
        print("Enter a number below to retrieve a menu.\n")
        print("1. US Vaccine Prediction Graph")
        print("2. State Vaccine Prediction Graph")
        print("3. Predict time left till herd immunity reached")
        print("4. Quit")
        choice = int(input())

        if choice == 1:
            print()
            # calculate_us_regression()
        elif choice == 2:
            state = get_state_input()
            x, y, regression_line = calculate_regression(state, state_vaccinations)
            graph_regression(x, y, regression_line, state)

        elif choice == 3:
            state = get_state_input()
            # calculate_us_regression()
            # make prediction using regression

if __name__ == '__main__':
    # Read data
    raw_vaccine_data = pd.read_csv("datasets/us_state_vaccinations0507.csv") # This spreadsheet has Covid-19 data up until May 1st
    vaccine_data = raw_vaccine_data[['date','location','people_vaccinated','people_fully_vaccinated']].copy() # Filter out some of the extra unnecessary columns
    state_populations = pd.read_csv("datasets/2019_Census_Data_Population.csv")
    states = get_states() # create a list of the names of all 50 states
    # current_date = date.today()

    state_vaccinations = {}

    # Clean data
    remove_non_states()
    order_by_state()

    # Display menu
    menu()

ext: py

reason: import numpy


text:
# Copyright (c) 2017, Lawrence Livermore National Security, LLC. Produced at
# the Lawrence Livermore National Laboratory.
# LLNL-CODE-743438
# All rights reserved.
# This file is part of MGmol. For details, see https://github.com/llnl/mgmol.
# Please also read this link https://github.com/llnl/mgmol/LICENSE
#
#source /usr/apps/python/default/setup.csh
#python averageMDpositions.py . > average_positions.dat
import sys, string, os
from numpy import *
from pdb_tools import printPDBfile
from collections import defaultdict

au2ps = 2.418885e-05

times=[]
#coords={}
coords=defaultdict(list)
nsteps=0
time=0.
dt=0.

def appendPositions(filename):
  global coords
  global nsteps
  global time
  global dt

  found=0
  found_force=0

  #print '#',filename
  file=open(filename,'r')
  L1=file.readlines()
  for line in L1: ## loop over lines of file
    words=string.split(line)
    if len(words)>1:
      #find positions
      if 'Stepper' in words and 'Forces:' in words:
        found_force=1
        found=0
        nsteps=nsteps+1
        time=time+dt*au2ps
        times.append(time)
      if words[0][0:2]=='##' and found_force:
        name=words[1]
        found=found+1
        valx=eval(words[2])
        valy=eval(words[3])
        valz=eval(words[4])
        if name not in coords:
          coords[name]=list()
        coords[name].append([valx,valy,valz])
    else:
        if found>0:
          found_force=0

  file.close()

#main
filesdir=sys.argv[1]
filenames=os.listdir(filesdir)

inputs=[]
for filename in filenames:
  if 'md_run' in filename:
    inputs.append(filename)

inputs.sort()

#read 'dt'
inputs[0]
file=open(inputs[0],'r')
L1=file.readlines()
for line in L1: ## loop over lines of file
  word=string.split(line)
  if len(word)>1:
    if word[0]=='Timestep':
      dt=eval(word[5])

for filename in inputs:
  appendPositions(filename)

#number of steps/ps
nf=int(ceil(1./(dt*au2ps)))

#sampling length in ps
number_ps=3

xyz=[]
names=[]
movables=[]

#loop over atoms
for c in coords.keys():
  #print coords[c]
  #print len(coords[c])
  npts=min(number_ps*nf,len(coords[c]))

  names.append(c)

  #build lists with latest npts values
  x=[]
  y=[]
  z=[]
  for i in range(len(coords[c])-npts,len(coords[c])-1):
    x.append(coords[c][i][0])
    y.append(coords[c][i][1])
    z.append(coords[c][i][2])

  #calculate averages with latest npts values
  avex = sum(x[-npts:])/npts
  avey = sum(y[-npts:])/npts
  avez = sum(z[-npts:])/npts

  xyz.append(str(avex)+' '+str(avey)+' '+str(avez))
  movables.append('1')

na=len(names)
printPDBfile(na,names,xyz,movables,'')

reason: from numpy


text:
import numpy as np

class Kelly_bot:

    def __init__(self,W0=1.0,r1=1.0,r2=-1.0,p0=None,dp0=None, strategy=None,lam=None):
        '''
        Parameters:
        p0: Initial coin bias
        dp0: Initial coin uncertainty
        lambda: The aggressiveness of the agent
        '''

        self.p0 = p0
        self.dp0 = dp0
        self.W0 = W0 # The agents initial wealth
        self.r1 = r1
        self.r2 = r2

        self.W_vec = []
        self.dW_vec = []
        self.f_vec = []
        self.bet_vec = []
        self.p_vec = []
        self.dp_vec = []

        # Compute the alpha and beta paramters for the conjugate priors

        if p0 is None or (p0==0.0):
            self.alpha = 0.0
            self.beta = 0.0
            self.p1 = 0.0
            self.p2 = 0.0
            self.p_est = 0.0
            self.dp_est = 0.0
        else:
            self.alpha = p0*( (p0*(1.0-p0)/dp0**2) -1)
            self.beta = (1-p0)*( (p0*(1.0-p0)/dp0**2) -1)
            self.p_est = self.alpha/(self.alpha+self.beta)
            self.dp_est = np.sqrt(self.p_est*(1.0-self.p_est)/(self.alpha+self.beta+1))
            self.p1 = self.alpha/(self.alpha+self.beta)
            self.p2 = 1.0-self.p1

        # The agents memory
        self.N_heads = 0
        self.N_tails = 0
        self.data = []

        if lam is None:
            self.lam = 0.0
        else:
            self.lam = lam

        # Store information about the agents strategy
        if strategy is not None:
            self.strategy = strategy
            self.strategy_name = strategy['name']

        return None

    def reset(self):
        '''
        Reset the memory of the kelly-bot
        '''

        self.N_heads = 0
        self.N_tails = 0
        self.p_est = 0.0
        self.dp_est = 0.0
        self.p1 = 0.0
        self.p2 = 0.0

        self.data = []
        self.dp_vec = []
        self.p_vec = []
        self.W_vec = []
        self.bet_vec = []
        self.f_vec = []

        return None

    def posterior_distribution(self):
        '''
        This function returns the agents current belief about the coin probability
        as a continous distribution
        '''

        def f(x):

            y = (x**(self.N_heads+self.alpha-1))*((1.0-x)**(self.N_tails+self.beta-1))
            norm = 1.0

            y /= norm

            return y

        return None

    def new_observation(self,obs):
        '''
        Observe a new data point and update all of the information

        '''

        # Add the new data to the agents memory
        self.data.append(obs)

        # Update the current estimate of p0 and its uncertainty
        self.N_heads += int(obs==1)
        self.N_tails += int(obs==0)

        # Update the mean value of the estimate
        self.p_est = (self.N_heads + self.alpha)/(self.N_heads+self.N_tails+self.alpha+self.beta)
        self.dp_est = np.sqrt(self.p_est*(1-self.p_est)/(self.alpha+self.beta+self.N_heads+self.N_tails+1))

        # Update the vector containing the estimates of the current values
        self.p_vec.append(self.p_est)
        self.dp_vec.append(self.dp_est)

        # New updated values
        self.p1 = self.p_est
        self.p2 = 1.0-self.p_est

        return None

    def update_wealth(self,result,wager):
        '''
        Update the agents wealth
        '''

        if result==1:
            dW = (wager/self.W0)
            self.W0+= wager
        elif result==0:
            dW = -1.0*(wager / self.W0)
            self.W0+= -1.0*wager

        if self.W0 <0:
            self.W0 = 0.0

        # Update the Agents wealth
        self.W_vec.append(self.W0)
        self.dW_vec.append(dW)

        return None

    def apply_strategy(self):
        '''

        '''

        if self.strategy_name == 'fixed_fraction':
            f,bet_amount = self.fixed_fraction(f_fixed=self.strategy['f_fixed'])
        elif self.strategy_name == 'fixed_bet':
            f,bet_amount = self.fixed_bet_size(bet_size=self.strategy['fixed_bet_size'])
        elif self.strategy_name == 'adaptive_kelly':
            f,bet_amount = self.kelly_bet_size(fraction=self.strategy['kelly_fraction'])

        return bet_amount

    def kelly_bet_size(self,fraction=1.0):
        '''
        Determine the optimal bet given all of the available information
        '''

        f = -(self.r1*self.p1+self.r2*self.p2)/(self.r1*self.r2)
        f += (self.lam/(1-self.lam))*((self.r1-self.r2)/(self.r1*self.r2))*self.dp_est
        f = fraction*f # Take a fraction of the kelly bet

        # Only Bet if the fraction is >0
        if f<0:
            f=0.0

        bet_amount = f*self.W0 # The size of the wager

        self.f_vec.append(f)
        self.bet_vec.append(bet_amount)

        return f,bet_amount

    def fixed_bet_size(self,bet_size):
        '''
        Place a bet of fixed amount
        '''

        f = (bet_size/self.W0)

        self.f_vec.append(f)
        self.bet_vec(bet_size)

        return f, bet_size

    def fixed_fraction(self,f_fixed):
        '''
        Determine the optimal bet given all of the available information
        '''

        bet_amount = f_fixed*self.W0 # The size of the wager

        self.f_vec.append(f_fixed)
        self.bet_vec.append(bet_amount)

        return f_fixed,bet_amount


text:
###
#
# Simple script for using world
#
# (Made for 2 agents with turn based movements, but can be extended for multiple agents)
# For Open Acess and Use
#
# Made by Willie McClinton at University of South Florida
#
###

from World import World
import numpy as np
import time

# Creates new World
test_world = World("Test")

# Display World in Terminal
test_world.render()
done = 0

while(not done):
    test_world.render()
    time.sleep(0.25)

    # Chose random actions
    random_action_agent1 = ["RIGHT","LEFT","UP","DOWN","ATTACK"][np.random.randint(5)]
    random_action_agent2 = ["RIGHT","LEFT","UP","DOWN","ATTACK"][np.random.randint(5)]

    # Take action
    state1, reward1, done = test_world.step(random_action_agent1,test_world.agent1.name)
    state2, reward2, done = test_world.step(random_action_agent2,test_world.agent2.name)

# Resets world just so you know how
test_world.reset()

# Prints Final Values
print(state1)
print(reward1)
print(done)


text:
import numpy as np
from scipy.linalg import qr
from scipy.optimize import least_squares
from tensorly.tenalg import khatri_rao
from traniest.tools.tensor_utils import get_symm_mat_pow

def VectorsToTensor(As, Ls = None):
	Dim = ()
	if Ls is None:
		Ls = np.ones((As[0].shape[1],))
	for A in As:
		iN = A.shape[0]
		Dim += (iN, )
	X = np.zeros(Dim)
	for r in range(As[0].shape[1]): # Rank of decomposition
		ais = []
		for A in As:
			ais.append(A[:, r])
		X += Ls[r] * np.prod(np.ix_(*ais))
	return X

def PARAFACError(X, As, Ls = None, ReturnVector = False):
	x = VectorsToTensor(As, Ls = Ls)
	if ReturnVector:
		return (X - x)
	return np.sqrt(((X - x)**2).sum())

def DIIS_Solve(Es):
	B = np.full((len(Es) + 1, len(Es) + 1), -1.0)
	for i in range(len(Es)):
		for j in range(len(Es)):
			B[i, j] = (Es[i] * Es[j]).sum()
	B[len(Es), len(Es)] = 0.0
	y = np.zeros((len(Es) + 1,))
	y[len(Es)] = -1.0
	return np.linalg.solve(B, y)[:len(Es)]

def DIIS(As, Es):
	cs = DIIS_Solve(Es)
	NewA = np.zeros(As[0].shape)
	for n in range(len(As)):
		NewA += cs[n] * As[n]
	return NewA

def Matricize(X, m):
	I = X.shape
	IArr = np.array(I)

	Dim1 = I[m]
	Dim2 = int(IArr.prod() / Dim1)
	Xm = X.copy()
	for i in range(m, 0, -1):
		Xm = np.swapaxes(Xm, i - 1, i)
	return Xm.reshape(Dim1, Dim2)

def UnMatricize(X, m, Dim):
	Xu = X.reshape(Dim)
	for i in range(m):
		Xu = np.swapaxes(Xu, i, i + 1)
	return Xu

'''
Input
	X (np.array): I1 x .. x IN
	U (np.array): J x In
Output
	UX (np.array): I1 x .. x J x .. x IN
'''
def nModeMatProduct(X, U, n):
	Dim = list(X.shape)
	Dim[n] = U.shape[0]
	Dim = tuple(Dim)
	Xn = Matricize(X, n)
	Yn = U @ Xn
	return UnMatricize(Yn, n, Dim)

def InitHOSVD(X, R):
	As = []
	for i in range(X.ndim):
		Xi = Matricize(X, i)
		U, S, V = np.linalg.svd(Xi)
		U = U[:, :R]
		As.append(U)
	return As

def InitA(I, r):
	As = []
	for i in I:
		A = np.eye(i, r)
		As.append(A)
	return As

def PARAFAC_Orth(X, R, max_iter = 1000, DIISSpace = 0, verbose = 0):
	I = X.shape
	As = InitHOSVD(X, R)
	def OptL(weights, factors, tensor):
		return PARAFACError(tensor, factors, Ls = weights)
	Ls = np.ones((R,))
	if DIISSpace > 0:
		ErrorVecs = []
		PrevAs = []
		for i in range(len(I)):
			ErrorVecs.append([])
			PrevAs.append([])
	for i in range(max_iter):
		for n in range(len(I)):
			W = khatri_rao(As, skip_matrix = n)
			Xn = Matricize(X, n)
			Y = get_symm_mat_pow(W.T @ Xn.T @ Xn @ W, -0.50, check_symm = False) # W @ np.linalg.pinv(V)
			As[n] = Xn @ W @ Y
			if DIISSpace > 0:
				Ls = least_squares(OptL, Ls, args = [As, X]).x
				E = PARAFACError(X, As, Ls = Ls, ReturnVector = True)
				ErrorVecs[n].append(E)
				ErrorVecs[n] = ErrorVecs[n][-DIISSpace:]
				PrevAs[n].append(As[n])
				PrevAs[n] = PrevAs[n][-DIISSpace:]
				if i > DIISSpace:
					As[n] = DIIS(PrevAs[n], ErrorVecs[n])
					Ls = least_squares(OptL, Ls, args = [As, X]).x
					E = PARAFACError(X, As, Ls = Ls, ReturnVector = True)
					PrevAs[n][-1] = As[n]
					ErrorVecs[n][-1] = E
		Ls = least_squares(OptL, Ls, args = [As, X]).x
		Err = PARAFACError(X, As, Ls = Ls)
		if verbose > 0:
			print("PARAFAC Iteration", i, "complete with error", Err)
		if Err < 1e-6:
			break
	print("PARAFAC complete with error", Err)
	return As, Ls

def PARAFAC_OrthQR(X, R, max_iter = 1000, Normalize = False):
	I = X.shape
	As = InitHOSVD(X, R)
	for i in range(max_iter):
		Ls = np.ones((R,))
		# Pre orthogonalize all factor matrices
		for n in range(len(As)):
			As[n] = qr(As[n], mode = 'economic')[0]
		for n in range(len(I)):
			W = khatri_rao(As, skip_matrix = n)
			Xn = Matricize(X, n)
			As[n] = Xn @ W @ np.linalg.inv(W.T @ W)
		for A in As:
			for j in range(A.shape[1]):
				norm = np.linalg.norm(A[:, j])
				A[:, j] /= norm
				Ls[j] *= norm
		Err = PARAFACError(X, As, Ls = Ls)
		print("PARAFAC Iteration", i, "complete with error", Err)
		if Err < 1e-6:
			break
	return As, Ls

def PARAFAC(X, R, max_iter = 1000, Normalize = False):
	I = X.shape
	As = InitHOSVD(X, R)
	for i in range(max_iter):
		if Normalize:
			Ls = np.ones((R,))
		else:
			Ls = None
		for n in range(len(I)):
			#V = np.ones((R, R))
			#for m, A in enumerate(As):
			#	if m == n:
			#		continue
			#	else:
			#		V *= A.T @ A

			W = khatri_rao(As, skip_matrix = n)
			Xn = Matricize(X, n)
			As[n] = Xn @ W @ np.linalg.inv(W.T @ W)
			#As[n] = Xn @ W @ np.linalg.pinv(V)
		if Normalize:
			for A in As:
				for j in range(A.shape[1]):
					norm = np.linalg.norm(A[:, j])
					A[:, j] /= norm
					Ls[j] *= norm
		Err = PARAFACError(X, As, Ls = Ls)
		if Err < 1e-6:
			break
	print("PARAFAC complete with error", Err)
	return As, Ls

if __name__ == "__main__":
	X = np.zeros((3, 4, 2))
	for i in range(X.shape[0]):
		for j in range(X.shape[1]):
			X[i, j, 0] = i + X.shape[0] * j + 1
			X[i, j, 1] = i + X.shape[0] * j + 13
	X = np.random.rand(10,10,10,10)

	#U = np.zeros((2, 3))
	#for i in range(U.shape[0]):
	#	for j in range(U.shape[1]):
	#		U[i, j] = i + U.shape[0] * j + 1

	#Y = nModeMatProduct(X, U, 0)
	#print(Y[:, :, 0])
	#print(Y[:, :, 1])

	As, Ls = PARAFAC_Orth(X, 2, max_iter = 100, DIISSpace = 3, verbose = 1)
	As, Ls = PARAFAC_Orth(X, 2, max_iter = 100, DIISSpace = 0, verbose = 1)
	from tensorly.decomposition import parafac
	S, Us = parafac(X, 2, n_iter_max = 100, orthogonalise = True, normalize_factors = False)
	print(PARAFACError(X, Us, Ls = S))

max_stars_repo_name: Sekie/TranIEST

max_issues_repo_name: Sekie/TranIEST

max_forks_repo_name: the-hktran/TranIEST

reason: import numpy,from scipy


text:
# Copyright 2018 Jörg Franke
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#    http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
import numpy as np
import tensorflow as tf
import pytest

from adnc.model.memory_units.base_cell import BaseMemoryUnitCell

@pytest.fixture(
    params=[{"seed": 123, "input_size": 13, "batch_size": 3, "memory_length": 4, "memory_width": 4, "read_heads": 3,
             "dnc_norm": True, "bypass_dropout": False},
            {"seed": 124, "input_size": 11, "batch_size": 3, "memory_length": 256, "memory_width": 23, "read_heads": 2,
             "dnc_norm": False, "bypass_dropout": False},
            {"seed": 125, "input_size": 5, "batch_size": 3, "memory_length": 4, "memory_width": 11, "read_heads": 8,
             "dnc_norm": True, "bypass_dropout": True},
            {"seed": 126, "input_size": 2, "batch_size": 3, "memory_length": 56, "memory_width": 9, "read_heads": 11,
             "dnc_norm": False, "bypass_dropout": True}
            ])
def memory_config(request):
    config = request.param
    return BaseMemoryUnitCell(input_size=config['input_size'], memory_length=config["memory_length"],
                              memory_width=config["memory_width"],
                              read_heads=config["read_heads"], seed=config["seed"],
                              reuse=False, name='test_mu'), config

@pytest.fixture()
def session():
    with tf.Session() as sess:
        yield sess
    tf.reset_default_graph()

@pytest.fixture()
def np_rng():
    seed = np.random.randint(1, 999)
    return np.random.RandomState(seed)

class TestDNCMemoryUnit():
    def test_init(self, memory_config):

        memory_unit, config = memory_config

        assert isinstance(memory_unit, object)
        assert isinstance(memory_unit.rng, np.random.RandomState)

        assert memory_unit.h_N == config["memory_length"]
        assert memory_unit.h_W == config["memory_width"]
        assert memory_unit.h_RH == config["read_heads"]

    def test_property_output_size(self, memory_config, session):
        memory_unit, config = memory_config

        output_size = memory_unit.output_size
        assert output_size == config['memory_width'] * config["read_heads"] + config['input_size']

    def test_calculate_content_weightings(self, memory_config, session, np_rng):
        memory_unit, config = memory_config

        np_memory = np_rng.normal(0, 1, (config['batch_size'], config['memory_length'], config['memory_width']))
        np_keys = np_rng.normal(0, 2, (config['batch_size'], 1, config['memory_width']))
        np_strengths = np_rng.uniform(1, 10, (config['batch_size'], 1))

        memory = tf.constant(np_memory, dtype=tf.float32)
        keys = tf.constant(np_keys, dtype=tf.float32)
        strengths = tf.constant(np_strengths, dtype=tf.float32)

        content_weightings = memory_unit._calculate_content_weightings(memory, keys, strengths)
        weightings = content_weightings.eval()

        np_similarity = np.empty([config['batch_size'], config['memory_length']])
        for b in range(config['batch_size']):
            for l in range(config['memory_length']):
                np_similarity[b, l] = np.dot(np_memory[b, l, :], np_keys[b, 0, :]) / (
                    np.sqrt(np.dot(np_memory[b, l, :], np_memory[b, l, :])) * np.sqrt(
                        np.dot(np_keys[b, 0, :], np_keys[b, 0, :])))

        def _weighted_softmax(x, s):
            e_x = np.exp(x * s)
            return e_x / e_x.sum(axis=1, keepdims=True)

        np_weightings = _weighted_softmax(np_similarity, np_strengths)

        assert weightings.shape == (config['batch_size'], config['memory_length'])
        assert 0 <= weightings.min() and weightings.max() <= 1 and weightings.sum(axis=1).all() <= 1
        assert np.allclose(weightings, np_weightings)

        np_memory = np_rng.uniform(0, 1, (config['batch_size'], config['memory_length'], config['memory_width']))
        np_keys = np_rng.normal(0, 2, (config['batch_size'], config['read_heads'], config['memory_width']))
        np_strengths = np_rng.uniform(1, 10, (config['batch_size'], config['read_heads'], 1))

        memory = tf.constant(np_memory, dtype=tf.float32)
        keys = tf.constant(np_keys, dtype=tf.float32)
        strengths = tf.constant(np_strengths, dtype=tf.float32)

        content_weightings = memory_unit._calculate_content_weightings(memory, keys, strengths)
        weightings = content_weightings.eval()

        np_similarity = np.empty([config['batch_size'], config['read_heads'], config['memory_length']])
        for b in range(config['batch_size']):
            for r in range(config['read_heads']):
                for l in range(config['memory_length']):
                    np_similarity[b, r, l] = np.dot(np_memory[b, l, :], np_keys[b, r, :]) / (
                        np.sqrt(np.dot(np_memory[b, l, :], np_memory[b, l, :])) * np.sqrt(
                            np.dot(np_keys[b, r, :], np_keys[b, r, :])))
        np_weightings = np.empty([config['batch_size'], config['read_heads'], config['memory_length']])

        def _weighted_softmax(x, s):
            e_x = np.exp(x * s)
            return e_x / e_x.sum(axis=1, keepdims=True)

        for r in range(config['read_heads']):
            np_weightings[:, r, :] = _weighted_softmax(np_similarity[:, r, :], np_strengths[:, r])

        assert weightings.shape == (config['batch_size'], config['read_heads'], config['memory_length'])
        assert 0 <= weightings.min() and weightings.max() <= 1 and weightings.sum(axis=2).all() <= 1
        assert np.allclose(weightings, np_weightings)

    def test_read_memory(self, memory_config, session, np_rng):
        memory_unit, config = memory_config

        np_memory = np_rng.normal(0, 1, [config['batch_size'], config['memory_length'], config['memory_width']])
        np_read_weightings = np_rng.uniform(0, 1 / config['memory_length'],
                                            [config['batch_size'], config['read_heads'], config['memory_length']])

        memory = tf.constant(np_memory, dtype=tf.float32)
        read_weightings = tf.constant(np_read_weightings, dtype=tf.float32)

        read_vectors = memory_unit._read_memory(memory, read_weightings)
        read_vectors = read_vectors.eval()

        np_read_vectors = np.empty([config['batch_size'], config['read_heads'], config['memory_width']])
        for b in range(config['batch_size']):
            for r in range(config['read_heads']):
                np_read_vectors[b, r, :] = np.matmul(np.expand_dims(np_read_weightings[b, r, :], 0), np_memory[b, :, :])

        assert read_vectors.shape == (config['batch_size'], config['read_heads'], config['memory_width'])
        assert np.allclose(read_vectors, np_read_vectors, atol=1e-06)

max_stars_repo_name: carusyte/ADNC

max_issues_repo_name: carusyte/ADNC

max_forks_repo_name: carusyte/ADNC