---
title: "Smooth (spline) terms in sccomp"
author: "Stefano Mangiola"
date: "`r Sys.Date()`"
package: sccomp
output:
  prettydoc::html_pretty:
    theme: cayman
    toc: yes
    toc_depth: 2
    number_sections: yes
    fig_caption: yes
    df_print: paged
abstract: >
  Splines let you model non-linear compositional effects of a continuous
  covariate. For example, pseudotime, age, dose, time-since-treatment, on cell-type
  proportions. 
vignette: >
  %\VignetteIndexEntry{Smooth (spline) terms in sccomp}
  %\VignetteEngine{knitr::rmarkdown}
  %\VignetteEncoding{UTF-8}
---

```{r setup, include=FALSE}
knitr::opts_chunk$set(
  collapse = TRUE,
  comment  = "#>",
  fig.width  = 7,
  fig.height = 4.5,
  dpi = 110
)
# Smooth vignette chunks need mgcv (Suggests) and CmdStan.
run_smooth_vignette <-
  instantiate::stan_cmdstan_exists()
```

# Why smooth terms?

Continuous factors in linear models are tricky, becausegenerally you need to make assumptions about their linearity of higher polinomial trend. 
A spline replaces a single linear column with a small, flexible basis that lets the curve bend, while a penalty keeps the bend conservative

`sccomp` adopts the same syntax as `brms` / `mgcv`:

```{r show-formulae, eval=FALSE}
~ type + s(pseudotime, k = 5)        
```

Under the hood every smooth is split into two pieces:

* an unpenalised **linear trend** block `Xf`, and
* a penalised **wiggly** block `Xr` (the non linear part).

# What is supported

| Feature                                            | Notes                                                |
|----------------------------------------------------|-------------------------------------------------------|
| `s(x, k = ...)` in `formula_composition`           |                                         |
| `s(x, z, bs = "fs", k = ...)` (factor smooth)      | |
| Multiple smooth terms in the same formula          |          |
| Mix smooths with parametric terms (`~ type + s(x)`)|                                                    |
| Mix smooths with explicit `(... | g)` random effects |                            |
| `s(x, by = factor)` (independent smooths per level) | Max 2 smooth terms allowed            |


```{r load-libs, message=FALSE, warning=FALSE}
library(dplyr)
library(tidyr)
library(ggplot2)
library(sccomp)
data("counts_obj")
```

# Continuous covariate


```{r make-pseudotime}
set.seed(1)
samples <- levels(counts_obj$sample)
pt_map <- tibble(
  sample     = samples,
  pseudotime = seq(0, 6, length.out = length(samples)) +
                 rnorm(length(samples), 0, 0.05)
)

counts_pt <- counts_obj |> left_join(pt_map, by = "sample")
counts_pt |> distinct(sample, type, pseudotime) |> head()
```

## 1. Fit `~ type + s(pseudotime, k = 5)`

```{r fit-smooth, message=FALSE, warning=FALSE, eval = run_smooth_vignette}
fit <- counts_pt |>
  sccomp_estimate(
    formula_composition = ~ type + s(pseudotime, k = 5),
    formula_variability = ~ 1,
    sample              = "sample",
    cell_group          = "cell_group",
    abundance           = "count",
    cores               = 1,
    verbose             = FALSE,
    max_sampling_iterations = 2000
  )
```

The output will include one new linear parameter

```
s(pseudotime, k = 5)__lin1
```

and three penalised basis 

```
s(pseudotime, k = 5)___basis01
s(pseudotime, k = 5)___basis02
s(pseudotime, k = 5)___basis03
```

```{r test-table, eval = run_smooth_vignette}
fit |>
  sccomp_test() |>
  select(cell_group, parameter, c_lower, c_effect, c_upper, c_FDR) 
```

These parameters cannot be interpreted directly, on the contrary of other parameter in the linear model. 

## 2. Visualise the curve 

