ppgtk 0.1.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 (41) hide show
  1. ppgtk-0.1.0/LICENSE +21 -0
  2. ppgtk-0.1.0/MANIFEST.in +2 -0
  3. ppgtk-0.1.0/PKG-INFO +108 -0
  4. ppgtk-0.1.0/README.md +80 -0
  5. ppgtk-0.1.0/ppgtk/__init__.py +6 -0
  6. ppgtk-0.1.0/ppgtk/calculate_frequencies/__init__.py +14 -0
  7. ppgtk-0.1.0/ppgtk/calculate_frequencies/calculate_frequencies.py +376 -0
  8. ppgtk-0.1.0/ppgtk/calculate_frequencies/impute.py +471 -0
  9. ppgtk-0.1.0/ppgtk/classify_ploidy/__init__.py +9 -0
  10. ppgtk-0.1.0/ppgtk/classify_ploidy/logistic_regression.py +891 -0
  11. ppgtk-0.1.0/ppgtk/conversion/__init__.py +0 -0
  12. ppgtk-0.1.0/ppgtk/conversion/structure.py +90 -0
  13. ppgtk-0.1.0/ppgtk/diversity_statistics/__init__.py +5 -0
  14. ppgtk-0.1.0/ppgtk/diversity_statistics/sfs.py +114 -0
  15. ppgtk-0.1.0/ppgtk/diversity_statistics/theta.py +85 -0
  16. ppgtk-0.1.0/ppgtk/fit_mixtures/__init__.py +6 -0
  17. ppgtk-0.1.0/ppgtk/fit_mixtures/fit_mixtures.py +81 -0
  18. ppgtk-0.1.0/ppgtk/fit_mixtures/gmm.py +124 -0
  19. ppgtk-0.1.0/ppgtk/fit_mixtures/gmm2.py +118 -0
  20. ppgtk-0.1.0/ppgtk/fit_mixtures/gmm_fixed_means.py +964 -0
  21. ppgtk-0.1.0/ppgtk/fit_mixtures/gmm_fixed_means_fixed_weights.py +971 -0
  22. ppgtk-0.1.0/ppgtk/fit_mixtures/lmm.py +91 -0
  23. ppgtk-0.1.0/ppgtk/fit_mixtures/plot_mixtures.py +160 -0
  24. ppgtk-0.1.0/ppgtk/math.py +11 -0
  25. ppgtk-0.1.0/ppgtk/sampling/__init__.py +7 -0
  26. ppgtk-0.1.0/ppgtk/sampling/sampling.py +228 -0
  27. ppgtk-0.1.0/ppgtk/utils.py +116 -0
  28. ppgtk-0.1.0/ppgtk/windowing/__init__.py +7 -0
  29. ppgtk-0.1.0/ppgtk/windowing/windowing.py +95 -0
  30. ppgtk-0.1.0/ppgtk.egg-info/PKG-INFO +108 -0
  31. ppgtk-0.1.0/ppgtk.egg-info/SOURCES.txt +39 -0
  32. ppgtk-0.1.0/ppgtk.egg-info/dependency_links.txt +1 -0
  33. ppgtk-0.1.0/ppgtk.egg-info/entry_points.txt +2 -0
  34. ppgtk-0.1.0/ppgtk.egg-info/requires.txt +8 -0
  35. ppgtk-0.1.0/ppgtk.egg-info/top_level.txt +2 -0
  36. ppgtk-0.1.0/ppgtk_cli.py +733 -0
  37. ppgtk-0.1.0/setup.cfg +4 -0
  38. ppgtk-0.1.0/setup.py +32 -0
  39. ppgtk-0.1.0/test/test_cli_estimate_theta_windowing.py +224 -0
  40. ppgtk-0.1.0/test/test_cli_individual_genotypes.py +52 -0
  41. ppgtk-0.1.0/test/test_cli_pop_freqs.py +67 -0
