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
+            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:

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.

python/rainbow/simulators/proximal_contact/blocking.py

@@ -0,0 +1,142 @@
+import numpy as np
+import numba as nb
+import networkx as nx
+from collections import defaultdict
+
+
+@nb.njit(parallel=True, nogil=True, cache=True)
+def _get_edges(J):
+    """ Get the edges of the contact graph.
+
+    Args:
+        J (ArrayLike): The contact jacobi matrix.
+
+    Returns:
+        (ArrayLike, ArrayLike): The source and target of the edges.
+    """
+    K = J.shape[0] // 4 # The number of contact points, here the each contact point has 4 rows of J.
+    step_size = 24 # Here the step_size is 24 , because each contact point has 8 non-zero blocks and each block has 3 columns.
+    cols = J.shape[1] # The number of columns of J.
+
+    # Preallocate arrays for source and target of edges
+    max_possible_edges = (cols // step_size) * K * K
+    source = np.full(max_possible_edges, -1, dtype=np.int64)
+    target = np.full(max_possible_edges, -1, dtype=np.int64)
+    edge_count = 0
+
+    for idx in nb.prange(0, cols//step_size):
+        col_idx = idx * step_size
+        contacts = np.where(J[:, col_idx: col_idx+step_size])[0]
+        contacts = np.unique(contacts // 4)
+        for i, contact_i in enumerate(contacts):
+            for j, contact_j in enumerate(contacts):
+                if i < j:
+                    source[edge_count] = contact_i
+                    target[edge_count] = contact_j
+                    edge_count += 1
+
+    # Trim the arrays to the actual size
+    valid_indices = source != -1
+    filtered_source = source[valid_indices]
+    filtered_source = target[valid_indices]
+
+    return filtered_source, filtered_source
+
+
+def build_contact_graph(J):
+    """ Create the contact graph based on the contact jacobi matrix.
+        The vertices of the graph are the contact points, and the edges of the graph are the value of the contact jacobi matrix is not zero between two contact points.
+
+    Args:
+        J (ArrayLike): The contact jacobi matrix.
+
+    Returns:
+        graph: The contact graph.
+    """
+    G = nx.Graph()
+    K = J.shape[0] // 4  # The number of contact points, here the each contact point has 4 rows of J.
+    G.add_nodes_from(np.arange(K)) # Add the contact points as the vertices of the graph.
+    sources, targets = _get_edges(J)
+    for s, t in zip(sources, targets):
+        G.add_edge(s, t)
+
+    return G
+
+
+def greedy_graph_coloring(G):
+    """ Greedy graph coloring algorithm
+
+    Args:
+        G (Graph): The contact grpah is created from the contact jacobi matrix.
+
+    Returns:
+        ArrayLike: 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"
+
+    Note: 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).
+
+    Args:
+        G (Graph): The contact grpah is created from the contact jacobi matrix.
+        max_iter (int, optional): The maximum iterations . Defaults to 200.
+
+    Returns:
+        dict: 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,124 @@
+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 __init__(self, J, WJT, b, mu, friction_solver, engine, stats, debug_on, prefix):
+        super().__init__(J, WJT, b, mu, friction_solver, engine, stats, debug_on, prefix)
+
+    def sweep(self):
+        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):
+    """ The worker function of the parallel gauss seidel algorithm.
+
+    Args:
+        color_group (ArrayLike): The color group, the value is  the block location, containing the start and end index.
+        J (ArrayLike): The contact jacobi matrix.
+        WJT (ArrayLike): The WJ^T matrix, here W = M^{-1}, and the M is the mass matrix.
+        b (ArrayLike): b = Ju^t + ∆t JM^{-1}h + EJu^t
+        mu (float): The coefficient of friction.
+        r (float): r-factor value.
+        x (ArrayLike): The current contact force.
+        w (ArrayLike): The WJT.dot(x).
+        delta_ws (ArrayLike): A array of the delta_w, each delta_w = WJT.dot(delta_x), here delta_x is the change of the contact force.
+        friction_solver (callable): The proximal operator of friction cone function.
+
+    Returns:
+        (ArrayLike, ArrayLike): 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 __init__(self, J, WJT, b, mu, friction_solver, engine, stats, debug_on, prefix):
+        super().__init__(J, WJT, b, mu, friction_solver, engine, stats, debug_on, prefix)
+
+    def sweep(self):
+        # Compute the color group
+        color_groups = defaultdict(list)
+        G = BLOCKING.build_contact_graph(self.J.toarray())
+        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)