【Rでブラインド音源分離】Second-Order Blind Identification（SOBI）入門

複数のセンサで観測した信号が、実は複数の独立した「音源」の混合であった、という状況は、信号処理の分野ではよく見られます。そのような問題はブラインド音源分離（Blind Source Separation, BSS）と呼ばれます。
　BSS の代表的な手法としては ICA（Independent Component Analysis）が広く知られていますが、時系列データを対象とする場合には、ICA よりも自然で、かつ強力に機能する方法が存在します。それが SOBI（Second-Order Blind Identification）です。
　我々のグループでは、これまで振動方向の異なるフラクタル変動を分解する方法として oriented fractal scaling components analysis (OFSCA) を提案してきました。しかし、実際比べてみると、SOBI でうまくいく場合が多いと感じました。そこで本稿では、SOBI を実際に用い、その有効性を確認してみたいと思います。

1. ブラインド音源分離の問題設定
2. SOBIとは
3. RでSOBIへの最短ルート
4. SOBIアルゴリズムのポイント
まとめ
[おまけ] n次元のSOBI

1. ブラインド音源分離の問題設定

　まずは、ブラインド音源分離の問題設定を説明します。
　ここでは、複数のセンサで観測された信号が、実は複数の「元の信号 (音源)」の線形混合として得られている、と仮定します。これを数式で書くと、次のようになります。

$\displaystyle{ \mathbf{x}(t) = \mathbf{A}\,\mathbf{s}(t) }$

ここで、

$\mathbf{s}(t)$ ：元の音源信号を成分にもつベクトルそれぞれの成分は互いに独立で、異なる物理的・統計的性質をもつと仮定します。
$\mathbf{A}$ ： 未知の混合行列 各音源が、どの程度ずつ各センサに影響を与えているかを表します。
$\mathbf{x}(t)$ ： 実際に観測される信号 センサで記録できるのは、この混合後の信号だけです。

　例として、音源が 2 つ、センサも 2 つの場合を考えてみましょう。このとき、各ベクトルと行列は次のように書けます。

$\displaystyle{ \mathbf{s}(t) = \begin{pmatrix} s_1(t) \\ s_2(t) \end{pmatrix}, \qquad \mathbf{x}(t) = \begin{pmatrix} x_1(t) \\ x_2(t) \end{pmatrix}, \qquad \mathbf{A} = \begin{pmatrix} a_{11} & a_{12} \\ a_{21} & a_{22} \end{pmatrix} }$

線形混合は

$\displaystyle{ \begin{pmatrix} x_1(t) \\ x_2(t) \end{pmatrix} = \begin{pmatrix} a_{11} & a_{12} \\ a_{21} & a_{22} \end{pmatrix} \begin{pmatrix} s_1(t) \\ s_2(t) \end{pmatrix} }$

となります。これを成分ごとに書き下すと、

$\displaystyle{ \begin{aligned} x_1(t) &= a_{11}\,s_1(t) + a_{12}\,s_2(t), \\ x_2(t) &= a_{21}\,s_1(t) + a_{22}\,s_2(t). \end{aligned} }$

となります。

　このモデルは、決して特殊なものではありません。たとえば、

複数のマイクで録音した音声
複数電極で測定した脳波や心電図
複数方向のセンサで記録した振動信号

などでは、それぞれのセンサが「一つの音源だけ」を見ていることはほとんどなく、複数の要因が重なった結果が観測されます。この「重なり方」が線形で近似できる、というのが $\mathbf{x}(t) = \mathbf{A}\, \mathbf{s}(t)$ という仮定です。
　「ブラインド」と呼ばれるのは「元の独立な音源成分 $\mathbf{s}(t)$ は未知」「混合の仕方 $\mathbf{A}$ も未知」だからです。分かっているのは、観測信号 $\mathbf{x}(t))$ だけです。つまり、事前情報がほとんど無い状態（ブラインド）で、元の音源を推定しなければなりません。これが「ブラインド音源分離（Blind Source Separation）」と呼ばれる理由です。
　この状況での目標は「観測信号 $\mathbf{x}(t)$ しか与えられていない状態で、元の音源 $\mathbf{s}(t)$ を（順序とスケールの違いを除いて）復元することです。
　上の 2 次元の例では、観測されている $x _ 1(t), x _ 2(t)$ の 2 つの信号だけから、混合行列 $\mathbf{A}$ を推定し、元の音源 $s _ 1(t), s _ 2(t)$ を抽出します。とはいえ、元の音源を完全に一意に復元することは原理的にできません。 $s _ 1(t), s _ 2(t)$ の順序と、値の大きさを、決めることはできません。それでも、元の音源と同じ時間的な振る舞いをもつ信号が得られることには大きな価値があります。

2. SOBIとは

　SOBI（Second-Order Blind Identification）は、「信号の時間的な自己相関の違いを手がかりに、混合信号を分離する方法」です。自己相関、あるいはそれと等価なパワースペクトルが、時系列データを扱う際に最も基本的で重要な特徴量であることは、本ブログでも繰り返し説明してきました。
　自己相関とは、信号の「現在」と「過去（あるいは未来）」との間にどの程度の関係性があるかを記述する統計量です。言い換えれば、信号がどれくらい時間的な記憶を持っているかを定量化する指標だと言えます。
　一方、パワースペクトルは、この自己相関の情報を周波数の観点から表現し直したものに過ぎません。両者はフーリエ変換を介して互いに一対一に対応しており、本質的には同じ情報を別の視点から見ているだけです。
　ICA が利用するのは、主に信号の分布の非ガウス性です。これは「時刻をシャッフルしても変わらない情報」であり、時間構造そのものには直接触れていません。
　それに対して SOBI は、「信号が時間的にどのように変動しているか」という、時系列として最も自然な特徴に注目します。その特徴量は、時間遅れを導入した共分散（2次統計量）です。SOBI という名前が示すとおり、この手法は Second-Order（2次）統計量だけを用いて分離を行います。

3. RでSOBIへの最短ルート

　私としては、理屈を説明したいですが、それを我慢でして、RでSOBIを実行してみましょう。
　SOBIを体験するために、以下のRスクリプトを実行してください。

############################################################
# Minimal SOBI demo in R (2 sources / 2 sensors)
#  - Sources: AR(1) and AR(2)
#  - Mixing:  x(t) = A s(t)
#  - Separation: SOBI using lagged covariance joint diagonalization (2x2 analytic)
############################################################

# -----------------------------
# 0) Helpers
# -----------------------------
zscore <- function(x) as.numeric((x - mean(x)) / sd(x))

# AR simulation (base R)
sim_ar1 <- function(n, a1) {
  as.numeric(arima.sim(model = list(ar = a1), n = n))
}
sim_ar2 <- function(n, a1, a2) {
  as.numeric(arima.sim(model = list(ar = c(a1, a2)), n = n))
}

# Autocorrelation estimation (robust extraction)
acf_est <- function(x, lag.max = 50) {
  a <- acf(x, lag.max = lag.max, plot = FALSE)
  # ---- FIX: robustly extract vectors regardless of array dims
  list(lag = as.numeric(a$lag), acf = as.numeric(a$acf))
}

# -----------------------------
# 1) SOBI for 2 channels (minimal)
# -----------------------------
sobi_2ch <- function(X, lags = 1:50) {
  # X: 2 x N matrix
  if (!is.matrix(X)) X <- as.matrix(X)
  if (nrow(X) != 2) stop("sobi_2ch: X must be 2 x N.")
  N <- ncol(X)

  # center
  Xc <- X - rowMeans(X)

  # whiten: Z = V Xc, Cov(Z)=I
  R0 <- (Xc %*% t(Xc)) / N
  eg <- eigen(R0, symmetric = TRUE)
  V  <- diag(1 / sqrt(eg$values), 2, 2) %*% t(eg$vectors)
  Z  <- V %*% Xc

  # symmetric lagged covariance in whitened space
  cov_lag_sym <- function(Z, tau) {
    Z1 <- Z[, (tau + 1):N, drop = FALSE]
    Z0 <- Z[, 1:(N - tau), drop = FALSE]
    R  <- (Z1 %*% t(Z0)) / (N - tau)
    0.5 * (R + t(R))
  }
  Rs <- lapply(lags, function(tau) cov_lag_sym(Z, tau))

  # 2x2 joint diagonalization by a single rotation:
  # For each R=[[a,b],[b,c]], define u=[b, (c-a)/2].
  U <- do.call(rbind, lapply(Rs, function(R) c(R[1,2], (R[2,2] - R[1,1]) / 2)))
  C <- t(U) %*% U

  # eigenvector of smallest eigenvalue gives direction minimizing off-diagonals
  egC <- eigen(C, symmetric = TRUE)
  v   <- egC$vectors[, which.min(egC$values)]
  phi <- atan2(v[2], v[1])
  th  <- phi / 2

  Rrot <- matrix(c(cos(th), -sin(th),
                   sin(th),  cos(th)), 2, 2, byrow = TRUE)

  # unmixing and estimated sources
  W    <- t(Rrot) %*% V
  Shat <- W %*% Xc

  list(Shat = Shat, W = W)
}

