def twocell_to_bz1(cell):
from ase.dft.bz import bz_vertices
# 2d in x-y plane
if len(cell) > 2:
assert all(abs(cell[2][0:2]) < 1e-6) and all(abs(cell.T[2][0:2]) < 1e-6)
else:
cell = [list(c) + [0] for c in cell] + [[0, 0, 1]]
icell = np.linalg.inv(cell).T
try:
bz1 = bz_vertices(icell[:3, :3], dim=2)
except TypeError:
bz1 = bz_vertices(icell[:3, :3])
return bz1, icell, cell
def calculate_bz_vertices_from_direct_cell(cell):
from ase.dft.bz import bz_vertices
if len(cell) > 2:
assert all(abs(cell[2][0:2]) < 1e-6) and all(
abs(cell.T[2][0:2]) < 1e-6)
else:
cell = [list(c) + [0] for c in cell] + [[0, 0, 1]]
icell = np.linalg.inv(cell).T
try:
bz1 = bz_vertices(icell[:3, :3], dim=2)
except TypeError:
bz1 = bz_vertices(icell[:3, :3])
return bz1
def calc_commensurate_moire_cell(underlayer_a,
overlayer_a,
relative_angle=0,
swap_angle=False):
"""
Calculates nearly commensurate moire unit cells for two hexagonal lattices
:return:
"""
from ase.dft.kpoints import get_special_points
from ase.dft.bz import bz_vertices
underlayer_direct = hex_cell_2d(a=underlayer_a)
overlayer_direct = hex_cell_2d(a=overlayer_a)
underlayer_direct = [list(c) + [0]
for c in underlayer_direct] + [[0, 0, 1]]
overlayer_direct = [list(c) + [0] for c in overlayer_direct] + [[0, 0, 1]]
underlayer_icell = np.linalg.inv(underlayer_direct).T
overlayer_icell = np.linalg.inv(overlayer_direct).T
underlayer_k = np.dot(underlayer_icell.T,
get_special_points(underlayer_direct)['K'])
overlayer_k = Rotation.from_rotvec([0, 0, relative_angle]).apply(
np.dot(overlayer_icell.T,
get_special_points(overlayer_direct)['K']))
moire_k = (underlayer_k - overlayer_k)
moire_a = underlayer_a * (np.linalg.norm(underlayer_k) /
np.linalg.norm(moire_k))
moire_angle = angle_between_vectors(underlayer_k, moire_k)
if swap_angle:
moire_angle = -moire_angle
moire_cell = hex_cell_2d(moire_a)
moire_cell = [list(c) + [0] for c in moire_cell] + [[0, 0, 1]]
moire_cell = Rotation.from_rotvec([0, 0, moire_angle]).apply(moire_cell)
moire_icell = np.linalg.inv(moire_cell).T
moire_bz_points = bz_vertices(moire_icell)
moire_bz_points = moire_bz_points[[len(p[0])
for p in moire_bz_points].index(6)][0]
return {
'k_points': (underlayer_k, overlayer_k, moire_k),
'moire_a': moire_a,
'moire_k': moire_k,
'moire_cell': moire_cell,
'moire_icell': moire_icell,
'moire_bz_points': moire_bz_points,
'moire_bz_angle': moire_angle,
}
def bz_points_for_hexagonal_lattice(a=1):
from ase.dft.bz import bz_vertices
cell = hex_cell_2d(a)
cell = [list(c) + [0] for c in cell] + [[0, 0, 1]]
icell = np.linalg.inv(cell).T
bz_vertices = bz_vertices(icell)
# get the first face which has six points, this is the top or bottom
# face of the cell
return as_2d(bz_vertices[[len(face[0])
for face in bz_vertices].index(6)][0])
def plot_plane_to_bz(cell, plane, ax, special_points=None, facecolor='red'):
from ase.dft.bz import bz_vertices
if isinstance(plane, str):
plane_points = process_kpath(plane, cell, special_points=special_points)[0]
else:
plane_points = plane
d1, d2 = plane_points[1] - plane_points[0], plane_points[2] - plane_points[0]
faces = [p[0] for p in bz_vertices(np.linalg.inv(cell).T)]
pts = polyhedron_intersect_plane(faces, np.cross(d1, d2), plane_points[0])
collection = Poly3DCollection([pts])
collection.set_facecolor(facecolor)
ax.add_collection3d(collection, zs='z')
def build_2dbz_poly(vertices=None, icell=None, cell=None):
"""
Converts brillouin zone or equivalent information to a polygon mask
that can be used to mask away data outside the zone boundary.
:param vertices:
:param icell:
:param cell:
:return:
"""
from arpes.analysis.mask import raw_poly_to_mask
from ase.dft.bz import bz_vertices # pylint: disable=import-error
assert (cell is not None or vertices is not None or icell is not None)
if vertices is None:
if icell is None:
icell = np.linalg.inv(cell).T
vertices = bz_vertices(icell)
points, _ = vertices[0] # points, normal
points_2d = [p[:2] for p in points]
return raw_poly_to_mask(points_2d)
def bz3d_plot(cell, vectors=False, paths=None, points=None, ax=None,
elev=None, scale=1, repeat=None, transformations=None,
hide_ax=True, **kwargs):
"""
For now this is lifted from ase.dft.bz.bz3d_plot with some modifications. All copyright and
licensing terms for this and bz2d_plot are those of the current release of ASE
(Atomic Simulation Environment).
:param cell:
:param vectors:
:param paths:
:param points:
:param elev:
:param scale:
:return:
"""
try:
from ase.dft.bz import bz_vertices # dynamic because we do not require ase
except ImportError:
warnings.warn('You will need to install ASE (Atomic Simulation Environment) to use this feature.')
raise ImportError('You will need to install ASE before using Brillouin Zone plotting')
class Arrow3D(FancyArrowPatch):
def __init__(self, xs, ys, zs, *args, **kwargs):
FancyArrowPatch.__init__(self, (0, 0), (0, 0), *args, **kwargs)
self._verts3d = xs, ys, zs
def draw(self, renderer):
xs3d, ys3d, zs3d = self._verts3d
xs, ys, zs = proj3d.proj_transform(xs3d, ys3d, zs3d, renderer.M)
self.set_positions((xs[0], ys[0]), (xs[1], ys[1]))
FancyArrowPatch.draw(self, renderer)
icell = np.linalg.inv(cell).T
kpoints = points
if isinstance(paths, str):
from ase.dft.kpoints import (get_special_points, special_paths, parse_path_string,
crystal_structure_from_cell)
special_points = get_special_points(cell)
structure = crystal_structure_from_cell(cell)
path_string = special_paths[structure] if paths == 'all' else paths
paths = []
for names in parse_path_string(path_string):
points = []
for name in names:
points.append(np.dot(icell.T, special_points[name]))
paths.append((names, points))
if ax is None:
fig = plt.figure(figsize=(5, 5))
ax = fig.gca(projection='3d')
azim = np.pi / 5
elev = elev or np.pi / 6
x = np.sin(azim)
y = np.cos(azim)
view = [x * np.cos(elev), y * np.cos(elev), np.sin(elev)]
bz1 = bz_vertices(icell)
maxp = 0.0
if repeat is None:
repeat = (1, 1, 1,)
dx, dy, dz = icell[0], icell[1], icell[2]
rep_x, rep_y, rep_z = repeat
if isinstance(rep_x, int):
rep_x = (0, rep_x)
if isinstance(rep_y, int):
rep_y = (0, rep_y)
if isinstance(rep_z, int):
rep_z = (0, rep_z)
c = kwargs.pop('c', 'k')
c = kwargs.pop('color', c)
for ix, iy, iz in itertools.product(range(*rep_x), range(*rep_y), range(*rep_z)):
delta = dx * ix + dy * iy + dz * iz
for points, normal in bz1:
color = c
if np.dot(normal, view) < 0:
ls = ':'
else:
ls = '-'
cosines = np.dot(icell, normal)/ np.linalg.norm(normal) / np.linalg.norm(icell, axis=1)
for idx, cosine in enumerate(cosines):
if np.abs(np.abs(cosine) - 1) < 1e-6 and False: # debugging this
tup = [rep_x, rep_y, rep_z][idx]
current = [ix, iy, iz][idx]
if cosine < 0:
current = current - 1
if tup[0] < current + 1 < tup[1]:
color = (1, 1, 1, 0)
if current + 1 != tup[1] and current != tup[0]:
ls = ':'
color='blue'
x, y, z = np.concatenate([points, points[:1]]).T
x, y, z = x + delta[0], y + delta[1], z + delta[2]
ax.plot(x, y, z, c=color, ls=ls, **kwargs)
maxp = max(maxp, points.max())
if vectors:
ax.add_artist(Arrow3D([0, icell[0, 0]],
[0, icell[0, 1]],
[0, icell[0, 2]],
mutation_scale=20, lw=1,
arrowstyle='-|>', color='k'))
ax.add_artist(Arrow3D([0, icell[1, 0]],
[0, icell[1, 1]],
[0, icell[1, 2]],
mutation_scale=20, lw=1,
arrowstyle='-|>', color='k'))
ax.add_artist(Arrow3D([0, icell[2, 0]],
[0, icell[2, 1]],
[0, icell[2, 2]],
mutation_scale=20, lw=1,
arrowstyle='-|>', color='k'))
maxp = max(maxp, 0.6 * icell.max())
if paths is not None:
for names, points in paths:
x, y, z = np.array(points).T
ax.plot(x, y, z, c='r', ls='-')
for name, point in zip(names, points):
x, y, z = point
if name == 'G':
name = '\\Gamma'
elif len(name) > 1:
name = name[0] + '_' + name[1]
ax.text(x, y, z, '$' + name + '$',
ha='center', va='bottom', color='r')
if kpoints is not None:
for p in kpoints:
ax.scatter(p[0], p[1], p[2], c='b')
if hide_ax:
ax.set_axis_off()
ax.autoscale_view(tight=True)
s = maxp / 0.5 * 0.45 * scale
ax.set_xlim(-s, s)
ax.set_ylim(-s, s)
ax.set_zlim(-s, s)
ax.set_aspect('equal')
ax.view_init(azim=azim / np.pi * 180, elev=elev / np.pi * 180)