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
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
from __future__ import print_function
import numpy as np
from mpi4py import MPI
from parallel import Parallel
try:
import h5py
use_hdf5 = True
except ImportError:
use_hdf5 = False
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
# Set the colormap
plt.rcParams['image.cmap'] = 'BrBG'
def write_field(field, parallel, step):
nx, ny = [i - 2 for i in field.shape]
nx_full = parallel.dims[0] * nx
ny_full = parallel.dims[1] * ny
if parallel.rank == 0:
field_full = np.zeros((nx_full, ny_full), float)
field_full[:nx, :ny] = field[1:-1, 1:-1]
# Receive data from other ranks
for p in range(1, parallel.size):
coords = parallel.comm.Get_coords(p)
idx = coords[0] * nx * ny_full + coords[1] * ny
parallel.comm.Recv((field_full.ravel()[idx:], 1,
parallel.subarraytype), source=p)
# Write out the data
plt.figure()
plt.gca().clear()
plt.imshow(field_full)
plt.axis('off')
plt.savefig('heat_{0:04d}.png'.format(step))
else:
# Send the data
parallel.comm.Ssend((field, 1, parallel.subarraytype), dest=0)
def read_field(filename):
# Read the initial temperature field from file
rank = MPI.COMM_WORLD.Get_rank()
if rank == 0:
with open(filename) as f:
line = f.readline().split()
shape = [int(n) for n in line[1:3]]
else:
shape = None
nx, ny = MPI.COMM_WORLD.bcast(shape, root=0)
parallel = Parallel(nx, ny)
nx_local = nx / parallel.dims[0]
ny_local = ny / parallel.dims[1]
field = np.zeros((nx_local + 2, ny_local + 2), dtype=float)
if parallel.rank == 0:
field_full = np.loadtxt(filename)
field[1:-1, 1:-1] = field_full[:nx_local, :ny_local]
# Send the data to other ranks
for p in range(1, parallel.size):
coords = parallel.comm.Get_coords(p)
idx = coords[0] * nx_local * ny + coords[1] * ny_local
parallel.comm.Send((field_full.ravel()[idx:], 1,
parallel.subarraytype), dest=p)
else:
parallel.comm.Recv((field, 1, parallel.subarraytype),
source=0)
# fo = np.empty((nx_local, ny_local), dtype=float)
# parallel.comm.Recv(fo, source=0)
# import pickle
# pickle.dump(fo, open('fo_f{0}.pckl'.format(parallel.rank), 'w'))
# Set the boundary values
field[0, :] = field[1, :]
field[-1, :] = field[-2, :]
field[:, 0] = field[:, 1]
field[:,-1] = field[:,-2]
field0 = field.copy()
return field, field0, parallel
def write_restart(field, parallel, iter):
nx, ny = [i - 2 for i in field.shape]
nx_full = parallel.dims[0] * nx
ny_full = parallel.dims[1] * ny
mode = MPI.MODE_CREATE | MPI.MODE_WRONLY
fh = MPI.File.Open(parallel.comm, "HEAT_RESTART.dat", mode)
if parallel.rank == 0:
fh.Write(np.array((nx_full, ny_full, iter)))
disp = 3 * 8 # Python ints are 8 bytes
fh.Set_view(disp, MPI.DOUBLE, parallel.filetype)
fh.Write_all((field, 1, parallel.restarttype))
fh.Close()
def read_restart():
mode = MPI.MODE_RDONLY
fh = MPI.File.Open(MPI.COMM_WORLD, "HEAT_RESTART.dat", mode)
buf = np.zeros(3, int)
fh.Read_all(buf)
nx, ny, iter = buf[0], buf[1], buf[2]
parallel = Parallel(nx, ny)
nx_local = nx / parallel.dims[0]
ny_local = ny / parallel.dims[1]
if parallel.rank == 0:
print("Restarting from iteration ", iter)
field = np.zeros((nx_local + 2, ny_local + 2), dtype=float)
disp = 3 * 8 # Python ints are 8 bytes
fh.Set_view(disp, MPI.DOUBLE, parallel.filetype)
fh.Read_all((field, 1, parallel.restarttype))
fh.Close()
field0 = field.copy()
return field, field0, parallel, iter