Rectangle packing

Rectangle Packing¶

Pack rectangles of given sizes to minimize overlap and overall bounding box area. The collision loss uses axis-aligned bounding-box (AABB) intersection area, computed with vizopt.components.multiple_bbox_intersections.

%config InlineBackend.figure_formats = ['svg']

import pathlib

import numpy as np
import pandas as pd
from IPython.display import SVG
from jax import numpy as jnp
from matplotlib import patches as mpatches
from matplotlib import pyplot as plt

from vizopt.animation import SnapshotCallback, snapshots_to_animated_svg
from vizopt.base import ObjectiveTerm, OptimConfig, OptimizationProblemTemplate
from vizopt.components import multiple_bbox_intersections

fig_folder = pathlib.Path("output") / "figs"
fig_folder.mkdir(parents=True, exist_ok=True)

rng = np.random.default_rng(0)
n_rects = 30
half_sizes = rng.uniform(0.1, 0.8, size=(n_rects, 2)).astype(np.float32)
total_scale = float(half_sizes.sum())
initial_positions = rng.uniform(0, total_scale / n_rects * 4, size=(n_rects, 2)).astype(np.float32)

def plot_rectangles(positions, half_sizes, title="Rectangle packing"):
    fig, ax = plt.subplots(figsize=(5, 5))
    colors = plt.cm.tab20.colors
    for i, (xy, hw) in enumerate(zip(positions, half_sizes)):
        rect = mpatches.FancyBboxPatch(
            xy - hw, 2 * hw[0], 2 * hw[1],
            boxstyle="square,pad=0",
            color=colors[i % len(colors)], alpha=0.6, ec="black", linewidth=0.5,
        )
        ax.add_patch(rect)
    margin = float(half_sizes.max())
    ax.set_xlim(positions[:, 0].min() - margin, positions[:, 0].max() + margin)
    ax.set_ylim(positions[:, 1].min() - margin, positions[:, 1].max() + margin)
    ax.set_aspect("equal")
    ax.set_title(title)
    plt.axis("off")
    plt.tight_layout()

plot_rectangles(initial_positions, half_sizes, title=f"Initial: {n_rects} rectangles")

output

Optimization problem¶

Two loss terms: - collision: sum of pairwise AABB intersection areas (should be zero when rectangles don't touch) - total_size: sum of the overall bounding box width and height

_TAB20 = [
    "#1f77b4", "#aec7e8", "#ff7f0e", "#ffbb78", "#2ca02c", "#98df8a",
    "#d62728", "#ff9896", "#9467bd", "#c5b0d5", "#8c564b", "#c49c94",
    "#e377c2", "#f7b6d2", "#7f7f7f", "#c7c7c7", "#bcbd22", "#dbdb8d",
    "#17becf", "#9edae5",
]


def _term_collision(optim_vars, input_params):
    positions = optim_vars["positions"]           # (n, 2)
    half_sizes = input_params["half_sizes"]       # (n, 2)
    bboxes = jnp.stack([positions - half_sizes, positions + half_sizes], axis=1)  # (n, 2, 2)
    overlaps = multiple_bbox_intersections(bboxes, bboxes)  # (n, n)
    n = half_sizes.shape[0]
    rows, cols = np.triu_indices(n, 1)
    return jnp.sum(overlaps[rows, cols])


def _term_total_size(optim_vars, input_params):
    positions = optim_vars["positions"]
    half_sizes = input_params["half_sizes"]
    max_coord = jnp.max(positions + half_sizes, axis=0)
    min_coord = jnp.min(positions - half_sizes, axis=0)
    return jnp.sum(max_coord - min_coord)


def _plot_configuration(optim_vars, input_params):
    positions = np.array(optim_vars["positions"])
    half_sizes = input_params["half_sizes"]
    plot_rectangles(positions, half_sizes, title=f"Optimized: {len(half_sizes)} rectangles")


def _svg_configuration(snapshots, input_params, size):
    half_sizes = input_params["half_sizes"]  # (n, 2)
    n = len(half_sizes)

    all_pos = np.concatenate([v["positions"] for _, v in snapshots], axis=0)
    margin = float(half_sizes.max())
    xmin = float(all_pos[:, 0].min()) - margin
    ymin = float(all_pos[:, 1].min()) - margin
    span = max(
        float(all_pos[:, 0].max()) + margin - xmin,
        float(all_pos[:, 1].max()) + margin - ymin,
    )

    def to_svg(x, y):
        return (x - xmin) / span * size, (span - (y - ymin)) / span * size

    elements = []
    for i in range(n):
        hw = half_sizes[i]  # (2,)
        w = hw[0] / span * size * 2
        h = hw[1] / span * size * 2
        xs, ys = [], []
        for _, v in snapshots:
            cx, cy = to_svg(float(v["positions"][i, 0]), float(v["positions"][i, 1]))
            xs.append(f"{cx - w / 2:.1f}")
            ys.append(f"{cy - h / 2:.1f}")
        elements.append({
            "tag": "rect",
            "width": f"{w:.1f}",
            "height": f"{h:.1f}",
            "fill": _TAB20[i % len(_TAB20)],
            "fill-opacity": "0.6",
            "stroke": "black",
            "stroke-width": "0.5",
            "x": xs,
            "y": ys,
        })
    return elements


def _initialize(input_params, seed):
    return {"positions": input_params["initial_positions"].copy()}


rectangle_packing_template = OptimizationProblemTemplate(
    terms=[
        ObjectiveTerm("total_size", _term_total_size, multiplier=1.0),
        ObjectiveTerm("collision", _term_collision, multiplier=5.0),
    ],
    initialize=_initialize,
    plot_configuration=_plot_configuration,
    svg_configuration=_svg_configuration,
)

input_parameters = {
    "half_sizes": half_sizes,
    "initial_positions": initial_positions,
}

problem = rectangle_packing_template.instantiate(input_parameters)
optim_vars, history = problem.optimize(
    OptimConfig(n_iters=5000, learning_rate=0.01), track_every=100
)

pd.DataFrame(history).set_index("iteration").plot(marker=".")
plt.ylabel("Loss value")

output

problem.plot_configuration(optim_vars, input_parameters)

output

Animation¶

cb = SnapshotCallback(every=50)
problem_anim = rectangle_packing_template.instantiate(input_parameters)
problem_anim.optimize(OptimConfig(n_iters=5000, learning_rate=0.01), callback=cb)

svg_str = snapshots_to_animated_svg(problem_anim, cb.snapshots, fps=10)
SVG(data=svg_str)

output