FV3-JEDI
fv3jedi_communication_mod.f90
Go to the documentation of this file.
1 module fv3jedi_communication_mod
2 
3 use mpi
4 use fv3jedi_geom_mod, only: fv3jedi_geom
5 use fv3jedi_kinds_mod, only: kind_real
6 
7 implicit none
8 
9 private
10 public gather_field, scatter_field
11 
12 !---------------------------------------------------------------------------------------------------
13 
14 contains
15 
16 !---------------------------------------------------------------------------------------------------
17 
18 subroutine gather_field(geom,comm,gproc,field_in,field_out)
19 
20  implicit none
21 
22  type(fv3jedi_geom), intent(in) :: geom !fv3-jedi geom
23  integer, intent(in) :: comm !MPI communicator
24  integer, intent(in) :: gproc !Processor on which to gather
25  real(kind=kind_real), intent(in) :: field_in(geom%isc:geom%iec,geom%jsc:geom%jec)
26  real(kind=kind_real), intent(out) :: field_out(1:geom%npx-1,1:geom%npy-1,6)
27 
28  integer :: comm_size, comm_rank, ierr, n, ji, jj, jc, npx_l, npy_l, npx_g, npy_g
29  integer, allocatable :: isc_l(:), iec_l(:), jsc_l(:), jec_l(:), til_l(:)
30  integer, allocatable :: counts(:), displs(:), vectorcounts(:), vectordispls(:)
31  real(kind=kind_real), allocatable :: vector_g(:), vector_l(:)
32 
33  ! Get comm size
34  ! -------------
35  call mpi_comm_size(comm, comm_size, ierr)
36  call mpi_comm_rank(comm, comm_rank, ierr)
37 
38  npx_l = geom%iec-geom%isc+1
39  npy_l = geom%jec-geom%jsc+1
40  npx_g = geom%npx-1
41  npy_g = geom%npy-1
42 
43  ! Array of counts and displacement
44  ! --------------------------------
45  allocate(counts(comm_size), displs(comm_size))
46  do n = 1,comm_size
47  displs(n) = n-1
48  counts(n) = 1
49  enddo
50 
51  ! Gather local dimensions
52  ! -----------------------
53  allocate(isc_l(comm_size), iec_l(comm_size), jsc_l(comm_size), jec_l(comm_size), til_l(comm_size))
54  call mpi_allgatherv(geom%isc, 1, mpi_int, isc_l, counts, displs, mpi_int, comm, ierr)
55  call mpi_allgatherv(geom%iec, 1, mpi_int, iec_l, counts, displs, mpi_int, comm, ierr)
56  call mpi_allgatherv(geom%jsc, 1, mpi_int, jsc_l, counts, displs, mpi_int, comm, ierr)
57  call mpi_allgatherv(geom%jec, 1, mpi_int, jec_l, counts, displs, mpi_int, comm, ierr)
58  call mpi_allgatherv(geom%ntile, 1, mpi_int, til_l, counts, displs, mpi_int, comm, ierr)
59  deallocate(counts,displs)
60 
61  ! Gather counts and displacement and pack local array into vector
62  ! ---------------------------------------------------------------
63 allocate(vectorcounts(comm_size), vectordispls(comm_size))
64 
65  n = 0
66  do jc = 1,comm_size
67  vectordispls(jc) = n
68  do jj = jsc_l(jc),jec_l(jc)
69  do ji = isc_l(jc),iec_l(jc)
70  n = n+1
71  enddo
72  enddo
73  vectorcounts(jc) = n - vectordispls(jc)
74  enddo
75 
76  allocate(vector_l(npx_l*npy_l))
77  n = 0
78  do jj = geom%jsc,geom%jec
79  do ji = geom%isc,geom%iec
80  n = n+1
81  vector_l(n) = field_in(ji,jj)
82  enddo
83  enddo
84 
85  ! Gather the full field
86  allocate(vector_g(npx_g*npy_g*6))
87  call mpi_gatherv( vector_l, npx_l*npy_l, mpi_double_precision, &
88  vector_g, vectorcounts, vectordispls, mpi_double_precision, &
89  gproc, comm, ierr)
90  deallocate(vector_l,vectorcounts,vectordispls)
91 
92 
93  ! Unpack global vector into array
94  ! -------------------------------
95 
96  if (comm_rank == gproc) then
97  n = 0
98  do jc = 1,comm_size
99  do jj = jsc_l(jc),jec_l(jc)
100  do ji = isc_l(jc),iec_l(jc)
101  n = n+1
102  field_out(ji,jj,til_l(jc)) = vector_g(n)
103  enddo
104  enddo
105  enddo
106  endif
107  deallocate(isc_l, iec_l, jsc_l, jec_l)
108 
109  deallocate(vector_g)
110 
111 end subroutine gather_field
112 
113 !---------------------------------------------------------------------------------------------------
114 
115 subroutine scatter_field(geom,comm,gproc,field_in,field_out)
116 
117  implicit none
118 
119  type(fv3jedi_geom), intent(in) :: geom !fv3-jedi geom
120  integer, intent(in) :: comm !MPI communicator
121  integer, intent(in) :: gproc !Processor on which to gather
122  real(kind=kind_real), intent(in) :: field_in(1:geom%npx-1,1:geom%npy-1,6)
123  real(kind=kind_real), intent(out) :: field_out(geom%isc:geom%iec,geom%jsc:geom%jec)
124 
125  integer :: comm_size, comm_rank, ierr, n, ji, jj, jc, npx_l, npy_l, npx_g, npy_g
126  integer, allocatable :: isc_l(:), iec_l(:), jsc_l(:), jec_l(:), til_l(:)
127  integer, allocatable :: counts(:), displs(:), vectorcounts(:), vectordispls(:)
128  real(kind=kind_real), allocatable :: vector_g(:), vector_l(:)
129 
130  ! Get comm size
131  ! -------------
132  call mpi_comm_size(comm, comm_size, ierr)
133  call mpi_comm_rank(comm, comm_rank, ierr)
134 
135  npx_l = geom%iec-geom%isc+1
136  npy_l = geom%jec-geom%jsc+1
137  npx_g = geom%npx-1
138  npy_g = geom%npy-1
139 
140  ! Array of counts and displacement
141  ! --------------------------------
142  allocate(counts(comm_size), displs(comm_size))
143  do n = 1,comm_size
144  displs(n) = n-1
145  counts(n) = 1
146  enddo
147 
148  ! Gather local dimensions
149  ! -----------------------
150  allocate(isc_l(comm_size), iec_l(comm_size), jsc_l(comm_size), jec_l(comm_size), til_l(comm_size))
151  call mpi_allgatherv(geom%isc, 1, mpi_int, isc_l, counts, displs, mpi_int, comm, ierr)
152  call mpi_allgatherv(geom%iec, 1, mpi_int, iec_l, counts, displs, mpi_int, comm, ierr)
153  call mpi_allgatherv(geom%jsc, 1, mpi_int, jsc_l, counts, displs, mpi_int, comm, ierr)
154  call mpi_allgatherv(geom%jec, 1, mpi_int, jec_l, counts, displs, mpi_int, comm, ierr)
155  call mpi_allgatherv(geom%ntile, 1, mpi_int, til_l, counts, displs, mpi_int, comm, ierr)
156  deallocate(counts,displs)
157 
158 
159  ! Pack whole tile array into vector
160  ! ---------------------------------
161  allocate(vector_g(npx_g*npy_g*6))
162  allocate(vectorcounts(comm_size), vectordispls(comm_size))
163 
164  n = 0
165  do jc = 1,comm_size
166  vectordispls(jc) = n
167  do jj = jsc_l(jc),jec_l(jc)
168  do ji = isc_l(jc),iec_l(jc)
169  n = n+1
170  enddo
171  enddo
172  vectorcounts(jc) = n - vectordispls(jc)
173  enddo
174 
175  if (comm_rank == gproc) then
176  n = 0
177  do jc = 1,comm_size
178  do jj = jsc_l(jc),jec_l(jc)
179  do ji = isc_l(jc),iec_l(jc)
180  n = n+1
181  vector_g(n) = field_in(ji,jj,til_l(jc))
182  enddo
183  enddo
184  enddo
185  endif
186 
187  deallocate(isc_l, iec_l, jsc_l, jec_l)
188 
189  ! Scatter tile array to processors
190  allocate(vector_l(npx_l*npy_l))
191 
192  call mpi_scatterv( vector_g, vectorcounts, vectordispls, mpi_double_precision, &
193  vector_l, npx_l*npy_l, mpi_double_precision, &
194  gproc, comm, ierr )
195 
196  deallocate(vector_g,vectorcounts,vectordispls)
197 
198  ! Unpack local vector into array
199  ! ------------------------------
200  n = 0
201  do jj = geom%jsc,geom%jec
202  do ji = geom%isc,geom%iec
203  n = n+1
204  field_out(ji,jj) = vector_l(n)
205  enddo
206  enddo
207 
208  deallocate(vector_l)
209 
210 end subroutine scatter_field
211 
212 !---------------------------------------------------------------------------------------------------
213 
214 end module fv3jedi_communication_mod
fv3jedi_geom_mod
Fortran module handling geometry for the FV3 model.
Definition: fv3jedi_geom_mod.f90:8
fv3jedi_communication_mod
Definition: fv3jedi_communication_mod.f90:1
fv3jedi_geom_mod::fv3jedi_geom
Fortran derived type to hold geometry data for the FV3JEDI model.
Definition: fv3jedi_geom_mod.f90:46
fv3jedi_communication_mod::scatter_field
subroutine, public scatter_field(geom, comm, gproc, field_in, field_out)
Definition: fv3jedi_communication_mod.f90:116
fv3jedi_communication_mod::gather_field
subroutine, public gather_field(geom, comm, gproc, field_in, field_out)
Definition: fv3jedi_communication_mod.f90:19
fv3jedi_kinds_mod::kind_real
integer, parameter, public kind_real
Definition: fv3jedi_kinds_mod.f90:14
fv3jedi_kinds_mod
Definition: fv3jedi_kinds_mod.f90:6