# Align permutation/sign for evaluation (2 sources)
align_2 <- function(Shat, Strue) {
  c11 <- abs(cor(Shat[1,], Strue[1,]))
  c12 <- abs(cor(Shat[1,], Strue[2,]))
  c21 <- abs(cor(Shat[2,], Strue[1,]))
  c22 <- abs(cor(Shat[2,], Strue[2,]))

  if ((c12 + c21) > (c11 + c22)) Shat <- Shat[c(2,1),,drop=FALSE]
  s1 <- sign(cor(Shat[1,], Strue[1,])); if (s1 == 0) s1 <- 1
  s2 <- sign(cor(Shat[2,], Strue[2,])); if (s2 == 0) s2 <- 1
  Shat[1,] <- s1 * Shat[1,]
  Shat[2,] <- s2 * Shat[2,]
  list(Shat = Shat,
       cor1 = cor(Shat[1,], Strue[1,]),
       cor2 = cor(Shat[2,], Strue[2,]))
}

# -----------------------------
# 2) Generate sources (AR(1), AR(2))
# -----------------------------
N <- 5000

s1 <- zscore(sim_ar1(N, a1 = 0.85))
# pick a stable AR(2) (roots outside unit circle)
s2 <- zscore(sim_ar2(N, a1 = 1.30, a2 = -0.75))

S <- rbind(s1, s2)  # 2 x N

# -----------------------------
# 3) Mix
# -----------------------------
A <- matrix(c(1.0, 0.9,
              0.7, 1.2), 2, 2, byrow = TRUE)
X <- A %*% S

# -----------------------------
# 4) SOBI separate
# -----------------------------
out  <- sobi_2ch(X, lags = 1:60)
algn <- align_2(out$Shat, S)

cat(sprintf("corr(true s1, est1) = %.4f\n", algn$cor1))
cat(sprintf("corr(true s2, est2) = %.4f\n", algn$cor2))

Shat <- algn$Shat

# -----------------------------
# 5) Quick visualization
# -----------------------------
idx <- 1:1000

lag.max <- 40   # ---- FIX: define lag.max before use

op <- par(no.readonly = TRUE)
par(mfrow = c(4, 2), mar = c(3, 3, 2, 1))

## 1st row: true sources
plot(S[1, idx], type = "l", col = "blue", main = "True s1: AR(1)", xlab = "", ylab = "")
plot(S[2, idx], type = "l", col = "blue", main = "True s2: AR(2)", xlab = "", ylab = "")

## 2nd row: observed mixtures
plot(X[1, idx], type = "l", main = "Observed x1 (mixture)", xlab = "", ylab = "")
plot(X[2, idx], type = "l", main = "Observed x2 (mixture)", xlab = "", ylab = "")

## 3rd row: SOBI estimates
plot(Shat[1, idx], type = "l", col = "red", main = "SOBI est1 (aligned)", xlab = "", ylab = "")
plot(Shat[2, idx], type = "l", col = "red", main = "SOBI est2 (aligned)", xlab = "", ylab = "")

## 4th row: ACF comparison
acf_s1  <- acf_est(S[1, ],    lag.max)
acf_e1  <- acf_est(Shat[1, ], lag.max)
acf_s2  <- acf_est(S[2, ],    lag.max)
acf_e2  <- acf_est(Shat[2, ], lag.max)

# s1 vs est1
plot(acf_s1$lag, acf_s1$acf,
     type = "n",
     main = "ACF comparison: s1",
     xlab = "Lag", ylab = "ACF",
     ylim = range(c(acf_s1$acf, acf_e1$acf), finite = TRUE))

# vertical lines (True)
segments(acf_s1$lag, 0, acf_s1$lag, acf_s1$acf,
         col = "black")

# vertical lines (SOBI)
segments(acf_e1$lag, 0, acf_e1$lag, acf_e1$acf,
         col = "red")

# points
points(acf_s1$lag, acf_s1$acf,
       pch = 16, cex = 0.8, col = "blue")
points(acf_e1$lag, acf_e1$acf,
       pch = 1, cex = 1.1, col = "red")

