Python numpy.number() Examples

The following are 30 code examples for showing how to use numpy.number(). 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: EXOSIMS   Author: dsavransky   File: SurveySimulation.py    License: BSD 3-Clause "New" or "Revised" License 6 votes vote down vote up
def revisitFilter(self,sInds,tmpCurrentTimeNorm):
        """Helper method for Overloading Revisit Filtering

        Args:
            sInds - indices of stars still in observation list
            tmpCurrentTimeNorm (MJD) - the simulation time after overhead was added in MJD form
        Returns:
            sInds - indices of stars still in observation list
        """

        tovisit = np.zeros(self.TargetList.nStars, dtype=bool)
        if len(sInds) > 0:
            tovisit[sInds] = ((self.starVisits[sInds] == min(self.starVisits[sInds])) \
                    & (self.starVisits[sInds] < self.nVisitsMax))#Checks that no star has exceeded the number of revisits and the indicies of all considered stars have minimum number of observations
            #The above condition should prevent revisits so long as all stars have not been observed
            if self.starRevisit.size != 0:
                dt_rev = np.abs(self.starRevisit[:,1]*u.day - tmpCurrentTimeNorm)
                ind_rev = [int(x) for x in self.starRevisit[dt_rev < self.dt_max,0] 
                        if x in sInds]
                tovisit[ind_rev] = (self.starVisits[ind_rev] < self.nVisitsMax)
            sInds = np.where(tovisit)[0]
        return sInds 
Example 2
Project: pyscf   Author: pyscf   File: fciqmc.py    License: Apache License 2.0 6 votes vote down vote up
def one_from_two_pdm(two_pdm, nelec):
    '''Return a 1-rdm, given a 2-rdm to contract.

    Args:
        two_pdm : ndarray
            A (spin-free) 2-particle reduced density matrix.
        nelec: int
            The number of electrons contributing to the RDMs.

    Returns:
        one_pdm : ndarray
            The (spin-free) 1-particle reduced density matrix.
    '''

    # Last two indices refer to middle two second quantized operators in the 2RDM
    one_pdm = numpy.einsum('ikjj->ik', two_pdm)
    one_pdm /= (numpy.sum(nelec)-1)
    return one_pdm 
Example 3
Project: pyscf   Author: pyscf   File: fciqmc.py    License: Apache License 2.0 6 votes vote down vote up
def one_from_two_pdm(two_pdm, nelec):
    '''Return a 1-rdm, given a 2-rdm to contract.

    Args:
        two_pdm : ndarray
            A (spin-free) 2-particle reduced density matrix.
        nelec: int
            The number of electrons contributing to the RDMs.

    Returns:
        one_pdm : ndarray
            The (spin-free) 1-particle reduced density matrix.
    '''

    # Last two indices refer to middle two second quantized operators in the 2RDM
    one_pdm = numpy.einsum('ikjj->ik', two_pdm)
    one_pdm /= (numpy.sum(nelec)-1)
    return one_pdm 
Example 4
Project: pyscf   Author: pyscf   File: direct_nosym.py    License: Apache License 2.0 6 votes vote down vote up
def kernel(self, h1e, eri, norb, nelec, ci0=None,
               tol=None, lindep=None, max_cycle=None, max_space=None,
               nroots=None, davidson_only=None, pspace_size=None,
               orbsym=None, wfnsym=None, ecore=0, **kwargs):
        if isinstance(nelec, (int, numpy.number)):
            nelecb = nelec//2
            neleca = nelec - nelecb
        else:
            neleca, nelecb = nelec
        link_indexa = cistring.gen_linkstr_index(range(norb), neleca)
        link_indexb = cistring.gen_linkstr_index(range(norb), nelecb)
        e, c = direct_spin1.kernel_ms1(self, h1e, eri, norb, nelec, ci0,
                                       (link_indexa,link_indexb),
                                       tol, lindep, max_cycle, max_space, nroots,
                                       davidson_only, pspace_size, ecore=ecore,
                                       **kwargs)
        self.eci, self.ci = e, c
        return e, c 
