Coverage for src/beamme/geometric_search/scipy.py: 91%

1# The MIT License (MIT) 

2# 

3# Copyright (c) 2018-2026 BeamMe Authors 

4# 

5# Permission is hereby granted, free of charge, to any person obtaining a copy 

6# of this software and associated documentation files (the "Software"), to deal 

7# in the Software without restriction, including without limitation the rights 

8# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 

9# copies of the Software, and to permit persons to whom the Software is 

10# furnished to do so, subject to the following conditions: 

11# 

12# The above copyright notice and this permission notice shall be included in 

13# all copies or substantial portions of the Software. 

14# 

15# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 

16# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 

17# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 

18# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 

19# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 

20# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN 

21# THE SOFTWARE. 

22"""This file defines the interface to the Scipy spatial geometric search 

23functionality.""" 

24 

25from scipy.spatial import KDTree as _KDTree 

26 

27 

28def pairs_to_partner_list(pairs, n_points): 

29 """Convert the pairs to a partner list.""" 

30 # Sort the pairs by the first column 

31 pairs = pairs[pairs[:, 0].argsort()] 

32 

33 n_partners = 0 

34 partner_index_list = [-1 for i in range(n_points)] 

35 for pair in pairs: 

36 pair_partners = [partner_index_list[pair[i]] for i in range(2)] 

37 if pair_partners[0] == -1 and pair_partners[1] == -1: 

38 # We have a new partner index 

39 for i in range(2): 

40 partner_index_list[pair[i]] = n_partners 

41 n_partners += 1 

42 elif pair_partners[0] == -1: 

43 partner_index_list[pair[0]] = pair_partners[1] 

44 elif pair_partners[1] == -1: 

45 partner_index_list[pair[1]] = pair_partners[0] 

46 else: 

47 # Both points already have a partner index. Since we order 

48 # the pairs this one has to be equal to each other or we 

49 # have multiple clusters. 

50 if not pair_partners[0] == pair_partners[1]: 

51 raise ValueError( 

52 "Two points are connected to different partner indices." 

53 ) 

54 return partner_index_list, n_partners 

55 

56 

57def find_close_points_scipy(point_coordinates, tol): 

58 """Call the Scipy implementation of find close_points.""" 

59 kd_tree = _KDTree(point_coordinates) 

60 pairs = kd_tree.query_pairs(r=tol, output_type="ndarray") 

61 return pairs_to_partner_list(pairs, len(point_coordinates))