def action_extend_grid(targets: List[str]):
del targets
# pylint: disable=protected-access
wfn = load_wave_function(self.path_initial)
tape = list(wfn._tape)
tape[8] += self.nleft + self.nright
wfn_new = WaveFunction(tape=tuple(tape))
spfs = get_spfs(wfn)
for i, spf in enumerate(spfs):
spfs[i] = numpy.pad(
spf,
(self.nleft, self.nright),
"constant",
constant_values=(self.value, self.value),
)
N = wfn._tape[1]
m = wfn._tape[3]
number_state = numpy.zeros(m, dtype=numpy.int64)
number_state[0] = N
wfn_new.init_coef_sing_spec_B(
number_state,
spfs,
1e-15,
1e-15,
full_spf=True,
)
wfn_new.PSI[0:wfn_new.tree._subnodes[0].
_z0] = wfn.PSI[0:wfn.tree._subnodes[0]._z0]
save_wave_function(self.path, wfn_new)
def action_add_momentum(targets: List[str]):
del targets
# pylint: disable=protected-access
wfn = load_wave_function(Path(self.path_initial))
wfn.tree._topNode._pgrid[0] = self.grid.get_x()
add_momentum(wfn, self.momentum)
save_wave_function(self.path, wfn)
def read_spfs(path: str) -> Tuple[numpy.ndarray, numpy.ndarray]:
with io.StringIO(read_first_frame(path)) as sio:
wfn = load_wave_function(sio)
_, times, psis = read_psi_ascii(path)
spfs = []
for psi in psis:
wfn.PSI = psi
spfs.append(numpy.array(get_spfs(wfn)))
return times, numpy.moveaxis(numpy.array(spfs), 1, 0)
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 main():
# pylint: disable=protected-access
parser = argparse.ArgumentParser()
parser.add_argument(
"basis",
type=Path,
help="wave function file containing the basis to project onto",
)
parser.add_argument(
"psi",
type=Path,
help="wave function file containing the wave function to analyse",
)
parser.add_argument(
"-o",
"--output",
type=Path,
help="name of the output file (optional)",
)
args = parser.parse_args()
output = args.output
if not output:
output = Path(f"{args.psi.stem}_{args.basis.stem}.fixed_ns.h5")
output = output.resolve()
with tempfile.TemporaryDirectory(dir=Path.cwd()) as tmpdir:
LOGGER.info("create a restart file from psi file")
with open(Path(tmpdir) / "restart", "w") as fp:
fp.write(read_first_frame(args.psi))
LOGGER.info("copy psi file")
shutil.copy2(args.psi, Path(tmpdir) / "psi")
LOGGER.info("copy basis")
shutil.copy2(args.basis, Path(tmpdir) / "basis")
with WorkingDir(tmpdir):
cmd = [
"qdtk_analysis.x",
"-fixed_ns",
"-rst_bra",
"basis",
"-rst_ket",
"restart",
"-psi",
"psi",
"-save",
"result",
]
LOGGER.info("run qdtk_analysis.x: %s", " ".join(cmd))
subprocess.check_output(cmd)
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],
)
LOGGER.info("copy result")
shutil.copy2("result.h5", output)
