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_write_wave_function(targets: List[str]):
with h5py.File(self.path_matrix, "r") as fptr:
matrix = fptr["matrix"][:, :]
energies, spfs = diagonalize_1b_operator(matrix, -1)
indices = numpy.nonzero(
(energies > self.threshold) if self.above_threshold else
(energies < self.threshold), )[0]
m = indices.max() - indices.min() + 1
if self.max_number > 0:
m = min(self.max_number, m)
self.logger.info("use m = %d spfs", m)
self.logger.info("spf indices: %d to %d", indices.min(),
indices.min() + m)
spfs = spfs[indices.min():indices.min() + m]
spfs_arr = numpy.array(spfs)
energies = energies[indices.min():indices.min() + m]
with h5py.File(targets[1], "w") as fptr:
dset = fptr.create_dataset(
"energies",
(len(energies), ),
dtype=numpy.float64,
)
dset[:] = energies
dset = fptr.create_dataset(
"spfs",
spfs_arr.shape,
dtype=numpy.complex128,
)
dset[:, :] = spfs_arr
grid_points = matrix.shape[0]
tape = (
-10,
self.number_of_particles,
+1,
len(energies),
-1,
1,
1,
0,
grid_points,
-2,
)
self.number_state = numpy.zeros(len(energies), dtype=numpy.int64)
self.number_state[0] = self.number_of_particles
wfn = WaveFunction(tape=tape)
wfn.init_coef_sing_spec_B(
self.number_state,
spfs,
1e-15,
1e-15,
full_spf=True,
)
save_wave_function(self.path, wfn)
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 action_write_wave_function(targets: List[str]):
del targets
all_spfs = []
for m, path_matrix, path_basis in zip(
self.number_of_spfs,
self.path_matrices,
self.path_bases,
):
with h5py.File(path_matrix, "r") as fptr:
matrix = fptr["matrix"][:, :]
energies, spfs = diagonalize_1b_operator(matrix, m)
spfs_arr = numpy.array(spfs)
all_spfs += spfs
with h5py.File(path_basis, "w") as fptr:
dset = fptr.create_dataset("energies", (m, ),
dtype=numpy.float64)
dset[:] = energies
dset = fptr.create_dataset(
"spfs",
spfs_arr.shape,
dtype=numpy.complex128,
)
dset[:, :] = spfs_arr[-1]
grid_points = matrix.shape[0]
tape = (
-10,
self.number_of_particles,
+1,
sum(self.number_of_spfs),
-1,
1,
1,
0,
grid_points,
-2,
)
wfn = WaveFunction(tape=tape)
wfn.init_coef_sing_spec_B(
self.number_state,
all_spfs,
1e-15,
1e-15,
full_spf=True,
)
save_wave_function(self.path, wfn)
def action_write_wave_function(targets: List[str]) -> Dict[str, Any]:
with h5py.File(self.path_matrix_A, "r") as fptr:
matrix_A = fptr["matrix"][:, :]
with h5py.File(self.path_matrix_B, "r") as fptr:
matrix_B = fptr["matrix"][:, :]
energies_A, spfs_A = diagonalize_1b_operator(
matrix_A, self.num_spfs_A)
spfs_arr_A = numpy.array(spfs_A)
energies_B, spfs_B = diagonalize_1b_operator(
matrix_B, self.num_spfs_B)
spfs_arr_B = numpy.array(spfs_B)
with h5py.File(self.path_basis_A, "w") as fptr:
fptr.create_dataset(
"energies",
(self.num_spfs_A, ),
dtype=numpy.float64,
)[:] = energies_A
fptr.create_dataset(
"spfs", spfs_arr_A.shape,
dtype=numpy.complex128)[:, :, ] = spfs_arr_A
with h5py.File(self.path_basis_B, "w") as fptr:
fptr.create_dataset(
"energies",
(self.num_spfs_B, ),
dtype=numpy.float64,
)[:] = energies_B
fptr.create_dataset(
"spfs", spfs_arr_B.shape,
dtype=numpy.complex128)[:, :, ] = spfs_arr_B
n_A = matrix_A.shape[0]
n_B = matrix_B.shape[0]
assert n_A == n_B
tape = (
-10,
2,
0,
self.num_sbs_A,
self.num_sbs_B,
-1,
1,
self.num_particles_A,
1,
self.num_spfs_A,
-1,
1,
1,
0,
n_A,
0,
0,
-1,
2,
self.num_particles_B,
1,
self.num_spfs_B,
-1,
1,
1,
0,
n_B,
-2,
)
wfn = WaveFunction(tape=tape)
wfn.init_coef_multi_spec(
2,
[self.ns_A, self.ns_B],
[spfs_A, spfs_B],
1e-15,
1e-15,
full_spf=True,
)
save_wave_function(self.path, wfn)