Do not traverse graph when extracting segments in path range.

9 files changed, 83 insertions, 26 deletions

diff --git a/doc/domagi-extract.dbk b/doc/domagi-extract.dbk
index a2b7c6b..d20afd9 100644
--- a/doc/domagi-extract.dbk
+++ b/doc/domagi-extract.dbk
@@ -9,23 +9,29 @@
       <title>Options</title>
       <xi:include href="input-db-argument.dbk" />
       <xi:include href="output-db-argument.dbk" />
+      <xi:include href="threads-argument.dbk" />
+      <xi:include href="help-option.dbk" />
+    </variablelist>
+    <variablelist>
+      <title>Traverse the graph from a segment</title>
       <varlistentry>
         <term><option>-n <replaceable>SEGMENT</replaceable></option></term>
         <term><option>--node=<replaceable>SEGMENT</replaceable></option></term>
         <listitem><para>Segment name from which to begin the traversal</para></listitem>
       </varlistentry>
       <varlistentry>
-        <term><option>-r <replaceable>PATH_RANGE</replaceable></option></term>
-        <term><option>--path-range=<replaceable>PATH_RANGE</replaceable></option></term>
-        <listitem><para>Path range, in <replaceable>path[:pos1[-pos2]]</replaceable> format, from which to begin the traversal</para></listitem>
-      </varlistentry>
-      <varlistentry>
         <term><option>-c <replaceable>STEPS</replaceable></option></term>
         <term><option>--context-steps=<replaceable>STEPS</replaceable></option></term>
         <listitem><para>The number of segments away from the initial segments to traverse</para></listitem>
       </varlistentry>
-      <xi:include href="threads-argument.dbk" />
-      <xi:include href="help-option.dbk" />
+    </variablelist>
+    <variablelist>
+      <title>Extract segments in path range</title>
+      <varlistentry>
+        <term><option>-r <replaceable>PATH_RANGE</replaceable></option></term>
+        <term><option>--path-range=<replaceable>PATH_RANGE</replaceable></option></term>
+        <listitem><para>Extract segments in <replaceable>PATH_RANGE</replaceable>, specified in the <replaceable>path[:pos1[-pos2]]</replaceable> format. <replaceable>pos1</replaceable> and <replaceable>pos2</replaceable> are 0-based coordinates. The extracted segments include <replaceable>pos1</replaceable> (inclusive) but not <replaceable>pos2</replaceable> (exclusive).</para></listitem>
+      </varlistentry>
     </variablelist>
   </refsection>
 </refentry>
diff --git a/domagi/domagi.py b/domagi/domagi.py
index 2aba383..e2bf457 100644
--- a/domagi/domagi.py
+++ b/domagi/domagi.py
@@ -226,34 +226,27 @@ def depth(con, graph_depth_table, paths, bed_input, threads, progress):
 @click.option("-c", "--context-steps", "steps",
               type=click.INT,
               # TODO: Add default=0
-              required=True,
               help="number of traversal steps")
 @common_options
 def extract(con, outfile, segment_name, path_range, steps, threads, progress):
     set_duckdb_threads(con, threads)
     with connect_duckdb(outfile, threads) as out_con:
         out_con.execute(read_sql("schema.sql"))
-    if segment_name:
-        con.execute("""
-        CREATE TEMPORARY TABLE initial_segment AS
-          SELECT id FROM segment WHERE segment.name=?
-        """,
-                    [segment_name])
-    elif path_range:
-        # TODO: We're assuming the interval is [start, end) rather
-        # than [start, end]. But check what odgi does.
+    if path_range:
+        path, start, end = re.match(r"^([^:]*):(\d+)-(\d+)", path_range).groups()
         con.execute("""
-        CREATE TEMPORARY TABLE initial_segment AS
+        CREATE TEMPORARY TABLE selected_segment AS
           SELECT segment_id AS id
           FROM path_segment
           INNER JOIN path ON path.id=path_segment.path_id
-          WHERE path.name=? AND start>=? AND start<?;
+          WHERE path.name=$1 AND path_segment.start<$3 AND $2<path_segment.end;
         """,
-        # TODO: Convert extracted strings to integers.
-        re.match(r"^([^:]*):(\d+)-(\d+)", path_range).groups())
+                    [path, int(start), int(end)])
+    elif segment_name:
+        con.execute(read_sql("extract-node-traversal.sql"),
+                    [segment_name, steps])
     else:
         raise ValueError("Neither --node and --path-range specified")
-    con.execute(read_sql("extract-node-traversal.sql"), [steps])
     con.execute(f"""
     ATTACH '{outfile}' AS subset_db (READ_WRITE);
     
@@ -278,7 +271,6 @@ def extract(con, outfile, segment_name, path_range, steps, threads, progress):
     GROUP BY id;
 
     DROP TABLE selected_segment;
-    DROP TABLE initial_segment;
     """)
 
 @main.command(short_help="Write graph in sparse matrix format")
diff --git a/domagi/extract-node-traversal.sql b/domagi/extract-node-traversal.sql
index 2624809..cfad925 100644
--- a/domagi/extract-node-traversal.sql
+++ b/domagi/extract-node-traversal.sql
@@ -1,6 +1,7 @@
 CREATE TEMPORARY TABLE selected_segment AS
