def _get_network_emissions(self, original_aa_string, pssm, token): aa = original_aa_string if self.use_aa else -1 evo = pssm if self.use_pssm else -1 tok = token if self.use_token else -1 packed_input_sequences = self.embed(aa, evo, tok) minibatch_size = int(packed_input_sequences[1][0]) self.init_hidden(minibatch_size) if self.use_gpu: packed_input_sequences = packed_input_sequences.cuda() (data, bi_lstm_batches, _, _), self.hidden_layer = self.bi_lstm(packed_input_sequences, self.hidden_layer) emissions_padded, batch_sizes = torch.nn.utils.rnn.pad_packed_sequence( torch.nn.utils.rnn.PackedSequence(self.hidden_to_labels(data), bi_lstm_batches)) emissions = emissions_padded.transpose(0, 1).transpose( 1, 2) # minibatch_size, self.mixture_size, -1 emissions = self.batch_norm(emissions) emissions = emissions.transpose(1, 2).transpose( 0, 1) # (minibatch_size, -1, self.mixture_size) output_angles = self._dehidrals(emissions, self.alphabet) backbone_atoms_padded, _ = \ get_backbone_positions_from_angular_prediction(output_angles, batch_sizes, self.use_gpu) return output_angles, backbone_atoms_padded, batch_sizes
def forward(self, latent_space): #decoding. takes in a single latent code from a single part of the batch. # where input is the teacher forcing or the predictions from the previous step. # batch_first = True for this. prev_out, hidden = self.decoder(latent_space) #predict sequences: pred_seqs = F.log_softmax(self.dense2_post_dec(prev_out), dim=2) # alternatives to weird angle mixture model. # # just predicting dihedrals directly. Then try to scale them to be between +/-pi # # then try using atan2 #np.pi* F.tanh( output_angles = self.latent_to_dihedral2( F.elu(self.latent_to_dihedral1(prev_out))).permute( [1, 0, 2]) # max size, minibatch size, 3 (angels) #print('output angles shape::: ', output_angles.shape) ###print('output angles::: ', output_angles) # weird angle mixture model thing. '''x = self.hidden_to_labels(prev_out) x = self.bn(x.permute([0,2,1])).permute([0,2,1]).contiguous() #x = x.transpose(1,2) #(minibatch_size, -1, self.mixture_size) p = torch.exp(self.soft(x)) output_angles = self.softmax_to_angle(p).transpose(0,1) # max size, minibatch size, 3 (angels)''' # used to feed in batch sizes here. could do so from encoder. but all I do I take the length of it... backbone_atoms_padded, batch_sizes_backbone = get_backbone_positions_from_angular_prediction( output_angles, latent_space.shape[0], self.device) if torch.isnan(backbone_atoms_padded).sum() > 0: print('angles are NOT VALID!!!++++++++++++++++++++++++') return pred_seqs, output_angles, backbone_atoms_padded, batch_sizes_backbone
def _get_network_emissions(self, original_aa_string): # (seqlen x features), (seqlen x batch x features ) packed_input_sequences, input_sequences, batch_sizes = self.embed( original_aa_string) # (batch x seqlen x features ) reshaped_packed_input_sequences = input_sequences.data.permute(1, 0, 2) # Turn (batch x seqlen x features) into (batch x features x seqlen) for CNN flipped_reshaped_packed_input_sequences = reshaped_packed_input_sequences.transpose( 1, 2) batch_size, embedding_size, sequence_length = flipped_reshaped_packed_input_sequences.shape # why init in every fwd pass? Is it stateless? # if self.minibatch_size != minibatch_size: # self.minibatch_size = minibatch_size # self.init_hidden(self.minibatch_size) # else: # self.retain_hidden() # stateless self.init_hidden(sequence_length) x_conv = self.conv_init(flipped_reshaped_packed_input_sequences) x_conv = x_conv.transpose(1, 2) x_lstm = torch.cat( (x_conv, reshaped_packed_input_sequences), 2) # Concatenate original input and output CNN to LSTM # performing the fwd pass data, self.hidden_layer = self.bi_lstm(x_lstm) x_total = torch.cat((data, x_conv), 2) emissions = self.hidden_to_labels(x_total) emissions = emissions.view(batch_size, self.mixture_size, -1) emissions = self.batch_norm(emissions) emissions = emissions.view(batch_size, -1, self.mixture_size) probabilities = torch.exp(self.soft(emissions)) output_angles = self.softmax_to_angle(probabilities) \ .transpose(0, 1) # max size, minibatch size, 3 (angles) assert not np.isnan(data.cpu().detach().numpy()).any() backbone_atoms_padded, _ = \ get_backbone_positions_from_angular_prediction(output_angles, batch_sizes, self.use_gpu) assert not np.isnan(backbone_atoms_padded.cpu().detach().numpy()).any() return output_angles, backbone_atoms_padded, batch_sizes
def _get_network_emissions(self, original_aa_string): packed_input_sequences, input_sequences, batch_sizes = self.embed( original_aa_string) minibatch_size = int(packed_input_sequences[1][0]) # why init in every fwd pass? Is it stateless? if self.minibatch_size != minibatch_size: self.minibatch_size = minibatch_size self.init_hidden(self.minibatch_size) else: self.retain_hidden() # stateless # self.init_hidden(minibatch_size) # performing the fwd pass (data, bi_lstm_batches, _, _), self.hidden_layer = self.bi_lstm(packed_input_sequences, self.hidden_layer) assert not np.isnan(data.cpu().detach().numpy()).any() emissions_padded, batch_sizes = torch.nn.utils.rnn.pad_packed_sequence( torch.nn.utils.rnn.PackedSequence(self.hidden_to_labels(data), bi_lstm_batches)) assert not np.isnan(emissions_padded.cpu().detach().numpy()).any() emissions = emissions_padded.transpose(0, 1)\ .transpose(1, 2) # minibatch_size, self.mixture_size, -1 emissions = self.batch_norm(emissions) emissions = emissions.transpose( 1, 2) # (minibatch_size, -1, self.mixture_size) probabilities = torch.exp(self.soft(emissions)) output_angles = self.softmax_to_angle(probabilities)\ .transpose(0, 1) # max size, minibatch size, 3 (angles) assert not np.isnan(output_angles.cpu().detach().numpy()).any() backbone_atoms_padded, _ = \ get_backbone_positions_from_angular_prediction(output_angles, batch_sizes, self.use_gpu) assert not np.isnan(backbone_atoms_padded.cpu().detach().numpy()).any() return output_angles, backbone_atoms_padded, batch_sizes
def _get_network_emissions(self, original_aa_string, pssm=-1): packed_input_sequences = self.embed(original_aa_string) minibatch_size = int(packed_input_sequences[1][0]) self.init_hidden(minibatch_size) (data, bi_lstm_batches, _, _), self.hidden_layer = self.bi_lstm( packed_input_sequences, self.hidden_layer) emissions_padded, batch_sizes = torch.nn.utils.rnn.pad_packed_sequence( torch.nn.utils.rnn.PackedSequence(self.hidden_to_labels(data), bi_lstm_batches)) emissions = emissions_padded.transpose(0, 1)\ .transpose(1, 2) # minibatch_size, self.mixture_size, -1 emissions = self.batch_norm(emissions) emissions = emissions.transpose(1, 2) # (minibatch_size, -1, self.mixture_size) probabilities = torch.exp(self.soft(emissions)) output_angles = self.softmax_to_angle(probabilities)\ .transpose(0, 1) # max size, minibatch size, 3 (angles) backbone_atoms_padded, _ = \ get_backbone_positions_from_angular_prediction(output_angles, batch_sizes, self.use_gpu) return output_angles, backbone_atoms_padded, batch_sizes
def _get_network_emissions(self, original_aa_string, pssm, primary_token): # set input aa = original_aa_string if self.use_aa else -1 evo = pssm if self.use_pssm else -1 tok = primary_token if self.use_token else -1 packed_input_sequences = self.embed(aa, evo, tok) minibatch_size = int(packed_input_sequences[1][0]) if self.src_mask is None or self.src_mask.size( 0) != packed_input_sequences[1].size(0): mask = self._generate_square_subsequent_mask( packed_input_sequences[1].size(0)).to(self.device) self.src_mask = mask # transfomer state = self.W(packed_input_sequences[0].to(self.device)) state = F.relu(state) state, batch_sizes = torch.nn.utils.rnn.pad_packed_sequence( torch.nn.utils.rnn.PackedSequence(state, packed_input_sequences[1])) positional_encodings = self.pos_encoder(state) output = self.transformer_encoder(positional_encodings, self.src_mask) # convert internal representation to label output_angles = self._dehidrals( output, self.alphabet, ) # coordinate backbone_atoms_padded, _ = get_backbone_positions_from_angular_prediction( output_angles, batch_sizes, self.use_gpu) return output_angles, backbone_atoms_padded, batch_sizes
def process_file(input_file, output_file, use_gpu): print("Processing raw data file", input_file) # create output file f = h5py.File(output_file, 'w') current_buffer_size = 1 current_buffer_allocation = 0 dset1 = f.create_dataset('primary',(current_buffer_size,MAX_SEQUENCE_LENGTH),maxshape=(None,MAX_SEQUENCE_LENGTH),dtype='int32') dset2 = f.create_dataset('tertiary',(current_buffer_size,MAX_SEQUENCE_LENGTH,9),maxshape=(None,MAX_SEQUENCE_LENGTH, 9),dtype='float') dset3 = f.create_dataset('mask',(current_buffer_size,MAX_SEQUENCE_LENGTH),maxshape=(None,MAX_SEQUENCE_LENGTH),dtype='uint8') input_file_pointer = open("raw_data/casp11/" + input_file, "r") while True: # while there's more proteins to process print(input_file_pointer) next_protein = read_protein_from_file(input_file_pointer) if next_protein is None: break sequence_length = len(next_protein['primary']) if sequence_length > MAX_SEQUENCE_LENGTH: print("Dropping protein as length too long:", sequence_length) continue if current_buffer_allocation >= current_buffer_size: current_buffer_size = current_buffer_size + 1 dset1.resize((current_buffer_size,MAX_SEQUENCE_LENGTH)) dset2.resize((current_buffer_size,MAX_SEQUENCE_LENGTH, 9)) dset3.resize((current_buffer_size,MAX_SEQUENCE_LENGTH)) primary_padded = np.zeros(MAX_SEQUENCE_LENGTH) tertiary_padded = np.zeros((9, MAX_SEQUENCE_LENGTH)) mask_padded = np.zeros(MAX_SEQUENCE_LENGTH) # masking and padding here happens so that the stored dataset is of the same size. # when the data is loaded in this padding is removed again. primary_padded[:sequence_length] = next_protein['primary'] t_transposed = np.ravel(np.array(next_protein['tertiary']).T) t_reshaped = np.reshape(t_transposed, (sequence_length,9)).T tertiary_padded[:,:sequence_length] = t_reshaped mask_padded[:sequence_length] = next_protein['mask'] mask = torch.Tensor(mask_padded).type(dtype=torch.uint8) prim = torch.masked_select(torch.Tensor(primary_padded).type(dtype=torch.long), mask) pos = torch.masked_select(torch.Tensor(tertiary_padded), mask).view(9, -1).transpose(0, 1).unsqueeze(1) / 100 if use_gpu: pos = pos.cuda() angles, batch_sizes = calculate_dihedral_angles_over_minibatch(pos, [len(prim)], use_gpu=use_gpu) tertiary, _ = get_backbone_positions_from_angular_prediction(angles, batch_sizes, use_gpu=use_gpu) tertiary = tertiary.squeeze(1) primary_padded = np.zeros(MAX_SEQUENCE_LENGTH) tertiary_padded = np.zeros((MAX_SEQUENCE_LENGTH, 9)) length_after_mask_removed = len(prim) primary_padded[:length_after_mask_removed] = prim.data.cpu().numpy() tertiary_padded[:length_after_mask_removed, :] = tertiary.data.cpu().numpy() mask_padded = np.zeros(MAX_SEQUENCE_LENGTH) mask_padded[:length_after_mask_removed] = np.ones(length_after_mask_removed) dset1[current_buffer_allocation] = primary_padded dset2[current_buffer_allocation] = tertiary_padded dset3[current_buffer_allocation] = mask_padded current_buffer_allocation += 1 print("Wrote output to", current_buffer_allocation, "proteins to", output_file)
def _get_network_emissions(self, original_aa_string): # (seqlen x features), (seqlen x batch x features ) packed_input_sequences, input_sequences, batch_sizes = self.embed( original_aa_string) # (batch x seqlen x features ) reshaped_packed_input_sequences = input_sequences.data.permute(1, 0, 2) # Turn (batch x seqlen x features) into (batch x features x seqlen) for CNN flipped_reshaped_packed_input_sequences = reshaped_packed_input_sequences.transpose( 1, 2) batch_size, embedding_size, sequence_length = flipped_reshaped_packed_input_sequences.shape # performing the fwd pass init = F.leaky_relu( self.conv_bn_init( self.conv_init(flipped_reshaped_packed_input_sequences))) res = self.conv1(F.leaky_relu(self.conv_bn1(init))) res = self.conv2(F.leaky_relu(self.conv_bn2(res))) x_conv = res + init res = self.conv4(F.leaky_relu(self.conv_bn4(x_conv))) res = self.conv5(F.leaky_relu(self.conv_bn5(res))) x_conv = self.conv6_expand(x_conv) x_conv += res res = self.conv7(F.leaky_relu(self.conv_bn7(x_conv))) res = self.conv8(F.leaky_relu(self.conv_bn8(res))) x_conv = self.conv9_expand(x_conv) x_conv += res res = self.conv10(F.leaky_relu(self.conv_bn10(x_conv))) res = self.conv11(F.leaky_relu(self.conv_bn11(res))) x_conv = self.conv12_expand(x_conv) x_conv += res res = self.conv13(F.leaky_relu(self.conv_bn13(x_conv))) res = self.conv14(F.leaky_relu(self.conv_bn14(res))) x_conv = self.conv15_expand(x_conv) x_conv += res res = self.conv16(F.leaky_relu(self.conv_bn16(x_conv))) res = self.conv17(F.leaky_relu(self.conv_bn17(res))) x_conv = self.conv18_expand(x_conv) x_conv += res res = self.conv19(F.leaky_relu(self.conv_bn19(x_conv))) res = self.conv20(F.leaky_relu(self.conv_bn20(res))) x_conv = self.conv21_expand(x_conv) x_conv += res res = self.conv22(F.leaky_relu(self.conv_bn22(x_conv))) res = self.conv23(F.leaky_relu(self.conv_bn23(res))) x_conv = self.conv24_expand(x_conv) x_conv += res x_conv = x_conv.transpose(1, 2) x_lstm = torch.cat( (x_conv, reshaped_packed_input_sequences), 2) # Concatenate original input and output CNN to LSTM # why init in every fwd pass? Is it stateless? # if self.minibatch_size != minibatch_size: # self.minibatch_size = minibatch_size # self.init_hidden(self.minibatch_size) # else: # self.retain_hidden() # stateless: self.init_hidden(sequence_length) data, self.hidden_layer = self.bi_lstm(x_lstm) x_total = torch.cat((data, x_conv), 2) emissions = self.hidden_to_labels(x_total) emissions = emissions.view(batch_size, self.mixture_size, -1) emissions = self.batch_norm(emissions) emissions = emissions.view(batch_size, -1, self.mixture_size) probabilities = torch.exp(self.soft(emissions)) output_angles = self.softmax_to_angle(probabilities) \ .transpose(0, 1) # max size, minibatch size, 3 (angles) assert not np.isnan(data.cpu().detach().numpy()).any() backbone_atoms_padded, _ = \ get_backbone_positions_from_angular_prediction(output_angles, batch_sizes, self.use_gpu) assert not np.isnan(backbone_atoms_padded.cpu().detach().numpy()).any() return output_angles, backbone_atoms_padded, batch_sizes
def process_file(input_file, output_file, use_gpu): print("Processing raw data file", input_file) # create output file f = h5py.File(output_file, 'w') current_buffer_size = 1 current_buffer_allocation = 0 dset1 = f.create_dataset( 'primary', (current_buffer_size, MAX_SEQUENCE_LENGTH), maxshape=(None, MAX_SEQUENCE_LENGTH), dtype='int32' ) # creates an empty dataset with given dimension, axes with none in maxshape are unlimited dset2 = f.create_dataset('tertiary', (current_buffer_size, MAX_SEQUENCE_LENGTH, 9), maxshape=(None, MAX_SEQUENCE_LENGTH, 9), dtype='float') dset3 = f.create_dataset('mask', (current_buffer_size, MAX_SEQUENCE_LENGTH), maxshape=(None, MAX_SEQUENCE_LENGTH), dtype='uint8') input_file_pointer = open("data/raw/" + input_file, "r") while True: # while there's more proteins to process next_protein = read_protein_from_file(input_file_pointer) if next_protein is None: break sequence_length = len(next_protein['primary']) if sequence_length > MAX_SEQUENCE_LENGTH: print("Dropping protein as length too long:", sequence_length) continue if current_buffer_allocation >= current_buffer_size: current_buffer_size = current_buffer_size + 1 dset1.resize((current_buffer_size, MAX_SEQUENCE_LENGTH)) dset2.resize((current_buffer_size, MAX_SEQUENCE_LENGTH, 9)) dset3.resize((current_buffer_size, MAX_SEQUENCE_LENGTH)) primary_padded = np.zeros(MAX_SEQUENCE_LENGTH) tertiary_padded = np.zeros((9, MAX_SEQUENCE_LENGTH)) mask_padded = np.zeros(MAX_SEQUENCE_LENGTH) # masking and padding here happens so that the stored dataset is of the same size. # when the data is loaded in this padding is removed again. primary_padded[:sequence_length] = next_protein['primary'] t_transposed = np.ravel(np.array( next_protein['tertiary']).T) # flattens the array into 1-D t_reshaped = np.reshape(t_transposed, (sequence_length, 9)).T tertiary_padded[:, :sequence_length] = t_reshaped mask_padded[:sequence_length] = next_protein['mask'] mask = torch.Tensor(mask_padded).type(dtype=torch.uint8) prim = torch.masked_select( torch.Tensor(primary_padded).type(dtype=torch.long), mask ) # only leaves those aa which have + (actually 0) in their mask pos = torch.masked_select(torch.Tensor(tertiary_padded), mask).view( 9, -1 ).transpose(0, 1).unsqueeze( 1 ) / 100 # divides by 100 because all values are artificially increased by 100 as specified in the proteinnet documentation, do not know yet why we need to add an additional dimension??? if use_gpu: pos = pos.cuda() angles, batch_sizes = calculate_dihedral_angles_over_minibatch( pos, [len(prim)], use_gpu=use_gpu) tertiary, _ = get_backbone_positions_from_angular_prediction( angles, batch_sizes, use_gpu=use_gpu) tertiary = tertiary.squeeze(1) primary_padded = np.zeros(MAX_SEQUENCE_LENGTH) tertiary_padded = np.zeros((MAX_SEQUENCE_LENGTH, 9)) length_after_mask_removed = len(prim) primary_padded[:length_after_mask_removed] = prim.data.cpu().numpy() tertiary_padded[:length_after_mask_removed, :] = tertiary.data.cpu( ).numpy() mask_padded = np.zeros(MAX_SEQUENCE_LENGTH) mask_padded[:length_after_mask_removed] = np.ones( length_after_mask_removed) dset1[current_buffer_allocation] = primary_padded dset2[current_buffer_allocation] = tertiary_padded dset3[current_buffer_allocation] = mask_padded current_buffer_allocation += 1 print("Wrote output to", current_buffer_allocation, "proteins to", output_file)
def process_file(input_file, output_file, use_gpu, max_sequence_length, use_mask=True, vocab='iupac'): print("Processing raw data file", input_file) # set tokenizer tokenizer = TAPETokenizer(vocab=vocab) # create output file file = h5py.File(output_file, 'w') current_buffer_size = 1 current_buffer_allocation = 0 dset1 = file.create_dataset('primary', (current_buffer_size, max_sequence_length), maxshape=(None, max_sequence_length), dtype='int32') dset2 = file.create_dataset('tertiary', (current_buffer_size, max_sequence_length, 9), maxshape=(None, max_sequence_length, 9), dtype='float') dset3 = file.create_dataset('mask', (current_buffer_size, max_sequence_length), maxshape=(None, max_sequence_length), dtype='uint8') dset4 = file.create_dataset('pssm', (current_buffer_size, max_sequence_length, 21), maxshape=(None, max_sequence_length, 21), dtype='float') dset5 = file.create_dataset('primary_token', (current_buffer_size, 2 * max_sequence_length), maxshape=(None, 2 * max_sequence_length), dtype='int32') input_file_pointer = open("data/raw/" + input_file, "r") while True: # while there's more proteins to process next_protein = read_protein_from_file(input_file_pointer) if next_protein is None: break sequence_length = len(next_protein['primary']) if sequence_length > max_sequence_length: # print("Dropping protein as length too long:", sequence_length) continue print("Process protein with length", sequence_length) if current_buffer_allocation >= current_buffer_size: current_buffer_size = current_buffer_size + 1 dset1.resize((current_buffer_size, max_sequence_length)) dset2.resize((current_buffer_size, max_sequence_length, 9)) dset3.resize((current_buffer_size, max_sequence_length)) dset4.resize((current_buffer_size, max_sequence_length, 21)) dset5.resize((current_buffer_size, 2 * max_sequence_length)) primary_padded = np.zeros(max_sequence_length) tertiary_padded = np.zeros((9, max_sequence_length)) mask_padded = np.zeros(max_sequence_length) pssm_padded = np.zeros((21, max_sequence_length)) primary_token_padded = np.zeros(2 * max_sequence_length) # masking and padding here happens so that the stored dataset is of the same size. # when the data is loaded in this padding is removed again. primary_padded[:sequence_length] = next_protein['primary'] t_transposed = np.ravel(np.array(next_protein['tertiary']).T) t_reshaped = np.reshape(t_transposed, (sequence_length, 9)).T tertiary_padded[:, :sequence_length] = t_reshaped mask_padded[:sequence_length] = next_protein['mask'] pssm_padded[:, :sequence_length] = np.array(next_protein['evolutionary']) if use_mask: mask = torch.Tensor(mask_padded).type(dtype=torch.bool) prim = torch.masked_select(torch.Tensor(primary_padded) .type(dtype=torch.long), mask) seq_token = torch.Tensor(tokenization(tokenizer, next_protein['seq'], next_protein['mask'])) pos = torch.masked_select(torch.Tensor(tertiary_padded), mask) \ .view(9, -1).transpose(0, 1).unsqueeze(1) / 100 pssm = torch.masked_select(torch.Tensor(pssm_padded), mask).view(21, -1).transpose(0, 1) if use_gpu: pos = pos.cuda() angles, batch_sizes = calculate_dihedral_angles_over_minibatch(pos, [len(prim)], use_gpu=use_gpu) tertiary, _ = get_backbone_positions_from_angular_prediction(angles, batch_sizes, use_gpu=use_gpu) tertiary = tertiary.squeeze(1) primary_padded = np.zeros(max_sequence_length) tertiary_padded = np.zeros((max_sequence_length, 9)) pssm_padded = np.zeros((max_sequence_length, 21)) length_after_mask_removed = len(prim) primary_padded[:length_after_mask_removed] = prim.data.cpu().numpy() primary_token_padded[:len(seq_token)] = seq_token.cpu().numpy() tertiary_padded[:length_after_mask_removed, :] = tertiary.data.cpu().numpy() pssm_padded[:length_after_mask_removed, :] = pssm.data.cpu().numpy() mask_padded = np.zeros(max_sequence_length) mask_padded[:length_after_mask_removed] = np.ones(length_after_mask_removed) dset1[current_buffer_allocation] = primary_padded dset2[current_buffer_allocation] = tertiary_padded dset3[current_buffer_allocation] = mask_padded dset4[current_buffer_allocation] = pssm_padded dset5[current_buffer_allocation] = primary_token_padded current_buffer_allocation += 1 print("Wrote output to", current_buffer_allocation, "proteins to", output_file)
def process_file(input_file, output_file, device, want_trimmed, want_pure): print("Processing raw data file", input_file) # create output file f = h5py.File(output_file, 'w') current_buffer_size = 1 current_buffer_allocation = 0 dset1 = f.create_dataset('primary', (current_buffer_size, MAX_SEQUENCE_LENGTH), maxshape=(None, MAX_SEQUENCE_LENGTH), dtype='int32') dset2 = f.create_dataset('tertiary', (current_buffer_size, MAX_SEQUENCE_LENGTH, 9), maxshape=(None, MAX_SEQUENCE_LENGTH, 9), dtype='float') dset3 = f.create_dataset('mask', (current_buffer_size, MAX_SEQUENCE_LENGTH), maxshape=(None, MAX_SEQUENCE_LENGTH), dtype='uint8') dset4 = f.create_dataset('padding_mask', (current_buffer_size, MAX_SEQUENCE_LENGTH), maxshape=(None, MAX_SEQUENCE_LENGTH), dtype='uint8') input_file_pointer = open("data/raw/" + input_file, "r") while True: # while there's more proteins to process next_protein = read_protein_from_file(input_file_pointer) if next_protein is None: break sequence_length = len(next_protein['primary']) if sequence_length > MAX_SEQUENCE_LENGTH: print("Dropping protein as length too long:", sequence_length) continue if want_pure: unpadded_mask = torch.Tensor( next_protein['mask']).type(dtype=torch.uint8) if unpadded_mask.sum() != unpadded_mask.shape[0]: print('dropping protein, mask has holes') continue elif want_trimmed: s = [str(i) for i in next_protein['mask']] s.append('0') # needed for those that dont end with a mask spot!. res = "".join(s) if len(res.split('10')) > 2: print('dropping protein, mask isnt just on edges') continue primary_padded = np.zeros(MAX_SEQUENCE_LENGTH) tertiary_padded = np.zeros((9, MAX_SEQUENCE_LENGTH)) mask_padded = np.zeros(MAX_SEQUENCE_LENGTH) # masking and padding here happens so that the stored dataset is of the same size. # when the data is loaded in this padding is removed again. primary_padded[:sequence_length] = next_protein['primary'] t_transposed = np.ravel(np.array(next_protein['tertiary']).T) t_reshaped = np.reshape(t_transposed, (sequence_length, 9)).T tertiary_padded[:, :sequence_length] = t_reshaped mask_padded[:sequence_length] = next_protein['mask'] mask = torch.Tensor(mask_padded).type(dtype=torch.uint8) prim = torch.masked_select( torch.Tensor(primary_padded).type(dtype=torch.long), mask) pos = torch.masked_select(torch.Tensor(tertiary_padded), mask).view( 9, -1).transpose(0, 1).unsqueeze(1) / 100 pos = pos.to(device) angles, batch_sizes = calculate_dihedral_angles_over_minibatch( pos, [len(prim)], device) # this must be what is creating the nans!! Not clear to me why... tertiary, _ = get_backbone_positions_from_angular_prediction( angles, batch_sizes, device) tertiary = tertiary.squeeze(1) if torch.isnan(tertiary).sum() > 0: print('there is a nan in tertiary! Dropping and printing mask') print(next_protein['mask']) continue primary_padded = np.zeros(MAX_SEQUENCE_LENGTH) tertiary_padded = np.zeros((MAX_SEQUENCE_LENGTH, 9)) length_after_mask_removed = len(prim) if sequence_length == 0: print( 'sequence length is zero after mask was applied. Dropping! ==========' ) continue #print('final size', length_after_mask_removed) #print('tertiary ', tertiary) primary_padded[:length_after_mask_removed] = prim.data.cpu().numpy() tertiary_padded[:length_after_mask_removed, :] = tertiary.data.cpu( ).numpy() mask_padded = np.zeros(MAX_SEQUENCE_LENGTH) mask_padded[:length_after_mask_removed] = np.ones( length_after_mask_removed) padding_mask_padded = np.zeros(MAX_SEQUENCE_LENGTH) # this mask has masking for both the padding and the AAs without angle data! # # THIS HAS BECOME COMPLETELY IRRELEVANT NOW THAT I AM GETTING RID OF ANY MISSING THINGS FIRST! padding_mask_padded[:length_after_mask_removed] = np.ones( length_after_mask_removed) if current_buffer_allocation >= current_buffer_size: current_buffer_size = current_buffer_size + 1 dset1.resize((current_buffer_size, MAX_SEQUENCE_LENGTH)) dset2.resize((current_buffer_size, MAX_SEQUENCE_LENGTH, 9)) dset3.resize((current_buffer_size, MAX_SEQUENCE_LENGTH)) dset4.resize((current_buffer_size, MAX_SEQUENCE_LENGTH)) dset1[current_buffer_allocation] = primary_padded dset2[current_buffer_allocation] = tertiary_padded dset3[current_buffer_allocation] = mask_padded dset4[current_buffer_allocation] = padding_mask_padded current_buffer_allocation += 1 print("Wrote output to", current_buffer_allocation, "proteins to", output_file)