Linear elasticity

src/dctkit/apps/linear_elasticity.py

@@ -0,0 +1,57 @@
+import numpy as np
+import numpy.typing as npt
+from dctkit.mesh.simplex import SimplicialComplex
+import dctkit.dec.cochain as C
+import dctkit.dec.vector as V
+from jax import Array
+import jax.numpy as jnp
+from typing import Tuple
+
+
+class LinearElasticity():
+    def __init__(self, S: SimplicialComplex, mu_: float, lambda_: float):
+        self.S = S
+        self.mu_ = mu_
+        self.lambda_ = lambda_
+
+    def __get_current_metric_2D(self, node_coords: npt.Array | Array):
+        """Compute the multiarray of shape (n, 2, 2) where n is the number of
+            2-simplices and any 2x2 matrix is the metric of the corresponding 2-simplex.
+        """
+        dim = self.S.dim
+        B = self.S.B[dim]
+        primal_edges = self.S.S[1]
+        # construct the matrix in which the i-th row corresponds to the vector
+        # of coordinates of the i-th primal edge
+        primal_edge_vectors = node_coords[primal_edges[:, 1], :] - \
+            node_coords[primal_edges[:, 0], :]
+        # construct the multiarray of shape (n, 3, 3) where any 3x3 matrix represents
+        # the coordinates of the edge vectors (arranged in rows) belonging to the
+        # corresponding 2-simplex
+        primal_edges_per_2_simplex = primal_edge_vectors[B]
+        # extract the first two rows, i.e. basis vectors, for each 3x3 matrix
+        basis_vectors = primal_edges_per_2_simplex[:, :-1, :]
+        metric = basis_vectors @ np.transpose(
+            basis_vectors, axes=(0, 2, 1))
+        return metric
+
+    def linear_elasticity_residual(self, node_coords: C.CochainP0,
+                                   f: C.CochainP2) -> C.CochainP2:
+        dim = self.S.dim
+        g = self.__get_current_metric_2D(node_coords=node_coords.coeffs)
+        g_tensor = V.DiscreteTensorFieldD(S=self.S, coeffs=g)
+        epsilon = 1/2 * (g_tensor - self.S.metric)
+        tr_epsilon = jnp.trace(epsilon, axis1=1, axis2=2)
+        stress = 2*self.mu*epsilon + self.lambda_*tr_epsilon[:, None, None] * \
+            np.stack([np.identity(2)]*dim)
+        stress_integrated = V.flat_DPD(stress)
+        residual = C.coboundary(C.star(stress_integrated)) + f
+        return residual
+
+    def obj_linear_elasticity(self, node_coords: npt.Array, f: npt.Array, gamma: float,
+                              boundary_values: Tuple[npt.NDArray, npt.NDArray]):
+        pos, value = boundary_values
+        node_coords_coch = C.CochainP0(S=self.S, coeffs=node_coords)
+        f_coch = C.CochainP2(complex=self.S, coeffs=f)
+        residual = self.linear_elasticity_residual(node_coords_coch, f_coch).coeffs
+        penalty = np.sum((node_coords[pos] - value)**2)

src/dctkit/dec/cochain.py

@@ -293,11 +293,12 @@ def star(c: Cochain) -> Cochain:
                      coeffs=jnp.empty_like(c.coeffs, dtype=dt.float_dtype))
 
     if c.is_primal:
-        star_c.coeffs = c.complex.hodge_star[c.dim]*c.coeffs
+        star_c.coeffs = (c.complex.hodge_star[c.dim]*c.coeffs.T).T
     else:
         # NOTE: this step only works with well-centered meshes!
         assert c.complex.is_well_centered
-        star_c.coeffs = c.complex.hodge_star_inverse[c.complex.dim - c.dim]*c.coeffs
+        star_c.coeffs = (
+            c.complex.hodge_star_inverse[c.complex.dim - c.dim]*c.coeffs.T).T
 
     return star_c
 

src/dctkit/dec/vector.py

