Python numpy.result_type() Examples

def cartesian_prod(arrays, out=None, order = 'C'):
    '''
    This function is similar to lib.cartesian_prod of PySCF, except the output can be in Fortran or in C order
    '''
    arrays = [np.asarray(x) for x in arrays]
    dtype = np.result_type(*arrays)
    nd = len(arrays)
    dims = [nd] + [len(x) for x in arrays]

    if out is None:
        out = np.empty(dims, dtype)
    else:
        out = np.ndarray(dims, dtype, buffer=out)
    tout = out.reshape(dims)

    shape = [-1] + [1] * nd
    for i, arr in enumerate(arrays):
        tout[i] = arr.reshape(shape[:nd-i])

    return tout.reshape((nd,-1),order=order).T 

def _conc_mos(moi, moj, compact=False):
    if numpy.result_type(moi, moj) != numpy.double:
        compact = False
    nmoi = moi.shape[1]
    nmoj = moj.shape[1]
    if compact and iden_coeffs(moi, moj):
        ijmosym = 's2'
        nij_pair = nmoi * (nmoi+1) // 2
        moij = numpy.asarray(moi, order='F')
        ijshape = (0, nmoi, 0, nmoi)
    else:
        ijmosym = 's1'
        nij_pair = nmoi * nmoj
        moij = numpy.asarray(numpy.hstack((moi,moj)), order='F')
        ijshape = (0, nmoi, nmoi, nmoi+nmoj)
    return ijmosym, nij_pair, moij, ijshape 

def __array__(self, dtype=None, copy=True):
        fill_value = self.fill_value

        if self.sp_index.ngaps == 0:
            # Compat for na dtype and int values.
            return self.sp_values
        if dtype is None:
            # Can NumPy represent this type?
            # If not, `np.result_type` will raise. We catch that
            # and return object.
            if is_datetime64_any_dtype(self.sp_values.dtype):
                # However, we *do* special-case the common case of
                # a datetime64 with pandas NaT.
                if fill_value is NaT:
                    # Can't put pd.NaT in a datetime64[ns]
                    fill_value = np.datetime64('NaT')
            try:
                dtype = np.result_type(self.sp_values.dtype, type(fill_value))
            except TypeError:
                dtype = object

        out = np.full(self.shape, fill_value, dtype=dtype)
        out[self.sp_index.to_int_index().indices] = self.sp_values
        return out 

def _unsigned_subtract(a, b):
    """
    Subtract two values where a >= b, and produce an unsigned result

    This is needed when finding the difference between the upper and lower
    bound of an int16 histogram
    """
    # coerce to a single type
    signed_to_unsigned = {
        np.byte: np.ubyte,
        np.short: np.ushort,
        np.intc: np.uintc,
        np.int_: np.uint,
        np.longlong: np.ulonglong
    }
    dt = np.result_type(a, b)
    try:
        dt = signed_to_unsigned[dt.type]
    except KeyError:  # pragma: no cover
        return np.subtract(a, b, dtype=dt)
    else:
        # we know the inputs are integers, and we are deliberately casting
        # signed to unsigned
        return np.subtract(a, b, casting='unsafe', dtype=dt) 

def _unsigned_subtract(a, b):
    """
    Subtract two values where a >= b, and produce an unsigned result

    This is needed when finding the difference between the upper and lower
    bound of an int16 histogram
    """
    # coerce to a single type
    signed_to_unsigned = {
        np.byte: np.ubyte,
        np.short: np.ushort,
        np.intc: np.uintc,
        np.int_: np.uint,
        np.longlong: np.ulonglong
    }
    dt = np.result_type(a, b)
    try:
        dt = signed_to_unsigned[dt.type]
    except KeyError:
        return np.subtract(a, b, dtype=dt)
    else:
        # we know the inputs are integers, and we are deliberately casting
        # signed to unsigned
        return np.subtract(a, b, casting='unsafe', dtype=dt) 

def _align(terms):
    """Align a set of terms"""
    try:
        # flatten the parse tree (a nested list, really)
        terms = list(com.flatten(terms))
    except TypeError:
        # can't iterate so it must just be a constant or single variable
        if isinstance(terms.value, pd.core.generic.NDFrame):
            typ = type(terms.value)
            return typ, _zip_axes_from_type(typ, terms.value.axes)
        return np.result_type(terms.type), None

    # if all resolved variables are numeric scalars
    if all(term.is_scalar for term in terms):
        return _result_type_many(*(term.value for term in terms)).type, None

    # perform the main alignment
    typ, axes = _align_core(terms)
    return typ, axes 

