<commit_before>import numpy as nm

try:
    import dask.array as da

except ImportError:
    da = None

try:
    import opt_einsum as oe

except ImportError:
    oe = None

try:
    from jax.config import config
    config.update("jax_enable_x64", True)
    import jax
    import jax.numpy as jnp

except ImportError:
    jnp = jax = None

from pyparsing import (Word, Suppress, oneOf, OneOrMore, delimitedList,
                       Combine, alphas, Literal)

from sfepy.base.base import output, Struct
from sfepy.base.timing import Timer
from sfepy.discrete import FieldVariable
from sfepy.mechanics.tensors import dim2sym
from sfepy.terms.terms import Term
from sfepy.terms import register_term

def _get_char_map(c1, c2):
    mm = {}
    for ic, char in enumerate(c1):
        if char in mm:
            print(char, '->eq?', mm[char], c2[ic])
            if mm[char] != c2[ic]:
                mm[char] += c2[ic]
        else:
            mm[char] = c2[ic]

    return mm

def append_all(seqs, item, ii=None):
    if ii is None:
        for seq in seqs:
            seq.append(item)

    else:
        seqs[ii].append(item)

def get_sizes(indices, operands):
    sizes = {}

    for iis, op in zip(indices, operands):
        for ii, size in zip(iis, op.shape):
            sizes[ii] = size

    return sizes

class ExpressionBuilder(Struct):
    letters = 'defgh'
    _aux_letters = 'rstuvwxyz'

    def __init__(self, n_add, dc_type, region, cache):
        self.n_add = n_add
        self.subscripts = [[] for ia in range(n_add)]
        self.operands = [[] for ia in range(n_add)]
        self.operand_names = [[] for ia in range(n_add)]
        self.out_subscripts = ['c' for ia in range(n_add)]
        self.ia = 0
        self.dc_type = dc_type
        self.region = region
        self.cache = cache
        self.aux_letters = iter(self._aux_letters)

    def make_eye(self, size):
        key = 'I{}'.format(size)
        ee = self.cache.get(key)
        if ee is None:
            ee = nm.eye(size)
            self.cache[key] = ee

        return ee

    def make_psg(self, dim):
        key = 'Psg{}'.format(dim)
        psg = self.cache.get(key)
        if psg is None:
            sym = dim2sym(dim)
            psg = nm.zeros((dim, dim, sym))
            if dim == 3:
                psg[0, [0,1,2], [0,3,4]] = 1
                psg[1, [0,1,2], [3,1,5]] = 1
                psg[2, [0,1,2], [4,5,2]] = 1

            elif dim == 2:
                psg[0, [0,1], [0,2]] = 1
                psg[1, [0,1], [2,1]] = 1

            self.cache[key] = psg

        return psg

    def add_constant(self, val, name):
        append_all(self.subscripts, 'cq')
        append_all(self.operands, val)
        append_all(self.operand_names, name)

    def add_bfg(self, iin, ein, qsbg, name):
        append_all(self.subscripts, 'cq{}{}'.format(ein[2], iin))
        append_all(self.operands, qsbg)
        append_all(self.operand_names, name)

    def add_bf(self, iin, ein, qsb, name):
        key = 'qsb{}'.format(id(qsb))
        _qsb  = self.cache.get(key)
        if qsb.shape[0] > 1: # cell-depending basis.
            append_all(self.subscripts, 'cq{}'.format(iin))
            if _qsb is None:
                _qsb = qsb[:, :, 0]
                self.cache[key] = _qsb

        else:
            append_all(self.subscripts, 'q{}'.format(iin))
            if _qsb is None:
                _qsb = qsb[0, :, 0]
                self.cache[key] = _qsb

        append_all(self.operands, _qsb)
        append_all(self.operand_names, name)

    def add_eye(self, iic, ein, eye, iia=None):
        append_all(self.subscripts, '{}{}'.format(ein[0], iic), ii=iia)
        append_all(self.operands, eye, ii=iia)
        append_all(self.operand_names, 'I', ii=iia)

    def add_psg(self, iic, ein, psg, iia=None):
        append_all(self.subscripts, '{}{}{}'.format(iic, ein[2], ein[0]),
                   ii=iia)
        append_all(self.operands, psg, ii=iia)
        append_all(self.operand_names, 'Psg', ii=iia)

    def add_arg_dofs(self, iin, ein, arg, iia=None):
        dofs = self.cache.get(arg.name)
        if dofs is None:
            conn = arg.field.get_econn(self.dc_type, self.region)
            dofs_vec = arg().reshape((-1, arg.n_components))
            # axis 0: cells, axis 1: node, axis 2: component
            dofs = dofs_vec[conn]
            self.cache[arg.name] = dofs

        if arg.n_components > 1:
            term = 'c{}{}'.format(iin, ein[0])

        else:
            dofs.shape = (dofs.shape[0], -1)
            term = 'c{}'.format(iin)

        append_all(self.subscripts, term, ii=iia)
        append_all(self.operands, dofs, ii=iia)
        append_all(self.operand_names, arg.name + '.dofs', ii=iia)

    def add_virtual_arg(self, arg, ii, ein, qsb, qsbg, modifier):
        iin = self.letters[ii] # node (qs basis index)<commit_msg>use faster(?) dofs layout in ExpressionBuilder.add_arg_dofs() (WIP)<commit_after>import numpy as nm

try:
    import dask.array as da

except ImportError:
    da = None

try:
    import opt_einsum as oe

except ImportError:
    oe = None

