"""This module provides basic mathematical functions."""
# 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
# Modified by Kristian B. Oelgaard, 2011
import cmath
import math
import numbers
import warnings
from ufl.constantvalue import (
ComplexValue,
ConstantValue,
FloatValue,
IntValue,
RealValue,
Zero,
as_ufl,
is_true_ufl_scalar,
)
from ufl.core.operator import Operator
from ufl.core.ufl_type import ufl_type
"""
TODO: Include additional functions available in <cmath> (need derivatives as well):
Exponential and logarithmic functions:
log10 Compute common logarithm (function)
TODO: Any other useful special functions?
About bessel functions:
http://en.wikipedia.org/wiki/Bessel_function
Portable implementations of bessel functions:
http://www.boost.org/doc/libs/1_47_0/libs/math/doc/sf_and_dist/html/math_toolkit/main_overview/tr1.html
Implementation in C++ std::tr1:: or boost::math::tr1::
- BesselK: cyl_bessel_k(nu, x)
- BesselI: cyl_bessel_i(nu, x)
- BesselJ: cyl_bessel_j(nu, x)
- BesselY: cyl_neumann(nu, x)
"""
# --- Function representations ---
[docs]
@ufl_type(is_abstract=True, is_scalar=True, num_ops=1)
class MathFunction(Operator):
"""Base class for all unary scalar math functions."""
# Freeze member variables for objects in this class
__slots__ = ("_name",)
def __init__(self, name, argument):
"""Initialise."""
Operator.__init__(self, (argument,))
if not is_true_ufl_scalar(argument):
raise ValueError("Expecting scalar argument.")
self._name = name
[docs]
def evaluate(self, x, mapping, component, index_values):
"""Evaluate."""
a = self.ufl_operands[0].evaluate(x, mapping, component, index_values)
try:
if isinstance(a, numbers.Real):
res = getattr(math, self._name)(a)
else:
res = getattr(cmath, self._name)(a)
except ValueError:
warnings.warn(
"Value error in evaluation of function %s with argument %s." % (self._name, a)
)
raise
return res
def __str__(self):
"""Format as a string."""
return "%s(%s)" % (self._name, self.ufl_operands[0])
[docs]
@ufl_type()
class Sqrt(MathFunction):
"""Square root."""
__slots__ = ()
def __new__(cls, argument):
"""Create a new Sqrt."""
if isinstance(argument, (RealValue, Zero, numbers.Real)):
if float(argument) < 0:
return ComplexValue(cmath.sqrt(complex(argument)))
else:
return FloatValue(math.sqrt(float(argument)))
if isinstance(argument, (ComplexValue, complex)):
return ComplexValue(cmath.sqrt(complex(argument)))
return MathFunction.__new__(cls)
def __init__(self, argument):
"""Initialise."""
MathFunction.__init__(self, "sqrt", argument)
[docs]
@ufl_type()
class Exp(MathFunction):
"""Exponentiation.."""
__slots__ = ()
def __new__(cls, argument):
"""Create a new Exp."""
if isinstance(argument, (RealValue, Zero)):
return FloatValue(math.exp(float(argument)))
if isinstance(argument, (ComplexValue)):
return ComplexValue(cmath.exp(complex(argument)))
return MathFunction.__new__(cls)
def __init__(self, argument):
"""Initialise."""
MathFunction.__init__(self, "exp", argument)
[docs]
@ufl_type()
class Ln(MathFunction):
"""Natural logarithm."""
__slots__ = ()
def __new__(cls, argument):
"""Create a new Ln."""
if isinstance(argument, (RealValue, Zero)):
return FloatValue(math.log(float(argument)))
if isinstance(argument, (ComplexValue)):
return ComplexValue(cmath.log(complex(argument)))
return MathFunction.__new__(cls)
def __init__(self, argument):
"""Initialise."""
MathFunction.__init__(self, "ln", argument)
[docs]
def evaluate(self, x, mapping, component, index_values):
"""Evaluate."""
a = self.ufl_operands[0].evaluate(x, mapping, component, index_values)
try:
return math.log(a)
except TypeError:
return cmath.log(a)
[docs]
@ufl_type()
class Cos(MathFunction):
"""Cosine."""
__slots__ = ()
def __new__(cls, argument):
"""Create a new Cos."""
if isinstance(argument, (RealValue, Zero)):
return FloatValue(math.cos(float(argument)))
if isinstance(argument, (ComplexValue)):
return ComplexValue(cmath.cos(complex(argument)))
return MathFunction.__new__(cls)
def __init__(self, argument):
"""Initialise."""
MathFunction.__init__(self, "cos", argument)
[docs]
@ufl_type()
class Sin(MathFunction):
"""Sine."""
__slots__ = ()
def __new__(cls, argument):
"""Create a new Sin."""
if isinstance(argument, (RealValue, Zero)):
return FloatValue(math.sin(float(argument)))
if isinstance(argument, (ComplexValue)):
return ComplexValue(cmath.sin(complex(argument)))
return MathFunction.__new__(cls)
def __init__(self, argument):
"""Initialise."""
MathFunction.__init__(self, "sin", argument)
[docs]
@ufl_type()
class Tan(MathFunction):
"""Tangent."""
__slots__ = ()
def __new__(cls, argument):
"""Create a new Tan."""
if isinstance(argument, (RealValue, Zero)):
return FloatValue(math.tan(float(argument)))
if isinstance(argument, (ComplexValue)):
return ComplexValue(cmath.tan(complex(argument)))
return MathFunction.__new__(cls)
def __init__(self, argument):
"""Initialise."""
MathFunction.__init__(self, "tan", argument)
[docs]
@ufl_type()
class Cosh(MathFunction):
"""Hyperbolic cosine."""
__slots__ = ()
def __new__(cls, argument):
"""Create a new Cosh."""
if isinstance(argument, (RealValue, Zero)):
return FloatValue(math.cosh(float(argument)))
if isinstance(argument, (ComplexValue)):
return ComplexValue(cmath.cosh(complex(argument)))
return MathFunction.__new__(cls)
def __init__(self, argument):
"""Initialise."""
MathFunction.__init__(self, "cosh", argument)
[docs]
@ufl_type()
class Sinh(MathFunction):
"""Hyperbolic sine."""
__slots__ = ()
def __new__(cls, argument):
"""Create a new Sinh."""
if isinstance(argument, (RealValue, Zero)):
return FloatValue(math.sinh(float(argument)))
if isinstance(argument, (ComplexValue)):
return ComplexValue(cmath.sinh(complex(argument)))
return MathFunction.__new__(cls)
def __init__(self, argument):
"""Initialise."""
MathFunction.__init__(self, "sinh", argument)
[docs]
@ufl_type()
class Tanh(MathFunction):
"""Hyperbolic tangent."""
__slots__ = ()
def __new__(cls, argument):
"""Create a new Tanh."""
if isinstance(argument, (RealValue, Zero)):
return FloatValue(math.tanh(float(argument)))
if isinstance(argument, (ComplexValue)):
return ComplexValue(cmath.tanh(complex(argument)))
return MathFunction.__new__(cls)
def __init__(self, argument):
"""Initialise."""
MathFunction.__init__(self, "tanh", argument)
[docs]
@ufl_type()
class Acos(MathFunction):
"""Inverse cosine."""
__slots__ = ()
def __new__(cls, argument):
"""Create a new Acos."""
if isinstance(argument, (RealValue, Zero)):
return FloatValue(math.acos(float(argument)))
if isinstance(argument, (ComplexValue)):
return ComplexValue(cmath.acos(complex(argument)))
return MathFunction.__new__(cls)
def __init__(self, argument):
"""Initialise."""
MathFunction.__init__(self, "acos", argument)
[docs]
@ufl_type()
class Asin(MathFunction):
"""Inverse sine."""
__slots__ = ()
def __new__(cls, argument):
"""Create a new Asin."""
if isinstance(argument, (RealValue, Zero)):
return FloatValue(math.asin(float(argument)))
if isinstance(argument, (ComplexValue)):
return ComplexValue(cmath.asin(complex(argument)))
return MathFunction.__new__(cls)
def __init__(self, argument):
"""Initialise."""
MathFunction.__init__(self, "asin", argument)
[docs]
@ufl_type()
class Atan(MathFunction):
"""Inverse tangent."""
__slots__ = ()
def __new__(cls, argument):
"""Create a new Atan."""
if isinstance(argument, (RealValue, Zero)):
return FloatValue(math.atan(float(argument)))
if isinstance(argument, (ComplexValue)):
return ComplexValue(cmath.atan(complex(argument)))
return MathFunction.__new__(cls)
def __init__(self, argument):
"""Initialise."""
MathFunction.__init__(self, "atan", argument)
[docs]
@ufl_type(is_scalar=True, num_ops=2)
class Atan2(Operator):
"""Inverse tangent with two inputs."""
__slots__ = ()
def __new__(cls, arg1, arg2):
"""Create a new Atan2."""
if isinstance(arg1, (RealValue, Zero)) and isinstance(arg2, (RealValue, Zero)):
return FloatValue(math.atan2(float(arg1), float(arg2)))
if isinstance(arg1, (ComplexValue)) or isinstance(arg2, (ComplexValue)):
raise TypeError("Atan2 does not support complex numbers.")
return Operator.__new__(cls)
def __init__(self, arg1, arg2):
"""Initialise."""
Operator.__init__(self, (arg1, arg2))
if isinstance(arg1, (ComplexValue, complex)) or isinstance(arg2, (ComplexValue, complex)):
raise TypeError("Atan2 does not support complex numbers.")
if not is_true_ufl_scalar(arg1):
raise ValueError("Expecting scalar argument 1.")
if not is_true_ufl_scalar(arg2):
raise ValueError("Expecting scalar argument 2.")
[docs]
def evaluate(self, x, mapping, component, index_values):
"""Evaluate."""
a = self.ufl_operands[0].evaluate(x, mapping, component, index_values)
b = self.ufl_operands[1].evaluate(x, mapping, component, index_values)
try:
res = math.atan2(a, b)
except TypeError:
raise ValueError("Atan2 does not support complex numbers.")
except ValueError:
warnings.warn(
"Value error in evaluation of function atan2 with arguments %s, %s." % (a, b)
)
raise
return res
def __str__(self):
"""Format as a string."""
return "atan2(%s,%s)" % (self.ufl_operands[0], self.ufl_operands[1])
[docs]
@ufl_type()
class Erf(MathFunction):
"""Erf function."""
__slots__ = ()
def __new__(cls, argument):
"""Create a new Erf."""
if isinstance(argument, (RealValue, Zero)):
return FloatValue(math.erf(float(argument)))
if isinstance(argument, (ConstantValue)):
return ComplexValue(math.erf(complex(argument)))
return MathFunction.__new__(cls)
def __init__(self, argument):
"""Initialise."""
MathFunction.__init__(self, "erf", argument)
[docs]
def evaluate(self, x, mapping, component, index_values):
"""Evaluate."""
a = self.ufl_operands[0].evaluate(x, mapping, component, index_values)
return math.erf(a)
[docs]
@ufl_type(is_abstract=True, is_scalar=True, num_ops=2)
class BesselFunction(Operator):
"""Base class for all bessel functions."""
__slots__ = "_name"
def __init__(self, name, nu, argument):
"""Initialise."""
if not is_true_ufl_scalar(nu):
raise ValueError("Expecting scalar nu.")
if not is_true_ufl_scalar(argument):
raise ValueError("Expecting scalar argument.")
# Use integer representation if suitable
fnu = float(nu)
inu = int(nu)
if fnu == inu:
nu = as_ufl(inu)
else:
nu = as_ufl(fnu)
Operator.__init__(self, (nu, argument))
self._name = name
[docs]
def evaluate(self, x, mapping, component, index_values):
"""Evaluate."""
a = self.ufl_operands[1].evaluate(x, mapping, component, index_values)
try:
import scipy.special
except ImportError:
raise ValueError(
"You must have scipy installed to evaluate bessel functions in python."
)
name = self._name[-1]
if isinstance(self.ufl_operands[0], IntValue):
nu = int(self.ufl_operands[0])
functype = "n" if name != "i" else "v"
else:
nu = self.ufl_operands[0].evaluate(x, mapping, component, index_values)
functype = "v"
func = getattr(scipy.special, name + functype)
return func(nu, a)
def __str__(self):
"""Format as a string."""
return "%s(%s, %s)" % (self._name, self.ufl_operands[0], self.ufl_operands[1])
[docs]
@ufl_type()
class BesselJ(BesselFunction):
"""Bessel J function."""
__slots__ = ()
def __init__(self, nu, argument):
"""Initialise."""
BesselFunction.__init__(self, "cyl_bessel_j", nu, argument)
[docs]
@ufl_type()
class BesselY(BesselFunction):
"""Bessel Y function."""
__slots__ = ()
def __init__(self, nu, argument):
"""Initialise."""
BesselFunction.__init__(self, "cyl_bessel_y", nu, argument)
[docs]
@ufl_type()
class BesselI(BesselFunction):
"""Bessel I function."""
__slots__ = ()
def __init__(self, nu, argument):
"""Initialise."""
BesselFunction.__init__(self, "cyl_bessel_i", nu, argument)
[docs]
@ufl_type()
class BesselK(BesselFunction):
"""Bessel K function."""
__slots__ = ()
def __init__(self, nu, argument):
"""Initialise."""
BesselFunction.__init__(self, "cyl_bessel_k", nu, argument)