def real_space_grid(v, grid_unit, density):
import sisl
# Create a temporary copy of the geometry
g = self.geometry.copy()
# Set new sc to create real-space grid
sc = sisl.SuperCell(
[self.xmax - self.xmin, self.ymax - self.ymin, 1000],
origo=[0, 0, -z])
g.set_sc(sc)
# Move geometry within the supercell
g = g.move([-self.xmin, -self.ymin, -np.amin(g.xyz[:, 2])])
# Make z~0 -> z = 0
g.xyz[np.where(np.abs(g.xyz[:, 2]) < 1e-3), 2] = 0
# Create the real-space grid
grid = sisl.Grid(grid_unit, sc=g.sc, geometry=g)
if density:
D = sisl.physics.DensityMatrix(g)
a = np.arange(len(D))
D.D[a, a] = v
D.density(grid)
else:
if isinstance(v, sisl.physics.electron.EigenstateElectron):
# Set parent geometry equal to the temporary one
v.parent = g
v.wavefunction(grid)
else:
# In case v is a vector
sisl.electron.wavefunction(v, grid, geometry=g)
del g
return grid
def wannier_on_grid(self, i, k=None, grid_prec=0.2, grid=None, geom=None):
'''
Projects the wannier function on a grid
'''
#all_coeffs = DataArray(self.ewf.wannR, dims=('k', 'orb', 'wannier'))
wannR = self.ewf.wannR
# Use the geometry of the hamiltonian if the user didn't provide one (usual case)
if geom is None:
geom = self.model.ham.geom
# Create a grid if the user didn't provide one
if grid is None:
grid = sisl.Grid(grid_prec, geometry=geom, dtype=complex)
# Get the coefficients of that we want
#coeffs = all_coeffs.sel(wannier=i)
coeffi = wannR[:, :, i]
#if k is None:
# coeffs = coeffs.mean('k')
#else:
# coeffs = coeffs.sel(k=k)
# Project the orbitals with the coefficients on the grid
wavefunction(coeffs, grid, geometry=geom)
return grid
def real_space_grid(geometry, sc, vector, shape, mode='wavefunction', **kwargs):
""" Create real space `sisl.Grid` object for a `sisl.Geometry` and a `sisl.Supercell`
Parameters
----------
g : sisl.Geometry
geometry
vector : numpy.array
vector to expand in the real space grid
shape: float or (3,) of int
the shape of the grid. A float specifies the grid spacing in Angstrom, while a list of integers specifies the exact grid size
(see `sisl.Grid`)
sc : sisl.SuperCell
supercell object of the grid
mode: str, optional
to build the grid from sisl.electron.wavefunction object (``mode='wavefunction'``)
or from the sisl.physics.DensityMatrix (``mode='charge'``), e.g. for charge-related plots
See Also
--------
sisl.Grid
"""
# Create a temporary copy of the geometry
g = geometry.copy()
# Set new sc to create real-space grid
g.set_sc(sc)
# Create the real-space grid
grid = sisl.Grid(shape, sc=sc, geometry=g)
if mode in ['wavefunction']:
if isinstance(vector, sisl.physics.electron.EigenstateElectron):
# Set parent geometry equal to the temporary one
vector.parent = g
vector.wavefunction(grid)
else:
# In case v is a vector
sisl.electron.wavefunction(vector, grid, geometry=g)
elif mode in ['charge']:
# The input vector v corresponds to charge-related quantities
# including spin-polarization understood as charge_up - charge_dn
D = sisl.physics.DensityMatrix(g)
a = np.arange(len(D))
D.D[a, a] = vector
D.density(grid)
if 'smooth' in kwargs:
# Smooth grid with gaussian function
if 'r_smooth' in kwargs:
r_smooth = kwargs['r_smooth']
else:
r_smooth = 0.7
grid = grid.smooth(method='gaussian', r=r_smooth)
del g
return grid
def init_func_and_attrs(self, request, siesta_test_files):
name = request.param
if name == "siesta_RHO":
init_func = sisl.get_sile(siesta_test_files("SrTiO3.RHO")).plot
attrs = {"grid_shape": (48, 48, 48)}
elif name == "complex_grid":
complex_grid_shape = (8, 10, 10)
values = np.random.random(complex_grid_shape).astype(
np.complex128) + np.random.random(complex_grid_shape) * 1j
complex_grid = sisl.Grid(complex_grid_shape, sc=1)
complex_grid.grid = values
init_func = complex_grid.plot
attrs = {"grid_shape": complex_grid_shape}
return init_func, attrs
def init_func_and_attrs(self, request, siesta_test_files, vasp_test_files):
name = request.param
if name == "siesta_RHO":
init_func = sisl.get_sile(siesta_test_files("SrTiO3.RHO")).plot
attrs = {"grid_shape": (48, 48, 48)}
if name == "VASP CHGCAR":
init_func = sisl.get_sile(
vasp_test_files(osp.join("graphene", "CHGCAR"))).plot.grid
attrs = {"grid_shape": (24, 24, 100)}
elif name == "complex_grid":
complex_grid_shape = (8, 10, 10)
np.random.seed(1)
values = np.random.random(complex_grid_shape).astype(
np.complex128) + np.random.random(complex_grid_shape) * 1j
complex_grid = sisl.Grid(complex_grid_shape, sc=1)
complex_grid.grid = values
init_func = complex_grid.plot
attrs = {"grid_shape": complex_grid_shape}
return init_func, attrs
def solve_poisson(geometry, shape, radius=2.0,
dtype=np.float64, tolerance=1e-12,
accel=None, boundary_fft=True,
device_val=None,
box=False, boundary=None, **elecs_V):
""" Solve Poisson equation """
error = False
elecs = []
for name in geometry.names:
if ('+' in name) or (name in ["Buffer", "Device"]):
continue
# This is actually an electrode
elecs.append(name)
error = error or (name not in elecs_V)
if len(elecs) == 0:
raise ValueError(f"{_script}: Could not find any electrodes in the geometry.")
error = error or len(elecs) != len(elecs_V)
if error:
for name in elecs:
if not name in elecs_V:
print(f" missing electrode bias: {name}")
raise ValueError(f"{_script}: Missing electrode arguments for specifying the bias.")
if boundary is None:
bc = [[si.Grid.PERIODIC, si.Grid.PERIODIC] for _ in range(3)]
else:
bc = []
def bc2bc(s):
return {'periodic': 'PERIODIC', 'dirichlet': 'DIRICHLET', 'neumann': 'NEUMANN',
'p': 'PERIODIC', 'd': 'DIRICHLET', 'n': 'NEUMANN'}.get(s.lower(), s.upper())
for bottom, top in boundary:
bc.append([getattr(si.Grid, bc2bc(bottom)), getattr(si.Grid, bc2bc(top))])
if len(bc) != 3:
raise ValueError(f"{_script}: Requires a 3x2 list input for the boundary conditions.")
def _create_shape_tree(xyz, A, B):
""" Takes two lists A and B which are forming a tree """
AA, BB = None, None
if len(A) == 1:
AA = si.Sphere(radius, xyz[A[0]])
if len(B) == 0:
return AA
if len(B) == 1:
BB = si.Sphere(radius, xyz[B[0]])
if len(A) == 0:
return BB
# Quick return if these are the final ones
if AA and BB:
return AA | BB
idx_A1, idx_A2 = np.array_split(A, 2)
idx_B1, idx_B2 = np.array_split(B, 2)
if not AA:
AA = _create_shape_tree(xyz, idx_A1, idx_A2)
if not BB:
BB = _create_shape_tree(xyz, idx_B1, idx_B2)
return AA | BB
# Create grid
grid = si.Grid(shape, geometry=geometry, bc=bc, dtype=dtype)
class _fake:
@property
def shape(self):
return shape
@property
def dtype(self):
return dtype
# Fake the grid to reduce memory requirement
grid.grid = _fake()
# Construct matrices we need to specify the boundary conditions on
A, b = grid.topyamg()
# Short-hand notation
xyz = geometry.xyz
if not device_val is None:
print(f"\nApplying device potential = {device_val}")
idx = geometry.names[elec]
idx_1, idx_2 = np.array_split(idx, 2)
device = _create_shape_tree(xyz, idx_1, idx_2)
del idx_1, idx_2
idx = grid.index_truncate(grid.index(device))
idx = grid.pyamg_index(idx)
grid.pyamg_fix(A, b, idx, device_val)
# Apply electrode constants
print("\nApplying electrode potentials")
for i, elec in enumerate(elecs):
V = elecs_V[elec]
print(f" - {elec} = {V}")
idx = geometry.names[elec]
idx_1, idx_2 = np.array_split(idx, 2)
elec_shape = _create_shape_tree(xyz, idx_1, idx_2)
idx = grid.index_truncate(grid.index(elec_shape))
idx = grid.pyamg_index(idx)
grid.pyamg_fix(A, b, idx, V)
del idx, idx_1, idx_2, elec_shape
# Now we have initialized both A and b with correct boundary conditions
# Lets solve the Poisson equation!
if box:
# No point in solving the boundary problem if requesting a box
boundary_fft = False
grid.grid = b.reshape(shape)
del A
else:
x = pyamg_solve(A, b, tolerance=tolerance, accel=accel)
grid.grid = x.reshape(shape)
del A, b
if boundary_fft:
# Change boundaries to always use dirichlet
# This ensures that once we set the boundaries we don't
# get any side-effects
periodic = grid.bc[:, 0] == grid.PERIODIC
bc = np.repeat(np.array([grid.DIRICHLET], np.int32), 6).reshape(3, 2)
for i in (0, 1, 2):
if periodic[i]:
bc[i, :] = grid.PERIODIC
grid.set_bc(bc)
A, b = grid.topyamg()
# Solve only for the boundary fixed
def sl2idx(grid, sl):
return grid.pyamg_index(grid.mgrid(sl))
# Create slices
sl = [slice(0, g) for g in grid.shape]
new_sl = sl[:]
# One boundary at a time
for i in (0, 1, 2):
if periodic[i]:
continue
new_sl = sl[:]
new_sl[i] = slice(0, 1)
idx = sl2idx(grid, new_sl)
grid.pyamg_fix(A, b, idx, grid.grid[new_sl[0], new_sl[1], new_sl[2]].reshape(-1))
new_sl[i] = slice(grid.shape[i] - 1, grid.shape[i])
idx = sl2idx(grid, new_sl)
grid.pyamg_fix(A, b, idx, grid.grid[new_sl[0], new_sl[1], new_sl[2]].reshape(-1))
grid.grid = _fake()
x = pyamg_solve(A, b, tolerance=tolerance, accel=accel)
grid.grid = x.reshape(shape)
del A, b
return grid
def solve_poisson(geometry, shape, radius=2.0, dtype=np.float64, tolerance=1e-12, boundary=None, **elecs_V):
""" Solve Poisson equation """
error = False
elecs = []
for name in geometry.names:
if ('+' in name) or (name in ["Buffer", "Device"]):
continue
# This is actually an electrode
elecs.append(name)
error = error or (name not in elecs_V)
if len(elecs) == 0:
raise ValueError("{}: Could not find any electrodes in the geometry.".format(_script))
error = error or len(elecs) != len(elecs_V)
if error:
for name in elecs:
if not name in elecs_V:
print(" missing electrode bias: {}".format(name))
raise ValueError("{}: Missing electrode arguments for specifying the bias.".format(_script))
if boundary is None:
bc = [[si.Grid.PERIODIC, si.Grid.PERIODIC] for _ in range(3)]
else:
bc = []
for bottom, top in boundary:
bottom = bottom.upper()
top = top.upper()
bc.append([getattr(si.Grid, bottom), getattr(si.Grid, top)])
if len(bc) != 3:
raise ValueError("{}: Requires a 3x2 list input for the boundary conditions.".format(_script))
def _create_shape_tree(xyz, A, B):
""" Takes two lists A and B which are forming a tree """
AA, BB = None, None
if len(A) == 1:
#add_A.append(A[0])
AA = si.Sphere(radius, xyz[A[0]])
if len(B) == 0:
return AA
if len(B) == 1:
#add_B.append(B[0])
BB = si.Sphere(radius, xyz[B[0]])
if len(A) == 0:
return BB
# Quick return if these are the final ones
if AA and BB:
return AA | BB
idx_A1, idx_A2 = np.array_split(A, 2)
idx_B1, idx_B2 = np.array_split(B, 2)
if not AA:
AA = _create_shape_tree(xyz, idx_A1, idx_A2)
if not BB:
BB = _create_shape_tree(xyz, idx_B1, idx_B2)
return AA | BB
#add_A = []
#add_B = []
print("Constructing electrode regions for fixing potential")
elec_shapes = []
xyz = geometry.xyz
for elec in elecs:
print(" - {}".format(elec))
# Get electrode indices and coordinates
idx_elec = geometry.names[elec]
idx_1, idx_2 = np.array_split(idx_elec, 2)
#add_A.clear()
#add_B.clear()
elec_shapes.append(_create_shape_tree(xyz, idx_1, idx_2))
#assert len(add_A) + len(add_B) == len(idx_elec)
# Create grid
grid = si.Grid(shape, geometry=geometry, bc=bc, dtype=dtype)
class _fake(object):
@property
def shape(self):
return shape
@property
def dtype(self):
return dtype
# Fake the grid to reduce memory requirement
grid.grid = _fake()
# Construct matrices we need to specify the boundary conditions on
A, b = grid.topyamg()
# Apply electrode constants
print("\nApplying electrode potentials")
for i, elec in enumerate(elecs):
print(" - {}".format(elec))
idx = grid.index_fold(grid.index(elec_shapes[i]))
idx = grid.pyamg_index(idx)
V = elecs_V[elec]
grid.pyamg_fix(A, b, idx, V)
del idx
del elec_shapes
# Now we have initialized both A and b with correct boundary conditions
# Lets solve the Poisson equation!
x = pyamg_solve(A, b, tolerance=tolerance)
grid.grid = x.reshape(shape)
del A, b
return grid
def solve_poisson(geometry, shape, radius="empirical",
dtype=np.float64, tolerance=1e-8,
accel=None, boundary_fft=True,
device_val=None, plot_boundary=False,
box=False, boundary=None, **elecs_V):
""" Solve Poisson equation """
error = False
elecs = []
for name in geometry.names:
if ('+' in name) or (name in ["Buffer", "Device"]):
continue
# This is actually an electrode
elecs.append(name)
error = error or (name not in elecs_V)
if len(elecs) == 0:
raise ValueError(f"{_script}: Could not find any electrodes in the geometry.")
error = error or len(elecs) != len(elecs_V)
if error:
for name in elecs:
if not name in elecs_V:
print(f" missing electrode bias: {name}")
raise ValueError(f"{_script}: Missing electrode arguments for specifying the bias.")
if boundary is None:
bc = [[si.Grid.PERIODIC, si.Grid.PERIODIC] for _ in range(3)]
else:
bc = []
def bc2bc(s):
return {'periodic': 'PERIODIC', 'p': 'PERIODIC', si.Grid.PERIODIC: 'PERIODIC',
'dirichlet': 'DIRICHLET', 'd': 'DIRICHLET', si.Grid.DIRICHLET: 'DIRICHLET',
'neumann': 'NEUMANN', 'n': 'NEUMANN',si.Grid.NEUMANN: 'NEUMANN',
}.get(s.lower(), s.upper())
for bottom, top in boundary:
bc.append([getattr(si.Grid, bc2bc(bottom)), getattr(si.Grid, bc2bc(top))])
if len(bc) != 3:
raise ValueError(f"{_script}: Requires a 3x2 list input for the boundary conditions.")
def _create_shape_tree(xyz, A, B=None):
""" Takes two lists A and B which returns a shape with a binary nesting
This makes further index handling much faster.
"""
if B is None or len(B) == 0:
return _create_shape_tree(xyz, *np.array_split(A, 2))
AA, BB = None, None
if len(A) == 1:
AA = si.Sphere(radius, xyz[A[0]])
if len(B) == 0:
return AA
if len(B) == 1:
BB = si.Sphere(radius, xyz[B[0]])
if len(A) == 0:
return BB
# Quick return if these are the final ones
if AA and BB:
return AA | BB
if not AA:
AA = _create_shape_tree(xyz, *np.array_split(A, 2))
if not BB:
BB = _create_shape_tree(xyz, *np.array_split(B, 2))
return AA | BB
# Create grid
grid = si.Grid(shape, geometry=geometry, bc=bc, dtype=dtype)
class _fake:
@property
def shape(self):
return shape
@property
def dtype(self):
return dtype
# Fake the grid to reduce memory requirement
grid.grid = _fake()
# Construct matrices we need to specify the boundary conditions on
A, b = grid.topyamg()
# Short-hand notation
xyz = geometry.xyz
if not device_val is None:
print(f"\nApplying device potential = {device_val}")
idx = geometry.names["Device"]
device = _create_shape_tree(xyz, idx)
idx = grid.index_truncate(grid.index(device))
idx = grid.pyamg_index(idx)
grid.pyamg_fix(A, b, idx, device_val)
# Apply electrode constants
print("\nApplying electrode potentials")
for i, elec in enumerate(elecs):
V = elecs_V[elec]
print(f" - {elec} = {V}")
idx = geometry.names[elec]
elec_shape = _create_shape_tree(xyz, idx)
idx = grid.index_truncate(grid.index(elec_shape))
idx = grid.pyamg_index(idx)
grid.pyamg_fix(A, b, idx, V)
del idx, elec_shape
# Now we have initialized both A and b with correct boundary conditions
# Lets solve the Poisson equation!
if box:
# No point in solving the boundary problem if requesting a box
boundary_fft = False
grid.grid = b.reshape(shape)
del A
else:
x = pyamg_solve(A, b, tolerance=tolerance, accel=accel,
title="solving electrode boundary conditions")
grid.grid = x.reshape(shape)
del A, b
if boundary_fft:
# Change boundaries to always use dirichlet
# This ensures that once we set the boundaries we don't
# get any side-effects
periodic = grid.bc[:, 0] == grid.PERIODIC
bc = np.repeat(np.array([grid.DIRICHLET], np.int32), 6).reshape(3, 2)
for i in (0, 1, 2):
if periodic[i]:
bc[i, :] = grid.PERIODIC
grid.set_bc(bc)
A, b = grid.topyamg()
# Solve only for the boundary fixed
def sl2idx(grid, sl):
return grid.pyamg_index(grid.mgrid(sl))
# Create slices
sl = [slice(0, g) for g in grid.shape]
new_sl = sl[:]
# One boundary at a time
for i in (0, 1, 2):
if periodic[i]:
continue
new_sl = sl[:]
new_sl[i] = slice(0, 1)
idx = sl2idx(grid, new_sl)
grid.pyamg_fix(A, b, idx, grid.grid[new_sl[0], new_sl[1], new_sl[2]].reshape(-1))
new_sl[i] = slice(grid.shape[i] - 1, grid.shape[i])
idx = sl2idx(grid, new_sl)
grid.pyamg_fix(A, b, idx, grid.grid[new_sl[0], new_sl[1], new_sl[2]].reshape(-1))
if plot_boundary:
dat = b.reshape(*grid.shape)
# now plot every plane
import matplotlib.pyplot as plt
slicex3 = np.index_exp[:] * 3
axs = [
np.linspace(0, grid.sc.length[ax], shape, endpoint=False)
for ax, shape in enumerate(grid.shape)
]
for i in (0, 1, 2):
idx = list(slicex3)
j = (i + 1) % 3
k = (i + 2) % 3
if i > j:
i, j = j, i
X, Y = np.meshgrid(axs[i], axs[j])
for v, head in ((0, "bottom"), (-1, "top")):
plt.figure()
plt.title(f"axis: {'ABC'[k]} ({head})")
idx[k] = v
plt.contourf(X, Y, dat[tuple(idx)].T)
plt.xlabel(f"Distance along {'ABC'[i]} [Ang]")
plt.ylabel(f"Distance along {'ABC'[j]} [Ang]")
plt.colorbar()
plt.show()
grid.grid = _fake()
x = pyamg_solve(A, b, tolerance=tolerance, accel=accel,
title="removing electrode boundaries and solving for edge fixing")
grid.grid = x.reshape(shape)
del A, b
return grid
plot.update_settings(axes=axes,
transforms=["cos"],
represent=representation)
# Check that transforms = ["cos"] applies np.cos
assert np.allclose(
self._get_plotted_values(),
np.cos(new_grid.grid).mean(axis=tuple(ax for ax in [0, 1, 2]
if ax not in axes)))
# Generate a complex-valued grid for testing
complex_grid_shape = (8, 10, 10)
values = np.random.random(complex_grid_shape).astype(
np.complex128) + np.random.random(complex_grid_shape) * 1j
complex_grid = sisl.Grid(complex_grid_shape, sc=1)
complex_grid.grid = values
class TestGridPlot(GridPlotTester):
run_for = {
# Test with a siesta RHO file
"siesta_RHO": {
"plot_file": osp.join(_dir, "SrTiO3.RHO"),
"grid_shape": (48, 48, 48)
},
"complex_grid": {
"init_func": complex_grid.plot,
"grid_shape": complex_grid_shape
}