Source code for boltz_data.draw.mol3d._layout

import math
from collections.abc import Generator

import numpy as np

from boltz_data.geom import BoundingBoxes

from ._primitive import Primitive


def enumerate_grid_indices(n: int) -> Generator[tuple[int, int]]:
    width = math.ceil(math.sqrt(n))
    height = math.ceil(n / width)
    for j in range(height):
        for i in range(width):
            if j * width + i == n:
                return
            yield (i, j)




def grid(*groups: list[Primitive], padding: float = 2) -> list[Primitive]:
    """
    Arrange multiple groups of primitives in a grid layout.

    Computes bounding boxes for each group and arranges them in a grid pattern
    (approximately square). All groups are sized to fit the largest bounding box
    plus padding.

    Args:
        *groups: Variable number of primitive groups to arrange.
        padding: Padding between grid cells.

    Returns:
        Flat list of all primitives with positions adjusted for grid layout.

    """
    bounding_boxes = BoundingBoxes.from_list(
        [BoundingBoxes.from_list([pri.bounding_box for pri in group]).bounding_box for group in groups]
    )
    largest_box = bounding_boxes.size.max(axis=0) + padding
    primitives: list[Primitive] = []
    for group_i, ((i, j), group) in enumerate(zip(enumerate_grid_indices(len(groups)), groups, strict=True)):
        offset = (
            np.array([i * largest_box[0], j * largest_box[1], 0]) - bounding_boxes[group_i].center + 0.5 * largest_box
        )
        primitives.extend(primitive + offset for primitive in group)
    return primitives