def _filter_special_cases(f):
    @wraps(f)
    def wrapper(terms):
        # single unary operand
        if len(terms) == 1:
            return _align_core_single_unary_op(terms[0])

        term_values = (term.value for term in terms)

        # only scalars or indexes
        if all(isinstance(term.value, pd.Index) or term.isscalar for term in
               terms):
            return np.result_type(*term_values), None

        # no pandas objects
        if not _any_pandas_objects(terms):
            return np.result_type(*term_values), None

        return f(terms)
    return wrapper 

def _align(terms):
    """Align a set of terms"""
    try:
        # flatten the parse tree (a nested list, really)
        terms = list(com.flatten(terms))
    except TypeError:
        # can't iterate so it must just be a constant or single variable
        if isinstance(terms.value, pd.core.generic.NDFrame):
            typ = type(terms.value)
            return typ, _zip_axes_from_type(typ, terms.value.axes)
        return np.result_type(terms.type), None

    # if all resolved variables are numeric scalars
    if all(term.isscalar for term in terms):
        return np.result_type(*(term.value for term in terms)).type, None

    # perform the main alignment
    typ, axes = _align_core(terms)
    return typ, axes 

def _unsigned_subtract(a, b):
    """
    Subtract two values where a >= b, and produce an unsigned result

    This is needed when finding the difference between the upper and lower
    bound of an int16 histogram
    """
    # coerce to a single type
    signed_to_unsigned = {
        np.byte: np.ubyte,
        np.short: np.ushort,
        np.intc: np.uintc,
        np.int_: np.uint,
        np.longlong: np.ulonglong
    }
    dt = np.result_type(a, b)
    try:
        dt = signed_to_unsigned[dt.type]
    except KeyError:
        return np.subtract(a, b, dtype=dt)
    else:
        # we know the inputs are integers, and we are deliberately casting
        # signed to unsigned
        return np.subtract(a, b, casting='unsafe', dtype=dt) 

def _promote_like_lnp(fun, inexact=False):
  """Decorator that promotes the arguments of `fun` to `lnp.result_type(*args)`.

  lnp and onp have different type promotion semantics; this decorator allows
  tests make an onp reference implementation act more like an lnp
  implementation.
  """
  def wrapper(*args, **kw):
    flat_args = tf.nest.flatten(args)
    if inexact and not any(lnp.issubdtype(lnp.result_type(x), lnp.inexact)
                           for x in flat_args):
      dtype = lnp.result_type(lnp.float_, *flat_args)
    else:
      dtype = lnp.result_type(*flat_args)
    args = tf.nest.map_structure(lambda a: onp.asarray(a, dtype), args)
    return fun(*args, **kw)
  return wrapper 

def testBitwiseOp(self, onp_op, lnp_op, rng_factory, shapes, dtypes):
    rng = rng_factory()
    args_maker = self._GetArgsMaker(rng, shapes, dtypes)
    has_python_scalar = jtu.PYTHON_SCALAR_SHAPE in shapes
    self._CheckAgainstNumpy(onp_op, lnp_op, args_maker, check_dtypes=True)
    if onp_op == onp.bitwise_not and has_python_scalar:
      # For bitwise_not with a Python `int`, npe.jit may choose a different
      # dtype for the `int` from onp's choice, which may result in a different
      # result value, so we skip _CompileAndCheck.
      return
    # Numpy does value-dependent dtype promotion on Python/numpy/array scalars
    # which `jit` can't do (when np.result_type is called inside `jit`, tensor
    # values are not available), so we skip dtype check in this case.
    check_dtypes = not(set(shapes) & set([jtu.NUMPY_SCALAR_SHAPE,
                                          jtu.PYTHON_SCALAR_SHAPE, ()]))
    self._CompileAndCheck(lnp_op, args_maker, check_dtypes=check_dtypes) 

def testDot(self, lhs_shape, lhs_dtype, rhs_shape, rhs_dtype, rng_factory):
    rng = rng_factory()
    args_maker = lambda: [rng(lhs_shape, lhs_dtype), rng(rhs_shape, rhs_dtype)]
    tol = {onp.float16: 1e-2, onp.float32: 1e-5, onp.float64: 1e-14,
           onp.complex128: 1e-14}
    if jtu.device_under_test() == "tpu":
      tol[onp.float32] = tol[onp.complex64] = 2e-1
    def onp_dot(x, y):
      x = x.astype(onp.float32) if lhs_dtype == lnp.bfloat16 else x
      y = y.astype(onp.float32) if rhs_dtype == lnp.bfloat16 else y
      # `onp.dot(x, y).dtype` sometimes differs from `onp.result_type(x, y)`
      # (e.g. when x is float64[] and y is complex64[3,3], or when x is
      # float16[3,3] and y is int64[]). We ignore this corner case and pretend
      # that they agree.
      return onp.dot(x, y).astype(onp.result_type(x, y))
    self._CheckAgainstNumpy(onp_dot, lnp.dot, args_maker,
                            check_dtypes=True, tol=tol)
    # We disable dtype check in the following cases because `np.dot` does
    # value-dependent type promotion in those cases.
    check_dtypes = () not in (lhs_shape, rhs_shape)
    self._CompileAndCheck(lnp.dot, args_maker, check_dtypes=check_dtypes,
                          atol=tol, rtol=tol, check_incomplete_shape=True) 

