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")