Example 5
Project: pyscf   Author: pyscf   File: direct_spin1_symm.py    License: Apache License 2.0 6 votes vote down vote up
def reorder_eri(eri, norb, orbsym):
    if orbsym is None:
        return [eri], numpy.arange(norb), numpy.zeros(norb,dtype=numpy.int32)
# map irrep IDs of Dooh or Coov to D2h, C2v
# see symm.basis.linearmole_symm_descent
    orbsym = numpy.asarray(orbsym) % 10
# irrep of (ij| pair
    trilirrep = (orbsym[:,None]^orbsym)[numpy.tril_indices(norb)]
# and the number of occurence for each irrep
    dimirrep = numpy.asarray(numpy.bincount(trilirrep), dtype=numpy.int32)
# we sort the irreps of (ij| pair, to group the pairs which have same irreps
# "order" is irrep-id-sorted index. The (ij| paired is ordered that the
# pair-id given by order[0] comes first in the sorted pair
# "rank" is a sorted "order". Given nth (ij| pair, it returns the place(rank)
# of the sorted pair
    old_eri_irrep = numpy.asarray(trilirrep, dtype=numpy.int32)
    rank_in_irrep = numpy.empty_like(old_eri_irrep)
    p0 = 0
    eri_irs = [numpy.zeros((0,0))] * TOTIRREPS
    for ir, nnorb in enumerate(dimirrep):
        idx = numpy.asarray(numpy.where(trilirrep == ir)[0], dtype=numpy.int32)
        rank_in_irrep[idx] = numpy.arange(nnorb, dtype=numpy.int32)
        eri_irs[ir] = lib.take_2d(eri, idx, idx)
        p0 += nnorb
    return eri_irs, rank_in_irrep, old_eri_irrep 
Example 6
Project: pyscf   Author: pyscf   File: wfn_format.py    License: Apache License 2.0 6 votes vote down vote up
def write_ci(fout, fcivec, norb, nelec, ncore=0):
    from pyscf import fci
    if isinstance(nelec, (int, numpy.number)):
        nelecb = nelec//2
        neleca = nelec - nelecb
    else:
        neleca, nelecb = nelec
    fout.write('NELACTIVE, NDETS, NORBCORE, NORBACTIVE\n')
    fout.write(' %5d %5d %5d %5d\n' % (neleca+nelecb, fcivec.size, ncore, norb))
    fout.write('COEFFICIENT/ OCCUPIED ACTIVE SPIN ORBITALS\n')

    nb = fci.cistring.num_strings(norb, nelecb)
    stringsa = fci.cistring.gen_strings4orblist(range(norb), neleca)
    stringsb = fci.cistring.gen_strings4orblist(range(norb), nelecb)
    def str2orbidx(string, ncore):
        bstring = bin(string)
        return [i+1+ncore for i,s in enumerate(bstring[::-1]) if s == '1']

    addrs = numpy.argsort(abs(fcivec.ravel()))
    for iaddr in reversed(addrs):
        addra, addrb = divmod(iaddr, nb)
        idxa = ['%3d' % x for x in str2orbidx(stringsa[addra], ncore)]
        idxb = ['%3d' % (-x) for x in str2orbidx(stringsb[addrb], ncore)]
        #TODO:add a cuttoff and a counter for ndets
        fout.write('%18.10E %s %s\n' % (fcivec[addra,addrb], ' '.join(idxa), ' '.join(idxb))) 
