Source code for qutip.bloch3d

__all__ = ['Bloch3d']

import numpy as np
from qutip.qobj import Qobj
from qutip.expect import expect
from qutip.operators import sigmax, sigmay, sigmaz


[docs]class Bloch3d: """Class for plotting data on a 3D Bloch sphere using mayavi. Valid data can be either points, vectors, or qobj objects corresponding to state vectors or density matrices. for a two-state system (or subsystem). Attributes ---------- fig : instance {None} User supplied Matplotlib Figure instance for plotting Bloch sphere. font_color : str {'black'} Color of font used for Bloch sphere labels. font_scale : float {0.08} Scale for font used for Bloch sphere labels. frame : bool {True} Draw frame for Bloch sphere frame_alpha : float {0.05} Sets transparency of Bloch sphere frame. frame_color : str {'gray'} Color of sphere wireframe. frame_num : int {8} Number of frame elements to draw. frame_radius : floats {0.005} Width of wireframe. point_color : list {['r', 'g', 'b', 'y']} List of colors for Bloch sphere point markers to cycle through. i.e. By default, points 0 and 4 will both be blue ('r'). point_mode : string {'sphere','cone','cube','cylinder','point'} Point marker shapes. point_size : float {0.075} Size of points on Bloch sphere. sphere_alpha : float {0.1} Transparency of Bloch sphere itself. sphere_color : str {'#808080'} Color of Bloch sphere. size : list {[500,500]} Size of Bloch sphere plot in pixels. Best to have both numbers the same otherwise you will have a Bloch sphere that looks like a football. vector_color : list {['r', 'g', 'b', 'y']} List of vector colors to cycle through. vector_width : int {3} Width of displayed vectors. view : list {[45,65]} Azimuthal and Elevation viewing angles. xlabel : list {``['|x>', '']``} List of strings corresponding to +x and -x axes labels, respectively. xlpos : list {[1.07,-1.07]} Positions of +x and -x labels respectively. ylabel : list {``['|y>', '']``} List of strings corresponding to +y and -y axes labels, respectively. ylpos : list {[1.07,-1.07]} Positions of +y and -y labels respectively. zlabel : list {``['|0>', '|1>']``} List of strings corresponding to +z and -z axes labels, respectively. zlpos : list {[1.07,-1.07]} Positions of +z and -z labels respectively. Notes ----- The use of mayavi for 3D rendering of the Bloch sphere comes with a few limitations: I) You can not embed a Bloch3d figure into a matplotlib window. II) The use of LaTex is not supported by the mayavi rendering engine. Therefore all labels must be defined using standard text. Of course you can post-process the generated figures later to add LaTeX using other software if needed. """ def __init__(self, fig=None): # ----check for mayavi----- try: from mayavi import mlab except: raise Exception("This function requires the mayavi module.") # ---Image options--- self.fig = None self.user_fig = None # check if user specified figure or axes. if fig: self.user_fig = fig # The size of the figure in inches, default = [500,500]. self.size = [500, 500] # Azimuthal and Elvation viewing angles, default = [45,65]. self.view = [45, 65] # Image background color self.bgcolor = 'white' # Image foreground color. Other options can override. self.fgcolor = 'black' # ---Sphere options--- # Color of Bloch sphere, default = #808080 self.sphere_color = '#808080' # Transparency of Bloch sphere, default = 0.1 self.sphere_alpha = 0.1 # ---Frame options--- # Draw frame? self.frame = True # number of lines to draw for frame self.frame_num = 8 # Color of wireframe, default = 'gray' self.frame_color = 'black' # Transparency of wireframe, default = 0.2 self.frame_alpha = 0.05 # Radius of frame lines self.frame_radius = 0.005 # --Axes--- # Axes color self.axes_color = 'black' # Transparency of axes self.axes_alpha = 0.4 # Radius of axes lines self.axes_radius = 0.005 # ---Labels--- # Labels for x-axis (in LaTex), default = ['$x$',''] self.xlabel = ['|x>', ''] # Position of x-axis labels, default = [1.2,-1.2] self.xlpos = [1.07, -1.07] # Labels for y-axis (in LaTex), default = ['$y$',''] self.ylabel = ['|y>', ''] # Position of y-axis labels, default = [1.1,-1.1] self.ylpos = [1.07, -1.07] # Labels for z-axis self.zlabel = ['|0>', '|1>'] # Position of z-axis labels, default = [1.05,-1.05] self.zlpos = [1.07, -1.07] # ---Font options--- # Color of fonts, default = 'black' self.font_color = 'black' # Size of fonts, default = 20 self.font_scale = 0.08 # ---Vector options--- # Object used for representing vectors on Bloch sphere. # List of colors for Bloch vectors, default = ['b','g','r','y'] self.vector_color = ['r', 'g', 'b', 'y'] # Width of Bloch vectors, default = 2 self.vector_width = 2.0 # Height of vector head self.vector_head_height = 0.15 # Radius of vector head self.vector_head_radius = 0.075 # ---Point options--- # List of colors for Bloch point markers, default = ['b','g','r','y'] self.point_color = ['r', 'g', 'b', 'y'] # Size of point markers self.point_size = 0.06 # Shape of point markers # Options: 'cone' or 'cube' or 'cylinder' or 'point' or 'sphere'. # Default = 'sphere' self.point_mode = 'sphere' # ---Data lists--- # Data for point markers self.points = [] # Data for Bloch vectors self.vectors = [] # Number of times sphere has been saved self.savenum = 0 # Style of points, 'm' for multiple colors, 's' for single color self.point_style = [] # Transparency of points self.point_alpha = [] # Transparency of vectors self.vector_alpha = [] def __str__(self): s = "" s += "Bloch3D data:\n" s += "-----------\n" s += "Number of points: " + str(len(self.points)) + "\n" s += "Number of vectors: " + str(len(self.vectors)) + "\n" s += "\n" s += "Bloch3D sphere properties:\n" s += "--------------------------\n" s += "axes_alpha: " + str(self.axes_alpha) + "\n" s += "axes_color: " + str(self.axes_color) + "\n" s += "axes_radius: " + str(self.axes_radius) + "\n" s += "bgcolor: " + str(self.bgcolor) + "\n" s += "fgcolor: " + str(self.fgcolor) + "\n" s += "font_color: " + str(self.font_color) + "\n" s += "font_scale: " + str(self.font_scale) + "\n" s += "frame: " + str(self.frame) + "\n" s += "frame_alpha: " + str(self.frame_alpha) + "\n" s += "frame_color: " + str(self.frame_color) + "\n" s += "frame_num: " + str(self.frame_num) + "\n" s += "frame_radius: " + str(self.frame_radius) + "\n" s += "point_color: " + str(self.point_color) + "\n" s += "point_mode: " + str(self.point_mode) + "\n" s += "point_size: " + str(self.point_size) + "\n" s += "sphere_alpha: " + str(self.sphere_alpha) + "\n" s += "sphere_color: " + str(self.sphere_color) + "\n" s += "size: " + str(self.size) + "\n" s += "vector_color: " + str(self.vector_color) + "\n" s += "vector_width: " + str(self.vector_width) + "\n" s += "vector_head_height: " + str(self.vector_head_height) + "\n" s += "vector_head_radius: " + str(self.vector_head_radius) + "\n" s += "view: " + str(self.view) + "\n" s += "xlabel: " + str(self.xlabel) + "\n" s += "xlpos: " + str(self.xlpos) + "\n" s += "ylabel: " + str(self.ylabel) + "\n" s += "ylpos: " + str(self.ylpos) + "\n" s += "zlabel: " + str(self.zlabel) + "\n" s += "zlpos: " + str(self.zlpos) + "\n" return s
[docs] def clear(self): """Resets the Bloch sphere data sets to empty. """ self.points = [] self.vectors = [] self.point_style = []
[docs] def add_points(self, points, meth='s', alpha=1.0): """Add a list of data points to bloch sphere. Parameters ---------- points : array/list Collection of data points. meth : str {'s','m'} Type of points to plot, use 'm' for multicolored. alpha : float, default=1. Transparency value for the vectors. Values between 0 and 1. """ if not isinstance(points[0], (list, np.ndarray)): points = [[points[0]], [points[1]], [points[2]]] points = np.array(points) if meth == 's': if len(points[0]) == 1: pnts = np.array( [[points[0][0]], [points[1][0]], [points[2][0]]]) pnts = np.append(pnts, points, axis=1) else: pnts = points self.points.append(pnts) self.point_style.append('s') else: self.points.append(points) self.point_style.append('m') self.point_alpha.append(alpha)
[docs] def add_states(self, state, kind='vector', alpha=1.0): """Add a state vector Qobj to Bloch sphere. Parameters ---------- state : qobj Input state vector. kind : str {'vector','point'} Type of object to plot. alpha : float, default=1. Transparency value for the vectors. Values between 0 and 1. """ if isinstance(state, Qobj): state = [state] for st in state: if kind == 'vector': vec = [expect(sigmax(), st), expect(sigmay(), st), expect(sigmaz(), st)] self.add_vectors(vec, alpha=alpha) elif kind == 'point': pnt = [expect(sigmax(), st), expect(sigmay(), st), expect(sigmaz(), st)] self.add_points(pnt, alpha=alpha)
[docs] def add_vectors(self, vectors, alpha=1.0): """Add a list of vectors to Bloch sphere. Parameters ---------- vectors : array/list Array with vectors of unit length or smaller. alpha : float, default=1. Transparency value for the vectors. Values between 0 and 1. """ if isinstance(vectors[0], (list, np.ndarray)): for vec in vectors: self.vectors.append(vec) self.vector_alpha.append(alpha) else: self.vectors.append(vectors) self.vector_alpha.append(alpha)
[docs] def plot_vectors(self): """ Plots vectors on the Bloch sphere. """ from mayavi import mlab from tvtk.api import tvtk import matplotlib.colors as colors ii = 0 for k in range(len(self.vectors)): vec = np.array(self.vectors[k]) norm = np.linalg.norm(vec) theta = np.arccos(vec[2] / norm) phi = np.arctan2(vec[1], vec[0]) vec -= 0.5 * self.vector_head_height * \ np.array([np.sin(theta) * np.cos(phi), np.sin(theta) * np.sin(phi), np.cos(theta)]) color = colors.colorConverter.to_rgb( self.vector_color[np.mod(k, len(self.vector_color))]) mlab.plot3d([0, vec[0]], [0, vec[1]], [0, vec[2]], name='vector' + str(ii), tube_sides=100, line_width=self.vector_width, opacity=self.vector_alpha[k], color=color) cone = tvtk.ConeSource(height=self.vector_head_height, radius=self.vector_head_radius, resolution=100) cone_mapper = tvtk.PolyDataMapper( input_connection=cone.output_port) prop = tvtk.Property(opacity=self.vector_alpha[k], color=color) cc = tvtk.Actor(mapper=cone_mapper, property=prop) cc.rotate_z(np.degrees(phi)) cc.rotate_y(-90 + np.degrees(theta)) cc.position = vec self.fig.scene.add_actor(cc)
[docs] def plot_points(self): """ Plots points on the Bloch sphere. """ from mayavi import mlab import matplotlib.colors as colors for k in range(len(self.points)): num = len(self.points[k][0]) dist = [np.sqrt(self.points[k][0][j] ** 2 + self.points[k][1][j] ** 2 + self.points[k][2][j] ** 2) for j in range(num)] if any(abs(dist - dist[0]) / dist[0] > 1e-12): # combine arrays so that they can be sorted together # and sort rates from lowest to highest zipped = sorted(zip(dist, range(num))) dist, indperm = zip(*zipped) indperm = np.array(indperm) else: indperm = range(num) if self.point_style[k] == 's': color = colors.colorConverter.to_rgb( self.point_color[np.mod(k, len(self.point_color))]) mlab.points3d( self.points[k][0][indperm], self.points[k][1][indperm], self.points[k][2][indperm], figure=self.fig, resolution=100, scale_factor=self.point_size, mode=self.point_mode, color=color, opacity=self.point_alpha[k]) elif self.point_style[k] == 'm': pnt_colors = np.array(self.point_color * np.ceil( num / float(len(self.point_color)))) pnt_colors = pnt_colors[0:num] pnt_colors = list(pnt_colors[indperm]) for kk in range(num): mlab.points3d( self.points[k][0][ indperm[kk]], self.points[k][1][indperm[kk]], self.points[k][2][ indperm[kk]], figure=self.fig, resolution=100, scale_factor=self.point_size, mode=self.point_mode, color=colors.colorConverter.to_rgb(pnt_colors[kk]), opacity=self.point_alpha[k])
[docs] def make_sphere(self): """ Plots Bloch sphere and data sets. """ # setup plot # Figure instance for Bloch sphere plot from mayavi import mlab import matplotlib.colors as colors if self.user_fig: self.fig = self.user_fig else: self.fig = mlab.figure( 1, size=self.size, bgcolor=colors.colorConverter.to_rgb(self.bgcolor), fgcolor=colors.colorConverter.to_rgb(self.fgcolor)) sphere = mlab.points3d( 0, 0, 0, figure=self.fig, scale_mode='none', scale_factor=2, color=colors.colorConverter.to_rgb(self.sphere_color), resolution=100, opacity=self.sphere_alpha, name='bloch_sphere') # Thse commands make the sphere look better sphere.actor.property.specular = 0.45 sphere.actor.property.specular_power = 5 sphere.actor.property.backface_culling = True # make frame for sphere surface if self.frame: theta = np.linspace(0, 2 * np.pi, 100) for angle in np.linspace(-np.pi, np.pi, self.frame_num): xlat = np.cos(theta) * np.cos(angle) ylat = np.sin(theta) * np.cos(angle) zlat = np.ones_like(theta) * np.sin(angle) xlon = np.sin(angle) * np.sin(theta) ylon = np.cos(angle) * np.sin(theta) zlon = np.cos(theta) mlab.plot3d( xlat, ylat, zlat, color=colors.colorConverter.to_rgb(self.frame_color), opacity=self.frame_alpha, tube_radius=self.frame_radius) mlab.plot3d( xlon, ylon, zlon, color=colors.colorConverter.to_rgb(self.frame_color), opacity=self.frame_alpha, tube_radius=self.frame_radius) # add axes axis = np.linspace(-1.0, 1.0, 10) other = np.zeros_like(axis) mlab.plot3d( axis, other, other, color=colors.colorConverter.to_rgb(self.axes_color), tube_radius=self.axes_radius, opacity=self.axes_alpha) mlab.plot3d( other, axis, other, color=colors.colorConverter.to_rgb(self.axes_color), tube_radius=self.axes_radius, opacity=self.axes_alpha) mlab.plot3d( other, other, axis, color=colors.colorConverter.to_rgb(self.axes_color), tube_radius=self.axes_radius, opacity=self.axes_alpha) # add data to sphere self.plot_points() self.plot_vectors() # #add labels mlab.text3d(0, 0, self.zlpos[0], self.zlabel[0], color=colors.colorConverter.to_rgb(self.font_color), scale=self.font_scale) mlab.text3d(0, 0, self.zlpos[1], self.zlabel[1], color=colors.colorConverter.to_rgb(self.font_color), scale=self.font_scale) mlab.text3d(self.xlpos[0], 0, 0, self.xlabel[0], color=colors.colorConverter.to_rgb(self.font_color), scale=self.font_scale) mlab.text3d(self.xlpos[1], 0, 0, self.xlabel[1], color=colors.colorConverter.to_rgb(self.font_color), scale=self.font_scale) mlab.text3d(0, self.ylpos[0], 0, self.ylabel[0], color=colors.colorConverter.to_rgb(self.font_color), scale=self.font_scale) mlab.text3d(0, self.ylpos[1], 0, self.ylabel[1], color=colors.colorConverter.to_rgb(self.font_color), scale=self.font_scale)
[docs] def show(self): """ Display the Bloch sphere and corresponding data sets. """ from mayavi import mlab self.make_sphere() mlab.view(azimuth=self.view[0], elevation=self.view[1], distance=5) if self.fig: mlab.show()
[docs] def save(self, name=None, format='png', dirc=None): """Saves Bloch sphere to file of type ``format`` in directory ``dirc``. Parameters ---------- name : str Name of saved image. Must include path and format as well. i.e. '/Users/Paul/Desktop/bloch.png' This overrides the 'format' and 'dirc' arguments. format : str Format of output image. Default is 'png'. dirc : str Directory for output images. Defaults to current working directory. Returns ------- File containing plot of Bloch sphere. """ from mayavi import mlab import os self.make_sphere() mlab.view(azimuth=self.view[0], elevation=self.view[1], distance=5) if dirc: if not os.path.isdir(os.getcwd() + "/" + str(dirc)): os.makedirs(os.getcwd() + "/" + str(dirc)) if name is None: if dirc: mlab.savefig(os.getcwd() + "/" + str(dirc) + '/bloch_' + str(self.savenum) + '.' + format) else: mlab.savefig(os.getcwd() + '/bloch_' + str(self.savenum) + '.' + format) else: mlab.savefig(name) self.savenum += 1 if self.fig: mlab.close(self.fig)