JDWavefunctionStepper.C 10.3 KB
Newer Older
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28
////////////////////////////////////////////////////////////////////////////////
//
// Copyright (c) 2008 The Regents of the University of California
//
// This file is part of Qbox
//
// Qbox is distributed under the terms of the GNU General Public License
// as published by the Free Software Foundation, either version 2 of
// the License, or (at your option) any later version.
// See the file COPYING in the root directory of this distribution
// or <http://www.gnu.org/licenses/>.
//
////////////////////////////////////////////////////////////////////////////////
//
// JDWavefunctionStepper.C
//
////////////////////////////////////////////////////////////////////////////////

#include "JDWavefunctionStepper.h"
#include "Wavefunction.h"
#include "SlaterDet.h"
#include "EnergyFunctional.h"
#include "Preconditioner.h"
#include <iostream>
using namespace std;

////////////////////////////////////////////////////////////////////////////////
JDWavefunctionStepper::JDWavefunctionStepper(Wavefunction& wf,
29 30 31
  Preconditioner& prec, EnergyFunctional& ef, TimerMap& tmap) :
  WavefunctionStepper(wf,tmap), prec_(prec), wft_(wf), dwft_(wf), ef_(ef)
{}
Francois Gygi committed
32 33 34

////////////////////////////////////////////////////////////////////////////////
JDWavefunctionStepper::~JDWavefunctionStepper(void)
35
{}
36 37 38 39 40 41 42

