tree_nine/summarize_changes_alt.py at main · aofarrel/tree_nine · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
import sys
import logging
from datetime import datetime, timezone
import polars as pl
today = datetime.now(timezone.utc)
pl.Config.set_tbl_rows(-1)
pl.Config.set_tbl_cols(-1)
pl.Config.set_tbl_width_chars(150)
pl.Config.set_fmt_str_lengths(5000)
pl.Config.set_fmt_table_cell_list_len(5000)

# pylint: disable=use-list-literal,duplicate-code

# polars refuses to output lists in TSV so we're going to cheat by claiming our sample IDs are actually strings
schema_overrides = {"sample_id": pl.Utf8}

df = pl.read_ndjson(sys.argv[1], ignore_errors=True, schema_overrides=schema_overrides)

df = df.select(['cluster_id', 'cluster_distance', 'n_samples', 'microreact_url', 'sample_id']).rename(
    {"cluster_distance": "distance",
    "sample_id": "samples"}
)
df.write_csv('all_clusters.tsv', separator='\t', include_header=True, quote_style="never")
df.filter(pl.col("distance") == pl.lit(20)).write_csv('all_20_clusters.tsv', separator='\t', include_header=True, quote_style="never")
df.filter(pl.col("distance") == pl.lit(10)).write_csv('all_10_clusters.tsv', separator='\t', include_header=True, quote_style="never")
df.filter(pl.col("distance") == pl.lit(5)).write_csv('all_5_clusters.tsv', separator='\t', include_header=True, quote_style="never")

# do it again, but without schema overrides (reading from the disk is faster trust)
df = pl.read_ndjson(sys.argv[1], ignore_errors=True)

change_report = list()

for row in df.iter_rows(named=True):
    logging.debug("Checking %s", row['cluster_id'])
    try:
        what_is = set(row["sample_id"])
    except TypeError:
        what_is = set()
    try:
        what_was = set(row["sample_id_previously"])
    except TypeError:
        what_was = set()
    gained = list(what_is - what_was)
    lost = list(what_was - what_is)
    logging.debug("what is: %s", what_is)
    logging.debug("what was: %s", what_was)
    logging.debug("gained: %s", gained)
    logging.debug("lost: %s", lost)
    change_report.append({"cluster": f"{row['cluster_id']}@{row['cluster_distance']}",
    	"gained": gained, "lost": lost, "kept": list(what_is.intersection(what_was)),
    	"microreact_url": row['microreact_url']})
change_report_df = pl.DataFrame(change_report).with_columns([
    pl.when(pl.col('gained').list.len() == 0).then(None).otherwise(pl.col('gained')).alias("gained"),
    pl.when(pl.col('lost').list.len() == 0).then(None).otherwise(pl.col('lost')).alias("lost"),
    pl.when(pl.col('kept').list.len() == 0).then(None).otherwise(pl.col('kept')).alias("kept"),
    pl.when(pl.col('microreact_url').is_null()).then(None).otherwise(pl.col('microreact_url')).alias("microreact_url"),
])
change_report_df = change_report_df.with_columns([
	pl.when(pl.col('gained').is_not_null()).then(pl.col('gained').list.len()).otherwise(pl.lit(0)).alias("n_gained"),
	pl.when(pl.col('lost').is_not_null()).then(pl.col('lost').list.len()).otherwise(pl.lit(0)).alias("n_lost"),
	pl.when(pl.col('kept').is_not_null()).then(pl.col('kept').list.len()).otherwise(pl.lit(0)).alias("n_kept"),
])
change_report_df = change_report_df.with_columns(n_now=pl.col('n_gained')-pl.col('n_lost')+pl.col('n_kept'))

abbreviated_change_report_df = change_report_df.filter(
    (pl.col('n_gained')).gt(pl.lit(0))
    .or_(pl.col('n_lost').gt(pl.lit(0)))
)

print("Here's how clusters have changed:")
print(abbreviated_change_report_df)
#change_report_df.write_ndjson(f'change_report{today.isoformat()}.json')

print("Existing clusters that lost samples (note: it's possible to gain and lose)")
lost_samples = change_report_df.filter(pl.col("lost").is_not_null()).select(['cluster', 'n_gained', 'n_lost', 'n_kept', 'microreact_url', 'lost'])
print(lost_samples)

print("Existing clusters that gained samples (note: it's possible to gain and lose)")
gained_samples = change_report_df.filter((pl.col("gained").is_not_null().and_(pl.col("kept").is_not_null()))).select(['cluster', 'n_gained', 'n_lost', 'n_kept', 'n_now', 'microreact_url', 'gained'])
print(gained_samples)

print("Brand new clusters")
new = change_report_df.filter((pl.col("gained").is_not_null().and_(pl.col("kept").is_null()))).select(['cluster', 'n_gained', 'microreact_url', 'gained'])
print(new)

print("Decimated clusters")
decimated = change_report_df.filter((pl.col("lost").is_null()).and_(pl.col("gained").is_null()).and_(pl.col("kept").is_null())).select(['cluster', 'n_gained', 'n_lost', 'n_kept', 'n_now', 'microreact_url'])
print(decimated)

#print("Unchanged clusters")
#unchanged = change_report_df.filter((pl.col("lost").is_null()).and_(pl.col("gained").is_null()).and_(pl.col("kept").is_not_null())).select(['cluster', 'n_now', 'microreact_url'])
#print(unchanged)