Skip to content

Commit 74f3f50

Browse files
committed
back in black
1 parent 2948ef4 commit 74f3f50

6 files changed

Lines changed: 107 additions & 114 deletions

File tree

mcdc/adapt.py

Lines changed: 0 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -106,11 +106,9 @@ def local_array(shape, dtype):
106106

107107
@numba.extending.type_callable(local_array)
108108
def type_local_array(context):
109-
110109
from numba.core.typing.npydecl import parse_dtype, parse_shape
111110

112111
if isinstance(context, numba.core.typing.context.Context):
113-
114112
# Function repurposed from Numba's ol_np_empty.
115113
def typer(shape, dtype):
116114
numba.np.arrayobj._check_const_str_dtype("empty", dtype)
@@ -148,10 +146,8 @@ def typer(shape, dtype):
148146
return typer
149147

150148
elif isinstance(context, numba.cuda.target.CUDATypingContext):
151-
152149
# Function repurposed from Numba's Cuda_array_decl.
153150
def typer(shape, dtype):
154-
155151
# Only integer literals and tuples of integer literals are valid
156152
# shapes
157153
if isinstance(shape, types.Integer):
@@ -200,14 +196,12 @@ def typer(shape, dtype):
200196

201197
@numba.extending.lower_builtin(local_array, types.IntegerLiteral, types.Any)
202198
def builtin_local_array(context, builder, sig, args):
203-
204199
shape, dtype = sig.args
205200

206201
from numba.core.typing.npydecl import parse_dtype, parse_shape
207202
import numba.np.arrayobj as arrayobj
208203

209204
if isinstance(context, numba.core.cpu.CPUContext):
210-
211205
# No default arguments.
212206
nb_dtype = parse_dtype(dtype)
213207
nb_shape = parse_shape(shape)
@@ -420,7 +414,6 @@ def nopython_mode(is_on):
420414

421415

422416
def gpu_forward_declare(args):
423-
424417
if args.gpu_rocm_path != None:
425418
harm.config.set_rocm_path(args.gpu_rocm_path)
426419

@@ -651,7 +644,6 @@ def finalize(prog: numba.uintp):
651644
final_fn(prog)
652645

653646
def step(prog: numba.uintp, arg: arg_type):
654-
655647
step_async()
656648

657649
(step_async,) = harm.RuntimeSpec.async_dispatch(step)

mcdc/kernel.py

Lines changed: 2 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -196,7 +196,6 @@ def dd_check_out(mcdc):
196196

197197
@trace.njit()
198198
def dd_signal_halt(mcdc):
199-
200199
with objmode():
201200
for rank in range(1, MPI.COMM_WORLD.Get_size()):
202201
dummy_buff = np.zeros((1,), dtype=np.int32)
@@ -207,7 +206,6 @@ def dd_signal_halt(mcdc):
207206

208207
@trace.njit()
209208
def dd_signal_block(mcdc):
210-
211209
with objmode(rank="int64"):
212210
rank = MPI.COMM_WORLD.Get_rank()
213211

@@ -240,7 +238,6 @@ def dd_signal_block(mcdc):
240238

241239
@trace.njit()
242240
def dd_signal_unblock(mcdc):
243-
244241
with objmode(rank="int64"):
245242
rank = MPI.COMM_WORLD.Get_rank()
246243

@@ -270,7 +267,6 @@ def dd_signal_unblock(mcdc):
270267

271268
@trace.njit()
272269
def dd_distribute_bank(mcdc, bank, dest_list):
273-
274270
with objmode(send_delta="int64"):
275271
dest_count = len(dest_list)
276272
send_delta = 0
@@ -334,7 +330,6 @@ def dd_particle_send(prog):
334330

335331
@trace.njit()
336332
def dd_get_recv_tag():
337-
338333
with objmode(tag="int64"):
339334
status = MPI.Status()
340335
MPI.COMM_WORLD.Probe(status=status)
@@ -345,7 +340,6 @@ def dd_get_recv_tag():
345340

