Randomly generate simply-typed lambda calculus expressions using Boltzmann samplers

gen_boltzmann.pl

:- ensure_loaded(library(apply)).
:- ensure_loaded(library(random)).

:- set_prolog_flag(optimise, true).
:- set_prolog_flag(optimise_unify, true).

% Randomly generate simply-typed lambda calculus terms using Boltzmann
% samplers.
%
% Based on the paper Random generation of closed simply-typed λ-terms: a
% synergy between logic programming and Boltzmann samplers
%
% https://arxiv.org/pdf/1612.07682.pdf
%
% An application of the algorithm shown here:
% https://github.com/pedrominicz/sugar
%
% Related:
% https://github.com/fredokun/arbogen
% https://github.com/maciej-bendkowski/boltzmann-brain
% https://github.com/maciej-bendkowski/lambda-sampler

% Note that there are multiple definitions for the size of a lambda term. In
% particular, this definition is different from the one in the paper A Hiking
% Trip Through the Orders of Magnitude: Deriving Efficient Generators for
% Closed Simply-Typed Lambda Terms and Normal Forms.
%
% Also note that there are two definitions for the size of a lambda term in the
% reference paper. This is the definition of size used for deriving the
% generating function (and thus the Boltzmann sampler) for plain lambda terms.
size(z, 0).
size(s(X), S) :- size(X, S0), S is S0 + 1.
size(l(X), S) :- size(X, S0), S is S0 + 1.
size(a(X, Y), S) :- size(X, S0), size(Y, S1), S is S0 + S1 + 2.

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%%  The Boltzmann sampler  %%
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

min_size(120).
boltzmann_variable(R) :- R < 0.3569605561766718.
boltzmann_variable_zero(R) :- R < 0.7044187641738427.
boltzmann_lambda(R) :- R < 0.6525417920028290.

next(R, Size1, Size2) :-
  random(R),
  Size2 is Size1 + 1.

random_lambda(X, A) :-
  random(R),
  random_lambda(X, A, [], R, 0, Size),
  min_size(MinSize),
  Size >= MinSize,
  !.
random_lambda(X, A) :- random_lambda(X, A).

random_lambda(X, A, Ctx, R) -->
  { boltzmann_variable(R), !, random(NewR) },
  random_variable(X, A, Ctx, NewR).
random_lambda(l(X), A -> B, Ctx, R) -->
  { boltzmann_lambda(R), ! },
  next(NewR),
  random_lambda(X, B, [A|Ctx], NewR).
random_lambda(a(X, Y), B, Ctx, _) -->
  next(R1),
  random_lambda(X, A -> B, Ctx, R1),
  next(R2),
  random_lambda(Y, A, Ctx, R2).

random_variable(z, A0, [A|_], R) -->
  { boltzmann_variable_zero(R), !, unify_with_occurs_check(A0, A) }.
random_variable(s(X), A, [_|Ctx], _) -->
  next(R),
  random_variable(X, A, Ctx, R).

multithreaded_random_lambda(X, A) :-
  prolog_flag(cpu_count, MaxThreads0),
  MaxThreads is MaxThreads0 - 1,
  Goal = random_lambda(X, A),
  length(Goals, MaxThreads),
  maplist(=(Goal), Goals),
  first_solution([X, A], Goals, []).

convert(X, NewX) :- convert(X, NewX, []).

convert(z, v(X), [X|_]).
convert(s(X), NewX, [_|Ctx]) :- convert(X, NewX, Ctx).
convert(l(Y), l(X, NewY), Ctx) :- convert(Y, NewY, [X|Ctx]).
convert(a(X, Y), a(NewX, NewY), Ctx) :-
  convert(X, NewX, Ctx),
  convert(Y, NewY, Ctx).

% Pretty print a term.
pretty(X0) :-
  convert(X0, X),
  numbervars(X, 0, _),
  pretty(X, Xs, []),
  maplist(write, Xs),
  nl.

pretty(v('$VAR'(I))) --> [x, I].
pretty(l('$VAR'(I), X)) --> ['(λ ', x, I, ', '], pretty(X), [')'].
pretty(a(X, Y)) --> ['('], pretty(X), [' '], pretty(Y), [')'].

