def _score_spectrum(self, precursor_mass, spectrum_original, state0_c, state0_h, candidate_list, model, model_output_logprob, model_lstm_state, session, direction): """TODO(nh2tran): docstring.""" #~ print("".join(["="] * 80)) # section-separating line #~ print("WorkerDB: _score()") # convert symbols into id candidate_list = [[deepnovo_config.vocab[x] for x in candidate] for candidate in candidate_list] # we shall group candidates into minibatches # === candidate_len === # s # i # z # e # ===================== minibatch_size = len(candidate_list) # number of candidates candidate_len = len(candidate_list[0]) # length of each candidate # candidates share the same state0, so repeat into [minibatch_size, 512] # the states will also be updated after every iteration state0_c = state0_c.reshape((1, -1)) # reshape to [1, 512] state0_h = state0_h.reshape((1, -1)) minibatch_state_c = np.repeat(state0_c, minibatch_size, axis=0) minibatch_state_h = np.repeat(state0_h, minibatch_size, axis=0) # mass of each candidate, will be accumulated everytime an AA is appended minibatch_prefix_mass = np.zeros(minibatch_size) # output is a list of candidate_len arrays of shape [minibatch_size, 26] # each row is log of probability distribution over 26 classes/symbols output_logprob_list = [] # recurrent iterations for position in range(candidate_len): # gather minibatch data minibatch_AA_id = np.zeros(minibatch_size) for index, candidate in enumerate(candidate_list): AA = candidate[position] minibatch_AA_id[index] = AA minibatch_prefix_mass[index] += deepnovo_config.mass_ID[AA] # this is the most time-consuming ~70-75% minibatch_intensity = [get_candidate_intensity(spectrum_original, precursor_mass, prefix_mass, direction) for prefix_mass in np.nditer(minibatch_prefix_mass)] # final shape [minibatch_size, 26, 8, 10] minibatch_intensity = np.array(minibatch_intensity) # model feed input_feed = {} input_feed[model.input_dict["AAid"][1].name] = minibatch_AA_id input_feed[model.input_dict["intensity"].name] = minibatch_intensity input_feed[model.input_dict["lstm_state"][0].name] = minibatch_state_c input_feed[model.input_dict["lstm_state"][1].name] = minibatch_state_h # and run output_feed = [model_output_logprob, model_lstm_state] output_logprob, (minibatch_state_c, minibatch_state_h) = session.run( fetches=output_feed, feed_dict=input_feed) output_logprob_list.append(output_logprob) return output_logprob_list
def _extend_peak(self, direction, session, model, spectrum_batch, peak_batch): """TODO(nh2tran): docstring. Inputs: spectrum_batch: a list of spectrum, each is a dictionary spectrum["scan"] spectrum["precursor_mass"] spectrum["spectrum_holder"] spectrum["spectrum_original_forward"] spectrum["spectrum_original_backward"] peak_batch: one peak for each spectrum, each peak is a dictionary peak["prefix_mass"] for extension in the forward direction peak["sufffix_mass"] for extension in the backward direction peak["mass_tolerance"] Outputs: top_path_batch: for every input spectrum, the output is a list of paths, each path is a dictionary path["AAid_list"] path["score_list"] path["score_sum"] """ print( "WorkerDenovo: _extend_peak(), direction={0:s}".format(direction)) # test running time and tensorflow time test_time_decode = 0.0 test_time_tf = 0.0 test_time = 0.0 start_time_decode = time.time() # for every input spectrum, the output is a list of paths, # each path is a dictionary # path["AAid_list"] # path["score_list"] # path["score_sum"] spectrum_batch_size = len(spectrum_batch) top_path_batch = [[] for x in range(spectrum_batch_size)] # forward/backward direction setting # the direction determines the model, the spectrum and the peak mass if direction == "forward": model_lstm_state0 = model.output_forward["lstm_state0"] model_output_log_prob = model.output_forward["logprob"] model_lstm_state = model.output_forward["lstm_state"] spectrum_original_name = "spectrum_original_forward" peak_mass_name = "prefix_mass" FIRST_LABEL = self.GO_ID LAST_LABEL = self.EOS_ID elif direction == "backward": model_lstm_state0 = model.output_backward["lstm_state0"] model_output_log_prob = model.output_backward["logprob"] model_lstm_state = model.output_backward["lstm_state"] spectrum_original_name = "spectrum_original_backward" peak_mass_name = "suffix_mass" FIRST_LABEL = self.EOS_ID LAST_LABEL = self.GO_ID # PEAK EXTENSION includes 4 steps: # STEP 1: initialize the lstm and the active_search_list. # STEP 2, 3, 4 are repeated until the active_search_list is empty. # STEP 2: gather data from active search entries and group into blocks. # STEP 3: run tensorflow model on data blocks to predict next AA. # STEP 4: retrieve data from blocks to update the active_search_list # with knapsack dynamic programming and beam search. start_time_tf = time.time() # STEP 1: initialize lstm spectrum_holder_array = np.array( [x["spectrum_holder"] for x in spectrum_batch]) input_feed = {} input_feed[model.input_dict["spectrum"].name] = spectrum_holder_array output_feed = model_lstm_state0 c_state0_array, h_state0_array = session.run(fetches=output_feed, feed_dict=input_feed) test_time_tf += time.time() - start_time_tf # STEP 1: initialize the active_search_list # active_search_list holds the info of search entries under processing # each search entry is a dictionary # search_entry["spectrum_id"] # search_entry["current_path_list"] # each path is also a dictionary # path["AAid_list"] # path["prefix_mass"] # path["score_list"] # path["score_sum"] # path["c_state"] # path["h_state"] active_search_list = [] for spectrum_id in range(spectrum_batch_size): search_entry = {} search_entry["spectrum_id"] = spectrum_id path = {} path["AAid_list"] = [FIRST_LABEL] path["prefix_mass"] = peak_batch[spectrum_id][peak_mass_name] path["score_list"] = [0.0] path["score_sum"] = 0.0 path["c_state"] = c_state0_array[spectrum_id] path["h_state"] = h_state0_array[spectrum_id] search_entry["current_path_list"] = [path] active_search_list.append(search_entry) # repeat STEP 2, 3, 4 until the active_search_list is empty. while True: # STEP 2: gather data from active search entries and group into blocks. # data blocks for the input feed of tensorflow model block_AAid_1 = [] # nobi block_AAid_2 = [] # nobi block_c_state = [] block_h_state = [] block_candidate_intensity = [] # data blocks to record the current status of search entries block_AAid_list = [] block_prefix_mass = [] block_score_list = [] block_score_sum = [] block_knapsack_candidate = [] # store the number of paths of each search entry in the big blocks # to retrieve the info of each search entry later in STEP 4. search_entry_size = [0] * len(active_search_list) # gather data into blocks through 2 nested loops over active_search_list # and over current_path_list of each search_entry for entry_index, search_entry in enumerate(active_search_list): spectrum_id = search_entry["spectrum_id"] current_path_list = search_entry["current_path_list"] precursor_mass = spectrum_batch[spectrum_id]["precursor_mass"] spectrum_original = spectrum_batch[spectrum_id][ spectrum_original_name] peak_mass_tolerance = peak_batch[spectrum_id]["mass_tolerance"] for path in current_path_list: # keep track of the AA predicted in the previous iteration # for nobi (short k-mer) model, we will need 2 previous AA AAid_list = path["AAid_list"] AAid_2 = AAid_list[-1] if len(AAid_list) > 1: AAid_1 = AAid_list[-2] else: AAid_1 = AAid_2 # nobi # the current status of this path prefix_mass = path["prefix_mass"] score_list = path["score_list"] score_sum = path["score_sum"] c_state = path["c_state"] h_state = path["h_state"] # when we reach LAST_LABEL, check if the mass of predicted sequence # is close to the given precursor_mass: # if yes, send the current path to output # if not, skip the current path if AAid_2 == LAST_LABEL: # nobi if (abs(prefix_mass - precursor_mass) <= peak_mass_tolerance): top_path_batch[spectrum_id].append({ "AAid_list": AAid_list, "score_list": score_list, "score_sum": score_sum }) continue start_time = time.time() # get CANDIDATE INTENSITY to predict next AA # TODO(nh2tran): change direction from 0/1 to "forward"/"backward" direction_id = 0 if direction == "forward" else 1 candidate_intensity = get_candidate_intensity( spectrum_original, precursor_mass, prefix_mass, direction_id) test_time += time.time() - start_time # use knapsack and SUFFIX MASS to filter next AA candidate suffix_mass = precursor_mass - prefix_mass - self.mass_ID[ LAST_LABEL] knapsack_tolerance = int( round(peak_mass_tolerance * self.KNAPSACK_AA_RESOLUTION)) knapsack_candidate = self._search_knapsack( suffix_mass, knapsack_tolerance) # if not possible to extend, add LAST_LABEL to end the sequence if not knapsack_candidate: knapsack_candidate.append(LAST_LABEL) # gather data blocks block_AAid_1.append(AAid_1) # nobi block_AAid_2.append(AAid_2) # nobi block_c_state.append(c_state) block_h_state.append(h_state) block_candidate_intensity.append(candidate_intensity) block_AAid_list.append(AAid_list) block_prefix_mass.append(prefix_mass) block_score_list.append(score_list) block_score_sum.append(score_sum) block_knapsack_candidate.append(knapsack_candidate) # record the size of each search entry in the blocks search_entry_size[entry_index] += 1 # STEP 3: run tensorflow model on data blocks to predict next AA. # output is stored in current_log_prob, current_c_state, current_h_state if block_AAid_1: start_time_tf = time.time() block_AAid_1 = np.array(block_AAid_1) # nobi block_AAid_2 = np.array(block_AAid_2) # nobi block_c_state = np.array(block_c_state) block_h_state = np.array(block_h_state) block_candidate_intensity = np.array(block_candidate_intensity) input_feed = {} input_feed[model.input_dict["AAid"] [0].name] = block_AAid_1 # nobi input_feed[model.input_dict["AAid"] [1].name] = block_AAid_2 # nobi input_feed[model.input_dict["lstm_state"] [0].name] = block_c_state input_feed[model.input_dict["lstm_state"] [1].name] = block_h_state input_feed[model.input_dict["intensity"]. name] = block_candidate_intensity output_feed = [model_output_log_prob, model_lstm_state] # lstm.len_full #~ output_feed = model_output_log_prob # nobi current_log_prob, (current_c_state, current_h_state) = session.run( output_feed, input_feed) # lstm.len_full #~ current_log_prob = session.run(output_feed,input_feed) # nobi test_time_tf += time.time() - start_time_tf # STEP 4: retrieve data from blocks to update the active_search_list # with knapsack dynamic programming and beam search. block_index = 0 for entry_index, search_entry in enumerate(active_search_list): # find all possible new paths within knapsack filter new_path_list = [] for index in range( block_index, block_index + search_entry_size[entry_index]): for AAid in block_knapsack_candidate[index]: new_path = {} new_path["AAid_list"] = block_AAid_list[index] + [AAid] new_path["prefix_mass"] = block_prefix_mass[ index] + self.mass_ID[AAid] if AAid > 2: # do NOT add score of GO, EOS, PAD new_path["score_list"] = ( block_score_list[index] + [current_log_prob[index][AAid]]) new_path["score_sum"] = ( block_score_sum[index] + current_log_prob[index][AAid]) else: new_path["score_list"] = block_score_list[index] new_path["score_sum"] = block_score_sum[index] new_path["c_state"] = current_c_state[ index] # lstm.len_full new_path["h_state"] = current_h_state[ index] # lstm.len_full #~ new_path["c_state"] = block_c_state[index] # nobi #~ new_path["h_state"] = block_h_state[index] # nobi new_path_list.append(new_path) # beam search to select top candidates if len(new_path_list) > self.beam_size: new_path_score = np.array( [x["score_sum"] for x in new_path_list]) top_k_index = np.argpartition(-new_path_score, self.beam_size)[:self.beam_size] # pylint: disable=line-too-long search_entry["current_path_list"] = [ new_path_list[top_k_index[x]] for x in range(self.beam_size) ] else: search_entry["current_path_list"] = new_path_list # update the accumulated block_index block_index += search_entry_size[entry_index] # update active_search_list by removing empty entries active_search_list = [ x for x in active_search_list if x["current_path_list"] ] # STOP the extension loop if active_search_list is empty if not active_search_list: break test_time_decode += time.time() - start_time_decode print(" test_time_tf = %.2f" % (test_time_tf)) print(" test_time_decode = %.2f" % (test_time_decode)) print(" test_time = %.2f" % (test_time)) return top_path_batch