Coverage for cosmolayer/store/subsampling.py: 100%

1"""Shrink a SegmentStore to a molecule subset without leaking splits. 

2 

3``SegmentStore.subsample`` shrinks a store to ``num_molecules`` while 

4keeping every molecule's ``split`` assignment (see ``splitting.py``) fixed: 

5molecules are only ever dropped from within a split, never moved to 

6another one, so a test molecule can never end up in train/val just 

7because the dataset got smaller. 

8""" 

9 

10from __future__ import annotations 

11 

12from typing import TYPE_CHECKING 

13 

14import numpy as np 

15from numpy.typing import NDArray 

16 

17if TYPE_CHECKING: 

18 from .segments import SegmentStore 

19 

20 

21def apportion_counts(sizes: NDArray[np.int64], total: int) -> NDArray[np.int64]: 

22 """Apportion ``total`` items across buckets, proportional to ``sizes``. 

23 

24 Each bucket's share is fixed via largest-remainder (Hamilton) 

25 apportionment: raw proportional shares are floored, then the leftover 

26 units are handed out one at a time, in order of largest fractional 

27 remainder, to whichever bucket still has room -- a bucket's own 

28 ``sizes[i]`` doubles as its capacity, so none is ever asked for more 

29 than it has. 

30 

31 Parameters 

32 ---------- 

33 sizes : np.ndarray, shape (k,) 

34 Bucket sizes (also used as capacities). 

35 total : int 

36 Total number of items to apportion; must not exceed ``sizes.sum()``. 

37 

38 Returns 

39 ------- 

40 np.ndarray, shape (k,) 

41 Non-negative integer counts, each clamped to its bucket's size, 

42 summing to ``total``. 

43 """ 

44 n_total = int(sizes.sum()) 

45 raw_counts = total * sizes / n_total 

46 counts = np.floor(raw_counts).astype(np.int64) 

47 remainder = total - int(counts.sum()) 

48 if remainder > 0: 

49 fractional = raw_counts - counts 

50 headroom = sizes - counts 

51 top_up_order = np.argsort(-fractional, kind="stable") 

52 for bucket in top_up_order: 

53 if remainder == 0: 

54 break 

55 if headroom[bucket] > 0: 

56 counts[bucket] += 1 

57 headroom[bucket] -= 1 

58 remainder -= 1 

59 return np.minimum(counts, sizes) 

60 

61 

62def restrict_to_molecules( 

63 store: SegmentStore, selected: NDArray[np.int64] 

64) -> SegmentStore: 

65 """Return a new, unsaved store restricted to ``selected`` molecules. 

66 

67 Segment-indexed arrays (``data``, each ``averaged_sigmas`` scheme) carry 

68 no molecule identity of their own, so they're sliced by a segment mask 

69 derived from ``selected``. ``atom_indices`` does carry atom identity 

70 (global, dataset-wide), so it's additionally rebased onto a new, 

71 compacted index space covering only the kept molecules' atoms; 

72 ``atoms_df`` is sliced by the same atom mask, preserving row order. 

73 

74 Parameters 

75 ---------- 

76 store : SegmentStore 

77 Store to restrict. 

78 selected : np.ndarray, shape (k,) 

79 Ascending-sorted row indices into ``store.molecules_df`` to keep. 

80 

81 Returns 

82 ------- 

83 SegmentStore 

84 A new store with the same ``storage_dir`` as ``store`` (a 

85 placeholder -- pass a real directory to ``save`` to persist it), 

86 holding only ``selected``'s molecules, segments, and atoms. 

87 """ 

88 # Deferred: segments.py imports this module at load time (for 

89 # SegmentStore.subsample), so importing SegmentStore back at module 

90 # level here would be circular. 

91 from .segments import SegmentStore, StoreMetadata # noqa: PLC0415 

92 

93 molecules_df = store.molecules_df 

94 n_mols_total = len(molecules_df) 

95 n_segs_total = len(store.data) 

96 full_segment_offsets = molecules_df["segment_offsets"].to_numpy().astype("int64") 

97 full_atom_offsets = molecules_df["atom_offsets"].to_numpy().astype("int64") 

98 full_num_atoms = molecules_df["num_atoms"].to_numpy().astype("int64") 

99 

100 segment_counts = np.diff(np.append(full_segment_offsets, n_segs_total)) 

101 segment_molecule = np.repeat(np.arange(n_mols_total), segment_counts) 

102 atom_molecule = np.repeat(np.arange(n_mols_total), full_num_atoms) 

103 

104 is_selected = np.zeros(n_mols_total, dtype=bool) 

105 is_selected[selected] = True 

106 segment_mask = is_selected[segment_molecule] 

107 atom_mask = is_selected[atom_molecule] 

108 

109 new_atom_offset_by_molecule = np.zeros(n_mols_total, dtype=np.int64) 

110 new_atom_offset_by_molecule[selected] = np.concatenate( 

111 [[0], np.cumsum(full_num_atoms[selected])[:-1]] 

112 ) 

113 

114 kept_molecule = segment_molecule[segment_mask] 

115 data = np.asarray(store.data)[segment_mask] 

116 atom_indices = ( 

117 np.asarray(store.atom_indices)[segment_mask] 

118 - full_atom_offsets[kept_molecule] 

119 + new_atom_offset_by_molecule[kept_molecule] 

120 ) 

121 atoms_df = store.atoms_df.iloc[atom_mask].reset_index(drop=True) 

122 averaged_sigmas = { 

123 name: np.asarray(arr)[segment_mask] 

124 for name, arr in store.averaged_sigmas.items() 

125 } 

126 

127 new_segment_offsets = np.concatenate( 

128 [[0], np.cumsum(segment_counts[selected])[:-1]] 

129 ) 

130 new_molecules_df = molecules_df.iloc[selected].copy() 

131 new_molecules_df["segment_offsets"] = new_segment_offsets 

132 new_molecules_df["atom_offsets"] = new_atom_offset_by_molecule[selected] 

133 

134 metadata = StoreMetadata( 

135 num_molecules=len(selected), 

136 num_cosmo_parse_failures=0, 

137 schemes=dict(store.metadata.schemes), 

138 ) 

139 return SegmentStore( 

140 store.storage_dir, 

141 data, 

142 atom_indices, 

143 new_molecules_df, 

144 atoms_df, 

145 metadata, 

146 averaged_sigmas, 

147 )