try:
    from jax.config import config
    config.update("jax_enable_x64", True)
    import jax
    import jax.numpy as jnp

except ImportError:
    jnp = jax = None

from pyparsing import (Word, Suppress, oneOf, OneOrMore, delimitedList,
                       Combine, alphas, Literal)

from sfepy.base.base import output, Struct
from sfepy.base.timing import Timer
from sfepy.discrete import FieldVariable
from sfepy.mechanics.tensors import dim2sym
from sfepy.terms.terms import Term
from sfepy.terms import register_term

def _get_char_map(c1, c2):
    mm = {}
    for ic, char in enumerate(c1):
        if char in mm:
            print(char, '->eq?', mm[char], c2[ic])
            if mm[char] != c2[ic]:
                mm[char] += c2[ic]
        else:
            mm[char] = c2[ic]

    return mm

def append_all(seqs, item, ii=None):
    if ii is None:
        for seq in seqs:
            seq.append(item)

    else:
        seqs[ii].append(item)

def get_sizes(indices, operands):
    sizes = {}

    for iis, op in zip(indices, operands):
        for ii, size in zip(iis, op.shape):
            sizes[ii] = size

    return sizes

class ExpressionBuilder(Struct):
    letters = 'defgh'
    _aux_letters = 'rstuvwxyz'

    def __init__(self, n_add, dc_type, region, cache):
        self.n_add = n_add
        self.subscripts = [[] for ia in range(n_add)]
        self.operands = [[] for ia in range(n_add)]
        self.operand_names = [[] for ia in range(n_add)]
        self.out_subscripts = ['c' for ia in range(n_add)]
        self.ia = 0
        self.dc_type = dc_type
        self.region = region
        self.cache = cache
        self.aux_letters = iter(self._aux_letters)

    def make_eye(self, size):
        key = 'I{}'.format(size)
        ee = self.cache.get(key)
        if ee is None:
            ee = nm.eye(size)
            self.cache[key] = ee

        return ee

    def make_psg(self, dim):
        key = 'Psg{}'.format(dim)
        psg = self.cache.get(key)
        if psg is None:
            sym = dim2sym(dim)
            psg = nm.zeros((dim, dim, sym))
            if dim == 3:
                psg[0, [0,1,2], [0,3,4]] = 1
                psg[1, [0,1,2], [3,1,5]] = 1
                psg[2, [0,1,2], [4,5,2]] = 1

            elif dim == 2:
                psg[0, [0,1], [0,2]] = 1
                psg[1, [0,1], [2,1]] = 1

            self.cache[key] = psg

        return psg

    def add_constant(self, val, name):
        append_all(self.subscripts, 'cq')
        append_all(self.operands, val)
        append_all(self.operand_names, name)

    def add_bfg(self, iin, ein, qsbg, name):
        append_all(self.subscripts, 'cq{}{}'.format(ein[2], iin))
        append_all(self.operands, qsbg)
        append_all(self.operand_names, name)

    def add_bf(self, iin, ein, qsb, name):
        key = 'qsb{}'.format(id(qsb))
        _qsb  = self.cache.get(key)
        if qsb.shape[0] > 1: # cell-depending basis.
            append_all(self.subscripts, 'cq{}'.format(iin))
            if _qsb is None:
                _qsb = qsb[:, :, 0]
                self.cache[key] = _qsb

        else:
            append_all(self.subscripts, 'q{}'.format(iin))
            if _qsb is None:
                _qsb = qsb[0, :, 0]
                self.cache[key] = _qsb

        append_all(self.operands, _qsb)
        append_all(self.operand_names, name)

    def add_eye(self, iic, ein, eye, iia=None):
        append_all(self.subscripts, '{}{}'.format(ein[0], iic), ii=iia)
        append_all(self.operands, eye, ii=iia)
        append_all(self.operand_names, 'I', ii=iia)

    def add_psg(self, iic, ein, psg, iia=None):
        append_all(self.subscripts, '{}{}{}'.format(iic, ein[2], ein[0]),
                   ii=iia)
        append_all(self.operands, psg, ii=iia)
        append_all(self.operand_names, 'Psg', ii=iia)

    def add_arg_dofs(self, iin, ein, arg, iia=None):
        dofs = self.cache.get(arg.name)
        if dofs is None:
            conn = arg.field.get_econn(self.dc_type, self.region)
            dofs_vec = arg().reshape((-1, arg.n_components))
            # # axis 0: cells, axis 1: node, axis 2: component
            # dofs = dofs_vec[conn]
            # axis 0: cells, axis 1: component, axis 2: node
            dofs = dofs_vec[conn].transpose((0, 2, 1))
            self.cache[arg.name] = dofs

        if arg.n_components > 1:
            #term = 'c{}{}'.format(iin, ein[0])
            term = 'c{}{}'.format(ein[0], iin)

        else:
            dofs.shape = (dofs.shape[0], -1)
            term = 'c{}'.format(iin)

        append_all(self.subscripts, term, ii=iia)
        append_all(self.operands, dofs, ii=iia)
        append_all(self.operand_names, arg.name + '.dofs', ii=iia)

    def add_virtual_arg(self, arg, ii, ein, qsb, qsbg, modifier):
        iin = self.letters[ii] # node (qs basis index)
{'fraction_non_alphanumeric': 0.13507014028056114, 'fraction_numerical': 0.014128256513026053, 'mean_word_length': 2.0494958753437214, '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': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '3908585', 'n_tokens_mistral': 3751, 'n_tokens_neox': 3611, 'n_words': 845}