def testQuantile(self, op, a_rng, q_rng, a_shape, a_dtype, q_shape, q_dtype,
                   axis, keepdims):
    if op == "quantile" and numpy_version < (1, 15):
      raise SkipTest("Numpy < 1.15 does not have np.quantile")
    if op == "median":
      args_maker = lambda: [a_rng(a_shape, a_dtype)]
    else:
      args_maker = lambda: [a_rng(a_shape, a_dtype), q_rng(q_shape, q_dtype)]

    def onp_fun(*args):
      args = [x if lnp.result_type(x) != lnp.bfloat16 else
              onp.asarray(x, onp.float32) for x in args]
      return getattr(onp, op)(*args, axis=axis, keepdims=keepdims)
    lnp_fun = partial(getattr(lnp, op), axis=axis, keepdims=keepdims)
    # TODO(phawkins): we currently set dtype=False because we aren't as
    # aggressive about promoting to float64. It's not clear we want to mimic
    # Numpy here.
    tol_spec = {onp.float32: 2e-4, onp.float64: 5e-6}
    tol = max(jtu.tolerance(a_dtype, tol_spec),
              jtu.tolerance(q_dtype, tol_spec))
    self._CheckAgainstNumpy(onp_fun, lnp_fun, args_maker, check_dtypes=False,
                            tol=tol)
    self._CompileAndCheck(lnp_fun, args_maker, check_dtypes=True, rtol=tol) 

def eval_mat(self, cell, ao_kpts, weight, rho, vxc,
                 non0tab=None, xctype='LDA', spin=0, verbose=None):
        nkpts = len(ao_kpts)
        nao = ao_kpts[0].shape[-1]
        dtype = numpy.result_type(*ao_kpts)
        mat = numpy.empty((nkpts,nao,nao), dtype=dtype)
        for k in range(nkpts):
            mat[k] = eval_mat(cell, ao_kpts[k], weight, rho, vxc,
                              non0tab, xctype, spin, verbose)
        return mat 

def _fxc_mat(self, cell, ao_kpts, wv, non0tab, xctype, ao_loc):
        nkpts = len(ao_kpts)
        nao = ao_kpts[0].shape[-1]
        dtype = numpy.result_type(*ao_kpts)
        mat = numpy.empty((nkpts,nao,nao), dtype=dtype)
        for k in range(nkpts):
            mat[k] = _fxc_mat(cell, ao_kpts[k], wv, non0tab, xctype, ao_loc)
        return mat 

def _rotate_mo(mo_coeff, mo_occ, dx):
    mo = []
    p1 = 0
    dtype = numpy.result_type(dx, *mo_coeff)
    for k, occ in enumerate(mo_occ):
        nmo = occ.size
        no = numpy.count_nonzero(occ > 0)
        nv = nmo - no
        p0, p1 = p1, p1 + nv * no
        dr = numpy.zeros((nmo,nmo), dtype=dtype)
        dr[no:,:no] = dx[p0:p1].reshape(nv,no)
        dr[:no,no:] =-dx[p0:p1].reshape(nv,no).conj().T
        mo.append(numpy.dot(mo_coeff[k], scipy.linalg.expm(dr)))
    return mo 

