nodes: 0, module; 1, function_definition; 2, function_name:sort_snps; 3, parameters; 4, block; 5, identifier:snps; 6, expression_statement; 7, expression_statement; 8, comment:# move PAR and MT to the end of the dataframe; 9, if_statement; 10, if_statement; 11, comment:# convert chrom column to category for sorting; 12, comment:# https://stackoverflow.com/a/26707444; 13, expression_statement; 14, comment:# sort based on ordered chromosome list and position; 15, expression_statement; 16, comment:# convert chromosome back to object; 17, expression_statement; 18, return_statement; 19, comment:""" Sort SNPs based on ordered chromosome list and position. """; 20, assignment; 21, comparison_operator:"PAR" in sorted_list; 22, block; 23, comparison_operator:"MT" in sorted_list; 24, block; 25, assignment; 26, assignment; 27, assignment; 28, identifier:snps; 29, identifier:sorted_list; 30, call; 31, string:"PAR"; 32, identifier:sorted_list; 33, expression_statement; 34, expression_statement; 35, string:"MT"; 36, identifier:sorted_list; 37, expression_statement; 38, expression_statement; 39, subscript; 40, call; 41, identifier:snps; 42, call; 43, subscript; 44, call; 45, identifier:sorted; 46, argument_list; 47, call; 48, call; 49, call; 50, call; 51, identifier:snps; 52, string:"chrom"; 53, attribute; 54, argument_list; 55, attribute; 56, argument_list; 57, identifier:snps; 58, string:"chrom"; 59, attribute; 60, argument_list; 61, call; 62, keyword_argument; 63, attribute; 64, argument_list; 65, attribute; 66, argument_list; 67, attribute; 68, argument_list; 69, attribute; 70, argument_list; 71, subscript; 72, identifier:astype; 73, call; 74, identifier:snps; 75, identifier:sort_values; 76, list; 77, subscript; 78, identifier:astype; 79, identifier:object; 80, attribute; 81, argument_list; 82, identifier:key; 83, identifier:_natural_sort_key; 84, identifier:sorted_list; 85, identifier:remove; 86, string:"PAR"; 87, identifier:sorted_list; 88, identifier:append; 89, string:"PAR"; 90, identifier:sorted_list; 91, identifier:remove; 92, string:"MT"; 93, identifier:sorted_list; 94, identifier:append; 95, string:"MT"; 96, identifier:snps; 97, string:"chrom"; 98, identifier:CategoricalDtype; 99, argument_list; 100, string:"chrom"; 101, string:"pos"; 102, identifier:snps; 103, string:"chrom"; 104, subscript; 105, identifier:unique; 106, keyword_argument; 107, keyword_argument; 108, identifier:snps; 109, string:"chrom"; 110, identifier:categories; 111, identifier:sorted_list; 112, identifier:ordered; 113, True
code: def sort_snps(snps):
      """ Sort SNPs based on ordered chromosome list and position. """
  
      sorted_list = sorted(snps["chrom"].unique(), key=_natural_sort_key)
  
      # move PAR and MT to the end of the dataframe
      if "PAR" in sorted_list:
          sorted_list.remove("PAR")
          sorted_list.append("PAR")
  
      if "MT" in sorted_list:
          sorted_list.remove("MT")
          sorted_list.append("MT")
  
      # convert chrom column to category for sorting
      # https://stackoverflow.com/a/26707444
      snps["chrom"] = snps["chrom"].astype(
          CategoricalDtype(categories=sorted_list, ordered=True)
      )
  
      # sort based on ordered chromosome list and position
      snps = snps.sort_values(["chrom", "pos"])
  
      # convert chromosome back to object
      snps["chrom"] = snps["chrom"].astype(object)
  
      return snps