10
10
11
11
import flopy
12
12
import numpy as np
13
- from common_regression import (COMPARE_PROGRAMS , adjust_htol ,
14
- get_mf6_comparison , get_mf6_files ,
15
- get_namefiles , get_rclose , get_regression_files ,
16
- setup_mf6 , setup_mf6_comparison )
13
+ from common_regression import (
14
+ COMPARE_PROGRAMS ,
15
+ adjust_htol ,
16
+ get_mf6_comparison ,
17
+ get_mf6_files ,
18
+ get_namefiles ,
19
+ get_rclose ,
20
+ get_regression_files ,
21
+ setup_mf6 ,
22
+ setup_mf6_comparison ,
23
+ )
17
24
from flopy .mbase import BaseModel
18
25
from flopy .mf6 import MFSimulation
19
26
from flopy .utils .compare import compare_heads
@@ -243,7 +250,7 @@ def __init__(
243
250
self .build = build
244
251
self .check = check
245
252
self .parallel = parallel
246
- self .ncpus = ncpus
253
+ self .ncpus = [ ncpus ] if isinstance ( ncpus , int ) else ncpus
247
254
self .api_func = api_func
248
255
self .compare = compare
249
256
self .outp = None
@@ -526,6 +533,7 @@ def run_sim_or_model(
526
533
workspace : Union [str , os .PathLike ],
527
534
target : Union [str , os .PathLike ],
528
535
xfail : bool = False ,
536
+ ncpus : int = 1 ,
529
537
) -> Tuple [bool , List [str ]]:
530
538
"""
531
539
Run a simulation or model with FloPy.
@@ -536,6 +544,8 @@ def run_sim_or_model(
536
544
The target executable to use
537
545
xfail : bool
538
546
Whether to expect failure
547
+ ncpus : int
548
+ The number of CPUs for a parallel run
539
549
"""
540
550
541
551
# make sure workspace exists
@@ -570,14 +580,12 @@ def run_sim_or_model(
570
580
# via MF6 executable
571
581
elif "mf6" in target .name :
572
582
# parallel test if configured
573
- if self .parallel :
583
+ if self .parallel and ncpus > 1 :
574
584
print (
575
585
f"Parallel test { self .name } on { self .ncpus } processes"
576
586
)
577
587
try :
578
- success , buff = run_parallel (
579
- workspace , target , self .ncpus
580
- )
588
+ success , buff = run_parallel (workspace , target , ncpus )
581
589
except Exception :
582
590
warn (
583
591
"MODFLOW 6 parallel test" ,
@@ -687,19 +695,31 @@ def run(self):
687
695
self .sims = sims
688
696
nsims = len (sims )
689
697
self .buffs = list (repeat (None , nsims ))
698
+
690
699
assert len (self .xfail ) in [
691
700
1 ,
692
701
nsims ,
693
702
], f"Invalid xfail: expected a single boolean or one for each model"
694
703
if len (self .xfail ) == 1 and nsims :
695
704
self .xfail = list (repeat (self .xfail [0 ], nsims ))
705
+
706
+ assert len (self .ncpus ) in [
707
+ 1 ,
708
+ nsims ,
709
+ ], f"Invalid ncpus: expected a single integer or one for each model"
710
+ if len (self .ncpus ) == 1 and nsims :
711
+ self .ncpus = list (repeat (self .ncpus [0 ], nsims ))
712
+
696
713
write_input (* sims , overwrite = self .overwrite , verbose = self .verbose )
697
714
else :
698
715
self .sims = [MFSimulation .load (sim_ws = self .workspace )]
699
716
self .buffs = [None ]
700
717
assert (
701
718
len (self .xfail ) == 1
702
- ), f"Invalid xfail: expected a single boolean or one for each model"
719
+ ), f"Invalid xfail: expected a single boolean"
720
+ assert (
721
+ len (self .ncpus ) == 1
722
+ ), f"Invalid ncpus: expected a single integer"
703
723
704
724
# run models/simulations
705
725
for i , sim_or_model in enumerate (self .sims ):
@@ -716,7 +736,10 @@ def run(self):
716
736
else tgts .get (exe_path .stem , tgts ["mf6" ])
717
737
)
718
738
xfail = self .xfail [i ]
719
- success , buff = self .run_sim_or_model (workspace , target , xfail )
739
+ ncpus = self .ncpus [i ]
740
+ success , buff = self .run_sim_or_model (
741
+ workspace , target , xfail , ncpus
742
+ )
720
743
self .buffs [i ] = buff # store model output for assertions later
721
744
assert success , (
722
745
f"{ 'Simulation' if 'mf6' in str (target ) else 'Model' } "
0 commit comments