////////////////////////////////////////////////////////////////////////////////
void JDWavefunctionStepper::update(Wavefunction& dwf)
{
  // save copy of wf in wft_ and dwf in dwft_
  wft_ = wf_;
  dwft_ = dwf;
Francois Gygi committed
43
  tmap_["jd_residual"].start();
44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63
  for ( int ispin = 0; ispin < wf_.nspin(); ispin++ )
  {
    for ( int ikp = 0; ikp < wf_.nkp(); ikp++ )
    {
      if ( wf_.sd(ispin,ikp)->basis().real() )
      {
        // compute A = Y^T H Y  and descent direction HY - YA
        // proxy real matrices c, cp
        DoubleMatrix c_proxy(wf_.sd(ispin,ikp)->c());
        DoubleMatrix cp_proxy(dwf.sd(ispin,ikp)->c());
        DoubleMatrix a(c_proxy.context(),c_proxy.n(),c_proxy.n(),
                       c_proxy.nb(),c_proxy.nb());

        // factor 2.0 in next line: G and -G
        a.gemm('t','n',2.0,c_proxy,cp_proxy,0.0);
        // rank-1 update correction
        a.ger(-1.0,c_proxy,0,cp_proxy,0);

        // cp = cp - c * a
        cp_proxy.gemm('n','n',-1.0,c_proxy,a,1.0);
Francois Gygi committed
64 65 66 67 68 69 70 71 72 73 74 75 76
      }
      else // real
      {
        // compute A = Y^T H Y  and descent direction HY - YA
        ComplexMatrix& c = wf_.sd(ispin,ikp)->c();
        ComplexMatrix& cp = dwf.sd(ispin,ikp)->c();
        ComplexMatrix a(c.context(),c.n(),c.n(),c.nb(),c.nb());

        // (Y,HY)
        a.gemm('c','n',1.0,c,cp,0.0);

        // cp = cp - c * a
        cp.gemm('n','n',-1.0,c,a,1.0);
77
        // dwf.sd->c() now contains the descent direction (HV-VA)
Francois Gygi committed
78 79 80 81
      }
    } // ikp
  } // ispin
  tmap_["jd_residual"].stop();
82

Francois Gygi committed
83
  // dwf.sd->c() now contains the descent direction (HV-VA)
84
  prec_.update(wf_);
Francois Gygi committed
85 86 87 88 89 90 91 92

  tmap_["jd_compute_z"].start();
  for ( int ispin = 0; ispin < wf_.nspin(); ispin++ )
  {
    for ( int ikp = 0; ikp < wf_.nkp(); ikp++ )
    {
      if ( wf_.sd(ispin,ikp)->basis().real() )
      {
93 94 95 96
        // Apply preconditioner K and store -K(HV-VA) in wf

        double* c = (double*) wf_.sd(ispin,ikp)->c().valptr();
        double* dc = (double*) dwf.sd(ispin,ikp)->c().valptr();
Francois Gygi committed
97 98 99 100
        DoubleMatrix c_proxy(wf_.sd(ispin,ikp)->c());
        DoubleMatrix a(c_proxy.context(),c_proxy.n(),c_proxy.n(),
                       c_proxy.nb(),c_proxy.nb());
        DoubleMatrix c_proxy_t(wft_.sd(ispin,ikp)->c());
101 102 103 104 105 106 107 108 109
        const int mloc = wf_.sd(ispin,ikp)->c().mloc();
        const int ngwl = wf_.sd(ispin,ikp)->basis().localsize();
        const int nloc = wf_.sd(ispin,ikp)->c().nloc();

        for ( int n = 0; n < nloc; n++ )
        {
          // note: double mloc length for complex<double> indices
          double* dcn = &dc[2*mloc*n];
          double* cn  =  &c[2*mloc*n];
Francois Gygi committed
110

111 112
          for ( int i = 0; i < ngwl; i++ )
          {
113
            const double fac = prec_.diag(ispin,ikp,n,i);
114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140
            const double f0 = -fac * dcn[2*i];
            const double f1 = -fac * dcn[2*i+1];
            cn[2*i] = f0;
            cn[2*i+1] = f1;
          }
        }
        // wf_ now contains the preconditioned descent
        // direction Z = -K(HY-YA)

        // orthogonalize Z to Y
        // Z = Z - YY^T Z
        // A = Y^T * Z
        // factor 2.0 in next line: G and -G
        a.gemm('t','n',2.0,c_proxy_t,c_proxy,0.0);
        // rank-1 update correction
        a.ger(-1.0,c_proxy_t,0,c_proxy,0);

        // Z = Z - Y * A
        c_proxy.gemm('n','n',-1.0,c_proxy_t,a,1.0);

        // orthogonalize Z: gram(Z)
        wf_.sd(ispin,ikp)->gram();

        // wf now contains Z, orthonormal
      } // if real
      else
      {
Francois Gygi committed
141
        // Apply preconditioner K and store -K(HV-VA) in wf
Francois Gygi committed
142
        ComplexMatrix& c = wf_.sd(ispin,ikp)->c();
Francois Gygi committed
143
        complex<double>* cv = c.valptr();
Francois Gygi committed
144
        ComplexMatrix& cp = dwf.sd(ispin,ikp)->c();
Francois Gygi committed
145
        complex<double>* cpv = cp.valptr();
Francois Gygi committed
146 147
        ComplexMatrix& ct = wft_.sd(ispin,ikp)->c();
        ComplexMatrix a(c.context(),c.n(),c.n(),c.nb(),c.nb());
Francois Gygi committed
148 149 150 151 152 153 154 155
        const int ngwl = wf_.sd(ispin,ikp)->basis().localsize();
        const int mloc = c.mloc();
        const int nloc = c.nloc();

        for ( int n = 0; n < nloc; n++ )
        {
          complex<double>* cpn = &cpv[mloc*n];
          complex<double>* cn  =  &cv[mloc*n];
Francois Gygi committed
156

Francois Gygi committed
157 158
          for ( int i = 0; i < ngwl; i++ )
          {
159
            const double fac = prec_.diag(ispin,ikp,n,i);
Francois Gygi committed
160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177
            cn[i] = -fac * cpn[i];
          }
        }
        // wf_ now contains the preconditioned descent
        // direction Z = -K(HY-YA)

        // orthogonalize Z to Y
        // Z = Z - YY^T Z
        // A = Y^T * Z
        a.gemm('c','n',1.0,ct,c,0.0);

        // Z = Z - Y * A
        c.gemm('n','n',-1.0,ct,a,1.0);

        // orthogonalize Z: gram(Z)
        wf_.sd(ispin,ikp)->gram();

        // wf now contains Z, orthonormal
178 179 180
      }
    } // ikp
  } // ispin
Francois Gygi committed
181
  tmap_["jd_compute_z"].stop();
182

Francois Gygi committed
183
  tmap_["jd_hz"].start();
184 185 186 187 188 189 190 191
  // compute HZ
  const bool compute_hpsi = true;
  const bool compute_forces = false;
  const bool compute_stress = false;
  vector<vector<double> > fion;
  valarray<double> sigma;
  ef_.energy(compute_hpsi,dwf,compute_forces,fion,
             compute_stress,sigma);
Francois Gygi committed
192
  tmap_["jd_hz"].stop();
193

Francois Gygi committed
194
  tmap_["jd_blocks"].start();
195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215
  for ( int ispin = 0; ispin < wf_.nspin(); ispin++ )
  {
    for ( int ikp = 0; ikp < wf_.nkp(); ikp++ )
    {
      if ( wf_.sd(ispin,ikp)->basis().real() )
      {
        // wf now contains Z
        // dwf now contains HZ
        // compute blocks (Y,HY), (Y,HZ), (Z,HZ)

        // proxy real matrices c, cp
        DoubleMatrix c_proxy(wf_.sd(ispin,ikp)->c());
        DoubleMatrix c_proxy_t(wft_.sd(ispin,ikp)->c());
        DoubleMatrix cp_proxy(dwf.sd(ispin,ikp)->c());
        DoubleMatrix cp_proxy_t(dwft_.sd(ispin,ikp)->c());

        DoubleMatrix a(c_proxy.context(),c_proxy.n(),c_proxy.n(),
                       c_proxy.nb(),c_proxy.nb());
        DoubleMatrix h(c_proxy.context(),2*c_proxy.n(),2*c_proxy.n(),
                       c_proxy.nb(),c_proxy.nb());

Francois Gygi committed
216
        // (Y,HY)
217 218 219 220 221 222 223
        // factor 2.0 in next line: G and -G
        a.gemm('t','n',2.0,c_proxy_t,cp_proxy_t,0.0);
        // rank-1 update correction
        a.ger(-1.0,c_proxy_t,0,cp_proxy_t,0);
        // a contains (Y,HY), copy to h11 block
        h.getsub(a,a.m(),a.n(),0,0,0,0);

Francois Gygi committed
224
        // (Z,HY)
225 226 227 228 229 230
        // factor 2.0 in next line: G and -G
        a.gemm('t','n',2.0,c_proxy,cp_proxy_t,0.0);
        // rank-1 update correction
        a.ger(-1.0,c_proxy,0,cp_proxy_t,0);
        // a contains (Z,HY), copy to h21 block
        h.getsub(a,a.m(),a.n(),0,0,a.m(),0);
Francois Gygi committed
231 232

        // (Z,HZ)
233 234 235 236 237 238
        // factor 2.0 in next line: G and -G
        a.gemm('t','n',2.0,c_proxy,cp_proxy,0.0);
        // rank-1 update correction
        a.ger(-1.0,c_proxy,0,cp_proxy,0);
        // a contains (Z,HZ), copy to h22 block
        h.getsub(a,a.m(),a.n(),0,0,a.m(),a.n());
Francois Gygi committed
239

240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264
        // diagonalize h
        // Note: we only need the first n eigenvectors of the (2n x 2n) matrix
        valarray<double> w(h.m());
        // q is (2n,2n)
        DoubleMatrix q(h.context(),h.n(),h.n(),h.nb(),h.nb());
        h.syev('l',w,q);

        // compute the first n eigenvectors and store in wf
        // Y = Z Q21 (store result in dwf)
        // get Q21 in a
        a.getsub(q,a.n(),a.n(),a.n(),0);
        cp_proxy.gemm('n','n',1.0,c_proxy,a,0.0);

        // Y = Y Q11 (store result in wf)
        // get Q11 in a
        a.getsub(q,a.n(),a.n(),0,0);
        c_proxy.gemm('n','n',1.0,c_proxy_t,a,0.0);

        // add two contributions
        c_proxy += cp_proxy;

        // wf now contains the corrected eigenvectors
      } // if real
      else
      {
Francois Gygi committed
265 266 267 268 269 270 271 272 273 274 275 276 277 278 279 280 281 282 283 284 285 286 287 288 289 290 291 292 293 294 295 296 297 298 299 300 301 302 303 304 305 306 307 308 309 310 311 312 313 314 315
        // wf now contains Z
        // dwf now contains HZ
        // compute blocks (Y,HY), (Y,HZ), (Z,HZ)

        ComplexMatrix& c = wf_.sd(ispin,ikp)->c();
        ComplexMatrix& ct = wft_.sd(ispin,ikp)->c();
        ComplexMatrix& cp = dwf.sd(ispin,ikp)->c();
        ComplexMatrix& cpt = dwft_.sd(ispin,ikp)->c();

        ComplexMatrix a(c.context(),c.n(),c.n(),c.nb(),c.nb());
        ComplexMatrix h(c.context(),2*c.n(),2*c.n(),c.nb(),c.nb());

        // (Y,HY)
        // factor 2.0 in next line: G and -G
        a.gemm('c','n',1.0,ct,cpt,0.0);
        // a contains (Y,HY), copy to h11 block
        h.getsub(a,a.m(),a.n(),0,0,0,0);

        // (Z,HY)
        a.gemm('c','n',1.0,c,cpt,0.0);
        // a contains (Z,HY), copy to h21 block
        h.getsub(a,a.m(),a.n(),0,0,a.m(),0);

        // (Z,HZ)
        a.gemm('c','n',1.0,c,cp,0.0);
        // a contains (Z,HZ), copy to h22 block
        h.getsub(a,a.m(),a.n(),0,0,a.m(),a.n());

        // diagonalize h
        // Note: we only need the first n eigenvectors of the (2n x 2n) matrix
        valarray<double> w(h.m());
        // q is (2n,2n)
        ComplexMatrix q(h.context(),h.n(),h.n(),h.nb(),h.nb());
        h.heev('l',w,q);

        // compute the first n eigenvectors and store in wf
        // Y = Z Q21 (store result in dwf)
        // get Q21 in a
        a.getsub(q,a.n(),a.n(),a.n(),0);
        cp.gemm('n','n',1.0,c,a,0.0);

        // Y = Y Q11 (store result in wf)
        // get Q11 in a
        a.getsub(q,a.n(),a.n(),0,0);
        c.gemm('n','n',1.0,ct,a,0.0);

        // add two contributions
        c += cp;

        // wf now contains the corrected eigenvectors
      } // if real
316 317
    } // ikp
  } // ispin
Francois Gygi committed
318
  tmap_["jd_blocks"].stop();
319
}