lib/pluck/src/bif/transpiler.zig
daab053ee43316e1809a84551d573ddd1e5bf3d2
1 const std = @import("std");
2 const Allocator = std.mem.Allocator;
3 const parser = @import("parser.zig");
4 const BayesianNetwork = parser.BayesianNetwork;
5 const Variable = parser.Variable;
6 const Probability = parser.Probability;
7 const CPTEntry = parser.CPTEntry;
8
9 pub const TranspileError = error{
10 VariableNotFound,
11 ProbabilityNotFound,
12 InvalidCPT,
13 OutOfMemory,
14 };
15
16 pub const Config = struct {
17 target_variable: []const u8,
18 include_header: bool = true,
19 network_name: ?[]const u8 = null,
20 };
21
22 pub const Transpiler = struct {
23 allocator: Allocator,
24 bn: *const BayesianNetwork,
25 output: std.ArrayListUnmanaged(u8),
26 config: Config,
27 };
28
29 fn init(allocator: Allocator, bn: *const BayesianNetwork, config: Config) Transpiler {
30 return .{
31 .allocator = allocator,
32 .bn = bn,
33 .output = .empty,
34 .config = config,
35 };
36 }
37
38 fn deinit(transpiler: *Transpiler) void {
39 transpiler.output.deinit(transpiler.allocator);
40 }
41
42 fn write(transpiler: *Transpiler, bytes: []const u8) !void {
43 try transpiler.output.appendSlice(transpiler.allocator, bytes);
44 }
45
46 fn print(transpiler: *Transpiler, comptime fmt: []const u8, args: anytype) !void {
47 try transpiler.output.print(transpiler.allocator, fmt, args);
48 }
49
50 fn writeIndent(transpiler: *Transpiler, level: usize) !void {
51 for (0..level) |_| {
52 try write(transpiler, " ");
53 }
54 }
55
56 fn isBinaryVariable(v: *const Variable) bool {
57 if (v.values.len != 2) return false;
58 const first = v.values[0];
59 const second = v.values[1];
60 return (std.mem.eql(u8, first, "True") and std.mem.eql(u8, second, "False")) or
61 (std.mem.eql(u8, first, "true") and std.mem.eql(u8, second, "false"));
62 }
63
64 fn toLowerFirst(transpiler: *Transpiler, name: []const u8) ![]const u8 {
65 if (name.len == 0) return name;
66
67 const result = try transpiler.allocator.alloc(u8, name.len);
68 @memcpy(result, name);
69 result[0] = std.ascii.toLower(name[0]);
70 return result;
71 }
72
73 fn toConstructorName(transpiler: *Transpiler, type_name: []const u8, value: []const u8) ![]const u8 {
74 const result = try transpiler.allocator.alloc(u8, type_name.len + value.len);
75
76 @memcpy(result[0..type_name.len], type_name);
77 result[0] = std.ascii.toUpper(type_name[0]);
78
79 @memcpy(result[type_name.len..], value);
80 result[type_name.len] = std.ascii.toUpper(value[0]);
81
82 return result;
83 }
84
85 fn writeTypeDecl(transpiler: *Transpiler, v: *const Variable) !void {
86 const type_name = try toLowerFirst(transpiler, v.name);
87 defer transpiler.allocator.free(type_name);
88
89 try print(transpiler, "(define-type {s}", .{type_name});
90 for (v.values) |val| {
91 const constructor = try toConstructorName(transpiler, v.name, val);
92 defer transpiler.allocator.free(constructor);
93 try print(transpiler, " ({s})", .{constructor});
94 }
95 try write(transpiler, ")\n");
96 }
97
98 fn findProbability(transpiler: *Transpiler, var_name: []const u8) ?*const Probability {
99 for (transpiler.bn.probabilities) |*p| {
100 if (std.mem.eql(u8, p.child, var_name)) {
101 return p;
102 }
103 }
104 return null;
105 }
106
107 fn writeDistribution(transpiler: *Transpiler, variable: *const Variable, probs: []const f64) !void {
108 if (isBinaryVariable(variable) and probs.len == 2) {
109 try print(transpiler, "(flip {d})", .{probs[0]});
110 } else {
111 try write(transpiler, "(discrete");
112 for (variable.values, probs) |val, p| {
113 const constructor = try toConstructorName(transpiler, variable.name, val);
114 defer transpiler.allocator.free(constructor);
115 try print(transpiler, " ({s} {d})", .{ constructor, p });
116 }
117 try write(transpiler, ")");
118 }
119 }
120
121 fn findEntryForValues(prob: *const Probability, values: []const []const u8) ?*const CPTEntry {
122 for (prob.entries) |*entry| {
123 if (entry.parent_values.len != values.len) continue;
124
125 var matches = true;
126 for (entry.parent_values, values) |ev, v| {
127 if (!std.mem.eql(u8, ev, v)) {
128 matches = false;
129 break;
130 }
131 }
132 if (matches) return entry;
133 }
134 return null;
135 }
136
137 fn writeConditionalExpr(
138 transpiler: *Transpiler,
139 variable: *const Variable,
140 prob: *const Probability,
141 parent_idx: usize,
142 current_values: *std.ArrayListUnmanaged([]const u8),
143 indent: usize,
144 ) !void {
145 if (parent_idx >= prob.parents.len) {
146 const entry = findEntryForValues(prob, current_values.items) orelse
147 return error.InvalidCPT;
148 try writeDistribution(transpiler, variable, entry.probabilities);
149 return;
150 }
151
152 const parent_name = prob.parents[parent_idx];
153 const parent_var = transpiler.bn.getVariable(parent_name) orelse return error.VariableNotFound;
154
155 if (isBinaryVariable(parent_var)) {
156 const def_name = try toLowerFirst(transpiler, parent_name);
157 defer transpiler.allocator.free(def_name);
158
159 try print(transpiler, "(if {s}\n", .{def_name});
160
161 try current_values.append(transpiler.allocator, "True");
162 try writeIndent(transpiler, indent + 1);
163 try writeConditionalExpr(transpiler, variable, prob, parent_idx + 1, current_values, indent + 1);
164 _ = current_values.pop();
165
166 try write(transpiler, "\n");
167
168 try current_values.append(transpiler.allocator, "False");
169 try writeIndent(transpiler, indent + 1);
170 try writeConditionalExpr(transpiler, variable, prob, parent_idx + 1, current_values, indent + 1);
171 _ = current_values.pop();
172 try write(transpiler, ")");
173 } else {
174 const def_name = try toLowerFirst(transpiler, parent_name);
175 defer transpiler.allocator.free(def_name);
176
177 try print(transpiler, "(case {s} of\n", .{def_name});
178
179 for (parent_var.values, 0..) |val, idx| {
180 const constructor = try toConstructorName(transpiler, parent_name, val);
181 defer transpiler.allocator.free(constructor);
182
183 try writeIndent(transpiler, indent + 1);
184 if (idx > 0) {
185 try write(transpiler, "| ");
186 }
187 try print(transpiler, "{s} => ", .{constructor});
188
189 try current_values.append(transpiler.allocator, val);
190 try writeConditionalExpr(transpiler, variable, prob, parent_idx + 1, current_values, indent + 2);
191 _ = current_values.pop();
192 if (idx + 1 < parent_var.values.len) {
193 try write(transpiler, "\n");
194 }
195 }
196 try write(transpiler, ")");
197 }
198 }
199
200 fn writeDefinition(transpiler: *Transpiler, variable: *const Variable, prob: *const Probability) !void {
201 const def_name = try toLowerFirst(transpiler, variable.name);
202 defer transpiler.allocator.free(def_name);
203
204 try print(transpiler, "(define {s} ", .{def_name});
205 if (prob.parents.len == 0) {
206 try writeDistribution(transpiler, variable, prob.entries[0].probabilities);
207 } else {
208 var current_values: std.ArrayListUnmanaged([]const u8) = .empty;
209 defer current_values.deinit(transpiler.allocator);
210
211 try writeConditionalExpr(transpiler, variable, prob, 0, ¤t_values, 0);
212 }
213 try write(transpiler, ")\n");
214 }
215
216 fn topologicalSort(transpiler: *Transpiler) ![][]const u8 {
217 var result: std.ArrayListUnmanaged([]const u8) = .empty;
218 errdefer result.deinit(transpiler.allocator);
219
220 var in_degree = std.StringHashMap(usize).init(transpiler.allocator);
221 defer in_degree.deinit();
222
223 for (transpiler.bn.variables) |v| {
224 try in_degree.put(v.name, 0);
225 }
226
227 for (transpiler.bn.probabilities) |p| {
228 const current = in_degree.get(p.child) orelse 0;
229 try in_degree.put(p.child, current + p.parents.len);
230 }
231
232 var queue: std.ArrayListUnmanaged([]const u8) = .empty;
233 defer queue.deinit(transpiler.allocator);
234
235 for (transpiler.bn.variables) |v| {
236 if ((in_degree.get(v.name) orelse 0) == 0) {
237 try queue.append(transpiler.allocator, v.name);
238 }
239 }
240
241 while (queue.items.len > 0) {
242 const node = queue.orderedRemove(0);
243 try result.append(transpiler.allocator, node);
244
245 for (transpiler.bn.probabilities) |p| {
246 for (p.parents) |parent| {
247 if (std.mem.eql(u8, parent, node)) {
248 const deg = in_degree.get(p.child) orelse 1;
249 if (deg > 0) {
250 try in_degree.put(p.child, deg - 1);
251 if (deg == 1) {
252 try queue.append(transpiler.allocator, p.child);
253 }
254 }
255 break;
256 }
257 }
258 }
259 }
260
261 return result.toOwnedSlice(transpiler.allocator);
262 }
263
264 fn render(transpiler: *Transpiler) ![]const u8 {
265 if (transpiler.config.include_header) {
266 if (transpiler.config.network_name) |name| {
267 try print(transpiler, ";; {s} Bayesian Network\n", .{name});
268 }
269 try print(transpiler, ";; Variables: {d}\n", .{transpiler.bn.variables.len});
270 try write(transpiler, ";; Auto-generated from BIF format\n\n");
271 }
272
273 for (transpiler.bn.variables) |*v| {
274 if (!isBinaryVariable(v)) {
275 try writeTypeDecl(transpiler, v);
276 }
277 }
278
279 const ordered = try topologicalSort(transpiler);
280 defer transpiler.allocator.free(ordered);
281
282 for (ordered) |var_name| {
283 const variable = transpiler.bn.getVariable(var_name) orelse return error.VariableNotFound;
284 const prob = findProbability(transpiler, var_name) orelse return error.ProbabilityNotFound;
285 try writeDefinition(transpiler, variable, prob);
286 try write(transpiler, "\n");
287 }
288
289 const target_lower = try toLowerFirst(transpiler, transpiler.config.target_variable);
290 defer transpiler.allocator.free(target_lower);
291
292 try print(transpiler, ";; Query: Marginal probability of {s}\n", .{transpiler.config.target_variable});
293 try print(transpiler, "(query (Marginal {s}))\n", .{target_lower});
294
295 return try transpiler.output.toOwnedSlice(transpiler.allocator);
296 }
297
298 pub fn transpile(
299 allocator: Allocator,
300 bn: *const BayesianNetwork,
301 config: Config,
302 ) TranspileError![]const u8 {
303 var transpiler = init(allocator, bn, config);
304 defer deinit(&transpiler);
305 return render(&transpiler);
306 }
307
308 pub fn transpileFromBif(
309 allocator: Allocator,
310 bif_source: []const u8,
311 target_variable: []const u8,
312 network_name: ?[]const u8,
313 ) ![]const u8 {
314 var bn = try parser.parseString(allocator, bif_source);
315 defer bn.deinit();
316
317 return transpile(allocator, &bn, .{
318 .target_variable = target_variable,
319 .network_name = network_name,
320 });
321 }
322
323 test "transpile simple binary network" {
324 const allocator = std.testing.allocator;
325 const source =
326 \\variable A {
327 \\ type discrete [ 2 ] { True, False };
328 \\}
329 \\probability ( A ) {
330 \\ table 0.3, 0.7;
331 \\}
332 ;
333
334 var bn = try parser.parseString(allocator, source);
335 defer bn.deinit();
336
337 const result = try transpile(allocator, &bn, .{
338 .target_variable = "A",
339 .include_header = false,
340 });
341 defer allocator.free(result);
342
343 try std.testing.expect(std.mem.indexOf(u8, result, "(define a (flip 0.3))") != null);
344 try std.testing.expect(std.mem.indexOf(u8, result, "(query (Marginal a))") != null);
345 }
346
347 test "transpile categorical network" {
348 const allocator = std.testing.allocator;
349 const source =
350 \\variable Weather {
351 \\ type discrete [ 3 ] { sunny, cloudy, rainy };
352 \\}
353 \\probability ( Weather ) {
354 \\ table 0.5, 0.3, 0.2;
355 \\}
356 ;
357
358 var bn = try parser.parseString(allocator, source);
359 defer bn.deinit();
360
361 const result = try transpile(allocator, &bn, .{
362 .target_variable = "Weather",
363 .include_header = false,
364 });
365 defer allocator.free(result);
366
367 try std.testing.expect(std.mem.indexOf(u8, result, "(define-type weather (WeatherSunny) (WeatherCloudy) (WeatherRainy))") != null);
368 try std.testing.expect(std.mem.indexOf(u8, result, "(discrete (WeatherSunny 0.5) (WeatherCloudy 0.3) (WeatherRainy 0.2))") != null);
369 }
370
371 test "transpile conditional network as S-expression conditionals" {
372 const allocator = std.testing.allocator;
373 const source =
374 \\variable A {
375 \\ type discrete [ 2 ] { True, False };
376 \\}
377 \\variable B {
378 \\ type discrete [ 2 ] { True, False };
379 \\}
380 \\probability ( A ) {
381 \\ table 0.3, 0.7;
382 \\}
383 \\probability ( B | A ) {
384 \\ (True) 0.9, 0.1;
385 \\ (False) 0.2, 0.8;
386 \\}
387 ;
388
389 var bn = try parser.parseString(allocator, source);
390 defer bn.deinit();
391
392 const result = try transpile(allocator, &bn, .{
393 .target_variable = "B",
394 .include_header = false,
395 });
396 defer allocator.free(result);
397
398 try std.testing.expect(std.mem.indexOf(u8, result, "(define a (flip 0.3))") != null);
399 try std.testing.expect(std.mem.indexOf(u8, result, "(define b (if a") != null);
400 try std.testing.expect(std.mem.indexOf(u8, result, "(flip 0.9)") != null);
401 try std.testing.expect(std.mem.indexOf(u8, result, "(flip 0.2)") != null);
402 try std.testing.expect(std.mem.indexOf(u8, result, "(query (Marginal b))") != null);
403 }