Skip to content

Commit f1a772e

Browse files
PhlyaPhlya
committed
split the main function
1 parent a839d34 commit f1a772e

1 file changed

Lines changed: 109 additions & 79 deletions

File tree

‎adjustText/__init__.py‎

Lines changed: 109 additions & 79 deletions
Original file line numberDiff line numberDiff line change
@@ -186,10 +186,7 @@ def apply_shifts(coords, shifts_x, shifts_y):
186186

187187

188188
def force_into_bbox(coords, bbox):
189-
xmin = bbox.xmin
190-
xmax = bbox.xmax
191-
ymin = bbox.ymin
192-
ymax = bbox.ymax
189+
xmin, xmax, ymin, ymax = bbox
193190
dx, dy = np.zeros((coords.shape[0])), np.zeros((coords.shape[0]))
194191
if np.any((coords[:, 0] < xmin) & (coords[:, 1] > xmax)):
195192
logging.warn("Some labels are too long, can't fit inside the X axis")
@@ -228,24 +225,113 @@ def explode(coords, static_coords, r=None):
228225
return xshifts, yshifts
229226

230227

228+
def iterate(
229+
coords,
230+
orig_coords,
231+
static_coords=None,
232+
force_text: tuple[float, float] = (0.1, 0.2),
233+
force_static: tuple[float, float] = (0.05, 0.1),
234+
force_pull: tuple[float, float] = (0.05, 0.1),
235+
expand: tuple[float, float] = (1.05, 1.1),
236+
bbox_to_contain=False,
237+
only_move={"text": "xy", "static": "xy", "explode": "xy", "pull": "xy"},
238+
):
239+
240+
text_shifts_x, text_shifts_y = get_shifts_texts(
241+
expand_coords(coords, expand[0], expand[1])
242+
)
243+
if static_coords.shape[0] > 0:
244+
static_shifts_x, static_shifts_y = get_shifts_extra(
245+
expand_coords(coords, expand[0], expand[1]), static_coords
246+
)
247+
else:
248+
static_shifts_x, static_shifts_y = np.zeros((1)), np.zeros((1))
249+
error_x = np.abs(text_shifts_x) + np.abs(static_shifts_x)
250+
error_y = np.abs(text_shifts_y) + np.abs(static_shifts_y)
251+
error = np.sum(np.append(error_x, error_y))
252+
253+
pull_x, pull_y = pull_back(coords, orig_coords)
254+
255+
text_shifts_x *= force_text[0]
256+
text_shifts_y *= force_text[1]
257+
static_shifts_x *= force_static[0]
258+
static_shifts_y *= force_static[1]
259+
# Pull is in the opposite direction, so need to negate it
260+
pull_x *= -force_pull[0]
261+
pull_y *= -force_pull[1]
262+
# pull_x[error_x != 0] = 0
263+
# pull_y[error_y != 0] = 0
264+
265+
if only_move:
266+
if "x" not in only_move["text"]:
267+
text_shifts_x = np.zeros_like(text_shifts_x)
268+
elif "x+" in only_move["text"]:
269+
text_shifts_x[text_shifts_x > 0] = 0
270+
elif "x-" in only_move["text"]:
271+
text_shifts_x[text_shifts_x < 0] = 0
272+
273+
if "y" not in only_move["text"]:
274+
text_shifts_y = np.zeros_like(text_shifts_y)
275+
elif "y+" in only_move["text"]:
276+
text_shifts_y[text_shifts_y > 0] = 0
277+
elif "y-" in only_move["text"]:
278+
text_shifts_y[text_shifts_y < 0] = 0
279+
280+
if "x" not in only_move["static"]:
281+
static_shifts_x = np.zeros_like(static_shifts_x)
282+
elif "x+" in only_move["static"]:
283+
static_shifts_x[static_shifts_x > 0] = 0
284+
elif "x-" in only_move["static"]:
285+
static_shifts_x[static_shifts_x < 0] = 0
286+
287+
if "y" not in only_move["static"]:
288+
static_shifts_y = np.zeros_like(static_shifts_y)
289+
elif "y+" in only_move["static"]:
290+
static_shifts_y[static_shifts_y > 0] = 0
291+
elif "y-" in only_move["static"]:
292+
static_shifts_y[static_shifts_y < 0] = 0
293+
294+
if "x" not in only_move["pull"]:
295+
pull_x = np.zeros_like(pull_x)
296+
elif "x+" in only_move["pull"]:
297+
pull_x[pull_x > 0] = 0
298+
elif "x-" in only_move["pull"]:
299+
pull_x[pull_x < 0] = 0
300+
301+
if "y" not in only_move["pull"]:
302+
pull_y = np.zeros_like(pull_y)
303+
elif "y+" in only_move["pull"]:
304+
pull_y[pull_y > 0] = 0
305+
elif "y-" in only_move["pull"]:
306+
pull_y[pull_y < 0] = 0
307+
308+
shifts_x = text_shifts_x + static_shifts_x + pull_x
309+
shifts_y = text_shifts_y + static_shifts_y + pull_y
310+
311+
coords = apply_shifts(coords, shifts_x, shifts_y)
312+
if bbox_to_contain:
313+
coords = force_into_bbox(coords, bbox_to_contain)
314+
return coords, error
315+
316+
231317
def adjust_text(
232318
texts,
233319
x=None,
234320
y=None,
235321
objects=None,
236322
avoid_self=True,
237323
force_text: tuple[float, float] = (0.1, 0.2),
238-
force_static: tuple[float, float] = (0.05, 0.1),
239-
force_pull: tuple[float, float] = (0.01, 0.001),
240-
force_explode: tuple[float, float] = (0.01, 0.01),
324+
force_static: tuple[float, float] = (0.1, 0.2),
325+
force_pull: tuple[float, float] = (0.01, 0.01),
326+
force_explode: tuple[float, float] = (0.01, 0.02),
241327
expand: tuple[float, float] = (1.05, 1.1),
242328
explode_radius="auto",
243329
ensure_inside_axes=True,
244330
expand_axes=False,
245331
only_move={"text": "xy", "static": "xy", "explode": "xy", "pull": "xy"},
246332
ax=None,
247333
min_arrow_len=5,
248-
time_lim=0.1,
334+
time_lim=0.5,
249335
*args,
250336
**kwargs,
251337
):
@@ -393,85 +479,29 @@ def adjust_text(
393479

394480
# expands = list(zip(np.linspace(expand_start[0], expand_end[0], expand_steps),
395481
# np.linspace(expand_start[1], expand_end[1], expand_steps)))
396-
ax_bbox = ax.patch.get_extents()
397482

398483
if expand_axes:
399484
expand_axes_to_fit(coords, ax, transform)
485+
if ensure_inside_axes:
486+
ax_bbox = ax.patch.get_extents()
487+
ax_bbox = ax_bbox.xmin, ax_bbox.xmax, ax_bbox.ymin, ax_bbox.ymax
488+
else:
489+
ax_bbox = False
400490

401491
# i = 0
402492
while error > 0:
403493
# expand = expands[min(i, expand_steps-1)]
404-
text_shifts_x, text_shifts_y = get_shifts_texts(
405-
expand_coords(coords, expand[0], expand[1])
494+
coords, error = iterate(
495+
coords,
496+
orig_xy_disp_coord,
497+
static_coords,
498+
force_text=force_text,
499+
force_static=force_static,
500+
force_pull=force_pull,
501+
expand=expand,
502+
bbox_to_contain=ax_bbox,
503+
only_move=only_move,
406504
)
407-
if static_coords.shape[0] > 0:
408-
static_shifts_x, static_shifts_y = get_shifts_extra(
409-
expand_coords(coords, expand[0], expand[1]), static_coords
410-
)
411-
else:
412-
static_shifts_x, static_shifts_y = np.zeros((1)), np.zeros((1))
413-
error_x = np.abs(text_shifts_x) + np.abs(static_shifts_x)
414-
error_y = np.abs(text_shifts_y) + np.abs(static_shifts_y)
415-
error = np.sum(np.append(error_x, error_y))
416-
pull_x, pull_y = pull_back(coords, orig_xy_disp_coord)
417-
418-
text_shifts_x *= force_text[0]
419-
text_shifts_y *= force_text[1]
420-
static_shifts_x *= force_static[0]
421-
static_shifts_y *= force_static[1]
422-
# Pull is in the opposite direction, so need to negate it
423-
pull_x *= -force_pull[0]
424-
pull_y *= -force_pull[1]
425-
426-
if only_move:
427-
if "x" not in only_move["text"]:
428-
text_shifts_x = np.zeros_like(text_shifts_x)
429-
elif "x+" in only_move["text"]:
430-
text_shifts_x[text_shifts_x > 0] = 0
431-
elif "x-" in only_move["text"]:
432-
text_shifts_x[text_shifts_x < 0] = 0
433-
434-
if "y" not in only_move["text"]:
435-
text_shifts_y = np.zeros_like(text_shifts_y)
436-
elif "y+" in only_move["text"]:
437-
text_shifts_y[text_shifts_y > 0] = 0
438-
elif "y-" in only_move["text"]:
439-
text_shifts_y[text_shifts_y < 0] = 0
440-
441-
if "x" not in only_move["static"]:
442-
static_shifts_x = np.zeros_like(static_shifts_x)
443-
elif "x+" in only_move["static"]:
444-
static_shifts_x[static_shifts_x > 0] = 0
445-
elif "x-" in only_move["static"]:
446-
static_shifts_x[static_shifts_x < 0] = 0
447-
448-
if "y" not in only_move["static"]:
449-
static_shifts_y = np.zeros_like(static_shifts_y)
450-
elif "y+" in only_move["static"]:
451-
static_shifts_y[static_shifts_y > 0] = 0
452-
elif "y-" in only_move["static"]:
453-
static_shifts_y[static_shifts_y < 0] = 0
454-
455-
if "x" not in only_move["pull"]:
456-
pull_x = np.zeros_like(pull_x)
457-
elif "x+" in only_move["pull"]:
458-
pull_x[pull_x > 0] = 0
459-
elif "x-" in only_move["pull"]:
460-
pull_x[pull_x < 0] = 0
461-
462-
if "y" not in only_move["pull"]:
463-
pull_y = np.zeros_like(pull_y)
464-
elif "y+" in only_move["pull"]:
465-
pull_y[pull_y > 0] = 0
466-
elif "y-" in only_move["pull"]:
467-
pull_y[pull_y < 0] = 0
468-
469-
shifts_x = text_shifts_x + static_shifts_x + pull_x
470-
shifts_y = text_shifts_y + static_shifts_y + pull_y
471-
472-
coords = apply_shifts(coords, shifts_x, shifts_y)
473-
if ensure_inside_axes:
474-
coords = force_into_bbox(coords, ax_bbox)
475505

476506
if timer() - start_time > time_lim:
477507
break

0 commit comments

Comments
 (0)