|
8 | 8 |
|
9 | 9 | import argparse
|
10 | 10 |
|
| 11 | + |
11 | 12 | def split(num):
|
12 | 13 | for a in range(math.isqrt(num), 0, -1):
|
13 | 14 | if num % a == 0:
|
@@ -53,14 +54,17 @@ def main():
|
53 | 54 | py = rank // xr
|
54 | 55 | print(f"Rank {rank}/{size} has partition ({px}, {py})/({xr}, {yr})")
|
55 | 56 | if rank == 0:
|
56 |
| - print(f"Each of {size} partitions has node size {pnx}x{pny} = {pnx*pny} for a total of {nx*ny} nodes on the base") |
57 |
| - |
| 57 | + print( |
| 58 | + f"Each of {size} partitions has node size {pnx}x{pny} = {pnx*pny} for a total of {nx*ny} nodes on the base" |
| 59 | + ) |
58 | 60 |
|
59 | 61 | def getMesh(nz):
|
60 |
| - basex = np.linspace(0, 1, nx)[px*pnx:(px+1)*pnx] |
61 |
| - basey = np.linspace(0, 1, ny)[py*pny:(py+1)*pny] |
| 62 | + basex = np.linspace(0, 1, nx)[px * pnx : (px + 1) * pnx] |
| 63 | + basey = np.linspace(0, 1, ny)[py * pny : (py + 1) * pny] |
62 | 64 | z = np.array(range(nz)) * dz
|
63 |
| - return np.stack(np.meshgrid(basex, basey, z, indexing="ij"), axis=-1).reshape(-1, 3) |
| 65 | + return np.stack(np.meshgrid(basex, basey, z, indexing="ij"), axis=-1).reshape( |
| 66 | + -1, 3 |
| 67 | + ) |
64 | 68 |
|
65 | 69 | def requiresEvent(tw):
|
66 | 70 | return tw % eventFrequency == 0
|
@@ -110,20 +114,24 @@ def dataAtTimeWindow(tw):
|
110 | 114 | coords = getMeshAtTimeWindow(tw)
|
111 | 115 | if rank == 0:
|
112 | 116 | print(
|
113 |
| - f"{participant_name}: Event grows local mesh from {oldCount} to { |
| 117 | + f"{participant_name}: Event grows local mesh from {oldCount} to { |
114 | 118 | len(coords)} and global mesh from {
|
115 | 119 | oldCount *
|
116 | 120 | size} to {
|
117 | 121 | len(coords) *
|
118 |
| - size}") |
| 122 | + size}" |
| 123 | + ) |
119 | 124 | participant.reset_mesh(mesh_name)
|
120 | 125 | vertex_ids = participant.set_mesh_vertices(mesh_name, coords)
|
121 | 126 |
|
122 |
| - participant.write_data(mesh_name, write_data_name, vertex_ids, dataAtTimeWindow(tw)) |
| 127 | + participant.write_data( |
| 128 | + mesh_name, write_data_name, vertex_ids, dataAtTimeWindow(tw) |
| 129 | + ) |
123 | 130 |
|
124 | 131 | participant.advance(dt)
|
125 | 132 | tw += 1
|
126 | 133 |
|
| 134 | + |
127 | 135 | if __name__ == "__main__":
|
128 | 136 | try:
|
129 | 137 | main()
|
|
0 commit comments