fw demod sim: add --prepare flag

1 files changed, 10 insertions, 1 deletions

diff --git a/controller/fw/tools/dsss_demod_test_runner.py b/controller/fw/tools/dsss_demod_test_runner.py
index a7acccd..c0f737d 100644
--- a/controller/fw/tools/dsss_demod_test_runner.py
+++ b/controller/fw/tools/dsss_demod_test_runner.py
@@ -1,6 +1,7 @@
 #!/usr/bin/env python3
 
 import os
+import sys
 from os import path
 import subprocess
 import json
@@ -146,6 +147,9 @@ if __name__ == '__main__':
     parser.add_argument('-d', '--dump', help='Write results to JSON file')
     parser.add_argument('-c', '--cachedir', default='dsss_test_cache', help='Directory to store build output and data in')
     parser.add_argument('-n', '--no-cache', action='store_true', help='Disable result cache')
+    parser.add_argument('-b', '--batches', type=int, default=1, help='Number of batches to split the computation into')
+    parser.add_argument('-i', '--index', type=int, default=0, help='Batch index to compute')
+    parser.add_argument('-p', '--prepare', action='store_true', help='Prepare mode: compile runners, then exit.')
     args = parser.parse_args()
 
     DecoderParams = namedtuple('DecoderParams', ['nbits', 'thf', 'decimation', 'symbols'])
@@ -175,6 +179,8 @@ if __name__ == '__main__':
             builder=build_test_binary, param_list=dec_paramses, desc='Building decoders',
             context=dict(cachedir=build_cache_dir)))
     print('Done building decoders.')
+    if args.prepare:
+        sys.exit(0)
 
     GeneratorParams = namedtuple('GeneratorParams', ['seed', 'amplitude_spec', 'background'])
     gen_params = [ GeneratorParams(rep, (5e-3, 1, 5), background)
@@ -195,6 +201,9 @@ if __name__ == '__main__':
     waveform_params = [ (*gp, *dp, thfs) for gp in gen_params for dp, thfs in dec_param_groups.items() ]
     print(f'Generated {len(waveform_params)} parameter sets')
 
+    # Separate out our batch
+    waveform_params = waveform_params[args.index::args.batches]
+
     def lookup_binary(*params):
         return path.join(build_cache_dir, builds[tuple(params)], 'tools/dsss_demod_test')
 
@@ -208,7 +217,7 @@ if __name__ == '__main__':
             ['seed', 'amplitude_spec', 'background', 'nbits', 'decimation', 'symbols', 'thresholds'],
             params_mapper=params_mapper,
             builder=run_test,
-            param_list=waveform_params, desc='Generating waveforms',
+            param_list=waveform_params, desc='Simulating demodulation',
             context=dict(cachedir=data_cache_dir, lookup_binary=lookup_binary),
             disable_cache=args.no_cache):
         results += chunk