To draw the fitted curve, predict at a regular grid of `pseudotime` values
(remember to also supply any other covariates the formula needs — here, `type`):

```{r predict-grid, message=FALSE, warning=FALSE, eval = run_smooth_vignette}
grid <- expand.grid(
  pseudotime = seq(min(pt_map$pseudotime), max(pt_map$pseudotime), length.out = 60),
  type       = levels(counts_obj$type)
) |>
  as_tibble() |>
  mutate(sample = sprintf("grid_%03d", row_number()))

pred <- fit |>
  sccomp_predict(new_data = grid, number_of_draws = 200)

head(pred)
```

`sccomp_predict()` already returns the `new_data` columns (`pseudotime`,
`type`, `sample`) alongside `proportion_mean` and the 95% credible interval
`proportion_lower` / `proportion_upper`.

We show the six cell groups with the largest pseudotime range in the
posterior mean.

```{r plot-curves, fig.width=8, fig.height=5, eval = run_smooth_vignette}
top_cgs <- pred |>
  group_by(cell_group) |>
  summarise(rng = max(proportion_mean) - min(proportion_mean)) |>
  arrange(desc(rng)) |>
  slice_head(n = 6) |>
  pull(cell_group)

raw <- counts_pt |>
  group_by(sample) |>
  mutate(prop = count / sum(count)) |>
  ungroup() |>
  filter(cell_group %in% top_cgs)

pred |>
  filter(cell_group %in% top_cgs) |>
  ggplot(aes(pseudotime, proportion_mean, colour = type, fill = type)) +
  geom_ribbon(aes(ymin = proportion_lower, ymax = proportion_upper),
              alpha = 0.15, linewidth = 0) +
  geom_line(linewidth = 0.7) +
  geom_point(data = raw, aes(y = prop), alpha = 0.5, size = 1.2) +
  facet_wrap(~ cell_group, scales = "free_y", ncol = 3) +
  theme_minimal(base_size = 10)
```

The blue/red curves are posterior means, ribbons are 95% credible intervals,
points are observed sample proportions. Note the wiggle — e.g. CD8 1 peaks
around pseudotime ≈ 4 — which a purely linear `~ type + pseudotime` model
could not capture.

## 3. Let's use a smaller basis (`k = 3`)

The same parametric + smooth layout works with a lower `k`. Fewer basis
functions mean a smoother curve and less risk of over-fitting when sample
size is modest. Here `k = 3` yields one unpenalised linear column and a
single wiggly basis column.

```{r fit-smooth-k3, message=FALSE, warning=FALSE, eval = run_smooth_vignette}
fit_k3 <- counts_pt |>
  sccomp_estimate(
    formula_composition = ~ type + s(pseudotime, k = 3),
    formula_variability = ~ 1,
    sample              = "sample",
    cell_group          = "cell_group",
    abundance           = "count",
    cores               = 1,
    verbose             = FALSE,
    max_sampling_iterations = 2000
  )
```

```{r predict-grid-k3, message=FALSE, warning=FALSE, eval = run_smooth_vignette}
pred_k3 <- fit_k3 |>
  sccomp_predict(new_data = grid, number_of_draws = 200)
```

```{r plot-curves-k3, fig.width=8, fig.height=5, eval = run_smooth_vignette}
pred_k3 |>
  filter(cell_group %in% top_cgs) |>
  ggplot(aes(pseudotime, proportion_mean, colour = type, fill = type)) +
  geom_ribbon(aes(ymin = proportion_lower, ymax = proportion_upper),
              alpha = 0.15, linewidth = 0) +
  geom_line(linewidth = 0.7) +
  geom_point(data = raw, aes(y = prop, colour = type), alpha = 0.5, size = 1.2) +
  facet_wrap(~ cell_group, scales = "free_y", ncol = 3) +
  theme_minimal(base_size = 10)
```

Compared with the `k = 5` figure above, the credible bands are often a
little tighter and the lines less wiggly: the model has less capacity to
bend, which is appropriate when you want a gentle pseudotime trend on top
of the `type` effect.