346341
@trace.njit()
347342
def dd_recv_particles(mcdc):
348-
349343
buff = np.zeros(
350344
mcdc["domain_decomp"]["bank_zp"]["particles"].shape[0],
351345
dtype=type_.particle_record,
@@ -373,7 +367,6 @@ def dd_recv_particles(mcdc):
373367

374368
@trace.njit()
375369
def dd_recv_turnstile(mcdc):
376-
377370
with objmode(busy_delta="int64", send_delta="int64"):
378371
event_buff = np.zeros((1,), dtype=type_.dd_turnstile_event)
379372
MPI.COMM_WORLD.Recv([event_buff, type_.dd_turnstile_event_mpi])
@@ -394,7 +387,6 @@ def dd_recv_turnstile(mcdc):
394387

395388
@trace.njit()
396389
def dd_recv_halt(mcdc):
397-
398390
with objmode():
399391
dummy_buff = np.zeros((1,), dtype=np.int32)
400392
MPI.COMM_WORLD.Recv(dummy_buff)
@@ -406,7 +398,6 @@ def dd_recv_halt(mcdc):
406398

407399
@trace.njit()
408400
def dd_recv(mcdc):
409-
410401
if mcdc["domain_decomp"]["rank_busy"]:
411402
dd_signal_block(mcdc)
412403
mcdc["domain_decomp"]["rank_busy"] = False
@@ -719,6 +710,8 @@ def distribute_work_dd_precursor(N, mcdc):
719710
mcdc["mpi_work_start_precursor"] = work_start
720711
mcdc["mpi_work_size_precursor"] = work_size
721712
mcdc["mpi_work_size_total_precursor"] = work_size_total
713+
714+
722715
# =============================================================================
723716
# Random sampling
724717
# =============================================================================
@@ -2117,7 +2110,6 @@ def score_cell_tally(P_arr, distance, tally, data, mcdc):
21172110
# Sweep through the distance
21182111
distance_swept = 0.0
21192112
while distance_swept < distance - COINCIDENCE_TOLERANCE:
2120-
21212113
# Find distance to mesh grids
21222114
dt = (min(mesh["t"][it + 1], t_final) - t) / ut
21232115

mcdc/loop.py

Lines changed: 0 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -62,7 +62,6 @@ def teardown_gpu(mcdc):
6262

6363
@trace.njit()
6464
def loop_fixed_source(data_arr, mcdc_arr):
65-
6665
# Ensure `data` and `mcdc` exist for the lifetime of the program
6766
# by intentionally leaking their memory
6867
adapt.leak(data_arr)
@@ -382,7 +381,6 @@ def source_dd_resolution(data, prog):
382381
if kernel.get_bank_size(mcdc["bank_active"]) > 0:
383382
# Loop until active bank is exhausted
384383
while kernel.get_bank_size(mcdc["bank_active"]) > 0:
385-
386384
kernel.get_particle(P_arr, mcdc["bank_active"], mcdc)
387385
if not kernel.particle_in_domain(P_arr, mcdc) and P["alive"] == True:
388386
print(f"recieved particle not in domain")
@@ -490,7 +488,6 @@ def step(prog: nb.uintp, P_input: adapt.particle_gpu):
490488

491489
@trace.njit(cache=caching)
492490
def gpu_loop_source(seed, data, mcdc):
493-
494491
# Progress bar indicator
495492
N_prog = 0
496493

@@ -514,7 +511,6 @@ def gpu_loop_source(seed, data, mcdc):
514511
phase_count = (full_work_size + phase_size - 1) // phase_size
515512

516513
for phase in range(phase_count):
517-
518514
mcdc["mpi_work_iter"][0] = phase_size * phase
519515
mcdc["mpi_work_size"] = min(phase_size * (phase + 1), full_work_size)
520516
mcdc["source_seed"] = seed
@@ -839,7 +835,6 @@ def step(prog: nb.uintp, P_input: adapt.particle_gpu):
839835

840836
@trace.njit(cache=caching)
841837
def gpu_loop_source_precursor(seed, data, mcdc):
842-
843838
# Progress bar indicator
844839
N_prog = 0
845840

@@ -898,7 +893,6 @@ def gpu_loop_source_precursor(seed, data, mcdc):
898893

899894

900895
def build_gpu_progs(input_deck, args):
901-
902896
STRAT = args.gpu_strat
903897

904898
src_spec = gpu_sources_spec()

mcdc/main.py

Lines changed: 30 additions & 22 deletions
Original file line numberDiff line numberDiff line change
@@ -417,7 +417,7 @@ def dd_mesh_bounds(idx):
417417
zlen = len(input_deck.mesh_tallies[idx].z)
418418
if (xlen == 2) and (ylen == 2) and (zlen == 2):
419419
return 0, 1, 0, 1, 0, 1
420-
420+
421421
xn = input_deck.technique["dd_mesh"]["x"][xmesh_idx]
422422
xp = input_deck.technique["dd_mesh"]["x"][xmesh_idx + 1]
423423
yn = input_deck.technique["dd_mesh"]["y"][ymesh_idx]
@@ -645,9 +645,9 @@ def prepare():
645645
for i in range(N_material):
646646
for name in type_.material.names:
647647
if name in ["nuclide_IDs", "nuclide_densities"]:
648-
mcdc["materials"][i][name][: mcdc["materials"][i]["N_nuclide"]] = (
649-
getattr(input_deck.materials[i], name)
650-
)
648+
mcdc["materials"][i][name][
649+
: mcdc["materials"][i]["N_nuclide"]
650+
] = getattr(input_deck.materials[i], name)
651651
else:
652652
copy_field(mcdc["materials"][i], input_deck.materials[i], name)
653653

@@ -877,9 +877,15 @@ def prepare():
877877
)
878878

