Python time.clock() Examples

def test_write_conf(self):
        rdd = self.sc.parallelize([{'key':i, 'text':i, 'int':i} for i in range(10)])
        save = partial(rdd.saveToCassandra, self.keyspace, self.table)

        save(batch_size=100)
        save(batch_buffer_size=100)
        save(batch_grouping_key='replica_set')
        save(batch_grouping_key='partition')
        save(consistency_level='ALL')
        save(consistency_level=ConsistencyLevel.LOCAL_QUORUM)
        save(parallelism_level=10)
        save(throughput_mibps=10)
        save(ttl=5)
        save(ttl=timedelta(minutes=30))
        save(timestamp=time.clock() * 1000 * 1000)
        save(timestamp=datetime.now())
        save(metrics_enabled=True)
        save(write_conf=WriteConf(ttl=3, metrics_enabled=True)) 

def _make_shared(self):
        cput0 = (time.clock(), time.time())

        kconserv = self.kconserv
        t1, t2, eris = self.t1, self.t2, self.eris

        self.Foo = imd.Foo(self._cc, t1, t2, eris, kconserv)
        self.Fvv = imd.Fvv(self._cc, t1, t2, eris, kconserv)
        self.Fov = imd.Fov(self._cc, t1, t2, eris, kconserv)

        # 2 virtuals
        self.Wovvo = imd.Wovvo(self._cc, t1, t2, eris, kconserv)
        self.Woovv = eris.oovv

        self._made_shared = True
        logger.timer_debug1(self, 'EOM-CCSD shared intermediates', *cput0)
        return self 

def make_ip(self):
        if not self._made_shared:
            self._make_shared()

        cput0 = (time.clock(), time.time())

        kconserv = self.kconserv
        t1, t2, eris = self.t1, self.t2, self.eris

        # 0 or 1 virtuals
        self.Woooo = imd.Woooo(self._cc, t1, t2, eris, kconserv)
        self.Wooov = imd.Wooov(self._cc, t1, t2, eris, kconserv)
        self.Wovoo = imd.Wovoo(self._cc, t1, t2, eris, kconserv)

        self.made_ip_imds = True
        logger.timer_debug1(self, 'EOM-CCSD IP intermediates', *cput0)
        return self 

def make_ea(self):
        if not self._made_shared:
            self._make_shared()

        cput0 = (time.clock(), time.time())

        kconserv = self.kconserv
        t1, t2, eris = self.t1, self.t2, self.eris

        # FIXME DELETE WOOOO
        # 0 or 1 virtuals
        self.Woooo = imd.Woooo(self._cc, t1, t2, eris, kconserv)
        # 3 or 4 virtuals
        self.Wvovv = imd.Wvovv(self._cc, t1, t2, eris, kconserv)
        self.Wvvvv = imd.Wvvvv(self._cc, t1, t2, eris, kconserv)
        self.Wvvvo = imd.Wvvvo(self._cc, t1, t2, eris, kconserv)

        self.made_ea_imds = True
        logger.timer_debug1(self, 'EOM-CCSD EA intermediates', *cput0)
        return self 

def make_t3p2_ea(self, cc):
        cput0 = (time.clock(), time.time())

        t1, t2, eris = cc.t1, cc.t2, self.eris
        delta_E_corr, pt1, pt2, Wovoo, Wvvvo = \
            imd.get_t3p2_imds_slow(cc, t1, t2, eris)
        self.t1 = pt1
        self.t2 = pt2

        self._made_shared = False  # Force update
        self.make_ea()  # Make after t1/t2 updated
        self.Wvvvo = self.Wvvvo + Wvvvo

        self.made_ea_imds = True
        logger.timer_debug1(self, 'EOM-CCSD(T)a EA intermediates', *cput0)
        return self 

def make_ee(self):
        if not self._made_shared:
            self._make_shared()

        cput0 = (time.clock(), time.time())

        kconserv = self.kconserv
        t1, t2, eris = self.t1, self.t2, self.eris

        if not self.made_ip_imds:
            # 0 or 1 virtuals
            self.Woooo = imd.Woooo(self._cc, t1, t2, eris, kconserv)
            self.Wooov = imd.Wooov(self._cc, t1, t2, eris, kconserv)
            self.Wovoo = imd.Wovoo(self._cc, t1, t2, eris, kconserv)
        if not self.made_ea_imds:
            # 3 or 4 virtuals
            self.Wvovv = imd.Wvovv(self._cc, t1, t2, eris, kconserv)
            self.Wvvvv = imd.Wvvvv(self._cc, t1, t2, eris, kconserv)
            self.Wvvvo = imd.Wvvvo(self._cc, t1, t2, eris, kconserv, self.Wvvvv)

        self.made_ee_imds = True
        logger.timer(self, 'EOM-CCSD EE intermediates', *cput0)
        return self 

def _make_shared_2e(self):
        cput0 = (time.clock(), time.time())
        log = logger.Logger(self.stdout, self.verbose)

        t1, t2, eris = self.t1, self.t2, self.eris
        kconserv = self.kconserv

        if self._fimd is not None:
            nkpts, nocc, nvir = t1.shape
            ovov_dest = self._fimd.create_dataset('ovov', (nkpts, nkpts, nkpts, nocc, nvir, nocc, nvir), t1.dtype.char)
            ovvo_dest = self._fimd.create_dataset('ovvo', (nkpts, nkpts, nkpts, nocc, nvir, nvir, nocc), t1.dtype.char)
        else:
            ovov_dest = ovvo_dest = None

        # 2 virtuals
        self.Wovov = imd.Wovov(t1, t2, eris, kconserv, ovov_dest)
        self.Wovvo = imd.Wovvo(t1, t2, eris, kconserv, ovvo_dest)
        self.Woovv = eris.oovv

        log.timer('EOM-CCSD shared two-electron intermediates', *cput0) 

def make_t3p2_ip(self, cc):
        cput0 = (time.clock(), time.time())

        t1, t2, eris = cc.t1, cc.t2, self.eris
        delta_E_tot, pt1, pt2, Wovoo, Wvvvo = \
            imd.get_t3p2_imds(cc, t1, t2, eris)
        self.t1 = pt1
        self.t2 = pt2

        self._made_shared_2e = False  # Force update
        self.make_ip()  # Make after t1/t2 updated
        self.Wovoo = self.Wovoo + Wovoo

        self.made_ip_imds = True
        logger.timer_debug1(self, 'EOM-CCSD(T)a IP intermediates', *cput0)
        return self 

def make_t3p2_ea(self, cc):
        cput0 = (time.clock(), time.time())

        t1, t2, eris = cc.t1, cc.t2, self.eris
        delta_E_tot, pt1, pt2, Wovoo, Wvvvo = \
            imd.get_t3p2_imds(cc, t1, t2, eris)
        self.t1 = pt1
        self.t2 = pt2

        self._made_shared_2e = False  # Force update
        self.make_ea()  # Make after t1/t2 updated
        self.Wvvvo = self.Wvvvo + Wvvvo

        self.made_ea_imds = True
        logger.timer_debug1(self, 'EOM-CCSD(T)a EA intermediates', *cput0)
        return self 

def _make_shared_2e(self):
        cput0 = (time.clock(), time.time())
        log = logger.Logger(self.stdout, self.verbose)

        t1,t2,eris = self.t1, self.t2, self.eris
        kconserv = self.kconserv

        # TODO: check whether to hold Wovov Wovvo in memory
        if self._fimd is None:
            self._fimd = lib.H5TmpFile()
        nkpts, nocc, nvir = t1.shape
        self._fimd.create_dataset('ovov', (nkpts,nkpts,nkpts,nocc,nvir,nocc,nvir), t1.dtype.char)
        self._fimd.create_dataset('ovvo', (nkpts,nkpts,nkpts,nocc,nvir,nvir,nocc), t1.dtype.char)

        # 2 virtuals
        self.Wovov = imd.Wovov(t1,t2,eris,kconserv, self._fimd['ovov'])
        self.Wovvo = imd.Wovvo(t1,t2,eris,kconserv, self._fimd['ovvo'])
        self.Woovv = eris.oovv

        log.timer('EOM-CCSD shared two-electron intermediates', *cput0) 

