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 }