JAL-2110 add 'products' parameter to filter results in madeCdsAlignment

[jalview.git] / src / jalview / analysis / AlignmentUtils.java

diff --git a/src/jalview/analysis/AlignmentUtils.java b/src/jalview/analysis/AlignmentUtils.java
index 74066d7..f94d393 100644 (file)
--- a/src/jalview/analysis/AlignmentUtils.java
+++ b/src/jalview/analysis/AlignmentUtils.java
@@ -1407,11 +1407,14 @@ public class AlignmentUtils
    *          aligned dna sequences
    * @param dataset
    *          - throws error if not given a dataset
+   * @param products
+   *          (optional) to restrict results to CDS that map to specified
+   *          protein products
    * @return an alignment whose sequences are the cds-only parts of the dna
    *         sequences (or null if no mappings are found)
    */
   public static AlignmentI makeCdsAlignment(SequenceI[] dna,
-          AlignmentI dataset)
+          AlignmentI dataset, AlignmentI products)
   {
     if (dataset.getDataset() != null)
     {
@@ -1420,7 +1423,16 @@ public class AlignmentUtils
     }
     List<SequenceI> cdsSeqs = new ArrayList<SequenceI>();
     List<AlignedCodonFrame> mappings = dataset.getCodonFrames();
-    
+    HashSet<SequenceI> productSeqs = null;
+    if (products != null)
+    {
+      productSeqs = new HashSet<SequenceI>();
+      for (SequenceI seq : products.getSequences())
+      {
+        productSeqs.add(seq.getDatasetSequence() == null ? seq : seq
+                .getDatasetSequence());
+      }
+    }
 
     /*
      * construct CDS sequences from the (cds-to-protein) mappings made earlier;
@@ -1453,9 +1465,20 @@ public class AlignmentUtils
            * the dna mapping's product
            */
           SequenceI cdsSeq = null;
+
           // TODO better mappings collection data model so we can do
-          // a table lookup instead of double loops to find mappings
+          // a direct lookup instead of double loops to find mappings
+
           SequenceI proteinProduct = aMapping.getTo();
+
+          /*
+           * skip if not mapped to one of a specified set of proteins
+           */
+          if (productSeqs != null && !productSeqs.contains(proteinProduct))
+          {
+            continue;
+          }
+
           for (AlignedCodonFrame acf : MappingUtils
                   .findMappingsForSequence(proteinProduct, mappings))
           {
@@ -1544,7 +1567,7 @@ public class AlignmentUtils
 
     AlignmentI cds = new Alignment(cdsSeqs.toArray(new SequenceI[cdsSeqs
             .size()]));
-    cds.setDataset((Alignment) dataset);
+    cds.setDataset(dataset);
 
     return cds;
   }