PlotCmd.C 15.4 KB
Newer Older
Francois Gygi committed
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 29 30 31
////////////////////////////////////////////////////////////////////////////////
//
// 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/>.
//
////////////////////////////////////////////////////////////////////////////////
//
// PlotCmd.C:
//
////////////////////////////////////////////////////////////////////////////////

#include "PlotCmd.h"
#include "isodate.h"
#include "release.h"

#include <iostream>
#include <fstream>
#include <sstream>
#include <complex>

#include "Context.h"
#include "Sample.h"
#include "SampleReader.h"
#include "Basis.h"
32
#include "EnergyFunctional.h"
Francois Gygi committed
33 34 35 36 37 38 39 40 41 42 43 44
#include "FourierTransform.h"
#include "SlaterDet.h"
#include "Matrix.h"
#include "Species.h"
#include "Atom.h"
#include "ChargeDensity.h"

using namespace std;

////////////////////////////////////////////////////////////////////////////////
int PlotCmd::action(int argc, char **argv)
{
45
  string usage("  Use: plot filename\n"
46 47
               "       plot -density  [-spin {1|2}] filename\n"
               "       plot -vlocal   [-spin {1|2}] filename\n"
48 49
               "       plot -wf n [-spin {1|2}] filename\n"
               "       plot -wfs nmin nmax [-spin {1|2}] filename");
Francois Gygi committed
50 51 52 53

  // parse arguments
  // plot filename               : plot atoms in xyz format
  // plot -density filename      : plot atoms and density in cube format
54
  // plot -vlocal  filename      : plot atoms and vlocal in cube format
Francois Gygi committed
55 56
  // plot -wf <n> filename       : plot atoms and wf <n> in cube format
  // plot -wf <n1> <n2> filename : plot atoms and wfs <n1> to <n2> in cube fmt
57 58
  // spin: 1 = first spin (up), 2 = second spin (down)
  if ( (argc < 2) || (argc > 7) )
Francois Gygi committed
59 60 61 62 63 64
  {
    if ( ui->onpe0() )
      cout << usage << endl;
    return 1;
  }

65 66
  bool plot_atoms = true;
  bool xyz = true;
Francois Gygi committed
67
  bool plot_density = false;
68
  bool plot_vlocal = false;
Francois Gygi committed
69
  bool plot_wf = false;
Francois Gygi committed
70
  bool plot_wfs = false;
Francois Gygi committed
71
  int nmin,nmax,nwf;
72 73 74 75 76
  // ispin = 0: plot both spins
  // ispin = 1: plot first spin
  // ispin = 2: plot second spin
  int ispin = 0;

Francois Gygi committed
77 78
  string filename;

79 80 81
  // process arguments
  int iarg = 1;
  while ( iarg < argc )
Francois Gygi committed
82
  {
83
    if ( !strcmp(argv[iarg],"-density") )
Francois Gygi committed
84
    {
85 86
      plot_density = true;
      xyz = false;
Francois Gygi committed
87
    }
88 89 90 91 92
    else if ( !strcmp(argv[iarg],"-vlocal") )
    {
      plot_vlocal = true;
      xyz = false;
    }
93
    else if ( !strcmp(argv[iarg],"-wf") )
Francois Gygi committed
94
    {
95 96 97 98 99 100 101
      plot_wf = true;
      xyz = false;
      // process argument: n
      iarg++;
      if ( iarg == argc )
      {
        if ( ui->onpe0() )
Francois Gygi committed
102
        cout << usage << endl;
103 104 105 106 107 108 109 110 111 112 113 114
        return 1;
      }
      nmin = atoi(argv[iarg]) - 1;
      nmax = nmin;
      nwf = 1;
      if ( nmin < 0 || nmax >= s->wf.nst() || nmin > nmax )
      {
        if ( ui->onpe0() )
          cout << " nmin or nmax incompatible with nst="
               << s->wf.nst() << endl;
        return 1;
      }
Francois Gygi committed
115
    }
116
    else if ( !strcmp(argv[iarg],"-wfs") )
Francois Gygi committed
117
    {
Francois Gygi committed
118
      plot_wfs = true;
119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145
      xyz = false;
      // process argument: nmin
      iarg++;
      if ( iarg==argc )
      {
        if ( ui->onpe0() )
        cout << usage << endl;
        return 1;
      }
      nmin = atoi(argv[iarg]) - 1;
      // process argument: nmax
      iarg++;
      if ( iarg==argc )
      {
        if ( ui->onpe0() )
        cout << usage << endl;
        return 1;
      }
      nmax = atoi(argv[iarg]) - 1;
      nwf = nmax-nmin+1;
      if ( nmin < 0 || nmax >= s->wf.nst() || nmin > nmax )
      {
        if ( ui->onpe0() )
          cout << " nmin or nmax incompatible with nst="
               << s->wf.nst() << endl;
        return 1;
      }
Francois Gygi committed
146
    }
147
    else if ( !strcmp(argv[iarg],"-spin") )
Francois Gygi committed
148
    {
Francois Gygi committed
149
      if ( !(plot_density || plot_wf || plot_wfs) )
150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175
      {
        if ( ui->onpe0() )
          cout << usage << endl;
        return 1;
      }
      if ( s->wf.nspin() != 2 )
      {
        if ( ui->onpe0() )
          cout << "nspin = 1, cannot select spin" << endl;
        return 1;
      }
      // process argument: ispin
      iarg++;
      if ( iarg==argc )
      {
        if ( ui->onpe0() )
          cout << usage << endl;
        return 1;
      }
      ispin = atoi(argv[iarg]);
      if ( ispin < 1 || ispin > 2 )
      {
        if ( ui->onpe0() )
          cout << " spin must be 1 or 2" <<  endl;
        return 1;
      }
Francois Gygi committed
176
    }
177
    else
Francois Gygi committed
178
    {
179 180
      // argument must be the file name
      filename = argv[iarg];
Francois Gygi committed
181
    }
182 183 184 185 186 187

    iarg++;

  } // while iarg

  // Must specify spin if plotting wave functions when nspin==2
188
  if ( s->wf.nspin()==2 && (plot_vlocal||plot_wf||plot_wfs) && ispin==0 )
189 190 191 192
  {
    if ( ui->onpe0() )
      cout << " must use -spin if nspin==2" <<  endl;
    return 1;
Francois Gygi committed
193 194 195 196
  }

  ofstream os;
  vector<double> tmpr;
197 198 199
  int np0 = 0;
  int np1 = 0;
  int np2 = 0;
Francois Gygi committed
200 201 202 203 204 205

  const Context& ctxt = *s->wf.spincontext();
  if ( plot_density )
  {
    ChargeDensity cd(s->wf);
    cd.update_density();
206
    tmpr.resize(cd.vft()->np012());
Francois Gygi committed
207 208 209
    np0 = cd.vft()->np0();
    np1 = cd.vft()->np1();
    np2 = cd.vft()->np2();
210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230
    if ( s->wf.nspin() == 1 )
    {
      for ( int i = 0; i < cd.vft()->np012loc(); i++ )
        tmpr[i] = cd.rhor[0][i];
    }
    else
    {
      if ( ispin == 0 )
      {
        // plot both spins
        for ( int i = 0; i < cd.vft()->np012loc(); i++ )
          tmpr[i] = cd.rhor[0][i] + cd.rhor[1][i];
      }
      else
      {
        // plot one spin only
        // ispin==1 or ispin==2
        for ( int i = 0; i < cd.vft()->np012loc(); i++ )
          tmpr[i] = cd.rhor[ispin-1][i];
      }
    }
Francois Gygi committed
231 232 233 234 235 236 237 238 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 265 266 267 268 269 270 271 272 273 274 275 276 277 278 279 280

    // send blocks of tmpr to pe0
    // send from first context column only
    for ( int i = 0; i < ctxt.nprow(); i++ )
    {
      bool iamsending = ctxt.mycol() == 0 && i == ctxt.myrow();

      // send size of tmpr block
      int size=-1;
      if ( ctxt.onpe0() )
      {
        if ( iamsending )
        {
          // sending to self, size not needed
        }
        else
          ctxt.irecv(1,1,&size,1,i,0);
      }
      else
      {
        if ( iamsending )
        {
          size = cd.vft()->np012loc();
          ctxt.isend(1,1,&size,1,0,0);
        }
      }

      // send tmpr block
      if ( ctxt.onpe0() )
      {
        if ( iamsending )
        {
          // do nothing, data is already in place
        }
        else
        {
          int istart = cd.vft()->np0() * cd.vft()->np1() *
                       cd.vft()->np2_first(i);
          ctxt.drecv(size,1,&tmpr[istart],1,i,0);
        }
      }
      else
      {
        if ( iamsending )
        {
          ctxt.dsend(size,1,&tmpr[0],1,0,0);
        }
      }
    }
  } // plot_density
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 316 317 318 319 320 321 322 323 324 325 326 327 328 329 330 331 332 333 334 335 336 337 338 339 340 341 342 343 344 345 346 347 348 349 350 351 352 353 354 355 356 357 358 359 360 361 362 363 364 365 366 367 368 369
  else if ( plot_vlocal )
  {
    ChargeDensity cd(s->wf);
    EnergyFunctional ef(*s,cd);
    cd.update_density();
    cd.update_rhor();
    bool compute_stress = false;
    ef.update_vhxc(compute_stress);

    FourierTransform *vft = cd.vft();
    Basis *vbasis = cd.vbasis();

    tmpr.resize(cd.vft()->np012());
    np0 = cd.vft()->np0();
    np1 = cd.vft()->np1();
    np2 = cd.vft()->np2();
    if ( s->wf.nspin() == 1 )
    {
      for ( int i = 0; i < cd.vft()->np012loc(); i++ )
        tmpr[i] = ef.v_r[0][i];
    }
    else
    {
      if ( ispin == 0 )
      {
        // should not get here
        // ispin must be set if nspin==2
        if ( ui->onpe0() )
          cout << usage << endl;
        return 1;
      }
      else
      {
        // plot one spin only
        // ispin==1 or ispin==2
        for ( int i = 0; i < cd.vft()->np012loc(); i++ )
          tmpr[i] = ef.v_r[ispin-1][i];
      }
    }

    // send blocks of tmpr to pe0
    // send from first context column only
    for ( int i = 0; i < ctxt.nprow(); i++ )
    {
      bool iamsending = ctxt.mycol() == 0 && i == ctxt.myrow();

      // send size of tmpr block
      int size=-1;
      if ( ctxt.onpe0() )
      {
        if ( iamsending )
        {
          // sending to self, size not needed
        }
        else
          ctxt.irecv(1,1,&size,1,i,0);
      }
      else
      {
        if ( iamsending )
        {
          size = cd.vft()->np012loc();
          ctxt.isend(1,1,&size,1,0,0);
        }
      }

      // send tmpr block
      if ( ctxt.onpe0() )
      {
        if ( iamsending )
        {
          // do nothing, data is already in place
        }
        else
        {
          int istart = cd.vft()->np0() * cd.vft()->np1() *
                       cd.vft()->np2_first(i);
          ctxt.drecv(size,1,&tmpr[istart],1,i,0);
        }
      }
      else
      {
        if ( iamsending )
        {
          ctxt.dsend(size,1,&tmpr[0],1,0,0);
        }
      }
    }
  } // plot_vlocal
Francois Gygi committed
370
  else if ( plot_wf || plot_wfs )
Francois Gygi committed
371
  {
Francois Gygi committed
372
    // compute wf or wf squared and store in tmpr
Francois Gygi committed
373 374 375 376 377 378 379 380 381 382 383 384 385 386 387 388 389 390 391
    if ( ctxt.onpe0() )
    {
      ctxt.ibcast_send(1,1,&nwf,1);
      ctxt.ibcast_send(1,1,&nmin,1);
      ctxt.ibcast_send(1,1,&nmax,1);
    }
    else
    {
      ctxt.ibcast_recv(1,1,&nwf,1,0,0);
      ctxt.ibcast_recv(1,1,&nmin,1,0,0);
      ctxt.ibcast_recv(1,1,&nmax,1,0,0);
    }

    if ( nwf > 0 && s->wf.nst() == 0 )
    {
      cout << " no states in sample" << endl;
      return 1;
    }

392 393 394 395 396 397 398 399 400 401 402 403 404 405 406 407 408 409 410 411
    int isp_min, isp_max;
    SlaterDet *sdp;
    if ( ispin == 0 )
    {
      // -spin was not specified:
      if ( s->wf.nspin() == 1 )
      {
        isp_min = 0; isp_max = 0;
      }
      else
      {
        isp_min = 0; isp_max = 1;
      }
    }
    else
    {
      isp_min = ispin-1; isp_max = ispin-1;
    }

    const Basis& basis = s->wf.sd(0,0)->basis();
Francois Gygi committed
412 413 414 415 416 417 418 419
    np0 = basis.np(0);
    np1 = basis.np(1);
    np2 = basis.np(2);
    FourierTransform ft(basis,np0,np1,np2);

    vector<complex<double> > wftmp(ft.np012loc());
    vector<double> wftmpr(ft.np012());
    tmpr.resize(ft.np012());
Francois Gygi committed
420
    tmpr.assign(ft.np012(),0.0);
Francois Gygi committed
421

422 423 424 425
    for ( int isp = isp_min; isp <= isp_max; isp++ )
    {
      sdp = s->wf.sd(isp,0);
      const ComplexMatrix& c = sdp->c();
Francois Gygi committed
426

427
      for ( int n = nmin; n <= nmax; n++ )
Francois Gygi committed
428
      {
429
        if ( n >= s->wf.nst(isp) )
Francois Gygi committed
430
        {
431 432 433 434
          if ( ui->onpe0() )
            cout << "invalid wave function index: " << n+1
                 << " > nst(ispin)" << endl;
          return 1;
Francois Gygi committed
435 436
        }

437
        // compute real-space wavefunction
Francois Gygi committed
438

439 440
        // transform wf on ctxt.mycol() hosting state n
        if ( c.pc(n) == c.context().mycol() )
Francois Gygi committed
441
        {
442 443 444 445 446 447 448
          //os << " state " << n << " is stored on column "
          //     << ctxt_.mycol() << " local index: " << c_.y(n) << endl;
          int nloc = c.y(n); // local index
          ft.backward(c.cvalptr(c.mloc()*nloc),&wftmp[0]);

          double *a = (double*) &wftmp[0];
          if ( basis.real() )
Francois Gygi committed
449
          {
450 451 452
            // real function: plot wf
            for ( int i = 0; i < ft.np012loc(); i++ )
              wftmpr[i] = a[2*i];
Francois Gygi committed
453 454 455
          }
          else
          {
456 457 458
            // complex function: plot modulus
            for ( int i = 0; i < ft.np012loc(); i++ )
              wftmpr[i] = sqrt(a[2*i]*a[2*i] + a[2*i+1]*a[2*i+1]);
Francois Gygi committed
459 460 461
          }
        }

462 463
        // send blocks of wftmpr to pe0
        for ( int i = 0; i < c.context().nprow(); i++ )
Francois Gygi committed
464
        {
465 466 467 468 469 470
          bool iamsending = c.pc(n) == c.context().mycol() &&
                            i == c.context().myrow();

          // send size of wftmpr block
          int size=-1;
          if ( c.context().onpe0() )
Francois Gygi committed
471
          {
472 473 474 475 476 477
            if ( iamsending )
            {
              // sending to self, size not needed
            }
            else
              c.context().irecv(1,1,&size,1,i,c.pc(n));
Francois Gygi committed
478 479 480
          }
          else
          {
481 482 483 484 485
            if ( iamsending )
            {
              size = ft.np012loc();
              c.context().isend(1,1,&size,1,0,0);
            }
Francois Gygi committed
486
          }
487 488 489

          // send wftmpr block
          if ( c.context().onpe0() )
Francois Gygi committed
490
          {
491 492 493 494 495 496 497 498 499 500 501 502 503 504 505 506
            if ( iamsending )
            {
              // do nothing, data is already in place
            }
            else
            {
              int istart = ft.np0() * ft.np1() * ft.np2_first(i);
              c.context().drecv(size,1,&wftmpr[istart],1,i,c.pc(n));
            }
          }
          else
          {
            if ( iamsending )
            {
              c.context().dsend(size,1,&wftmpr[0],1,0,0);
            }
Francois Gygi committed
507 508 509
          }
        }

510 511
        // process the data on task 0
        if ( c.context().onpe0() )
Francois Gygi committed
512
        {
513
          // wftmpr is now complete on task 0
Francois Gygi committed
514
          if ( plot_wfs )
Francois Gygi committed
515
          {
Francois Gygi committed
516
            // multiple wfs, accumulate square
517 518
            for ( int i = 0; i < ft.np012(); i++ )
            {
Francois Gygi committed
519
              tmpr[i] += wftmpr[i]*wftmpr[i];
520
            }
Francois Gygi committed
521
          }
522
          else
Francois Gygi committed
523
          {
Francois Gygi committed
524
            // plot individual wf
525 526
            for ( int i = 0; i < ft.np012(); i++ )
            {
Francois Gygi committed
527
              tmpr[i] = wftmpr[i];
528
            }
Francois Gygi committed
529 530
          }
        }
531 532
      } // for n
    } // for isp
Francois Gygi committed
533
  } // if plot_wf || plot_wfs
Francois Gygi committed
534 535 536 537 538 539 540 541 542 543 544 545 546 547 548 549 550 551 552 553 554 555 556 557 558 559 560 561 562 563

  // tmpr now contains the function to plot on task 0

  if ( ctxt.onpe0() )
    os.open(filename.c_str());

  if ( plot_atoms )
  {
    if ( ctxt.onpe0() )
    {
      if ( xyz )
      {
        const double a0 = 0.529177;
        int natoms = s->atoms.size();
        os << natoms << endl;
        os << "Created " << isodate() << " by qbox-" << release() << endl;
        const int nsp = s->atoms.nsp();
        for ( int is = 0; is < nsp; is++ )
        {
          Species* sp = s->atoms.species_list[is];
          string symbol = sp->symbol();
          const int na = s->atoms.na(is);
          for ( int ia = 0; ia < na; ia++ )
          {
            Atom *ap = s->atoms.atom_list[is][ia];
            os << setprecision(5);
            os << symbol << " " << a0*ap->position() << endl;
          }
        }
      }
Francois Gygi committed
564
      else
Francois Gygi committed
565 566 567 568 569 570 571 572 573 574 575 576 577 578 579 580 581 582 583 584 585 586 587 588 589 590 591 592 593 594 595 596
      {
        // write header and atoms
        os << "Created " << isodate() << " by qbox-" << release() << endl;
        os << endl;

        int natoms = s->atoms.size();
        D3vector a0 = s->atoms.cell().a(0);
        D3vector a1 = s->atoms.cell().a(1);
        D3vector a2 = s->atoms.cell().a(2);
        os << natoms << " " << -0.5*(a0+a1+a2) << endl;

        // write unit cell
        os << np0 << " " << a0/np0 << endl;
        os << np1 << " " << a1/np1 << endl;
        os << np2 << " " << a2/np2 << endl;
        const int nsp = s->atoms.nsp();
        for ( int is = 0; is < nsp; is++ )
        {
          Species* sp = s->atoms.species_list[is];
          const int z = sp->atomic_number();
          const int na = s->atoms.na(is);
          for ( int ia = 0; ia < na; ia++ )
          {
            Atom *ap = s->atoms.atom_list[is][ia];
            os << setprecision(5);
            os << z << " " << ((double) z) << " " << ap->position() << endl;
          }
        }
      }
    }
  } // if plot_atoms

597
  if ( plot_density || plot_vlocal || plot_wf || plot_wfs )
Francois Gygi committed
598 599 600 601 602 603 604 605 606 607 608 609 610 611 612 613 614 615 616 617 618 619 620 621
  {
    // process the function in tmpr
    if ( ctxt.onpe0() )
    {
      os.setf(ios::scientific,ios::floatfield);
      os << setprecision(5);
      for ( int i = 0; i < np0; i++ )
      {
        const int ip = (i + np0/2 ) % np0;
        for ( int j = 0; j < np1; j++ )
        {
          const int jp = (j + np1/2 ) % np1;
          for ( int k = 0; k < np2; k++ )
          {
            const int kp = (k + np2/2 ) % np2;
            os << setw(13) << tmpr[ip+np0*(jp+np1*kp)];
            if ( ( k % 6 ) == 5 )
              os << endl;
          }
          if ( ( np2 % 6 ) != 0 )
            os << endl;
        }
      }
    }
622
  } // if plot_density || ...
Francois Gygi committed
623 624 625 626 627

  os.close();

  return 0;
}