def make_ip(self, ip_partition=None):
        self._make_shared_1e()
        if self._made_shared_2e is False and ip_partition != 'mp':
            self._make_shared_2e()
            self._made_shared_2e = True

        cput0 = (time.clock(), time.time())
        log = logger.Logger(self.stdout, self.verbose)

        t1,t2,eris = self.t1, self.t2, self.eris
        kconserv = self.kconserv

        nkpts, nocc, nvir = t1.shape
        self._fimd.create_dataset('oooo', (nkpts,nkpts,nkpts,nocc,nocc,nocc,nocc), t1.dtype.char)
        self._fimd.create_dataset('ooov', (nkpts,nkpts,nkpts,nocc,nocc,nocc,nvir), t1.dtype.char)
        self._fimd.create_dataset('ovoo', (nkpts,nkpts,nkpts,nocc,nvir,nocc,nocc), t1.dtype.char)

        # 0 or 1 virtuals
        if ip_partition != 'mp':
            self.Woooo = imd.Woooo(t1,t2,eris,kconserv, self._fimd['oooo'])
        self.Wooov = imd.Wooov(t1,t2,eris,kconserv, self._fimd['ooov'])
        self.Wovoo = imd.Wovoo(t1,t2,eris,kconserv, self._fimd['ovoo'])
        self.made_ip_imds = True
        log.timer('EOM-CCSD IP intermediates', *cput0) 

def _make_shared(self):
        cput0 = (time.clock(), time.time())

        t1, t2, eris = self.t1, self.t2, self.eris
        self.Foo, self.FOO = kintermediates_uhf.Foo(self._cc, t1, t2, eris)
        self.Fvv, self.FVV = kintermediates_uhf.Fvv(self._cc, t1, t2, eris)
        self.Fov, self.FOV = kintermediates_uhf.Fov(self._cc, t1, t2, eris)

        # 2 virtuals
        self.Wovvo, self.WovVO, self.WOVvo, self.WOVVO = kintermediates_uhf.Wovvo(self._cc, t1, t2, eris)
        self.Woovv, self.WooVV, self.WOOvv, self.WOOVV = kintermediates_uhf.Woovv(self._cc, t1, t2, eris)
        self.Wovov = eris.ovov - np.asarray(eris.ovov).transpose(2,1,0,5,4,3,6)
        self.WOVOV = eris.OVOV - np.asarray(eris.OVOV).transpose(2,1,0,5,4,3,6)
        self.WovOV = eris.ovOV
        self.WOVov = None

        self._made_shared = True
        logger.timer_debug1(self, 'EOM-KCCSD shared intermediates', *cput0)
        return self 

def make_ip(self):
        if not self._made_shared:
            self._make_shared()

        kconserv = self.kconserv
        cput0 = (time.clock(), time.time())

        t1, t2, eris = self.t1, self.t2, self.eris

        # 0 or 1 virtuals
        self.Woooo, self.WooOO, _         , self.WOOOO = kintermediates_uhf.Woooo(self._cc, t1, t2, eris)
        self.Wooov, self.WooOV, self.WOOov, self.WOOOV = kintermediates_uhf.Wooov(self._cc, t1, t2, eris, kconserv)  # TODO
        self.Woovo, self.WooVO, self.WOOvo, self.WOOVO = kintermediates_uhf.Woovo(self._cc, t1, t2, eris)  # TODO

        self.made_ip_imds = True
        logger.timer_debug1(self, 'EOM-KUCCSD IP intermediates', *cput0)
        return self 

def make_ea(self):
        if not self._made_shared:
            self._make_shared()

        cput0 = (time.clock(), time.time())

        t1, t2, eris = self.t1, self.t2, self.eris

        # 3 or 4 virtuals
        #self.Wvovv, self.WvoVV, self.WVOvv, self.WVOVV = kintermediates_uhf.Wvovv(self._cc, t1, t2, eris)
        self.Wvvov, self.WvvOV, self.WVVov, self.WVVOV = kintermediates_uhf.Wvvov(self._cc, t1, t2, eris)
        if eris.vvvv is not None:
            self.Wvvvv, self.WvvVV, self.WVVVV = Wvvvv = kintermediates_uhf.Wvvvv(self._cc, t1, t2, eris)
        else:
            self.Wvvvv = self.WvvVV = self.WVVVV = None
        self.Wvvvo, self.WvvVO, self.WVVvo, self.WVVVO = kintermediates_uhf.Wvvvo(self._cc, t1, t2, eris)

        self.made_ea_imds = True
        logger.timer_debug1(self, 'EOM-KUCCSD EA intermediates', *cput0)
        return self 

def kernel(self, mo_energy=None, mo_coeff=None, mo_occ=None, atmlst=None):
        cput0 = (time.clock(), time.time())
        if mo_energy is None: mo_energy = self.base.mo_energy
        if mo_coeff is None: mo_coeff = self.base.mo_coeff
        if mo_occ is None: mo_occ = self.base.mo_occ
        if atmlst is None:
            atmlst = self.atmlst
        else:
            self.atmlst = atmlst

        if self.verbose >= logger.INFO:
            self.dump_flags()

        de = self.grad_elec(mo_energy, mo_coeff, mo_occ, atmlst)
        self.de = de + self.grad_nuc(atmlst=atmlst)
        logger.timer(self, 'SCF gradients', *cput0)
        self._finalize()
        return self.de 

def kernel(self, mo_energy=None, mo_coeff=None, td_e=None, td_xy=None,
               eris=None, orbs=None):
        if mo_coeff is None:
            mo_coeff = self._scf.mo_coeff
        if mo_energy is None:
            mo_energy = self._scf.mo_energy
        if td_e is None:
            td_e = self._tdscf.e
        if td_xy is None:
            td_xy = self._tdscf.xy

        cput0 = (time.clock(), time.time())
        self.dump_flags()
        self.converged, self.mo_energy, self.mo_coeff = \
                kernel(self, mo_energy, mo_coeff, td_e, td_xy,
                       eris=eris, orbs=orbs, verbose=self.verbose)

        logger.timer(self, 'GW', *cput0)
        return self.mo_energy 

def _make_eris_incore(mycc, mo_coeff=None, ao2mofn=None):
    cput0 = (time.clock(), time.time())
    eris = _ChemistsERIs()
    eris._common_init_(mycc, mo_coeff)
    nocc = eris.nocc
    nmo = eris.fock.shape[0]

    if callable(ao2mofn):
        eri1 = ao2mofn(eris.mo_coeff).reshape([nmo]*4)
    else:
        eri1 = ao2mo.incore.full(mycc._scf._eri, eris.mo_coeff)
        eri1 = ao2mo.restore(1, eri1, nmo)
    eris.oooo = eri1[:nocc,:nocc,:nocc,:nocc].copy()
    eris.ovoo = eri1[:nocc,nocc:,:nocc,:nocc].copy()
    eris.ovov = eri1[:nocc,nocc:,:nocc,nocc:].copy()
    eris.oovv = eri1[:nocc,:nocc,nocc:,nocc:].copy()
    eris.ovvo = eri1[:nocc,nocc:,nocc:,:nocc].copy()
    eris.ovvv = eri1[:nocc,nocc:,nocc:,nocc:].copy()
    logger.timer(mycc, 'GW integral transformation', *cput0)
    return eris 

def get_jk(self, mol=None, dm=None, hermi=1, with_j=True, with_k=True,
               omega=None):
        if mol is None: mol = self.mol
        if dm is None: dm = self.make_rdm1()
        cpu0 = (time.clock(), time.time())
        if self.direct_scf and self.opt is None:
            self.opt = self.init_direct_scf(mol)

        if with_j and with_k:
            vj, vk = get_jk(mol, dm, hermi, self.opt, with_j, with_k, omega)
        else:
            if with_j:
                prescreen = 'CVHFnrs8_vj_prescreen'
            else:
                prescreen = 'CVHFnrs8_vk_prescreen'
            with lib.temporary_env(self.opt, prescreen=prescreen):
                vj, vk = get_jk(mol, dm, hermi, self.opt, with_j, with_k, omega)

        logger.timer(self, 'vj and vk', *cpu0)
        return vj, vk 

def solve_nos1(fvind, mo_energy, mo_occ, h1,
               max_cycle=20, tol=1e-9, hermi=False, verbose=logger.WARN):
    '''For field independent basis. First order overlap matrix is zero'''
    log = logger.new_logger(verbose=verbose)
    t0 = (time.clock(), time.time())

    e_a = mo_energy[mo_occ==0]
    e_i = mo_energy[mo_occ>0]
    e_ai = 1 / lib.direct_sum('a-i->ai', e_a, e_i)
    mo1base = h1 * -e_ai

    def vind_vo(mo1):
        v = fvind(mo1.reshape(h1.shape)).reshape(h1.shape)
        v *= e_ai
        return v.ravel()
    mo1 = lib.krylov(vind_vo, mo1base.ravel(),
                     tol=tol, max_cycle=max_cycle, hermi=hermi, verbose=log)
    log.timer('krylov solver in CPHF', *t0)
    return mo1.reshape(h1.shape), None

# h1 shape is (:,nocc+nvir,nocc) 

def make_ip(self):
        if self._made_shared is False:
            self._make_shared()
            self._made_shared = True

        cput0 = (time.clock(), time.time())
        log = logger.Logger(self.cc.stdout, self.cc.verbose)

        t1,t2,eris = self.cc.t1, self.cc.t2, self.cc.eris

        # 0 or 1 virtuals
        self.Woooo = imd.Woooo(t1,t2,eris)
        self.Wooov = imd.Wooov(t1,t2,eris)
        self.Wovoo = imd.Wovoo(t1,t2,eris)

        self.made_ip_imds = True
        log.timer('EOM-CCSD IP intermediates', *cput0) 

def make_ea(self):
        if self._made_shared is False:
            self._make_shared()
            self._made_shared = True

        cput0 = (time.clock(), time.time())
        log = logger.Logger(self.cc.stdout, self.cc.verbose)

        t1,t2,eris = self.cc.t1, self.cc.t2, self.cc.eris

        # 3 or 4 virtuals
        self.Wvovv = imd.Wvovv(t1,t2,eris)
        self.Wvvvv = imd.Wvvvv(t1,t2,eris)
        self.Wvvvo = imd.Wvvvo(t1,t2,eris,self.Wvvvv)

        self.made_ea_imds = True
        log.timer('EOM-CCSD EA intermediates', *cput0) 

def make_ip(self, ip_partition=None):
        self._make_shared_1e()
        if self._made_shared_2e is False and ip_partition != 'mp':
            self._make_shared_2e()
            self._made_shared_2e = True

        cput0 = (time.clock(), time.time())
        log = logger.Logger(self.stdout, self.verbose)

        t1,t2,eris = self.t1, self.t2, self.eris

        # 0 or 1 virtuals
        if ip_partition != 'mp':
            self.Woooo = imd.Woooo(t1,t2,eris)
        self.Wooov = imd.Wooov(t1,t2,eris)
        self.Wovoo = imd.Wovoo(t1,t2,eris)
        self.made_ip_imds = True
        log.timer('EOM-CCSD IP intermediates', *cput0) 

def make_ea(self, ea_partition=None):
        self._make_shared_1e()
        if self._made_shared_2e is False and ea_partition != 'mp':
            self._make_shared_2e()
            self._made_shared_2e = True

        cput0 = (time.clock(), time.time())
        log = logger.Logger(self.stdout, self.verbose)

        t1,t2,eris = self.t1, self.t2, self.eris

        # 3 or 4 virtuals
        self.Wvovv = imd.Wvovv(t1,t2,eris)
        if ea_partition == 'mp':
            self.Wvvvo = imd.Wvvvo(t1,t2,eris)
        else:
            self.Wvvvv = imd.Wvvvv(t1,t2,eris)
            self.Wvvvo = imd.Wvvvo(t1,t2,eris,self.Wvvvv)
        self.made_ea_imds = True
        log.timer('EOM-CCSD EA intermediates', *cput0) 

def _make_shared(self):
        cput0 = (time.clock(), time.time())

        t1, t2, eris = self.t1, self.t2, self.eris
        self.Foo, self.FOO = uintermediates.Foo(t1, t2, eris)
        self.Fvv, self.FVV = uintermediates.Fvv(t1, t2, eris)
        self.Fov, self.FOV = uintermediates.Fov(t1, t2, eris)

        # 2 virtuals
        self.Wovvo, self.WovVO, self.WOVvo, self.WOVVO, self.WoVVo, self.WOvvO = \
                uintermediates.Wovvo(t1, t2, eris)
        Wovov = np.asarray(eris.ovov)
        WOVOV = np.asarray(eris.OVOV)
        Wovov = Wovov - Wovov.transpose(0,3,2,1)
        WOVOV = WOVOV - WOVOV.transpose(0,3,2,1)
        self.Wovov = Wovov
        self.WovOV = eris.ovOV
        self.WOVov = None
        self.WOVOV = WOVOV

        self._made_shared = True
        logger.timer_debug1(self, 'EOM-CCSD shared intermediates', *cput0)
        return self 

def make_ip(self):
        if not self._made_shared:
            self._make_shared()

        cput0 = (time.clock(), time.time())

        t1, t2, eris = self.t1, self.t2, self.eris

        # 0 or 1 virtuals
        self.Woooo, self.WooOO, _         , self.WOOOO = uintermediates.Woooo(t1, t2, eris)
        self.Wooov, self.WooOV, self.WOOov, self.WOOOV = uintermediates.Wooov(t1, t2, eris)
        self.Woovo, self.WooVO, self.WOOvo, self.WOOVO = uintermediates.Woovo(t1, t2, eris)

        self.made_ip_imds = True
        logger.timer_debug1(self, 'EOM-UCCSD IP intermediates', *cput0)
        return self 

def tic(self, name='default'):
        # using time.time instead of time.clock because time time.clock
        # does not normalize for multithreading
        torch.cuda.synchronize()
        self._start_time[name] = time.time() 

def _log_data(self, net_info_type='active_only'):
        log_list = [self.num_gen, self.num_eval, time.clock(), self.pop[0].eval, self.pop[0].count_active_node()]
        if net_info_type == 'active_only':
            log_list.append(self.pop[0].active_net_list())
        elif net_info_type == 'full':
            log_list += self.pop[0].gene.flatten().tolist()
        else:
            pass
        return log_list 

def benchmark(mod, dry_run=10, iterations=10):
    if len(mod._context) == 1:
        ctx = mod._context[0]
    else:
        ctx = mx.cpu()
    data = [mx.random.uniform(-1.0, 1.0, shape=shape, ctx=ctx) for _, shape in mod.data_shapes]
    label = [mx.nd.array(np.random.randint(1, 100, size=shape), ctx=ctx) for _, shape in mod.label_shapes]
    batch = mx.io.DataBatch(data, label)

    # dry run
    for i in range(dry_run):
        mod.forward(batch, is_train=True)
        mod.backward()
        for output in mod.get_outputs(merge_multi_context=False)[0]:
            output.wait_to_read()
        mod.update()

    t0 = time.clock()

    profiler.set_state('run')
    # real run
    for i in range(iterations):
        mod.forward(batch, is_train=True)
        mod.backward()
        mod.update()
        for output in mod.get_outputs(merge_multi_context=False)[0]:
            output.wait_to_read()
    profiler.set_state('stop')

    t1 = time.clock()
    return (t1 - t0)*1000.0 / iterations 

