def _apply_new_correction(data): """Applies the new 2015 correction to the data. Arguments: data -- Raw data without any correction applied (old correction should be already reverted at this point). Returns: return_data -- Corrected data using 2015 2D correction polynomial where possible. Where only data of one plane is available 1D fallback is used. """ return_data = copy.deepcopy( data ) # need a deepcopy of data here because we want to modify values but still need the old ones for plane in data: for bpm in data[plane]: bpm_type = bpm.split(".")[0] if bpm_type == "BPMS": continue # we will skip skewed BPMS type BPMs as they need special treatment and are only a few use_1d_fallback = True if bpm not in data["X"] or bpm not in data["Y"] else False if use_1d_fallback: print("Warning: No corresponding BPM for %s found, using 1D fallback correction!" % bpm) for i in range(len(data[plane][bpm])): if not use_1d_fallback: if plane == "X": u_raw = data["X"][bpm][i] v_raw = data["Y"][bpm][i] else: # bit confusing but we just want to replace x by y and vice versa for the correction or the other plane u_raw = data["Y"][bpm][i] v_raw = data["X"][bpm][i] coefficients = polynomial_correction.get_coefficients_of_bpm_type( bpm_type, polynomial_correction.new_coefficients ) else: u_raw = data[plane][bpm][i] v_raw = 0 coefficients = polynomial_correction.get_coefficients_of_bpm_type( bpm_type, polynomial_correction.new_coefficients_1D_fallback ) return_data[plane][bpm][i] = polynomial_correction.new_P5(u_raw, v_raw, coefficients) return return_data
def _create_fake_test_data(self): import textwrap, time, random import polynomial_correction self.sdds_header = textwrap.dedent('''\ #SDDSASCIIFORMAT v1 #Beam: %s #Created: %s By: Python unittest (NL correction) #bunchid :0 #number of turns :%d #number of monitors :%d\n''' % ( self.beam_name, time.strftime('%Y-%m-%d#%H-%M-%S', time.localtime()), self.number_of_turns, self.number_of_bpms ) ) # first create raw data and write to sdds file with open(self.valid_raw_path, 'w') as sdds_raw: # write header sdds_raw.write(self.sdds_header) # write data for bpm_num in range(self.number_of_bpms): random_bpm_position = random.random() * 27e3 # about 27km random_bpm_type = random.choice(self.bpm_types) bpm_name = '%s.%d.%s' % (random_bpm_type, bpm_num, self.beam_name) self.bpm_positions[bpm_name] = random_bpm_position for plane in ('X', 'Y'): if plane == 'X': plane_num = 0 else: plane_num = 1 if bpm_name not in self.raw_data[plane]: self.raw_data[plane][bpm_name] = [] sdds_raw.write('%d %s %.5f ' % (plane_num, bpm_name, random_bpm_position)) for turn_num in range(self.number_of_turns): random_amp = (random.random()*2 - 1)*self.max_raw_amp self.raw_data[plane][bpm_name].append(random_amp) format_string = '%.5f ' if turn_num == self.number_of_turns - 1: format_string = '%.5f' # to avoid ' ' at the end of each line sdds_raw.write(format_string % random_amp) # random number in (-self.max_raw_amp, +self.max_raw_amp) sdds_raw.write('\n') # now we create old correction file with this data with open(self.valid_old_correction_path, 'w') as sdds_old: # write header sdds_old.write(self.sdds_header) # write data for plane in self.raw_data: if plane == 'X': plane_num = 0 else: plane_num = 1 for bpm_name in self.raw_data[plane]: bpm_type = bpm_name.split('.')[0] bpm_position = self.bpm_positions[bpm_name] sdds_old.write('%d %s %.5f ' % (plane_num, bpm_name, bpm_position)) for turn_num in range(len(self.raw_data[plane][bpm_name])): coefficients = polynomial_correction.get_coefficients_of_bpm_type(bpm_type, polynomial_correction.old_coefficients[self.year_for_old_correction]) turn_amp = self.raw_data[plane][bpm_name][turn_num] turn_amp *= coefficients[0] turn_amp = polynomial_correction._old_P5(turn_amp, coefficients) format_string = '%.5f ' if turn_num == len(self.raw_data[plane][bpm_name]) - 1: format_string = '%.5f' # to avoid ' ' at the end of each line sdds_old.write(format_string % turn_amp) # u_old = P5_old(kf * u_raw) sdds_old.write('\n') # new correction file with this data with open(self.valid_new_correction_path, 'w') as sdds_new: # write header sdds_new.write(self.sdds_header) # write data for plane in self.raw_data: if plane == 'X': plane_num = 0 else: plane_num = 1 for bpm_name in self.raw_data[plane]: bpm_type = bpm_name.split('.')[0] bpm_position = self.bpm_positions[bpm_name] sdds_new.write('%d %s %.5f ' % (plane_num, bpm_name, bpm_position)) for turn_num in range(len(self.raw_data[plane][bpm_name])): coefficients = polynomial_correction.get_coefficients_of_bpm_type(bpm_type, polynomial_correction.new_coefficients) turn_amp = self.raw_data[plane][bpm_name][turn_num] if plane == 'X': v_raw = self.raw_data['Y'][bpm_name][turn_num] else: v_raw = self.raw_data['X'][bpm_name][turn_num] turn_amp = polynomial_correction.new_P5(self.raw_data[plane][bpm_name][turn_num], v_raw, coefficients) format_string = '%.5f ' if turn_num == len(self.raw_data[plane][bpm_name]) - 1: format_string = '%.5f' # to avoid ' ' at the end of each line sdds_new.write(format_string % turn_amp) # u_new = P5_new(u_raw) sdds_new.write('\n')