legend("topright",
       legend = c("True", "SOBI"),
       col = c("black", "red"),
       pch = c(16, 1),
       pt.cex = c(0.8, 1.1),
       bty = "n")

# s2 vs est2
plot(acf_s2$lag, acf_s2$acf,
     type = "n",
     main = "ACF comparison: s2",
     xlab = "Lag", ylab = "ACF",
     ylim = range(c(acf_s2$acf, acf_e2$acf), finite = TRUE))

# vertical lines (True)
segments(acf_s2$lag, 0, acf_s2$lag, acf_s2$acf,
         col = "black")

# vertical lines (SOBI)
segments(acf_e2$lag, 0, acf_e2$lag, acf_e2$acf,
         col = "red")

# points
points(acf_s2$lag, acf_s2$acf,
       pch = 16, cex = 0.8, col = "blue")
points(acf_e2$lag, acf_e2$acf,
       pch = 1, cex = 1.1, col = "red")

legend("topright",
       legend = c("True", "SOBI"),
       col = c("black", "red"),
       pch = c(16, 1),
       pt.cex = c(0.8, 1.1),
       bty = "n")
par(op)

この R スクリプトを実行すると、下図が描かれます。

　上図の最上段の 2 つのパネルは、今回用意した元の音源信号（青実線）です。

左：True s1: AR(1) 1 次の自己回帰過程で生成した信号で、強い正の自己相関をもち、時間的に「なめらかな」変動が見られます。
右：True s2: AR(2) 2 次の自己回帰過程で生成した信号で、自己相関に振動成分が含まれ、やや周期的な構造が現れています。

この 2 つの信号は、分布の形だけを見るとどちらもガウス（正規分布）ですが、時間構造（自己相関）は異なっています。
　上図の2 段目は、これら 2 つの音源を未知の混合行列で線形に混合した、 観測信号です。

Observed x1 (mixture)
Observed x2 (mixture)

どちらの信号も、見た目だけでは「どの成分が AR(1) 由来で、どの成分が AR(2) 由来か」を判別することは困難です。これがブラインド音源分離問題の出発点です。
　上図の3 段目は、SOBI によって推定された 分離信号です。

SOBI est1 (aligned)
SOBI est2 (aligned)

順序とスケールは調整してありますが、波形の時間的な振る舞いを見ると、

est1 は AR(1) 型の変動
est2 は AR(2) 型の変動

をそれぞれよく再現していることが分かります。ただし、SOBI が本当に正しく分離できているかどうかは、波形の見た目だけでは判断できません。
　そこで、最も重要なのが、上図最下段の 自己相関関数（ACF）の比較です。

黒 ●：元の音源（True）
赤 ○：SOBI による推定信号

上図最下段左のプロットには、AR(1) に特徴的な「ラグとともに単調に減衰する自己相関」が、SOBI 推定信号でもほぼ完全に再現されています。黒と赤の丸が、ほぼ重なっていることから、SOBI により、元信号のAR(1) の時間構造が忠実に再現されていることが分かります。上図最下段右のプロットには、AR(2) に特有の「振動を伴いながら減衰する自己相関構造」も、SOBI によって正確に再現されています。

4. SOBIアルゴリズムのポイント

　今回は詳細には踏み込みませんが、最低限、SOBI アルゴリズムのキーアイデアを整理しておきます。

4.1 ポイント1：白色化による問題の単純化

　SOBI の最初の重要なステップは、白色化（whitening）です。
　観測信号 $\mathbf{x}(t)$ は、一般に成分間で相関をもっています。そこで、共分散行列を用いた線形変換によって、

平均 0
分散 1
相互に無相関

という状態に正規化します。この操作によって、未知の混合行列 $\mathbf{A}$ は、 直交行列（回転行列）で書けることになります。
　たとえば２次元の場合、観測信号を

$\displaystyle{ \mathbf{x}(t)= \begin{pmatrix} x_1(t)\\ x_2(t) \end{pmatrix} }$

とし、平均を引いた信号を $\mathbf{x} _ c(t)$ とします。
　共分散行列は

$\displaystyle{ \mathbf{R}_x(0) = \mathbb{E}\!\left[\mathbf{x}_c(t)\, \mathbf{x}_c(t)^{\mathsf T}\right] }$

で与えられ、対称行列なので固有値分解できます。

$\displaystyle{ \mathbf{R}_x(0) = \mathbf{E}\mathbf{\Lambda}\mathbf{E}^{\mathsf T}, \qquad \mathbf{\Lambda}=\mathrm{diag}(\lambda_1,\lambda_2) }$

