← Files ClaraARCHIVED FILE
modules/period-comparison/vendor/modules/charting/draw_bubble.py
38.2 KB · Oct 2, 2026 · 00:29 UTC
import copy
import logging
import math
import re
import numpy as np
import plotly.express as px
import plotly.graph_objects as go
import polars as pl
from modules.charting.chart_primitives import (
get_color_array,
get_color_dictionary,
get_hightlight_color,
millify,
)
from modules.charting.draw_charts_utils import get_polars_value_at_index
from modules.charting.polars_helpers import unique_values_lazy
from modules.utilities.config import (
get_config_params,
get_metric_array_params,
get_naming_params,
)
from modules.utilities.helpers import (
duplicate_dataframe,
unique,
)
from modules.utilities.ui_notifier import ui
from modules.utilities.utils import (
ensure_lazyframe,
ensure_polars_df,
get_schema_and_column_names,
)
def add_split_lines(fig):
fig.add_hline(
y=0.5,
opacity=1,
line_width=0.5,
line_color="lightgrey",
yref="y domain",
)
fig.add_vline(
x=0.5,
opacity=1,
line_width=0.5,
line_color="lightgrey",
xref="x domain",
)
return fig
def _safe_float(value: object) -> float | None:
try:
result = float(value)
except (TypeError, ValueError):
return None
if not math.isfinite(result):
return None
return result
def _bubble_axis_range(values: list[float]) -> tuple[float, float]:
if not values:
return 0.0, 1.0
minimum = min(values)
maximum = max(values)
if minimum == maximum:
padding = max(abs(minimum) * 0.1, 1.0)
return minimum - padding, maximum + padding
minimum_multiplier = 1.1 if minimum < 0 else 0.9
maximum_multiplier = 1.1
return minimum * minimum_multiplier, maximum * maximum_multiplier
def _estimate_text_box(text: str, font_size: int | float) -> tuple[float, float]:
label = str(text or "")
longest_line = max(label.split("<br>") or [""], key=len)
width = max(24.0, len(longest_line) * float(font_size) * 0.58)
height = max(14.0, (label.count("<br>") + 1) * float(font_size) * 1.25)
return width, height
def _boxes_overlap(
left: tuple[float, float, float, float],
right: tuple[float, float, float, float],
padding: float = 3.0,
) -> bool:
return not (
left[2] + padding <= right[0]
or right[2] + padding <= left[0]
or left[3] + padding <= right[1]
or right[3] + padding <= left[1]
)
def _bubble_text_box(
x_px: float,
y_px: float,
x_shift: float,
y_shift: float,
width: float,
height: float,
) -> tuple[float, float, float, float]:
center_x = x_px + x_shift
center_y = y_px + y_shift
return (
center_x - (width / 2),
center_y - (height / 2),
center_x + (width / 2),
center_y + (height / 2),
)
def _bubble_label_candidates(
label_width: float,
radius_px: float,
) -> list[tuple[int, int]]:
gap = max(radius_px + 10.0, 14.0)
side_gap = max(radius_px + (label_width / 2) + 8.0, 28.0)
return [
(0, int(round(gap))),
(-int(round(label_width * 0.35)), int(round(gap + 10))),
(int(round(label_width * 0.35)), int(round(gap + 10))),
(-int(round(side_gap)), 0),
(int(round(side_gap)), 0),
(0, -int(round(gap))),
(-int(round(label_width * 0.35)), -int(round(gap + 10))),
(int(round(label_width * 0.35)), -int(round(gap + 10))),
]
def _is_aggregate_other_label(value: object) -> bool:
label = str(value or "").strip().lower()
return label.startswith("other rank >") or label.startswith("others rank >")
def _metric_tokens(metric: object) -> set[str]:
return set(re.findall(r"[a-z0-9]+", str(metric or "").lower()))
def _configured_non_summable_metrics() -> set[str]:
namingParams = get_naming_params()
metricArrayParams = get_metric_array_params()
metric_groups = [
namingParams["noSumMetricsArray"],
namingParams["growthMetricArray"],
namingParams["percentMetricsArray"],
namingParams["priceMetricsArray"],
]
return {
str(metric).strip().lower()
for group in metric_groups
for metric in metricArrayParams[group]
}
def _axis_metric_is_non_summable(metric: object, chartDict: dict) -> bool:
# Total bubble coordinates are averages, so the marker is only meaningful
# when both plotted axes are non-additive metrics.
metric_label = str(metric or "").strip().lower()
if metric_label in _configured_non_summable_metrics():
return True
namingParams = get_naming_params()
countMetricsAvgArray = namingParams["countMetricsAvgArray"]
if metric_label in {
str(item).strip().lower() for item in chartDict.get(countMetricsAvgArray, [])
}:
return True
tokens = _metric_tokens(metric)
return bool(
tokens
& {
"avg",
"average",
"change",
"coverage",
"cwd",
"distribution",
"growth",
"index",
"pct",
"percent",
"percentage",
"price",
"rate",
"ratio",
"share",
"variance",
}
)
def _can_plot_total_bubble(chartDict: dict) -> bool:
namingParams = get_naming_params()
plotTotalBubble = namingParams["plotTotalBubble"]
xAxisMetric = namingParams["xAxisMetric"]
yAxisMetric = namingParams["yAxisMetric"]
if not chartDict.get(plotTotalBubble, False):
return False
return _axis_metric_is_non_summable(
chartDict[xAxisMetric], chartDict
) and _axis_metric_is_non_summable(chartDict[yAxisMetric], chartDict)
def _select_bubble_label_positions(
rows: list[dict[str, object]],
*,
x_range: tuple[float, float],
y_range: tuple[float, float],
size_ref: float | bool | None,
font_size: int | float,
show_label: bool,
show_value: bool,
plot_width: int = 900,
plot_height: int = 600,
collision_padding: float = 3.0,
allowed_edge_overflow: float = 24.0,
) -> dict[int, tuple[int, int]]:
"""Select readable bubble labels using mechanically verifiable box overlap."""
x_span = max(x_range[1] - x_range[0], 1.0)
y_span = max(y_range[1] - y_range[0], 1.0)
numeric_size_ref = (
float(size_ref)
if isinstance(size_ref, (int, float)) and math.isfinite(float(size_ref))
else 0.0
)
size_values = [
row["size_float"]
for row in rows
if isinstance(row.get("size_float"), (int, float))
]
size_min = min(size_values) if size_values else 0.0
size_span = max((max(size_values) - size_min) if size_values else 0.0, 1.0)
scored_rows: list[tuple[float, int, dict[str, object]]] = []
for order, row in enumerate(rows):
x_float = row.get("x_float")
y_float = row.get("y_float")
size_float = row.get("size_float")
if not isinstance(x_float, (int, float)) or not isinstance(
y_float, (int, float)
):
continue
x_norm = (float(x_float) - x_range[0]) / x_span
y_norm = (float(y_float) - y_range[0]) / y_span
size_norm = (
(float(size_float) - size_min) / size_span
if isinstance(size_float, (int, float))
else 0.0
)
edge_score = abs(x_norm - 0.5) + abs(y_norm - 0.5)
aggregate_other_bonus = (
4.0 if _is_aggregate_other_label(row.get("label")) else 0.0
)
score = (size_norm * 2.0) + (edge_score * 0.65) + aggregate_other_bonus
scored_rows.append((score, -order, row))
occupied: list[tuple[float, float, float, float]] = []
selected: dict[int, tuple[int, int]] = {}
for _score, _order, row in sorted(scored_rows, reverse=True):
label_text = str(row.get("label", ""))
value_text = str(row.get("value", ""))
primary_text = label_text if show_label else value_text
if not primary_text:
continue
x_float = float(row["x_float"])
y_float = float(row["y_float"])
x_px = ((x_float - x_range[0]) / x_span) * plot_width
y_px = ((y_float - y_range[0]) / y_span) * plot_height
label_width, label_height = _estimate_text_box(primary_text, font_size)
value_width, value_height = _estimate_text_box(value_text, font_size)
size_float = row.get("size_float")
radius_px = 0.0
if numeric_size_ref > 0 and isinstance(size_float, (int, float)):
radius_px = max(
2.0, math.sqrt(max(float(size_float), 0.0) / numeric_size_ref) / 2.0
)
if show_label:
candidates = _bubble_label_candidates(label_width, radius_px)
else:
candidates = [(0, 0)]
for x_shift, y_shift in candidates:
candidate_boxes: list[tuple[float, float, float, float]] = []
if show_value and value_text:
candidate_boxes.append(
_bubble_text_box(x_px, y_px, 0, 0, value_width, value_height)
)
if show_label and label_text:
candidate_boxes.append(
_bubble_text_box(
x_px,
y_px,
x_shift,
y_shift,
label_width,
label_height,
)
)
if not candidate_boxes:
continue
if any(
box[0] < -allowed_edge_overflow
or box[1] < -allowed_edge_overflow
or box[2] > plot_width + allowed_edge_overflow
or box[3] > plot_height + allowed_edge_overflow
for box in candidate_boxes
):
continue
if any(
_boxes_overlap(
candidate_box,
placed_box,
padding=collision_padding,
)
for candidate_box in candidate_boxes
for placed_box in occupied
):
continue
row_index = row.get("row_index")
if isinstance(row_index, int):
selected[row_index] = (x_shift, y_shift)
occupied.extend(candidate_boxes)
break
return selected
def _bubble_axis_average(
df: pl.DataFrame | pl.LazyFrame,
x_dimension: str,
y_dimension: str,
) -> tuple[float | None, float | None]:
lf = ensure_lazyframe(df)
columns, _ = get_schema_and_column_names(lf)
if x_dimension not in columns or y_dimension not in columns:
return None, None
rows = (
lf.select(
pl.col(x_dimension).cast(pl.Float64, strict=False).mean().alias("x_avg"),
pl.col(y_dimension).cast(pl.Float64, strict=False).mean().alias("y_avg"),
)
.collect(engine="streaming")
.to_dicts()
)
if not rows:
return None, None
return _safe_float(rows[0]["x_avg"]), _safe_float(rows[0]["y_avg"])
def start_bubble_axes_from_zero(
fig: go.Figure,
df: pl.DataFrame | pl.LazyFrame,
column: str,
chartDict: dict,
countRows: int,
countCols: int,
) -> go.Figure:
"""Adjust axis ranges for bubble charts with LazyFrame support."""
namingParams = get_naming_params()
startAxesFromZero = namingParams["startAxesFromZero"]
yAxisMetric = namingParams["yAxisMetric"]
xAxisMetric = namingParams["xAxisMetric"]
totalName = namingParams["totalName"]
minXDimensionKey = namingParams["minXDimension"]
maxXDimensionKey = namingParams["maxXDimension"]
minYDimensionKey = namingParams["minYDimension"]
maxYDimensionKey = namingParams["maxYDimension"]
xDimension = chartDict[xAxisMetric]
yDimension = chartDict[yAxisMetric]
yTitleText = None
xTitleText = None
if countRows == 1 and countCols == 1:
yTitleText = yDimension
xTitleText = xDimension
if startAxesFromZero in chartDict and chartDict[startAxesFromZero]:
row_arg = countRows if getattr(fig, "_grid_ref", None) is not None else None
col_arg = countCols if getattr(fig, "_grid_ref", None) is not None else None
fig.update_xaxes(
title_text=xTitleText, rangemode="tozero", row=row_arg, col=col_arg
)
fig.update_yaxes(
title_text=yTitleText, rangemode="tozero", row=row_arg, col=col_arg
)
else:
lf = ensure_lazyframe(df)
min_max = (
lf.select(
pl.col(xDimension).min().alias("min_x"),
pl.col(xDimension).max().alias("max_x"),
pl.col(yDimension).min().alias("min_y"),
pl.col(yDimension).max().alias("max_y"),
)
.collect(engine="streaming")
.to_dicts()[0]
)
minXDimension = min_max["min_x"]
maxXDimension = min_max["max_x"]
minYDimension = min_max["min_y"]
maxYDimension = min_max["max_y"]
if minXDimension is None:
minXDimension = chartDict.get(minXDimensionKey, 0)
if maxXDimension is None:
maxXDimension = chartDict.get(maxXDimensionKey, 0)
if minYDimension is None:
minYDimension = chartDict.get(minYDimensionKey, 0)
if maxYDimension is None:
maxYDimension = chartDict.get(maxYDimensionKey, 0)
minXMultiplier, maxXMultiplier = 0.9, 1.1
minYMultiplier, maxYMultiplier = 0.9, 1.1
if minXDimension < 0:
minXMultiplier = 1.1
if minYDimension < 0:
minYMultiplier = 1.1
if maxXDimension < 0:
maxXMultiplier = 1.1
if maxYDimension < 0:
maxYMultiplier = 1.1
row_arg = countRows if getattr(fig, "_grid_ref", None) is not None else None
col_arg = countCols if getattr(fig, "_grid_ref", None) is not None else None
fig.update_xaxes(
title_text=xTitleText,
range=[minXDimension * minXMultiplier, maxXDimension * maxXMultiplier],
rangemode="normal",
row=row_arg,
col=col_arg,
)
fig.update_yaxes(
title_text=yTitleText,
range=[minYDimension * minYMultiplier, maxYDimension * maxYMultiplier],
rangemode="normal",
row=row_arg,
col=col_arg,
)
return fig
def draw_motion_chart(df, paramDict, periodOrder, chartDict):
"""Actually draw a motion chart.
Accepts both :class:`polars.DataFrame` and :class:`polars.LazyFrame` and
performs most operations lazily, collecting only when necessary.
"""
df_lazy = ensure_lazyframe(df)
namingParams = get_naming_params()
configParams = get_config_params()
font = configParams[namingParams["fontChoice"]]
fontSize = configParams[namingParams["fontSizeText"]]
annotationDict = configParams[namingParams["annotationDict"]]
monetaryName = namingParams["monetaryLocalCurrencyName"]
yAxisMetric = namingParams["yAxisMetric"]
xAxisMetric = namingParams["xAxisMetric"]
xAxisDimension = namingParams["xAxisDimension"]
yAxisDimension = namingParams["yAxisDimension"]
bubbleSize = namingParams["bubbleSize"]
totalName = namingParams["totalName"]
periodName = namingParams["periodName"]
dateName = namingParams["dateName"]
plName = namingParams["plName"]
fcName = namingParams["fcName"]
acName = namingParams["acName"]
selectedPeriods = namingParams["selectedPeriods"]
plotValuesAsChoice = namingParams["plotValuesAsChoice"]
absolute = namingParams["absolute"]
chosenChart = namingParams["chosenChart"]
labelColor = namingParams["labelColor"]
blackLabelChoice = namingParams["blackLabelChoice"]
whiteLabelChoice = namingParams["whiteLabelChoice"]
greyLabelChoice = namingParams["greyLabelChoice"]
showBubbleLabel = namingParams["showBubbleLabel"]
showBoth = namingParams["showBoth"]
showLabelsOnly = namingParams["showLabelsOnly"]
showNothing = namingParams["showNothing"]
colorDict = get_color_dictionary(chartDict)
colorArray = get_color_array(colorDict, chartDict)
chosenChart = chartDict[chosenChart]
periodOrder = chartDict[selectedPeriods]
xDimension = chartDict[xAxisMetric]
yDimension = chartDict[yAxisMetric]
bubbleSizeDimension = chartDict[bubbleSize]
chosenDimension = chartDict[xAxisDimension]
bubbleColorDimension = chartDict[yAxisDimension]
if chartDict[plotValuesAsChoice] == absolute:
chartFormat = ",.3s"
else:
chartFormat = ",.0f"
for element in [bubbleSizeDimension, xDimension, yDimension]:
df_lazy = df_lazy.with_columns(
pl.when(pl.col(element) <= 0)
.then(0)
.otherwise(pl.col(element))
.alias(element)
)
other_name = namingParams["otherName"]
undefined_name = namingParams["undefinedName"]
highlighted_dimension = namingParams["highlightedDimension"]
highlight_color = get_hightlight_color(chartDict, colorDict)
unique_colors = unique_values_lazy(bubbleColorDimension, df_lazy)
for i, val in enumerate(unique_colors):
if val in [other_name, undefined_name] or (
isinstance(val, str) and other_name.lower() in val.lower()
):
if i < len(colorArray):
colorArray[i] = colorDict["veryLightGreyColor"]
elif (
highlighted_dimension in chartDict
and val in chartDict[highlighted_dimension]
):
if i < len(colorArray):
colorArray[i] = highlight_color
min_max = (
df_lazy.select(
pl.col(xDimension).min().alias("min_x"),
pl.col(xDimension).max().alias("max_x"),
pl.col(yDimension).min().alias("min_y"),
pl.col(yDimension).max().alias("max_y"),
)
.collect(engine="streaming")
.to_dicts()[0]
)
min_x, max_x, min_y, max_y = (
min_max["min_x"],
min_max["max_x"],
min_max["min_y"],
min_max["max_y"],
)
rangeX = [min_x * 0.8, max_x * 1.2]
rangeY = [min_y * 0.8, max_y * 1.2]
unique_dates = unique_values_lazy(dateName, df_lazy)
color_map = {
val: colorArray[i % len(colorArray)] for i, val in enumerate(unique_colors)
}
max_size = (
df_lazy.select(pl.col(bubbleSizeDimension).max())
.collect(engine="streaming")
.item()
)
mode = "markers+text" if chartDict[showBubbleLabel] != showNothing else "markers"
frames = []
for date in unique_dates:
df_frame = (
df_lazy.filter(pl.col(dateName) == date)
.select(
xDimension,
yDimension,
bubbleSizeDimension,
chosenDimension,
bubbleColorDimension,
)
.collect(engine="streaming")
)
frames.append(
go.Frame(
data=[
go.Scatter(
x=df_frame[xDimension].to_list(),
y=df_frame[yDimension].to_list(),
mode=mode,
text=(
df_frame[chosenDimension].to_list()
if chartDict[showBubbleLabel] != showNothing
else None
),
marker=dict(
size=df_frame[bubbleSizeDimension].to_list(),
color=[
color_map[val]
for val in df_frame[bubbleColorDimension].to_list()
],
sizemode="area",
sizeref=2.0 * max_size / (70.0**2),
sizemin=4,
opacity=0.8,
),
hovertext=df_frame[chosenDimension].to_list(),
showlegend=False,
)
],
name=str(date),
)
)
fig = go.Figure(data=frames[0].data if frames else [], frames=frames)
fig.update_xaxes(range=rangeX)
fig.update_yaxes(range=rangeY)
fig.for_each_trace(
lambda t: t.update(
textfont_color=colorDict["blackColor"], textposition="top center"
)
)
df_full = df_lazy.collect()
rangeArray = range(df_full.height)
fig = add_split_lines(fig)
if chartDict[showBubbleLabel] != showNothing:
for element in rangeArray:
if bubbleColorDimension == periodName:
if df_full[bubbleColorDimension][element] in [plName, fcName]:
color = colorDict["blackColor"]
elif df_full[bubbleColorDimension][element] == periodOrder[0]:
color = colorDict["blackColor"]
else:
color = colorDict["whiteColor"]
text = df_full[chosenDimension][element]
bubbleValue = millify(df_full[bubbleSizeDimension][element], 1)
bubbleLabel = str(df_full[chosenDimension][element])
if chartDict[showBubbleLabel] == showLabelsOnly:
bubbleText = bubbleLabel
elif chartDict[showBubbleLabel] == showValuesOnly:
bubbleText = bubbleValue
else:
bubbleText = bubbleLabel + "<br>" + bubbleValue
return fig
def add_labels_and_values_to_bubbles(
fig, df, chartDict, colorDict, sizeRef, countRows, countCols
):
namingParams = get_naming_params()
configParams = get_config_params()
font = configParams[namingParams["fontChoice"]]
fontSize = configParams[namingParams["fontSizeText"]]
bubbleSize = namingParams["bubbleSize"]
adjustBubbleLabels = namingParams["adjustBubbleLabels"]
showBoth = namingParams["showBoth"]
showLabelsOnly = namingParams["showLabelsOnly"]
showValuesOnly = namingParams["showValuesOnly"]
showNothing = namingParams["showNothing"]
showBubbleLabel = namingParams["showBubbleLabel"]
xAxisDimension = namingParams["xAxisDimension"]
yAxisMetric = namingParams["yAxisMetric"]
xAxisMetric = namingParams["xAxisMetric"]
scaledSize = namingParams["scaledSizeName"]
bubbleSizeDimension = chartDict[bubbleSize]
chosenDimension = chartDict[xAxisDimension]
xDimension = chartDict[xAxisMetric]
yDimension = chartDict[yAxisMetric]
label_mode = chartDict[showBubbleLabel]
if label_mode == showNothing:
return fig
show_label = label_mode in [showBoth, showLabelsOnly]
show_value = label_mode in [showBoth, showValuesOnly]
lf = (
ensure_lazyframe(df)
.sort(xDimension, descending=True)
.with_row_index("_idx")
.with_columns(
((pl.col(bubbleSizeDimension) / sizeRef / 100) + 100).alias(scaledSize)
)
.with_columns((pl.col(scaledSize) / 100).alias(scaledSize))
.with_columns(((pl.len() < 6).cast(pl.Int64)).alias("__round"))
.select(
pl.col("_idx"),
pl.col(xDimension).alias("__x"),
pl.col(yDimension).alias("__y"),
pl.col(scaledSize).alias("__scale"),
pl.col(bubbleSizeDimension).alias("__size"),
pl.struct([pl.col(bubbleSizeDimension), pl.col("__round")])
.map_elements(
lambda s: millify(s[bubbleSizeDimension], s["__round"]),
return_dtype=pl.Utf8,
)
.alias("__value"),
pl.col(chosenDimension).cast(pl.Utf8).alias("__label"),
)
)
ann_df = lf.collect(engine="streaming")
rows: list[dict[str, object]] = []
x_values: list[float] = []
y_values: list[float] = []
for row in ann_df.iter_rows(named=True):
x_float = _safe_float(row["__x"])
y_float = _safe_float(row["__y"])
size_float = _safe_float(row["__size"])
if x_float is not None:
x_values.append(x_float)
if y_float is not None:
y_values.append(y_float)
rows.append(
{
"row_index": int(row["_idx"]),
"value": row["__value"],
"label": row["__label"],
"x": row["__x"],
"y": row["__y"],
"scale": row["__scale"],
"x_float": x_float,
"y_float": y_float,
"size_float": size_float,
}
)
if chartDict[adjustBubbleLabels]:
selected_positions = _select_bubble_label_positions(
rows,
x_range=_bubble_axis_range(x_values),
y_range=_bubble_axis_range(y_values),
size_ref=sizeRef,
font_size=fontSize,
show_label=show_label,
show_value=show_value,
)
else:
selected_positions = {
row["row_index"]: (0, 0)
for row in rows
if isinstance(row.get("row_index"), int)
}
for row in rows:
row_index = row["row_index"]
if not isinstance(row_index, int) or row_index not in selected_positions:
continue
if show_value:
fig.add_annotation(
text=row["value"],
showarrow=False,
align="center",
yshift=0,
y=row["y"],
yref="y",
x=row["x"],
ax=0,
xref="x",
font=dict(size=fontSize, color=colorDict["blackColor"], family=font),
row=countRows,
col=countCols,
)
if show_label and chartDict[adjustBubbleLabels]:
x_shift, y_shift = selected_positions[row_index]
fig.add_annotation(
text=row["label"],
showarrow=False,
align="center",
yshift=y_shift,
xshift=x_shift,
y=row["y"],
yref="y",
x=row["x"],
ax=0,
xref="x",
font=dict(size=fontSize, color=colorDict["blackColor"], family=font),
row=countRows,
col=countCols,
)
return fig
def add_legend_title_to_bubble_chart(
fig, plotLegend, chartDict, colorDict, countRows, countCols
):
namingParams = get_naming_params()
yAxisDimension = namingParams["yAxisDimension"]
bubbleColorDimension = chartDict[yAxisDimension]
if plotLegend:
fig.add_annotation(
text=bubbleColorDimension,
showarrow=False,
x=1.1,
y=1.03,
xref="paper",
yref="paper",
font=dict(
color=colorDict["blackColor"],
),
)
return fig
def get_colors_for_bubble(
fig,
df,
column,
chartDict,
colorDimensionArray,
countCols,
countRows,
aggregateOtherItemsName,
colorArray,
):
namingParams = get_naming_params()
yAxisDimension = namingParams["yAxisDimension"]
xAxisDimension = namingParams["xAxisDimension"]
nothingFilteredName = namingParams["nothingFilteredName"]
smallMultiplesColumn = namingParams["smallMultiplesColumn"]
totalName = namingParams["totalName"]
notMetConditionValue = namingParams["notMetConditionValue"]
smallMultiplesColumn = chartDict[smallMultiplesColumn]
chosenDimension = chartDict[xAxisDimension]
bubbleColorDimension = chartDict[yAxisDimension]
lf = ensure_lazyframe(df)
plotLegend = True
if column == totalName:
if bubbleColorDimension in [None, nothingFilteredName, notMetConditionValue]:
colorDimensionArray = [0]
lf = lf.with_columns(pl.lit(0).alias(bubbleColorDimension))
plotLegend = False
elif chosenDimension == bubbleColorDimension:
colorDimensionArray = unique_values_lazy(bubbleColorDimension, lf)
plotLegend = False
else:
colorDimensionArray = unique_values_lazy(bubbleColorDimension, lf)
else:
if bubbleColorDimension in [None, nothingFilteredName]:
lf = lf.with_columns(pl.lit(0).alias(bubbleColorDimension))
plotLegend = False
elif (
chosenDimension == bubbleColorDimension
and chosenDimension != smallMultiplesColumn
):
plotLegend = False
if countCols == 1 and countRows == 1:
plotLegend = True
elif bubbleColorDimension == smallMultiplesColumn:
chartDict[yAxisDimension] = nothingFilteredName
colorDimensionArray = [0]
lf = lf.with_columns(pl.lit(0).alias(chartDict[yAxisDimension]))
plotLegend = False
elif bubbleColorDimension != smallMultiplesColumn:
plotLegend = False
if countCols == 1 and countRows == 1:
plotLegend = True
else:
pass
if bubbleColorDimension in [None, nothingFilteredName]:
lf = lf.with_columns(
pl.when(pl.col(chosenDimension) == aggregateOtherItemsName)
.then(1)
.otherwise(pl.col(bubbleColorDimension))
.alias(bubbleColorDimension)
)
plotLegend = False
if len(colorDimensionArray) == 1:
colorDimensionArray.append(1)
colorDict = get_color_dictionary(chartDict)
colorArray.insert(1, colorDict["veryLightGreyColor"])
df_out = ensure_polars_df(lf)
return df_out, colorDimensionArray, plotLegend, colorArray, chartDict
def color_other_bubbles_in_grey(
df: pl.DataFrame | pl.LazyFrame,
colorArray: list[str],
bubbleColorDimension: str,
colorDict: dict,
chartDict: dict,
) -> list[str]:
"""Update ``colorArray`` for special bubble colors lazily."""
namingParams = get_naming_params()
otherName = namingParams["otherName"]
undefinedName = namingParams["undefinedName"]
highlightedDimension = namingParams["highlightedDimension"]
highlightColor = get_hightlight_color(chartDict, colorDict)
columns, _ = get_schema_and_column_names(df)
if bubbleColorDimension in columns:
lf = ensure_lazyframe(df)
bubbleColorArray = unique_values_lazy(bubbleColorDimension, lf)
count = 0
for element in bubbleColorArray:
if element in [otherName, undefinedName] or (
isinstance(element, str) and otherName.lower() in element.lower()
):
colorArray[count] = colorDict["veryLightGreyColor"]
elif (
highlightedDimension in chartDict
and element in chartDict[highlightedDimension]
):
colorArray[count] = highlightColor
count += 1
return colorArray
def add_total_bubble_to_bubble_chart(
fig, df, chartDict, dfTotals, sizeRef, countRows, countCols
):
namingParams = get_naming_params()
plotTotalBubble = namingParams["plotTotalBubble"]
yAxisMetric = namingParams["yAxisMetric"]
xAxisMetric = namingParams["xAxisMetric"]
bubbleSize = namingParams["bubbleSize"]
totalName = namingParams["totalName"]
xDimension = chartDict[xAxisMetric]
yDimension = chartDict[yAxisMetric]
bubbleSizeDimension = chartDict[bubbleSize]
df_lazy_totals = ensure_lazyframe(dfTotals)
totals_columns, _ = get_schema_and_column_names(df_lazy_totals)
has_totals = bubbleSizeDimension in totals_columns and (
df_lazy_totals.select([bubbleSizeDimension])
.head(1)
.collect(engine="streaming")
.height
> 0
)
if _can_plot_total_bubble(chartDict) and has_totals:
xTotal, yTotal = _bubble_axis_average(df, xDimension, yDimension)
if xTotal is None and xDimension in totals_columns:
xTotal = get_polars_value_at_index(df_lazy_totals, xDimension, 0)
if yTotal is None and yDimension in totals_columns:
yTotal = get_polars_value_at_index(df_lazy_totals, yDimension, 0)
sizeTotal = get_polars_value_at_index(df_lazy_totals, bubbleSizeDimension, 0)
if xTotal is None or yTotal is None or _safe_float(sizeTotal) is None:
return fig
fig.add_trace(
go.Scatter(
x=[xTotal],
y=[yTotal],
mode="markers",
marker=dict(
size=[sizeTotal],
color="rgba(0,0,0,0)",
line=dict(
color="Red",
width=1,
),
sizemode="area",
sizeref=sizeRef,
sizemin=4,
opacity=1,
),
showlegend=False, # do not show this in the legend
),
row=countRows,
col=countCols,
)
totalBubbleValue = millify(sizeTotal, 1)
fig.add_annotation(
text=totalName + ": " + str(totalBubbleValue),
showarrow=False,
align="center",
yshift=0,
y=yTotal,
yref="y",
x=xTotal,
ax=0,
xref="x",
font=dict(
size=10,
color="Red",
),
row=countRows,
col=countCols,
)
return fig
def add_bubbles_to_bubble_chart(
fig,
df,
plotLegend,
chartDict,
colorArray,
colorDimensionArray,
sizeRef,
countRows,
countCols,
):
namingParams = get_naming_params()
yAxisDimension = namingParams["yAxisDimension"]
yAxisMetric = namingParams["yAxisMetric"]
xAxisMetric = namingParams["xAxisMetric"]
bubbleSize = namingParams["bubbleSize"]
xAxisDimension = namingParams["xAxisDimension"]
adjustBubbleLabels = namingParams["adjustBubbleLabels"]
bubbleSizeDimension = chartDict[bubbleSize]
xDimension = chartDict[xAxisMetric]
yDimension = chartDict[yAxisMetric]
chosenDimension = chartDict[xAxisDimension]
bubbleColorDimension = chartDict[yAxisDimension]
countColors = 0
for item in colorDimensionArray:
if (
xDimension is not None
and not isinstance(xDimension, bool)
and yDimension is not None
and not isinstance(yDimension, bool)
):
# Polars-style row filtering; ensure eager DataFrame for plotting
dfFiltered = ensure_lazyframe(df).filter(
pl.col(bubbleColorDimension) == item
)
dfFiltered = ensure_polars_df(dfFiltered)
bubbleLabel = dfFiltered[chosenDimension]
mode = "markers+text"
if chartDict[adjustBubbleLabels]:
bubbleLabel = False
mode = "markers"
fig.add_trace(
go.Scatter(
x=dfFiltered[xDimension],
y=dfFiltered[yDimension],
mode=mode,
textposition="top center",
marker=dict(
size=dfFiltered[bubbleSizeDimension],
color=colorArray[countColors],
sizemode="area",
sizeref=sizeRef,
sizemin=4,
opacity=0.1,
),
text=bubbleLabel, # this will appear when you hover over a point
showlegend=False, # do not show this in the legend
),
row=countRows,
col=countCols,
)
dfCopy = duplicate_dataframe(dfFiltered)
dfCopy_df = ensure_polars_df(dfCopy)
if dfCopy_df.height == 0 and plotLegend:
empty_df = pl.DataFrame(
{
xDimension: [None],
yDimension: [None],
bubbleSizeDimension: [1],
}
)
dfCopy = pl.concat([ensure_lazyframe(dfCopy), empty_df.lazy()])
dfCopy_df = empty_df
else:
dfCopy = ensure_lazyframe(dfCopy)
fig.add_trace(
go.Scatter(
x=dfCopy_df[xDimension],
y=dfCopy_df[yDimension],
mode="markers",
marker=dict(
size=dfCopy_df[bubbleSizeDimension],
color="rgba(0,0,0,0)", # transparent color
line=dict(color=colorArray[countColors], width=2.5),
sizemode="area",
sizeref=sizeRef,
sizemin=4,
),
name=item,
showlegend=plotLegend,
),
row=countRows,
col=countCols,
)
countColors = countColors + 1
return fig
def get_size_def_for_bubbles(df, sizeRef, chartDict, column, count):
namingParams = get_naming_params()
bubbleSize = namingParams["bubbleSize"]
totalName = namingParams["totalName"]
bubbleSizeDimension = chartDict[bubbleSize]
lf = ensure_lazyframe(df)
try:
max_size = (
lf.select(pl.col(bubbleSizeDimension).max())
.collect(engine="streaming")
.item()
)
except Exception as e:
logging.exception(e)
ui.error("Something went wrong while computing bubble max size.")
max_size = None
if max_size is None:
return False
if column == totalName:
return 2.0 * max_size / (40.0**2)
return 2 * max_size / (40.0**2)
def draw_bubble_chart(
fig,
df,
colorDimensionArray,
plotLegend,
chartDict,
colorDict,
colorArray,
dfTotals,
column,
count,
countRows,
countCols,
sizeRef,
):
"""
actually draws bubble chart
"""
sizeRef = get_size_def_for_bubbles(df, sizeRef, chartDict, column, count)
fig = add_bubbles_to_bubble_chart(
fig,
df,
plotLegend,
chartDict,
colorArray,
colorDimensionArray,
sizeRef,
countRows,
countCols,
)
fig = add_total_bubble_to_bubble_chart(
fig, df, chartDict, dfTotals, sizeRef, countRows, countCols
)
fig = add_labels_and_values_to_bubbles(
fig, df, chartDict, colorDict, sizeRef, countRows, countCols
)
fig = add_legend_title_to_bubble_chart(
fig, plotLegend, chartDict, colorDict, countRows, countCols
)
fig = start_bubble_axes_from_zero(fig, df, column, chartDict, countRows, countCols)
fig = add_split_lines(fig)
return fig, sizeRef
SHA-256: da11bef92f80038c86a80aa009835d6bb27d20323cc9c813b023546078aaab16