Example 7
Project: pyscf   Author: pyscf   File: cisd.py    License: Apache License 2.0 6 votes vote down vote up
def from_fcivec(ci0, norb, nelec, frozen=None):
    '''Extract CISD coefficients from FCI coefficients'''
    if not (frozen is None or frozen == 0):
        raise NotImplementedError

    if isinstance(nelec, (int, numpy.number)):
        nelecb = nelec//2
        neleca = nelec - nelecb
    else:
        neleca, nelecb = nelec
    nocc = neleca
    nvir = norb - nocc
    t1addr, t1sign = t1strs(norb, nocc)

    c0 = ci0[0,0]
    c1 = ci0[0,t1addr] * t1sign
    c2 = numpy.einsum('i,j,ij->ij', t1sign, t1sign, ci0[t1addr[:,None],t1addr])
    c1 = c1.reshape(nocc,nvir)
    c2 = c2.reshape(nocc,nvir,nocc,nvir).transpose(0,2,1,3)
    return amplitudes_to_cisdvec(c0, c1, c2) 
Example 8
Project: python-control   Author: python-control   File: iosys.py    License: BSD 3-Clause "New" or "Revised" License 6 votes vote down vote up
def __rmul__(sys1, sys2):
        """Pre-multiply an input/output systems by a scalar/matrix"""
        if isinstance(sys2, (int, float, np.number)):
            # TODO: Scale the output
            raise NotImplemented("Scalar multiplication not yet implemented")
        elif isinstance(sys2, np.ndarray):
            # TODO: Post-multiply by a matrix
            raise NotImplemented("Matrix multiplication not yet implemented")
        elif isinstance(sys1, StateSpace) and isinstance(sys2, StateSpace):
            # Special case: maintain linear systems structure
            new_ss_sys = StateSpace.__rmul__(sys1, sys2)
            # TODO: set input and output names
            new_io_sys = LinearIOSystem(new_ss_sys)

            return new_io_sys
        elif not isinstance(sys2, InputOutputSystem):
            raise ValueError("Unknown I/O system object ", sys1)
        else:
            # Both systetms are InputOutputSystems => use __mul__
            return InputOutputSystem.__mul__(sys2, sys1) 
Example 9
Project: python-control   Author: python-control   File: iosys.py    License: BSD 3-Clause "New" or "Revised" License 6 votes vote down vote up
def set_outputs(self, outputs, prefix='y'):
        """Set the number/names of the system outputs.

        Parameters
        ----------
        outputs : int, list of str, or None
            Description of the system outputs.  This can be given as an integer
            count or as a list of strings that name the individual signals.
            If an integer count is specified, the names of the signal will be
            of the form `u[i]` (where the prefix `u` can be changed using the
            optional prefix parameter).
        prefix : string, optional
            If `outputs` is an integer, create the names of the states using
            the given prefix (default = 'y').  The names of the input will be
            of the form `prefix[i]`.

        """
        self.noutputs, self.output_index = \
            self._process_signal_list(outputs, prefix=prefix) 
Example 10
Project: python-control   Author: python-control   File: iosys.py    License: BSD 3-Clause "New" or "Revised" License 6 votes vote down vote up
def set_states(self, states, prefix='x'):
        """Set the number/names of the system states.

        Parameters
        ----------
        states : int, list of str, or None
            Description of the system states.  This can be given as an integer
            count or as a list of strings that name the individual signals.
            If an integer count is specified, the names of the signal will be
            of the form `u[i]` (where the prefix `u` can be changed using the
            optional prefix parameter).
        prefix : string, optional
            If `states` is an integer, create the names of the states using
            the given prefix (default = 'x').  The names of the input will be
            of the form `prefix[i]`.

        """
        self.nstates, self.state_index = \
            self._process_signal_list(states, prefix=prefix) 
Example 11
Project: python-control   Author: python-control   File: iosys.py    License: BSD 3-Clause "New" or "Revised" License 6 votes vote down vote up
def set_input_map(self, input_map):
        """Set the input map for an interconnected I/O system.

        Parameters
        ----------
        input_map : 2D array
             Specify the matrix that will be used to multiply the vector of
             system inputs to obtain the vector of subsystem inputs.  These
             values are added to the inputs specified in the connection map.

        """
        # Figure out the number of internal inputs
        ninputs = sum(sys.ninputs for sys in self.syslist)

        # Make sure the input map is the right size
        if input_map.shape[0] != ninputs:
            ValueError("Input map is not the right shape")
        self.input_map = input_map
        self.ninputs = input_map.shape[1] 
