def request(self, parameters: Parameters, compute: bool = True) -> Path: self.logger.info("request wave function for parameters %s", repr(parameters)) with cwd.WorkingDir(self.working_dir.parent): common_parameter_names = parameters.get_common_parameter_names( self.prototype, ) path = self.get_path(parameters) if path: self.logger.info("wave function is present") return path self.logger.info("wave function is not present") p = copy.deepcopy(self.prototype) for name in common_parameter_names: p[name] = parameters[name] if p not in self.missing_wave_functions: self.store_missing_wave_function(p) self.combinations = (self.stored_wave_functions + self.missing_wave_functions) if not compute: raise MissingWfnError(p) self._compute(p) return self.request(parameters, compute)
def action_compute(targets: List[str]): del targets operator = self.operator.resolve() psi = self.psi.resolve() expval = self.expval.resolve() wave_function = self.wave_function.resolve() with tempfile.TemporaryDirectory() as tmpdir: with cwd.WorkingDir(tmpdir): self.logger.info("compute expectation value") copy_file(operator, "operator") copy_file(psi, "psi") copy_file(wave_function, "restart") cmd = [ "qdtk_expect.x", "-opr", "operator", "-rst", "restart", "-psi", "psi", "-save", "expval", ] self.logger.info("command: %s", " ".join(cmd)) env = os.environ.copy() env["OMP_NUM_THREADS"] = env.get("OMP_NUM_THREADS", "1") subprocess.run(cmd, env=env) write_expval_hdf5("expval.h5", *read_expval_ascii("expval")) copy_file("expval.h5", expval)
def link_simulations(self): if not self.combinations: return self.logger.info("create simulation symlinks") self.create_working_dir() # link by index with cwd.WorkingDir(self.working_dir): path_by_index = Path("by_index") if not path_by_index.exists(): os.makedirs(path_by_index) for index, combination in enumerate(self.combinations): path = self.compute_working_dir(combination) link = path_by_index / str(index) if link.exists(): link.unlink() subprocess.run(["ln", "-s", "-f", "-r", path, link]) # link by parameters with cwd.WorkingDir(self.working_dir): path_by_param = Path("by_param") if not path_by_param.exists(): os.makedirs(path_by_param) variables, constants = mlxtk.parameters.get_variables( self.combinations) for combination in self.combinations: try: if not variables: name = "_".join(constant + "=" + str(combination[constant]) for constant in constants) else: name = "_".join(variable + "=" + str(combination[variable]) for variable in variables) path = self.compute_working_dir(combination) link = path_by_param / name if link.exists(): link.unlink() subprocess.run(["ln", "-s", "-f", "-r", path, link]) except OSError: continue
def load_db(path: Path, variable_name: str = "db") -> WaveFunctionDB: with cwd.WorkingDir(path.resolve().parent): spec = importlib.util.spec_from_file_location("db", str(path)) db_module = importlib.util.module_from_spec(spec) spec.loader.exec_module(db_module) db = getattr(db_module, variable_name) db.working_dir = path.parent / db.name return db
def load_stored_wave_functions(self): self.create_working_dir() p = Path("stored_wave_functions.pickle") with cwd.WorkingDir(self.working_dir): if not p.exists(): return with open("stored_wave_functions.pickle", "rb") as fptr: self.stored_wave_functions = pickle.load(fptr)
def unlink_simulations(self): if not self.working_dir.exists(): return with cwd.WorkingDir(self.working_dir): path_by_index = Path("by_index") if path_by_index.exists(): with cwd.WorkingDir(path_by_index): for entry in os.listdir(os.getcwd()): path_entry = Path(entry) if path_entry.is_symlink(): path_entry.unlink() path_by_param = Path("by_param") if path_by_param.exists(): with cwd.WorkingDir(path_by_param): for entry in os.listdir(os.getcwd()): path_entry = Path(entry) if path_entry.is_symlink(): path_entry.unlink()
def action_compute(targets: List[str]): del targets wave_function = self.wave_function.resolve() psi = self.psi.resolve() basis = self.basis.resolve() result = self.result.resolve() with tempfile.TemporaryDirectory() as tmpdir: with cwd.WorkingDir(tmpdir): self.logger.info("perform number state analysis") copy_file(wave_function, "restart") copy_file(basis, "basis") copy_file(psi, "psi") cmd = [ "qdtk_analysis.x", "-fixed_ns", "-rst_bra", "basis", "-rst_ket", "restart", "-save", "result", "-psi", "psi", ] self.logger.info("command: %s", " ".join(cmd)) env = os.environ.copy() env["OMP_NUM_THREADS"] = env.get("OMP_NUM_THREADS", "1") subprocess.run(cmd, env=env) times, real, imag = inout.read_fixed_ns_ascii("result") wfn = load_wave_function("basis") inout.write_fixed_ns_hdf5( "result.h5", times, real, imag, wfn._tape[1], wfn._tape[3], ) with h5py.File("result.h5", "a") as fptr: dset = fptr["fixed_ns"].create_dataset( "total_magnitude", shape=times.shape, dtype=numpy.float64, ) dset[:] = numpy.sum((real**2) + (imag**2), axis=1) copy_file("result.h5", result)
def store_parameters(self): self.logger.info("storing scan parameters") self.create_working_dir() if not self.combinations: return with cwd.WorkingDir(self.working_dir): for combination in self.combinations: working_dir = self.compute_working_dir(combination) working_dir.mkdir(exist_ok=True, parents=True) with cwd.WorkingDir(working_dir): with open("parameters.pickle", "wb") as fptr: pickle.dump(combination, fptr, protocol=3) with open("parameters.json", "w") as fptr: fptr.write(combination.to_json() + "\n") with open("scan.pickle", "wb") as fptr: pickle.dump( [combination for combination in self.combinations], fptr, protocol=3, )
def action_compute(targets: list[str]): del targets wave_function = self.wave_function.resolve() psi = self.psi.resolve() outpath = self.output.resolve() with tempfile.TemporaryDirectory() as tmpdir: with cwd.WorkingDir(tmpdir): copy_file(psi, "psi") copy_file(wave_function, "restart") results: None | numpy.ndarray = None cmd = [ "qdtk_expect.x", "-opr", "operator", "-rst", "restart", "-psi", "psi", "-save", "expval", ] for i, combination in enumerate(self.indices): with open("operator", "w") as fptr: self.func(*combination).get_operator( ).createOperatorFile(fptr, ) self.logger.info(f'command: {" ".join(cmd)}') env = os.environ.copy() env["OMP_NUM_THREADS"] = env.get( "OMP_NUM_THREADS", "1") subprocess.run(cmd, env=env) time, values = read_expval_ascii("expval") if results is None: results = numpy.zeros( [len(time), len(self.indices)], dtype=numpy.complex128, ) results[:, i] = values with h5py.File(outpath, "w") as fptr: fptr.create_dataset("time", data=time) fptr.create_dataset("indices", data=self.indices) fptr.create_dataset("values", data=results)
def store_wave_function(self, parameters: Parameters): self.create_working_dir() p = Path("stored_wave_functions.pickle") with cwd.WorkingDir(self.working_dir): entries = [] # type: List[Parameters] if p.exists(): with open(p, "rb") as fptr: entries = pickle.load(fptr) if parameters in entries: return entries.append(parameters) with open(p, "wb") as fptr: pickle.dump(entries, fptr, protocol=3) self.stored_wave_functions.append(parameters)
def remove_wave_function(self, parameters: Parameters): self.create_working_dir() p = Path("stored_wave_functions.pickle") with cwd.WorkingDir(self.working_dir): if not p.exists(): return with open(p, "rb") as fptr: entries = pickle.load(fptr) if parameters not in entries: return entries.remove(parameters) with open(p, "wb") as fptr: pickle.dump(entries, fptr, protocol=3) self.stored_wave_functions.remove(parameters)
def action_compute(targets: List[str]): wave_function = self.wave_function.resolve() basis = self.basis.resolve() result = self.result.resolve() with tempfile.TemporaryDirectory() as tmpdir: with cwd.WorkingDir(tmpdir): self.logger.info( "compute MCTDHB wave function overlap (static)") copy_file(wave_function, "restart") copy_file(basis, "basis") cmd = [ "qdtk_analysis.x", "-fixed_ns", "-rst_bra", "basis", "-rst_ket", "restart", "-save", "result", ] self.logger.info("command: %s", " ".join(cmd)) env = os.environ.copy() env["OMP_NUM_THREADS"] = env.get("OMP_NUM_THREADS", "1") subprocess.run(cmd, env=env) _, real, imag = inout.read_fixed_ns_ascii("result") self.logger.info(real.shape) coeffs = load_wave_function("basis").PSI[:real.shape[1]] overlap = (numpy.abs( numpy.sum( numpy.conjugate(coeffs) * (real + 1j * imag)), )**2) self.logger.info("overlap = %f", overlap) with h5py.File(result, "w") as fptr: fptr.create_dataset("overlap", shape=(1, ), dtype=numpy.float64)[:] = overlap
def action_run(targets: List[str]): del targets if not self.path_name.exists(): self.path_name.mkdir() path_temp = self.path_name.resolve().with_name("." + self.path_name.name) if path_temp.exists(): path_temp_operator = path_temp / "hamiltonian" if not path_temp_operator.exists(): self.logger.debug( 'temporary operator "%s" does not exist', str(path_temp_operator), ) shutil.rmtree(path_temp) else: hash_a = hash_file(path_temp_operator) hash_b = hash_file(self.path_hamiltonian.resolve()) if hash_a != hash_b: self.logger.debug("operator hashes do not match") shutil.rmtree(path_temp) if path_temp.exists(): path_temp_wfn = path_temp / "initial" if not path_temp_wfn.exists(): self.logger.debug( 'temporary wave function "%s" does not exist', str(path_temp_wfn), ) shutil.rmtree(path_temp) else: hash_a = hash_file(path_temp_wfn) hash_b = hash_file(self.path_wave_function.resolve()) if hash_a != hash_b: self.logger.debug("wave function hashes do not match") shutil.rmtree(path_temp) if path_temp.exists(): self.logger.info("attempt continuation run") self.flags["cont"] = True self.flags["rstzero"] = False self.flag_list.append("-cont") try: self.flag_list.remove("-rstzero") except ValueError: pass with TemporaryDir(path_temp) as tmpdir: operator = self.path_hamiltonian.resolve() wave_function = self.path_wave_function.resolve() output_dir = self.path_name.resolve() hdf5_file = self.path_hdf5.resolve() diag_gauge_oper = (self.path_diag_gauge_oper.resolve() if self.path_diag_gauge_oper else None) with cwd.WorkingDir(tmpdir.path): self.logger.info("propagate wave function") if not self.flags["cont"]: copy_file(operator, "hamiltonian") copy_file(wave_function, "initial") copy_file(wave_function, "restart") if self.flags.get("gauge", "standard") == "diagonalization": if not diag_gauge_oper: raise ValueError( "no operator specified for diagonalization gauge", ) copy_file(diag_gauge_oper, "gauge_oper") cmd = ["qdtk_propagate.x"] + self.flag_list env = os.environ.copy() env["OMP_NUM_THREADS"] = env.get("OMP_NUM_THREADS", "1") self.logger.info("command: %s", " ".join(cmd)) result = subprocess.run(cmd, env=env) if result.returncode != 0: raise RuntimeError("Failed to run qdtk_propagate.x") if self.flags["exact_diag"]: with h5py.File("result.h5", "w") as fptr: add_eigenbasis_to_hdf5(fptr, *read_eigenbasis_ascii(".")) else: shutil.move("restart", "final.wfn") with h5py.File("result.h5", "w") as fptr: add_gpop_to_hdf5(fptr.create_group("gpop"), *read_gpop_ascii("gpop")) add_natpop_to_hdf5( fptr.create_group("natpop"), *read_natpop_ascii("natpop"), ) add_output_to_hdf5( fptr.create_group("output"), *read_output_ascii("output"), ) if "gauge" in self.flags: if self.flags["gauge"] != "standard": time, error = numpy.loadtxt( "constraint_error.txt", unpack=True, ) group = fptr.create_group( "constraint_error") group.create_dataset( "time", time.shape, dtype=time.dtype, )[:] = time group.create_dataset( "error", error.shape, dtype=error.dtype, )[:] = error shutil.move("result.h5", hdf5_file) if self.flags["exact_diag"]: with open("eigenvectors") as fptr: lines = fptr.readlines() if not lines[0].startswith("$tape"): with open("eigenvectors", "w") as fptr: fptr.write("$tape\n") fptr.writelines(lines) if self.flags["eig_tot"]: self.logger.info( "exact diagonalization with one eigenvector, remove last two lines in eigenvectors file", ) with open("eigenvectors") as fptr: lines = fptr.readlines()[:-1] with open("eigenvectors", "w") as fptr: fptr.writelines(lines) for fname in self.qdtk_files: copy_file(fname, output_dir / fname) tmpdir.complete = True
def compress_folder(path: Union[str, Path], compression: int = 9, jobs: int = 1): path = make_path(path) folder = path.name archive = path.with_suffix(".tar.gz").name exe_pv = shutil.which("pv") exe_tar = shutil.which("tar") exe_gzip = shutil.which("tar") exe_pigz = shutil.which("pigz") with cwd.WorkingDir(path.parent): if exe_pigz: if exe_pv: with open(archive, "wb") as fptr: size = get_folder_size(folder) process_tar = subprocess.Popen( [exe_tar, "cf", "-", folder], stdout=subprocess.PIPE, ) process_pv = subprocess.Popen( [exe_pv, "-s", str(size)], stdin=process_tar.stdout, stdout=subprocess.PIPE, ) process_pigz = subprocess.Popen( [exe_pigz, "-" + str(compression), "-p", str(jobs)], stdin=process_pv.stdout, stdout=fptr, ) process_tar.wait() process_pv.wait() process_pigz.wait() else: LOGGER.warning("cannot find pv, no progress will be displayed") with open(archive, "wb") as fptr: process_tar = subprocess.Popen( [exe_tar, "cf", "-", folder], stdout=subprocess.PIPE, ) process_pigz = subprocess.Popen( [exe_pigz, "-" + str(compression), "-p", str(jobs)], stdin=process_tar.stdout, stdout=fptr, ) process_tar.wait() process_pigz.wait() elif exe_gzip: if jobs > 1: LOGGER.warning( "gzip does not support parallel compression, using one thread only", ) if exe_pv: with open(archive, "wb") as fptr: size = get_folder_size(folder) process_tar = subprocess.Popen( [exe_tar, "cf", "-", folder], stdout=subprocess.PIPE, ) process_pv = subprocess.Popen( [exe_pv, "-s", str(size)], stdin=process_tar.stdout, stdout=subprocess.PIPE, ) process_gzip = subprocess.Popen( [exe_gzip, "-" + str(compression)], stdin=process_pv.stdout, stdout=fptr, ) process_tar.wait() process_pv.wait() process_gzip.wait() else: LOGGER.warning("cannot find pv, no progress will be displayed") with open(archive, "wb") as fptr: process_tar = subprocess.Popen( [exe_tar, "cf", "-", folder], stdout=subprocess.PIPE, ) process_gzip = subprocess.Popen( [exe_gzip, "-" + str(compression)], stdin=process_gzip.stdout, stdout=fptr, ) process_tar.wait() process_gzip.wait() else: raise RuntimeError("Cannot find either pigz or gzip")