update

Author: srush

chalk/array_types.py

@@ -48,7 +48,6 @@ def vmap(fn: Callable[[Array], Array]) -> Callable[[Array], Array]:
         "Fake jax vmap for numpy as a for loop."
         if JAX_MODE:
             return vmap(fn)
-
         def vmap2(x: Array) -> Array:
             if isinstance(x, tuple):
                 size = x[-1].size()  # type: ignore

chalk/backend/cairo.py

@@ -124,7 +124,7 @@ def animate(
     path_frame = "/tmp/frame-{:d}.png"
     import imageio
 
-    with imageio.get_writer(path, fps=20, loop=0) as writer:
+    with imageio.get_writer(path, fps=10, loop=0) as writer:
         for i in range(shape[0]):
             path = path_frame.format(i)
             patches_to_file(patches, path, h, w, (i,))

chalk/backend/patch.py

@@ -51,7 +51,8 @@ def order_patches(
 ) -> List[Tuple[Tuple[int, ...], Patch, Dict[str, Any]]]:
     import numpy as onp
 
-    patches = tx.tree_map(onp.asarray, patches)
+    if tx.JAX_MODE:
+        patches = tx.tree_map(onp.asarray, patches)
 
     d = {}
     for patch in patches:
@@ -99,6 +100,7 @@ def from_path(
         for loc_trail in path.loc_trails:
             p = loc_trail.location
             segments = loc_trail.located_segments()
+            print("SEG", transform.shape, transform.strides, loc_trail.trail.segments.transform.strides)
             vert = segment_to_curve(segments.transform, segments.angles)
             if path.scale_invariant is not None:
                 scale = height / 20
@@ -147,10 +149,11 @@ def patch_from_prim(
     style = prim.style if prim.style is not None else style
     assert isinstance(prim.prim_shape, Path)
     assert prim.order is not None
+    in_style = style.to_mpl() #tx.multi_vmap(style.to_mpl.__func__, len(size))(style),  # type: ignore
     patch = Patch.from_path(
         prim.prim_shape,
         prim.transform,
-        tx.multi_vmap(style.to_mpl.__func__, len(size))(style),  # type: ignore
+        in_style,
         prim.order,
         height,
     )

chalk/backend/svg.py

@@ -12,27 +12,27 @@ def to_svg(patch: Patch, ind: Tuple[int, ...]) -> str:
     v, c = patch.vert[ind], patch.command[ind]
     if v.shape[0] == 0:
         return "<g></g>"
-    line = ""
+    parts = []
     i = 0
     while i < c.shape[0]:
         if c[i] == chalk.backend.patch.Command.MOVETO.value:
-            line += f"M {v[i, 0]} {v[i, 1]}"
+            parts.append( f"M {v[i, 0]:.2f} {v[i, 1]:.2f}")
             i += 1
         elif c[i] == chalk.backend.patch.Command.LINETO.value:
-            line += f"L {v[i, 0]} {v[i, 1]}"
+            parts.append(f"L {v[i, 0]} {v[i, 1]}")
             i += 1
         elif c[i] == chalk.backend.patch.Command.CURVE3.value:
-            line += f"Q {v[i, 0]} {v[i, 1]} {v[i+1, 0]} {v[i+1, 1]}"
+            parts.append(f"Q {v[i, 0]} {v[i, 1]} {v[i+1, 0]} {v[i+1, 1]}")
             i += 2
         elif c[i] == chalk.backend.patch.Command.CLOSEPOLY.value:
-            line += "Z"
+            parts.append("Z")
             i += 1
         elif c[i] == chalk.backend.patch.Command.SKIP.value:
             i += 1
         elif c[i] == chalk.backend.patch.Command.CURVE4.value:
-            line += f"C {v[i, 0]} {v[i, 1]} {v[i+1, 0]} {v[i+1, 1]} {v[i+2, 0]} {v[i+2, 1]}"
+            parts.append(f"C {v[i, 0]:.2f} {v[i, 1]:.2f} {v[i+1, 0]:.2f} {v[i+1, 1]:.2f} {v[i+2, 0]:.2f} {v[i+2, 1]:.2f}")
             i += 3
-    return line
+    return " ".join(parts)
 
 
 def write_style(d: Dict[str, Any]) -> Dict[str, str]:
@@ -51,48 +51,68 @@ def write_style(d: Dict[str, Any]) -> Dict[str, str]:
     return out
 
 
-def render_svg_patches(patches: List[Patch], animate:bool =False, time_steps=0) -> str:
-    out = ""
-    patches = [chalk.backend.patch.order_patches(patches, (step,)) 
-               for step in range(time_steps)]
-    for v in zip(*patches):
-        out += f"""
-    <path>
-"""     
-        lines = []
-        css = {}
-        for ind, patch, style_new in v:
-            lines.append(to_svg(patch, ind))
-            s = write_style(style_new)
-            for k, v in s.items():
-                css.setdefault(k, []).append(v)
-        
-        values = ";".join(lines)
-        out += f"""
-        <animate attributeName="d" values="{values}" dur="2s" repeatCount="indefinite"/>
+def render_svg_patches(patches: List[Patch], animate:bool=False, time_steps:int=0) -> str:
+    if animate:
+        out = ""
+        patches = [chalk.backend.patch.order_patches(patches, (step,)) 
+                for step in range(time_steps)]
+        for v in zip(*patches):
+            out += "\n\n <path>\n"
+            lines = []
+            css = {}
+            
+
+            for ind, patch, style_new in v:
+
+                lines.append(to_svg(patch, ind))
+                s = write_style(style_new)
+                for k, v in s.items():
+                    css.setdefault(k, []).append(v)
+            s = set(lines)
+            if len(s) == 1:
+                out += f"""
+                <set attributeName="d" to="{list(s)[0]}"/>
+                """
+            else:
+                values = ";".join(lines)
+                out += f"""
+                <animate attributeName="d" values="{values}" dur="2s" repeatCount="indefinite"/>
+                """
+            for k, v in css.items():
+                s = set(v)
+                if len(s) == 1:
+                    out += f"""<set attributeName="{k}" to="{list(s)[0]}"/>"""
+
+                else:
+                    out += f"""
+                <animate attributeName="{k}" values="{';'.join(v)}" dur="2s" repeatCount="indefinite"/>
         """
-        for k, v in css.items():
+            out += "</path>\n\n"
+        return out
+    else:
+        out = ""
+        for ind, patch, style_new in chalk.backend.patch.order_patches(patches):
+            inner = to_svg(patch, ind)
+            style_t = ";".join([f"{k}:{v}" for k, v in write_style(style_new).items()] )
             out += f"""
-        <animate attributeName="{k}" values="{';'.join(v)}" dur="2s" repeatCount="indefinite"/>
-"""
-        out += """
-    </path>
-    """
-    return out
-
+            <g style="{style_t}">
+                <path d="{inner}" />
+            </g>"""
+        return out
 def patches_to_file(
     patches: List[Patch], path: str, height: tx.IntLike, 
     width: tx.IntLike,
     animate: bool = False,
-    time_steps=0
+    time_steps: int =0
 ) -> None:
-    dwg = f"""<?xml version="1.0" encoding="utf-8" ?>
-<svg baseProfile="full" height="{int(height)}" version="1.1" width="{int(width)}" xmlns="http://www.w3.org/2000/svg" xmlns:ev="http://www.w3.org/2001/xml-events" xmlns:xlink="http://www.w3.org/1999/xlink">
-    """
-    dwg += render_svg_patches(patches, animate, time_steps)
-    dwg += "</svg>"
     with open(path, "w") as f:
-        f.write(dwg)
+        f.write(f"""<?xml version="1.0" encoding="utf-8" ?>
+<svg baseProfile="full" height="{int(height)}" version="1.1" width="{int(width)}" xmlns="http://www.w3.org/2000/svg" xmlns:ev="http://www.w3.org/2001/xml-events" xmlns:xlink="http://www.w3.org/1999/xlink">
+    """)
+        
+        out = render_svg_patches(patches, animate, time_steps)
+        f.write(out)
+        f.write("</svg>")
 
 
 def render(

chalk/broadcast.py

@@ -55,6 +55,14 @@ def check(x: tx.Array) -> None:
 V1 = TypeVar("V1", bound=Diagram)
 V2 = TypeVar("V2", bound=Diagram)
 
+def broadcast_to(tree, old_shape: Tuple[int, ...], new_shape: Tuple[int, ...]):
+    if old_shape == new_shape: return tree
+    def reshape(x: tx.Array) -> tx.Array:
+        shape = x.shape 
+        return tx.np.broadcast_to(x, new_shape + shape[len(old_shape):])
+
+    return tx.tree_map(reshape, tree)
+
 
 def broadcast_diagrams(self: V1, other: V2) -> Tuple[V1, V2]:
     """
@@ -65,20 +73,23 @@ def broadcast_diagrams(self: V1, other: V2) -> Tuple[V1, V2]:
     other_size = other.size()
     if size == other_size:
         return self, other
-    check(size, other_size, str(type(self)), str(type(other)))
-    ml = max(len(size), len(other_size))
-    for i in range(ml):
-        off = -1 - i
-        if i > len(other_size) - 1:
-            other = other.add_axis(size[off])  # type: ignore
-        elif i > len(size) - 1:
-            self = self.add_axis(other_size[off])  # type: ignore
-        elif size[off] == 1 and other_size[off] != 1:
-            self = self.repeat_axis(other_size[off], len(size) + off)  # type: ignore
-        elif size[off] != 1 and other_size[off] == 1:
-            other = other.repeat_axis(size[off], len(other_size) + off)  # type: ignore
-    check_consistent(self)
-    check_consistent(other)
+    new_shape = tx.np.broadcast_shapes(size, other_size)
+    self = broadcast_to(self, size, new_shape)
+    other = broadcast_to(other, other_size, new_shape) 
+
+    # ml = max(len(size), len(other_size))
+    # for i in range(ml):
+    #     off = -1 - i
+    #     if i > len(other_size) - 1:
+    #         other = other.add_axis(size[off])  # type: ignore
+    #     elif i > len(size) - 1:
+    #         self = self.add_axis(other_size[off])  # type: ignore
+    #     elif size[off] == 1 and other_size[off] != 1:
+    #         self = self.repeat_axis(other_size[off], len(size) + off)  # type: ignore
+    #     elif size[off] != 1 and other_size[off] == 1:
+    #         other = other.repeat_axis(size[off], len(other_size) + off)  # type: ignore
+    # check_consistent(self)
+    # check_consistent(other)
     assert (
         self.size() == other.size()
     ), f"{size} {other_size} {self.size()} {other.size()}"

chalk/combinators.py

@@ -102,14 +102,13 @@ def call_scan(diagram: Diagram) -> Diagram:
         env = diagram.get_envelope()
         right = env(v)
         left = env(-v)
-        off = tx.np.roll(right, 1) + left + sep
-        off = tx.index_update(off, 0, 0)
-        off = tx.np.cumsum(off, axis=0)
+        off = tx.np.roll(right, 1, axis=axis) + left + sep
+        off = tx.index_update(off, (Ellipsis, 0), 0)
+        off = tx.np.cumsum(off, axis=axis)
         t = v * off[..., None, None]
         return diagram.translate_by(t)
-    
 
-    call_scan = tx.multi_vmap(call_scan, axis)  # type: ignore
+    # call_scan = tx.multi_vmap(call_scan, axis)  # type: ignore
     return call_scan(diagram).compose_axis()
 
 

chalk/core.py

@@ -85,7 +85,7 @@ def apply_style(self: B1, style: StyleHolder) -> B:  # type: ignore
         return ApplyStyle(new_diagram.style, self)
 
     def __repr__(self) -> str:
-        return f"Diagram[self.shape]"
+        return f"Diagram[{self.shape}]"
 
     def __tree_pp__(self, **kwargs):  # type: ignore
         import jax._src.pretty_printer as pp

chalk/path.py

@@ -86,6 +86,7 @@ def from_array(points: P2_t, closed: bool = False) -> Path:
 
 
 def from_points(points: List[P2_t], closed: bool = False) -> Path:
+    points = tx.np.broadcast_arrays(*points)
     return Path.from_array(tx.np.stack(points, axis=-3))
 
 

chalk/style.py

@@ -131,9 +131,9 @@ def shape(self) -> Tuple[int, ...]:
         return self.base.shape[:-1]
 
     def get(self, key: str) -> tx.Scalars:
-        v = self.base[slice(*STYLE_LOCATIONS[key])]
+        v = self.base[..., slice(*STYLE_LOCATIONS[key])]
         return tx.np.where(
-            self.mask[slice(*STYLE_LOCATIONS[key])], v, DEFAULTS[key]
+            self.mask[..., slice(*STYLE_LOCATIONS[key])], v, DEFAULTS[key]
         )
 
     @property
@@ -193,8 +193,8 @@ def to_mpl(self) -> Dict[str, Any]:
 
         # Set by observation
         lw = self.line_width_
-        style["linewidth"] = lw.reshape(-1)[0]
-        style["alpha"] = self.fill_opacity_[0]
+        style["linewidth"] = lw[...,0]
+        style["alpha"] = self.fill_opacity_[...,0]
         return style
 
 

chalk/transform.py

@@ -83,6 +83,14 @@ def P2(x: Floating, y: Floating) -> P2_t:
     s: P2_t = np.stack([x, y, o], axis=-1)[..., None]
     return s
 
+@jit
+# @partial(vectorize, signature="(),()->(3,1)")
+def to_P2(x: Float[Array, "*B 2"]) -> P2_t:
+    "Map a standard vector to a point"
+    _, o = np.broadcast_arrays(x[..., :1], ftos(1.0))
+    s: P2_t = np.concatenate([x, o], axis=-1)[..., None]
+    return s
+
 
 @jit
 def norm(v: V2_t) -> V2_t: