from __future__ import division, print_function import functools import math import os import numbers import re import sys import time import h5py import numpy as np import meep as mp from . import mode_solver, with_hermitian_epsilon from meep.simulation import get_num_args try: basestring except NameError: basestring = str U_MIN = 0 U_PROD = 1 U_SUM = 2 class MPBArray(np.ndarray): def __new__(cls, input_array, lattice, kpoint=None, bloch_phase=False): # Input array is an already formed ndarray instance # We first cast to be our class type obj = np.asarray(input_array).view(cls) # add the new properties to the created instance obj.lattice = lattice obj.kpoint = kpoint obj.bloch_phase = bloch_phase # Finally, we must return the newly created object: return obj def __array_finalize__(self, obj): # ``self`` is a new object resulting from # ndarray.__new__(MPBArray, ...), therefore it only has # attributes that the ndarray.__new__ constructor gave it - # i.e. those of a standard ndarray. # We could have got to the ndarray.__new__ call in 3 ways: # From an explicit constructor - e.g. MPBArray(lattice): # obj is None # (we're in the middle of the MPBArray.__new__ # constructor, and self.lattice will be set when we return to # MPBArray.__new__) if obj is None: return # From view casting - e.g arr.view(MPBArray): # obj is arr # (type(obj) can be MPBArray) # From new-from-template - e.g mpbarr[:3] # type(obj) is MPBArray # # Note that it is here, rather than in the __new__ method, # that we set the default value for 'lattice', because this # method sees all creation of default objects - with the # MPBArray.__new__ constructor, but also with # arr.view(MPBArray). self.lattice = getattr(obj, 'lattice', None) self.kpoint = getattr(obj, 'kpoint', None) self.bloch_phase = getattr(obj, 'bloch_phase', False) class ModeSolver(object): def __init__(self, resolution=10, is_negative_epsilon_ok=False, eigensolver_flops=0, eigensolver_flags=68, use_simple_preconditioner=False, force_mu=False, mu_input_file='', epsilon_input_file='', mesh_size=3, target_freq=0.0, tolerance=1.0e-7, num_bands=1, k_points=[], ensure_periodicity=True, geometry=[], geometry_lattice=mp.Lattice(), geometry_center=mp.Vector3(0, 0, 0), default_material=mp.Medium(epsilon=1), dimensions=3, random_fields=False, filename_prefix='', deterministic=False, verbose=False, optimize_grid_size=True, eigensolver_nwork=3, eigensolver_block_size=-11): self.resolution = resolution self.is_negative_epsilon_ok = is_negative_epsilon_ok self.eigensolver_flops = eigensolver_flops self.eigensolver_nwork = eigensolver_nwork self.eigensolver_block_size = eigensolver_block_size self.eigensolver_flags = eigensolver_flags self.use_simple_preconditioner = use_simple_preconditioner self.force_mu = force_mu self.mu_input_file = mu_input_file self.epsilon_input_file = epsilon_input_file self.mesh_size = mesh_size self.target_freq = target_freq self.tolerance = tolerance self.num_bands = num_bands self.k_points = k_points self.ensure_periodicity = ensure_periodicity self.geometry = geometry self.geometry_lattice = geometry_lattice self.geometry_center = geometry_center self.default_material = default_material self.dimensions = dimensions self.random_fields = random_fields self.filename_prefix = filename_prefix self.deterministic = deterministic self.verbose = verbose self.optimize_grid_size = optimize_grid_size self.eigensolver_nwork = eigensolver_nwork self.eigensolver_block_size = eigensolver_block_size self.parity = '' self.iterations = 0 self.all_freqs = None self.freqs = [] self.band_range_data = [] self.eigensolver_flops = 0 self.total_run_time = 0 self.current_k = mp.Vector3() self.k_split_num = 1 self.k_split_index = 0 self.eigensolver_iters = [] self.mode_solver = None @property def resolution(self): return self._resolution @resolution.setter def resolution(self, val): if isinstance(val, numbers.Number): self._resolution = [val, val, val] elif isinstance(val, mp.Vector3): self._resolution = [val.x, val.y, val.z] else: t = type(val) raise TypeError("resolution must be a number or a Vector3: Got {}".format(t)) def allow_negative_epsilon(self): self.is_negative_epsilon_ok = True self.target_freq = 1 / mp.inf def get_filename_prefix(self): if self.filename_prefix: return self.filename_prefix + '-' else: _, filename = os.path.split(sys.argv[0]) if filename == 'ipykernel_launcher.py' or filename == '__main__.py': return '' else: return re.sub(r'\.py$', '', filename) + '-' def get_freqs(self): return self.mode_solver.get_freqs() def multiply_bloch_phase(self, arr): dims = arr.shape arr = arr.ravel() self.mode_solver.multiply_bloch_phase(arr) return np.reshape(arr, dims) def get_poynting(self, which_band): e = self.get_efield(which_band, False).ravel() h = self.get_hfield(which_band, False).ravel() # Reshape into rows of vector3s e = e.reshape((int(e.shape[0] / 3), 3)) h = h.reshape((int(h.shape[0] / 3), 3)) res = np.zeros(e.shape, dtype=np.complex128) def ExH(e, h): ev = mp.Vector3(e[0], e[1], e[2]) hv = mp.Vector3(h[0], h[1], h[2]) return ev.conj().cross(hv) for i in range(e.shape[0]): res[i] = np.array(ExH(e[i], h[i])) flat_res = res.ravel() self.mode_solver.set_curfield_cmplx(flat_res) self.mode_solver.set_curfield_type('v') return MPBArray(res, self.get_lattice, self.current_k) def get_epsilon(self): self.mode_solver.get_epsilon() return self.get_curfield_as_array(False) def get_mu(self): self.mode_solver.get_mu() return self.get_curfield_as_array(False) def get_bfield(self, which_band, bloch_phase=True): return self._get_field('b', which_band, bloch_phase) def get_efield(self, which_band, bloch_phase=True): return self._get_field('e', which_band, bloch_phase) def get_dfield(self, which_band, bloch_phase=True): return self._get_field('d', which_band, bloch_phase) def get_hfield(self, which_band, bloch_phase=True): return self._get_field('h', which_band, bloch_phase) def get_charge_density(self, which_band): self.get_efield(which_band) self.mode_solver.compute_field_divergence() def _get_field(self, f, band, bloch_phase): if self.mode_solver is None: raise ValueError("Must call a run function before attempting to get a field") if f == 'b': self.mode_solver.get_bfield(band) elif f == 'd': self.mode_solver.get_dfield(band) elif f == 'e': self.mode_solver.get_efield(band) elif f == 'h': self.mode_solver.get_hfield(band) dims = self.mode_solver.get_dims() while len(dims) < 3: dims += [1] dims += [3] arr = np.zeros(np.prod(dims), np.complex128) if bloch_phase: self.mode_solver.multiply_bloch_phase() self.mode_solver.get_curfield_cmplx(arr) arr = np.reshape(arr, dims) res = MPBArray(arr, self.get_lattice(), self.current_k, bloch_phase=bloch_phase) return res def get_curfield_as_array(self, bloch_phase=True): dims = self.mode_solver.get_dims() arr = np.zeros(np.prod(dims)) self.mode_solver.get_curfield(arr) arr = np.reshape(arr, dims) return MPBArray(arr, self.get_lattice(), self.current_k, bloch_phase=bloch_phase) def get_dpwr(self, band): self.get_dfield(band, False) self.compute_field_energy() return self.get_curfield_as_array(False) def get_bpwr(self, band): self.get_bfield(band, False) self.compute_field_energy() return self.get_curfield_as_array(False) def fix_field_phase(self): self.mode_solver.fix_field_phase() def get_epsilon_point(self, p): return self.mode_solver.get_epsilon_point(p) def get_epsilon_inverse_tensor_point(self, p): return self.mode_solver.get_epsilon_inverse_tensor_point(p) def get_energy_point(self, p): return self.mode_solver.get_energy_point(p) def get_field_point(self, p): return self.mode_solver.get_field_point(p) def get_bloch_field_point(self, p): return self.mode_solver.get_bloch_field_point(p) def get_tot_pwr(self, which_band): epwr = self.get_dpwr(which_band) hpwr = self.get_bpwr(which_band) tot_pwr = epwr + hpwr self.mode_solver.set_curfield(tot_pwr.ravel()) self.mode_solver.set_curfield_type('R') return MPBArray(tot_pwr, self.get_lattice(), self.current_k, bloch_phase=False) def get_eigenvectors(self, first_band, num_bands): dims = self.mode_solver.get_eigenvectors_slice_dims(num_bands) ev = np.zeros(np.prod(dims), dtype=np.complex128) self.mode_solver.get_eigenvectors(first_band - 1, num_bands, ev) return MPBArray(ev.reshape(dims), self.get_lattice(), self.current_k) def set_eigenvectors(self, ev, first_band): self.mode_solver.set_eigenvectors(first_band - 1, ev.flatten()) def save_eigenvectors(self, filename): with h5py.File(filename, 'w') as f: ev = self.get_eigenvectors(1, self.num_bands) f['rawdata'] = ev def load_eigenvectors(self, filename): with h5py.File(filename, 'r') as f: ev = f['rawdata'].value self.set_eigenvectors(ev, 1) self.mode_solver.curfield_reset() # The band-range-data is a list of tuples, each consisting of a (min, k-point) # tuple and a (max, k-point) tuple, with each min/max pair describing the # frequency range of a band and the k-points where it achieves its minimum/maximum. # Here, we update this data with a new list of band frequencies, and return the new # data. If band-range-data is null or too short, the needed entries will be created. def update_band_range_data(self, brd, freqs, kpoint): def update_brd(brd, freqs, br_start): if not freqs: return br_start + brd else: br = ((mp.inf, -1), (-mp.inf, -1)) if not brd else brd[0] br_rest = [] if not brd else brd[1:] newmin = (freqs[0], kpoint) if freqs[0] < br[0][0] else br[0] newmax = (freqs[0], kpoint) if freqs[0] > br[1][0] else br[1] new_start = br_start + [(newmin, newmax)] return update_brd(br_rest, freqs[1:], new_start) return update_brd(brd, freqs, []) def output_band_range_data(self, br_data): for tup, band in zip(br_data, range(1, len(br_data) + 1)): fmt = "Band {} range: {} at {} to {} at {}" min_band, max_band = tup min_freq, min_kpoint = min_band max_freq, max_kpoint = max_band print(fmt.format(band, min_freq, min_kpoint, max_freq, max_kpoint)) # Output any gaps in the given band ranges, and return a list of the gaps as # a list of (percent, freq-min, freq-max) tuples. def output_gaps(self, br_data): def ogaps(br_cur, br_rest, i, gaps): if not br_rest: ordered_gaps = [] gaps = list(reversed(gaps)) for i in range(0, len(gaps), 3): ordered_gaps.append((gaps[i + 2], gaps[i + 1], gaps[i])) return ordered_gaps else: br_rest_min_f = br_rest[0][0][0] br_cur_max_f = br_cur[1][0] if br_cur_max_f >= br_rest_min_f: return ogaps(br_rest[0], br_rest[1:], i + 1, gaps) else: gap_size = ((200 * (br_rest_min_f - br_cur_max_f)) / (br_rest_min_f + br_cur_max_f)) fmt = "Gap from band {} ({}) to band {} ({}), {}%" print(fmt.format(i, br_cur_max_f, i + 1, br_rest_min_f, gap_size)) return ogaps(br_rest[0], br_rest[1:], i + 1, [gap_size, br_cur_max_f, br_rest_min_f] + gaps) if not br_data: return [] else: return ogaps(br_data[0], br_data[1:], 1, []) # Return the frequency gap from the band #lower-band to the band # #(lower-band+1), as a percentage of mid-gap frequency. The "gap" # may be negative if the maximum of the lower band is higher than the # minimum of the upper band. (The gap is computed from the # band-range-data of the previous run.) def retrieve_gap(self, lower_band): if lower_band + 1 > len(self.band_range_data): raise ValueError("retrieve-gap called for higher band than was calculated") f1 = self.band_range_data[lower_band - 1][1][0] f2 = self.band_range_data[lower_band][0][0] return (f2 - f1) / (0.005 * (f1 + f2)) # Split a list L into num more-or-less equal pieces, returning the piece # given by index (in 0..num-1), along with the index in L of the first # element of the piece, as a list: [first-index, piece-of-L] def list_split(self, l, num, index): def list_sub(l, start, length, index, rest): if not l: return list(reversed(rest)) if index >= start and index < (start + length): return list_sub(l[1:], start, length, index + 1, [l[0]] + rest) else: return list_sub(l[1:], start, length, index + 1, rest) if index >= num or index < 0: return (len(l), []) else: block_size = (len(l) + num - 1) // num start = index * block_size length = min(block_size, (len(l) - index * block_size)) return (start, list_sub(l, start, length, 0, [])) def get_lattice(self): if self.mode_solver is None: raise RuntimeError("Must call ModeSolver.run before getting the lattice.") lattice = np.zeros((3, 3)) self.mode_solver.get_lattice(lattice) return lattice def output_field(self): self.output_field_to_file(mp.ALL, self.get_filename_prefix()) def output_field_x(self): self.output_field_to_file(0, self.get_filename_prefix()) def output_field_y(self): self.output_field_to_file(1, self.get_filename_prefix()) def output_field_z(self): self.output_field_to_file(2, self.get_filename_prefix()) def output_epsilon(self): self.mode_solver.get_epsilon() self.output_field_to_file(mp.ALL, self.get_filename_prefix()) def output_mu(self): self.mode_solver.get_mu() self.output_field_to_file(mp.ALL, self.get_filename_prefix()) def output_field_to_file(self, component, fname_prefix): curfield_type = self.mode_solver.get_curfield_type() output_k = self.mode_solver.get_output_k() if curfield_type in 'Rv': # Generic scalar/vector field. Don't know k output_k = [0, 0, 0] if curfield_type in 'dhbecv': self._output_vector_field(curfield_type, fname_prefix, output_k, component) elif curfield_type == 'C': self._output_complex_scalar_field(fname_prefix, output_k) elif curfield_type in 'DHBnmR': self._output_scalar_field(curfield_type, fname_prefix) else: raise ValueError("Unkown field type: {}".format(curfield_type)) self.mode_solver.curfield_reset() def _output_complex_scalar_field(self, fname_prefix, output_k): curfield_type = 'C' kpoint_index = self.mode_solver.get_kpoint_index() curfield_band = self.mode_solver.curfield_band fname = "{}.k{:02d}.b{:02d}".format(curfield_type, kpoint_index, curfield_band) description = "{} field, kpoint {}, band {}, freq={:.6g}".format( curfield_type, kpoint_index, curfield_band, self.freqs[curfield_band - 1] ) fname = self._create_fname(fname, fname_prefix, True) print("Outputting complex scalar field to {}...".format(fname)) with h5py.File(fname, 'w') as f: f['description'] = description.encode() f['Bloch wavevector'] = np.array(output_k) self._write_lattice_vectors(f) dims = self.mode_solver.get_dims() field = np.empty(np.prod(dims), np.complex128) self.mode_solver.get_curfield_cmplx(field) reshaped_field = field.reshape(dims) f['c.r'] = np.real(reshaped_field) f['c.i'] = np.imag(reshaped_field) def _output_vector_field(self, curfield_type, fname_prefix, output_k, component): components = ['x', 'y', 'z'] kpoint_index = self.mode_solver.get_kpoint_index() curfield_band = self.mode_solver.curfield_band fname = "{}.k{:02d}.b{:02d}".format(curfield_type, kpoint_index, curfield_band) if component >= 0: fname += ".{}".format(components[component]) description = "{} field, kpoint {}, band {}, freq={:.6g}".format( curfield_type, kpoint_index, curfield_band, self.freqs[curfield_band - 1] ) fname = self._create_fname(fname, fname_prefix, True) print("Outputting fields to {}...".format(fname)) with h5py.File(fname, 'w') as f: f['description'] = description.encode() f['Bloch wavevector'] = np.array(output_k) self._write_lattice_vectors(f) if curfield_type != 'v': self.mode_solver.multiply_bloch_phase() for c_idx, c in enumerate(components): if component >= 0 and c_idx != component: continue dims = self.mode_solver.get_dims() field = np.empty(np.prod(dims) * 3, np.complex128) self.mode_solver.get_curfield_cmplx(field) component_field = field[c_idx::3].reshape(dims) name = "{}.r".format(c) f[name] = np.real(component_field) name = "{}.i".format(c) f[name] = np.imag(component_field) def _output_scalar_field(self, curfield_type, fname_prefix): components = ['x', 'y', 'z'] if curfield_type == 'n': fname = 'epsilon' description = 'dielectric function, epsilon' elif curfield_type == 'm': fname = 'mu' description = 'permeability mu' else: kpoint_index = self.mode_solver.get_kpoint_index() curfield_band = self.mode_solver.curfield_band fname = "{}pwr.k{:02d}.b{:02d}".format(curfield_type.lower(), kpoint_index, curfield_band) descr_fmt = "{} field energy density, kpoint {}, band {}, freq={:.6g}" description = descr_fmt.format(curfield_type, kpoint_index, curfield_band, self.freqs[curfield_band - 1]) parity_suffix = False if curfield_type in 'mn' else True fname = self._create_fname(fname, fname_prefix, parity_suffix) print("Outputting {}...".format(fname)) with h5py.File(fname, 'w') as f: f['description'] = description.encode() self._create_h5_dataset(f, 'data') self._write_lattice_vectors(f) if curfield_type == 'n': for inv in [False, True]: inv_str = 'epsilon_inverse' if inv else 'epsilon' for c1 in range(3): for c2 in range(c1, 3): self.mode_solver.get_epsilon_tensor(c1, c2, 0, inv) dataname = "{}.{}{}".format(inv_str, components[c1], components[c2]) self._create_h5_dataset(f, dataname) if with_hermitian_epsilon() and c1 != c2: self.mode_solver.get_epsilon_tensor(c1, c2, 1, inv) dataname += '.i' self._create_h5_dataset(f, dataname) def _write_lattice_vectors(self, h5file): lattice = np.zeros((3, 3)) self.mode_solver.get_lattice(lattice) h5file['lattice vectors'] = lattice def _create_h5_dataset(self, h5file, key): h5file[key] = self.get_curfield_as_array(False) def _create_fname(self, fname, prefix, parity_suffix): parity_str = self.mode_solver.get_parity_string() if parity_suffix and parity_str: suffix = ".{}".format(parity_str) else: suffix = '' return prefix + fname + suffix + '.h5' def compute_field_energy(self): return self.mode_solver.compute_field_energy() def compute_field_divergence(self): mode_solver.compute_field_divergence() def compute_energy_in_objects(self, objs): return self.mode_solver.compute_energy_in_objects(objs) def compute_energy_in_dielectric(self, eps_low, eps_high): return self.mode_solver.compute_energy_in_dielectric(eps_low, eps_high) def compute_energy_integral(self, f): return self.mode_solver.compute_energy_integral(f) def compute_field_integral(self, f): return self.mode_solver.compute_field_integral(f) def compute_group_velocities(self): xarg = mp.cartesian_to_reciprocal(mp.Vector3(1), self.geometry_lattice) vx = self.mode_solver.compute_group_velocity_component(xarg) yarg = mp.cartesian_to_reciprocal(mp.Vector3(y=1), self.geometry_lattice) vy = self.mode_solver.compute_group_velocity_component(yarg) zarg = mp.cartesian_to_reciprocal(mp.Vector3(z=1), self.geometry_lattice) vz = self.mode_solver.compute_group_velocity_component(zarg) return [mp.Vector3(x, y, z) for x, y, z in zip(vx, vy, vz)] def compute_group_velocity_component(self, direction): return self.mode_solver.compute_group_velocity_component(direction) def compute_one_group_velocity(self, which_band): return self.mode_solver.compute_1_group_velocity(which_band) def compute_one_group_velocity_component(self, direction, which_band): return self.mode_solver.compute_1_group_velocity_component(direction, which_band) def compute_zparities(self): return self.mode_solver.compute_zparities() def compute_yparities(self): return self.mode_solver.compute_yparities() def randomize_fields(self): self.mode_solver.randomize_fields() def display_kpoint_data(self, name, data): k_index = self.mode_solver.get_kpoint_index() print("{}{}:, {}".format(self.parity, name, k_index), end='') for d in data: print(", {}".format(d), end='') print() def display_eigensolver_stats(self): num_runs = len(self.eigensolver_iters) if num_runs <= 0: return min_iters = min(self.eigensolver_iters) max_iters = max(self.eigensolver_iters) mean_iters = np.mean(self.eigensolver_iters) fmt = "eigensolver iterations for {} kpoints: {}-{}, mean = {}" print(fmt.format(num_runs, min_iters, max_iters, mean_iters), end='') sorted_iters = sorted(self.eigensolver_iters) idx1 = num_runs // 2 idx2 = ((num_runs + 1) // 2) - 1 median_iters = 0.5 * (sorted_iters[idx1] + sorted_iters[idx2]) print(", median = {}".format(median_iters)) mean_flops = self.eigensolver_flops / (num_runs * mean_iters) print("mean flops per iteration = {}".format(mean_flops)) mean_time = self.total_run_time / (mean_iters * num_runs) print("mean time per iteration = {} s".format(mean_time)) def _get_grid_size(self): grid_size = mp.Vector3(self.resolution[0] * self.geometry_lattice.size.x, self.resolution[1] * self.geometry_lattice.size.y, self.resolution[2] * self.geometry_lattice.size.z) grid_size.x = max(math.ceil(grid_size.x), 1) grid_size.y = max(math.ceil(grid_size.y), 1) grid_size.z = max(math.ceil(grid_size.z), 1) return grid_size def _optimize_grid_size(self): self.grid_size.x = self.next_factor2357(self.grid_size.x) self.grid_size.y = self.next_factor2357(self.grid_size.y) self.grid_size.z = self.next_factor2357(self.grid_size.z) def next_factor2357(self, n): def is_factor2357(n): def divby(n, p): if n % p == 0: return divby(n // p, p) return n return divby(divby(divby(divby(n, 2), 3), 5), 7) == 1 if is_factor2357(n): return n return self.next_factor2357(n + 1) def init_params(self, p, reset_fields): self.grid_size = self._get_grid_size() if self.optimize_grid_size: self._optimize_grid_size() self.mode_solver = mode_solver( self.num_bands, p, self.resolution, self.geometry_lattice, self.tolerance, self.mesh_size, self.default_material, self.geometry, True if reset_fields else False, self.deterministic, self.target_freq, self.dimensions, self.verbose, self.ensure_periodicity, self.eigensolver_flops, self.is_negative_epsilon_ok, self.epsilon_input_file, self.mu_input_file, self.force_mu, self.use_simple_preconditioner, self.grid_size, self.eigensolver_nwork, self.eigensolver_block_size, ) def set_parity(self, p): self.mode_solver.set_parity(p) def solve_kpoint(self, k): self.mode_solver.solve_kpoint(k) def run_parity(self, p, reset_fields, *band_functions): if self.random_fields and self.randomize_fields not in band_functions: band_functions.append(self.randomize_fields) start = time.time() self.all_freqs = np.zeros((len(self.k_points), self.num_bands)) self.band_range_data = [] init_time = time.time() print("Initializing eigensolver data") print("Computing {} bands with {} tolerance".format(self.num_bands, self.tolerance)) if type(self.default_material) is not mp.Medium and callable(self.default_material): # TODO: Support epsilon_function user materials like meep? self.default_material.eps = False self.init_params(p, reset_fields) if isinstance(reset_fields, basestring): self.load_eigenvectors(reset_fields) print("{} k-points".format(len(self.k_points))) for kp in self.k_points: print(" {}".format(kp)) print("elapsed time for initialization: {}".format(time.time() - init_time)) # TODO: Split over multiple processes # k_split = list_split(self.k_points, self.k_split_num, self.k_split_index) k_split = (0, self.k_points) self.mode_solver.set_kpoint_index(k_split[0]) if self.num_bands > 0: for i, k in enumerate(k_split[1]): self.current_k = k solve_kpoint_time = time.time() self.mode_solver.solve_kpoint(k) self.iterations = self.mode_solver.get_iterations() print("elapsed time for k point: {}".format(time.time() - solve_kpoint_time)) self.freqs = self.get_freqs() self.all_freqs[i, :] = np.array(self.freqs) self.band_range_data = self.update_band_range_data(self.band_range_data, self.freqs, k) self.eigensolver_iters += [self.iterations / self.num_bands] for f in band_functions: num_args = get_num_args(f) if num_args == 1: f(self) elif num_args == 2: band = 1 while band <= self.num_bands: f(self, band) band += 1 else: raise ValueError("Band function should take 1 or 2 arguments. " "The first must be a ModeSolver instance") if len(k_split[1]) > 1: self.output_band_range_data(self.band_range_data) self.gap_list = self.output_gaps(self.band_range_data) else: self.gap_list = [] end = time.time() - start print("total elapsed time for run: {}".format(end)) self.total_run_time += end self.eigensolver_flops = self.mode_solver.get_eigensolver_flops() self.parity = self.mode_solver.get_parity_string() print("done") def run(self, *band_functions): self.run_parity(mp.NO_PARITY, True, *band_functions) def run_zeven(self, *band_functions): self.run_parity(mp.EVEN_Z, True, *band_functions) def run_zodd(self, *band_functions): self.run_parity(mp.ODD_Z, True, *band_functions) def run_yeven(self, *band_functions): self.run_parity(mp.EVEN_Y, True, *band_functions) def run_yodd(self, *band_functions): self.run_parity(mp.ODD_Y, True, *band_functions) def run_yeven_zeven(self, *band_functions): self.run_parity(mp.EVEN_Y + mp.EVEN_Z, True, *band_functions) def run_yeven_zodd(self, *band_functions): self.run_parity(mp.EVEN_Y + mp.ODD_Z, True, *band_functions) def run_yodd_zeven(self, *band_functions): self.run_parity(mp.ODD_Y + mp.EVEN_Z, True, *band_functions) def run_yodd_zodd(self, *band_functions): self.run_parity(mp.ODD_Y + mp.ODD_Z, True, *band_functions) run_te = run_zeven run_tm = run_zodd run_te_yeven = run_yeven_zeven run_te_yodd = run_yodd_zeven run_tm_yeven = run_yeven_zodd run_tm_yodd = run_yodd_zodd 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 def first_brillouin_zone(self, k): """ Function to convert a k-point k into an equivalent point in the first Brillouin zone (not necessarily the irreducible Brillouin zone) """ def n(k): return mp.reciprocal_to_cartesian(k, self.geometry_lattice).norm() def try_plus(k, v): if n(k + v) < n(k): return try_plus(k + v, v) else: return k def _try(k, v): return try_plus(try_plus(k, v), mp.Vector3() - v) try_list = [ mp.Vector3(1, 0, 0), mp.Vector3(0, 1, 0), mp.Vector3(0, 0, 1), mp.Vector3(0, 1, 1), mp.Vector3(1, 0, 1), mp.Vector3(1, 1, 0), mp.Vector3(0, 1, -1), mp.Vector3(1, 0, -1), mp.Vector3(1, -1, 0), mp.Vector3(1, 1, 1), mp.Vector3(-1, 1, 1), mp.Vector3(1, -1, 1), mp.Vector3(1, 1, -1), ] def try_all(k): return functools.reduce(_try, try_list, k) def try_all_and_repeat(k): knew = try_all(k) return try_all_and_repeat(knew) if n(knew) < n(k) else k k0 = k - mp.Vector3(*[round(x) for x in k]) return try_all_and_repeat(k0) if n(k0) < n(k) else try_all_and_repeat(k) def get_dominant_planewave(self, band): return self.mode_solver.get_dominant_planewave(band) # Predefined output functions (functions of the band index), for passing to `run` def output_hfield(ms, which_band): ms.get_hfield(which_band, False) ms.output_field() def output_hfield_x(ms, which_band): ms.get_hfield(which_band, False) ms.output_field_x() def output_hfield_y(ms, which_band): ms.get_hfield(which_band, False) ms.output_field_y() def output_hfield_z(ms, which_band): ms.get_hfield(which_band, False) ms.output_field_z() def output_bfield(ms, which_band): ms.get_bfield(which_band, False) ms.output_field() def output_bfield_x(ms, which_band): ms.get_bfield(which_band, False) ms.output_field_x() def output_bfield_y(ms, which_band): ms.get_bfield(which_band, False) ms.output_field_y() def output_bfield_z(ms, which_band): ms.get_bfield(which_band, False) ms.output_field_z() def output_dfield(ms, which_band): ms.get_dfield(which_band, False) ms.output_field() def output_dfield_x(ms, which_band): ms.get_dfield(which_band, False) ms.output_field_x() def output_dfield_y(ms, which_band): ms.get_dfield(which_band, False) ms.output_field_y() def output_dfield_z(ms, which_band): ms.get_dfield(which_band, False) ms.output_field_z() def output_efield(ms, which_band): ms.get_efield(which_band, False) ms.output_field() def output_efield_x(ms, which_band): ms.get_efield(which_band, False) ms.output_field_x() def output_efield_y(ms, which_band): ms.get_efield(which_band, False) ms.output_field_y() def output_efield_z(ms, which_band): ms.get_efield(which_band, False) ms.output_field_z() def output_bpwr(ms, which_band): ms.get_bfield(which_band, False) ms.compute_field_energy() ms.output_field() def output_dpwr(ms, which_band): ms.get_dfield(which_band, False) ms.compute_field_energy() ms.output_field() def output_tot_pwr(ms, which_band): ms.get_tot_pwr(which_band) ms.output_field_to_file(-1, ms.get_filename_prefix() + 'tot.') def output_dpwr_in_objects(output_func, min_energy, objects=[]): """ The following function returns an output function that calls output_func for bands with D energy in objects > min-energy. For example, output_dpwr_in_objects(output_dfield, 0.20, some_object) would return an output function that would spit out the D field for bands with at least %20 of their D energy in some-object. """ def _output(ms, which_band): ms.get_dfield(which_band, False) ms.compute_field_energy() energy = ms.compute_energy_in_objects(objects) fmt = "dpwr:, {}, {}, {} " print(fmt.format(which_band, ms.freqs[which_band - 1], energy)) if energy >= min_energy: apply_band_func(ms, output_func, which_band) return _output def output_charge_density(ms, which_band): ms.get_charge_density(which_band) ms.output_field_to_file(-1, ms.get_filename_prefix()) def output_poynting(ms, which_band): ms.get_poynting(which_band) ms.output_field_to_file(-1, ms.get_filename_prefix() + 'flux.') def output_poynting_x(ms, which_band): ms.get_poynting(which_band) ms.output_field_to_file(0, ms.get_filename_prefix() + 'flux.') def output_poynting_y(ms, which_band): ms.get_poynting(which_band) ms.output_field_to_file(1, ms.get_filename_prefix() + 'flux.') def output_poynting_z(ms, which_band): ms.get_poynting(which_band) ms.output_field_to_file(2, ms.get_filename_prefix() + 'flux.') def display_yparities(ms): ms.display_kpoint_data('yparity', ms.mode_solver.compute_yparities()) def display_zparities(ms): ms.display_kpoint_data('zparity', ms.mode_solver.compute_zparities()) def display_group_velocities(ms): ms.display_kpoint_data('velocity', ms.compute_group_velocities()) # Band functions to pick a canonical phase for the eigenstate of the # given band based upon the spatial representation of the given field def fix_hfield_phase(ms, which_band): ms.get_hfield(which_band, False) ms.mode_solver.fix_field_phase() def fix_bfield_phase(ms, which_band): ms.get_bfield(which_band, False) ms.mode_solver.fix_field_phase() def fix_dfield_phase(ms, which_band): ms.get_dfield(which_band, False) ms.mode_solver.fix_field_phase() def fix_efield_phase(ms, which_band): ms.get_efield(which_band, False) ms.mode_solver.fix_field_phase() def apply_band_func_thunk(ms, band_func, which_band, eval_thunk): """ We need a special function to evaluate band functions, since band functions can either be a function of the band number or a thunk (function of no arguments, evaluated once per k-point). """ if get_num_args(band_func) == 1: if eval_thunk: band_func(ms) # evaluate thunks once per k-point else: band_func(ms, which_band) def apply_band_func(ms, band_func, which_band): apply_band_func_thunk(ms, band_func, which_band, which_band == 1) def combine_band_functions(*band_funcs): """Combines zero or more band functions into one""" def _combine(ms, which_band): for f in band_funcs: apply_band_func(ms, f, which_band) return _combine def output_at_kpoint(kpoint, *band_funcs): """Only invoke the given band functions for the specified k-point""" band_func = combine_band_functions(*band_funcs) def _output_at_kpoint(ms, which_band): if ms.current_k.close(kpoint, tol=1e-8 * kpoint.norm()): band_func(ms, which_band) return _output_at_kpoint