Example 12
Project: python-control   Author: python-control   File: frdata.py    License: BSD 3-Clause "New" or "Revised" License 6 votes vote down vote up
def __truediv__(self, other):
        """Divide two LTI objects."""

        if isinstance(other, (int, float, complex, np.number)):
            return FRD(self.fresp * (1/other), self.omega,
                       smooth=(self.ifunc is not None))
        else:
            other = _convertToFRD(other, omega=self.omega)

        if (self.inputs > 1 or self.outputs > 1 or
            other.inputs > 1 or other.outputs > 1):
            raise NotImplementedError(
                "FRD.__truediv__ is currently only implemented for SISO "
                "systems.")

        return FRD(self.fresp/other.fresp, self.omega,
                   smooth=(self.ifunc is not None) and
                          (other.ifunc is not None))

    # TODO: Remove when transition to python3 complete 
Example 13
Project: python-control   Author: python-control   File: frdata.py    License: BSD 3-Clause "New" or "Revised" License 6 votes vote down vote up
def evalfr(self, omega):
        """Evaluate a transfer function at a single angular frequency.

        self._evalfr(omega) returns the value of the frequency response
        at frequency omega.

        Note that a "normal" FRD only returns values for which there is an
        entry in the omega vector. An interpolating FRD can return
        intermediate values.

        """
        warn("FRD.evalfr(omega) will be deprecated in a future release "
             "of python-control; use sys.eval(omega) instead",
             PendingDeprecationWarning)         # pragma: no coverage
        return self._evalfr(omega)

    # Define the `eval` function to evaluate an FRD at a given (real)
    # frequency.  Note that we choose to use `eval` instead of `evalfr` to
    # avoid confusion with :func:`evalfr`, which takes a complex number as its
    # argument.  Similarly, we don't use `__call__` to avoid confusion between
    # G(s) for a transfer function and G(omega) for an FRD object. 
Example 14
Project: python-control   Author: python-control   File: lti.py    License: BSD 3-Clause "New" or "Revised" License 6 votes vote down vote up
def issiso(sys, strict=False):
    """
    Check to see if a system is single input, single output

    Parameters
    ----------
    sys : LTI system
        System to be checked
    strict: bool (default = False)
        If strict is True, do not treat scalars as SISO
    """
    if isinstance(sys, (int, float, complex, np.number)) and not strict:
        return True
    elif not isinstance(sys, LTI):
        raise ValueError("Object is not an LTI system")

    # Done with the tricky stuff...
    return sys.issiso()

# Return the timebase (with conversion if unspecified) 
Example 15
Project: python-control   Author: python-control   File: lti.py    License: BSD 3-Clause "New" or "Revised" License 6 votes vote down vote up
def timebase(sys, strict=True):
    """Return the timebase for an LTI system

    dt = timebase(sys)

    returns the timebase for a system 'sys'.  If the strict option is
    set to False, dt = True will be returned as 1.
    """
    # System needs to be either a constant or an LTI system
    if isinstance(sys, (int, float, complex, np.number)):
        return None
    elif not isinstance(sys, LTI):
        raise ValueError("Timebase not defined")

    # Return the sample time, with converstion to float if strict is false
    if (sys.dt == None):
        return None
    elif (strict):
        return float(sys.dt)

    return sys.dt

