jbr/#252 allow calculating intersection with flanking (#27)

* Preparation to add region selection to overlap table
* Add flanking support to intersectRanges method
* Part of JetBrains-Research/jbr#252

src/main/kotlin/org/jetbrains/bio/genome/containers/GenomeMaps.kt

@@ -110,6 +110,10 @@ interface GenomeStrandMapLike<T> : Iterable<T> {
         Strand.values().map { strand -> this[chromosome, strand] }
     }.iterator()
 
+    fun iteratorWithKeys(): Iterator<Pair<Pair<Chromosome,Strand>, T>> = genomeQuery.get().flatMap { chromosome ->
+        Strand.values().map { strand -> (chromosome to strand) to this[chromosome, strand] }
+    }.iterator()
+
     fun entries(): Iterable<Triple<Chromosome, Strand, T>> = genomeQuery.get().flatMap { chromosome ->
         Strand.values().map { strand ->
             Triple(chromosome, strand, this[chromosome, strand])

src/main/kotlin/org/jetbrains/bio/genome/containers/LocationsList.kt

@@ -77,6 +77,27 @@ abstract class LocationsList<T : RangesList> : GenomeStrandMapLike<List<Location
         return acc.result()
     }
 
+    /**
+     * 'inline' is important here otherwise each method usage will
+     * create new instance of anonymous 'metric' class
+     */
+    inline fun calcMarkedLocations(
+        other: LocationsList<out RangesList>,
+        metric: (RangesList, RangesList) -> List<Range>
+    ): List<Location> {
+        require(genomeQuery == other.genomeQuery)
+
+        val acc = mutableListOf<Location>()
+        genomeQuery.get().forEach { chromosome ->
+            Strand.values().forEach { strand ->
+                val ranges = metric(rangeLists[chromosome, strand], other. rangeLists[chromosome, strand])
+                acc.addAll(ranges.map { Location(it.startOffset, it.endOffset, chromosome, strand) })
+            }
+        }
+
+        return acc
+    }
+
     /**
      * 'inline' is important here otherwise each method usage will
      * create new instance of anonymous 'metric' class

src/main/kotlin/org/jetbrains/bio/genome/containers/RangesList.kt

@@ -49,7 +49,7 @@ interface RangesList : Iterable<Range> {
      * Intersect each list interval with requested [range], empty intervals not reported
      */
     fun intersectRanges(range: Range) = intersectRanges(range.startOffset, range.endOffset)
-    infix fun intersectRanges(other: RangesList): List<Range>
+    fun intersectRanges(other: RangesList, flankBothSides: Int = 0): List<Range>
 
     /**
      * If intersected ranges not needed use this function to do less GS
@@ -85,11 +85,17 @@ abstract class BaseRangesList(
         return acc
     }
 
-    override infix fun intersectRanges(other: RangesList): List<Range> {
+    override fun intersectRanges(other: RangesList, flankBothSides: Int): List<Range> {
         val acc = ArrayList<Range>()
 
+
         for (idx in 0 until size) {
-            acc.addAll(other.intersectRanges(startOffsets[idx], endOffsets[idx]))
+            val startOffset = startOffsets[idx]
+            val endOffset = endOffsets[idx]
+
+            val startOffsetFlnk = kotlin.math.max(0, startOffset - flankBothSides)
+            val endOffsetFlnk = endOffset + flankBothSides
+            acc.addAll(other.intersectRanges(startOffsetFlnk, endOffsetFlnk))
         }
 
         return acc

src/main/kotlin/org/jetbrains/bio/genome/containers/intersection/IntersectionMetrics.kt

@@ -28,6 +28,10 @@ object IntersectionMetrics {
 
         override fun calcMetric(ra: RangesList, rb: RangesList): Double =
             ra.intersectRanges(rb).sumOf { it.length().toDouble() }
+
+        //TODO implement to calculate proper regions
+        override fun calcRegions(a: LocationsList<out RangesList>, b: LocationsList<out RangesList>) =
+            a.calcMarkedLocations(b) {ra , rb -> ra.intersectRanges(rb)}
     }
 
     val OVERLAP_FRACTION = object : RegionsMetric {
@@ -37,6 +41,8 @@ object IntersectionMetrics {
             OVERLAP.calcMetric(a, b) / a.size
 
         override fun calcMetric(ra: RangesList, rb: RangesList): Double = OVERLAP.calcMetric(ra, rb)
+
+        override fun calcRegions(a: LocationsList<out RangesList>, b: LocationsList<out RangesList>) = OVERLAP.calcRegions(a, b)
     }
 
     fun parse(options: OptionSet) = (options.valueOf("metric") as String).let { text ->

src/main/kotlin/org/jetbrains/bio/genome/containers/intersection/IntersectionNumberMetric.kt

@@ -15,4 +15,8 @@ class IntersectionNumberMetric(val aSetFlankedBothSides: Int = 0) : RegionsMetri
 
     override fun calcMetric(ra: RangesList, rb: RangesList): Double =
         ra.intersectRangesNumber(rb, flankBothSides = aSetFlankedBothSides).toDouble()
+
+    //TODO implement to calculate proper regions
+    override fun calcRegions(a: LocationsList<out RangesList>, b: LocationsList<out RangesList>) =
+        a.calcMarkedLocations(b) {ra , rb -> ra.intersectRanges(rb, aSetFlankedBothSides)}
 }

src/main/kotlin/org/jetbrains/bio/genome/containers/intersection/OverlapNumberMetric.kt

@@ -1,5 +1,6 @@
 package org.jetbrains.bio.genome.containers.intersection
 
+import org.jetbrains.bio.genome.Location
 import org.jetbrains.bio.genome.containers.LocationsList
 import org.jetbrains.bio.genome.containers.RangesList
 
@@ -15,4 +16,7 @@ class OverlapNumberMetric(val aSetFlankedBothSides: Int = 0) : RegionsMetric {
 
     override fun calcMetric(ra: RangesList, rb: RangesList): Double =
         ra.overlapRangesNumber(rb, flankBothSides = aSetFlankedBothSides).toDouble()
+
+    override fun calcRegions(a: LocationsList<out RangesList>, b: LocationsList<out RangesList>) =
+        a.calcMarkedLocations(b) { ra, rb -> ra.overlapRanges(rb, aSetFlankedBothSides) }
 }

src/main/kotlin/org/jetbrains/bio/genome/containers/intersection/RegionsMetric.kt

@@ -1,5 +1,6 @@
 package org.jetbrains.bio.genome.containers.intersection
 
+import org.jetbrains.bio.genome.Location
 import org.jetbrains.bio.genome.containers.LocationsList
 import org.jetbrains.bio.genome.containers.RangesList
 
@@ -8,4 +9,5 @@ interface RegionsMetric {
 
     fun calcMetric(a: LocationsList<out RangesList>, b: LocationsList<out RangesList>): Double
     fun calcMetric(ra: RangesList, rb: RangesList): Double
+    fun calcRegions(a: LocationsList<out RangesList>, b: LocationsList<out RangesList>): List<Location>
 }

src/test/kotlin/org/jetbrains/bio/genome/containers/RangesMergingListTest.kt

@@ -91,16 +91,16 @@ class RangesMergingListTest {
         val rl1 = rangeMergingList(Range(0, 10))
         val rl2 = rangeMergingList(Range(20, 30))
         assertEquals(
-            rangeMergingList(Range(0, 10), Range(20, 30)).toList(),
-            (rl1 or rl2).toList()
+                rangeMergingList(Range(0, 10), Range(20, 30)).toList(),
+                (rl1 or rl2).toList()
         )
     }
 
     @Test
     fun orWithIntersections() {
         assertEquals(
-            rangeMergingList(Range(0, 30)).toList(),
-            (rangeMergingList(Range(0, 10)) or rangeMergingList(Range(5, 30))).toList()
+                rangeMergingList(Range(0, 30)).toList(),
+                (rangeMergingList(Range(0, 10)) or rangeMergingList(Range(5, 30))).toList()
         )
     }
 
@@ -122,25 +122,40 @@ class RangesMergingListTest {
         val rl2 = rangeMergingList(Range(5, 30))
         doCheckIntersectRanges(rangeMergingList(Range(5, 10)).toList(), rl1, rl2)
         doCheckIntersectRanges(
-            rangeMergingList(Range(5, 10), Range(25, 30)).toList(),
-            rl1 or rangeMergingList(Range(25, 30)),
-            rl2
+                rangeMergingList(Range(5, 10), Range(25, 30)).toList(),
+                rl1 or rangeMergingList(Range(25, 30)),
+                rl2
         )
         doCheckIntersectRanges(
-            rangeMergingList(Range(5, 10), Range(25, 30)).toList(),
-            rl2,
-            rl1 or rangeMergingList(Range(25, 30))
+                rangeMergingList(Range(5, 10), Range(25, 30)).toList(),
+                rl2,
+                rl1 or rangeMergingList(Range(25, 30))
         )
 
         doCheckIntersectRanges(rangeMergingList(Range(6, 8)).toList(), rangeMergingList(Range(0, 2), Range(6, 8)), rl2)
     }
 
+    @Test
+    fun intersectRangesFlnk() {
+        val rl1 = rangeMergingList(Range(0, 10), Range(20, 25))
+        val rl2 = rangeMergingList(Range(0, 30))
+        assertEquals(listOf(Range(0, 12), Range(18, 27)), (rl1.intersectRanges(rl2, 2)).toList())
+        assertEquals(listOf(Range(11, 12), Range(18, 20)), (rl1.intersectRanges(rangeMergingList(Range(11, 20)), 2)))
+    }
+
+    @Test
+    fun intersectRangesFlnkOverlapAfterFlanking() {
+        val rl1 = rangeMergingList(Range(0, 11), Range(13, 25))
+        val rl2 = rangeMergingList(Range(0, 30))
+        assertEquals(listOf(Range(0, 14), Range(10, 28)), (rl1.intersectRanges(rl2, 3)).toList())
+    }
+
     private fun doCheckIntersectRanges(
-        expected: List<Range>,
-        rl1: RangesMergingList,
-        rl2: RangesMergingList
+            expected: List<Range>,
+            rl1: RangesMergingList,
+            rl2: RangesMergingList
     ) {
-        assertEquals(expected, (rl1 intersectRanges rl2).toList())
+        assertEquals(expected, (rl1.intersectRanges(rl2)).toList())
         assertEquals(expected.size, (rl1.intersectRangesNumber(rl2)))
     }
 
@@ -186,10 +201,10 @@ class RangesMergingListTest {
     }
 
     private fun checkOverlap(
-        expected: List<Range>,
-        a: RangesMergingList,
-        b: RangesMergingList,
-        flankBothSides: Int = 0
+            expected: List<Range>,
+            a: RangesMergingList,
+            b: RangesMergingList,
+            flankBothSides: Int = 0
     ) {
         assertEquals(expected, (a.overlapRanges(b, flankBothSides)).toList())
         assertEquals(expected.size, a.overlapRangesNumber(b, flankBothSides))
@@ -201,14 +216,14 @@ class RangesMergingListTest {
         val universe = Range(0, 100)
         assertEquals(listOf(universe), universe - emptyList())
         assertEquals(
-            listOf(Range(0, 20), Range(40, 100)),
-            universe - listOf(Range(20, 40))
+                listOf(Range(0, 20), Range(40, 100)),
+                universe - listOf(Range(20, 40))
         )
         assertEquals(listOf(Range(40, 100)), universe - listOf(Range(0, 40)))
         assertEquals(listOf(Range(0, 40)), universe - listOf(Range(40, 100)))
         assertEquals(
-            emptyList<Range>(),
-            universe - listOf(Range(0, 40), Range(40, 100))
+                emptyList<Range>(),
+                universe - listOf(Range(0, 40), Range(40, 100))
         )
     }
 
@@ -218,21 +233,21 @@ class RangesMergingListTest {
         assertEquals(listOf(Range(10, 90)), rangeMergingList(Range(0, 100)).intersectRanges(Range(10, 90)))
         assertEquals(listOf(Range(10, 90)), rangeMergingList(Range(10, 90)).intersectRanges(Range(0, 100)))
         assertEquals(
-            listOf(Range(20, 40), Range(50, 70)),
-            rangeMergingList(Range(0, 40), Range(50, 100)).intersectRanges(Range(20, 70))
+                listOf(Range(20, 40), Range(50, 70)),
+                rangeMergingList(Range(0, 40), Range(50, 100)).intersectRanges(Range(20, 70))
         )
     }
 
     @Test
     fun lookup() {
         val data = listOf(
-            Range(1, 2),
-            Range(4, 6),
-            Range(4, 9),
-            Range(4, 7),
-            Range(6, 10),
-            Range(20, 30),
-            Range(40, 50)
+                Range(1, 2),
+                Range(4, 6),
+                Range(4, 9),
+                Range(4, 7),
+                Range(6, 10),
+                Range(20, 30),
+                Range(40, 50)
         ).toRangeMergingList()
         assertEquals(-1, data.internalLookup(0))
         assertEquals(0, data.internalLookup(1))
@@ -247,34 +262,34 @@ class RangesMergingListTest {
     @Test
     fun complementaryRanges() {
         checkComplement(
-            listOf(),
-            listOf(Range(0, 100))
+                listOf(),
+                listOf(Range(0, 100))
         )
         checkComplement(
-            listOf(Range(4, 6)),
-            listOf(Range(0, 4), Range(6, 100))
+                listOf(Range(4, 6)),
+                listOf(Range(0, 4), Range(6, 100))
         )
         checkComplement(
-            listOf(Range(3, 6), Range(14, 16), Range(40, 90)),
-            listOf(Range(0, 3), Range(6, 14), Range(16, 40), Range(90, 100))
+                listOf(Range(3, 6), Range(14, 16), Range(40, 90)),
+                listOf(Range(0, 3), Range(6, 14), Range(16, 40), Range(90, 100))
         )
         checkComplement(
-            listOf(Range(1, 2), Range(4, 6), Range(10, 99)),
-            listOf(Range(0, 1), Range(2, 4), Range(6, 10), Range(99, 100))
+                listOf(Range(1, 2), Range(4, 6), Range(10, 99)),
+                listOf(Range(0, 1), Range(2, 4), Range(6, 10), Range(99, 100))
         )
         checkComplement(
-            listOf(Range(0, 2), Range(4, 6), Range(10, 100)),
-            listOf(Range(2, 4), Range(6, 10))
+                listOf(Range(0, 2), Range(4, 6), Range(10, 100)),
+                listOf(Range(2, 4), Range(6, 10))
         )
         checkComplement(
-            listOf(Range(0, 100)),
-            listOf()
+                listOf(Range(0, 100)),
+                listOf()
         )
     }
 
     private fun checkComplement(
-        data: List<Range>,
-        expected: List<Range>
+            data: List<Range>,
+            expected: List<Range>
     ) {
         val actual = data.toRangeMergingList().complementaryRanges(100).toList()
         assertEquals(expected, actual)