Find x in this pyramid by Pythagorean theorem (3D drawing by matplotlib)

pyramid.py

"""Find x in this pyramid
by Pythagorean theorem

3D drawing by matplotlib"""

import matplotlib as mpl
from mpl_toolkits.mplot3d import Axes3D
import numpy as np
import matplotlib.pyplot as plt

mpl.rcParams['legend.fontsize'] = 10

fig = plt.figure()
ax = fig.add_subplot(111, projection='3d')
ax.plot([0,4],[0,0], [0,0], 'b')
ax.plot([0,4],[2,2], [0,0], 'b')
ax.plot([0,0],[0,2], [0,0], 'b')
ax.plot([4,4],[0,2], [0,0], 'b')

ax.plot([0,2],[0,1], [0,1], 'b')
ax.plot([4,2],[0,1], [0,1], 'b')
ax.plot([0,2],[2,1], [0,1], 'b')
ax.plot([4,2],[2,1], [0,1], 'b')
ax.text(3, 1.5, 0.5, 'x')

ax.plot([2,2],[1,1], [0,1], 'r')
ax.legend()

plt.show()