1#!/usr/bin/env python
2# Copyright 2014-2018 The PySCF Developers. All Rights Reserved.
3#
4# Licensed under the Apache License, Version 2.0 (the "License");
5# you may not use this file except in compliance with the License.
6# You may obtain a copy of the License at
7#
8#     http://www.apache.org/licenses/LICENSE-2.0
9#
10# Unless required by applicable law or agreed to in writing, software
11# distributed under the License is distributed on an "AS IS" BASIS,
12# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13# See the License for the specific language governing permissions and
14# limitations under the License.
15
16import numpy
17from pyscf import lib
18
19def spatial2spin(tx, orbspin=None):
20    '''Convert T1/T2 of spatial orbital representation to T1/T2 of
21    spin-orbital representation
22    '''
23    if isinstance(tx, numpy.ndarray) and tx.ndim == 2:
24        # RCCSD t1 amplitudes
25        return spatial2spin((tx,tx), orbspin)
26    elif isinstance(tx, numpy.ndarray) and tx.ndim == 4:
27        # RCCSD t2 amplitudes
28        t2aa = tx - tx.transpose(0,1,3,2)
29        return spatial2spin((t2aa,tx,t2aa), orbspin)
30    elif len(tx) == 2:  # t1
31        t1a, t1b = tx
32        nocc_a, nvir_a = t1a.shape
33        nocc_b, nvir_b = t1b.shape
34    else:
35        t2aa, t2ab, t2bb = tx
36        nocc_a, nocc_b, nvir_a, nvir_b = t2ab.shape
37
38    if orbspin is None:
39        orbspin = numpy.zeros((nocc_a+nvir_a)*2, dtype=int)
40        orbspin[1::2] = 1
41
42    nocc = nocc_a + nocc_b
43    nvir = nvir_a + nvir_b
44    idxoa = numpy.where(orbspin[:nocc] == 0)[0]
45    idxob = numpy.where(orbspin[:nocc] == 1)[0]
46    idxva = numpy.where(orbspin[nocc:] == 0)[0]
47    idxvb = numpy.where(orbspin[nocc:] == 1)[0]
48
49    if len(tx) == 2:  # t1
50        t1 = numpy.zeros((nocc,nvir), dtype=t1a.dtype)
51        lib.takebak_2d(t1, t1a, idxoa, idxva)
52        lib.takebak_2d(t1, t1b, idxob, idxvb)
53        t1 = lib.tag_array(t1, orbspin=orbspin)
54        return t1
55
56    else:
57        t2 = numpy.zeros((nocc**2,nvir**2), dtype=t2aa.dtype)
58        idxoaa = idxoa[:,None] * nocc + idxoa
59        idxoab = idxoa[:,None] * nocc + idxob
60        idxoba = idxob[:,None] * nocc + idxoa
61        idxobb = idxob[:,None] * nocc + idxob
62        idxvaa = idxva[:,None] * nvir + idxva
63        idxvab = idxva[:,None] * nvir + idxvb
64        idxvba = idxvb[:,None] * nvir + idxva
65        idxvbb = idxvb[:,None] * nvir + idxvb
66        t2aa = t2aa.reshape(nocc_a*nocc_a,nvir_a*nvir_a)
67        t2ab = t2ab.reshape(nocc_a*nocc_b,nvir_a*nvir_b)
68        t2bb = t2bb.reshape(nocc_b*nocc_b,nvir_b*nvir_b)
69        lib.takebak_2d(t2, t2aa, idxoaa.ravel()  , idxvaa.ravel()  )
70        lib.takebak_2d(t2, t2bb, idxobb.ravel()  , idxvbb.ravel()  )
71        lib.takebak_2d(t2, t2ab, idxoab.ravel()  , idxvab.ravel()  )
72        lib.takebak_2d(t2, t2ab, idxoba.T.ravel(), idxvba.T.ravel())
73        abba = -t2ab
74        lib.takebak_2d(t2, abba, idxoab.ravel()  , idxvba.T.ravel())
75        lib.takebak_2d(t2, abba, idxoba.T.ravel(), idxvab.ravel()  )
76        t2 = lib.tag_array(t2, orbspin=orbspin)
77        return t2.reshape(nocc,nocc,nvir,nvir)
78
79spatial2spinorb = spatial2spin
80
81def spin2spatial(tx, orbspin):
82    if tx.ndim == 2:  # t1
83        nocc, nvir = tx.shape
84    else:
85        nocc, nvir = tx.shape[1:3]
86
87    idxoa = numpy.where(orbspin[:nocc] == 0)[0]
88    idxob = numpy.where(orbspin[:nocc] == 1)[0]
89    idxva = numpy.where(orbspin[nocc:] == 0)[0]
90    idxvb = numpy.where(orbspin[nocc:] == 1)[0]
91    nocc_a = len(idxoa)
92    nocc_b = len(idxob)
93    nvir_a = len(idxva)
94    nvir_b = len(idxvb)
95
96    if tx.ndim == 2:  # t1
97        t1a = lib.take_2d(tx, idxoa, idxva)
98        t1b = lib.take_2d(tx, idxob, idxvb)
99        return t1a, t1b
100    else:
101        idxoaa = idxoa[:,None] * nocc + idxoa
102        idxoab = idxoa[:,None] * nocc + idxob
103        idxobb = idxob[:,None] * nocc + idxob
104        idxvaa = idxva[:,None] * nvir + idxva
105        idxvab = idxva[:,None] * nvir + idxvb
106        idxvbb = idxvb[:,None] * nvir + idxvb
107        t2 = tx.reshape(nocc**2,nvir**2)
108        t2aa = lib.take_2d(t2, idxoaa.ravel(), idxvaa.ravel())
109        t2bb = lib.take_2d(t2, idxobb.ravel(), idxvbb.ravel())
110        t2ab = lib.take_2d(t2, idxoab.ravel(), idxvab.ravel())
111        t2aa = t2aa.reshape(nocc_a,nocc_a,nvir_a,nvir_a)
112        t2bb = t2bb.reshape(nocc_b,nocc_b,nvir_b,nvir_b)
113        t2ab = t2ab.reshape(nocc_a,nocc_b,nvir_a,nvir_b)
114        return t2aa,t2ab,t2bb
115
116def convert_to_uccsd(mycc):
117    from pyscf import scf
118    from pyscf.cc import uccsd, gccsd
119    if isinstance(mycc, uccsd.UCCSD):
120        return mycc
121    elif isinstance(mycc, gccsd.GCCSD):
122        raise NotImplementedError
123
124    mf = scf.addons.convert_to_uhf(mycc._scf)
125    ucc = uccsd.UCCSD(mf)
126    assert(mycc._nocc is None)
127    assert(mycc._nmo is None)
128    ucc.__dict__.update(mycc.__dict__)
129    ucc._scf = mf
130    ucc.mo_coeff = mf.mo_coeff
131    ucc.mo_occ = mf.mo_occ
132    if not (mycc.frozen is None or isinstance(mycc.frozen, (int, numpy.integer))):
133        raise NotImplementedError
134    ucc.t1, ucc.t2 = uccsd.amplitudes_from_rccsd(mycc.t1, mycc.t2)
135    return ucc
136
137def convert_to_gccsd(mycc):
138    from pyscf import scf
139    from pyscf.cc import gccsd
140    if isinstance(mycc, gccsd.GCCSD):
141        return mycc
142
143    mf = scf.addons.convert_to_ghf(mycc._scf)
144    gcc = gccsd.GCCSD(mf)
145    assert(mycc._nocc is None)
146    assert(mycc._nmo is None)
147    gcc.__dict__.update(mycc.__dict__)
148    gcc._scf = mf
149    gcc.mo_coeff = mf.mo_coeff
150    gcc.mo_occ = mf.mo_occ
151    if isinstance(mycc.frozen, (int, numpy.integer)):
152        gcc.frozen = mycc.frozen * 2
153    elif not (mycc.frozen is None or mycc.frozen == 0):
154        raise NotImplementedError
155    gcc.t1 = spatial2spin(mycc.t1, mf.mo_coeff.orbspin)
156    gcc.t2 = spatial2spin(mycc.t2, mf.mo_coeff.orbspin)
157    return gcc
158