Optimized :meth:.NumberLine.number_to_point

[chopan050]

Overview: What does this pull request change?

I've made some changes to NumberLine in order to speedup its number_to_point method, which is relevant when plotting or updating many graphs in a coordinate system:

• Changed np.vstack(alphas) to alphas.reshape(-1, 1), which removes a major bottleneck when using vectorization.
• Recalculated the range of the portion of the line excluding the tips, to use it instead of repeatedly calling VMobject.get_start and VMobject.get_end, which are the most important bottleneck in general.
• Added a new class UnitLinearBase in bases.py, with methods function and inverse_function that return the same value that is passed, for an extra minor speedup.

Motivation and Explanation: Why and how do your changes improve the library?

This is the code I've been using for testing:

from manim import *

class LaplaceScene(Scene):
    def construct(self):
        axes = Axes(x_range=[-1, 7], y_range=[-1, 4], x_length=16.32, y_length=10.20).add_coordinates()
        s = ValueTracker(0)
    
        def get_min_x(k):
            if k > 3:
                return np.floor(2*np.log(k/3.5)/0.5) / 2 - 0.5
            if k < 0:
                return np.floor(2*np.log(-k/0.5)/0.5) / 2 - 0.5
            return -0.5

        draw_horizontal = lambda k: axes.plot(
            lambda t: k*np.exp(-s.get_value()*t), x_range=[get_min_x(k), 6.5, 0.5],
            stroke_width=1.0, stroke_color=RED_A,
            use_vectorized=True,
        )
        horizontals = {k: draw_horizontal(k) for k in range(-15, 96)}
        del horizontals[0]

        verticals = {
            k: Line(
                axes.c2p(k, -1), axes.c2p(k, 5),
                stroke_width=1.0, stroke_color=BLUE_A,
            )
            for k in range(1, 7)
        }

        draw_func = lambda: axes.plot(
            lambda t: t*np.exp(-s.get_value()*t), x_range=[0, 6.5, 0.5],
            stroke_color=YELLOW,
            use_vectorized=True,
        )
        func = draw_func()

        self.play(
            *[Create(verticals[key]) for key in verticals],
            *[Create(horizontals[key]) for key in horizontals],
            DrawBorderThenFill(axes),
        )
        self.play(Create(func), run_time=2)

        func.add_updater(lambda plot: plot.become(draw_func()))
        for k in horizontals:
            (lambda k: horizontals[k].add_updater(
                lambda line: line.become(draw_horizontal(k))
            ))(k)

        self.play(s.animate.set_value(0.5), run_time=5)

        func.clear_updaters()
        for k in horizontals:
            horizontals[k].clear_updaters()

        self.wait()
        self.play(VGroup(axes, func, *horizontals.values(), *verticals.values()).animate.scale(0.5))
        self.wait()
        
with tempconfig({"preview": True, "quality": "medium_quality", "disable_caching": True}):
    scene = LaplaceScene()
    scene.render()

There are 109 ParametricFunctions being updated for 5 seconds, which is costly. A big portion of the runtime is spent on ParametricFunction.generate_points, and there are 3 main culprits: VMobject.make_smooth, VMobject.add_points_as_corners, and Axes.coords_to_point. I'm addressing the first 2 of them in other PRs.

In this PR I shall address Axes.coords_to_point, or more specifically NumberLine.number_to_point (the pale green block which is two blocks below Axes.coords_to_point):

Here is the original code for NumberLine.number_to_point:

def number_to_point(self, number: float | np.ndarray) -> np.ndarray:
        number = np.asarray(number)
        scalar = number.ndim == 0
        number = self.scaling.inverse_function(number)
        alphas = (number - self.x_range[0]) / (self.x_range[1] - self.x_range[0])
        alphas = float(alphas) if scalar else np.vstack(alphas)
        val = interpolate(self.get_start(), self.get_end(), alphas)
        return val

np.vstack

NOTE: this problem only happened when an array (rather than a single float) is passed to NumberLine.number_to_line.

The major bottleneck is np.vstack(alphas). What np.vstack does in this case is to transform alphas into a column, i.e. if alphas was [0.4, 0.8, 0.1], it becomes [[0.4], [0.8], [0.1]]. The thing is, np.vstack does a lot of preprocessing behind the curtains, and creates a new array to hold these values. In this case, that preprocessing is not necessary, and we don't actually need to copy the array. A simple view generated with alphas.reshape(-1, 1) is sufficient:

        alphas = float(alphas) if scalar else alphas.reshape(-1, 1)

And this bottleneck is now gone.

TipableVMobject.get_start and TipableVMobject.get_end

In the last line:

        val = interpolate(self.get_start(), self.get_end(), alphas)

there are calls to TipableVMobject's methods get_start and get_end. They check if the line has a tip (calling has_tip which does return hasattr(self, tip) and self.tip in self, which is expensive) and, if it does, call its get_start method, or else call TipableVMobject.get_start. They call VMobject.throw_error_if_no_points, which is another verification process. All of that is unnecessarily expensive for this case, where we do know the NumberLine should have points, and there should be a way to get the x range of the line whether it has tips or not without recurring to these expensive calculations, right?

