def construct_eigvect_DOS_plot(xy, fig, DOS_ax, eig_ax, eigval, eigvect, en, sim_type, Ni, Nk, marker_num=0, color_scheme='default', sub_lattice=-1, cmap_lines='BlueBlackRed', line_climv=None, cmap_patches='isolum_rainbow', draw_strain=False, lw=3, bondval_matrix=None, dotsz=.04, normalization=0.9, xycolor=lecmaps.green(), wzcolor=lecmaps.yellow()): """puts together lattice and DOS plots and draws normal mode ellipsoids on top Parameters ---------- xy: array 2N x 3 Equilibrium position of the gyroscopes fig : figure with lattice and DOS drawn DOS_ax: axis for the DOS plot eig_ax axis for the eigenvalue plot eigval : array of dimension 2nx1 Eigenvalues of matrix for system eigvect : array of dimension 2nx2n Eigenvectors of matrix for system. Eigvect is stored as NModes x NP*2 array, with x and y components alternating, like: x0, y0, x1, y1, ... xNP, yNP. en: int Number of the eigenvalue you are plotting marker_num : int the index of the 'timestep' (or phase of the eigvect) at which we place the t=0 line/dot for each particle color_scheme : str (default='default') sub_lattice : int cmap_lines : str line_climv : tuple of floats or None cmap_patches : str draw_strain : bool lw : float bondval_matrix : (NP x max#NN) float array or None a color specification for the bonds in the network Returns ---------- fig : completed figure for normal mode [scat_fg, p, f_mark] : things to be cleared before next normal mode is drawn """ # ppu = leplt.get_points_per_unit() s = leplt.absolute_sizer() plt.sca(DOS_ax) ev = eigval[en] ev1 = ev # Show where current eigenvalue is in DOS plot (f_mark, ) = plt.plot([np.real(ev), np.real(ev)], plt.ylim(), '-r') NP = len(xy[:, 0]) im1 = np.imag(ev) re1 = np.real(ev) plt.sca(eig_ax) plt.title('Mode %d; $\Omega=( %0.6f + %0.6f i)$' % (en, re1, im1)) # Preallocate ellipsoid plot vars angles_arr = np.zeros(NP) tangles_arr = np.zeros(NP) # patch = [] polygons, tiltgons = [], [] colors = np.zeros(2 * NP + 2) # x_mag = np.zeros(NP) # y_mag = np.zeros(NP) x0s, y0s = np.zeros(NP), np.zeros(NP) w0s, z0s = np.zeros(NP), np.zeros(NP) mag1 = eigvect[en] # Eigvect is stored as NModes x NP*2 array, with x and y components alternating, like: # x0, y0, x1, y1, ... xNP, yNP ... w0, z0, ... wNP, zNP. mag1x = np.array([mag1[2 * i] for i in range(NP)]) mag1y = np.array([mag1[2 * i + 1] for i in range(NP)]) mag1w = np.array([mag1[2 * i] for i in np.arange(NP, 2 * NP)]) mag1z = np.array([mag1[2 * i + 1] for i in np.arange(NP, 2 * NP)]) # Pick a series of times to draw out the ellipsoids if abs(ev1) > 0: time_arr = np.arange(21) * 2 * np.pi / (np.abs(ev1) * 20) exp1 = np.exp(1j * ev1 * time_arr) else: time_arr = np.arange(21) * 2 * np.pi / 20 exp1 = np.exp(1j * time_arr) # Normalization for the ellipsoids lim_mag1 = np.max( np.array([ np.sqrt( np.abs(exp1 * mag1x[i])**2 + np.abs(exp1 * mag1y[i])**2 + np.abs(exp1 * mag1w[i])**2 + np.abs(exp1 * mag1z[i])**2) for i in range(len(mag1x)) ]).flatten()) if np.isnan(lim_mag1): print 'found nan for limiting magnitude, replacing lim_mag1 with 1.0' lim_mag1 = 1. mag1x *= normalization / lim_mag1 mag1y *= normalization / lim_mag1 mag1w *= normalization / lim_mag1 mag1z *= normalization / lim_mag1 # sys.exit() cw = [] ccw = [] lines_stretch = [] lines_twist = [] for i in range(NP): # Draw COM movement of each node as an ellipse x_disps = 0.5 * (exp1 * mag1x[i]).real y_disps = 0.5 * (exp1 * mag1y[i]).real x_vals = xy[i, 0] + x_disps y_vals = xy[i, 1] + y_disps # Draw movement of the orientation of each node as an ellipse w_disps = 0.5 * (exp1 * mag1w[i]).real z_disps = 0.5 * (exp1 * mag1z[i]).real w_vals = xy[i, 0] + w_disps z_vals = xy[i, 1] + z_disps poly_points = np.array([x_vals, y_vals]).T tilt_points = np.array([w_vals, z_vals]).T # polygon = Polygon(poly_points, True, lw=lw, ec='g') # tiltgon = Polygon(tilt_points, True, lw=lw, ec='r') # polygon = plt.plot(poly_points[:, 0], poly_points[:, 1], 'g-', lw=lw) # tiltgon = plt.plot(tilt_points[:, 0], tilt_points[:, 1], 'r-', lw=lw) npolypts = len(poly_points[:, 0]) for ii in range(npolypts): lines_s = [[ poly_points[ii, 0], poly_points[(ii + 1) % npolypts, 0] ], [poly_points[ii, 1], poly_points[(ii + 1) % npolypts, 1]]] lines_t = [[ tilt_points[ii, 0], tilt_points[(ii + 1) % npolypts, 0] ], [tilt_points[ii, 1], tilt_points[(ii + 1) % npolypts, 1]]] lines_stretch.append(lines_s) lines_twist.append(lines_t) # x0 is the marker_num^th element of x_disps x0 = x_disps[marker_num] y0 = y_disps[marker_num] w0 = w_disps[marker_num] z0 = z_disps[marker_num] # x0s is the position (global pos, not relative) of each gyro at time = marker_num(out of 81) x0s[i] = x_vals[marker_num] y0s[i] = y_vals[marker_num] w0s[i] = w_vals[marker_num] z0s[i] = z_vals[marker_num] # Get angle for xy mag = np.sqrt(x0**2 + y0**2) if mag > 0: anglez = np.arccos(x0 / mag) else: anglez = 0 if y0 < 0: anglez = 2 * np.pi - anglez # Get angle for wz tmag = np.sqrt(w0**2 + z0**2) if tmag > 0: tanglez = np.arccos(w0 / tmag) else: tanglez = 0 if y0 < 0: tanglez = 2 * np.pi - tanglez angles_arr[i] = anglez tangles_arr[i] = tanglez # polygons.append(polygon) # tiltgons.append(tiltgon) # Do Fast Fourier Transform (FFT) # ff = abs(fft.fft(x_disps + 1j*y_disps))**2 # ff_freq = fft.fftfreq(len(x_vals), 1) # mm_f = ff_freq[ff == max(ff)][0] if color_scheme == 'default': colors[2 * i] = anglez colors[2 * i + 1] = anglez # add two more colors to ensure clims go from 0 to 2pi, I think... colors[2 * NP] = 0 colors[2 * NP + 1] = 2 * np.pi plt.yticks([]) plt.xticks([]) # this is the part that puts a dot a t=0 point scat_fg = eig_ax.scatter(x0s, y0s, s=s(dotsz), edgecolors='k', facecolors='none') scat_fg2 = eig_ax.scatter(w0s, z0s, s=s(dotsz), edgecolors='gray', facecolors='none') try: NN = np.shape(Ni)[1] except IndexError: NN = 0 Rnorm = np.array([x0s, y0s]).T # Bond Stretches if draw_strain: inc = 0 stretches = np.zeros(4 * len(xy)) for i in range(len(xy)): if NN > 0: for j, k in zip(Ni[i], Nk[i]): if i < j and abs(k) > 0: n1 = float(np.linalg.norm(Rnorm[i] - Rnorm[j])) n2 = np.linalg.norm(xy[i] - xy[j]) stretches[inc] = (n1 - n2) inc += 1 stretch = np.array(stretches[0:inc]) else: # simply get length of BL (len(BL) = inc) by iterating over all bondsssa inc = 0 for i in range(len(xy)): if NN > 0: for j, k in zip(Ni[i], Nk[i]): if i < j and abs(k) > 0: inc += 1 # For particles with neighbors, get list of bonds to draw. # If bondval_matrix is not None, color by the elements of that matrix if bondval_matrix is not None or draw_strain: test = list(np.zeros([inc, 1])) bondvals = list(np.ones([inc, 1])) inc = 0 xy = np.array([x0s, y0s]).T for i in range(len(xy)): if NN > 0: for j, k in zip(Ni[i], Nk[i]): if i < j and abs(k) > 0: test[inc] = [xy[(i, j), 0], xy[(i, j), 1]] if bondval_matrix is not None: bondvals[inc] = bondval_matrix[i, j] inc += 1 # lines connect sites (bonds), while lines_12 draw the black lines from the pinning to location sites lines = [zip(x, y) for x, y in test] # Check that we have all the cmaps if cmap_lines not in plt.colormaps() or cmap_patches not in plt.colormaps( ): lecmaps.register_cmaps() # Add lines colored by strain here if bondval_matrix is not None: lines_st = LineCollection(lines, array=bondvals, cmap=cmap_lines, linewidth=0.8) if line_climv is None: maxk = np.max(np.abs(bondvals)) mink = np.min(np.abs(bondvals)) if (bondvals - bondvals[0] < 1e-8).all(): lines_st.set_clim([mink - 1., maxk + 1.]) else: lines_st.set_clim([mink, maxk]) lines_st.set_zorder(2) eig_ax.add_collection(lines_st) else: if draw_strain: lines_st = LineCollection(lines, array=stretch, cmap=cmap_lines, linewidth=0.8) if line_climv is None: maxstretch = np.max(np.abs(stretch)) if maxstretch < 1e-8: line_climv = 1.0 else: line_climv = maxstretch lines_st.set_clim([-line_climv, line_climv]) lines_st.set_zorder(2) eig_ax.add_collection(lines_st) # Draw lines for movement lines_stretch = [zip(x, y) for x, y in lines_stretch] polygons = LineCollection(lines_stretch, linewidth=lw, linestyles='solid', color=xycolor) polygons.set_zorder(1) eig_ax.add_collection(polygons) lines_twist = [zip(x, y) for x, y in lines_twist] tiltgons = LineCollection(lines_twist, linewidth=lw, linestyles='solid', color=wzcolor) tiltgons.set_zorder(1) eig_ax.add_collection(tiltgons) # Draw polygons # polygons = PatchCollection(polygons, cmap=cmap_patches, alpha=0.6) # # p.set_array(np.array(colors)) # polygons.set_clim([0, 2 * np.pi]) # polygons.set_zorder(1) # eig_ax.add_collection(polygons) # tiltgons = PatchCollection(tiltgons, cmap=cmap_patches, alpha=0.6) # # p.set_array(np.array(colors)) # tiltgons.set_clim([0, 2 * np.pi]) # tiltgons.set_zorder(1000) # eig_ax.add_collection(tiltgons) eig_ax.set_aspect('equal') # erased ev/(2*pi) here npm 2016 cw_ccw = [cw, ccw, ev] # print cw_ccw[1] # If on a virtualenv, check it here # if not hasattr(sys, 'real_prefix'): # plt.show() # eig_ax.set_facecolor('#000000') # print 'leplt: construct_eigvect_DOS_plot() exiting' return fig, [scat_fg, scat_fg2, f_mark, polygons, tiltgons], cw_ccw
def plot_eigvect_excitation(xy, fig, dos_ax, eig_ax, eigval, eigvect, en, marker_num=0, draw_strain=False, black_t0lines=False, mark_t0=True, title='auto', normalization=1., alpha=0.6, lw=1, zorder=10, cmap='isolum_rainbow', color_scheme='default', color='phase', theta=None, t0_ptsize=.02, bondval_matrix=None, dotsz=.04, xycolor=lecmaps.green(), wzcolor=lecmaps.yellow()): """Draws normal mode ellipsoids on axis eig_ax. If black_t0lines is true, draws the black line from pinning site to positions Parameters ---------- xy: array 2N x 3 Equilibrium position of the gyroscopes fig : figure with lattice and DOS drawn dos_ax: matplotlib axis instance or None axis for the DOS plot. If None, ignores this input eig_ax : matplotlib axis instance axis for the eigenvalue plot eigval : array of dimension 2nx1 Eigenvalues of matrix for system eigvect : array of dimension 2nx2n Eigenvectors of matrix for system. Eigvect is stored as NModes x NP*2 array, with x and y components alternating, like: x0, y0, x1, y1, ... xNP, yNP. en: int Number of the eigenvalue you are plotting marker_num : int in (0, 80) where in the phase (0 to 80) to call t=t0. This sets "now" for drawing where in the normal mode to draw black_t0lines : bool Draw black lines extending from the pinning site to the current site (where 'current' is determined by marker_num) color : str ('phase' or 'displacement') Whether to color the excitations by their instantaneous (relative) phase, or by the magnitude of the mean displacement theta : float or None angle by which to rotate the initial displacement plotted by black lines, dots, etc bondval_matrix : (NP x max#NN) float array or None a color specification for the bonds in the network Returns ---------- fig : matplotlib figure instance completed figure for normal mode [scat_fg, pp, f_mark, lines12_st] : things to be cleared before next normal mode is drawn """ # ensure that colormap is registered lecmaps.register_cmap(cmap) # ppu = get_points_per_unit() s = leplt.absolute_sizer() ev = eigval[en] # Show where current eigenvalue is in DOS plot if dos_ax is not None: (f_mark, ) = dos_ax.plot([abs(ev), abs(ev)], dos_ax.get_ylim(), '-r') else: f_mark = None NP = len(xy) print 'twisty.plotting.plotting: NP = ', NP print 'twisty.plotting.plotting: np.shape(xy) = ', np.shape(xy) re1 = np.real(ev) plt.sca(eig_ax) if title == 'auto': eig_ax.set_title('$\omega = %0.6f$' % re1) elif title is not None and title not in ['', 'none']: eig_ax.set_title(title) # Preallocate ellipsoid plot vars angles_arr = np.zeros(NP) tangles_arr = np.zeros(NP) # patch = [] polygons, tiltgons = [], [] colors = np.zeros(2 * NP + 2) # x_mag = np.zeros(NP) # y_mag = np.zeros(NP) x0s, y0s = np.zeros(NP), np.zeros(NP) w0s, z0s = np.zeros(NP), np.zeros(NP) mag1 = eigvect[en] print '' # Eigvect is stored as NModes x NP*2 array, with x and y components alternating, like: # x0, y0, x1, y1, ... xNP, yNP ... w0, z0, ... wNP, zNP. mag1x = np.array([mag1[2 * i] for i in range(NP)]) mag1y = np.array([mag1[2 * i + 1] for i in range(NP)]) mag1w = np.array([mag1[2 * i] for i in np.arange(NP, 2 * NP)]) mag1z = np.array([mag1[2 * i + 1] for i in np.arange(NP, 2 * NP)]) # Pick a series of times to draw out the ellipsoids if abs(ev) > 0: time_arr = np.arange(21) * 2 * np.pi / (np.abs(ev) * 20) exp1 = np.exp(1j * ev * time_arr) else: time_arr = np.arange(21) * 2 * np.pi / 20 exp1 = np.exp(1j * time_arr) # Normalization for the ellipsoids lim_mag1 = np.max( np.array([ np.sqrt( np.abs(exp1 * mag1x[i])**2 + np.abs(exp1 * mag1y[i])**2 + np.abs(exp1 * mag1w[i])**2 + np.abs(exp1 * mag1z[i])**2) for i in range(len(mag1x)) ]).flatten()) if np.isnan(lim_mag1): print 'found nan for limiting magnitude, replacing lim_mag1 with 1.0' lim_mag1 = 1. mag1x *= normalization / lim_mag1 mag1y *= normalization / lim_mag1 mag1w *= normalization / lim_mag1 mag1z *= normalization / lim_mag1 # sys.exit() cw = [] ccw = [] lines_stretch = [] lines_twist = [] for i in range(NP): # Draw COM movement of each node as an ellipse x_disps = 0.5 * (exp1 * mag1x[i]).real y_disps = 0.5 * (exp1 * mag1y[i]).real x_vals = xy[i, 0] + x_disps y_vals = xy[i, 1] + y_disps # Draw movement of the orientation of each node as an ellipse w_disps = 0.5 * (exp1 * mag1w[i]).real z_disps = 0.5 * (exp1 * mag1z[i]).real w_vals = xy[i, 0] + w_disps z_vals = xy[i, 1] + z_disps poly_points = np.array([x_vals, y_vals]).T tilt_points = np.array([w_vals, z_vals]).T # polygon = Polygon(poly_points, True, lw=lw, ec='g') # tiltgon = Polygon(tilt_points, True, lw=lw, ec='r') # polygon = plt.plot(poly_points[:, 0], poly_points[:, 1], 'g-', lw=lw) # tiltgon = plt.plot(tilt_points[:, 0], tilt_points[:, 1], 'r-', lw=lw) npolypts = len(poly_points[:, 0]) for ii in range(npolypts): lines_s = [[ poly_points[ii, 0], poly_points[(ii + 1) % npolypts, 0] ], [poly_points[ii, 1], poly_points[(ii + 1) % npolypts, 1]]] lines_t = [[ tilt_points[ii, 0], tilt_points[(ii + 1) % npolypts, 0] ], [tilt_points[ii, 1], tilt_points[(ii + 1) % npolypts, 1]]] lines_stretch.append(lines_s) lines_twist.append(lines_t) # x0 is the marker_num^th element of x_disps x0 = x_disps[marker_num] y0 = y_disps[marker_num] w0 = w_disps[marker_num] z0 = z_disps[marker_num] # x0s is the position (global pos, not relative) of each gyro at time = marker_num(out of 81) x0s[i] = x_vals[marker_num] y0s[i] = y_vals[marker_num] w0s[i] = w_vals[marker_num] z0s[i] = z_vals[marker_num] # Get angle for xy mag = np.sqrt(x0**2 + y0**2) if mag > 0: anglez = np.arccos(x0 / mag) else: anglez = 0 if y0 < 0: anglez = 2 * np.pi - anglez # Get angle for wz tmag = np.sqrt(w0**2 + z0**2) if tmag > 0: tanglez = np.arccos(w0 / tmag) else: tanglez = 0 if y0 < 0: tanglez = 2 * np.pi - tanglez angles_arr[i] = anglez tangles_arr[i] = tanglez # polygons.append(polygon) # tiltgons.append(tiltgon) # Do Fast Fourier Transform (FFT) # ff = abs(fft.fft(x_disps + 1j*y_disps))**2 # ff_freq = fft.fftfreq(len(x_vals), 1) # mm_f = ff_freq[ff == max(ff)][0] if color_scheme == 'default': colors[2 * i] = anglez colors[2 * i + 1] = anglez # add two more colors to ensure clims go from 0 to 2pi, I think... colors[2 * NP] = 0 colors[2 * NP + 1] = 2 * np.pi plt.yticks([]) plt.xticks([]) # this is the part that puts a dot a t=0 point scat_fg = eig_ax.scatter(x0s, y0s, s=s(dotsz), edgecolors='k', facecolors='none') scat_fg2 = eig_ax.scatter(w0s, z0s, s=s(dotsz), edgecolors='gray', facecolors='none') # try: # NN = np.shape(tlat.lattice.NL)[1] # except IndexError: # NN = 0 # # Rnorm = np.array([x0s, y0s]).T # # # Bond Stretches # if draw_strain: # inc = 0 # stretches = np.zeros(4 * len(xy)) # for i in range(len(xy)): # if NN > 0: # for j, k in zip(Ni[i], Nk[i]): # if i < j and abs(k) > 0: # n1 = float(np.linalg.norm(Rnorm[i] - Rnorm[j])) # n2 = np.linalg.norm(xy[i] - xy[j]) # stretches[inc] = (n1 - n2) # inc += 1 # # stretch = np.array(stretches[0:inc]) # else: # # simply get length of BL (len(BL) = inc) by iterating over all bondsssa # inc = 0 # for i in range(len(xy)): # if NN > 0: # for j, k in zip(Ni[i], Nk[i]): # if i < j and abs(k) > 0: # inc += 1 # # # For particles with neighbors, get list of bonds to draw. # # If bondval_matrix is not None, color by the elements of that matrix # if bondval_matrix is not None or draw_strain: # test = list(np.zeros([inc, 1])) # bondvals = list(np.ones([inc, 1])) # inc = 0 # xy = np.array([x0s, y0s]).T # for i in range(len(xy)): # if NN > 0: # for j, k in zip(Ni[i], Nk[i]): # if i < j and abs(k) > 0: # test[inc] = [xy[(i, j), 0], xy[(i, j), 1]] # if bondval_matrix is not None: # bondvals[inc] = bondval_matrix[i, j] # inc += 1 # # # lines connect sites (bonds), while lines_12 draw the black lines from the pinning to location sites # lines = [zip(x, y) for x, y in test] # # # Check that we have all the cmaps # if cmap_lines not in plt.colormaps() or cmap_patches not in plt.colormaps(): # lecmaps.register_cmaps() # # # Add lines colored by strain here # if bondval_matrix is not None: # lines_st = LineCollection(lines, array=bondvals, cmap=cmap_lines, linewidth=0.8) # if line_climv is None: # maxk = np.max(np.abs(bondvals)) # mink = np.min(np.abs(bondvals)) # if (bondvals - bondvals[0] < 1e-8).all(): # lines_st.set_clim([mink - 1., maxk + 1.]) # else: # lines_st.set_clim([mink, maxk]) # # lines_st.set_zorder(2) # eig_ax.add_collection(lines_st) # else: # if draw_strain: # lines_st = LineCollection(lines, array=stretch, cmap=cmap_lines, linewidth=0.8) # if line_climv is None: # maxstretch = np.max(np.abs(stretch)) # if maxstretch < 1e-8: # line_climv = 1.0 # else: # line_climv = maxstretch # # lines_st.set_clim([-line_climv, line_climv]) # lines_st.set_zorder(2) # eig_ax.add_collection(lines_st) # Draw lines for movement lines_stretch = [zip(x, y) for x, y in lines_stretch] polygons = LineCollection(lines_stretch, linewidth=lw, linestyles='solid', color=xycolor) polygons.set_zorder(1) eig_ax.add_collection(polygons) lines_twist = [zip(x, y) for x, y in lines_twist] tiltgons = LineCollection(lines_twist, linewidth=lw, linestyles='solid', color=wzcolor) tiltgons.set_zorder(1) eig_ax.add_collection(tiltgons) # Draw polygons # polygons = PatchCollection(polygons, cmap=cmap_patches, alpha=0.6) # # p.set_array(np.array(colors)) # polygons.set_clim([0, 2 * np.pi]) # polygons.set_zorder(1) # eig_ax.add_collection(polygons) # tiltgons = PatchCollection(tiltgons, cmap=cmap_patches, alpha=0.6) # # p.set_array(np.array(colors)) # tiltgons.set_clim([0, 2 * np.pi]) # tiltgons.set_zorder(1000) # eig_ax.add_collection(tiltgons) eig_ax.set_aspect('equal') # erased ev/(2*pi) here npm 2016 cw_ccw = [cw, ccw, ev] # print cw_ccw[1] return fig, [scat_fg, scat_fg2, f_mark, polygons, tiltgons], cw_ccw
def construct_eigenvalue_DOS_plot_haldane(xy, fig, dos_ax, eig_ax, eigval, eigvect, en, Ni, Nk, marker_num=0, color_scheme='default', sub_lattice=-1, PVx=[], PVy=[], black_t0lines=False, mark_t0=True, title='auto', normalization=1., alpha=0.6, lw=1, zorder=10): """puts together lattice and DOS plots and draws normal mode magitudes as circles on top Parameters ---------- xy: array 2nx3 Equilibrium position of the gyroscopes fig : figure with lattice and DOS drawn dos_ax: axis for the DOS plot eig_ax axis for the eigenvalue plot eigval : array of dimension 2nx1 Eigenvalues of matrix for system eigvect : array of dimension 2nx2n Eigenvectors of matrix for system. Eigvect is stored as NModes x NP array, like: mode0_psi0, mode0_psi1, ... / mode1_psi0, ... en: int Number of the eigenvalue you are plotting Returns ---------- fig : completed figure for normal mode [scat_fg, p, f_mark] : things to be cleared before next normal mode is drawn """ ppu = leplt.get_points_per_unit() s = leplt.absolute_sizer() # re_eigvals = sum(abs(real(eigval))) # im_eigvals = sum(abs(imag(eigval))) ev = eigval[en] ev1 = ev # Show where current eigenvalue is in DOS plot (red line ticking current eigval) if dos_ax is not None: (f_mark, ) = dos_ax.plot([ev.real, ev.real], P.ylim(), '-r') plt.sca(dos_ax) NP = len(xy) im1 = np.imag(ev) re1 = np.real(ev) P.sca(eig_ax) if title == 'auto': eig_ax.set_title('Mode %d: $\omega=( %0.6f + %0.6f i)$' % (en, re1, im1)) elif title is not None and title not in ['', 'none']: eig_ax.set_title(title) # Preallocate ellipsoid plot vars shap = eigvect.shape angles_arr = np.zeros(NP) major_Ax = np.zeros(NP) patch = [] polygon = [] colors = np.zeros(NP + 2) # x_mag = np.zeros(NP) # y_mag = np.zeros(NP) x0s = np.zeros(NP) y0s = np.zeros(NP) mag1 = eigvect[en] # Eigvect is stored as NModes x NP*2 array, with x and y components alternating, like: # x0, y0, x1, y1, ... xNP, yNP. mag1x = np.array([mag1[i] for i in range(NP)]) mag1y = np.array([mag1[i] for i in range(NP)]) # Pick a series of times to draw out the ellipsoid time_arr = np.arange(81) * 2 * pi / (abs(ev1) * 80) exp1 = np.exp(1j * ev1 * time_arr) # Normalization for the ellipsoids lim_mag1 = max( np.array([ np.sqrt(2 * abs(exp1 * mag1x[i])**2) for i in range(len(mag1x)) ]).flatten()) mag1x /= lim_mag1 mag1y /= lim_mag1 mag1x *= normalization mag1y *= normalization cw = [] ccw = [] lines_1 = [] for i in range(NP): unit = mag1x[i] x_disps = 0.5 * (exp1 * unit).real y_disps = 0.5 * (exp1 * unit).imag x_vals = xy[i, 0] + x_disps y_vals = xy[i, 1] + y_disps # x_mag[i] = max(x_vals-xy[i,0]).real # y_mag[i] = max(y_vals-xy[i,1]).real poly_points = array([x_vals, y_vals]).T polygon = Polygon(poly_points, True) # x0 is the marker_num^th element of x_disps x0 = x_disps[marker_num] y0 = y_disps[marker_num] x0s[i] = x_vals[marker_num] y0s[i] = y_vals[marker_num] # These are the black lines protruding from pivot point to current position lines_1.append([[xy[i, 0], x_vals[marker_num]], [xy[i, 1], y_vals[marker_num]]]) mag = sqrt(x0**2 + y0**2) if mag > 0: anglez = np.arccos(x0 / mag) else: anglez = 0 if y0 < 0: anglez = 2 * np.pi - anglez # testangle = arctan2(y0,x0) # print ' x0 - x_disps[0] =', x0-x_disps[marker_num] angles_arr[i] = anglez # print 'polygon = ', poly_points patch.append(polygon) # Do Fast Fourier Transform (FFT) # ff = abs(fft.fft(x_disps + 1j*y_disps))**2 # ff_freq = fft.fftfreq(len(x_vals), 1) # mm_f = ff_freq[ff == max(ff)][0] if color_scheme == 'default': colors[i] = anglez else: if sub_lattice[i] == 0: colors[i] = 0 else: colors[i] = pi # if mm_f > 0: # colors[i] = 0 # else: # colors[i] = pi colors[NP] = 0 colors[NP + 1] = 2 * pi plt.yticks([]) plt.xticks([]) # this is the part that puts a dot a t=0 point if mark_t0: scat_fg = eig_ax.scatter(x0s[cw], y0s[cw], s=s(.02), c='k') scat_fg2 = eig_ax.scatter(x0s[cw], y0s[cw], s=s(.02), c='r') scat_fg = [scat_fg, scat_fg2] else: scat_fg = [] NP = len(xy) try: NN = shape(Ni)[1] except IndexError: NN = 0 z = np.zeros(NP) Rnorm = np.array([x0s, y0s, z]).T # Bond Stretches inc = 0 stretches = zeros(3 * len(xy)) if PVx == [] and PVy == []: '''There are no periodic boundaries supplied''' for i in range(len(xy)): if NN > 0: for j, k in zip(Ni[i], Nk[i]): if i < j and abs(k) > 0: n1 = float(linalg.norm(Rnorm[i] - Rnorm[j])) n2 = linalg.norm(xy[i] - xy[j]) stretches[inc] = (n1 - n2) inc += 1 else: '''There are periodic boundaries supplied''' # get boundary particle indices KLabs = np.zeros_like(Nk, dtype='int') KLabs[Nk > 0] = 1 boundary = extract_boundary_from_NL(xy, Ni, KLabs) for i in range(len(xy)): if NN > 0: for j, k in zip(Ni[i], Nk[i]): # if i in boundary and j in boundary: # col = np.where( Ni[i] == j)[0][0] # print 'col = ', col # n1 = float( np.linalg.norm(Rnorm[i]-Rnorm[j]) ) # n2 = np.linalg.norm(R[i] - R[j] ) # stretches[inc] = (n1 - n2) # inc += 1 # # #test[inc] = [R[i], np.array([R[j,0]+PVx[i,col], R[j,1] + PVy[i,col], 0])] # else: if i < j and abs(k) > 0: n1 = float(np.linalg.norm(Rnorm[i] - Rnorm[j])) n2 = np.linalg.norm(xy[i] - xy[j]) stretches[inc] = (n1 - n2) inc += 1 stretch = np.array(stretches[0:inc]) # For particles with neighbors, get list of bonds to draw by stretches test = list(np.zeros([inc, 1])) inc = 0 xy = np.array([x0s, y0s, z]).T if PVx == [] and PVy == []: '''There are no periodic boundaries supplied''' for i in range(len(xy)): if NN > 0: for j, k in zip(Ni[i], Nk[i]): if i < j and abs(k) > 0: test[inc] = [xy[(i, j), 0], xy[(i, j), 1]] inc += 1 else: '''There are periodic boundaries supplied''' # get boundary particle indices KLabs = np.zeros_like(Nk, dtype='int') KLabs[Nk > 0] = 1 boundary = extract_boundary_from_NL(xy, Ni, KLabs) for i in range(len(xy)): if NN > 0: for j, k in zip(Ni[i], Nk[i]): # if i in boundary and j in boundary: # col = np.where( Ni[i] == j)[0][0] # print 'i,j = (', i,j, ')' # print 'PVx[i,col] = ', PVx[i,col] # print 'PVy[i,col] = ', PVy[i,col] # test[inc] = [xy[i], np.array([xy[j,0]+PVx[i,col], xy[j,1] + PVy[i,col], 0])] # #plt.plot([ xy[i,0], xy[j,0]+PVx[i,col]], [xy[i,1], xy[j,1] + PVy[i,col] ], 'k-') # #plt.plot(xy[:,0], xy[:,1],'b.') # #plt.show() # print 'test = ', test # inc += 1 # else: if i < j and abs(k) > 0: test[inc] = [xy[(i, j), 0], xy[(i, j), 1]] inc += 1 lines = [zip(x, y) for x, y in test] # angles[-1] = 0 # angles[-2] = 2*pi lines_st = LineCollection(lines, array=stretch, cmap='seismic', linewidth=8) lines_st.set_clim([-1. * 0.25, 1 * 0.25]) lines_st.set_zorder(2) if black_t0lines: lines_12 = [zip(x, y) for x, y in lines_1] lines_12_st = LineCollection(lines_12, linewidth=0.8) lines_12_st.set_color('k') eig_ax.add_collection(lines_12_st) else: lines_12_st = [] p = PatchCollection(patch, cmap='hsv', lw=lw, alpha=alpha, zorder=zorder) p.set_array(array(colors)) p.set_clim([0, 2 * pi]) p.set_zorder(1) # eig_ax.add_collection(lines_st) eig_ax.add_collection(p) eig_ax.set_aspect('equal') s = leplt.absolute_sizer() # erased ev/(2*pi) here npm 2016 cw_ccw = [cw, ccw, ev] # print cw_ccw[1] return fig, [scat_fg, p, f_mark, lines_12_st], cw_ccw
def plot_eigvect_excitation_haldane(xy, fig, dos_ax, eig_ax, eigval, eigvect, en, marker_num=0, black_t0lines=False, mark_t0=True, title='auto', normalization=1., alpha=0.6, lw=1, zorder=10): """Draws normal mode ellipsoids on axis eig_ax. If black_t0lines is true, draws the black line from pinning site to positions. The difference from hlatpfns.construct_haldane_eigvect_DOS_plot() is doesn't draw lattice. Parameters ---------- xy: array 2N x 3 Equilibrium position of the gyroscopes fig : figure with lattice and DOS drawn dos_ax: matplotlib axis instance or None axis for the DOS plot. If None, ignores this input eig_ax : matplotlib axis instance axis for the eigenvalue plot eigval : array of dimension 2nx1 Eigenvalues of matrix for system eigvect : array of dimension 2nx2n Eigenvectors of matrix for system. Eigvect is stored as NModes x NP*2 array, with x and y components alternating, like: x0, y0, x1, y1, ... xNP, yNP. en: int Number of the eigenvalue you are plotting marker_num : int in (0, 80) where in the phase (0 to 80) to call t=t0. This sets "now" for drawing where in the normal mode to draw black_t0lines : bool Draw black lines extending from the pinning site to the current site (where 'current' is determined by marker_num) Returns ---------- fig : matplotlib figure instance completed figure for normal mode [scat_fg, pp, f_mark, lines12_st] : things to be cleared before next normal mode is drawn """ s = leplt.absolute_sizer() ev = eigval[en] ev1 = ev # Show where current eigenvalue is in DOS plot if dos_ax is not None: (f_mark, ) = dos_ax.plot([np.real(ev), np.real(ev)], dos_ax.get_ylim(), '-r') NP = len(xy) im1 = np.real(ev) plt.sca(eig_ax) if title == 'auto': eig_ax.set_title('$\omega = %0.6f$' % im1) elif title is not None and title not in ['', 'none']: eig_ax.set_title(title) # Preallocate ellipsoid plot vars angles_arr = np.zeros(NP, dtype=float) patch = [] colors = np.zeros(NP) x0s = np.zeros(NP, dtype=float) y0s = np.zeros(NP, dtype=float) mag1 = eigvect[en] # Pick a series of times to draw out the ellipsoid time_arr = np.arange(81) * 2 * np.pi / (np.abs(ev1) * 80) exp1 = np.exp(1j * ev1 * time_arr) # Normalization for the ellipsoids mag1 /= np.max(np.abs(mag1)) mag1 *= normalization if black_t0lines: lines_1 = [] else: lines_12_st = [] for i in range(NP): x_disps = 0.5 * (exp1 * mag1[i]).real y_disps = 0.5 * (exp1 * mag1[i]).imag x_vals = xy[i, 0] + x_disps y_vals = xy[i, 1] + y_disps poly_points = np.array([x_vals, y_vals]).T polygon = Polygon(poly_points, True) # x0 is the marker_num^th element of x_disps x0 = x_disps[marker_num] y0 = y_disps[marker_num] x0s[i] = x_vals[marker_num] y0s[i] = y_vals[marker_num] if black_t0lines: # These are the black lines protruding from pivot point to current position lines_1.append([[xy[i, 0], x_vals[marker_num]], [xy[i, 1], y_vals[marker_num]]]) mag = np.sqrt(x0**2 + y0**2) if mag > 0: anglez = np.arccos(x0 / mag) else: anglez = 0 if y0 < 0: anglez = 2 * np.pi - anglez angles_arr[i] = anglez patch.append(polygon) colors[i] = anglez # this is the part that puts a dot a t=0 point if mark_t0: scat_fg = eig_ax.scatter(x0s, y0s, s=s(.02), c='k') else: scat_fg = [] pp = PatchCollection(patch, cmap='hsv', lw=lw, alpha=alpha, zorder=zorder) pp.set_array(np.array(colors)) pp.set_clim([0, 2 * np.pi]) pp.set_zorder(1) eig_ax.add_collection(pp) if black_t0lines: lines_12 = [zip(x, y) for x, y in lines_1] lines_12_st = LineCollection(lines_12, linewidth=0.8) lines_12_st.set_color('k') eig_ax.add_collection(lines_12_st) eig_ax.set_aspect('equal') return fig, [scat_fg, pp, f_mark, lines_12_st]
def construct_haldane_eigvect_DOS_plot(xy, fig, DOS_ax, eig_ax, eigval, eigvect, en, NL, KL, marker_num=0, color_scheme='default', sub_lattice=-1, normalization=None): """puts together lattice and DOS plots and draws normal mode ellipsoids on top Parameters ---------- xy: array 2N x 3 Equilibrium position of the gyroscopes fig : figure with lattice and DOS drawn DOS_ax: axis for the DOS plot eig_ax axis for the eigenvalue plot eigval : array of dimension 2nx1 Eigenvalues of matrix for system eigvect : array of dimension 2nx2n Eigenvectors of matrix for system. Eigvect is stored as NModes x NP*2 array, with x and y components alternating, like: x0, y0, x1, y1, ... xNP, yNP. en: int Number of the eigenvalue you are plotting Returns ---------- fig : completed figure for normal mode [scat_fg, p, f_mark] : things to be cleared before next normal mode is drawn """ s = leplt.absolute_sizer() plt.sca(DOS_ax) ev = eigval[en] ev1 = ev # Show where current eigenvalue is in DOS plot (f_mark, ) = plt.plot([ev, ev], plt.ylim(), '-r') NP = len(xy) im1 = np.imag(ev) re1 = np.real(ev) plt.sca(eig_ax) plt.title('Mode %d; $\Omega=( %0.6f + %0.6f i)$' % (en, re1, im1)) # Preallocate ellipsoid plot vars angles_arr = np.zeros(NP) patch = [] colors = np.zeros(NP + 2) x0s = np.zeros(NP) y0s = np.zeros(NP) mag1 = eigvect[en] if normalization is None: mag1 /= np.max(np.abs(mag1)) else: mag1 *= normalization * float(len(xy)) # Pick a series of times to draw out the ellipsoid time_arr = np.arange(81.0) * 2. * np.pi / float(abs(ev1) * 80) exp1 = np.exp(1j * ev1 * time_arr) cw = [] ccw = [] lines_1 = [] for i in range(NP): x_disps = 0.5 * (exp1 * mag1[i]).real y_disps = 0.5 * (exp1 * mag1[i]).imag x_vals = xy[i, 0] + x_disps y_vals = xy[i, 1] + y_disps poly_points = np.array([x_vals, y_vals]).T polygon = Polygon(poly_points, True) # x0 is the marker_num^th element of x_disps x0 = x_disps[marker_num] y0 = y_disps[marker_num] x0s[i] = x_vals[marker_num] y0s[i] = y_vals[marker_num] # These are the black lines protruding from pivot point to current position lines_1.append([[xy[i, 0], x_vals[marker_num]], [xy[i, 1], y_vals[marker_num]]]) mag = np.sqrt(x0**2 + y0**2) if mag > 0: anglez = np.arccos(x0 / mag) else: anglez = 0 if y0 < 0: anglez = 2 * np.pi - anglez angles_arr[i] = anglez patch.append(polygon) if color_scheme == 'default': colors[i] = anglez else: if sub_lattice[i] == 0: colors[i] = 0 else: colors[i] = np.pi ccw.append(i) colors[NP] = 0 colors[NP + 1] = 2 * np.pi plt.yticks([]) plt.xticks([]) # this is the part that puts a dot a t=0 point scat_fg = eig_ax.scatter(x0s[cw], y0s[cw], s=s(.02), c='DodgerBlue') scat_fg2 = eig_ax.scatter(x0s[ccw], y0s[ccw], s=s(.02), c='Red', zorder=3) NP = len(xy) try: NN = np.shape(NL)[1] except IndexError: NN = 0 z = np.zeros(NP) Rnorm = np.array([x0s, y0s, z]).T # Bond Stretches inc = 0 stretches = np.zeros(4 * len(xy)) for i in range(len(xy)): if NN > 0: for j, k in zip(NL[i], KL[i]): if i < j and abs(k) > 0: n1 = float(linalg.norm(Rnorm[i] - Rnorm[j])) n2 = linalg.norm(xy[i] - xy[j]) stretches[inc] = (n1 - n2) inc += 1 # For particles with neighbors, get list of bonds to draw by stretches test = list(np.zeros([inc, 1])) inc = 0 xy = np.array([x0s, y0s, z]).T for i in range(len(xy)): if NN > 0: for j, k in zip(NL[i], KL[i]): if i < j and abs(k) > 0: test[inc] = [xy[(i, j), 0], xy[(i, j), 1]] inc += 1 stretch = np.array(stretches[0:inc]) # lines connect sites (bonds), while lines_12 draw the black lines from the pinning to location sites lines = [zip(x, y) for x, y in test] lines_12 = [zip(x, y) for x, y in lines_1] lines_st = LineCollection(lines, array=stretch, cmap='seismic', linewidth=8) lines_st.set_clim([-1. * 0.25, 1 * 0.25]) lines_st.set_zorder(2) lines_12_st = LineCollection(lines_12, linewidth=0.8) lines_12_st.set_color('k') p = PatchCollection(patch, cmap='hsv', alpha=0.6) p.set_array(np.array(colors)) p.set_clim([0, 2 * np.pi]) p.set_zorder(1) # eig_ax.add_collection(lines_st) eig_ax.add_collection(lines_12_st) eig_ax.add_collection(p) eig_ax.set_aspect('equal') # erased ev/(2*pi) here npm 2016 cw_ccw = [cw, ccw, ev] # print cw_ccw[1] return fig, [scat_fg, scat_fg2, p, f_mark, lines_12_st], cw_ccw
def movie_plot_2D(xy, BL, bs=None, fname='none', title='', NL=[], KL=[], BLNNN=[], NLNNN=[], KLNNN=[], PVx=[], PVy=[], PVxydict={}, nljnnn=None, kljnnn=None, klknnn=None, ax=None, fig=None, axcb='auto', cbar_ax=None, cbar_orientation='vertical', xlimv='auto', ylimv='auto', climv=0.1, colorz=True, ptcolor=None, figsize='auto', colorpoly=False, bondcolor=None, colormap='seismic', bgcolor=None, axis_off=False, axis_equal=True, text_topleft=None, lw=-1., ptsize=10, negative_NNN_arrows=False, show=False, arrow_alpha=1.0, fontsize=8, cax_label='Strain', zorder=0, rasterized=False): """Plots (and saves if fname is not 'none') a 2D image of the lattice with colored bonds and particles colored by coordination (both optional). Parameters ---------- xy : array of dimension nx2 2D lattice of points (positions x,y) BL : array of dimension #bonds x 2 Each row is a bond and contains indices of connected points bs : array of dimension #bonds x 1 or None Strain in each bond fname : string Full path including name of the file (.png, etc), if None, will not save figure title : string The title of the frame NL : NP x NN int array (optional, for speed) Specify to speed up computation, if colorz or colorpoly or if periodic boundary conditions KL : NP x NN int array (optional, for speed) Specify to speed up computation, if colorz or colorpoly or if periodic boundary conditions BLNNN : NLNNN : KLNNN : PVx : NP x NN float array (optional, for periodic lattices and speed) ijth element of PVx is the x-component of the vector taking NL[i,j] to its image as seen by particle i If PVx and PVy are specified, PVxydict need not be specified. PVy : NP x NN float array (optional, for periodic lattices and speed) ijth element of PVy is the y-component of the vector taking NL[i,j] to its image as seen by particle i If PVx and PVy are specified, PVxydict need not be specified. PVxydict : dict (optional, for periodic lattices) dictionary of periodic bonds (keys) to periodic vectors (values) nljnnn : #pts x max(#NNN) int array or None nearest neighbor array matching NLNNN and KLNNN. nljnnn[i, j] gives the neighbor of i such that NLNNN[i, j] is the next nearest neighbor of i through the particle nljnnn[i, j] kljnnn : #pts x max(#NNN) int array or None bond array describing periodicity of bonds matching NLNNN and KLNNN. kljnnn[i, j] describes the bond type (bulk -> +1, periodic --> -1) of bond connecting i to nljnnn[i, j] klknnn : #pts x max(#NNN) int array or None bond array describing periodicity of bonds matching NLNNN and KLNNN. klknnn[i, j] describes the bond type (bulk -> +1, periodic --> -1) of bond connecting nljnnn[i, j] to NLNNN[i, j] ax: matplotlib figure axis instance Axis on which to draw the network fig: matplotlib figure instance Figure on which to draw the network axcb: matplotlib colorbar instance Colorbar to use for strains in bonds cbar_ax : axis instance Axis to use for colorbar. If colorbar instance is not already defined, use axcb instead. cbar_orientation : str ('horizontal' or 'vertical') Orientation of the colorbar xlimv: float or tuple of floats ylimv: float or tuple of floats climv : float or tuple Color limit for coloring bonds by bs colorz: bool whether to color the particles by their coordination number ptcolor: string color spec or tuple color spec or None color specification for coloring the points, if colorz is False. Default is None (no coloring of points) figsize : tuple w,h tuple in inches colorpoly : bool Whether to color in polygons formed by bonds according to the number of sides bondcolor : color specification (hexadecimal or RGB) colormap : if bondcolor is None, uses bs array to color bonds bgcolor : hex format string, rgb color spec, or None If not None, sets the bgcolor. Often used is '#d9d9d9' axis_off : bool Turn off the axis border and canvas axis_equal : bool text_topleft : str or None lw: float line width for plotting bonds. If lw == -1, then uses automatic line width to adjust for bond density.. ptsize: float size of points passed to absolute_sizer negative_NNN_arrows : bool make positive and negative NNN hoppings different color show : bool whether to show the plot after creating it arrow_alpha : float opacity of the arrow fontsize : int (default=8) fontsize for all labels cax_label : int (default='Strain') Label for the colorbar zorder : int z placement on axis (higher means bringing network to the front, lower is to the back Returns ---------- [ax,axcb] : stuff to clear after plotting """ if fig is None or fig == 'none': fig = plt.gcf() if ax is None or ax == 'none': if figsize == 'auto': plt.clf() else: fig = plt.figure(figsize=figsize) ax = plt.axes() if bs is None: bs = np.zeros_like(BL[:, 0], dtype=float) if colormap not in plt.colormaps(): lecmaps.register_colormaps() NP = len(xy) if lw == -1: if NP < 10000: lw = 0.5 s = leplt.absolute_sizer() else: lw = (10 / np.sqrt(len(xy))) if NL == [] and KL == []: if colorz or colorpoly: NL, KL = le.BL2NLandKL(BL, NP=NP, NN='min') if (BL < 0).any(): if len(PVxydict) == 0: raise RuntimeError( 'PVxydict must be supplied to display_lattice_2D() when periodic BCs exist, ' + 'if NL and KL not supplied!') else: PVx, PVy = le.PVxydict2PVxPVy(PVxydict, NL, KL) if colorz: zvals = (KL != 0).sum(1) zmed = np.median(zvals) # print 'zmed = ', zmed under1 = np.logical_and(zvals < zmed - 0.5, zvals > zmed - 1.5) over1 = np.logical_and(zvals > zmed + 0.5, zvals < zmed + 1.5) Cz = np.zeros((len(xy), 3), dtype=int) # far under black // under blue // equal white // over red // far over green Cz[under1] = [0. / 255, 50. / 255, 255. / 255] Cz[zvals == zmed] = [100. / 255, 100. / 255, 100. / 255] Cz[over1] = [255. / 255, 0. / 255, 0. / 255] Cz[zvals > zmed + 1.5] = [0. / 255, 255. / 255, 50. / 255] # Cz[zvals<zmed-1.5] = [0./255,255./255,150./255] #leave these black s = leplt.absolute_sizer() sval = min([.005, .12 / np.sqrt(len(xy))]) sizes = np.zeros(NP, dtype=float) sizes[zvals > zmed + 0.5] = sval sizes[zvals == zmed] = sval * 0.5 sizes[zvals < zmed - 0.5] = sval # topinds = zvals!=zmed ax.scatter(xy[:, 0], xy[:, 1], s=s(sizes), c=Cz, edgecolor='none', zorder=10, rasterized=rasterized) ax.axis('equal') elif ptcolor is not None and ptcolor != 'none' and ptcolor != '': if NP < 10000: # if smallish #pts, plot them # print 'xy = ', xy # plt.plot(xy[:,0],xy[:,1],'k.') s = leplt.absolute_sizer() ax.scatter(xy[:, 0], xy[:, 1], s=ptsize, alpha=0.5, facecolor=ptcolor, edgecolor='none', rasterized=rasterized) if colorpoly: # Color the polygons based on # sides # First extract polygons. To do that, if there are periodic boundaries, we need to supply as dict if PVxydict == {} and len(PVx) > 0: PVxydict = le.PVxy2PVxydict(PVx, PVy, NL, KL=KL) polygons = le.extract_polygons_lattice(xy, BL, NL=NL, KL=KL, viewmethod=True, PVxydict=PVxydict) PolyPC = le.polygons2PPC(polygons) # number of polygon sides Pno = np.array([len(polyg) for polyg in polygons], dtype=int) print 'nvis: Pno = ', Pno print 'nvis: medPno = ', np.floor(np.median(Pno)) medPno = np.floor(np.median(Pno)) uIND = np.where(Pno == medPno - 1)[0] mIND = np.where(Pno == medPno)[0] oIND = np.where(Pno == medPno + 1)[0] loIND = np.where(Pno < medPno - 1.5)[0] hiIND = np.where(Pno > medPno + 1.5)[0] print ' uIND = ', uIND print ' oIND = ', oIND print ' loIND = ', loIND print ' hiIND = ', hiIND if len(uIND) > 0: PPCu = [PolyPC[i] for i in uIND] pu = PatchCollection(PPCu, color='b', alpha=0.5) ax.add_collection(pu) if len(mIND) > 0: PPCm = [PolyPC[i] for i in mIND] pm = PatchCollection(PPCm, color=[0.5, 0.5, 0.5], alpha=0.5) ax.add_collection(pm) if len(oIND) > 0: PPCo = [PolyPC[i] for i in oIND] po = PatchCollection(PPCo, color='r', alpha=0.5) ax.add_collection(po) if len(loIND) > 0: PPClo = [PolyPC[i] for i in loIND] plo = PatchCollection(PPClo, color='k', alpha=0.5) ax.add_collection(plo) if len(hiIND) > 0: PPChi = [PolyPC[i] for i in hiIND] phi = PatchCollection(PPChi, color='g', alpha=0.5) ax.add_collection(phi) # Efficiently plot many lines in a single set of axes using LineCollection # First check if there are periodic bonds if BL.size > 0: if (BL < 0).any(): if PVx == [] or PVy == [] or PVx is None or PVy is None: raise RuntimeError( 'PVx and PVy must be supplied to display_lattice_2D when periodic BCs exist!' ) else: # get indices of periodic bonds perINDS = np.unique(np.where(BL < 0)[0]) perBL = np.abs(BL[perINDS]) # # Check # print 'perBL = ', perBL # plt.plot(xy[:,0], xy[:,1],'b.') # for i in range(len(xy)): # plt.text(xy[i,0]+0.05, xy[i,1],str(i)) # plt.show() # define the normal bonds which are not periodic normINDS = np.setdiff1d(np.arange(len(BL)), perINDS) BLtmp = BL[normINDS] bstmp = bs[normINDS] lines = [ zip(xy[BLtmp[i, :], 0], xy[BLtmp[i, :], 1]) for i in range(len(BLtmp)) ] xy_add = np.zeros((4, 2)) # Build new strain list bs_out by storing bulk lines first, then recording the strain twice # for each periodic bond since the periodic bond is at least two lines in the plot, get bs_out # ready for appending # bs_out = np.zeros(len(normINDS) + 5 * len(perINDS), dtype=float) # bs_out[0:len(normINDS)] = bstmp bs_out = bstmp.tolist() # Add periodic bond lines to image # Note that we have to be careful that a single particle can be connected to another both in the bulk # and through a periodic boundary, and perhaps through more than one periodic boundary # kk indexes perINDS to determine what the strain of each periodic bond should be # draw_perbond_count counts the number of drawn linesegments that are periodic bonds kk, draw_perbond_count = 0, 0 for row in perBL: colA = np.argwhere(NL[row[0]] == row[1]).ravel() colB = np.argwhere(NL[row[1]] == row[0]).ravel() if len(colA) > 1 or len(colB) > 1: # Look for where KL < 0 to pick out just the periodic bond(s) -- ie there # were both bulk and periodic bonds connecting row[0] to row[1] a_klneg = np.argwhere(KL[row[0]] < 0) colA = np.intersect1d(colA, a_klneg) b_klneg = np.argwhere(KL[row[1]] < 0) colB = np.intersect1d(colB, b_klneg) # print 'colA = ', colA # print 'netvis here' # sys.exit() if len(colA) > 1 or len(colB) > 1: # there are multiple periodic bonds connecting one particle to another (in different # directions). Plot each of them. for ii in range(len(colA)): print 'netvis: colA = ', colA print 'netvis: colB = ', colB # columns a and b for this ii index caii, cbii = colA[ii], colB[ii] # add xy points to the network to plot to simulate image particles xy_add[0] = xy[row[0]] xy_add[1] = xy[row[1]] + np.array( [PVx[row[0], caii], PVy[row[0], caii]]) xy_add[2] = xy[row[1]] xy_add[3] = xy[row[0]] + np.array( [PVx[row[1], cbii], PVy[row[1], cbii]]) # Make the lines to draw (dashed lines for this periodic case) lines += zip(xy_add[0:2, 0], xy_add[0:2, 1]), zip( xy_add[2:4, 0], xy_add[2:4, 1]) bs_out.append(bs[perINDS[kk]]) bs_out.append(bs[perINDS[kk]]) draw_perbond_count += 1 kk += 1 else: # print 'row = ', row # print 'NL = ', NL # print 'KL = ', KL # print 'colA, colB = ', colA, colB colA, colB = colA[0], colB[0] xy_add[0] = xy[row[0]] xy_add[1] = xy[row[1]] + np.array( [PVx[row[0], colA], PVy[row[0], colA]]) xy_add[2] = xy[row[1]] xy_add[3] = xy[row[0]] + np.array( [PVx[row[1], colB], PVy[row[1], colB]]) lines += zip(xy_add[0:2, 0], xy_add[0:2, 1]), zip( xy_add[2:4, 0], xy_add[2:4, 1]) # bs_out[2 * kk + len(normINDS)] = bs[perINDS[kk]] # bs_out[2 * kk + 1 + len(normINDS)] = bs[perINDS[kk]] bs_out.append(bs[perINDS[kk]]) bs_out.append(bs[perINDS[kk]]) draw_perbond_count += 1 kk += 1 else: colA, colB = colA[0], colB[0] xy_add[0] = xy[row[0]] xy_add[1] = xy[row[1]] + np.array( [PVx[row[0], colA], PVy[row[0], colA]]) xy_add[2] = xy[row[1]] xy_add[3] = xy[row[0]] + np.array( [PVx[row[1], colB], PVy[row[1], colB]]) lines += zip(xy_add[0:2, 0], xy_add[0:2, 1]), zip(xy_add[2:4, 0], xy_add[2:4, 1]) # bs_out[2 * kk + len(normINDS)] = bs[perINDS[kk]] # bs_out[2 * kk + 1 + len(normINDS)] = bs[perINDS[kk]] bs_out.append(bs[perINDS[kk]]) bs_out.append(bs[perINDS[kk]]) draw_perbond_count += 1 kk += 1 # replace bs by the new bs (bs_out) bs = np.array(bs_out) # store number of bulk bonds nbulk_bonds = len(normINDS) else: if len(np.shape(BL)) > 1: lines = [ zip(xy[BL[i, :], 0], xy[BL[i, :], 1]) for i in range(np.shape(BL)[0]) ] # store number of bulk bonds nbulk_bonds = len(lines) else: lines = [ zip(xy[BL[i][0]], xy[BL[i][1]]) for i in range(np.shape(BL)[0]) ] # store number of bulk bonds nbulk_bonds = 1 if isinstance(climv, tuple): cmin = climv[0] cmax = climv[1] elif isinstance(climv, float): cmin = -climv cmax = climv elif climv is None: cmin = None cmax = None if bondcolor is None: # draw the periodic bonds as dashed, regular bulk bonds as solid line_segments = LineCollection( lines[0:nbulk_bonds], # Make a sequence of x,y pairs linewidths=lw, # could iterate over list linestyles='solid', cmap=colormap, norm=plt.Normalize(vmin=cmin, vmax=cmax), zorder=zorder, rasterized=rasterized) line_segments.set_array(bs[0:nbulk_bonds]) # draw the periodic bonds as dashed, if there are any periodic_lsegs = LineCollection( lines[nbulk_bonds:], # Make a sequence of x,y pairs linewidths=lw, # could iterate over list linestyles='dashed', cmap=colormap, norm=plt.Normalize(vmin=cmin, vmax=cmax), zorder=zorder, rasterized=rasterized) periodic_lsegs.set_array(bs[nbulk_bonds:]) else: line_segments = LineCollection(lines[0:nbulk_bonds], linewidths=lw, linestyles='solid', colors=bondcolor, zorder=zorder, rasterized=rasterized) # draw the periodic bonds as dashed, if there are any periodic_lsegs = LineCollection(lines[nbulk_bonds:], linewidths=lw, linestyles='dashed', colors=bondcolor, zorder=zorder, rasterized=rasterized) ax.add_collection(line_segments) if periodic_lsegs: ax.add_collection(periodic_lsegs) # If there is only a single bond color, ignore the colorbar specification if bondcolor is None or isinstance(bondcolor, np.ndarray): if axcb == 'auto': if cbar_ax is None: print 'nvis: Instantiating colorbar...' axcb = fig.colorbar(line_segments) else: print 'nvis: Using cbar_ax to instantiate colorbar' axcb = fig.colorbar(line_segments, cax=cbar_ax, orientation=cbar_orientation) if axcb != 'none' and axcb is not None: print 'nvis: Creating colorbar...' axcb.set_label(cax_label, fontsize=fontsize) axcb.set_clim(vmin=cmin, vmax=cmax) else: # Ignore colorbar axis specification axcb = 'none' else: axcb = 'none' if len(BLNNN) > 0: # todo: add functionality for periodic NNN connections # Efficiently plot many lines in a single set of axes using LineCollection lines = [ zip(xy[BLNNN[i, :], 0], xy[BLNNN[i, :], 1]) for i in range(len(BLNNN)) ] linesNNN = LineCollection( lines, # Make a sequence of x,y pairs linewidths=lw, # could iterate over list linestyles='dashed', color='blue', zorder=100) ax.add_collection(linesNNN, rasterized=rasterized) elif len(NLNNN) > 0 and len(KLNNN) > 0: factor = 0.8 if (BL < 0).any(): print 'nvis: plotting periodic NNN...' if nljnnn is None: raise RuntimeError( 'Must supply nljnnn to plot NNN hoppings/connections') for i in range(NP): todo = np.where(KLNNN[i, :] > 1e-12)[0] for index in todo: kk = NLNNN[i, index] # Ascribe the correct periodic vector based on both PVx[i, NNind] and PVx[NNind, ind] # Note : nljnnn is # nearest neighbor array matching NLNNN and KLNNN. nljnnn[i, j] gives the neighbor of i such that # NLNNN[i, j] is the next nearest neighbor of i through the particle nljnnn[i, j] jj = nljnnn[i, index] if kljnnn[i, index] < 0 or klknnn[i, index] < 0: jind = np.where(NL[i, :] == jj)[0][0] kind = np.where(NL[jj, :] == kk)[0][0] # print 'jj = ', jj # print 'kk = ', kk # print 'jind = ', jind # print 'kind = ', kind # print 'NL[i, :] =', NL[i, :] # print 'NL[jj, :] =', NL[jj, :] # print 'PVx[i, jind] = ', PVx[i, jind] # print 'PVy[i, jind] = ', PVy[i, jind] # print 'PVx[jj, kind] = ', PVx[jj, kind] # print 'PVy[jj, kind] = ', PVy[jj, kind] dx = (xy[kk, 0] + PVx[i, jind] + PVx[jj, kind] - xy[i, 0]) * factor dy = (xy[kk, 1] + PVy[i, jind] + PVy[jj, kind] - xy[i, 1]) * factor else: dx = (xy[kk, 0] - xy[i, 0]) * factor dy = (xy[kk, 1] - xy[i, 1]) * factor ax.arrow(xy[i, 0], xy[i, 1], dx, dy, head_width=0.1, head_length=0.2, fc='b', ec='b', linestyle='dashed') # Check # print 'dx = ', dx # print 'dy = ', dy # for ind in range(NP): # plt.text(xy[ind, 0]-0.2, xy[ind, 1]-0.2, str(ind)) # plt.show() # sys.exit() else: # amount to offset clockwise nnn arrows for i in range(NP): todo = np.where(KLNNN[i, :] > 1e-12)[0] # Allow for both blue and red arrows (forward/backward), or just blue. If just blue, use no offset and # full scale factor if negative_NNN_arrows: scalef = 0.3 else: scalef = 0.8 offset = np.array([0.0, 0.0]) for ind in NLNNN[i, todo]: if negative_NNN_arrows: offset = (xy[ind, :] - xy[i, :]) * 0.5 ax.arrow(xy[i, 0] + offset[0], xy[i, 1] + offset[1], (xy[ind, 0] - xy[i, 0]) * scalef, (xy[ind, 1] - xy[i, 1]) * scalef, head_width=0.1, head_length=0.2, fc='b', ec='b', alpha=arrow_alpha) if negative_NNN_arrows: todo = np.where(KLNNN[i, :] < -1e-12)[0] for ind in NLNNN[i, todo]: offset = (xy[ind, :] - xy[i, :]) * 0.5 ax.arrow(xy[i, 0] + offset[0], xy[i, 1] + offset[1], (xy[ind, 0] - xy[i, 0]) * 0.3, (xy[ind, 1] - xy[i, 1]) * 0.3, head_width=0.1, head_length=0.2, fc='r', ec='r', alpha=arrow_alpha) if bgcolor is not None: ax.set_axis_bgcolor(bgcolor) # set limits ax.axis('scaled') if xlimv != 'auto' and xlimv is not None: if isinstance(xlimv, tuple): ax.set_xlim(xlimv[0], xlimv[1]) else: print 'nvis: setting xlimv' ax.set_xlim(-xlimv, xlimv) else: ax.set_xlim(np.min(xy[:, 0]) - 2.5, np.max(xy[:, 0]) + 2.5) if ylimv != 'auto' and ylimv is not None: if isinstance(ylimv, tuple): print 'nvis: setting ylimv to tuple' ax.set_ylim(ylimv[0], ylimv[1]) else: ax.set_ylim(-ylimv, ylimv) else: print 'nvis: setting', min(xy[:, 1]), max(xy[:, 1]) ax.set_ylim(np.min(xy[:, 1]) - 2, np.max(xy[:, 1]) + 2) if title is not None: ax.set_title(title, fontsize=fontsize) if text_topleft is not None: ax.text(0.05, .98, text_topleft, horizontalalignment='right', verticalalignment='top', transform=ax.transAxes) if axis_off: ax.axis('off') if fname != 'none' and fname != '' and fname is not None: print 'nvis: saving figure: ', fname plt.savefig(fname) if show: plt.show() return [ax, axcb]