Source code for ufl.algorithms.renumbering

"""Algorithms for renumbering of counted objects, currently variables and indices."""
# Copyright (C) 2008-2024 Martin Sandve Alnæs, Anders Logg, Jørgen S. Dokken and Lawrence Mitchell
#
# This file is part of UFL (https://www.fenicsproject.org)
#
# SPDX-License-Identifier:    LGPL-3.0-or-later

from collections import defaultdict
from itertools import count as _count

from ufl.algorithms.map_integrands import map_integrand_dags
from ufl.core.multiindex import Index
from ufl.corealg.multifunction import MultiFunction


class IndexRelabeller(MultiFunction):
    """Renumber indices to have a consistent index numbering starting from 0."""

    def __init__(self):
        """Initialize index relabeller with a zero count."""
        super().__init__()
        count = _count()
        self.index_cache = defaultdict(lambda: Index(next(count)))

    expr = MultiFunction.reuse_if_untouched

    def multi_index(self, o):
        """Apply to multi-indices."""
        return type(o)(
            tuple(self.index_cache[i] if isinstance(i, Index) else i for i in o.indices())
        )

    def zero(self, o):
        """Apply to zero."""
        fi = o.ufl_free_indices
        fid = o.ufl_index_dimensions
        new_indices = [self.index_cache[Index(i)].count() for i in fi]
        if fi == () and fid == ():
            return o
        new_fi, new_fid = zip(*sorted(zip(new_indices, fid), key=lambda x: x[0]))
        return type(o)(o.ufl_shape, tuple(new_fi), tuple(new_fid))


def renumber_indices(form):
    """Renumber indices to have a consistent index numbering starting from 0.

    This is useful to avoid multiple kernels for the same integrand,
    but with different subdomain ids.

    Args:
        form: A UFL form, integral or expression.

    Returns:
        A new form, integral or expression with renumbered indices.
    """
    reindexer = IndexRelabeller()
    return map_integrand_dags(reindexer, form)