Merge pull request #1494 from milaboratory/add-clones-to-shmt

Add clones to shmt
milaboratory · Jan 2, 2024 · e42a6b2 · e42a6b2
2 parents 5e29ae5 + 32c8794
commit e42a6b2
Show file tree

Hide file tree

Showing 4 changed files with 23 additions and 5 deletions.
diff --git a/build.gradle.kts b/build.gradle.kts
@@ -134,7 +134,7 @@ val toObfuscate: Configuration by configurations.creating {
 val obfuscationLibs: Configuration by configurations.creating
 
 
-val mixcrAlgoVersion = "4.6.0-17-consensus-by-chain"
+val mixcrAlgoVersion = "4.6.0-20-add-clones-to-shmt"
 // may be blank (will be inherited from mixcr-algo)
 val milibVersion = ""
 // may be blank (will be inherited from mixcr-algo or milib)

diff --git a/regression/list b/regression/list
@@ -4,6 +4,7 @@ V06_base_build_trees.shmt
 V07_01f3146f1b_base_build_trees.shmt
 V08_4c5f869d5d_base_build_trees.shmt
 V09_85ab04f88a_base_build_trees.shmt
+V10_3c89f63e85_base_build_trees.shmt
 V10_baseSingleCell_vdjcontigs.clna
 V10_case3.clna
 V11_base_single_cell.vdjcontigs.clna

diff --git a/src/main/kotlin/com/milaboratory/mixcr/cli/CommandFindShmTrees.kt b/src/main/kotlin/com/milaboratory/mixcr/cli/CommandFindShmTrees.kt
@@ -18,6 +18,7 @@ import cc.redberry.pipe.util.asOutputPort
 import cc.redberry.pipe.util.asSequence
 import cc.redberry.pipe.util.drainToAndClose
 import cc.redberry.pipe.util.flatMap
+import cc.redberry.pipe.util.flatten
 import cc.redberry.pipe.util.map
 import cc.redberry.pipe.util.toList
 import com.milaboratory.app.InputFileType
@@ -61,6 +62,7 @@ import com.milaboratory.util.TempFileDest
 import com.milaboratory.util.TempFileManager
 import com.milaboratory.util.cached
 import com.milaboratory.util.groupByOnDisk
+import com.milaboratory.util.withExpectedSize
 import io.repseq.core.GeneFeature
 import io.repseq.core.GeneFeature.CDR3
 import io.repseq.core.GeneFeatures
@@ -363,8 +365,19 @@ class CommandFindShmTrees : MiXCRCommandWithOutputs() {
                     processSingleCell
                 )
 
-                val writer = shmTreesWriter.treesWriter()
-                writeResults(writer, singleCellTrees, datasets, shmTreeBuilder)
+                shmTreesWriter.treesWriter { writer ->
+                    writeResults(writer, singleCellTrees, datasets, shmTreeBuilder)
+                }
+                shmTreesWriter.clonesWriter { writer ->
+                    datasets
+                        .mapIndexed { datasetId, clnsReader ->
+                            clnsReader.readClones().map { clone -> CloneWithDatasetId(clone, datasetId) }
+                        }
+                        .flatten()
+                        .withExpectedSize(datasets.sumOf { it.numberOfClones().toLong() })
+                        .reportProgress("Writing clones")
+                        .drainToAndClose(writer)
+                }
             }
             reportBuilder.setFinishMillis(System.currentTimeMillis())
             report = reportBuilder.buildReport()

diff --git a/src/main/kotlin/com/milaboratory/mixcr/cli/CommandSlice.kt b/src/main/kotlin/com/milaboratory/mixcr/cli/CommandSlice.kt
@@ -13,11 +13,11 @@ package com.milaboratory.mixcr.cli
 
 import cc.redberry.pipe.util.asOutputPort
 import cc.redberry.pipe.util.asSequence
+import cc.redberry.pipe.util.drainToAndClose
 import cc.redberry.pipe.util.filter
 import cc.redberry.pipe.util.it
 import cc.redberry.pipe.util.map
 import cc.redberry.pipe.util.toList
-import cc.redberry.pipe.util.use
 import com.milaboratory.app.InputFileType
 import com.milaboratory.app.ValidationException
 import com.milaboratory.mixcr.basictypes.ClnAReader
@@ -238,7 +238,7 @@ class CommandSlice : MiXCRCommandWithOutputs() {
             SHMTreesWriter(out).use { writer ->
                 writer.copyHeaderFrom(reader)
 
-                writer.treesWriter().use { treesWriter ->
+                writer.treesWriter { treesWriter ->
                     reader.readTrees()
                         .filter { it.treeId.toLong() in ids }
                         .asSequence().withIndex()
@@ -252,6 +252,10 @@ class CommandSlice : MiXCRCommandWithOutputs() {
                         }
                 }
 
+                writer.clonesWriter { clonesWriter ->
+                    reader.readAllClones().drainToAndClose(clonesWriter)
+                }
+
                 writer.setFooter(reader.footer)
             }
         }