taps.apps.physics

taps/apps/physics.py

TerrainConfig dataclass

TerrainConfig(
    width: int,
    height: float,
    resolution: int,
    scale: float,
    octaves: int,
    persistence: float,
    lacunarity: float,
    filter_size: int,
)

Terrain configuration.

SimulationConfig dataclass

SimulationConfig(
    ball_diameter: float,
    ball_mass: float,
    tick_rate: int,
    total_time: int,
    real_time: bool,
)

Simulation configuration.

PhysicsApp

PhysicsApp(
    num_simulations: int,
    terrain: TerrainConfig,
    simulation: SimulationConfig,
    seed: int | None = None,
)

Physics simulation application.

Simulate the physics of golf balls landing and rolling on a golf green.

Parameters:

• num_simulations (int) –

  Number of balls to simulate.

• terrain (TerrainConfig) –

  Terrain configuration.

• simulation (SimulationConfig) –

  Simulation configuration.

• seed (int | None, default: None ) –

  Random seed.

Source code in taps/apps/physics.py

def __init__(
    self,
    num_simulations: int,
    terrain: TerrainConfig,
    simulation: SimulationConfig,
    seed: int | None = None,
) -> None:
    self.num_simulations = num_simulations
    self.terrain = terrain
    self.simulation = simulation
    self.seed = seed

close

close() -> None

Close the application.

Source code in taps/apps/physics.py

def close(self) -> None:
    """Close the application."""
    pass

run

run(engine: Engine, run_dir: Path) -> None

Run the application.

Parameters:

• engine (Engine) –

  Application execution engine.

• run_dir (Path) –

  Run directory.

Source code in taps/apps/physics.py

def run(self, engine: Engine, run_dir: pathlib.Path) -> None:
    """Run the application.

    Args:
        engine: Application execution engine.
        run_dir: Run directory.
    """
    terrain_heightmap = generate_noisemap(self.terrain, seed=self.seed)
    terrain_mesh = generate_vertices(terrain_heightmap, self.terrain)
    logger.info('Generated terrain mesh')

    initial_positions = generate_initial_positions(
        self.num_simulations,
        self.terrain,
        seed=self.seed,
    )
    logger.info(f'Generated {len(initial_positions)} initial position(s)')

    logger.info('Submitting simulations...')
    futures = [
        engine.submit(
            simulate,
            terrain_mesh,
            [position],
            sim_config=self.simulation,
            terrain_config=self.terrain,
        )
        for position in initial_positions
    ]
    logger.info('Simulations submitted')

    results = [future.result() for future in futures]
    final_positions = [pos for result in results for pos in result]
    logger.info(f'Received {len(final_positions)} final position(s)')

    contour_plot = run_dir / 'images' / 'contour.png'
    create_contour_plot(
        initial_positions,
        final_positions,
        terrain_heightmap,
        self.terrain,
        contour_plot,
    )
    logger.info(f'Saved contour map to {contour_plot}')

    terrain_plot = run_dir / 'images' / 'terrain.png'
    create_terrain_plot(
        initial_positions,
        final_positions,
        terrain_heightmap,
        self.terrain,
        terrain_plot,
    )
    logger.info(f'Saved terrain map to {terrain_plot}')

create_contour_plot

create_contour_plot(
    initial_positions: list[Position],
    final_positions: list[Position],
    heightmap: NDArray[float64],
    config: TerrainConfig,
    filepath: str | Path,
) -> None

Write contour plot with initial/final ball positions.

Parameters:

• initial_positions (list[Position]) –

  Initial ball positions.

• final_positions (list[Position]) –

  Final ball positions.

• heightmap (NDArray[float64]) –

  Terrain heightmap.

• config (TerrainConfig) –

  Terrain configuration.

• filepath (str | Path) –

  Output file.

Source code in taps/apps/physics.py

def create_contour_plot(
    initial_positions: list[Position],
    final_positions: list[Position],
    heightmap: NDArray[numpy.float64],
    config: TerrainConfig,
    filepath: str | pathlib.Path,
) -> None:
    """Write contour plot with initial/final ball positions.

    Args:
        initial_positions: Initial ball positions.
        final_positions: Final ball positions.
        heightmap: Terrain heightmap.
        config: Terrain configuration.
        filepath: Output file.
    """
    fig, axs = plt.subplots(1, 2, sharey=True, figsize=(8, 4))

    x = numpy.linspace(0, config.width, num=config.width * config.resolution)
    y = numpy.linspace(0, config.width, num=config.width * config.resolution)
    for ax, positions in zip(axs, (initial_positions, final_positions)):
        handle = ax.contour(x, y, heightmap, levels=10)
        plt.clabel(handle, inline=True)
        px, py = (
            # Clamp balls that may have rolled off the surface to the
            # edge of the plot.
            [max(0, min(p[0], config.width)) for p in positions],
            [max(0, min(p[1], config.width)) for p in positions],
        )
        ax.scatter(px, py, s=16, c='#FFFFFF', zorder=100)

    for i, (ax, title) in enumerate(zip(axs, ('Initial', 'Final'))):
        ax.set_title(title)
        ax.set_xlim(0, config.width)
        ax.set_ylim(0, config.width)
        ax.set_xlabel('X Position (m)')
        if i == 0:
            ax.set_ylabel('Y Position (m)')
        ax.set_facecolor('#58a177')

    fig.tight_layout(w_pad=2)

    pathlib.Path(filepath).parent.mkdir(parents=True, exist_ok=True)
    fig.savefig(filepath, pad_inches=0.05, dpi=300)

