Benchmark

Author: kngwyu

benchmark/bench_phyjax2d.py

@@ -0,0 +1,130 @@
+import csv
+from datetime import datetime, timedelta
+from pathlib import Path
+
+import jax
+import jax.numpy as jnp
+import numpy as np
+import typer
+
+from phyjax2d import SpaceBuilder, Vec2d, step
+from phyjax2d.moderngl_vis import MglVisualizer
+
+
+def ball_fall_phyjax2d(
+    n_balls: int,
+    debug_vis: bool,
+    n_iter: int = 1000,
+) -> timedelta:
+    """
+    Simulates n_balls falling using phyjax2d.
+    If debug_vis is True, uses MglVisualizer for rendering.
+    """
+    builder = SpaceBuilder(
+        gravity=(0.0, -900.0),
+        dt=0.002,
+        jacobi_damping=0.5,
+        n_velocity_iter=10,
+        n_position_iter=2,
+        bias_factor=0.02,
+        bounce_threshold=4,
+    )
+
+    for _ in range(n_balls):
+        builder.add_circle(
+            radius=4.0,
+            density=1.0 / (16 * np.pi),
+            elasticity=0.5,
+            friction=0.5,
+        )
+
+    # Container setup
+    builder.add_segment(
+        p1=Vec2d(50.0, 50.0),
+        p2=Vec2d(550.0, 50.0),
+        elasticity=0.4,
+        friction=0.5,
+    )
+    builder.add_segment(
+        p1=Vec2d(50.0, 50.0),
+        p2=Vec2d(50.0, 800.0),
+        elasticity=0.4,
+        friction=0.5,
+    )
+    builder.add_segment(
+        p1=Vec2d(550.0, 50.0),
+        p2=Vec2d(550.0, 800.0),
+        elasticity=0.4,
+        friction=0.5,
+    )
+
+    space = builder.build()
+
+    # 2. Initialize State
+    rng = np.random.default_rng()
+    x_coords = rng.uniform(70, 530, n_balls)
+    y_coords = rng.uniform(400, 1000, n_balls)
+    pos_array = jnp.stack([jnp.array(x_coords), jnp.array(y_coords)], axis=-1)
+
+    sd = space.zeros_state().nested_replace("circle.p.xy", pos_array)
+    vs = space.init_solver()
+
+    # 3. Initialize Visualizer
+    visualizer = None
+    if debug_vis:
+        # We define the range based on the window size/container
+        visualizer = MglVisualizer(
+            x_range=600.0,
+            y_range=1000.0,
+            space=space,
+            stated=sd,
+            title=f"Phyjax2D Debug: {n_balls} balls",
+            figsize=(600, 1000),
+        )
+
+    # 4. Simulation Loop
+    jit_step = jax.jit(step, static_argnums=(0,))
+
+    # Warm-up (Optional: only if you want to exclude first-run JIT from benchmark)
+    sd, _, _ = jit_step(space, sd, vs)
+
+    start = datetime.now()
+    for _ in range(n_iter):
+        sd, _, _ = jit_step(space, sd, vs)
+
+        if visualizer is not None:
+            visualizer.render(state=sd)
+            visualizer.show()
+
+    duration = datetime.now() - start
+
+    if visualizer is not None:
+        visualizer.close()
+
+    return duration
+
+
+def main(
+    count_min: int = 1000,
+    count_max: int = 100000,
+    count_interval: int = 1000,
+    debug_vis: bool = False,
+    filename: Path = Path("bench.csv"),
+) -> None:
+    results = []
+
+    for count in [x for x in range(count_min, count_max + 1, count_interval)]:
+        duration = ball_fall_phyjax2d(count, debug_vis)
+        # Convert timedelta to total seconds as a float for the CSV
+        seconds = duration.total_seconds()
+        results.append((count, seconds))
+
+    if not debug_vis:
+        with open(filename, mode="w", newline="") as file:
+            writer = csv.writer(file)
+            writer.writerow(["n_balls", "duration_seconds"])  # Header
+            writer.writerows(results)
+
+
+if __name__ == "__main__":
+    typer.run(main)

benchmark/bench_pymunk.py

