1"""Element formulations: the Discrete Kirchhoff Triangle for bending and the constant-strain
2triangle for the Marguerre-von Karman membrane strains.
4Nodal degrees of freedom are (u, v, w, theta_x, theta_y) with theta_x = dw/dy and
5theta_y = -dw/dx (the DKT convention of Batoz, Bathe and Ho, 1980). Bending enters through a
6precomputed 9 x 9 element stiffness acting on (w, theta_x, theta_y) of the three nodes; the
7membrane energy is evaluated from the constant gradients of u, v and w on the element together
8with the slope of the initial surface. See docs/model.md, section "Spatial discretization".
9"""
10from __future__ import annotations
12import jax
13import jax.numpy as jnp
16def dkt_stiffness(xe, ye, D, nu):
17 """9 x 9 DKT bending stiffness for one triangle.
19 ``xe``, ``ye``: nodal coordinates (3,), ``D`` the bending rigidity E h^3 / 12 (1 - nu^2).
20 DOF order is (w1, thx1, thy1, w2, thx2, thy2, w3, thx3, thy3). The curvature operator is
21 integrated with the three-point rule, which is exact for the quadratic rotation field.
22 """
23 i, j = jnp.array([1, 2, 0]), jnp.array([2, 0, 1]) # sides k = 4, 5, 6 <-> ij = 23, 31, 12
24 xij, yij = xe[i] - xe[j], ye[i] - ye[j]
25 l2 = xij ** 2 + yij ** 2
26 a = -xij / l2
27 b = 0.75 * xij * yij / l2
28 c = (0.25 * xij ** 2 - 0.5 * yij ** 2) / l2
29 d = -yij / l2
30 e = (0.25 * yij ** 2 - 0.5 * xij ** 2) / l2
31 a4, a5, a6 = a; b4, b5, b6 = b; c4, c5, c6 = c; d4, d5, d6 = d; e4, e5, e6 = e
33 def H(s):
34 """Rotation shape functions (beta_x, beta_y) = H(xi, eta) . U, shape (2, 9)."""
35 xi, eta = s
36 lam = 1 - xi - eta
37 N1, N2, N3 = lam * (2 * lam - 1), xi * (2 * xi - 1), eta * (2 * eta - 1)
38 N4, N5, N6 = 4 * xi * eta, 4 * eta * lam, 4 * xi * lam
39 Hx = jnp.stack([1.5 * (a6 * N6 - a5 * N5), b5 * N5 + b6 * N6, N1 - c5 * N5 - c6 * N6,
40 1.5 * (a4 * N4 - a6 * N6), b6 * N6 + b4 * N4, N2 - c4 * N4 - c6 * N6,
41 1.5 * (a5 * N5 - a4 * N4), b4 * N4 + b5 * N5, N3 - c4 * N4 - c5 * N5])
42 Hy = jnp.stack([1.5 * (d6 * N6 - d5 * N5), -N1 + e5 * N5 + e6 * N6, -b5 * N5 - b6 * N6,
43 1.5 * (d4 * N4 - d6 * N6), -N2 + e4 * N4 + e6 * N6, -b6 * N6 - b4 * N4,
44 1.5 * (d5 * N5 - d4 * N4), -N3 + e4 * N4 + e5 * N5, -b4 * N4 - b5 * N5])
45 return jnp.stack([Hx, Hy])
47 x21, x31 = xe[1] - xe[0], xe[2] - xe[0]
48 y21, y31 = ye[1] - ye[0], ye[2] - ye[0]
49 det = x21 * y31 - x31 * y21 # = 2 A
50 xi_x, xi_y, eta_x, eta_y = y31 / det, -x31 / det, -y21 / det, x21 / det
52 def B_at(s):
53 """Curvature operator (kappa_xx, kappa_yy, 2 kappa_xy) = B(s) . U, shape (3, 9)."""
54 J = jax.jacfwd(H)(s) # (2, 9, 2): d(Hx, Hy)/d(xi, eta)
55 Hx_xi, Hx_eta, Hy_xi, Hy_eta = J[0, :, 0], J[0, :, 1], J[1, :, 0], J[1, :, 1]
56 Hx_x = Hx_xi * xi_x + Hx_eta * eta_x
57 Hx_y = Hx_xi * xi_y + Hx_eta * eta_y
58 Hy_x = Hy_xi * xi_x + Hy_eta * eta_x
59 Hy_y = Hy_xi * xi_y + Hy_eta * eta_y
60 return jnp.stack([Hx_x, Hy_y, Hx_y + Hy_x])
62 Db = D * jnp.array([[1, nu, 0], [nu, 1, 0], [0, 0, 0.5 * (1 - nu)]])
63 gps = jnp.array([[1 / 6, 1 / 6], [2 / 3, 1 / 6], [1 / 6, 2 / 3]])
64 Bs = jax.vmap(B_at)(gps) # (3, 3, 9)
65 return (0.5 * det / 3) * jnp.einsum("gki,kl,glj->ij", Bs, Db, Bs)
68def membrane_strains(qe, Gx, Gy, z0x, z0y, nl):
69 """Marguerre-von Karman membrane strains (e_xx, e_yy, gamma_xy) on one element.
71 ``qe``: nodal (3, 5) values [u, v, w, thx, thy]; ``Gx``, ``Gy``: CST gradient operators (3,);
72 ``z0x``, ``z0y``: slope of the initial surface at the element; ``nl``: 1.0 keeps the quadratic
73 terms, 0.0 gives the linear shallow shell with the same geometry.
74 """
75 ux, uy = Gx @ qe[:, 0], Gy @ qe[:, 0]
76 vx, vy = Gx @ qe[:, 1], Gy @ qe[:, 1]
77 wx, wy = Gx @ qe[:, 2], Gy @ qe[:, 2]
78 exx = ux + nl * 0.5 * wx * wx + z0x * wx
79 eyy = vy + nl * 0.5 * wy * wy + z0y * wy
80 gxy = uy + vx + nl * wx * wy + z0x * wy + z0y * wx
81 return exx, eyy, gxy
84def element_energy_parts(qe, Gx, Gy, A, z0x, z0y, Cm, Ke, nl):
85 """(membrane, bending) strain energy of one element.
87 ``Cm`` = (C11, C12, C33) of the membrane stiffness E h / (1 - nu^2) [[1, nu, 0], [nu, 1, 0],
88 [0, 0, (1 - nu)/2]]; ``Ke`` the 9 x 9 DKT stiffness.
89 """
90 exx, eyy, gxy = membrane_strains(qe, Gx, Gy, z0x, z0y, nl)
91 C11, C12, C33 = Cm
92 Em = 0.5 * A * (C11 * (exx * exx + eyy * eyy) + 2 * C12 * exx * eyy + C33 * gxy * gxy)
93 U = qe[:, 2:].reshape(9)
94 return Em, 0.5 * (U @ Ke @ U)
97def element_energy(*args):
98 Em, Eb = element_energy_parts(*args)
99 return Em + Eb