Minor fixes to adam-python API

fnothaft · Aug 22, 2017 · 9b51df5 · 9b51df5
1 parent 209f997
commit 9b51df5
Show file tree

Hide file tree

Showing 2 changed files with 31 additions and 2 deletions.
diff --git a/adam-python/src/bdgenomics/adam/rdd.py b/adam-python/src/bdgenomics/adam/rdd.py
@@ -283,7 +283,11 @@ def toCoverage(self, collapse = True):
         :rtype: bdgenomics.adam.rdd.CoverageRDD
         """
 
-        return CoverageRDD(self._jvmRdd.toCoverage(collapse), self.sc)
+        coverageRDD = CoverageRDD(self._jvmRdd.toCoverage(), self.sc)
+        if (collapse):
+            return coverageRDD.collapse()
+        else:
+            return coverageRDD
 
 
     def save(self, filePath, isSorted = False):
@@ -745,7 +749,7 @@ def _replaceRdd(self, newRdd):
         return FragmentRDD(newRdd, self.sc)
 
 
-    def __init__(self, jvmRdd):
+    def __init__(self, jvmRdd, sc):
         """
         Constructs a Python FragmentRDD from a JVM FragmentRDD.
         Should not be called from user code; instead, go through

diff --git a/adam-python/src/bdgenomics/adam/test/alignmentRecordRdd_test.py b/adam-python/src/bdgenomics/adam/test/alignmentRecordRdd_test.py
@@ -138,3 +138,28 @@ def test_transmute_to_coverage(self):
 
         assert(isinstance(readsAsCoverage, CoverageRDD))
         self.assertEquals(readsAsCoverage.toDF().count(), 5)
+
+
+    def test_to_coverage(self):
+
+        readsPath = self.resourceFile("unsorted.sam")
+        ac = ADAMContext(self.sc)
+
+        reads = ac.loadAlignments(readsPath)
+
+        coverage = reads.toCoverage()
+        self.assertEquals(coverage.toDF().count(), 42)
+
+        coverage = reads.toCoverage(collapse = False)
+        self.assertEquals(coverage.toDF().count(), 46)
+
+
+    def test_to_fragments(self):
+
+        readsPath = self.resourceFile("unsorted.sam")
+        ac = ADAMContext(self.sc)
+
+        reads = ac.loadAlignments(readsPath)
+
+        fragments = reads.toFragments()
+        self.assertEquals(fragments.toDF().count(), 5)