Support AA abundance plot - #25
Conversation
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request adds support for per-position amino acid export to facilitate AA abundance plotting. The changes include adding a toggle in the UI, updating the block model, and modifying the anarci-numbering script and Tengo workflow to generate and process the position data. Feedback was provided regarding the Python implementation's memory efficiency and a data type mismatch in the count column, suggesting a refactor using the polars library.
| if args.positions_tsv: | ||
| chain_domain = {"H": "IGHeavy", "KL": "IGLight"} | ||
| pos_rows: List[List[str]] = [] | ||
| unique_keys = list(dict.fromkeys(keys)) | ||
| for key in unique_keys: | ||
| for chain in chains: | ||
| anarci_rows, pos_labels = anarci_by_chain[chain] | ||
| residues = anarci_rows.get(key) if anarci_rows and pos_labels else None | ||
| if residues is None: | ||
| continue | ||
| ranges = REGION_RANGES[args.scheme][chain] | ||
| cdr3_start, cdr3_end = ranges["CDR3"] | ||
| for pos_label, residue in zip(pos_labels, residues): | ||
| num = position_number(pos_label) | ||
| if num is None or num < cdr3_start or num > cdr3_end: | ||
| continue | ||
| residue = (residue or "").strip() | ||
| if residue in {"", "-", "."}: | ||
| continue | ||
| pos_rows.append([key, chain_domain[chain], pos_label, residue, "1"]) | ||
| pl.DataFrame( | ||
| pos_rows, schema=["clonotypeKey", "chain", "position", "aminoacid", "count"], orient="row" | ||
| ).write_csv(args.positions_tsv, separator="\t") |
There was a problem hiding this comment.
The current implementation for generating the per-position CDR3 amino acid data builds a list of lists in memory, which can be inefficient and consume a large amount of memory for datasets with many clonotypes. Additionally, the count column is being written as a string "1", but the downstream workflow expects a numeric type (Long).
I suggest refactoring this section to use a more idiomatic polars approach. This will improve performance and memory efficiency by leveraging polars' lazy evaluation and optimized backend operations. The suggested change also corrects the data type of the count column.
if args.positions_tsv:
chain_domain = {"H": "IGHeavy", "KL": "IGLight"}
all_pos_dfs = []
for chain in chains:
anarci_rows, pos_labels = anarci_by_chain[chain]
if not anarci_rows or not pos_labels:
continue
ranges = REGION_RANGES[args.scheme][chain]
cdr3_start, cdr3_end = ranges["CDR3"]
df = pl.DataFrame(
list(anarci_rows.items()),
schema=["clonotypeKey", "residues"],
)
df = df.with_columns(pl.lit(pos_labels).alias("pos_labels"))
df = df.explode(["residues", "pos_labels"])
df = df.rename({"residues": "aminoacid", "pos_labels": "position"})
df = df.with_columns(
pl.col("position").str.extract(r"^(\d+)", 1).cast(pl.Int64).alias("pos_num")
)
df = df.filter(
(pl.col("pos_num") >= cdr3_start)
& (pl.col("pos_num") <= cdr3_end)
& (~pl.col("aminoacid").str.strip().is_in(["", "-", "."]))
)
df = df.with_columns(
pl.lit(chain_domain[chain]).alias("chain"),
pl.lit(1, dtype=pl.Int64).alias("count"),
)
all_pos_dfs.append(df.select("clonotypeKey", "chain", "position", "aminoacid", "count"))
if all_pos_dfs:
final_df = pl.concat(all_pos_dfs)
final_df.write_csv(args.positions_tsv, separator="\t")
else:
pl.DataFrame(
schema={
"clonotypeKey": pl.Utf8,
"chain": pl.Utf8,
"position": pl.Utf8,
"aminoacid": pl.Utf8,
"count": pl.Int64,
}
).write_csv(args.positions_tsv, separator="\t")
2c9b947 to
2e7cfc6
Compare
1 participant
This PR adds a new export to the redefine-clonotype block that allow to build per position AA abundance plot. 2 new columns are exported with the following specs:
New PColumns Spec
Export is added only if some of the numbering schemes is selected, column names vary depending on the selected numbering scheme, e.g.
IMGT CDR3 position,Kabat CDR3 position,Chothia CDR3 position.To view the plot using Graph Maker please use these settings: add "Stacked Bar Chart" and set up the data mapping:
Note 1: new export requires fix in clonotype browser, here is the PR: platforma-open/clonotype-browser#17
Note 2: plot is also added to cdr3-spectratype block in this PR platforma-open/cdr3-spectratype#37