combine_mozaic.py

from math import gcd
from functools import reduce

def lcm(a, b):
    return a * b // gcd(a, b) if a and b else max(a, b)

def _lcm_list(lst):
    return reduce(lcm, lst, 1)

def _repeat_chars(line, times):
    return "".join(c * times for c in line)

def _transpose(block):
    if not block:
        return []
    max_len = max(len(row) for row in block)
    block = [row.ljust(max_len) for row in block]
    return ["".join(block[r][c] for r in range(len(block))) for c in range(max_len)]

def _check_unique_letters(*blocks):
    """
    Ensure all blocks have unique letters across blocks.
    Raises an AssertionError if any letter appears in more than one block.
    """
    unique = set()
    for i, block in enumerate(blocks, 1):
        letters = set(block.replace('\n', ''))
        assert not (letters & unique), f"Duplicate letters found in block {i}: {letters & unique}"
        unique.update(letters)

def combine_mosaics(*blocks, ratio=None, orient='v'):
    if len(blocks) < 2:
        raise ValueError("Need at least two blocks to combine")

    _check_unique_letters(*blocks)
    blocks = [_format_block(block) for block in blocks]

    # Normalize input
    blocks_lines = [block.split("\n") for block in blocks]

    # Normalize ratio
    if ratio is None:
        ratios = [1.0] * len(blocks_lines)
    else:
        try:
            ratios = list(ratio)
            if len(ratios) != len(blocks_lines):
                raise ValueError
        except Exception:
            ratios = [float(ratio)] * len(blocks_lines)

    # Transpose if horizontal
    transposed = False
    if orient == 'v':
        blocks_lines = [_transpose(b) for b in blocks_lines]
        transposed = True

    # Horizontal expansion (columns)
    cols_list = [max(len(line) for line in b) if b else 0 for b in blocks_lines]
    Lw = _lcm_list(cols_list)
    blocks_expanded = []
    for b, c, r in zip(blocks_lines, cols_list, ratios):
        b = [line.ljust(c) for line in b]
        h = max(1, int(round(Lw / c * r)))
        blocks_expanded.append([_repeat_chars(line, h) for line in b])

    # Vertical expansion (rows)
    rows_list = [len(b) for b in blocks_expanded]
    Lh = _lcm_list(rows_list)
    blocks_tiled = []
    for b, r in zip(blocks_expanded, ratios):
        v = max(1, int(round(Lh / len(b))))
        blocks_tiled.append([line for line in b for _ in range(v)])

    # Combine all blocks
    combined = ["".join(lines) for lines in zip(*blocks_tiled)]

    # Transpose back if needed
    if transposed:
        combined = _transpose(combined)

    return _format_block("\n".join(combined))


def _format_block(mosaic:str) -> str:
    return mosaic.replace(' ', '').lstrip('\n').rstrip('\n')

def shrink_ax(ax, shrink=0.03):
      # shrink 3% each subplot
    pos = ax.get_position()
    # shrink from all sides
    new_pos = [
        pos.x0 + shrink/2,
        pos.y0 + shrink/2,
        pos.width - shrink,
        pos.height - shrink
    ]
    ax.set_position(new_pos)

def label_ax(ax, label, width=0.12, height=0.12, x=-0.15, y=1.05, fontsize=14, fontweight='bold', facecolor='none', edgecolor='none'):
    # Create inset axes
    label_ax = ax.inset_axes([x, y, width, height])
    label_ax.text(0.5, 0.5, label, fontsize=fontsize, fontweight=fontweight,
                  ha='center', va='center')
    label_ax.set_facecolor(facecolor)
    for spine in label_ax.spines.values():
        spine.set_edgecolor(edgecolor)
    label_ax.set_xticks([])
    label_ax.set_yticks([])


import matplotlib.pyplot as plt

if True:
    top = """
    AB
    CD
    """
    bottom = "EFGHI"
    
    mosaic = combine_mosaics(top, bottom, orient='h')
    expect = """
    AAAAABBBBBEEFFGGHHII
    CCCCCDDDDDEEFFGGHHII
    """
    assert mosaic == _format_block(expect)
    
    mosaic = combine_mosaics(top, bottom, ratio=[1., .5], orient='h')
    expect = """
    AAAAABBBBBEFGHI
    CCCCCDDDDDEFGHI
    """
    assert mosaic == _format_block(expect)
    
    mosaic = combine_mosaics(top, bottom, orient='v')
    expect = """
    AAAAABBBBB
    CCCCCDDDDD
    EEFFGGHHII
    EEFFGGHHII
    """
    assert mosaic == _format_block(expect)
    
    mosaic = combine_mosaics(top, bottom, ratio=[2., 1.])
    expect = """
    AAAAABBBBB
    AAAAABBBBB
    CCCCCDDDDD
    CCCCCDDDDD
    EEFFGGHHII
    EEFFGGHHII
    """
    assert mosaic == _format_block(expect)
    fig, axes = plt.subplot_mosaic(mosaic)
    axes['A'].plot(range(10), range(10))
    axes['I'].plot(range(10), range(10))