enable array job submission

observation_sim/ObservationSim.py

@@ -131,7 +131,7 @@ class Observation(object):
                              chip.chipID, os.getpid(), (psutil.Process(os.getpid()).memory_info().rss / 1024 / 1024 / 1024)))
         del chip.img
 
-    def runExposure_MPI_PointingList(self, pointing_list, chips=None):
+    def runExposure_MPI_PointingList(self, pointing_list, chips=None, task_id=None):
         comm = MPI.COMM_WORLD
         ind_thread = comm.Get_rank()
         num_thread = comm.Get_size()
@@ -184,7 +184,10 @@ class Observation(object):
 
             for ichip in range(nchips_per_fp):
                 i_process = process_counter + ichip
-                if i_process % num_thread != ind_thread:
+                if task_id is not None:
+                    if i_process != task_id:
+                        continue
+                elif i_process % num_thread != ind_thread:
                     continue
                 pid = os.getpid()
 

observation_sim/_util.py

@@ -16,13 +16,15 @@ def parse_args():
     parser.add_argument('--config_file', type=str, required=True,
                         help='.yaml config file for simulation settings.')
     parser.add_argument('--catalog', type=str,
-                        help='name of the catalog interface class to be loaded.')
+                        help='Name of the catalog interface class to be loaded.')
     parser.add_argument('-c', '--config_dir', type=str,
                         help='Directory that houses the .yaml config file.')
     parser.add_argument('-d', '--data_dir', type=str,
                         help='Directory that houses the input data.')
     parser.add_argument('-w', '--work_dir', type=str,
                         help='The path for output.')
+    parser.add_argument('--task_id', type=int,
+                        help='The task id to run a particular chip, used when running with job array')
     return parser.parse_args()
 
 

run_sim.py

@@ -103,7 +103,8 @@ def run_sim():
                           work_dir=config['work_dir'], data_dir=config['data_dir'])
 
     # Run simulation
-    obs.runExposure_MPI_PointingList(pointing_list=pointing_list)
+    obs.runExposure_MPI_PointingList(
+        pointing_list=pointing_list, task_id=args.task_id)
 
 
 if __name__ == '__main__':