## 4. Under the hood: What `s()` actually decomposes into

The smooth metadata used to fit the model is stored on the result, so we
can re-evaluate the basis at any new pseudotime value. This is exactly what
`sccomp_predict()` does internally.

```{r basis-decomp, fig.width=8, fig.height=6, eval = run_smooth_vignette}
smooth_meta    <- sccomp:::get_smooth_results(fit)
smooth_specs   <- smooth_meta$smooth_specs
smooth_re_objs <- smooth_meta$smooth_re_objs

grid_basis <- tibble(
  pseudotime = seq(min(pt_map$pseudotime), max(pt_map$pseudotime),  length.out = 200),
  type = levels(counts_obj$type)[1]  # any constant works; smooth ignores it
)

basis <- sccomp:::predict_smooth_at_newdata(
  smooth_specs[[1]], smooth_re_objs[[1]], grid_basis
)
# Plain s() has one penalised block — for multi-penalty smooths
# (e.g. bs = "fs", t2(...)) `Xr_blocks` would contain several matrices.
Xr_single <- basis$Xr_blocks[[1]]

basis_long <- bind_rows(
  as.data.frame(basis$Xf) |>
    setNames(sprintf("lin%d", seq_len(ncol(basis$Xf)))) |>
    mutate(pseudotime = grid_basis$pseudotime) |>
    pivot_longer(-pseudotime, names_to = "column", values_to = "value") |>
    mutate(block = "Xf  (unpenalised, joins fixed effects)"),
  as.data.frame(Xr_single) |>
    setNames(sprintf("basis%d", seq_len(ncol(Xr_single)))) |>
    mutate(pseudotime = grid_basis$pseudotime) |>
    pivot_longer(-pseudotime, names_to = "column", values_to = "value") |>
    mutate(block = "Xr  (penalised, joins one RE slot, shared sds)")
)

ggplot(basis_long, aes(pseudotime, value, colour = column)) +
  geom_line(linewidth = 0.7) +
  facet_wrap(~ block, scales = "free_y", ncol = 1) +
  labs(x = "pseudotime", y = "basis value",
       title = "Decomposition of s(pseudotime, k = 5)",
       subtitle = "Linear null-space (Xf) is a fixed effect; wiggly basis (Xr) is one RE slot") +
  theme_minimal(base_size = 10)
```

The top panel is the **linear trend**, the bottom panel is the **wiggly basis**: 
each of these columns gets one coefficient per cell group, and a
*single* smoothing standard deviation `sds` per cell group controls how
much wiggle is allowed. Increasing `k` would add more wiggly basis columns;
the penalty would then keep the extra capacity from over-fitting.

# Hierarchical smooths: per-group curves with shared structure

A common ask is: "fit a smooth per tissue (or per condition, per donor, ...)
but partial-pool across groups". The shorthand for this (also used in `mgcv`) is
`s(x, factor, bs = "fs")`. It builds one common
wiggly curve and a per-level deviation, governed jointly by a small set of
smoothing parameters. 

## 1. Fit ~ s(pseudotime, tissue, bs = "fs", k = 5)

We extend `counts_obj` with a synthetic `tissue` factor (3 levels) on top of
the existing `pseudotime`:

```{r make-tissue, message=FALSE}
set.seed(11)
samples <- levels(counts_obj$sample)
covs <- tibble::tibble(
  sample     = samples,
  pseudotime = seq(0, 6, length.out = length(samples)) +
                 rnorm(length(samples), 0, 0.05),
  tissue     = factor(rep(c("blood", "lymph", "tumor"),
                          length.out = length(samples)))
)
counts_tissue <- counts_obj |> dplyr::left_join(covs, by = "sample")
counts_tissue |> dplyr::distinct(sample, type, pseudotime, tissue) |> head()
```

