def test_pyr_4mode(multi_e, dvr):
basis, ham_terms = construct_vibronic_model(multi_e, dvr)
model = Model(basis, ham_terms)
mpo = Mpo(model)
logger.info(f"mpo_bond_dims:{mpo.bond_dims}")
# same form whether multi_e is True or False
init_condition = {"s2": 1}
if dvr:
for dof in model.v_dofs:
idx = model.order[dof]
init_condition[dof] = basis[idx].dvr_v[0]
mps = Mps.hartree_product_state(model, condition=init_condition)
compress_config = CompressConfig(CompressCriteria.fixed, max_bonddim=10)
evolve_config = EvolveConfig(EvolveMethod.tdvp_ps)
job = VibronicModelDynamics(model,
mps0=mps,
h_mpo=mpo,
compress_config=compress_config,
evolve_config=evolve_config,
expand=True)
time_step_fs = 2
job.evolve(evolve_dt=time_step_fs * fs2au, nsteps=60)
from renormalizer.vibronic.tests.mctdh_data import mctdh_data
assert np.allclose(mctdh_data[::round(time_step_fs / 0.5)][:61, 1:],
job.e_occupations_array,
atol=2e-2)
def test_pyr_4mode(multi_e, translator):
order, basis, vibronic_model = construct_vibronic_model(multi_e)
if translator is ModelTranslator.vibronic_model:
model = vibronic_model
elif translator is ModelTranslator.general_model:
model = vibronic_to_general(vibronic_model)
else:
assert False
mol_list2 = MolList2(order, basis, model, model_translator=translator)
mpo = Mpo(mol_list2)
logger.info(f"mpo_bond_dims:{mpo.bond_dims}")
mps = Mps.hartree_product_state(mol_list2, condition={"e_1": 1})
# for multi-e case the `expand bond dimension` routine is currently not working
# because creation operator is not defined yet
mps.use_dummy_qn = True
mps.build_empty_qn()
compress_config = CompressConfig(CompressCriteria.fixed, max_bonddim=10)
evolve_config = EvolveConfig(EvolveMethod.tdvp_ps)
job = VibronicModelDynamics(mol_list2,
mps0=mps,
h_mpo=mpo,
compress_config=compress_config,
evolve_config=evolve_config)
time_step_fs = 2
job.evolve(evolve_dt=time_step_fs * fs2au, nsteps=59)
from renormalizer.vibronic.tests.mctdh_data import mctdh_data
assert np.allclose(mctdh_data[::round(time_step_fs / 0.5)][:, 1:],
job.e_occupations_array,
atol=5e-2)
def init_mps(self):
if self.mps0 is None:
assert self.init_condition is not None
init_mp = Mps.hartree_product_state(self.model, self.init_condition)
self.mps0 = init_mp.copy()
else:
init_mp = self.mps0.copy()
init_mp.compress_config = self.compress_config
init_mp.evolve_config = self.evolve_config
init_mp.model = self.model
if self.evolve_config.is_tdvp and self.expand:
init_mp = init_mp.expand_bond_dimension(self.h_mpo, include_ex=False)
return init_mp
def init_mps(self):
if self.mps0 is None:
assert self.init_condition is not None
init_mp = Mps.hartree_product_state(self.mol_list,
self.init_condition)
self.mps0 = init_mp.copy()
else:
init_mp = self.mps0.copy()
init_mp.compress_config = self.compress_config
init_mp.evolve_config = self.evolve_config
init_mp.mol_list = self.mol_list
if self.evolve_config.is_tdvp:
init_mp = init_mp.expand_bond_dimension(self.h_mpo)
return init_mp
def analysis_dominant_config(self, thresh=0.8, alias=None, tda_m_trunc=20,
return_compressed_mps=False):
r""" analyze the dominant configuration of each tda root.
The algorithm is to compress the tda wavefunction to a rank-1 Hartree
state and get the ci coefficient of the largest configuration.
Then, the configuration is subtracted from the tda wavefunction and
redo the first step to get the second largest configuration. The
two steps continue until the thresh is achieved.
Parameters
----------
thresh: float, optional
the threshold to stop the analysis procedure of each root.
:math:`\sum_i |c_i|^2 > thresh`. Default is 0.8.
alias: dict, optional
The alias of each site. For example, ``alias={0:"v_0", 1:"v_2",
2:"v_1"}``. Default is `None`.
tda_m_trunc: int, optional
the ``m`` to compress a tda wavefunction. Default is 20.
return_compressed_mps: bool, optional
If ``True``, return the tda excited state as a single compressed
mps. Default is `False`.
Returns
-------
configs: dict
The dominant configration of each root.
``configs = {0:[(config0, config_name0, ci_coeff0),(config1,
config_name1, ci_coeff1),...], 1:...}``
compressed_mps: List[renormalizer.mps.Mps]
see the description in ``return_compressed_mps``.
Note
----
The compressed_mps is an approximation of the tda wavefunction with
``m=tda_m_trunc``.
"""
mps_l_cano, mps_r_cano, tangent_u, tda_coeff_list = self.wfn
if alias is not None:
assert len(alias) == mps_l_cano.site_num
compressed_mps = []
for iroot in range(self.nroots):
logger.info(f"iroot: {iroot}")
tda_coeff = tda_coeff_list[iroot]
mps_tangent_list = []
weight = []
for ims in range(mps_l_cano.site_num):
if tangent_u[ims] is None:
assert tda_coeff[ims] is None
continue
weight.append(np.sum(tda_coeff[ims]**2))
mps_tangent = merge(mps_l_cano, mps_r_cano, ims+1)
mps_tangent[ims] = asnumpy(tensordot(tangent_u[ims],
tda_coeff[ims],[-1,0]))
mps_tangent_list.append(mps_tangent)
assert np.allclose(np.sum(weight), 1)
# sort the mps_tangent from large weight to small weight
mps_tangent_list = [mps_tangent_list[i] for i in np.argsort(weight,axis=None)[::-1]]
coeff_square_sum = 0
mps_delete = None
config_visited = []
while coeff_square_sum < thresh:
if mps_delete is None:
# first compress it to M=tda_m_trunc
mps_rank1 = compressed_sum(mps_tangent_list, batchsize=5,
temp_m_trunc=tda_m_trunc)
else:
mps_rank1 = compressed_sum([mps_delete] + mps_tangent_list,
batchsize=5, temp_m_trunc=tda_m_trunc)
if coeff_square_sum == 0 and return_compressed_mps:
compressed_mps.append(mps_rank1.copy())
mps_rank1 = mps_rank1.canonicalise().compress(temp_m_trunc=1)
# get config with the largest coeff
config = []
for ims, ms in enumerate(mps_rank1):
ms = ms.array.flatten()**2
quanta = int(np.argmax(ms))
config.append(quanta)
# check if the config has been visited
if config in config_visited:
break
config_visited.append(config)
ci_coeff_list = []
for mps_tangent in mps_tangent_list:
sentinel = xp.ones((1,1))
for ims, ms in enumerate(mps_tangent):
sentinel = sentinel.dot(asxp(ms[:,config[ims],:]))
ci_coeff_list.append(float(sentinel[0,0]))
ci_coeff = np.sum(ci_coeff_list)
coeff_square_sum += ci_coeff**2
if alias is not None:
config_name = [f"{quanta}"+f"{alias[isite]}" for isite, quanta
in enumerate(config) if quanta != 0]
config_name = " ".join(config_name)
self.configs[iroot].append((config, config_name, ci_coeff))
logger.info(f"config: {config}, {config_name}")
else:
self.configs[iroot].append((config, ci_coeff))
logger.info(f"config: {config}")
logger.info(f"ci_coeff: {ci_coeff}, weight:{ci_coeff**2}")
condition = {dof:config[idof] for idof, dof in
enumerate(self.model.dofs)}
mps_delete_increment = Mps.hartree_product_state(self.model, condition).scale(-ci_coeff)
if mps_delete is None:
mps_delete = mps_delete_increment
else:
mps_delete = mps_delete + mps_delete_increment
logger.info(f"coeff_square_sum: {coeff_square_sum}")
return self.configs, compressed_mps
