Specify the focal length for 3d plots

scottshambaugh · scottshambaugh · commit e8a3f57a9238 · 2022-01-04T13:17:48.000-07:00
Create a gallery example showing the different proj_type options

Typo

Test fix, example fix, and linting

What's new, example fix

Merge conflict

Offset zooming of focal length

Code review comments, consolidate projection functions

Try and fix zooming

Try to fix the focal length zooming

Update example

Cleanup

example cleanup

more example tweaks

more example tweak

swap plot order

Enforce a positive focal length

focal lentgh tests

flake8 linting

docstring tweak
diff --git a/doc/users/next_whats_new/3d_plot_focal_length.rst b/doc/users/next_whats_new/3d_plot_focal_length.rst
@@ -0,0 +1,27 @@
+Give the 3D camera a custom focal length
+----------------------------------------
+
+Users can now better mimic real-world cameras by specifying the focal length of
+the virtual camera in 3D plots. An increasing focal length between 1 and
++infinity "flattens" the image, while a decreasing focal length between 1 and 0
+exaggerates the perspective and gives the image more apparent depth. The
+default focal length of 1 corresponds to a Field of View (FOV) of 90 deg, and
+is backwards-compatible with existing 3D plots.
+
+The focal length can be calculated from a desired FOV via the equation:
+| ``focal_length = 1/tan(FOV/2)``
+
+.. plot::
+    :include-source: true
+
+    from mpl_toolkits.mplot3d import axes3d
+    import matplotlib.pyplot as plt
+    fig, axs = plt.subplots(1, 3, subplot_kw={'projection': '3d'})
+    X, Y, Z = axes3d.get_test_data(0.05)
+    focal_lengths = [0.25, 1, 4]
+    for ax, fl in zip(axs, focal_lengths):
+        ax.plot_wireframe(X, Y, Z, rstride=10, cstride=10)
+        ax.set_proj_type('persp', focal_length=fl)
+        ax.set_title(f"focal_length = {fl}")
+    plt.tight_layout()
+    plt.show()
diff --git a/examples/mplot3d/projections.py b/examples/mplot3d/projections.py
@@ -0,0 +1,48 @@
+"""
+========================
+3D plot projection types
+========================
+
+Demonstrates the different camera projections for 3D plots, and the effects of
+changing the focal length for a perspective projection. Note that matplotlib
+corrects for the 'zoom' effect of changing the focal length.
+
+An increasing focal length between 1 and +infinity "flattens" the image, while
+a decreasing focal length between 1 and 0 exaggerates the perspective and gives
+the image more apparent depth. The default focal length of 1 corresponds to a
+Field of View (FOV) of 90 deg. In the limiting case, a focal length of
++infinity corresponds to an orthographic projection after correction of the
+zoom effect.
+
+You can calculate focal length from a FOV via the equation:
+focal_length = 1/tan(FOV/2)
+
+Or vice versa:
+FOV = 2*atan(1/focal_length)
+"""
+
+from mpl_toolkits.mplot3d import axes3d
+import matplotlib.pyplot as plt
+
+
+fig, axs = plt.subplots(1, 3, subplot_kw={'projection': '3d'})
+
+# Get the test data
+X, Y, Z = axes3d.get_test_data(0.05)
+
+# Plot the data
+for ax in axs:
+    ax.plot_wireframe(X, Y, Z, rstride=10, cstride=10)
+
+# Set the orthographic projection.
+axs[0].set_proj_type('ortho')  # FOV = 0 deg
+axs[0].set_title("'ortho'\nfocal_length = +∞", fontsize=10)
+
+# Set the perspective projections
+axs[1].set_proj_type('persp')  # FOV = 90 deg
+axs[1].set_title("'persp'\nfocal_length = 1 (default)", fontsize=10)
+
+axs[2].set_proj_type('persp', focal_length=0.2)  # FOV = 157.4 deg
+axs[2].set_title("'persp'\nfocal_length = 0.2", fontsize=10)
+
+plt.show()
diff --git a/lib/mpl_toolkits/mplot3d/axes3d.py b/lib/mpl_toolkits/mplot3d/axes3d.py
@@ -56,7 +56,7 @@ class Axes3D(Axes):
     def __init__(
             self, fig, rect=None, *args,
             elev=30, azim=-60, roll=0, sharez=None, proj_type='persp',
-            box_aspect=None, computed_zorder=True,
+            box_aspect=None, computed_zorder=True, focal_length=None,
             **kwargs):
         """
         Parameters
@@ -103,6 +103,11 @@ def __init__(
             This behavior is deprecated in 3.4, the default will
             change to False in 3.5.  The keyword will be undocumented
             and a non-False value will be an error in 3.6.
+        focal_length : float, default: None
+            For a projection type of 'persp', the focal length of the virtual
+            camera. Must be > 0. If None, defaults to 1.
+            The focal length can be computed from a desired Field Of View via
+            the equation: focal_length = 1/tan(FOV/2)
 
         **kwargs
             Other optional keyword arguments:
@@ -116,7 +121,7 @@ def __init__(
         self.initial_azim = azim
         self.initial_elev = elev
         self.initial_roll = roll
-        self.set_proj_type(proj_type)
+        self.set_proj_type(proj_type, focal_length)
         self.computed_zorder = computed_zorder
 
         self.xy_viewLim = Bbox.unit()
@@ -1027,18 +1032,36 @@ def view_init(self, elev=None, azim=None, roll=None, vertical_axis="z"):
             dict(x=0, y=1, z=2), vertical_axis=vertical_axis
         )
 
-    def set_proj_type(self, proj_type):
+    def set_proj_type(self, proj_type, focal_length=None):
         """
         Set the projection type.
 
         Parameters
         ----------
         proj_type : {'persp', 'ortho'}