def test_profiler():
    iter_num = 5
    begin_profiling_iter = 2
    end_profiling_iter = 4

    enable_profiler('test_profiler.json', False, False)

    A = mx.sym.Variable('A')
    B = mx.sym.Variable('B')
    C = mx.symbol.dot(A, B)

    executor = C.simple_bind(mx.cpu(1), 'write', A=(4096, 4096), B=(4096, 4096))

    a = mx.random.uniform(-1.0, 1.0, shape=(4096, 4096))
    b = mx.random.uniform(-1.0, 1.0, shape=(4096, 4096))

    a.copyto(executor.arg_dict['A'])
    b.copyto(executor.arg_dict['B'])

    print("execution begin")
    for i in range(iter_num):
        print("Iteration {}/{}".format(i + 1, iter_num))
        if i == begin_profiling_iter:
            t0 = time.clock()
            profiler.set_state('run')
        if i == end_profiling_iter:
            t1 = time.clock()
            profiler.set_state('stop')
        executor.forward()
        c = executor.outputs[0]
        c.wait_to_read()
    print("execution end")
    duration = t1 - t0
    print('duration: {0}s'.format(duration))
    print('          {0}ms/operator'.format(duration*1000/iter_num))
    profiler.dump(True)
    profiler.set_state('stop') 

def run_one(genNewPlanets=True, rewindPlanets=True, outpath='.'):    
    # wrap the run_sim in a try/except loop
    nbmax = 10
    for attempt in range(nbmax):
        try:
            # run one survey simulation
            SS.run_sim()
            DRM = SS.DRM[:]
            systems = SS.SimulatedUniverse.dump_systems()
            systems['MsTrue'] = SS.TargetList.MsTrue
            systems['MsEst'] = SS.TargetList.MsEst
            seed = SS.seed
        except Exception as e:
            # if anything goes wrong, log the error and reset simulation
            with open(os.path.join(outpath,'log.err'), 'ab') as f:
                f.write(repr(e))
                f.write('\n')
                f.write(traceback.format_exc())
                f.write('\n\n')
            
            SS.reset_sim()
        else:
            break
    else:
        raise ValueError("Unsuccessful run_sim after %s reset_sim attempts"%nbmax)
    
    # reset simulation at the end of each simulation
    SS.reset_sim(genNewPlanets=genNewPlanets, rewindPlanets=rewindPlanets)
    
    pklname = 'run'+str(int(time.clock()*100))+''.join(["%s" % random.randint(0, 9) for num in range(5)]) + '.pkl'
    pklpath = os.path.join(outpath, pklname)
    with open(pklpath, 'wb') as f:
        pickle.dump({'DRM':DRM,'systems':systems,'seed':seed}, f)
        
    return 0 

def run_one(genNewPlanets=True, rewindPlanets=True, outpath='.'):
    # wrap the run_sim in a try/except loop
    # reset simulation at the end of each simulation
    #NOTE: Methods are imported from IPClusterEnsemble sync_imports function line
    SS.reset_sim(genNewPlanets=genNewPlanets, rewindPlanets=rewindPlanets)
    nbmax = 10
    for attempt in numpy.arange(nbmax):
        try:
            # run one survey simulation
            SS.run_sim()
            DRM = SS.DRM[:]
            systems = SS.SimulatedUniverse.dump_systems()
            systems['MsTrue'] = SS.TargetList.MsTrue
            systems['MsEst'] = SS.TargetList.MsEst
            seed = SS.seed
        except Exception as e:
            # if anything goes wrong, log the error and reset simulation
            with open(os.path.join(outpath,'log.err'), 'ab') as f:
                f.write(repr(e))
                f.write('\n')
                f.write(traceback.format_exc())
                f.write('\n\n')
            
            SS.reset_sim(genNewPlanets=genNewPlanets, rewindPlanets=rewindPlanets)
        else:
            break
    else:
        raise ValueError("Unsuccessful run_sim after %s reset_sim attempts"%nbmax)
    
    # reset simulation at the end of each simulation
    SS.reset_sim(genNewPlanets=genNewPlanets, rewindPlanets=rewindPlanets)  

    pklname = 'run'+str(int(time.clock()*100))+''.join(["%s" % random.randint(0, 9) for num in numpy.arange(5)]) + '.pkl'
    pklpath = os.path.join(outpath, pklname)
    with open(pklpath, 'wb') as f:
        pickle.dump({'DRM':DRM,'systems':systems,'seed':seed}, f)
        
    return 0 

def get_jk(self, mol=None, dm=None, hermi=1, with_j=True, with_k=True,
               omega=None):
        if mol is None: mol = self.mol
        if dm is None: dm = self.make_rdm1()
        t0 = (time.clock(), time.time())
        if self.direct_scf and self.opt is None:
            self.opt = self.init_direct_scf(mol)
        vj, vk = get_jk(mol, dm, hermi, self.opt, with_j, with_k)
        logger.timer(self, 'vj and vk', *t0)
        return vj, vk 

def get_jk(self, cell=None, dm_kpts=None, hermi=1, kpts=None, kpts_band=None,
               with_j=True, with_k=True, omega=None, **kwargs):
        if cell is None: cell = self.cell
        if kpts is None: kpts = self.kpts
        if dm_kpts is None: dm_kpts = self.make_rdm1()
        cpu0 = (time.clock(), time.time())
        vj, vk = self.with_df.get_jk(dm_kpts, hermi, kpts, kpts_band,
                                     with_j, with_k, omega, exxdiv=self.exxdiv)
        logger.timer(self, 'vj and vk', *cpu0)
        return vj, vk 

def _make_shared_1e(self):
        cput0 = (time.clock(), time.time())
        log = logger.Logger(self.stdout, self.verbose)

        t1, t2, eris = self.t1, self.t2, self.eris
        kconserv = self.kconserv
        self.Loo = imd.Loo(t1, t2, eris, kconserv)
        self.Lvv = imd.Lvv(t1, t2, eris, kconserv)
        self.Fov = imd.cc_Fov(t1, t2, eris, kconserv)

        log.timer('EOM-CCSD shared one-electron intermediates', *cput0) 

def make_ip(self, ip_partition=None):
        self._make_shared_1e()
        if self._made_shared_2e is False and ip_partition != 'mp':
            self._make_shared_2e()
            self._made_shared_2e = True

        cput0 = (time.clock(), time.time())
        log = logger.Logger(self.stdout, self.verbose)

        t1, t2, eris = self.t1, self.t2, self.eris
        kconserv = self.kconserv

        if self._fimd is not None:
            nkpts, nocc, nvir = t1.shape
            oooo_dest = self._fimd.create_dataset('oooo', (nkpts, nkpts, nkpts, nocc, nocc, nocc, nocc), t1.dtype.char)
            ooov_dest = self._fimd.create_dataset('ooov', (nkpts, nkpts, nkpts, nocc, nocc, nocc, nvir), t1.dtype.char)
            ovoo_dest = self._fimd.create_dataset('ovoo', (nkpts, nkpts, nkpts, nocc, nvir, nocc, nocc), t1.dtype.char)
        else:
            oooo_dest = ooov_dest = ovoo_dest = None

        # 0 or 1 virtuals
        if ip_partition != 'mp':
            self.Woooo = imd.Woooo(t1, t2, eris, kconserv, oooo_dest)
        self.Wooov = imd.Wooov(t1, t2, eris, kconserv, ooov_dest)
        self.Wovoo = imd.Wovoo(t1, t2, eris, kconserv, ovoo_dest)
        self.made_ip_imds = True
        log.timer('EOM-CCSD IP intermediates', *cput0) 

def make_ea(self, ea_partition=None):
        self._make_shared_1e()
        if self._made_shared_2e is False and ea_partition != 'mp':
            self._make_shared_2e()
            self._made_shared_2e = True

        cput0 = (time.clock(), time.time())
        log = logger.Logger(self.stdout, self.verbose)

        t1, t2, eris = self.t1, self.t2, self.eris
        kconserv = self.kconserv

        if self._fimd is not None:
            nkpts, nocc, nvir = t1.shape
            vovv_dest = self._fimd.create_dataset('vovv', (nkpts, nkpts, nkpts, nvir, nocc, nvir, nvir), t1.dtype.char)
            vvvo_dest = self._fimd.create_dataset('vvvo', (nkpts, nkpts, nkpts, nvir, nvir, nvir, nocc), t1.dtype.char)
            vvvv_dest = self._fimd.create_dataset('vvvv', (nkpts, nkpts, nkpts, nvir, nvir, nvir, nvir), t1.dtype.char)
        else:
            vovv_dest = vvvo_dest = vvvv_dest = None

        # 3 or 4 virtuals
        self.Wvovv = imd.Wvovv(t1, t2, eris, kconserv, vovv_dest)
        if ea_partition == 'mp' and np.all(t1 == 0):
            self.Wvvvo = imd.Wvvvo(t1, t2, eris, kconserv, vvvo_dest)
        else:
            self.Wvvvv = imd.Wvvvv(t1, t2, eris, kconserv, vvvv_dest)
            self.Wvvvo = imd.Wvvvo(t1, t2, eris, kconserv, self.Wvvvv, vvvo_dest)
        self.made_ea_imds = True
        log.timer('EOM-CCSD EA intermediates', *cput0) 

