[WIP] First working test using petsc4py.

src/dctkit/experimental/petsc_test.py

@@ -0,0 +1,100 @@
+import petsc4py
+from petsc4py import PETSc
+import numpy as np
+from jax import grad, jit, value_and_grad
+from dctkit.mesh import util
+import dctkit as dt
+from dctkit.dec import cochain as C
+import time
+
+dt.config()
+
+lc = 0.008
+
+mesh, _ = util.generate_square_mesh(lc)
+S = util.build_complex_from_mesh(mesh)
+S.get_hodge_star()
+bnodes = mesh.cell_sets_dict["boundary"]["line"]
+node_coord = S.node_coords
+
+# NOTE: exact solution of Delta u + f = 0
+u_true = np.array(node_coord[:, 0]**2 + node_coord[:, 1] ** 2, dtype=dt.float_dtype)
+b_values = u_true[bnodes]
+
+boundary_values = (np.array(bnodes, dtype=dt.int_dtype), b_values)
+
+num_nodes = S.num_nodes
+print(num_nodes)
+f_vec = -4.*np.ones(num_nodes, dtype=dt.float_dtype)
+
+u_0 = np.zeros(num_nodes, dtype=dt.float_dtype)
+
+gamma = 1000.
+
+
+def energy_poisson(x, f, boundary_values, gamma):
+    pos, value = boundary_values
+    f = C.Cochain(0, True, S, f)
+    u = C.Cochain(0, True, S, x)
+    du = C.coboundary(u)
+    norm_grad = 1/2.*C.inner_product(du, du)
+    bound_term = -C.inner_product(u, f)
+    penalty = 0.5*gamma*dt.backend.sum((x[pos] - value)**2)
+    energy = norm_grad + bound_term + penalty
+    return energy
+
+
+args = (f_vec, boundary_values, gamma)
+
+energy_jit = jit(energy_poisson)
+objgrad = jit(grad(energy_poisson))
+objandgrad = jit(value_and_grad(energy_poisson))
+# wrappers for the objective function and its gradient following the signature of
+# TAOObjectiveFunction
+
+
+def objective_function(tao, x, f, boundary_values, gamma):
+    return energy_jit(x.getArray(), f, boundary_values, gamma)
+
+
+def gradient_function(tao, x, g, f, boundary_values, gamma):
+    g_jax = objgrad(x.getArray(), f, boundary_values, gamma)
+    g.setArray(g_jax)
+
+
+def objective_and_gradient(tao, x, g, f, boundary_values, gamma):
+    fval, grad_jax = objandgrad(x.getArray(), f, boundary_values, gamma)
+    g.setArray(grad_jax)
+    return fval
+
+
+# Initialize petsc4py
+petsc4py.init()
+
+# Create a PETSc vector to hold the optimization variables
+x = PETSc.Vec().createWithArray(u_0)
+
+# Create a PETSc TAO object for the solver
+tao = PETSc.TAO().create()
+tao.setType(PETSc.TAO.Type.LMVM)  # Specify the solver type
+tao.setSolution(x)
+# tao.setObjective(objective_function, args=args)  # Set the objective function
+g = PETSc.Vec().createSeq(num_nodes)
+# tao.setGradient(gradient_function, g, args=args)  # Set the gradient function
+tao.setObjectiveGradient(objective_and_gradient, g, args=args)
+tao.setMaximumIterations(500)
+# tao.setTolerances(gatol=1e-3)
+tao.setFromOptions()  # Set options for the solver
+
+tic = time.time()
+# Minimize the function using the Nonlinear CG method
+tao.solve()
+toc = time.time()
+tao.view()
+print("Elapsed time = ", toc - tic)
+
+# Get the solution and the objective value
+u = tao.getSolution()
+objective_value = tao.getObjectiveValue()
+
+assert np.allclose(u, u_true, atol=1e-2)

src/dctkit/mesh/circumcenter.py

