Skip to content

Commit 5ba7d9c

Browse files
davidmeijerdavidmeijer
committed
WIP: cross-modal retrieval
1 parent 41cd220 commit 5ba7d9c

7 files changed

Lines changed: 800 additions & 49 deletions

File tree

‎alembic/versions/0006_rev.py‎

Lines changed: 20 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,20 @@
1+
import sqlalchemy as sa
2+
from alembic import op
3+
from pgvector.sqlalchemy import Vector
4+
from sqlalchemy.dialects.postgresql import BIT
5+
6+
revision = "0006_rev"
7+
down_revision = "0005_rev"
8+
branch_labels = None
9+
depends_on = None
10+
11+
12+
def upgrade():
13+
op.execute("CREATE EXTENSION IF NOT EXISTS vector") # ensure pgvector is available (no-op if already installed)
14+
op.drop_column("retromol_compound", "fp_retro")
15+
op.add_column("retromol_compound", sa.Column("coverage", sa.Float(), nullable=True))
16+
17+
18+
def downgrade():
19+
op.add_column("retromol_compound", sa.Column("fp_retro", BIT(512), nullable=True))
20+
op.drop_column("retromol_compound", "coverage")

‎alembic/versions/0007_rev.py‎

Lines changed: 141 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,141 @@
1+
import sqlalchemy as sa
2+
from alembic import op
3+
from pgvector.sqlalchemy import Vector
4+
from sqlalchemy.dialects.postgresql import BIT, JSONB
5+
6+
revision = "0007_rev"
7+
down_revision = "0006_rev"
8+
branch_labels = None
9+
depends_on = None
10+
11+
12+
def upgrade():
13+
op.execute("CREATE EXTENSION IF NOT EXISTS vector;")
14+
op.create_table(
15+
"biocracker_genbank",
16+
sa.Column("id", sa.BigInteger, primary_key=True),
17+
sa.Column(
18+
"genbank_region_id",
19+
sa.BigInteger,
20+
sa.ForeignKey("genbank_region.id", ondelete="CASCADE"),
21+
nullable=False,
22+
),
23+
sa.Column("result_json", JSONB, nullable=False),
24+
sa.Column(
25+
"created_at",
26+
sa.DateTime(timezone=True),
27+
server_default=sa.text("NOW()"),
28+
nullable=False,
29+
),
30+
sa.Column(
31+
"updated_at",
32+
sa.DateTime(timezone=True),
33+
server_default=sa.text("NOW()"),
34+
nullable=False,
35+
),
36+
sa.UniqueConstraint(
37+
"genbank_region_id",
38+
name="uq_biocracker_genbank_genbank_region"
39+
),
40+
)
41+
42+
# Explicit index for clarity
43+
op.create_index(
44+
"ix_biocracker_genbank_genbank_region_id",
45+
"biocracker_genbank",
46+
["genbank_region_id"]
47+
)
48+
49+
op.create_table(
50+
"retrofingerprint",
51+
sa.Column("id", sa.BigInteger, primary_key=True),
52+
sa.Column(
53+
"retromol_compound_id",
54+
sa.BigInteger,
55+
sa.ForeignKey("retromol_compound.id", ondelete="CASCADE"),
56+
nullable=True,
57+
),
58+
sa.Column(
59+
"biocracker_genbank_id",
60+
sa.BigInteger,
61+
sa.ForeignKey("biocracker_genbank.id", ondelete="CASCADE"),
62+
nullable=True,
63+
),
64+
sa.Column("fp_retro_b512_bit", BIT(512), nullable=False),
65+
sa.Column("fp_retro_b512_pop", sa.SmallInteger, nullable=False),
66+
sa.Column("fp_retro_b512_vec_binary", Vector(512), nullable=False),
67+
sa.Column("fp_retro_b512_vec_counted", Vector(512), nullable=False),
68+
sa.Column(
69+
"created_at",
70+
sa.DateTime(timezone=True),
71+
server_default=sa.text("NOW()"),
72+
nullable=False,
73+
),
74+
sa.Column(
75+
"updated_at",
76+
sa.DateTime(timezone=True),
77+
server_default=sa.text("NOW()"),
78+
nullable=False,
79+
),
80+
)
81+
82+
# Explicit indexes for clarity
83+
op.create_index(
84+
"ix_retrofingerprint_retromol_compound_id",
85+
"retrofingerprint",
86+
["retromol_compound_id"],
87+
)
88+
op.create_index(
89+
"ix_retrofingerprint_biocracker_genbank_id",
90+
"retrofingerprint",
91+
["biocracker_genbank_id"],
92+
)
93+
94+
# Timestamp trigger (table-scoped function name to avoid collisions)
95+
op.execute("""
96+
CREATE OR REPLACE FUNCTION public.set_timestamp_biocracker_genbank()
97+
RETURNS TRIGGER AS $$
98+
BEGIN
99+
NEW.updated_at = NOW(); -- use clock_timestamp() if you want wall-clock per statement
100+
RETURN NEW;
101+
END;
102+
$$ LANGUAGE plpgsql;
103+
""")
104+
op.execute("""
105+
CREATE TRIGGER biocracker_genbank_set_timestamp
106+
BEFORE UPDATE ON public.biocracker_genbank
107+
FOR EACH ROW
108+
EXECUTE FUNCTION public.set_timestamp_biocracker_genbank();
109+
""")
110+
111+
op.execute("""
112+
CREATE OR REPLACE FUNCTION public.set_timestamp_retrofingerprint()
113+
RETURNS TRIGGER AS $$
114+
BEGIN
115+
NEW.updated_at = NOW(); -- use clock_timestamp() if you want wall-clock per statement
116+
RETURN NEW;
117+
END;
118+
$$ LANGUAGE plpgsql;
119+
""")
120+
op.execute("""
121+
CREATE TRIGGER retrofingerprint_set_timestamp
122+
BEFORE UPDATE ON public.retrofingerprint
123+
FOR EACH ROW
124+
EXECUTE FUNCTION public.set_timestamp_retrofingerprint();
125+
""")
126+
127+
128+
def downgrade():
129+
# drop trigger and its function first
130+
op.execute("DROP TRIGGER IF EXISTS retrofingerprint_set_timestamp ON public.retrofingerprint;")
131+
op.execute("DROP FUNCTION IF EXISTS public.set_timestamp_retrofingerprint;")
132+
133+
op.execute("DROP TRIGGER IF EXISTS biocracker_genbank_set_timestamp ON public.biocracker_genbank;")
134+
op.execute("DROP FUNCTION IF EXISTS public.set_timestamp_biocracker_genbank;")
135+
136+
op.drop_index("ix_retrofingerprint_biocracker_genbank_id", table_name="retrofingerprint")
137+
op.drop_index("ix_retrofingerprint_retromol_compound_id", table_name="retrofingerprint")
138+
op.drop_table("retrofingerprint")
139+
140+
op.drop_index("ix_biocracker_genbank_genbank_region_id", table_name="biocracker_genbank")
141+
op.drop_table("biocracker_genbank")

