<commit_before>#     Copyright 2017 TVB-HPC contributors
#
#     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 time
import logging
import unittest
import loopy as lp
from loopy.target.c import CTarget
import pymbolic as pm
import numpy as np
import numpy.testing
from scipy.stats import kstest
from .bold import BalloonWindkessel
from .compiler import Compiler, CppCompiler, CompiledKernel, Spec
from .coupling import (Linear as LCf, Diff, Sigmoidal, Kuramoto as KCf,
    PostSumStat, BaseCoupling)
from .model import BaseModel, _TestModel, HMJE, RWW, JansenRit, Linear, G2DO
from .model import Kuramoto
from .network import DenseNetwork, DelayNetwork
from .rng import RNG
from .scheme import euler_maruyama_logp, EulerStep, EulerMaryuyamaStep
# from .harness import SimpleTimeStep
from .utils import getLogger, VarSubst


LOG = logging.getLogger(__name__)


class TestCase(unittest.TestCase):

    def setUp(self):
        super().setUp()
        self.tic = time.time()
        self.logger = getLogger(self.id())
        self._time_limit = 0.1

    def timeit(self, fn, *args, **kwds):
        niter = 0
        tic = toc = time.time()
        while ((toc - tic) < self._time_limit):
            fn(*args, **kwds)
            toc = time.time()
            niter += 1
        per_iter = self._time_limit / niter
        self.logger.info('%r requires %.3f ms / iter',
                         fn, per_iter * 1e3)

    def tearDown(self):
        super().tearDown()
        msg = 'required %.3fs'
        self.logger.info(msg, time.time() - self.tic)


class TestUtils(TestCase):

    def test_var_subst(self):
        subst = VarSubst(b=pm.parse('b[i, j]'))
        expr = subst(pm.parse('a + b * pre_syn[i, j]'))
        self.assertEqual(str(expr), 'a + b[i, j]*pre_syn[i, j]')


class TestLoopTransforms(TestCase):
    """
    These are more tests to check that our use of Loopy is correct.

    """

    def setUp(self):
        super().setUp()
        from loopy.target.ispc import ISPCTarget
        target = ISPCTarget()
        self.knl = lp.make_kernel('{[i]:0<=i<n}', "out[i] = in[i]",
                                  target=target)

    def _dtype_and_code(self, knl):
        knl = lp.add_dtypes(knl, {'in': np.float32, 'out': np.float32})
        code, _ = lp.generate_code(knl)
        return code

    def test_chunk_iname(self):
        "Chunk useful to split work for e.g. omp par for"
        knl = lp.chunk_iname(self.knl, 'i', 8)
        print(self._dtype_and_code(knl))

    def test_split_iname2(self):
        "Split useful for omp simd inner loop"
        knl = lp.split_iname(self.knl, 'i', 8)
        knl = lp.tag_inames(knl, [('i_inner', 'ilp.unr',)])
        print(self._dtype_and_code(knl))

    def test_wrap_loop(self):
        "Take kernel, place in larger loop, offsetting certain vars"
        knl = lp.make_kernel("{[i,j]:0<=i,j<n}",
                             "out[i] = sum(j, (i/j)*in[i, j])",
                             target=CTarget())
        # in will depend on t
        knl2 = lp.to_batched(knl, 'T', ['in'], 't')
        print(self._dtype_and_code(knl2))

    def test_split_iname2(self):
        "Split one of two inames."
        from loopy.target.ispc import ISPCTarget as CTarget
        knl = lp.make_kernel("{[i,j]:0<=i,j<n}",
                             "out[i, j] = in[i, j]",
                             target=CTarget())
        knl = lp.split_iname(knl, 'i', 8)
        knl = lp.prioritize_loops(knl, ['i_outer', 'j', 'i_inner'])
        print(self._dtype_and_code(knl))

    def test_sparse_matmul(self):
        knl = lp.make_kernel(
            [
                '{[i]: 0   <= i <   n}',
                '{[j]: jlo <= j < jhi}'
            ],
            """
            <> jlo = row[i]
            <> jhi = row[i + 1]
            out[i] = sum(j, dat[j] * vec[col[j]])
            """,
            'n nnz row col dat vec out'.split(),
            target=CTarget())
        knl = lp.add_dtypes(knl, {
            'out': np.float32,
            'dat': np.float32,
            'vec': np.float32,
            'col': np.uintc,
            'row': np.uintc,
            'n': np.uintc,
            'nnz': np.uintc,
        })
        # col and dat have uninferrable shape
        knl.args[3].shape = pm.var('nnz'),
        knl.args[4].shape = pm.var('nnz'),
        cknl = CompiledKernel(knl)
        from scipy.sparse import csr_matrix
        n = 64
        mat = csr_matrix(np.ones((64, 64)) * (np.random.rand(64, 64) < 0.1))
        row = mat.indptr.astype(np.uintc)
        col = mat.indices.astype(np.uintc)
        dat = mat.data.astype(np.float32)
        out, vec = np.random.rand(2, n).astype(np.float32)
        nnz = mat.nnz
        cknl(n, nnz, row, col, dat, vec, out)
        np.testing.assert_allclose(out, mat * vec, )