@@ -13,13 +13,13 @@ def circumcenter(S: npt.NDArray,
         (Reference: Bell, Hirani, PyDEC: Software and Algorithms for Discretization
         of Exterior Calculus, 2012, Section 10.1).
 
-        Args:
-            S: matrix containing the IDs of the nodes belonging to each simplex.
-            node_coords: coordinates (cols) of each node of the complex to which the
-                simplices belongs.
-        Returns:
-            a tuple consisting of the Cartesian coordinates of the circumcenter and the
-                barycentric coordinates.
+    Args:
+        S: matrix containing the IDs of the nodes belonging to each simplex.
+        node_coords: coordinates (cols) of each node of the complex to which the
+            simplices belongs.
+    Returns:
+        a tuple consisting of the Cartesian coordinates of the circumcenter and the
+        barycentric coordinates.
     """
     # get global data type
     float_dtype = dctkit.float_dtype

src/dctkit/mesh/simplex.py

@@ -26,8 +26,8 @@ class SimplicialComplex:
         circ (list): list where each entry p is a matrix containing the
             coordinates of the circumcenters (cols) of all the p-simplexes (rows).
         boundary (list): list of the boundary matrices at all dimensions (0..dim-1).
-        node_coords (float np.array): Cartesian coordinates (cols) of the
-            nodes (rows) of the simplicial complex.
+        node_coords (npt.NDArray): Cartesian coordinates (cols) of the nodes (rows) of
+            the simplicial complex.
         primal_volumes (list): list where each entry p is an array containing all the
             volumes of the primal p-simplices.
         dual_volumes (list): list where each entry p is an array containing all
@@ -77,7 +77,8 @@ def get_boundary_operators(self):
 
     def get_complex_boundary_faces_indices(self):
         """Find the IDs of the boundary faces of the complex, i.e. the row indices of
-        the boundary faces in the matrix S[dim-1]."""
+        the boundary faces in the matrix S[dim-1].
+        """
         # boundary faces IDs appear only once in the matrix simplices_faces[dim]
         unique_elements, counts = np.unique(
             self.simplices_faces[self.dim], return_counts=True)
@@ -107,7 +108,6 @@ def get_primal_volumes(self):
 
     def get_dual_volumes(self):
         """Compute all the dual volumes."""
-
         if not hasattr(self, "circ"):
             self.get_circumcenters()
 
@@ -267,8 +267,8 @@ def get_current_metric_2D(self, node_coords: npt.NDArray | Array) -> Array:
 
             Returns:
                 the multiarray of shape (n, 2, 2), where n is the number of 2-simplices
-                    and each 2x2 matrix is the current metric of the corresponding
-                    2-simplex.
+                and each 2x2 matrix is the current metric of the corresponding
+                2-simplex.
         """
         # NOTATION:
         # a_i, reference covariant basis (pairs of edge vectors of a primal 2-simplex)
@@ -360,9 +360,9 @@ def compute_boundary_COO(S: npt.NDArray) -> Tuple[list, npt.NDArray, npt.NDArray
 
     Returns:
         a tuple containing a list with the COO representation of the boundary, the
-            matrix of node IDs belonging to each (p-1)-face ordered lexicographically,
-            and a matrix containing the IDs of the nodes (cols) belonging to
-            each p-simplex (rows) counted with repetition and ordered lexicographically.
+        matrix of node IDs belonging to each (p-1)-face ordered lexicographically,
+        and a matrix containing the IDs of the nodes (cols) belonging to
+        each p-simplex (rows) counted with repetition and ordered lexicographically.
 
     """
     # number of p-simplices
@@ -427,7 +427,7 @@ def compute_simplices_faces(S: npt.NDArray, faces_ordered:
 
     Returns:
         a matrix containing the IDs of the (p-1)-simplices (cols) belonging
-            to each p-simplex (rows).
+        to each p-simplex (rows).
 
     """