‎src/bionexus/cli.py‎

Lines changed: 98 additions & 47 deletions
Original file line numberDiff line numberDiff line change
@@ -251,15 +251,6 @@ def cmd_parse_compounds(args: argparse.Namespace) -> None:
251251
console.print(f"Used RetroMol to parse {n_parsed} compounds")
252252

253253

254-
def cmd_parse_bgcs(args: argparse.Namespace) -> None:
255-
"""
256-
Parse BGCs from the database using RetroMol.
257-
258-
:param args: command-line arguments
259-
"""
260-
raise NotImplementedError("BGC parsing with RetroMol is not yet implemented")
261-
262-
263254
def cmd_compute_fp_morgan(args: argparse.Namespace) -> None:
264255
"""
265256
Compute Morgan fingerprints for compounds with SMILES.
@@ -272,14 +263,16 @@ def cmd_compute_fp_morgan(args: argparse.Namespace) -> None:
272263
console.print(f"Computed fingerprints for {done} compounds (batch={args.batch})")
273264

274265

275-
def cmd_compute_fp_retro(args: argparse.Namespace) -> None:
266+
def cmd_compute_fp_retro_compound(args: argparse.Namespace) -> None:
276267
"""
277268
Compute biosynthetic fingerprints for compounds and/or GenBank records.
278269
279270
:param args: command-line arguments
280271
"""
281-
console.print("[yellow]Biosynthetic fingerprint computation is not yet implemented[/]")
282-
exit(1)
272+
from bionexus.etl.retromol import backfill_retro_fingerprints
273+
274+
done = backfill_retro_fingerprints(cache_dir=args.cache_dir, batch=args.batch, recompute=args.recompute)
275+
console.print(f"Computed RetroMol fingerprints for {done} entries (batch={args.batch})")
283276

284277

285278
def cmd_dump_db(args: argparse.Namespace) -> None:
@@ -353,7 +346,7 @@ def cmd_search_morgan(args: argparse.Namespace) -> None:
353346
if not args.out:
354347
# console.print(df[["source", "name", "jacc"]])
355348
table = Table(show_header=True, header_style="bold magenta")
356-
table.add_column("id", style="dim", width=6)
349+
table.add_column("compound_id", style="dim", width=6)
357350
table.add_column("source", style="dim", width=10)
358351
table.add_column("name", style="white", width=30)
359352
table.add_column("jacc", justify="right")
@@ -371,14 +364,81 @@ def cmd_search_morgan(args: argparse.Namespace) -> None:
371364
console.print(f"Wrote {len(df)} results to [green]{args.out}[/]")
372365

373366

374-
def cmd_search_retro(args: argparse.Namespace) -> None:
367+
def cmd_search_retro_compound(args: argparse.Namespace) -> None:
368+
"""
369+
Search compounds by biosynthetic similarity to a given RetroMol fingerprint.
370+
371+
:param args: command-line arguments
372+
"""
373+
from bionexus.db.search import retro_search_compound
374+
375+
rows = retro_search_compound(
376+
smiles=args.smiles,
377+
top_k=args.top_k,
378+
counted=getattr(args, "counted", False),
379+
)
380+
df = pd.DataFrame(rows)
381+
382+
if not args.out:
383+
# console.print(df[["source", "name", "jacc"]])
384+
table = Table(show_header=True, header_style="bold magenta")
385+
table.add_column("compound_id", style="dim", width=6)
386+
table.add_column("source", style="dim", width=10)
387+
table.add_column("name", style="white", width=30)
388+
table.add_column("cosine", justify="right")
389+
for _, row in df.iterrows():
390+
table.add_row(
391+
f"{row['id']}",
392+
str(row["source"]),
393+
f"[cyan]{row['name']}",
394+
f"{row['cosine']:.3f}",
395+
)
396+
console.print(table)
397+
else:
398+
sep = "," if args.out.lower().endswith(".csv") else "\t"
399+
df.to_csv(args.out, index=False, sep=sep)
400+
console.print(f"Wrote {len(df)} results to [green]{args.out}[/]")
401+
402+
403+
def cmd_search_retro_gbk(args: argparse.Namespace) -> None:
375404
"""
376-
Search compounds and/or BGCs by biosynthetic fingerprint similarity.
405+
Search compounds by biosynthetic similarity to a GenBank file.
377406
378407
:param args: command-line arguments
379408
"""
380-
console.print("[yellow]Searching by biosynthetic fingerprint is not yet implemented[/]")
381-
exit(1)
409+
from bionexus.db.search import retro_search_gbk
410+
411+
rows = retro_search_gbk(
412+
path=args.path,
413+
top_k=args.top_k,
414+
readout_toplevel=args.readout_toplevel,
415+
readout_sublevel=args.readout_sublevel,
416+
counted=getattr(args, "counted", False),
417+
cache_dir=getattr(args, "cache_dir", None),
418+
)
419+
df = pd.DataFrame(rows)
420+
421+
if not args.out:
422+
# console.print(df[["source", "name", "jacc"]])
423+
table = Table(show_header=True, header_style="bold magenta")
424+
table.add_column("record", style="dim", width=30)
425+
table.add_column("compound_id", style="dim", width=6)
426+
table.add_column("source", style="dim", width=10)
427+
table.add_column("name", style="white", width=30)
428+
table.add_column("cosine", justify="right")
429+
for _, row in df.iterrows():
430+
table.add_row(
431+
f"{row['record']}",
432+
f"{row['id']}",
433+
str(row["source"]),
434+
f"[cyan]{row['name']}",
435+
f"{row['cosine']:.3f}",
436+
)
437+
console.print(table)
438+
else:
439+
sep = "," if args.out.lower().endswith(".csv") else "\t"
440+
df.to_csv(args.out, index=False, sep=sep)
441+
console.print(f"Wrote {len(df)} results to [green]{args.out}[/]")
382442

383443

384444
def build_parser() -> argparse.ArgumentParser:
@@ -461,33 +521,16 @@ def build_parser() -> argparse.ArgumentParser:
461521
p_parse_compounds.add_argument("--workers", type=int, default=1, help="Number of parallel workers to use")
462522
p_parse_compounds.set_defaults(func=cmd_parse_compounds)
463523

464-
p_parse_bgcs = sub.add_parser("parse-bgcs", help="Parse BGCs from database with RetroMol")
465-
p_parse_bgcs.add_argument("--batch", type=int, default=2000)
466-
p_parse_bgcs.add_argument("--recompute", action="store_true", help="Force recomputation for all BGCs")
467-
p_parse_bgcs.set_defaults(func=cmd_parse_bgcs)
468-
469524
p_fp = sub.add_parser("compute-fp-morgan", help="Compute fingerprints for compounds with SMILES")
470525
p_fp.add_argument("--batch", type=int, default=2000)
471526
p_fp.add_argument("--recompute", action="store_true", help="Force recomputation for all compounds")
472527
p_fp.set_defaults(func=cmd_compute_fp_morgan)
473528

474-
p_do = sub.add_parser(
475-
"compute-fp-retro",
476-
help="Compute biosynthetic fingerprints for compounds and/or GenBank records",
477-
)
478-
p_do.add_argument(
479-
"--for",
480-
choices=["compounds", "gbks", "both"],
481-
default="both",
482-
help="What to compute fingerprints for",
483-
)
529+
p_do = sub.add_parser("compute-fp-retro-compound",help="Compute biosynthetic fingerprints for compounds")
530+
p_do.add_argument("--cache-dir", default=None, help="Cache/work dir for RetroMol")
484531
p_do.add_argument("--batch", type=int, default=2000)
485-
p_do.add_argument(
486-
"--recompute",
487-
action="store_true",
488-
help="Force recomputation for all compounds/records",
489-
)
490-
p_do.set_defaults(func=cmd_compute_fp_retro)
532+
p_do.add_argument("--recompute", action="store_true", help="Force recomputation")
533+
p_do.set_defaults(func=cmd_compute_fp_retro_compound)
491534

492535
p_dump = sub.add_parser("dump-db", help="Write pg_dump custom format")
493536
p_dump.add_argument("--out", default="dumps/bionexus.dump")
@@ -504,16 +547,24 @@ def build_parser() -> argparse.ArgumentParser:
504547
p_search_m.add_argument("--out", default=None, help="Optional output file (TSV/CSV)")
505548
p_search_m.set_defaults(func=cmd_search_morgan)
506549

507-
p_search_r = sub.add_parser("search-retro", help="Search compoundsa and BGCs to a biosynthetic fingerprint")
508-
p_search_r.add_argument("--for", required=True, choices=["compound", "gbk"], help="Input type")
509-
p_search_r.add_argument(
510-
"--input",
511-
required=True,
512-
type=str,
513-
help="SMILES for 'compound' input type or path to GBK region file for 'gbk' input type",
514-
)
550+
p_search_r = sub.add_parser("search-retro-compound", help="Search compounds by biosynthetic similarity to a RetroMol fingerprint")
551+
p_search_r.add_argument("--smiles", required=True)
552+
p_search_r.add_argument("--top-k", type=int, default=20)
515553
p_search_r.add_argument("--out", default=None, help="Optional output file (TSV/CSV)")
516-
p_search_r.set_defaults(func=cmd_search_retro)
554+
p_search_r.add_argument("--counted", action="store_true", help="Use counted fingerprint for search")
555+
p_search_r.set_defaults(func=cmd_search_retro_compound)
556+
557+
p_search_r_gbk = sub.add_parser("search-retro-gbk", help="Search compounds by biosynthetic similarity to a GenBank file")
558+
p_search_r_gbk.add_argument("--path", required=True, help="Path to GenBank file")
559+
p_search_r_gbk.add_argument("--top-k", type=int, default=50)
560+
p_search_r_gbk.add_argument("--readout-toplevel", choices=["region", "cand_cluster"], default="region",
561+
help="Read out fingerprints at 'region' or 'cand_cluster' level")
562+
p_search_r_gbk.add_argument("--readout-sublevel", choices=["rec", "gene"], default="rec",
563+
help="Read out fingerprints at 'rec' or 'gene' level")
564+
p_search_r_gbk.add_argument("--out", default=None, help="Optional output file (TSV/CSV)")
565+
p_search_r_gbk.add_argument("--counted", action="store_true", help="Use counted fingerprint for search")
566+
p_search_r_gbk.add_argument("--cache-dir", default=None, help="Cache/work dir for RetroMol")
567+
p_search_r_gbk.set_defaults(func=cmd_search_retro_gbk)
517568

518569
return p
519570

0 commit comments

Comments
 (0)