def amplitudes_to_vector_singlet(r1, r2, kconserv):
    '''Transform 3- and 7-dimensional arrays, r1 and r2, to a 1-dimensional
    array with unique indices.

    For example:
        r1: t_{i k_i}^{a k_a}
        r2: t_{i k_i, j k_j}^{a k_a, b k_b}
        return: a vector with all r1 elements, and r2 elements whose indices
    satisfy (i k_i a k_a) >= (j k_j b k_b)
    '''
    cput0 = (time.clock(), time.time())
    log = logger.Logger(sys.stdout, logger.DEBUG)
    # r1 indices: k_i, i, a
    nkpts, nocc, nvir = np.asarray(r1.shape)[[0, 1, 2]]
    nov = nocc * nvir

    # r2 indices (old): k_i, k_J, k_a, i, J, a, B
    # r2 indices (new): (k_i, k_a), (k_J), (i, a), (J, B)
    r2 = r2.transpose(0,2,1,3,5,4,6).reshape(nkpts**2, nkpts, nov, nov)

    idx, idy = np.tril_indices(nov)
    nov2_tril = nov * (nov + 1) // 2
    nov2 = nov * nov

    vector = np.empty(r2.size, dtype=r2.dtype)
    offset = 0
    for ki, ka, kj in kpts_helper.loop_kkk(nkpts):
        kb = kconserv[ki, ka, kj]
        kika = ki * nkpts + ka
        kjkb = kj * nkpts + kb
        r2ovov = r2[kika, kj]
        if kika == kjkb:
            vector[offset:offset+nov2_tril] = r2ovov[idx, idy]
            offset += nov2_tril
        elif kika > kjkb:
            vector[offset:offset+nov2] = r2ovov.ravel()
            offset += nov2

    vector = np.hstack((r1.ravel(), vector[:offset]))
    log.timer("amplitudes_to_vector_singlet", *cput0)
    return vector 

def eeccsd_cis_approx_slow(eom, kshift, nroots=1, imds=None, **kwargs):
    '''Build initial R vector through diagonalization of <r1|Hbar|r1>

    This method evaluates the matrix elements of Hbar in r1 space in the following way:
    - 1st col of Hbar = matvec(r1_col1) where r1_col1 = [1, 0, 0, 0, ...]
    - 2nd col of Hbar = matvec(r1_col2) where r1_col2 = [0, 1, 0, 0, ...]
    - and so on

    Note that such evaluation has N^3 cost, but error free (because matvec() has been proven correct).
    '''
    cput0 = (time.clock(), time.time())
    log = logger.Logger(eom.stdout, eom.verbose)

    if imds is None: imds = eom.make_imds()
    nkpts, nocc, nvir = imds.t1.shape
    dtype = imds.t1.dtype
    r1_size = nkpts * nocc * nvir

    H1 = np.zeros([r1_size, r1_size], dtype=dtype)
    for col in range(r1_size):
        vec = np.zeros(r1_size, dtype=dtype)
        vec[col] = 1.0
        H1[:, col] = eeccsd_matvec_singlet_Hr1(eom, vec, kshift, imds=imds)

    eigval, eigvec = np.linalg.eig(H1)
    idx = eigval.argsort()[:nroots]
    eigval = eigval[idx]
    eigvec = eigvec[:, idx]

    log.timer("EOMEE CIS approx", *cput0)

    return eigval, eigvec 

def make_ip(self, ip_partition=None):
        self._make_shared_1e()
        if self._made_shared_2e is False and ip_partition != 'mp':
            self._make_shared_2e()
            self._made_shared_2e = True

        cput0 = (time.clock(), time.time())
        log = logger.Logger(self.stdout, self.verbose)

        t1, t2, eris = self.t1, self.t2, self.eris
        kconserv = self.kconserv

        if self._fimd is not None:
            nkpts, nocc, nvir = t1.shape
            oooo_dest = self._fimd.create_dataset('oooo', (nkpts, nkpts, nkpts, nocc, nocc, nocc, nocc), t1.dtype.char)
            ooov_dest = self._fimd.create_dataset('ooov', (nkpts, nkpts, nkpts, nocc, nocc, nocc, nvir), t1.dtype.char)
            ovoo_dest = self._fimd.create_dataset('ovoo', (nkpts, nkpts, nkpts, nocc, nvir, nocc, nocc), t1.dtype.char)
        else:
            oooo_dest = ooov_dest = ovoo_dest = None

        # 0 or 1 virtuals
        if ip_partition != 'mp':
            self.Woooo = imd.Woooo(t1, t2, eris, kconserv, oooo_dest)
        self.Wooov = imd.Wooov(t1, t2, eris, kconserv, ooov_dest)
        self.Wovoo = imd.Wovoo(t1, t2, eris, kconserv, ovoo_dest)
        self.made_ip_imds = True
        log.timer('EOM-CCSD IP intermediates', *cput0) 

def make_ea(self, ea_partition=None):
        self._make_shared_1e()
        if self._made_shared_2e is False and ea_partition != 'mp':
            self._make_shared_2e()
            self._made_shared_2e = True

        cput0 = (time.clock(), time.time())
        log = logger.Logger(self.stdout, self.verbose)

        t1, t2, eris = self.t1, self.t2, self.eris
        kconserv = self.kconserv

        if self._fimd is not None:
            nkpts, nocc, nvir = t1.shape
            vovv_dest = self._fimd.create_dataset('vovv', (nkpts, nkpts, nkpts, nvir, nocc, nvir, nvir), t1.dtype.char)
            vvvo_dest = self._fimd.create_dataset('vvvo', (nkpts, nkpts, nkpts, nvir, nvir, nvir, nocc), t1.dtype.char)
            if eris.vvvv is not None:
                vvvv_dest = self._fimd.create_dataset('vvvv', (nkpts, nkpts, nkpts, nvir, nvir, nvir, nvir), t1.dtype.char)
        else:
            vovv_dest = vvvo_dest = vvvv_dest = None

        # 3 or 4 virtuals
        self.Wvovv = imd.Wvovv(t1, t2, eris, kconserv, vovv_dest)
        if ea_partition == 'mp' and np.all(t1 == 0):
            self.Wvvvo = imd.Wvvvo(t1, t2, eris, kconserv, vvvo_dest)
        else:
            if eris.vvvv is None:
                self.Wvvvv = None
            else:
                self.Wvvvv = imd.Wvvvv(t1, t2, eris, kconserv, vvvv_dest)
            self.Wvvvo = imd.Wvvvo(t1, t2, eris, kconserv, self.Wvvvv, vvvo_dest)
        self.made_ea_imds = True
        log.timer('EOM-CCSD EA intermediates', *cput0) 

