Source code for qmctorch.wavefunction.pooling.orbital_configurations

import torch


[docs]class OrbitalConfigurations: def __init__(self, mol): # self.mol = mol self.nup = mol.nup self.ndown = mol.ndown self.nelec = self.nup + self.ndown self.norb = mol.basis.nmo
[docs] def get_configs(self, configs): """Get the configuratio in the CI expansion Args: configs (str): name of the configs we want mol (mol object): molecule object Returns: tuple(torch.LongTensor,torch.LongTensor): the spin up/spin down electronic confs """ if isinstance(configs, str): configs = configs.lower() if isinstance(configs, torch.Tensor): return configs elif configs == 'ground_state': return self._get_ground_state_config() elif configs.startswith('cas('): nelec, norb = eval(configs.lstrip("cas")) self.sanity_check(nelec, norb) nocc, nvirt = self._get_orb_number(nelec, norb) return self._get_cas_config(nocc, nvirt, nelec) elif configs.startswith('single('): nelec, norb = eval(configs.lstrip("single")) self.sanity_check(nelec, norb) nocc, nvirt = self._get_orb_number(nelec, norb) return self._get_single_config(nocc, nvirt) elif configs.startswith('single_double('): nelec, norb = eval(configs.lstrip("single_double")) self.sanity_check(nelec, norb) nocc, nvirt = self._get_orb_number(nelec, norb) return self._get_single_double_config(nocc, nvirt) else: print(configs, " not recognized as valid configuration") print('Options are : ground_state') print(' single(nelec,norb)') print(' single_double(nelec,norb)') print(' cas(nelec,norb)') raise ValueError("Config error")
[docs] def sanity_check(self, nelec, norb): """Check if the number of elec/orb is consistent with the properties of the molecule Args: nelec (int): required number of electrons in config norb (int): required number of orb in config """ if nelec > self.nelec: raise ValueError( 'required number of electron in config too large') if norb > self.norb: raise ValueError( 'required number of orbitals in config too large')
def _get_ground_state_config(self): """Return only the ground state configuration Args: mol (mol): mol object Returns: tuple(torch.LongTensor,torch.LongTensor): the spin up/spin down electronic confs """ _gs_up = list(range(self.nup)) _gs_down = list(range(self.ndown)) cup, cdown = [_gs_up], [_gs_down] return (torch.LongTensor(cup), torch.LongTensor(cdown)) def _get_single_config(self, nocc, nvirt): """Get the confs of the singlet conformations Args: mol (mol): mol object nocc (int): number of occupied orbitals in the active space nvirt (int): number of virtual orbitals in the active space """ _gs_up = list(range(self.nup)) _gs_down = list(range(self.ndown)) cup, cdown = [_gs_up], [_gs_down] for iocc in range( self.nup - 1, self.nup - 1 - nocc, -1): for ivirt in range(self.nup, self.nup + nvirt, 1): _xt = self._create_excitation( _gs_up.copy(), iocc, ivirt) cup, cdown = self._append_excitations( cup, cdown, _xt, _gs_down) _xt = self._create_excitation( _gs_down.copy(), iocc, ivirt) cup, cdown = self._append_excitations( cup, cdown, _gs_up, _xt) return (torch.LongTensor(cup), torch.LongTensor(cdown)) def _get_single_double_config(self, nocc, nvirt): """Get the confs of the single + double Args: nelec (int): number of electrons in the active space norb (int) : number of orbitals in the active space """ _gs_up = list(range(self.nup)) _gs_down = list(range(self.ndown)) cup, cdown = self._get_single_config(nocc, nvirt) cup = cup.tolist() cdown = cdown.tolist() idx_occ_up = list( range(self.nup - 1, self.nup - 1 - nocc, -1)) idx_vrt_up = list(range(self.nup, self.nup + nvirt, 1)) idx_occ_down = list(range( self.ndown - 1, self.ndown - 1 - nocc, -1)) idx_vrt_down = list(range(self.ndown, self.ndown + nvirt, 1)) # ground, single and double with 1 elec excited per spin for iocc_up in idx_occ_up: for ivirt_up in idx_vrt_up: for iocc_down in idx_occ_down: for ivirt_down in idx_vrt_down: _xt_up = self._create_excitation( _gs_up.copy(), iocc_up, ivirt_up) _xt_down = self._create_excitation( _gs_down.copy(), iocc_down, ivirt_down) cup, cdown = self._append_excitations( cup, cdown, _xt_up, _xt_down) # double with 2elec excited on spin up for occ1, occ2 in torch.combinations(torch.as_tensor(idx_occ_up), r=2): for vrt1, vrt2 in torch.combinations(torch.as_tensor(idx_vrt_up), r=2): _xt_up = self._create_excitation( _gs_up.copy(), occ1, vrt2) _xt_up = self._create_excitation(_xt_up, occ2, vrt1) cup, cdown = self._append_excitations( cup, cdown, _xt_up, _gs_down) # double with 2elec excited per spin for occ1, occ2 in torch.combinations(torch.as_tensor(idx_occ_down), r=2): for vrt1, vrt2 in torch.combinations(torch.as_tensor(idx_vrt_down), r=2): _xt_down = self._create_excitation( _gs_down.copy(), occ1, vrt2) _xt_down = self._create_excitation( _xt_down, occ2, vrt1) cup, cdown = self._append_excitations( cup, cdown, _gs_up, _xt_down) return (torch.LongTensor(cup), torch.LongTensor(cdown)) def _get_cas_config(self, nocc, nvirt, nelec): """get confs of the CAS Args: nocc (int): number of occupied orbitals in the CAS nvirt ([type]): number of virt orbitals in the CAS """ from itertools import combinations, product idx_low, idx_high = self.nup - nocc, self.nup + nvirt orb_index_up = range(idx_low, idx_high) idx_frz = list(range(idx_low)) _cup = [idx_frz + list(l) for l in list(combinations(orb_index_up, nelec // 2))] idx_low, idx_high = self.nup - nocc - 1, self.nup + nvirt - 1 _cdown = [ idx_frz + list(l) for l in list( combinations(orb_index_up, nelec // 2))] confs = list(product(_cup, _cdown)) cup, cdown = [], [] for c in confs: cup.append(c[0]) cdown.append(c[1]) return (torch.LongTensor(cup), torch.LongTensor(cdown)) def _get_orb_number(self, nelec, norb): """compute the number of occupied and virtual orbital __ PER SPIN __ __ ONLY VALID For spin up/down ___ Args: nelec (int): total number of elec in the CAS norb (int): total number of orb in the CAS Returns: [int,int]: number of occpuied/virtual orb per spi """ nocc = nelec // 2 nvirt = norb - nocc return nocc, nvirt def _create_excitation(self, conf, iocc, ivirt): return self._create_excitation_replace(conf, iocc, ivirt) @staticmethod def _create_excitation_ordered(conf, iocc, ivirt): """promote an electron from iocc to ivirt Args: conf (list): index of the occupied orbitals iocc (int): index of the occupied orbital ivirt (int): index of the virtual orbital Returns: list: new configuration by increasing order e.g: 4->6 leads to : [0,1,2,3,5,6] Note: if that method is used to define the exciation index permutation must be accounted for when computing the determinant as det(A[:,[0,1,2,3]]) = -det(A[:,[0,1,3,2]]) see : ExcitationMask.get_index_unique_single() in oribtal_projector.py """ conf.pop(iocc) conf += [ivirt] return conf @staticmethod def _create_excitation_replace(conf, iocc, ivirt): """promote an electron from iocc to ivirt Args: conf (list): index of the occupied orbitals iocc (int): index of the occupied orbital ivirt (int): index of the virtual orbital Returns: list: new configuration not ordered e.g.: 4->6 leads tpo : [0,1,2,3,6,5] """ conf[iocc] = ivirt return conf @staticmethod def _append_excitations(cup, cdown, new_cup, new_cdown): """Append new excitations Args: cup (list): configurations of spin up cdown (list): configurations of spin down new_cup (list): new spin up confs new_cdown (list): new spin down confs """ cup.append(new_cup) cdown.append(new_cdown) return cup, cdown
[docs]def get_excitation(configs): """get the excitation data Args: configs (tuple): configuratin of the electrons Returns: exc_up, exc_down : index of the obitals in the excitaitons [i,j],[l,m] : excitation i -> l, j -> l """ exc_up, exc_down = [], [] for ic, (cup, cdown) in enumerate(zip(configs[0], configs[1])): set_cup = set(tuple(cup.tolist())) set_cdown = set(tuple(cdown.tolist())) if ic == 0: set_gs_up = set_cup set_gs_down = set_cdown else: exc_up.append([list(set_gs_up.difference(set_cup)), list(set_cup.difference(set_gs_up))]) exc_down.append([list(set_gs_down.difference(set_cdown)), list(set_cdown.difference(set_gs_down))]) return (exc_up, exc_down)
[docs]def get_unique_excitation(configs): """get the unique excitation data Args: configs (tuple): configuratin of the electrons Returns: exc_up, exc_down : index of the obitals in the excitaitons [i,j],[l,m] : excitation i -> l, j -> l index_up, index_down : index map for the unique exc [0,0,...], [0,1,...] means that 1st : excitation is composed of unique_up[0]*unique_down[0] 2nd : excitation is composed of unique_up[0]*unique_down[1] .... """ uniq_exc_up, uniq_exc_down = [], [] index_uniq_exc_up, index_uniq_exc_down = [], [] for ic, (cup, cdown) in enumerate(zip(configs[0], configs[1])): set_cup = set(tuple(cup.tolist())) set_cdown = set(tuple(cdown.tolist())) if ic == 0: set_gs_up = set_cup set_gs_down = set_cdown exc_up = [list(set_gs_up.difference(set_cup)), list(set_cup.difference(set_gs_up))] exc_down = [list(set_gs_down.difference(set_cdown)), list(set_cdown.difference(set_gs_down))] if exc_up not in uniq_exc_up: uniq_exc_up.append(exc_up) if exc_down not in uniq_exc_down: uniq_exc_down.append(exc_down) index_uniq_exc_up.append(uniq_exc_up.index(exc_up)) index_uniq_exc_down.append( uniq_exc_down.index(exc_down)) return (uniq_exc_up, uniq_exc_down), (index_uniq_exc_up, index_uniq_exc_down)