Updated sorcha docs on Parallelization (#1183)

* Updated sorcha docs on Parallelization

In multi_sorcha.py, fixed the order of parameters config and stats in the function run_sorcha, and removed path_inputs as it wasn't being used.  Also updated the code so that --norbits is no longer required. Now the code will split the ObjIDs provided in the input file and --chunksize into different cores as equally as possible using the np.array_split function.

Updated multi_sorcha.sh and hpc.rst to match these changes to multi_sorcha.py

* Used black

Author: Little-Ryugu

docs/example_files/multi_sorcha.py

@@ -3,8 +3,9 @@
 from multiprocessing import Pool
 import pandas as pd
 import sqlite3
+import numpy as np
 
-def run_sorcha(i, args, path_inputs, pointings, instance,stats, config):
+def run_sorcha(i, args, pointings, instance, config, stats):
     print(f"sorcha run -c {config} --pd {pointings} -o {args.path}{instance}/ -t {instance}_{i} --ob  {args.path}{instance}/orbits_{i}.csv -p {args.path}{instance}/physical_{i}.csv --st {stats}_{i}", flush=True)
     os.system(f"sorcha run -c {config} --pd {pointings} -o {args.path}{instance}/ -t {instance}_{i} --ob  {args.path}{instance}/orbits_{i}.csv -p {args.path}{instance}/physical_{i}.csv --st {stats}_{i}")
 
@@ -16,43 +17,39 @@ def run_sorcha(i, args, path_inputs, pointings, instance,stats, config):
         parser.add_argument('--input_physical', type=str)
         parser.add_argument('--path', type=str)
         parser.add_argument('--chunksize', type=int)
-        parser.add_argument('--norbits', type=int)
         parser.add_argument('--cores', type=int)
         parser.add_argument('--instance', type=int)
         parser.add_argument('--cleanup',  action='store_true')
-        parser.add_argument('--copy_inputs', action='store_true')
         parser.add_argument('--pointings', type=str)
         parser.add_argument('--stats', type=str)
         parser.add_argument('--config', type=str)
         args = parser.parse_args()
         chunk = args.chunksize
         instance = args.instance
-        norbits = args.norbits
         pointings = args.pointings
         path = args.path
         config = args.config
         stats=args.stats
 
         orbits = tb.Table.read(args.input_orbits)
         orbits = orbits[instance*chunk:(instance+1)*chunk]
+        orb_splits = np.array_split(range(len(orbits)), args.cores)
+        
         physical = tb.Table.read(args.input_physical)
         physical = physical[instance*chunk:(instance+1)*chunk]
+        phys_splits = np.array_split(range(len(physical)), args.cores)
 
         os.system(f'mkdir {instance}')
-
-
-        if args.copy_inputs:
-                os.system(f'cp {pointings} {instance}/')
-                path_inputs = f'{instance}'
+        os.system(f'cp {pointings} {instance}/')
 
         for i in range(args.cores):
-                sub_orb = orbits[i*norbits:(i+1)*norbits]
-                sub_phys = physical[i*norbits:(i+1)*norbits]
+                sub_orb = orbits[orb_splits[i]]
+                sub_phys = physical[phys_splits[i]]
                 sub_orb.write(f"{args.path}{instance}/orbits_{i}.csv", overwrite=True)
                 sub_phys.write(f"{args.path}{instance}/physical_{i}.csv", overwrite=True)
 
         with Pool(processes=args.cores) as pool:
-            pool.starmap(run_sorcha, [(i, args, path_inputs, pointings, instance, config, stats) for i in range(args.cores)])
+            pool.starmap(run_sorcha, [(i, args, pointings, instance, config, stats) for i in range(args.cores)])
 
         data = [] 
         for i in range(args.cores):

docs/example_files/multi_sorcha.sh

@@ -6,4 +6,4 @@
 #SBATCH --time=24:00:00
 #SBATCH --output=log-%a.log
 
-python3 multi_sorcha.py --config my_config.ini --input_orbits my_orbits.csv --input_physical my_colors.csv --pointings my_pointings.db --path ./ --chunksize $(($1 * $2)) --norbits $1 --cores $2 --instance ${SLURM_ARRAY_TASK_ID} --cleanup --copy_inputs
+python3 multi_sorcha.py --config my_config.ini --input_orbits my_orbits.csv --input_physical my_colors.csv --pointings my_pointings.db --path ./ --chunksize $1 --cores $2 --instance ${SLURM_ARRAY_TASK_ID} --cleanup 

docs/hpc.rst

@@ -43,13 +43,13 @@ Below is a more complex example of a Slurm script. Here, multi_sorcha.sh calls m
 .. note::
   We provide these here for you to copy, paste, and edit as needed. You might have to some slight modifications to both the Slurm script and multi_sorcha.py, for example if you're using ``Sorcha`` without calling the stats file.   
 
-``multi_sorcha.sh`` requests many parallel Slurm jobs of ``multi_sorcha.py``, feeding each a different --instance parameter. After changing ‘my_orbits.csv’, ‘my_colors.csv’, ‘my_pointings.db’, ‘my_config.ini’, and the various Slurm parameters to match the above, you could generate 10 jobs, each with 4 cores running 25 orbits each, as follows::
+``multi_sorcha.sh`` requests many parallel Slurm jobs of ``multi_sorcha.py``, feeding each a different --instance parameter. After changing ‘my_orbits.csv’, ‘my_colors.csv’, ‘my_pointings.db’, ‘my_config.ini’, and the various Slurm parameters to match the above, for a file of 1000 objects you could generate 10 jobs with 4 cores running 25 orbits each, as follows::
 
-   sbatch --array=0-9 multi_sorcha.sh 25 4
+   sbatch --array=0-9 multi_sorcha.sh 100 4
 
 You can run multi_sorcha.py on the command line as well::
 
-   python multi_sorcha.py --config sorcha_config_demo.ini --input_orbits mba_sample_1000_orbit.csv --input_physical mba_sample_1000_physical.csv --pointings baseline_v2.0_1yr.db --path ./ --chunksize 1000 --norbits 250 --cores 4 --instance 0 --stats mbastats --cleanup --copy_inputs 
+   python multi_sorcha.py --config sorcha_config_demo.ini --input_orbits mba_sample_1000_orbit.csv --input_physical mba_sample_1000_physical.csv --pointings baseline_v2.0_1yr.db --path ./ --chunksize 1000 --cores 4 --instance 0 --stats mbastats --cleanup
 
 This will generate a single output file. It should work fine on a laptop, and be a bit (but not quite 4x) faster than the single-core equivalent due to overheads.