Refactored coordinate_systems.py, fixed bugs, added NumberPlane test (#1092)

* Add function to update a dict recursively with other dicts

* Refactor Axes and NumberPlane

* Fix NumberLine's rotation

* Rotate NumberLine just before adding numbers

* More refactoring

* Add graphical unit test to NumberPlane

* More cleanup and fix range attributes issue(#840)

* Remove overlooked print statements and run black

* Make graphical unit test of NumberPlane more robust

* Shift z_axis to XY-plane's origin and minor refactor in its rotation

* Apply review suggestions

Author: abhijithmuthyala

manim/mobject/coordinate_systems.py

@@ -14,7 +14,7 @@
 from ..mobject.number_line import NumberLine
 from ..mobject.svg.tex_mobject import MathTex
 from ..mobject.types.vectorized_mobject import VGroup
-from ..utils.config_ops import merge_dicts_recursively
+from ..utils.config_ops import merge_dicts_recursively, update_dict_recursively
 from ..utils.simple_functions import binary_search
 from ..utils.space_ops import angle_of_vector
 from ..utils.color import LIGHT_GREY, WHITE, BLUE_D, BLUE
@@ -28,16 +28,12 @@ class CoordinateSystem:
     Abstract class for Axes and NumberPlane
     """
 
-    def __init__(self, dim=2):
+    def __init__(self, x_min=None, x_max=None, y_min=None, y_max=None, dim=2):
         self.dimension = dim
-        if not hasattr(self, "x_min"):
-            self.x_min = -config["frame_x_radius"]
-        if not hasattr(self, "x_max"):
-            self.x_max = config["frame_x_radius"]
-        if not hasattr(self, "y_min"):
-            self.y_min = -config["frame_y_radius"]
-        if not hasattr(self, "y_max"):
-            self.y_max = config["frame_y_radius"]
+        self.x_min = -config["frame_x_radius"] if x_min is None else x_min
+        self.x_max = config["frame_x_radius"] if x_max is None else x_max
+        self.y_min = -config["frame_y_radius"] if y_min is None else y_min
+        self.y_max = config["frame_y_radius"] if y_max is None else y_max
 
     def coords_to_point(self, *coords):
         raise NotImplementedError()
@@ -132,44 +128,56 @@ def input_to_graph_point(self, x, graph):
 class Axes(VGroup, CoordinateSystem):
     def __init__(
         self,
+        x_min=None,
+        x_max=None,
+        y_min=None,
+        y_max=None,
         axis_config=None,
         x_axis_config=None,
         y_axis_config=None,
         center_point=ORIGIN,
         **kwargs
     ):
-        if axis_config is None:
-            axis_config = {
-                "color": LIGHT_GREY,
-                "include_tip": True,
-                "exclude_zero_from_default_numbers": True,
-            }
-        if y_axis_config is None:
-            y_axis_config = {"label_direction": LEFT, "rotation": 90 * DEGREES}
-        self.axis_config = axis_config
-        if x_axis_config is None:
-            x_axis_config = {}
-        self.x_axis_config = x_axis_config
-        self.y_axis_config = y_axis_config
-        self.center_point = center_point
-        CoordinateSystem.__init__(self)
+        CoordinateSystem.__init__(
+            self, x_min=x_min, x_max=x_max, y_min=y_min, y_max=y_max
+        )
         VGroup.__init__(self, **kwargs)
-        self.x_axis = self.create_axis(self.x_min, self.x_max, self.x_axis_config)
-        self.y_axis = self.create_axis(self.y_min, self.y_max, self.y_axis_config)
+
+        self.axis_config = {
+            "color": LIGHT_GREY,
+            "include_tip": True,
+            "exclude_zero_from_default_numbers": True,
+        }
+        self.x_axis_config = {"x_min": self.x_min, "x_max": self.x_max}
+        self.y_axis_config = {
+            "x_min": self.y_min,
+            "x_max": self.y_max,
+            "label_direction": LEFT,
+            "rotation": 90 * DEGREES,
+        }
+
+        self.update_default_configs(
+            (self.axis_config, self.x_axis_config, self.y_axis_config),
+            (axis_config, x_axis_config, y_axis_config),
+        )
+        self.center_point = center_point
+        self.x_axis = self.create_axis(self.x_axis_config)
+        self.y_axis = self.create_axis(self.y_axis_config)
         # Add as a separate group in case various other
         # mobjects are added to self, as for example in
         # NumberPlane below
         self.axes = VGroup(self.x_axis, self.y_axis, dim=self.dim)
         self.add(*self.axes)
         self.shift(self.center_point)
 
-    def create_axis(self, min_val, max_val, axis_config):
-        new_config = merge_dicts_recursively(
-            self.axis_config,
-            {"x_min": min_val, "x_max": max_val},
-            axis_config,
-        )
-        return NumberLine(**new_config)
+    @staticmethod
+    def update_default_configs(default_configs, passed_configs):
+        for default_config, passed_config in zip(default_configs, passed_configs):
+            if passed_config is not None:
+                update_dict_recursively(default_config, passed_config)
+
+    def create_axis(self, axis_config):
+        return NumberLine(**merge_dicts_recursively(self.axis_config, axis_config))
 
     def coords_to_point(self, *coords):
         origin = self.x_axis.number_to_point(0)
@@ -209,37 +217,33 @@ def add_coordinates(self, x_vals=None, y_vals=None):
 class ThreeDAxes(Axes):
     def __init__(
         self,
-        z_axis_config=None,
-        z_min=-3.5,
-        z_max=3.5,
         x_min=-5.5,
         x_max=5.5,
         y_min=-5.5,
         y_max=5.5,
+        z_min=-3.5,
+        z_max=3.5,
+        z_axis_config=None,
         z_normal=DOWN,
         num_axis_pieces=20,
         light_source=9 * DOWN + 7 * LEFT + 10 * OUT,
         **kwargs
     ):
-        if z_axis_config is None:
-            z_axis_config = {}
-        self.z_axis_config = z_axis_config
+        Axes.__init__(
+            self, x_min=x_min, x_max=x_max, y_min=y_min, y_max=y_max, **kwargs
+        )
         self.z_min = z_min
         self.z_max = z_max
+        self.z_axis_config = {"x_min": self.z_min, "x_max": self.z_max}
+        self.update_default_configs((self.z_axis_config,), (z_axis_config,))
         self.z_normal = z_normal
         self.num_axis_pieces = num_axis_pieces
         self.light_source = light_source
-        self.x_min = x_min
-        self.x_max = x_max
-        self.y_min = y_min
-        self.y_max = y_max
-        Axes.__init__(self, **kwargs)
         self.dimension = 3
-        z_axis = self.z_axis = self.create_axis(
-            self.z_min, self.z_max, self.z_axis_config
-        )
-        z_axis.rotate(-np.pi / 2, UP, about_point=ORIGIN)
-        z_axis.rotate(angle_of_vector(self.z_normal), OUT, about_point=ORIGIN)
+        z_axis = self.z_axis = self.create_axis(self.z_axis_config)
+        z_axis.shift(self.center_point)
+        z_axis.rotate_about_zero(-np.pi / 2, UP)
+        z_axis.rotate_about_zero(angle_of_vector(self.z_normal))
         self.axes.add(z_axis)
         self.add(z_axis)
 
@@ -281,28 +285,26 @@ def __init__(
         make_smooth_after_applying_functions=True,
         **kwargs
     ):
-        if axis_config is None:
-            axis_config = {
-                "stroke_color": WHITE,
-                "stroke_width": 2,
-                "include_ticks": False,
-                "include_tip": False,
-                "line_to_number_buff": SMALL_BUFF,
-                "label_direction": DR,
-                "number_scale_val": 0.5,
-            }
-        self.axis_config = axis_config
-        if y_axis_config is None:
-            y_axis_config = {"label_direction": DR}
-        self.y_axis_config = y_axis_config
-
-        if background_line_style is None:
-            background_line_style = {
-                "stroke_color": BLUE_D,
-                "stroke_width": 2,
-                "stroke_opacity": 1,
-            }
-        self.background_line_style = background_line_style
+        self.axis_config = {
+            "stroke_color": WHITE,
+            "stroke_width": 2,
+            "include_ticks": False,
+            "include_tip": False,
+            "line_to_number_buff": SMALL_BUFF,
+            "label_direction": DR,
+            "number_scale_val": 0.5,
+        }
+        self.y_axis_config = {"label_direction": DR}
+        self.background_line_style = {
+            "stroke_color": BLUE_D,
+            "stroke_width": 2,
+            "stroke_opacity": 1,
+        }
+
+        self.update_default_configs(
+            (self.axis_config, self.y_axis_config, self.background_line_style),
+            (axis_config, y_axis_config, background_line_style),
+        )
 
         # Defaults to a faded version of line_config
         self.faded_line_style = faded_line_style

manim/mobject/number_line.py

@@ -94,10 +94,16 @@ def __init__(
             self.add_tip()
         if self.include_ticks:
             self.add_tick_marks()
-        self.rotate(self.rotation)
+        self.rotate_about_zero(self.rotation)
         if self.include_numbers:
             self.add_numbers()
 
+    def rotate_about_zero(self, angle, axis=OUT, **kwargs):
+        return self.rotate_about_number(0, angle, axis, **kwargs)
+
+    def rotate_about_number(self, number, angle, axis=OUT, **kwargs):
+        return self.rotate(angle, axis, about_point=self.n2p(number), **kwargs)
+
     def init_leftmost_tick(self):
         if self.leftmost_tick is None:
             self.leftmost_tick = op.mul(

manim/utils/config_ops.py

@@ -2,6 +2,7 @@
 
 __all__ = [
     "merge_dicts_recursively",
+    "update_dict_recursively",
     "DictAsObject",
 ]
 
@@ -29,6 +30,11 @@ def merge_dicts_recursively(*dicts):
     return result
 
 
+def update_dict_recursively(current_dict, *others):
+    updated_dict = merge_dicts_recursively(current_dict, *others)
+    current_dict.update(updated_dict)
+
+
 # Occasionally convenient in order to write dict.x instead of more laborious
 # (and less in keeping with all other attr accesses) dict["x"]
 

tests/control_data/graphical_units_data/coordinate_systems/NumberPlaneTest.npz

@@ -0,0 +1,27 @@
+import pytest
+
+from manim import *
+from ..utils.testing_utils import get_scenes_to_test
+from ..utils.GraphicalUnitTester import GraphicalUnitTester
+
+
+class NumberPlaneTest(Scene):
+    def construct(self):
+        plane = NumberPlane(
+            axis_config={"include_numbers": True, "x_min": -8},
+            x_min=-4,
+            x_max=6,
+            x_axis_config={"unit_size": 1.2},
+            y_min=-2,
+            y_axis_config={"x_max": 5, "width": 6, "label_direction": UL},
+            center_point=2 * DL,
+        )
+        self.play(Animation(plane))
+
+
+MODULE_NAME = "coordinate_systems"
+
+
+@pytest.mark.parametrize("scene_to_test", get_scenes_to_test(__name__), indirect=False)
+def test_scene(scene_to_test, tmpdir, show_diff):
+    GraphicalUnitTester(scene_to_test[1], MODULE_NAME, tmpdir).test(show_diff=show_diff)