From 7c8b995edfceac9d59ca07baf3d700be1bc97ed0 Mon Sep 17 00:00:00 2001
From: Seth Axen <seth@sethaxen.com>
Date: Tue, 4 Jun 2024 16:58:13 +0200
Subject: [PATCH] Add multi-logit-normal implementation

Adapted from mjhajharia/transforms#74
---
 .../functions/multi-logit-normal.stanfunction | 28 +++++++++++++++++++
 1 file changed, 28 insertions(+)
 create mode 100644 targets/functions/multi-logit-normal.stanfunction

diff --git a/targets/functions/multi-logit-normal.stanfunction b/targets/functions/multi-logit-normal.stanfunction
new file mode 100644
index 0000000..934e40c
--- /dev/null
+++ b/targets/functions/multi-logit-normal.stanfunction
@@ -0,0 +1,28 @@
+/**
+ * Return the multivariate logistic normal density for the specified simplex.
+ *
+ * See: https://en.wikipedia.org/wiki/Logit-normal_distribution#Multivariate_generalization
+ * 
+ * @param theta a simplex (N rows)
+ * @param mu location of normal (N-1 rows)
+ * @param L_Sigma Cholesky factor of covariance (N-1 rows, N-1 cols)
+ */
+real multi_logit_normal_cholesky_lpdf(vector theta, vector mu, matrix L_Sigma) {
+  vector[rows(theta)] log_theta = log(theta);
+  return sum(-log_theta)
+    + multi_log_logit_normal_cholesky_lpdf(log_theta | mu, L_Sigma);
+}
+
+/**
+ * Return the multivariate logistic normal density for the specified log simplex.
+ *
+ * See: https://en.wikipedia.org/wiki/Logit-normal_distribution#Multivariate_generalization
+ * 
+ * @param theta a log simplex (N rows)
+ * @param mu location of normal (N-1 rows)
+ * @param L_Sigma Cholesky factor of covariance (N-1 rows, N-1 cols)
+ */
+real multi_log_logit_normal_cholesky_lpdf(vector log_theta, vector mu, matrix L_Sigma) {
+  int N = rows(log_theta);     
+  return multi_normal_cholesky_lpdf(log_theta[1:N-1] - log_theta[N] | mu, L_Sigma);
+}