strkit 0.15.0a22__tar.gz → 0.16.0__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (69) hide show
  1. {strkit-0.15.0a22/strkit.egg-info → strkit-0.16.0}/PKG-INFO +7 -8
  2. {strkit-0.15.0a22 → strkit-0.16.0}/README.md +5 -6
  3. {strkit-0.15.0a22 → strkit-0.16.0}/setup.py +1 -1
  4. strkit-0.16.0/strkit/VERSION +1 -0
  5. {strkit-0.15.0a22 → strkit-0.16.0}/strkit/call/allele.py +11 -5
  6. {strkit-0.15.0a22 → strkit-0.16.0}/strkit/call/call_locus.py +88 -126
  7. {strkit-0.15.0a22 → strkit-0.16.0}/strkit/call/call_sample.py +51 -37
  8. strkit-0.16.0/strkit/call/consensus.py +57 -0
  9. strkit-0.16.0/strkit/call/output/__init__.py +12 -0
  10. strkit-0.16.0/strkit/call/output/json_report.py +83 -0
  11. {strkit-0.15.0a22 → strkit-0.16.0}/strkit/call/output/vcf.py +5 -3
  12. {strkit-0.15.0a22 → strkit-0.16.0}/strkit/call/params.py +5 -1
  13. strkit-0.16.0/strkit/call/realign.py +123 -0
  14. {strkit-0.15.0a22 → strkit-0.16.0}/strkit/call/repeats.py +7 -71
  15. {strkit-0.15.0a22 → strkit-0.16.0}/strkit/call/types.py +62 -8
  16. {strkit-0.15.0a22 → strkit-0.16.0}/strkit/call/utils.py +16 -1
  17. {strkit-0.15.0a22 → strkit-0.16.0}/strkit/entry.py +13 -4
  18. {strkit-0.15.0a22 → strkit-0.16.0}/strkit/json.py +13 -6
  19. {strkit-0.15.0a22 → strkit-0.16.0}/strkit/mi/result.py +2 -2
  20. {strkit-0.15.0a22 → strkit-0.16.0}/strkit/viz/server.py +9 -3
  21. {strkit-0.15.0a22 → strkit-0.16.0}/strkit/viz/templates/browser.html +9 -2
  22. {strkit-0.15.0a22 → strkit-0.16.0/strkit.egg-info}/PKG-INFO +7 -8
  23. {strkit-0.15.0a22 → strkit-0.16.0}/strkit.egg-info/requires.txt +1 -1
  24. strkit-0.15.0a22/strkit/VERSION +0 -1
  25. strkit-0.15.0a22/strkit/call/consensus.py +0 -67
  26. strkit-0.15.0a22/strkit/call/output/__init__.py +0 -10
  27. strkit-0.15.0a22/strkit/call/output/json_report.py +0 -42
  28. strkit-0.15.0a22/strkit/call/realign.py +0 -58
  29. {strkit-0.15.0a22 → strkit-0.16.0}/LICENSE +0 -0
  30. {strkit-0.15.0a22 → strkit-0.16.0}/MANIFEST.in +0 -0
  31. {strkit-0.15.0a22 → strkit-0.16.0}/pyproject.toml +0 -0
  32. {strkit-0.15.0a22 → strkit-0.16.0}/setup.cfg +0 -0
  33. {strkit-0.15.0a22 → strkit-0.16.0}/strkit/__init__.py +0 -0
  34. {strkit-0.15.0a22 → strkit-0.16.0}/strkit/call/__init__.py +0 -0
  35. {strkit-0.15.0a22 → strkit-0.16.0}/strkit/call/align_matrix.py +0 -0
  36. {strkit-0.15.0a22 → strkit-0.16.0}/strkit/call/cigar.py +0 -0
  37. {strkit-0.15.0a22 → strkit-0.16.0}/strkit/call/non_daemonic_pool.py +0 -0
  38. {strkit-0.15.0a22 → strkit-0.16.0}/strkit/call/output/tsv.py +0 -0
  39. {strkit-0.15.0a22 → strkit-0.16.0}/strkit/call/snvs.py +0 -0
  40. {strkit-0.15.0a22 → strkit-0.16.0}/strkit/catalog/__init__.py +0 -0
  41. {strkit-0.15.0a22 → strkit-0.16.0}/strkit/catalog/combine.py +0 -0
  42. {strkit-0.15.0a22 → strkit-0.16.0}/strkit/constants.py +0 -0
  43. {strkit-0.15.0a22 → strkit-0.16.0}/strkit/convert/__init__.py +0 -0
  44. {strkit-0.15.0a22 → strkit-0.16.0}/strkit/convert/_bed_4.py +0 -0
  45. {strkit-0.15.0a22 → strkit-0.16.0}/strkit/convert/converter.py +0 -0
  46. {strkit-0.15.0a22 → strkit-0.16.0}/strkit/convert/expansionhunter.py +0 -0
  47. {strkit-0.15.0a22 → strkit-0.16.0}/strkit/convert/gangstr.py +0 -0
  48. {strkit-0.15.0a22 → strkit-0.16.0}/strkit/convert/hipstr.py +0 -0
  49. {strkit-0.15.0a22 → strkit-0.16.0}/strkit/convert/straglr.py +0 -0
  50. {strkit-0.15.0a22 → strkit-0.16.0}/strkit/convert/tandem_genotypes.py +0 -0
  51. {strkit-0.15.0a22 → strkit-0.16.0}/strkit/exceptions.py +0 -0
  52. {strkit-0.15.0a22 → strkit-0.16.0}/strkit/logger.py +0 -0
  53. {strkit-0.15.0a22 → strkit-0.16.0}/strkit/mi/__init__.py +0 -0
  54. {strkit-0.15.0a22 → strkit-0.16.0}/strkit/mi/base.py +0 -0
  55. {strkit-0.15.0a22 → strkit-0.16.0}/strkit/mi/expansionhunter.py +0 -0
  56. {strkit-0.15.0a22 → strkit-0.16.0}/strkit/mi/gangstr.py +0 -0
  57. {strkit-0.15.0a22 → strkit-0.16.0}/strkit/mi/repeathmm.py +0 -0
  58. {strkit-0.15.0a22 → strkit-0.16.0}/strkit/mi/straglr.py +0 -0
  59. {strkit-0.15.0a22 → strkit-0.16.0}/strkit/mi/strkit.py +0 -0
  60. {strkit-0.15.0a22 → strkit-0.16.0}/strkit/mi/tandem_genotypes.py +0 -0
  61. {strkit-0.15.0a22 → strkit-0.16.0}/strkit/mi/vcf_utils.py +0 -0
  62. {strkit-0.15.0a22 → strkit-0.16.0}/strkit/utils.py +0 -0
  63. {strkit-0.15.0a22 → strkit-0.16.0}/strkit/viz/__init__.py +0 -0
  64. {strkit-0.15.0a22 → strkit-0.16.0}/strkit/viz/static/logo.png +0 -0
  65. {strkit-0.15.0a22 → strkit-0.16.0}/strkit.egg-info/SOURCES.txt +0 -0
  66. {strkit-0.15.0a22 → strkit-0.16.0}/strkit.egg-info/dependency_links.txt +0 -0
  67. {strkit-0.15.0a22 → strkit-0.16.0}/strkit.egg-info/entry_points.txt +0 -0
  68. {strkit-0.15.0a22 → strkit-0.16.0}/strkit.egg-info/top_level.txt +0 -0
  69. {strkit-0.15.0a22 → strkit-0.16.0}/tests/test_caller_utils.py +0 -0
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.1
2
2
  Name: strkit
