Skip to content

Commit baa338c

Browse files
max-modelsmax-models
committed
Added more tikz primitives
1 parent 38fc954 commit baa338c

4 files changed

Lines changed: 588 additions & 13 deletions

File tree

‎src/maxplotlib/canvas/canvas.py‎

Lines changed: 176 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -18,7 +18,13 @@
1818
from maxplotlib.backends.plotext import PlotextFigure, create_plotext_figure
1919
from maxplotlib.colors.colors import Color
2020
from maxplotlib.linestyle.linestyle import Linestyle
21-
from maxplotlib.subfigure.line_plot import LinePlot
21+
from maxplotlib.subfigure.line_plot import (
22+
LinePlot,
23+
_TIKZ_SUPPORTED_PLOT_TYPES,
24+
_tikz_error_bounds,
25+
_tikz_step_coordinates,
26+
_tikz_style_kwargs,
27+
)
2228
from maxplotlib.utils.options import Backends
2329

2430

@@ -2269,7 +2275,12 @@ def plot_tikzfigure(
22692275

22702276
# Add each plot line to the subfigure
22712277
for line_data in line_plot.line_data:
2272-
if line_data.get("plot_type") == "plot":
2278+
plot_type = line_data.get("plot_type")
2279+
if plot_type not in _TIKZ_SUPPORTED_PLOT_TYPES:
2280+
raise NotImplementedError(
2281+
f"{plot_type} is not supported by the tikzfigure backend"
2282+
)
2283+
if plot_type == "plot":
22732284
# Extract and transform x, y data
22742285
x = (line_data["x"] + line_plot._xshift) * line_plot._xscale
22752286
y = (line_data["y"] + line_plot._yshift) * line_plot._yscale
@@ -2280,11 +2291,170 @@ def plot_tikzfigure(
22802291
ax.add_plot(
22812292
x=x,
22822293
y=y,
2283-
# label=kwargs.get("label", ""),
2284-
color=kwargs.get("color", "black"),
2285-
line_width=kwargs.get("linewidth", 1.0),
2294+
**_tikz_style_kwargs(kwargs),
2295+
)
2296+
elif plot_type == "scatter":
2297+
x = (line_data["x"] + line_plot._xshift) * line_plot._xscale
2298+
y = (line_data["y"] + line_plot._yshift) * line_plot._yscale
2299+
kwargs = _tikz_style_kwargs(line_data.get("kwargs", {}))
2300+
kwargs.setdefault("mark", "*")
2301+
kwargs["line_width"] = 0
2302+
ax.add_plot(x=x, y=y, **kwargs)
2303+
elif plot_type in {"bar", "barh"}:
2304+
source_kwargs = line_data.get("kwargs", {})
2305+
kwargs = _tikz_style_kwargs(source_kwargs)
2306+
kwargs["fill"] = source_kwargs.get("color", "blue")
2307+
kwargs["fill_opacity"] = source_kwargs.get("alpha", 1.0)
2308+
kwargs["line_width"] = source_kwargs.get("linewidth", 0)
2309+
if plot_type == "bar":
2310+
width = source_kwargs.get("width", 0.8)
2311+
for x, height in zip(line_data["x"], line_data["height"]):
2312+
ax.add_plot(
2313+
x=[x - width / 2, x + width / 2, x + width / 2, x - width / 2],
2314+
y=[0, 0, height, height],
2315+
cycle=True,
2316+
**kwargs,
2317+
)
2318+
else:
2319+
height = source_kwargs.get("height", 0.8)
2320+
for y, width in zip(line_data["y"], line_data["width"]):
2321+
ax.add_plot(
2322+
x=[0, width, width, 0],
2323+
y=[y - height / 2, y - height / 2, y + height / 2, y + height / 2],
2324+
cycle=True,
2325+
**kwargs,
2326+
)
2327+
elif plot_type == "fill_between":
2328+
x = line_data["x"]
2329+
y1 = np.asarray(line_data["y1"])
2330+
y2 = np.broadcast_to(line_data["y2"], y1.shape)
2331+
source_kwargs = line_data.get("kwargs", {})
2332+
kwargs = _tikz_style_kwargs(source_kwargs)
2333+
kwargs["fill"] = source_kwargs.get("color", "blue")
2334+
kwargs["fill_opacity"] = source_kwargs.get("alpha", 0.25)
2335+
ax.add_plot(
2336+
x=list(x) + list(x[::-1]),
2337+
y=list(y1) + list(y2[::-1]),
2338+
cycle=True,
2339+
**kwargs,
2340+
)
2341+
elif plot_type == "errorbar":
2342+
x = line_data["x"]
2343+
y = line_data["y"]
2344+
kwargs = _tikz_style_kwargs(line_data.get("kwargs", {}))
2345+
ax.add_plot(x=x, y=y, **kwargs)
2346+
y_bounds = _tikz_error_bounds(line_data["yerr"], y)
2347+
if y_bounds is not None:
2348+
lower, upper = y_bounds
2349+
for xi, low, high in zip(x, y - lower, y + upper):
2350+
ax.add_plot(x=[xi, xi], y=[low, high], **kwargs)
2351+
x_bounds = _tikz_error_bounds(line_data["xerr"], x)
2352+
if x_bounds is not None:
2353+
lower, upper = x_bounds
2354+
for yi, low, high in zip(y, x - lower, x + upper):
2355+
ax.add_plot(x=[low, high], y=[yi, yi], **kwargs)
2356+
elif plot_type in {"step", "stairs"}:
2357+
source_kwargs = line_data.get("kwargs", {})
2358+
if plot_type == "step":
2359+
x = line_data["x"]
2360+
y = line_data["y"]
2361+
where = source_kwargs.get("where", "pre")
2362+
else:
2363+
values = line_data["values"]
2364+
edges = line_data["edges"]
2365+
if edges is None:
2366+
edges = np.arange(len(values) + 1)
2367+
x = edges
2368+
y = np.r_[values, values[-1]]
2369+
where = "post"
2370+
x, y = _tikz_step_coordinates(x, y, where=where)
2371+
ax.add_plot(
2372+
x=x,
2373+
y=y,
2374+
**_tikz_style_kwargs(source_kwargs),
22862375
)
2287-
elif line_data.get("plot_type") == "gantt":
2376+
elif plot_type == "stem":
2377+
x = line_data["x"]
2378+
y = line_data["y"]
2379+
source_kwargs = line_data.get("kwargs", {})
2380+
style = _tikz_style_kwargs(source_kwargs)
2381+
marker_style = dict(style)
2382+
marker_style.update(mark=source_kwargs.get("marker", "*"), line_width=0)
2383+
ax.add_plot(x=x, y=y, **marker_style)
2384+
for xi, yi in zip(x, y):
2385+
ax.add_plot(x=[xi, xi], y=[0, yi], **style)
2386+
elif plot_type in {"hlines", "vlines"}:
2387+
style = _tikz_style_kwargs(line_data.get("kwargs", {}))
2388+
if plot_type == "hlines":
2389+
for yi, left, right in zip(
2390+
np.atleast_1d(line_data["y"]),
2391+
np.atleast_1d(line_data["xmin"]),
2392+
np.atleast_1d(line_data["xmax"]),
2393+
):
2394+
ax.add_plot(x=[left, right], y=[yi, yi], **style)
2395+
else:
2396+
for xi, bottom, top in zip(
2397+
np.atleast_1d(line_data["x"]),
2398+
np.atleast_1d(line_data["ymin"]),
2399+
np.atleast_1d(line_data["ymax"]),
2400+
):
2401+
ax.add_plot(x=[xi, xi], y=[bottom, top], **style)
2402+
elif plot_type in {"axvspan", "axhspan"}:
2403+
source_kwargs = line_data.get("kwargs", {})
2404+
style = _tikz_style_kwargs(source_kwargs)
2405+
style["fill"] = source_kwargs.get("color", "blue")
2406+
style["fill_opacity"] = source_kwargs.get("alpha", 0.2)
2407+
if plot_type == "axvspan":
2408+
xmin, xmax = line_data["xmin"], line_data["xmax"]
2409+
ymin, ymax = line_plot._ymin or 0, line_plot._ymax or 1
2410+
x = [xmin, xmax, xmax, xmin]
2411+
y = [ymin, ymin, ymax, ymax]
2412+
else:
2413+
ymin, ymax = line_data["ymin"], line_data["ymax"]
2414+
xmin, xmax = line_plot._xmin or 0, line_plot._xmax or 1
2415+
x = [xmin, xmax, xmax, xmin]
2416+
y = [ymin, ymin, ymax, ymax]
2417+
ax.add_plot(x=x, y=y, cycle=True, **style)
2418+
elif plot_type == "fill":
2419+
if len(line_data["args"]) < 2:
2420+
raise ValueError("tikzfigure fill requires x and y coordinates")
2421+
x, y = line_data["args"][:2]
2422+
source_kwargs = line_data.get("kwargs", {})
2423+
style = _tikz_style_kwargs(source_kwargs)
2424+
style["fill"] = source_kwargs.get("color", "blue")
2425+
style["fill_opacity"] = source_kwargs.get("alpha", 0.25)
2426+
ax.add_plot(x=x, y=y, cycle=True, **style)
2427+
elif plot_type == "flame_chart":
2428+
labels = line_data["labels"]
2429+
parents = line_data["parents"]
2430+
values = line_data["values"] * line_plot._xscale
2431+
start_times = line_data["start_times"]
2432+
depths = np.zeros(len(labels), dtype=int)
2433+
if start_times is None:
2434+
start_times = np.zeros(len(labels))
2435+
else:
2436+
start_times = (
2437+
start_times + line_plot._xshift
2438+
) * line_plot._xscale
2439+
for index, parent in enumerate(parents):
2440+
if parent is not None:
2441+
parent_index = (
2442+
parent
2443+
if isinstance(parent, int)
2444+
else labels.index(parent)
2445+
)
2446+
depths[index] = depths[parent_index] + 1
2447+
colors = ["red", "blue", "green", "orange", "purple", "cyan"]
2448+
for index, (start, value) in enumerate(zip(start_times, values)):
2449+
y = depths[index]
2450+
ax.add_plot(
2451+
x=[start, start + value, start + value, start],
2452+
y=[y - 0.4, y - 0.4, y + 0.4, y + 0.4],
2453+
cycle=True,
2454+
fill=colors[y % len(colors)],
2455+
line_width=0,
2456+
)
2457+
elif plot_type == "gantt":
22882458
tasks = line_data["tasks"]
22892459
start_times = (
22902460
line_data["start_times"] + line_plot._xshift

0 commit comments

Comments
 (0)