class TestCompiledKernel(TestCase):

    def test_simple_kernel(self):
        knl = lp.make_kernel("{ [i]: 0<=i<n }", "out[i] = 2*a[i]", target=CTarget())
        typed = lp.add_dtypes(knl, {'a': np.float32})
        code, _ = lp.generate_code(typed)
        fn = CompiledKernel(typed)
        a, out = np.zeros((2, 10), np.float32)
        a[:] = np.r_[:a.size]
        fn(a, 10, out)
        np.testing.assert_allclose(out, a * 2)


class TestLogProb(TestCase):

    def setUp(self):
        super().setUp()
        self.model = _TestModel()<commit_msg>loosen prec on sp mat test<commit_after>#     Copyright 2017 TVB-HPC contributors
#
#     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 time
import logging
import unittest
import loopy as lp
from loopy.target.c import CTarget
import pymbolic as pm
import numpy as np
import numpy.testing
from scipy.stats import kstest
from .bold import BalloonWindkessel
from .compiler import Compiler, CppCompiler, CompiledKernel, Spec
from .coupling import (Linear as LCf, Diff, Sigmoidal, Kuramoto as KCf,
    PostSumStat, BaseCoupling)
from .model import BaseModel, _TestModel, HMJE, RWW, JansenRit, Linear, G2DO
from .model import Kuramoto
from .network import DenseNetwork, DelayNetwork
from .rng import RNG
from .scheme import euler_maruyama_logp, EulerStep, EulerMaryuyamaStep
# from .harness import SimpleTimeStep
from .utils import getLogger, VarSubst


LOG = logging.getLogger(__name__)


class TestCase(unittest.TestCase):

    def setUp(self):
        super().setUp()
        self.tic = time.time()
        self.logger = getLogger(self.id())
        self._time_limit = 0.1

    def timeit(self, fn, *args, **kwds):
        niter = 0
        tic = toc = time.time()
        while ((toc - tic) < self._time_limit):
            fn(*args, **kwds)
            toc = time.time()
            niter += 1
        per_iter = self._time_limit / niter
        self.logger.info('%r requires %.3f ms / iter',
                         fn, per_iter * 1e3)

    def tearDown(self):
        super().tearDown()
        msg = 'required %.3fs'
        self.logger.info(msg, time.time() - self.tic)


class TestUtils(TestCase):

    def test_var_subst(self):
        subst = VarSubst(b=pm.parse('b[i, j]'))
        expr = subst(pm.parse('a + b * pre_syn[i, j]'))
        self.assertEqual(str(expr), 'a + b[i, j]*pre_syn[i, j]')