3
- Version: 0.15.0a22
3
+ Version: 0.16.0
4
4
  Summary: A toolkit for analyzing variation in short(ish) tandem repeats.
5
5
  Home-page: https://github.com/davidlougheed/strkit
6
6
  Author: David Lougheed
@@ -20,7 +20,7 @@ Requires-Dist: parasail<1.4,>=1.2.4
20
20
  Requires-Dist: scikit-learn<1.5,>=1.2.1
21
21
  Requires-Dist: scipy<1.14,>=1.10
22
22
  Requires-Dist: statsmodels<0.15,>=0.14.0
23
- Requires-Dist: strkit_rust_ext==0.12.2
23
+ Requires-Dist: strkit_rust_ext==0.13.0
24
24
 
25
25
  # STRkit - short tandem repeat genotyping with long reads
26
26
 
@@ -88,10 +88,10 @@ You may need to install the [Rust toolchain](https://www.rust-lang.org/tools/ins
88
88
  to compile the `strkit_rust_ext` wheels, although prebuilt wheels for this module are available
89
89
  for some platforms.
90
90
 
91
- On Digital Research Alliance of Canada/Compute Canada clusters, this involves loading a module:
91
+ On Digital Research Alliance of Canada/Compute Canada clusters, this involves loading a few modules:
92
92
 
93
93
  ```bash
94
- module load rust/1.76.0
94
+ module load rust/1.76.0 clang/17.0.6 python/3.11 scipy-stack/2023b parasail/2.6.2
95
95
  python -m pip install strkit
96
96
  ```
97
97
 
@@ -138,8 +138,8 @@ long read data should still work.
138
138
  * Whole-genome and targeted genotyping modes to adjust this re-weighting.
139
139
  * Incorporation of single-nucleotide variation (SNVs) for better and faster calling plus
140
140
  additional downstream analysis possibilities.
141
- * Recommended for **HiFi data only**. In my testing, this worsens runtime and call quality for
142
- ONT-UL data, but speeds up the tool and improves call quality for HiFi data.
141
+ * Recommended for **HiFi data and ONT duplex data only**. In my testing, this worsens runtime and call quality for
142
+ ONT ultra-long-read data, but speeds up the tool and improves call quality for HiFi/ONT duplex data.
143
143
  * Parallelized for faster computing on clusters and for ad-hoc fast analysis of single samples.
144
144
  * 95% confidence intervals on calls via a user-configurable optional parametric bootstrapping process.
145
145
 
@@ -149,7 +149,7 @@ long read data should still work.
149
149
  ```bash
150
150
  strkit call \
151
151
  path/to/read/file.bam \ # [REQUIRED] One indexed read file (BAM/CRAM)
152
- --hq \ # If using PacBio HiFi reads, enable this to get better genotyping & more robust expansion detection
152
+ --hq \ # If using PacBio HiFi or ONT duplex reads, enable this to get better genotyping & more robust expansion detection
153
153
  --realign \ # If using PacBio HiFi reads, enable this to enable local realignment / read recovery. Good for detecting expansions, but slows down calling.
154
154
  --ref path/to/reference.fa.gz \ # [REQUIRED] Indexed FASTA-formatted reference genome
155
155
  --loci path/to/loci.bed \ # [REQUIRED] TRF-formatted (or 4-col, with motif as last column) sorted list of loci to genotype
@@ -187,7 +187,6 @@ If you want to output a full call report, you can use the `--json output-file.js
187
187
  specify a path to output a more detailed JSON document to. This document contains 99% CIs, peak
188
188
  labels, and some other information that isn't included in the normal TSV file. If you want this
189
189
  file to be indented and human-readable, use the `--indent-json` flag in addition to `--json ...`.
190
- Note that **memory usage will be higher** when JSON output is enabled.
191
190
 
192
191
  If you want to output a VCF file (STRs and SNVs if called; currently not phased), use the
193
192
  `--vcf ...` argument. If you pass `--vcf stdout`, the VCF will be written to `stdout` instead of a
@@ -64,10 +64,10 @@ You may need to install the [Rust toolchain](https://www.rust-lang.org/tools/ins
64
64
  to compile the `strkit_rust_ext` wheels, although prebuilt wheels for this module are available
65
65
  for some platforms.
66
66
 
67
- On Digital Research Alliance of Canada/Compute Canada clusters, this involves loading a module:
67
+ On Digital Research Alliance of Canada/Compute Canada clusters, this involves loading a few modules:
68
68
 
69
69
  ```bash
70
- module load rust/1.76.0
70
+ module load rust/1.76.0 clang/17.0.6 python/3.11 scipy-stack/2023b parasail/2.6.2
71
71
  python -m pip install strkit
72
72
  ```
73
73
 
@@ -114,8 +114,8 @@ long read data should still work.
114
114
  * Whole-genome and targeted genotyping modes to adjust this re-weighting.
115
115
  * Incorporation of single-nucleotide variation (SNVs) for better and faster calling plus
116
116
  additional downstream analysis possibilities.
117
- * Recommended for **HiFi data only**. In my testing, this worsens runtime and call quality for
118
- ONT-UL data, but speeds up the tool and improves call quality for HiFi data.
117
+ * Recommended for **HiFi data and ONT duplex data only**. In my testing, this worsens runtime and call quality for
118
+ ONT ultra-long-read data, but speeds up the tool and improves call quality for HiFi/ONT duplex data.
119
119
  * Parallelized for faster computing on clusters and for ad-hoc fast analysis of single samples.
120
120
  * 95% confidence intervals on calls via a user-configurable optional parametric bootstrapping process.
121
121
 
@@ -125,7 +125,7 @@ long read data should still work.
125
125
  ```bash
126
126
  strkit call \
127
127
  path/to/read/file.bam \ # [REQUIRED] One indexed read file (BAM/CRAM)
128
- --hq \ # If using PacBio HiFi reads, enable this to get better genotyping & more robust expansion detection
128
+ --hq \ # If using PacBio HiFi or ONT duplex reads, enable this to get better genotyping & more robust expansion detection
129
129
  --realign \ # If using PacBio HiFi reads, enable this to enable local realignment / read recovery. Good for detecting expansions, but slows down calling.
130
130
  --ref path/to/reference.fa.gz \ # [REQUIRED] Indexed FASTA-formatted reference genome
131
131
  --loci path/to/loci.bed \ # [REQUIRED] TRF-formatted (or 4-col, with motif as last column) sorted list of loci to genotype
@@ -163,7 +163,6 @@ If you want to output a full call report, you can use the `--json output-file.js
163
163
  specify a path to output a more detailed JSON document to. This document contains 99% CIs, peak
164
164
  labels, and some other information that isn't included in the normal TSV file. If you want this
165
165
  file to be indented and human-readable, use the `--indent-json` flag in addition to `--json ...`.
166
- Note that **memory usage will be higher** when JSON output is enabled.
167
166
 
168
167
  If you want to output a VCF file (STRs and SNVs if called; currently not phased), use the
169
168
  `--vcf ...` argument. If you pass `--vcf stdout`, the VCF will be written to `stdout` instead of a
@@ -22,7 +22,7 @@ setup(
22
22
  "scikit-learn>=1.2.1,<1.5",
23
23
  "scipy>=1.10,<1.14",
24
24
  "statsmodels>=0.14.0,<0.15",
25
- "strkit_rust_ext==0.12.2",
25
+ "strkit_rust_ext==0.13.0",
26
26
  ],
27
27
 
28
28
  description="A toolkit for analyzing variation in short(ish) tandem repeats.",
@@ -0,0 +1 @@
1
+ 0.16.0
@@ -58,7 +58,7 @@ def _array_as_int(n: Union[NDArray[np.int_], NDArray[np.float_]]) -> NDArray[np.
58
58
  return np.rint(n).astype(np.int32)
59
59
 
60
60
 
61
- def _calculate_cis(samples, ci: str = Literal["95", "99"]) -> Union[NDArray[np.int32], NDArray[np.float_]]:
61
+ def _calculate_cis(samples, ci: str = Literal["95", "99"]) -> NDArray[np.int32]:
62
62
  percentiles = np.percentile(
63
63
  samples, CI_PERCENTILE_RANGES[ci], axis=1, method="interpolated_inverted_cdf"
64
64
  ).transpose()
@@ -145,7 +145,7 @@ def fit_gmm(
145
145
  return g
146
146
 
147
147
 
148
- class CallDict(TypedDict):
148
+ class BaseCallDict(TypedDict):
149
149
  call: Union[NDArray[np.int32], NDArray[np.float_]]
150
150
  call_95_cis: Union[NDArray[np.int32], NDArray[np.float_]] # 2D arrays
151
151
  call_99_cis: Union[NDArray[np.int32], NDArray[np.float_]] # 2D arrays
@@ -155,6 +155,10 @@ class CallDict(TypedDict):
155
155
  modal_n_peaks: int
156
156
 
157
157
 
158
+ class CallDict(BaseCallDict, total=False):
159
+ ps: int
160
+
161
+
158
162
  def make_read_weights(read_weights: Optional[Iterable[float]], num_reads: int) -> NDArray[np.float_]:
159
163
  return np.array(
160
164
  read_weights if read_weights is not None else np.array(([1/num_reads] * num_reads) if num_reads else []))
@@ -262,7 +266,7 @@ def call_alleles(
262
266
  gmm_cache = {}
263
267
 
264
268
  def _get_fitted_gmm(s: Union[NDArray[np.int_], NDArray[np.float_]]) -> Optional[object]:
265
- if (s_t := tuple(s)) not in gmm_cache:
269
+ if (s_t := s.tobytes()) not in gmm_cache:
266
270
  # Fit Gaussian mixture model to the resampled data
267
271
  gmm_cache[s_t] = fit_gmm(rng, s, n_alleles, allele_filter, params.hq, gm_filter_factor)
268
272
 
@@ -335,17 +339,19 @@ def call_alleles(
335
339
  median_idx = allele_samples.shape[1] // 2 #
336
340
  medians_of_means = allele_samples[:, median_idx]
337
341
  medians_of_means_final = np.rint(medians_of_means).astype(np.int32)
338
- peak_weights = allele_weight_samples[:, median_idx]
342
+ peak_weights = allele_weight_samples[:, median_idx].flatten()
339
343
  peak_stdevs = allele_stdev_samples[:, median_idx]
340
344
  modal_n_peaks: int = statistics.mode(sample_peaks).item()
341
345
 
346
+ peak_weights /= peak_weights.sum() # re-normalize weights
347
+
342
348
  return {
343
349
  "call": medians_of_means_final.flatten(),
344
350
  "call_95_cis": allele_cis_95,
345
351
  "call_99_cis": allele_cis_99,
346
352
 
347
353
  "peaks": medians_of_means.flatten(), # Don't round, so we can recover original Gaussian model
348
- "peak_weights": peak_weights.flatten(),
354
+ "peak_weights": peak_weights,
349
355
  "peak_stdevs": peak_stdevs.flatten(),
350
356
  # TODO: should be ok to use this, because resample gets put at end, vertically (3rd allele in a 3-ploid case)
351
357
  # so taking the first 2 alleles still works in terms of stdev/mean estimates? I think?
@@ -2,12 +2,9 @@ from __future__ import annotations
2
2
 
3
3
  import functools
4
4
  import logging
5
- import multiprocessing as mp
6
5
  import multiprocessing.managers as mmg
7
6
  import numpy as np
8
7
  import operator
9
- import queue
10
- import time
11
8
  import threading
12
9
 
13
10
  from collections import Counter
@@ -15,22 +12,23 @@ from collections.abc import Sequence
15
12
  from datetime import datetime
16
13
  from pysam import FastaFile
17
14
  from sklearn.cluster import AgglomerativeClustering
15
+ from sklearn.mixture import GaussianMixture
18
16
 
19
17
  from numpy.typing import NDArray
20
- from typing import Iterable, Literal, Optional, Union
18
+ from typing import Iterable, Optional, Union
21
19
 
22
20
  from strkit_rust_ext import (
23
21
  get_pairs_and_tr_read_coords, STRkitBAMReader, STRkitAlignedSegment, STRkitVCFReader, CandidateSNVs
24
22
  )
25
23
 
26
24
  from strkit.call.allele import CallDict, call_alleles
27
- from strkit.utils import cat_strs, apply_or_none
25
+ from strkit.utils import apply_or_none
28
26
 
29
27
  from .align_matrix import match_score
30
28
  from .cigar import decode_cigar_np
31
- from .consensus import ConsensusMethod, consensus_seq
29
+ from .consensus import consensus_seq
32
30
  from .params import CallParams
33
- from .realign import realign_read
31
+ from .realign import perform_realign
34
32
  from .repeats import get_repeat_count, get_ref_repeat_count
35
33
  from .snvs import (
36
34
  SNV_GAP_CHAR,
@@ -38,8 +36,10 @@ from .snvs import (
38
36
  call_and_filter_useful_snvs,
39
37
  process_read_snvs_for_locus_and_calculate_useful_snvs,
40
38
  )
41
- from .types import ReadDict, ReadDictExtra, CalledSNV
42
- from .utils import idx_0_getter, find_pair_by_ref_pos, normalize_contig, get_new_seed
39
+ from .types import (
40
+ VCFContigFormat, AssignMethod, AssignMethodWithHP, ConsensusMethod, ReadDict, ReadDictExtra, CalledSNV, LocusResult
41
+ )
42
+ from .utils import idx_0_getter, find_pair_by_ref_pos, normalize_contig, get_new_seed, calculate_seq_with_wildcards
43
43
 
44
44
 
45
45
  __all__ = [
@@ -56,7 +56,6 @@ default_ref_max_iters: int = 250
56
56
  ref_max_iters_to_be_slow: int = 100
57
57
 
58
58
  max_rc_iters = 50
59
- min_read_score = 0.9 # TODO: parametrize
60
59
  # TODO: Parametrize - if very low alignment or very slow, then "bad read"
61
60
  extremely_low_read_adj_score: float = 0.1 # fraction
62
61
  bad_read_alignment_time: int = 15 # seconds
@@ -65,6 +64,8 @@ max_bad_reads: int = 3
65
64
 
66
65
  base_wildcard_threshold = 3
67
66
 
67
+ min_hp_read_coverage: int = 8 # TODO: parametrize
68
+
68
69
  roughly_equiv_stdev_dist = 1
69
70
 
70
71
  force_realign = False
@@ -73,7 +74,8 @@ many_realigns_threshold = 2
73
74
  significant_clip_threshold = 100
74
75
  significant_clip_snv_take_in = 250
75
76
 
76
- max_poa_length = 5000 # maximum number of bases before we can't use POA for consensus anymore due to performance
77
+ # maximum median number of bases before we can't use POA for consensus anymore due to performance:
78
+ max_mdn_poa_length = 5000
77
79
 
78
80
 
79
81
  # property getters & other partials
@@ -83,17 +85,6 @@ not_snv_out_of_range_char = functools.partial(operator.ne, SNV_OUT_OF_RANGE_CHAR
83
85
  eq_0 = functools.partial(operator.eq, 0)
84
86
 
85
87
 
86
- @functools.cache
87
- def _mask_low_q_base(base_and_qual: tuple[str, int]) -> str:
88
- return base_and_qual[0] if base_and_qual[1] > base_wildcard_threshold else "X"
89
-
90
-
91
- def calculate_seq_with_wildcards(qs: str, quals: Optional[NDArray[np.uint8]]) -> str:
92
- if quals is None:
93
- return qs # No quality information, so don't do anything
94
- return cat_strs(map(_mask_low_q_base, zip(qs, quals)))
95
-
96
-
97
88
  def get_read_coords_from_matched_pairs(
98
89
  left_flank_coord: int,
99
90
  left_coord: int,
@@ -131,8 +122,8 @@ def get_read_coords_from_matched_pairs(
131
122
  # ------------------------------------------------------------------------------------------------------------------
132
123
 
133
124
  for i in range(lhs + 1, len(q_coords)):
134
- query_coord = q_coords[i]
135
- ref_coord = r_coords[i]
125
+ query_coord = q_coords[i].item()
126
+ ref_coord = r_coords[i].item()
136
127
 
137
128
  # Skip gaps on either side to find mapped flank indices
138
129
  if ref_coord < left_coord:
@@ -514,8 +505,8 @@ def call_alleles_with_incorporated_snvs(
514
505
  rng: np.random.Generator,
515
506
  logger_: logging.Logger,
516
507
  locus_log_str: str,
517
- ) -> tuple[Literal["dist", "snv", "snv+dist", "single"], Optional[tuple[dict, list[dict]]]]:
518
- assign_method: Literal["dist", "snv", "snv+dist", "single"] = "dist"
508
+ ) -> tuple[AssignMethod, Optional[tuple[dict, list[CalledSNV]]]]:
509
+ assign_method: AssignMethod = "dist"
519
510
 
520
511
  # TODO: parametrize min 'enough to do pure SNV haplotyping' thresholds
521
512
 
@@ -524,8 +515,6 @@ def call_alleles_with_incorporated_snvs(
524
515
  read_dict_items_with_at_least_one_snv: list[tuple[str, ReadDict]] = []
525
516
  read_dict_items_with_no_snvs: list[tuple[str, ReadDict]] = []
526
517
 
527
- print_snvs = False
528
-
529
518
  # TODO: Check that we don't have disjoint groups - how would we actually ascertain this?
530
519
  # - each read as a set; merge sets when there is overlap, make sure we have 1 set at the end
531
520
 
@@ -556,10 +545,6 @@ def call_alleles_with_incorporated_snvs(
556
545
 
557
546
  read["snvu"] = read_useful_snv_bases # Store read-level 'useful' SNVs
558
547
 
559
- if print_snvs:
560
- print(
561
- rn, f"\t{read['cn']:.0f}", "\t", read_useful_snv_bases, n_non_blank_hq_read_useful_snv_bases)
562
-
563
548
  n_reads_with_many_snvs: int = len(read_dict_items_with_many_snvs)
564
549
  n_reads_with_at_least_one_snv: int = len(read_dict_items_with_at_least_one_snv)
565
550
  n_reads_with_no_snvs: int = len(read_dict_items_with_no_snvs)
@@ -738,7 +723,7 @@ def call_alleles_with_incorporated_snvs(
738
723
  peak_weights_pre_adj = np.concatenate(tuple(cc["peak_weights"] for cc in cdd), axis=0)
739
724
 
740
725
  # All call_datas are truth-y; all arrays should be ordered by peak_order
741
- call_data = {
726
+ call_data: CallDict = {
742
727
  "call": np.concatenate(tuple(cc["call"] for cc in cdd_ordered), axis=0),
743
728
  "call_95_cis": np.concatenate(tuple(cc["call_95_cis"] for cc in cdd_ordered), axis=0),
744
729
  "call_99_cis": np.concatenate(tuple(cc["call_99_cis"] for cc in cdd_ordered), axis=0),
@@ -797,20 +782,20 @@ def call_locus(
797
782
  # ---
798
783
  snv_vcf_file: Optional[STRkitVCFReader] = None,
799
784
  snv_vcf_contigs: tuple[str, ...] = (),
800
- snv_vcf_file_format: Literal["chr", "num", "acc", ""] = "",
785
+ snv_vcf_file_format: VCFContigFormat = "",
801
786
  # ---
802
787
  read_file_has_chr: bool = True,
803
788
  ref_file_has_chr: bool = True,
804
- ) -> Optional[dict]:
789
+ ) -> Optional[LocusResult]:
805
790
  call_timer = datetime.now()
806
791
 
807
792
  # params de-structuring ------------
808
793
  consensus = params.consensus
809
794
  count_kmers = params.count_kmers
810
795
  flank_size = params.flank_size
811
- log_level = params.log_level
812
796
  realign = params.realign
813
797
  respect_ref = params.respect_ref
798
+ min_read_align_score = params.min_read_align_score
814
799
  snv_min_base_qual = params.snv_min_base_qual
815
800
  # ----------------------------------
816
801
 
@@ -830,7 +815,7 @@ def call_locus(
830
815
  ref_seq: str = ""
831
816
  raised: bool = False
832
817
 
833
- call_dict_base = {
818
+ locus_result: LocusResult = {
834
819
  "locus_index": t_idx,
835
820
  "contig": contig,
836
821
  "start": left_coord,
@@ -872,15 +857,15 @@ def call_locus(
872
857
  raised = True
873
858
 
874
859
  if raised:
875
- return None # don't even return call_dict_base, this locus has an invalid position
860
+ return None # don't even return locus_result, this locus has an invalid position
876
861
 
877
862
  if len(ref_left_flank_seq) < params.flank_size or len(ref_right_flank_seq) < params.flank_size:
878
863
  logger_.warning(f"{locus_log_str} - skipping locus, reference flank size too small")
879
- return call_dict_base
864
+ return locus_result
880
865
 
881
866
  if ref_left_flank_seq.endswith("N" * motif_size) or ref_right_flank_seq.startswith("N" * motif_size):
882
867
  logger_.warning(f"{locus_log_str} - skipping locus, reference has flanking N[...] sequence")
883
- return call_dict_base
868
+ return locus_result
884
869
 
885
870
  # Get reference repeat count by our method, so we can calculate offsets from reference
886
871
  # - Replace flanking/ref TR sequences with adjusted sequences
@@ -917,12 +902,12 @@ def call_locus(
917
902
  local_search_range=ref_local_search_range,
918
903
  step_size=ref_step_size,
919
904
  )
920
- call_dict_base["ref_cn"] = ref_cn # tag call dictionary with ref_cn
905
+ locus_result["ref_cn"] = ref_cn # tag call dictionary with ref_cn
921
906
 
922
907
  slow_ref_count = any(x > ref_max_iters_to_be_slow for x in r_n_is)
923
908
 
924
909
  logger_.debug(
925
- f"{locus_log_str} - got ref. copy number: {ref_cn} ({l_offset=}; {r_offset=}; iters={r_n_is})")
910
+ f"{locus_log_str} - {ref_cn=} (l.o.={l_offset}; r.o.={r_offset}; #i={r_n_is})")
926
911
 
927
912
  if slow_ref_count:
928
913
  logger_.warning(f"{locus_log_str} - slow reference copy number counting ({motif=}; {ref_cn=}; iters={r_n_is})")
@@ -933,11 +918,9 @@ def call_locus(
933
918
  left_coord_adj = left_coord if respect_ref else left_coord - max(0, l_offset)
934
919
  right_coord_adj = right_coord if respect_ref else right_coord + max(0, r_offset)
935
920
 
936
- if not respect_ref: # tag call_dict_base with adjusted start/end
937
- call_dict_base.update({
938
- "start_adj": left_coord_adj,
939
- "end_adj": right_coord_adj,
940
- })
921
+ if not respect_ref: # tag locus_result with adjusted start/end
922
+ locus_result["start_adj"] = left_coord_adj
923
+ locus_result["end_adj"] = right_coord_adj
941
924
 
942
925
  ref_time = (datetime.now() - ref_timer).total_seconds()
943
926
 
@@ -965,9 +948,8 @@ def call_locus(
965
948
 
966
949
  logger_.debug(f"{locus_log_str} - got {n_overlapping_reads} overlapping aligned segments")
967
950
 
968
- if n_overlapping_reads > params.max_reads:
969
- logger_.warning(f"{locus_log_str} - skipping locus; too many overlapping reads")
970
- return call_dict_base
951
+ if n_overlapping_reads > params.max_reads: # TODO: sample across full set instead?
952
+ logger_.warning(f"{locus_log_str} - locus has excess reads, using the first {params.max_reads}; misalignment?")
971
953
 
972
954
  sorted_read_lengths = np.sort(read_lengths)
973
955
 
@@ -1034,40 +1016,21 @@ def call_locus(
1034
1016
  # Run the realignment in a separate process, to give us a timeout mechanism.
1035
1017
  # This means we're spawning a second process for this job, just momentarily, beyond the pool size.
1036
1018
 
1037
- q: mp.Queue = mp.Queue()
1038
- proc = mp.Process(target=realign_read, daemon=False, kwargs=dict(
1039
- # fetch an extra base for the right flank coordinate check later (needs to be >= the exclusive coord)
1040
- ref_seq=ref_total_seq, # TODO: with the plus 1, really?
1041
- query_seq=calculate_seq_with_wildcards(qs, fqqs),
1042
- left_flank_coord=left_flank_coord,
1043
- flank_size=flank_size,
1044
- rn=rn,
1045
- t_idx=t_idx,
1046
- always_realign=force_realign,
1047
- q=q,
1048
- log_level=log_level,
1049
- ))
1050
- proc.start()
1051
- pairs_new = None
1052
- try:
1053
- pairs_new = q.get(timeout=realign_timeout)
1054
- proc.join()
1055
- except queue.Empty:
1056
- logger_.warning(
1057
- f"{locus_log_str} - experienced timeout while re-aligning read {rn}. Reverting to initial "
1058
- f"alignment.")
1059
- proc.terminate()
1060
- time.sleep(0.1) # wait a little for the process to terminate
1061
- finally:
1062
- wait_count: int = 0
1063
- while proc.is_alive():
1064
- logger_.warning(f"{locus_log_str} - realign job has still not exited. Waiting 0.5 seconds...")
1065
- time.sleep(0.5)
1066
- wait_count += 1
1067
- if wait_count > 5:
1068
- logger_.fatal(f"{locus_log_str} - realign job never exited. Terminating...")
1069
- exit(1)
1070
- proc.close()
1019
+ pairs_new = perform_realign(
1020
+ t_idx,
1021
+ left_flank_coord,
1022
+ ref_total_seq,
1023
+ rn,
1024
+ qs,
1025
+ fqqs,
1026
+ # ---
1027
+ params,
1028
+ realign_timeout,
1029
+ force_realign,
1030
+ # ---
1031
+ logger_,
1032
+ locus_log_str,
1033
+ )
1071
1034
 
1072
1035
  if pairs_new is not None:
1073
1036
  q_coords, r_coords = pairs_new
@@ -1166,8 +1129,8 @@ def call_locus(
1166
1129
  rc_time = (datetime.now() - rc_timer).total_seconds()
1167
1130
 
1168
1131
  logger_.debug(
1169
- f"{locus_log_str} - {rn} | start: {read_sc}, offset frac.: {read_offset_frac_from_starting_guess:.4f} | "
1170
- f"rc time: {rc_time}, copy number: {read_cn}, score: {read_cn_score}")
1132
+ f"{locus_log_str} - {rn} | start={read_sc}, o.f.={read_offset_frac_from_starting_guess:.4f} | "
1133
+ f"rct={rc_time}s, cn={read_cn}, s={read_cn_score}")
1171
1134
 
1172
1135
  if n_read_cn_iters >= max_rc_iters:
1173
1136
  logger_.debug(f"{locus_log_str} - locus repeat counting exceeded maximum # iterations ({n_read_cn_iters})")
@@ -1178,7 +1141,7 @@ def call_locus(
1178
1141
  if rc_time >= really_bad_read_alignment_time:
1179
1142
  logger_.debug(
1180
1143
  f"{locus_log_str} - not calling locus due to a pathologically-poorly-aligning read ({rn}; "
1181
- f"repeat count alignment scored {read_adj_score:.2f} < {min_read_score}; get_repeat_count time: "
1144
+ f"repeat count alignment scored {read_adj_score:.2f} < {min_read_align_score}; get_repeat_count time: "
1182
1145
  f"{rc_time:.3f}s; # get_repeat_count iters: {n_read_cn_iters}; {motif=}; {ref_cn=})")
1183
1146
  logger_.debug(f"{locus_log_str} - ref left flank: {ref_left_flank_seq}")
1184
1147
  logger_.debug(f"{locus_log_str} - read left flank: {flank_left_seq}")
@@ -1187,16 +1150,16 @@ def call_locus(
1187
1150
  logger_.debug(f"{locus_log_str} - ref right flank: {ref_right_flank_seq}")
1188
1151
  logger_.debug(f"{locus_log_str} - read right flank: {flank_right_seq}")
1189
1152
  return {
1190
- **call_dict_base,
1153
+ **locus_result,
1191
1154
  "peaks": None,
1192
1155
  "read_peaks_called": False,
1193
1156
  "time": (datetime.now() - call_timer).total_seconds(),
1194
1157
  }
1195
1158
 
1196
- if read_adj_score < min_read_score:
1159
+ if read_adj_score < min_read_align_score:
1197
1160
  logger_.debug(
1198
1161
  f"{locus_log_str} - skipping read {rn} (repeat count alignment scored {read_adj_score:.2f} < "
1199
- f"{min_read_score}; get_repeat_count time: {rc_time:.3f}s; # get_repeat_count iters: "
1162
+ f"{min_read_align_score}; get_repeat_count time: {rc_time:.3f}s; # get_repeat_count iters: "
1200
1163
  f"{n_read_cn_iters})")
1201
1164
 
1202
1165
  if read_adj_score < extremely_low_read_adj_score or rc_time >= bad_read_alignment_time:
@@ -1206,7 +1169,7 @@ def call_locus(
1206
1169
  f"{locus_log_str} - not calling locus due to >3 extremely poor-aligning reads (TR seq: "
1207
1170
  f"{tr_read_seq_wc[:20]}...)")
1208
1171
  return {
1209
- **call_dict_base,
1172
+ **locus_result,
1210
1173
  "peaks": None,
1211
1174
  "read_peaks_called": False,
1212
1175
  "time": (datetime.now() - call_timer).total_seconds(),
@@ -1263,7 +1226,7 @@ def call_locus(
1263
1226
 
1264
1227
  phase_set_lock.release()
1265
1228
 
1266
- read_dict[rn]["hp"] = hp
1229
+ read_dict[rn]["hp"] = hp # not none inside this if-statement
1267
1230
  read_dict[rn]["ps"] = ps_remapped
1268
1231
  haplotags.add(hp)
1269
1232
  haplotagged_reads_count += 1
@@ -1294,7 +1257,7 @@ def call_locus(
1294
1257
 
1295
1258
  n_reads_in_dict: int = len(read_dict)
1296
1259
 
1297
- call_dict_base.update({
1260
+ locus_result.update({
1298
1261
  # TODO: alt anchors:
1299
1262
  **({"ref_start_anchor": ref_left_flank_seq[-1], "ref_seq": ref_seq} if consensus else {}),
1300
1263
  "reads": read_dict,
@@ -1304,7 +1267,7 @@ def call_locus(
1304
1267
  if n_reads_in_dict < params.min_reads:
1305
1268
  logger_.debug(f"{locus_log_str} - not enough reads to make a call ({n_reads_in_dict} < {params.min_reads})")
1306
1269
  return {
1307
- **call_dict_base,
1270
+ **locus_result,
1308
1271
  "peaks": None,
1309
1272
  "read_peaks_called": False,
1310
1273
  "time": (datetime.now() - call_timer).total_seconds(),
@@ -1316,11 +1279,7 @@ def call_locus(
1316
1279
  # noinspection PyTypeChecker
1317
1280
  read_dict_items: tuple[tuple[str, ReadDict], ...] = tuple(read_dict.items())
1318
1281
 
1319
- assign_method: Literal["dist", "snv", "snv+dist", "single", "hp"] = "dist"
1320
- if n_alleles < 2:
1321
- assign_method = "single"
1322
-
1323
- min_hp_read_coverage: int = 8 # TODO: parametrize
1282
+ assign_method: AssignMethodWithHP = "single" if n_alleles == 1 else "dist"
1324
1283
 
1325
1284
  # Realigns are missing significant amounts of flanking information since the realignment only uses a portion of the
1326
1285
  # reference genome. If we have a rare realignment (e.g., a large expansion), we cannot use SNVs.
@@ -1413,7 +1372,7 @@ def call_locus(
1413
1372
  assign_method = am
1414
1373
  if call_res is not None:
1415
1374
  call_data = call_res[0] # Call data dictionary
1416
- call_dict_base["snvs"] = call_res[1] # Called useful SNVs
1375
+ locus_result["snvs"] = call_res[1] # Called useful SNVs
1417
1376
 
1418
1377
  single_or_dist_assign: bool = assign_method in ("single", "dist")
1419
1378
 
@@ -1477,16 +1436,22 @@ def call_locus(
1477
1436
  # Hack: add small value to stdevs if we are exactly sure to make the below code work okay
1478
1437
  stdevs += 0.00001
1479
1438
 
1480
- sd_dist = np.abs((peaks - cn) / stdevs)
1481
- weighted_dist = np.abs(((peaks - cn) / stdevs) * weights)
1439
+ # np.abs((peaks - cn) / stdevs) is copy number distance from each peak, in # standard deviations
1440
+ peak: int
1441
+ if call_modal_n == 2 and np.all(np.abs((peaks - cn) / stdevs) < roughly_equiv_stdev_dist):
1442
+ # Hack: if both peaks are 1 stdev away, pretend we aren't sure and fill in whichever allele has less
1443
+ peak = int(len(allele_reads[0]) > len(allele_reads[1]))
1444
+ else:
1445
+ # Create an already-fitted Gaussian mixture instance and use it to predict the peak
1482
1446
 
1483
- # Hack: if both peaks are 1 stdev away, pretend we aren't sure and fill in whichever allele has less
1484
- peak: int = (
1485
- # bool to int conversion: 1 if we add to allele_reads[1]
1486
- int(len(allele_reads[0]) > len(allele_reads[1]))
1487
- if call_modal_n == 2 and np.all(sd_dist < roughly_equiv_stdev_dist)
1488
- else np.argmin(weighted_dist).item()
1489
- )
1447
+ g: GaussianMixture = GaussianMixture(n_components=call_modal_n, covariance_type="spherical")
1448
+ g.means_ = peaks.reshape(-1, 1)
1449
+ g.covariances_ = stdevs ** 2
1450
+ g.weights_ = weights
1451
+ # see https://github.com/scikit-learn/scikit-learn/blob/1.5.0/sklearn/mixture/_gaussian_mixture.py#L347
1452
+ g.precisions_cholesky_ = 1.0 / stdevs
1453
+
1454
+ peak = g.predict([[cn]])[0].item()
1490
1455
 
1491
1456
  allele_reads[peak].append(r)
1492
1457
  rd["p"] = peak
@@ -1515,16 +1480,16 @@ def call_locus(
1515
1480
  lambda a: consensus_seq(
1516
1481
  map(lambda rr: read_dict_extra[rr]["_tr_seq"], a),
1517
1482
  logger_,
1518
- max_poa_length=max_poa_length,
1483
+ max_mdn_poa_length=max_mdn_poa_length,
1519
1484
  ),
1520
1485
  allele_reads,
1521
1486
  )
1522
1487
  )
1523
1488
 
1524
1489
  peak_data = {
1525
- "means": call_peaks.tolist(), # from np.ndarray
1526
- "weights": call_weights.tolist(), # from np.ndarray
1527
- "stdevs": call_stdevs.tolist(), # from np.ndarray
1490
+ "means": call_peaks,
1491
+ "weights": call_weights,
1492
+ "stdevs": call_stdevs,
1528
1493
  "modal_n": call_modal_n,
1529
1494
  "n_reads": call_peak_n_reads,
1530
1495
  **({"kmers": list(map(dict, peak_kmers))} if count_kmers in ("peak", "both") else {}),
@@ -1552,16 +1517,13 @@ def call_locus(
1552
1517
  f"{locus_log_str} - got call: {call_val} (95% CIs: {call_95_cis_val}); peak assign method={assign_method}; "
1553
1518
  f"# reads={call_peak_n_reads}")
1554
1519
 
1555
- return {
1556
- **call_dict_base,
1557
- "assign_method": assign_method,
1558
- "call": call_val,
1559
- "call_95_cis": call_95_cis_val,
1560
- "call_99_cis": apply_or_none(_nested_ndarray_serialize, call_99_cis),
1561
- "peaks": peak_data,
1562
- "ps": call_ps,
1563
- # make typecheck happy above by checking all of these are not None (even though if call is false-y, all of them
1564
- # should be None and otherwise none of them should).
1565
- "read_peaks_called": read_peaks_called,
1566
- "time": call_time,
1567
- }
1520
+ locus_result["assign_method"] = assign_method
1521
+ locus_result["call"] = call_val
1522
+ locus_result["call_95_cis"] = call_95_cis_val
1523
+ locus_result["call_99_cis"] = apply_or_none(_nested_ndarray_serialize, call_99_cis)
1524
+ locus_result["peaks"] = peak_data
1525
+ locus_result["ps"] = call_ps
1526
+ locus_result["read_peaks_called"] = read_peaks_called
1527
+ locus_result["time"] = call_time
1528
+
1529
+ return locus_result