Fix issue with plotting functions when nuts_sampler = "nutpie"

NEWS.md

@@ -1,5 +1,9 @@
 # coevolve - development version
 
+### Bug Fixes
+
+* Fixed issue with plotting functions when `nuts_sampler = "nutpie"` (#114)
+
 ### Other Changes
 
 * Update documentation (#111)

inst/python/coev_jax_model.py

@@ -506,6 +506,7 @@ def _expand_var_list(self):
         vs.append(("Q_sigma", (J,)))
         vs.append(("b", (J,)))
         vs.append(("eta_anc", (self.N_tree, J)))
+        vs.append(("eta", (self.N_tree, self.N_seg, J)))
         if self.estimate_correlated_drift:
             vs.append(("cor_R", (J, J)))
         for j1, nc in zip(self.ordered_j, self.ordered_ncuts):
@@ -1199,13 +1200,26 @@ def _expand_jax(x):
             Q, L_R, _ = self._build_Q_matrix(params)
             vals.append(A_mat)  # A
             vals.append(Q)      # Q
+
+            Q_inf = ksolve(A_mat, Q, self.J)
+            A_delta_cache, L_VCV_cache, A_solve_cache, b_delta_cache = (
+                self._compute_caches(A_mat, Q_inf, params)
+            )
+
+            eta, tip_L_VCV = self._tree_traversal(
+                params,
+                A_delta_cache,
+                L_VCV_cache,
+                b_delta_cache,
+            )
 
             vals.append(params["A_diag"])
             if self.n_offdiag > 0:
                 vals.append(params["A_offdiag"])
             vals.append(params["Q_sigma"])
             vals.append(params["b"])
             vals.append(params["eta_anc"])
+            vals.append(eta)
 
             if self.estimate_correlated_drift:
                 vals.append(L_R @ L_R.T)  # cor_R

tests/testthat/test-coev_fit_nutpie.R

@@ -402,3 +402,38 @@ test_that("JAX fit with HSGP returns coevfit with GP params", {
   expect_true(any(grepl("^sigma_dist", vars)))
   expect_true(any(grepl("^rho_dist", vars)))
 })
+
+test_that("Plotting functions work with JAX fit", {
+  skip_if_not(
+    coevolve:::check_jax_available(),
+    message = "JAX not available - skipping JAX tests"
+  )
+  fit <- coev_fit(
+    data = authority$data,
+    variables = list(
+      political_authority = "ordered_logistic",
+      religious_authority = "ordered_logistic"
+    ),
+    id = "language",
+    tree = authority$phylogeny,
+    prior = list(A_offdiag = "normal(0, 2)"),
+    nuts_sampler = "nutpie",
+    chains = 1,
+    iter_warmup = 50,
+    iter_sampling = 50,
+    seed = 1
+  )
+  expect_no_error(coev_plot_trait_values(fit))
+  expect_no_error(coev_plot_delta_theta(fit))
+  expect_no_error(coev_plot_pred_series(fit))
+  expect_no_error(
+    coev_plot_flowfield(
+      fit, var1 = "political_authority", var2 = "religious_authority"
+    )
+  )
+  expect_no_error(
+    coev_plot_selection_gradient(
+      fit, var1 = "political_authority", var2 = "religious_authority"
+    )
+  )
+})