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

1"""Cluster-preserving train/val/test splitting of molecules. 

2 

3Used by ``SegmentStore.assign_splits`` to partition stored molecules into 

4named splits (e.g. ``"train"``/``"val"``/``"test"``) without letting 

5near-duplicate structures -- as identified by ``cluster_id`` (see 

6``clustering.py``) -- leak across splits. 

7""" 

8 

9from collections.abc import Mapping 

10 

11import numpy as np 

12from numpy.typing import NDArray 

13 

14from cosmolayer.store._chalcedon.greedy_cluster_split import ( 

15 greedy_cluster_split as _chalcedon_greedy_cluster_split, 

16) 

17 

18 

19def greedy_cluster_split( 

20 cluster_ids: NDArray[np.int64], fractions: Mapping[str, float] 

21) -> NDArray[np.str_]: 

22 """Assign each molecule a split name, keeping every cluster intact. 

23 

24 Delegates to chalcedon's greedy LPT-scheduling split (vendored in 

25 ``cosmolayer.store._chalcedon``): clusters are assigned, largest first, 

26 to whichever split is currently furthest below its target fraction, so 

27 a whole cluster always lands in a single split. 

28 

29 Parameters 

30 ---------- 

31 cluster_ids : np.ndarray, shape (n,) 

32 Cluster id per molecule, as produced by ``clustering.butina_cluster``. 

33 fractions : Mapping[str, float] 

34 Target fraction per split name, e.g. ``{"train": 0.8, "val": 0.1, 

35 "test": 0.1}`` or ``{"train": 0.8, "test": 0.2}``. Values must be 

36 positive and sum to 1.0. 

37 

38 Returns 

39 ------- 

40 np.ndarray, shape (n,) 

41 Split name assigned to each molecule, in ``cluster_ids`` order. 

42 

43 Raises 

44 ------ 

45 ValueError 

46 If ``fractions`` is empty, contains non-positive values, or doesn't 

47 sum to 1.0 within 1e-6. 

48 """ 

49 labels = np.empty(len(cluster_ids), dtype=object) 

50 splits = _chalcedon_greedy_cluster_split(cluster_ids, dict(fractions)) 

51 for name, indices in splits.items(): 

52 labels[indices] = name 

53 return labels.astype(str)