I have this program that plot two surfaces, my problem is that when I try to add a vector to the program it doesn't work.I tried to use quiver but my program returned two different plots, I want the vector and the surface in the same plot.
This is the code:
import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits import mplot3d
from sympy import symbols, diff, sqrt
# Mirror parameters
q = 0.5
# Light source coordinates
a, b, c = 0, 0, 2
# Z0 plane
Z0 = np.full((11, 11), 5)
# Define the domain
x = np.linspace(-1, 1, 11)
y = np.linspace(-1, 1, 11)
# Create a meshgrid for X and Y coordinates
X, Y = np.meshgrid(x, y)
# Equation for the mirror surface (Z)
Z = q * (X ** 2 + Y ** 2) / (1 + np.sqrt(1 - q ** 2 * (X ** 2 + Y ** 2)))
# Partial derivatives of Z with respect to X and Y
Zx = (1.0 * X * q ** 3 * (X ** 2 + Y ** 2) /
((-q ** 2 * (X ** 2 + Y ** 2) + 1) ** 0.5 * ((-q ** 2 * (X ** 2 + Y ** 2) + 1) ** 0.5 + 1) ** 2) +
2 * X * q / ((-q ** 2 * (X ** 2 + Y ** 2) + 1) ** 0.5 + 1))
Zy = (1.0 * Y * q ** 3 * (X ** 2 + Y ** 2) /
((-q ** 2 * (X ** 2 + Y ** 2) + 1) ** 0.5 * ((-q ** 2 * (X ** 2 + Y ** 2) + 1) ** 0.5 + 1) ** 2) +
2 * Y * q / ((-q ** 2 * (X ** 2 + Y ** 2) + 1) ** 0.5 + 1))
# Calculate transformed coordinates (x0, y0) based on mirror equation
x0 = X + (Z0 - Z) * ((X - a) * (1 - Zx ** 2 + Zy ** 2) - 2 * Zx * (Zy * (Y - b) + (c - Z))) / (
(c - Z) * (1 - Zx ** 2 - Zy ** 2) + 2 * (Zx * (X - a) + Zy * (Y - b)))
y0 = Y + (Z0 - Z) * ((Y - b) * (1 + Zx ** 2 - Zy ** 2) - 2 * Zy * (Zx * (X - a) + (c - Z))) / (
(c - Z) * (1 - Zx ** 2 - Zy ** 2) + 2 * (Zx * (X - a) + Zy * (Y - b)))
# Create a meshgrid for the plane Z0
X0, Y0 = np.meshgrid(x0[0, :], y0[:, 0])
# Create a 3D plot
fig = plt.figure()
ax = plt.axes(projection='3d')
# Plot the mirror surface and the Z0 plane
ax.plot_surface(X, Y, Z, cmap="plasma", linewidth=0, antialiased=False, alpha=0.5)
ax.plot_surface(X0, Y0, Z0, cmap="viridis", linewidth=0, antialiased=False, alpha=0.5)
# Set labels for axes
ax.set_xlabel('x')
ax.set_ylabel('y')
ax.set_zlabel('z')
# Set the view angle for better visualization
ax.view_init(30, 60)
# Save and display the plot
plt.savefig('3Dplot1.png', dpi=600)
plt.show()