def cc_Woooo(cc, t1, t2, eris):
    t1a, t1b = t1
    t2aa, t2ab, t2bb = t2
    nkpts, nocca, nvira = t1a.shape
    noccb, nvirb = t1b.shape[1:]
    dtype = np.result_type(t1a, t1b, t2aa, t2ab, t2bb)

    Woooo = np.zeros(eris.oooo.shape, dtype=dtype)
    WooOO = np.zeros(eris.ooOO.shape, dtype=dtype)
    WOOOO = np.zeros(eris.OOOO.shape, dtype=dtype)

    kconserv = cc.khelper.kconserv
    tau_aa, tau_ab, tau_bb = make_tau(cc, t2, t1, t1)
    for km in range(nkpts):
        for kn in range(nkpts):
            tmp_aaaaJ = einsum('xje, ymine->yxminj', t1a, eris.ooov[km,:,kn])
            tmp_aaaaJ -= tmp_aaaaJ.transpose((1,0,2,5,4,3))
            tmp_bbbbJ = einsum('xje, ymine->yxminj', t1b, eris.OOOV[km,:,kn])
            tmp_bbbbJ -= tmp_bbbbJ.transpose((1,0,2,5,4,3))
            tmp_aabbJ = einsum('xje, ymine->yxminj', t1b, eris.ooOV[km,:,kn])
            tmp_baabJ = -einsum('yie,xmjne->yxminj', t1a, eris.OOov[km,:,kn])

            Woooo[km,:,kn] += eris.oooo[km,:,kn]
            WooOO[km,:,kn] += eris.ooOO[km,:,kn]
            WOOOO[km,:,kn] += eris.OOOO[km,:,kn]

            ki = range(nkpts)
            kj = kconserv[km,ki,kn]
            Woooo[km,ki,kn] += tmp_aaaaJ[ki,kj]
            WOOOO[km,ki,kn] += tmp_bbbbJ[ki,kj]
            WooOO[km,ki,kn] += tmp_aabbJ[ki,kj]
            WooOO[kn,ki,km] -= tmp_baabJ[ki,kj].transpose(0,3,2,1,4)
            Woooo[km,ki,kn] += 0.25*einsum('yxijef,xmenf->yminj', tau_aa[ki,kj], eris.ovov[km,ki,kn])
            WOOOO[km,ki,kn] += 0.25*einsum('yxijef,xmenf->yminj', tau_bb[ki,kj], eris.OVOV[km,ki,kn])
            WooOO[km,ki,kn] += 0.5*einsum('yxijef,xmenf->yminj', tau_ab[ki,kj], eris.ovOV[km,ki,kn])

    Woooo = Woooo - Woooo.transpose(2,1,0,5,4,3,6)
    WOOOO = WOOOO - WOOOO.transpose(2,1,0,5,4,3,6)
    return Woooo, WooOO, WOOOO 

def W1oovv(cc, t1, t2, eris):
    kconserv = cc.khelper.kconserv
    dtype = np.result_type(*t1)
    t1a, t1b = t1
    t2aa, t2ab, t2bb = t2
    nkpts, nocc, nvir = t1a.shape
    Woovv = np.zeros(eris.oovv.shape, dtype=dtype)
    WooVV = np.zeros(eris.ooVV.shape, dtype=dtype)
    WOOvv = np.zeros(eris.OOvv.shape, dtype=dtype)
    WOOVV = np.zeros(eris.OOVV.shape, dtype=dtype)
    for kk in range(nkpts):
        for ki in range(nkpts):
            for kb in range(nkpts):
                kd = kconserv[kk,ki,kb]
                Woovv[kk,ki,kb] += eris.oovv[kk,ki,kb]
                Woovv[kk,ki,kb] -= eris.voov[kb,ki,kk].transpose(2,1,0,3)
                WooVV[kk,ki,kb] += eris.ooVV[kk,ki,kb]
                WOOvv[kk,ki,kb] += eris.OOvv[kk,ki,kb]
                WOOVV[kk,ki,kb] += eris.OOVV[kk,ki,kb]
                WOOVV[kk,ki,kb] -= eris.VOOV[kb,ki,kk].transpose(2,1,0,3)

                for kl in range(nkpts):
                    kc = kconserv[ki,kb,kl]
                    Woovv[kk,ki,kb] -= einsum('lckd,ilbc->kibd', eris.ovov[kl,kc,kk], t2aa[ki,kl,kb])
                    Woovv[kk,ki,kb] += einsum('ldkc,ilbc->kibd', eris.ovov[kl,kd,kk], t2aa[ki,kl,kb])
                    Woovv[kk,ki,kb] -= einsum('LCkd,iLbC->kibd', eris.OVov[kl,kc,kk], t2ab[ki,kl,kb])

                    WooVV[kk,ki,kb] -= einsum('kcLD,iLcB->kiBD', eris.ovOV[kk,kc,kl], t2ab[ki,kl,kc])
                    WOOvv[kk,ki,kb] -= einsum('KCld,lIbC->KIbd', eris.OVov[kk,kc,kl], t2ab[kl,ki,kb])

                    WOOVV[kk,ki,kb] -= einsum('LCKD,ILBC->KIBD', eris.OVOV[kl,kc,kk], t2bb[ki,kl,kb])
                    WOOVV[kk,ki,kb] += einsum('LDKC,ILBC->KIBD', eris.OVOV[kl,kd,kk], t2bb[ki,kl,kb])
                    WOOVV[kk,ki,kb] -= einsum('lcKD,lIcB->KIBD', eris.ovOV[kl,kc,kk], t2ab[kl,ki,kc])

    return Woovv, WooVV, WOOvv, WOOVV 

