Python numpy.result_type() Examples

The following are 30 code examples for showing how to use numpy.result_type(). These examples are extracted from open source projects. You can vote up the ones you like or vote down the ones you don't like, and go to the original project or source file by following the links above each example.

You may check out the related API usage on the sidebar.

You may also want to check out all available functions/classes of the module numpy , or try the search function .

Example 1
Project: pyscf   Author: pyscf   File: pywannier90.py    License: Apache License 2.0 6 votes vote down vote up
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 
Example 2
Project: pyscf   Author: pyscf   File: incore.py    License: Apache License 2.0 6 votes vote down vote up
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 
Example 3
Project: recruit   Author: Frank-qlu   File: sparse.py    License: Apache License 2.0 6 votes vote down vote up
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 
Example 4
Project: mars   Author: mars-project   File: histogram.py    License: Apache License 2.0 6 votes vote down vote up
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) 
Example 5
Project: lambda-packs   Author: ryfeus   File: histograms.py    License: MIT License 6 votes vote down vote up
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) 
Example 6
Project: vnpy_crypto   Author: birforce   File: align.py    License: MIT License 6 votes vote down vote up
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 
Example 7
Project: Computable   Author: ktraunmueller   File: align.py    License: MIT License 6 votes vote down vote up
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 
Example 8
Project: Computable   Author: ktraunmueller   File: align.py    License: MIT License 6 votes vote down vote up
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 
Example 9
Project: Mastering-Elasticsearch-7.0   Author: PacktPublishing   File: histograms.py    License: MIT License 6 votes vote down vote up
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) 
Example 10
Project: trax   Author: google   File: lax_numpy_test.py    License: Apache License 2.0 6 votes vote down vote up
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 
Example 11
Project: trax   Author: google   File: lax_numpy_test.py    License: Apache License 2.0 6 votes vote down vote up
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) 
Example 12
Project: trax   Author: google   File: lax_numpy_test.py    License: Apache License 2.0 6 votes vote down vote up
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) 
Example 13
Project: trax   Author: google   File: lax_numpy_test.py    License: Apache License 2.0 6 votes vote down vote up
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) 
Example 14
Project: pyscf   Author: pyscf   File: numint.py    License: Apache License 2.0 5 votes vote down vote up
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 
Example 15
Project: pyscf   Author: pyscf   File: numint.py    License: Apache License 2.0 5 votes vote down vote up
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 
Example 16
Project: pyscf   Author: pyscf   File: stability.py    License: Apache License 2.0 5 votes vote down vote up
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 
Example 17
Project: pyscf   Author: pyscf   File: kintermediates_uhf.py    License: Apache License 2.0 5 votes vote down vote up
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 
Example 18
Project: pyscf   Author: pyscf   File: kintermediates_uhf.py    License: Apache License 2.0 5 votes vote down vote up
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 
Example 19
Project: pyscf   Author: pyscf   File: kintermediates_uhf.py    License: Apache License 2.0 5 votes vote down vote up
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 
Example 20
Project: pyscf   Author: pyscf   File: kccsd_rhf.py    License: Apache License 2.0 5 votes vote down vote up
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 
Example 21
Project: pyscf   Author: pyscf   File: fft_ao2mo.py    License: Apache License 2.0 5 votes vote down vote up
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 
Example 22
Project: pyscf   Author: pyscf   File: eom_rccsd_hybrid.py    License: Apache License 2.0 5 votes vote down vote up
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 
Example 23
Project: pyscf   Author: pyscf   File: uintermediates.py    License: Apache License 2.0 5 votes vote down vote up
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 
Example 24
Project: pyscf   Author: pyscf   File: ccsd.py    License: Apache License 2.0 5 votes vote down vote up
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) 
Example 25
Project: pyscf   Author: pyscf   File: uccsd_t.py    License: Apache License 2.0 5 votes vote down vote up
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 
Example 26
Project: pyscf   Author: pyscf   File: eom_rccsd.py    License: Apache License 2.0 5 votes vote down vote up
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 
Example 27
Project: pyscf   Author: pyscf   File: eom_rccsd.py    License: Apache License 2.0 5 votes vote down vote up
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 
Example 28
Project: pyscf   Author: pyscf   File: linalg_helper.py    License: Apache License 2.0 5 votes vote down vote up
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 
Example 29
Project: recruit   Author: Frank-qlu   File: histograms.py    License: Apache License 2.0 5 votes vote down vote up
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) 
Example 30
Project: recruit   Author: Frank-qlu   File: recfunctions.py    License: Apache License 2.0 5 votes vote down vote up
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)