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 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 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 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 main():
parser = argparse.ArgumentParser()
parser.add_argument(
"--restart",
type=Path,
default=Path("restart"),
help="path to the restart file",
)
parser.add_argument(
"--operator",
type=Path,
default=Path("operator"),
help="path to the hamiltonian operator",
)
parser.add_argument(
"--psi",
type=Path,
default=Path("psi"),
help="path to the psi file",
)
parser.add_argument("--dof1",
type=int,
default=1,
help="first degree of freedom")
parser.add_argument("--dof2",
type=int,
default=1,
help="second degree of freedom")
parser.add_argument("-o",
"--output",
type=Path,
help="path for the output file")
args = parser.parse_args()
restart_file = Path(args.restart).resolve()
operator_file = Path(args.operator).resolve()
psi_file = Path(args.psi).resolve()
basename = f"dmat2_dof{args.dof1}_dof{args.dof2}"
if not args.output:
output_file = Path.cwd() / (basename + "_gridrep.h5")
else:
output_file = Path(args.output).resolve()
with tempfile.TemporaryDirectory() as tmpdir:
with WorkingDir(tmpdir):
copy_file(restart_file, "restart")
copy_file(operator_file, "operator")
copy_file(psi_file, "psi")
subprocess.run([
"qdtk_analysis.x",
"-dmat2",
"-rst",
"restart",
"-opr",
"operator",
"-psi",
"psi",
"-dof",
str(args.dof1),
"-dofB",
str(args.dof2),
], )
with h5py.File(output_file, "w") as fptr:
add_dmat2_gridrep_to_hdf5(
fptr,
read_dmat2_gridrep_ascii(basename + "_grid"),
)
def main():
parser = argparse.ArgumentParser()
parser.add_argument("operator",
type=Path,
help="path of the operator file")
parser.add_argument("restart", type=Path, help="path of the restart file")
parser.add_argument("psi", type=Path, help="path of the psi file")
parser.add_argument("--node", type=int, default=0, help="node to use")
parser.add_argument("--dof",
type=int,
default=1,
help="degree of freedom to use")
parser.add_argument(
"-o",
"--output",
type=Path,
default=Path.cwd() / "dmat.h5",
help="output file",
)
parser.add_argument("--slice", type=str)
group_diag = parser.add_mutually_exclusive_group()
group_diag.add_argument(
"--diagonalize",
action="store_true",
help="whether to diagonalize the one-body density matrix",
)
group_diag.add_argument(
"--only-diagonalize",
action="store_true",
help="whether to only diagonalize the one-body density matrix",
)
parser.add_argument(
"--only-eigenvalues",
action="store_true",
help="whether to only output the eigenvalues",
)
group_rep = parser.add_mutually_exclusive_group()
group_rep.add_argument(
"--spfrep",
action="store_true",
help="use spf representation",
)
group_rep.add_argument(
"--gridrep",
action="store_true",
help="use grid representation",
)
args = parser.parse_args()
if args.only_eigenvalues and not (args.diagonalize
or args.only_diagonalize):
parser.error(
"--only-eigenvalues only applies when --diagonalize or --only-diagonalize is specified",
)
output = args.output.resolve()
opr = args.operator.resolve()
rst = args.restart.resolve()
psi = args.psi.resolve()
with tempfile.TemporaryDirectory() as tempdir:
with WorkingDir(tempdir):
copy_file(opr, "opr")
copy_file(rst, "rst")
if args.slice:
tape, times, psi = read_psi_ascii(psi)
m = RE_SLICE.match(args.slice)
if not m:
raise RuntimeError(f'Invalid slice format "{args.slice}"')
start = 0
end = len(times)
step = 1
if m.group(1) != "":
start = int(m.group(1))
if m.group(2) != "":
stop = int(m.group(2))
try:
if m.group(3) != "":
step = int(m.group(3))
except IndexError:
pass
indices = [i for i in range(start, end, step)]
write_psi_ascii("psi", (tape, times[indices], psi[indices]))
else:
copy_file(psi, "psi")
gridrep = args.gridrep or ((not args.gridrep) and
(not args.spfrep))
cmd = [
"qdtk_analysis.x",
"-opr",
"opr",
"-rst",
"rst",
"-psi",
"psi",
"-dmat",
"-nd",
str(args.node),
"-dof",
str(args.dof),
]
if not gridrep:
cmd.append("-spfrep")
if gridrep:
cmd.append("-gridrep")
if args.diagonalize:
cmd.append("-diagonalize")
if args.only_diagonalize:
cmd += ["-diagonalize", "-onlydiag"]
if args.only_eigenvalues:
cmd.append("-onlyeigval")
LOGGER.info("cmd: %s", " ".join(cmd))
subprocess.run(cmd)
with h5py.File(output, "a") as fptr:
if not args.only_diagonalize:
if gridrep:
dmat.add_dmat_gridrep_to_hdf5(
fptr,
dmat.read_dmat_gridrep_ascii(
f"dmat_dof{args.dof}_grid"),
)
else:
dmat.add_dmat_spfrep_to_hdf5(
fptr,
*dmat.read_dmat_spfrep_ascii(
f"dmat_dof{args.dof}_spf"),
)
if args.diagonalize or args.only_diagonalize:
dmat.add_dmat_evals_to_hdf5(
fptr,
*dmat.read_dmat_evals_ascii(
f"eval_dmat_dof{args.dof}"),
)
if not args.only_eigenvalues:
if gridrep:
dmat.add_dmat_evecs_grid_to_hdf5(
fptr,
*dmat.read_dmat_evecs_grid_ascii(
f"evec_dmat_dof{args.dof}_grid", ),
)
else:
dmat.add_dmat_evecs_spf_to_hdf5(
fptr,
*dmat.read_dmat_evecs_spf_ascii(
f"evec_dmat_dof{args.dof}_spf", ),
)
def compute_reduced_density_matrix(
wave_function: PathLike,
output_file: PathLike,
dvrs: Sequence[DVRSpecification],
dofs_A: Sequence[int],
basis_states: Mapping[int, Sequence[ArrayLike]] | None = None,
threads: int = 1,
):
# convert paths
wave_function = Path(wave_function).resolve()
output_file = Path(output_file).resolve()
# function to generate basis states for DoF
def generate_basis_states(dof: int) -> list[ArrayLike]:
states: list[ArrayLike] = []
dim = dvrs[dof].get().npoints
for i in range(dim):
states.append(numpy.zeros((dim, )))
states[-1][i] = 1
return states
# check/prepare basis states for each DoF
if basis_states is None:
# generate basis_states from DVR
basis_states: dict[int, list[ArrayLike]] = {}
for dof in dofs_A:
basis_states[dof] = generate_basis_states(dof)
else:
# check that basis_states are consitent with DVR
for dof in dofs_A:
# no basis states for DoF -> generate them
if dof not in basis_states:
basis_states[dof] = generate_basis_states(dof)
dim = dvrs[dof].get().npoints
if len(basis_states[dof]) != dim:
raise ValueError(
f"basis for DoF {dof} is overcomplete or incomplete (got {len(basis_states[dof])} states, expected {dim})",
)
for i, state in enumerate(basis_states[dof]):
if len(state) != dim:
raise ValueError(
f"basis state {i} of DoF {dof} has wrong size (got {len(state)}, expected {dim})",
)
# compute dimension of reduced density operator
dim_A = reduce(mul, (len(basis_states[dof]) for dof in basis_states))
# create storage for results
results: dict[tuple[int, int], complex] = {}
with TemporaryDir(
Path(output_file).parent /
("." + Path(output_file).name + ".tmp"), ) as tmpdir:
with WorkingDir(tmpdir.path):
# copy wave function file
copy_file(Path(wave_function), Path.cwd() / "psi")
# generate a restart file
tape, times, psi = read_psi_frame_ascii("psi", 0)
write_psi_ascii("restart", [tape, [times], psi[numpy.newaxis, :]])
# generate the tasks for the multiprocessing pool
tasks = ([dvrs, dofs_A, a, state_a, b, state_b]
for a, state_a in enumerate(
product(*(basis_states[dof] for dof in dofs_A)), )
for b, state_b in enumerate(
product(*(basis_states[dof] for dof in dofs_A)), )
if b <= a)
with multiprocessing.Pool(threads) as pool:
for a, b, time, values in pool.starmap(
compute_reduced_density_matrix_element,
tasks,
1,
):
results[(a, b)] = values
tmpdir.complete = True
# convert results to an array
steps = len(results[(0, 0)])
results_arr = numpy.zeros((steps, dim_A, dim_A), dtype=numpy.complex128)
for a in range(dim_A):
for b in range(a + 1):
results_arr[:, a, b] = results[(a, b)]
for b in range(a + 1, dim_A):
results_arr[:, a, b] = numpy.conjugate(results[(b, a)])
# create parent directory of output file
Path(output_file).parent.mkdir(exist_ok=True, parents=True)
# make sure that results array is a 3rd order tensor
# first index denotes time
if len(results_arr.shape) == 2:
results_arr = results_arr[numpy.newaxis, :, :]
# store results in a HDF5 file
with h5py.File(output_file, "w") as fptr:
fptr.create_dataset("time", data=time)
fptr.create_dataset("rho_A", data=results_arr)
fptr.create_dataset("trace",
data=numpy.trace(results_arr, axis1=1, axis2=2))
for dof in dofs_A:
fptr.create_dataset(
f"basis_states_{dof}",
data=numpy.array(basis_states[dof]),
)
def main():
parser = argparse.ArgumentParser()
parser.add_argument("input", type=Path, help="path for the input file")
parser.add_argument(
"-o",
"--output",
type=Path,
default=None,
help="path for the video file",
)
parser.add_argument(
"-f",
"--fps",
type=float,
default=5,
help="frames per second for the video",
)
parser.add_argument(
"--dpi",
type=int,
default=600,
help="resolution (dpi) of the individual frames",
)
parser.add_argument(
"-j",
"--jobs",
type=int,
default=os.cpu_count(),
help="number of threads to use to create the individual frames",
)
args = parser.parse_args()
times, x1, x2, dmat2 = read_dmat2_gridrep_hdf5(args.input)
X2, X1 = numpy.meshgrid(x2, x1)
num_digits = len(str(len(times)))
if args.output:
output = args.output.resolve()
else:
output = args.input.with_suffix(".mp4").resolve()
with tempfile.TemporaryDirectory() as tmpdir:
with WorkingDir(tmpdir):
with multiprocessing.Pool(args.jobs) as pool:
pool.map(
partial(
render_frame,
times=times,
dmat2=dmat2,
X1=X1,
X2=X2,
num_digits=num_digits,
dpi=args.dpi,
),
[i for i, _ in enumerate(times)],
)
cmd = [
"ffmpeg",
"-y",
"-framerate",
str(args.fps),
"-i",
"%0" + str(num_digits) + "d.png",
"out.mp4",
]
LOGGER.info("ffmpeg command: %s", " ".join(cmd))
subprocess.run(cmd)
LOGGER.info("copy file: out.mp4 -> %s", str(output))
copy_file("out.mp4", output)