# Check to see if two timebases are equal 
Example 16
Project: recruit   Author: Frank-qlu   File: test_regression.py    License: Apache License 2.0 6 votes vote down vote up
def test_ticket_1539(self):
        dtypes = [x for x in np.typeDict.values()
                  if (issubclass(x, np.number)
                      and not issubclass(x, np.timedelta64))]
        a = np.array([], np.bool_)  # not x[0] because it is unordered
        failures = []

        for x in dtypes:
            b = a.astype(x)
            for y in dtypes:
                c = a.astype(y)
                try:
                    np.dot(b, c)
                except TypeError:
                    failures.append((x, y))
        if failures:
            raise AssertionError("Failures: %r" % failures) 
Example 17
Project: EXOSIMS   Author: dsavransky   File: SurveySimulation.py    License: BSD 3-Clause "New" or "Revised" License 5 votes vote down vote up
def initializeStorageArrays(self):
        """
        Initialize all storage arrays based on # of stars and targets
        """

        self.DRM = []
        self.fullSpectra = np.zeros(self.SimulatedUniverse.nPlans, dtype=int)
        self.partialSpectra = np.zeros(self.SimulatedUniverse.nPlans, dtype=int)
        self.propagTimes = np.zeros(self.TargetList.nStars)*u.d
        self.lastObsTimes = np.zeros(self.TargetList.nStars)*u.d
        self.starVisits = np.zeros(self.TargetList.nStars, dtype=int)#contains the number of times each star was visited
        self.starRevisit = np.array([])
        self.starExtended = np.array([], dtype=int)
        self.lastDetected = np.empty((self.TargetList.nStars, 4), dtype=object) 
Example 18
Project: esmlab   Author: NCAR   File: core.py    License: Apache License 2.0 5 votes vote down vote up
def decode_arbitrary_time(num_time_var, units, calendar):
        """Decode an arbitrary time var of type number
        """
        # Check if num_time_var is already decoded:
        if not issubclass(num_time_var.dtype.type, np.number):
            raise ValueError('cannot decode non-numeric time')
        return cftime.num2date(num_time_var, units=units, calendar=calendar) 
Example 19
Project: esmlab   Author: NCAR   File: test_core.py    License: Apache License 2.0 5 votes vote down vote up
def test_time_year_to_midyeardate(dset, time_coord_name='time'):
    assert isinstance(dset[time_coord_name].values[0], np.number)
    dset = dset.esmlab.set_time(time_coord_name=time_coord_name).time_year_to_midyeardate()
    assert isinstance(dset[time_coord_name].values[0], cftime.datetime) 
Example 20
Project: esmlab   Author: NCAR   File: test_core.py    License: Apache License 2.0 5 votes vote down vote up
def test_uncompute_time_var(dset, time_coord_name='time'):
    esm = dset.esmlab.set_time(time_coord_name=time_coord_name)
    xr.testing._assert_internal_invariants(esm._ds_time_computed)
    ds = esm.compute_time_var()
    assert ds[time_coord_name].dtype == np.dtype('O')
    dset_with_uncomputed_time = esm.uncompute_time_var()
    assert np.issubdtype(dset_with_uncomputed_time[time_coord_name].dtype, np.number)


# For some strange reason, this case fails when using pytest parametrization 
Example 21
Project: pyscf   Author: pyscf   File: fciqmc.py    License: Apache License 2.0 5 votes vote down vote up
def dump_flags(self, verbose=None):
        log = logger.new_logger(self, verbose)
        log.info('******** FCIQMC options ********')
        log.info('Number of walkers = %s', self.maxwalkers)
        log.info('Maximum number of iterations = %d', self.maxIter) 
Example 22
Project: pyscf   Author: pyscf   File: fciqmc.py    License: Apache License 2.0 5 votes vote down vote up
def read_neci_one_pdm(fciqmcci, filename, norb, nelec, directory='.'):
    '''Obtain the spin-free 1RDM from neci by reading in the spin free 2RDM.
    If core orbitals have been indicated as frozen in neci, this core contribution
    will be explicitly added back in to the RDM. Therefore, the norb parameter
    should be the total number of orbitals passed to neci (inc. frozen), while
    nelec is the total number of electrons (inc. frozen), but not inactive if running
    through CASSCF.
    '''

    two_pdm = read_neci_two_pdm(fciqmcci, filename, norb, directory)
    one_pdm = one_from_two_pdm(two_pdm, nelec)
    return one_pdm 