@@ -22,14 +22,29 @@ def __init__(self, S: spx.SimplicialComplex, is_primal: bool,
         self.coeffs = coeffs
 
 
+class DiscreteTensorField():
+    def __init__(self, S: spx.SimplicialComplex, is_primal: bool,
+                 coeffs: npt.NDArray | Array):
+        self.S = S
+        self.is_primal = is_primal
+        self.coeffs = coeffs
+
+
 class DiscreteVectorFieldD(DiscreteVectorField):
     """Inherited class for dual discrete vector fields."""
 
     def __init__(self, S: spx.SimplicialComplex, coeffs: npt.NDArray | Array):
         super().__init__(S, False, coeffs)
 
 
-def flat_DPD(v: DiscreteVectorFieldD) -> CochainD1:
+class DiscreteTensorFieldD(DiscreteTensorField):
+    """Inherited class for dual discrete tensor fields."""
+
+    def __init__(self, S: spx.SimplicialComplex, coeffs: npt.NDArray | Array):
+        super().__init__(S, False, coeffs)
+
+
+def flat_DPD(v: DiscreteVectorFieldD | DiscreteTensorFieldD) -> CochainD1:
     """Implements the flat operator for dual discrete vector fields.
 
     Args:
@@ -40,7 +55,17 @@ def flat_DPD(v: DiscreteVectorFieldD) -> CochainD1:
     dedges = v.S.dual_edges_vectors
     flat_matrix = v.S.flat_weights
     # multiply weights of each dual edge by the vectors associated to the dual nodes
-    # belonging to the edge and then perform dot product row-wise with the edge vectors
-    # of the dual edges (see definition of DPD in Hirani, pag. 54).
-    coch_coeffs = jnp.sum((v.coeffs @ flat_matrix).T * dedges, axis=1)
+    # belonging to the edge
+    weighted_v = v.coeffs @ flat_matrix
+    weighted_v_T = weighted_v.T
+    if v.coeffs.ndim == 2:
+        # vector field case
+        # perform dot product row-wise with the edge vectors
+        # of the dual edges (see definition of DPD in Hirani, pag. 54).
+        coch_coeffs = jnp.einsum("ij, ij -> i", weighted_v_T, dedges)
+    elif v.coeffs.ndim == 3:
+        # tensor field case
+        # apply each matrix (rows of the multiarray weighted_v_T fixing the first axis)
+        # to the edge vector of the corresponding dual edge
+        coch_coeffs = jnp.einsum("ijk, ik -> ij", weighted_v_T, dedges)
     return CochainD1(v.S, coch_coeffs)

src/dctkit/mesh/simplex.py

@@ -256,6 +256,11 @@ def get_flat_weights(self):
                 dual_edges_indices] = np.linalg.norm(diff_circs, axis=1)
 
         self.flat_weights = self.dual_edges_fractions_lengths/self.dual_edges_lengths
+        # in the case of non-well centered mesh an entry of the flat weights matrix
+        # can be NaN. In this case, the corresponding dual edge is the null vector,
+        # hence we shouldn't take in account dot product with it. We then substitute
+        # any NaN with 0.
+        self.flat_weights = np.nan_to_num(self.flat_weights)
 
     def get_metric_2D(self):
         """Compute the multiarray of shape (n, 2, 2) where n is the number of

tests/test_cochain.py

@@ -123,3 +123,30 @@ def test_cochain(setup_test):
 
     assert inner_product.dtype == dctkit.float_dtype
     assert np.allclose(inner_product, inner_product_true)
+
+    # vector-valued cochain test
+    c0_v_coeffs = np.arange(15).reshape((5, 3))
+    c1_v_coeffs = np.arange(24).reshape((8, 3))
+    c0_v = cochain.CochainP0(cpx, c0_v_coeffs)
+    c1_v = cochain.CochainP1(cpx, c1_v_coeffs)
+    dc0_v = cochain.coboundary(c0_v)
+    dc1_v = cochain.coboundary(c1_v)
+    dc0_v_true = np.array([[3, 3, 3],
+                           [6,  6,  6],
+                           [12, 12, 12],
+                           [6, 6, 6],
+                           [9, 9, 9],
+                           [3, 3, 3],
+                           [6, 6, 6],
+                           [3, 3, 3]], dtype=dctkit.float_dtype)
+    dc1_v_true = np.array([[6, 7, 8],
+                           [-15, -16, -17],
+                           [18, 19, 20],
+                           [-18, -19, -20]], dtype=dctkit.float_dtype)
+    assert np.allclose(dc0_v.coeffs, dc0_v_true)
+    assert np.allclose(dc1_v.coeffs, dc1_v_true)
+
+    star_c0_v = cochain.star(c0_v)
+    star_c0_v_true = 0.125*c0_v_coeffs
+    star_c0_v_true[-1, :] = np.array([6, 6.5, 7])
+    assert np.allclose(star_c0_v.coeffs, star_c0_v_true)

tests/test_vector.py

@@ -1,4 +1,5 @@
 import numpy as np
+import dctkit as dt
 from dctkit.mesh import simplex, util
 import dctkit.dec.vector as V
 import jax.numpy as jnp
@@ -15,9 +16,16 @@ def test_vector_cochain(setup_test):
     S.get_flat_weights()
 
     # test flat operator
-    v_coeffs = np.ones((S.embedded_dim, S.S[2].shape[0]))
+    v_coeffs = np.ones((S.embedded_dim, S.S[2].shape[0]), dtype=dt.float_dtype)
+    T_coeffs = np.ones((S.embedded_dim, S.embedded_dim,
+                       S.S[2].shape[0]), dtype=dt.float_dtype)
     v = V.DiscreteVectorFieldD(S, v_coeffs)
-    c = V.flat_DPD(v)
-    c_true_coeffs = S.dual_edges_vectors.sum(axis=1)
+    T = V.DiscreteTensorFieldD(S, T_coeffs)
+    c_v = V.flat_DPD(v)
+    c_T = V.flat_DPD(T)
+    c_v_true_coeffs = S.dual_edges_vectors.sum(axis=1)
+    c_T_true_coeffs = np.ones((12, 3), dtype=dt.float_dtype)
+    c_T_true_coeffs = c_v_true_coeffs[:, None]*c_T_true_coeffs
 
-    assert jnp.allclose(c.coeffs, c_true_coeffs)
+    assert jnp.allclose(c_v.coeffs, c_v_true_coeffs)
+    assert jnp.allclose(c_T.coeffs, c_T_true_coeffs)