% Pretty print a type.
pretty_type(A) :-
  numbervars(A, 0, _),
  pretty_type(A, As, []),
  maplist(write, As),
  nl.

pretty_type(A -> B) -->
  ['('], pretty_type(A), !, [' → '], pretty_type(B), [')'].
pretty_type(A) --> [A].

main :-
  multithreaded_random_lambda(X, A),
  pretty(X),
  pretty_type(A).

gen_boltzmann.py

#!/usr/bin/env python3

from decimal import Decimal, getcontext

getcontext().prec = 100

d1 = Decimal(1)
d2 = Decimal(2)
d3 = Decimal(3)
d4 = Decimal(4)
d5 = Decimal(5)
d7 = Decimal(7)

d3_div_d2 = d3 / d2

def A0(z):
    #
    #         sqrt(1 - 3z - z^2 - z^3)
    # 1 - z - ------------------------
    #                sqrt(1 - z)
    # --------------------------------
    #               2*z^2
    #
    z = Decimal(z)

    tmp1 = d1 - z

    tmp2 = (d1 - d3 * z - z ** d2 - z ** d3).sqrt()
    tmp3 = (d1 - z).sqrt()
    tmp4 = tmp2 / tmp3

    dividend = tmp1 - tmp4

    divisor = d2 * (z ** d2)

    return dividend / divisor

# First derivative of `A0`.
def A1(z):
    z = Decimal(z)
    d1_sub_z = d1 - z

    aux = (d1 - (z ** d3) - (z ** d2) - (3 * z)).sqrt()

    # tmp3 = ((1-z)^(3/2)*z - 2*(1-z)^(3/2)) * sqrt(1 - z^3 - z^2 - 3*z)
    tmp1 = (d1_sub_z ** d3_div_d2) * z
    tmp2 = d2 * (d1_sub_z ** d3_div_d2)
    tmp3 = (tmp1 - tmp2) * aux

    # tmp8 = z^4 + z^3 + 5*z^2 - 7*z+2
    tmp4 = z ** d4
    tmp5 = z ** d3
    tmp6 = d5 * (z ** d2)
    tmp7 = d7 * z
    tmp8 = tmp4 + tmp5 + tmp6 - tmp7 + d2

    dividend = tmp3 + tmp8

    # divisor = 2*(1-z)^(3/2) * z^3 * sqrt(-z^3-z^2-3*z+1)
    tmp9 = d2 * (d1_sub_z ** d3_div_d2)
    tmp10 = z ** d3 # `tmp5 = tmp10`, but repeating is probably less confusing.
    divisor = tmp9 * tmp10 * aux

    return dividend / divisor

# Expected size.
def E(x):
    x = Decimal(x)
    return x * (A1(x) / A0(x))

# https://arxiv.org/pdf/1506.02367.pdf
rho = Decimal('0.2955977425220839')

def parameter(size):
    size = Decimal(size)

    divisor = d2
    guess = rho / divisor

    while True:
        expected_size = E(guess)

        if abs(size - expected_size) < Decimal('0.0001'):
            return guess

        divisor *= d2
        guess += rho / divisor if size > expected_size else -rho / divisor

def probability_variable(x):
    x = Decimal(x)
    return (d1 / (d1 - x)) / A0(x)

def probability_variable_zero(x):
    x = Decimal(x)
    return d1 - x

def probability_lambda(x):
    x = Decimal(x)
    return x

size = 120

# The numbers in `https://arxiv.org/pdf/1612.07682.pdf` are:
#
#   x = Decimal('0.29558095907')
#   print('min_size(120).')
x = parameter(size)
print(f'min_size({size}).')

p1 = probability_variable(x)
print(f'boltzmann_variable(R) :- R < {p1:.16}.')

p2 = probability_variable_zero(x)
print(f'boltzmann_variable_zero(R) :- R < {p2:.16}.')

p3 = probability_variable(x) + probability_lambda(x)
print(f'boltzmann_lambda(R) :- R < {p3:.16}.')