-  WITH RECURSIVE cte (id, distance) AS (
-      SELECT id, 0 FROM initial_segment
+  WITH RECURSIVE
+    cte (id, distance) AS (
+      SELECT id, 0 FROM segment WHERE segment.name=?
       UNION ALL
       SELECT DISTINCT to_segment, distance+1 FROM cte
       INNER JOIN (
@@ -13,5 +14,5 @@ CREATE TEMPORARY TABLE selected_segment AS
       )
       ON from_segment=id
       WHERE distance<?
-  )
+    )
   SELECT DISTINCT id FROM cte;
diff --git a/test-data/expected-output/test1-extract-path-7-10.gfa b/test-data/expected-output/test1-extract-path-7-10.gfa
new file mode 100644
index 0000000..327c140
--- /dev/null
+++ b/test-data/expected-output/test1-extract-path-7-10.gfa
@@ -0,0 +1,8 @@
+H	VN:Z:1.0
+S	1	CAAATAAG
+L	1	+	3	+	0M
+S	3	G
+L	3	+	5	+	0M
+S	5	C
+P	x	1+,3+,5+	*
+P	y	1+	*
diff --git a/test-data/expected-output/test1-extract-path-7-8.gfa b/test-data/expected-output/test1-extract-path-7-8.gfa
new file mode 100644
index 0000000..6ac10e2
--- /dev/null
+++ b/test-data/expected-output/test1-extract-path-7-8.gfa
@@ -0,0 +1,4 @@
+H	VN:Z:1.0
+S	1	CAAATAAG
+P	x	1+	*
+P	y	1+	*
diff --git a/test-data/expected-output/test1-extract-path-7-9.gfa b/test-data/expected-output/test1-extract-path-7-9.gfa
new file mode 100644
index 0000000..d013efd
--- /dev/null
+++ b/test-data/expected-output/test1-extract-path-7-9.gfa
@@ -0,0 +1,6 @@
+H	VN:Z:1.0
+S	1	CAAATAAG
+L	1	+	3	+	0M
+S	3	G
+P	x	1+,3+	*
+P	y	1+	*
diff --git a/test-data/expected-output/test1-extract-path-8-10.gfa b/test-data/expected-output/test1-extract-path-8-10.gfa
new file mode 100644
index 0000000..703d856
--- /dev/null
+++ b/test-data/expected-output/test1-extract-path-8-10.gfa
@@ -0,0 +1,5 @@
+H	VN:Z:1.0
+S	3	G
+L	3	+	5	+	0M
+S	5	C
+P	x	3+,5+	*
diff --git a/test-data/expected-output/test1-extract-path-8-9.gfa b/test-data/expected-output/test1-extract-path-8-9.gfa
new file mode 100644
index 0000000..243a879
--- /dev/null
+++ b/test-data/expected-output/test1-extract-path-8-9.gfa
@@ -0,0 +1,3 @@
+H	VN:Z:1.0
+S	3	G
+P	x	3+	*
diff --git a/tests/test_domagi.py b/tests/test_domagi.py
index d7190a0..621c8c9 100644
--- a/tests/test_domagi.py
+++ b/tests/test_domagi.py
@@ -233,6 +233,38 @@ def test_domagi_extract_node(tmp_path, request, domagi_db_name, context_steps, e
     with open(expected_output) as file:
         assert_gfa_equal(file, io.StringIO(result.stdout))
 
+@pytest.mark.parametrize("domagi_db_name, path_range, expected_output",
+                         [("domagi_db_test1",
+                           "x:7-8",
+                           Path("test-data/expected-output/test1-extract-path-7-8.gfa")),
+                          ("domagi_db_test1",
+                           "x:7-9",
+                           Path("test-data/expected-output/test1-extract-path-7-9.gfa")),
+                          ("domagi_db_test1",
+                           "x:7-10",
+                           Path("test-data/expected-output/test1-extract-path-7-10.gfa")),
+                          ("domagi_db_test1",
+                           "x:8-9",
+                           Path("test-data/expected-output/test1-extract-path-8-9.gfa")),
+                          ("domagi_db_test1",
+                           "x:8-10",
+                           Path("test-data/expected-output/test1-extract-path-8-10.gfa"))])
+def test_domagi_extract_path(tmp_path, request, domagi_db_name, path_range, expected_output):
+    domagi_db = request.getfixturevalue(domagi_db_name)
+    output_duckdb_path = tmp_path / f"{domagi_db.stem}-output.db"
+    runner = CliRunner()
+    result = runner.invoke(main, ["extract",
+                                  "--db", domagi_db,
+                                  "--path-range", path_range,
+                                  "--out", output_duckdb_path])
+    assert result.exit_code == 0
+    result = runner.invoke(main, ["view",
+                                  "--to-gfa",
+                                  "--db", output_duckdb_path])
+    assert result.exit_code == 0
+    with open(expected_output) as file:
+        assert_gfa_equal(file, io.StringIO(result.stdout))
+
 @pytest.mark.parametrize("domagi_db_name, expected_output",
                          [("domagi_db_test1",
                            Path("test-data/expected-output/test1-matrix")),