@@ -0,0 +1,98 @@
+import csv
+from datetime import datetime, timedelta
+from pathlib import Path
+
+import numpy as np
+import pygame
+import pymunk
+import pymunk.pygame_util
+import typer
+
+
+def ball_fall(n_balls: int, debug_vis: bool, n_iter: int = 1000) -> timedelta:
+    space = pymunk.Space()
+    # 1. Flip Gravity: Positive Y pulls "down" in PyGame coordinates
+    space.gravity = (0, 900)
+
+    static_body = space.static_body
+    # 2. Invert Container: Floor is now at Y=800, walls go up toward Y=50
+    segments = [
+        pymunk.Segment(static_body, (50, 800), (550, 800), 5),  # Floor (Bottom)
+        pymunk.Segment(static_body, (50, 800), (50, 50), 5),  # Left Wall
+        pymunk.Segment(static_body, (550, 800), (550, 50), 5),  # Right Wall
+    ]
+    for seg in segments:
+        seg.elasticity = 0.4
+        seg.friction = 0.5
+    space.add(*segments)
+
+    radius = 4
+    mass = 1
+    moment = pymunk.moment_for_circle(mass, 0, radius)
+
+    rng = np.random.default_rng()
+    # 3. Invert Spawn: Balls start near the top (Y=100 to 400)
+    x_coords = rng.uniform(70, 530, n_balls)
+    y_coords = rng.uniform(50, 400, n_balls)
+
+    for i in range(n_balls):
+        body = pymunk.Body(mass, moment)
+        body.position = (float(x_coords[i]), float(y_coords[i]))
+        shape = pymunk.Circle(body, radius)
+        shape.elasticity = 0.5
+        shape.friction = 0.5
+        space.add(body, shape)
+
+    if debug_vis:
+        pygame.init()
+        screen = pygame.display.set_mode((600, 850))
+        pygame.display.set_caption(f"Inverted Gravity Benchmark: {n_balls} balls")
+        draw_options = pymunk.pygame_util.DrawOptions(screen)
+    else:
+        screen = None
+        draw_options = None
+
+    start = datetime.now()
+    for _ in range(n_iter):
+        if screen is not None and draw_options is not None:
+            for event in pygame.event.get():
+                if event.type == pygame.QUIT:
+                    pygame.quit()
+                    return datetime.now() - start
+
+            screen.fill((255, 255, 255))
+            space.debug_draw(draw_options)
+            pygame.display.flip()
+
+        space.step(0.002)
+
+    if debug_vis:
+        pygame.quit()
+
+    return datetime.now() - start
+
+
+def main(
+    count_min: int = 1000,
+    count_max: int = 100000,
+    count_interval: int = 1000,
+    debug_vis: bool = False,
+    filename: Path = Path("bench.csv"),
+) -> None:
+    results = []
+
+    for count in [x for x in range(count_min, count_max + 1, count_interval)]:
+        duration = ball_fall(count, debug_vis)
+        # Convert timedelta to total seconds as a float for the CSV
+        seconds = duration.total_seconds()
+        results.append((count, seconds))
+
+    if not debug_vis:
+        with open(filename, mode="w", newline="") as file:
+            writer = csv.writer(file)
+            writer.writerow(["n_balls", "duration_seconds"])  # Header
+            writer.writerows(results)
+
+
+if __name__ == "__main__":
+    typer.run(main)

pyproject.toml

@@ -21,6 +21,7 @@ requires-python = ">= 3.10"
 dependencies = [
     "chex >= 0.1.86",
     "jax >= 0.4.26",
+    "pygame>=2.6.1",
 ]
 dynamic = ["version"]
 

src/phyjax2d/impl.py

@@ -861,6 +861,7 @@ class Space:
     dt: float = 0.1
     linear_damping: float = 0.95
     angular_damping: float = 0.95
+    jacobi_damping: float = 1.0
     bias_factor: float = 0.2
     n_velocity_iter: int = 6
     n_position_iter: int = 2
@@ -1014,7 +1015,7 @@ def __str__(self) -> str:
             f"  Timestep (dt): {self.dt}",
             f"  Damping: Linear={self.linear_damping}, Angular={self.angular_damping}",
             f"  Solver Iterations: Velocity={self.n_velocity_iter}, Position={self.n_position_iter}",
-            f"  Safety: Linear Slop={self.linear_slop}, Max Velocity={self.max_velocity}"
+            f"  Safety: Linear Slop={self.linear_slop}, Max Velocity={self.max_velocity}",
         ]
         return "\n".join(lines)
 
@@ -1282,8 +1283,13 @@ def solve_constraints(
     """Resolve collisions by Sequential Impulse method"""
     idx1, idx2 = space._ci_total.index1, space._ci_total.index2
 
-    def gather(a: jax.Array, b: jax.Array, orig: jax.Array) -> jax.Array:
-        return orig.at[idx1].add(a).at[idx2].add(b)
+    def gather(
+        a: jax.Array,
+        b: jax.Array,
+        orig: jax.Array,
+        damping: float = 1.0,
+    ) -> jax.Array:
+        return orig.at[idx1].add(a * damping).at[idx2].add(b * damping)
 
     p1, p2 = p.get_slice(idx1), p.get_slice(idx2)
     v1, v2 = v.get_slice(idx1), v.get_slice(idx2)
@@ -1317,7 +1323,7 @@ def vstep(
         (v.into_axy(), solver),
     )
     bv1, bv2 = apply_bounce(contact, helper, solver)
-    v_axy = gather(bv1, bv2, v_axy)
+    v_axy = gather(bv1, bv2, v_axy, damping=space.jacobi_damping)
 
     def pstep(
         _: int,

src/phyjax2d/utils.py

@@ -195,6 +195,7 @@ class SpaceBuilder:
     dt: float = 0.1
     linear_damping: float = 0.9
     angular_damping: float = 0.9
+    jacobi_damping: float = 1.0
     bias_factor: float = 0.2
     n_velocity_iter: int = 6
     n_position_iter: int = 2
@@ -525,10 +526,12 @@ def build(self) -> Space:
             jnp.inf if self.max_angular_velocity is None else self.max_angular_velocity
         )
         return Space(
+            dt=self.dt,
             gravity=jnp.array(self.gravity),
             shaped=shaped,
             linear_damping=linear_damping,
             angular_damping=angular_damping,
+            jacobi_damping=self.jacobi_damping,
             bias_factor=self.bias_factor,
             n_velocity_iter=self.n_velocity_iter,
             n_position_iter=self.n_position_iter,