Source code for ufl.indexsum

"""This module defines the IndexSum class."""

# Copyright (C) 2008-2016 Martin Sandve Alnæs
#
# This file is part of UFL (https://www.fenicsproject.org)
#
# SPDX-License-Identifier:    LGPL-3.0-or-later

from ufl.algebra import Product
from ufl.constantvalue import Zero
from ufl.core.expr import Expr, ufl_err_str
from ufl.core.multiindex import MultiIndex
from ufl.core.operator import Operator
from ufl.core.ufl_type import ufl_type
from ufl.indexed import Indexed
from ufl.precedence import parstr

# --- Sum over an index ---




@ufl_type(num_ops=2)
class IndexSum(Operator):
    """Index sum."""

    __slots__ = ("_dimension", "_initialised", "ufl_free_indices", "ufl_index_dimensions")

    def __new__(cls, summand, index):
        """Create a new IndexSum."""
        # Error checks
        if not isinstance(summand, Expr):
            raise ValueError(f"Expecting Expr instance, got {ufl_err_str(summand)}")
        if not isinstance(index, MultiIndex):
            raise ValueError(f"Expecting MultiIndex instance, got {ufl_err_str(index)}")
        if len(index) != 1:
            raise ValueError(f"Expecting a single Index but got {len(index)}.")

        (j,) = index
        # Simplification to zero
        if isinstance(summand, Zero):
            sh = summand.ufl_shape
            fi = summand.ufl_free_indices
            fid = summand.ufl_index_dimensions
            pos = fi.index(j.count())
            fi = fi[:pos] + fi[pos + 1 :]
            fid = fid[:pos] + fid[pos + 1 :]
            return Zero(sh, fi, fid)

        # Factor out common factors
        if isinstance(summand, Product):
            a, b = summand.ufl_operands
            if j.count() not in a.ufl_free_indices:
                return Product(a, IndexSum(b, index))
            elif j.count() not in b.ufl_free_indices:
                return Product(b, IndexSum(a, index))

        self = Operator.__new__(cls)
        self._initialised = False
        return self

    def __init__(self, summand, index):
        """Initialise."""
        if self._initialised:
            return
        (j,) = index
        fi = summand.ufl_free_indices
        fid = summand.ufl_index_dimensions
        pos = fi.index(j.count())
        self._dimension = fid[pos]
        self.ufl_free_indices = fi[:pos] + fi[pos + 1 :]
        self.ufl_index_dimensions = fid[:pos] + fid[pos + 1 :]
        Operator.__init__(self, (summand, index))
        self._initialised = True



    def index(self):
        """Get index."""
        return self.ufl_operands[1][0]




    def dimension(self):
        """Get dimension."""
        return self._dimension


    @property
    def ufl_shape(self):
        """Get UFL shape."""
        return self.ufl_operands[0].ufl_shape

    def _simplify_indexed(self, multiindex):
        """Return a simplified Expr used in the constructor of Indexed(self, multiindex)."""
        A, i = self.ufl_operands
        return IndexSum(Indexed(A, multiindex), i)



    def evaluate(self, x, mapping, component, index_values):
        """Evaluate."""
        (i,) = self.ufl_operands[1]
        tmp = 0
        for k in range(self._dimension):
            index_values.push(i, k)
            tmp += self.ufl_operands[0].evaluate(x, mapping, component, index_values)
            index_values.pop()
        return tmp


    def __str__(self):
        """Format as a string."""
        return f"sum_{{{self.ufl_operands[1]!s}}} {parstr(self.ufl_operands[0], self)} "