diff --git a/gwemopt/args.py b/gwemopt/args.py index 8f547da..db46488 100644 --- a/gwemopt/args.py +++ b/gwemopt/args.py @@ -128,6 +128,7 @@ def parse_args(args): parser.add_argument("--doParallel", action="store_true", default=False) parser.add_argument("--Ncores", default=4, type=int) + parser.add_argument("--parallelBackend", default="threading") parser.add_argument("--doBlocks", action="store_true", default=False) parser.add_argument("--Nblocks", default=4, type=int) diff --git a/gwemopt/moc.py b/gwemopt/moc.py index 3739124..6ddfa6f 100644 --- a/gwemopt/moc.py +++ b/gwemopt/moc.py @@ -10,6 +10,7 @@ import gwemopt.tiles from gwemopt.chipgaps import get_decam_quadrant_moc, get_ztf_quadrant_moc from gwemopt.paths import CONFIG_DIR +from gwemopt.utils.parallel import tqdm_joblib from gwemopt.utils.pixels import getCirclePixels, getRegionPixels, getSquarePixels @@ -23,16 +24,19 @@ def create_moc(params, map_struct=None): moc_struct = {} if params["doParallel"]: - moclists = Parallel( - n_jobs=params["Ncores"], - backend="threading", - batch_size=int(len(tesselation) / params["Ncores"]) + 1, - )( - delayed(Fov2Moc)( - params, config_struct, telescope, tess[1], tess[2], nside + with tqdm_joblib( + tqdm(desc="MOC creation", total=len(tesselation)) + ) as progress_bar: + moclists = Parallel( + n_jobs=params["Ncores"], + backend=params["parallelBackend"], + batch_size=int(len(tesselation) / params["Ncores"]) + 1, + )( + delayed(Fov2Moc)( + params, config_struct, telescope, tess[1], tess[2], nside + ) + for tess in tesselation ) - for tess in tqdm(tesselation) - ) for ii, tess in tqdm(enumerate(tesselation), total=len(tesselation)): index, ra, dec = tess[0], tess[1], tess[2] if (telescope == "ZTF") and params["doUsePrimary"] and (index > 880): diff --git a/gwemopt/params.py b/gwemopt/params.py index 07f04be..1fe0461 100644 --- a/gwemopt/params.py +++ b/gwemopt/params.py @@ -162,4 +162,8 @@ def params_struct(opts): json.loads(opts.milpOptions) if hasattr(opts, "milpOptions") else {} ) + params["parallelBackend"] = ( + opts.parallelBackend if hasattr(opts, "parallelBackend") else "threading" + ) + return params diff --git a/gwemopt/plotting/plot_coverage.py b/gwemopt/plotting/plot_coverage.py index 3b85db4..f8f9a80 100644 --- a/gwemopt/plotting/plot_coverage.py +++ b/gwemopt/plotting/plot_coverage.py @@ -276,7 +276,10 @@ def plot_coverage_scaled(params, map_struct, coverage_struct, plot_sun_moon, max moviedir = params["outputDir"].joinpath("movie") moviedir.mkdir(exist_ok=True, parents=True) - for jj in tqdm(range(len(mjds))): + def make_plot(jj, params, map_struct, coverage_struct): + hdu = map_struct["hdu"] + columns = [col.name for col in hdu.columns] + mjd = mjds[jj] plot_name = moviedir.joinpath(f"coverage-{jj:04d}.png") title = f"Coverage Map: {mjd:.2f}" @@ -310,6 +313,9 @@ def plot_coverage_scaled(params, map_struct, coverage_struct, plot_sun_moon, max plt.savefig(plot_name, dpi=200) plt.close() + for jj in tqdm(range(len(mjds))): + make_plot(jj, params, map_struct, coverage_struct) + moviefiles = moviedir.joinpath("coverage-%04d.png") filename = params["outputDir"].joinpath("coverage.mpg") diff --git a/gwemopt/scheduler.py b/gwemopt/scheduler.py index b60ee2e..17d2bc5 100644 --- a/gwemopt/scheduler.py +++ b/gwemopt/scheduler.py @@ -122,7 +122,10 @@ def get_order_heuristic( if type(tile_struct[key]["exposureTime"]) in [float, np.float64]: tileexpdur[jj] = tile_struct[key]["exposureTime"] elif type(tile_struct[key]["exposureTime"]) in [list, np.ndarray]: - tileexpdur[jj] = tile_struct[key]["exposureTime"][0] + if len(tile_struct[key]["exposureTime"]) > 0: + tileexpdur[jj] = tile_struct[key]["exposureTime"][0] + else: + tileexpdur[jj] = 0.0 else: tileexpdur[jj] = 0.0 diff --git a/gwemopt/segments.py b/gwemopt/segments.py index acfb29f..a51bfc2 100644 --- a/gwemopt/segments.py +++ b/gwemopt/segments.py @@ -11,6 +11,7 @@ from gwemopt.utils.geometry import angular_distance from gwemopt.utils.misc import get_exposures +from gwemopt.utils.parallel import tqdm_joblib from gwemopt.utils.sidereal_time import greenwich_sidereal_time # conversion between MJD (tt) and DJD (what ephem uses) @@ -288,16 +289,17 @@ def get_segments_tiles(params, config_struct, tile_struct): moon_radecs = get_moon_radecs(segmentlist, observer) if params["doParallel"]: - tilesegmentlists = Parallel( - n_jobs=params["Ncores"], - backend="threading", - batch_size=int(len(radecs) / params["Ncores"]) + 1, - )( - delayed(get_segments_tile)( - config_struct, radec, segmentlist, moon_radecs, params["airmass"] + with tqdm_joblib(tqdm(desc="MOC creation", total=len(radecs))) as progress_bar: + tilesegmentlists = Parallel( + n_jobs=params["Ncores"], + backend=params["parallelBackend"], + batch_size=int(len(radecs) / params["Ncores"]) + 1, + )( + delayed(get_segments_tile)( + config_struct, radec, segmentlist, moon_radecs, params["airmass"] + ) + for radec in radecs ) - for radec in tqdm(radecs) - ) for ii, key in enumerate(keys): tile_struct[key]["segmentlist"] = tilesegmentlists[ii] else: diff --git a/gwemopt/tiles.py b/gwemopt/tiles.py index 0616fb1..00c43ee 100644 --- a/gwemopt/tiles.py +++ b/gwemopt/tiles.py @@ -9,6 +9,7 @@ from astropy import units as u from astropy.coordinates import ICRS, SkyCoord from astropy.time import Time +from joblib import Parallel, delayed from mocpy import MOC from scipy.stats import norm from shapely.geometry import MultiPoint @@ -18,6 +19,7 @@ import gwemopt.moc import gwemopt.segments from gwemopt.utils.geometry import angular_distance +from gwemopt.utils.parallel import tqdm_joblib LEVEL = MOC.MAX_ORDER HPX = ah.HEALPix(nside=ah.level_to_nside(LEVEL), order="nested", frame=ICRS()) @@ -1083,8 +1085,27 @@ def compute_tiles_map( keys = tile_struct.keys() ntiles = len(keys) vals = np.nan * np.ones((ntiles,)) - for ii, key in tqdm(enumerate(keys), total=len(keys)): - vals[ii] = tile_struct[key]["moc"].probability_in_multiordermap(skymap) + + if params["doParallel"]: + + def calculate_probability(moc, skymap): + return moc.probability_in_multiordermap(skymap) + + with tqdm_joblib( + tqdm(desc="Probability calcuation", total=len(keys)) + ) as progress_bar: + vals = Parallel( + n_jobs=params["Ncores"], + backend=params["parallelBackend"], + batch_size=int(len(keys) / params["Ncores"]) + 1, + )( + delayed(calculate_probability)(tile_struct[key]["moc"], skymap) + for key in keys + ) + vals = np.array(vals) + else: + for ii, key in tqdm(enumerate(keys), total=len(keys)): + vals[ii] = tile_struct[key]["moc"].probability_in_multiordermap(skymap) return vals