　白色化変換は

$\displaystyle{ \mathbf{z}(t) = \mathbf{V}\mathbf{x}_c(t), \qquad \mathbf{V} = \mathbf{\Lambda}^{-1/2}\mathbf{E}^{\mathsf T} }$

と定義します。このとき

$\displaystyle{ \mathbb{E}\!\left[\mathbf{z}(t)\, \mathbf{z}(t)^{\mathsf T}\right] = \mathbf{I} }$

が成り立ち、信号は白色化されています。
　元の混合モデル

$\displaystyle{ \mathbf{x}(t)=\mathbf{A}\, \mathbf{s}(t) }$

は、白色化後に

$\displaystyle{ \mathbf{z}(t) = \mathbf{V}\mathbf{A}\, \mathbf{s}(t) \equiv \mathbf{B}\, \mathbf{s}(t) }$

と書けます。このとき、2 次元では混合行列 $\mathbf{B}$ は直交行列として扱え、

$\displaystyle{ \mathbf{B} = \begin{pmatrix} \cos\theta & -\sin\theta\\ \sin\theta & \cos\theta \end{pmatrix} }$

すなわち、未知量は回転角 $\theta$ だけに簡約されます。
　このように、白色化は SOBI の問題を「未知行列を推定する問題」から「回転角を決める問題」へ落とす重要な前処理です。

4.2 ポイント2：複数ラグの自己相関を同時に使う

　白色化された信号に対して、SOBI は 時間遅れを導入した共分散行列を計算します。白色化された信号を

$\displaystyle{ \mathbf{z}(t)= \begin{pmatrix} z_1(t)\\ z_2(t) \end{pmatrix}, \qquad \mathbb{E}\!\left[\mathbf{z}(t)\mathbf{z}(t)^{\mathsf T}\right] =\mathbf{I} }$

とします。SOBI では、この信号に対して 時間遅れ $\tau$ を導入した共分散行列

$\displaystyle{ \mathbf{R}_z(\tau) = \mathbb{E}\!\left[ \mathbf{z}(t)\mathbf{z}(t-\tau)^{\mathsf T} \right] }$

を計算します。具体的には、

$\displaystyle{ \tau = 1, 2, 3, \ldots }$

といった複数のラグについて、 $\mathbf{R} _ z(\tau)$ を求めます。
　 $\mathbf{z}(t)$ は、前処理で白色化されているので、

$\displaystyle{ \mathbf{R}_z(0)=\mathbf{I} }$

ですが、 $\displaystyle{ \tau\gt0 }$ では一般に

$\displaystyle{ \mathbf{R}_z(\tau) = \begin{pmatrix} r_{11}(\tau) & r_{12}(\tau)\\ r_{21}(\tau) & r_{22}(\tau) \end{pmatrix} }$

となり、非対角成分が残ります。
　SOBI の基本的な考え方は「すべての $\tau$ に対して、 $\mathbf{R} _ z(\tau)$ を同時に対角化する回転を求める」というものです。

4.3 ポイント3：同時対角化

　SOBI の核心は、同時対角化（joint diagonalization）にあります。白色化後の混合モデルは

$\displaystyle{ \mathbf{z}(t)=\mathbf{B}\, \mathbf{s}(t) }$

と書け、ここで $\mathbf{B}$ は直交行列です。このとき、

$\displaystyle{ \mathbf{R}_z(\tau) = \mathbf{B}\, \mathbf{R}_s(\tau)\, \mathbf{B}^{\mathsf T} }$

が成り立ちます。音源 $\mathbf{s}(t)$ が互いに独立であれば、 $\mathbf{R} _ s(\tau)$ はすべての $\tau$ に対して対角行列になります。
　ここで重要なのは、単一の時間遅れに注目しただけでは回転行列は決まらないという点です。ある 1 つのラグに対して共分散行列を対角化する回転は、一般にいくつも存在します。したがって、単一の自己相関情報だけでは、音源を一意に分離することはできません。
　しかし、複数の時間遅れに対する共分散行列を同時に対角化することを要求すると状況は大きく変わります。すべてのラグで同時に対角となる回転は、通常ただ 1 つに定まるため、回転行列が事実上一意に決まります。
　このように、「同時に対角化する」という強い制約を課すことによって、混合を打ち消す回転が特定され、音源分離が可能になります。

4.4 ポイント4：最適な回転行列の求め方

