Add ABD embedding mesh visualization (#6)

* [add] add a draw abd embedding mesh option for the reduced dof sim script

* added separator and method checks

---------

Co-authored-by: Minchen Li <minchernl@gmail.com>

Author: dwxrycb123

9_reduced_DOF/simulator.py

@@ -4,19 +4,21 @@
 import pygame       # pygame for visualization
 pygame.init()
 
+import math 
 import utils
 import square_mesh   # square mesh
 import time_integrator
 
 # simulation setup
 side_len = 1
-rho = 1000      # density of square
-E = 2e4         # Young's modulus
-nu = 0.4        # Poisson's ratio
-n_seg = 10       # num of segments per side of the square
-h = 0.01        # time step size in s
-DBC = []        # no nodes need to be fixed
-y_ground = -1   # height of the planar ground
+rho = 1000          # density of square
+E = 2e4             # Young's modulus
+nu = 0.4            # Poisson's ratio
+n_seg = 10          # num of segments per side of the square
+h = 0.01            # time step size in s
+DBC = []            # no nodes need to be fixed
+y_ground = -1       # height of the planar ground
+draw_abd = True     # whether to draw the embedding mesh for ABD 
 
 # initialize simulation
 [x, e] = square_mesh.generate(side_len, n_seg)  # node positions and edge node indices
@@ -41,7 +43,10 @@
 contact_area = [side_len / n_seg] * len(x)     # perimeter split to each node
 # ANCHOR_END: contact_area
 # compute reduced basis using 0: no reduction; 1: polynomial functions; 2: modal reduction
-reduced_basis = utils.compute_reduced_basis(x, e, vol, IB, mu_lame, lam, method=1, order=2)
+method = 1
+order = 1
+reduced_basis = utils.compute_reduced_basis(x, e, vol, IB, mu_lame, lam, method, order)
+abd_anchor_basis = utils.compute_abd_anchor_basis(x)
 
 # simulation with visualization
 resolution = np.array([900, 900])
@@ -65,6 +70,15 @@ def screen_projection(x):
     # fill the background and draw the square
     screen.fill((255, 255, 255))
     pygame.draw.aaline(screen, (0, 0, 255), screen_projection([-2, y_ground]), screen_projection([2, y_ground]))   # ground
+    # draw abd 
+    if draw_abd and method == 1 and order == 1: 
+        for i in range(3):
+            reduced_vars = np.linalg.solve(np.transpose(reduced_basis) @ reduced_basis, np.transpose(reduced_basis) @ x.reshape(-1))
+            abd_anchors = abd_anchor_basis @ reduced_vars
+            pygame.draw.circle(screen, (255, 0, 0), screen_projection(abd_anchors[2 * i: 2 * i + 2]), 0.1 * side_len / n_seg * scale)
+        pygame.draw.aaline(screen, (255, 0, 0), screen_projection(abd_anchors[0:2]), screen_projection(abd_anchors[2:4]))
+        pygame.draw.aaline(screen, (255, 0, 0), screen_projection(abd_anchors[2:4]), screen_projection(abd_anchors[4:6]))
+        pygame.draw.aaline(screen, (255, 0, 0), screen_projection(abd_anchors[4:6]), screen_projection(abd_anchors[0:2]))
     for eI in e:
         # ANCHOR: draw_tri
         pygame.draw.aaline(screen, (0, 0, 255), screen_projection(x[eI[0]]), screen_projection(x[eI[1]]))

9_reduced_DOF/utils.py

@@ -1,3 +1,4 @@
+import math
 import numpy as np
 import numpy.linalg as LA
 
@@ -33,6 +34,21 @@ def smallest_positive_real_root_quad(a, b, c, tol = 1e-6):
     return t
 # ANCHOR_END: find_positive_real_root
 
+def compute_abd_anchor_basis(x):
+    c = x.mean(axis=0)
+    diag = np.linalg.norm(x.max(axis=0) - x.min(axis=0))
+    scale = diag / math.sqrt(2)
+    c -= 1/3 * scale
+    anchors = np.stack([c, c + np.asarray([scale, 0]), c + np.asarray([0, scale])], axis=0)
+
+    basis = np.zeros((len(anchors) * 2, 6)) # 1, x, y for both x- and y-displacements
+    for i in range(len(anchors)):
+        for d in range(2):
+            basis[i * 2 + d][d * 3] = 1
+            basis[i * 2 + d][d * 3 + 1] = anchors[i][0]
+            basis[i * 2 + d][d * 3 + 2] = anchors[i][1]
+    return basis
+
 def compute_reduced_basis(x, e, vol, IB, mu_lame, lam, method, order):
     if method == 0: # full basis, no reduction
         basis = np.zeros((len(x) * 2, len(x) * 2))