add multiplier keyword to unit vector

In https://discourse.mc-stan.org/t/a-better-unit-vector/26989/30 Seth Axen lays out the reasoning for adding an additional parameter which repels values away from 0.

The current implementation is equivalent to

data {
  int<lower=0> N;
}
parameters {
  vector[N] u_raw;
}
transformed parameters {
  real r = dot_self(u_raw);
  vector[N] u = u_raw / sqrt(r);
}
model {
  target += -0.5 * r;
}

The proposal is to parameterize the unit vector as

data {
  int<lower=0> N;
}
parameters {
  vector[N] u_raw;
}
transformed parameters {
  real r = dot_self(u_raw);
  vector[N] u = u_raw / sqrt(r);
}
model {
 target += 0.5 * (a * log(r) - r);
}

where a >= 0 and is given as a multiplier keyword. As Seth notes in the post, it corresponds to the chi distribution with a + 1 degrees of freedom. The user would see:

unit_vector<multiplier=a>[N] u;

Setting a = 0 is equivalent to what is currently in Stan and will still be the default when the multiplier keyword is not used.

How to select a?

From the same thread the issue manifests with smaller dimension sizes of the unit vector. Heuristically, I'm finding that setting a = N^(6/N) where N is the dimension of the unit-vector works well.

         N a
 [1,]    1 1.000000
 [2,]    2 8.000000
 [3,]    3 9.000000
 [4,]    4 8.000000
 [5,]    5 6.898648
 [6,]    6 6.000000
 [7,]    7 5.301146
 [8,]    8 4.756828
 [9,]    9 4.326749
[10,]   10 3.981072

Here's a test model for 2d unit vectors where divergences occur in the current parameterization. When the data size, M, is small there are fewer divergences then when it's larger (M > 50).

data {
  int<lower=0> M;
  vector<lower=-pi(), upper=pi()>[M] y;
  real a;
}
parameters {
  vector[2] u_raw;
  real<lower=0> kappa;
}
transformed parameters {
  real r = dot_self(u_raw);
  vector[2] u = u_raw / sqrt(r);
  real mu = atan2(u[2], u[1]);
}
model {
  target += 0.5 * (a * log(r) - r);
  kappa ~ exponential(1);
  y ~ von_mises(mu, kappa);
}

The data for this can be created as

M <- 100
y <- runif(M, min = -pi/4, max = pi/4)
true_mean <- 0

mod_unit <- cmdstan_model("unit.stan")

mod_unit_out <- mod_unit$sample(
  data = list(M = M,
              y = y,
              a = 2^(6/2),
  seed = 2309423,
  parallel_chains = 4
)