-
Notifications
You must be signed in to change notification settings - Fork 23
/
utils.py
319 lines (269 loc) · 13.4 KB
/
utils.py
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
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
import anndata
try:
from anndata.base import Raw
except ImportError:
from anndata import Raw
import batchglm.api as glm
import numpy as np
import pandas as pd
import patsy
import scipy.sparse
from typing import List, Tuple, Union
# Relay util functions for diffxpy api.
# design_matrix, preview_coef_names and constraint_system_from_star are redefined here.
from batchglm.data import constraint_matrix_from_string, constraint_matrix_from_dict
from batchglm.data import view_coef_names
def parse_gene_names(
data: Union[anndata.AnnData, Raw, np.ndarray, scipy.sparse.csr_matrix, glm.typing.InputDataBase],
gene_names: Union[list, np.ndarray, None]
):
if gene_names is None:
if anndata is not None and (isinstance(data, anndata.AnnData) or isinstance(data, Raw)):
gene_names = data.var_names
elif isinstance(data, glm.typing.InputDataBase):
gene_names = data.features
else:
raise ValueError("Missing gene names")
return np.asarray(gene_names)
def parse_sample_description(
data: Union[anndata.AnnData, Raw, np.ndarray, scipy.sparse.csr_matrix, glm.typing.InputDataBase],
sample_description: Union[pd.DataFrame, None]
) -> pd.DataFrame:
"""
Parse sample description from input.
:param data: Input data matrix (observations x features) or (cells x genes).
:param sample_description: pandas.DataFrame containing sample annotations, can be None.
:return: Assembled sample annotations.
"""
if sample_description is None:
if anndata is not None and isinstance(data, anndata.AnnData):
sample_description = data.obs
else:
raise ValueError(
"Please specify `sample_description` or provide `data` as anndata.AnnData " +
"with corresponding sample annotations"
)
if sample_description is not None:
if anndata is not None and isinstance(data, Raw):
# Raw does not have attribute shape.
assert data.X.shape[0] == sample_description.shape[0], \
"data matrix and sample description must contain same number of cells: %i, %i" % \
(data.X.shape[0], sample_description.shape[0])
elif isinstance(data, glm.typing.InputDataBase):
assert data.x.shape[0] == sample_description.shape[0], \
"data matrix and sample description must contain same number of cells: %i, %i" % \
(data.x.shape[0], sample_description.shape[0])
else:
assert data.shape[0] == sample_description.shape[0], \
"data matrix and sample description must contain same number of cells: %i, %i" % \
(data.shape[0], sample_description.shape[0])
return sample_description
def parse_size_factors(
size_factors: Union[np.ndarray, pd.core.series.Series, np.ndarray],
data: Union[anndata.AnnData, Raw, np.ndarray, scipy.sparse.csr_matrix, glm.typing.InputDataBase],
sample_description: pd.DataFrame
) -> Union[np.ndarray, None]:
"""
Parse size-factors from input.
:param size_factors: 1D array of transformed library size factors for each cell in the
same order as in data or string-type column identifier of size-factor containing
column in sample description.
:param data: Input data matrix (observations x features) or (cells x genes).
:param sample_description: optional pandas.DataFrame containing sample annotations
:return: Assebled size-factors.
"""
if size_factors is not None:
if isinstance(size_factors, pd.core.series.Series):
size_factors = size_factors.values
elif isinstance(size_factors, str):
assert size_factors in sample_description.columns, ""
size_factors = sample_description[size_factors].values
if anndata is not None and isinstance(data, Raw):
data_shape = data.X.shape
elif isinstance(data, glm.typing.InputDataBase):
data_shape = data.x.shape
else:
data_shape = data.shape
assert size_factors.shape[0] == data_shape[0], "data matrix and size factors must contain same number of cells"
assert np.all(size_factors > 0), "size_factors <= 0 found, please remove these cells"
return size_factors
def parse_grouping(data, sample_description, grouping):
if isinstance(grouping, str):
sample_description = parse_sample_description(data, sample_description)
grouping = sample_description[grouping]
return np.squeeze(np.asarray(grouping))
def split_x(data, grouping):
grouping = np.asarray(grouping)
groups = np.unique(grouping)
x0 = data[np.where(grouping == groups[0])[0]]
x1 = data[np.where(grouping == groups[1])[0]]
return x0, x1
def dmat_unique(dmat, sample_description):
dmat, idx = np.unique(dmat, axis=0, return_index=True)
sample_description = sample_description.iloc[idx].reset_index(drop=True)
return dmat, sample_description
def design_matrix(
data: Union[anndata.AnnData, Raw, np.ndarray, scipy.sparse.csr_matrix] = None,
sample_description: Union[None, pd.DataFrame] = None,
formula: Union[None, str] = None,
as_numeric: Union[List[str], Tuple[str], str] = (),
dmat: Union[pd.DataFrame, None] = None,
return_type: str = "patsy"
) -> Union[patsy.design_info.DesignMatrix, pd.DataFrame]:
""" Create a design matrix from some sample description.
This function defaults to perform formatting if dmat is directly supplied as a pd.DataFrame.
This function relays batchglm.data.design_matrix() to behave like the other wrappers in diffxpy.
:param data: Input data matrix (observations x features) or (cells x genes).
:param sample_description: pandas.DataFrame of length "num_observations" containing explanatory variables as columns
:param formula: model formula as string, describing the relations of the explanatory variables.
E.g. '~ 1 + batch + confounder'
:param as_numeric:
Which columns of sample_description to treat as numeric and
not as categorical. This yields columns in the design matrix
which do not correpond to one-hot encoded discrete factors.
This makes sense for number of genes, time, pseudotime or space
for example.
:param dmat: a model design matrix as a pd.DataFrame
:param return_type: type of the returned value.
- "patsy": return plain patsy.design_info.DesignMatrix object
- "dataframe": return pd.DataFrame with observations as rows and params as columns
:param dmat: model design matrix
"""
if data is None and sample_description is None and dmat is None:
raise ValueError("supply either data or sample_description or dmat")
if dmat is None and formula is None:
raise ValueError("supply either dmat or formula")
if dmat is None:
sample_description = parse_sample_description(data, sample_description)
if sample_description is not None:
as_categorical = [False if x in as_numeric else True for x in sample_description.columns.values]
else:
as_categorical = True
return glm.data.design_matrix(
sample_description=sample_description,
formula=formula,
as_categorical=as_categorical,
dmat=dmat,
return_type=return_type
)
def preview_coef_names(
sample_description: pd.DataFrame,
formula: str,
as_numeric: Union[List[str], Tuple[str], str] = ()
) -> np.ndarray:
"""
Return coefficient names of model.
Use this to preview what the model would look like.
This function relays batchglm.data.preview_coef_names() to behave like the other wrappers in diffxpy.
:param sample_description: pandas.DataFrame of length "num_observations" containing explanatory variables as columns
:param formula: model formula as string, describing the relations of the explanatory variables.
E.g. '~ 1 + batch + confounder'
:param as_numeric:
Which columns of sample_description to treat as numeric and
not as categorical. This yields columns in the design matrix
which do not correpond to one-hot encoded discrete factors.
This makes sense for number of genes, time, pseudotime or space
for example.
:return: A list of coefficient names.
"""
if isinstance(as_numeric, str):
as_numeric = [as_numeric]
if isinstance(as_numeric, tuple):
as_numeric = list(as_numeric)
return glm.data.preview_coef_names(
sample_description=sample_description,
formula=formula,
as_categorical=[False if x in as_numeric else True for x in sample_description.columns.values]
)
def constraint_system_from_star(
dmat: Union[None, patsy.design_info.DesignMatrix] = None,
sample_description: Union[None, pd.DataFrame] = None,
formula: Union[None, str] = None,
as_numeric: Union[List[str], Tuple[str], str] = (),
constraints: dict = {},
return_type: str = "patsy",
) -> Tuple:
"""
Create a design matrix and a constraint matrix.
This function relays batchglm.data.constraint_matrix_from_star() to behave like the other wrappers in diffxpy.
:param dmat: Pre-built model design matrix.
:param sample_description: pandas.DataFrame of length "num_observations" containing explanatory variables as columns
:param formula: model formula as string, describing the relations of the explanatory variables.
E.g. '~ 1 + batch + confounder'
:param as_numeric:
Which columns of sample_description to treat as numeric and
not as categorical. This yields columns in the design matrix
which do not correspond to one-hot encoded discrete factors.
:param constraints: Grouped factors to enfore equality constraints on. Every element of
the dictionary corresponds to one set of equality constraints. Each set has to be
be an entry of the form {..., x: y, ...} where x is the factor to be constrained and y is
a factor by which levels of x are grouped and then constrained. Set y="1" to constrain
all levels of x to sum to one, a single equality constraint.
E.g.: {"batch": "condition"} Batch levels within each condition are constrained to sum to
zero. This is applicable if repeats of a an experiment within each condition
are independent so that the set-up ~1+condition+batch is perfectly confounded.
Can only group by non-constrained effects right now, use constraint_matrix_from_string
for other cases.
:param return_type: type of the returned value.
- "patsy": return plain patsy.design_info.DesignMatrix object
- "dataframe": return pd.DataFrame with observations as rows and params as columns
This option is overridden if constraints are supplied as dict.
:return: a model design matrix and a constraint matrix
"""
if isinstance(as_numeric, str):
as_numeric = [as_numeric]
if isinstance(as_numeric, tuple):
as_numeric = list(as_numeric)
if sample_description is not None:
as_categorical = [False if x in as_numeric else True for x in sample_description.columns.values]
else:
as_categorical = True
return glm.data.constraint_system_from_star(
dmat=dmat,
sample_description=sample_description,
formula=formula,
as_categorical=as_categorical,
constraints=constraints,
return_type=return_type
)
def bin_continuous_covariate(
factor_to_bin: str,
bins: Union[int, list, np.ndarray, Tuple],
data: Union[None, anndata.AnnData] = None,
sample_description: Union[None, pd.DataFrame] = None
):
r"""
Bin a continuous covariate.
Adds the binned covariate to the table. If data is supplied, the covariate is added in place in data.obs, otherwise
the covariate is added in the sample_description and the new sample_description is returned.
Binning is performed on quantiles of the distribution.
:param factor_to_bin: Name of columns of factor to bin.
:param bins: Number of bins or iteratable with bin borders. If given as integer, the bins are defined on the
quantiles of the covariate, ie the bottom 20% of observations are in the first bin if bins==5.
:param data: Anndata object that contains sample description table in .obs.
:param sample_description: Sample description table.
:return: Sample description table with binned covariate added if sample_description was supplied, otherwise None is
returned as the new column was added in place.
"""
if data is None and sample_description is not None:
sd = sample_description
elif data is not None and sample_description is None:
sd = data.obs
else:
raise ValueError("supply either data or sample_description")
if isinstance(bins, list) or isinstance(bins, np.ndarray) or isinstance(bins, Tuple):
bins = np.asarray(bins)
else:
bins = np.arange(0, 1, 1 / bins)
fac_binned = glm.data.bin_continuous_covariate(
sample_description=sd,
factor_to_bin=factor_to_bin,
bins=bins
)
if data is None and sample_description is not None:
sd[factor_to_bin + "_binned"] = fac_binned
return sample_description
elif data is not None and sample_description is None:
data.obs[factor_to_bin + "_binned"] = fac_binned
else:
assert False