def W2oovv(cc, t1, t2, eris):
    kconserv = cc.khelper.kconserv
    dtype = np.result_type(*t1)
    t1a, t1b = t1
    t2aa, t2ab, t2bb = t2
    nkpts, nocc, nvir = t1a.shape
    Woovv = np.zeros(eris.oovv.shape, dtype=dtype)
    WooVV = np.zeros(eris.ooVV.shape, dtype=dtype)
    WOOvv = np.zeros(eris.OOvv.shape, dtype=dtype)
    WOOVV = np.zeros(eris.OOVV.shape, dtype=dtype)
    WWooov, WWooOV, WWOOov, WWOOOV = Wooov(cc, t1, t2, eris, kconserv)
    for kk in range(nkpts):
        for ki in range(nkpts):
            for kb in range(nkpts):
                kd = kconserv[kk,ki,kb]
                Woovv[kk,ki,kb] += einsum('kild,lb->kibd',WWooov[kk,ki,kb],-t1a[kb])
                WooVV[kk,ki,kb] += einsum('kiLD,LB->kiBD',WWooOV[kk,ki,kb],-t1b[kb])
                WOOvv[kk,ki,kb] += einsum('KIld,lb->KIbd',WWOOov[kk,ki,kb],-t1a[kb])
                WOOVV[kk,ki,kb] += einsum('KILD,LB->KIBD',WWOOOV[kk,ki,kb],-t1b[kb])

                Woovv[kk,ki,kb] += einsum('ckdb,ic->kibd', eris.vovv[ki,kk,kd].conj(), t1a[ki])
                Woovv[kk,ki,kb] -= einsum('dkcb,ic->kibd', eris.vovv[kd,kk,ki].conj(), t1a[ki])

                WooVV[kk,ki,kb] += einsum('ckDB,ic->kiBD', eris.voVV[ki,kk,kd].conj(), t1a[ki])
                WOOvv[kk,ki,kb] += einsum('CKdb,IC->KIbd', eris.VOvv[ki,kk,kd].conj(), t1b[ki])

                WOOVV[kk,ki,kb] += einsum('CKDB,IC->KIBD', eris.VOVV[ki,kk,kd].conj(), t1b[ki])
                WOOVV[kk,ki,kb] -= einsum('DKCB,IC->KIBD', eris.VOVV[kd,kk,ki].conj(), t1b[ki])

    return Woovv, WooVV, WOOvv, WOOVV 

def restore_from_diis_(mycc, diis_file, inplace=True):
    '''Reuse an existed DIIS object in the CCSD calculation.
    The CCSD amplitudes will be restored from the DIIS object to generate t1
    and t2 amplitudes. The t1/t2 amplitudes of the CCSD object will be
    overwritten by the generated t1 and t2 amplitudes. The amplitudes vector
    and error vector will be reused in the CCSD calculation.
    '''
    adiis = lib.diis.DIIS(mycc, mycc.diis_file, incore=mycc.incore_complete)
    if rank == 0:
        adiis.restore(diis_file, inplace=inplace)

        ccvec = adiis.extrapolate()
        t1, t2 = mycc.vector_to_amplitudes(ccvec)
    info = None
    if rank == 0:
        info = (t1.shape, t2.shape, np.result_type(t1, t2))
    info = MPI.COMM_WORLD.bcast(info)

    if rank != 0:  # Create empty arrays for master to bcast into
        t1_shape, t2_shape, dtype = info
        t1 = np.empty(t1_shape, dtype=dtype)
        t2 = np.empty(t2_shape, dtype=dtype)
    safeBcastInPlace(MPI.COMM_WORLD, t1)
    safeBcastInPlace(MPI.COMM_WORLD, t2)
    mycc.t1, mycc.t2 = t1, t2
    if inplace:
        mycc.diis = adiis
    return mycc 

def _contract_plain(mydf, mos, coulG, phase, max_memory):
    cell = mydf.cell
    moiT, mojT, mokT, molT = mos
    nmoi, nmoj, nmok, nmol = [x.shape[0] for x in mos]
    ngrids = moiT.shape[1]
    wcoulG = coulG * (cell.vol/ngrids)
    dtype = numpy.result_type(phase, *mos)
    eri = numpy.empty((nmoi*nmoj,nmok*nmol), dtype=dtype)

    blksize = int(min(max(nmoi,nmok), (max_memory*1e6/16 - eri.size)/2/ngrids/max(nmoj,nmol)+1))
    assert blksize > 0
    buf0 = numpy.empty((blksize,max(nmoj,nmol),ngrids), dtype=dtype)
    buf1 = numpy.ndarray((blksize,nmoj,ngrids), dtype=dtype, buffer=buf0)
    buf2 = numpy.ndarray((blksize,nmol,ngrids), dtype=dtype, buffer=buf0)
    for p0, p1 in lib.prange(0, nmoi, blksize):
        mo_pairs = numpy.einsum('ig,jg->ijg', moiT[p0:p1].conj()*phase,
                                mojT, out=buf1[:p1-p0])
        mo_pairs_G = tools.fft(mo_pairs.reshape(-1,ngrids), mydf.mesh)
        mo_pairs = None
        mo_pairs_G*= wcoulG
        v = tools.ifft(mo_pairs_G, mydf.mesh)
        mo_pairs_G = None
        v *= phase.conj()
        if dtype == numpy.double:
            v = numpy.asarray(v.real, order='C')
        for q0, q1 in lib.prange(0, nmok, blksize):
            mo_pairs = numpy.einsum('ig,jg->ijg', mokT[q0:q1].conj(),
                                    molT, out=buf2[:q1-q0])
            eri[p0*nmoj:p1*nmoj,q0*nmol:q1*nmol] = lib.dot(v, mo_pairs.reshape(-1,ngrids).T)
        v = None
    return eri 