　音源が 2 つの場合、白色化後の回転は、1 つの回転角 (\theta) だけで表されます。このとき、各ラグにおける「非対角成分の大きさ」は、(\theta) の三角関数として解析的に書くことができます。
　SOBI では「すべてのラグにおける非対角成分の二乗和を最小にする」という 1 変数の最適化問題を解きますが、この問題は 固有値問題に帰着でき、最適な回転角は 閉形式で一意に求まります。
　そのため、2 次元の SOBI では、

探索不要
初期値不要
局所解の心配なし

という、数値的に非常に安定した計算が可能です。
　一方、3 次元以上になると、

回転の自由度が増える
単一の解析解は存在しない

ため、反復的な同時対角化アルゴリズムが用いられます。
　具体的には、

成分のペアを 1 組ずつ選び
その 2 次元部分空間で最適な回転（←これは解析的）
これを全ペアについて繰り返す

という Jacobi 型の反復回転が使われます。
　重要なのは、各 2×2 回転自体は、2 次元の場合と同じ解析解を用いているという点です。つまり、探索的に見えても、各ステップは理論的に明確な計算から成り立っています。

まとめ

　正直に言うと、oriented fractal scaling components analysis（OFSCA）をさらに洗練させるために SOBI を改めて勉強し直しましたが、その過程で、少なくとも実用的な音源分離という観点では、OFSCA は SOBI に太刀打ちできないのではないかと感じるようになりました。
　それでも、SOBI の限界を丁寧に見極めた上で、OFSCA ならではの強みや独自の価値を見いだせないか、引き続き検討していきたいと考えています。

※ もし記事の中で「ここ違うよ」という点や気になるところがあれば、気軽に指摘していただけると助かります。質問や「このテーマも取り上げてほしい」といったリクエストも大歓迎です。必ず対応するとは約束できませんが、できるだけ今後の記事で扱いたいと思います。それと、下のはてなブログランキングはあまり信用できる指標ではなさそうですが (私のブログを読んでいる人は、実際とても少ないです)、押してもらえるとシンプルに励みになります。気が向いたときにポチッとしていただけたら嬉しいです。

[おまけ] n次元のSOBI

　以下の R スクリプトは、複数の時系列が線形に混合された信号から、元の音源を SOBI によって分離するための最小限の実装例です。ここでは、3 つの音源（AR(1)、AR(2)、AR(2)）を用いた数値例を示しています。
　特別なパッケージは必要ありません。R コンソールまたは RStudio 上で、スクリプト全体をそのまま実行してください。実行すると、

3 つの AR 過程からなる元信号を生成
それらを未知の混合行列で線形に混合
SOBI により音源分離を実行
結果を図として可視化

という一連の処理が自動的に行われます。
　描画される図は 4 行 × 3 列の構成になっています。

1 行目：True 元の音源信号（AR(1)、AR(2)、AR(2)）
2 行目：Observed 混合後の観測信号（各チャンネル）
3 行目：SOBI estimated SOBI によって推定された音源（順序と符号は自動整列）
4 行目：ACF comparison 元の音源（●）と推定音源（○）の自己相関の比較 → 時間構造が正しく復元されているかを確認できます

　以下の部分を変更することで、挙動を簡単に調整できます。

lags = 1:60：SOBI に用いる時間遅れの範囲
n_sweeps = 20：同時対角化の反復回数
AR 係数（sim_ar() の引数）：音源の時間構造を変更

############################################################
# SOBI in R for n channels (n sources / n sensors)
#  - Sources demo: AR(1), AR(2), AR(2)  (n = 3)
#  - Mixing: X(t) = A S(t)
#  - SOBI: whitening + joint diagonalization (Jacobi sweeps)
#  - Plot: 4 rows x n columns (True / Observed / Estimated / ACF compare)
############################################################

# -----------------------------
# 0) Helpers
# -----------------------------
zscore <- function(x) as.numeric((x - mean(x)) / sd(x))

sim_ar <- function(n, ar) {
  as.numeric(arima.sim(model = list(ar = ar), n = n))
}

acf_est <- function(x, lag.max = 50) {
  a <- acf(x, lag.max = lag.max, plot = FALSE)
  list(lag = as.numeric(a$lag), acf = as.numeric(a$acf))
}

