Coverage for cosmolayer/store/_chalcedon/greedy_cluster_split.py: 93%

1"""Dataset splitting strategies for molecular datasets. 

2 

3Vendored from https://github.com/rowansci/chalcedon at commit 

492da3cc5bd6ffb0d397cb49ea556f168d1d38b7e (MIT license, see ``NOTICE`` in 

5this directory), unmodified from upstream. 

6""" 

7 

8from __future__ import annotations 

9 

10from typing import TYPE_CHECKING 

11 

12import numpy as np 

13 

14if TYPE_CHECKING: 

15 from numpy.typing import NDArray 

16 

17 

18def greedy_cluster_split( 

19 cluster_ids: NDArray[np.integer], 

20 fractions: dict[str, float], 

21) -> dict[str, NDArray[np.intp]]: 

22 """Split points into named groups while keeping each cluster intact. 

23 

24 Walks clusters from largest to smallest and drops each one into whichever 

25 split is currently furthest below its target fraction. Because whole 

26 clusters stay together, points in different splits never share a cluster, 

27 which is what keeps the resulting train/val/test sets dissimilar. 

28 

29 The underlying algorithm is Longest Processing Time (LPT) scheduling, from 

30 Graham, R. L. (1969), "Bounds on Multiprocessing Timing Anomalies", SIAM 

31 Journal on Applied Mathematics 17(2):416-429, doi:10.1137/0117039. 

32 

33 Args: 

34 cluster_ids: cluster label per point, shape `(n,)`. 

35 fractions: mapping from split name to target fraction. Values must be 

36 positive and sum to 1.0. Iteration order breaks ties when multiple 

37 splits share the maximum deficit. 

38 

39 Returns: 

40 Mapping from split name to ascending-sorted point indices. 

41 

42 Raises: 

43 ValueError: if `fractions` is empty, contains non-positive values, 

44 or does not sum to 1.0 within 1e-6. 

45 

46 Examples: 

47 >>> import numpy as np 

48 >>> ids = np.array([0, 0, 0, 1, 1, 2, 3]) 

49 >>> result = greedy_cluster_split(ids, {"train": 0.6, "test": 0.4}) 

50 >>> result["train"].tolist() 

51 [0, 1, 2, 5] 

52 >>> result["test"].tolist() 

53 [3, 4, 6] 

54 """ 

55 if not fractions: 

56 raise ValueError("fractions must be non-empty") 

57 if any(value <= 0 for value in fractions.values()): 

58 raise ValueError("all target fractions must be positive") 

59 total = sum(fractions.values()) 

60 if abs(total - 1.0) > 1e-6: 

61 raise ValueError(f"target fractions must sum to 1.0, got {total}") 

62 

63 point_count = len(cluster_ids) 

64 if point_count == 0: 

65 return {name: np.empty(0, dtype=np.intp) for name in fractions} 

66 

67 _, inverse = np.unique(cluster_ids, return_inverse=True) 

68 inverse = inverse.ravel() 

69 sizes = np.bincount(inverse) 

70 order = np.argsort(-sizes, kind="stable") 

71 

72 split_names = list(fractions) 

73 targets = np.array([fractions[name] for name in split_names], dtype=np.float64) 

74 counts = np.zeros(len(split_names), dtype=np.float64) 

75 members_by_split: dict[str, list[NDArray[np.intp]]] = {name: [] for name in split_names} 

76 

77 for cluster_index in order: 

78 deficits = targets - counts / point_count 

79 chosen = int(np.argmax(deficits)) 

80 members = np.flatnonzero(inverse == cluster_index) 

81 members_by_split[split_names[chosen]].append(members) 

82 counts[chosen] += sizes[cluster_index] 

83 

84 return { 

85 name: np.sort(np.concatenate(arrays)) if arrays else np.empty(0, dtype=np.intp) 

86 for name, arrays in members_by_split.items() 

87 }