Custom likelihood with if and i

Modeling
,
1

Hello everyone, I’m trying to fit a new model with a custom likelihood. I’m sharing a toy example to get some tips on improving efficiency.

  1. Is there a more efficient way to have a different likelihood for each chosen option rather than using “if”?
  2. Is my usage of the complex number is the way to go?

Thanks a lot!
Ido

functions {

  real custom_likelihood(real rt, real lambda, real d, real sigma,

                         real theta, int choice) {

    real P_L = 1 - normal_cdf(lambda - d | 0, sigma);

    

    real P_S = normal_cdf(-lambda - d | 0, sigma);

    

    real P_C = 1 - P_L - P_S;

    

    real result;

    

    complex z = 0 + 1i;

    

    if (choice == 1) {

      result = get_real((P_S / (theta * sqrt(P_C))) * exp(-rt / theta)

                        * ((exp(rt * sqrt(P_C) / theta))

                           - cos(z * rt * sqrt(P_C) / theta)));

    } else if (choice == 2) {

      result = get_real((P_L / (theta * sqrt(P_C))) * exp(-rt / theta)

                        * ((exp(rt * sqrt(P_C) / theta))

                           - cos(z * rt * sqrt(P_C) / theta)));

    } else {

      return 0; // Default return in case choice is invalid

    }

    

    return result;

  }

}

data {

  int<lower=1> Nsubjects; //number of subjects

  

  int<lower=1> Nblocks;

  

  int<lower=1> Ntrials; // Number of observations

  

  array[Nsubjects, Ntrials] real<lower=0> rts; // Response times

  

  array[Nsubjects, Ntrials] real<lower=0> d; // Constant d

  

  array[Nsubjects, Ntrials] real<lower=0> sigma; // Constant sigma

  

  array[Nsubjects, Ntrials] real<lower=0> theta; // Constant theta

  

  array[Nsubjects, Ntrials] int<lower=1, upper=2> choices; // Choice array (1 or 2 for each observation)

}

parameters {

  array[Nsubjects] real<lower=0> lambda; // Parameter to estimate

}

model {

  lambda ~ uniform(0, 4);

  

  for (subject in 1 : Nsubjects) {

    for (trial in 1 : Ntrials) {

      target += custom_likelihood(rts[subject, trial], lambda[subject],

                                  d[subject, trial], sigma[subject, trial],

                                  theta[subject, trial],

                                  choices[subject, trial]);

    }

  }

}
2

No idea what happened to introduce all the spaces, but here’s a more readable form

functions {
  real custom_likelihood(real rt, real lambda, real d, real sigma,
                         real theta, int choice) {
    real P_L = 1 - normal_cdf(lambda - d | 0, sigma);
    real P_S = normal_cdf(-lambda - d | 0, sigma);
    real P_C = 1 - P_L - P_S;
    real result;
    complex z = 0 + 1i;
    if (choice == 1) {
      result = get_real((P_S / (theta * sqrt(P_C))) * exp(-rt / theta)
                        * ((exp(rt * sqrt(P_C) / theta))
                           - cos(z * rt * sqrt(P_C) / theta)));
    } else if (choice == 2) {
      result = get_real((P_L / (theta * sqrt(P_C))) * exp(-rt / theta)
                        * ((exp(rt * sqrt(P_C) / theta))
                           - cos(z * rt * sqrt(P_C) / theta)));
    } else {
      return 0; // Default return in case choice is invalid
    }
    return result;
  }
}
data {
  int<lower=1> Nsubjects; //number of subjects
  int<lower=1> Nblocks;
  int<lower=1> Ntrials; // Number of observations
  array[Nsubjects, Ntrials] real<lower=0> rts; // Response times
  array[Nsubjects, Ntrials] real<lower=0> d; // Constant d
  array[Nsubjects, Ntrials] real<lower=0> sigma; // Constant sigma
  array[Nsubjects, Ntrials] real<lower=0> theta; // Constant theta
  array[Nsubjects, Ntrials] int<lower=1, upper=2> choices; // Choice array (1 or 2 for each observation)
}
parameters {
  array[Nsubjects] real<lower=0> lambda; // Parameter to estimate
}
model {
  lambda ~ uniform(0, 4);
  for (subject in 1 : Nsubjects) {
    for (trial in 1 : Ntrials) {
      target += custom_likelihood(rts[subject, trial], lambda[subject],
                                  d[subject, trial], sigma[subject, trial],
                                  theta[subject, trial],
                                  choices[subject, trial]);
    }
  }
}

First of all, you are adding the result of custom_likelihood to target, which is confusing. The target accumulates log likelihood. So make sure that’s the log likelihood you are adding. I would suggest renaming the function to log_likelihood to make this clear (the prefix custom_ isn’t adding any information—the reader knows it’s custom because it’s defined in the functions block).

Getting rid of the conditional won’t help with efficiency here as the rest of the code does so much more computation. You can clean it all up in a couple of ways. First, you don’t need the else because the choice variable is constrained in the data declaration to be 1 or 2. If it’s not, the data won’t load. You don’t need to check things that have already been checked elsewhere.

I’d replace the end of your custom function, starting with real result; and replace it with this:

return get_real(((choice == 1 ? P_S : P_L) / (theta * sqrt(P_C))) * exp(-rt / theta)
                        * ((exp(rt * sqrt(P_C) / theta))
                           - cos(z * rt * sqrt(P_C) / theta)));

Or define intermediates:

real P = choice == 1 ? P_S : PL;
complex z = ((P / (theta * sqrt(P_C))) * exp(-rt / theta)
                        * ((exp(rt * sqrt(P_C) / theta))
                           - cos(z * rt * sqrt(P_C) / theta)));
return get_real(z);

Rather than writing 1 - normal_cdf(...), you can write normal_ccdf(...). It should be more numerically table. But int he case of defining P_C, you can take it to be

real P_C = normal_cdf(lambda - d | 0, sigma) - normal_cdf(-lambda - d | 0, sigma);

One thing you can do to save efficiency is cache redundant calculations. Rather than taking sqrt and dividing the same things multiple times, you can do this:

real sqrt_P_C = sqrt(P_C);
real rt_over_theta;
complex z = ((P / (theta * sqrt_P_C)) * exp(-rt_over_theta)
                        * ((exp(rt_over_theta * sqrt_P_C))
                           - cos(1i * rt_over_theta * sqrt_P_C)));

I also just used 1i in place of z. This can get noticeable when functions like sqrt are used a lot.

I’m afraid I don’t know what you’re trying to do, and even if I did, I’m not very good at complex arithmetic, much less complex calculus. It’s fine as far as Stan code is concerned as long as the sqrt isn’t being applied to negative numbers. Our sqrt function just applies to real numbers and will throw an exception if it gets a negative real number. If you want complex sqrt, you need to cast it to complex first—easiest way to do that is assign to a complex variable. You can also add 0 * 1i.

Also, I think the get_real can be moved down to just scope over the cos as I think the rest is going to return a real number, so that subtracting a complex number is component wise.

1 Like
3

Thanks a lot, Bob!