"""
Variables are used to define the relationship between trial functions, in both time and approximations, to test
functions and the mapping between both functions and their FEniCS subdomain mapping.
"""
from functools import wraps
from typing import Callable, Iterable, List, NamedTuple, Optional, Union
import dolfin as fem
from ufl.indexed import Indexed # type: ignore
from mtnlion.domain import Domain
[docs]class UndefinedDomain:
"""
Used to represent objects with an undefined domain.
"""
# pylint: disable=too-few-public-methods
def __init__(self, domain: str):
self.domain = domain
def __repr__(self):
return f"{UndefinedDomain.__qualname__}({self.domain})"
def __eq__(self, other) -> bool:
return isinstance(other, UndefinedDomain) and self.domain == other.domain
[docs]class SingleDomainError(Exception):
"""This exception is raised when the number of domains should have been reduced to one for the operation."""
[docs]class UVMap(NamedTuple):
"""
Mapping between the trial functions (u), test functions (v), and the mapping of elements to the trial functions.
Trial functions are represented as a list of functions in time. Each element of the list is a dictionary of the
applicable domains containing a list of related functions. Test functions are not effected by time stepping schemes,
so there is no need to provide a list in time. Similarly, the element types are not allowed to change with each step
in time.
"""
trial_function: List[Domain[str, List[Union[Indexed, UndefinedDomain]]]]
test_function: Domain[str, List[Union[Indexed, UndefinedDomain]]]
element_map: Domain[str, List[Union[fem.FiniteElement, UndefinedDomain]]]
[docs]class UVMapSingleDomain(NamedTuple):
"""
Mapping between the trial functions (u), test functions (v), and the mapping of elements to the trial functions
restricted to a single domain.
Trial functions are represented as a list of functions in time. Each element of the list is a list of related
functions. Test functions are not effected by time stepping schemes, so there is no need to provide a list in time.
Similarly, the element types are not allowed to change with each step in time.
"""
trial_function: List[List[Union[Indexed, UndefinedDomain]]]
test_function: List[Union[Indexed, UndefinedDomain]]
element_map: List[Union[fem.FiniteElement, UndefinedDomain]]
[docs]def check_single(func: Callable) -> Callable:
"""
Used to wrap operations of `Variable` in order to allow variables to be used as parameters in formulas.
:param func: operation to wrap
:return: wrapped function
"""
@wraps(func)
def new_func(*args) -> Callable:
"""
Convert arguments into FEniCS recognized types. The first argument is assumed to be "this" while the second
argument is assumed to be the other object.
If the argument is a `Variable`, retrieve the trial function at the current time step. The `Variable` must be
domain restricted to be usable so the domain doesn't have to be parsed.
:param args: args[0] must be "this" object, while args[0] must be the other.
:return: Callable with parsed/converted arguments
"""
new_args = list(args)
try:
if isinstance(args[1], Variable):
if isinstance(args[1].uv_map, UVMapSingleDomain):
new_args[1] = args[1].uv_map.trial_function[0][0]
else:
raise SingleDomainError(f"{repr(args[1])} must be reduced to a single domain.")
except IndexError:
pass
if isinstance(args[0].uv_map, UVMapSingleDomain):
return func(*new_args)
raise SingleDomainError(f"{repr(args[0])} must be reduced to a single domain.")
return new_func
[docs]class Variable:
"""
This class allows the partial definition of trial functions and test functions allowing them to be defined and used
before the relationships in time, function approximations, and element mappings are defined.
`Variable`'s are able to then define the relationships in formulations without worrying about the specifics of
time discretization schemes or function approximations.
"""
def __init__(self, name: str, domains: Iterable[str], num_functions: int = 1):
"""
Create a `Variable` object with a given name, a list of applicable domains, and the number of functions required
to represent the `Variable`.
:param name: Name of the `Variable` the name must be unique amongst used variables
:param domains: list of domains the `Variable` is defined
:param num_functions: number of functions used to represent this `Variable`
"""
self.name = name
self.domains = domains
self.num_functions = num_functions
self.uv_map: Union[UVMap, UVMapSingleDomain] = UVMap(
[self._init_mapping(self.domains, self.num_functions)],
self._init_mapping(self.domains, self.num_functions),
self._init_mapping(self.domains, self.num_functions),
)
self._index = 0
@staticmethod
def _init_mapping(domains: Iterable[str], num_functions: int) -> Domain[str, List[UndefinedDomain]]:
"""
Initialize the UV maps to None to indicate the variable has not been fully specified.
:param domains: Applicable domains
:param num_functions: Number of functions representing the variable
:return: Mappings to None
"""
return Domain({domain: [UndefinedDomain("None") for _ in range(num_functions)] for domain in domains})
@property
def is_defined(self) -> bool:
"""
Indicates whether or not the variable has been defined by checking the uv map element map.
:return: true if the elements have been defined
"""
if isinstance(self.uv_map, UVMapSingleDomain):
return UndefinedDomain not in (type(element) for element in self.uv_map.element_map)
return UndefinedDomain not in (
type(element) for elements in self.uv_map.element_map.values() for element in elements
)
[docs] def get_domain(self, domain: str) -> "Variable":
"""
Retrieve a `Variable` object defined only on the given domain using `UVMapSingleDomain`.
:param domain: Domain to restrict the variable to
:return: `Variable` defined on a single domain
"""
new_variable = Variable(self.name, [domain], self.num_functions)
if isinstance(self.uv_map, UVMapSingleDomain):
if domain not in self.domains:
new_variable.uv_map = UVMapSingleDomain(
[[UndefinedDomain(domain)]], [UndefinedDomain(domain)], [UndefinedDomain(domain)]
)
else:
new_variable.uv_map = self.uv_map
else:
new_variable.uv_map = UVMapSingleDomain(
[data.get(domain, UndefinedDomain(domain)) for data in self.uv_map.trial_function],
self.uv_map.test_function.get(domain, [UndefinedDomain(domain)]),
self.uv_map.element_map.get(domain, [UndefinedDomain(domain)]),
)
return new_variable
[docs] def trial(
self, history: Optional[int] = None, subfunction: int = 0, all_funcs: bool = False
) -> Union[Optional[Indexed], List[Union[Indexed, UndefinedDomain]], List[List[Union[Indexed, UndefinedDomain]]]]:
"""
Retrieve the trial function for this `Variable`.
History and subfunction constraints may be applied simultaneously. If no constraints are applied, the entire
trial function will be returned.
:param history: Index in time to retrieve the trial function, 0 represents the current time, one is the
previous time step, etc.
:param subfunction: Index of the desired subfunction
:param all_funcs: return all subfunctions
:return: Desired trial function
"""
if not isinstance(self.uv_map, UVMapSingleDomain):
raise SingleDomainError("Variable domain has not been reduced.")
if history is not None:
if all_funcs:
return self.uv_map.trial_function[history]
return self.uv_map.trial_function[history][subfunction]
return [time if all_funcs else time[subfunction] for time in self.uv_map.trial_function]
[docs] def test(self, subfunction: int = 0, all_funcs: bool = False) -> Union[List[Optional[Indexed]], Optional[Indexed]]:
"""
Retrieve the test function for this `Variable`.
If no constraints are applied, the entire test function will be returned.
:param subfunction: Index of the desired subfunction
:param all_funcs: return all subfunctions
:return: Desired test function
"""
if not isinstance(self.uv_map, UVMapSingleDomain):
raise SingleDomainError("Variable domain has not been reduced.")
if all_funcs:
return self.uv_map.test_function
return self.uv_map.test_function[subfunction]
def __eq__(self, other):
if not isinstance(other, Variable):
# don't attempt to compare against unrelated types
raise TypeError
return (
self.name == other.name
and self.domains == other.domains
and self.num_functions == other.num_functions
and self.uv_map == other.uv_map
)
def __iter__(self) -> "Variable":
"""
Allow the `Variable` to be iterable so that operations such as sum(variable) are possible.
:return: self
"""
self._index = 0
return self
def __next__(self) -> Optional[Indexed]:
"""
When iterating through the `Variable`, it is assumed that the first time index is used.
:return: trial function subfunction
"""
if not isinstance(self.uv_map, UVMapSingleDomain):
raise SingleDomainError("Variable domain has not been reduced.")
self._index += 1
try:
return self.uv_map.trial_function[0][self._index - 1] # Only works for first moment of time
except IndexError:
raise StopIteration
@check_single
def __truediv__(self, other):
return self.uv_map.trial_function[0][0] / other # Only works for first moment of time and first subfunction
@check_single
def __rtruediv__(self, other):
return other / self.uv_map.trial_function[0][0] # Only works for first moment of time and first subfunction
@check_single
def __sub__(self, other):
return self.uv_map.trial_function[0][0] - other
@check_single
def __rsub__(self, other):
return other - self.uv_map.trial_function[0][0]
@check_single
def __add__(self, other):
return self.uv_map.trial_function[0][0] + other
@check_single
def __radd__(self, other):
return self.__add__(other)
@check_single
def __mul__(self, other):
return self.uv_map.trial_function[0][0] * other
@check_single
def __rmul__(self, other):
return self.__mul__(other)
@check_single
def __pow__(self, other):
return self.uv_map.trial_function[0][0] ** other
@check_single
def __rpow__(self, other):
return other ** self.uv_map.trial_function[0][0]
def __repr__(self):
return f"{self.name}{self.domains}"