Example 23
Project: pyscf   Author: pyscf   File: fciqmc.py    License: Apache License 2.0 5 votes vote down vote up
def dump_flags(self, verbose=None):
        log = logger.new_logger(self, verbose)
        log.info('******** FCIQMC options ********')
        log.info('Number of walkers = %s', self.maxwalkers)
        log.info('Maximum number of iterations = %d', self.maxIter) 
Example 24
Project: pyscf   Author: pyscf   File: fciqmc.py    License: Apache License 2.0 5 votes vote down vote up
def read_neci_one_pdm(fciqmcci, filename, norb, nelec, directory='.'):
    '''Obtain the spin-free 1RDM from neci by reading in the spin free 2RDM.
    If core orbitals have been indicated as frozen in neci, this core contribution
    will be explicitly added back in to the RDM. Therefore, the norb parameter
    should be the total number of orbitals passed to neci (inc. frozen), while
    nelec is the total number of electrons (inc. frozen), but not inactive if running
    through CASSCF.
    '''

    two_pdm = read_neci_two_pdm(fciqmcci, filename, norb, directory)
    one_pdm = one_from_two_pdm(two_pdm, nelec)
    return one_pdm 
Example 25
Project: pyscf   Author: pyscf   File: direct_spin1.py    License: Apache License 2.0 5 votes vote down vote up
def absorb_h1e(h1e, eri, norb, nelec, fac=1):
    '''Modify 2e Hamiltonian to include 1e Hamiltonian contribution.
    '''
    if h1e.dtype == numpy.complex or eri.dtype == numpy.complex:
        raise NotImplementedError('Complex Hamiltonian')

    if not isinstance(nelec, (int, numpy.number)):
        nelec = sum(nelec)
    h2e = ao2mo.restore(1, eri.copy(), norb)
    f1e = h1e - numpy.einsum('jiik->jk', h2e) * .5
    f1e = f1e * (1./(nelec+1e-100))
    for k in range(norb):
        h2e[k,k,:,:] += f1e
        h2e[:,:,k,k] += f1e
    return ao2mo.restore(4, h2e, norb) * fac 
Example 26
Project: pyscf   Author: pyscf   File: direct_nosym.py    License: Apache License 2.0 5 votes vote down vote up
def absorb_h1e(h1e, eri, norb, nelec, fac=1):
    '''Modify 2e Hamiltonian to include 1e Hamiltonian contribution.
    '''
    if not isinstance(nelec, (int, numpy.number)):
        nelec = sum(nelec)
    h2e = ao2mo.restore(1, eri.copy(), norb)
    f1e = h1e - numpy.einsum('jiik->jk', h2e) * .5
    f1e = f1e * (1./(nelec+1e-100))
    for k in range(norb):
        h2e[k,k,:,:] += f1e
        h2e[:,:,k,k] += f1e
    return h2e * fac 
Example 27
Project: pyscf   Author: pyscf   File: direct_nosym.py    License: Apache License 2.0 5 votes vote down vote up
def _unpack(norb, nelec, link_index):
    if link_index is None:
        if isinstance(nelec, (int, numpy.number)):
            nelecb = nelec//2
            neleca = nelec - nelecb
        else:
            neleca, nelecb = nelec
        link_indexa = cistring.gen_linkstr_index(range(norb), neleca)
        link_indexb = cistring.gen_linkstr_index(range(norb), nelecb)
        return link_indexa, link_indexb
    else:
        return link_index 
