1!
2! Copyright (C) 1996-2016	The SIESTA group
3!  This file is distributed under the terms of the
4!  GNU General Public License: see COPYING in the top directory
5!  or http://www.gnu.org/copyleft/gpl.txt.
6! See Docs/Contributors.txt for a list of contributors.
7!
8module iowfs_netcdf
9
10#ifdef CDF
11use netcdf
12use sys, only: die
13use precision, only: dp
14use parallel, only: Node, Nodes
15use alloc, only: re_alloc, de_alloc
16
17#ifdef MPI
18use mpi_siesta
19use parallelsubs, only : GLobalToLocalOrb, WhichNodeOrb
20#endif
21
22implicit none
23
24! These module variables should be put into a derived type, and maybe
25! not all of them are really necessary
26!
27integer ncid, norbs_id, nspin_id, nk_id, ncomplex_id, nvectors_id
28integer wf_id    , eigval_id
29
30public :: setup_wfs_netcdf_file, write_wfs_netcdf, get_wfs_netcdf
31public :: open_wfs_netcdf_file, close_wfs_netcdf_file
32public :: get_wfs_block_netcdf
33
34private
35
36CONTAINS
37
38subroutine setup_wfs_netcdf_file( norbs, nk, nspin, nvectors)
39
40!  To be called only by the master node
41
42      integer, intent(in)   ::    norbs  ! Number of atomic orbitals
43      integer, intent(in)   ::    nspin  ! Number of spins
44      integer, intent(in)   ::    nk     ! Number of kpoints
45      integer, intent(in)   ::    nvectors   ! Number of eigenvectors stored per spin and k-point
46
47      call check( nf90_create('WFS.nc',NF90_CLOBBER,ncid))
48!
49!     Dimensions
50!
51      call check( nf90_def_dim(ncid,'norbs',norbs,norbs_id))  !"Number of basis orbitals"
52      call check( nf90_def_dim(ncid,'nspin',nspin,nspin_id))   !"Number of spin components"
53      call check( nf90_def_dim(ncid,'nk',nk,nk_id))     !"Number of k-points"
54      call check( nf90_def_dim(ncid,'nvectors',nvectors,nvectors_id))     !"Number of k-points"
55      call check( nf90_def_dim(ncid,'ncomplex',2,ncomplex_id))     !"Real/imaginary"
56
57!
58!     Variables
59!
60      call check( nf90_def_var(ncid,'eigval',nf90_float,(/nvectors_id,nk_id,nspin_id/),eigval_id))
61      call check( nf90_put_att(ncid,eigval_id,'Description',"Eigenvalues in Ryd"))
62      call check( nf90_def_var(ncid,'wf',nf90_float,    &
63                            (/ncomplex_id,norbs_id,nvectors_id,nk_id,nspin_id/),wf_id))
64      call check( nf90_put_att(ncid,wf_id,'Description',"Wavefunctions"))
65
66      call check( nf90_enddef(ncid))
67
68end subroutine setup_wfs_netcdf_file
69
70subroutine write_wfs_netcdf( norbs, no_l, ik, ispin, psi, nvectors, eigvals )
71
72use precision, only : dp
73
74integer, intent(in)   ::    norbs  ! Total number of orbitals in unit cell (components of wf)
75integer, intent(in)   ::    no_l   ! Number of eigenvectors in this node
76integer, intent(in)   ::    ik     ! k-point index
77integer, intent(in)   ::    ispin  ! spin index
78integer, intent(in)   ::    nvectors ! total number of eigenvectors to be stored for this k-point and spin
79
80real(dp), intent(in)  :: psi(2,norbs,no_l)    ! wave-functions stored in this node
81real(dp), intent(in)  :: eigvals(nvectors)    ! wave-functions stored in this node
82
83#ifdef MPI
84      real(dp), dimension(:,:), pointer  ::  psi_buf  => null()
85      integer  :: MPIerror, stat(MPI_STATUS_SIZE), count, BNode, tag, ivec_local
86#endif
87
88      integer  :: ivec
89
90#ifdef MPI
91
92if (Node == 0) then
93   nullify( psi_buf )
94   call re_alloc( psi_buf, 1, 2, 1, norbs, name='psi_buf', routine='iowfs_netcdf' )
95   call check( nf90_put_var(ncid,eigval_id,eigvals(1:nvectors),         &
96                            start=(/1, ik, ispin/),  count=(/nvectors, 1, 1 /) ) )
97endif
98
99    do ivec = 1, nvectors
100       call WhichNodeOrb(ivec,Nodes,Bnode)
101       call GlobalToLocalOrb(ivec,BNode,Nodes,ivec_local)
102       tag = ivec
103
104       if (Node == 0) then
105          if (BNode == 0) then
106             psi_buf(1:2,1:norbs) = psi(1:2,1:norbs,ivec_local)     ! Could do with pointer
107          else
108             call MPI_Recv(psi_buf(1,1),2*norbs,MPI_double_precision,BNode,tag,  &
109                                                        MPI_Comm_World,stat,mpierror)
110          endif
111          call check( nf90_put_var(ncid,wf_id,psi_buf(1:2,1:norbs),         &
112                                       start=(/1, 1, ivec, ik, ispin/),        &
113                                       count=(/2, norbs, 1, 1, 1 /) ) )
114       else      ! Non-master nodes
115
116          if (Node == BNode) then
117             call MPI_Send(psi(1,1,ivec_local),2*norbs, &
118                           MPI_double_precision,0,tag,MPI_Comm_World,mpierror)
119          endif
120
121       endif
122
123    enddo  ! loop over vectors
124
125    call de_alloc(psi_buf,name="psi_buf", routine="iowfs_netcdf")
126
127#else
128    call check( nf90_put_var(ncid,eigval_id,eigvals(1:nvectors),         &
129                             start=(/1, ik, ispin/),  count=(/nvectors, 1, 1 /) ) )
130    do ivec = 1, nvectors
131       call check( nf90_put_var(ncid, wf_id, psi(1:2,1:norbs,ivec),  &
132                              start = (/1, 1, ivec, ik, ispin /), &
133                              count = (/2, norbs, 1, 1, 1 /) ))
134    enddo
135#endif
136
137if (Node == 0) then
138   call check( nf90_sync(ncid))
139endif
140
141end subroutine write_wfs_netcdf
142
143subroutine close_wfs_netcdf_file()
144if (Node == 0) then
145   call check( nf90_close(ncid))
146endif
147end subroutine close_wfs_netcdf_file
148
149subroutine open_wfs_netcdf_file()
150if (Node == 0) then
151   call check( nf90_open("WFS.nc",NF90_NOWRITE,ncid))
152   call check( nf90_inq_varid(ncid,"wf",wf_id))
153endif
154end subroutine open_wfs_netcdf_file
155
156subroutine get_wfs_netcdf(ivec,ik,ispin,norbs,psi_aux)
157
158use precision, only : dp
159
160integer, intent(in)   ::    ivec   ! Eigenvector index
161integer, intent(in)   ::    ik     ! k-point index
162integer, intent(in)   ::    ispin  ! spin index
163integer, intent(in)   ::    norbs  ! Total number of orbitals in unit cell (components of wf)
164
165real(dp), intent(out)  :: psi_aux(2,norbs)    ! wave-function
166
167#ifdef MPI
168      integer  :: MPIerror
169#endif
170
171if (Node == 0) then
172       call check( nf90_get_var(ncid, wf_id, psi_aux(1:2,1:norbs),  &
173                              start = (/1, 1, ivec, ik, ispin /), &
174                              count = (/2, norbs, 1, 1, 1 /) ))
175endif
176#ifdef MPI
177call MPI_Bcast(psi_aux(1,1),2*norbs,MPI_Double_Precision,0,MPI_Comm_World, MPIerror)
178#endif
179end subroutine get_wfs_netcdf
180
181subroutine get_wfs_block_netcdf(ivec_start,nvecs,ik,ispin,norbs,psi_block)
182
183use precision, only : dp
184
185integer, intent(in)   ::    ivec_start ! Index of first eigenvector to get
186integer, intent(in)   ::    nvecs  ! Number of eigenvectors to get
187integer, intent(in)   ::    ik     ! k-point index
188integer, intent(in)   ::    ispin  ! spin index
189integer, intent(in)   ::    norbs  ! Total number of orbitals in unit cell (components of wf)
190
191real(dp), intent(out)  :: psi_block(2,norbs,nvecs)    ! wave-function block
192
193#ifdef MPI
194      integer  :: MPIerror
195#endif
196
197if (Node == 0) then
198       call check( nf90_get_var(ncid, wf_id, psi_block(1:2,1:norbs,1:nvecs),  &
199                              start = (/1, 1, ivec_start, ik, ispin /), &
200                              count = (/2, norbs, nvecs, 1, 1 /) ))
201endif
202
203#ifdef MPI
204call MPI_Bcast(psi_block(1,1,1),2*norbs*nvecs,MPI_Double_Precision,0,MPI_Comm_World, MPIerror)
205#endif
206
207end subroutine get_wfs_block_netcdf
208
209subroutine check(code)
210use netcdf, only: nf90_noerr, nf90_strerror
211integer, intent(in) :: code
212if (code /= nf90_noerr) call die("netCDF error: " // NF90_STRERROR(code))
213end subroutine check
214
215
216#endif
217end module iowfs_netcdf
218