def make_ee(self, ee_partition=None):
        self._make_shared_1e()
        if self._made_shared_2e is False:
            self._make_shared_2e()
            self._made_shared_2e = True

        cput0 = (time.clock(), time.time())
        log = logger.Logger(self.stdout, self.verbose)

        t1, t2, eris = self.t1, self.t2, self.eris
        kconserv = self.kconserv

        # Rename imds to match the notations in pyscf.cc.eom_rccsd
        self.Foo = self.Loo
        self.Fvv = self.Lvv
        self.woOvV = self.Woovv
        self.woVvO = self.Wovvo
        self.woVoV = self.Wovov

        if not self.made_ip_imds:
            # 0 or 1 virtuals
            self.woOoO = imd.Woooo(t1, t2, eris, kconserv)
            self.woOoV = imd.Wooov(t1, t2, eris, kconserv)
            self.woVoO = imd.Wovoo(t1, t2, eris, kconserv)
        else:
            self.woOoO = self.Woooo
            self.woOoV = self.Wooov
            self.woVoO = self.Wovoo

        if not self.made_ea_imds:
            # 3 or 4 virtuals
            self.wvOvV = imd.Wvovv(t1, t2, eris, kconserv)
            self.wvVvV = imd.Wvvvv(t1, t2, eris, kconserv)
            self.wvVvO = imd.Wvvvo(t1, t2, eris, kconserv, self.wvVvV)
        else:
            self.wvOvV = self.Wvovv
            self.wvVvV = self.Wvvvv
            self.wvVvO = self.Wvvvo

        self.made_ee_imds = True
        log.timer('EOM-CCSD EE intermediates', *cput0) 

def init_amps(self, eris):
        time0 = time.clock(), time.time()
        nocc = self.nocc
        nvir = self.nmo - nocc
        nkpts = self.nkpts
        mo_e_o = [eris.mo_energy[k][:nocc] for k in range(nkpts)]
        mo_e_v = [eris.mo_energy[k][nocc:] for k in range(nkpts)]
        t1 = numpy.zeros((nkpts, nocc, nvir), dtype=numpy.complex128)
        t2 = numpy.zeros((nkpts, nkpts, nkpts, nocc, nocc, nvir, nvir), dtype=numpy.complex128)
        self.emp2 = 0
        eris_oovv = eris.oovv.copy()

        # Get location of padded elements in occupied and virtual space
        nonzero_opadding, nonzero_vpadding = padding_k_idx(self, kind="split")

        kconserv = kpts_helper.get_kconserv(self._scf.cell, self.kpts)
        for ki, kj, ka in kpts_helper.loop_kkk(nkpts):
            kb = kconserv[ki, ka, kj]
            # For LARGE_DENOM, see t1new update above
            eia = LARGE_DENOM * numpy.ones((nocc, nvir), dtype=eris.mo_energy[0].dtype)
            n0_ovp_ia = numpy.ix_(nonzero_opadding[ki], nonzero_vpadding[ka])
            eia[n0_ovp_ia] = (mo_e_o[ki][:,None] - mo_e_v[ka])[n0_ovp_ia]

            ejb = LARGE_DENOM * numpy.ones((nocc, nvir), dtype=eris.mo_energy[0].dtype)
            n0_ovp_jb = numpy.ix_(nonzero_opadding[kj], nonzero_vpadding[kb])
            ejb[n0_ovp_jb] = (mo_e_o[kj][:,None] - mo_e_v[kb])[n0_ovp_jb]
            eijab = eia[:, None, :, None] + ejb[:, None, :]

            t2[ki, kj, ka] = eris_oovv[ki, kj, ka] / eijab

        t2 = numpy.conj(t2)
        self.emp2 = 0.25 * numpy.einsum('pqrijab,pqrijab', t2, eris_oovv).real
        self.emp2 /= nkpts

        logger.info(self, 'Init t2, MP2 energy = %.15g', self.emp2.real)
        logger.timer(self, 'init mp2', *time0)
        return self.emp2, t1, t2 

def _make_shared_1e(self):
        cput0 = (time.clock(), time.time())
        log = logger.Logger(self.stdout, self.verbose)

        t1,t2,eris = self.t1, self.t2, self.eris
        kconserv = self.kconserv
        self.Loo = imd.Loo(t1,t2,eris,kconserv)
        self.Lvv = imd.Lvv(t1,t2,eris,kconserv)
        self.Fov = imd.cc_Fov(t1,t2,eris,kconserv)

        log.timer('EOM-CCSD shared one-electron intermediates', *cput0) 

def ipccsd(self, nroots=2*4, kptlist=None):
        time0 = time.clock(), time.time()
        log = logger.Logger(self.stdout, self.verbose)
        nocc = self.nocc
        nvir = self.nmo - nocc
        nkpts = self.nkpts
        if kptlist is None:
            kptlist = range(nkpts)
        size = self.vector_size_ip()
        for k,kshift in enumerate(kptlist):
            self.kshift = kshift
            nfrozen = np.sum(self.mask_frozen_ip(np.zeros(size, dtype=int), kshift, const=1))
            nroots = min(nroots, size - nfrozen)
        evals = np.zeros((len(kptlist),nroots), np.float)
        evecs = np.zeros((len(kptlist),nroots,size), np.complex)

        for k,kshift in enumerate(kptlist):
            self.kshift = kshift
            diag = self.ipccsd_diag()
            diag = self.mask_frozen_ip(diag, kshift, const=LARGE_DENOM)
            precond = lambda dx, e, x0: dx/(diag-e)
            # Initial guess from file
            amplitude_filename = "__ripccsd" + str(kshift) + "__.hdf5"
            rsuccess, x0 = read_eom_amplitudes((nroots,size),filename=amplitude_filename)
            if x0 is not None:
                x0 = x0.T
            #if not rsuccess:
            #    x0 = np.zeros_like(diag)
            #    x0[np.argmin(diag)] = 1.0

            conv, evals_k, evecs_k = eigs(self.ipccsd_matvec, size, nroots, x0=x0, Adiag=diag, verbose=self.verbose)

            evals_k = evals_k.real
            evals[k] = evals_k
            evecs[k] = evecs_k.T

            write_eom_amplitudes(evecs[k],filename=amplitude_filename)
        time0 = log.timer_debug1('converge ip-ccsd', *time0)
        comm.Barrier()
        return evals.real, evecs 

def eaccsd(self, nroots=2*4, kptlist=None):
        time0 = time.clock(), time.time()
        log = logger.Logger(self.stdout, self.verbose)
        nocc = self.nocc
        nvir = self.nmo - nocc
        nkpts = self.nkpts
        if kptlist is None:
            kptlist = range(nkpts)
        size = self.vector_size_ea()
        for k,kshift in enumerate(kptlist):
            self.kshift = kshift
            nfrozen = np.sum(self.mask_frozen_ea(np.zeros(size, dtype=int), kshift, const=1))
            nroots = min(nroots, size - nfrozen)
        evals = np.zeros((len(kptlist),nroots), np.float)
        evecs = np.zeros((len(kptlist),nroots,size), np.complex)

        for k,kshift in enumerate(kptlist):
            self.kshift = kshift
            diag = self.eaccsd_diag()
            diag = self.mask_frozen_ea(diag, kshift, const=LARGE_DENOM)
            precond = lambda dx, e, x0: dx/(diag-e)
            # Initial guess from file
            amplitude_filename = "__reaccsd" + str(kshift) + "__.hdf5"
            rsuccess, x0 = read_eom_amplitudes((nroots,size),filename=amplitude_filename)
            if x0 is not None:
                x0 = x0.T
            #if not rsuccess:
            #    x0 = np.zeros_like(diag)
            #    x0[np.argmin(diag)] = 1.0

            conv, evals_k, evecs_k = eigs(self.eaccsd_matvec, size, nroots, x0=x0, Adiag=diag, verbose=self.verbose)

            evals_k = evals_k.real
            evals[k] = evals_k
            evecs[k] = evecs_k.T

            write_eom_amplitudes(evecs[k],filename=amplitude_filename)
        time0 = log.timer_debug1('converge ea-ccsd', *time0)
        comm.Barrier()
        return evals.real, evecs 

