PERF Call KDTree in mutual_info to reduce memory footprint - #17878

Thank you @noelano

Please add an |Efficiency| entry to the change log at doc/whats_new/v0.24.rst. Like the other entries there, please reference this pull request with :pr: and credit yourself (and other contributors if applicable) with :user:

I applied suggestions, but now I see that we still need to have the "what's new" entry to merge.

I wonder if this metric can take advantage of n_jobs + chunking with a threading backend.

We could investigate, but maybe let's merge this PR first..

@noelano Would you have benchmark results before/after this change?

In terms of benchmark I can't share the data I was using to test but here's a quick illustration with 1m points using memory_profiler:

x = np.random.random(size=1000000)
y = np.random.choice(['a', 'b'], size=1000000)
    
_compute_mi_cd(x, y, 3)

before:

152.3 MiB 0.0 MiB nn.set_params(algorithm='kd_tree')
154.3 MiB 2.0 MiB nn.fit(c)
323.5 MiB 169.2 MiB ind = nn.radius_neighbors(radius=radius, return_distance=False)
331.8 MiB 0.1 MiB m_all = np.array([i.size for i in ind])

after:
157.7 MiB 0.0 MiB nn = KDTree(c)
165.4 MiB 7.6 MiB m_all = nn.query_radius(c, radius, count_only=True, return_distance=False)
165.4 MiB 0.0 MiB m_all = np.array(m_all) - 1.0

Execution time is slight improvement also

x = np.random.random(size=1000000)
y = np.random.choice(['a', 'b'], size=1000000)
    
print(timeit.timeit('mi = _compute_mi_cd(x, y, 3)', number=100, setup="from __main__ import x, y, _compute_mi_cd"))
print(timeit.timeit('mi = _compute_mi_cd(x, y, 3, "new")', number=100, setup="from __main__ import x, y, _compute_mi_cd"))

gives
71.61950279999999
64.1891042

In terms of n_jobs and chunking I hadn't looked into it since the public functions don't expose those anyway. I guess it'd be a separate task to look into adding that as an enhancement

Thank you for the benchmark!

In terms of n_jobs and chunking I hadn't looked into it since the public functions don't expose those anyway. I guess it'd be a separate task to look into adding that as an enhancement

This PR should not consider chunking. It was a passing thought when reviewing the code.

Thank you @noelano and reviewers!