FastaExtract

#!/usr/bin/perl

=head1 NAME

 FastaExtract.pl
 Tool to extract sequences from a fasta file in different ways

=cut

=head1 SYPNOSIS

 FastaExtract.pl [-h] -f <input_fasta_file> [-o output_basename]
                       [-i <extract_by_id>] [-l <extract_by_length>] 
                       [-r <extract_by_regexp>]
                       [-c <extract_by_composition>]
                       [-v] [-C] [-O] [-R] [-I] [-S] [-U]

=head2 I<Flags:>

=over


=item -f

B<input_fasta_file>       input fasta file (mandatory)

=item -o

B<output_basename>        output basename (by default it will printed as stdout)

=item -i

B<extract_by_id>          filename to extract sequences by id

=item -l 

B<extract_by_length>      integer to extract sequences by length     

=item -r

B<regexp>                 regular expression to extract the sequences

=item -c

B<extract_by_composition> by nucleotide content (for example: GC=50)

=item -C

B<cut_sequence>           cut a part of the sequence extracted (coordinates)
                          from the input filename (-i)

=item -O

B<orientate_strand>       orientate the sequence based in the strand specified
                          in the ID list filename (-i)

=item -I

B<replace_ids>            replace ID for the id specified in the ID list
                          file (-i)

=item -R

B<revcom>                 return reverse complemetary of the sequence

=item -S

B<one_file_per_read>      split the output to one file per read

=item -U

B<select_unique>          select unique (remove redundancy if there is any)

=item -v

B<invert_selection>       invert the selection

=item -h

B<help>                   print the help

=back

=cut

=head1 DESCRIPTION

 This script extract a list of sequences from a fasta/qual file. 

 There are three ways to supply the list of sequences:

   1) As list of ids in a file (-i <id_list_file>)

   2) As length value (-l <length>)

   3) As regular expression to match with the id (-r <regexp>)

   4) By minimum nucleotide/aminoacid content in percentage (-c <composition>)

   Examples:

      fasta_extract.pl -f test.fasta -i id_file.txt > extract.fasta

      fasta_extract.pl -f test.fasta -l 100 -v > seq_less100.fasta

      fasta_extract.pl -f test.fasta -r 'Sl' > sl_seq.fasta

      fasta_extract.pl -f test.fasta -r 'AAAAAAAAAA$' > seq_with_poliA.fasta

      fasta_extract.pl -f test.fasta -c GC=50 > seq_GC50_omore.fasta

      fasta_extract.pl -f test.fasta -c N=100 -v > seq_with_noNs.fasta

   To create one file per read, use -S option, in this case -o will be 
   overwrite with sequence id name

   -i <id_list> can be a file with 1 column (just the IDs), 3 columns (ID plus
   start and end coordinates to use with -C), 4 columns (ID + start + end + 
   strand, to get the specific strand -O) or 5 columns (ID + start + end + 
   strand + NewID, use with -I option)

=cut

=head1 AUTHORS

  Aureliano Bombarely Gomez.
  (ab782@cornell.edu).

=cut

=head1 METHODS

 FastaExtract.pl


=cut

use strict;
use warnings;

use Getopt::Std;
use Bio::SeqIO;
use Bio::SeqUtils;
use Bio::Tools::IUPAC;

our ($opt_f, $opt_o, $opt_i, $opt_l, $opt_c, $opt_r, $opt_v, $opt_C, $opt_O,
     $opt_I, $opt_R, $opt_S, $opt_U, $opt_h);
getopts("f:o:i:l:c:r:vCOIRSUh");
if (!$opt_f && !$opt_o && !$opt_i && !$opt_l && !$opt_c && !$opt_r && !$opt_v 
    && !$opt_C && !$opt_O && !$opt_I && !$opt_R && !$opt_S && !$opt_U && 
    !$opt_h) {
    print "There are n\'t any tags. Print help\n\n";
    help();
}
if ($opt_h) {
    help();
}

## Get the arguments and check them

my $infasta_file = $opt_f || 
    die("INPUT ARG. ERROR: -f <input_fastafile> argument was not supplied.\n");
my $idfile = $opt_i;
my $length_cf = $opt_l;
my $regexp;
if (defined $opt_r) {
    $regexp = qr/$opt_r/;
}
if ($opt_C) {
    unless (defined $idfile) {
	die("ERROR: -C option only can be used with -i option.\n");
    }
}

my %seqcond = ();
if (defined $opt_c) {

    my @conds = split(/,/, $opt_c);
    foreach my $cond (@conds) {
    
	if ($cond =~ m/^(\w+)=(\d+)$/) {
	
	    $seqcond{$1} = $2;
	}
    }
}

if (!$idfile && !$length_cf && !$regexp && scalar(keys(%seqcond)) == 0) {
    die("EXTRACTION ARG. ERROR: No extract options were used (-i,-l,-r,-c)\n");
}

my $out = $opt_o;

## Extract the list of ids from the id file if exists

my %ids;
my %trims;
my $trimsc = 0;

if (defined $idfile) {

    print STDERR "\nEXTRACT BY IDS option [Enabled].\n";

    open my $idfh, '<', $idfile || 
	die("OPEN FILE ERROR: File=$idfile can not be openned (error:$!)\n");

    my $L = `cut -f1 $idfile | wc -l`;
    chomp($L);
    my $l = 0;

    while (<$idfh>) {
	chomp($_);

	## It will take only the first column and remove the rest
	$l++;
	
	my @data = split(/\t/, $_);
	my $fc = scalar(@data);

	## It will ignore the duplications overwritting the ids

	if (defined $data[0]) {
	    $ids{$data[0]} = 1;

	    if ($opt_C) {
	    
		my $start = $data[1];
		my $end = $data[2];
		my %metadata = ( start => $start, end => $end);

		if ($start > $end) {
		    die("ERROR: $l line, start=$start > end=$end\n");
		}

		if ($opt_O) {
		
		    my $orient = $data[3];
		    if ($orient !~ m/^(\+|-)$/) {
		    
			die("ERROR: $l line, strand has wrong format (+|-)\n");
		    }
		    else {
		    
			$metadata{strand} = $orient;
		    }

		    if ($opt_I) {
		    
			if (scalar(@data) != 5) {
			    
			    print STDERR "\n\nOptions -C -O -I Enabled.\n";
			    die("ERROR: $l line doesnt have 5 fields");
			}

			$metadata{newid} = $data[4];
		    }
		    else {
		    
			unless (scalar(@data) >= 4) {
			
			    print STDERR "\n\nOptions -C -O Enabled.\n";
			    die("ERROR: $l line doesnt have 4 or more fields");
			}
		    }
		}
		else {
		
		    unless (scalar(@data) >= 3) {
			
			print STDERR "\n\nOptions -C Enabled.\n";
			die("ERROR: $l line doesnt have 3 or more fields");
		    }
		}
	    
		if (exists $trims{$data[0]}) {
		    
		    push @{$trims{$data[0]}}, \%metadata;
		}
		else {
		    $trims{$data[0]} = [\%metadata];
		}
		$trimsc++;
	    
	    }
	    else {
	    
		unless (scalar(@data) >= 1) {
		    
		    print STDERR "\n\nOptions -C Enabled.\n";
		    die("ERROR: $l line doesnt have 1 or more fields");
		}		
	    }		
	    
	}
    }

    my $ids_count = scalar(keys %ids);
    print STDERR "\t$ids_count ids ($trimsc trims) were parsed from $idfile\n";
}




## Open the input and the output fasta file


print STDERR "\nEXTRACTING SEQUENCES FROM $infasta_file.\n";

my $c = 0;

my $inio = Bio::SeqIO->new(
                             -file   => $infasta_file,
                             -format => 'fasta'
                          );


## It will print the sequences in STDOUT by default

my $outio;
unless (defined $opt_S) {
    if (defined $out) {

	my $outfasta = $out;
	
	$outio = Bio::SeqIO->new( 
	    -file   => ">$outfasta",
	    -format => 'fasta'
	    );	
    }
    else {
	$outio = Bio::SeqIO->new( -format => 'fasta' );
    }
}

my $s = 0;
my %selected = ();

while (my $seqobj = $inio->next_seq()) {
    
    $s++;
    my $id = $seqobj->primary_id();
    my $lg = $seqobj->length();

    print STDERR "Processing id=$id (sequence $s)    \r";

    ## Select only if it have one of the extract option

    my $sel = 0;
	
    if (defined $idfile) {
	if (exists $ids{$id}) {
	    $sel = 1;
	}
    }
    if (defined $length_cf) {
	if ($lg > $length_cf) {
	    $sel = 1;
	}
    }
    if (defined $regexp) {
	if ($id =~ $regexp) {
	    $sel = 1;
	}
	elsif ($seqobj->seq() =~ $regexp) {
	    $sel = 1;
	}
    }
    if (scalar(keys %seqcond) > 0 && $lg > 0) {
    
	my %seqcomp = ();

	foreach my $ch (split(//, $seqobj->seq())) {
	    
	    if (exists $seqcomp{uc($ch)}) {
	    
		$seqcomp{uc($ch)}++;
	    }
	    else {
	    
		$seqcomp{uc($ch)} = 1;
	    }
	}

	my @composition = ();
	foreach my $chq (sort {$a cmp $b} keys %seqcomp) {
	    
	    $seqcomp{$chq} = $seqcomp{$chq} * 100 / $lg;
	    push @composition, $chq . '=' . $seqcomp{$chq};
	}

	my $cond_n = scalar(keys %seqcond);
	my $cond_t = 0;

	foreach my $cond (sort {$a cmp $b} keys %seqcond) {
	
	    my $sumval = 0;
	    my @sumcond = split(//, $cond);
	    foreach my $singcond (@sumcond) {
		
		if (exists $seqcomp{uc($singcond)}) {
		    $sumval += $seqcomp{uc($singcond)};
		}
	    }
	    if ($sumval >= $seqcond{$cond}) {
		
		$cond_t++;
	    }	
	}
	if ($cond_t == $cond_n) {
	    $sel = 1;
	}

	#my $test = join(",", @composition);
	#print STDERR "\nTESTCOMP (SEL:$sel): $id\t$test\n";
    }

    if ($opt_v) {
	
	## Invert the selection 
	$sel =~ s/0/2/;
	$sel =~ s/1/0/;
	$sel =~ s/2/1/;
    }

    if ($opt_U && exists $selected{$id}) {
    
	$sel = 0;
    }

    ## If the sequence was selected, it will printed

    if ($sel == 1) {
	
	if ($opt_C) {
	    if (exists $trims{$id}) {

		my @coords = @{$trims{$id}};

		foreach my $coord (@coords) { 

		    my $st = $coord->{start};
		    my $en = $coord->{end};

		    if ($st < $en) {

			my $newid = $id . '_subseq_from' . $st . '_to' . $en;

			if ($opt_I) {
		    
			    $newid = $coord->{newid};
			}
			my $newseq = $seqobj->trunc($st,$en);
			$newseq->id($newid);
			
			if ($opt_O && $coord->{strand} =~ /^-$/) {
			    
			    $newseq = $newseq->revcom();
			}
			if ($opt_R) {
			    $newseq = $newseq->revcom();
			}

			if ($opt_S) {
		    
			    my $outs_fasta = $id . '.fasta';
			    if (defined $out) {
				$outs_fasta = $out . '.' . $id . '.fasta';
			    }
			    my $s_outio = Bio::SeqIO->new( 
				-file   => ">$outs_fasta",
				-format => 'fasta'
				);
			    $s_outio->write_seq($newseq);
			}
			else {
			    
			    $outio->write_seq($newseq);
			}
			$c++;
		    }
		}
	    }
	}
	else {
	
	    if ($opt_R) {
		$seqobj = $seqobj->revcom();
	    }

	    if ($opt_S) {
		    
		my $outs_fasta = $id . '.fasta';
		if (defined $out) {
		    $outs_fasta = $out . '.' . $id . '.fasta';
		}
		my $s_outio = Bio::SeqIO->new( 
		    -file   => ">$outs_fasta",
		    -format => 'fasta'
		    );
		$s_outio->write_seq($seqobj);
	    }
	    else {

		$outio->write_seq($seqobj);
	    }
	    $c++;
	}

	$selected{$id} = 1;
    }
}

print STDERR "\n\nDONE.\n\t$c sequences were extracted from the $infasta_file.\n\n";



=head2 help

  Usage: help()
  Desc: print help of this script
  Ret: none
  Args: none
  Side_Effects: exit of the script
  Example: if (!@ARGV) {
               help();
           }

=cut

sub help {
  print STDERR <<EOF;
  $0:

    Description:

      This script extract a list of sequences from a fasta/qual file. 

      There are three ways to supply the list of sequences:

	  1) As list of ids in a file (-i <id_list_file>)
    
          2) As length value (-l <length>)

          3) As regular expression to match with the id (-r <regexp>)

          4) By minimum nucleotide/aminoacid content in percentage 
             (-c <composition>)

       To create one file per read, use -S option, in this case -o will be 
       overwrite with sequence id name

      -i <id_list> can be a file with: 
         1 column (just the IDs), 
         3 columns (ID plus start and end coordinates to use with -C), 
         4 columns (ID + start + end + strand, to get the specific strand -O) 
         5 columns (ID + start + end + strand + NewID, use with -I option)

    Usage:
     
      FastaExtract.pl [-h] -f <input_fasta_file> [-o output_basename]
                       [-i <extract_by_id>] [-l <extract_by_length>] 
                       [-r <extract_by_regexp>]
                       [-c <extract_by_composition>]
                       [-v] [-C] [-O] [-R] [-I] [-S]

      
    Examples:

      fasta_extract.pl -f test.fasta -i id_file.txt > extract.fasta

      fasta_extract.pl -f test.fasta -l 100 -v > seq_less100.fasta

      fasta_extract.pl -f test.fasta -r 'Sl' > sl_seq.fasta

      fasta_extract.pl -f test.fasta -r 'AAAAAAAAAA$' > seq_with_poliA.fasta
 
      fasta_extract.pl -f test.fasta -c GC=50 > seq_GC50_omore.fasta

      fasta_extract.pl -f test.fasta -c N=100 -v > seq_with_noNs.fasta

    Flags:

      -f <input_fasta_file>       input fasta file (mandatory)
      -o <output_basename>        output basename (by default it will printed 
                                  as stdout)
      -i <extract_by_id>          filename to extract sequences by id
      -l <extract_by_length>      integer to extract sequences by length     
      -c <extract_by_composition> nt/aa min composition
      -r <regexp>                 regular expression to extract the sequences
      -C <enable_trimming>        trim seqs according -i file.
      -O <orientate_strand>       orientate the sequence based in the strand 
                                  specified in the ID list filename (-i)
      -I <replace_ids>            replace ID for the id specified in the ID list
                                  file (-i)
      -R <revcom>                 get reverse complementary
      -S <onefile_per_read>       split the output as one file per read
      -v <invert_selection>       invert the selection
      -h <help>                   print the help
     

EOF
exit (1);
}