dragonnfruit 0.3.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.
- dragonnfruit-0.3.0/LICENSE +21 -0
- dragonnfruit-0.3.0/PKG-INFO +9 -0
- dragonnfruit-0.3.0/README.md +2 -0
- dragonnfruit-0.3.0/cmd/dragonnfruit +99 -0
- dragonnfruit-0.3.0/dragonnfruit/__init__.py +4 -0
- dragonnfruit-0.3.0/dragonnfruit/attributions.py +66 -0
- dragonnfruit-0.3.0/dragonnfruit/io.py +267 -0
- dragonnfruit-0.3.0/dragonnfruit/models.py +605 -0
- dragonnfruit-0.3.0/dragonnfruit/preprocessing.py +384 -0
- dragonnfruit-0.3.0/dragonnfruit.egg-info/PKG-INFO +9 -0
- dragonnfruit-0.3.0/dragonnfruit.egg-info/SOURCES.txt +14 -0
- dragonnfruit-0.3.0/dragonnfruit.egg-info/dependency_links.txt +1 -0
- dragonnfruit-0.3.0/dragonnfruit.egg-info/requires.txt +5 -0
- dragonnfruit-0.3.0/dragonnfruit.egg-info/top_level.txt +1 -0
- dragonnfruit-0.3.0/setup.cfg +4 -0
- dragonnfruit-0.3.0/setup.py +20 -0
|
@@ -0,0 +1,21 @@
|
|
|
1
|
+
MIT License
|
|
2
|
+
|
|
3
|
+
Copyright (c) 2022 Jacob Schreiber
|
|
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,9 @@
|
|
|
1
|
+
Metadata-Version: 2.1
|
|
2
|
+
Name: dragonnfruit
|
|
3
|
+
Version: 0.3.0
|
|
4
|
+
Summary: dragonnfruit is a method for analyzing scATAC-seq experiments.
|
|
5
|
+
Home-page: https://github.com/jmschrei/dragonnfruit
|
|
6
|
+
Author: Jacob Schreiber
|
|
7
|
+
Author-email: jmschreiber91@gmail.com
|
|
8
|
+
License: LICENSE.txt
|
|
9
|
+
License-File: LICENSE
|
|
@@ -0,0 +1,99 @@
|
|
|
1
|
+
#!/usr/bin/env python
|
|
2
|
+
# dragonnfruit command-line tool
|
|
3
|
+
# Author: Jacob Schreiber <jmschreiber91@gmail.com>
|
|
4
|
+
|
|
5
|
+
import sys
|
|
6
|
+
import numpy
|
|
7
|
+
import torch
|
|
8
|
+
import pyfaidx
|
|
9
|
+
import argparse
|
|
10
|
+
|
|
11
|
+
import deepdish as dd
|
|
12
|
+
|
|
13
|
+
from dragonnfruit.preprocessing import fragments_to_sparse
|
|
14
|
+
from dragonnfruit.preprocessing import preprocess_atac
|
|
15
|
+
|
|
16
|
+
import json
|
|
17
|
+
|
|
18
|
+
desc = """dragonnfruit is an extension of ChromBPNet for single-cell data.
|
|
19
|
+
At a high level, it takes in nucleotide sequence and some aspect of cell
|
|
20
|
+
state, such as an LSI matrix from scATAC-seq data or gene expression data
|
|
21
|
+
from multimodal data, and makes predictions for dynamically pseudobulked
|
|
22
|
+
scATAC-seq signal. It works by running the cell state representation through
|
|
23
|
+
a small multi-layer perceptron and outputting the convolution parameters
|
|
24
|
+
in the accessibility component of a ChromBPNet model. Put another way:
|
|
25
|
+
after training a dragonnfruit model, one can produce a ChromBPNet model
|
|
26
|
+
for each cell in the experiment."""
|
|
27
|
+
|
|
28
|
+
# Read in the arguments
|
|
29
|
+
parser = argparse.ArgumentParser(description=desc)
|
|
30
|
+
subparsers = parser.add_subparsers(help="Must be either 'train', 'predict', 'interpret', or 'marginalize'.", required=True, dest='cmd')
|
|
31
|
+
|
|
32
|
+
preprocess_help = """Preprocess a set of fragment files, and optionally
|
|
33
|
+
inclusion/exclusion matrices, into a set of sparse matrices used in training
|
|
34
|
+
the dragonnfruit model."""
|
|
35
|
+
preprocess_parser = subparsers.add_parser("preprocess", help=preprocess_help)
|
|
36
|
+
preprocess_parser.add_argument("-p", "--parameters", type=str, required=True,
|
|
37
|
+
help="A JSON file containing the parameters for preprocessing the data.")
|
|
38
|
+
|
|
39
|
+
|
|
40
|
+
# Pull the arguments
|
|
41
|
+
args = parser.parse_args()
|
|
42
|
+
|
|
43
|
+
if args.cmd == "preprocess":
|
|
44
|
+
with open(args.parameters, "r") as infile:
|
|
45
|
+
parameters = json.load(infile)
|
|
46
|
+
|
|
47
|
+
default_parameters = {
|
|
48
|
+
'fragments': None,
|
|
49
|
+
'loci': None,
|
|
50
|
+
'h5_filename':'dragonnfruit_data.h5',
|
|
51
|
+
'rd_filename':'dragonnfruit_read_depths.npz',
|
|
52
|
+
'chrom_sizes': None,
|
|
53
|
+
'include_cells': None,
|
|
54
|
+
'exclude_cells': None,
|
|
55
|
+
'chroms': ['chr1', 'chr2', 'chr3', 'chr4', 'chr5', 'chr6', 'chr7',
|
|
56
|
+
'chr8', 'chr9', 'chr10', 'chr11', 'chr12', 'chr13', 'chr14',
|
|
57
|
+
'chr15', 'chr16', 'chr17', 'chr18', 'chr19', 'chr20', 'chr21',
|
|
58
|
+
'chr22', 'chrX'],
|
|
59
|
+
'cell_name_prefixes': None,
|
|
60
|
+
'max_fragment_length': 1000,
|
|
61
|
+
'start_offset': 4,
|
|
62
|
+
'end_offset': -5,
|
|
63
|
+
'verbose': False
|
|
64
|
+
}
|
|
65
|
+
|
|
66
|
+
for parameter, value in default_parameters.items():
|
|
67
|
+
if parameter not in parameters:
|
|
68
|
+
if value is None and parameter in ('fragments', 'chrom_sizes'):
|
|
69
|
+
raise ValueError("Must provide value for '{}'".format(parameter))
|
|
70
|
+
|
|
71
|
+
parameters[parameter] = value
|
|
72
|
+
|
|
73
|
+
###
|
|
74
|
+
|
|
75
|
+
'''
|
|
76
|
+
X_cscs, read_depths = fragments_to_sparse(
|
|
77
|
+
fragments=parameters['fragments'],
|
|
78
|
+
chrom_sizes=parameters['chrom_sizes'],
|
|
79
|
+
chroms=parameters['chroms'],
|
|
80
|
+
include_cells=parameters['include_cells'],
|
|
81
|
+
exclude_cells=parameters['exclude_cells'],
|
|
82
|
+
cell_name_prefixes=parameters['cell_name_prefixes'],
|
|
83
|
+
max_fragment_length=parameters['max_fragment_length'],
|
|
84
|
+
start_offset=parameters['start_offset'],
|
|
85
|
+
end_offset=parameters['end_offset'],
|
|
86
|
+
verbose=parameters['verbose']
|
|
87
|
+
)
|
|
88
|
+
|
|
89
|
+
dd.io.save(parameters['h5_filename'], X_cscs, compression='blosc')
|
|
90
|
+
numpy.savez_compressed(parameters['rd_filename'], read_depths)
|
|
91
|
+
'''
|
|
92
|
+
|
|
93
|
+
X_cscs = dd.io.load(parameters['h5_filename'])
|
|
94
|
+
|
|
95
|
+
X_pca, X_umap, neighbors = preprocess_atac(X_cscs,
|
|
96
|
+
peaks=parameters['loci'], chroms=parameters['chroms'])
|
|
97
|
+
|
|
98
|
+
numpy.savez_compressed("X_pca.npz", X_pca)
|
|
99
|
+
numpy.savez_compressed("neighbors.npz", neighbors)
|
|
@@ -0,0 +1,66 @@
|
|
|
1
|
+
|
|
2
|
+
# attributions.py
|
|
3
|
+
# Author: Jacob Schreiber <jmschreiber91@gmail.com>
|
|
4
|
+
|
|
5
|
+
import numpy
|
|
6
|
+
import torch
|
|
7
|
+
|
|
8
|
+
from tqdm import trange
|
|
9
|
+
from captum.attr import DeepLiftShap
|
|
10
|
+
|
|
11
|
+
from bpnetlite.attributions import dinucleotide_shuffle
|
|
12
|
+
from bpnetlite.attributions import hypothetical_attributions
|
|
13
|
+
|
|
14
|
+
class ProfileWrapper(torch.nn.Module):
|
|
15
|
+
"""A wrapper class that returns transformed profiles.
|
|
16
|
+
|
|
17
|
+
This class takes in a trained model and returns the weighted softmaxed
|
|
18
|
+
outputs of the first dimension. Specifically, it takes the predicted
|
|
19
|
+
"logits" and takes the dot product between them and the softmaxed versions
|
|
20
|
+
of those logits. This is for convenience when using captum to calculate
|
|
21
|
+
attribution scores.
|
|
22
|
+
|
|
23
|
+
Parameters
|
|
24
|
+
----------
|
|
25
|
+
model: torch.nn.Module
|
|
26
|
+
A torch model to be wrapped.
|
|
27
|
+
"""
|
|
28
|
+
|
|
29
|
+
def __init__(self, model):
|
|
30
|
+
super(ProfileWrapper, self).__init__()
|
|
31
|
+
self.model = model
|
|
32
|
+
|
|
33
|
+
def forward(self, X, cell_states):
|
|
34
|
+
logits = self.model(X, cell_states)
|
|
35
|
+
logits = logits - torch.mean(logits, dim=-1, keepdims=True)
|
|
36
|
+
l = torch.clone(logits).detach()
|
|
37
|
+
|
|
38
|
+
y = torch.exp(l - torch.logsumexp(l, dim=-1, keepdims=True))
|
|
39
|
+
return (logits * y).sum(axis=-1, keepdims=True)
|
|
40
|
+
|
|
41
|
+
|
|
42
|
+
def calculate_attributions(model, X, cell_states, n_shuffles=10, batch_size=1):
|
|
43
|
+
wrapper = ProfileWrapper(model.cuda())
|
|
44
|
+
ig = DeepLiftShap(wrapper)
|
|
45
|
+
|
|
46
|
+
reference = dinucleotide_shuffle(X, n_shuffles=n_shuffles).cuda()
|
|
47
|
+
reference = reference.type(X.dtype)
|
|
48
|
+
|
|
49
|
+
X = X.unsqueeze(0).cuda()
|
|
50
|
+
_cell_states = cell_states.cuda()
|
|
51
|
+
|
|
52
|
+
attributions = []
|
|
53
|
+
with torch.no_grad():
|
|
54
|
+
for start in trange(0, len(cell_states), batch_size):
|
|
55
|
+
c_ = _cell_states[start:start+batch_size]
|
|
56
|
+
X_ = X.expand(len(c_), -1, -1)
|
|
57
|
+
|
|
58
|
+
attr = ig.attribute(X_, reference, target=0,
|
|
59
|
+
additional_forward_args=(c_,),
|
|
60
|
+
custom_attribution_func=hypothetical_attributions)
|
|
61
|
+
|
|
62
|
+
attr = (attr * X_).cpu()
|
|
63
|
+
attributions.append(attr)
|
|
64
|
+
|
|
65
|
+
attributions = torch.cat(attributions)
|
|
66
|
+
return attributions
|
|
@@ -0,0 +1,267 @@
|
|
|
1
|
+
# io.py
|
|
2
|
+
# Author: Jacob Schreiber <jmschreiber91@gmail.com>
|
|
3
|
+
|
|
4
|
+
import time
|
|
5
|
+
import numpy
|
|
6
|
+
import scipy
|
|
7
|
+
import torch
|
|
8
|
+
import pandas
|
|
9
|
+
|
|
10
|
+
from bpnetlite.io import one_hot_encode
|
|
11
|
+
|
|
12
|
+
|
|
13
|
+
def _extract_example(self, chrom, mid, cell_idx, idx):
|
|
14
|
+
start, end = mid - self.window // 2, mid + self.window // 2
|
|
15
|
+
neighbs = self.neighbors[cell_idx]
|
|
16
|
+
|
|
17
|
+
X = self.sequence[chrom][start:end].T.astype('float32')
|
|
18
|
+
y = self.signal[chrom][:, start+self.trimming:end-self.trimming]
|
|
19
|
+
y = numpy.array(y[neighbs].sum(axis=0))[0]
|
|
20
|
+
|
|
21
|
+
c = self.cell_states[cell_idx]
|
|
22
|
+
r = self.read_depths[cell_idx]
|
|
23
|
+
|
|
24
|
+
if self.reverse_complement and idx % 2 == 0:
|
|
25
|
+
X = X[::-1][:, ::-1].copy()
|
|
26
|
+
y = y[::-1].copy()
|
|
27
|
+
|
|
28
|
+
X = torch.from_numpy(X)
|
|
29
|
+
y = torch.from_numpy(y)
|
|
30
|
+
c = torch.from_numpy(c)
|
|
31
|
+
r = torch.from_numpy(r)
|
|
32
|
+
return X, y, c, r
|
|
33
|
+
|
|
34
|
+
|
|
35
|
+
class LocusGenerator(torch.utils.data.Dataset):
|
|
36
|
+
"""A data generator for dragonnfruit inputs. Adapted from bpnet-lite.
|
|
37
|
+
|
|
38
|
+
This generator takes in a set of sequences and output signals
|
|
39
|
+
and will return a single element with random jitter and reverse-complement
|
|
40
|
+
augmentation applied. Because the data is single-cell where the output
|
|
41
|
+
signals differ across cells, each returned element is a random locus
|
|
42
|
+
in a random cell.
|
|
43
|
+
|
|
44
|
+
A conceptual difference between this DataGenerator and the one implemented
|
|
45
|
+
in bpnet-lite is that the bpnet-lite one assumes that you can extract all
|
|
46
|
+
loci into an array. Here, because there are hundreds of thousands of peaks
|
|
47
|
+
and potentially thousands of cells, it is actually more efficient to
|
|
48
|
+
|
|
49
|
+
Parameters
|
|
50
|
+
----------
|
|
51
|
+
sequences: dict of torch.tensors, shape=(n, 4), dtype=torch.float32
|
|
52
|
+
A dictionary of the nucleotide sequences to use.
|
|
53
|
+
|
|
54
|
+
signals: dict of torch.tensors, shape=(n, n_cells), dtype=torch.float32
|
|
55
|
+
A dictionary of the cell signals
|
|
56
|
+
|
|
57
|
+
loci: str or pandas.DataFrame
|
|
58
|
+
A set of loci to use.
|
|
59
|
+
|
|
60
|
+
cell_states:
|
|
61
|
+
"""
|
|
62
|
+
|
|
63
|
+
def __init__(self, sequence, signal, loci_file, neighbors, cell_states,
|
|
64
|
+
read_depths, trimming, window, chroms, n_cells_per_locus=1,
|
|
65
|
+
reverse_complement=True, max_jitter=128, random_state=None):
|
|
66
|
+
self.trimming = trimming
|
|
67
|
+
self.window = window
|
|
68
|
+
self.chroms = chroms
|
|
69
|
+
self.reverse_complement = reverse_complement
|
|
70
|
+
self.max_jitter = max_jitter
|
|
71
|
+
self.random_state = numpy.random.RandomState(random_state)
|
|
72
|
+
|
|
73
|
+
self.signal = {chrom: signal[chrom] for chrom in chroms}
|
|
74
|
+
self.sequence = {chrom: sequence[chrom] for chrom in chroms}
|
|
75
|
+
self.neighbors = neighbors
|
|
76
|
+
self.cell_states = cell_states
|
|
77
|
+
self.read_depths = read_depths
|
|
78
|
+
self.n_cells_per_locus = n_cells_per_locus
|
|
79
|
+
|
|
80
|
+
if not isinstance(loci_file, list):
|
|
81
|
+
loci_file = [loci_file]
|
|
82
|
+
|
|
83
|
+
loci = []
|
|
84
|
+
names = ['chrom', 'start', 'end']
|
|
85
|
+
for filename in loci_file:
|
|
86
|
+
loci_ = pandas.read_csv(filename, sep='\t', usecols=[0, 1, 2],
|
|
87
|
+
header=None, index_col=False, names=names)
|
|
88
|
+
loci_['mid'] = (loci_['end'] - loci_['start']) // 2 + loci_['start']
|
|
89
|
+
loci_ = loci_[numpy.isin(loci_['chrom'], chroms)]
|
|
90
|
+
loci_ = loci_.sample(frac=1, random_state=self.random_state)
|
|
91
|
+
loci_ = loci_.reset_index(drop=True)
|
|
92
|
+
loci.append(loci_)
|
|
93
|
+
|
|
94
|
+
self.loci = pandas.concat(loci).sort_index().reset_index(drop=True) # interleave
|
|
95
|
+
print("Loci: {}".format(self.loci.shape[0]))
|
|
96
|
+
|
|
97
|
+
def __len__(self):
|
|
98
|
+
return self.loci.shape[0]
|
|
99
|
+
|
|
100
|
+
def __getitem__(self, idx):
|
|
101
|
+
chrom, _, _, mid = self.loci.iloc[idx // self.n_cells_per_locus]
|
|
102
|
+
mid += self.random_state.randint(-self.max_jitter, self.max_jitter+1)
|
|
103
|
+
cell_idx = self.random_state.randint(self.cell_states.shape[0])
|
|
104
|
+
return _extract_example(self, chrom, mid, cell_idx, idx)
|
|
105
|
+
|
|
106
|
+
|
|
107
|
+
class GWGenerator(torch.utils.data.Dataset):
|
|
108
|
+
"""A data generator for dragonnfruit inputs. Adapted from bpnet-lite.
|
|
109
|
+
|
|
110
|
+
This generator takes in a set of sequences and output signals
|
|
111
|
+
and will return a single element with random jitter and reverse-complement
|
|
112
|
+
augmentation applied. Because the data is single-cell where the output
|
|
113
|
+
signals differ across cells, each returned element is a random locus
|
|
114
|
+
in a random cell.
|
|
115
|
+
|
|
116
|
+
A conceptual difference between this DataGenerator and the one implemented
|
|
117
|
+
in bpnet-lite is that the bpnet-lite one assumes that you can extract all
|
|
118
|
+
loci into an array. Here, because there are hundreds of thousands of peaks
|
|
119
|
+
and potentially thousands of cells, it is actually more efficient to
|
|
120
|
+
|
|
121
|
+
Parameters
|
|
122
|
+
----------
|
|
123
|
+
sequences: dict of torch.tensors, shape=(n, 4), dtype=torch.float32
|
|
124
|
+
A dictionary of the nucleotide sequences to use.
|
|
125
|
+
|
|
126
|
+
signals: dict of torch.tensors, shape=(n, n_cells), dtype=torch.float32
|
|
127
|
+
A dictionary of the cell signals
|
|
128
|
+
|
|
129
|
+
loci: str or pandas.DataFrame
|
|
130
|
+
A set of loci to use.
|
|
131
|
+
|
|
132
|
+
cell_states:
|
|
133
|
+
"""
|
|
134
|
+
|
|
135
|
+
def __init__(self, sequence, signal, neighbors, cell_states,
|
|
136
|
+
read_depths, trimming, window, chroms, cells_per_loci=1,
|
|
137
|
+
reverse_complement=True, random_state=None):
|
|
138
|
+
self.trimming = trimming
|
|
139
|
+
self.window = window
|
|
140
|
+
self.chroms = chroms
|
|
141
|
+
self.cells_per_loci = cells_per_loci
|
|
142
|
+
self.reverse_complement = reverse_complement
|
|
143
|
+
self.random_state = numpy.random.RandomState(random_state)
|
|
144
|
+
|
|
145
|
+
self.signal = {chrom: signal[chrom] for chrom in chroms}
|
|
146
|
+
self.sequence = {chrom: sequence[chrom] for chrom in chroms}
|
|
147
|
+
self.neighbors = neighbors
|
|
148
|
+
self.cell_states = cell_states
|
|
149
|
+
self.read_depths = read_depths
|
|
150
|
+
self._lengths = numpy.array([seq.shape[0] for seq in self.sequence.values()])
|
|
151
|
+
print(self._lengths)
|
|
152
|
+
|
|
153
|
+
def __len__(self):
|
|
154
|
+
return sum(self._lengths)
|
|
155
|
+
|
|
156
|
+
def __getitem__(self, idx):
|
|
157
|
+
c_idx = numpy.random.choice(len(self._lengths),
|
|
158
|
+
p=self._lengths / self._lengths.sum())
|
|
159
|
+
chrom = self.chroms[c_idx]
|
|
160
|
+
|
|
161
|
+
mid = numpy.random.randint(10000, self._lengths[c_idx]-10000)
|
|
162
|
+
cell_idx = numpy.random.randint(self.cell_states.shape[0])
|
|
163
|
+
return _extract_example(self, chrom, mid, cell_idx, idx)
|
|
164
|
+
|
|
165
|
+
|
|
166
|
+
class GWBGenerator(torch.utils.data.Dataset):
|
|
167
|
+
"""A data generator for dragonnfruit inputs. Adapted from bpnet-lite.
|
|
168
|
+
|
|
169
|
+
This generator takes in a set of sequences and output signals
|
|
170
|
+
and will return a single element with random jitter and reverse-complement
|
|
171
|
+
augmentation applied. Because the data is single-cell where the output
|
|
172
|
+
signals differ across cells, each returned element is a random locus
|
|
173
|
+
in a random cell.
|
|
174
|
+
|
|
175
|
+
A conceptual difference between this DataGenerator and the one implemented
|
|
176
|
+
in bpnet-lite is that the bpnet-lite one assumes that you can extract all
|
|
177
|
+
loci into an array. Here, because there are hundreds of thousands of peaks
|
|
178
|
+
and potentially thousands of cells, it is actually more efficient to
|
|
179
|
+
|
|
180
|
+
Parameters
|
|
181
|
+
----------
|
|
182
|
+
sequences: dict of torch.tensors, shape=(n, 4), dtype=torch.float32
|
|
183
|
+
A dictionary of the nucleotide sequences to use.
|
|
184
|
+
|
|
185
|
+
signals: dict of torch.tensors, shape=(n, n_cells), dtype=torch.float32
|
|
186
|
+
A dictionary of the cell signals
|
|
187
|
+
|
|
188
|
+
loci: str or pandas.DataFrame
|
|
189
|
+
A set of loci to use.
|
|
190
|
+
|
|
191
|
+
cell_states:
|
|
192
|
+
"""
|
|
193
|
+
|
|
194
|
+
def __init__(self, sequence, signal, neighbors, cell_states,
|
|
195
|
+
read_depths, trimming, window, chroms, batch_size=1024,
|
|
196
|
+
cells_per_loci=1, reverse_complement=True, random_state=None):
|
|
197
|
+
self.trimming = trimming
|
|
198
|
+
self.window = window
|
|
199
|
+
self.chroms = chroms
|
|
200
|
+
self.reverse_complement = reverse_complement
|
|
201
|
+
self.batch_size = batch_size
|
|
202
|
+
self.cells_per_loci = cells_per_loci
|
|
203
|
+
self.random_state = numpy.random.RandomState(random_state)
|
|
204
|
+
|
|
205
|
+
print(signal['chr21'].dtype, sequence['chr21'].dtype)
|
|
206
|
+
|
|
207
|
+
self.signal = {chrom: signal[chrom] for chrom in chroms}
|
|
208
|
+
self.sequence = {chrom: sequence[chrom] for chrom in chroms}
|
|
209
|
+
self.neighbors = neighbors
|
|
210
|
+
self.cell_states = cell_states
|
|
211
|
+
self.read_depths = read_depths
|
|
212
|
+
self._lengths = numpy.array([seq.shape[0] for seq in self.sequence.values()])
|
|
213
|
+
self._probs = self._lengths / self._lengths.sum()
|
|
214
|
+
|
|
215
|
+
print(self._lengths)
|
|
216
|
+
|
|
217
|
+
def __len__(self):
|
|
218
|
+
return sum(self._lengths)
|
|
219
|
+
|
|
220
|
+
def __getitem__(self, idx):
|
|
221
|
+
tic = time.time()
|
|
222
|
+
X = numpy.empty((self.batch_size, 4, self.window), dtype=numpy.int8)
|
|
223
|
+
y = numpy.empty((self.batch_size, self.window - self.trimming*2), dtype=numpy.float32)
|
|
224
|
+
c = numpy.empty((self.batch_size, self.cell_states.shape[1]), dtype=numpy.float32)
|
|
225
|
+
r = numpy.empty((self.batch_size, 1), dtype=numpy.float32)
|
|
226
|
+
|
|
227
|
+
cell_idxs = self.random_state.randint(self.cell_states.shape[0], size=self.batch_size)
|
|
228
|
+
chrom_idxs = self.random_state.choice(len(self._lengths), p=self._probs, size=self.batch_size)
|
|
229
|
+
|
|
230
|
+
nc = self.cells_per_loci
|
|
231
|
+
a = time.time() - tic
|
|
232
|
+
b, d, e = 0, 0, 0
|
|
233
|
+
|
|
234
|
+
for i, (cell_idx, chrom_idx) in enumerate(zip(cell_idxs, chrom_idxs)):
|
|
235
|
+
tic = time.time()
|
|
236
|
+
if i % nc == 0:
|
|
237
|
+
chrom = self.chroms[chrom_idx]
|
|
238
|
+
|
|
239
|
+
mid = self.random_state.randint(10000, self._lengths[chrom_idx]-10000)
|
|
240
|
+
start, end = mid - self.window // 2, mid + self.window // 2
|
|
241
|
+
|
|
242
|
+
lidx = (i // nc) * nc
|
|
243
|
+
X[lidx:lidx+nc] = self.sequence[chrom][start:end].T
|
|
244
|
+
y_ = self.signal[chrom][:, start+self.trimming:end-self.trimming].tocsr()
|
|
245
|
+
b += time.time() - tic
|
|
246
|
+
|
|
247
|
+
tic = time.time()
|
|
248
|
+
neighbs = self.neighbors[cell_idx]
|
|
249
|
+
y[i] = numpy.array(y_[neighbs].sum(axis=0))[0]
|
|
250
|
+
d += time.time() - tic
|
|
251
|
+
|
|
252
|
+
tic = time.time()
|
|
253
|
+
c[i] = self.cell_states[cell_idx]
|
|
254
|
+
r[i] = self.read_depths[cell_idx]
|
|
255
|
+
e += time.time() - tic
|
|
256
|
+
|
|
257
|
+
if self.reverse_complement and idx % 2 == 0:
|
|
258
|
+
X[i] = X[i][::-1][:, ::-1].copy()
|
|
259
|
+
y[i] = y[i][::-1].copy()
|
|
260
|
+
|
|
261
|
+
tic = time.time()
|
|
262
|
+
X = torch.from_numpy(X.astype('float32'))
|
|
263
|
+
y = torch.from_numpy(y)
|
|
264
|
+
c = torch.from_numpy(c)
|
|
265
|
+
r = torch.from_numpy(r)
|
|
266
|
+
#print(idx, a, b, d, e, time.time() - tic)
|
|
267
|
+
return X, y, c, r
|