Merge pull request #117 from Jax922/feature/proximal-parallel

Feature/proximal parallel

.github/workflows/python-package-conda.yml

@@ -33,6 +33,8 @@ jobs:
             conda install pyparsing
             conda install -c anaconda ipython_genutils 
             conda install -c conda-forge meshplot
+            conda install numba=0.58.1
+            conda install -c anaconda networkx
             pip install usd-core
             coverage run  -m unittest python/unit_tests/test_*.py
             coverage report

python/rainbow/simulators/prox_rigid_bodies/solver.py

@@ -5,7 +5,7 @@
 import rainbow.geometry.surface_mesh as MESH
 import rainbow.simulators.prox_rigid_bodies.mass as MASS
 import rainbow.simulators.prox_rigid_bodies.collision_detection as CD
-import rainbow.simulators.prox_rigid_bodies.gauss_seidel as GS
+import rainbow.simulators.proximal_contact.prox_solvers as CONTACT_SOLVERS
 from rainbow.simulators.prox_rigid_bodies.types import *
 from rainbow.util.timer import Timer
 import numpy as np
@@ -421,8 +421,10 @@ def apply_post_stabilization(J, WJT, x, engine, stats: dict, debug_on) -> dict:
     if not g.any():
         return stats
     mu = np.zeros(K, dtype=np.float64)