def leaccsd(self, nroots=2*4, kptlist=None):
        time0 = time.clock(), time.time()
        log = logger.Logger(self.stdout, self.verbose)
        nocc = self.nocc
        nvir = self.nmo - nocc
        nkpts = self.nkpts
        if kptlist is None:
            kptlist = range(nkpts)
        size = self.vector_size_ea()
        for k,kshift in enumerate(kptlist):
            self.kshift = kshift
            nfrozen = np.sum(self.mask_frozen_ea(np.zeros(size, dtype=int), kshift, const=1))
            nroots = min(nroots, size - nfrozen)
        evals = np.zeros((len(kptlist),nroots), np.float)
        evecs = np.zeros((len(kptlist),nroots,size), np.complex)

        for k,kshift in enumerate(kptlist):
            self.kshift = kshift
            diag = self.eaccsd_diag()
            precond = lambda dx, e, x0: dx/(diag-e)
            # Initial guess from file
            amplitude_filename = "__leaccsd" + str(kshift) + "__.hdf5"
            rsuccess, x0 = read_eom_amplitudes((nroots,size),filename=amplitude_filename)
            if x0 is not None:
                x0 = x0.T
            #if not rsuccess:
            #    x0 = np.zeros_like(diag)
            #    x0[np.argmin(diag)] = 1.0

            conv, evals_k, evecs_k = eigs(self.leaccsd_matvec, size, nroots, x0=x0, Adiag=diag, verbose=self.verbose)

            evals_k = evals_k.real
            evals[k] = evals_k
            evecs[k] = evecs_k.T

            write_eom_amplitudes(evecs[k],amplitude_filename)
        time0 = log.timer_debug1('converge lea-ccsd', *time0)
        comm.Barrier()
        return evals.real, evecs 

def get_veff(ks_grad, dm=None, kpts=None):
    mf = ks_grad.base
    cell = ks_grad.cell
    if dm is None: dm = mf.make_rdm1()
    if kpts is None: kpts = mf.kpts
    t0 = (time.clock(), time.time())

    ni = mf._numint
    if ks_grad.grids is not None:
        grids = ks_grad.grids
    else:
        grids = mf.grids
    if grids.coords is None:
        grids.build(with_non0tab=True)

    omega, alpha, hyb = ni.rsh_and_hybrid_coeff(mf.xc, spin=cell.spin)

    mem_now = lib.current_memory()[0]
    max_memory = max(2000, ks_grad.max_memory*.9-mem_now)
    if ks_grad.grid_response:
        raise NotImplementedError
    else:
        vxc =  get_vxc(ni, cell, grids, mf.xc, dm, kpts,
                           max_memory=max_memory, verbose=ks_grad.verbose)
    t0 = logger.timer(ks_grad, 'vxc', *t0)

    if abs(hyb) < 1e-10 and abs(alpha) < 1e-10:
        vj = ks_grad.get_j(dm, kpts)
        vxc += vj[:,0][:,None] + vj[:,1][:,None]
    else:
        vj, vk = ks_grad.get_jk(dm, kpts)
        vk *= hyb
        if abs(omega) > 1e-10:  # For range separated Coulomb operator
            with cell.with_range_coulomb(omega):
                vk += ks_grad.get_k(dm, kpts) * (alpha - hyb)
        vxc += vj[:,0][:,None] + vj[:,1][:,None] - vk

    return vxc 

def grad_elec(mf_grad, mo_energy=None, mo_coeff=None, mo_occ=None, atmlst=None):
    '''
    Electronic part of KRHF/KRKS gradients
    Args:
        mf_grad : pbc.grad.krhf.Gradients or pbc.grad.krks.Gradients object
    '''
    mf = mf_grad.base
    cell = mf_grad.cell
    kpts = mf.kpts
    nkpts = len(kpts)
    if mo_energy is None: mo_energy = mf.mo_energy
    if mo_occ is None:    mo_occ = mf.mo_occ
    if mo_coeff is None:  mo_coeff = mf.mo_coeff
    if atmlst is None: atmlst = range(cell.natm)

    log = logger.Logger(mf_grad.stdout, mf_grad.verbose)
    hcore_deriv = mf_grad.hcore_generator(cell, kpts)
    s1 = mf_grad.get_ovlp(cell, kpts)
    dm0 = mf.make_rdm1(mo_coeff, mo_occ)
    t0 = (time.clock(), time.time())
    log.debug('Computing Gradients of NR-HF Coulomb repulsion')
    vhf = mf_grad.get_veff(dm0, kpts)
    log.timer('gradients of 2e part', *t0)
    dme0 = mf_grad.make_rdm1e(mo_energy, mo_coeff, mo_occ)
    aoslices = cell.aoslice_by_atom()
    de = np.zeros([len(atmlst),3])
    for x, ia in enumerate(atmlst):
        p0, p1 = aoslices[ia,2:]
        h1ao = hcore_deriv(ia)
        de[x] += np.einsum('xkij,kji->x', h1ao, dm0).real
        # nabla was applied on bra in vhf, *2 for the contributions of nabla|ket>
        de[x] += np.einsum('xkij,kji->x', vhf[:,:,p0:p1], dm0[:,:,p0:p1]).real * 2
        de[x] -= np.einsum('kxij,kji->x', s1[:,:,p0:p1], dme0[:,:,p0:p1]).real * 2
        de[x] /= nkpts
        de[x] += mf_grad.extra_force(ia, locals())

    if log.verbose > logger.DEBUG:
        log.debug('gradients of electronic part')
        mf_grad._write(log, cell, de, atmlst)
    return de 

def get_jk(self, dm=None, kpts=None):
        if kpts is None: kpts = self.kpts
        if dm is None: dm = self.base.make_rdm1()
        exxdiv = self.base.exxdiv
        cpu0 = (time.clock(), time.time())
        vj, vk = self.base.with_df.get_jk_e1(dm, kpts, exxdiv=exxdiv)
        logger.timer(self, 'vj and vk', *cpu0)
        return vj, vk 

def get_j(self, dm=None, kpts=None):
        if kpts is None: kpts = self.kpts
        if dm is None: dm = self.base.make_rdm1()
        cpu0 = (time.clock(), time.time())
        vj = self.base.with_df.get_j_e1(dm, kpts)
        logger.timer(self, 'vj', *cpu0)
        return vj 

def get_k(self, dm=None, kpts=None, kpts_band=None):
        if kpts is None: kpts = self.kpts
        if dm is None: dm = self.base.make_rdm1()
        exxdiv = self.base.exxdiv
        cpu0 = (time.clock(), time.time())
        vk = self.base.with_df.get_k_e1(dm, kpts, kpts_band, exxdiv)
        logger.timer(self, 'vk', *cpu0)
        return vk 

def scf(self, dm0=None, **kwargs):
        '''SCF main driver

        Kwargs:
            dm0 : ndarray
                If given, it will be used as the initial guess density matrix

        Examples:

        >>> import numpy
        >>> from pyscf import gto, scf
        >>> mol = gto.M(atom='H 0 0 0; F 0 0 1.1')
        >>> mf = scf.hf.SCF(mol)
        >>> dm_guess = numpy.eye(mol.nao_nr())
        >>> mf.kernel(dm_guess)
        converged SCF energy = -98.5521904482821
        -98.552190448282104
        '''
        cput0 = (time.clock(), time.time())

        self.dump_flags()
        self.build(self.mol)

        if self.max_cycle > 0 or self.mo_coeff is None:
            self.converged, self.e_tot, \
                    self.mo_energy, self.mo_coeff, self.mo_occ = \
                    kernel(self, self.conv_tol, self.conv_tol_grad,
                           dm0=dm0, callback=self.callback,
                           conv_check=self.conv_check, **kwargs)
        else:
            # Avoid to update SCF orbitals in the non-SCF initialization
            # (issue #495).  But run regular SCF for initial guess if SCF was
            # not initialized.
            self.e_tot = kernel(self, self.conv_tol, self.conv_tol_grad,
                                dm0=dm0, callback=self.callback,
                                conv_check=self.conv_check, **kwargs)[1]

        logger.timer(self, 'SCF', *cput0)
        self._finalize()
        return self.e_tot 

def get_jk(self, mol=None, dm=None, hermi=1, with_j=True, with_k=True,
               omega=None):
        if mol is None: mol = self.mol
        if dm is None: dm = self.make_rdm1()
        t0 = (time.clock(), time.time())
        log = logger.new_logger(self)
        if self.direct_scf and self.opt is None:
            self.opt = self.init_direct_scf(mol)
        opt_llll, opt_ssll, opt_ssss, opt_gaunt = self.opt

        vj, vk = get_jk_coulomb(mol, dm, hermi, self._coulomb_now,
                                opt_llll, opt_ssll, opt_ssss, omega, log)

        if self.with_breit:
            if 'SSSS' in self._coulomb_now.upper():
                vj1, vk1 = _call_veff_gaunt_breit(mol, dm, hermi, opt_gaunt, True)
                log.debug('Add Breit term')
                vj += vj1
                vk += vk1
        elif self.with_gaunt and 'SS' in self._coulomb_now.upper():
            log.debug('Add Gaunt term')
            vj1, vk1 = _call_veff_gaunt_breit(mol, dm, hermi, opt_gaunt, False)
            vj += vj1
            vk += vk1

        log.timer('vj and vk', *t0)
        return vj, vk 

