schematic: add various options

Author: jcmgray

docs/changelog.md

@@ -10,6 +10,7 @@ Release notes for `quimb`.
 - add new example on tracing tensor network functions {ref}`ex_tracing_tn_functions`
 - update infrastructure for TEBD and SimpleUpdate based algorithms.
 - [`schematic.Drawing`](quimb.schematic.drawing): add [`grid`](quimb.schematic.drawing.grid), [`grid3d`](quimb.schematic.drawing.grid3d), [`bezier`](quimb.schematic.drawing.bezier), [`star`](quimb.schematic.drawing.star), [`cross`](quimb.schematic.drawing.cross) and [`zigzag`](quimb.schematic.drawing.zigzag) methods.
+- [`schematic.Drawing`](quimb.schematic.drawing): add `relative` option to [`arrowhead`](quimb.schematic.drawing.arrowhead), `shorten` option to [`text_between`](quimb.schematic.drawing.text_between) and `text_left` and `text_right` options to [`line`](quimb.schematic.drawing.line).
 - refactor [`TEBDGen`](quimb.tensor.tensor_arbgeom_tebd.TEBDGen) and [`SimpleUpdateGen`](quimb.tensor.tensor_arbgeom_tebd.SimpleUpdateGen)
 
 (whats-new-1-11-2)=

quimb/schematic.py

@@ -230,6 +230,10 @@ def text_between(self, cooa, coob, text, preset=None, **kwargs):
         center : float, optional
             The position of the text along the line, where 0.0 is the start and
             1.0 is the end. Default is 0.5.
+        shorten : float or tuple[float, float], optional
+            Shorten the implicit line by this *absolute* amount at each end.
+            If a tuple, the first value is the start shortening, the second the
+            end shortening.
         preset : str, optional
             A preset style to use for the text.
         kwargs
@@ -241,6 +245,7 @@ def text_between(self, cooa, coob, text, preset=None, **kwargs):
         style.setdefault("verticalalignment", "center")
         style.setdefault("clip_on", False)
         center = style.pop("center", 0.5)
+        shorten = style.pop("shorten", None)
         zorder_delta = style.pop("zorder_delta", 0.02)
 
         if len(cooa) == 2:
@@ -256,6 +261,18 @@ def text_between(self, cooa, coob, text, preset=None, **kwargs):
             xa, ya = self._3d_project(*cooa)
             xb, yb = self._3d_project(*coob)
 
+        if shorten is not None:
+            forward, inverse = get_rotator_and_inverse((xa, ya), (xb, yb))
+            R = forward(xb, yb)[0]
+            try:
+                start, end = shorten
+            except TypeError:
+                start = end = shorten
+            ra = (start, 0.0)
+            rb = (R - end, 0.0)
+            xa, ya = inverse(*ra)
+            xb, yb = inverse(*rb)
+
         # compute midpoint
         x = xa * (1 - center) + xb * center
         y = ya * (1 - center) + yb * center
@@ -591,6 +608,8 @@ def line(self, cooa, coob, preset=None, **kwargs):
         shorten = style.pop("shorten", None)
         arrowhead = style.pop("arrowhead", None)
         text = style.pop("text")
+        text_left = style.pop("text_left", None)
+        text_right = style.pop("text_right", None)
         zorder_delta = style.pop("zorder_delta", 0.0)
 
         if len(cooa) == 2:
@@ -624,6 +643,8 @@ def line(self, cooa, coob, preset=None, **kwargs):
             xa, ya = inverse(*ra)
             xb, yb = inverse(*rb)
 
+        cooa, coob = (xa, ya), (xb, yb)
+
         if arrowhead is not None:
             if arrowhead is True:
                 arrowhead = {}
@@ -641,9 +662,27 @@ def line(self, cooa, coob, preset=None, **kwargs):
                 text = dict(text)
 
             # don't want to pass full style dict to text_between
-            text.setdefault("zorder", style["zorder"])
+            text.setdefault("zorder", style["zorder"] + 0.02)
             self.text_between(cooa, coob, **text)
 
+        if text_left:
+            if isinstance(text_left, str):
+                text_left = {"text": text_left}
+            else:
+                text_left = dict(text_left)
+
+            text_left.setdefault("zorder", style["zorder"] + 0.02)
+            self.text_between(cooa, coob, center=0.0, **text_left)
+
+        if text_right:
+            if isinstance(text_right, str):
+                text_right = {"text": text_right}
+            else:
+                text_right = dict(text_right)
+
+            text_right.setdefault("zorder", style["zorder"] + 0.02)
+            self.text_between(cooa, coob, center=1.0, **text_right)
+
         self._adjust_lims(xa, ya)
         self._adjust_lims(xb, yb)
 
@@ -690,6 +729,8 @@ def line_offset(
         style.setdefault("text", None)
         arrowhead = style.pop("arrowhead")
         text = style.pop("text")
+        text_left = style.pop("text_left", None)
+        text_right = style.pop("text_right", None)
         zorder_delta = style.pop("zorder_delta", 0.0)
 
         if len(cooa) == 2:
@@ -741,6 +782,24 @@ def line_offset(
             text.setdefault("zorder", style["zorder"])
             self.text_between(cooml, coomr, **text)
 
+        if text_left:
+            if isinstance(text_left, str):
+                text_left = {"text": text_left}
+            else:
+                text_left = dict(text_left)
+
+            text_left.setdefault("zorder", style["zorder"])
+            self.text_between(cooml, coomr, center=0.0, **text_left)
+
+        if text_right:
+            if isinstance(text_right, str):
+                text_right = {"text": text_right}
+            else:
+                text_right = dict(text_right)
+
+            text_right.setdefault("zorder", style["zorder"])
+            self.text_between(cooml, coomr, center=1.0, **text_right)
+
         for coo in curve_pts:
             self._adjust_lims(*coo)
 
@@ -836,6 +895,9 @@ def arrowhead(self, cooa, coob, preset=None, **kwargs):
             The width of the arrowhead. Default is 0.05.
         length : float, optional
             The length of the arrowhead. Default is 0.1.
+        relative : bool or float, optional
+            If True, the the `width` and `length` parameters are scaled as
+            `total_line_length ** relative`.
         preset : str, optional
             A preset style to use for the arrowhead, including the above
             options.
@@ -875,10 +937,13 @@ def arrowhead(self, cooa, coob, preset=None, **kwargs):
         center = style.pop("center")
         width = style.pop("width")
         length = style.pop("length")
+        relative = style.pop("relative", 1.0)
+
         lab = rb[0]
+        scale = lab ** relative
         xc = center * lab
-        arrow_x = xc - length * lab
-        arrow_y = width * lab
+        arrow_x = xc - length * scale
+        arrow_y = width * scale
 
         aa, ab, ac = [
             inverse(arrow_x, arrow_y),