def eaccsd_diag(eom, imds=None):
    if imds is None: imds = eom.make_imds()
    t1, t2 = imds.t1, imds.t2
    dtype = np.result_type(t1, t2)
    nocc, nvir = t1.shape

    fock = imds.eris.fock
    foo = fock[:nocc,:nocc]
    fvv = fock[nocc:,nocc:]

    Hr1 = np.diag(imds.Lvv)
    Hr2 = np.zeros((nocc,nvir,nvir), dtype)
    for a in range(nvir):
        if eom.partition != 'mp':
            _Wvvvva = np.array(imds.Wvvvv[a])
        for b in range(nvir):
            for j in range(nocc):
                if eom.partition == 'mp':
                    Hr2[j,a,b] += fvv[a,a]
                    Hr2[j,a,b] += fvv[b,b]
                    Hr2[j,a,b] += -foo[j,j]
                else:
                    Hr2[j,a,b] += imds.Lvv[a,a]
                    Hr2[j,a,b] += imds.Lvv[b,b]
                    Hr2[j,a,b] += -imds.Loo[j,j]
                    Hr2[j,a,b] += 2*imds.Wovvo[j,b,b,j]
                    Hr2[j,a,b] += -imds.Wovov[j,b,j,b]
                    Hr2[j,a,b] += -imds.Wovov[j,a,j,a]
                    Hr2[j,a,b] += -imds.Wovvo[j,b,b,j]*(a==b)
                    Hr2[j,a,b] += _Wvvvva[b,a,b]
                    Hr2[j,a,b] += -2*np.dot(imds.Woovv[:,j,a,b], t2[:,j,a,b])
                    Hr2[j,a,b] += np.dot(imds.Woovv[:,j,b,a], t2[:,j,a,b])

    vector = amplitudes_to_vector_ea(Hr1,Hr2)
    return vector 

def Wooov(t1, t2, eris):
    t1a, t1b = t1
    t2aa, t2ab, t2bb = t2
    dtype = np.result_type(t1a, t1b, t2aa, t2ab, t2bb)

    eris_ovoo = np.asarray(eris.ovoo)
    eris_OVOO = np.asarray(eris.OVOO)
    eris_OVoo = np.asarray(eris.OVoo)
    eris_ovOO = np.asarray(eris.ovOO)
    ovoo = eris_ovoo - eris_ovoo.transpose(2,1,0,3)
    OVOO = eris_OVOO - eris_OVOO.transpose(2,1,0,3)
    wooov = np.array(     ovoo.transpose(2,3,0,1), dtype=dtype)
    wOOOV = np.array(     OVOO.transpose(2,3,0,1), dtype=dtype)
    wooOV = np.array(eris_OVoo.transpose(2,3,0,1), dtype=dtype)
    wOOov = np.array(eris_ovOO.transpose(2,3,0,1), dtype=dtype)
    eris_ovoo = eris_OVOO = eris_ovOO = eris_OVoo = None

    eris_ovov = np.asarray(eris.ovov)
    eris_OVOV = np.asarray(eris.OVOV)
    eris_ovOV = np.asarray(eris.ovOV)
    ovov = eris_ovov - eris_ovov.transpose(0,3,2,1)
    OVOV = eris_OVOV - eris_OVOV.transpose(0,3,2,1)

    wooov += lib.einsum('if,mfne->mine', t1a,      ovov)
    wOOOV += lib.einsum('if,mfne->mine', t1b,      OVOV)
    wooOV += lib.einsum('if,mfNE->miNE', t1a, eris_ovOV)
    wOOov += lib.einsum('IF,neMF->MIne', t1b, eris_ovOV)
    return wooov, wooOV, wOOov, wOOOV 

