PERF Call KDTree in mutual_info to reduce memory footprint

[noelano]

What does this implement/fix? Explain your changes.

The mutual_info functions can end up requiring high amounts of memory which makes it particularly tricky to perform multiple MI computations for features in parallel.
The biggest impact is the nn.radius_neighbours call which loads the full nested arrays of all neighbours into memory before getting their sizes.

Since the algorithm is already being set to kd_tree anyway we can explicity call the KDTree class in order to use the query_radius method to get the count of neighbours without having to first store them.

[jnothman]

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:

[thomasjpfan]

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

Thank you for the PR @noelano! LGTM

[rth]

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?

[thomasjpfan]

Suggested change

	for counting nearest neighbors
	for counting nearest neighbors.

[noelano]

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

[thomasjpfan]

May need to rewrap to make this pass linting.

[thomasjpfan]

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.

[rth]

Thank you @noelano and reviewers!