main.bend
mlbend-split@0.1.0.0 · 179 lines · 0 definitions
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
20import Base
21import 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
26type 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.
33law 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.
41def 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.
50def 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.
55def 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.
60def 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.
70def 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.
74def 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.
82def 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.
86def 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.
94def 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.
102def 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.
106def 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.
110def 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.
118def 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.
126def 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.
135def 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.
140def 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.
149def 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.
153def 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.
157def 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.
161def 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.
169def 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
173def 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 {==}