text: import os
  import torch
  import torchvision.transforms as transforms
  import numpy as np
  from torch.utils.data import Dataset
  
  from srcOld.dataloader_a3m import Dataset_a3m
  from srcOld.dataloader_lmdb import Dataset_lmdb
  from srcOld.dataloader_pnet import Dataset_pnet
  from srcOld.dataloader_utils import SeqFlip, ListToNumpy, ConvertPnetFeaturesTo2D, ConvertCoordToDists, \
      ConvertDistAnglesToBins, ConvertPnetFeaturesTo1D
  
  
  def select_dataset(path_train,path_test,type='2D',batch_size=1, network=None):
      '''
      This is a wrapper routine for various dataloaders.
      Currently supports:
          folders with a3m files/pdb.
          pnet files.
      #TODO implement a3m files without pdb for testing
      :param path: path of input folder/file
      :return:
      '''
  
      Flip = SeqFlip()
      if os.path.isdir(path_train):
          dataset_train = Dataset_a3m(path_train)
      elif os.path.isfile(path_train) and os.path.splitext(path_train)[1].lower() == '.pnet':
          if type == '2D':
              transform_train = transforms.Compose([Flip, ListToNumpy(), ConvertPnetFeaturesTo2D()])
              transform_target_train = transforms.Compose([Flip, ListToNumpy(), ConvertCoordToDists()])
              transform_mask_train = transforms.Compose([Flip, ListToNumpy()])
          elif type == '1D':
              transform_train = transforms.Compose([Flip, ListToNumpy(), ConvertPnetFeaturesTo1D()])
              transform_target_train = transforms.Compose([Flip, ListToNumpy(), ConvertCoordToDists()])
              transform_mask_train = transforms.Compose([Flip, ListToNumpy()])
  
  
          dataset_train = Dataset_pnet(path_train, transform=transform_train,transform_target=transform_target_train,transform_mask=transform_mask_train)
      elif os.path.isfile(path_train) and os.path.splitext(path_train)[1].lower() == '.lmdb':
          if type == '2D':
              transform_train = transforms.Compose([Flip, ConvertPnetFeaturesTo2D()])
          elif type == '1D':
              transform_train = transforms.Compose([Flip, ConvertPnetFeaturesTo1D()])
          transform_target_train = transforms.Compose([Flip, ConvertCoordToDists()])
          transform_mask_train = transforms.Compose([Flip])
  
          dataset_train = Dataset_lmdb(path_train, transform=transform_train, target_transform=transform_target_train, mask_transform=transform_mask_train)
      else:
          raise NotImplementedError("dataset not implemented yet.")
      if os.path.isdir(path_test):
          dataset_test = Dataset_a3m(path_test)
      elif os.path.isfile(path_test) and os.path.splitext(path_test)[1].lower() == '.pnet':
          if type == '2D':
              transform_test = transforms.Compose([ListToNumpy(), ConvertPnetFeaturesTo2D()])
              transform_target_test = transforms.Compose([ListToNumpy(), ConvertCoordToDists()])
              transform_mask_test = transforms.Compose([ListToNumpy()])
          elif type == '1D':
              transform_test = transforms.Compose([Flip, ListToNumpy(), ConvertPnetFeaturesTo1D()])
              transform_target_test = transforms.Compose([Flip, ListToNumpy(), ConvertCoordToDists()])
              transform_mask_test = transforms.Compose([Flip, ListToNumpy()])
  
          dataset_test = Dataset_pnet(path_test, transform=transform_test,transform_target=transform_target_test, transform_mask=transform_mask_test)
      elif os.path.isfile(path_test) and os.path.splitext(path_test)[1].lower() == '.lmdb':
          if type == '2D':
              transform_test = transforms.Compose([ConvertPnetFeaturesTo2D()])
          elif type == '1D':
              transform_test = transforms.Compose([ConvertPnetFeaturesTo1D()])
          transform_target_test = transforms.Compose([ConvertCoordToDists()])
          dataset_test = Dataset_lmdb(path_test, transform=transform_test, target_transform=transform_target_test)
      else:
          raise NotImplementedError("dataset not implemented yet.")
  
      assert len(dataset_train) >= batch_size
  
      if network.lower() == 'vnet':
          pad_modulo = 8
      else:
          pad_modulo = 1
  
  
      dl_train = torch.utils.data.DataLoader(dataset_train, batch_size=batch_size, shuffle=True, num_workers=0, collate_fn=PadCollate(pad_modulo=pad_modulo),
                                             drop_last=True)
      dl_test = torch.utils.data.DataLoader(dataset_test, batch_size=min(batch_size,len(dataset_test)), shuffle=False, num_workers=0, collate_fn=PadCollate(pad_modulo=pad_modulo),
                                             drop_last=False)
  
      return dl_train, dl_test
  
  
  
  def pad_numpy_array_to_tensor(vec, pad, dim):
      """
      args:
          vec - tensor to pad
          pad - the size to pad to
          dim - dimension to pad
  
      return:
          a new tensor padded to 'pad' in dimension 'dim'
      """
      pad_size = list(vec.shape)
      pad_size[dim] = pad - vec.shape[dim]
      return torch.cat([torch.from_numpy(vec), torch.zeros(*pad_size)], dim=dim)
  
  
  class PadCollate:
      """
      a variant of callate_fn that pads according to the longest sequence in
      a batch of sequences
      """
  
      def __init__(self, dim=1, pad_modulo=1):
          """
          args:
              dim - the dimension to be padded (dimension of time in sequences)
          """
          self.dim = dim
          self.pad_mod = pad_modulo
  
      def pad_collate(self, batch):
          """
          args:
              batch - list of (tensor, label)
  
          return:
              xs - a tensor of all examples in 'batch' after padding
              ys - a LongTensor of all labels in batch
          """
          # find longest sequence
          nb = len(batch)
          nf = batch[0][0].shape[0]
          max_len = max(map(lambda x: x[0].shape[self.dim], batch))
          max_len = int(self.pad_mod * np.ceil(max_len / self.pad_mod))
  
          distances = ()
          nt = len(batch[0][1]) # Number of distograms
          tt = torch.empty((nt,nb,max_len,max_len),dtype=torch.float32)
  
          coords = ()
          nc = len(batch[0][2]) # Number of coordinates
          cc = torch.empty((nc,nb,3,max_len),dtype=torch.float32)
  
          features = torch.empty((nb,nf,max_len),dtype=torch.float32)
          masks = torch.ones((nb,max_len),dtype=torch.int64)
          for i,batchi in enumerate(batch):
              feature = batchi[0]
              pad_size = list(feature.shape)
              pad_size[self.dim] = max_len - pad_size[self.dim]
              features[i,:,:] = torch.cat([torch.from_numpy(feature), torch.zeros(*pad_size)],dim=self.dim)
  
              for j in range(nt):
                  distance = batchi[1][j]
                  pad_size = list(distance.shape)
                  pad_size[0] = max_len - pad_size[0]
                  tt[j,i,:,:] = torch.cat([torch.cat([torch.from_numpy(distance), torch.zeros(*pad_size)],dim=0),torch.zeros((max_len,pad_size[0]))],dim=1)
  
              for j in range(nc):
                  coord = torch.from_numpy(batchi[2][j].copy())
                  pad_size = list(coord.shape)
                  pad_size[1] = max_len - pad_size[1]
                  cc[j,i,:,:] = torch.cat([coord, torch.zeros(*pad_size)], dim=1)
  
              masks[i,feature.shape[self.dim]:] = 0
          for i in range(nt):
              distances += (tt[i,:,:,:],)
          for i in range(nc):
              coords += (cc[i,:,:,:],)
          return features, distances, masks, coords
  
      def __call__(self, batch):
          return self.pad_collate(batch)
  