From 894c02a42db446c3f454a3b6a19b1091a84b41b5 Mon Sep 17 00:00:00 2001 From: Dong She Date: Tue, 14 Nov 2023 19:08:30 +0100 Subject: [PATCH] update proximal algorithm --- .github/workflows/python-package-conda.yml | 2 +- .../simulators/prox_rigid_bodies/solver.py | 14 +- .../simulators/prox_rigid_bodies/types.py | 1 + .../prox_soft_bodies/gauss_seidel_parallel.py | 381 ------------------ .../prox_soft_bodies/jacobi_parallel.py | 143 ------- .../simulators/prox_soft_bodies/solver.py | 28 +- .../simulators/prox_soft_bodies/types.py | 6 +- .../simulators/proximal_contact/blocking.py | 142 +++++++ .../proximal_contact/gauss_seidel_solver.py | 124 ++++++ .../proximal_contact/jacobi_solver.py | 82 ++++ .../proximal_contact/prox_solvers.py | 41 ++ .../proximal_operators.py} | 16 +- .../proximal_contact/solver_interface.py | 133 ++++++ .../unit_tests/test_random_graph_coloring.py | 34 ++ 14 files changed, 584 insertions(+), 563 deletions(-) delete mode 100644 python/rainbow/simulators/prox_soft_bodies/gauss_seidel_parallel.py delete mode 100644 python/rainbow/simulators/prox_soft_bodies/jacobi_parallel.py create mode 100644 python/rainbow/simulators/proximal_contact/blocking.py create mode 100644 python/rainbow/simulators/proximal_contact/gauss_seidel_solver.py create mode 100644 python/rainbow/simulators/proximal_contact/jacobi_solver.py create mode 100644 python/rainbow/simulators/proximal_contact/prox_solvers.py rename python/rainbow/simulators/{prox_soft_bodies/prox.py => proximal_contact/proximal_operators.py} (89%) create mode 100644 python/rainbow/simulators/proximal_contact/solver_interface.py create mode 100644 python/unit_tests/test_random_graph_coloring.py diff --git a/.github/workflows/python-package-conda.yml b/.github/workflows/python-package-conda.yml index 0ccd95a..4fb83e6 100644 --- a/.github/workflows/python-package-conda.yml +++ b/.github/workflows/python-package-conda.yml @@ -34,7 +34,7 @@ jobs: conda install -c anaconda ipython_genutils conda install -c conda-forge meshplot conda install numba - pip install networkx + conda install -c anaconda networkx pip install usd-core coverage run -m unittest python/unit_tests/test_*.py coverage report diff --git a/python/rainbow/simulators/prox_rigid_bodies/solver.py b/python/rainbow/simulators/prox_rigid_bodies/solver.py index e434202..8b97121 100644 --- a/python/rainbow/simulators/prox_rigid_bodies/solver.py +++ b/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) diff --git a/python/rainbow/simulators/prox_rigid_bodies/types.py b/python/rainbow/simulators/prox_rigid_bodies/types.py index 607939f..bb344f1 100644 --- a/python/rainbow/simulators/prox_rigid_bodies/types.py +++ b/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: diff --git a/python/rainbow/simulators/prox_soft_bodies/gauss_seidel_parallel.py b/python/rainbow/simulators/prox_soft_bodies/gauss_seidel_parallel.py deleted file mode 100644 index 516bb3a..0000000 --- a/python/rainbow/simulators/prox_soft_bodies/gauss_seidel_parallel.py +++ /dev/null @@ -1,381 +0,0 @@ -import numpy as np -import numba as nb -import networkx as nx -from collections import defaultdict -from rainbow.util.timer import Timer - - -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 - G.add_nodes_from(np.arange(K)) - sources, targets = _get_edges(J) - for s, t in zip(sources, targets): - G.add_edge(s, t) - - return G - - -@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 - step_size = 24 - cols = J.shape[1] - - # 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 _greedy_graph_coloring(G): - """ Greedy graph coloring algorithm - - Args: - G (Graph): The contact grpah is created by on 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 - block_end = block_start + 4 - 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 by on 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 - - return C - - -@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 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 - - -def sweep_parallel(K, J, WJT, b, mu, r, x, friction_solver): - """ Parallel Gauss-Seidel algorithm: - This method aims to parallelize the Gauss-Seidel algorithm by dividing contact points into different groups based on graph coloring. Each group, representing a unique color, is processed in parallel within itself, but serially across groups. Here are the detailed steps: - - 1. Construct the contact graph: Vertices of this graph represent contact points, and edges indicate non-zero values in the contact Jacobian matrix between two such points. For more information, refer to: https://dspace5.zcu.cz/bitstream/11025/11058/1/Lazarevych.pdf. - - 2. Graph Coloring: Utilize the graph coloring algorithm to segregate contact points into distinct groups (colors). We experimented with both the Random Graph Coloring and Greedy Graph Coloring algorithms. However, the Greedy Graph Coloring proved superior, hence our choice. - - 3. Parallel Gauss Seidel computation: Each color group's contact force can be computed in parallel. However, computations for each color group are handled serially. - - - Why we can parallel the gauss seidel algorithm? - The origin idea is easy to understand why we can parallel the gauss seidel algorithm, if the two rows of the contact jacobi matrix are independent, then we can compute the new contact force of these two rows in parallel, the indepandent rows defined the value of the contact jacobi matrix is zero between these two contact points. We can move to a easy case of linear system as follows: - Ax = b - [[1, 1, 0, 0], [x_1, [1, - [2, 2, 2, 0], x_2, = 2, - [0, 1, 1, 1], x_3, 3, - [0, 0, 2, 2]] x_4] 4] - We can find A_{0,3} = 0 and A_{2,0} = 0, so we can compute x_1 and x_3 in parallel. - - - Args: - K (int): Contact points number. - J (ArrayLike): The contact jacobi matrix. - WJT (ArrayLike): The WJ^T matrix, here W = M^{-1}. - 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. - friction_solver (callable): The friction cone function. - - Returns: - ArrayLike: The new contact force. - """ - color_groups = defaultdict(list) - G = _build_contact_graph(J.toarray()) - color_groups = _greedy_graph_coloring(G) - nb.set_num_threads(4 if (nb.config.NUMBA_NUM_THREADS // 2) >= 4 else nb.config.NUMBA_NUM_THREADS // 2) - - w0 = WJT.dot(x) - delta_ws = np.zeros((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] - x, delta_ws = sweep_worker(np.array(color_group), J.toarray(), WJT.toarray(), b, mu, r, x, w, delta_ws, friction_solver) - - return x - - -def sweep(K, J, WJT, b, mu, r, x, friction_solver): - """ Serial Gauss Seidel algorithm. - - Args: - K (int): Contact points number. - J (ArrayLike): The contact jacobi matrix. - WJT (ArrayLike): The WJ^T matrix, here W = M^{-1}. - 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. - friction_solver (callable): The friction cone function. - - Returns: - ArrayLike: The new contact force. - """ - w = WJT.dot(x) - for k in range(K): - block = range(4 * k, 4 * k + 4) - mu_k = mu[k] # Only isotropic Coulomb friction - x_b = x[block] - delta = ( - x_b.copy() - ) # Used to keep the old values and compute the change in values - r_b = r[block] - b_b = 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, (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] = friction_solver(z_b[1], z_b[2], z_b[3], mu_k, x_b[0]) - # Put updated contact forces back into solution vector - x[block] = x_b - # Get the change in the x_block - np.subtract(x_b, delta, delta) - # Update w - w += WJT.tocsr()[:, block].dot(delta) - return x - - -def solve(J, WJT, b, mu, friction_solver, engine, stats, debug_on, prefix): - """ Gauss Seidel Proximal Solver, which supports parallel computing and serial computing. - - Args: - J (ArrayLike): The contact jacobi matrix. - WJT (ArrayLike): The WJ^T matrix, here W = M^{-1}. - b (ArrayLike): b = Ju^t + ∆t JM^{-1}h + EJu^t - mu (float): The coefficient of friction. - friction_solver (callable): The friction cone function. - engine (Object): The engine object. - stats (dict): The statistics information. - debug_on (boolean): Whether to debug. - prefix (string): The prefix of the statistics information. - - Returns: - (ArrayLike, dict): The new contact force and the statistics information. - """ - timer = None - if debug_on: - timer = Timer("Gauss Seidel") - stats[prefix + "residuals"] = ( - np.ones(engine.params.max_iterations, dtype=np.float64) * np.inf - ) - stats[prefix + "lambda"] = np.zeros( - [engine.params.max_iterations] + list(b.shape), dtype=np.float64 - ) - stats[prefix + "reject"] = np.zeros(engine.params.max_iterations, dtype=bool) - stats[prefix + "exitcode"] = 0 - stats[prefix + "iterations"] = engine.params.max_iterations - timer.start() - - K = len(engine.contact_points) - x = np.zeros(b.shape, dtype=np.float64) # The current iterate - sol = np.zeros( - b.shape, dtype=np.float64 - ) # The last best known solution, used for restarting if divergence - error = np.zeros(b.shape, dtype=np.float64) # The residual vector - - # Compute initial r-factor value - delassus_diag = np.sum(J.multiply(WJT.T), axis=1).A1 - delassus_diag[delassus_diag == 0] = 1 - r = 0.1 / delassus_diag - - # Extract parameter values for controlling the adaptive r-factor strategy - nu_reduce = engine.params.nu_reduce - nu_increase = engine.params.nu_increase - too_small_merit_change = engine.params.too_small_merit_change - - last_merit = np.Inf - - for iteration in range(engine.params.max_iterations): - if engine.params.proximal_solver['parallel']: - x = sweep_parallel(K, J, WJT, b, mu, r, x, friction_solver) - else: - x = sweep(K, J, WJT, b, mu, r, x, friction_solver) - - np.subtract(x, sol, error) - merit = np.linalg.norm(error, np.inf) - if debug_on: - stats[prefix + "lambda"][iteration] = x - stats[prefix + "residuals"][iteration] = merit - # Test stopping criteria - if merit < engine.params.absolute_tolerance: - if debug_on: - stats[prefix + "iterations"] = iteration - stats[prefix + "exitcode"] = 1 - timer.end() - stats[prefix + "solver_time"] = timer.elapsed - return x, stats - if np.abs(merit - last_merit) < engine.params.relative_tolerance * last_merit: - if debug_on: - stats[prefix + "iterations"] = iteration - stats[prefix + "exitcode"] = 2 - timer.end() - stats[prefix + "solver_time"] = timer.elapsed - return x, stats - - # Update r-factors - if merit > last_merit: - # Divergence detected: reduce R-factor and roll-back solution to last known good iterate! - np.multiply(nu_reduce, r, r) - np.copyto(x, sol) - if debug_on: - stats[prefix + "reject"][iteration] = True - else: - if last_merit - merit < too_small_merit_change: - # Convergence is slow: increase r-factor - np.multiply(nu_increase, r, r) - # Convergence detected: accept x as better solution - last_merit = merit - np.copyto(sol, x) - - # If this point of the code is reached then it means the method did not converge within the given iterations. - if debug_on: - timer.end() - stats[prefix + "solver_time"] = timer.elapsed - return sol, stats diff --git a/python/rainbow/simulators/prox_soft_bodies/jacobi_parallel.py b/python/rainbow/simulators/prox_soft_bodies/jacobi_parallel.py deleted file mode 100644 index fac4dfb..0000000 --- a/python/rainbow/simulators/prox_soft_bodies/jacobi_parallel.py +++ /dev/null @@ -1,143 +0,0 @@ -import numpy as np -import numba as nb -from rainbow.util.timer import Timer - - -@nb.njit(parallel=True, nogil=True, cache=True) -def sweep(K, w, z, b, mu, x, friction_solver, hybrid = True): - """ Parallel Jacobi Proximal Algorithm. - - Args: - K (int): Contact points number. - w (ArrayLike): The WJT.dot(x). - z (ArrayLike): The current contact force. - b (ArrayLike): b = Ju^t + ∆t JM^{-1}h + EJu^t - mu (float): The coefficient of friction. - x (ArrayLike): The current contact force. - friction_solver (callable): The friction cone function. - hybrid (bool): If True, use the hybrid parallel Jacobi algorithm. - - Returns: - ArrayLike: The new contact force. - """ - if not hybrid: - x_old = x.copy() - for k in nb.prange(K): - block = np.arange(4 * k, 4 * k + 4) - mu_k = mu[k] # Only isotropic Coulomb friction - x_b = x[block] - b_b = b[block] - z_b = z[block] - - # Solve: x_n = prox_{R^+}( x_n - r (A x_n + b) ) - x_b[0] = np.max(np.array([0.0, z_b[0]])) - - # Solve: x_f = prox_C( x_f - r (A x_f + b)) - x_b_0 = x_b[0] if hybrid else x_old[block][0] - 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) - # Put updated contact forces back into solution vector - x[block] = x_b - return x - - -def solve(J, WJT, b, mu, friction_solver, engine, stats, debug_on, prefix): - """ Jacobi Proximal Solver, which just supports parallel computing. - - Args: - J (ArrayLike): The contact jacobi matrix. - WJT (ArrayLike): The WJ^T matrix, here W = M^{-1}. - b (ArrayLike): b = Ju^t + ∆t JM^{-1}h + EJu^t - mu (float): The coefficient of friction. - friction_solver (callable): The friction cone function. - engine (Object): The engine object. - stats (dict): The statistics information. - debug_on (boolean): Whether to debug. - prefix (string): The prefix of the statistics information. - - Returns: - (ArrayLike, dict): The new contact force and the statistics information. - """ - timer = None - if debug_on: - timer = Timer("Gauss Seidel") - stats[prefix + "residuals"] = ( - np.ones(engine.params.max_iterations, dtype=np.float64) * np.inf - ) - stats[prefix + "lambda"] = np.zeros( - [engine.params.max_iterations] + list(b.shape), dtype=np.float64 - ) - stats[prefix + "reject"] = np.zeros(engine.params.max_iterations, dtype=bool) - stats[prefix + "exitcode"] = 0 - stats[prefix + "iterations"] = engine.params.max_iterations - timer.start() - - K = len(engine.contact_points) - x = np.zeros(b.shape, dtype=np.float64) # The current iterate - sol = np.zeros( - b.shape, dtype=np.float64 - ) # The last best known solution, used for restarting if divergence - error = np.zeros(b.shape, dtype=np.float64) # The residual vector - - # Compute initial r-factor value - delassus_diag = np.sum(J.multiply(WJT.T), axis=1).A1 - delassus_diag[delassus_diag == 0] = 1 - r = 0.1 / delassus_diag - - # Extract parameter values for controlling the adaptive r-factor strategy - nu_reduce = engine.params.nu_reduce - nu_increase = engine.params.nu_increase - too_small_merit_change = engine.params.too_small_merit_change - last_merit = np.Inf - - # Set number of threads for Numba - nb.set_num_threads(4 if (nb.config.NUMBA_NUM_THREADS // 2) >= 4 else nb.config.NUMBA_NUM_THREADS // 2) - - for iteration in range(engine.params.max_iterations): - w = WJT.dot(x) - z = x - np.multiply(r, (J.dot(w) + b)) - hybrid = True if engine.params.proximal_solver['scheme'] == 'jacboi_hybrid' else False - x = sweep(K, w, z, b, mu, x, friction_solver, hybrid) - - np.subtract(x, sol, error) - merit = np.linalg.norm(error, np.inf) - if debug_on: - stats[prefix + "lambda"][iteration] = x - stats[prefix + "residuals"][iteration] = merit - # Test stopping criteria - if merit < engine.params.absolute_tolerance: - if debug_on: - stats[prefix + "iterations"] = iteration+1 - stats[prefix + "exitcode"] = 1 - timer.end() - stats[prefix + "solver_time"] = timer.elapsed - # print("convergence_iteration1 = ", iteration+1) - return x, stats - if np.abs(merit - last_merit) < engine.params.relative_tolerance * last_merit: - if debug_on: - stats[prefix + "iterations"] = iteration+1 - stats[prefix + "exitcode"] = 2 - timer.end() - stats[prefix + "solver_time"] = timer.elapsed - # print("convergence_iteration2 = ", iteration+1) - return x, stats - - # Update r-factors - if merit > last_merit: - # Divergence detected: reduce R-factor and roll-back solution to last known good iterate! - np.multiply(nu_reduce, r, r) - np.copyto(x, sol) - if debug_on: - stats[prefix + "reject"][iteration] = True - else: - if last_merit - merit < too_small_merit_change: - # Convergence is slow: increase r-factor - np.multiply(nu_increase, r, r) - # Convergence detected: accept x as better solution - last_merit = merit - np.copyto(sol, x) - - # If this point of the code is reached then it means the method did not converge within the given iterations. - if debug_on: - timer.end() - stats[prefix + "solver_time"] = timer.elapsed - return sol, stats \ No newline at end of file diff --git a/python/rainbow/simulators/prox_soft_bodies/solver.py b/python/rainbow/simulators/prox_soft_bodies/solver.py index 4016205..0e9bd5c 100644 --- a/python/rainbow/simulators/prox_soft_bodies/solver.py +++ b/python/rainbow/simulators/prox_soft_bodies/solver.py @@ -1,8 +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_soft_bodies.prox as PROX -import rainbow.simulators.prox_soft_bodies.jacobi_parallel as JP -import rainbow.simulators.prox_soft_bodies.gauss_seidel_parallel 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 @@ -976,14 +974,10 @@ def apply_post_stabilization(J, WJT, engine, stats, debug_on): return stats mu = np.zeros(K, dtype=np.float64) - if engine.params.proximal_solver['scheme'] == "jacobi": - sol, stats = JP.solve( - J, WJT, g, mu, PROX.prox_origin, engine, stats, debug_on, "post_stabilization_" - ) - if engine.params.proximal_solver['scheme'] == "gauss_seidel": - sol, stats = GS.solve( - J, WJT, g, mu, PROX.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 ----------------- @@ -1121,15 +1115,9 @@ def step(self, dt: float, engine: Engine, debug_on: bool) -> None: mu = get_friction_coefficient_vector(engine) b = J.dot(u_prime) - if engine.params.proximal_solver['scheme'] == "jacobi": - sol, stats = JP.solve( - J, WJT, b, mu, PROX.prox_sphere, engine, stats, debug_on, "jacobi_" - ) - if engine.params.proximal_solver['scheme'] == "gauss_seidel": - sol, stats = GS.solve( - J, WJT, b, mu, PROX.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 ----------------- diff --git a/python/rainbow/simulators/prox_soft_bodies/types.py b/python/rainbow/simulators/prox_soft_bodies/types.py index 3e3d70d..e10a4bf 100644 --- a/python/rainbow/simulators/prox_soft_bodies/types.py +++ b/python/rainbow/simulators/prox_soft_bodies/types.py @@ -252,11 +252,7 @@ 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 = { - "scheme": "gauss_seidel", # or "jacobi", "jacboi_hybrid" - "parallel": False, # This a boolean flag that indicates if we should use parallelization or not. - } - + self.proximal_solver = "gauss_seidel" # or "gauss_seidel", "parallel_gauss_seidel", "parallel_jacobi", "parallel_jacboi_hybrid" class Engine: """ diff --git a/python/rainbow/simulators/proximal_contact/blocking.py b/python/rainbow/simulators/proximal_contact/blocking.py new file mode 100644 index 0000000..5582210 --- /dev/null +++ b/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 diff --git a/python/rainbow/simulators/proximal_contact/gauss_seidel_solver.py b/python/rainbow/simulators/proximal_contact/gauss_seidel_solver.py new file mode 100644 index 0000000..9f1bd20 --- /dev/null +++ b/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) diff --git a/python/rainbow/simulators/proximal_contact/jacobi_solver.py b/python/rainbow/simulators/proximal_contact/jacobi_solver.py new file mode 100644 index 0000000..000544d --- /dev/null +++ b/python/rainbow/simulators/proximal_contact/jacobi_solver.py @@ -0,0 +1,82 @@ +import numpy as np +import numba as nb +from rainbow.simulators.proximal_contact.solver_interface import SolverInterface + + +@nb.njit(parallel=True, nogil=True, cache=True) +def sweep_worker(K, w, z, b, mu, x, friction_solver, hybrid): + """ Parallel Jacobi Proximal Algorithm. + + Args: + K (int): Contact points number. + w (ArrayLike): The WJT.dot(x). + z (ArrayLike): The current contact force. + b (ArrayLike): b = Ju^t + ∆t JM^{-1}h + EJu^t + mu (float): The coefficient of friction. + x (ArrayLike): The current contact force. + friction_solver (callable): The proximal operator of friction cone function. + hybrid (bool): If True, use the hybrid parallel Jacobi algorithm. + + Returns: + ArrayLike: The new contact force. + """ + if not hybrid: + x_old = x.copy() + for k in nb.prange(K): + block = np.arange(4 * k, 4 * k + 4) + mu_k = mu[k] # Only isotropic Coulomb friction + x_b = x[block] + b_b = b[block] + z_b = z[block] + + # Solve: x_n = prox_{R^+}( x_n - r (A x_n + b) ) + x_b[0] = np.max(np.array([0.0, z_b[0]])) + + # Solve: x_f = prox_C( x_f - r (A x_f + b)) + x_b_0 = x_b[0] if hybrid else x_old[block][0] + 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) + # Put updated contact forces back into solution vector + x[block] = x_b + return x + + +class ParallelJacobiSolver(SolverInterface): + """ Parallel Jacobi Proximal Algorithm. + """ + 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) + z = self.x - np.multiply(self.r, (self.J.dot(w) + self.b)) + self.x = sweep_worker( + self.K, + w, + z, + self.b, + self.mu, + self.x, + self.friction_solver, + False, + ) + + +class ParallelJacobiHybridSolver(SolverInterface): + """ Parallel Jacobi Hybrid Proximal Algorithm. + """ + 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) + z = self.x - np.multiply(self.r, (self.J.dot(w) + self.b)) + self.x = sweep_worker( + self.K, + w, + z, + self.b, + self.mu, + self.x, + self.friction_solver, + True, + ) \ No newline at end of file diff --git a/python/rainbow/simulators/proximal_contact/prox_solvers.py b/python/rainbow/simulators/proximal_contact/prox_solvers.py new file mode 100644 index 0000000..a3ea73d --- /dev/null +++ b/python/rainbow/simulators/proximal_contact/prox_solvers.py @@ -0,0 +1,41 @@ +from rainbow.simulators.proximal_contact.proximal_operators import * +from rainbow.simulators.proximal_contact.gauss_seidel_solver import GaussSeidelSolver, ParallelGaussSeidelSolver +from rainbow.simulators.proximal_contact.jacobi_solver import ParallelJacobiSolver, ParallelJacobiHybridSolver + + +def solve(J, WJT, b, mu, friction_solver, engine, stats, debug_on, prefix="", scheme="gauss_seidel"): + """_summary_ + + Args: + J (ArrayLike): The contact jacobi matrix. + WJT (ArrayLike): The WJ^T matrix, here W = M^{-1}. + b (ArrayLike): b = Ju^t + ∆t JM^{-1}h + EJu^t + mu (float): The coefficient of friction. + friction_solver (callable): The proximal operator of friction cone function. + engine (Object): The engine object. + stats (dict): The statistics information. + debug_on (boolean): Whether to debug. + prefix (string): The prefix of the statistics information. + scheme (str, optional): The scheme of proximal solver. Defaults to "gauss_seidel", you can use one of 'gauss_seidel', 'parallel_gauss_seidel', 'parallel_jacobi', 'parallel_jacboi_hybrid'. + + Raises: + NotImplementedError: Unknown solver scheme. + + Returns: The new contact force and the statistics information. + """ + + # If the prefix is empty, use the scheme as the prefix + prefix = prefix if prefix != "" else scheme + "_" + + if scheme == "gauss_seidel": + solver = GaussSeidelSolver(J, WJT, b, mu, friction_solver, engine, stats, debug_on, prefix) + elif scheme == "parallel_gauss_seidel": + solver = ParallelGaussSeidelSolver(J, WJT, b, mu, friction_solver, engine, stats, debug_on, prefix) + elif scheme == "parallel_jacobi": + solver = ParallelJacobiSolver(J, WJT, b, mu, friction_solver, engine, stats, debug_on, prefix) + elif scheme == "parallel_jacboi_hybrid": + solver = ParallelJacobiHybridSolver(J, WJT, b, mu, friction_solver, engine, stats, debug_on, prefix) + else: + raise NotImplementedError("Unknown solver scheme: {}, you can use one of 'gauss_seidel', 'parallel_gauss_seidel', 'parallel_jacobi', 'parallel_jacboi_hybrid'.".format(scheme)) + + return solver.solve() \ No newline at end of file diff --git a/python/rainbow/simulators/prox_soft_bodies/prox.py b/python/rainbow/simulators/proximal_contact/proximal_operators.py similarity index 89% rename from python/rainbow/simulators/prox_soft_bodies/prox.py rename to python/rainbow/simulators/proximal_contact/proximal_operators.py index fe1bec9..e44c652 100644 --- a/python/rainbow/simulators/prox_soft_bodies/prox.py +++ b/python/rainbow/simulators/proximal_contact/proximal_operators.py @@ -47,14 +47,14 @@ def prox_ellipsoid(engine, z_s, z_t, z_tau, mu_s, mu_t, mu_tau, x_n): """ :param engine: - :param z_s: - :param z_t: - :param z_tau: - :param mu_s: - :param mu_t: - :param mu_tau: - :param x_n: - :return: + :param z_s: s-component of current z-point. + :param z_t: t-component of current z-point. + :param z_tau: tau-component of current z-point. + :param mu_s: The coefficient of friction. + :param mu_t: The coefficient of friction. + :param mu_tau: The coefficient of friction. + :param x_n: The current normal force magnitude. + :return: The x_s, x_t, x_tau which will be the closest point to the ellipsoid with radius my*x_n. """ if x_n <= 0.0: return 0.0, 0.0, 0.0 diff --git a/python/rainbow/simulators/proximal_contact/solver_interface.py b/python/rainbow/simulators/proximal_contact/solver_interface.py new file mode 100644 index 0000000..531c671 --- /dev/null +++ b/python/rainbow/simulators/proximal_contact/solver_interface.py @@ -0,0 +1,133 @@ +import numpy as np +from abc import ABC, abstractmethod +from rainbow.util.timer import Timer + + +class SolverInterface(ABC): + def __init__(self, J, WJT, b, mu, friction_solver, engine, stats, debug_on, prefix) -> None: + self.J = J # The contact jacobi matrix + self.WJT = WJT # The WJ^T matrix, here W = M^{-1}. + self.b = b # b = Ju^t + ∆t JM^{-1}h + EJu^t + self.mu = mu # The coefficient of friction. + self.friction_solver = friction_solver # The proximal operator of the friction cone. + self.engine = engine # The engine object. + self.stats = stats # The statistics information. + self.debug_on = debug_on # Whether to debug (Ture or False). + self.prefix = prefix # The prefix of the statistics information. + self.timer = None # The timer used to record the time + self.K = len(engine.contact_points) # Number of contact points. + self.iteration = 0 # The current iteration. + self.x = np.zeros(self.b.shape, dtype=np.float64) # The current iterate + self.sol = np.zeros( + self.b.shape, dtype=np.float64 + ) # The last best known solution, used for restarting if divergence + self.error = np.zeros(self.b.shape, dtype=np.float64) # The residual vector + self.last_merit = np.Inf # The last merit value + + def compute_initial_r_factor(self): + """ Compute the initial r-factor value. + """ + delassus_diag = np.sum(self.J.multiply(self.WJT.T), axis=1).A1 + delassus_diag[delassus_diag == 0] = 1 + self.r = 0.1 / delassus_diag + + def update_r_factor(self): + """ Update the r-factor value. + """ + nu_reduce = self.engine.params.nu_reduce + nu_increase = self.engine.params.nu_increase + too_small_merit_change = self.engine.params.too_small_merit_change + + if self.merit > self.last_merit: + # Divergence detected: reduce R-factor and roll-back solution to last known good iterate! + np.multiply(nu_reduce, self.r, self.r) + np.copyto(self.x, self.sol) + if self.debug_on: + self.stats[self.prefix + "reject"][self.iteration] = True + else: + if self.last_merit - self.merit < too_small_merit_change: + # Convergence is slow: increase r-factor + np.multiply(nu_increase, self.r, self.r) + # Convergence detected: accept x as better solution + self.last_merit = self.merit + np.copyto(self.sol, self.x) + + def initialize_stats(self, timer): + """ Initialize the statistics information, when debug_on is True. + + Args: + timer ('Timer'): The timer used to record the time. + """ + self.timer = timer + self.stats[self.prefix + "residuals"] = ( + np.ones(self.engine.params.max_iterations, dtype=np.float64) * np.inf + ) + self.stats[self.prefix + "lambda"] = np.zeros( + [self.engine.params.max_iterations] + list(self.b.shape), dtype=np.float64 + ) + self.stats[self.prefix + "reject"] = np.zeros(self.engine.params.max_iterations, dtype=bool) + self.stats[self.prefix + "exitcode"] = 0 + self.stats[self.prefix + "iterations"] = self.engine.params.max_iterations + self.timer.start() + + def check_convergence(self): + """ Check the convergence. + + Returns: + bool: if True, the solver converges, otherwise(False), the solver does not converge. + """ + np.subtract(self.x, self.sol, self.error) + self.merit = np.linalg.norm(self.error, np.inf) + + if self.debug_on: + self.stats[self.prefix + "lambda"][self.iteration] = self.x + self.stats[self.prefix + "residuals"][self.iteration] = self.merit + + if self.merit < self.engine.params.absolute_tolerance: + if self.debug_on: + self.stats[self.prefix + "iterations"] = self.iteration + self.stats[self.prefix + "exitcode"] = 1 + self.timer.end() + self.stats[self.prefix + "solver_time"] = self.timer.elapsed + return True + if np.abs(self.merit - self.last_merit) < self.engine.params.relative_tolerance * self.last_merit: + if self.debug_on: + self.stats[self.prefix + "iterations"] = self.iteration + self.stats[self.prefix + "exitcode"] = 2 + self.timer.end() + self.stats[self.prefix + "solver_time"] = self.timer.elapsed + return True + return False + + @abstractmethod + def sweep(self): + """ Sweep the contact points to compute the new contact force. + This method should be implemented in the derived class. + """ + pass + + def solve(self): + """ Solve the contact problem. + + Returns: The new contact force and the statistics information. + """ + if self.debug_on: + self.initialize_stats(Timer("CONTACT_SOLVER")) + + # Compute initial r-factor value + self.compute_initial_r_factor() + + for iteration in range(self.engine.params.max_iterations): + self.iteration = iteration + self.sweep() + # Check convergence + if self.check_convergence(): + return self.x, self.stats + # Update r-factors + self.update_r_factor() + + # If this point of the code is reached then it means the method did not converge within the given iterations. + if self.debug_on: + self.timer.end() + self.stats[self.prefix + "solver_time"] = self.timer.elapsed + return self.sol, self.stats \ No newline at end of file diff --git a/python/unit_tests/test_random_graph_coloring.py b/python/unit_tests/test_random_graph_coloring.py new file mode 100644 index 0000000..31d8f92 --- /dev/null +++ b/python/unit_tests/test_random_graph_coloring.py @@ -0,0 +1,34 @@ +import unittest +import os +import sys +import numpy as np +import networkx as nx + +sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) + +import rainbow.simulators.proximal_contact.blocking as BLOCKING +import rainbow.util.test_tools as TEST + + +class TestBlocking(unittest.TestCase): + + def setUp(self): + self.K = 5 + self.G = nx.Graph() + self.G.add_nodes_from(np.arange(self.K)) + + def test_random_graph_coloring(self): + # Create a graph: 0-1-2-3-4, here the number is the node id. + self.G.add_edges_from([(0, 1), (1, 2), (2, 3), (3, 4)]) + color_groups = BLOCKING.random_graph_coloring(self.G) + + # Check if the color groups are valid + # The color groups should be: {1,2,4}, {1, 3} + expect_color_groups = {0: [0, 2, 4], 1: [1, 3]} + for k, v in color_groups.items(): + self.assertIn(k, expect_color_groups) + TEST.is_array_equal(v, expect_color_groups[k]) + + +if __name__ == "__main__": + unittest.main() \ No newline at end of file