def terminate(*args): # pragma: nocover log.info('Terminating...') try: client = Client.current() client.close() client.cluster.close() except: # noqa: E722 pass sys.exit(1)
def __init__(self, client=None, storage=None): try: self.client = client or Client.current() except ValueError: # Initialise new client self.client = get_worker().client self.storage = storage if self.client.asynchronous or getattr(thread_state, "on_event_loop_thread", False): async def _register(): await self.client.run_on_scheduler(register_with_scheduler) self.client.loop.add_callback(_register) else: self.client.run_on_scheduler(register_with_scheduler)
def find_offsets(self): # noqa: C901 """Calculate offsets between all pairs of overlapping images""" client = Client.current() # convert to find_shifts results = [] logger.info("Processing section %s", self._section.name) # debug # img_cube = self.get_img_section(0, Settings.channel_to_use - 1) img_cube = self.get_img_section(0, -1) # Calculate confidence map. Only needs to be done once per section dist_conf = self.get_distconf() flat = read_calib(Settings.flat_file)[Settings.channel_to_use - 1] dark = (read_calib(Settings.dark_file)[Settings.channel_to_use - 1] / Settings.norm_val) flat = flat.persist() dark = dark.persist() dx0, dy0, delta_x, delta_y = self.find_grid() dx_mos, dy_mos = self.get_mos_pos() # We calculate the ref_img here, for if the slice # has no data (very low max. avg flux), we skip calculations too img_cube_stack = da.stack(img_cube) # ref image is the one with the max integrated flux cube_totals = img_cube_stack.sum(axis=(1, 2)) cube_totals = cube_totals.compute() self.cube_means = cube_totals / 2080.0**2 self.absolute_ref_img = self.cube_means.argmax() self.mean_ref_img = self.cube_means.max() logger.info("Max mean: {0:.5f}".format(self.mean_ref_img)) # If the default displacements are too large # to for overlaps, stick with the default scale self.offset_mode = "sampled" if delta_x / Settings.mosaic_scale > 1950.0: logger.info("Displacement in X too large: {0:.1f}".format( delta_x / Settings.mosaic_scale)) self.offset_mode = "default" if delta_y / Settings.mosaic_scale > 1950.0: logger.info("Displacement in Y too large: {0:.1f}".format( delta_y / Settings.mosaic_scale)) self.offset_mode = "default" if self.mean_ref_img < 0.05: logger.info("Avg. flux too low: {0:.3f}<0.05".format( self.mean_ref_img)) self.offset_mode = "default" for i, img in enumerate(img_cube): r = np.sqrt((dx0 - dx0[i])**2 + (dy0 - dy0[i])**2) # including no diagonals i_t = np.where((r <= np.sqrt(1)) & (r > 0))[0].tolist() im_i = img for this_img in i_t: if i > this_img: continue desp = [ (dx_mos[this_img] - dx_mos[i]) / Settings.mosaic_scale, (dy_mos[this_img] - dy_mos[i]) / Settings.mosaic_scale, ] # trying to do always positive displacements if (desp[0] < -100) or (desp[1] < -100): desp[0] *= -1 desp[1] *= -1 obj_img = i im_obj = im_i ref_img = this_img im_ref = img_cube[this_img] sign = -1 else: obj_img = this_img im_obj = img_cube[this_img] ref_img = i im_ref = im_i sign = 1 if self.offset_mode == "sampled": im1 = apply_geometric_transform( im_ref.data / Settings.norm_val, dark, flat, Settings.cof_dist) im2 = apply_geometric_transform( im_obj.data / Settings.norm_val, dark, flat, Settings.cof_dist) res = delayed(find_overlap_conf)(im1, im2, dist_conf, dist_conf, desp) results.append(_sink(ref_img, obj_img, res, sign)) else: logger.debug( "Initial offsets i: %d j: %d dx: %d dy: %d", ref_img, obj_img, desp[0], desp[1], ) results.append([ref_img, obj_img, desp[0], desp[1], sign]) if self.offset_mode == "sampled": futures = client.compute(results) offsets = [] for fut in as_completed(futures): i, j, res, sign = fut.result() dx, dy, mi, avf = res.x, res.y, res.mi, res.avg_flux offsets.append([i, j, res, sign]) logger.debug( "Section %s offsets i: %d j: %d dx: %d dy: %d mi: %f avg_f: %f sign: %d", self._section.name, i, j, dx, dy, mi, avf, sign, ) else: offsets = [] for t in results: i, j, dx, dy, sign = t offsets.append([i, j, [dx, dy, 0.0, 0.0], sign]) logger.debug( "Section %s offsets i: %d j: %d dx: %d dy: %d mi: %f avg_f: %f sign: %d", self._section.name, i, j, dx, dy, 0.0, 0.0, sign, ) self._offsets = offsets