```{r fs-fit, message=FALSE, warning=FALSE, eval = run_smooth_vignette}
fit_fs <- counts_tissue |>
  sccomp_estimate(
    formula_composition = ~ s(pseudotime, tissue, bs = "fs", k = 5),
    formula_variability = ~ 1,
    sample              = "sample",
    cell_group          = "cell_group",
    abundance           = "count",
    cores               = 1,
    verbose             = FALSE,
    max_sampling_iterations = 2000
  )

# This one smooth occupies 3 RE slots (one per mgcv penalty block);
# you can confirm via:
attr(fit_fs, "model_input")$ncol_X_random_eff
```

The three non-zero entries correspond to the wiggly main block (9 columns
for `k = 5` × 3 tissues) and two smaller per-level deviation blocks.

## 2. Predict and plot per-tissue curves

`sccomp_predict()` takes a grid that crosses `pseudotime` with `tissue`:

```{r fs-predict, message=FALSE, warning=FALSE, eval = run_smooth_vignette}
grid_fs <- expand.grid(
  pseudotime = seq(min(covs$pseudotime), max(covs$pseudotime),
                   length.out = 60),
  tissue     = levels(covs$tissue)
) |>
  as_tibble() |>
  mutate(sample = sprintf("fs_grid_%03d", row_number()))

pred_fs <- fit_fs |>
  sccomp_predict(new_data = grid_fs, number_of_draws = 200)

head(pred_fs)
```

Pick the most pseudotime-responsive cell groups and plot one curve per tissue:

```{r fs-plot, fig.width=8.5, fig.height=5.5, eval = run_smooth_vignette}
top_cgs_fs <- pred_fs |>
  group_by(cell_group) |>
  summarise(rng = max(proportion_mean) - min(proportion_mean)) |>
  arrange(desc(rng)) |>
  slice_head(n = 6) |>
  pull(cell_group)

raw_fs <- counts_tissue |>
  group_by(sample) |>
  mutate(prop = count / sum(count)) |>
  ungroup() |>
  filter(cell_group %in% top_cgs_fs)

pred_fs |>
  filter(cell_group %in% top_cgs_fs) |>
  ggplot(aes(pseudotime, proportion_mean,
             colour = tissue, fill = tissue)) +
  geom_ribbon(aes(ymin = proportion_lower, ymax = proportion_upper),
              alpha = 0.15, linewidth = 0) +
  geom_line(linewidth = 0.7) +
  geom_point(data = raw_fs, aes(y = prop, colour = tissue),
             alpha = 0.5, size = 1.2) +
  facet_wrap(~ cell_group, scales = "free_y", ncol = 3) +
  labs(x = "pseudotime", y = "proportion",
       title = "Factor smooth s(pseudotime, tissue, bs = 'fs', k = 5)",
       subtitle = "Per-tissue curves sharing one factor-smooth structure") +
  theme_minimal(base_size = 10)
```

The three coloured curves are the same smooth applied to different tissues,
allowed to deviate at each pseudotime point. Where the deviations are small
(e.g. CD8 1) the curves track each other; where they are large
(e.g. CD8 2) one tissue diverges. Whether the deviations are "real" is
expressed by the credible-interval ribbons — overlap means the data
don't support a difference.

## fs vs `s(x, by = tissue)`

`s(x, by = tissue)` is an alternative spelling that *also* gives a per-tissue
curve, but it treats each tissue's smooth as independent: each level gets
its own smoothing parameter, and the curves can have totally unrelated
shape. 
Use `bs = "fs"` when you want a single shared smoothness scale and partial
pooling across levels; use `by =` when the tissues are clearly heterogeneous
and a shared smoothness assumption would be misleading.


# Multiple smooths and mixing with random effects

`sccomp` can model up to two splines together with random effects. 

```{r multi-smooth-formula, eval=FALSE}
~ type +
  s(pseudotime, k = 5) +
  s(age,        k = 4) +
  (1 | donor)
```


# Session info

```{r session-info, echo=FALSE}
sessionInfo()
```