# For n=3 demo: best permutation by brute force + sign fix
align_n_bruteforce <- function(Shat, Strue) {
  n <- nrow(Strue)
  if (n != nrow(Shat)) stop("align_n_bruteforce: dimension mismatch.")
  if (n > 8) stop("This brute-force aligner is intended for small n (<=8).")

  perms <- function(v) {
    if (length(v) == 1) return(list(v))
    out <- list()
    for (i in seq_along(v)) {
      rest <- v[-i]
      for (p in perms(rest)) out[[length(out) + 1]] <- c(v[i], p)
    }
    out
  }

  P <- perms(1:n)
  best_score <- -Inf
  best_perm <- 1:n

  for (perm in P) {
    score <- 0
    for (i in 1:n) score <- score + abs(cor(Shat[i, ], Strue[perm[i], ]))
    if (score > best_score) {
      best_score <- score
      best_perm <- perm
    }
  }

  Sh2 <- Shat
  # reorder true to match Sh2 rows
  S2 <- Strue[best_perm, , drop = FALSE]

  # sign fix per component
  for (i in 1:n) {
    sgn <- sign(cor(Sh2[i, ], S2[i, ]))
    if (sgn == 0) sgn <- 1
    Sh2[i, ] <- sgn * Sh2[i, ]
  }

  cors <- sapply(1:n, function(i) cor(Sh2[i, ], S2[i, ]))
  list(Shat = Sh2, Strue_reordered = S2, perm = best_perm, cors = cors)
}

# -----------------------------
# 1) Joint diagonalization (Jacobi) for symmetric matrices
#    Input: list of K symmetric n x n matrices (lagged covariances in whitened space)
#    Output: Q (orthogonal) s.t. Q^T R_k Q are "as diagonal as possible"
# -----------------------------
joint_diag_jacobi <- function(Rs, n, n_sweeps = 15, eps = 1e-12) {
  K <- length(Rs)
  Q <- diag(n)

  for (sweep in 1:n_sweeps) {
    changed <- FALSE

    for (p in 1:(n - 1)) {
      for (q in (p + 1):n) {

        # Build 2x2 criterion matrix C = Σ u_k u_k^T, u_k = [b_k, d_k]
        # where b_k = R_k[p,q], d_k = (R_k[q,q]-R_k[p,p])/2
        c11 <- 0; c12 <- 0; c22 <- 0
        for (k in 1:K) {
          Rk <- Rs[[k]]
          b  <- Rk[p, q]
          d  <- (Rk[q, q] - Rk[p, p]) / 2
          c11 <- c11 + b * b
          c12 <- c12 + b * d
          c22 <- c22 + d * d
        }
        C <- matrix(c(c11, c12, c12, c22), 2, 2)

        # Smallest-eigenvector gives direction minimizing off-diagonal energy
        eg <- eigen(C, symmetric = TRUE)
        v  <- eg$vectors[, which.min(eg$values)]
        phi <- atan2(v[2], v[1])
        th  <- 0.5 * phi

        # If rotation is negligible, skip
        if (abs(sin(th)) < eps) next

        cth <- cos(th); sth <- sin(th)

        # Apply Jacobi rotation to all matrices: R <- G^T R G
        # where G acts only on indices (p,q)
        for (k in 1:K) {
          Rk <- Rs[[k]]

          # Update rows p,q (right-multiply by G)
          Rp <- cth * Rk[, p] + sth * Rk[, q]
          Rq <- -sth * Rk[, p] + cth * Rk[, q]
          Rk[, p] <- Rp
          Rk[, q] <- Rq

          # Update cols p,q (left-multiply by G^T)
          Cp <- cth * Rk[p, ] + sth * Rk[q, ]
          Cq <- -sth * Rk[p, ] + cth * Rk[q, ]
          Rk[p, ] <- Cp
          Rk[q, ] <- Cq

          # Symmetrize to suppress numerical drift
          Rs[[k]] <- 0.5 * (Rk + t(Rk))
        }

        # Accumulate Q <- Q G
        G <- diag(n)
        G[p, p] <- cth; G[q, q] <- cth
        G[p, q] <- -sth; G[q, p] <- sth
        Q <- Q %*% G

        changed <- TRUE
      }
    }

    if (!changed) break
  }

  list(Q = Q, Rs_diag = Rs)
}