Well, this issue was trickier to solve, because I couldn't just plug in self.points[0] and self.points[-1] instead of self.get_start() and self.get_end(). This is because the actual "line" in NumberLine becomes shorter when tips are added, and thus the NumberLine.x_range attribute does no longer represent the range of the line without the tips, but the range of the full line including the tips. I had to do something more.

To solve this, I overrode TipableVMobject's methods add_tip and pop_tips, where I calculated the range of the actual portion of the line excluding the tips. For this, I created new attributes for NumberLine: x_range_no_tips, x_min_no_tips and x_max_no_tips.

class NumberLine(Line):
    def __init__(self, ...):
        # ...some preprocessing...
        # turn into a NumPy array to scale by just applying the function
        self.x_range = np.array(x_range, dtype=float)
        self.x_min, self.x_max, self.x_step = scaling.function(self.x_range)
        self.x_range_no_tips = self.x_range.copy()
        self.x_min_no_tips = self.x_min
        self.x_max_no_tips = self.x_max
        # ...init some attributes...
        if self.include_tip:
            self.add_tip(
                tip_length=self.tip_height,
                tip_width=self.tip_width,
                tip_shape=tip_shape,
            )
            self.tip.set_stroke(self.stroke_color, self.stroke_width)
        # ...some more processes...

    def add_tip(
        self, tip=None, tip_shape=None, tip_length=None, tip_width=None, at_start=False
    ):
        old_start = self.points[0].copy()
        old_end = self.points[-1].copy()
        super().add_tip(tip, tip_shape, tip_length, tip_width, at_start)
        new_start = self.points[0]
        new_end = self.points[-1]

        direction = old_end - old_start
        length2 = np.dot(direction, direction)
        new_start_proportion = np.dot(new_start - old_start, direction) / length2
        new_end_proportion = np.dot(new_end - old_start, direction) / length2

        range_min, range_max, _ = self.x_range_no_tips
        diff = range_max - range_min
        self.x_range_no_tips[0] = range_min + new_start_proportion * diff
        self.x_range_no_tips[1] = range_min + new_end_proportion * diff
        self.x_min_no_tips, self.x_max_no_tips, _ = self.scaling.function(
            self.x_range_no_tips
        )

        return self

    def pop_tips(self):
        result = super().pop_tips()
        self.x_range_no_tips[:] = self.x_range
        self.x_min_no_tips = self.x_min
        self.x_max_no_tips = self.x_max
        return result

With that done, now I can just plug in self.points[0] and self.points[-1] in NumberLine.number_to_point, if I use x_range_no_tips instead of x_range:

def number_to_point(self, number: float | np.ndarray) -> np.ndarray:
        number = np.asarray(number)
        scalar = number.ndim == 0
        number = self.scaling.inverse_function(number)
        range_min, range_max, _ = self.x_range_no_tips
        alphas = (number - range_min) / (range_max - range_min)
        alphas = float(alphas) if scalar else alphas.reshape(-1, 1)
        val = interpolate(self.points[0], self.points[-1], alphas)
        return val

IMPORTANT NOTE: With those parameters, I also modified some other functions to optimize them: see point_to_number, get_unit_vector and get_unit_size.

NumberLine.scaling.inverse_function and the adding of UnitLinearBase

Finally, there's a very small portion of NumberLine.number_to_point (the small paler green block at the right of NumberLine.number_to_line) where we call NumberLine.scaling.inverse_function. Now, the default base of NumberLine is a LinearBase, whose default scale factor is 1, and its function and inverse function consist of return self.scale_factor * value and return value / self.scale_factor respectively. If the default scale factor is 1, this just returns the value itself, but multiplying and dividing by 1 still creates some small overhead.

It is truly a small detail in comparison to the previous two points, but as I wanted to optimize NumberLine.number_to_point as much as possible for my use case, I decided to create a new UnitLinearBase whose function and inverse function consist solely of returning the same value passed as parameter:

class UnitLinearBase(LinearBase):
    def __init__(self):
        super().__init__(scale_factor=1.0)

    def function(self, value: float) -> float:
        return value

    def inverse_function(self, value: float) -> float:
        return value

Then I imported it in number_line.py and used it as the default NumberLine.scaling attribute. It's around 20x faster than the original "multiply/divide by 1" approach.

Results

I managed a speedup of around 3.5x in NumberLine.number_to_line! Compare the results:

Before	After

Links to added or changed documentation pages

Further Information and Comments

Reviewer Checklist

• The PR title is descriptive enough for the changelog, and the PR is labeled correctly
• If applicable: newly added non-private functions and classes have a docstring including a short summary and a PARAMETERS section
• If applicable: newly added functions and classes are tested

[chopan050]

The tests have failed, but only by errors of really small orders of magnitude (around 1e-16 or less). Can these be ignored (as these errors can be negligible) or not really?

[MrDiver]

I would still suggest taking a look at the currently failing tests. There maybe is a way to make it perfectly centered again with some small nudging.

[chopan050]

Converted to draft while I try to figure out what to do with this. I tried some stuff, and it still doesn't pass the tests:

• round the result returned by n2p to "fix" the floating point errors
• reverted the calculation of endpoints when there are tips
• reverted almost everything

and it still errors out 😅 I'll figure out later what to do with this.