def scf(self, dm0=None):
        cput0 = (time.clock(), time.time())

        self.build()
        self.dump_flags()
        self.converged, self.e_tot, \
                self.mo_energy, self.mo_coeff, self.mo_occ \
                = kernel(self, self.conv_tol, self.conv_tol_grad,
                         dm0=dm0, callback=self.callback,
                         conv_check=self.conv_check)

        logger.timer(self, 'SCF', *cput0)
        self._finalize()
        return self.e_tot 

def _iterative_kernel(mp, eris, verbose=None):
    cput1 = cput0 = (time.clock(), time.time())
    log = logger.new_logger(mp, verbose)

    emp2, t2 = mp.init_amps(eris=eris)
    log.info('Init E(MP2) = %.15g', emp2)

    adiis = lib.diis.DIIS(mp)

    conv = False
    for istep in range(mp.max_cycle):
        t2new = mp.update_amps(t2, eris)

        if isinstance(t2new, numpy.ndarray):
            normt = numpy.linalg.norm(t2new - t2)
            t2 = None
            t2new = adiis.update(t2new)
        else: # UMP2
            normt = numpy.linalg.norm([numpy.linalg.norm(t2new[i] - t2[i])
                                       for i in range(3)])
            t2 = None
            t2shape = [x.shape for x in t2new]
            t2new = numpy.hstack([x.ravel() for x in t2new])
            t2new = adiis.update(t2new)
            t2new = lib.split_reshape(t2new, t2shape)

        t2, t2new = t2new, None
        emp2, e_last = mp.energy(t2, eris), emp2
        log.info('cycle = %d  E_corr(MP2) = %.15g  dE = %.9g  norm(t2) = %.6g',
                 istep+1, emp2, emp2 - e_last, normt)
        cput1 = log.timer('MP2 iter', *cput1)
        if abs(emp2-e_last) < mp.conv_tol and normt < mp.conv_tol_normt:
            conv = True
            break
    log.timer('MP2', *cput0)
    return conv, emp2, t2 

def kernel(self, mo_coeff=None, mo_occ=None, dm0=None):
        cput0 = (time.clock(), time.time())
        if dm0 is not None:
            if isinstance(dm0, str):
                sys.stderr.write('Newton solver reads density matrix from chkfile %s\n' % dm0)
                dm0 = self.from_chk(dm0)

        elif mo_coeff is not None and mo_occ is None:
            logger.warn(self, 'Newton solver expects mo_coeff with '
                        'mo_occ as initial guess but mo_occ is not found in '
                        'the arguments.\n      The given '
                        'argument is treated as density matrix.')
            dm0 = mo_coeff
            mo_coeff = mo_occ = None

        else:
            if mo_coeff is None: mo_coeff = self.mo_coeff
            if mo_occ is None: mo_occ = self.mo_occ

            # TODO: assert mo_coeff orth-normality. If not orth-normal,
            # build dm from mo_coeff and mo_occ then unset mo_coeff and mo_occ.

        self.build(self.mol)
        self.dump_flags()

        self.converged, self.e_tot, \
                self.mo_energy, self.mo_coeff, self.mo_occ = \
                kernel(self, mo_coeff, mo_occ, dm0, conv_tol=self.conv_tol,
                       conv_tol_grad=self.conv_tol_grad,
                       max_cycle=self.max_cycle,
                       callback=self.callback, verbose=self.verbose)

        logger.timer(self, 'Second order SCF', *cput0)
        self._finalize()
        return self.e_tot 

def init_amps(self, eris=None):
        time0 = time.clock(), time.time()
        if eris is None:
            eris = self.ao2mo(self.mo_coeff)
        nocca, noccb = self.nocc

        fova = eris.focka[:nocca,nocca:]
        fovb = eris.fockb[:noccb,noccb:]
        mo_ea_o = eris.mo_energy[0][:nocca]
        mo_ea_v = eris.mo_energy[0][nocca:]
        mo_eb_o = eris.mo_energy[1][:noccb]
        mo_eb_v = eris.mo_energy[1][noccb:]
        eia_a = lib.direct_sum('i-a->ia', mo_ea_o, mo_ea_v)
        eia_b = lib.direct_sum('i-a->ia', mo_eb_o, mo_eb_v)

        t1a = fova.conj() / eia_a
        t1b = fovb.conj() / eia_b

        eris_ovov = np.asarray(eris.ovov)
        eris_OVOV = np.asarray(eris.OVOV)
        eris_ovOV = np.asarray(eris.ovOV)
        t2aa = eris_ovov.transpose(0,2,1,3) / lib.direct_sum('ia+jb->ijab', eia_a, eia_a)
        t2ab = eris_ovOV.transpose(0,2,1,3) / lib.direct_sum('ia+jb->ijab', eia_a, eia_b)
        t2bb = eris_OVOV.transpose(0,2,1,3) / lib.direct_sum('ia+jb->ijab', eia_b, eia_b)
        t2aa = t2aa - t2aa.transpose(0,1,3,2)
        t2bb = t2bb - t2bb.transpose(0,1,3,2)
        e  =      np.einsum('iJaB,iaJB', t2ab, eris_ovOV)
        e += 0.25*np.einsum('ijab,iajb', t2aa, eris_ovov)
        e -= 0.25*np.einsum('ijab,ibja', t2aa, eris_ovov)
        e += 0.25*np.einsum('ijab,iajb', t2bb, eris_OVOV)
        e -= 0.25*np.einsum('ijab,ibja', t2bb, eris_OVOV)
        self.emp2 = e.real
        logger.info(self, 'Init t2, MP2 energy = %.15g', self.emp2)
        logger.timer(self, 'init mp2', *time0)
        return self.emp2, (t1a,t1b), (t2aa,t2ab,t2bb) 

def init_amps(self, eris):
        time0 = time.clock(), time.time()
        mo_e = eris.fock.diagonal()
        nocc = self.nocc
        eia = mo_e[:nocc,None] - mo_e[None,nocc:]
        eijab = lib.direct_sum('ia,jb->ijab',eia,eia)
        t1 = eris.fock[:nocc,nocc:] / eia
        eris_oovv = np.array(eris.oovv)
        t2 = eris_oovv/eijab
        self.emp2 = 0.25*einsum('ijab,ijab',t2,eris_oovv.conj()).real
        logger.info(self, 'Init t2, MP2 energy = %.15g', self.emp2)
        logger.timer(self, 'init mp2', *time0)
        return self.emp2, t1, t2 

def _make_shared(self):
        cput0 = (time.clock(), time.time())
        log = logger.Logger(self.cc.stdout, self.cc.verbose)

        t1,t2,eris = self.cc.t1, self.cc.t2, self.cc.eris
        self.Foo = imd.Foo(t1,t2,eris)
        self.Fvv = imd.Fvv(t1,t2,eris)
        self.Fov = imd.Fov(t1,t2,eris)

        # 2 virtuals
        self.Wovvo = imd.Wovvo(t1,t2,eris)
        self.Woovv = eris.oovv

        log.timer('EOM-CCSD shared intermediates', *cput0) 

def _make_shared_1e(self):
        cput0 = (time.clock(), time.time())
        log = logger.Logger(self.stdout, self.verbose)

        t1,t2,eris = self.t1, self.t2, self.eris
        self.Loo = imd.Loo(t1,t2,eris)
        self.Lvv = imd.Lvv(t1,t2,eris)
        self.Fov = imd.cc_Fov(t1,t2,eris)

        log.timer('EOM-CCSD shared one-electron intermediates', *cput0)