chalk icon indicating copy to clipboard operation
chalk copied to clipboard

[WIP] Chalk-NumPy

Open srush opened this issue 1 month ago • 0 comments

This is an experimental rewrite of chalk (maybe like a v3) which I thought would be kind of interesting.

Philosophy here is to make any coordinate operation potentially be a batch of coordinates. So the type V2_t could be a vector or 100 vector. What is nice about this is that you can apply many different affine transformations simultaneously with matrix operations. To make this play nicely with Chalk, I rewrote the Trail class so that you can apply Affine simultaneously to all the different arcs in the trail. This should in theory speed up big diagrams and possibly homogenous animations, with a speed cost to smaller ones. The rest of the API should be the same.

Since chalk is mostly an exercise in type systems, one neat component here is that I used the JaxTyping library which lets you do all this in a type safe / dimension way with run time checks: https://docs.kidger.site/jaxtyping/api/array/. I think the types are not that much more complex but it is helpful that it type checks.

Code Changes:

  • This eliminates the need for Planar library dependency.
  • Removed line segments, now everything is an arc.
  • Thinking of removing the Tikz backend. It's annoying to maintain, and I don't use it.

Issues:

  • [ ] - Need to figure out how Trace works here. The fact that it is variable length doesn't play well with numpy
  • [ ] - Currently the implementation is in Jax, but I want it to work with raw numpy as well (shouldn't be too hard)
  • [ ] - Add back some other trail functions

Todos:

  • batch example
  • "PathGroup" for rendering lots of shapes that differ only in transfor
  • differentiation
  • Numpy for style?

srush avatar May 28 '24 21:05 srush