create_terrain_plot

create_terrain_plot(
    initial_positions: list[Position],
    final_positions: list[Position],
    heightmap: NDArray[float64],
    config: TerrainConfig,
    filepath: str | Path,
) -> None

Write 3D terrain mesh with heatmap of final ball positions.

Parameters:

• initial_positions (list[Position]) –

  Initial ball positions.

• final_positions (list[Position]) –

  Final ball positions.

• heightmap (NDArray[float64]) –

  Terrain heightmap.

• config (TerrainConfig) –

  Terrain configuration.

• filepath (str | Path) –

  Output file.

Source code in taps/apps/physics.py

def create_terrain_plot(
    initial_positions: list[Position],
    final_positions: list[Position],
    heightmap: NDArray[numpy.float64],
    config: TerrainConfig,
    filepath: str | pathlib.Path,
) -> None:
    """Write 3D terrain mesh with heatmap of final ball positions.

    Args:
        initial_positions: Initial ball positions.
        final_positions: Final ball positions.
        heightmap: Terrain heightmap.
        config: Terrain configuration.
        filepath: Output file.
    """
    fig = plt.figure(figsize=(8, 4))
    axs = [
        fig.add_subplot(1, 2, 1, projection='3d'),
        fig.add_subplot(1, 2, 2, projection='3d'),
    ]
    plt.subplots_adjust(wspace=0.5, hspace=0.5)

    x = numpy.arange(0, config.width, 1 / config.resolution)
    y = numpy.arange(0, config.width, 1 / config.resolution)
    x, y = numpy.meshgrid(x, y)

    for ax, positions in zip(axs, (initial_positions, final_positions)):
        px, py, _ = zip(*positions)
        xy = numpy.vstack([px, py])
        kde = gaussian_kde(xy)(xy)
        kde_grid = numpy.zeros_like(heightmap)

        for i, (xi, yi) in enumerate(zip(px, py)):
            # Clamp positions to be in [0, config.width]. If a ball rolled
            # off the edge, it's position would be outside the mesh map.
            max_index = (config.width * config.resolution) - 1
            xi_scaled = int(xi * config.resolution)
            yi_scaled = int(yi * config.resolution)
            xi_scaled = max(0, min(xi_scaled, max_index))
            yi_scaled = max(0, min(yi_scaled, max_index))
            kde_grid[yi_scaled, xi_scaled] = kde[i]

        kde_smoothed = gaussian_filter(kde_grid, sigma=1)
        ax.plot_surface(
            x,
            y,
            heightmap,
            facecolors=plt.cm.jet(kde_smoothed / kde_smoothed.max()),
            alpha=0.9,
        )

    for ax, title in zip(axs, ('Initial', 'Final')):
        ax.set_title(title, pad=-20)
        ax.set_xlim(0, config.width)
        ax.set_ylim(0, config.width)
        ax.set_zlim(0, 2 * heightmap.max())
        ax.set_xlabel('X Position (m)')
        ax.set_ylabel('Y Position (m)')
        ax.set_zlabel('Z Position (m)')

    fig.tight_layout()
    fig.subplots_adjust(wspace=0.15, left=0, right=0.92, bottom=0.05, top=0.98)

    pathlib.Path(filepath).parent.mkdir(parents=True, exist_ok=True)
    fig.savefig(filepath, dpi=300)

generate_initial_positions

generate_initial_positions(
    num_balls: int,
    config: TerrainConfig,
    seed: int | None = None,
) -> list[Position]

Generate initial ball positions.

Parameters:

• num_balls (int) –

  Number of balls.

• config (TerrainConfig) –

  Terrain config used to bound locations of balls.

• seed (int | None, default: None ) –

  Random seed for initial positions.

Returns:

• list[Position] –

  List of ball positions.

Source code in taps/apps/physics.py

def generate_initial_positions(
    num_balls: int,
    config: TerrainConfig,
    seed: int | None = None,
) -> list[Position]:
    """Generate initial ball positions.

    Args:
        num_balls: Number of balls.
        config: Terrain config used to bound locations of balls.
        seed: Random seed for initial positions.

    Returns:
        List of ball positions.
    """
    buffer = 0.2 * config.width
    min_width, max_width = buffer, config.width - buffer

    random.seed(seed)

    def _generate() -> Position:
        return (
            random.uniform(min_width, max_width),
            random.uniform(min_width, max_width),
            2 * config.height,
        )

    return [_generate() for _ in range(num_balls)]

generate_noisemap

generate_noisemap(
    config: TerrainConfig, seed: int | None = None
) -> NDArray[float64]

Generate Perlin noise map for terrain generation.

Parameters:

• config (TerrainConfig) –

  Terrain configuration.

• seed (int | None, default: None ) –

  Random seed.

Returns:

• NDArray[float64] –

  Noise map.

Source code in taps/apps/physics.py

def generate_noisemap(
    config: TerrainConfig,
    seed: int | None = None,
) -> NDArray[numpy.float64]:
    """Generate Perlin noise map for terrain generation.

    Args:
        config: Terrain configuration.
        seed: Random seed.

    Returns:
        Noise map.
    """
    dimension = config.width * config.resolution
    heightmap = numpy.zeros((dimension, dimension))

    offset = seed if seed is not None else random.randint(0, 1_000_000)

    for i in range(dimension):
        for j in range(dimension):
            x = (i / config.resolution) / config.scale
            y = (j / config.resolution) / config.scale
            heightmap[i][j] = pnoise2(
                x + offset,
                y + offset,
                octaves=config.octaves,
                persistence=config.persistence,
                lacunarity=config.lacunarity,
            )

    # Smooth terrain with gaussian filter
    heightmap = gaussian_filter(heightmap, config.filter_size)
    # Scale terrain height to be [0, config.height]
    old_min, old_max = heightmap.min(), heightmap.max()
    return (heightmap - old_min) * config.height / old_max

generate_vertices

generate_vertices(
    heightmap: NDArray[float64], config: TerrainConfig
) -> Terrain

Generate vertex mesh from heighmap.

Parameters:

• heightmap (NDArray[float64]) –

  Terrain heighmap.

• config (TerrainConfig) –

  Terrain config.

Returns:

• Terrain –

  Vertex mesh.

Source code in taps/apps/physics.py

def generate_vertices(
    heightmap: NDArray[numpy.float64],
    config: TerrainConfig,
) -> Terrain:
    """Generate vertex mesh from heighmap.

    Args:
        heightmap: Terrain heighmap.
        config: Terrain config.

    Returns:
        Vertex mesh.
    """
    width, height = heightmap.shape[0], heightmap.shape[1]
    vertices = width * height
    parts = math.ceil(vertices / MAX_VERTICES_PER_MESH)
    max_width_per_part = math.ceil(width / parts)

    terrain: Terrain = []
    for x1 in range(0, width, max_width_per_part):
        # Add one to right index of slice so that the terrain parts
        # overlap each other slightly.
        x2 = min(width, x1 + max_width_per_part + 1)
        # For now, we only split on axis 0 (referred to as x-axis).
        y = 0 / config.resolution
        heightmap_part = heightmap[x1:x2, :]
        terrain_part = _generate_vertices(heightmap_part, config.resolution)
        terrain.append((x1 / config.resolution, y, terrain_part))

    return terrain

simulate

simulate(
    terrain: Terrain,
    positions: list[Position],
    sim_config: SimulationConfig,
    terrain_config: TerrainConfig,
) -> list[Position]

Simulate balls landing on terrain.

Parameters:

• terrain (Terrain) –

  Terrain.

• positions (list[Position]) –

  Initial ball positions.

• sim_config (SimulationConfig) –

  Simulation configuration.

• terrain_config (TerrainConfig) –

  Terrain configuration.

Returns:

• list[Position] –

  List of final ball positions.

Source code in taps/apps/physics.py

@task()
def simulate(
    terrain: Terrain,
    positions: list[Position],
    sim_config: SimulationConfig,
    terrain_config: TerrainConfig,
) -> list[Position]:
    """Simulate balls landing on terrain.

    Args:
        terrain: Terrain.
        positions: Initial ball positions.
        sim_config: Simulation configuration.
        terrain_config: Terrain configuration.

    Returns:
        List of final ball positions.
    """
    pybullet.connect(pybullet.DIRECT)
    pybullet.setAdditionalSearchPath(pybullet_data.getDataPath())
    pybullet.setTimeStep(1 / sim_config.tick_rate)
    pybullet.setGravity(0, 0, -9.81)

    _create_terrain_body(terrain)

    ball_ids = [
        _create_ball_body(
            position,
            radius=sim_config.ball_diameter / 2,
            mass=sim_config.ball_mass,
            max_velocity=0.1 * terrain_config.width,
        )
        for position in positions
    ]

    logger.debug(
        f'Simulating for {sim_config.total_time} '
        f'({sim_config.tick_rate} steps per seconds)',
    )

    for _ in range(sim_config.total_time * sim_config.tick_rate):
        pybullet.stepSimulation()
        if sim_config.real_time:
            time.sleep(1 / sim_config.tick_rate)

    logger.debug('Simulation completed')

    final_positions = [
        pybullet.getBasePositionAndOrientation(ball_id)[0]
        for ball_id in ball_ids
    ]

    pybullet.disconnect()

    return final_positions