ppgtk-0.1.0/LICENSE ADDED
@@ -0,0 +1,21 @@
1
+ MIT License
2
+
3
+ Copyright (c) 2024 George Tiley
4
+
5
+ Permission is hereby granted, free of charge, to any person obtaining a copy
6
+ of this software and associated documentation files (the "Software"), to deal
7
+ in the Software without restriction, including without limitation the rights
8
+ to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
9
+ copies of the Software, and to permit persons to whom the Software is
10
+ furnished to do so, subject to the following conditions:
11
+
12
+ The above copyright notice and this permission notice shall be included in all
13
+ copies or substantial portions of the Software.
14
+
15
+ THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
16
+ IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
17
+ FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
18
+ AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
19
+ LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
20
+ OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
21
+ SOFTWARE.
@@ -0,0 +1,2 @@
1
+ include LICENSE
2
+ include README.md
ppgtk-0.1.0/PKG-INFO ADDED
@@ -0,0 +1,108 @@
1
+ Metadata-Version: 2.4
2
+ Name: ppgtk
3
+ Version: 0.1.0
4
+ Summary: A package for polyploid popoulation genomics analyses and data exploration
5
+ Home-page: https://github.com/gtiley/ppgtk
6
+ Author: George P. Tiley
7
+ License: MIT
8
+ Requires-Python: >=3.11
9
+ Description-Content-Type: text/markdown
10
+ License-File: LICENSE
11
+ Requires-Dist: cyvcf2>=0.31.1
12
+ Requires-Dist: pandas>=2.2.2
13
+ Requires-Dist: scikit-learn<1.8,>=1.5.1
14
+ Requires-Dist: click>=8.1.7
15
+ Requires-Dist: matplotlib>=3.10.5
16
+ Requires-Dist: scipy>=1.15.2
17
+ Requires-Dist: seaborn>=0.13.2
18
+ Requires-Dist: statsmodels>=0.14.5
19
+ Dynamic: author
20
+ Dynamic: description
21
+ Dynamic: description-content-type
22
+ Dynamic: home-page
23
+ Dynamic: license
24
+ Dynamic: license-file
25
+ Dynamic: requires-dist
26
+ Dynamic: requires-python
27
+ Dynamic: summary
28
+
29
+ # Polyploid Population Genomics Tool Kit (ppgtk)
30
+ [![CI](https://github.com/tileylab/ppgtk/actions/workflows/ci-install.yml/badge.svg?branch=main)](https://github.com/tileylab/ppgtk/actions/workflows/ci-install.yml)
31
+
32
+ Python package for polyploid population genomics analyses and data exploration
33
+
34
+ ## Disclaimer
35
+ This package is in early developmental stages and a lot of functionality and specific commands are still a work in progress. The *classify-ploidy* method does work and has been vetted against several test datasets. Publication of the methods are anticipated in the near future but made available in case it is helpful for some exploratory analyses.
36
+
37
+ ## Installation
38
+
39
+ ### Conda Install
40
+ ```python
41
+ conda env create -f environment.yml
42
+ conda activate ppgtk
43
+ ```
44
+ The *environment.yml* file was created from the ppgtk development environment with `conda env export --from-history > environment.yml`.
45
+
46
+ ### Pip Install
47
+ A pip installation can be done system-wide or within a new conda/venv environment. The recommended install path is to install from package metadata so dependency resolution stays platform-aware. This install path is tested on Python 3.11 and 3.12 across Ubuntu and macOS. The Windows installation is currently not working due to a limitation of a dependency for reading VCF files, but we are working towards a native parser to allow windows compatibility. If Windows is necessary for you, please submit an issue and it will happen faster.
48
+
49
+ ```python
50
+ python -m pip install --upgrade pip
51
+ python -m pip install --upgrade setuptools wheel
52
+ python -m pip install .
53
+ ```
54
+
55
+ If you need a fully pinned developer environment, `requirements.txt` is still available, but pinned transitive dependencies can be less portable across operating systems.
56
+
57
+ ## Ploidy Classification
58
+
59
+ The ploidy classification method can be implemented as a single command-line program after successful installation. To bring up the program options, you can look at the help menu:
60
+
61
+ ```bash
62
+ ppgtk classify-ploidy --help
63
+ ```
64
+
65
+ And an example analysis would look like this:
66
+ ```bash
67
+ ppgtk classify-ploidy -v input.vcf -o results metadata.csv
68
+ ```
69
+
70
+ Note that the metadata file comes at the end without a flag. This is by design. All commands in **ppgtk** require a metadata file, which has a minimal format. Only three columns are needed:
71
+ ```
72
+ individual, population, ploidy
73
+ sample_1, pop_A, 2
74
+ sample_2, pop_A, 2
75
+ sample_3, pop_A, NA
76
+ ...
77
+ sample_n, pop_B, 4
78
+ ```
79
+ Missing values are accepted for ploidy. Missing values should only be encoded as "NA". Other numbers that are not an expected ploidy will be treated as a seperate class in the regression model.
80
+
81
+ The population field is used when calculating population genetic statistics, but it does not matter for the *classify-ploidy* method. If you do not have *a priori* or data-driven population assignments, that is fine. You could simply duplicate the individual names or assign all individuals to a single placeholder population, such as "pop_A".
82
+
83
+ ## Classifier accuracy metrics
84
+
85
+ The ploidy classifier (logistic regression) prints and writes several performance artifacts when run via the `logistic_regression` entrypoint. These are produced by the `cross_validate_logistic_regression` and `evaluate_logistic_regression_model` helpers and can be found in the output directory when provided.
86
+
87
+ - `logistic_regression_cv_accuracy.csv` — a single-row table with summary metrics:
88
+ - `accuracy`: overall fraction of correctly classified samples
89
+ - `balanced_accuracy`: per-class recall averaged equally across classes (useful for imbalanced label sets)
90
+ - `macro_f1`: unweighted mean of per-class F1 scores
91
+ - `weighted_f1`: mean of per-class F1 scores weighted by support (class prevalence)
92
+
93
+ - `logistic_regression_cv_report.csv` — the full sklearn classification report (precision / recall / f1 / support) for each class. Use this to inspect per-ploidy performance.
94
+
95
+ - `logistic_regression_cv_confusion_matrix.csv` — confusion matrix (rows = true labels, columns = predicted labels) showing raw classification counts.
96
+
97
+ - `logistic_regression_predictions.csv` — per-sample results including:
98
+ - `known_ploidy` (when available), `predicted_ploidy`
99
+ - probability columns named `prob_<class_label>` for class-wise predicted probabilities
100
+
101
+ Interpreting the metrics:
102
+ - Use `balanced_accuracy` and `macro_f1` when classes are imbalanced to avoid dominance by the largest class.
103
+ - `weighted_f1` is helpful when you want a single F1 that reflects the dataset's class proportions.
104
+ - The confusion matrix and per-class precision/recall help identify which ploidy levels are being confused by the model.
105
+
106
+ <!--
107
+ When adding or removing features (for example the new `mean_ab` / `median_ab` summaries), re-run cross-validation and compare these artifacts to track improvements or regressions.
108
+ -->
ppgtk-0.1.0/README.md ADDED
@@ -0,0 +1,80 @@
1
+ # Polyploid Population Genomics Tool Kit (ppgtk)
2
+ [![CI](https://github.com/tileylab/ppgtk/actions/workflows/ci-install.yml/badge.svg?branch=main)](https://github.com/tileylab/ppgtk/actions/workflows/ci-install.yml)
3
+
4
+ Python package for polyploid population genomics analyses and data exploration
5
+
6
+ ## Disclaimer
7
+ This package is in early developmental stages and a lot of functionality and specific commands are still a work in progress. The *classify-ploidy* method does work and has been vetted against several test datasets. Publication of the methods are anticipated in the near future but made available in case it is helpful for some exploratory analyses.
8
+
9
+ ## Installation
10
+
11
+ ### Conda Install
12
+ ```python
13
+ conda env create -f environment.yml
14
+ conda activate ppgtk
15
+ ```
16
+ The *environment.yml* file was created from the ppgtk development environment with `conda env export --from-history > environment.yml`.
17
+
18
+ ### Pip Install
19
+ A pip installation can be done system-wide or within a new conda/venv environment. The recommended install path is to install from package metadata so dependency resolution stays platform-aware. This install path is tested on Python 3.11 and 3.12 across Ubuntu and macOS. The Windows installation is currently not working due to a limitation of a dependency for reading VCF files, but we are working towards a native parser to allow windows compatibility. If Windows is necessary for you, please submit an issue and it will happen faster.
20
+
21
+ ```python
22
+ python -m pip install --upgrade pip
23
+ python -m pip install --upgrade setuptools wheel
24
+ python -m pip install .
25
+ ```
26
+
27
+ If you need a fully pinned developer environment, `requirements.txt` is still available, but pinned transitive dependencies can be less portable across operating systems.
28
+
29
+ ## Ploidy Classification
30
+
31
+ The ploidy classification method can be implemented as a single command-line program after successful installation. To bring up the program options, you can look at the help menu:
32
+
33
+ ```bash
34
+ ppgtk classify-ploidy --help
35
+ ```
36
+
37
+ And an example analysis would look like this:
38
+ ```bash
39
+ ppgtk classify-ploidy -v input.vcf -o results metadata.csv
40
+ ```
41
+
42
+ Note that the metadata file comes at the end without a flag. This is by design. All commands in **ppgtk** require a metadata file, which has a minimal format. Only three columns are needed:
43
+ ```
44
+ individual, population, ploidy
45
+ sample_1, pop_A, 2
46
+ sample_2, pop_A, 2
47
+ sample_3, pop_A, NA
48
+ ...
49
+ sample_n, pop_B, 4
50
+ ```
51
+ Missing values are accepted for ploidy. Missing values should only be encoded as "NA". Other numbers that are not an expected ploidy will be treated as a seperate class in the regression model.
52
+
53
+ The population field is used when calculating population genetic statistics, but it does not matter for the *classify-ploidy* method. If you do not have *a priori* or data-driven population assignments, that is fine. You could simply duplicate the individual names or assign all individuals to a single placeholder population, such as "pop_A".
54
+
55
+ ## Classifier accuracy metrics
56
+
57
+ The ploidy classifier (logistic regression) prints and writes several performance artifacts when run via the `logistic_regression` entrypoint. These are produced by the `cross_validate_logistic_regression` and `evaluate_logistic_regression_model` helpers and can be found in the output directory when provided.
58
+
59
+ - `logistic_regression_cv_accuracy.csv` — a single-row table with summary metrics:
60
+ - `accuracy`: overall fraction of correctly classified samples
61
+ - `balanced_accuracy`: per-class recall averaged equally across classes (useful for imbalanced label sets)
62
+ - `macro_f1`: unweighted mean of per-class F1 scores
63
+ - `weighted_f1`: mean of per-class F1 scores weighted by support (class prevalence)
64
+
65
+ - `logistic_regression_cv_report.csv` — the full sklearn classification report (precision / recall / f1 / support) for each class. Use this to inspect per-ploidy performance.
66
+
67
+ - `logistic_regression_cv_confusion_matrix.csv` — confusion matrix (rows = true labels, columns = predicted labels) showing raw classification counts.
68
+
69
+ - `logistic_regression_predictions.csv` — per-sample results including:
70
+ - `known_ploidy` (when available), `predicted_ploidy`
71
+ - probability columns named `prob_<class_label>` for class-wise predicted probabilities
72
+
73
+ Interpreting the metrics:
74
+ - Use `balanced_accuracy` and `macro_f1` when classes are imbalanced to avoid dominance by the largest class.
75
+ - `weighted_f1` is helpful when you want a single F1 that reflects the dataset's class proportions.
76
+ - The confusion matrix and per-class precision/recall help identify which ploidy levels are being confused by the model.
77
+
78
+ <!--
79
+ When adding or removing features (for example the new `mean_ab` / `median_ab` summaries), re-run cross-validation and compare these artifacts to track improvements or regressions.
80
+ -->
@@ -0,0 +1,6 @@
1
+ from ppgtk.calculate_frequencies.calculate_frequencies import get_ind_genotypes, get_ind_ab, get_pop_freqs
2
+ from ppgtk.calculate_frequencies.impute import average_missing, remove_missing
3
+
4
+ """
5
+ ppgtk takes a vcf as input to create a matrix of allele frequencies
6
+ """
@@ -0,0 +1,14 @@
1
+ from ppgtk.calculate_frequencies.calculate_frequencies import get_ind_genotypes
2
+ from ppgtk.calculate_frequencies.calculate_frequencies import get_ind_ab
3
+ from ppgtk.calculate_frequencies.calculate_frequencies import get_pop_freqs
4
+ from ppgtk.calculate_frequencies.impute import apply_missing_imputation
5
+ from ppgtk.calculate_frequencies.impute import average_missing
6
+ from ppgtk.calculate_frequencies.impute import randomly_impute_missing
7
+ from ppgtk.calculate_frequencies.impute import learn_impute_missing
8
+ from ppgtk.calculate_frequencies.impute import zero_impute_missing
9
+ from ppgtk.calculate_frequencies.impute import remove_missing
10
+ from ppgtk.calculate_frequencies.impute import remove_missing
11
+
12
+ """
13
+ Calculate allele balance for variants from a multisample vcf
14
+ """
@@ -0,0 +1,376 @@
1
+ import re
2
+ import numpy as np
3
+ import pandas as pd
4
+ import logging
5
+
6
+
7
+ def _safe_int(value, default=0):
8
+ try:
9
+ return int(value)
10
+ except (TypeError, ValueError):
11
+ return default
12
+
13
+
14
+ def _extract_sample_fields(format_string, sample_string):
15
+ """Extract GT/AD/GQ fields from a sample column using FORMAT keys."""
16
+ format_keys = format_string.split(':')
17
+ sample_values = sample_string.split(':')
18
+ field_map = {}
19
+ for idx, key in enumerate(format_keys):
20
+ if idx < len(sample_values):
21
+ field_map[key] = sample_values[idx]
22
+
23
+ gt_field = field_map.get('GT', sample_values[0] if len(sample_values) > 0 else '')
24
+ ad_field = field_map.get('AD', sample_values[1] if len(sample_values) > 1 else '')
25
+ gq_field = field_map.get('GQ', '')
26
+ return gt_field, ad_field, gq_field, ('GQ' in field_map)
27
+
28
+
29
+ def get_ind_genotypes(n_sites, n_tax, ind_map, vcf_file, min_depth, min_count, min_qual, pass_flag, output_dir, return_site_data=False):
30
+ '''
31
+ Returns per-individual genotype dosage and quality layers across sites from a multisample VCF.
32
+
33
+ Parameters:
34
+ n_sites (int): the number of sites to process from the VCF
35
+ n_tax (int): the number of individuals to process from the VCF
36
+ ind_map (dict): a dictionary mapping individuals in the VCF to a population or other identifier
37
+ vcf_file (string): a multisample vcf file uncompressed
38
+ min_depth (int): the minimum depth of a site to be considered high-quality
39
+ min_count (int): the minimum number of reads supporting the minor allele to be considered high-quality
40
+ min_qual (int): the minimum phred-scaled genotype likelihood to be considered high-quality
41
+ pass_flag (bool): Does the VCF have an PASS field or something else to condider like a "."
42
+ output_dir (string): the directory where all results will be written
43
+
44
+ Returns:
45
+ tax_list (list): A list of individual labels
46
+ genotype_dat: a numpy array with shape (4, n_sites, n_tax) where layers are genotype dosage,
47
+ depth, genotype quality, and pass-filter indicator
48
+ site_df (pd.DataFrame, optional): site-level coordinates and chromosome ids when return_site_data=True
49
+ '''
50
+
51
+ genotype_data = np.empty((n_sites, n_tax), dtype=np.int8)
52
+ site_depth_data = np.empty((n_sites, n_tax), dtype=np.uint16)
53
+ genotype_quality_data = np.empty((n_sites, n_tax), dtype=np.uint8)
54
+ passing_filter_data = np.empty((n_sites, n_tax), dtype=np.bool_)
55
+ chromosome_data = np.empty((n_sites,), dtype=np.uint16)
56
+ site_position_data = np.empty((n_sites,), dtype=np.uint32)
57
+ chromosome_label_data = np.empty((n_sites,), dtype=object)
58
+ ref_allele_data = np.empty((n_sites,), dtype='U1')
59
+ alt_allele_data = np.empty((n_sites,), dtype='U1')
60
+ vcf_map = {}
61
+ vcf_index = {}
62
+ tax_list = []
63
+ n_tax = 0
64
+ n_sites = 0
65
+ n_variants = {}
66
+ skip_header = 1
67
+
68
+ chromosome_map = {}
69
+ chromosome_count = 0
70
+
71
+ with open(vcf_file,'r') as fh:
72
+ for line in fh:
73
+ line = line.strip()
74
+ if '#CHROM' in line:
75
+ temp = line.split()
76
+ #print(line)
77
+ for i in range(9, len(temp)):
78
+ this_tax = ''
79
+ if '/' in temp[i]:
80
+ tax_path = temp[i].split('/')
81
+ this_tax = tax_path[-1]
82
+ else:
83
+ this_tax = temp[i]
84
+ #print(this_tax)
85
+ if this_tax in ind_map.keys():
86
+ tax_list.append(this_tax)
87
+ n_variants[this_tax] = 0
88
+ vcf_index[this_tax] = n_tax
89
+ n_tax = n_tax + 1
90
+ vcf_map[i] = this_tax
91
+ skip_header = 0
92
+ else:
93
+ if skip_header == 0:
94
+ temp = line.split()
95
+ if (temp[6] == pass_flag):
96
+ format_string = temp[8]
97
+ for i in range(9, len(temp)): ####Continue fixing here
98
+ if i in vcf_map.keys():
99
+ genotypes = []
100
+ genotype_sum = -1
101
+ ref_counts = 0
102
+ alt_counts = 0
103
+ total_count = 0
104
+ genotype_quality = 0
105
+ indicator = 0
106
+ genotype_string = temp[i]
107
+ gt_field, ad_field, gq_field, has_gq = _extract_sample_fields(format_string, genotype_string)
108
+ if re.match(r'\d+,\d+', ad_field):
109
+ allele_counts = ad_field.split(',')
110
+ ref_counts = _safe_int(allele_counts[0], default=0)
111
+ alt_counts = _safe_int(allele_counts[1], default=0)
112
+ total_count = ref_counts + alt_counts
113
+ if re.match(r'\d+\S+\d+', gt_field):
114
+ if '/' in gt_field:
115
+ genotypes = gt_field.split('/')
116
+ elif '|' in gt_field:
117
+ genotypes = gt_field.split('|')
118
+ genotype_sum = int(genotypes.count('1'))
119
+ if re.match(r'\d+', gq_field):
120
+ genotype_quality = min(_safe_int(gq_field, default=0), np.iinfo(np.uint8).max)
121
+ qual_ok = (genotype_quality >= min_qual) if has_gq else True
122
+ if ((total_count >= min_depth) and (ref_counts >= 1) and (alt_counts >= min_count) and qual_ok):
123
+ indicator = 1
124
+
125
+ genotype_data[n_sites, vcf_index[vcf_map[i]]] = genotype_sum
126
+ site_depth_data[n_sites, vcf_index[vcf_map[i]]] = total_count
127
+ genotype_quality_data[n_sites, vcf_index[vcf_map[i]]] = genotype_quality
128
+ passing_filter_data[n_sites, vcf_index[vcf_map[i]]] = indicator
129
+ n_variants[vcf_map[i]] = n_variants[vcf_map[i]] + 1
130
+ chromosome = temp[0]
131
+ position = temp[1]
132
+ ref_allele_data[n_sites] = temp[3]
133
+ alt_allele_data[n_sites] = temp[4].split(',')[0]
134
+ if chromosome not in chromosome_map.keys():
135
+ chromosome_map[chromosome] = chromosome_count
136
+ chromosome_count = chromosome_count + 1
137
+ chromosome_data[n_sites] = int(chromosome_map[chromosome])
138
+ elif chromosome in chromosome_map.keys():
139
+ chromosome_data[n_sites] = int(chromosome_map[chromosome])
140
+ chromosome_label_data[n_sites] = chromosome
141
+ site_position_data[n_sites] = int(position)
142
+ n_sites = n_sites + 1
143
+
144
+ if (output_dir != 'dummy'):
145
+ for i in range(0, len(tax_list)):
146
+ output_file = f'{output_dir}/{tax_list[i]}.txt'
147
+ outfile = open(output_file, 'w')
148
+ outfile.write('genotype\tdepth\tgenotype_quality\tpass_filters\n')
149
+ for j in range(0, n_variants[tax_list[i]]):
150
+ outstring = f'{genotype_data[j,vcf_index[tax_list[i]]]}\t{site_depth_data[j,vcf_index[tax_list[i]]]}\t{genotype_quality_data[j,vcf_index[tax_list[i]]]}\t{passing_filter_data[j,vcf_index[tax_list[i]]]}\n'
151
+ outfile.write(outstring)
152
+ outfile.close()
153
+
154
+ genotype_dat = np.array([
155
+ genotype_data,
156
+ site_depth_data,
157
+ genotype_quality_data,
158
+ passing_filter_data
159
+ ])
160
+
161
+ logging.info(f'Array shape: {genotype_dat.shape}')
162
+ logging.info(f'Memory usage: {genotype_dat.nbytes / 1024 / 1024:.2f} MB')
163
+ logging.info(f'Processed VCF of {n_sites} for {n_tax}\n')
164
+ if return_site_data:
165
+ site_df = pd.DataFrame(
166
+ {
167
+ 'site_index': np.arange(0, n_sites, dtype=np.int32),
168
+ 'chromosome': chromosome_label_data[:n_sites],
169
+ 'chromosome_id': chromosome_data[:n_sites],
170
+ 'position': site_position_data[:n_sites],
171
+ 'ref_allele': ref_allele_data[:n_sites],
172
+ 'alt_allele': alt_allele_data[:n_sites],
173
+ }
174
+ )
175
+ site_df['chromosome'] = site_df['chromosome'].astype('category')
176
+ return(tax_list, genotype_dat, site_df)
177
+ return(tax_list, genotype_dat)
178
+
179
+ def get_ind_ab(n_sites, n_tax, ind_map, vcf_file, min_depth, min_count, min_qual, pass_flag, output_dir, return_site_data=False):
180
+ '''
181
+ Returns per-individual allele-balance and quality layers across sites from a multisample VCF.
182
+
183
+ Parameters:
184
+ n_sites (int): the number of sites to process from the VCF
185
+ n_tax (int): the number of individuals to process from the VCF
186
+ ind_map (dict): a dictionary mapping individuals in the VCF to a population or other identifier
187
+ vcf_file (string): a multisample vcf file uncompressed
188
+ min_depth (int): the minimum depth of a site to be considered high-quality
189
+ min_count (int): the minimum number of reads supporting the minor allele to be considered high-quality
190
+ min_qual (int): the minimum phred-scaled genotype likelihood to be considered high-quality
191
+ pass_flag (bool): Does the VCF have an PASS field or something else to condider like a "."
192
+ output_dir (string): the directory where all results will be written
193
+
194
+ Returns:
195
+ tax_list (list): A list of individual labels
196
+ ab_dat: a numpy array with shape (4, n_sites, n_tax) where layers are allele balance,
197
+ depth, genotype quality, and pass-filter indicator
198
+ site_df (pd.DataFrame, optional): site-level coordinates and chromosome ids when return_site_data=True
199
+ '''
200
+
201
+ # Goal - these all need to be typed as arrays to keep the memory from exploding
202
+ # The individual files can be written out using pandas from array
203
+ allele_balance_data = np.empty((n_sites, n_tax), dtype=np.float32)
204
+ site_depth_data = np.empty((n_sites, n_tax), dtype=np.uint16)
205
+ genotype_quality_data = np.empty((n_sites, n_tax), dtype=np.uint8)
206
+ passing_filter_data = np.empty((n_sites, n_tax), dtype=np.bool_)
207
+ chromosome_data = np.empty((n_sites,), dtype=np.uint16)
208
+ site_position_data = np.empty((n_sites,), dtype=np.uint32)
209
+ chromosome_label_data = np.empty((n_sites,), dtype=object)
210
+ vcf_map = {}
211
+ vcf_index = {}
212
+ tax_list = []
213
+ n_tax = 0
214
+ n_sites = 0
215
+ n_variants = {}
216
+ skip_header = 1
217
+
218
+ chromosome_map = {}
219
+ chromosome_count = 0
220
+
221
+ with open(vcf_file,'r') as fh:
222
+ for line in fh:
223
+ line = line.strip()
224
+ if '#CHROM' in line:
225
+ temp = line.split()
226
+ #print(line)
227
+ for i in range(9, len(temp)):
228
+ this_tax = ''
229
+ if '/' in temp[i]:
230
+ tax_path = temp[i].split('/')
231
+ this_tax = tax_path[-1]
232
+ else:
233
+ this_tax = temp[i]
234
+ #print(this_tax)
235
+ if this_tax in ind_map.keys():
236
+ tax_list.append(this_tax)
237
+ n_variants[this_tax] = 0
238
+ vcf_index[this_tax] = n_tax
239
+ n_tax = n_tax + 1
240
+ vcf_map[i] = this_tax
241
+ skip_header = 0
242
+ else:
243
+ if skip_header == 0:
244
+ temp = line.split()
245
+ if (temp[6] == pass_flag):
246
+ format_string = temp[8]
247
+ for i in range(9, len(temp)): ####Continue fixing here
248
+ if i in vcf_map.keys():
249
+ ref_counts = 0
250
+ alt_counts = 0
251
+ total_count = 0
252
+ total_count = 0
253
+ allele_balance = 0
254
+ genotype_quality = 0
255
+ indicator = 0
256
+ genotype_string = temp[i]
257
+ _, ad_field, gq_field, has_gq = _extract_sample_fields(format_string, genotype_string)
258
+ if re.match(r'\d+,\d+', ad_field):
259
+ allele_counts = ad_field.split(',')
260
+ ref_counts = _safe_int(allele_counts[0], default=0)
261
+ alt_counts = _safe_int(allele_counts[1], default=0)
262
+ total_count = ref_counts + alt_counts
263
+ if re.match(r'\d+', gq_field):
264
+ genotype_quality = min(_safe_int(gq_field, default=0), np.iinfo(np.uint8).max)
265
+ if (total_count > 0):
266
+ allele_balance = alt_counts / total_count
267
+ qual_ok = (genotype_quality >= min_qual) if has_gq else True
268
+ if ((total_count >= min_depth) and (ref_counts >= 1) and (alt_counts >= min_count) and qual_ok):
269
+ indicator = 1
270
+ else:
271
+ print(f'WARNING: Incorrectly formatted VCF fields!\n--> {vcf_map[i]} at variant {n_variants[vcf_map[i]]}\n-->{temp[0]}: {temp[1]}\n')
272
+
273
+ allele_balance_data[n_sites, vcf_index[vcf_map[i]]] = allele_balance
274
+ site_depth_data[n_sites, vcf_index[vcf_map[i]]] = total_count
275
+ genotype_quality_data[n_sites, vcf_index[vcf_map[i]]] = genotype_quality
276
+ passing_filter_data[n_sites, vcf_index[vcf_map[i]]] = indicator
277
+ n_variants[vcf_map[i]] = n_variants[vcf_map[i]] + 1
278
+ chromosome = temp[0]
279
+ position = temp[1]
280
+ if chromosome not in chromosome_map.keys():
281
+ chromosome_map[chromosome] = chromosome_count
282
+ chromosome_count = chromosome_count + 1
283
+ chromosome_data[n_sites] = int(chromosome_map[chromosome])
284
+ elif chromosome in chromosome_map.keys():
285
+ chromosome_data[n_sites] = int(chromosome_map[chromosome])
286
+ chromosome_label_data[n_sites] = chromosome
287
+ site_position_data[n_sites] = int(position)
288
+ n_sites = n_sites + 1
289
+
290
+ if (output_dir != 'dummy'):
291
+ for i in range(0, len(tax_list)):
292
+ output_file = f'{output_dir}/{tax_list[i]}.txt'
293
+ outfile = open(output_file, 'w')
294
+ outfile.write('allele_balance\tdepth\tgenotype_quality\tpass_filters\n')
295
+ for j in range(0, n_variants[tax_list[i]]):
296
+ outstring = f'{allele_balance_data[j,vcf_index[tax_list[i]]]}\t{site_depth_data[j,vcf_index[tax_list[i]]]}\t{genotype_quality_data[j,vcf_index[tax_list[i]]]}\t{passing_filter_data[j,vcf_index[tax_list[i]]]}\n'
297
+ outfile.write(outstring)
298
+ outfile.close()
299
+
300
+ ab_dat = np.array([
301
+ allele_balance_data,
302
+ site_depth_data,
303
+ genotype_quality_data,
304
+ passing_filter_data
305
+ ])
306
+
307
+ #ab_df = pd.DataFrame(allele_balance_data)
308
+ #print(ab_df)
309
+ logging.info(f'Array shape: {ab_dat.shape}')
310
+ logging.info(f'Memory usage: {ab_dat.nbytes / 1024 / 1024:.2f} MB')
311
+ logging.info(f'Processed VCF of {n_sites} for {n_tax}\n')
312
+ if return_site_data:
313
+ site_df = pd.DataFrame(
314
+ {
315
+ 'site_index': np.arange(0, n_sites, dtype=np.int32),
316
+ 'chromosome': chromosome_label_data[:n_sites],
317
+ 'chromosome_id': chromosome_data[:n_sites],
318
+ 'position': site_position_data[:n_sites],
319
+ }
320
+ )
321
+ site_df['chromosome'] = site_df['chromosome'].astype('category')
322
+ return(tax_list, ab_dat, site_df)
323
+ return(tax_list, ab_dat)
324
+
325
+
326
+ def get_pop_freqs(genotype_dat, tax_list, ind_map, site_df):
327
+ '''
328
+ Calculate population-level allele frequencies from individual genotype dosages.
329
+
330
+ Parameters:
331
+ genotype_dat (np.ndarray): shape (4, n_sites, n_taxa), layer 0 is dosage, layer 3 is pass-filter
332
+ tax_list (list): individual IDs aligned to genotype_dat axis 2
333
+ ind_map (dict): individual metadata with 'population' and 'ploidy' keys
334
+ site_df (pd.DataFrame): site-level coordinates with 'chromosome' and 'position' columns
335
+
336
+ Returns:
337
+ freq_df (pd.DataFrame): wide matrix with populations as rows and locus labels as columns
338
+ '''
339
+ from ppgtk.utils import assign_populations
340
+
341
+ tax_index = {tax: idx for idx, tax in enumerate(tax_list)}
342
+ populations = assign_populations(ind_map)
343
+ sorted_pops = sorted(populations.keys())
344
+
345
+ n_sites = genotype_dat.shape[1]
346
+ dosage_layer = genotype_dat[0]
347
+ pass_layer = genotype_dat[3]
348
+
349
+ locus_labels = [
350
+ f"{site_df.iloc[i]['chromosome']}_{site_df.iloc[i]['position']}"
351
+ for i in range(n_sites)
352
+ ]
353
+
354
+ freq_data = {}
355
+ for pop in sorted_pops:
356
+ members = populations[pop]
357
+ member_indices = np.array([tax_index[tax] for tax in members if tax in tax_index])
358
+ ploidies = np.array([int(ind_map[tax]['ploidy']) for tax in members if tax in tax_index])
359
+
360
+ pop_dosages = dosage_layer[:, member_indices].astype(np.float64)
361
+ pop_passing = pass_layer[:, member_indices].astype(np.float64)
362
+
363
+ # Only count individuals that pass filters at each site
364
+ valid_dosages = np.where(pop_passing > 0, pop_dosages, 0.0)
365
+ valid_ploidies = pop_passing * ploidies[np.newaxis, :]
366
+
367
+ sum_dosages = valid_dosages.sum(axis=1)
368
+ sum_ploidies = valid_ploidies.sum(axis=1)
369
+
370
+ with np.errstate(invalid='ignore'):
371
+ freqs = np.where(sum_ploidies > 0, sum_dosages / sum_ploidies, np.nan)
372
+ freq_data[pop] = freqs
373
+
374
+ freq_df = pd.DataFrame(freq_data, index=locus_labels).T
375
+ freq_df.index.name = 'population'
376
+ return freq_df