feat: allow neutral colors and postprocessing

[gdalle]

Fixes #122, fixes #158

Done:

• Add a postprocessing kwarg to GreedyColoringAlgorithm which allows removing useless colors and replacing them with the "neutral" color 0. This only applies to symmetric or bidirectional colorings.
• Explain postprocessing in the relevant docstrings.
• Implement postprocessing by:
  1. mirroring the decompression code to detect which colors are needed
  2. zero out unneeded colors
  3. renumber colors to get a lower maximum color
• Improve the StarSet and TreeSet storages:
  • Clarify the negative hub trick for star coloring, and undo it before building the StarSet
  • Move the BFS computation for acyclic coloring from the result constructor to the TreeSet constructor, so that it is available for postprocessing
• Modify color grouping to ignore neutral colors, adjust stacking to handle empty color groups
• Adapt tests.
• Adapt visualization.

To do:

• Add tests on small instance with known results (More tests for bicoloring with preprocessing #169)
  • Rectangle
  • Arrowhead

[codecov]

Codecov Report

All modified and coverable lines are covered by tests ✅

Project coverage is 100.00%. Comparing base (0794c68) to head (974e38b).
Report is 1 commits behind head on main.

Additional details and impacted files

@@            Coverage Diff            @@
##              main      #166   +/-   ##
=========================================
  Coverage   100.00%   100.00%           
=========================================
  Files           13        13           
  Lines         1349      1422   +73     
=========================================
+ Hits          1349      1422   +73     

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

[amontoison]

We could easily check if a neutral color was assigned before calling this function.
We just need to update a counter in the "if condition" line 564.

[amontoison]

Do we already know at this step which columns / rows have the same color?
It could help to determine if the neutral color helped to remove a color or just impact the seed.

Talking about the seed, we didn't discussed about that but even if have the same number of colors, having seeds with more zeros (because of rows / columns of neutral color) could not help AD tools where they walk through an expression tree?

[gdalle]

We could easily check if a neutral color was assigned before calling this function.

Good idea, I'm doing it right now

Do we already know at this step which columns / rows have the same color?
It could help to determine if the neutral color helped to remove a color or just impact the seed.

The present code doesn't remove a color only from some vertices, it removes it all at once if that color was found to be useless.

Talking about the seed, we didn't discussed about that but even if have the same number of colors, having seeds with more zeros (because of rows / columns of neutral color) could not help AD tools where they walk through an expression tree?

No, the AD tools are agnostic to what's inside the seed. You can think of it as matrix-vector multiplication: doing A * rand(10) doesn't take more time than A * zeros(10)

[amontoison]

@gdalle In acyclic coloring, we can change the color of some columns to the neutral color if they are leaves in certain special trees (trivial trees and stars).
It's the only cases where a given color only appears as leaves.
Since we use 2-colored trees, all other kind of trees will require both colors to be non-neutral.

[gdalle]

@gdalle In acyclic coloring, we can change the color of some columns to the neutral color if they are leaves in certain special trees (trivial trees and stars).
It's the only cases where a given color only appears as leaves.
Since we use 2-colored trees, all other kind of trees will require both colors to be non-neutral.

Good point, this can definitely go in the paper! Functionally, the current code does the same thing, except it is slightly obscured because I chose to stick to the structure of the decompression routine. But I opened #168 to keep track of the idea

[amontoison]

LGTM
I propose to add more unit tests (#169) before a new release.
The rectangle problem is good example for checking that the option postprocessing doesn't break something.

[amontoison]

@gdalle In acyclic coloring, we can change the color of some columns to the neutral color if they are leaves in certain special trees (trivial trees and stars).
It's the only cases where a given color only appears as leaves.
Since we use 2-colored trees, all other kind of trees will require both colors to be non-neutral.

Good point, this can definitely go in the paper! Functionally, the current code does the same thing, except it is slightly obscured because I chose to stick to the structure of the decompression routine. But I opened #168 to keep track of the idea

To add to my answer, some vertices could only be leaves of "normal" trees and could be neutralized. However, because other vertices of the same initial colors are internal nodes of the trees, this will not lead to a reduction in the number of colors.
Only seeds could be impacted, but as discussed earlier today, this will not provide any advantage for most common AD tools.

I briefly discussed this question of gain with modified seeds (which lead to full rows of zeros in the seed matrix) with Michel Schanen (@michel2323) this morning.
Expression trees could potentially be simplified / pruned in code that dynamically generates or specializes the differentiation code.
An example could be packages with their own DSL, like JuMP, which could simplify the expression tree for forward or reverse passes.
That's the theory—in practice, not many people take advantage of that.