From 524286d599a870b909ec8ad0775cd92f2e7e2392 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?J=C3=B8rgen=20S=2E=20Dokken?= Date: Fri, 27 Jun 2025 15:18:04 +0000 Subject: [PATCH] Slight fix dolfin adaptors and add dolfinx adaptors --- moola/adaptors/__init__.py | 19 ++- moola/adaptors/adaptor.py | 11 +- moola/adaptors/dolfin_vector.py | 4 +- moola/adaptors/dolfin_vector_set.py | 2 +- moola/adaptors/dolfinx_vector.py | 206 +++++++++++++++++++++++++++ moola/adaptors/dolfinx_vector_set.py | 192 +++++++++++++++++++++++++ 6 files changed, 426 insertions(+), 8 deletions(-) create mode 100644 moola/adaptors/dolfinx_vector.py create mode 100644 moola/adaptors/dolfinx_vector_set.py diff --git a/moola/adaptors/__init__.py b/moola/adaptors/__init__.py index 639b3b7..549d13f 100644 --- a/moola/adaptors/__init__.py +++ b/moola/adaptors/__init__.py @@ -1,14 +1,25 @@ from .numpy_vector import NumpyPrimalVector from .numpy_vector import NumpyDualVector +from .dolfinx_vector import DolfinxPrimalVector +from .dolfinx_vector import DolfinxDualVector +from .dolfinx_vector_set import DolfinxPrimalVectorSet +from .dolfinx_vector_set import DolfinxDualVectorSet + from .dolfin_vector import DolfinPrimalVector from .dolfin_vector import DolfinDualVector -from .dolfin_vector import RieszMap +from .dolfin_vector_set import DolfinPrimalVectorSet +from .dolfin_vector_set import DolfinDualVectorSet +# FIXME: Make these functions backend agnostic. +try: + import dolfinx + from .dolfinx_vector import RieszMap + from .dolfinx_vector_set import RieszMapSet +except ModuleNotFoundError: + from .dolfin_vector import RieszMap + from .dolfin_vector_set import RieszMapSet -from .dolfin_vector_set import DolfinPrimalVectorSet -from .dolfin_vector_set import DolfinDualVectorSet -from .dolfin_vector_set import RieszMapSet from .adaptor import convert_to_moola_dual_vector \ No newline at end of file diff --git a/moola/adaptors/adaptor.py b/moola/adaptors/adaptor.py index 3219899..4156de6 100644 --- a/moola/adaptors/adaptor.py +++ b/moola/adaptors/adaptor.py @@ -1,5 +1,7 @@ from .dolfin_vector import DolfinPrimalVector, DolfinDualVector from .dolfin_vector_set import DolfinDualVectorSet +from .dolfinx_vector import DolfinxPrimalVector, DolfinxDualVector +from .dolfinx_vector_set import DolfinxDualVectorSet from .numpy_vector import NumpyPrimalVector, NumpyDualVector @@ -9,10 +11,17 @@ def convert_to_moola_dual_vector(x, y): """ if isinstance(y, DolfinPrimalVector): r = DolfinDualVector(x, riesz_map=y.riesz_map) + elif isinstance(y, DolfinxPrimalVector): + r = DolfinxDualVector(x, riesz_map=y.riesz_map) elif isinstance(y, NumpyPrimalVector): r = NumpyDualVector(x) else: - r = DolfinDualVectorSet( + try: + r = DolfinxDualVectorSet( + [DolfinxDualVector(xi, riesz_map=yi.riesz_map) for (xi, yi) in zip(x, y.vector_list)], + riesz_map=y.riesz_map) + except: + r = DolfinDualVectorSet( [DolfinDualVector(xi, riesz_map=yi.riesz_map) for (xi, yi) in zip(x, y.vector_list)], riesz_map=y.riesz_map) diff --git a/moola/adaptors/dolfin_vector.py b/moola/adaptors/dolfin_vector.py index e2f9f9c..f02f3bc 100644 --- a/moola/adaptors/dolfin_vector.py +++ b/moola/adaptors/dolfin_vector.py @@ -17,7 +17,7 @@ def __init__(self, V, inner_product="L2", map_operator=None, inverse = "default" self.V = V import dolfin - if inner_product is not "custom": + if inner_product != "custom": u = dolfin.TrialFunction(V) v = dolfin.TestFunction(V) @@ -108,7 +108,7 @@ def __setitem__(self, index, value): def array(self): ''' Returns the vector as a numpy.array object. If local=False, the global array must be returned in a distributed environment. ''' - return self.data.vector().array() + return self.data.vector()[:] def scale(self, s): v = self.data.vector() diff --git a/moola/adaptors/dolfin_vector_set.py b/moola/adaptors/dolfin_vector_set.py index 9a20f74..9c1e45c 100644 --- a/moola/adaptors/dolfin_vector_set.py +++ b/moola/adaptors/dolfin_vector_set.py @@ -34,7 +34,7 @@ def primal_map(self, x, b): else: b[:] = self.riesz_inv * x.vector_list - def dual_map(self, x): + def dual_map(self, x, b): if isinstance(self.riesz_map, ndarray): b[:] = self.riesz_map.dot(x.vector_list) else: diff --git a/moola/adaptors/dolfinx_vector.py b/moola/adaptors/dolfinx_vector.py new file mode 100644 index 0000000..e085cdb --- /dev/null +++ b/moola/adaptors/dolfinx_vector.py @@ -0,0 +1,206 @@ +from moola.linalg import Vector +from moola.misc import events +from math import sqrt + + +class IdentityMap(): + + def primal_map(self, x, b): + x.zero() + x.axpy(1, b) + + def dual_map(self, x): + return x.copy() + +class RieszMap(): + + def __init__(self, V, inner_product="L2", map_operator=None, inverse = "default", + jit_options=None, form_compiler_options=None): + self.V = V + from petsc4py import PETSc + import ufl + import dolfinx.fem.petsc + if inner_product!="custom": + u = ufl.TrialFunction(V) + v = ufl.TestFunction(V) + + + default_forms = {"L2": ufl.inner(u, v)*ufl.dx, + "H0_1": ufl.inner(ufl.grad(u), ufl.grad(v))*ufl.dx, + "H1": (ufl.inner(u, v) + ufl.inner(ufl.grad(u), + ufl.grad(v)))*ufl.dx, + } + + form = dolfinx.fem.form(default_forms[inner_product], jit_options=jit_options,form_compiler_options=form_compiler_options) + map_operator = dolfinx.fem.petsc.assemble_matrix(form) + map_operator.assemble() + self.map_operator = map_operator + self._P = None + if inverse in ("default", "lu"): + self.petsc_options = {"ksp_type": "preonly", "pc_type": "lu", + "pc_factor_mat_solver_type": "mumps", + "ksp_error_if_not_converged": True} + elif inverse == "jacobi": + self.petsc_options = {"ksp_type": "preonly", "pc_type": "jacobi", + "ksp_error_if_not_converged": True} + elif inverse == "sor": + self.petsc_options = {"ksp_type": "preonly", "pc_type": "sor", + "ksp_error_if_not_converged": True} + + elif inverse == "amg": + self.petsc_options = {"ksp_type": "preonly", "pc_type": "hypre"} + + elif isinstance(inverse, PETSc.Mat): + self.petsc_options = {"skp_type": "preonly", "pc_type": "mat"} + self._P = inverse + else: + raise RuntimeError(f"Unknown inverse type {inverse}. ") + self.solver_type = inverse + def primal_map(self, x, b): + from dolfinx_adjoint.petsc_utils import solve_linear_problem + solve_linear_problem(self.map_operator, x, b, petsc_options=self.petsc_options, P=self._P) + + + def dual_map(self, x): + return self.map_operator * x + + +class DolfinxVector(Vector): + ''' An implementation for vectors based on Dolfin data types. ''' + + def __init__(self, data, inner_product="L2", riesz_map=None): + ''' Wraps the Dolfinx object `data` (Function) in a moola.DolfinxVector object. ''' + + if riesz_map is None: + if inner_product=="l2": + riesz_map = IdentityMap() + else: + fn_space = data.function_space + riesz_map = RieszMap(fn_space, inner_product) + self.riesz_map = riesz_map + self.data = data + self.version = 0 + + def bump_version(self): + self.version += 1 + + def __getitem__(self, index): + ''' Returns the value of the (local) index. ''' + return self.data.x.array[index] + + def __setitem__(self, index, value): + ''' Sets the value of the (local) index. ''' + self.data.x.array[index] = value + + def array(self): + ''' Returns the vector as a numpy.array object. If local=False, the + global array must be returned in a distributed environment. ''' + return self.data.x.array + + def scale(self, s): + v = self.data.x.array + v *= s + self.bump_version() + + def __hash__(self): + ''' Returns a hash of the vector ''' + return hash(self.data)*100 + self.version + + def axpy(self, a, x): + ''' Adds a*x to the function. ''' + assert self.__class__ == x.__class__ + v = self.data.x.array + v[:] += a * x.data.x.array[:] + self.bump_version() + + def zero(self): + ''' Zeros the function. ''' + self.data.x.array[:] = 0 + self.bump_version() + return self + + def local_size(self): + ''' Returns the (local) size of the vector. ''' + return self.data.index_map.size_local* self.data.x.block_size + + def size(self): + ''' Returns the (gobal) size of the vector. ''' + return self.data.index_map.size_global* self.data.x.block_size + + def copy(self): + # FIXME: As we don't have a copy method for the dolfinx Function, this is a bit tedious. + from dolfinx_adjoint import Function + from dolfinx_adjoint.blocks.assembly import _vector + func = Function(self.data.function_space, annotate=True) + func.x.array[:] = self.data.x.array.copy() + tt = self.__class__(func, + riesz_map=self.riesz_map) + return tt + + +class DolfinxPrimalVector(DolfinxVector): + """ A class for representing primal vectors. """ + + + def dual(self): + """ Returns the dual representation. """ + + events.increment("Primal -> dual map") + import dolfinx.fem + if isinstance(self.data, dolfinx.fem.Function): + V = self.data.function_space + + dual_vec = self.riesz_map.dual_map(self.data.x) + dual = dolfinx.fem.Function(V, dual_vec) + + return DolfinxDualVector(dual, riesz_map=self.riesz_map) + else: + return self + + def inner(self, vec): + """ Computes the inner product with vec. """ + assert isinstance(vec, DolfinxPrimalVector) + events.increment("Inner product") + + v = self.riesz_map.dual_map(self.data.x) + import dolfinx.cpp + return dolfinx.cpp.la.inner_product(v._cpp_object, self.data.x._cpp_object) + + def norm(self): + """ Computes the vector norm induced by the inner product. """ + + return sqrt(self.inner(self)) + primal_norm = norm + + +class DolfinxDualVector(DolfinxVector): + """ A class for representing dual vectors. """ + + def apply(self, primal): + """ Applies the dual vector to a primal vector. """ + assert isinstance(primal, DolfinxPrimalVector) + import dolfinx.cpp + return dolfinx.cpp.la.inner_product(self.data.x._cpp_object, primal.data.x._cpp_object) + + + + def primal(self): + """ Returns the primal representation. """ + events.increment("Dual -> primal map") + import dolfinx + if isinstance(self.data, dolfinx.fem.Function): + V = self.data.function_space + + dual = dolfinx.fem.Function(V) + self.riesz_map.primal_map(dual.x, self.data.x) + + return DolfinxPrimalVector(dual, riesz_map=self.riesz_map) + else: + return self + + def primal_norm(self): + """ Computes the norm of the primal representation. """ + + return sqrt(self.apply(self.primal())) + +DolfinxLinearFunctional = DolfinxDualVector diff --git a/moola/adaptors/dolfinx_vector_set.py b/moola/adaptors/dolfinx_vector_set.py new file mode 100644 index 0000000..8929959 --- /dev/null +++ b/moola/adaptors/dolfinx_vector_set.py @@ -0,0 +1,192 @@ +from .dolfinx_vector import DolfinxVector, DolfinxPrimalVector, DolfinxDualVector +from moola.linalg import Vector +from moola.misc import events +from math import sqrt +from numpy import zeros, ndarray + +class RieszMapSet(): + + def __init__(self, inner_product="l2", map_operator=None, inverse = "default"): + + if inner_product == "l2": + self.riesz_map = 1 + self.riesz_inv = 1 + + elif inner_product == "custom": + from numpy import ndarray, equal + from scipy.sparse.linalg import LinearOperator, aslinearoperator + if not isinstance(map_operator, (ndarray, LinearOperator)) or not equal(*map_operator.shape): + raise TypeError("only square numpy arrays are currently supported") + self.riesz_map = aslinearoperator(map_operator) + + if inverse in ("default", "lu") and isinstance(map_operator, ndarray): + from numpy.linalg import inv + self.riesz_inv = aslinearoperator(inv(map_operator)) + + else: + self.riesz_inv = inverse + + self.inner_product = inner_product + + def primal_map(self, x, b): + if isinstance(self.riesz_inv, ndarray): + b[:] = self.riesz_inv.dot(x.vector_list) + else: + b[:] = self.riesz_inv * x.vector_list + + def dual_map(self, x, b): + if isinstance(self.riesz_map, ndarray): + b[:] = self.riesz_map.dot(x.vector_list) + else: + b[:] = self.riesz_map * x.vector_list + +class IdentityMapSet(RieszMapSet): + def __init__(self): + RieszMapSet.__init__(self, inner_product = "l2") + +class DolfinxVectorSet(Vector): + + def __init__(self, vector_list, riesz_map = None): + '''An implementation for set of vectors based on FEniCS data types. + Currently supports only simple Riesz maps of numpy.ndarry type. + + Args: + vector_list (list): A list with moola.DolfinVector + + riesz_map (array): An operator to be applied to the vector + of controls prior to the individual control Riesz maps. + + ''' + + for vec in vector_list: + if not isinstance(vec, DolfinxVector): + raise ValueError("vector_list must be a list of DolfinxVectors") + self.vector_list = zeros(len(vector_list), dtype = object) + self.vector_list[:] = vector_list + + if riesz_map == None: + riesz_map = RieszMapSet("l2") + + self.riesz_map = riesz_map + + def __len__(self): + return len(self.vector_list) + + def __getitem__(self, index): + ''' Returns the subvector with given index. ''' + return self.vector_list[index] + + def __setitem__(self, index, value): + ''' Sets the subvector with the given index. ''' + self.vector_list[index] = value + + @property + def data(self): + return [vec.data for vec in self.vector_list] + + def array(self): + ''' Returns the vector as a numpy.array object. If local=False, the + global array must be returned in a distributed environment. ''' + #return self.data.vector().array() + raise NotImplementedError + + def scale(self, s): + [vec.scale(s) for vec in self.vector_list] + + def __hash__(self): + ''' Returns a hash of the vector ''' + hashes = tuple([hash(vec) for vec in self.vector_list]) + return hash(hashes) + + def axpy(self, a, x): + ''' Adds a*x to the function. ''' + assert self.__class__ == x.__class__ + return [vec.axpy(a, xx) for vec, xx in zip(self.vector_list, x.vector_list)] + + + def zero(self): + ''' Zeros the function. ''' + [vec.zero() for vec in self.vector_list] + return self + + def local_size(self): + ''' Returns the (local) size of the vector. ''' + #return self.data.vector().local_size() + raise NotImplementedError + + def size(self): + ''' Returns the (gobal) size of the vector. ''' + return sum([vec.size() for vec in self.vector_list]) + + def copy(self): + vector_list_cpy = [vec.copy() for vec in self.vector_list] + if hasattr(self, "riesz_map"): + return self.__class__(vector_list_cpy, riesz_map = self.riesz_map) + else: + return self.__class__(vector_list_cpy) + + +class DolfinxPrimalVectorSet(DolfinxVectorSet): + """ A class for representing primal vectors. """ + def __init__(self, vector_list, riesz_map = None): + for i, v in enumerate(vector_list): + if not isinstance(v, DolfinxPrimalVector): + raise TypeError("Vector with index {} is not a DolfinPrimalVector.".format(i)) + DolfinxVectorSet.__init__(self, vector_list, riesz_map = riesz_map) + + def dual(self): + """ Returns the dual representation. """ + if self.riesz_map.inner_product == "l2": + return DolfinxDualVectorSet([vec.dual() for vec in self.vector_list], + riesz_map = self.riesz_map) + else: + return DolfinxDualVectorSet([vec.dual() for vec in self.riesz_map.riesz_map * self.vector_list], + riesz_map = self.riesz_map) + + + def inner(self, vec): + """ Computes the inner product with vec. """ + assert isinstance(vec, DolfinxPrimalVectorSet) + events.increment("Inner product") + + return sum([ss.inner(vec) for ss, vec in zip(self.vector_list, vec.vector_list)]) + + def norm(self): + """ Computes the vector norm induced by the inner product. """ + + return sqrt(self.inner(self)) + primal_norm = norm + + +class DolfinxDualVectorSet(DolfinxVectorSet): + """ A class for representing dual vectors. """ + def __init__(self, vector_list, riesz_map = None): + for i, v in enumerate(vector_list): + if not isinstance(v, DolfinxDualVector): + raise TypeError("Vector with index {} is not a DolfinDualVector.".format(i)) + DolfinxVectorSet.__init__(self, vector_list, riesz_map) + + def apply(self, primal): + """ Applies the dual vector to a primal vector. """ + assert isinstance(primal, DolfinxPrimalVectorSet) + return sum([ss.apply(vec) for ss, vec in zip(self.vector_list, primal.vector_list)]) + + def primal(self): + """ Returns the primal representation. """ + #events.increment("Dual -> primal map") + if self.riesz_map.inner_product == "l2": + return DolfinxPrimalVectorSet([vec.primal() for vec in self.vector_list], + riesz_map = self.riesz_map) + else: + primal_vecs = zeros(len(self), dtype = "object") + primal_vecs[:] = [v.primal() for v in self.vector_list] + return DolfinxPrimalVectorSet(self.riesz_map.riesz_inv * primal_vecs, + riesz_map = self.riesz_map) + + + def primal_norm(self): + """ Computes the norm of the primal representation. """ + + return sqrt(self.apply(self.primal())) + +DolfinxLinearFunctional = DolfinxDualVector