def find_k(self, p, omega, band_min, band_max, korig_and_kdir, tol, kmag_guess, kmag_min, kmag_max, *band_funcs): num_bands_save = self.num_bands kpoints_save = self.k_points nb = band_max - band_min + 1 kdir = korig_and_kdir[1] if type( korig_and_kdir) is list else korig_and_kdir lat = self.geometry_lattice kdir1 = mp.cartesian_to_reciprocal( mp.reciprocal_to_cartesian(kdir, lat).unit(), lat) if type(korig_and_kdir) is list: korig = korig_and_kdir[0] else: korig = mp.Vector3() # k0s is a list caching the best k value found for each band: if type(kmag_guess) is list: k0s = kmag_guess else: k0s = [kmag_guess] * (band_max - band_min + 1) # dict to memoize all "band: k" results bktab = {} def rootfun(b): def _rootfun(k): # First, look in the cached table tab_val = bktab.get((b, k), None) if tab_val: print("find-k {} at {}: {} (cached)".format( b, k, tab_val[0])) return tab_val # Otherwise, compute bands and cache results else: self.num_bands = b self.k_points = [korig + kdir1.scale(k)] self.run_parity(p, False) v = self.mode_solver.compute_group_velocity_component( kdir1) # Cache computed values for _b, _f, _v in zip(range(band_min, b - band_min + 1, 1), self.freqs[band_min - 1:], v[band_min - 1:]): tabval = bktab.get((_b, k0s[_b - band_min]), None) if not tabval or abs(_f - omega) < abs(tabval[0]): k0s[_b - band_min + 1] = k bktab[(_b, k)] = (_f - omega, _v) fun = self.freqs[-1] - omega print("find-k {} at {}: {}".format(b, k, fun)) return (fun, v[-1]) return _rootfun # Don't let previous computations interfere if self.mode_solver: self.randomize_fields() ks = [] for b in range(band_max, band_max - nb, -1): ks.append( mp.find_root_deriv(rootfun(b), tol, kmag_min, kmag_max, k0s[b - band_min])) if band_funcs: for b, k in zip(range(1, band_max + 1), reversed(ks)): self.num_bands = b self.k_points = [korig + kdir1.scale(k)] def bfunc(ms, b_prime): if b_prime == b: for f in band_funcs: apply_band_func_thunk(ms, f, b, True) self.run_parity(p, False, bfunc) self.num_bands = num_bands_save self.k_points = kpoints_save ks = list(reversed(ks)) print("{}kvals:, {}, {}, {}".format(self.parity, omega, band_min, band_max), end='') for k in korig: print(", {}".format(k), end='') for k in kdir1: print(", {}".format(k), end='') for k in ks: print(", {}".format(k), end='') print() return ks
def find_k(self, p, omega, band_min, band_max, korig_and_kdir, tol, kmag_guess, kmag_min, kmag_max, *band_funcs): num_bands_save = self.num_bands kpoints_save = self.k_points nb = band_max - band_min + 1 kdir = korig_and_kdir[1] if type(korig_and_kdir) is list else korig_and_kdir lat = self.geometry_lattice kdir1 = mp.cartesian_to_reciprocal(mp.reciprocal_to_cartesian(kdir, lat).unit(), lat) if type(korig_and_kdir) is list: korig = korig_and_kdir[0] else: korig = mp.Vector3() # k0s is a list caching the best k value found for each band: if type(kmag_guess) is list: k0s = kmag_guess else: k0s = [kmag_guess] * (band_max - band_min + 1) # dict to memoize all "band: k" results bktab = {} def rootfun(b): def _rootfun(k): # First, look in the cached table tab_val = bktab.get((b, k), None) if tab_val: print("find-k {} at {}: {} (cached)".format(b, k, tab_val[0])) return tab_val # Otherwise, compute bands and cache results else: self.num_bands = b self.k_points = [korig + kdir1.scale(k)] self.run_parity(p, False) v = self.mode_solver.compute_group_velocity_component(kdir1) # Cache computed values for _b, _f, _v in zip(range(band_min, b - band_min + 1, 1), self.freqs[band_min - 1:], v[band_min - 1:]): tabval = bktab.get((_b, k0s[_b - band_min]), None) if not tabval or abs(_f - omega) < abs(tabval[0]): k0s[_b - band_min + 1] = k bktab[(_b, k)] = (_f - omega, _v) fun = self.freqs[-1] - omega print("find-k {} at {}: {}".format(b, k, fun)) return (fun, v[-1]) return _rootfun # Don't let previous computations interfere if self.mode_solver: self.randomize_fields() ks = [] for b in range(band_max, band_max - nb, -1): ks.append(mp.find_root_deriv(rootfun(b), tol, kmag_min, kmag_max, k0s[b - band_min])) if band_funcs: for b, k in zip(range(band_max, band_max - nb, -1), reversed(ks)): self.num_bands = b self.k_points = [korig + kdir1.scale(k)] def bfunc(ms, b_prime): if b_prime == b: for f in band_funcs: apply_band_func_thunk(ms, f, b, True) self.run_parity(p, False, bfunc) self.num_bands = num_bands_save self.k_points = kpoints_save ks = list(reversed(ks)) print("{}kvals:, {}, {}, {}".format(self.parity, omega, band_min, band_max), end='') for k in korig: print(", {}".format(k), end='') for k in kdir1: print(", {}".format(k), end='') for k in ks: print(", {}".format(k), end='') print() return ks