Skip to content

perf: parallel downsample#4004

Open
flying-sheep wants to merge 43 commits intomainfrom
parallel-downsample
Open

perf: parallel downsample#4004
flying-sheep wants to merge 43 commits intomainfrom
parallel-downsample

Conversation

@flying-sheep
Copy link
Member

@flying-sheep flying-sheep commented Mar 18, 2026

needed to implement np.random.Generator.choice for this.

@codecov
Copy link

codecov bot commented Mar 18, 2026

❌ 2 Tests Failed:

Tests completed Failed Passed Skipped
1675 2 1673 1120
View the top 2 failed test(s) by shortest run time
tests/test_neighbors.py::test_neighbors_distance_equivalence
Stack Traces | 0.944s run time
#x1B[0m#x1B[94mdef#x1B[39;49;00m#x1B[90m #x1B[39;49;00m#x1B[92mtest_neighbors_distance_equivalence#x1B[39;49;00m() -> #x1B[94mNone#x1B[39;49;00m:#x1B[90m#x1B[39;49;00m
        adata = pbmc68k_reduced()#x1B[90m#x1B[39;49;00m
        adata_d = adata.copy()#x1B[90m#x1B[39;49;00m
    #x1B[90m#x1B[39;49;00m
        sc.pp.neighbors(adata)#x1B[90m#x1B[39;49;00m
        #x1B[90m# reusing the same distances#x1B[39;49;00m#x1B[90m#x1B[39;49;00m
        sc.pp.neighbors(adata_d, distances=adata.obsp[#x1B[33m"#x1B[39;49;00m#x1B[33mdistances#x1B[39;49;00m#x1B[33m"#x1B[39;49;00m])#x1B[90m#x1B[39;49;00m
    #x1B[90m#x1B[39;49;00m
        np.testing.assert_allclose(#x1B[90m#x1B[39;49;00m
            adata.obsp[#x1B[33m"#x1B[39;49;00m#x1B[33mconnectivities#x1B[39;49;00m#x1B[33m"#x1B[39;49;00m].toarray(),#x1B[90m#x1B[39;49;00m
            adata_d.obsp[#x1B[33m"#x1B[39;49;00m#x1B[33mconnectivities#x1B[39;49;00m#x1B[33m"#x1B[39;49;00m].toarray(),#x1B[90m#x1B[39;49;00m
            rtol=#x1B[94m1e-5#x1B[39;49;00m,#x1B[90m#x1B[39;49;00m
        )#x1B[90m#x1B[39;49;00m
        np.testing.assert_allclose(#x1B[90m#x1B[39;49;00m
            adata.obsp[#x1B[33m"#x1B[39;49;00m#x1B[33mdistances#x1B[39;49;00m#x1B[33m"#x1B[39;49;00m].toarray(),#x1B[90m#x1B[39;49;00m
            adata_d.obsp[#x1B[33m"#x1B[39;49;00m#x1B[33mdistances#x1B[39;49;00m#x1B[33m"#x1B[39;49;00m].toarray(),#x1B[90m#x1B[39;49;00m
            rtol=#x1B[94m1e-5#x1B[39;49;00m,#x1B[90m#x1B[39;49;00m
        )#x1B[90m#x1B[39;49;00m
        p, p_d = (ad.uns[#x1B[33m"#x1B[39;49;00m#x1B[33mneighbors#x1B[39;49;00m#x1B[33m"#x1B[39;49;00m][#x1B[33m"#x1B[39;49;00m#x1B[33mparams#x1B[39;49;00m#x1B[33m"#x1B[39;49;00m].copy() #x1B[94mfor#x1B[39;49;00m ad #x1B[95min#x1B[39;49;00m (adata, adata_d))#x1B[90m#x1B[39;49;00m
        #x1B[94massert#x1B[39;49;00m p.pop(#x1B[33m"#x1B[39;49;00m#x1B[33mmetric#x1B[39;49;00m#x1B[33m"#x1B[39;49;00m) == #x1B[33m"#x1B[39;49;00m#x1B[33meuclidean#x1B[39;49;00m#x1B[33m"#x1B[39;49;00m#x1B[90m#x1B[39;49;00m
        #x1B[94massert#x1B[39;49;00m p_d.pop(#x1B[33m"#x1B[39;49;00m#x1B[33mmetric#x1B[39;49;00m#x1B[33m"#x1B[39;49;00m) #x1B[95mis#x1B[39;49;00m #x1B[94mNone#x1B[39;49;00m#x1B[90m#x1B[39;49;00m
>       #x1B[94massert#x1B[39;49;00m p == p_d#x1B[90m#x1B[39;49;00m
#x1B[1m#x1B[31mE       AssertionError: assert {'method': 'u...dom_state': 0} == {'method': 'u...eighbors': 15}#x1B[0m
#x1B[1m#x1B[31mE         #x1B[0m
#x1B[1m#x1B[31mE         Omitting 2 identical items, use -vv to show#x1B[0m
#x1B[1m#x1B[31mE         Left contains 1 more item:#x1B[0m
#x1B[1m#x1B[31mE         #x1B[0m{#x1B[33m'#x1B[39;49;00m#x1B[33mrandom_state#x1B[39;49;00m#x1B[33m'#x1B[39;49;00m: #x1B[94m0#x1B[39;49;00m}#x1B[90m#x1B[39;49;00m#x1B[0m
#x1B[1m#x1B[31mE         #x1B[0m
#x1B[1m#x1B[31mE         Full diff:#x1B[0m
#x1B[1m#x1B[31mE         #x1B[0m#x1B[90m #x1B[39;49;00m {#x1B[90m#x1B[39;49;00m#x1B[0m
#x1B[1m#x1B[31mE         #x1B[90m #x1B[39;49;00m     'method': 'umap',#x1B[90m#x1B[39;49;00m#x1B[0m
#x1B[1m#x1B[31mE         #x1B[90m #x1B[39;49;00m     'n_neighbors': 15,#x1B[90m#x1B[39;49;00m#x1B[0m
#x1B[1m#x1B[31mE         #x1B[92m+     'random_state': 0,#x1B[39;49;00m#x1B[90m#x1B[39;49;00m#x1B[0m
#x1B[1m#x1B[31mE         #x1B[90m #x1B[39;49;00m }#x1B[90m#x1B[39;49;00m#x1B[0m

#x1B[1m#x1B[31mtests/test_neighbors.py#x1B[0m:303: AssertionError
tests/test_preprocessing.py::test_downsample_counts_per_cell_multiple_targets[csc_matrix-float64-replace]
Stack Traces | 1.08s run time
#x1B[0m#x1B[37m@pytest#x1B[39;49;00m.mark.parametrize(#x1B[33m"#x1B[39;49;00m#x1B[33mreplace#x1B[39;49;00m#x1B[33m"#x1B[39;49;00m, [#x1B[94mTrue#x1B[39;49;00m, #x1B[94mFalse#x1B[39;49;00m], ids=[#x1B[33m"#x1B[39;49;00m#x1B[33mreplace#x1B[39;49;00m#x1B[33m"#x1B[39;49;00m, #x1B[33m"#x1B[39;49;00m#x1B[33mno_replace#x1B[39;49;00m#x1B[33m"#x1B[39;49;00m])#x1B[90m#x1B[39;49;00m
    #x1B[94mdef#x1B[39;49;00m#x1B[90m #x1B[39;49;00m#x1B[92mtest_downsample_counts_per_cell_multiple_targets#x1B[39;49;00m(#x1B[90m#x1B[39;49;00m
        *, count_matrix_format: _MatrixFormat, replace: #x1B[96mbool#x1B[39;49;00m, dtype: DTypeLike#x1B[90m#x1B[39;49;00m
    ) -> #x1B[94mNone#x1B[39;49;00m:#x1B[90m#x1B[39;49;00m
        rng = np.random.default_rng()#x1B[90m#x1B[39;49;00m
        targets = rng.integers(#x1B[94m500#x1B[39;49;00m, #x1B[94m1500#x1B[39;49;00m, #x1B[94m1000#x1B[39;49;00m)#x1B[90m#x1B[39;49;00m
        x = rng.integers(#x1B[94m0#x1B[39;49;00m, #x1B[94m100#x1B[39;49;00m, (#x1B[94m1000#x1B[39;49;00m, #x1B[94m100#x1B[39;49;00m)) * rng.binomial(#x1B[94m1#x1B[39;49;00m, #x1B[94m0.3#x1B[39;49;00m, (#x1B[94m1000#x1B[39;49;00m, #x1B[94m100#x1B[39;49;00m))#x1B[90m#x1B[39;49;00m
        x = x.astype(dtype)#x1B[90m#x1B[39;49;00m
        adata = AnnData(X=count_matrix_format(x).astype(dtype))#x1B[90m#x1B[39;49;00m
        initial_totals = np.ravel(adata.X.sum(axis=#x1B[94m1#x1B[39;49;00m))#x1B[90m#x1B[39;49;00m
        #x1B[94mwith#x1B[39;49;00m pytest.raises(#x1B[96mValueError#x1B[39;49;00m, match=#x1B[33mr#x1B[39;49;00m#x1B[33m"#x1B[39;49;00m#x1B[33mcounts_per_cell.*length as number of obs#x1B[39;49;00m#x1B[33m"#x1B[39;49;00m):#x1B[90m#x1B[39;49;00m
            sc.pp.downsample_counts(adata, counts_per_cell=[#x1B[94m40#x1B[39;49;00m, #x1B[94m10#x1B[39;49;00m], replace=replace)#x1B[90m#x1B[39;49;00m
>       adata = sc.pp.downsample_counts(#x1B[90m#x1B[39;49;00m
            adata, counts_per_cell=targets, replace=replace, copy=#x1B[94mTrue#x1B[39;49;00m#x1B[90m#x1B[39;49;00m
        )#x1B[90m#x1B[39;49;00m

#x1B[1m#x1B[31mtests/test_preprocessing.py#x1B[0m:575: 
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ 
#x1B[1m#x1B[.../scanpy/_utils/random.py#x1B[0m:201: in wrapper
    #x1B[0m#x1B[94mreturn#x1B[39;49;00m func(*args, **kwargs)#x1B[90m#x1B[39;49;00m
           ^^^^^^^^^^^^^^^^^^^^^#x1B[90m#x1B[39;49;00m
#x1B[1m#x1B[.../scanpy/preprocessing/_simple.py#x1B[0m:1044: in downsample_counts
    #x1B[0madata.X = _downsample_per_cell(#x1B[90m#x1B[39;49;00m
#x1B[1m#x1B[.../scanpy/preprocessing/_simple.py#x1B[0m:1079: in _downsample_per_cell
    #x1B[0m_downsample_arrays(#x1B[90m#x1B[39;49;00m
#x1B[1m#x1B[31msrc/scanpy/_compat.py#x1B[0m:149: in wrapper
    #x1B[0m#x1B[94mreturn#x1B[39;49;00m fns[parallel](*args, **kwargs)#x1B[90m#x1B[39;49;00m
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^#x1B[90m#x1B[39;49;00m
#x1B[1m#x1B[31m../../../..../scanpy/B9PcT7QG/hatch-test.few-extras/lib/python3.14.../numba/core/dispatcher.py#x1B[0m:443: in _compile_for_args
    #x1B[0m#x1B[94mraise#x1B[39;49;00m e#x1B[90m#x1B[39;49;00m
#x1B[1m#x1B[31m../../../..../scanpy/B9PcT7QG/hatch-test.few-extras/lib/python3.14.../numba/core/dispatcher.py#x1B[0m:376: in _compile_for_args
    #x1B[0mreturn_val = #x1B[96mself#x1B[39;49;00m.compile(#x1B[96mtuple#x1B[39;49;00m(argtypes))#x1B[90m#x1B[39;49;00m
                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^#x1B[90m#x1B[39;49;00m
#x1B[1m#x1B[31m../../../..../scanpy/B9PcT7QG/hatch-test.few-extras/lib/python3.14.../numba/core/dispatcher.py#x1B[0m:908: in compile
    #x1B[0mcres = #x1B[96mself#x1B[39;49;00m._compiler.compile(args, return_type)#x1B[90m#x1B[39;49;00m
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^#x1B[90m#x1B[39;49;00m
#x1B[1m#x1B[31m../../../..../scanpy/B9PcT7QG/hatch-test.few-extras/lib/python3.14.../numba/core/dispatcher.py#x1B[0m:80: in compile
    #x1B[0mstatus, retval = #x1B[96mself#x1B[39;49;00m._compile_cached(args, return_type)#x1B[90m#x1B[39;49;00m
                     ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^#x1B[90m#x1B[39;49;00m
#x1B[1m#x1B[31m../../../..../scanpy/B9PcT7QG/hatch-test.few-extras/lib/python3.14.../numba/core/dispatcher.py#x1B[0m:94: in _compile_cached
    #x1B[0mretval = #x1B[96mself#x1B[39;49;00m._compile_core(args, return_type)#x1B[90m#x1B[39;49;00m
             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^#x1B[90m#x1B[39;49;00m
#x1B[1m#x1B[31m../../../..../scanpy/B9PcT7QG/hatch-test.few-extras/lib/python3.14.../numba/core/dispatcher.py#x1B[0m:107: in _compile_core
    #x1B[0mcres = compiler.compile_extra(#x1B[96mself#x1B[39;49;00m.targetdescr.typing_context,#x1B[90m#x1B[39;49;00m
#x1B[1m#x1B[31m../../../..../scanpy/B9PcT7QG/hatch-test.few-extras/lib/python3.14.../numba/core/compiler.py#x1B[0m:739: in compile_extra
    #x1B[0m#x1B[94mreturn#x1B[39;49;00m pipeline.compile_extra(func)#x1B[90m#x1B[39;49;00m
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^#x1B[90m#x1B[39;49;00m
#x1B[1m#x1B[31m../../../..../scanpy/B9PcT7QG/hatch-test.few-extras/lib/python3.14.../numba/core/compiler.py#x1B[0m:439: in compile_extra
    #x1B[0m#x1B[94mreturn#x1B[39;49;00m #x1B[96mself#x1B[39;49;00m._compile_bytecode()#x1B[90m#x1B[39;49;00m
           ^^^^^^^^^^^^^^^^^^^^^^^^#x1B[90m#x1B[39;49;00m
#x1B[1m#x1B[31m../../../..../scanpy/B9PcT7QG/hatch-test.few-extras/lib/python3.14.../numba/core/compiler.py#x1B[0m:505: in _compile_bytecode
    #x1B[0m#x1B[94mreturn#x1B[39;49;00m #x1B[96mself#x1B[39;49;00m._compile_core()#x1B[90m#x1B[39;49;00m
           ^^^^^^^^^^^^^^^^^^^^#x1B[90m#x1B[39;49;00m
#x1B[1m#x1B[31m../../../..../scanpy/B9PcT7QG/hatch-test.few-extras/lib/python3.14.../numba/core/compiler.py#x1B[0m:481: in _compile_core
    #x1B[0m#x1B[94mraise#x1B[39;49;00m e#x1B[90m#x1B[39;49;00m
#x1B[1m#x1B[31m../../../..../scanpy/B9PcT7QG/hatch-test.few-extras/lib/python3.14.../numba/core/compiler.py#x1B[0m:473: in _compile_core
    #x1B[0mpm.run(#x1B[96mself#x1B[39;49;00m.state)#x1B[90m#x1B[39;49;00m
#x1B[1m#x1B[31m../../../..../scanpy/B9PcT7QG/hatch-test.few-extras/lib/python3.14.../numba/core/compiler_machinery.py#x1B[0m:363: in run
    #x1B[0m#x1B[94mraise#x1B[39;49;00m e#x1B[90m#x1B[39;49;00m
#x1B[1m#x1B[31m../../../..../scanpy/B9PcT7QG/hatch-test.few-extras/lib/python3.14.../numba/core/compiler_machinery.py#x1B[0m:356: in run
    #x1B[0m#x1B[96mself#x1B[39;49;00m._runPass(idx, pass_inst, state)#x1B[90m#x1B[39;49;00m
#x1B[1m#x1B[31m../../../..../scanpy/B9PcT7QG/hatch-test.few-extras/lib/python3.14.../numba/core/compiler_lock.py#x1B[0m:35: in _acquire_compile_lock
    #x1B[0m#x1B[94mreturn#x1B[39;49;00m func(*args, **kwargs)#x1B[90m#x1B[39;49;00m
           ^^^^^^^^^^^^^^^^^^^^^#x1B[90m#x1B[39;49;00m
#x1B[1m#x1B[31m../../../..../scanpy/B9PcT7QG/hatch-test.few-extras/lib/python3.14.../numba/core/compiler_machinery.py#x1B[0m:311: in _runPass
    #x1B[0mmutated |= check(pss.run_pass, internal_state)#x1B[90m#x1B[39;49;00m
               ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^#x1B[90m#x1B[39;49;00m
#x1B[1m#x1B[31m../../../..../scanpy/B9PcT7QG/hatch-test.few-extras/lib/python3.14.../numba/core/compiler_machinery.py#x1B[0m:272: in check
    #x1B[0mmangled = func(compiler_state)#x1B[90m#x1B[39;49;00m
              ^^^^^^^^^^^^^^^^^^^^#x1B[90m#x1B[39;49;00m
#x1B[1m#x1B[31m../../../..../scanpy/B9PcT7QG/hatch-test.few-extras/lib/python3.14.../numba/core/typed_passes.py#x1B[0m:472: in run_pass
    #x1B[0mlower.lower()#x1B[90m#x1B[39;49;00m
#x1B[1m#x1B[31m../../../..../scanpy/B9PcT7QG/hatch-test.few-extras/lib/python3.14.../numba/core/lowering.py#x1B[0m:193: in lower
    #x1B[0m#x1B[96mself#x1B[39;49;00m.lower_normal_function(#x1B[96mself#x1B[39;49;00m.fndesc)#x1B[90m#x1B[39;49;00m
#x1B[1m#x1B[31m../../../..../scanpy/B9PcT7QG/hatch-test.few-extras/lib/python3.14.../numba/core/lowering.py#x1B[0m:232: in lower_normal_function
    #x1B[0mentry_block_tail = #x1B[96mself#x1B[39;49;00m.lower_function_body()#x1B[90m#x1B[39;49;00m
                       ^^^^^^^^^^^^^^^^^^^^^^^^^^#x1B[90m#x1B[39;49;00m
#x1B[1m#x1B[31m../../../..../scanpy/B9PcT7QG/hatch-test.few-extras/lib/python3.14.../numba/core/lowering.py#x1B[0m:262: in lower_function_body
    #x1B[0m#x1B[96mself#x1B[39;49;00m.lower_block(block)#x1B[90m#x1B[39;49;00m
#x1B[1m#x1B[31m../../../..../scanpy/B9PcT7QG/hatch-test.few-extras/lib/python3.14.../numba/core/lowering.py#x1B[0m:276: in lower_block
    #x1B[0m#x1B[96mself#x1B[39;49;00m.lower_inst(inst)#x1B[90m#x1B[39;49;00m
#x1B[1m#x1B[31m../../../..../scanpy/B9PcT7QG/hatch-test.few-extras/lib/python3.14.../numba/parfors/parfor_lowering.py#x1B[0m:52: in lower_inst
    #x1B[0m_lower_parfor_parallel(#x1B[96mself#x1B[39;49;00m, inst)#x1B[90m#x1B[39;49;00m
#x1B[1m#x1B[31m../../../..../scanpy/B9PcT7QG/hatch-test.few-extras/lib/python3.14.../numba/parfors/parfor_lowering.py#x1B[0m:68: in _lower_parfor_parallel
    #x1B[0m#x1B[94mreturn#x1B[39;49;00m _lower_parfor_parallel_std(lowerer, parfor)#x1B[90m#x1B[39;49;00m
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^#x1B[90m#x1B[39;49;00m
#x1B[1m#x1B[31m../../../..../scanpy/B9PcT7QG/hatch-test.few-extras/lib/python3.14.../numba/parfors/parfor_lowering.py#x1B[0m:350: in _lower_parfor_parallel_std
    #x1B[0mexp_name_to_tuple_var) = _create_gufunc_for_parfor_body(#x1B[90m#x1B[39;49;00m
#x1B[1m#x1B[31m../../../..../scanpy/B9PcT7QG/hatch-test.few-extras/lib/python3.14.../numba/parfors/parfor_lowering.py#x1B[0m:1609: in _create_gufunc_for_parfor_body
    #x1B[0mkernel_func = compiler.compile_ir(#x1B[90m#x1B[39;49;00m
#x1B[1m#x1B[31m../../../..../scanpy/B9PcT7QG/hatch-test.few-extras/lib/python3.14.../numba/core/compiler.py#x1B[0m:802: in compile_ir
    #x1B[0m#x1B[94mreturn#x1B[39;49;00m pipeline.compile_ir(func_ir=func_ir, lifted=lifted,#x1B[90m#x1B[39;49;00m
#x1B[1m#x1B[31m../../../..../scanpy/B9PcT7QG/hatch-test.few-extras/lib/python3.14.../numba/core/compiler.py#x1B[0m:449: in compile_ir
    #x1B[0m#x1B[94mreturn#x1B[39;49;00m #x1B[96mself#x1B[39;49;00m._compile_ir()#x1B[90m#x1B[39;49;00m
           ^^^^^^^^^^^^^^^^^^#x1B[90m#x1B[39;49;00m
#x1B[1m#x1B[31m../../../..../scanpy/B9PcT7QG/hatch-test.few-extras/lib/python3.14.../numba/core/compiler.py#x1B[0m:512: in _compile_ir
    #x1B[0m#x1B[94mreturn#x1B[39;49;00m #x1B[96mself#x1B[39;49;00m._compile_core()#x1B[90m#x1B[39;49;00m
           ^^^^^^^^^^^^^^^^^^^^#x1B[90m#x1B[39;49;00m
#x1B[1m#x1B[31m../../../..../scanpy/B9PcT7QG/hatch-test.few-extras/lib/python3.14.../numba/core/compiler.py#x1B[0m:481: in _compile_core
    #x1B[0m#x1B[94mraise#x1B[39;49;00m e#x1B[90m#x1B[39;49;00m
#x1B[1m#x1B[31m../../../..../scanpy/B9PcT7QG/hatch-test.few-extras/lib/python3.14.../numba/core/compiler.py#x1B[0m:473: in _compile_core
    #x1B[0mpm.run(#x1B[96mself#x1B[39;49;00m.state)#x1B[90m#x1B[39;49;00m
#x1B[1m#x1B[31m../../../..../scanpy/B9PcT7QG/hatch-test.few-extras/lib/python3.14.../numba/core/compiler_machinery.py#x1B[0m:363: in run
    #x1B[0m#x1B[94mraise#x1B[39;49;00m e#x1B[90m#x1B[39;49;00m
#x1B[1m#x1B[31m../../../..../scanpy/B9PcT7QG/hatch-test.few-extras/lib/python3.14.../numba/core/compiler_machinery.py#x1B[0m:356: in run
    #x1B[0m#x1B[96mself#x1B[39;49;00m._runPass(idx, pass_inst, state)#x1B[90m#x1B[39;49;00m
#x1B[1m#x1B[31m../../../..../scanpy/B9PcT7QG/hatch-test.few-extras/lib/python3.14.../numba/core/compiler_lock.py#x1B[0m:35: in _acquire_compile_lock
    #x1B[0m#x1B[94mreturn#x1B[39;49;00m func(*args, **kwargs)#x1B[90m#x1B[39;49;00m
           ^^^^^^^^^^^^^^^^^^^^^#x1B[90m#x1B[39;49;00m
#x1B[1m#x1B[31m../../../..../scanpy/B9PcT7QG/hatch-test.few-extras/lib/python3.14.../numba/core/compiler_machinery.py#x1B[0m:311: in _runPass
    #x1B[0mmutated |= check(pss.run_pass, internal_state)#x1B[90m#x1B[39;49;00m
               ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^#x1B[90m#x1B[39;49;00m
#x1B[1m#x1B[31m../../../..../scanpy/B9PcT7QG/hatch-test.few-extras/lib/python3.14.../numba/core/compiler_machinery.py#x1B[0m:272: in check
    #x1B[0mmangled = func(compiler_state)#x1B[90m#x1B[39;49;00m
              ^^^^^^^^^^^^^^^^^^^^#x1B[90m#x1B[39;49;00m
#x1B[1m#x1B[31m../../../..../scanpy/B9PcT7QG/hatch-test.few-extras/lib/python3.14.../numba/core/typed_passes.py#x1B[0m:114: in run_pass
    #x1B[0mtypemap, return_type, calltypes, errs = type_inference_stage(#x1B[90m#x1B[39;49;00m
#x1B[1m#x1B[31m../../../..../scanpy/B9PcT7QG/hatch-test.few-extras/lib/python3.14.../numba/core/typed_passes.py#x1B[0m:80: in type_inference_stage
    #x1B[0m#x1B[94mwith#x1B[39;49;00m callstack_ctx, warnings:#x1B[90m#x1B[39;49;00m
                        ^^^^^^^^#x1B[90m#x1B[39;49;00m
#x1B[1m#x1B[31m../../../..../scanpy/B9PcT7QG/hatch-test.few-extras/lib/python3.14.../numba/core/errors.py#x1B[0m:530: in __exit__
    #x1B[0m#x1B[96mself#x1B[39;49;00m.flush()#x1B[90m#x1B[39;49;00m
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ 

self = <numba.core.errors.WarningsFixer object at 0x7fd91481d790>

    #x1B[0m#x1B[94mdef#x1B[39;49;00m#x1B[90m #x1B[39;49;00m#x1B[92mflush#x1B[39;49;00m(#x1B[96mself#x1B[39;49;00m):#x1B[90m#x1B[39;49;00m
    #x1B[90m    #x1B[39;49;00m#x1B[33m"""#x1B[39;49;00m
    #x1B[33m    Emit all stored warnings.#x1B[39;49;00m
    #x1B[33m    """#x1B[39;49;00m#x1B[90m#x1B[39;49;00m
        #x1B[94mdef#x1B[39;49;00m#x1B[90m #x1B[39;49;00m#x1B[92mkey#x1B[39;49;00m(arg):#x1B[90m#x1B[39;49;00m
            #x1B[90m# It is possible through codegen to create entirely identical#x1B[39;49;00m#x1B[90m#x1B[39;49;00m
            #x1B[90m# warnings, this leads to comparing types when sorting which breaks#x1B[39;49;00m#x1B[90m#x1B[39;49;00m
            #x1B[90m# on Python 3. Key as str() and if the worse happens then `id`#x1B[39;49;00m#x1B[90m#x1B[39;49;00m
            #x1B[90m# creates some uniqueness#x1B[39;49;00m#x1B[90m#x1B[39;49;00m
            #x1B[94mreturn#x1B[39;49;00m #x1B[96mstr#x1B[39;49;00m(arg) + #x1B[96mstr#x1B[39;49;00m(#x1B[96mid#x1B[39;49;00m(arg))#x1B[90m#x1B[39;49;00m
    #x1B[90m#x1B[39;49;00m
        #x1B[94mfor#x1B[39;49;00m (filename, lineno, category), messages #x1B[95min#x1B[39;49;00m #x1B[96msorted#x1B[39;49;00m(#x1B[90m#x1B[39;49;00m
                #x1B[96mself#x1B[39;49;00m._warnings.items(), key=key):#x1B[90m#x1B[39;49;00m
            #x1B[94mfor#x1B[39;49;00m msg #x1B[95min#x1B[39;49;00m #x1B[96msorted#x1B[39;49;00m(messages):#x1B[90m#x1B[39;49;00m
>               warnings.warn_explicit(msg, category, filename, lineno)#x1B[90m#x1B[39;49;00m
#x1B[1m#x1B[31mE               numba.core.errors.NumbaTypeSafetyWarning: unsafe cast from uint64 to int64. Precision may be lost.#x1B[0m

#x1B[1m#x1B[31m../../../..../scanpy/B9PcT7QG/hatch-test.few-extras/lib/python3.14.../numba/core/errors.py#x1B[0m:523: NumbaTypeSafetyWarning

To view more test analytics, go to the Test Analytics Dashboard
📋 Got 3 mins? Take this short survey to help us improve Test Analytics.

@flying-sheep
Copy link
Member Author

flying-sheep commented Mar 20, 2026

local timings

downsample_per_cell

[ 0.00%] · For scanpy commit a5e57615 <main>:
[25.00%] ··· ================= ===== ==============
                  dataset       rng   random_state 
             ================= ===== ==============
              pbmc68k_reduced   n/a   7.02±0.05ms  
                   pbmc3k       n/a     117±10ms   
             ================= ===== ==============
[33.33%] · For scanpy commit d8d0e804 <pa/rng>:
[58.33%] ··· ================= ============ ==============
                  dataset          rng       random_state 
             ================= ============ ==============
              pbmc68k_reduced   20.5±0.6ms    16.8±0.3ms  
                   pbmc3k        92.4±5ms      140±5ms    
             ================= ============ ==============
[66.67%] · For scanpy commit 41cd6a65 <parallel-downsample>:
[91.67%] ··· ================= ============= ==============
                  dataset           rng       random_state 
             ================= ============= ==============
              pbmc68k_reduced   1.22±0.02ms   1.52±0.05ms  
                   pbmc3k        14.3±0.8ms     28.1±1ms   
             ================= ============= ==============

downsample_total

[ 0.00%] · For scanpy commit a5e57615 <main>:
[33.33%] ··· preprocessing_counts.PreprocessingCountsRngSuite.time_downsample_total
[33.33%] ··· ================= ===== ==============
                  dataset       rng   random_state 
             ================= ===== ==============
              pbmc68k_reduced   n/a    9.69±0.4ms  
                   pbmc3k       n/a     375±6ms    
             ================= ===== ==============
[33.33%] · For scanpy commit d8d0e804 <pa/rng>:
[66.67%] ··· preprocessing_counts.PreprocessingCountsRngSuite.time_downsample_total
[66.67%] ··· ================= ============ ==============
                  dataset          rng       random_state 
             ================= ============ ==============
              pbmc68k_reduced   4.79±0.2ms    8.92±0.3ms  
                   pbmc3k        59.0±2ms      277±6ms    
             ================= ============ ==============
[66.67%] · For scanpy commit 41cd6a65 <parallel-downsample>:
[100.00%] ··· preprocessing_counts.PreprocessingCountsRngSuite.time_downsample_total
[100.00%] ··· ================= ============ ==============
              --                           layer           
              ----------------- ---------------------------
                   dataset          rng       random_state 
              ================= ============ ==============
               pbmc68k_reduced   5.47±0.2ms    10.0±0.1ms  
                    pbmc3k        145±2ms       424±20ms   
              ================= ============ ==============

@flying-sheep flying-sheep requested a review from ilan-gold March 20, 2026 10:21
@flying-sheep flying-sheep marked this pull request as ready for review March 20, 2026 10:21
@scverse-benchmark
Copy link

scverse-benchmark bot commented Mar 20, 2026

Benchmark changes

Change Before [9de11b1] After [d748f0b] Ratio Benchmark (Parameter)
+ 431M 477M 1.11 preprocessing_counts.PreprocessingCountsRngSuite.peakmem_downsample_total('pbmc3k', 'rng')
- 205±0.5ms 28.0±6ms 0.14 preprocessing_counts.PreprocessingCountsRngSuite.time_downsample_per_cell('pbmc3k', 'random_state')
- 71.8±0.7ms 20.8±6ms 0.29 preprocessing_counts.PreprocessingCountsRngSuite.time_downsample_per_cell('pbmc3k', 'rng')
- 25.1±0.2ms 1.48±0.2ms 0.06 preprocessing_counts.PreprocessingCountsRngSuite.time_downsample_per_cell('pbmc68k_reduced', 'random_state')
+ 192±8ms 306±4ms 1.59 preprocessing_counts.PreprocessingCountsRngSuite.time_downsample_total('pbmc3k', 'random_state')
+ 63.3±0.3ms 182±20ms 2.88 preprocessing_counts.PreprocessingCountsRngSuite.time_downsample_total('pbmc3k', 'rng')
+ 13.1±0.4ms 16.7±2ms 1.27 preprocessing_counts.PreprocessingCountsRngSuite.time_downsample_total('pbmc68k_reduced', 'random_state')
+ 7.14±0.3ms 11.1±0.6ms 1.55 preprocessing_counts.PreprocessingCountsRngSuite.time_downsample_total('pbmc68k_reduced', 'rng')

Comparison: https://github.com/scverse/scanpy/compare/9de11b14a0b4cac36c256806ce906d9cbbe1a96a..d748f0b5640e22732beb6ab169ff47b1ffbc83e0
Last changed: Fri, 20 Mar 2026 14:39:36 +0000

More details: https://github.com/scverse/scanpy/pull/4004/checks?check_run_id=67912471934

Comment on lines +1122 to +1133
for i in numba.prange(len(rows)):
if not mask[i]:
continue
thread = numba.get_thread_id()
_downsample_array_inner(
rows[i],
targets[i],
rng=rngs[thread] if rngs is not None else None,
seed=seed,
replace=replace,
inplace=True,
)
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

cc @selmanozleyen re: my comments on scverse/squidpy#1125

This is also a good option - let numba handle the threading for you

replace: bool,
) -> T:
total_counts = int(total_counts)
total = x.sum()
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

should use fau to prepare for scverse/fast-array-utils#155 maybe?

Base automatically changed from pa/rng to main March 20, 2026 14:01
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants