| 1003 | run_tm_yodd = run_yodd_zodd |
| 1004 | |
| 1005 | def find_k( |
| 1006 | self, |
| 1007 | p, |
| 1008 | omega, |
| 1009 | band_min, |
| 1010 | band_max, |
| 1011 | korig_and_kdir, |
| 1012 | tol, |
| 1013 | kmag_guess, |
| 1014 | kmag_min, |
| 1015 | kmag_max, |
| 1016 | *band_funcs, |
| 1017 | ): |
| 1018 | |
| 1019 | num_bands_save = self.num_bands |
| 1020 | kpoints_save = self.k_points |
| 1021 | nb = band_max - band_min + 1 |
| 1022 | kdir = korig_and_kdir[1] if type(korig_and_kdir) is list else korig_and_kdir |
| 1023 | lat = self.geometry_lattice |
| 1024 | kdir1 = mp.cartesian_to_reciprocal( |
| 1025 | mp.reciprocal_to_cartesian(kdir, lat).unit(), lat |
| 1026 | ) |
| 1027 | |
| 1028 | if type(korig_and_kdir) is list: |
| 1029 | korig = korig_and_kdir[0] |
| 1030 | else: |
| 1031 | korig = mp.Vector3() |
| 1032 | |
| 1033 | # k0s is a list caching the best k value found for each band: |
| 1034 | if type(kmag_guess) is list: |
| 1035 | k0s = kmag_guess |
| 1036 | else: |
| 1037 | k0s = [kmag_guess] * (band_max - band_min + 1) |
| 1038 | |
| 1039 | # dict to memoize all "band: k" results |
| 1040 | bktab = {} |
| 1041 | |
| 1042 | def rootfun(b): |
| 1043 | def _rootfun(k): |
| 1044 | # First, look in the cached table |
| 1045 | tab_val = bktab.get((b, k), None) |
| 1046 | if tab_val: |
| 1047 | if verbosity.mpb > 0: |
| 1048 | print(f"find-k {b} at {k}: {tab_val[0]} (cached)") |
| 1049 | return tab_val |
| 1050 | else: |
| 1051 | self.num_bands = b |
| 1052 | self.k_points = [korig + kdir1.scale(k)] |
| 1053 | self.run_parity(p, False) |
| 1054 | v = self.mode_solver.compute_group_velocity_component(kdir1) |
| 1055 | |
| 1056 | # Cache computed values |
| 1057 | for _b, _f, _v in zip( |
| 1058 | range(band_min, b - band_min + 1), |
| 1059 | self.freqs[band_min - 1 :], |
| 1060 | v[band_min - 1 :], |
| 1061 | ): |
| 1062 | tabval = bktab.get((_b, k0s[_b - band_min]), None) |