main.bend

mlbend-split@0.1.1.0 · 203 lines · 27 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      {==}
180
181# The rows of the fold f: the ones that test.
182def 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.
190def 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.
194def 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.
202def train_of(fs: List<&2, Nat>, f: Nat) -> List<&2, Nat>:
203  train_of_go(fs, f, 0n, Nil{})