[WIP] Chalk-NumPy

[srush]

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]

Huh, for some reason I have been really getting into this again. Thought this was really cool.

With numpy Trace is now super fast because you can batch over both the many rays and many segments. This lets you run a ton of traces simultaneously. So I thought it would be fun to write a rasterizer in Chalk. Basically you just trace from the left corner, and turn on the pixels in between the traces. You can do multiple traces per pixel and get basic antialiasing. And you can use envelopes to limit the size of the scan.

The top figure is the traces in Chalk, and the bottom is the "rendered" output!

[danoneata]

Hey! Thanks for the mega effort @srush! A quick heads up: I'm on holiday until the end of next week; I'll check the PR when I get back. 😉

[srush]

Of course, enjoy your holiday!

[danoneata]

Hello Sasha! I finally had some time to look at the PR! Thanks again for the impressive work! I will need more time to go through everything, but I'll leave some feedback on some parts that I've looked at.

A difficulty that I'm facing is that the code runs much slower than before: for example, even drawing a simple shape like a circle takes five seconds. I'm not sure whether I'm doing something wrong: are there some Jax settings that I need to specify?

[srush]

Yeah, I'm sorry this all got much too complex, I'm simplifying it down a bit and trying to make sure all the examples work. I'm going to put together a blog post first.

The tradeoff is if you use Numpy it will be as fast as before, but if you use Jax it has a startup cost, but then scales really well for things like animations.

[srush]

Fun Jax example:

def draw(t):
    red = to_color("red")
    green = to_color("green")

    @jax.vmap
    def multi(i):
        i_s = i / 100
        s = (0.05 * i ) / 2
        rot = tx.rotation(i * t + i)
        return rectangle(s, s).fill_color(i_s * red + (1-i_s) *green).line_width(0).apply_transform(rot)
    return multi(np.arange(100, 1, -1))

[danoneata]

Yeah, I'm sorry this all got much too complex, I'm simplifying it down a bit and trying to make sure all the examples work. I'm going to put together a blog post first.

No worries! The nature of the change makes it touch many parts of the codebase. Looking forward to the blog post! 🙂

The tradeoff is if you use Numpy it will be as fast as before, but if you use Jax it has a startup cost, but then scales really well for things like animations.

Ah, I missed the flag in transform.py that controls which library to use. This flag might be useful to expose to the user in some way.

Fun Jax example:

Very cool example indeed! 😵‍💫🙂 I've also used the same primitive (a function from time to diagram) to specify an animation in the animation repository.

[srush]

Yeah, the animation lib looks super neat! Was playing around with it a bit, it's very clean. (Thinking a bit about how these two systems might complement each other)

Here's a draft of the blog post: https://srush.github.io/DiffRast/ . Haven't done a proofread of the math yet so there might be some bugs. Would love any feedback!

[danoneata]

Wow, this is unbelievably cool @srush! I'll have a closer read tomorrow (especially the last two sections), but here are a few small observations:

1. In the equation block at the start of the article:

   • second line: should be probably be image = rasterizer(vector) instead of image = rasterizer(x)?
   • third line: should probably use L(x, target) instead of L(png, target)?

2. The matplotlib images seem to be missing from Section 4?

3. The table of contents links to sections 3–5 do not work.

4. To do rasterization, we will use a simple scanline algorithm.

The link to "scaline" redirects to the same page.

5.

Here search for the outer angle that leads to the largest trace.

Instead of "largest" shouldn't it be "smallest" or "thinnest"?

6. Maybe it's a bit more clearer if the title or the beginning of the article would explicitly mention the task of interest: image vectorization (or raster-to-vector conversion)? Differentiable rasterization seems to be the means to this end.

[danoneata]

A couple more small comments:

• $θ$ could be removed from $f(x - u; θ)$ since it's not defined nor used anywhere else?
• Typo: "To do step, we are going [...]"

Regarding the technical part from section 5, unfortunately that's a bit over my head 🙁 So I don't think I have any useful comments regarding that. I do have many questions, but I don't know if they are well posed or useful to improve the presentation of the blog-post; for example:

• Why is it sufficient to use the 1D integral rule; I see that in the paper they mention the 2D generalisation?
• Why is it accurate to approximate $f(x−\text{split})k(\text{split})$ with $(f(x−\text{split}+ϵ)−f(x−\text{split}−ϵ))k(\text{split})$? It looks to me that for $\epsilon\to 0$, the second term vanishes?
• When defining the backward pass for the boundary function, why is okay to omit the second term in Leibniz's integral rule (the integral)? Edit: Ah, is it because the shape's derivative takes care of that?
• Why does the integral start from 0 (and not from an arbitrary boundary)?

[srush]

These are all very helpful and well posed. Let me try to answer them.

Why is it sufficient to use the 1D integral rule; I see that in the paper they mention the 2D generalisation?

It's not really sufficient per se. It only works because I'm doing two scanlines, one left-to-right and one top-to-bottom. Therefore I get an approximation of 2d with the simpler 1d formulation.

In their paper they're using a much more complex rasterizer, but they also require normals, path sampling, nearby paths, etc, which we don't have acess to in Chalk.

Why is it accurate to approximate $f(x−\text{split})k(\text{split})$ with $(f(x−\text{split}+ϵ)−f(x−\text{split}−ϵ))k(\text{split})$? It looks to me that for $\epsilon\to 0$, the second term vanishes?

Hmm, you're right this is weird and seems wrong. I just took it from their paper.... I wonder if what is really going on is what you mentioned below 👇 i.e. it is + for the right boundary of A_1 and - for the left boundary of A_2.

Seems like this section needs to be more clear.

When defining the backward pass for the boundary function, why is okay to omit the second term in Leibniz's integral rule (the integral)? Edit: Ah, is it because the shape's derivative takes care of that?

Yeah so I think it's because that term is 0 wrt to the trace, but non-zero wrt color. So we are already handling that.

Why does the integral start from 0 (and not from an arbitrary boundary

I'll make this more clear. I was just showing the general one sided version of Leibniz. You are right that in practice we have both.

[danoneata]

Thanks a lot for sharing your thoughts, Sasha!