879879
else: # decomposed mesh filters
880-
mcdc["technique"]["dd_xsum"] = max(mcdc["technique"]["dd_xsum"],len(input_deck.mesh_tallies[i].x) - 1)
881-
mcdc["technique"]["dd_ysum"] = max(mcdc["technique"]["dd_ysum"],len(input_deck.mesh_tallies[i].y) - 1)
882-
mcdc["technique"]["dd_zsum"] = max(mcdc["technique"]["dd_zsum"],len(input_deck.mesh_tallies[i].z) - 1)
880+
mcdc["technique"]["dd_xsum"] = max(
881+
mcdc["technique"]["dd_xsum"], len(input_deck.mesh_tallies[i].x) - 1
882+
)
883+
mcdc["technique"]["dd_ysum"] = max(
884+
mcdc["technique"]["dd_ysum"], len(input_deck.mesh_tallies[i].y) - 1
885+
)
886+
mcdc["technique"]["dd_zsum"] = max(
887+
mcdc["technique"]["dd_zsum"], len(input_deck.mesh_tallies[i].z) - 1
888+
)
883889

884890
mxn, mxp, myn, myp, mzn, mzp = dd_mesh_bounds(i)
885891

@@ -1404,9 +1410,9 @@ def prepare():
14041410
for i in range(M):
14051411
idm = input_deck.uq_deltas["materials"][i].ID
14061412
mcdc["technique"]["uq_"]["materials"][i]["info"]["ID"] = idm
1407-
mcdc["technique"]["uq_"]["materials"][i]["info"]["distribution"] = (
1408-
input_deck.uq_deltas["materials"][i].distribution
1409-
)
1413+
mcdc["technique"]["uq_"]["materials"][i]["info"][
1414+
"distribution"
1415+
] = input_deck.uq_deltas["materials"][i].distribution
14101416
for name in input_deck.uq_deltas["materials"][i].flags:
14111417
mcdc["technique"]["uq_"]["materials"][i]["flags"][name] = True
14121418
mcdc["technique"]["uq_"]["materials"][i]["delta"][name] = getattr(
@@ -1420,15 +1426,15 @@ def prepare():
14201426
flags["nu_f"] = True
14211427
if mcdc["materials"][idm]["N_nuclide"] > 1:
14221428
for name in type_.uq_mat.names:
1423-
mcdc["technique"]["uq_"]["materials"][i]["mean"][name] = (
1424-
input_deck.materials[idm][name]
1425-
)
1429+
mcdc["technique"]["uq_"]["materials"][i]["mean"][
1430+
name
1431+
] = input_deck.materials[idm][name]
14261432

14271433
N = len(input_deck.uq_deltas["nuclides"])
14281434
for i in range(N):
1429-
mcdc["technique"]["uq_"]["nuclides"][i]["info"]["distribution"] = (
1430-
input_deck.uq_deltas["nuclides"][i].distribution
1431-
)
1435+
mcdc["technique"]["uq_"]["nuclides"][i]["info"][
1436+
"distribution"
1437+
] = input_deck.uq_deltas["nuclides"][i].distribution
14321438
idn = input_deck.uq_deltas["nuclides"][i].ID
14331439
mcdc["technique"]["uq_"]["nuclides"][i]["info"]["ID"] = idn
14341440
for name in type_.uq_nuc.names:
@@ -1806,7 +1812,8 @@ def dd_mergemesh(mcdc):
18061812

18071813
if d_Nz > 1:
18081814
sendcounts = np.array(
1809-
MPI.COMM_WORLD.gather(len(mcdc["mesh_tallies"][0]["filter"]["z"]), root=0))
1815+
MPI.COMM_WORLD.gather(len(mcdc["mesh_tallies"][0]["filter"]["z"]), root=0)
1816+
)
18101817
if mcdc["mpi_master"]:
18111818
z_filter = np.zeros((mcdc["mesh_tallies"].shape[0], sum(sendcounts)))
18121819
else:
@@ -1932,7 +1939,7 @@ def generate_hdf5(data, mcdc):
19321939
# Set tally shape
19331940
N_score = tally["N_score"]
19341941
if mcdc["technique"]["domain_decomposition"]:
1935-
if Nx or Ny or Nz: # check if spatial mesh exists
1942+
if Nx or Ny or Nz: # check if spatial mesh exists
19361943
Nx = mcdc["technique"]["dd_xsum"]
19371944
Ny = mcdc["technique"]["dd_ysum"]
19381945
Nz = mcdc["technique"]["dd_zsum"]
@@ -2386,7 +2393,6 @@ def recombine_tallies(file="output.h5"):
23862393

23872394

23882395
def closeout(mcdc):
2389-
23902396
loop.teardown_gpu(mcdc)
23912397

23922398
# Runtime
@@ -2523,9 +2529,11 @@ def visualize(
25232529
particle["t"] = t
25242530

25252531
# Random direction
2526-
particle["ux"], particle["uy"], particle["uz"] = (
2527-
kernel.sample_isotropic_direction(particle_container)
2528-
)
2532+
(
2533+
particle["ux"],
2534+
particle["uy"],
2535+
particle["uz"],
2536+
) = kernel.sample_isotropic_direction(particle_container)
25292537

25302538
# RGB color data for each pixels
25312539
data = np.zeros(pixels + (3,))

0 commit comments

Comments
 (0)