-    sol, stats = GS.solve(
-        J, WJT, g, mu, GS.prox_origin, engine, stats, debug_on, "post_stabilization_"
+    sol, stats = CONTACT_SOLVERS.solve(
+        J, WJT, g, mu, CONTACT_SOLVERS.prox_origin, engine, stats, debug_on,
+        prefix="post_stabilization_", 
+        scheme=engine.params.proximal_solver
     )
     vector_positional_update = WJT.dot(sol)
     position_update(x, vector_positional_update, 1, engine)
@@ -537,8 +539,10 @@ def step(self, dt: float, engine: Engine, debug_on: bool) -> None:
 
             mu = get_friction_coefficient_vector(engine)
             b = np.multiply(1 + e, v) + J.dot(du_ext) + g
-            sol, stats = GS.solve(
-                J, WJT, b, mu, GS.prox_sphere, engine, stats, debug_on, ""
+            sol, stats = CONTACT_SOLVERS.solve(
+                J, WJT, b, mu, CONTACT_SOLVERS.prox_sphere, engine, stats, debug_on, 
+                prefix="",
+                scheme=engine.params.proximal_solver
             )
             du_contact = WJT.dot(sol)
 

python/rainbow/simulators/prox_rigid_bodies/types.py

@@ -313,6 +313,7 @@ def __init__(self):
         self.resolution = (
             64  # The number of grid cells along each axis in the signed distance fields
         )
+        self.proximal_solver = "gauss_seidel" # or "gauss_seidel", "parallel_gauss_seidel", "parallel_jacobi", "parallel_jacboi_hybrid"
 
 
 class Engine:
@@ -331,6 +332,7 @@ def __init__(self):
         """
         Create a default empty engine instance.
         """
+        self.simulator_type = 'rigid_body' # simulation type for the engine
         self.bodies = dict()
         self.forces = dict()
         self.shapes = dict()

python/rainbow/simulators/prox_soft_bodies/solver.py

@@ -1,6 +1,6 @@
 import rainbow.simulators.prox_soft_bodies.mechanics as MECH
 import rainbow.simulators.prox_soft_bodies.collision_detection as CD
-import rainbow.simulators.prox_rigid_bodies.gauss_seidel as GS
+import rainbow.simulators.proximal_contact.prox_solvers as CONTACT_SOLVERS
 import rainbow.math.vector3 as V3
 import rainbow.math.matrix3 as M3
 import numpy as np
@@ -974,9 +974,10 @@ def apply_post_stabilization(J, WJT, engine, stats, debug_on):
         return stats
 
     mu = np.zeros(K, dtype=np.float64)
-    sol, stats = GS.solve(
-        J, WJT, g, mu, GS.prox_origin, engine, stats, debug_on, "post_stabilization_"
-    )
+
+    sol, stats = CONTACT_SOLVERS.solve(J, WJT, g, mu, CONTACT_SOLVERS.prox_origin, engine, stats, debug_on, 
+                                       prefix="post_stabilization_", 
+                                       scheme=engine.params.proximal_solver)
     delta_x = WJT.dot(sol)
 
     # --- Convert from 3N-by-1 into N-by-3 vector format -----------------
@@ -1114,9 +1115,9 @@ def step(self, dt: float, engine: Engine, debug_on: bool) -> None:
 
             mu = get_friction_coefficient_vector(engine)
             b = J.dot(u_prime)
-            sol, stats = GS.solve(
-                J, WJT, b, mu, GS.prox_sphere, engine, stats, debug_on, "gauss_seidel_"
-            )
+            sol, stats = CONTACT_SOLVERS.solve(J, WJT, b, mu, CONTACT_SOLVERS.prox_sphere, engine, stats, debug_on, 
+                                               prefix="", 
+                                               scheme=engine.params.proximal_solver)
             WPc = WJT.dot(sol)
 
         # --- Convert from 3N-by-1 into N-by-3 vector format -----------------

python/rainbow/simulators/prox_soft_bodies/types.py

@@ -253,10 +253,10 @@ def __init__(self):
             0.1  # Any geometry within this distance generates a contact point.
         )
         self.resolution = 64  # The number of grid cells along each axis in the signed distance fields.
+        self.proximal_solver = "gauss_seidel" # or "gauss_seidel", "parallel_gauss_seidel", "parallel_jacobi", "parallel_jacboi_hybrid"
         self.use_spatial_hashing = True  # Boolean flag that indicates if spatial hashing should be used instead of the BVH or not.
         self.time_stamp = 0 # The time step to use when simulating forward.
 
-
 class Engine:
     """
     This class holds all data that describes the world that we are currently simulating.
@@ -269,6 +269,7 @@ def __init__(self):
         This initializes an engine instance which holds all the configuration data of the world that is to
         be simulated.
         """
+        self.simulator_type = 'soft_body' # simulation type for the engine
         self.materials = dict()  # All materials that exist in the world.
         self.materials[
             "default"
@@ -288,3 +289,4 @@ def __init__(self):
         self.number_of_nodes = 0  # The total number of nodes in the world.
         self.stepper = None  # A reference to the time-stepper used to simulator forward.
         self.hash_grid = SH.HashGird()
+

python/rainbow/simulators/proximal_contact/blocking.py

@@ -0,0 +1,111 @@
+import numpy as np
+import networkx as nx
+from collections import defaultdict
+
+
+def build_contact_graph(contact_points, type="rigid_body"):
+    """ Create the contact graph based on the contact points for rigid body or soft body.
+        The vertices of the graph are the contact points, and the edges of the graph are :
+        1. If the type is rigid body, the edges are the contact points which belong to the same rigid body.
+        2. If the type is soft body, the edges are the contact points which belong to the same tetrahedron.
+
+    :param contact_points: The contact points array of the engine.
+    :param type: The type of the body, 'rigid body' or 'soft body', defaults to "rigid_body"
+    :raises ValueError: If the type of body is not supported.
+    :return: The contact graph.
+    """
+    G = nx.Graph()
+    K = len(contact_points) # The number of contact points.
+    G.add_nodes_from(np.arange(K)) # Add the contact points as the vertices of the graph.
+
+    for k in range(K):
+        cp1 = contact_points[k]
+        if type == "rigid_body":
+            for i in range(k+1, K):
+                cp2 = contact_points[i]
+                if cp1.bodyA.idx == cp2.bodyA.idx or cp1.bodyA.idx == cp2.bodyB.idx or cp1.bodyB.idx == cp2.bodyA.idx or cp1.bodyB.idx == cp2.bodyB.idx:
+                    G.add_edge(k, i)
+        elif type == "soft_body":
+            cp1_iA, cp1_jA, cp1_kA, cp1_mA = cp1.bodyA.T[cp1.idx_tetA] + cp1.bodyA.offset
+            cp1_iB, cp1_jB, cp1_kB, cp1_mB = cp1.bodyB.T[cp1.idx_tetB] + cp1.bodyB.offset
+            set1 = set([cp1_iA, cp1_jA, cp1_kA, cp1_mA, cp1_iB, cp1_jB, cp1_kB, cp1_mB])
+            for i in range(k+1, K):
+                cp2 = contact_points[i]
+                cp2_iA, cp2_jA, cp2_kA, cp2_mA = cp2.bodyA.T[cp2.idx_tetA] + cp2.bodyA.offset
+                cp2_iB, cp2_jB, cp2_kB, cp2_mB = cp2.bodyB.T[cp2.idx_tetB] + cp2.bodyB.offset
+                set2 = set([cp2_iA, cp2_jA, cp2_kA, cp2_mA, cp2_iB, cp2_jB, cp2_kB, cp2_mB])
+                if len(set1.intersection(set2)) > 0:
+                    G.add_edge(k, i)
+        else:
+            raise ValueError("The type of body is not supported.")
+
+    return G
+
+
+def greedy_graph_coloring(G):
+    """ Greedy graph coloring algorithm
+
+    :param G: The contact grpah.
+    :return: The color dictionary, the key is the color, the value is a array of the block location.
+    """
+    C = nx.coloring.greedy_color(G)
+    color_groups = defaultdict(list)
+    sorted_C = dict(sorted(C.items()))
+
+    for k, color in sorted_C.items():
+        block_start = 4 * k # The start index of a block
+        block_end = block_start + 4 # The end index of a block
+        color_groups[color].append((block_start, block_end))
+    return color_groups
+
+
+def random_graph_coloring(G, max_iter = 200):
+    """ Random graph coloring algorithm, which is posted as this paper: "Vivace: a Practical Gauss-Seidel Method for Stable Soft Body Dynamics". According to this paper, random graph coloring algorithm is faster than greedy graph coloring algorithm, and it can get a better result than greedy graph coloring algorithm; but I did not got the same result as the paper. In my test, greedy graph coloring algorithm is better than random graph coloring algorithm. I think the reason is that the random graph coloring algorithm is not stable, it can get different result in different test. So I did not use this algorithm in the final version of the code, however we can save this algorithm for future use (maybe can achieve better performance in different scene).
+
+    :param G: The contact grpah.
+    :param max_iter: The maximum iterations, defaults to 200.
+    :return:  A color dictionary, the key is the contact index, the value is the color index.  
+
+    """
+    degrees = np.array([degree for _, degree in G.degree()])
+    colors = (degrees // 1).astype(int)
+    palettes = [set(np.arange(c)) for c in colors]
+    palette_dict = {node: palette for node, palette in zip(G.nodes(), palettes)}
+
+    U = set(G.nodes())
+    C = defaultdict(int)
+    iter_count = 0
+    while len(U) > 0 and iter_count < max_iter:
+        for v in U:
+            if palette_dict[v]:
+                C[v] = np.random.choice(list(palette_dict[v]))
+
+        I = set()
+        for v in U:
+            neighbors_colors = {C[u] for u in G.neighbors(v) if u in U}
+            if C[v] not in neighbors_colors:
+                I.add(v)
+                for u in G.neighbors(v):
+                    if u in U:
+                        palette_dict[u].discard(C[v])
+        U -= I
+        max_used_color = max(C.values(), default=-1)
+        for v in U:
+            if not palette_dict[v]:
+                max_used_color += 1
+                palette_dict[v].add(max_used_color)
+
+        iter_count += 1
+
+    if U:
+        max_used_color = max(C.values(), default=-1)
+        for v in U:
+            max_used_color += 1
+            C[v] = max_used_color  
+
+    color_groups = defaultdict(list)
+    sorted_C = dict(sorted(C.items()))
+    for key, value in sorted_C.items():
+        color_groups[value].append(key)
+
+    return color_groups

python/rainbow/simulators/proximal_contact/gauss_seidel_solver.py

@@ -0,0 +1,122 @@
+import numpy as np
+import numba as nb
+from collections import defaultdict
+import rainbow.simulators.proximal_contact.blocking as BLOCKING
+from rainbow.simulators.proximal_contact.solver_interface import SolverInterface
+
+
+class GaussSeidelSolver(SolverInterface):
+    """ Serial Gauss Seidel Solver.
+    """
+
+    def sweep(self):
+        """ Sweep the contact force.
+        """
+        w = self.WJT.dot(self.x)
+        for k in range(self.K):
+            block = range(4 * k, 4 * k + 4)
+            mu_k = self.mu[k]  # Only isotropic Coulomb friction
+            x_b = self.x[block]
+            delta = (
+                x_b.copy()
+            )  # Used to keep the old values and compute the change in values
+            r_b = self.r[block]
+            b_b = self.b[block]
+
+            # By definition
+            #       z = x - r (J WJ^T x  + b)
+            #         = x - r ( A x  + b)
+            # We use
+            #        w =  WJ^T x
+            # so
+            #       z  = x - r ( J w  + b)
+            z_b = x_b - np.multiply(r_b, (self.J.dot(w)[block] + b_b))
+
+            # Solve:         x_n = prox_{R^+}( x_n - r (A x_n + b) )
+            x_b[0] = np.max([0.0, z_b[0]])
+
+            # Solve:         x_f = prox_C( x_f - r (A x_f + b))
+            x_b[1], x_b[2], x_b[3] = self.friction_solver(z_b[1], z_b[2], z_b[3], mu_k, x_b[0])
+            # Put updated contact forces back into solution vector
+            self.x[block] = x_b
+            # Get the change in the x_block
+            np.subtract(x_b, delta, delta)
+            # Update w
+            w += self.WJT.tocsr()[:, block].dot(delta)
+
+
+@nb.njit(parallel=True, nogil=True, cache=True)
+def sweep_worker(color_group, J, WJT, b, mu, r, x, w, delta_ws, friction_solver):
+    """ Sweep worker for parallel Gauss Seidel solver.
+
+    :param color_group: The color group, the value is  the block location, containing the start and end index.
+    :param J: The contact jacobi matrix.
+    :param WJT: The WJ^T matrix, here W = M^{-1}, and the M is the mass matrix.
+    :param b: b = Ju^t + ∆t JM^{-1}h + EJu^t
+    :param mu: The coefficient of friction.
+    :param r: r-factor value.
+    :param x: The current contact force.
+    :param w: The WJT.dot(x).
+    :param delta_ws: A array of the delta_w, each delta_w = WJT.dot(delta_x), here delta_x is the change of the contact force.
+    :param friction_solver: The proximal operator of friction cone function.
+    :return: The new contact force and the new delta_ws.
+    """
+
+    for i in nb.prange(len(color_group)):
+        block_start, block_end = color_group[i]
+        block = np.arange(block_start, block_end)
+        x_b = x[block]
+        r_b = r[block]
+        b_b = b[block]
+
+        delta = x_b.copy()
+        z_b = x_b - np.multiply(r_b, (J.dot(w)[block] + b_b))
+
+        x_b[0] = np.max(np.array([0.0, z_b[0]]))
+
+        mu_k = mu[block_start // 4]
+        x_b[1], x_b[2], x_b[3] = friction_solver(z_b[1], z_b[2], z_b[3], mu_k, x_b[0])
+
+        np.subtract(x_b, delta, delta)
+        delta_w = WJT[:, block].dot(delta)
+
+        x[block] = x_b
+        delta_ws[block_start//4] = delta_w
+
+    return x, delta_ws
+
+
+class ParallelGaussSeidelSolver(SolverInterface):
+    """ Parallel Gauss Seidel Solver.
+    """
+
+    def sweep(self):
+        """ Sweep the contact force.
+        """
+        # Compute the color group
+        color_groups = defaultdict(list)
+        G = BLOCKING.build_contact_graph(self.engine.contact_points, self.engine.simulator_type)
+        color_groups = BLOCKING.greedy_graph_coloring(G)
+
+        # Set the number of threads
+        nb.set_num_threads(4 if (nb.config.NUMBA_NUM_THREADS // 2) >= 4 else nb.config.NUMBA_NUM_THREADS // 2)
+
+        # Compute the new contact force
+        w0 = self.WJT.dot(self.x)
+        delta_ws = np.zeros((self.WJT.shape[1]//4, w0.shape[0]), dtype=np.float64)
+        for color_group in color_groups.values():
+            w = w0.copy()
+            block_start, _ = color_group[0]
+            for i in range(block_start // 4):
+                w += delta_ws[i]
+
+            self.x, delta_ws = sweep_worker(np.array(color_group), 
+                                            self.J.toarray(), 
+                                            self.WJT.toarray(), 
+                                            self.b, 
+                                            self.mu, 
+                                            self.r, 
+                                            self.x, 
+                                            w,
+                                            delta_ws, 
+                                            self.friction_solver)