# -----------------------------
# 2) SOBI (n channels)
# -----------------------------
sobi_n <- function(X, lags = 1:60, n_sweeps = 15) {
  # X: n x N
  if (!is.matrix(X)) X <- as.matrix(X)
  n <- nrow(X); N <- ncol(X)

  # center
  Xc <- X - rowMeans(X)

  # whiten: Z = V Xc, Cov(Z)=I
  R0 <- (Xc %*% t(Xc)) / N
  eg <- eigen(R0, symmetric = TRUE)

  # Guard against tiny/negative eigenvalues from numerical issues
  vals <- pmax(eg$values, 1e-15)
  V <- diag(1 / sqrt(vals), n, n) %*% t(eg$vectors)
  Z <- V %*% Xc

  # lagged covariances in whitened domain (symmetric)
  cov_lag_sym <- function(Z, tau) {
    Z1 <- Z[, (tau + 1):N, drop = FALSE]
    Z0 <- Z[, 1:(N - tau), drop = FALSE]
    R  <- (Z1 %*% t(Z0)) / (N - tau)
    0.5 * (R + t(R))
  }
  Rs <- lapply(lags, function(tau) cov_lag_sym(Z, tau))

  # joint diagonalization -> Q
  jd <- joint_diag_jacobi(Rs, n = n, n_sweeps = n_sweeps)
  Q <- jd$Q

  # unmixing and estimated sources
  W <- t(Q) %*% V
  Shat <- W %*% Xc

  list(Shat = Shat, W = W, V = V, Q = Q)
}

# -----------------------------
# 3) Demo: n=3 sources (AR(1), AR(2), AR(2))
# -----------------------------
N <- 20000
n <- 3

# Sources (choose AR params so time-structure differs)
s1 <- zscore(sim_ar(N, ar = 0.85))                  # AR(1)
s2 <- zscore(sim_ar(N, ar = c(1.30, -0.75)))        # AR(2) oscillatory-ish
s3 <- zscore(sim_ar(N, ar = c(0.90, -0.20)))        # AR(2) different shape

S <- rbind(s1, s2, s3)  # 3 x N

# Mixing matrix (invertible)
A <- matrix(c(1.0,  0.6,  0.2,
              0.3,  1.2, -0.4,
             -0.2,  0.5,  1.1), 3, 3, byrow = TRUE)

X <- A %*% S

# SOBI
out <- sobi_n(X, lags = 1:60, n_sweeps = 20)

# Align estimated sources to true (perm + sign)
al <- align_n_bruteforce(out$Shat, S)
Shat <- al$Shat
S_aligned <- al$Strue_reordered

cat("Correlations after best alignment:\n")
for (i in 1:n) cat(sprintf("  corr(true s%d, est%d) = %.4f\n", i, i, al$cors[i]))

# -----------------------------
# 4) Plot: 4 rows x 3 columns (same structure as before)
# -----------------------------
idx <- 1:2000
lag.max <- 40

op <- par(no.readonly = TRUE)
par(mfrow = c(4, 3), mar = c(3, 3, 2, 1))

# Row 1: true sources
for (i in 1:n) {
  plot(S_aligned[i, idx], type = "l", col = "blue",
       main = sprintf("True s%d", i), xlab = "", ylab = "")
}

# Row 2: observed mixtures
for (i in 1:n) {
  plot(X[i, idx], type = "l",
       main = sprintf("Observed x%d (mixture)", i), xlab = "", ylab = "")
}

# Row 3: SOBI estimates (aligned)
for (i in 1:n) {
  plot(Shat[i, idx], type = "l", col = "red",
       main = sprintf("SOBI est%d (aligned)", i), xlab = "", ylab = "")
}

# Row 4: ACF comparison (True ● vs SOBI ○) with vertical stems
for (i in 1:n) {
  acf_s <- acf_est(S_aligned[i, ], lag.max)
  acf_e <- acf_est(Shat[i, ],      lag.max)

  ylim <- range(c(acf_s$acf, acf_e$acf), finite = TRUE)

  plot(acf_s$lag, acf_s$acf, type = "n",
       main = sprintf("ACF comparison: s%d", i),
       xlab = "Lag", ylab = "ACF", ylim = ylim)

  segments(acf_s$lag, 0, acf_s$lag, acf_s$acf, col = "black")
  segments(acf_e$lag, 0, acf_e$lag, acf_e$acf, col = "red")

  points(acf_s$lag, acf_s$acf, pch = 16, cex = 0.8, col = "blue")
  points(acf_e$lag, acf_e$acf, pch = 1,  cex = 1.1, col = "red")

  legend("topright",
         legend = c("True", "SOBI"),
         col = c("black", "red"),
         pch = c(16, 1),
         pt.cex = c(0.8, 0.9),
         bty = "n")
}

par(op)