Check if custom functions give desired outcome

Hi all,

I have a number of functions written from Python to custom functions in Stan. Now I want to test whether they give the same output before I continue.
How can I test the functions in a simple Stan model?

This is my implementation

functions {

    int calculate_K(int a, int b, int m) {
        return to_int(floor((a + b) / m));
    }

    int calculate_R(int a, int b, int m, int K) {
        return to_int(a + b - K * m);
    }

    real p_a_wins_R_is_zero_and_K_is_even(int x, int a, real p_a, real p_b, int K, int m) {
        int binom_coeff_a;
        int binom_coeff_b;
        int failed_serves_a;
        int failed_serves_b;

        binom_coeff_a = choose(to_int(K / 2 * m), x);
        binom_coeff_b = choose(to_int(K / 2 * m - 1), a - x - 1);
        failed_serves_a = K / 2 * m - x;
        failed_serves_b = K / 2 * m - a + x;

        return (binom_coeff_a * p_a ^ x) * ((1 - p_a) ^ failed_serves_a) *
               (binom_coeff_b * (1 - p_b) ^ (a - x)) * (p_b ^ failed_serves_b);
    }

    real p_b_wins_R_is_zero_and_K_is_even(int x, int a, real p_a, real p_b, int K, int m) {
        int binom_coeff_a;
        int binom_coeff_b;
        int failed_serves_a;
        int failed_serves_b;

        binom_coeff_a = choose(K / 2 * m, x);
        binom_coeff_b = choose(K / 2 * m - 1, a - x);
        failed_serves_a = K / 2 * m - x;
        failed_serves_b = K / 2 * m - a + x;

        return (binom_coeff_a * p_a ^ x) * ((1 - p_a) ^ failed_serves_a) *
               (binom_coeff_b * (1 - p_b) ^ (a - x)) * (p_b ^ failed_serves_b);
    }

    real p_a_wins_R_is_zero_and_K_is_odd(int x, int a, real p_a, real p_b, int K, int m) {
        int binom_coeff_a;
        int binom_coeff_b;
        int failed_serves_a;
        int failed_serves_b;

        binom_coeff_a = choose(to_int(ceil(K / 2) * m - 1), x - 1);
        binom_coeff_b = choose(to_int(floor(K / 2) * m), a - x);
        failed_serves_a = to_int(ceil(K / 2) * m - x);
        failed_serves_b = to_int(floor(K / 2) * m - a + x);

        return (binom_coeff_a * p_a ^ x) * ((1 - p_a) ^ failed_serves_a) *
               (binom_coeff_b * (1 - p_b) ^ (a - x)) * (p_b ^ failed_serves_b);
    }

    real p_b_wins_R_is_zero_and_K_is_odd(int x, int a, real p_a, real p_b, int K, int m) {
        int binom_coeff_a;
        int binom_coeff_b;
        int failed_serves_a;
        int failed_serves_b;

        binom_coeff_a = choose(to_int(ceil(K / 2) * m - 1), x);
        binom_coeff_b = choose(to_int(floor(K / 2) * m), a - x);
        failed_serves_a = to_int(ceil(K / 2) * m - x);
        failed_serves_b = to_int(floor(K / 2) * m - a + x);

        return (binom_coeff_a * p_a ^ x) * ((1 - p_a) ^ failed_serves_a) *
               (binom_coeff_b * (1 - p_b) ^ (a - x)) * (p_b ^ failed_serves_b);
    }

    real p_a_wins_R_is_greater_than_zero_and_K_is_even(int x, int a, real p_a, real p_b, int K, int R, int m) {
        int binom_coeff_a;
        int binom_coeff_b;
        int failed_serves_a;
        int failed_serves_b;

        binom_coeff_a = choose(K / 2 * m + R - 1, x - 1);
        binom_coeff_b = choose(K / 2 * m, a - x);
        failed_serves_a = K / 2 * m + R - x;
        failed_serves_b = K / 2 * m - a + x;

        return (binom_coeff_a * p_a ^ x) * ((1 - p_a) ^ failed_serves_a) *
               (binom_coeff_b * (1 - p_b) ^ (a - x)) * (p_b ^ failed_serves_b);
    }

    real p_b_wins_R_is_greater_than_zero_and_K_is_even(int x, int a, real p_a, real p_b, int K, int R, int m) {
        int binom_coeff_a;
        int binom_coeff_b;
        int failed_serves_a;
        int failed_serves_b;

        binom_coeff_a = choose(K / 2 * m + R - 1, x);
        binom_coeff_b = choose(K / 2 * m, a - x);
        failed_serves_a = K / 2 * m + R - x;
        failed_serves_b = K / 2 * m - a + x;

        return choose(K / 2 * m + R - 1, x) * p_a ^ x * (1 - p_a) ^ (K / 2 * m + R - x) *
               choose(K / 2 * m, a - x) * (1 - p_b) ^ (a - x) * p_b ^ (K / 2 * m - a + x);
    }

    real p_a_wins_R_is_greater_than_zero_and_K_is_odd(int x, int a, real p_a, real p_b, int K, int R, int m) {
        int binom_coeff_a;
        int binom_coeff_b;
        int failed_serves_a;
        int failed_serves_b;

        binom_coeff_a = choose(to_int(ceil(K / 2) * m), x);
        binom_coeff_b = choose(to_int(floor(K / 2) * m + R - 1), a - x - 1);
        failed_serves_a = to_int(ceil(K / 2) * m - x);
        failed_serves_b = to_int(floor(K / 2) * m + R - a + x);

        return (binom_coeff_a * p_a ^ x) * ((1 - p_a) ^ failed_serves_a) *
               (binom_coeff_b * (1 - p_b) ^ (a - x)) * (p_b ^ failed_serves_b);
    }

    real p_b_wins_R_is_greater_than_zero_and_K_is_odd(int x, int a, real p_a, real p_b, int K, int R, int m) {
        int binom_coeff_a;
        int binom_coeff_b;
        int failed_serves_a;
        int failed_serves_b;

        binom_coeff_a = choose(to_int(ceil(K / 2) * m), x);
        binom_coeff_b = choose(to_int(floor(K / 2) * m + R - 1), a - x);
        failed_serves_a = to_int(ceil(K / 2) * m - x);
        failed_serves_b = to_int(floor(K / 2) * m + R - a + x);

        return (binom_coeff_a * p_a ^ x) * ((1 - p_a) ^ failed_serves_a) *
               (binom_coeff_b * (1 - p_b) ^ (a - x)) * (p_b ^ failed_serves_b);
    }

    real p_a_wins_set(int x, int a, int b, real p_a, real p_b, int m) {
        int K;
        int R;

        K = calculate_K(a, b, m);
        R = calculate_R(a, b, m, K);

        if (K % 2 == 0) {
            if (R > 0) {
                return p_a_wins_R_is_greater_than_zero_and_K_is_even(x, a, p_a, p_b, K, R, m);
            } else {
                return p_a_wins_R_is_zero_and_K_is_even(x, a, p_a, p_b, K, m);
            }
        } else {
            if (R > 0) {
                return p_a_wins_R_is_greater_than_zero_and_K_is_odd(x, a, p_a, p_b, K, R, m);
            } else {
                return p_a_wins_R_is_zero_and_K_is_odd(x, a, p_a, p_b, K, m);
            }
        }
    }

    real p_b_wins_set(int x, int a, int b, real p_a, real p_b, int m) {
        int K;
        int R;

        K = calculate_K(a, b, m);
        R = calculate_R(a, b, m, K);

        if (K % 2 == 0) {
            if (R > 0) {
                return p_b_wins_R_is_greater_than_zero_and_K_is_even(x, a, p_a, p_b, K, R, m);
            } else {
                return p_b_wins_R_is_zero_and_K_is_even(x, a, p_a, p_b, K, m);
            }
        } else {
            if (R > 0) {
                return p_b_wins_R_is_greater_than_zero_and_K_is_odd(x, a, p_a, p_b, K, R, m);
            } else {
                return p_b_wins_R_is_zero_and_K_is_odd(x, a, p_a, p_b, K, m);
            }
        }
    }

    real p_a_wins_after_tie(real p_a, real p_b) {
        return p_a * (1 - p_b) / (1 - ((1 - p_a) * (1 - p_b) + p_a * p_b));
    }

    int[] get_xs(int turn, int a, int b, int m) {
        int K;
        int R;
        int xs[2];

        K = calculate_K(a, b, m);
        R = calculate_R(a, b, m, K);

        if (turn == 1) {
            if (K % 2 == 0) {
                if (R > 0) {
                    xs[1] = max(1, a - K / 2 * m);
                    xs[2] = min(a, K / 2 * m + R) + 1;
                } else {
                    xs[1] = max(0, a - K / 2 * m);
                    xs[2] = min(a - 1, K / 2 * m) + 1;
                }
            } else {
                if (R > 0) {
                    xs[1] = max(0, to_int(a - floor(K / 2) * m - R));
                    xs[2] = min(a - 1, to_int(ceil(K / 2) * m)) + 1;
                } else {
                    xs[1] = max(1, to_int(a - floor(K / 2) * m));
                    xs[2] = min(a, to_int(ceil(K / 2) * m)) + 1;
                }
            }
        } else {
            if (K % 2 == 0) {
                if (R > 0) {
                    xs[1] = max(0, a - K / 2 * m);
                    xs[2] = min(a, K / 2 * m + R - 1) + 1;
                } else {
                    xs[1] = max(0, a - K / 2 * m + 1);
                    xs[2] = min(a, K / 2 * m) + 1;
                }
            } else {
                if (R > 0) {
                    xs[1] = max(0, to_int(a - floor(K / 2) * m - R + 1));
                    xs[2] = min(a, to_int(ceil(K / 2) * m)) + 1;
                } else {
                    xs[1] = max(0, to_int(a - floor(K / 2) * m));
                    xs[2] = min(a, to_int(ceil(K / 2) * m - 1)) + 1;
                }
            }
        }

        return xs;
    }

    real calculate_p_a_wins(real p_a, real p_b, int m, int n) {
        real p_a_win = 0;
        real p_a_tie = 0;
        real p_b_tie = 0;
        real p_a_after_tie;

        for (k in 0:n-2) {
            for (x in get_xs(1, n, k-n, m)) {
                p_a_win += p_a_wins_set(x, n, k, p_a, p_b, m);
            }
        }

        for (x in get_xs(1, n-1, n-1, m)) {
            p_a_tie += p_a_wins_set(x, n-1, n-1, p_a, p_b, m);
        }

        for (x in get_xs(0, n-1, n-1, m)) {
            p_b_tie += p_b_wins_set(x, n-1, n-1, p_a, p_b, m);
        }

        p_a_after_tie = p_a_wins_after_tie(p_a, p_b);
        p_a_win += (p_a_tie + p_b_tie) * p_a_after_tie;

        return p_a_win;
    }
}

data {
  real<lower=0, upper=1> p_a;
  real<lower=0,upper=1> p_b;
}


transformed parameters {
  real<lower=0,upper=1> p_a_wins;
  p_a_wins = calculate_p_a_wins(p_a, p_b, 2, 11);
}


model {
  target += p_a_wins;
}

If p_a is 0.3 and p_b is 0.1, p_a_wins should be:

0.9048684360973515

Running the above Stan model gives the following error:

Initialization failed. Rejecting initial value: Error evaluating the log probability at the initial value.

How can I test my custom function easily? Basically I don’t need the model part but Stan needs it, so how to deal with this?

If you’re using CmdStanPy (and you’re not on Windows*) I have a tool which should be fairly easy to use that lets you call Stan functions directly from Python. I use it to do exactly what you describe:

  • using it on Windows is possible, just not as automatic
2 Likes