Optimized :meth:`.NumberLine.number_to_point`
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)toalphas.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_startandVMobject.get_end, which are the most important bottleneck in general. - Added a new class
UnitLinearBaseinbases.py, with methodsfunctionandinverse_functionthat 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
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?
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.
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
n2pto "fix" the floating point errors - reverted the calculation of endpoints when there are tips
- reverted almost everything
and it still errors out :sweat_smile: I'll figure out later what to do with this.