<commit_before>"""An interpreter for FInAT recipes. This is not expected to be
performant, but rather to provide a test facility for FInAT code."""
import pymbolic.primitives as p
from pymbolic.mapper.evaluator import FloatEvaluationMapper, UnknownVariableError
from ast import IndexSum, ForAll, LeviCivita
import numpy as np
import copy


def _as_range(e):
    """Convert a slice to a range."""

    return range(e.start or 0, e.stop, e.step or 1)


class FinatEvaluationMapper(FloatEvaluationMapper):

    def __init__(self, context={}):
        """
        :arg context: a mapping from variable names to values
        """
        super(FinatEvaluationMapper, self).__init__(context)

        self.indices = {}

        # Storage for wave variables while they are out of scope.
        self.wave_vars = {}

    def map_variable(self, expr):
        try:
            var = self.context[expr.name]
            if isinstance(var, p.Expression):
                return self.rec(var)
            else:
                return var
        except KeyError:
            raise UnknownVariableError(expr.name)

    def map_recipe(self, expr):
        """Evaluate expr for all values of free indices"""

        d, b, p = expr.indices

        return self.rec(ForAll(d + b + p, expr.body))

    def map_index_sum(self, expr):

        indices, body = expr.children

        # Sum over multiple indices recursively.
        if len(indices) > 1:
            expr = IndexSum(indices[1:], body)
        else:
            expr = body

        idx = indices[0]

        e = idx.extent

        total = 0.0
        for i in _as_range(e):
            self.indices[idx] = i
            total += self.rec(expr)

        self.indices.pop(idx)

        return total

    def map_index(self, expr):

        return self.indices[expr]

    def map_for_all(self, expr):

        indices, body = expr.children

        # Execute over multiple indices recursively.
        if len(indices) > 1:
            expr = ForAll(indices[1:], body)
        else:
            expr = body

        idx = indices[0]

        e = idx.extent

        total = []
        for i in _as_range(e):
            self.indices[idx] = i
            total.append(self.rec(expr))

        self.indices.pop(idx)

        return np.array(total)

    def map_wave(self, expr):

        (var, index, base, update, body) = expr.children

        try:
            self.context[var] = self.wave_vars[var]
            self.context[var] = self.rec(update)
        except KeyError:
            # We're at the start of the loop over index.
            assert self.rec(index) == index.extent.start
            self.context[var] = self.rec(base)

        self.wave_vars[var] = self.context[var]

        # Execute the body.
        result = self.rec(body)

        # Remove the wave variable from scope.
        self.context.pop(var)
        if self.rec(index) >= index.extent.stop - 1:
            self.wave_vars.pop(var)

        return result

    def map_levi_civita(self, expr):

        free, bound, body = expr.children

        if len(bound) == 3:
            return self.rec(IndexSum(bound[:1],
                                     LeviCivita(bound[:1], bound[1:], body)))
        elif len(bound) == 2:

            self.indices[bound[0]] = (self.indices[free[0]] + 1) % 3
            self.indices[bound[1]] = (self.indices[free[0]] + 2) % 3
            tmp = self.rec(body)
            self.indices[bound[1]] = (self.indices[free[0]] + 1) % 3
            self.indices[bound[0]] = (self.indices[free[0]] + 2) % 3
            tmp -= self.rec(body)

            self.indices.pop(bound[0])
            self.indices.pop(bound[1])
<commit_msg>draft implementation of interpreter for Let<commit_after>"""An interpreter for FInAT recipes. This is not expected to be
performant, but rather to provide a test facility for FInAT code."""
import pymbolic.primitives as p
from pymbolic.mapper.evaluator import FloatEvaluationMapper, UnknownVariableError
from ast import IndexSum, ForAll, LeviCivita
import numpy as np
import copy


def _as_range(e):
    """Convert a slice to a range."""

    return range(e.start or 0, e.stop, e.step or 1)


class FinatEvaluationMapper(FloatEvaluationMapper):

    def __init__(self, context={}):
        """
        :arg context: a mapping from variable names to values
        """
        super(FinatEvaluationMapper, self).__init__(context)

        self.indices = {}

        # Storage for wave variables while they are out of scope.
        self.wave_vars = {}

    def map_variable(self, expr):
        try:
            var = self.context[expr.name]
            if isinstance(var, p.Expression):
                return self.rec(var)
            else:
                return var
        except KeyError:
            raise UnknownVariableError(expr.name)

    def map_recipe(self, expr):
        """Evaluate expr for all values of free indices"""

        d, b, p = expr.indices

        return self.rec(ForAll(d + b + p, expr.body))

    def map_index_sum(self, expr):

        indices, body = expr.children

        # Sum over multiple indices recursively.
        if len(indices) > 1:
            expr = IndexSum(indices[1:], body)
        else:
            expr = body

        idx = indices[0]

        e = idx.extent

        total = 0.0
        for i in _as_range(e):
            self.indices[idx] = i
            total += self.rec(expr)

        self.indices.pop(idx)

        return total

    def map_index(self, expr):

        return self.indices[expr]

    def map_for_all(self, expr):

        indices, body = expr.children

        # Execute over multiple indices recursively.
        if len(indices) > 1:
            expr = ForAll(indices[1:], body)
        else:
            expr = body

        idx = indices[0]

        e = idx.extent

        total = []
        for i in _as_range(e):
            self.indices[idx] = i
            total.append(self.rec(expr))

        self.indices.pop(idx)

        return np.array(total)

    def map_wave(self, expr):

        (var, index, base, update, body) = expr.children

        try:
            self.context[var] = self.wave_vars[var]
            self.context[var] = self.rec(update)
        except KeyError:
            # We're at the start of the loop over index.
            assert self.rec(index) == index.extent.start
            self.context[var] = self.rec(base)

        self.wave_vars[var] = self.context[var]

        # Execute the body.
        result = self.rec(body)

        # Remove the wave variable from scope.
        self.context.pop(var)
        if self.rec(index) >= index.extent.stop - 1:
            self.wave_vars.pop(var)

        return result

    def map_let(self, expr):

        for var, value in expr.bindings:
            if var in self.context:
                raise ValueError("Let variable %s was already in scope."
                                 % var.name)
            self.context[var] = self.rec(value)

        result = self.rec(expr.body)

        for var, value in expr.bindings:
            self.context.pop(var)

        return result

    def map_levi_civita(self, expr):

        free, bound, body = expr.children

        if len(bound) == 3:
            return self.rec(IndexSum(bound[:1],
                                     LeviCivita(bound[:1], bound[1:], body)))
        elif len(bound) == 2:

            self.indices[bound[0]] = (self.indices[free[0]] + 1) % 3
            self.indices[bound[1]] = (self.indices[free[0]] + 2) % 3
            tmp = self.rec(body)
            self.indices[bound[1]] = (self.indices[free[0]] + 1) % 3
            self.indices[bound[0]] = (self.indices[free[0]] + 2) % 3
            tmp -= self.rec(body)

            self.indices.pop(bound[0])
            self.indices.pop(bound[1])

{'fraction_non_alphanumeric': 0.11021402024862233, 'fraction_numerical': 0.008714596949891068, 'mean_word_length': 1.8575613328451117, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 9, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 3, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '5962517', 'n_tokens_mistral': 2406, 'n_tokens_neox': 2257, 'n_words': 624}