class TestLoopTransforms(TestCase):
    """
    These are more tests to check that our use of Loopy is correct.

    """

    def setUp(self):
        super().setUp()
        from loopy.target.ispc import ISPCTarget
        target = ISPCTarget()
        self.knl = lp.make_kernel('{[i]:0<=i<n}', "out[i] = in[i]",
                                  target=target)

    def _dtype_and_code(self, knl):
        knl = lp.add_dtypes(knl, {'in': np.float32, 'out': np.float32})
        code, _ = lp.generate_code(knl)
        return code

    def test_chunk_iname(self):
        "Chunk useful to split work for e.g. omp par for"
        knl = lp.chunk_iname(self.knl, 'i', 8)
        print(self._dtype_and_code(knl))

    def test_split_iname2(self):
        "Split useful for omp simd inner loop"
        knl = lp.split_iname(self.knl, 'i', 8)
        knl = lp.tag_inames(knl, [('i_inner', 'ilp.unr',)])
        print(self._dtype_and_code(knl))

    def test_wrap_loop(self):
        "Take kernel, place in larger loop, offsetting certain vars"
        knl = lp.make_kernel("{[i,j]:0<=i,j<n}",
                             "out[i] = sum(j, (i/j)*in[i, j])",
                             target=CTarget())
        # in will depend on t
        knl2 = lp.to_batched(knl, 'T', ['in'], 't')
        print(self._dtype_and_code(knl2))

    def test_split_iname2(self):
        "Split one of two inames."
        from loopy.target.ispc import ISPCTarget as CTarget
        knl = lp.make_kernel("{[i,j]:0<=i,j<n}",
                             "out[i, j] = in[i, j]",
                             target=CTarget())
        knl = lp.split_iname(knl, 'i', 8)
        knl = lp.prioritize_loops(knl, ['i_outer', 'j', 'i_inner'])
        print(self._dtype_and_code(knl))

    def test_sparse_matmul(self):
        knl = lp.make_kernel(
            [
                '{[i]: 0   <= i <   n}',
                '{[j]: jlo <= j < jhi}'
            ],
            """
            <> jlo = row[i]
            <> jhi = row[i + 1]
            out[i] = sum(j, dat[j] * vec[col[j]])
            """,
            'n nnz row col dat vec out'.split(),
            target=CTarget())
        knl = lp.add_dtypes(knl, {
            'out': np.float32,
            'dat': np.float32,
            'vec': np.float32,
            'col': np.uintc,
            'row': np.uintc,
            'n': np.uintc,
            'nnz': np.uintc,
        })
        # col and dat have uninferrable shape
        knl.args[3].shape = pm.var('nnz'),
        knl.args[4].shape = pm.var('nnz'),
        cknl = CompiledKernel(knl)
        from scipy.sparse import csr_matrix
        n = 64
        mat = csr_matrix(np.ones((64, 64)) * (np.random.rand(64, 64) < 0.1))
        row = mat.indptr.astype(np.uintc)
        col = mat.indices.astype(np.uintc)
        dat = mat.data.astype(np.float32)
        out, vec = np.random.rand(2, n).astype(np.float32)
        nnz = mat.nnz
        cknl(n, nnz, row, col, dat, vec, out)
        np.testing.assert_allclose(out, mat * vec, 1e-5, 1e-6)




class TestCompiledKernel(TestCase):

    def test_simple_kernel(self):
        knl = lp.make_kernel("{ [i]: 0<=i<n }", "out[i] = 2*a[i]", target=CTarget())
        typed = lp.add_dtypes(knl, {'a': np.float32})
        code, _ = lp.generate_code(typed)
        fn = CompiledKernel(typed)
        a, out = np.zeros((2, 10), np.float32)
        a[:] = np.r_[:a.size]
        fn(a, 10, out)
        np.testing.assert_allclose(out, a * 2)


class TestLogProb(TestCase):

    def setUp(self):
        super().setUp()
        self.model = _TestModel()
{'fraction_non_alphanumeric': 0.13399440630561912, 'fraction_numerical': 0.012204424103737605, 'mean_word_length': 2.1677852348993287, 'pattern_counts': {'":': 0, '<': 35, '<?xml version=': 0, '>': 7, 'https://': 0, 'lorem ipsum': 0, 'www.': 2, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 4, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '4355212', 'n_tokens_mistral': 4166, 'n_tokens_neox': 3997, 'n_words': 1024}