M3Drop 0.4.55__tar.gz → 0.4.56__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.
- {m3drop-0.4.55 → m3drop-0.4.56/M3Drop.egg-info}/PKG-INFO +1 -1
- {m3drop-0.4.55/M3Drop.egg-info → m3drop-0.4.56}/PKG-INFO +1 -1
- m3drop-0.4.56/m3Drop/NormalizationCPU.py +323 -0
- {m3drop-0.4.55 → m3drop-0.4.56}/setup.py +1 -1
- m3drop-0.4.55/m3Drop/NormalizationCPU.py +0 -207
- {m3drop-0.4.55 → m3drop-0.4.56}/LICENSE +0 -0
- {m3drop-0.4.55 → m3drop-0.4.56}/M3Drop.egg-info/SOURCES.txt +0 -0
- {m3drop-0.4.55 → m3drop-0.4.56}/M3Drop.egg-info/dependency_links.txt +0 -0
- {m3drop-0.4.55 → m3drop-0.4.56}/M3Drop.egg-info/requires.txt +0 -0
- {m3drop-0.4.55 → m3drop-0.4.56}/M3Drop.egg-info/top_level.txt +0 -0
- {m3drop-0.4.55 → m3drop-0.4.56}/README.md +0 -0
- {m3drop-0.4.55 → m3drop-0.4.56}/m3Drop/ControlDeviceCPU.py +0 -0
- {m3drop-0.4.55 → m3drop-0.4.56}/m3Drop/ControlDeviceGPU.py +0 -0
- {m3drop-0.4.55 → m3drop-0.4.56}/m3Drop/CoreCPU.py +0 -0
- {m3drop-0.4.55 → m3drop-0.4.56}/m3Drop/CoreGPU.py +0 -0
- {m3drop-0.4.55 → m3drop-0.4.56}/m3Drop/DiagnosticsCPU.py +0 -0
- {m3drop-0.4.55 → m3drop-0.4.56}/m3Drop/DiagnosticsGPU.py +0 -0
- {m3drop-0.4.55 → m3drop-0.4.56}/m3Drop/NormalizationGPU.py +0 -0
- {m3drop-0.4.55 → m3drop-0.4.56}/m3Drop/__init__.py +0 -0
- {m3drop-0.4.55 → m3drop-0.4.56}/pyproject.toml +0 -0
- {m3drop-0.4.55 → m3drop-0.4.56}/setup.cfg +0 -0
|
@@ -0,0 +1,323 @@
|
|
|
1
|
+
import pickle
|
|
2
|
+
import time
|
|
3
|
+
import sys
|
|
4
|
+
import numpy as np
|
|
5
|
+
import h5py
|
|
6
|
+
import anndata
|
|
7
|
+
import pandas as pd
|
|
8
|
+
import os
|
|
9
|
+
import matplotlib.pyplot as plt
|
|
10
|
+
import seaborn as sns
|
|
11
|
+
from scipy import sparse
|
|
12
|
+
|
|
13
|
+
try:
|
|
14
|
+
from numba import jit, prange
|
|
15
|
+
except ImportError:
|
|
16
|
+
print("CRITICAL ERROR: 'numba' not found. Please install it (pip install numba).")
|
|
17
|
+
sys.exit(1)
|
|
18
|
+
|
|
19
|
+
# Strict Relative Import
|
|
20
|
+
from .ControlDeviceCPU import ControlDevice
|
|
21
|
+
|
|
22
|
+
# ==========================================
|
|
23
|
+
# NUMBA KERNELS (CPU)
|
|
24
|
+
# ==========================================
|
|
25
|
+
|
|
26
|
+
@jit(nopython=True, parallel=True, fastmath=True)
|
|
27
|
+
def pearson_residual_kernel_cpu(counts, tj, ti, theta, total, out_matrix):
|
|
28
|
+
rows = counts.shape[0]
|
|
29
|
+
cols = counts.shape[1]
|
|
30
|
+
for r in prange(rows):
|
|
31
|
+
ti_val = ti[r]
|
|
32
|
+
for c in range(cols):
|
|
33
|
+
count_val = counts[r, c]
|
|
34
|
+
mu = (tj[c] * ti_val) / total
|
|
35
|
+
theta_val = theta[c]
|
|
36
|
+
denom_sq = mu + ((mu * mu) / theta_val)
|
|
37
|
+
denom = np.sqrt(denom_sq)
|
|
38
|
+
if denom < 1e-12:
|
|
39
|
+
out_matrix[r, c] = 0.0
|
|
40
|
+
else:
|
|
41
|
+
out_matrix[r, c] = (count_val - mu) / denom
|
|
42
|
+
|
|
43
|
+
@jit(nopython=True, parallel=True, fastmath=True)
|
|
44
|
+
def pearson_approx_kernel_cpu(counts, tj, ti, total, out_matrix):
|
|
45
|
+
rows = counts.shape[0]
|
|
46
|
+
cols = counts.shape[1]
|
|
47
|
+
for r in prange(rows):
|
|
48
|
+
ti_val = ti[r]
|
|
49
|
+
for c in range(cols):
|
|
50
|
+
count_val = counts[r, c]
|
|
51
|
+
mu = (tj[c] * ti_val) / total
|
|
52
|
+
denom = np.sqrt(mu)
|
|
53
|
+
if denom < 1e-12:
|
|
54
|
+
out_matrix[r, c] = 0.0
|
|
55
|
+
else:
|
|
56
|
+
out_matrix[r, c] = (count_val - mu) / denom
|
|
57
|
+
|
|
58
|
+
# ==========================================
|
|
59
|
+
# NORMALIZATION FUNCTION
|
|
60
|
+
# ==========================================
|
|
61
|
+
|
|
62
|
+
def NBumiPearsonResidualsCombinedCPU(
|
|
63
|
+
raw_filename: str,
|
|
64
|
+
mask_filename: str,
|
|
65
|
+
fit_filename: str,
|
|
66
|
+
stats_filename: str,
|
|
67
|
+
output_filename_full: str,
|
|
68
|
+
output_filename_approx: str,
|
|
69
|
+
plot_summary_filename: str = None,
|
|
70
|
+
plot_detail_filename: str = None,
|
|
71
|
+
mode: str = "auto",
|
|
72
|
+
manual_target: int = 3000
|
|
73
|
+
):
|
|
74
|
+
"""
|
|
75
|
+
CPU-Optimized: Calculates Full and Approximate residuals in a SINGLE PASS.
|
|
76
|
+
Includes "Sidecar" Visualization logic (Streaming Stats + Subsampling).
|
|
77
|
+
"""
|
|
78
|
+
start_time = time.perf_counter()
|
|
79
|
+
print(f"FUNCTION: NBumiPearsonResidualsCombinedCPU() | FILE: {raw_filename}")
|
|
80
|
+
|
|
81
|
+
# 1. Load Mask
|
|
82
|
+
with open(mask_filename, 'rb') as f: mask = pickle.load(f)
|
|
83
|
+
ng_filtered = int(np.sum(mask))
|
|
84
|
+
|
|
85
|
+
# 2. Init Device
|
|
86
|
+
with h5py.File(raw_filename, 'r') as f: indptr_cpu = f['X']['indptr'][:]; total_rows = len(indptr_cpu) - 1
|
|
87
|
+
device = ControlDevice(indptr=indptr_cpu, total_rows=total_rows, n_genes=ng_filtered, mode=mode, manual_target=manual_target)
|
|
88
|
+
nc = device.total_rows
|
|
89
|
+
|
|
90
|
+
print("Phase [1/2]: Initializing parameters...")
|
|
91
|
+
with open(fit_filename, 'rb') as f: fit = pickle.load(f)
|
|
92
|
+
|
|
93
|
+
total = fit['vals']['total']
|
|
94
|
+
tjs = fit['vals']['tjs'].values.astype(np.float64)
|
|
95
|
+
tis = fit['vals']['tis'].values.astype(np.float64)
|
|
96
|
+
sizes = fit['sizes'].values.astype(np.float64)
|
|
97
|
+
|
|
98
|
+
# Setup Output Files
|
|
99
|
+
adata_in = anndata.read_h5ad(raw_filename, backed='r')
|
|
100
|
+
filtered_var = adata_in.var[mask]
|
|
101
|
+
|
|
102
|
+
adata_out_full = anndata.AnnData(obs=adata_in.obs, var=filtered_var)
|
|
103
|
+
adata_out_full.write_h5ad(output_filename_full, compression=None)
|
|
104
|
+
|
|
105
|
+
adata_out_approx = anndata.AnnData(obs=adata_in.obs, var=filtered_var)
|
|
106
|
+
adata_out_approx.write_h5ad(output_filename_approx, compression=None)
|
|
107
|
+
|
|
108
|
+
# --- VISUALIZATION SETUP (THE SIDECAR) ---
|
|
109
|
+
# 1. Sampling Rate (Strict Cap to prevent CPU RAM explosion)
|
|
110
|
+
TARGET_SAMPLES = 5_000_000
|
|
111
|
+
total_points = nc * ng_filtered
|
|
112
|
+
|
|
113
|
+
if total_points <= TARGET_SAMPLES:
|
|
114
|
+
sampling_rate = 1.0
|
|
115
|
+
else:
|
|
116
|
+
sampling_rate = TARGET_SAMPLES / total_points
|
|
117
|
+
|
|
118
|
+
print(f" > Visualization Sampling Rate: {sampling_rate*100:.4f}% (Target: {TARGET_SAMPLES:,} points)")
|
|
119
|
+
|
|
120
|
+
# 2. Accumulators (Numpy Arrays - Small memory footprint)
|
|
121
|
+
acc_raw_sum = np.zeros(ng_filtered, dtype=np.float64)
|
|
122
|
+
acc_approx_sum = np.zeros(ng_filtered, dtype=np.float64)
|
|
123
|
+
acc_approx_sq = np.zeros(ng_filtered, dtype=np.float64)
|
|
124
|
+
acc_full_sum = np.zeros(ng_filtered, dtype=np.float64)
|
|
125
|
+
acc_full_sq = np.zeros(ng_filtered, dtype=np.float64)
|
|
126
|
+
|
|
127
|
+
# 3. Lists for Plots (Sampled Only)
|
|
128
|
+
viz_approx_samples = []
|
|
129
|
+
viz_full_samples = []
|
|
130
|
+
# -----------------------------------------
|
|
131
|
+
|
|
132
|
+
storage_chunk_rows = int(1_000_000_000 / (ng_filtered * 8))
|
|
133
|
+
if storage_chunk_rows > nc: storage_chunk_rows = nc
|
|
134
|
+
if storage_chunk_rows < 1: storage_chunk_rows = 1
|
|
135
|
+
|
|
136
|
+
with h5py.File(output_filename_full, 'a') as f_full, h5py.File(output_filename_approx, 'a') as f_approx:
|
|
137
|
+
if 'X' in f_full: del f_full['X']
|
|
138
|
+
if 'X' in f_approx: del f_approx['X']
|
|
139
|
+
|
|
140
|
+
out_x_full = f_full.create_dataset('X', shape=(nc, ng_filtered), chunks=(storage_chunk_rows, ng_filtered), dtype='float64')
|
|
141
|
+
out_x_approx = f_approx.create_dataset('X', shape=(nc, ng_filtered), chunks=(storage_chunk_rows, ng_filtered), dtype='float64')
|
|
142
|
+
|
|
143
|
+
with h5py.File(raw_filename, 'r') as f_in:
|
|
144
|
+
h5_indptr = f_in['X']['indptr']
|
|
145
|
+
h5_data = f_in['X']['data']
|
|
146
|
+
h5_indices = f_in['X']['indices']
|
|
147
|
+
|
|
148
|
+
current_row = 0
|
|
149
|
+
while current_row < nc:
|
|
150
|
+
end_row = device.get_next_chunk(current_row, mode='dense', overhead_multiplier=3.0)
|
|
151
|
+
if end_row is None or end_row <= current_row: break
|
|
152
|
+
|
|
153
|
+
chunk_size = end_row - current_row
|
|
154
|
+
print(f"Phase [2/2]: Processing rows {end_row} of {nc} | Chunk: {chunk_size}", end='\r')
|
|
155
|
+
|
|
156
|
+
start_idx, end_idx = h5_indptr[current_row], h5_indptr[end_row]
|
|
157
|
+
|
|
158
|
+
data = np.array(h5_data[start_idx:end_idx], dtype=np.float64)
|
|
159
|
+
indices = np.array(h5_indices[start_idx:end_idx])
|
|
160
|
+
indptr = np.array(h5_indptr[current_row:end_row+1] - h5_indptr[current_row])
|
|
161
|
+
|
|
162
|
+
chunk_csr = sparse.csr_matrix((data, indices, indptr), shape=(chunk_size, len(mask)))
|
|
163
|
+
chunk_csr = chunk_csr[:, mask]
|
|
164
|
+
chunk_csr.data = np.ceil(chunk_csr.data)
|
|
165
|
+
|
|
166
|
+
# Numba needs dense
|
|
167
|
+
counts_dense = chunk_csr.toarray()
|
|
168
|
+
|
|
169
|
+
# --- VIZ ACCUMULATION 1: RAW MEAN ---
|
|
170
|
+
acc_raw_sum += np.sum(counts_dense, axis=0)
|
|
171
|
+
|
|
172
|
+
# --- VIZ SAMPLING: GENERATE INDICES ---
|
|
173
|
+
chunk_total_items = chunk_size * ng_filtered
|
|
174
|
+
n_samples_chunk = int(chunk_total_items * sampling_rate)
|
|
175
|
+
sample_indices = None
|
|
176
|
+
|
|
177
|
+
if n_samples_chunk > 0:
|
|
178
|
+
sample_indices = np.random.randint(0, int(chunk_total_items), size=n_samples_chunk)
|
|
179
|
+
|
|
180
|
+
# --- CALC 1: APPROX ---
|
|
181
|
+
approx_out = np.empty_like(counts_dense)
|
|
182
|
+
pearson_approx_kernel_cpu(
|
|
183
|
+
counts_dense,
|
|
184
|
+
tjs,
|
|
185
|
+
tis[current_row:end_row],
|
|
186
|
+
total,
|
|
187
|
+
approx_out
|
|
188
|
+
)
|
|
189
|
+
|
|
190
|
+
# Accumulate
|
|
191
|
+
acc_approx_sum += np.sum(approx_out, axis=0)
|
|
192
|
+
|
|
193
|
+
# Sample
|
|
194
|
+
if sample_indices is not None:
|
|
195
|
+
# Ravel creates a view, take copies the data. Safe.
|
|
196
|
+
viz_approx_samples.append(np.take(approx_out.ravel(), sample_indices))
|
|
197
|
+
|
|
198
|
+
# Write
|
|
199
|
+
out_x_approx[current_row:end_row, :] = approx_out
|
|
200
|
+
|
|
201
|
+
# Square (Explicit multiplication for safety)
|
|
202
|
+
approx_out = approx_out * approx_out
|
|
203
|
+
acc_approx_sq += np.sum(approx_out, axis=0)
|
|
204
|
+
del approx_out
|
|
205
|
+
|
|
206
|
+
# --- CALC 2: FULL (In-place on counts_dense) ---
|
|
207
|
+
pearson_residual_kernel_cpu(
|
|
208
|
+
counts_dense,
|
|
209
|
+
tjs,
|
|
210
|
+
tis[current_row:end_row],
|
|
211
|
+
sizes,
|
|
212
|
+
total,
|
|
213
|
+
counts_dense # Overwrite input
|
|
214
|
+
)
|
|
215
|
+
|
|
216
|
+
# Accumulate
|
|
217
|
+
acc_full_sum += np.sum(counts_dense, axis=0)
|
|
218
|
+
|
|
219
|
+
# Sample
|
|
220
|
+
if sample_indices is not None:
|
|
221
|
+
viz_full_samples.append(np.take(counts_dense.ravel(), sample_indices))
|
|
222
|
+
|
|
223
|
+
# Write
|
|
224
|
+
out_x_full[current_row:end_row, :] = counts_dense
|
|
225
|
+
|
|
226
|
+
# Square
|
|
227
|
+
counts_dense = counts_dense * counts_dense
|
|
228
|
+
acc_full_sq += np.sum(counts_dense, axis=0)
|
|
229
|
+
|
|
230
|
+
current_row = end_row
|
|
231
|
+
|
|
232
|
+
print(f"\nPhase [2/2]: COMPLETE{' '*50}")
|
|
233
|
+
|
|
234
|
+
# ==========================================
|
|
235
|
+
# VIZ GENERATION (POST-PROCESS)
|
|
236
|
+
# ==========================================
|
|
237
|
+
if plot_summary_filename and plot_detail_filename:
|
|
238
|
+
print("Phase [Viz]: Generating Diagnostics (CPU)...")
|
|
239
|
+
|
|
240
|
+
# 1. Finalize Variance Stats
|
|
241
|
+
mean_raw = acc_raw_sum / nc
|
|
242
|
+
|
|
243
|
+
mean_approx = acc_approx_sum / nc
|
|
244
|
+
mean_sq_approx = acc_approx_sq / nc
|
|
245
|
+
var_approx = mean_sq_approx - (mean_approx**2)
|
|
246
|
+
|
|
247
|
+
mean_full = acc_full_sum / nc
|
|
248
|
+
mean_sq_full = acc_full_sq / nc
|
|
249
|
+
var_full = mean_sq_full - (mean_full**2)
|
|
250
|
+
|
|
251
|
+
# 2. Finalize Samples
|
|
252
|
+
if viz_approx_samples:
|
|
253
|
+
flat_approx = np.concatenate(viz_approx_samples)
|
|
254
|
+
flat_full = np.concatenate(viz_full_samples)
|
|
255
|
+
else:
|
|
256
|
+
flat_approx = np.array([])
|
|
257
|
+
flat_full = np.array([])
|
|
258
|
+
|
|
259
|
+
print(f" > Samples Collected: {len(flat_approx):,} points")
|
|
260
|
+
|
|
261
|
+
# --- FILE 1: SUMMARY (1080p) ---
|
|
262
|
+
print(f" > Saving Summary Plot: {plot_summary_filename}")
|
|
263
|
+
fig1, ax1 = plt.subplots(1, 2, figsize=(16, 7))
|
|
264
|
+
|
|
265
|
+
# Plot 1: Variance Stabilization
|
|
266
|
+
ax = ax1[0]
|
|
267
|
+
ax.scatter(mean_raw, var_approx, s=2, alpha=0.5, color='red', label='Approx (Poisson)')
|
|
268
|
+
ax.scatter(mean_raw, var_full, s=2, alpha=0.5, color='blue', label='Full (NB Pearson)')
|
|
269
|
+
ax.axhline(1.0, color='black', linestyle='--', linewidth=1)
|
|
270
|
+
ax.set_xscale('log')
|
|
271
|
+
ax.set_yscale('log')
|
|
272
|
+
ax.set_title("Variance Stabilization Check")
|
|
273
|
+
ax.set_xlabel("Mean Raw Expression (log)")
|
|
274
|
+
ax.set_ylabel("Variance of Residuals (log)")
|
|
275
|
+
ax.legend()
|
|
276
|
+
ax.grid(True, which='both', linestyle='--', alpha=0.5)
|
|
277
|
+
|
|
278
|
+
# Plot 2: Distribution (Histogram + KDE Overlay)
|
|
279
|
+
ax = ax1[1]
|
|
280
|
+
if len(flat_approx) > 100:
|
|
281
|
+
mask_kde = (flat_approx > -10) & (flat_approx < 10)
|
|
282
|
+
bins = np.linspace(-5, 5, 100)
|
|
283
|
+
ax.hist(flat_approx[mask_kde], bins=bins, color='red', alpha=0.2, density=True, label='_nolegend_')
|
|
284
|
+
ax.hist(flat_full[mask_kde], bins=bins, color='blue', alpha=0.2, density=True, label='_nolegend_')
|
|
285
|
+
|
|
286
|
+
sns.kdeplot(flat_approx[mask_kde], fill=False, color='red', linewidth=2, label='Approx', ax=ax, warn_singular=False)
|
|
287
|
+
sns.kdeplot(flat_full[mask_kde], fill=False, color='blue', linewidth=2, label='Full', ax=ax, warn_singular=False)
|
|
288
|
+
|
|
289
|
+
ax.set_yscale('log')
|
|
290
|
+
ax.set_ylim(bottom=0.001)
|
|
291
|
+
ax.set_xlim(-5, 5)
|
|
292
|
+
ax.set_title("Distribution of Residuals (Log Scale)")
|
|
293
|
+
ax.set_xlabel("Residual Value")
|
|
294
|
+
ax.legend()
|
|
295
|
+
ax.grid(True, alpha=0.3)
|
|
296
|
+
|
|
297
|
+
plt.tight_layout()
|
|
298
|
+
plt.savefig(plot_summary_filename, dpi=120)
|
|
299
|
+
plt.close()
|
|
300
|
+
|
|
301
|
+
# --- FILE 2: DETAIL (4K) ---
|
|
302
|
+
print(f" > Saving Detail Plot: {plot_detail_filename}")
|
|
303
|
+
fig2, ax2 = plt.subplots(figsize=(20, 11))
|
|
304
|
+
|
|
305
|
+
if len(flat_approx) > 0:
|
|
306
|
+
ax2.scatter(flat_approx, flat_full, s=1, alpha=0.5, color='purple')
|
|
307
|
+
lims = [
|
|
308
|
+
np.min([ax2.get_xlim(), ax2.get_ylim()]),
|
|
309
|
+
np.max([ax2.get_xlim(), ax2.get_ylim()]),
|
|
310
|
+
]
|
|
311
|
+
ax2.plot(lims, lims, 'k-', alpha=0.75, zorder=0)
|
|
312
|
+
|
|
313
|
+
ax2.set_title("Residual Shrinkage (Sampled)")
|
|
314
|
+
ax2.set_xlabel("Approx Residuals")
|
|
315
|
+
ax2.set_ylabel("Full Residuals")
|
|
316
|
+
ax2.grid(True, alpha=0.3)
|
|
317
|
+
|
|
318
|
+
plt.tight_layout()
|
|
319
|
+
plt.savefig(plot_detail_filename, dpi=200)
|
|
320
|
+
plt.close()
|
|
321
|
+
|
|
322
|
+
if hasattr(adata_in, "file") and adata_in.file is not None: adata_in.file.close()
|
|
323
|
+
print(f"Total time: {time.perf_counter() - start_time:.2f} seconds.\n")
|
|
@@ -5,7 +5,7 @@ with open("README.md", "r", encoding="utf-8") as fh:
|
|
|
5
5
|
|
|
6
6
|
setuptools.setup(
|
|
7
7
|
name="M3Drop", # Name for pip (pip install M3Drop)
|
|
8
|
-
version="0.4.
|
|
8
|
+
version="0.4.56",
|
|
9
9
|
author="Tallulah Andrews",
|
|
10
10
|
author_email="tandrew6@uwo.ca",
|
|
11
11
|
description="A Python implementation of the M3Drop single-cell RNA-seq analysis tool.",
|
|
@@ -1,207 +0,0 @@
|
|
|
1
|
-
import pickle
|
|
2
|
-
import time
|
|
3
|
-
import sys
|
|
4
|
-
import numpy as np
|
|
5
|
-
import h5py
|
|
6
|
-
import anndata
|
|
7
|
-
import pandas as pd
|
|
8
|
-
import os
|
|
9
|
-
from scipy import sparse
|
|
10
|
-
|
|
11
|
-
try:
|
|
12
|
-
from numba import jit, prange
|
|
13
|
-
except ImportError:
|
|
14
|
-
print("CRITICAL ERROR: 'numba' not found. Please install it (pip install numba).")
|
|
15
|
-
sys.exit(1)
|
|
16
|
-
|
|
17
|
-
# [FIX] Strict Relative Import
|
|
18
|
-
from .ControlDeviceCPU import ControlDevice
|
|
19
|
-
|
|
20
|
-
# ==========================================
|
|
21
|
-
# NUMBA KERNELS (CPU)
|
|
22
|
-
# ==========================================
|
|
23
|
-
|
|
24
|
-
@jit(nopython=True, parallel=True, fastmath=True)
|
|
25
|
-
def pearson_residual_kernel_cpu(counts, tj, ti, theta, total, out_matrix):
|
|
26
|
-
"""
|
|
27
|
-
Calculates Pearson residuals using Negative Binomial logic.
|
|
28
|
-
Parallelized across CPU cores.
|
|
29
|
-
"""
|
|
30
|
-
rows = counts.shape[0]
|
|
31
|
-
cols = counts.shape[1]
|
|
32
|
-
|
|
33
|
-
for r in prange(rows):
|
|
34
|
-
ti_val = ti[r]
|
|
35
|
-
for c in range(cols):
|
|
36
|
-
count_val = counts[r, c]
|
|
37
|
-
mu = (tj[c] * ti_val) / total
|
|
38
|
-
|
|
39
|
-
# theta is vector of size cols (genes)
|
|
40
|
-
theta_val = theta[c]
|
|
41
|
-
|
|
42
|
-
denom_sq = mu + ((mu * mu) / theta_val)
|
|
43
|
-
denom = np.sqrt(denom_sq)
|
|
44
|
-
|
|
45
|
-
if denom < 1e-12:
|
|
46
|
-
out_matrix[r, c] = 0.0
|
|
47
|
-
else:
|
|
48
|
-
out_matrix[r, c] = (count_val - mu) / denom
|
|
49
|
-
|
|
50
|
-
@jit(nopython=True, parallel=True, fastmath=True)
|
|
51
|
-
def pearson_approx_kernel_cpu(counts, tj, ti, total, out_matrix):
|
|
52
|
-
"""
|
|
53
|
-
Calculates Approximate Pearson residuals (Poisson limit).
|
|
54
|
-
"""
|
|
55
|
-
rows = counts.shape[0]
|
|
56
|
-
cols = counts.shape[1]
|
|
57
|
-
|
|
58
|
-
for r in prange(rows):
|
|
59
|
-
ti_val = ti[r]
|
|
60
|
-
for c in range(cols):
|
|
61
|
-
count_val = counts[r, c]
|
|
62
|
-
mu = (tj[c] * ti_val) / total
|
|
63
|
-
|
|
64
|
-
denom = np.sqrt(mu)
|
|
65
|
-
|
|
66
|
-
if denom < 1e-12:
|
|
67
|
-
out_matrix[r, c] = 0.0
|
|
68
|
-
else:
|
|
69
|
-
out_matrix[r, c] = (count_val - mu) / denom
|
|
70
|
-
|
|
71
|
-
# ==========================================
|
|
72
|
-
# NORMALIZATION FUNCTION
|
|
73
|
-
# ==========================================
|
|
74
|
-
|
|
75
|
-
def NBumiPearsonResidualsCombinedCPU(
|
|
76
|
-
raw_filename: str,
|
|
77
|
-
mask_filename: str,
|
|
78
|
-
fit_filename: str,
|
|
79
|
-
stats_filename: str,
|
|
80
|
-
output_filename_full: str,
|
|
81
|
-
output_filename_approx: str,
|
|
82
|
-
mode: str = "auto",
|
|
83
|
-
manual_target: int = 3000
|
|
84
|
-
):
|
|
85
|
-
"""
|
|
86
|
-
CPU-Optimized: Calculates Full and Approximate residuals in a SINGLE PASS.
|
|
87
|
-
Uses Numba for acceleration on L3-sized dense chunks.
|
|
88
|
-
"""
|
|
89
|
-
start_time = time.perf_counter()
|
|
90
|
-
print(f"FUNCTION: NBumiPearsonResidualsCombinedCPU() | FILE: {raw_filename}")
|
|
91
|
-
|
|
92
|
-
# 1. Load Mask
|
|
93
|
-
with open(mask_filename, 'rb') as f: mask = pickle.load(f)
|
|
94
|
-
ng_filtered = int(np.sum(mask))
|
|
95
|
-
|
|
96
|
-
# 2. Init Device
|
|
97
|
-
with h5py.File(raw_filename, 'r') as f: indptr_cpu = f['X']['indptr'][:]; total_rows = len(indptr_cpu) - 1
|
|
98
|
-
device = ControlDevice(indptr=indptr_cpu, total_rows=total_rows, n_genes=ng_filtered, mode=mode, manual_target=manual_target)
|
|
99
|
-
nc = device.total_rows
|
|
100
|
-
|
|
101
|
-
print("Phase [1/2]: Initializing parameters...")
|
|
102
|
-
# Load parameters
|
|
103
|
-
with open(fit_filename, 'rb') as f: fit = pickle.load(f)
|
|
104
|
-
with open(stats_filename, 'rb') as f: stats = pickle.load(f)
|
|
105
|
-
|
|
106
|
-
# Common params (Numpy Arrays)
|
|
107
|
-
total = fit['vals']['total']
|
|
108
|
-
tjs = fit['vals']['tjs'].values.astype(np.float64)
|
|
109
|
-
tis = fit['vals']['tis'].values.astype(np.float64)
|
|
110
|
-
|
|
111
|
-
# Specific params
|
|
112
|
-
sizes = fit['sizes'].values.astype(np.float64) # For Full
|
|
113
|
-
|
|
114
|
-
# Setup Output Files
|
|
115
|
-
adata_in = anndata.read_h5ad(raw_filename, backed='r')
|
|
116
|
-
filtered_var = adata_in.var[mask]
|
|
117
|
-
|
|
118
|
-
# Create skeletons
|
|
119
|
-
adata_out_full = anndata.AnnData(obs=adata_in.obs, var=filtered_var)
|
|
120
|
-
adata_out_full.write_h5ad(output_filename_full, compression=None)
|
|
121
|
-
|
|
122
|
-
adata_out_approx = anndata.AnnData(obs=adata_in.obs, var=filtered_var)
|
|
123
|
-
adata_out_approx.write_h5ad(output_filename_approx, compression=None)
|
|
124
|
-
|
|
125
|
-
# --- CHUNK SIZE FIX ---
|
|
126
|
-
# Calculate appropriate H5 storage chunks
|
|
127
|
-
storage_chunk_rows = int(1_000_000_000 / (ng_filtered * 8))
|
|
128
|
-
|
|
129
|
-
# [CRITICAL FIX] Clamp chunk size to total rows (nc)
|
|
130
|
-
if storage_chunk_rows > nc:
|
|
131
|
-
storage_chunk_rows = nc
|
|
132
|
-
|
|
133
|
-
if storage_chunk_rows < 1:
|
|
134
|
-
storage_chunk_rows = 1
|
|
135
|
-
# ----------------------
|
|
136
|
-
|
|
137
|
-
# Open both files for writing simultaneously
|
|
138
|
-
with h5py.File(output_filename_full, 'a') as f_full, h5py.File(output_filename_approx, 'a') as f_approx:
|
|
139
|
-
if 'X' in f_full: del f_full['X']
|
|
140
|
-
if 'X' in f_approx: del f_approx['X']
|
|
141
|
-
|
|
142
|
-
# Float64 output
|
|
143
|
-
out_x_full = f_full.create_dataset(
|
|
144
|
-
'X', shape=(nc, ng_filtered), chunks=(storage_chunk_rows, ng_filtered), dtype='float64'
|
|
145
|
-
)
|
|
146
|
-
out_x_approx = f_approx.create_dataset(
|
|
147
|
-
'X', shape=(nc, ng_filtered), chunks=(storage_chunk_rows, ng_filtered), dtype='float64'
|
|
148
|
-
)
|
|
149
|
-
|
|
150
|
-
with h5py.File(raw_filename, 'r') as f_in:
|
|
151
|
-
h5_indptr = f_in['X']['indptr']
|
|
152
|
-
h5_data = f_in['X']['data']
|
|
153
|
-
h5_indices = f_in['X']['indices']
|
|
154
|
-
|
|
155
|
-
current_row = 0
|
|
156
|
-
while current_row < nc:
|
|
157
|
-
# Dense mode is faster for Numba
|
|
158
|
-
end_row = device.get_next_chunk(current_row, mode='dense', overhead_multiplier=3.0)
|
|
159
|
-
if end_row is None or end_row <= current_row: break
|
|
160
|
-
|
|
161
|
-
chunk_size = end_row - current_row
|
|
162
|
-
print(f"Phase [2/2]: Processing rows {end_row} of {nc} | Chunk: {chunk_size}", end='\r')
|
|
163
|
-
|
|
164
|
-
start_idx, end_idx = h5_indptr[current_row], h5_indptr[end_row]
|
|
165
|
-
|
|
166
|
-
# Load & Filter
|
|
167
|
-
data = np.array(h5_data[start_idx:end_idx], dtype=np.float64)
|
|
168
|
-
indices = np.array(h5_indices[start_idx:end_idx])
|
|
169
|
-
indptr = np.array(h5_indptr[current_row:end_row+1] - h5_indptr[current_row])
|
|
170
|
-
|
|
171
|
-
chunk_csr = sparse.csr_matrix((data, indices, indptr), shape=(chunk_size, len(mask)))
|
|
172
|
-
chunk_csr = chunk_csr[:, mask]
|
|
173
|
-
chunk_csr.data = np.ceil(chunk_csr.data)
|
|
174
|
-
|
|
175
|
-
# Convert to Dense for Numba (faster than sparse iteration for dense ops)
|
|
176
|
-
counts_dense = chunk_csr.toarray()
|
|
177
|
-
|
|
178
|
-
# --- CALC 1: APPROX ---
|
|
179
|
-
approx_out = np.empty_like(counts_dense)
|
|
180
|
-
pearson_approx_kernel_cpu(
|
|
181
|
-
counts_dense,
|
|
182
|
-
tjs,
|
|
183
|
-
tis[current_row:end_row],
|
|
184
|
-
total,
|
|
185
|
-
approx_out
|
|
186
|
-
)
|
|
187
|
-
out_x_approx[current_row:end_row, :] = approx_out
|
|
188
|
-
del approx_out
|
|
189
|
-
|
|
190
|
-
# --- CALC 2: FULL (In-place on counts_dense) ---
|
|
191
|
-
# We can reuse the counts_dense buffer for output to save RAM
|
|
192
|
-
pearson_residual_kernel_cpu(
|
|
193
|
-
counts_dense,
|
|
194
|
-
tjs,
|
|
195
|
-
tis[current_row:end_row],
|
|
196
|
-
sizes,
|
|
197
|
-
total,
|
|
198
|
-
counts_dense # Overwrite input
|
|
199
|
-
)
|
|
200
|
-
out_x_full[current_row:end_row, :] = counts_dense
|
|
201
|
-
|
|
202
|
-
current_row = end_row
|
|
203
|
-
|
|
204
|
-
print(f"\nPhase [2/2]: COMPLETE{' '*50}")
|
|
205
|
-
|
|
206
|
-
if hasattr(adata_in, "file") and adata_in.file is not None: adata_in.file.close()
|
|
207
|
-
print(f"Total time: {time.perf_counter() - start_time:.2f} seconds.\n")
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|