Source code for modelly.viz.render.drawers.voronoi

"""Drawer for Voronoi diagram spaces, with Cohen-Sutherland line clipping."""

from itertools import pairwise

import altair as alt
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
from matplotlib.collections import LineCollection

from modelly.space import VoronoiGrid
from modelly.viz.render.drawers.base import (
    BaseSpaceDrawer,
    chart_properties,
    line_mark_kwargs,
)

# Cohen-Sutherland region codes
_INSIDE, _LEFT, _RIGHT, _BOTTOM, _TOP = 0, 1, 2, 4, 8


[docs] class VoronoiSpaceDrawer(BaseSpaceDrawer): """Drawer for Voronoi diagram spaces.""" def __init__(self, space: VoronoiGrid): """Initialize the Voronoi space drawer. Args: space: The Voronoi grid space to draw """ super().__init__(space) # Use the Cell.position property for calculations positions = np.array([cell.position for cell in self.space.all_cells]) if len(positions) > 0: x_min, y_min = positions.min(axis=0) x_max, y_max = positions.max(axis=0) else: x_max, x_min, y_max, y_min = 1, 0, 1, 0 self._set_padded_limits((x_min, x_max), (y_min, y_max)) def _clip_line(self, p1, p2, box): """Clips a line segment using the Cohen-Sutherland algorithm. Returns the clipped line segment (p1, p2) or None if it's outside. """ return clip_line(p1, p2, box) def _get_clipped_segments(self): """Helper method to perform the segment extraction, de-duplication and clipping logic.""" clip_box = ( self.viz_xmin, self.viz_ymin, self.viz_xmax, self.viz_ymax, ) unique_segments = set() for cell in self.space.all_cells.cells: vertices = [tuple(v) for v in cell.properties["polygon"]] for p1, p2 in pairwise([*vertices, vertices[0]]): # Sort to avoid duplicate segments going in opposite directions unique_segments.add(tuple(sorted((p1, p2)))) # Clip each unique segment final_segments = [] for p1, p2 in unique_segments: clipped_segment = self._clip_line(p1, p2, clip_box) if clipped_segment: final_segments.append(clipped_segment) return final_segments, clip_box
[docs] def draw_matplotlib(self, ax=None, **draw_space_kwargs): """Draw the Voronoi diagram using matplotlib. Args: ax: Matplotlib axes object to draw on **draw_space_kwargs: Keyword arguments passed to matplotlib's LineCollection. Examples: lw=2, alpha=0.5, colors='red' Returns: The modified axes object """ if ax is None: _, ax = plt.subplots() final_segments, clip_box = self._get_clipped_segments() ax.set_xlim(clip_box[0], clip_box[2]) ax.set_ylim(clip_box[1], clip_box[3]) if final_segments: # Define default styles for the plot style_args = {"colors": "k", "linestyle": "dotted", "lw": 1} style_args.update(draw_space_kwargs) # Create the LineCollection with the final styles lc = LineCollection(final_segments, **style_args) ax.add_collection(lc) return ax
[docs] def draw_altair(self, chart_width=450, chart_height=350, **draw_chart_kwargs): """Draw the Voronoi diagram using Altair. Args: chart_width: Width for the shown chart chart_height: Height for the shown chart **draw_chart_kwargs: Additional keyword arguments for styling the chart. Examples: * Line properties like color, strokeDash, strokeWidth, opacity. * Other kwargs (e.g., width, title) apply to the chart. Returns: An Altair Chart object representing the Voronoi diagram. """ final_segments, clip_box = self._get_clipped_segments() df = _segment_frame(final_segments) mark_kwargs = line_mark_kwargs(draw_chart_kwargs) chart_props = chart_properties(chart_width, chart_height, draw_chart_kwargs) chart = ( alt.Chart(df) .mark_line(**mark_kwargs) .encode( x=alt.X( "x:Q", scale=alt.Scale(domain=[clip_box[0], clip_box[2]]), axis=None ), y=alt.Y( "y:Q", scale=alt.Scale(domain=[clip_box[1], clip_box[3]]), axis=None ), detail="line_id:N", ) .properties(**chart_props) ) return chart
def _segment_frame(segments): """Return one row per segment end point, keyed by line id.""" final_data = [] for i, (p1, p2) in enumerate(segments): final_data.append({"x": p1[0], "y": p1[1], "line_id": i}) final_data.append({"x": p2[0], "y": p2[1], "line_id": i}) return pd.DataFrame(final_data) def _outcode(x, y, box): """Return the Cohen-Sutherland region code of a point.""" min_x, min_y, max_x, max_y = box code = _INSIDE if x < min_x: code |= _LEFT elif x > max_x: code |= _RIGHT if y < min_y: code |= _BOTTOM elif y > max_y: code |= _TOP return code def _boundary_point(start, end, outcode_out, box): """Return where the segment crosses the box edge named by the outside code.""" (x1, y1), (x2, y2) = start, end min_x, min_y, max_x, max_y = box x, y = 0.0, 0.0 # Check for horizontal line if y1 != y2: if outcode_out & _TOP: x = x1 + (x2 - x1) * (max_y - y1) / (y2 - y1) y = max_y elif outcode_out & _BOTTOM: x = x1 + (x2 - x1) * (min_y - y1) / (y2 - y1) y = min_y # Check for vertical line if x1 != x2: if outcode_out & _RIGHT: y = y1 + (y2 - y1) * (max_x - x1) / (x2 - x1) x = max_x elif outcode_out & _LEFT: y = y1 + (y2 - y1) * (min_x - x1) / (x2 - x1) x = min_x return x, y
[docs] def clip_line(p1, p2, box): """Clip a line segment to a (min_x, min_y, max_x, max_y) box (Cohen-Sutherland). Returns the clipped line segment (p1, p2) or None if it's outside. """ x1, y1 = p1 x2, y2 = p2 outcode1 = _outcode(x1, y1, box) outcode2 = _outcode(x2, y2, box) while True: if not (outcode1 | outcode2): # Both points inside return (x1, y1), (x2, y2) if outcode1 & outcode2: # Both points share an outside region return None outcode_out = outcode1 if outcode1 else outcode2 x, y = _boundary_point((x1, y1), (x2, y2), outcode_out, box) if outcode_out == outcode1: x1, y1 = x, y outcode1 = _outcode(x1, y1, box) else: x2, y2 = x, y outcode2 = _outcode(x2, y2, box)