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