def _contract_s1vvvv_t2(mycc, mol, vvvv, t2, out=None, verbose=None):
    '''Ht2 = numpy.einsum('ijcd,acdb->ijab', t2, vvvv)
    where vvvv can be real or complex and no permutation symmetry is available in vvvv.

    Args:
        vvvv : None or integral object
            if vvvv is None, contract t2 to AO-integrals using AO-direct algorithm
    '''
    # vvvv == None means AO-direct CCSD. It should redirect to
    # _contract_s4vvvv_t2(mycc, mol, vvvv, t2, out, verbose)
    assert(vvvv is not None)

    time0 = time.clock(), time.time()
    log = logger.new_logger(mycc, verbose)

    nvira, nvirb = t2.shape[-2:]
    x2 = t2.reshape(-1,nvira,nvirb)
    nocc2 = x2.shape[0]
    dtype = numpy.result_type(t2, vvvv)
    Ht2 = numpy.ndarray(x2.shape, dtype=dtype, buffer=out)

    max_memory = mycc.max_memory - lib.current_memory()[0]
    unit = nvirb**2*nvira*2 + nocc2*nvirb + 1
    blksize = min(nvira, max(BLKMIN, int(max_memory*1e6/8/unit)))

    for p0,p1 in lib.prange(0, nvira, blksize):
        Ht2[:,p0:p1] = lib.einsum('xcd,acbd->xab', x2, vvvv[p0:p1])
        time0 = log.timer_debug1('vvvv [%d:%d]' % (p0,p1), *time0)
    return Ht2.reshape(t2.shape) 

def _gen_contract_baa(ts, vooo, fock, mo_energy, orbsym, log):
    t1aT, t1bT, t2aaT, t2abT = ts
    focka, fockb = fock
    vooo, vOoO, VoOo = vooo
    nvira, nocca = t1aT.shape
    nvirb, noccb = t1bT.shape
    mo_ea = numpy.asarray(mo_energy[0], order='C')
    mo_eb = numpy.asarray(mo_energy[1], order='C')
    fvo = focka[nocca:,:nocca].copy()
    fVO = fockb[noccb:,:noccb].copy()

    cpu2 = [time.clock(), time.time()]
    dtype = numpy.result_type(t2aaT.dtype, vooo.dtype)
    if dtype == numpy.complex:
        drv = _ccsd.libcc.CCuccsd_t_zbaa
    else:
        drv = _ccsd.libcc.CCuccsd_t_baa
    def contract(et_sum, a0, a1, b0, b1, cache):
        cache_row_a, cache_col_a, cache_row_b, cache_col_b = cache
        drv(et_sum.ctypes.data_as(ctypes.c_void_p),
            mo_ea.ctypes.data_as(ctypes.c_void_p),
            mo_eb.ctypes.data_as(ctypes.c_void_p),
            t1aT.ctypes.data_as(ctypes.c_void_p),
            t1bT.ctypes.data_as(ctypes.c_void_p),
            t2aaT.ctypes.data_as(ctypes.c_void_p),
            t2abT.ctypes.data_as(ctypes.c_void_p),
            vooo.ctypes.data_as(ctypes.c_void_p),
            vOoO.ctypes.data_as(ctypes.c_void_p),
            VoOo.ctypes.data_as(ctypes.c_void_p),
            fvo.ctypes.data_as(ctypes.c_void_p),
            fVO.ctypes.data_as(ctypes.c_void_p),
            ctypes.c_int(nocca), ctypes.c_int(noccb),
            ctypes.c_int(nvira), ctypes.c_int(nvirb),
            ctypes.c_int(a0), ctypes.c_int(a1),
            ctypes.c_int(b0), ctypes.c_int(b1),
            cache_row_a.ctypes.data_as(ctypes.c_void_p),
            cache_col_a.ctypes.data_as(ctypes.c_void_p),
            cache_row_b.ctypes.data_as(ctypes.c_void_p),
            cache_col_b.ctypes.data_as(ctypes.c_void_p))
        cpu2[:] = log.timer_debug1('contract %d:%d,%d:%d'%(a0,a1,b0,b1), *cpu2)
    return contract 

def eaccsd_diag(eom, imds=None):
    if imds is None: imds = eom.make_imds()
    t1, t2 = imds.t1, imds.t2
    dtype = np.result_type(t1, t2)
    nocc, nvir = t1.shape

    fock = imds.eris.fock
    foo = fock[:nocc,:nocc]
    fvv = fock[nocc:,nocc:]

    Hr1 = np.diag(imds.Lvv)
    Hr2 = np.zeros((nocc,nvir,nvir), dtype)
    for a in range(nvir):
        if eom.partition != 'mp':
            _Wvvvva = np.array(imds.Wvvvv[a])
        for b in range(nvir):
            for j in range(nocc):
                if eom.partition == 'mp':
                    Hr2[j,a,b] += fvv[a,a]
                    Hr2[j,a,b] += fvv[b,b]
                    Hr2[j,a,b] += -foo[j,j]
                else:
                    Hr2[j,a,b] += imds.Lvv[a,a]
                    Hr2[j,a,b] += imds.Lvv[b,b]
                    Hr2[j,a,b] += -imds.Loo[j,j]
                    Hr2[j,a,b] += 2*imds.Wovvo[j,b,b,j]
                    Hr2[j,a,b] += -imds.Wovov[j,b,j,b]
                    Hr2[j,a,b] += -imds.Wovov[j,a,j,a]
                    Hr2[j,a,b] += -imds.Wovvo[j,b,b,j]*(a==b)
                    Hr2[j,a,b] += _Wvvvva[b,a,b]
                    Hr2[j,a,b] += -2*np.dot(imds.Woovv[:,j,a,b], t2[:,j,a,b])
                    Hr2[j,a,b] += np.dot(imds.Woovv[:,j,b,a], t2[:,j,a,b])

    vector = amplitudes_to_vector_ea(Hr1,Hr2)
    return vector 