-        """
-        self._projection = _api.check_getitem({
-            'persp': proj3d.persp_transformation,
-            'ortho': proj3d.ortho_transformation,
-        }, proj_type=proj_type)
+            The projection type.
+        focal_length : float, default: None
+            For a projection type of 'persp', the focal length of the virtual
+            camera. Must be > 0. If None, defaults to 1.
+            The focal length can be computed from a desired Field Of View via
+            the equation: focal_length = 1/tan(FOV/2)
+        """
+        if proj_type == 'persp':
+            if focal_length is None:
+                self.focal_length = 1
+            else:
+                if focal_length <= 0:
+                    raise ValueError(f"focal_length = {focal_length} must be" +
+                                     " greater than 0")
+                self.focal_length = focal_length
+        elif proj_type == 'ortho':
+            if focal_length not in (None, np.inf):
+                raise ValueError(f"focal_length = {focal_length} must be" +
+                                 f"None for proj_type = {proj_type}")
+            self.focal_length = np.inf
+        else:
+            raise ValueError(f"proj_type = {proj_type} must be in" +
+                             f"{'persp', 'ortho'}")
 
     def _roll_to_vertical(self, arr):
         """Roll arrays to match the different vertical axis."""
@@ -1094,8 +1117,21 @@ def get_proj(self):
         V = np.zeros(3)
         V[self._vertical_axis] = -1 if abs(elev_rad) > 0.5 * np.pi else 1
 
-        viewM = proj3d.view_transformation(eye, R, V, roll_rad)
-        projM = self._projection(-self.dist, self.dist)
+        # Generate the view and projection transformation matrices
+        if self.focal_length == np.inf:
+            # Orthographic projection
+            viewM = proj3d.view_transformation(eye, R, V, roll_rad)
+            projM = proj3d.ortho_transformation(-self.dist, self.dist)
+        else:
+            # Perspective projection
+            # Scale the eye dist to compensate for the focal length zoom effect
+            eye_focal = R + self.dist * ps * self.focal_length
+            viewM = proj3d.view_transformation(eye_focal, R, V, roll_rad)
+            projM = proj3d.persp_transformation(-self.dist,
+                                                self.dist,
+                                                self.focal_length)
+
+        # Combine all the transformation matrices to get the final projection
         M0 = np.dot(viewM, worldM)
         M = np.dot(projM, M0)
         return M
@@ -1158,7 +1194,7 @@ def cla(self):
                 pass
 
         self._autoscaleZon = True
-        if self._projection is proj3d.ortho_transformation:
+        if self.focal_length == np.inf:
             self._zmargin = rcParams['axes.zmargin']
         else:
             self._zmargin = 0.
diff --git a/lib/mpl_toolkits/mplot3d/proj3d.py b/lib/mpl_toolkits/mplot3d/proj3d.py
@@ -90,23 +90,27 @@ def view_transformation(E, R, V, roll):
     return np.dot(Mr, Mt)
 
 
-def persp_transformation(zfront, zback):
-    a = (zfront+zback)/(zfront-zback)
-    b = -2*(zfront*zback)/(zfront-zback)
-    return np.array([[1, 0, 0, 0],
-                     [0, 1, 0, 0],
-                     [0, 0, a, b],
-                     [0, 0, -1, 0]])
+def persp_transformation(zfront, zback, focal_length):
+    e = focal_length
+    a = 1  # aspect ratio
+    b = (zfront+zback)/(zfront-zback)
+    c = -2*(zfront*zback)/(zfront-zback)
+    proj_matrix = np.array([[e,   0,  0, 0],
+                            [0, e/a,  0, 0],
+                            [0,   0,  b, c],
+                            [0,   0, -1, 0]])
+    return proj_matrix
 
 
 def ortho_transformation(zfront, zback):
     # note: w component in the resulting vector will be (zback-zfront), not 1
     a = -(zfront + zback)
     b = -(zfront - zback)
-    return np.array([[2, 0, 0, 0],
-                     [0, 2, 0, 0],
-                     [0, 0, -2, 0],
-                     [0, 0, a, b]])
+    proj_matrix = np.array([[2, 0,  0, 0],
+                            [0, 2,  0, 0],
+                            [0, 0, -2, 0],
+                            [0, 0,  a, b]])
+    return proj_matrix
 
 
 def _proj_transform_vec(vec, M):
diff --git a/lib/mpl_toolkits/tests/baseline_images/test_mplot3d/axes3d_focal_length.png b/lib/mpl_toolkits/tests/baseline_images/test_mplot3d/axes3d_focal_length.png
diff --git a/lib/mpl_toolkits/tests/test_mplot3d.py b/lib/mpl_toolkits/tests/test_mplot3d.py
@@ -872,7 +872,7 @@ def _test_proj_make_M():
     V = np.array([0, 0, 1])
     roll = 0
     viewM = proj3d.view_transformation(E, R, V, roll)
-    perspM = proj3d.persp_transformation(100, -100)
+    perspM = proj3d.persp_transformation(100, -100, 1)
     M = np.dot(perspM, viewM)
     return M
 
@@ -1037,6 +1037,13 @@ def test_unautoscale(axis, auto):
     np.testing.assert_array_equal(get_lim(), (-0.5, 0.5))
 
 
+@mpl3d_image_comparison(['axes3d_focal_length.png'], remove_text=False)
+def test_axes3d_focal_length():
+    fig, axs = plt.subplots(1, 2, subplot_kw={'projection': '3d'})
+    axs[0].set_proj_type('persp', focal_length=np.inf)
+    axs[1].set_proj_type('persp', focal_length=0.15)
+
+
 @mpl3d_image_comparison(['axes3d_ortho.png'], remove_text=False)
 def test_axes3d_ortho():
     fig = plt.figure()
