visualization.py

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


def cavity_vis(cavity_x, cavity_y, cavity_z, transmons):
    cylinder_radius = 5e-4

    fig = plt.figure()
    ax = fig.add_subplot(111, projection='3d')

    max_dim = max(cavity_x, cavity_y, cavity_z)
    max_dim *= 1.1

    # Plotting the cavity
    rect_prism = [(-cavity_x / 2, -cavity_y / 2, -cavity_z / 2), (cavity_x / 2, -cavity_y / 2, -cavity_z / 2),
                  (cavity_x / 2, cavity_y / 2, -cavity_z / 2), (-cavity_x / 2, cavity_y / 2, -cavity_z / 2),
                  (-cavity_x / 2, -cavity_y / 2, cavity_z / 2), (cavity_x / 2, -cavity_y / 2, cavity_z / 2),
                  (cavity_x / 2, cavity_y / 2, cavity_z / 2), (-cavity_x / 2, cavity_y / 2, cavity_z / 2)]

    r = np.array(rect_prism)
    ax.plot3D(*zip(*r[[0, 1, 2, 3, 0]]), color='b')  # Bottom rectangle
    ax.plot3D(*zip(*r[[4, 5, 6, 7, 4]]), color='b')  # Top rectangle
    for i in range(4):
        ax.plot3D(*zip(*r[[i, i + 4]]), color='b')   # Vertical lines

    # Plotting transmons
    for i in range(len(transmons)):
        x = transmons[i][0]
        y = transmons[i][1]
        z = transmons[i][2]
        theta = transmons[i][3]
        phi = transmons[i][4]
        length = transmons[i][5]
        ax.plot([x - length * np.sin(np.radians(phi)) * np.cos(np.radians(theta)) / 2, x + length * np.sin(np.radians(phi)) * np.cos(np.radians(theta)) / 2],
                [y - length * np.sin(np.radians(phi)) * np.sin(np.radians(theta)) / 2, y + length * np.sin(np.radians(phi)) * np.sin(np.radians(theta)) / 2],
                [z - length * np.cos(np.radians(phi)) / 2, z + length * np.cos(np.radians(phi)) / 2])

    ax.set_xlim(- max_dim / 2, max_dim / 2)
    ax.set_ylim(- max_dim / 2, max_dim / 2)
    ax.set_zlim(- max_dim / 2, max_dim / 2)
    ax.set_xlabel('X')
    ax.set_ylabel('Y')
    ax.set_zlabel('Z')

    plt.savefig('test_cavity_img.jpg', dpi=350)
    return


cavity_vis(0.02, 0.015, 0.01, [[0, 0, 0, 90, 90, 6e-3], [0.005, 0, 0, 45, 135, 6e-3], [-0.005, 0, 0, 0, 0, 6e-3]])