Single-Nucleus Transcriptomic Mapping Reveals Correlative Microglial Changes Associated with rTMS in the Motor Cortex After Spinal Cord Injury.

1. #!/usr/bin/env python
2. #
3. # Copyright (c) 2018 10x Genomics, Inc. All rights reserved.
4. #
5. """Library functions for performing batch correction."""
8. from __future__ import annotations
10. import struct
11. from collections import Counter
13. import numpy as np
14. import sklearn.neighbors as sk_neighbors
15. from sklearn.metrics.pairwise import rbf_kernel
17. DEFAULT_BALLTREE_LEAFSIZE = 40
20. def batch_effect_score(
21. dimred_matrix: np.ndarray[tuple[int, int], np.dtype[np.float64]],
22. batch_ids: np.ndarray[int, np.dtype[np.bytes_]],
23. knn_neighbors: int | None = None,
24. knn_frac: float | None = 0.01,
25. max_num_bcs: int | None = 10000,
26. ):
27. """Compute batch effect score on an aggregated dimension-reduced matrix.
29. The batch effect score quantifies the degree of separation between groups of barcodes (batches).
30. For each barcode, the fraction of its k-nearest neighbors with the same batch id is calculated.
31. Local batch scores are computed by scaling and shifting the same-batch fractions to be strictly
32. less than the number of batches and have expectation equal to 1 under identical batches.
33. Finally, the average of these local batch scores is computed to obtain the overall batch effect
34. score. A batch score of 1 indicates no separation between batches, and a batch score equal to
35. the number of batches indicates perfect separation. A batch score less than 1 may occur in rare
36. cases and is consistent with no batch separation.
38. Args:
39. dimred_matrix (np.ndarray[tuple[int, int], np.dtype[np.float64]]): dimension reduced matrix.
40. Rows are samples and columns are features (e.g. principal components)
41. batch_ids (np.ndarray[int, np.dtype[np.string_]]): 1-D array of batch IDs. Must have at
42. least 2 observations per batch
43. knn_frac (Optional[float], optional): Sets k in the k-nearest neighbors calculation to a
44. fraction of the total barcodes (after subsampling to max_num_bcs). Defaults to 0.01.
45. max_num_bcs (Optional[int], optional): Maximum number of barcodes to use. Larger data sets
46. will be subsampled. Defaults to 10000.
47. knn_neighbors (Optional[int], optional): Use a fixed k rather than knn_frac
48. """
49. if knn_neighbors is None and knn_frac is None:
50. raise ValueError("One of knn_neighbors or knn_frac must be specified")
52. num_bcs = dimred_matrix.shape[0]
53. if num_bcs != len(batch_ids):
54. raise ValueError("Length of batch_ids must equal number of rows in dimred_matrix")
56. batch_counts_orig = Counter(batch_ids)
58. # subsample barcodes if greater than the specified max_num_bcs
59. if max_num_bcs is not None and num_bcs > max_num_bcs:
60. np.random.seed(0)
61. select_bc_idx = np.random.choice(num_bcs, max_num_bcs)
62. select_bc_idx.sort()
63. dimred_matrix = dimred_matrix[select_bc_idx]
64. batch_ids = batch_ids[select_bc_idx]
65. num_bcs = dimred_matrix.shape[0]
67. # Return NaN if we completely dropped any batch or if not enough data
68. # NOTE: this can happen due to subsampling but only for severely imbalanced cases,
69. # e.g. 10,000x more cells in batch 1 vs 2)
70. batch_counts = Counter(batch_ids)
71. if len(batch_counts) != len(batch_counts_orig) or min(batch_counts.values()) < 2:
72. return np.nan
74. if knn_neighbors is not None:
75. num_neighbors = knn_neighbors
76. else:
77. num_neighbors = int(np.ceil(knn_frac * num_bcs))
79. # For a cell in a given batch, what fraction of other cells share the same batch?
80. # This is the expectation of same_batch_frac below, given the null of identical batch mixing
81. batch_counts = Counter(batch_ids)
82. num_batches = len(batch_counts)
83. batch_to_frac = {batch: (count - 1) / (num_bcs - 1) for batch, count in batch_counts.items()}
84. null_same_batch_frac = np.fromiter((batch_to_frac[i] for i in batch_ids), dtype=np.float64)
85. # What is largest that same_batch_frac can be? Correction only relevant for very small batches
86. # (e.g. 10k total cells, 100 or fewer cells in a single batch)
87. batch_to_max_frac = {
88. batch: min(count - 1, num_neighbors) / num_neighbors
89. for batch, count in batch_counts.items()
90. }
91. max_same_batch_frac = np.fromiter((batch_to_max_frac[i] for i in batch_ids), dtype=np.float64)
93. balltree = sk_neighbors.BallTree(dimred_matrix, leaf_size=DEFAULT_BALLTREE_LEAFSIZE)
94. knn_idx = balltree.query(dimred_matrix, k=num_neighbors + 1, return_distance=False)
96. same_batch_frac = np.mean(batch_ids[:, None] == batch_ids[knn_idx[:, 1:]], axis=1)
97. # for identical batches, local_batch_score is equal to null_same_batch_frac in expectation
98. # we rescale to be between 1 and num_batches
99. local_batch_score = 1 + (num_batches - 1) * (same_batch_frac - null_same_batch_frac) / (
100. max_same_batch_frac - null_same_batch_frac
101. )
103. return np.mean(local_batch_score)
106. def find_knn(curr_matrix, ref_matrix, knn):
107. """For each row in curr_matrix, find k nearest neighbors in ref_matrix,.
109. return an array of shape=[curr_matrix.shape[0] * knn, ], which stores
110. the index of nearest neighbors in ref_matrix
111. """
112. balltree = sk_neighbors.BallTree(ref_matrix, leaf_size=DEFAULT_BALLTREE_LEAFSIZE)
113. num_neighbors = min(ref_matrix.shape[0], knn)
114. nn_idx = balltree.query(curr_matrix, k=num_neighbors, return_distance=False)
115. return nn_idx.ravel().astype(int)
118. def serialize_batch_nearest_neighbor(fp, batch_nearest_neighbor):
119. for (a, b), s in batch_nearest_neighbor.items():
120. fp.write(struct.pack("qqQ", a, b, len(s)))
121. for i, j in s:
122. fp.write(struct.pack("qq", i, j))
125. def deserialize_batch_nearest_neighbor(fp):
126. """>>> from cStringIO import StringIO.
128. >>> batch1 = dict()
129. >>> batch1[(0, 1)] = set([(1, 2), (3, 4), (5, 6)])
130. >>> batch1[(1, 2)] = set([(7, 8), (9, 10)])
131. >>> batch1[(3, 4)] = set([(11, 12)])
132. >>> fp = StringIO()
133. >>> serialize_batch_nearest_neighbor(fp, batch1)
134. >>> fp.seek(0)
135. >>> batch2 = deserialize_batch_nearest_neighbor(fp)
136. >>> batch1 == batch2
137. True
138. """
139. batch_nearest_neighbor = {}
140. while True:
141. fmt = "qqQ"
142. sz = struct.calcsize("qqQ")
143. buf = fp.read(sz)
144. if len(buf) == 0:
145. break
146. elif len(buf) != sz:
147. raise RuntimeError("corrupted batch_nearest_neighbor stream (key)")
148. a, b, slen = struct.unpack(fmt, buf)
149. fmt = "qq"
150. sz = struct.calcsize("qq")
151. s = set()
152. for _ in range(slen):
153. buf = fp.read(sz)
154. if len(buf) != sz:
155. raise RuntimeError("corrupted batch_nearest_neighbor stream (set)")
156. i, j = struct.unpack(fmt, buf)
157. s.add((i, j))
158. batch_nearest_neighbor[(a, b)] = s
159. return batch_nearest_neighbor
162. def correction_vector(dimred_matrix, cur_submatrix_idx, mnn_cur_idx, mnn_ref_idx, sigma):
163. """Compute the batch-correction vector.
165. 1. For each MNN pair in current dataset and the reference, a pair-specific
166. batch-correction vector is computed as the vector difference between the
167. paired cells.
168. 2. For each barcode in cur dataset, a batch-correction vector is calculated
169. as a weighted average of these pair-specific vectors, as computed with a
170. Gaussian kernel.
171. """
172. num_pcs = dimred_matrix.shape[1]
173. corr_vector = np.zeros((0, num_pcs))
175. # the number of mnn and submatrix dim might be very large, process by chunk to save memory
176. cur_submatrix_size = len(cur_submatrix_idx)
177. mnn_size = len(mnn_cur_idx)
178. # based on empirical testing
179. cur_submatrix_chunk_size = int(1e6 / num_pcs)
180. mnn_chunk_size = int(2e7 / num_pcs)
182. for i in range(0, cur_submatrix_size, cur_submatrix_chunk_size):
183. cur_submatrix_chunk = cur_submatrix_idx[i : i + cur_submatrix_chunk_size]
184. cur_submatrix = dimred_matrix[cur_submatrix_chunk]
186. weighted_sum, weights_sum = np.zeros(cur_submatrix.shape), np.zeros(cur_submatrix.shape)
188. for j in range(0, mnn_size, mnn_chunk_size):
189. mnn_cur_chunk = mnn_cur_idx[j : j + mnn_chunk_size]
190. mnn_ref_chunk = mnn_ref_idx[j : j + mnn_chunk_size]
192. mnn_cur = dimred_matrix[mnn_cur_chunk]
193. weights = rbf_kernel(cur_submatrix, mnn_cur, gamma=0.5 * sigma)
194. bias = dimred_matrix[mnn_ref_chunk] - mnn_cur
195. weighted_sum += np.dot(weights, bias)
196. weights_sum += np.tile(np.sum(weights, axis=1), (num_pcs, 1)).T
198. corr_vector = np.vstack((corr_vector, weighted_sum / weights_sum))
200. return corr_vector

batch_correction.py at commit 669395e, under MIT · at the source