Example 28
Project: pyscf   Author: pyscf   File: direct_spin0_symm.py    License: Apache License 2.0 5 votes vote down vote up
def dump_flags(self, verbose=None):
        direct_spin0.FCISolver.dump_flags(self, verbose)
        log = logger.new_logger(self, verbose)
        if isinstance(self.wfnsym, str):
            log.info('specified CI wfn symmetry = %s', self.wfnsym)
        elif isinstance(self.wfnsym, (int, numpy.number)):
            log.info('specified CI wfn symmetry = %s',
                     symm.irrep_id2name(self.mol.groupname, self.wfnsym)) 
Example 29
Project: pyscf   Author: pyscf   File: addons.py    License: Apache License 2.0 5 votes vote down vote up
def des_b(ci0, norb, neleca_nelecb, ap_id):
    r'''Construct (N-1)-electron wavefunction by removing a beta electron from
    N-electron wavefunction.

    Args:
        ci0 : 2D array
            CI coefficients, row for alpha strings and column for beta strings.
        norb : int
            Number of orbitals.
        (neleca,nelecb) : (int,int)
            Number of (alpha, beta) electrons of the input CI function
        ap_id : int
            Orbital index (0-based), for the annihilation operator

    Returns:
        2D array, row for alpha strings and column for beta strings. Note it
        has different number of columns to the input CI coefficients.
    '''
    neleca, nelecb = neleca_nelecb
    if nelecb <= 0:
        return numpy.zeros_like(ci0)
    if ci0.ndim == 1:
        ci0 = ci0.reshape(cistring.num_strings(norb, neleca),
                          cistring.num_strings(norb, nelecb))
    des_index = cistring.gen_des_str_index(range(norb), nelecb)
    nb_ci1 = cistring.num_strings(norb, nelecb-1)
    ci1 = numpy.zeros((ci0.shape[0], nb_ci1))

    entry_has_ap = (des_index[:,:,1] == ap_id)
    addr_ci0 = numpy.any(entry_has_ap, axis=1)
    addr_ci1 = des_index[entry_has_ap,2]
    sign = des_index[entry_has_ap,3]
    # This sign prefactor accounts for interchange of operators with alpha and beta spins
    if neleca % 2 == 1:
        sign *= -1
    ci1[:,addr_ci1] = ci0[:,addr_ci0] * sign
    return ci1 
Example 30
Project: pyscf   Author: pyscf   File: addons.py    License: Apache License 2.0 5 votes vote down vote up
def cre_a(ci0, norb, neleca_nelecb, ap_id):
    r'''Construct (N+1)-electron wavefunction by adding an alpha electron in
    the N-electron wavefunction.

    ... math::

        |N+1\rangle = \hat{a}^+_p |N\rangle

    Args:
        ci0 : 2D array
            CI coefficients, row for alpha strings and column for beta strings.
        norb : int
            Number of orbitals.
        (neleca,nelecb) : (int,int)
            Number of (alpha, beta) electrons of the input CI function
        ap_id : int
            Orbital index (0-based), for the creation operator

    Returns:
        2D array, row for alpha strings and column for beta strings. Note it
        has different number of rows to the input CI coefficients.
    '''
    neleca, nelecb = neleca_nelecb
    if neleca >= norb:
        return numpy.zeros_like(ci0)
    if ci0.ndim == 1:
        ci0 = ci0.reshape(cistring.num_strings(norb, neleca),
                          cistring.num_strings(norb, nelecb))
    cre_index = cistring.gen_cre_str_index(range(norb), neleca)
    na_ci1 = cistring.num_strings(norb, neleca+1)
    ci1 = numpy.zeros((na_ci1, ci0.shape[1]))

    entry_has_ap = (cre_index[:,:,0] == ap_id)
    addr_ci0 = numpy.any(entry_has_ap, axis=1)
    addr_ci1 = cre_index[entry_has_ap,2]
    sign = cre_index[entry_has_ap,3]
    ci1[addr_ci1] = sign.reshape(-1,1) * ci0[addr_ci0]
    return ci1

# construct (N+1)-electron wavefunction by adding a beta electron to
# N-electron wavefunction: