"""Basic algebra operations."""
# 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
#
# Modified by Anders Logg, 2008
from ufl.checks import is_true_ufl_scalar, is_ufl_scalar
from ufl.constantvalue import ComplexValue, IntValue, ScalarValue, Zero, as_ufl, zero
from ufl.core.expr import ufl_err_str
from ufl.core.operator import Operator
from ufl.core.ufl_type import ufl_type
from ufl.index_combination_utils import merge_unique_indices
from ufl.precedence import parstr
from ufl.sorting import sorted_expr
# --- Algebraic operators ---
[docs]@ufl_type(
num_ops=2,
inherit_shape_from_operand=0,
inherit_indices_from_operand=0,
binop="__add__",
rbinop="__radd__",
)
class Sum(Operator):
"""Sum."""
__slots__ = ()
def __new__(cls, a, b):
"""Create a new Sum."""
# Make sure everything is an Expr
a = as_ufl(a)
b = as_ufl(b)
# Assert consistent tensor properties
sh = a.ufl_shape
fi = a.ufl_free_indices
fid = a.ufl_index_dimensions
if b.ufl_shape != sh:
raise ValueError("Can't add expressions with different shapes.")
if b.ufl_free_indices != fi:
raise ValueError("Can't add expressions with different free indices.")
if b.ufl_index_dimensions != fid:
raise ValueError("Can't add expressions with different index dimensions.")
# Skip adding zero
if isinstance(a, Zero):
return b
elif isinstance(b, Zero):
return a
# Handle scalars specially and sort operands
sa = isinstance(a, ScalarValue)
sb = isinstance(b, ScalarValue)
if sa and sb:
# Apply constant propagation
return as_ufl(a._value + b._value)
elif sa:
# Place scalar first
# operands = (a, b)
pass # a, b = a, b
elif sb:
# Place scalar first
# operands = (b, a)
a, b = b, a
# elif a == b:
# # Replace a+b with 2*foo
# return 2*a
else:
# Otherwise sort operands in a canonical order
# operands = (b, a)
a, b = sorted_expr((a, b))
# construct and initialize a new Sum object
self = Operator.__new__(cls)
self._init(a, b)
return self
def _init(self, a, b):
"""Initialise."""
self.ufl_operands = (a, b)
def __init__(self, a, b):
"""Initialise."""
Operator.__init__(self)
[docs] def evaluate(self, x, mapping, component, index_values):
"""Evaluate."""
return sum(o.evaluate(x, mapping, component, index_values) for o in self.ufl_operands)
def __str__(self):
"""Format as a string."""
return " + ".join([parstr(o, self) for o in self.ufl_operands])
[docs]@ufl_type(num_ops=2, binop="__mul__", rbinop="__rmul__")
class Product(Operator):
"""The product of two or more UFL objects."""
__slots__ = ("ufl_free_indices", "ufl_index_dimensions")
def __new__(cls, a, b):
"""Create a new product."""
# Conversion
a = as_ufl(a)
b = as_ufl(b)
# Type checking
# Make sure everything is scalar
if a.ufl_shape or b.ufl_shape:
raise ValueError(
"Product can only represent products of scalars, "
f"got\n {ufl_err_str(a)}\nand\n {ufl_err_str(b)}"
)
# Simplification
if isinstance(a, Zero) or isinstance(b, Zero):
# Got any zeros? Return zero.
fi, fid = merge_unique_indices(
a.ufl_free_indices,
a.ufl_index_dimensions,
b.ufl_free_indices,
b.ufl_index_dimensions,
)
return Zero((), fi, fid)
sa = isinstance(a, ScalarValue)
sb = isinstance(b, ScalarValue)
if sa and sb: # const * const = const
# FIXME: Handle free indices like with zero? I think
# IntValue may be index annotated now?
return as_ufl(a._value * b._value)
elif sa: # 1 * b = b
if a._value == 1:
return b
# a, b = a, b
elif sb: # a * 1 = a
if b._value == 1:
return a
a, b = b, a
# elif a == b: # a * a = a**2 # TODO: Why? Maybe just remove this?
# if not a.ufl_free_indices:
# return a**2
else: # a * b = b * a
# Sort operands in a semi-canonical order
# (NB! This is fragile! Small changes here can have large effects.)
a, b = sorted_expr((a, b))
# Construction
self = Operator.__new__(cls)
self._init(a, b)
return self
def _init(self, a, b):
"""Constructor, called by __new__ with already checked arguments."""
self.ufl_operands = (a, b)
# Extract indices
fi, fid = merge_unique_indices(
a.ufl_free_indices, a.ufl_index_dimensions, b.ufl_free_indices, b.ufl_index_dimensions
)
self.ufl_free_indices = fi
self.ufl_index_dimensions = fid
def __init__(self, a, b):
"""Initialise."""
Operator.__init__(self)
ufl_shape = ()
[docs] def evaluate(self, x, mapping, component, index_values):
"""Evaluate."""
ops = self.ufl_operands
sh = self.ufl_shape
if sh:
if sh != ops[-1].ufl_shape:
raise ValueError(
"Expecting nonscalar product operand to be the last by convention."
)
tmp = ops[-1].evaluate(x, mapping, component, index_values)
ops = ops[:-1]
else:
tmp = 1
for o in ops:
tmp *= o.evaluate(x, mapping, (), index_values)
return tmp
def __str__(self):
"""Format as a string."""
a, b = self.ufl_operands
return " * ".join((parstr(a, self), parstr(b, self)))
[docs]@ufl_type(num_ops=2, inherit_indices_from_operand=0, binop="__div__", rbinop="__rdiv__")
class Division(Operator):
"""Division."""
__slots__ = ()
def __new__(cls, a, b):
"""Create a new Division."""
# Conversion
a = as_ufl(a)
b = as_ufl(b)
# Type checking
# TODO: Enabled workaround for nonscalar division in __div__,
# so maybe we can keep this assertion. Some algorithms may
# need updating.
if not is_ufl_scalar(a):
raise ValueError("Expecting scalar nominator in Division.")
if not is_true_ufl_scalar(b):
raise ValueError("Division by non-scalar is undefined.")
if isinstance(b, Zero):
raise ValueError("Division by zero!")
# Simplification
# Simplification a/b -> a
if isinstance(a, Zero) or (isinstance(b, ScalarValue) and b._value == 1):
return a
# Simplification "literal a / literal b" -> "literal value of
# a/b". Avoiding integer division by casting to float
if isinstance(a, ScalarValue) and isinstance(b, ScalarValue):
try:
return as_ufl(float(a._value) / float(b._value))
except TypeError:
return as_ufl(complex(a._value) / complex(b._value))
# Simplification "a / a" -> "1"
# if not a.ufl_free_indices and not a.ufl_shape and a == b:
# return as_ufl(1)
# Construction
self = Operator.__new__(cls)
self._init(a, b)
return self
def _init(self, a, b):
"""Initialise."""
self.ufl_operands = (a, b)
def __init__(self, a, b):
"""Initialise."""
Operator.__init__(self)
ufl_shape = () # self.ufl_operands[0].ufl_shape
[docs] def evaluate(self, x, mapping, component, index_values):
"""Evaluate."""
a, b = self.ufl_operands
a = a.evaluate(x, mapping, component, index_values)
b = b.evaluate(x, mapping, component, index_values)
# Avoiding integer division by casting to float
try:
e = float(a) / float(b)
except TypeError:
e = complex(a) / complex(b)
return e
def __str__(self):
"""Format as a string."""
return f"{parstr(self.ufl_operands[0], self)} / {parstr(self.ufl_operands[1], self)}"
[docs]@ufl_type(num_ops=2, inherit_indices_from_operand=0, binop="__pow__", rbinop="__rpow__")
class Power(Operator):
"""Power."""
__slots__ = ()
def __new__(cls, a, b):
"""Create new Power."""
# Conversion
a = as_ufl(a)
b = as_ufl(b)
# Type checking
if not is_true_ufl_scalar(a):
raise ValueError(f"Cannot take the power of a non-scalar expression {ufl_err_str(a)}.")
if not is_true_ufl_scalar(b):
raise ValueError(f"Cannot raise an expression to a non-scalar power {ufl_err_str(b)}.")
# Simplification
if isinstance(a, ScalarValue) and isinstance(b, ScalarValue):
return as_ufl(a._value**b._value)
if isinstance(b, Zero):
return IntValue(1)
if isinstance(a, Zero) and isinstance(b, ScalarValue):
if isinstance(b, ComplexValue):
raise ValueError("Cannot raise zero to a complex power.")
bf = float(b)
if bf < 0:
raise ValueError("Division by zero, cannot raise 0 to a negative power.")
else:
return zero()
if isinstance(b, ScalarValue) and b._value == 1:
return a
# Construction
self = Operator.__new__(cls)
self._init(a, b)
return self
def _init(self, a, b):
"""Initialise."""
self.ufl_operands = (a, b)
def __init__(self, a, b):
"""Initialise."""
Operator.__init__(self)
ufl_shape = ()
[docs] def evaluate(self, x, mapping, component, index_values):
"""Evalute."""
a, b = self.ufl_operands
a = a.evaluate(x, mapping, component, index_values)
b = b.evaluate(x, mapping, component, index_values)
return a**b
def __str__(self):
"""Format as a string."""
a, b = self.ufl_operands
return f"{parstr(a, self)} ** {parstr(b, self)}"
[docs]@ufl_type(num_ops=1, inherit_shape_from_operand=0, inherit_indices_from_operand=0, unop="__abs__")
class Abs(Operator):
"""Absolute value."""
__slots__ = ()
def __new__(cls, a):
"""Create a new Abs."""
a = as_ufl(a)
# Simplification
if isinstance(a, (Zero, Abs)):
return a
if isinstance(a, Conj):
return Abs(a.ufl_operands[0])
if isinstance(a, ScalarValue):
return as_ufl(abs(a._value))
return Operator.__new__(cls)
def __init__(self, a):
"""Initialise."""
Operator.__init__(self, (a,))
[docs] def evaluate(self, x, mapping, component, index_values):
"""Evaluate."""
a = self.ufl_operands[0].evaluate(x, mapping, component, index_values)
return abs(a)
def __str__(self):
"""Format as a string."""
(a,) = self.ufl_operands
return f"|{parstr(a, self)}|"
[docs]@ufl_type(num_ops=1, inherit_shape_from_operand=0, inherit_indices_from_operand=0)
class Conj(Operator):
"""Complex conjugate."""
__slots__ = ()
def __new__(cls, a):
"""Creatr a new Conj."""
a = as_ufl(a)
# Simplification
if isinstance(a, (Abs, Real, Imag, Zero)):
return a
if isinstance(a, Conj):
return a.ufl_operands[0]
if isinstance(a, ScalarValue):
return as_ufl(a._value.conjugate())
return Operator.__new__(cls)
def __init__(self, a):
"""Initialise."""
Operator.__init__(self, (a,))
[docs] def evaluate(self, x, mapping, component, index_values):
"""Evaluate."""
a = self.ufl_operands[0].evaluate(x, mapping, component, index_values)
return a.conjugate()
def __str__(self):
"""Format as a string."""
(a,) = self.ufl_operands
return f"conj({parstr(a, self)})"
[docs]@ufl_type(num_ops=1, inherit_shape_from_operand=0, inherit_indices_from_operand=0)
class Real(Operator):
"""Real part."""
__slots__ = ()
def __new__(cls, a):
"""Create a new Real."""
a = as_ufl(a)
# Simplification
if isinstance(a, Conj):
a = a.ufl_operands[0]
if isinstance(a, Zero):
return a
if isinstance(a, ScalarValue):
return as_ufl(a.real())
if isinstance(a, Real):
a = a.ufl_operands[0]
return Operator.__new__(cls)
def __init__(self, a):
"""Initialise."""
Operator.__init__(self, (a,))
[docs] def evaluate(self, x, mapping, component, index_values):
"""Evaluate."""
a = self.ufl_operands[0].evaluate(x, mapping, component, index_values)
return a.real
def __str__(self):
"""Format as a string."""
(a,) = self.ufl_operands
return f"Re[{parstr(a, self)}]"
[docs]@ufl_type(num_ops=1, inherit_shape_from_operand=0, inherit_indices_from_operand=0)
class Imag(Operator):
"""Imaginary part."""
__slots__ = ()
def __new__(cls, a):
"""Create a new Imag."""
a = as_ufl(a)
# Simplification
if isinstance(a, Zero):
return a
if isinstance(a, (Real, Imag, Abs)):
return Zero(a.ufl_shape, a.ufl_free_indices, a.ufl_index_dimensions)
if isinstance(a, ScalarValue):
return as_ufl(a.imag())
return Operator.__new__(cls)
def __init__(self, a):
"""Initialise."""
Operator.__init__(self, (a,))
[docs] def evaluate(self, x, mapping, component, index_values):
"""Evaluate."""
a = self.ufl_operands[0].evaluate(x, mapping, component, index_values)
return a.imag
def __str__(self):
"""Format as a string."""
(a,) = self.ufl_operands
return f"Im[{parstr(a, self)}]"