def amplitudes_to_vector_triplet(t1, t2, out=None):
    t2aa, t2ab = t2
    dtype = np.result_type(t1, t2aa, t2ab)
    nocc, nvir = t1.shape
    nov = nocc * nvir
    size1 = nov + nocc*(nocc-1)//2*nvir*(nvir-1)//2
    size = size1 + nov*(nov+1)//2
    vector = np.ndarray(size, dtype, buffer=out)
    ccsd.amplitudes_to_vector_s4(t1, t2[0], out=vector)
    t2ab = t2[1].transpose(0,2,1,3).reshape(nov,nov)
    lib.pack_tril(t2ab, out=vector[size1:])
    return vector 

def dsolve(aop, b, precond, tol=1e-12, max_cycle=30, dot=numpy.dot,
           lindep=DSOLVE_LINDEP, verbose=0, tol_residual=None):
    '''Davidson iteration to solve linear equation.  It works bad.
    '''

    if tol_residual is None:
        toloose = numpy.sqrt(tol)
    else:
        toloose = tol_residual

    assert callable(precond)

    xs = [precond(b)]
    ax = [aop(xs[-1])]

    dtype = numpy.result_type(ax[0], xs[0])
    aeff = numpy.zeros((max_cycle,max_cycle), dtype=dtype)
    beff = numpy.zeros((max_cycle), dtype=dtype)
    for istep in range(max_cycle):
        beff[istep] = dot(xs[istep], b)
        for i in range(istep+1):
            aeff[istep,i] = dot(xs[istep], ax[i])
            aeff[i,istep] = dot(xs[i], ax[istep])

        v = scipy.linalg.solve(aeff[:istep+1,:istep+1], beff[:istep+1])
        xtrial = dot(v, xs)
        dx = b - dot(v, ax)
        rr = numpy_helper.norm(dx)
        if verbose:
            print('davidson', istep, rr)
        if rr < toloose:
            break
        xs.append(precond(dx))
        ax.append(aop(xs[-1]))

    if verbose:
        print(istep)

    return xtrial 

def _unsigned_subtract(a, b):
    """
    Subtract two values where a >= b, and produce an unsigned result

    This is needed when finding the difference between the upper and lower
    bound of an int16 histogram
    """
    # coerce to a single type
    signed_to_unsigned = {
        np.byte: np.ubyte,
        np.short: np.ushort,
        np.intc: np.uintc,
        np.int_: np.uint,
        np.longlong: np.ulonglong
    }
    dt = np.result_type(a, b)
    try:
        dt = signed_to_unsigned[dt.type]
    except KeyError:
        return np.subtract(a, b, dtype=dt)
    else:
        # we know the inputs are integers, and we are deliberately casting
        # signed to unsigned
        return np.subtract(a, b, casting='unsafe', dtype=dt) 

def apply_along_fields(func, arr):
    """
    Apply function 'func' as a reduction across fields of a structured array.

    This is similar to `apply_along_axis`, but treats the fields of a
    structured array as an extra axis. The fields are all first cast to a
    common type following the type-promotion rules from `numpy.result_type`
    applied to the field's dtypes.

    Parameters
    ----------
    func : function
       Function to apply on the "field" dimension. This function must
       support an `axis` argument, like np.mean, np.sum, etc.
    arr : ndarray
       Structured array for which to apply func.

    Returns
    -------
    out : ndarray
       Result of the recution operation

    Examples
    --------

    >>> b = np.array([(1, 2, 5), (4, 5, 7), (7, 8 ,11), (10, 11, 12)],
    ...              dtype=[('x', 'i4'), ('y', 'f4'), ('z', 'f8')])
    >>> apply_along_fields(np.mean, b)
    array([ 2.66666667,  5.33333333,  8.66666667, 11.        ])
    >>> apply_along_fields(np.mean, b[['x', 'z']])
    array([ 3. ,  5.5,  9. , 11. ])

    """
    if arr.dtype.names is None:
        raise ValueError('arr must be a structured array')

    uarr = structured_to_unstructured(arr)
    return func(uarr, axis=-1)
    # works and avoids axis requirement, but very, very slow:
    #return np.apply_along_axis(func, -1, uarr)