diff --git a/DESCRIPTION b/DESCRIPTION
index 7cb0903..bfc79d7 100644
--- a/DESCRIPTION
+++ b/DESCRIPTION
@@ -1,7 +1,7 @@
 Package: bmstate
 Type: Package
 Title: Bayesian multistate modeling
-Version: 0.3.0
+Version: 0.3.1
 Authors@R:
     c(person(given = "Juho",
            family = "Timonen",
diff --git a/NAMESPACE b/NAMESPACE
index 7e19011..b39c1d1 100644
--- a/NAMESPACE
+++ b/NAMESPACE
@@ -23,7 +23,9 @@ export(ensure_exposed_stan_functions)
 export(example_sim_setup_illnessdeath)
 export(fit_stan)
 export(generate_paths)
+export(msmfit_check_xpsr_norm)
 export(msmfit_exposure)
+export(msmfit_exposure_df)
 export(msmfit_inst_hazard_param_draws)
 export(msmfit_log_hazard_multipliers)
 export(msmfit_pk_params)
diff --git a/R/DosingData.R b/R/DosingData.R
index 259ac68..5a76b73 100644
--- a/R/DosingData.R
+++ b/R/DosingData.R
@@ -104,23 +104,14 @@ PSSDosingData <- R6::R6Class(
     #'
     #' @param t A vector of output times for each subject (a list).
     #' @param theta A matrix of parameters.
-    #' @param skip_assert Skip most assertions and call exposed Stan directly,
-    #' assuming that it exists?
     #' @param MAX_CONC concentration upper bound
     #' @return a \code{data.frame}
-    simulate_pk = function(t, theta, MAX_CONC, skip_assert = FALSE) {
+    simulate_pk = function(t, theta, MAX_CONC) {
       checkmate::assert_list(t, len = self$num_subjects())
-      checkmate::assert_logical(skip_assert, len = 1)
       checkmate::assert_number(MAX_CONC, lower = 0)
-      if (skip_assert) {
-        out <- pop_2cpt_partly_ss(
-          t, self$dose_ss, self$times, self$doses, theta, self$tau_ss, MAX_CONC
-        )
-      } else {
-        out <- pk_2cpt_pss(
-          t, self$dose_ss, self$times, self$doses, theta, self$tau_ss, MAX_CONC
-        )
-      }
+      out <- pk_2cpt_pss(
+        t, self$dose_ss, self$times, self$doses, theta, self$tau_ss, MAX_CONC
+      )
 
       time <- as.numeric(unlist(t))
       val <- unlist(out)
diff --git a/R/MultistateModel.R b/R/MultistateModel.R
index d4e6667..3a357a9 100644
--- a/R/MultistateModel.R
+++ b/R/MultistateModel.R
@@ -474,24 +474,3 @@ place_internal_knots <- function(t_max, num_knots, t_event) {
   knots <- stats::quantile(t_event, probs = seq(0, 1, h))
   knots[2:(length(knots) - 1)]
 }
-
-# Normalize columns of matrix A
-normalize_columns <- function(A) {
-  K <- ncol(A)
-  cn <- colnames(A)
-  for (j in seq_len(K)) {
-    A[, j] <- normalize(A[, j], cn[j])
-  }
-  A
-}
-
-# Normalize
-normalize <- function(a, name) {
-  sdd <- stats::sd(a)
-  if (sdd == 0) {
-    stop("error in normalization of ", name, ", zero variance")
-  }
-  x_norm <- (a - mean(a)) / sdd
-  check_normalized_covariate(x_norm, name)
-  x_norm
-}
diff --git a/R/MultistateModelFit.R b/R/MultistateModelFit.R
index 9ef5e60..6a8f654 100644
--- a/R/MultistateModelFit.R
+++ b/R/MultistateModelFit.R
@@ -152,7 +152,7 @@ MultistateModelFit <- R6::R6Class("MultistateModelFit",
       DF <- list()
       for (s in seq_len(S)) {
         pb$tick()
-        df_s <- data$dosing$simulate_pk(ts, pkpar[[s]], MC, skip_assert = TRUE)
+        df_s <- data$dosing$simulate_pk(ts, pkpar[[s]], MC)
         df_s$.draw_idx <- s
         DF[[s]] <- df_s
       }
@@ -423,6 +423,50 @@ msmfit_exposure <- function(fit, oos = FALSE, data = NULL) {
   out
 }
 
+#' Compute exposure and return as rvar in df
+#'
+#' @export
+#' @inheritParams msmfit_exposure
+msmfit_exposure_df <- function(fit, oos = FALSE, data = NULL) {
+  check_oos(oos, data)
+  xp <- msmfit_exposure(fit, oos, data)
+  xp <- posterior::rvar(t(sapply(xp, function(x) x)))
+  if (is.null(data)) {
+    data <- fit$data
+  }
+  df <- data$paths$subject_df
+  df$xpsr_estimate <- xp
+  df
+}
+
+#' Check exposure normalization
+#'
+#' @export
+#' @inheritParams msmfit_exposure
+msmfit_check_xpsr_norm <- function(fit) {
+  df <- msmfit_exposure_df(fit)
+  r <- mean(df$xpsr_estimate)
+  xn <- fit$model$get_xpsr_normalizers()
+  msg1 <- paste0(
+    "xpsr normalization loc = ", round(xn$loc, 5),
+    ", mean estimated xpsr = ",
+    round(mean(r), 5)
+  )
+  msg2 <- paste0(
+    "xpsr normalization scale = ", round(xn$scale, 5),
+    ", estimated xpsr sd = ",
+    round(stats::sd(r), 5)
+  )
+  message(
+    "If these don't roughly match, consider refitting after setting ",
+    "xpsr normalizer loc and scale closer to estimated mean and sd. ",
+    "Otherwise interpret baseline hazards and xpsr effect size accodingly."
+  )
+  message(" - ", msg1)
+  message(" - ", msg2)
+}
+
+
 # Helper
 check_oos <- function(oos, data) {
   checkmate::assert_logical(oos, len = 1)
diff --git a/R/PKModel.R b/R/PKModel.R
index f919ef9..96bec24 100644
--- a/R/PKModel.R
+++ b/R/PKModel.R
@@ -112,8 +112,9 @@ PKModel <- R6::R6Class("PKModel",
       CL <- theta[2]
       V2 <- theta[3]
       checkmate::assert_number(CL, lower = 0)
+      checkmate::assert_number(V2, lower = 0)
       checkmate::assert_number(dose, lower = 0)
-      log(dose) - log(CL * V2)
+      log(dose) - log(CL) - log(V2)
     },
 
     #' @description Simulate data with many subjects
diff --git a/R/stan-data.R b/R/stan-data.R
index 3c2bc1d..b67bd64 100644
--- a/R/stan-data.R
+++ b/R/stan-data.R
@@ -184,12 +184,6 @@ create_stan_data_pk <- function(data, model) {
   }
   t_obs_pk <- pk_obs[, 1:2, drop = FALSE]
   conc_pk <- pk_obs[, 3:4, drop = FALSE]
-  if (max(conc_pk) > 1e5) {
-    stop(
-      "Rescale concentration measurements to larger units so that maximum",
-      "concentration has smaller numeric value"
-    )
-  }
   pk_lloq <- as.numeric(pk_obs[, 5])
 
   # Return
@@ -343,11 +337,11 @@ time_since_last_dose <- function(t_obs_pk, last_two_times) {
   n_before_post <- sum(last_two_times <= t_obs_pk[2])
   if (n_before_pre == 0) {
     n_before_pre <- 1
-    stop("invalid data")
+    stop("invalid data, , 0 doses before pre-dose time")
   }
   if (n_before_post == 0) {
     n_before_pre <- 1
-    stop("invalid data")
+    stop("invalid data, 0 doses before post-dose time")
   }
   t_before_pre <- last_two_times[n_before_pre]
   t_before_post <- last_two_times[n_before_post]
diff --git a/R/stan.R b/R/stan.R
index 6166a64..00cb5b6 100644
--- a/R/stan.R
+++ b/R/stan.R
@@ -150,6 +150,9 @@ fit_stan <- function(model, data,
     time = stan_fit$time()
   )
   fit <- MultistateModelFit$new(data, sd, model, draws, info)
+  if (fit$model$has_pk()) {
+    msmfit_check_xpsr_norm(fit)
+  }
   if (!return_stanfit) {
     return(fit)
   }
diff --git a/R/utils.R b/R/utils.R
index 4bbd93d..3c97474 100644
--- a/R/utils.R
+++ b/R/utils.R
@@ -175,3 +175,25 @@ prefit_checks <- function(model, data) {
   }
   TRUE
 }
+
+
+# Normalize columns of matrix A
+normalize_columns <- function(A) {
+  K <- ncol(A)
+  cn <- colnames(A)
+  for (j in seq_len(K)) {
+    A[, j] <- normalize(A[, j], cn[j])
+  }
+  A
+}
+
+# Normalize
+normalize <- function(a, name) {
+  sdd <- stats::sd(a)
+  if (sdd == 0) {
+    stop("error in normalization of ", name, ", zero variance")
+  }
+  x_norm <- (a - mean(a)) / sdd
+  check_normalized_covariate(x_norm, name)
+  x_norm
+}
diff --git a/inst/stan/msm.stan b/inst/stan/msm.stan
index 89308c0..2ae604b 100644
--- a/inst/stan/msm.stan
+++ b/inst/stan/msm.stan
@@ -2,7 +2,7 @@ functions {
 
   // log SS AUC
   vector log_ss_area_under_conc(vector dose_ss, matrix log_theta_pk){
-    return(log(dose_ss) - log_theta_pk[:,2] - log_theta_pk[:,3]) ; // D/(CL*V2)
+    return(log(dose_ss) - log_theta_pk[:,2] - log_theta_pk[:,3]) ; // log D/(CL*V2)
   }
 
   // log of hazard multiplier
@@ -304,7 +304,7 @@ data {
 
   // PK options
   int<lower=0,upper=1> I_xpsr;
-  real<lower=0> xpsr_loc;
+  real xpsr_loc;
   real<lower=0> xpsr_scale;
   int<lower=0> nc_ka; // num of predictors for ka
   int<lower=0> nc_CL; // num of predictors for CL
@@ -448,7 +448,7 @@ model {
   if(do_haz == 1){
     if(nc_haz > 0){
       for(k in 1:nc_haz){
-        beta_oth[1, k] ~ normal(0, 1);
+        beta_oth[1, k] ~ normal(0, 2);
       }
     }
     if(I_xpsr==1){
@@ -471,7 +471,7 @@ model {
     for(n in 1:size(log_z_pk[1])){
       log_z_pk[1, n] ~ normal(0, 1);
     }
-    log_mu_pk[1] ~ normal(0, 3);
+    log_mu_pk[1] ~ normal(0, 2);
     log_sig_pk[1] ~ normal(0, 1);
     beta_ka[1] ~ normal(0, 1);
     beta_CL[1] ~ normal(0, 1);
diff --git a/man/PSSDosingData.Rd b/man/PSSDosingData.Rd
index 6ff85ea..5601725 100644
--- a/man/PSSDosingData.Rd
+++ b/man/PSSDosingData.Rd
@@ -89,7 +89,7 @@ As data frame
 \subsection{Method \code{simulate_pk()}}{
 Simulate PK dynamics
 \subsection{Usage}{
-\if{html}{\out{<div class="r">}}\preformatted{PSSDosingData$simulate_pk(t, theta, MAX_CONC, skip_assert = FALSE)}\if{html}{\out{</div>}}
+\if{html}{\out{<div class="r">}}\preformatted{PSSDosingData$simulate_pk(t, theta, MAX_CONC)}\if{html}{\out{</div>}}
 }
 
 \subsection{Arguments}{
@@ -100,9 +100,6 @@ Simulate PK dynamics
 \item{\code{theta}}{A matrix of parameters.}
 
 \item{\code{MAX_CONC}}{concentration upper bound}
-
-\item{\code{skip_assert}}{Skip most assertions and call exposed Stan directly,
-assuming that it exists?}
 }
 \if{html}{\out{</div>}}
 }
diff --git a/man/msmfit_check_xpsr_norm.Rd b/man/msmfit_check_xpsr_norm.Rd
new file mode 100644
index 0000000..5763d74
--- /dev/null
+++ b/man/msmfit_check_xpsr_norm.Rd
@@ -0,0 +1,14 @@
+% Generated by roxygen2: do not edit by hand
+% Please edit documentation in R/MultistateModelFit.R
+\name{msmfit_check_xpsr_norm}
+\alias{msmfit_check_xpsr_norm}
+\title{Check exposure normalization}
+\usage{
+msmfit_check_xpsr_norm(fit)
+}
+\arguments{
+\item{fit}{A \code{\link{MultistateModelFit}} object}
+}
+\description{
+Check exposure normalization
+}
diff --git a/man/msmfit_exposure_df.Rd b/man/msmfit_exposure_df.Rd
new file mode 100644
index 0000000..3fc951a
--- /dev/null
+++ b/man/msmfit_exposure_df.Rd
@@ -0,0 +1,21 @@
+% Generated by roxygen2: do not edit by hand
+% Please edit documentation in R/MultistateModelFit.R
+\name{msmfit_exposure_df}
+\alias{msmfit_exposure_df}
+\title{Compute exposure and return as rvar in df}
+\usage{
+msmfit_exposure_df(fit, oos = FALSE, data = NULL)
+}
+\arguments{
+\item{fit}{A \code{\link{MultistateModelFit}} object}
+
+\item{oos}{Out-of-sample mode? If \code{FALSE}, the possible subject-specific
+fitted parameters are used. If \code{TRUE}, acting
+as if the subjects are new.}
+
+\item{data}{A \code{\link{JointData}} object. If \code{NULL}, the
+data used to fit the model is used.}
+}
+\description{
+Compute exposure and return as rvar in df
+}
diff --git a/tests/testthat/test-workflow.R b/tests/testthat/test-workflow.R
index 859e973..452cd73 100644
--- a/tests/testthat/test-workflow.R
+++ b/tests/testthat/test-workflow.R
@@ -139,7 +139,7 @@ test_that("entire workflow works (with PK)", {
     iter_warmup = options$iter_warmup,
     iter_sampling = options$iter_sampling,
     chains = options$chains,
-    refresh = 5,
+    refresh = 20,
     adapt_delta = 0.95,
     init = 0.1
   )
@@ -179,28 +179,28 @@ test_that("entire workflow works (with PK)", {
 })
 
 test_that("entire workflow works (with single transition PK)", {
-  # Setup
-  h0_base <- 1e-3
-
   # Create model
   sn <- c("Alive", "Death")
   tm <- transmat_survival(state_names = sn)
   t3yr <- 3 * 365.25
-  covs <- c("age")
+  covs <- c("age", "sex")
   pk_covs <- list(
     CL = "CrCL",
     V2 = c("weight", "sex")
   )
-  mod <- create_msm(tm, covs, pk_covs, num_knots = 5, t_max = t3yr)
-  bh_true <- matrix(0, 1, 2)
+  mod <- create_msm(tm, covs, pk_covs,
+    num_knots = 5, t_max = t3yr,
+    categ_covs = "sex"
+  )
+  bh_true <- matrix(0, 1, 3)
   bh_true[1, 1] <- 0.5 # age effect on death
-  bh_true[1, 2] <- -0.5 # dose effect on death
+  bh_true[1, 2] <- -0.5 # xpsr effect on death
   sn <- "Death"
   rownames(bh_true) <- paste0("Effect on ", sn)
   colnames(bh_true) <- mod$covs()
 
   h0_true <- 1e-3
-  beta_pk <- list(CL = -0.1, V2 = c(0, 0.1))
+  beta_pk <- list(CL = -0.3, V2 = c(0.2, 0.3))
 
   # Simulate data
   simdat <- mod$simulate_data(
@@ -219,7 +219,7 @@ test_that("entire workflow works (with single transition PK)", {
     iter_warmup = options$iter_warmup,
     iter_sampling = options$iter_sampling,
     chains = options$chains,
-    refresh = 5,
+    refresh = 20,
     adapt_delta = 0.95,
     init = 0.1
   )
@@ -250,22 +250,43 @@ test_that("PK-only works", {
   h0_base <- 1e-3
 
   # Simulate data
-  sim <- simulate_example_data(N = options$N_subject, pk_only = TRUE)
-  mod <- sim$model
-  jd <- sim$data
-
-  # Split
-  jd <- split_data(jd)
+  pk_covs <- list(
+    V2 = "weight"
+  )
+  tm <- transmat_survival()
+  tmax <- 3 * 365.25
+  mod <- create_msm(tm, hazard_covs = NULL, pk_covs, t_max = tmax, pk_only = TRUE)
+  jd <- mod$simulate_data(options$N_subject)
 
   # Fit the model
-  fit <- fit_stan(mod, jd$train,
+  fit <- fit_stan(mod, jd,
     iter_warmup = options$iter_warmup,
     iter_sampling = options$iter_sampling,
     chains = options$chains,
-    refresh = 5,
+    refresh = 20,
     adapt_delta = 0.95,
     init = 0.1
   )
+
+  # Plot pk fit
+  sid <- unique(jd$paths$subject_df$subject_id)[1:6]
+  plt <- fit$plot_pk(subject_ids = sid)
+  ttt <- seq(100, 300, by = 1)
+  tl <- list()
+  for (j in 1:options$N_subject) {
+    tl[[j]] <- seq(180, 220, by = 1)
+  }
+  dd <- fit$data$dosing
+  TH_true <- as.matrix(fit$data$paths$subject_df[, c("ka", "CL", "V2")])
+  TH_fit <- as.matrix(msmfit_pk_params(fit$mean_fit())[[1]])
+  pks1 <- dd$simulate_pk(tl, TH_true, 10^5)
+  pks1 <- pks1 |> dplyr::filter(.data$subject_id %in% sid)
+  pks2 <- dd$simulate_pk(tl, TH_fit, 10^5)
+  pks2 <- pks2 |> dplyr::filter(.data$subject_id %in% sid)
+  plt2 <- plt + geom_line(data = pks1, aes(x = time, y = val, group = subject_id)) +
+    geom_line(data = pks2, aes(x = time, y = val, group = subject_id), color = "red")
+  expect_true(is_ggplot(plt2))
+
   fit <- fit$mean_fit()
   expect_true(inherits(fit, "MultistateModelFit"))