1 # mlbend-split: a seeded, deterministic split of the rows, stratified by label. 2 # 3 # import mlbend-split@0.1.0.0/main.bend as Split 4 # 5 # The rows of a class are shuffled and then dealt out over the folds, so 6 # every fold holds the same proportion of every class — the rule 7 # scikit-learn's StratifiedKFold uses. The shuffle is bend-kit-random's 8 # xoshiro128**, seeded: the same seed gives the same split, and another 9 # seed gives another one. Nothing here generates numbers of its own. 10 # 11 # The labels are class indices, 0 to the number of classes. 12 # 13 # What is proved: every row of the dataset is placed in exactly one fold 14 # (`split_covers`) — the fold list holds one entry per row, so the folds 15 # are a partition of the rows: the test rows of a fold and its training 16 # rows are disjoint, and together they hold every row. What is tested: 17 # that the sizes and the classes of the folds agree with scikit-learn's 18 # (reference/test_split.py). 19 20 import Base 21 import bend-kit-random@0.1.0.1/random.bend as Rand 22 23 # The split of the rows: the training rows and the test rows, as their 24 # indices in the dataset. 25 26 type Split is Data: 27 Split{train: List<&2, Nat>, test: List<&2, Nat>} 28 29 # LAW (split_covers): every row of the dataset is placed in a fold — the 30 # fold list holds one entry per row, so the folds are a partition of the 31 # rows: the test rows of a fold and its training rows are disjoint, and 32 # together they hold every row. 33 law split_covers: 34 for labels: List<&2, Nat> 35 for +orders: List<&2, List<&2, Nat>> 36 for +k: Nat 37 for +i: Nat 38 {count_rows(folds_from(labels, orders, k, i)) == List.length(&2, Nat, labels) : Nat} 39 40 # The rows of one class, with their indices. 41 def rows_of_class(labels: List<&2, Nat>, +c: Nat, +i: Nat, +acc: List<&2, Nat>) -> List<&2, Nat>: 42 match labels: 43 case Nil{}: 44 acc 45 case Con{h, t}: 46 rows_of_class(t, c, Nat.add(i, 1n), 47 Bool.pick(List<&2, Nat>, Nat.is_eq(h, c), List.append(&2, Nat, acc, Con{i, Nil{}}), acc)) 48 49 # The list of a pair: the list, and the generator that comes after it. 50 def shuffled_list(pair: List<&2, Nat> & Rand.Rng) -> List<&2, Nat>: 51 (xs, r) = pair 52 xs 53 54 # The rows of one class, shuffled by the generator. 55 def shuffled_class(labels: List<&2, Nat>, +c: Nat, +seed: U32) -> List<&2, Nat>: 56 shuffled_list(Rand.shuffle(~Nat, rows_of_class(labels, c, 0n, Nil{}), 57 Rand.seed(U32.add(seed, U32.from_nat(c))))) 58 59 # The shuffled rows of every class, one class after the other. 60 def class_orders_go(+labels: List<&2, Nat>, +seed: U32, remaining: Nat, +next: Nat, 61 +acc: List<&2, List<&2, Nat>>) -> List<&2, List<&2, Nat>>: 62 match remaining: 63 case 0n: 64 acc 65 case 1n+rest: 66 class_orders_go(labels, seed, rest, Nat.add(next, 1n), 67 List.append(&2, List<&2, Nat>, acc, Con{shuffled_class(labels, next, seed), Nil{}})) 68 69 # The shuffled rows of every class, one class after the other. 70 def class_orders(labels: List<&2, Nat>, seed: U32, c: Nat) -> List<&2, List<&2, Nat>>: 71 class_orders_go(labels, seed, c, 0n, Nil{}) 72 73 # The largest label of the labels. 74 def largest_of(labels: List<&2, Nat>, +acc: Nat) -> Nat: 75 match labels: 76 case Nil{}: 77 acc 78 case Con{+h, t}: 79 largest_of(t, Bool.pick(Nat, Nat.is_gt(h, acc), h, acc)) 80 81 # The number of classes of the labels: one more than the largest label. 82 def classes_of(labels: List<&2, Nat>) -> Nat: 83 Nat.add(1n, largest_of(labels, 0n)) 84 85 # The number of rows a list of folds holds. 86 def count_rows(fs: List<&2, Nat>) -> Nat: 87 match fs: 88 case Nil{}: 89 0n 90 case Con{h, t}: 91 (1n + count_rows(t) : Nat) 92 93 # The place of a row among the shuffled rows of its class. 94 def position_go(order: List<&2, Nat>, +i: Nat, +at: Nat) -> Nat: 95 match order: 96 case Nil{}: 97 at 98 case Con{h, t}: 99 Bool.pick(Nat, Nat.is_eq(h, i), at, position_go(t, i, Nat.add(at, 1n))) 100 101 # The place of a row among the shuffled rows of its class. 102 def position_of(order: List<&2, Nat>, i: Nat) -> Nat: 103 position_go(order, i, 0n) 104 105 # The fold of a row, given the shuffled rows of its class. 106 def fold_in_of(order: List<&2, Nat>, k: Nat, i: Nat) -> Nat: 107 Nat.mod(position_of(order, i), k) 108 109 # The first of the shuffled rows of every class. 110 def order_head(orders: List<&2, List<&2, Nat>>) -> List<&2, Nat>: 111 match orders: 112 case Nil{}: 113 Nil{} 114 case Con{h, t}: 115 h 116 117 # The orders of every class but the first. 118 def order_tail(+orders: List<&2, List<&2, Nat>>) -> List<&2, List<&2, Nat>>: 119 match orders: 120 case Nil{}: 121 Nil{} 122 case Con{h, t}: 123 t 124 125 # The shuffled rows of the class of a label: the n-th of the orders. 126 def order_at(n: Nat, +orders: List<&2, List<&2, Nat>>) -> List<&2, Nat>: 127 match n: 128 case 0n: 129 order_head(orders) 130 case 1n+rest: 131 order_at(rest, order_tail(orders)) 132 133 # The fold of a row: its place among the shuffled rows of its class, 134 # dealt out over the folds in turn. 135 def fold_in(orders: List<&2, List<&2, Nat>>, k: Nat, label: Nat, i: Nat) -> Nat: 136 fold_in_of(order_at(label, orders), k, i) 137 138 # The fold of every row from the row i on, given the shuffled rows of 139 # every class. 140 def folds_from(labels: List<&2, Nat>, +orders: List<&2, List<&2, Nat>>, +k: Nat, +i: Nat) -> 141 List<&2, Nat>: 142 match labels: 143 case Nil{}: 144 Nil{} 145 case Con{h, t}: 146 Con{fold_in(orders, k, h, i), folds_from(t, orders, k, Nat.add(i, 1n))} 147 148 # The fold of every row of the dataset, in the order of the rows. 149 def folds(+labels: List<&2, Nat>, seed: U32, k: Nat) -> List<&2, Nat>: 150 folds_from(labels, class_orders(labels, seed, classes_of(labels)), k, 0n) 151 152 # The row, in the fold. 153 def test_row(fold: Nat, f: Nat, +test: List<&2, Nat>, i: Nat) -> List<&2, Nat>: 154 Bool.pick(List<&2, Nat>, Nat.is_eq(fold, f), List.append(&2, Nat, test, Con{i, Nil{}}), test) 155 156 # The row, out of the fold. 157 def train_row(fold: Nat, f: Nat, +train: List<&2, Nat>, i: Nat) -> List<&2, Nat>: 158 Bool.pick(List<&2, Nat>, Nat.is_eq(fold, f), train, List.append(&2, Nat, train, Con{i, Nil{}})) 159 160 # The rows of one fold: the ones that test and the ones that train. 161 def split_of(fs: List<&2, Nat>, +f: Nat, +i: Nat, +test: List<&2, Nat>, +train: List<&2, Nat>) -> Split: 162 match fs: 163 case Nil{}: 164 Split{train, test} 165 case Con{+h, t}: 166 split_of(t, f, Nat.add(i, 1n), test_row(h, f, test, i), train_row(h, f, train, i)) 167 168 # The rows of one fold: the ones that test and the ones that train. 169 def split(labels: List<&2, Nat>, seed: U32, k: Nat, f: Nat) -> Split: 170 split_of(folds(labels, seed, k), f, 0n, Nil{}, Nil{}) 171 172 # every row of the dataset is placed in a fold 173 def split_covers(labels, orders, k, i): 174 match labels: 175 case Nil{}: 176 {==} 177 case Con{h, +t}: 178 %split_covers(t, orders, k, Nat.add(i, 1n)) : {1n+count_rows(folds_from(t, orders, k, Nat.add(i, 1n))) == 1n+_ : Nat} 179 {==} 180 181 # The rows of the fold f: the ones that test. 182 def test_of_go(+fs: List<&2, Nat>, +f: Nat, +i: Nat, +acc: List<&2, Nat>) -> List<&2, Nat>: 183 match fs: 184 case Nil{}: 185 acc 186 case Con{+h, t}: 187 test_of_go(t, f, Nat.add(i, 1n), test_row(h, f, acc, i)) 188 189 # The rows of the fold f: the ones that test. 190 def test_of(fs: List<&2, Nat>, f: Nat) -> List<&2, Nat>: 191 test_of_go(fs, f, 0n, Nil{}) 192 193 # The rows of the fold f: the ones that train. 194 def train_of_go(+fs: List<&2, Nat>, +f: Nat, +i: Nat, +acc: List<&2, Nat>) -> List<&2, Nat>: 195 match fs: 196 case Nil{}: 197 acc 198 case Con{+h, t}: 199 train_of_go(t, f, Nat.add(i, 1n), train_row(h, f, acc, i)) 200 201 # The rows of the fold f: the ones that train. 202 def train_of(fs: List<&2, Nat>, f: Nat) -> List<&2, Nat>: 203 train_of_go(fs, f, 0n, Nil{})