diff --git a/README.ja.md b/README.ja.md new file mode 100644 --- /dev/null +++ b/README.ja.md @@ -0,0 +1,88 @@ +# hanalyze-bayes + +[`hanalyze`](../README.ja.md) の**ベイズ推論層**。 階層ベイズモデル +(HBM) の DSL・MCMC サンプラ群・自動微分・ベイズ的モデル比較を担う。 + +依存は `hanalyze-core` のみ (dataframe には依存しない)。 `-frame` と +並んで core の直上に位置するので、 **データ表現を持ち込まずに +サンプリング library としてだけ使える**のが特徴。 + +## 主要 module (全 26 module) + +### HBM DSL (`Hanalyze.Model.HBM.*`) + +| Module | 役割 | +|---|---| +| `Model.HBM` | 下位 8 module (Util / Distribution / Sampling / Model / Track / Eval / IR / Gradient) を束ねる facade。 通常はこれだけを import する | +| `Model.HBM.Model` | Free monad による多相モデル DSL 記述層 (`sample` / `observe` / `dataNamed*`) | +| `Model.HBM.Distribution` | 多相確率分布 ADT と密度・CDF | +| `Model.HBM.Eval` | log-joint / 尤度インタープリタ + DAG 構築 | +| `Model.HBM.Gradient` / `Model.HBM.IR` / `Model.HBM.VecAD` | AD 勾配コンパイラ層・中間表現・自作 reverse-mode AD (NUTS の per-draw ホット経路) | +| `Model.HBM.Ast` / `Model.HBM.Interp` | dialog DSL の AST + JSON decoder と評価系 | + +### サンプラ (`Hanalyze.MCMC.*`) + +| Module | 役割 | +|---|---| +| `MCMC.NUTS` | No-U-Turn Sampler — Hoffman & Gelman (2014) Algorithm 3 実装。 主力サンプラ | +| `MCMC.HMC` | Hamiltonian Monte Carlo (AD による exact gradient) | +| `MCMC.MH` / `MCMC.Slice` | Random-Walk Metropolis-Hastings / Slice sampler (Neal 2003) | +| `MCMC.Gibbs` | 共役事前分布向け Gibbs sampler (解析的フル条件付き) | +| `MCMC.SMC` | Tempered target による Sequential Monte Carlo | +| `MCMC.BayesianTest` | Bayesian A/B test — 2 群の平均差を NUTS でサンプルし ROPE / HDI で判定 | +| `MCMC.Progress` | 全 chain 集計の進捗表示 (stderr) | + +### 推論・モデル比較 (`Hanalyze.Stat.*`) + +| Module | 役割 | +|---|---| +| `Stat.AD` | HMC / NUTS を支える自動微分層 | +| `Stat.VI` | 変分推論 (ADVI) | +| `Stat.BridgeSampling` | Bridge Sampling による周辺尤度 log p(y) 推定 (Meng & Wong 1996) | +| `Stat.BayesFactor` | Bridge Sampling ベースの Bayes factor (Kass & Raftery 1995) | +| `Stat.BayesianModelAveraging` | log marginal を用いた真の BMA | +| `Stat.PosteriorPredictive` | 事前 / 事後予測サンプリング (PyMC の `sample_*_predictive` 相当) | + +## 単体で使う + +```cabal +-- Chain / 事後統計量は core 側の型なので、直接触るなら core も明示する +build-depends: hanalyze-bayes, hanalyze-core, containers +``` + +```haskell +{-# LANGUAGE OverloadedStrings #-} +import qualified Data.Map.Strict as Map +import Hanalyze.Model.HBM (ModelP, sample, observe, Distribution (..)) +import Hanalyze.MCMC.NUTS (nutsPure, defaultNUTSConfig) +import Hanalyze.MCMC.Core (posteriorMean, posteriorSD) + +myModel :: ModelP () +myModel = do + mu <- sample "mu" (Normal 0 10) + observe "y" (Normal mu 2) [1.2, 2.3, 3.1, 2.8, 1.9] -- observe は [Double] + +main = do + -- nutsPure は seed (Word32) を取り、純粋・決定的に Chain を返す + let chain = nutsPure myModel defaultNUTSConfig (Map.fromList [("mu", 0.0)]) 42 + print (posteriorMean "mu" chain, posteriorSD "mu" chain) + -- (Just 2.2673259586131507,Just 0.8340460004499451) +``` + +`Chain` は core の `Hanalyze.MCMC.Core` の型なので、 事後統計量 +(`posteriorMean` / `posteriorSD` / `posteriorQuantile`) と診断 +(`Hanalyze.Stat.MCMC` の `rhat` / `ess` / `hdi`) は core 側の API で扱える。 +HTML レポートや DAG 図が要る場合は `-viz` 層 (`Hanalyze.Viz.*`) を追加する。 + +## 関連 docs + +- 確率モデルの書き方: [docs/bayesian/02-probabilistic-model.ja.md](../docs/bayesian/02-probabilistic-model.ja.md) +- サンプラの選び方: [docs/bayesian/03-mcmc-samplers.ja.md](../docs/bayesian/03-mcmc-samplers.ja.md) +- Gibbs: [docs/bayesian/04-gibbs.ja.md](../docs/bayesian/04-gibbs.ja.md) / + 変分推論: [docs/bayesian/05-vi.ja.md](../docs/bayesian/05-vi.ja.md) +- モデル比較 (WAIC / LOO): [docs/bayesian/06-model-comparison.ja.md](../docs/bayesian/06-model-comparison.ja.md) +- 周辺尤度 / Bayes factor / BMA: [docs/bayesian/07-advanced-marginal-likelihood.ja.md](../docs/bayesian/07-advanced-marginal-likelihood.ja.md) +- Bayesian A/B test: [docs/bayesian/usage-bayesian-ab-test.ja.md](../docs/bayesian/usage-bayesian-ab-test.ja.md) +- 理論: [docs/bayesian/theory-hmc-nuts.ja.md](../docs/bayesian/theory-hmc-nuts.ja.md) + +← [repository README](../README.ja.md) diff --git a/README.md b/README.md new file mode 100644 --- /dev/null +++ b/README.md @@ -0,0 +1,91 @@ +# hanalyze-bayes + +The **Bayesian inference layer** of [`hanalyze`](../README.md): the +hierarchical Bayesian model (HBM) DSL, the MCMC samplers, automatic +differentiation, and Bayesian model comparison. + +It depends on `hanalyze-core` only — not on dataframe. Sitting next to +`-frame` directly above core, it can be used as a **standalone sampling library +without pulling in any data representation**. + +## Main modules (26 in total) + +### HBM DSL (`Hanalyze.Model.HBM.*`) + +| Module | Role | +|---|---| +| `Model.HBM` | Facade over 8 submodules (Util / Distribution / Sampling / Model / Track / Eval / IR / Gradient). Normally the only import you need | +| `Model.HBM.Model` | The polymorphic free-monad model DSL (`sample` / `observe` / `dataNamed*`) | +| `Model.HBM.Distribution` | Polymorphic distribution ADT with densities and CDFs | +| `Model.HBM.Eval` | log-joint / likelihood interpreter and DAG construction | +| `Model.HBM.Gradient` / `Model.HBM.IR` / `Model.HBM.VecAD` | AD gradient compiler, intermediate representation, and the in-house reverse-mode AD (the per-draw hot path of NUTS) | +| `Model.HBM.Ast` / `Model.HBM.Interp` | AST + JSON decoder of the dialog DSL and its interpreter | + +### Samplers (`Hanalyze.MCMC.*`) + +| Module | Role | +|---|---| +| `MCMC.NUTS` | No-U-Turn Sampler — Hoffman & Gelman (2014) Algorithm 3. The workhorse sampler | +| `MCMC.HMC` | Hamiltonian Monte Carlo with exact AD gradients | +| `MCMC.MH` / `MCMC.Slice` | Random-walk Metropolis-Hastings / slice sampler (Neal 2003) | +| `MCMC.Gibbs` | Gibbs sampler for conjugate priors (analytic full conditionals) | +| `MCMC.SMC` | Sequential Monte Carlo over a tempered target | +| `MCMC.BayesianTest` | Bayesian A/B test — samples the group mean difference with NUTS and decides via ROPE / HDI | +| `MCMC.Progress` | Progress display aggregated across chains (stderr) | + +### Inference and model comparison (`Hanalyze.Stat.*`) + +| Module | Role | +|---|---| +| `Stat.AD` | The automatic-differentiation layer behind HMC / NUTS | +| `Stat.VI` | Variational inference (ADVI) | +| `Stat.BridgeSampling` | Marginal likelihood log p(y) via bridge sampling (Meng & Wong 1996) | +| `Stat.BayesFactor` | Bayes factors on top of bridge sampling (Kass & Raftery 1995) | +| `Stat.BayesianModelAveraging` | True BMA from the log marginal likelihoods | +| `Stat.PosteriorPredictive` | Prior / posterior predictive sampling (PyMC's `sample_*_predictive`) | + +## Using it standalone + +```cabal +-- Chain and posterior statistics are core types, so depend on core too +-- if you touch them directly +build-depends: hanalyze-bayes, hanalyze-core, containers +``` + +```haskell +{-# LANGUAGE OverloadedStrings #-} +import qualified Data.Map.Strict as Map +import Hanalyze.Model.HBM (ModelP, sample, observe, Distribution (..)) +import Hanalyze.MCMC.NUTS (nutsPure, defaultNUTSConfig) +import Hanalyze.MCMC.Core (posteriorMean, posteriorSD) + +myModel :: ModelP () +myModel = do + mu <- sample "mu" (Normal 0 10) + observe "y" (Normal mu 2) [1.2, 2.3, 3.1, 2.8, 1.9] -- observe takes [Double] + +main = do + -- nutsPure takes a seed (Word32) and returns a Chain purely and deterministically + let chain = nutsPure myModel defaultNUTSConfig (Map.fromList [("mu", 0.0)]) 42 + print (posteriorMean "mu" chain, posteriorSD "mu" chain) + -- (Just 2.2673259586131507,Just 0.8340460004499451) +``` + +`Chain` is a core type (`Hanalyze.MCMC.Core`), so posterior summaries +(`posteriorMean` / `posteriorSD` / `posteriorQuantile`) and diagnostics +(`rhat` / `ess` / `hdi` in `Hanalyze.Stat.MCMC`) come from core. +Add the `-viz` layer (`Hanalyze.Viz.*`) when you need HTML reports or +DAG figures. + +## Related docs + +- Writing probabilistic models: [docs/bayesian/02-probabilistic-model.md](../docs/bayesian/02-probabilistic-model.md) +- Choosing a sampler: [docs/bayesian/03-mcmc-samplers.md](../docs/bayesian/03-mcmc-samplers.md) +- Gibbs: [docs/bayesian/04-gibbs.md](../docs/bayesian/04-gibbs.md) / + variational inference: [docs/bayesian/05-vi.md](../docs/bayesian/05-vi.md) +- Model comparison (WAIC / LOO): [docs/bayesian/06-model-comparison.md](../docs/bayesian/06-model-comparison.md) +- Marginal likelihood / Bayes factor / BMA: [docs/bayesian/07-advanced-marginal-likelihood.ja.md](../docs/bayesian/07-advanced-marginal-likelihood.ja.md) (ja only for now) +- Bayesian A/B test: [docs/bayesian/usage-bayesian-ab-test.md](../docs/bayesian/usage-bayesian-ab-test.md) +- Theory: [docs/bayesian/theory-hmc-nuts.md](../docs/bayesian/theory-hmc-nuts.md) + +← [repository README](../README.md) diff --git a/hanalyze-bayes.cabal b/hanalyze-bayes.cabal new file mode 100644 --- /dev/null +++ b/hanalyze-bayes.cabal @@ -0,0 +1,81 @@ +cabal-version: 3.0 +name: hanalyze-bayes +version: 0.2.0.1 +synopsis: Bayesian layer of hanalyze: HBM DSL, MCMC samplers, model comparison +description: + The Bayesian inference layer of the hanalyze toolkit, depending on + hanalyze-core only (no dataframe), so it can be used as a + standalone sampling library. A free-monad hierarchical Bayesian model DSL + with in-house reverse-mode AD, samplers (NUTS after Hoffman & Gelman 2014, + HMC, Metropolis-Hastings, slice, Gibbs, SMC), variational inference + (ADVI), posterior predictive sampling, and model comparison via bridge + sampling (marginal likelihood, Bayes factors, Bayesian model averaging). + . + Module names match the umbrella package hanalyze, which re-exports + everything, so downstream imports stay identical. See README.md for the + module map and a standalone usage example. +license: BSD-3-Clause +author: Toshiaki Honda +maintainer: frenzieddoll@gmail.com +copyright: 2026 Aelysce Project (Toshiaki Honda) +category: Math, Statistics, Numeric, Machine Learning +build-type: Simple +tested-with: GHC == 9.6.7 +extra-source-files: + README.md + README.ja.md + +common warnings + ghc-options: -Wall -Wcompat -Widentities -Wredundant-constraints + +-- -O2 は分割前と同一 (性能変更と構造変更を混ぜない、 層別 -O 調整は 106.5 後の別 Phase) +common opt + ghc-options: -O2 -funbox-strict-fields + +library + import: warnings, opt + hs-source-dirs: src + default-language: GHC2021 + exposed-modules: + Hanalyze.MCMC.BayesianTest + Hanalyze.MCMC.Gibbs + Hanalyze.MCMC.HMC + Hanalyze.MCMC.MH + Hanalyze.MCMC.NUTS + Hanalyze.MCMC.Progress + Hanalyze.MCMC.SMC + Hanalyze.MCMC.Slice + Hanalyze.Model.HBM + Hanalyze.Model.HBM.Ast + Hanalyze.Model.HBM.Distribution + Hanalyze.Model.HBM.Eval + Hanalyze.Model.HBM.Gradient + Hanalyze.Model.HBM.IR + Hanalyze.Model.HBM.Interp + Hanalyze.Model.HBM.Model + Hanalyze.Model.HBM.Sampling + Hanalyze.Model.HBM.Track + Hanalyze.Model.HBM.Util + Hanalyze.Model.HBM.VecAD + Hanalyze.Stat.AD + Hanalyze.Stat.BayesFactor + Hanalyze.Stat.BayesianModelAveraging + Hanalyze.Stat.BridgeSampling + Hanalyze.Stat.PosteriorPredictive + Hanalyze.Stat.VI + build-depends: + base >= 4.14 && < 5 + , array >= 0.5 && < 0.6 + , async >= 2.2 && < 2.3 + , containers >= 0.6 && < 0.8 + , hmatrix >= 0.20 && < 0.22 + , mwc-random >= 0.15 && < 0.16 + , primitive >= 0.7 && < 0.10 + , deepseq >= 1.4 && < 1.6 + , parallel >= 3.2 && < 3.3 + , text >= 1.2 && < 2.2 + , aeson >= 2.0 && < 2.3 + , ad >= 4.4 && < 4.6 + , reflection >= 2.1 && < 2.2 + , vector >= 0.12 && < 0.14 + , hanalyze-core == 0.2.0.1 diff --git a/src/Hanalyze/MCMC/BayesianTest.hs b/src/Hanalyze/MCMC/BayesianTest.hs new file mode 100644 --- /dev/null +++ b/src/Hanalyze/MCMC/BayesianTest.hs @@ -0,0 +1,277 @@ +-- | +-- Module : Hanalyze.MCMC.BayesianTest +-- Description : Bayesian A/B test — 2 群間の平均差を NUTS でサンプルし ROPE/HDI で判定 +-- Copyright : (c) 2026 Aelysce Project (Toshiaki Honda) +-- License : BSD-3-Clause +-- +-- [日本語]: Bayesian A/B test helper — 2 群間の平均差を NUTS でサンプル、 +-- ROPE / HDI に基づき決定。 +-- +-- Spotfire 風 "Good vs Bad" の Bayesian 版。 既存の頻度論版 +-- ('Hanalyze.Stat.GroupComparison.goodVsBad') が Welch t + Cohen's d で +-- 並列比較するのに対し、 本モジュールは __2 群の平均差の posterior__ を +-- 得て、 HDI (highest density interval) + ROPE (region of practical +-- equivalence) で意思決定する。 +-- +-- モデル: +-- +-- @ +-- μ_A ~ Normal(0, priorScale) +-- μ_B ~ Normal(0, priorScale) +-- σ_A ~ HalfNormal(sigmaScale) +-- σ_B ~ HalfNormal(sigmaScale) +-- y_A ~ Normal(μ_A, σ_A) +-- y_B ~ Normal(μ_B, σ_B) +-- diff = μ_B - μ_A +-- @ +-- +-- 決定ルール (`ROPEDecision lo hi`): +-- +-- - HDI が ROPE [lo, hi] と __重ならず HDI 全体が ROPE の外__ → 'RejectH0' +-- - HDI が ROPE 内に __完全に含まれる__ → 'AcceptH0' +-- - それ以外 → 'Inconclusive' +-- [English]: Bayesian A/B test helper — samples the difference in means +-- between 2 groups via NUTS and decides using ROPE / HDI. +-- +-- A Bayesian counterpart to the Spotfire-style "Good vs Bad". Where the +-- existing frequentist version +-- ('Hanalyze.Stat.GroupComparison.goodVsBad') compares in parallel +-- with Welch's t-test + Cohen's d, this module obtains the +-- __posterior of the difference in means between the 2 groups__ and +-- makes a decision using the HDI (highest density interval) + ROPE +-- (region of practical equivalence). +-- +-- Model: +-- +-- @ +-- μ_A ~ Normal(0, priorScale) +-- μ_B ~ Normal(0, priorScale) +-- σ_A ~ HalfNormal(sigmaScale) +-- σ_B ~ HalfNormal(sigmaScale) +-- y_A ~ Normal(μ_A, σ_A) +-- y_B ~ Normal(μ_B, σ_B) +-- diff = μ_B - μ_A +-- @ +-- +-- Decision rule (`ROPEDecision lo hi`): +-- +-- - HDI __does not overlap__ ROPE [lo, hi] and lies entirely outside it → +-- 'RejectH0' +-- - HDI is __entirely contained__ within ROPE → 'AcceptH0' +-- - Otherwise → 'Inconclusive' +{-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE BangPatterns #-} +{-# LANGUAGE RankNTypes #-} +module Hanalyze.MCMC.BayesianTest + ( -- * 入力 + BayesianABConfig (..) + , DecisionRule (..) + , defaultBayesianABConfig + -- * 出力 + , BayesianABResult (..) + , ABDecision (..) + -- * 実行 + , bayesianAB + -- * 補助 + , highestDensityInterval + ) where + +import qualified Data.Map.Strict as Map +import Data.List (sort) +import qualified System.Random.MWC as MWC + +import qualified Hanalyze.MCMC.Core as MC +import qualified Hanalyze.MCMC.NUTS as NUTS +import qualified Hanalyze.Model.HBM as HBM + +-- =========================================================================== +-- 型 +-- =========================================================================== + +-- | [日本語]: 意思決定ルール。 [English]: The decision rule. +data DecisionRule + = HDIOnly + -- ^ [日本語]: HDI を計算するのみ、 自動判定しない。 + -- [English]: Only computes the HDI; does not decide automatically. + | ROPEDecision !Double !Double + -- ^ [日本語]: @ROPEDecision lo hi@ で「実用上 0 と区別不能な区間 @[lo, hi]@」 を指定。 + -- [English]: @ROPEDecision lo hi@ specifies the interval @[lo, hi]@ + -- that is practically indistinguishable from 0. + deriving (Show, Eq) + +-- | [日本語]: A/B 試験の入力設定。 [English]: The A/B test's input configuration. +data BayesianABConfig = BayesianABConfig + { babCredible :: !Double -- ^ [日本語]: HDI の信頼水準 (例 0.95)。 [English]: The HDI's credible level (e.g. 0.95). + , babRule :: !DecisionRule + , babPriorScale :: !Double -- ^ [日本語]: μ_A, μ_B の prior σ (default 10)。 [English]: The prior σ for μ_A, μ_B (default 10). + , babSigmaScale :: !Double -- ^ [日本語]: HalfNormal σ の scale (default 5)。 [English]: The scale of the HalfNormal σ (default 5). + , babNUTS :: !NUTS.NUTSConfig + } deriving (Show) + +defaultBayesianABConfig :: BayesianABConfig +defaultBayesianABConfig = BayesianABConfig + { babCredible = 0.95 + , babRule = HDIOnly + , babPriorScale = 10.0 + , babSigmaScale = 5.0 + , babNUTS = NUTS.defaultNUTSConfig + { NUTS.nutsIterations = 1000 + , NUTS.nutsBurnIn = 500 + } + } + +-- | [日本語]: 自動判定の結果。 [English]: The result of the automatic decision. +data ABDecision + = AcceptH0 -- ^ [日本語]: HDI が ROPE 内 → 「実用上 0」 と判定。 [English]: HDI is within ROPE → judged "practically 0". + | RejectH0 -- ^ [日本語]: HDI が ROPE の外 → 「明確に差がある」 と判定。 [English]: HDI is outside ROPE → judged "clearly different". + | Inconclusive -- ^ [日本語]: HDI が ROPE と部分的に重なる → 「データ不足」。 [English]: HDI partially overlaps ROPE → "insufficient data". + | NoRuleApplied -- ^ [日本語]: 'HDIOnly' 指定で判定なし。 [English]: No decision made because 'HDIOnly' was specified. + deriving (Show, Eq) + +-- | [日本語]: A/B 試験の出力。 [English]: The A/B test's output. +data BayesianABResult = BayesianABResult + { babPosteriorDiff :: ![Double] + -- ^ [日本語]: 平均差 (μ_B − μ_A) の post-burn-in サンプル。 + -- [English]: The post-burn-in samples of the difference in means (μ_B − μ_A). + , babMeanDiff :: !Double + -- ^ [日本語]: posterior mean (μ_B − μ_A)。 [English]: The posterior mean (μ_B − μ_A). + , babHDI :: !(Double, Double) + -- ^ [日本語]: @babCredible@ 信頼水準の HDI。 [English]: The HDI at the @babCredible@ credible level. + , babDecision :: !ABDecision + , babProbDiffPos :: !Double + -- ^ [日本語]: @P(μ_B > μ_A)@ の posterior 確率。 [English]: The posterior probability @P(μ_B > μ_A)@. + , babChain :: !MC.Chain + -- ^ [日本語]: 生 chain (μ_A / μ_B / σ_A / σ_B / diff の post-burn-in サンプル)。 + -- [English]: The raw chain (post-burn-in samples of μ_A / μ_B / σ_A / σ_B / diff). + } deriving (Show) + +-- =========================================================================== +-- 公開関数 +-- =========================================================================== + +-- | [日本語]: 2 群のデータから Bayesian A/B 試験を実行。 +-- +-- 内部で HBM モデルを組み立て、 NUTS で posterior をサンプル、 +-- 平均差の HDI と決定を返す。 +-- +-- 失敗条件: いずれかの群が空 → @error@ (canvas backend では事前に検査)。 +-- [English]: Runs a Bayesian A/B test from the data of 2 groups. +-- +-- Internally builds an HBM model, samples the posterior via NUTS, and +-- returns the HDI of the difference in means along with the decision. +-- +-- Failure condition: either group is empty → @error@ (checked beforehand +-- in the canvas backend). +bayesianAB + :: BayesianABConfig + -> [Double] -- ^ [日本語]: 群 A の観測値。 [English]: The observations of group A. + -> [Double] -- ^ [日本語]: 群 B の観測値。 [English]: The observations of group B. + -> MWC.GenIO + -> IO BayesianABResult +bayesianAB cfg ysA ysB gen + | null ysA || null ysB = + error "Hanalyze.MCMC.BayesianTest.bayesianAB: both groups must be non-empty" + | otherwise = do + let priorScale_ = babPriorScale cfg + sigmaScale_ = babSigmaScale cfg + model :: HBM.ModelP () + model = do + muA <- HBM.sample "mu_a" (HBM.Normal 0 (realToFrac priorScale_)) + muB <- HBM.sample "mu_b" (HBM.Normal 0 (realToFrac priorScale_)) + sA <- HBM.sample "sigma_a" (HBM.HalfNormal (realToFrac sigmaScale_)) + sB <- HBM.sample "sigma_b" (HBM.HalfNormal (realToFrac sigmaScale_)) + HBM.observe "ya" (HBM.Normal muA sA) ysA + HBM.observe "yb" (HBM.Normal muB sB) ysB + _ <- HBM.deterministic "diff" (muB - muA) + pure () + initParams = Map.fromList + [ ("mu_a", mean ysA) + , ("mu_b", mean ysB) + , ("sigma_a", max 0.1 (stddev ysA)) + , ("sigma_b", max 0.1 (stddev ysB)) + ] + rawChain <- NUTS.nuts model (babNUTS cfg) initParams gen + -- deterministic 値 "diff" は raw chain に入っていないため augment で注入 + let chain = HBM.augmentChainWithDeterministic model rawChain + diffs = MC.chainVals "diff" chain + n = length diffs + mu = if n == 0 then 0 else sum diffs / fromIntegral n + hdi = highestDensityInterval (babCredible cfg) diffs + probP = if n == 0 + then 0 + else fromIntegral (length (filter (> 0) diffs)) + / fromIntegral n + decision = case babRule cfg of + HDIOnly -> NoRuleApplied + ROPEDecision lo hi -> classifyROPE hdi lo hi + pure BayesianABResult + { babPosteriorDiff = diffs + , babMeanDiff = mu + , babHDI = hdi + , babDecision = decision + , babProbDiffPos = probP + , babChain = chain + } + +-- =========================================================================== +-- 補助 +-- =========================================================================== + +-- | [日本語]: サンプル列の __highest density interval (HDI)__。 +-- +-- ソート後、 窓幅 @floor(n · level)@ で全 sliding window を試し、 +-- 最も狭い窓を返す。 unimodal な posterior では HDI = 最短連続区間。 +-- +-- @level ∈ (0, 1)@、 例: 0.95 で 95% HDI。 +-- [English]: The __highest density interval (HDI)__ of a sample sequence. +-- +-- After sorting, tries every sliding window of width @floor(n · level)@ +-- and returns the narrowest one. For a unimodal posterior, the HDI is +-- the shortest contiguous interval. +-- +-- @level ∈ (0, 1)@, e.g. 0.95 for the 95% HDI. +highestDensityInterval :: Double -> [Double] -> (Double, Double) +highestDensityInterval level xs + | null xs = (0, 0) + | level <= 0 || level >= 1 = error "HDI: level must be in (0, 1)" + | otherwise = + let sorted = sort xs + n = length sorted + k = max 1 (floor (fromIntegral n * level :: Double)) + -- 全 sliding windows (start = 0 .. n-k) + arr = case sorted of + [] -> [] + _ -> sorted + windows = [ (arr !! i, arr !! (i + k - 1)) + | i <- [0 .. n - k] ] + -- 最も狭い窓 + best = head $ foldr keepNarrower [head windows] (tail windows) + in best + where + keepNarrower w (b:_) = + if (snd w - fst w) < (snd b - fst b) then [w] else [b] + keepNarrower w [] = [w] + +-- | [日本語]: HDI と ROPE [lo, hi] から ABDecision を分類。 +-- [English]: Classifies the ABDecision from the HDI and ROPE [lo, hi]. +classifyROPE :: (Double, Double) -> Double -> Double -> ABDecision +classifyROPE (hdiLo, hdiHi) ropeLo ropeHi + | hdiHi < ropeLo || hdiLo > ropeHi = RejectH0 -- HDI 全体が ROPE 外 + | hdiLo >= ropeLo && hdiHi <= ropeHi = AcceptH0 -- HDI 全体が ROPE 内 + | otherwise = Inconclusive -- 部分重複 + +-- =========================================================================== +-- 統計 helper +-- =========================================================================== + +mean :: [Double] -> Double +mean [] = 0 +mean xs = sum xs / fromIntegral (length xs) + +stddev :: [Double] -> Double +stddev xs + | length xs < 2 = 1 + | otherwise = + let n = fromIntegral (length xs) :: Double + m = mean xs + in sqrt (sum [ (x - m) ** 2 | x <- xs ] / (n - 1)) diff --git a/src/Hanalyze/MCMC/Gibbs.hs b/src/Hanalyze/MCMC/Gibbs.hs new file mode 100644 --- /dev/null +++ b/src/Hanalyze/MCMC/Gibbs.hs @@ -0,0 +1,548 @@ +-- | +-- Module : Hanalyze.MCMC.Gibbs +-- Description : 共役事前分布向け Gibbs sampler (解析的フル条件付きサンプリング) +-- Copyright : (c) 2026 Aelysce Project (Toshiaki Honda) +-- License : BSD-3-Clause +-- +-- Gibbs sampler — analytic full-conditional sampling for conjugate priors. +-- +-- Each 'GibbsUpdate' draws a single parameter directly from its full +-- conditional distribution, so no Metropolis rejection step is needed and +-- every sample is accepted. When non-conjugate parameters are mixed in, +-- combine with Metropolis-Hastings ('gibbsMH'). +{-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE RankNTypes #-} +module Hanalyze.MCMC.Gibbs + ( -- * 共役アップデートブロック + GibbsUpdate + , normalNormal + , betaBinomial + , gammaPoisson + , sampleBetaBB + -- * Samplers + , GibbsConfig (..) + , defaultGibbsConfig + , gibbs + , gibbsBetaBinomial + , gibbsChains + , gibbsPure + , gibbsChainsPure + , gibbsBetaBinomialPure + -- * HBM-DSL integration: conjugacy auto-detection + , gibbsFromModel + -- * Hybrid Gibbs+MH sampler + , gibbsMH + , gibbsMHChains + , gibbsMHPure + , gibbsMHChainsPure + ) where + +import Control.Concurrent.Async (mapConcurrently) +import Control.Monad (foldM, replicateM, when) +import Control.Monad.Primitive (PrimMonad, PrimState) +import Control.Monad.ST (runST) +import Control.Parallel.Strategies (parList, rdeepseq, using) +import Data.Primitive.MutVar +import Data.List (nub) +import Data.Maybe (listToMaybe) +import Data.Word (Word32) +import qualified Data.Map.Strict as Map +import Data.Map.Strict (Map) +import Data.Text (Text) +import qualified Data.Vector as V +import qualified Data.Vector.Storable as VS +import System.Random.MWC (Gen, GenIO, uniform, initialize) +import System.Random.MWC.Distributions (gamma, normal) + +import Hanalyze.MCMC.Core (Chain (..), spawnGen) +import Hanalyze.Model.HBM (ModelP, Params, Distribution (..), + Node (..), NodeKind (..), collectNodes, + logJoint, runObserveDists, priorList) + +-- --------------------------------------------------------------------------- +-- 型 +-- --------------------------------------------------------------------------- + +-- | [日本語]: Gibbs update block。 現在のパラメータ集合を受け取り、 割り当てられた +-- パラメータのフル条件付き分布からサンプルした新しい @(name, value)@ を +-- 1 個返す。 +-- +-- monad パラメータ化 (@m@) により 'IO' でも @ST s@ でも走らせられる +-- (純粋 Gibbs に必要)。 rank-N alias にすると @Maybe@/list 構築が impredicative に +-- なるため、 alias を kind @* -> *@ にして関数側 ('gibbsFromModel' 等) を多相にする。 +-- [English]: A Gibbs update block. Receives the current parameter set and +-- returns a single fresh @(name, value)@ sampled from the assigned +-- parameter's full conditional distribution. +-- +-- Parameterizing over the monad (@m@) lets this run under either 'IO' +-- or @ST s@ (needed for pure Gibbs). Making the alias rank-N would make +-- @Maybe@/list construction impredicative, so instead the alias is kept +-- at kind @* -> *@ and the functions that use it (e.g. +-- 'gibbsFromModel') are made polymorphic. +type GibbsUpdate m = Params -> Gen (PrimState m) -> m (Text, Double) + +-- --------------------------------------------------------------------------- +-- 共役アップデート (モデル非依存) +-- --------------------------------------------------------------------------- + +-- | Conjugate update for a Normal prior × Normal likelihood with known +-- @σ@. +normalNormal + :: PrimMonad m => Text -> Double -> Double -> [Double] -> Double -> GibbsUpdate m +normalNormal paramName mu0 sig0 ys sigLik _ps gen = do + let n = fromIntegral (length ys) :: Double + ybar = if n == 0 then 0 else sum ys / n + prec0 = 1 / sig0 ^ (2::Int) + precLik = 1 / sigLik ^ (2::Int) + precPost = prec0 + n * precLik + sigPost = sqrt (1 / precPost) + muPost = (mu0 * prec0 + n * ybar * precLik) / precPost + val <- normal muPost sigPost gen + return (paramName, val) + +-- | Conjugate update for a Beta prior × Binomial likelihood. +betaBinomial + :: PrimMonad m => Text -> Double -> Double -> Int -> Int -> GibbsUpdate m +betaBinomial paramName alpha0 beta0 n k _ps gen = do + val <- sampleBeta (alpha0 + fromIntegral k) + (beta0 + fromIntegral (n - k)) + gen + return (paramName, val) + +-- | Conjugate update for a Gamma prior × Poisson likelihood +-- (rate parameterization). +gammaPoisson + :: PrimMonad m => Text -> Double -> Double -> [Double] -> GibbsUpdate m +gammaPoisson paramName alpha0 beta0 ys _ps gen = do + let n = fromIntegral (length ys) :: Double + aPost = alpha0 + sum ys + bPost = beta0 + n + val <- gamma aPost (1 / bPost) gen + return (paramName, val) + +-- | Sample @Beta(a, b)@. Implemented as @X / (X + Y)@ with +-- @X ~ Gamma(a)@, @Y ~ Gamma(b)@, since @mwc-random@ has no Beta sampler. +sampleBeta :: PrimMonad m => Double -> Double -> Gen (PrimState m) -> m Double +sampleBeta a b gen + | a > 1 && b > 1 = sampleBetaBB a b gen -- Cheng's BB, much faster + | otherwise = sampleBetaGamma a b gen +{-# INLINE sampleBeta #-} + +-- | Generic fallback: @X / (X + Y)@ with @X ~ Gamma(a)@, @Y ~ Gamma(b)@. +-- Used when the Cheng-BB precondition @a, b > 1@ is violated; the BB +-- algorithm's @λ = √((α − 2) / (2 a b − α))@ becomes imaginary at +-- @a, b ≤ 1@ so a different branch (BC) would be required there. +sampleBetaGamma :: PrimMonad m => Double -> Double -> Gen (PrimState m) -> m Double +sampleBetaGamma a b gen = do + x <- gamma a 1 gen + y <- gamma b 1 gen + return (x / (x + y)) + +-- | R. C. H. Cheng's BB algorithm (1978), valid for @min(a, b) > 1@. +-- Direct Beta sampler that avoids the two Gamma calls + division +-- ("X / (X+Y)") used by @sampleBetaGamma@. +-- +-- P37 (2026-05-07): the n=10000 Gibbs Beta-Binomial bench is 78% +-- @sampleBetaGamma@ (1.38 ms / 1.76 ms total). Each gamma call uses +-- mwc-random's Marsaglia-Tsang squeeze, which needs ~3 uniforms + log +-- + cube on average — heavier than Cheng-BB which is ~1.5 uniforms + +-- log + exp per accepted sample on this regime. +-- +-- Reference: Cheng (1978), "Generating Beta variates with non-integral +-- shape parameters", CACM 21(4):317-322. Algorithm BB on p. 319. +sampleBetaBB :: forall m. PrimMonad m => Double -> Double -> Gen (PrimState m) -> m Double +sampleBetaBB a b gen = do + let !alpha = a + b + !beta_ = sqrt ((alpha - 2) / (2 * a * b - alpha)) + !gamma_ = a + 1 / beta_ + !logFour = log 4 + !log5 = log 5 + !logAlpha = log alpha + let loop = do + u1 <- uniform gen :: m Double + u2 <- uniform gen :: m Double + let !v = beta_ * log (u1 / (1 - u1)) + !w = a * exp v + !z = u1 * u1 * u2 + !r = gamma_ * v - logFour + !s = a + r - w + -- Cheng-BB's three accept tests (numpy randomkit.c): + -- Step 4 (squeeze): s + 1 + log(5) ≥ 5 z + -- Step 5a: s ≥ log z + -- Step 5b: r + α · log(α / (b + w)) ≥ log z + if s + 1 + log5 >= 5 * z + then return (w / (b + w)) + else do + let !t = log z + if s >= t + then return (w / (b + w)) + else if r + alpha * (logAlpha - log (b + w)) >= t + then return (w / (b + w)) + else loop + loop +{-# INLINE sampleBetaBB #-} + +-- --------------------------------------------------------------------------- +-- Gibbs サンプラー (汎用ランナー、モデル非依存) +-- --------------------------------------------------------------------------- + +-- | Gibbs configuration. +data GibbsConfig = GibbsConfig + { gibbsIterations :: Int -- ^ Total iterations (burn-in included). + , gibbsBurnIn :: Int -- ^ Burn-in iterations to discard. + } deriving (Show) + +-- | Default configuration: 2000 iterations, 500 burn-in. +defaultGibbsConfig :: GibbsConfig +defaultGibbsConfig = GibbsConfig + { gibbsIterations = 2000 + , gibbsBurnIn = 500 + } + +-- | Apply each update in @updates@ once per iteration, in order. Every +-- Gibbs step is accepted by construction, so @chainAccepted@ equals +-- @(length updates) × iterations@. +gibbs :: PrimMonad m => [GibbsUpdate m] -> GibbsConfig -> Params -> Gen (PrimState m) -> m Chain +gibbs updates cfg initP gen = do + let total = gibbsBurnIn cfg + gibbsIterations cfg + nUpd = length updates + samplesRef <- newMutVar [] + acceptedRef <- newMutVar (0 :: Int) + let step current = foldM applyOne current updates + where + applyOne ps upd = do + (name, val) <- upd ps gen + return (Map.insert name val ps) + let loop 0 current = return current + loop i current = do + next <- step current + modifyMutVar' acceptedRef (+ nUpd) + when (i <= gibbsIterations cfg) $ + modifyMutVar' samplesRef (next :) + loop (i - 1) next + _ <- loop total initP + samples <- fmap reverse (readMutVar samplesRef) + accepted <- readMutVar acceptedRef + return Chain + { chainSamples = samples + , chainAccepted = accepted + , chainTotal = total * nUpd + , chainEnergy = [] + , chainDivergences = [] + , chainTreeDepths = [] + } + +-- | Specialised Beta-Binomial conjugate sampler. Equivalent to +-- @gibbs [betaBinomial p a0 b0 n k] cfg (Map.singleton p init) gen@ +-- but bypasses the generic loop's per-iteration overhead. +-- +-- P37 (2026-05-07): per-iteration profile of the n=10000 bench: +-- @sampleBeta@ itself (two @gamma@ draws + division) is ~80 ns — +-- well over half the 180 ns/iter budget. The remaining ~100 ns is +-- bookkeeping that the generic runner has to do because it doesn't +-- know whether updates depend on @ps@: +-- +-- * @Map.insert paramName val ps@ — fresh Map allocation every iter +-- (size-1 Map but still a tree node + Text key + boxed Double) +-- * @modifyMutVar' acceptedRef (+ nUpd)@ — counter that's exactly +-- @total@ at the end (every Gibbs step is unconditionally accepted) +-- * @modifyMutVar' samplesRef (next :)@ + final @reverse@ — list cons +-- of every kept iteration plus a 10000-element reverse +-- * @foldM applyOne current updates@ — closure construction even +-- though the update list has length 1 +-- +-- For Beta-Binomial in isolation the conjugate posterior is +-- /independent/ of the previous draw, so we replace the entire loop +-- with @VS.replicateM total (sampleBeta postA postB gen)@. The +-- @Params@ Map is constructed only at the chain-construction +-- boundary (lazily, one entry per kept sample), avoiding the per-iter +-- allocation while keeping the public 'Chain' shape intact. +-- +-- numpy.random.beta(a, b, size=10000) does the same thing in C with +-- SIMD; this brings the Haskell side to within ~2× of it without FFI. +gibbsBetaBinomial + :: PrimMonad m + => Text -- ^ Parameter name (= sample-Map key). + -> Double -- ^ Beta prior @α@. + -> Double -- ^ Beta prior @β@. + -> Int -- ^ Binomial @n@. + -> Int -- ^ Observed successes @k@. + -> GibbsConfig + -> Gen (PrimState m) + -> m Chain +gibbsBetaBinomial paramName alpha0 beta0 n k cfg gen = do + let !total = gibbsBurnIn cfg + gibbsIterations cfg + !keep = gibbsIterations cfg + !postA = alpha0 + fromIntegral k + !postB = beta0 + fromIntegral (n - k) + -- Storable Vector keeps the n=10000 doubles in 80 KB of contiguous + -- memory rather than as a linked list of boxed thunks, and avoids + -- the @reverse@ pass at the end of the generic loop. + vals <- VS.replicateM total (sampleBeta postA postB gen) + let kept = VS.drop (total - keep) vals + samples = [Map.singleton paramName v | v <- VS.toList kept] + return Chain + { chainSamples = samples + , chainAccepted = total -- every Gibbs step is accepted + , chainTotal = total + , chainEnergy = [] + , chainDivergences = [] + , chainTreeDepths = [] + } + +-- | Run 'gibbs' on @numChains@ parallel chains. +gibbsChains :: [GibbsUpdate IO] -> GibbsConfig -> Int -> Params -> GenIO -> IO [Chain] +gibbsChains updates cfg numChains initP baseGen = do + gens <- replicateM numChains (spawnGen baseGen) + mapConcurrently (\g -> gibbs updates cfg initP g) gens + +-- --------------------------------------------------------------------------- +-- HBM DSL 統合: 共役構造の自動検出 +-- --------------------------------------------------------------------------- + +distParams :: Distribution Double -> [Double] +distParams (Normal mu sig) = [mu, sig] +distParams (Binomial n p) = [fromIntegral n, p] +distParams (Poisson lam) = [lam] +distParams (Exponential r) = [r] +distParams (Gamma a b) = [a, b] +distParams (Beta a b) = [a, b] +distParams (Uniform lo hi) = [lo, hi] +distParams (StudentT df mu s) = [df, mu, s] +distParams (Cauchy loc s) = [loc, s] +distParams (HalfNormal s) = [s] +distParams (HalfCauchy s) = [s] +distParams (LogNormal mu s) = [mu, s] +distParams (Bernoulli p) = [p] +distParams (Categorical ps) = ps +distParams (Mixture ws _) = ws -- 共役検出には使えない (重みのみ) +distParams (Truncated _ _ _) = [] -- 共役検出対象外 +distParams (Censored _ _ _) = [] -- 共役検出対象外 +distParams MvNormal{} = [] -- 共役検出対象外 (観測専用) +distParams MvNormalChol{} = [] -- 共役検出対象外 (観測専用) +distParams (NegativeBinomial mu a) = [mu, a] +distParams (Multinomial _ ps) = ps +distParams (ZeroInflatedPoisson psi lam) = [psi, lam] +distParams (ZeroInflatedBinomial _ psi p) = [psi, p] +distParams (InverseGamma a b) = [a, b] +distParams (Weibull k l) = [k, l] +distParams (Pareto a xm) = [a, xm] +distParams (BetaBinomial _ a b) = [a, b] +distParams (VonMises mu k) = [mu, k] + +-- 各潜在変数が Observe ノードのどの (obsIndex, slotIndex) に影響するかを検出。 +detectObsDeps :: ModelP r -> [Text] -> Map Text [(Int, Int)] +detectObsDeps m latNames = + let baseline = map (\(_, d, _) -> distParams d) (runObserveDists m Map.empty) + perturb v = map (\(_, d, _) -> distParams d) + (runObserveDists m (Map.singleton v 1.0)) + in Map.fromList + [ (v, nub + [ (oi, si) + | let pp = perturb v + , (oi, (bp, pp')) <- zip [0..] (zip baseline pp) + , (si, (bv, pv)) <- zip [0..] (zip bp pp') + , bv /= pv + ]) + | v <- latNames + ] + +-- | Inspect an HBM model's structure and synthesise the conjugate +-- 'GibbsUpdate' steps automatically. +-- +-- Detected conjugate pairs: +-- +-- * @Gamma(α,β)@ + @Poisson(λ)@ → 'gammaPoisson' +-- * @Beta(α,β)@ + @Binomial(n,p)@ → 'betaBinomial' +-- * @Normal(μ₀,σ₀)@ + @Normal(μ,σ)@ → 'normalNormal' +-- +-- Returns @(updates, remaining)@: the synthesised updates and the names +-- of parameters that still need an MH step. +gibbsFromModel :: forall r m. PrimMonad m => ModelP r -> ([GibbsUpdate m], [Text]) +gibbsFromModel m = + let nodes = collectNodes m + latNames = [ nodeName n | n <- nodes, nodeKind n == LatentN ] + priorMap = Map.fromList (priorList m) + obsList = runObserveDists m Map.empty + indexedObs = zip [0 :: Int ..] obsList + deps = detectObsDeps m latNames + + obsAt i = listToMaybe [ (d, xs) | (j, (_, d, xs)) <- indexedObs, i == j ] + + buildUpd v = + let priorD = Map.findWithDefault (Normal 0 1) v priorMap + vDeps = Map.findWithDefault [] v deps + in case (priorD, vDeps) of + (Gamma a b, [(obsIdx, 0)]) -> + case obsAt obsIdx of + Just (Poisson _, xs) -> Just (gammaPoisson v a b xs) + _ -> Nothing + + (Beta a b, [(obsIdx, 1)]) -> + case obsAt obsIdx of + Just (Binomial nPerObs _, xs) -> + let k = round (sum xs) :: Int + n = nPerObs * length xs + in Just (betaBinomial v a b n k) + _ -> Nothing + + (Normal mu0 sig0, [(obsIdx, 0)]) -> + case obsAt obsIdx of + Just (Normal _ _, xs) -> + let sigmaVar = listToMaybe + [ w | (w, wDeps) <- Map.toList deps + , any (\(oi, si) -> oi == obsIdx && si == 1) wDeps + , w /= v + ] + in Just $ \ps gen -> + let sigLik = maybe 1.0 (\sv -> Map.findWithDefault 1.0 sv ps) sigmaVar + in normalNormal v mu0 sig0 xs sigLik ps gen + _ -> Nothing + + _ -> Nothing + + results = map buildUpd latNames + updates = [ u | Just u <- results ] + remaining = [ v | (v, Nothing) <- zip latNames results ] + in (updates, remaining) + +-- --------------------------------------------------------------------------- +-- ハイブリッド Gibbs+MH +-- --------------------------------------------------------------------------- + +hybridStep + :: PrimMonad m + => [GibbsUpdate m] + -> [Text] + -> Map Text Double + -> ModelP r + -> Params -> Gen (PrimState m) + -> m (Params, Bool) +hybridStep gibbsUpds mhNames mhSteps model current gen = do + afterGibbs <- foldM (\ps upd -> do + (name, val) <- upd ps gen + return (Map.insert name val ps)) current gibbsUpds + if null mhNames + then return (afterGibbs, True) + else do + proposed <- foldM (\ps n -> do + let s = Map.findWithDefault 1.0 n mhSteps + cv = Map.findWithDefault 0.0 n ps + eps <- normal 0 s gen + return (Map.insert n (cv + eps) ps)) afterGibbs mhNames + let logA = logJoint model proposed - logJoint model afterGibbs + u <- uniform gen + let accepted = log (u :: Double) < logA + return (if accepted then proposed else afterGibbs, accepted) + +-- | Hybrid sampler: Gibbs-update conjugate parameters and use Random-Walk +-- Metropolis on the rest. +gibbsMH + :: PrimMonad m + => ModelP r + -> GibbsConfig + -> Map Text Double -- ^ MH step size per non-conjugate parameter. + -> Params + -> Gen (PrimState m) + -> m Chain +gibbsMH model cfg mhSteps initP gen = do + let (gibbsUpds, mhNames) = gibbsFromModel model + total = gibbsBurnIn cfg + gibbsIterations cfg + samplesRef <- newMutVar [] + acceptedRef <- newMutVar (0 :: Int) + let loop 0 current = return current + loop i current = do + (next, acc) <- hybridStep gibbsUpds mhNames mhSteps model current gen + when acc $ modifyMutVar' acceptedRef (+1) + when (i <= gibbsIterations cfg) $ + modifyMutVar' samplesRef (next :) + loop (i - 1) next + _ <- loop total initP + samples <- fmap reverse (readMutVar samplesRef) + accepted <- readMutVar acceptedRef + return Chain + { chainSamples = samples + , chainAccepted = accepted + , chainTotal = total + , chainEnergy = [] + , chainDivergences = [] + , chainTreeDepths = [] + } + +gibbsMHChains + :: ModelP r + -> GibbsConfig + -> Map Text Double + -> Int + -> Params + -> GenIO + -> IO [Chain] +gibbsMHChains model cfg mhSteps numChains initP baseGen = do + gens <- replicateM numChains (spawnGen baseGen) + mapConcurrently (\g -> gibbsMH model cfg mhSteps initP g) gens + +-- --------------------------------------------------------------------------- +-- Phase 50: 純粋 (ST + seed) ラッパ +-- --------------------------------------------------------------------------- + +-- | [日本語]: 純粋・決定的な hybrid Gibbs+MH (モデルから共役 update を内部導出)。 +-- seed → 確定 Chain。 +-- [English]: A pure, deterministic hybrid Gibbs+MH sampler (internally +-- derives the conjugate updates from the model). Maps seed → a +-- deterministic Chain. +gibbsMHPure :: ModelP r -> GibbsConfig -> Map Text Double -> Params -> Word32 -> Chain +gibbsMHPure model cfg mhSteps initP seed = + runST (initialize (V.singleton seed) >>= gibbsMH model cfg mhSteps initP) + +-- | [日本語]: 純粋・決定的な multi-chain hybrid Gibbs+MH。 子 seed を純粋導出し +-- @parList rdeepseq@ で並列。 +-- [English]: A pure, deterministic multi-chain hybrid Gibbs+MH sampler. +-- Derives the child seeds purely and parallelizes with +-- @parList rdeepseq@. +gibbsMHChainsPure :: ModelP r -> GibbsConfig -> Map Text Double -> Int -> Params -> Word32 -> [Chain] +gibbsMHChainsPure model cfg mhSteps numChains initP seed = + let childSeeds :: [Word32] + childSeeds = runST $ do + g <- initialize (V.singleton seed) + replicateM numChains (uniform g) + chains = [ gibbsMHPure model cfg mhSteps initP s | s <- childSeeds ] + in chains `using` parList rdeepseq + +-- | [日本語]: 純粋・決定的な Beta-Binomial 共役 Gibbs (seed → 確定 Chain)。 +-- [English]: A pure, deterministic Beta-Binomial conjugate Gibbs sampler +-- (seed → deterministic Chain). +gibbsBetaBinomialPure + :: Text -> Double -> Double -> Int -> Int -> GibbsConfig -> Word32 -> Chain +gibbsBetaBinomialPure paramName alpha0 beta0 n k cfg seed = + runST (initialize (V.singleton seed) + >>= gibbsBetaBinomial paramName alpha0 beta0 n k cfg) + +-- | [日本語]: 純粋・決定的な汎用 Gibbs (seed → 確定 Chain)。 update 群は __rank-N__ +-- (@forall m. PrimMonad m => [GibbsUpdate m]@) で渡す = リストリテラルを直接渡せば +-- 多相のまま通る (@let updates = …@ で束縛すると単相化するので注意・直接渡しが楽)。 +-- [English]: A pure, deterministic generic Gibbs sampler (seed → +-- deterministic Chain). The update set is passed as __rank-N__ +-- (@forall m. PrimMonad m => [GibbsUpdate m]@) — passing a list literal +-- directly keeps it polymorphic (binding it with @let updates = …@ +-- monomorphizes it, so beware; passing it directly is easiest). +gibbsPure + :: (forall m. PrimMonad m => [GibbsUpdate m]) + -> GibbsConfig -> Params -> Word32 -> Chain +gibbsPure updates cfg initP seed = + runST (initialize (V.singleton seed) >>= gibbs updates cfg initP) + +-- | [日本語]: 純粋・決定的な汎用 multi-chain Gibbs。 子 seed を純粋導出し +-- @parList rdeepseq@ で並列。 +-- [English]: A pure, deterministic generic multi-chain Gibbs sampler. +-- Derives the child seeds purely and parallelizes with +-- @parList rdeepseq@. +gibbsChainsPure + :: (forall m. PrimMonad m => [GibbsUpdate m]) + -> GibbsConfig -> Int -> Params -> Word32 -> [Chain] +gibbsChainsPure updates cfg numChains initP seed = + let childSeeds :: [Word32] + childSeeds = runST $ do + g <- initialize (V.singleton seed) + replicateM numChains (uniform g) + chains = [ gibbsPure updates cfg initP s | s <- childSeeds ] + in chains `using` parList rdeepseq diff --git a/src/Hanalyze/MCMC/HMC.hs b/src/Hanalyze/MCMC/HMC.hs new file mode 100644 --- /dev/null +++ b/src/Hanalyze/MCMC/HMC.hs @@ -0,0 +1,340 @@ +-- | +-- Module : Hanalyze.MCMC.HMC +-- Description : Hamiltonian Monte Carlo (HMC) サンプラー +-- Copyright : (c) 2026 Aelysce Project (Toshiaki Honda) +-- License : BSD-3-Clause +-- +-- Hamiltonian Monte Carlo (HMC) sampler. +-- +-- Computes exact gradients of polymorphic 'Hanalyze.Model.HBM' models ('ModelP') via +-- 'Numeric.AD.Mode.Reverse.Double'. Constrained parameters (@PositiveT@, +-- @UnitIntervalT@) are detected automatically from the prior distribution. +-- +-- @ +-- import Hanalyze.Model.HBM +-- import Hanalyze.MCMC.HMC +-- +-- myModel :: ModelP () +-- myModel = do +-- mu <- sample "mu" (Normal 0 10) +-- sigma <- sample "sigma" (Exponential 1) +-- observe "y" (Normal mu sigma) [1.5, 2.0, 1.8] +-- +-- chain <- hmc myModel defaultHMCConfig (Map.fromList [("mu",0),("sigma",1)]) gen +-- @ +{-# LANGUAGE StrictData #-} +{-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE RankNTypes #-} +module Hanalyze.MCMC.HMC + ( -- * Configuration + HMCConfig (..) + , defaultHMCConfig + -- * Constraint-transform helpers + , toUnconstrainedParams + , fromUnconstrainedParams + , logJointU + , leapfrogWith + , leapfrogWithM + , leapfrogWithMVS + -- * Basic utilities + , kinetic + , kineticM + , kineticMVS + , paramsToVec + , vecToParams + -- * Sampler + , hmc + , hmcChains + , hmcPure + , hmcChainsPure + ) where + +import Control.Concurrent.Async (mapConcurrently) +import Control.Monad (forM, replicateM, when) +import Control.Monad.Primitive (PrimMonad, PrimState) +import Control.Monad.ST (runST) +import Control.Parallel.Strategies (parList, rdeepseq, using) +import Data.Primitive.MutVar +import Data.Word (Word32) +import qualified Data.Vector as V +import qualified Data.Map.Strict as Map +import Data.Map.Strict (Map) +import Data.Text (Text) +import qualified Data.Vector.Storable as VS +import System.Random.MWC (Gen, GenIO, uniform, initialize) +import System.Random.MWC.Distributions (standard) + +import Hanalyze.Model.HBM (ModelP, Params, sampleNames, getTransforms, + logJointUnconstrained, gradADU) +import Hanalyze.MCMC.Core (Chain (..), spawnGen) +import Hanalyze.Stat.Distribution (Transform, toUnconstrained, fromUnconstrained) + +-- --------------------------------------------------------------------------- +-- Configuration +-- --------------------------------------------------------------------------- + +-- | HMC configuration. +data HMCConfig = HMCConfig + { hmcIterations :: Int -- ^ Total iterations (burn-in included). + , hmcBurnIn :: Int -- ^ Burn-in iterations to discard. + , hmcStepSize :: Double -- ^ Leapfrog step size @ε@. + , hmcLeapfrogSteps :: Int -- ^ Number of leapfrog steps per HMC iteration. + } deriving (Show) + +-- | Default HMC configuration: 2000 iterations, 500 burn-in, +-- @ε = 0.1@, 10 leapfrog steps. +defaultHMCConfig :: HMCConfig +defaultHMCConfig = HMCConfig + { hmcIterations = 2000 + , hmcBurnIn = 500 + , hmcStepSize = 0.1 + , hmcLeapfrogSteps = 10 + } + +-- --------------------------------------------------------------------------- +-- パラメータ変換ユーティリティ +-- --------------------------------------------------------------------------- + +-- | Pack parameters into a flat vector in the given name order. +paramsToVec :: [Text] -> Params -> [Double] +paramsToVec names params = map (\n -> Map.findWithDefault 0.0 n params) names + +-- | Inverse of 'paramsToVec': pair names with values. +vecToParams :: [Text] -> [Double] -> Params +vecToParams names vals = Map.fromList (zip names vals) + +-- | Apply 'toUnconstrained' to every named parameter; unmapped names are +-- left untouched. +toUnconstrainedParams :: Map Text Transform -> Params -> Params +toUnconstrainedParams transforms = + Map.mapWithKey (\k v -> maybe v (`toUnconstrained` v) (Map.lookup k transforms)) + +-- | Apply 'fromUnconstrained' to every named parameter. +fromUnconstrainedParams :: Map Text Transform -> Params -> Params +fromUnconstrainedParams transforms = + Map.mapWithKey (\k u -> maybe u (`fromUnconstrained` u) (Map.lookup k transforms)) + +-- --------------------------------------------------------------------------- +-- unconstrained 空間での log-joint (Jacobian 補正付き) +-- --------------------------------------------------------------------------- + +-- | Log-joint of a polymorphic model in the unconstrained space (shared +-- with VI and NUTS). +logJointU :: ModelP r -> Map Text Transform -> Params -> Double +logJointU model transforms paramsU = + let names = sampleNames model + transList = [Map.findWithDefault errT n transforms | n <- names] + errT = error "logJointU: transform missing" + in logJointUnconstrained model names transList paramsU + +-- --------------------------------------------------------------------------- +-- リープフロッグ積分 +-- --------------------------------------------------------------------------- + +-- | Kinetic energy @0.5 ‖r‖²@ for unit-mass momentum @r@. +kinetic :: [Double] -> Double +kinetic r = 0.5 * sum (map (^ (2 :: Int)) r) + +-- | Kinetic energy with a diagonal mass matrix: +-- @½ rᵀ M⁻¹ r = ½ Σ M⁻¹_ii · r_i²@. +-- +-- Used by NUTS (B11) when running with diagonal mass-matrix adaptation. +-- @kinetic = kineticM (repeat 1)@ recovers the identity-mass case. +kineticM :: [Double] -> [Double] -> Double +kineticM mInv r = 0.5 * sum (zipWith (\m_inv ri -> m_inv * ri * ri) mInv r) + +-- | Storable-Vector variant of 'kineticM'. +kineticMVS :: VS.Vector Double -> VS.Vector Double -> Double +kineticMVS mInv r = + 0.5 * VS.sum (VS.zipWith (\m_inv ri -> m_inv * ri * ri) mInv r) +{-# INLINE kineticMVS #-} + +-- | Leapfrog integrator with a user-supplied gradient function. Takes +-- the gradient function, parameter names, step size @ε@, number of +-- steps, initial @θ@ and momentum @r@, and returns the updated pair. +leapfrogWith + :: ([Text] -> Params -> [Double]) + -> [Text] + -> Double + -> Int + -> Params + -> [Double] + -> (Params, [Double]) +leapfrogWith gradFn names eps steps theta0 r0 = go steps theta0 r0 + where + go 0 theta r = (theta, r) + go n theta r = + let g = gradFn names theta + rHalf = zipWith (\ri gi -> ri - (eps / 2) * gi) r g + tVec' = zipWith (\ti ri -> ti + eps * ri) (paramsToVec names theta) rHalf + theta' = vecToParams names tVec' + g' = gradFn names theta' + r' = zipWith (\ri gi -> ri - (eps / 2) * gi) rHalf g' + in go (n - 1) theta' r' + +-- | Leapfrog integrator with a diagonal mass matrix. +-- +-- * Position update: @θ' = θ + ε · M⁻¹ · r@ (so smaller @M_ii@ +-- ⇒ slower per-step move along that coordinate, matching the +-- intent that posterior-narrow directions get smaller steps). +-- * Momentum update: @r' = r − (ε/2) · ∇U(θ)@ (unchanged). +-- +-- @leapfrogWith = leapfrogWithM (repeat 1)@. +leapfrogWithM + :: ([Text] -> Params -> [Double]) + -> [Text] + -> [Double] -- ^ Diagonal @M⁻¹@ (length = number of params). + -> Double -- ^ Step size @ε@. + -> Int -- ^ Number of leapfrog steps. + -> Params + -> [Double] + -> (Params, [Double]) +leapfrogWithM gradFn names mInv eps steps theta0 r0 = go steps theta0 r0 + where + go 0 theta r = (theta, r) + go n theta r = + let g = gradFn names theta + rHalf = zipWith (\ri gi -> ri - (eps / 2) * gi) r g + -- θ' = θ + ε · M⁻¹ · r + tVec' = zipWith3 (\ti m_inv ri -> ti + eps * m_inv * ri) + (paramsToVec names theta) mInv rHalf + theta' = vecToParams names tVec' + g' = gradFn names theta' + r' = zipWith (\ri gi -> ri - (eps / 2) * gi) rHalf g' + in go (n - 1) theta' r' + +-- | Storable-Vector–native variant of 'leapfrogWithM'. Position, +-- momentum, gradient, and the diagonal @M⁻¹@ all live on +-- @VS.Vector Double@ throughout the integration; no @Map@ or +-- @[Double]@ traversal occurs in the inner loop. +-- +-- Used by 'Hanalyze.MCMC.NUTS' where each leapfrog step is invoked up to +-- @2¹⁰@ times per iteration: the previous form went +-- @[Double] → Map → [Double]@ at every step (per-name @Map.lookup@ +-- ×p plus list cell allocation for @zipWith3@), which dominated the +-- profile after the algorithmic improvements were in place. +leapfrogWithMVS + :: (VS.Vector Double -> VS.Vector Double) -- ^ Gradient (Vector → Vector). + -> VS.Vector Double -- ^ Diagonal @M⁻¹@. + -> Double -- ^ Step size @ε@. + -> Int -- ^ Steps. + -> VS.Vector Double -- ^ Initial @θ@. + -> VS.Vector Double -- ^ Initial @r@. + -> (VS.Vector Double, VS.Vector Double) +leapfrogWithMVS gradFn mInv eps steps theta0 r0 = go steps theta0 r0 + where + !halfEps = eps * 0.5 + -- Phase 54.7b: SCC で勾配呼出 (= compiled カーネル) と VS 更新を分離計測。 + go !n theta r + | n <= 0 = (theta, r) + | otherwise = + let g = {-# SCC "leapfrog_grad1" #-} gradFn theta + rHalf = {-# SCC "leapfrog_vec_rhalf" #-} + VS.zipWith (\ri gi -> ri - halfEps * gi) r g + theta' = {-# SCC "leapfrog_vec_pos" #-} + VS.zipWith3 (\ti m_inv ri -> ti + eps * m_inv * ri) + theta mInv rHalf + g' = {-# SCC "leapfrog_grad2" #-} gradFn theta' + r' = {-# SCC "leapfrog_vec_rfull" #-} + VS.zipWith (\ri gi -> ri - halfEps * gi) rHalf g' + in go (n - 1) theta' r' + +-- --------------------------------------------------------------------------- +-- HMC サンプラー (AD 勾配版) +-- --------------------------------------------------------------------------- + +-- | HMC sampler for a polymorphic HBM model ('ModelP'). +-- +-- Uses AD gradients ('Numeric.AD.Mode.Reverse.Double'), so it is more accurate +-- and faster than numeric differentiation. Constraint transforms are +-- detected automatically from the priors via @getTransforms@. +hmc :: PrimMonad m => ModelP r -> HMCConfig -> Params -> Gen (PrimState m) -> m Chain +hmc m cfg initC gen = do + let names = sampleNames m + trMap = getTransforms m + transList = [Map.findWithDefault errT n trMap | n <- names] + errT = error "hmc: missing transform (should not happen)" + + initU = Map.fromList + [ (n, toUnconstrained t v) + | (n, t) <- zip names transList + , Just v <- [Map.lookup n initC] ] + + total = hmcBurnIn cfg + hmcIterations cfg + + logJU :: Params -> Double + logJU paramsU = logJointUnconstrained m names transList paramsU + + gradFn :: [Text] -> Params -> [Double] + gradFn ns paramsU = + let xs = [Map.findWithDefault 0 n paramsU | n <- ns] + in map negate (gradADU m names transList xs) + + samplesRef <- newMutVar [] + energyRef <- newMutVar ([] :: [Double]) + acceptedRef <- newMutVar (0 :: Int) + + let step currentU = do + r <- forM names (\_ -> standard gen) + let h0 = -(logJU currentU) + kinetic r + (proposedU, rFinal) = + leapfrogWith gradFn names + (hmcStepSize cfg) (hmcLeapfrogSteps cfg) + currentU r + logAlpha = (logJU proposedU - kinetic rFinal) + - (logJU currentU - kinetic r) + u <- uniform gen + nextU <- if log (u :: Double) < logAlpha + then do modifyMutVar' acceptedRef (+1); return proposedU + else return currentU + return (nextU, h0) + + let toConstrained pu = Map.fromList + [ (n, fromUnconstrained t (Map.findWithDefault 0 n pu)) + | (n, t) <- zip names transList ] + + let loop 0 currentU = return currentU + loop i currentU = do + (nextU, h0) <- step currentU + when (i <= hmcIterations cfg) $ do + modifyMutVar' samplesRef (toConstrained nextU :) + modifyMutVar' energyRef (h0 :) + loop (i - 1) nextU + + _ <- loop total initU + samples <- fmap reverse (readMutVar samplesRef) + energies <- fmap reverse (readMutVar energyRef) + accepted <- readMutVar acceptedRef + return Chain + { chainSamples = samples + , chainAccepted = accepted + , chainTotal = total + , chainEnergy = energies + , chainDivergences = [] + , chainTreeDepths = [] + } + +-- | Run 'hmc' on @numChains@ parallel chains (use @+RTS -N@ for CPU +-- parallelism). +hmcChains :: ModelP r -> HMCConfig -> Int -> Params -> GenIO -> IO [Chain] +hmcChains m cfg numChains initC baseGen = do + gens <- replicateM numChains (spawnGen baseGen) + mapConcurrently (\g -> hmc m cfg initC g) gens + +-- | [日本語]: 純粋・決定的な HMC (seed → 確定 Chain)。 +-- [English]: A pure, deterministic HMC run (seed → a fixed Chain). +hmcPure :: ModelP r -> HMCConfig -> Params -> Word32 -> Chain +hmcPure m cfg initC seed = + runST (initialize (V.singleton seed) >>= hmc m cfg initC) + +-- | [日本語]: 純粋・決定的な multi-chain HMC。 子 seed を純粋導出し @parList rdeepseq@ で並列。 +-- [English]: A pure, deterministic multi-chain HMC run. Child seeds are +-- derived purely, and chains run in parallel via @parList rdeepseq@. +hmcChainsPure :: ModelP r -> HMCConfig -> Int -> Params -> Word32 -> [Chain] +hmcChainsPure m cfg numChains initC seed = + let childSeeds :: [Word32] + childSeeds = runST $ do + g <- initialize (V.singleton seed) + replicateM numChains (uniform g) + chains = [ hmcPure m cfg initC s | s <- childSeeds ] + in chains `using` parList rdeepseq diff --git a/src/Hanalyze/MCMC/MH.hs b/src/Hanalyze/MCMC/MH.hs new file mode 100644 --- /dev/null +++ b/src/Hanalyze/MCMC/MH.hs @@ -0,0 +1,134 @@ +-- | +-- Module : Hanalyze.MCMC.MH +-- Description : Random-Walk Metropolis-Hastings サンプラー +-- Copyright : (c) 2026 Aelysce Project (Toshiaki Honda) +-- License : BSD-3-Clause +-- +-- Random-Walk Metropolis-Hastings sampler. +-- +-- Tune the per-parameter step sizes ('mcmcStepSizes') so the acceptance rate +-- lands in the 20-50% range. Pair 'Hanalyze.MCMC.Core.Chain' with +-- 'Hanalyze.Viz.Report.renderReport' to produce diagnostic plots. +{-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE RankNTypes #-} +module Hanalyze.MCMC.MH + ( MCMCConfig (..) + , defaultMCMCConfig + , metropolis + , metropolisChains + , metropolisPure + , metropolisChainsPure + ) where + +import Control.Concurrent.Async (mapConcurrently) +import Control.Monad (forM, replicateM) +import Control.Monad.Primitive (PrimMonad, PrimState) +import Control.Monad.ST (runST) +import Control.Parallel.Strategies (parList, rdeepseq, using) +import Data.Primitive.MutVar +import Data.Word (Word32) +import qualified Data.Map.Strict as Map +import qualified Data.Vector as V +import Data.Text (Text) +import System.Random.MWC (Gen, GenIO, uniform, initialize) +import System.Random.MWC.Distributions (normal) + +import Hanalyze.Model.HBM (ModelP, Params, logJoint, sampleNames) +import Hanalyze.MCMC.Core (Chain (..), spawnGen) + +-- --------------------------------------------------------------------------- +-- Configuration +-- --------------------------------------------------------------------------- + +-- | Random-Walk Metropolis configuration. +data MCMCConfig = MCMCConfig + { mcmcIterations :: Int -- ^ Total iterations (burn-in included). + , mcmcBurnIn :: Int -- ^ Burn-in iterations to discard. + , mcmcStepSizes :: Map.Map Text Double -- ^ Per-parameter proposal step. + } deriving (Show) + +-- | Default configuration: 2000 iterations, 500 burn-in, step size 1.0 +-- for every parameter. +defaultMCMCConfig :: [Text] -> MCMCConfig +defaultMCMCConfig names = MCMCConfig + { mcmcIterations = 2000 + , mcmcBurnIn = 500 + , mcmcStepSizes = Map.fromList [(n, 1.0) | n <- names] + } + +-- --------------------------------------------------------------------------- +-- Random Walk Metropolis +-- --------------------------------------------------------------------------- + +-- | Run Random-Walk Metropolis. Uses a joint proposal that updates all +-- latent variables simultaneously. +metropolis :: PrimMonad m => ModelP r -> MCMCConfig -> Params -> Gen (PrimState m) -> m Chain +metropolis model cfg init_ gen = do + let names = sampleNames model + total = mcmcBurnIn cfg + mcmcIterations cfg + steps = mcmcStepSizes cfg + + samplesRef <- newMutVar [] + acceptedRef <- newMutVar (0 :: Int) + + let step current = do + proposed <- fmap Map.fromList $ forM names $ \n -> do + let s = Map.findWithDefault 1.0 n steps + cur = Map.findWithDefault 0.0 n current + eps <- normal 0 s gen + return (n, cur + eps) + let logA = logJoint model proposed - logJoint model current + u <- uniform gen + if log (u :: Double) < logA + then do modifyMutVar' acceptedRef (+1) + return proposed + else return current + + let loop 0 current = return current + loop i current = do + next <- step current + if i <= mcmcIterations cfg + then modifyMutVar' samplesRef (next :) + else return () + loop (i - 1) next + + _ <- loop total init_ + samples <- fmap reverse (readMutVar samplesRef) + accepted <- readMutVar acceptedRef + return Chain + { chainSamples = samples + , chainAccepted = accepted + , chainTotal = total + , chainEnergy = [] + , chainDivergences = [] + , chainTreeDepths = [] + } + +-- | Run 'metropolis' on @numChains@ parallel chains, each with an +-- independent RNG (use @+RTS -N@ to run on multiple cores). +metropolisChains :: ModelP r -> MCMCConfig -> Int -> Params -> GenIO -> IO [Chain] +metropolisChains model cfg numChains initP baseGen = do + gens <- replicateM numChains (spawnGen baseGen) + mapConcurrently (\g -> metropolis model cfg initP g) gens + +-- | [日本語]: 純粋・決定的な Metropolis (seed → 確定 Chain・IO 不要)。 +-- [English]: A pure, deterministic Metropolis run (seed → a fixed Chain, +-- no IO needed). +metropolisPure :: ModelP r -> MCMCConfig -> Params -> Word32 -> Chain +metropolisPure model cfg initP seed = + runST (initialize (V.singleton seed) >>= metropolis model cfg initP) + +-- | [日本語]: 純粋・決定的な multi-chain Metropolis。 親 seed から子 seed を純粋導出し +-- 各 chain 別 'runST' → @parList rdeepseq@ で chain 横断を並列評価 (決定性は seed 由来)。 +-- [English]: A pure, deterministic multi-chain Metropolis run. Child +-- seeds are derived purely from the parent seed; each chain gets its own +-- 'runST', and chains are evaluated in parallel across chains via +-- @parList rdeepseq@ (determinism comes from the seed). +metropolisChainsPure :: ModelP r -> MCMCConfig -> Int -> Params -> Word32 -> [Chain] +metropolisChainsPure model cfg numChains initP seed = + let childSeeds :: [Word32] + childSeeds = runST $ do + g <- initialize (V.singleton seed) + replicateM numChains (uniform g) + chains = [ metropolisPure model cfg initP s | s <- childSeeds ] + in chains `using` parList rdeepseq diff --git a/src/Hanalyze/MCMC/NUTS.hs b/src/Hanalyze/MCMC/NUTS.hs new file mode 100644 --- /dev/null +++ b/src/Hanalyze/MCMC/NUTS.hs @@ -0,0 +1,963 @@ +-- | +-- Module : Hanalyze.MCMC.NUTS +-- Description : No-U-Turn Sampler (NUTS) — Hoffman & Gelman (2014) Algorithm 3 実装 +-- Copyright : (c) 2026 Aelysce Project (Toshiaki Honda) +-- License : BSD-3-Clause +-- +-- No-U-Turn Sampler (NUTS). +-- +-- Implements Hoffman & Gelman (2014) Algorithm 3, with Nesterov dual +-- averaging for step-size adaptation (Stan's strategy). Gradients are +-- exact, computed via 'Numeric.AD.Mode.Reverse.Double' ([日本語]: reverse +-- モードで勾配を latent 数非依存の ~1 sweep に。 旧 forward は O(p) だった。 +-- [English]: reverse-mode autodiff makes the gradient cost ~1 sweep +-- independent of the number of latent parameters; the previous +-- forward-mode implementation was O(p)). +-- +-- Constrained parameters (@PositiveT@, @UnitIntervalT@) are detected +-- automatically from the prior distribution. +-- +-- @ +-- import Hanalyze.Model.HBM +-- import Hanalyze.MCMC.NUTS +-- +-- chain <- nuts myModel defaultNUTSConfig +-- (Map.fromList [("mu",0),("sigma",1)]) gen +-- @ +{-# LANGUAGE StrictData #-} +{-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE RankNTypes #-} +module Hanalyze.MCMC.NUTS + ( NUTSConfig (..) + , defaultNUTSConfig + , nuts + , nutsStream + , nutsChains + , nutsPure + , nutsChainsPure + , nutsChainsStream + , chainSeeds + , SampleEvent (..) + ) where + +import Control.Concurrent.Async (mapConcurrently) +import Control.Monad (foldM, replicateM, when) +import Control.Monad.ST (ST, runST) +import Control.Parallel.Strategies (parList, rdeepseq, using) +import Data.Primitive.MutVar +import Data.Word (Word32) +import qualified Data.Map.Strict as Map +import qualified Data.Vector as V +import qualified Data.Vector.Storable as VS +import System.Random.MWC (Gen, GenIO, uniform, initialize) +import Control.Monad.Primitive (PrimMonad, PrimState, RealWorld) +import System.Random.MWC.Distributions (standard) + +import Hanalyze.MCMC.Core (Chain (..), spawnGen) +import Hanalyze.MCMC.HMC (kineticMVS, leapfrogWithMVS) +import Hanalyze.Model.HBM (ModelP, Params, sampleNames, getTransforms, + compileGradUV, compileGradValUVM, compileLogPUV) +import Hanalyze.Stat.Distribution (toUnconstrained, fromUnconstrained) + +-- --------------------------------------------------------------------------- +-- Configuration +-- --------------------------------------------------------------------------- + +-- | NUTS configuration. +data NUTSConfig = NUTSConfig + { nutsIterations :: Int -- ^ Post-burn-in draws to keep (the loop runs + -- @nutsBurnIn + nutsIterations@ total). + , nutsBurnIn :: Int -- ^ Burn-in iterations to discard. + , nutsStepSize :: Double -- ^ Initial leapfrog step size @ε@. + , nutsMaxDepth :: Int -- ^ Maximum tree depth (typically 10). + , nutsAdaptStepSize :: Bool -- ^ Enable Nesterov dual-averaging step-size adaptation. + , nutsTargetAccept :: Double -- ^ Target acceptance rate (0.8 typical, 0.95 for hard problems). + , nutsWarmupInitMaxDepth :: Maybe Int + -- ^ [日本語]: 質量行列の__初回更新前__ (init + -- buffer + 第 1 window・M=I 期間) に適用する + -- tree depth 上限 (opt-in・既定 'Nothing' = + -- 無効)。 M=I では幾何が合わず dual averaging + -- の ε 鋸歯で depth 7-10 の木を掘り radon 実測 + -- で warmup leapfrog の 68% を浪費するため、 + -- 'Just 6' 等で抑制できる。 ただし参照実装 + -- (Stan/PyMC) に無いヒューリスティックゆえ + -- 既定 OFF — 原理側の対策は 'nutsInitEpsSearch'。 + -- 'nutsAdaptMass' が False のときは不適用。 + -- [English]: The tree-depth cap applied + -- __before the mass matrix's first update__ + -- (during the init buffer + first window, + -- the M=I period; opt-in, default + -- 'Nothing' = disabled). Under M=I the + -- geometry doesn't match, so dual + -- averaging's ε saws back and forth, + -- digging depth-7-10 trees; radon + -- measurements showed this wastes 68% of + -- warmup leapfrogs, which can be curbed + -- with e.g. 'Just 6'. However, since this + -- heuristic has no counterpart in the + -- reference implementations (Stan/PyMC), + -- it defaults to OFF — the principled fix + -- is 'nutsInitEpsSearch'. Not applied when + -- 'nutsAdaptMass' is False. + , nutsInitEpsSearch :: Bool -- ^ [日本語]: Stan (Hoffman–Gelman + -- Algorithm 4) の ε 倍加探索を (i) サンプリング + -- 開始前と (ii) 質量行列の各 window 末更新直後 + -- (Stan adapt_diag_e_nuts の + -- init_stepsize+restart と同順) に行う (既定 + -- True)。 DA anchor (μ = log 10ε) が幾何と + -- 乖離すると ε が鋸歯振動して深い木を掘るため、 + -- ε を 1 step leapfrog の受容率 ~1/2 になる値へ + -- 都度較正する (Stan と同じ標準機構)。 + -- 'nutsAdaptStepSize' が True のときのみ有効。 + -- [English]: Runs Stan's (Hoffman–Gelman + -- Algorithm 4) ε doubling search + -- (i) before sampling starts and + -- (ii) right after each mass-matrix window + -- ends (in the same order as Stan's + -- @adapt_diag_e_nuts@ @init_stepsize@ + + -- restart; default True). If the DA + -- anchor (μ = log 10ε) drifts away from the + -- geometry, ε oscillates and digs deep + -- trees, so ε is recalibrated each time to + -- the value that gives a ~1/2 acceptance + -- rate for one leapfrog step (the same + -- standard mechanism as Stan). Only takes + -- effect when 'nutsAdaptStepSize' is True. + , nutsAdaptMass :: Bool -- ^ Enable diagonal mass-matrix adaptation (B11). + -- Stan-style multi-window: init buffer (15% / + -- ≥75 iter, M=I) → doubling windows + -- 25→50→100→200→… (M updated + dual avg + -- restarted at each window end) → term buffer + -- (10% / ≥50 iter, M frozen, ε converges). + -- Recommended for posteriors with strongly + -- varying scales across parameters. + , nutsInitJitter :: Double -- ^ [日本語]: 各 chain の初期位置 (unconstrained) + -- に加える一様 jitter 半幅 (PyMC jitter+adapt_diag + -- 相当)。 chain ごとに独立に @U(-j, +j)@ を各成分へ + -- 加算し、 funnel 首での whole-chain 崩壊 (全 chain + -- 同一 init 由来) を減らす。 @0@ = 無操作 (= 従来 + -- 挙動・単一 chain 再現性テスト非影響)。 多 chain + -- 経路 (@hbmNutsConfig@) で 1.0 を設定。 + -- [English]: The half-width of the uniform + -- jitter added to each chain's initial + -- position (unconstrained space; the same + -- idea as PyMC's @jitter+adapt_diag@). An + -- independent @U(-j, +j)@ is added to each + -- coordinate per chain, reducing + -- whole-chain collapse at a funnel's neck + -- (which happens when all chains share the + -- same init). @0@ = no-op (= the previous + -- behavior; does not affect single-chain + -- reproducibility tests). The multi-chain + -- path (@hbmNutsConfig@) sets this to 1.0. + } deriving (Show) + +-- | Default NUTS configuration: 2000 post-burn-in draws, 500 burn-in +-- (2500 total), @ε = 0.1@, +-- max depth 10, dual averaging enabled, target acceptance 0.8, +-- diagonal mass-matrix adaptation off (opt-in via 'nutsAdaptMass'). +defaultNUTSConfig :: NUTSConfig +defaultNUTSConfig = NUTSConfig + { nutsIterations = 2000 + , nutsBurnIn = 500 + , nutsStepSize = 0.1 + , nutsMaxDepth = 10 + , nutsAdaptStepSize = True + , nutsTargetAccept = 0.8 + , nutsWarmupInitMaxDepth = Nothing + , nutsInitEpsSearch = True + , nutsAdaptMass = False + , nutsInitJitter = 0.0 + } + +-- --------------------------------------------------------------------------- +-- Dual averaging +-- --------------------------------------------------------------------------- + +-- | Internal state for Nesterov's dual-averaging step-size adaptation. +data DualAvgState = DualAvgState + { daLogEps :: Double -- ^ Current @log ε@ used for sampling. + , daLogEpsBar :: Double -- ^ Running smoothed @log ε̄@ (post-adaptation value). + , daH :: Double -- ^ Running average of (target − accept-stat). + , daMu :: Double -- ^ Anchor @μ = log(10 ε₀)@. + , daM :: Int -- ^ Iteration counter. + } + +-- | Initialize 'DualAvgState' from an initial step size @ε₀@. +initDualAvg :: Double -> DualAvgState +initDualAvg eps0 = DualAvgState + { daLogEps = log eps0 + , daLogEpsBar = log eps0 + , daH = 0.0 + , daMu = log (10 * eps0) + , daM = 0 + } + +-- | [日本語]: Stan (@base_hmc::init_stepsize@) 準拠の ε 探索。 +-- 与えられた ε を起点に、 1 step leapfrog の受容比が 0.8 を跨ぐまで倍加/半減 +-- する (運動量は試行ごとに再サンプル)。 dual averaging の anchor μ = log(10 ε₀) +-- が幾何に合った値になり、 ε 鋸歯振動 (radon で depth 7-10 の深掘り) を防ぐ。 +-- ★Hoffman–Gelman 2014 Alg.4 (起点 1.0・閾値 1/2・運動量 1 本固定) でなく +-- Stan 実装 (起点 = 現在 ε・閾値 0.8・毎試行再サンプル) に合わせる — +-- window 末の再較正では既適応の ε 近傍から保守的に探す必要がある +-- (起点 1.0/閾値 0.5 は radon 実測で新 M 下の trajectory が支えない大きな ε を +-- 返し、 次 window の深掘りを招いた)。 非有限 (発散) は比 −∞ 扱い = 半減方向。 +-- 反復と ε は安全側に有界。 +-- +-- [English]: An ε search following Stan's @base_hmc::init_stepsize@. +-- Starting from a given ε, it doubles or halves until one leapfrog +-- step's acceptance ratio crosses 0.8 (momentum is resampled on every +-- trial). This brings dual averaging's anchor μ = log(10 ε₀) into line +-- with the geometry, preventing ε sawtooth oscillation (which digs +-- depth-7-10 trees, per radon measurements). ★This follows Stan's +-- implementation (starting point = the current ε, threshold 0.8, +-- resampling momentum on every trial), not Hoffman–Gelman 2014's +-- Algorithm 4 (starting point 1.0, threshold 1/2, a single fixed +-- momentum draw) — the recalibration at a window's end needs to search +-- conservatively near the already-adapted ε (a starting point of 1.0 / +-- threshold of 0.5 returned, per radon measurements, an ε too large for +-- the trajectory to sustain under the new M, triggering deep digging in +-- the next window). Non-finite (divergent) is treated as a ratio of +-- −∞, i.e. the halving direction. Both the iteration count and ε are +-- bounded on the safe side. +findReasonableEpsilon + :: PrimMonad m + => (VS.Vector Double -> VS.Vector Double) -- ^ [日本語]: gradFn (−∇ logπ・NUTS と同じ向き)。 [English]: gradFn (−∇ logπ; same sign convention as NUTS). + -> (VS.Vector Double -> Double) -- ^ logπ (unconstrained) + -> VS.Vector Double -- ^ [日本語]: M⁻¹ 対角。 [English]: The diagonal of M⁻¹. + -> Double -- ^ [日本語]: 探索起点 ε (現在の nominal ε)。 [English]: The search's starting ε (the current nominal ε). + -> VS.Vector Double -- ^ [日本語]: 初期位置 θ (unconstrained)。 [English]: The initial position θ (unconstrained). + -> Gen (PrimState m) + -> m Double +findReasonableEpsilon gradFn logPiFn mInv eps0 theta gen = do + dH0 <- trial epsInit + let dir = if dH0 > thresh then 1 else -1 :: Int + loop dir epsInit (50 :: Int) + where + epsInit = max 1e-10 (min 1e7 eps0) + thresh = log 0.8 + -- 1 step leapfrog の log 受容比 (Stan と同じく運動量を都度引き直す)。 + trial eps = do + r0 <- sampleMomentum mInv gen + let h0 = negate (logPiFn theta) + kineticMVS mInv r0 + (th', r') = leapfrogWithMVS gradFn mInv eps 1 theta r0 + h' = negate (logPiFn th') + kineticMVS mInv r' + pure (if isNaN h' || isInfinite h' then (-1) / 0 else h0 - h') + loop dir !eps !k + | k <= 0 = pure eps + | otherwise = do + dH <- trial eps + let keepGoing = if dir == 1 then dH > thresh else dH < thresh + eps' = if dir == 1 then eps * 2 else eps / 2 + if not keepGoing then pure eps + else if eps' > 1e7 || eps' < 1e-10 then pure eps + else loop dir eps' (k - 1) + +-- | Apply one dual-averaging update given the target acceptance @δ@ and +-- the observed acceptance statistic @α@ for the iteration. +updateDualAvg :: Double -> Double -> DualAvgState -> DualAvgState +updateDualAvg delta alpha da = + let m = daM da + 1 + gamma = 0.05 + t0 = 10.0 + kappa = 0.75 + hNew = (1 - 1 / (fromIntegral m + t0)) * daH da + + (1 / (fromIntegral m + t0)) * (delta - alpha) + logEps = daMu da - sqrt (fromIntegral m) / gamma * hNew + logEpsClip = max (-7) (min 5 logEps) + logEpsBar = (fromIntegral m ** (-kappa)) * logEpsClip + + (1 - fromIntegral m ** (-kappa)) * daLogEpsBar da + in da { daLogEps = logEpsClip, daLogEpsBar = logEpsBar, daH = hNew, daM = m } + +-- --------------------------------------------------------------------------- +-- 内部ツリー +-- --------------------------------------------------------------------------- + +-- | Internal NUTS tree node. All position/momentum are 'VS.Vector +-- Double' rather than 'Params' (= @Map@) / @[Double]@: the +-- @doubleTree@ recursion creates up to @2¹⁰@ intermediate trees per +-- iteration, and the previous Map / list representation paid an +-- order-of-magnitude in allocation that swamped the actual leapfrog +-- arithmetic. +data NUTSTree = NUTSTree + { ntThMinus :: VS.Vector Double + , ntRMinus :: VS.Vector Double + , ntGMinus :: VS.Vector Double + -- ^ [日本語]: minus 端点の ∇U = −∇logπ (leapfrog 勾配キャッシュ)。 + -- 同方向の次の葉が始点勾配を再計算せずに済む (Stan の z_.g と同じ)。 + -- [English]: The ∇U = −∇logπ at the minus endpoint (a leapfrog + -- gradient cache). This lets the next leaf in the same direction + -- avoid recomputing the starting gradient (same idea as Stan's z_.g). + , ntThPlus :: VS.Vector Double + , ntRPlus :: VS.Vector Double + , ntGPlus :: VS.Vector Double + -- ^ [日本語]: plus 端点の ∇U (同上)。 + -- [English]: The ∇U at the plus endpoint (same idea as above). + , ntThPrime :: VS.Vector Double + , ntN :: Int + , ntS :: Bool + , ntDiv :: Bool + -- ^ [日本語]: サブツリー中で divergent (|ΔH| > deltaMax) が発生したか + -- [English]: Whether a divergence (|ΔH| > deltaMax) occurred anywhere + -- in the subtree. + , ntASum :: !Double + -- ^ [日本語]: Σ min(1, exp(H0 − H_leaf)) — Stan の accept_stat 蓄積。 + -- dual averaging はこの平均 ᾱ を学習する (旧: 1-step probe = 毎 draw + -- 余分な leapfrog+エネルギー評価を払う非標準の独自実装だった)。 + -- [English]: The accumulated Σ min(1, exp(H0 − H_leaf)) — Stan's + -- accept_stat. Dual averaging learns from the mean ᾱ of this (the + -- previous 1-step probe was a non-standard bespoke implementation that + -- paid for an extra leapfrog + energy evaluation on every draw). + , ntANum :: !Int + -- ^ [日本語]: ᾱ の分母 (サブツリーの葉数・棄却葉も含む)。 + -- [English]: The denominator of ᾱ (the number of leaves in the + -- subtree, including rejected leaves). + } + +deltaMax :: Double +deltaMax = 1000.0 + +-- | U-turn check on Storable Vectors. @(θ⁺ − θ⁻) · r⁻ < 0@ or +-- @(θ⁺ − θ⁻) · r⁺ < 0@ ⇒ trajectory has begun to retrace itself. +-- +-- [日本語]: 旧実装は @delta@ の共有 binding で stream fusion が切れ delta +-- ベクトルを毎回実体化していた (prof 実測: nuts_uturn が総 alloc の +-- 23.5%)。 2 つの内積を単一パス・確保なしで融合する。 加算順序は旧 +-- 'VS.sum' (左畳み込み) と同一 = ビット同一。 +-- [English]: The previous implementation lost stream fusion at a shared +-- @delta@ binding and materialized the delta vector on every call +-- (profiling showed nuts_uturn accounted for 23.5% of total allocation). +-- This version fuses the two dot products into a single pass with no +-- allocation. The summation order is identical to the previous 'VS.sum' +-- (left fold), so results are bit-identical. +uTurnVS + :: VS.Vector Double -> VS.Vector Double + -> VS.Vector Double -> VS.Vector Double -> Bool +uTurnVS thMinus rMinus thPlus rPlus = go 0 0 0 + where + !n = VS.length thMinus + go !d1 !d2 !j + | j >= n = d1 < 0 || d2 < 0 + | otherwise = + let d = thPlus `VS.unsafeIndex` j - thMinus `VS.unsafeIndex` j + in go (d1 + d * (rMinus `VS.unsafeIndex` j)) + (d2 + d * (rPlus `VS.unsafeIndex` j)) + (j + 1) +{-# INLINE uTurnVS #-} + +-- | Sample momentum @r ~ N(0, M)@ from the diagonal mass matrix +-- represented as @M⁻¹@. Per coordinate: @r_i = z / sqrt(M⁻¹_i)@, +-- @z ~ N(0,1)@. Storable-vector tight loop, no list allocation. +sampleMomentum :: PrimMonad m => VS.Vector Double -> Gen (PrimState m) -> m (VS.Vector Double) +sampleMomentum mInv gen = do + let n = VS.length mInv + VS.generateM n $ \i -> do + z <- standard gen + return (z / sqrt (mInv `VS.unsafeIndex` i)) +{-# INLINE sampleMomentum #-} + +-- --------------------------------------------------------------------------- +-- ツリービルダー +-- --------------------------------------------------------------------------- + +buildTree + :: forall m. PrimMonad m + => (VS.Vector Double -> m (Double, VS.Vector Double)) + -- ^ [日本語]: 融合評価: θ ↦ (logπ(θ), ∇U(θ) = −∇logπ(θ))。 chain 閉包に + -- 確保した arena/adj を再利用するため monadic。 + -- [English]: A fused evaluation: θ ↦ (logπ(θ), ∇U(θ) = −∇logπ(θ)). + -- Monadic so it can reuse the arena/adjoint buffers allocated in the + -- chain's closure. + -> VS.Vector Double -- ^ Diagonal M⁻¹. + -> Double -- ^ Step size @ε@. + -> VS.Vector Double -- ^ Position. + -> VS.Vector Double -- ^ Momentum. + -> VS.Vector Double -- ^ [日本語]: position での ∇U (キャッシュ)。 [English]: ∇U at position (cached). + -> Double -- ^ @log u@ slice. + -> Double -- ^ [日本語]: 初期エネルギー @H0@ (ᾱ 用)。 [English]: The initial energy @H0@ (used for ᾱ). + -> Int -- ^ Direction (±1). + -> Int -- ^ Recursion depth. + -> Gen (PrimState m) + -> m NUTSTree +-- Phase 54.7b: PrimMonad 多相 (Phase 50) は SPECIALIZE が無いと dictionary 渡しで +-- mwc の uniform/standard が unbox されない (prof 実測で RNG 系 11%/alloc 25%)。 +-- IO / ST の両具体型に特殊化して Phase 50 以前の機械語品質に戻す。 +{-# SPECIALIZE buildTree + :: (VS.Vector Double -> IO (Double, VS.Vector Double)) + -> VS.Vector Double -> Double -> VS.Vector Double -> VS.Vector Double + -> VS.Vector Double + -> Double -> Double -> Int -> Int -> Gen RealWorld -> IO NUTSTree #-} +{-# SPECIALIZE buildTree + :: (VS.Vector Double -> ST s (Double, VS.Vector Double)) + -> VS.Vector Double -> Double -> VS.Vector Double -> VS.Vector Double + -> VS.Vector Double + -> Double -> Double -> Int -> Int -> Gen s -> ST s NUTSTree #-} +buildTree gradValU mInv eps theta r gU logU h0 dir depth gen + | depth == 0 = do + -- Phase 87.2b: 1-step leapfrog を融合評価でインライン化。 始点勾配は + -- 端点キャッシュ (gU) を使い、 終点は (logπ, ∇U) を 1 回の融合評価で + -- 取得 (旧: 葉ごとに grad 2 回 + logπ 1 回 = 始点勾配の再計算と + -- エネルギー用 forward の重複を払っていた)。 + let !epsD = fromIntegral dir * eps + !halfEps = 0.5 * epsD + rHalf = {-# SCC "nuts_leapfrog_kick1" #-} + VS.zipWith (\ri gi -> ri - halfEps * gi) r gU + theta' = {-# SCC "nuts_leapfrog_drift" #-} + VS.zipWith3 (\ti m_inv ri -> ti + epsD * m_inv * ri) + theta mInv rHalf + (v', g') <- {-# SCC "nuts_gradval" #-} gradValU theta' + let r' = {-# SCC "nuts_leapfrog_kick2" #-} + VS.zipWith (\ri gi -> ri - halfEps * gi) rHalf g' + h' = {-# SCC "nuts_energy" #-} (negate v' + kineticMVS mInv r') + n' = if logU <= -h' then 1 else 0 + s' = logU < deltaMax - h' + divergent = not s' + -- Phase 87.2: Stan の accept_stat = min(1, exp(H0 − H')) を葉ごとに + -- 蓄積 (非有限は 0 = 棄却扱い)。 + a' = let d = h0 - h' + in if isNaN d then 0 else min 1 (exp (min 0 d)) + return NUTSTree + { ntThMinus = theta', ntRMinus = r', ntGMinus = g' + , ntThPlus = theta', ntRPlus = r', ntGPlus = g' + , ntThPrime = theta', ntN = n', ntS = s' + , ntDiv = divergent + , ntASum = a', ntANum = 1 + } + | otherwise = do + t1 <- buildTree gradValU mInv eps theta r gU logU h0 dir (depth - 1) gen + if not (ntS t1) then return t1 + else do + let (th0, r0, g0) = if dir == -1 + then (ntThMinus t1, ntRMinus t1, ntGMinus t1) + else (ntThPlus t1, ntRPlus t1, ntGPlus t1) + t2 <- buildTree gradValU mInv eps th0 r0 g0 logU h0 dir (depth - 1) gen + let n1 = ntN t1; n2 = ntN t2 + thPrime' <- + if n1 == 0 then return (ntThPrime t2) + else if n2 == 0 then return (ntThPrime t1) + else do + u <- {-# SCC "nuts_rng_uniform" #-} (uniform gen :: m Double) + return $ if u < min 1.0 (fromIntegral n2 / fromIntegral n1) + then ntThPrime t2 + else ntThPrime t1 + let (minus', rMinus', gMinus', plus', rPlus', gPlus') = if dir == -1 + then (ntThMinus t2, ntRMinus t2, ntGMinus t2, + ntThPlus t1, ntRPlus t1, ntGPlus t1) + else (ntThMinus t1, ntRMinus t1, ntGMinus t1, + ntThPlus t2, ntRPlus t2, ntGPlus t2) + s' = ntS t2 && not ({-# SCC "nuts_uturn" #-} uTurnVS minus' rMinus' plus' rPlus') + return NUTSTree + { ntThMinus = minus', ntRMinus = rMinus', ntGMinus = gMinus' + , ntThPlus = plus', ntRPlus = rPlus', ntGPlus = gPlus' + , ntThPrime = thPrime', ntN = n1 + n2, ntS = s' + , ntDiv = ntDiv t1 || ntDiv t2 + , ntASum = ntASum t1 + ntASum t2 + , ntANum = ntANum t1 + ntANum t2 + } + +-- --------------------------------------------------------------------------- +-- Streaming hook +-- --------------------------------------------------------------------------- + +-- | Per-iteration sample event emitted by @nutsStream@. +-- +-- Used by callers that want to observe MCMC progress as it happens +-- (e.g. live trace plots, real-time R-hat / ESS updates over the wire). +-- The callback receives one event per iteration of the outer loop, +-- including burn-in iterations (distinguished by 'seIsBurnIn'). +-- +-- The 'seParams' values are in the __constrained__ parameter space, +-- matching the convention used in 'chainSamples'. Burn-in events are +-- /not/ included in 'chainSamples', but are still streamed via the +-- callback so the UI can show warmup progress and adaptation. +data SampleEvent = SampleEvent + { seIter :: !Int -- ^ 0-based iteration index (burn-in inclusive). + -- Ranges over @[0 .. nutsBurnIn + nutsIterations - 1]@. + , seIsBurnIn :: !Bool -- ^ True if @seIter < nutsBurnIn@. + , seParams :: !Params -- ^ Current sample (constrained space). + , seEnergy :: !Double -- ^ Hamiltonian H0 at the start of this iteration. + , seDivergent :: !Bool -- ^ Whether this iteration's trajectory diverged. + , seAccepted :: !Bool -- ^ Whether the proposal was accepted + -- (@proposedU /= currentU@). + , seStepSize :: !Double -- ^ Current ε (after this iteration's adaptation). + , seTreeDepth :: !Int -- ^ [日本語]: この draw で実行された doubling 回数 + -- (leapfrog 数 ≈ 2^depth・warmup 固定費の診断用)。 + -- [English]: The number of doublings performed + -- for this draw (leapfrog count ≈ 2^depth; + -- used to diagnose fixed warmup cost). + , seAcceptStat :: !Double -- ^ [日本語]: この draw の mean accept-stat α + -- (dual averaging が target と比較する統計・ + -- 'seAccepted' の bool とは別物)。ε̄ 収束診断用。 + -- [English]: This draw's mean accept-stat α + -- (the statistic dual averaging compares + -- against the target — distinct from the + -- 'seAccepted' bool). Used to diagnose ε̄ + -- convergence. + } + +-- --------------------------------------------------------------------------- +-- NUTS サンプラー +-- --------------------------------------------------------------------------- + +-- | NUTS sampler for a polymorphic HBM model ('ModelP'). +-- [日本語]: 軌道長は U-Turn 判定で自動決定。 +-- [English]: The trajectory length is determined automatically by the +-- U-turn check. +-- +-- This is a thin wrapper around @nutsStream@ with a no-op callback. +-- Use @nutsStream@ directly if you want per-iteration progress +-- (e.g. for live UI updates over a WebSocket / SSE channel). +nuts :: PrimMonad m => ModelP r -> NUTSConfig -> Params -> Gen (PrimState m) -> m Chain +{-# SPECIALIZE nuts :: ModelP r -> NUTSConfig -> Params -> Gen RealWorld -> IO Chain #-} +{-# SPECIALIZE nuts :: ModelP r -> NUTSConfig -> Params -> Gen s -> ST s Chain #-} +nuts m cfg initC gen = nutsStream m cfg initC gen (\_ -> pure ()) + +-- | NUTS sampler with a per-iteration callback. Identical to 'nuts' +-- semantically; in addition, calls @onSample event@ once per outer +-- loop iteration (burn-in inclusive). The callback runs synchronously +-- inside the sampler loop, so it should return quickly (push events to +-- a queue rather than do IO of unbounded latency). +nutsStream :: forall r m. PrimMonad m + => ModelP r -> NUTSConfig -> Params -> Gen (PrimState m) + -> (SampleEvent -> m ()) + -> m Chain +{-# SPECIALIZE nutsStream + :: ModelP r -> NUTSConfig -> Params -> Gen RealWorld + -> (SampleEvent -> IO ()) -> IO Chain #-} +{-# SPECIALIZE nutsStream + :: ModelP r -> NUTSConfig -> Params -> Gen s + -> (SampleEvent -> ST s ()) -> ST s Chain #-} +nutsStream m cfg initC gen onSample = do + let names = sampleNames m + trMap = getTransforms m + transList = [Map.findWithDefault errT n trMap | n <- names] + errT = error "nuts: missing transform" + + -- Initial unconstrained position as a Storable Vector. The hot + -- loop never touches 'Params' (= Map); we only convert at the + -- boundary to record samples. + initUV0 :: VS.Vector Double + initUV0 = VS.fromList + [ toUnconstrained t (Map.findWithDefault 0 n initC) + | (n, t) <- zip names transList ] + + total = nutsBurnIn cfg + nutsIterations cfg + doAdapt = nutsAdaptStepSize cfg && nutsBurnIn cfg > 0 + + -- Vector-native log target density. Phase 54.4d/54.6: エネルギー評価も + -- 'compileLogPUV' で静的部分 (名前→index 解決込み) を 1 度だけ前処理した + -- compiled closure を全 tree node で再利用する (旧: 毎回 + -- 'logJointUnconstrained' の Free walk + per-obs スカラ logDensityObs)。 + logPiFn :: VS.Vector Double -> Double + logPiFn = compileLogPUV m names transList + + -- Vector-native gradient. Phase 54.4b/54.6: モデル構造は draw 間で不変 + -- ゆえ @compileGradUV@ で静的部分を **1 度だけ**前処理し、 返った + -- vector-native クロージャを全 leapfrog で再利用する (VS↔list 変換なし)。 + gradV :: VS.Vector Double -> VS.Vector Double + gradV = compileGradUV m names transList + gradFn :: VS.Vector Double -> VS.Vector Double + gradFn uv = VS.map negate (gradV uv) + + toConstrained :: VS.Vector Double -> Params + toConstrained uv = Map.fromList + [ (n, fromUnconstrained t (uv `VS.unsafeIndex` i)) + | (i, (n, t)) <- zip [0..] (zip names transList) ] + + -- Phase 87.2b: 値+勾配の融合評価 (JAX value_and_grad 相当)。 tree の葉が + -- leapfrog 最終勾配とエネルギーを同一点で二重評価していた重複を除去。 + -- Phase 90 A11-4①: 'compileGradValUVM' は forward/随伴 arena を **この chain + -- 閉包生成時に 1 度だけ**確保して全 leapfrog で再利用する (per-call 34k×2 + -- セル確保 + GC churn を除去)。 chain ごとに別 @nutsStream@ 呼出 = 別バッファ + -- ゆえ chain 横断並列 (@nutsChainsPure@/'nutsChainsStream') と非干渉。 + -- Phase 94 A4-2: 各 chain の初期位置に一様 jitter (funnel 首の whole-chain 崩壊対策)。 + -- j=0 なら initUV0 をそのまま (従来挙動)。 gen は chain 固有ゆえ chain ごと独立。 + initUV <- let j = nutsInitJitter cfg + in if j <= 0 then pure initUV0 + else VS.mapM (\x -> do u <- uniform gen + pure (x + (u * 2 - 1) * j)) initUV0 + gradValV <- compileGradValUVM m names transList + let gradValU :: VS.Vector Double -> m (Double, VS.Vector Double) + gradValU uv = do + (v, g) <- gradValV uv + pure (v, VS.map negate g) + + samplesRef <- newMutVar [] + energyRef <- newMutVar ([] :: [Double]) + divergenceRef <- newMutVar ([] :: [Int]) + depthRef <- newMutVar ([] :: [Int]) -- Phase 85.3: per-draw tree depth + acceptedRef <- newMutVar (0 :: Int) + -- Phase 85.6c: 初期 ε の較正 (Stan Algorithm 4・doAdapt 時のみ)。 + eps0 <- if doAdapt && nutsInitEpsSearch cfg + then findReasonableEpsilon gradFn logPiFn + (VS.replicate (length names) 1.0) (nutsStepSize cfg) initUV gen + else pure (nutsStepSize cfg) + daRef <- newMutVar (initDualAvg eps0) + + -- B11: Stan-style multi-window diagonal mass-matrix adaptation. + -- + -- Schedule (warmup W): + -- * init buffer (max 75 / W÷7 iters): step-size adapt only, M = I + -- * window phase: doubling windows 25 → 50 → 100 → 200 → ... + -- At the end of each window: update M⁻¹ from window's + -- Welford-accumulated diagonal variance, restart dual averaging. + -- * term buffer (max 50 / W÷10 iters): M frozen, step-size adapt + -- continues to converge ε under the final geometry. + let nParams = length names + adaptM = nutsAdaptMass cfg && nutsBurnIn cfg > 0 + (windowEnds, initBuf, _termBuf) = stanWindows (nutsBurnIn cfg) + windowPhaseEnd = if null windowEnds then 0 else last windowEnds + mInvRef <- newMutVar (VS.replicate nParams 1.0) + welfordRef <- newMutVar (emptyWelford nParams) + + let step :: VS.Vector Double -> Double -> Int -> VS.Vector Double + -> m (VS.Vector Double, Double, Double, Bool, Int) + step mInv eps maxDep currentU = do + -- r ~ N(0, M) ⇔ r_i = sqrt(M_ii) * z = z / sqrt(M⁻¹_ii) + r0 <- {-# SCC "nuts_sampleMomentum" #-} sampleMomentum mInv gen + u0 <- {-# SCC "nuts_rng_uniform" #-} (uniform gen :: m Double) + -- Phase 87.2b: 始点の (logπ, ∇U) を融合評価 1 回で取得。 値は H0 に、 + -- 勾配は両方向の最初の葉の始点キャッシュに使う。 + (v0, gU0) <- {-# SCC "nuts_gradval0" #-} gradValU currentU + let h0 = negate v0 + kineticMVS mInv r0 + logU = log u0 - h0 + let tree0 = NUTSTree + { ntThMinus = currentU, ntRMinus = r0, ntGMinus = gU0 + , ntThPlus = currentU, ntRPlus = r0, ntGPlus = gU0 + , ntThPrime = currentU, ntN = 1, ntS = True + , ntDiv = False + , ntASum = 0, ntANum = 0 + } + let doubleTree tree j = + if not (ntS tree) then return tree + else do + u <- {-# SCC "nuts_rng_uniform" #-} (uniform gen :: m Double) + let dir = if u < 0.5 then -1 else 1 :: Int + (th0, r0', g0') = if dir == -1 + then (ntThMinus tree, ntRMinus tree, ntGMinus tree) + else (ntThPlus tree, ntRPlus tree, ntGPlus tree) + subtree <- {-# SCC "nuts_buildTree" #-} + buildTree gradValU mInv eps th0 r0' g0' logU h0 dir j gen + let n1 = ntN tree; n2 = ntN subtree + thPrime' <- + if not (ntS subtree) || n2 == 0 + then return (ntThPrime tree) + else do + u2 <- {-# SCC "nuts_rng_uniform" #-} (uniform gen :: m Double) + return $ if u2 < min 1.0 (fromIntegral n2 / fromIntegral n1) + then ntThPrime subtree + else ntThPrime tree + let (minus', rMinus', gMinus', plus', rPlus', gPlus') = if dir == -1 + then (ntThMinus subtree, ntRMinus subtree, ntGMinus subtree, + ntThPlus tree, ntRPlus tree, ntGPlus tree) + else (ntThMinus tree, ntRMinus tree, ntGMinus tree, + ntThPlus subtree, ntRPlus subtree, ntGPlus subtree) + s' = ntS subtree && not ({-# SCC "nuts_uturn" #-} uTurnVS minus' rMinus' plus' rPlus') + return NUTSTree + { ntThMinus = minus', ntRMinus = rMinus', ntGMinus = gMinus' + , ntThPlus = plus', ntRPlus = rPlus', ntGPlus = gPlus' + , ntThPrime = thPrime', ntN = n1 + n2, ntS = s' + , ntDiv = ntDiv tree || ntDiv subtree + , ntASum = ntASum tree + ntASum subtree + , ntANum = ntANum tree + ntANum subtree + } + -- Phase 85.3: 実行された doubling 回数 = tree depth (PyMC の + -- tree_depth 相当・leapfrog 数 ≈ 2^depth) を数える。 + let doubleTreeD (tree, !dep) j = + if not (ntS tree) then return (tree, dep) + else do + t' <- doubleTree tree j + return (t', dep + 1 :: Int) + (finalTree, treeDepth) <- + foldM doubleTreeD (tree0, 0) [0 .. maxDep - 1] + -- Phase 87.2: alpha は Stan の accept_stat = tree 全葉の + -- min(1, exp(H0−H')) 平均 (buildTree で蓄積)。 旧 1-step probe + -- (毎 draw 余分な leapfrog + エネルギー評価・非標準の独自実装) を廃止。 + let proposedU = ntThPrime finalTree + alpha = if ntANum finalTree > 0 + then ntASum finalTree / fromIntegral (ntANum finalTree) + else 0 + when (proposedU /= currentU) $ modifyMutVar' acceptedRef (+1) + return (proposedU, alpha, h0, ntDiv finalTree, treeDepth) + + let loop 0 currentU _eps = return currentU + loop i currentU eps = do + mInv <- readMutVar mInvRef + let isBurnIn = i > nutsIterations cfg + -- iteration index from start (1-based); total counts down. + iterIdx = total - i + 1 + -- Phase 85.6: M 初回更新前 (M=I) は深い木を掘らない (init 期の + -- draw は捨てる区間・radon で warmup leapfrog の 68% を占めた)。 + firstMUpd = case windowEnds of { (w : _) -> w; [] -> 0 } + maxDep + | adaptM && isBurnIn && iterIdx <= firstMUpd + , Just cap <- nutsWarmupInitMaxDepth cfg = + min (nutsMaxDepth cfg) cap + | otherwise = nutsMaxDepth cfg + -- Inside the window phase: collect samples for Welford. + inWindowPhase = adaptM && isBurnIn + && iterIdx > initBuf + && iterIdx <= windowPhaseEnd + -- This iteration ends a window: update M, restart DA. + isWindowEnd = adaptM && isBurnIn && iterIdx `elem` windowEnds + (nextU, alpha, h0, divergent, treeDepth) <- step mInv eps maxDep currentU + when inWindowPhase $ + modifyMutVar' welfordRef (\w -> {-# SCC "nuts_welford" #-} welfordAddVS w nextU) + -- Phase 86: window 末に M を更新したら、 Stan (adapt_diag_e_nuts の + -- init_stepsize + restart) と同じく**新 metric の下で ε を再較正** + -- (Algorithm 4) して DA を restart する。 旧実装は鋸歯振動中の瞬間値 + -- ε を anchor (μ = log 10ε) にしており、 M 更新直後に ε が幾何と桁で + -- 乖離すると次 window 丸ごと深掘りする (radon seed=1 実測で window + -- [150,250) が depth 9.9・101k leapfrog = warmup 全体の 76%)。 + recalEps <- if isWindowEnd + then do + w <- readMutVar welfordRef + -- Reset Welford for the next window (window-local variance). + writeMutVar welfordRef (emptyWelford nParams) + if wN w >= 5 -- need a few samples to be meaningful + then do + let mInv' = welfordMInvVS w + writeMutVar mInvRef mInv' + if doAdapt && nutsInitEpsSearch cfg + then + -- Phase 87.1: **最終 window 末 (term buffer 直前) は restart + -- しない** (M のみ更新・DA 継続 = PyMC の連続 DA と同じ挙動)。 + -- restart すると DA が m=1 の暴れ期からやり直しになり、 50 + -- draw の term buffer では ε̄ が鋸歯の暴れを拾って過小に着地 + -- する (radon 実測: ε 振動 0.035-1.37・ε̄=0.18・sampling α + -- 0.95/depth 5。 PyMC は restart なしで振動 0.14-0.50・ + -- ε 0.23-0.34・depth 4/α 0.80-0.88)。 中間 window 末の + -- recal+restart (Phase 86・爆発対策) は維持 — この時点まで + -- に M はほぼ収束しており継続 DA の ε がそのまま通用する。 + if iterIdx == windowPhaseEnd + then pure Nothing + else do + epsNew <- findReasonableEpsilon gradFn logPiFn mInv' eps nextU gen + writeMutVar daRef (initDualAvg epsNew) + pure (Just epsNew) + else do + -- 旧挙動 (opt-out 時): 現 ε anchor で restart。 + writeMutVar daRef (initDualAvg eps) + pure Nothing + else pure Nothing + else pure Nothing + eps' <- case recalEps of + -- Stan と同じく restart 直後はこの draw の accept 統計を学習しない + -- (旧 metric 下の α で較正済 anchor を汚さない)。 + Just epsNew -> pure epsNew + Nothing + | doAdapt && isBurnIn -> do + da <- readMutVar daRef + let da' = {-# SCC "nuts_dualavg" #-} updateDualAvg (nutsTargetAccept cfg) alpha da + writeMutVar daRef da' + return (exp (daLogEps da')) + | otherwise -> do + da <- readMutVar daRef + let epsBar = if doAdapt && not isBurnIn && i == nutsIterations cfg + then exp (daLogEpsBar da) + else eps + return epsBar + let nextParams = {-# SCC "nuts_toConstrained" #-} toConstrained nextU + if not isBurnIn + then do + modifyMutVar' samplesRef (nextParams :) + modifyMutVar' energyRef (h0 :) + modifyMutVar' depthRef (treeDepth :) + when divergent $ + modifyMutVar' divergenceRef + ((nutsIterations cfg - i) :) + else return () + -- Phase 9.1a: per-iteration callback for streaming UIs. + -- 0-based iter index running 0 .. total-1; isBurnIn for first nutsBurnIn. + onSample SampleEvent + { seIter = total - i + , seIsBurnIn = isBurnIn + , seParams = nextParams + , seEnergy = h0 + , seDivergent = divergent + , seAccepted = nextU /= currentU + , seStepSize = eps' + , seTreeDepth = treeDepth + , seAcceptStat = alpha + } + loop (i - 1) nextU eps' + + _ <- loop total initUV eps0 + samples <- fmap reverse (readMutVar samplesRef) + energies <- fmap reverse (readMutVar energyRef) + divs <- fmap reverse (readMutVar divergenceRef) + depths <- fmap reverse (readMutVar depthRef) + accepted <- readMutVar acceptedRef + return Chain + { chainSamples = samples + , chainAccepted = accepted + , chainTotal = total + , chainEnergy = energies + , chainDivergences = divs + , chainTreeDepths = depths + } + +-- --------------------------------------------------------------------------- +-- B11: Mass-matrix adaptation helpers +-- --------------------------------------------------------------------------- + +-- | Welford online accumulator for diagonal sample variance. +-- +-- Per-coordinate one-pass mean / M2; variance = M2 / (n − 1). +-- Used by Stan-style window adaptation to estimate posterior variance +-- without keeping the raw samples around. +-- | Plain (non-record) constructor: @Welford n mean m2@. Kept positional +-- because the @m2@ field is only ever pattern-matched, never read via a +-- selector (record syntax would generate an unused-binding warning). +-- | Storable-Vector Welford. The previous list-based form allocated +-- four @[Double]@ vectors per add (warmup ~500 iters × 4 cells = 10K +-- list cells per fit) and was hot during the mass-matrix adaptation +-- window phase. +data Welford = Welford !Int !(VS.Vector Double) !(VS.Vector Double) + +wN :: Welford -> Int +wN (Welford n _ _) = n + +emptyWelford :: Int -> Welford +emptyWelford p = Welford 0 (VS.replicate p 0) (VS.replicate p 0) + +welfordAddVS :: Welford -> VS.Vector Double -> Welford +welfordAddVS (Welford n mean m2) x = + let !n' = n + 1 + !nD = fromIntegral n' :: Double + !d = VS.zipWith (-) x mean + !mean' = VS.zipWith (\me di -> me + di / nD) mean d + !d2 = VS.zipWith (-) x mean' + !m2' = VS.zipWith3 (\m2i d1 d22 -> m2i + d1 * d22) m2 d d2 + in Welford n' mean' m2' + +-- | Stan-regularised diagonal @M⁻¹@ from a Welford accumulator. +-- +-- @σ̂² = (n / (n+5)) · sample_var + 1e-3 · (5 / (n+5))@. +-- The 1e-3 shrinkage target keeps the estimator non-degenerate when +-- @n@ is tiny; for moderate @n@ it reduces to the sample variance. +-- +-- /Convention/: following Stan/blackjax, @M⁻¹@ stores the posterior +-- covariance directly (so @M⁻¹_ii = σ̂²_i@). With kinetic energy +-- @½ rᵀ M⁻¹ r@ and @r ~ N(0, M)@, this gives a per-leapfrog position +-- step @ε · σ̂_i@ in absolute units (i.e. @ε@ in posterior-sd units), +-- which is what NUTS needs for tree depth ~ @1/ε@. +welfordMInvVS :: Welford -> VS.Vector Double +welfordMInvVS (Welford n _ m2) + | n < 2 = VS.replicate (VS.length m2) 1.0 + | otherwise = + let nD = fromIntegral n :: Double + k = 5.0 :: Double + weight = nD / (nD + k) + target = 1e-3 + in VS.map (\v -> let raw = v / (nD - 1) + in max 1e-12 (weight * raw + (1 - weight) * target)) + m2 + +-- | Stan-style adaptation schedule for a warmup of @W@ iterations. +-- +-- Returns @(windowEndIters, initBuffer, termBuffer)@ where +-- @windowEndIters@ are 1-based iteration indices at which to update the +-- mass matrix, and @initBuffer@ / @termBuffer@ are the no-update +-- prefix / suffix lengths (Stan defaults: 15% / 10%, with floors of 75 +-- and 50 iters respectively). Windows double in size starting from 25; +-- the last window absorbs any remainder. +-- +-- For @W = 500@: @initBuffer = 75@, @termBuffer = 50@, middle = 375, +-- windows = @[100, 150, 250, 450]@. +stanWindows :: Int -> ([Int], Int, Int) +stanWindows w + | w < 50 = ([], w, 0) + | otherwise = + let initB = max 75 (w `div` 7) + termB = max 50 (w `div` 10) + midLen = w - initB - termB + in if midLen < 25 + then ([], w, 0) + else (genW (initB + 1) midLen 25, initB, termB) + where + genW _ 0 _ = [] + genW start rest wsize + | wsize * 2 > rest = + -- Next doubled window wouldn't fit; absorb the remainder. + [start + rest - 1] + | otherwise = + let endIter = start + wsize - 1 + in endIter : genW (endIter + 1) (rest - wsize) (wsize * 2) + +-- | Run 'nuts' on @numChains@ parallel chains. +nutsChains :: ModelP r -> NUTSConfig -> Int -> Params -> GenIO -> IO [Chain] +nutsChains m cfg numChains initC baseGen = do + gens <- replicateM numChains (spawnGen baseGen) + mapConcurrently (\g -> nuts m cfg initC g) gens + +-- --------------------------------------------------------------------------- +-- Phase 50: 純粋 (ST + seed) ラッパ +-- +-- 'nuts' を @ST@ で走らせ 'runST' で閉じることで、 **seed → 確定 'Chain'** の +-- 純粋関数にする (同 seed → ビット同一・IO 不要)。 mwc は 'PrimMonad' 汎用ゆえ +-- ロジックは 50.2 で一般化した 'nuts' をそのまま使う。 +-- --------------------------------------------------------------------------- + +-- | [日本語]: 純粋・決定的な単一 NUTS chain。 同じ @seed@ なら必ず同じ 'Chain' を返す。 +-- [English]: A pure, deterministic single NUTS chain. The same @seed@ +-- always returns the same 'Chain'. +nutsPure :: ModelP r -> NUTSConfig -> Params -> Word32 -> Chain +nutsPure m cfg initC seed = + runST (initialize (V.singleton seed) >>= nuts m cfg initC) + +-- | [日本語]: 親 @seed@ から chain ごとの child seed 列を純粋に導出する +-- (@nutsChainsPure@ から抽出)。 pure 経路と IO 経路 ('nutsChainsStream') が +-- __同じ seed 列__を共有することで両経路のビット一致を保証する (複製すると drift)。 +-- [English]: Purely derives a per-chain child seed sequence from a parent +-- @seed@ (factored out of @nutsChainsPure@). The pure path and the IO +-- path ('nutsChainsStream') sharing the __same seed sequence__ guarantees +-- bit-identical results between the two paths (duplicating the logic +-- would let them drift apart). +chainSeeds :: Word32 -> Int -> [Word32] +chainSeeds seed numChains = runST $ do + g <- initialize (V.singleton seed) + replicateM numChains (uniform g) + +-- | [日本語]: 純粋・決定的な multi-chain。 親 @seed@ から子 seed を純粋に導出 +-- (各 chain は別 'runST') し、 chain 横断を @parList rdeepseq@ で__最初から__ +-- 並列評価する (純粋性と並列性は直交。 @+RTS -N@ でマルチコア。 結果は +-- spark/コア数に依らずビット同一)。 +-- [English]: A pure, deterministic multi-chain run. Child seeds are +-- derived purely from the parent @seed@ (each chain gets its own +-- 'runST'), and the chains are evaluated in parallel __from the start__ +-- via @parList rdeepseq@ (purity and parallelism are orthogonal here; +-- @+RTS -N@ enables multiple cores. The result is bit-identical +-- regardless of spark count / core count). +nutsChainsPure :: ModelP r -> NUTSConfig -> Int -> Params -> Word32 -> [Chain] +nutsChainsPure m cfg numChains initC seed = + let chains = [ nutsPure m cfg initC s | s <- chainSeeds seed numChains ] + in chains `using` parList rdeepseq + +-- | [日本語]: @nutsChainsPure@ の IO 版: 同じ child seed 規約 +-- (@chainSeeds@) で chain ごとに @nutsStream@ を回し、 chain index 付き +-- callback で進捗を観測できるようにする。 chain 横断は 'mapConcurrently' +-- (既存 'nutsChains' と同様・実 OS スレッド並列には @-threaded +RTS -N@)。 +-- +-- mwc の 'PrimMonad' 汎用性 + 実証済の ST/IO ビット同一により、 +-- no-op callback なら結果は @nutsChainsPure m cfg n initC seed@ と +-- __ビット一致__する (回帰テストで固定)。 +-- [English]: The IO counterpart of @nutsChainsPure@: runs @nutsStream@ +-- per chain under the same child-seed convention (@chainSeeds@), and lets +-- progress be observed via a chain-index-aware callback. Chains run in +-- parallel via 'mapConcurrently' (same as the existing 'nutsChains'; use +-- @-threaded +RTS -N@ for true OS-thread parallelism). +-- +-- Thanks to mwc's 'PrimMonad' polymorphism plus the proven bit-identity +-- of the ST and IO paths, a no-op callback makes the result +-- __bit-identical__ to @nutsChainsPure m cfg n initC seed@ (pinned down +-- by a regression test). +nutsChainsStream :: ModelP r -> NUTSConfig -> Int -> Params -> Word32 + -> (Int -> SampleEvent -> IO ()) + -> IO [Chain] +nutsChainsStream m cfg numChains initC seed onSample = + mapConcurrently + (\(i, s) -> do + g <- initialize (V.singleton s) + nutsStream m cfg initC g (onSample i)) + (zip [0 ..] (chainSeeds seed numChains)) diff --git a/src/Hanalyze/MCMC/Progress.hs b/src/Hanalyze/MCMC/Progress.hs new file mode 100644 --- /dev/null +++ b/src/Hanalyze/MCMC/Progress.hs @@ -0,0 +1,204 @@ +-- | +-- Module : Hanalyze.MCMC.Progress +-- Description : MCMC サンプリングの進捗表示 (全 chain 集計を stderr に描画) +-- Copyright : (c) 2026 Aelysce Project (Toshiaki Honda) +-- License : BSD-3-Clause +-- +-- [日本語]: MCMC サンプリングの進捗表示。 +-- +-- 'Hanalyze.MCMC.NUTS.nutsChainsStream' の chain index 付き callback に +-- 接続して、 全 chain 集計の進捗 1 行を stderr に描画する: +-- +-- > chains 2/4 done | draw 3400/8000 (warmup) | div 12 | 380.0 it/s +-- +-- 設計 (計画の柱): +-- +-- - 表示は「現在の chain」 でなく__全 chain 集計__ (chain は mapConcurrently +-- 並列で同時進行するため「現在」 が無い)。 +-- - callback はサンプラループ内で__同期実行__される (@nutsStream@ doc 明記) +-- ので、 描画はカウンタ先行の間引き (全体の ~0.5% 刻み) を通過した時だけ +-- 時刻取得 + 描画する。 ホットパスに乗るのはカウンタ更新のみ。 +-- - TTY (対話端末) では @\\r@ 上書きの 1 行、 非 TTY (CI ログ等) では +-- 10% 刻みの行出力。 +-- - 並列 chain からの stderr 競合は @MVar@ の単一描画権で回避 +-- (取れなければ描画 skip = 次の間引き通過で追いつく)。 +-- [English]: Progress display for MCMC sampling. +-- +-- This hooks into the chain-indexed callback of +-- 'Hanalyze.MCMC.NUTS.nutsChainsStream' and renders one progress +-- line, aggregated across all chains, to stderr: +-- +-- > chains 2/4 done | draw 3400/8000 (warmup) | div 12 | 380.0 it/s +-- +-- Design (pillars of the plan): +-- +-- - The display shows __aggregated progress across all chains__, not "the +-- current chain" (since chains run concurrently via mapConcurrently, +-- there is no single "current" chain). +-- - The callback runs __synchronously__ inside the sampler loop (as +-- documented on @nutsStream@), so rendering only fetches the time and +-- draws once the counter passes a throttling threshold (~0.5% of the +-- total). Only the counter update sits on the hot path. +-- - On a TTY (interactive terminal), a single line is overwritten with +-- @\\r@; on a non-TTY (e.g. CI logs), a line is emitted every 10%. +-- - Contention on stderr from parallel chains is avoided via a single +-- drawing right held in an @MVar@ (if it can't be taken, the render is +-- skipped and catches up at the next throttled pass). +{-# LANGUAGE StrictData #-} +{-# LANGUAGE OverloadedStrings #-} +module Hanalyze.MCMC.Progress + ( ProgressSnapshot (..) + , formatProgress + , newProgressRenderer + ) where + +import Control.Concurrent.MVar (newMVar, tryTakeMVar, putMVar) +import Control.Monad (when) +import Data.IORef (IORef, newIORef, readIORef, atomicModifyIORef') +import Data.Text (Text) +import qualified Data.IntMap.Strict as IM +import qualified Data.Text as T +import qualified Data.Text.IO as TIO +import GHC.Clock (getMonotonicTime) +import Numeric (showFFloat) +import System.IO (stderr, hIsTerminalDevice, hFlush) + +import Hanalyze.MCMC.NUTS (SampleEvent (..)) + +-- =========================================================================== +-- スナップショット + 純粋フォーマッタ +-- =========================================================================== + +-- | [日本語]: 全 chain 集計の進捗スナップショット (描画と独立な純粋データ)。 +-- [English]: A progress snapshot aggregated across all chains (pure data, +-- independent of rendering). +data ProgressSnapshot = ProgressSnapshot + { psChains :: Int -- ^ [日本語]: 総 chain 数。 [English]: Total number of chains. + , psChainsDone :: Int -- ^ [日本語]: 完了 chain 数。 [English]: Number of completed chains. + , psDraw :: Int -- ^ [日本語]: 全 chain 合算の消化 iteration 数 (burn-in 込み)。 [English]: Iterations consumed, summed across all chains (including burn-in). + , psTotal :: Int -- ^ [日本語]: 全 chain 合算の総 iteration 数。 [English]: Total iterations, summed across all chains. + , psWarmup :: Bool -- ^ [日本語]: いずれかの chain が warmup (burn-in) 中か。 [English]: Whether any chain is currently in warmup (burn-in). + , psDivergent :: Int -- ^ [日本語]: divergence 累計 (全 chain)。 [English]: Cumulative divergence count (all chains). + , psItersPerSec :: Double -- ^ [日本語]: 開始からの平均スループット (iteration/s)。 [English]: Average throughput since start (iterations/s). + } deriving (Show, Eq) + +-- | [日本語]: 進捗 1 行の純粋フォーマッタ。 例: +-- +-- @ +-- formatProgress (ProgressSnapshot 4 2 3400 8000 True 12 380.0) +-- == "chains 2\/4 done | draw 3400\/8000 (warmup) | div 12 | 380.0 it\/s" +-- @ +-- [English]: A pure formatter for one progress line. Example: +-- +-- @ +-- formatProgress (ProgressSnapshot 4 2 3400 8000 True 12 380.0) +-- == "chains 2\/4 done | draw 3400\/8000 (warmup) | div 12 | 380.0 it\/s" +-- @ +formatProgress :: ProgressSnapshot -> Text +formatProgress ps = T.intercalate " | " + [ "chains " <> tshow (psChainsDone ps) <> "/" <> tshow (psChains ps) <> " done" + , "draw " <> tshow (psDraw ps) <> "/" <> tshow (psTotal ps) + <> (if psWarmup ps then " (warmup)" else "") + , "div " <> tshow (psDivergent ps) + , T.pack (showFFloat (Just 1) (psItersPerSec ps) "") <> " it/s" + ] + where tshow = T.pack . show + +-- =========================================================================== +-- stderr レンダラ +-- =========================================================================== + +-- | [日本語]: レンダラ内部の可変状態 (chain ごとの消化数 / warmup フラグ / div 累計)。 +-- [English]: The renderer's internal mutable state (per-chain consumed +-- count / warmup flag / cumulative divergence count). +data RState = RState + { rsDraws :: !(IM.IntMap Int) -- ^ [日本語]: chain index → 消化 iteration 数。 [English]: chain index → number of consumed iterations. + , rsWarm :: !(IM.IntMap Bool) -- ^ [日本語]: chain index → 直近 event が burn-in か。 [English]: chain index → whether the most recent event was burn-in. + , rsDiv :: !Int -- ^ [日本語]: divergence 累計。 [English]: Cumulative divergence count. + } + +-- | [日本語]: stderr 進捗レンダラを作る。 返り値 = (chain index 付き callback, 終了処理)。 +-- +-- 終了処理は最終スナップショットを描画して行を閉じる (TTY では改行を補う)。 +-- 'Hanalyze.MCMC.NUTS.nutsChainsStream' に渡す想定: +-- +-- @ +-- (onSample, finish) <- newProgressRenderer chains (burnIn + iters) +-- chains <- nutsChainsStream m cfg chains initC seed onSample +-- finish +-- @ +-- [English]: Creates an stderr progress renderer. Returns +-- (a chain-index-aware callback, a finalizer). +-- +-- The finalizer renders the final snapshot and closes the line (adding a +-- newline on a TTY). Intended to be passed to +-- 'Hanalyze.MCMC.NUTS.nutsChainsStream': +-- +-- @ +-- (onSample, finish) <- newProgressRenderer chains (burnIn + iters) +-- chains <- nutsChainsStream m cfg chains initC seed onSample +-- finish +-- @ +newProgressRenderer :: Int -- ^ [日本語]: 総 chain 数 [English]: Total number of chains + -> Int -- ^ [日本語]: chain あたりの総 iteration 数 (burn-in 込み) [English]: Total iterations per chain (including burn-in) + -> IO (Int -> SampleEvent -> IO (), IO ()) +newProgressRenderer nChains perChain = do + isTTY <- hIsTerminalDevice stderr + t0 <- getMonotonicTime + stRef <- newIORef (RState IM.empty IM.empty 0) + lastPct <- newIORef (-1 :: Int) -- 非 TTY の 10% 刻み判定 + drawLock <- newMVar () -- 単一描画権 + let totalAll = nChains * perChain + stride = max 1 (totalAll `div` 200) -- ~0.5% 刻みで描画候補 + + snapshot :: RState -> Double -> ProgressSnapshot + snapshot st now = + let drawn = sum (IM.elems (rsDraws st)) + done = IM.size (IM.filter (>= perChain) (rsDraws st)) + warm = or (IM.elems (rsWarm st)) + dt = max 1e-9 (now - t0) + in ProgressSnapshot + { psChains = nChains, psChainsDone = done + , psDraw = drawn, psTotal = totalAll + , psWarmup = warm, psDivergent = rsDiv st + , psItersPerSec = fromIntegral drawn / dt + } + + -- 描画権が取れた時だけ描画 (競合時は skip・次の間引きで追いつく)。 + render :: Bool -> IO () + render final = do + got <- tryTakeMVar drawLock + case got of + Nothing -> pure () + Just () -> do + st <- readIORef stRef + now <- getMonotonicTime + let snap = snapshot st now + line = formatProgress snap + if isTTY + then do + TIO.hPutStr stderr ("\r" <> line) + when final (TIO.hPutStr stderr "\n") + hFlush stderr + else do + -- 非 TTY: 10% 境界を跨いだ時 (or 終了時) だけ 1 行出す。 + let pct10 = (10 * psDraw snap) `div` max 1 totalAll + prev <- readIORef lastPct + when (pct10 > prev || final) $ do + atomicModifyIORef' lastPct (\p -> (max p pct10, ())) + TIO.hPutStrLn stderr line + hFlush stderr + putMVar drawLock () + + onSample :: Int -> SampleEvent -> IO () + onSample i ev = do + n <- atomicModifyIORef' stRef $ \st -> + let st' = RState + { rsDraws = IM.insertWith (+) i 1 (rsDraws st) + , rsWarm = IM.insert i (seIsBurnIn ev) (rsWarm st) + , rsDiv = rsDiv st + (if seDivergent ev then 1 else 0) + } + in (st', sum (IM.elems (rsDraws st'))) + when (n `mod` stride == 0) (render False) + + pure (onSample, render True) diff --git a/src/Hanalyze/MCMC/SMC.hs b/src/Hanalyze/MCMC/SMC.hs new file mode 100644 --- /dev/null +++ b/src/Hanalyze/MCMC/SMC.hs @@ -0,0 +1,329 @@ +-- | +-- Module : Hanalyze.MCMC.SMC +-- Description : Tempered target による Sequential Monte Carlo (SMC) サンプラー +-- Copyright : (c) 2026 Aelysce Project (Toshiaki Honda) +-- License : BSD-3-Clause +-- +-- Sequential Monte Carlo (SMC) sampler with tempered targets. +-- +-- Implements a particle-based sampler that bridges from a broad initial +-- distribution to the full posterior @π(θ) ∝ p(θ) · L(θ)@ via a sequence +-- of intermediate targets @π_t(θ) ∝ p(θ) · L(θ)^β_t@, where +-- @β_0 = 0 → β_T = 1@. +-- +-- Reference: Del Moral, Doucet, Jasra (2006) "Sequential Monte Carlo +-- samplers". JRSSB 68:411-436. +-- +-- [日本語]: +-- +-- ## アルゴリズム概要 +-- +-- 1. __Init__: N 個の粒子を @init_@ を中心とする広い Gaussian cloud から +-- サンプル (= 近似 prior) +-- 2. __Tempering loop__ (t = 1..T): +-- a. __Weight__: 重み更新 @w_i ∝ exp((β_t − β_{t-1}) · logL(θ_i))@ +-- b. __log marginal contribution__: @log(mean w_i)@ を累積 +-- c. __Resample__: ESS = @(Σw)² / Σw²@ が閾値以下なら systematic resampling +-- d. __Move__: 各粒子に対し K 回の MH 移動 (target = π_t、 random walk +-- proposal) +-- 3. __Output__: 最終粒子集合を 'Chain' として返す + log marginal likelihood +-- の推定値 +-- +-- ## NUTS / MH との位置付け +-- +-- SMC の advantage: +-- +-- - 並列性が高い (粒子間は独立、 移動が並列化可能) +-- - 多峰分布で chain がはまりにくい (= temperature annealing) +-- - __log marginal likelihood の副産物推定__: Bridge Sampling より +-- 軽量で取れる (= Bayes Factor / BMA の前処理に使える) +-- +-- SMC の disadvantage: +-- +-- - 単峰分布なら NUTS の方が effective sample size / 時間 で有利 +-- - temperature schedule の選択が結果に影響 +-- +-- Bridge Sampling は本 SMC の log marginal 推定の __独立な検証手段__ として +-- 使う (両者で 5% 以内一致なら確からしい)。 +-- +-- [English]: +-- +-- ## Algorithm overview +-- +-- 1. __Init__: sample N particles from a broad Gaussian cloud centered on +-- @init_@ (= an approximate prior) +-- 2. __Tempering loop__ (t = 1..T): +-- a. __Weight__: update weights @w_i ∝ exp((β_t − β_{t-1}) · logL(θ_i))@ +-- b. __log marginal contribution__: accumulate @log(mean w_i)@ +-- c. __Resample__: if ESS = @(Σw)² / Σw²@ falls below the threshold, +-- systematic resampling +-- d. __Move__: K MH moves per particle (target = π_t, random walk +-- proposal) +-- 3. __Output__: returns the final particle set as a 'Chain' plus an +-- estimate of the log marginal likelihood +-- +-- ## Positioning relative to NUTS / MH +-- +-- SMC's advantages: +-- +-- - Highly parallel (particles are independent; moves parallelize) +-- - Chains are less likely to get stuck in a multimodal distribution +-- (= temperature annealing) +-- - __A by-product estimate of the log marginal likelihood__: obtained +-- more cheaply than Bridge Sampling (= usable as a preprocessing step +-- for a Bayes Factor / model averaging) +-- +-- SMC's disadvantages: +-- +-- - For unimodal distributions, NUTS is favorable in effective sample +-- size per unit time +-- - The choice of temperature schedule affects the result +-- +-- Bridge Sampling is used as an __independent verification method__ for +-- this SMC's log marginal estimate (agreement within 5% between the two +-- is taken as reassurance). +{-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE RankNTypes #-} +{-# LANGUAGE BangPatterns #-} +module Hanalyze.MCMC.SMC + ( SMCConfig (..) + , defaultSMCConfig + , SMCResult (..) + , smc + , smcPure + ) where + +import Control.Monad (forM, replicateM, foldM) +import Control.Monad.Primitive (PrimMonad, PrimState) +import Control.Monad.ST (runST) +import qualified Data.Map.Strict as Map +import Data.List (sort) +import Data.Text (Text) +import Data.Word (Word32) +import qualified Data.Vector as V +import qualified Data.Vector.Unboxed as VU +import System.Random.MWC (Gen, uniform, initialize) +import System.Random.MWC.Distributions (normal) + +import Hanalyze.Model.HBM (ModelP, Params, logPrior, logLikelihood, sampleNames) +import Hanalyze.MCMC.Core (Chain (..)) + +-- --------------------------------------------------------------------------- +-- Configuration +-- --------------------------------------------------------------------------- + +-- | SMC configuration. +data SMCConfig = SMCConfig + { smcNParticles :: !Int -- ^ [日本語]: N: 粒子数 (典型 500-2000) [English]: N: the number of particles (typically 500-2000) + , smcNSteps :: !Int -- ^ [日本語]: T: temperature step 数 (典型 10-50) [English]: T: the number of temperature steps (typically 10-50) + , smcMHIterations :: !Int -- ^ [日本語]: K: 各 temperature 内の MH 移動 回数 (典型 5-20) [English]: K: the number of MH moves within each temperature (typically 5-20) + , smcMHStepSize :: !(Map.Map Text Double) -- ^ [日本語]: Random walk MH の per-param std [English]: The per-parameter std for the random-walk MH proposal + , smcInitJitter :: !Double -- ^ [日本語]: 初期粒子を init_ から散らす Gaussian σ (typical 2-5) [English]: The Gaussian σ used to scatter initial particles around init_ (typically 2-5) + , smcESSThreshold :: !Double -- ^ [日本語]: 0..1、 ESS < N · threshold で resample (typical 0.5) [English]: In 0..1; resamples when ESS < N · threshold (typically 0.5) + } deriving (Show) + +-- | [日本語]: Default: N=500、 T=20、 K=10、 step=0.5、 jitter σ=3、 ESS threshold=0.5。 +-- [English]: Default: N=500, T=20, K=10, step=0.5, jitter σ=3, ESS threshold=0.5. +defaultSMCConfig :: [Text] -> SMCConfig +defaultSMCConfig names = SMCConfig + { smcNParticles = 500 + , smcNSteps = 20 + , smcMHIterations = 10 + , smcMHStepSize = Map.fromList [(n, 0.5) | n <- names] + , smcInitJitter = 3.0 + , smcESSThreshold = 0.5 + } + +-- | [日本語]: SMC の結果。 粒子を Chain 形に詰めた posterior 推定 + log marginal + +-- temperature step ごとの ESS 履歴。 +-- +-- __重要__: 'smcLogMarginal' は「初期粒子が prior からサンプルされている」 +-- __ことを仮定した推定値__。 本実装は init_ を中心とする jittered Gaussian +-- から初期粒子を作るため、 prior が広いと bias する。 厳密な log marginal が +-- 必要な場合は 'Hanalyze.Stat.BridgeSampling.bridgeSampling' を +-- 使用すること (SMC chain を入力に独立に推定する)。 SMC の primary 用途は +-- __多峰 posterior の効率的なサンプリング__。 +-- [English]: The result of SMC. A posterior estimate packed into a Chain +-- shape, plus the log marginal, plus a per-temperature-step ESS history. +-- +-- __Important__: 'smcLogMarginal' is an estimate that assumes the +-- __initial particles were sampled from the prior__. This implementation +-- creates initial particles from a jittered Gaussian centered on init_, +-- so it is biased when the prior is broad. If a rigorous log marginal is +-- needed, use +-- 'Hanalyze.Stat.BridgeSampling.bridgeSampling' instead (it +-- estimates independently, taking the SMC chain as input). SMC's primary +-- use case is __efficient sampling of multimodal posteriors__. +data SMCResult = SMCResult + { smcChain :: !Chain + , smcLogMarginal :: !Double + , smcESSHistory :: ![Double] + } deriving (Show) + +-- --------------------------------------------------------------------------- +-- 公開 API +-- --------------------------------------------------------------------------- + +-- | [日本語]: SMC を実行。 @init_@ を中心に initial particles を散らし、 linear +-- temperature schedule (β_t = t/T) で posterior に温めていく。 +-- [English]: Runs SMC. Scatters initial particles around @init_@ and +-- tempers toward the posterior along a linear temperature schedule +-- (β_t = t/T). +smc :: forall r m. PrimMonad m => ModelP r -> SMCConfig -> Params -> Gen (PrimState m) -> m SMCResult +smc model cfg init_ gen = do + let n = smcNParticles cfg + tT = smcNSteps cfg + names = sampleNames model + steps = smcMHStepSize cfg + jitter = smcInitJitter cfg + -- 1. Init particles: init_ + N(0, jitter · stepSizes_i) + particles0 <- replicateM n (jitterInit names jitter steps init_ gen) + let betas = [ fromIntegral t / fromIntegral tT | t <- [0 .. tT] ] -- [0, 1/T, .., 1] + betaSteps = zip betas (tail betas) -- [(β_{t-1}, β_t)] + + -- 2. Tempering loop + (finalParticles, logMarg, essHist) <- + foldM (stepTemper model steps (smcMHIterations cfg) (smcESSThreshold cfg) gen n) + (particles0, 0.0 :: Double, []) + betaSteps + + let accepted = chainAcceptedAcc (length finalParticles * tT * smcMHIterations cfg) + total = length finalParticles * tT * smcMHIterations cfg + pure SMCResult + { smcChain = Chain + { chainSamples = finalParticles + , chainAccepted = accepted + , chainTotal = total + , chainEnergy = [] + , chainDivergences = [] + , chainTreeDepths = [] + } + , smcLogMarginal = logMarg + , smcESSHistory = reverse essHist + } + where + -- 受理数は本実装では追跡しない (= 0 を入れて acceptanceRate は意味なし) + chainAcceptedAcc _ = 0 + +-- | [日本語]: 純粋・決定的な SMC (seed → 確定 SMCResult・IO 不要)。 'smc' の ST/seed 版。 +-- [English]: A pure, deterministic SMC (seed → a fixed SMCResult, no IO +-- needed). The ST/seed counterpart of 'smc'. +smcPure :: ModelP r -> SMCConfig -> Params -> Word32 -> SMCResult +smcPure model cfg initP seed = + runST (initialize (V.singleton seed) >>= smc model cfg initP) + +-- | [日本語]: 1 ステップの tempering: +-- - 重み計算 + log marginal 累積 +-- - ESS 判定して resample +-- - K 回の MH 移動 (target = π_t = p(θ) · L(θ)^β_t) +-- [English]: One tempering step: +-- - Compute weights + accumulate the log marginal +-- - Check ESS and resample +-- - K MH moves (target = π_t = p(θ) · L(θ)^β_t) +stepTemper + :: forall r m. PrimMonad m => ModelP r + -> Map.Map Text Double -- ^ step sizes + -> Int -- ^ K + -> Double -- ^ ESS threshold + -> Gen (PrimState m) + -> Int -- ^ [日本語]: N (元の粒子数、 resample で N keep) [English]: N (the original particle count; resampling keeps N) + -> ([Params], Double, [Double]) -- ^ [日本語]: (粒子、 累積 log marginal、 ESS 履歴) [English]: (particles, cumulative log marginal, ESS history) + -> (Double, Double) -- ^ (β_{t-1}, β_t) + -> m ([Params], Double, [Double]) +stepTemper model steps k essThr gen n (particles, logMarg, essHist) (b0, b1) = do + let dbeta = b1 - b0 + logLs = map (logLikelihood model) particles + logWs = map (dbeta *) logLs -- log incremental weights + logSumW = logSumExp logWs + logMean = logSumW - log (fromIntegral (length particles)) + ws = map (\lw -> exp (lw - logSumW)) logWs -- normalized weights + ess = if sum (map (** 2) ws) == 0 then 0 + else 1 / sum (map (** 2) ws) + logMarg' = logMarg + logMean + + -- Resample if ESS < threshold · N + resampled <- + if ess < essThr * fromIntegral n + then systematicResample particles ws n gen + else pure particles + + -- Move with K MH iterations + moved <- moveK model steps b1 k resampled gen + pure (moved, logMarg', ess : essHist) + +-- | [日本語]: systematic resampling (= particle filter standard)。 +-- [English]: Systematic resampling (the standard particle-filter method). +systematicResample + :: forall m. PrimMonad m => [Params] -> [Double] -> Int -> Gen (PrimState m) -> m [Params] +systematicResample particles ws n gen = do + u0 <- uniform gen :: m Double + let total = sum ws + ws' = map (/ total) ws -- normalize + cdf = scanl1 (+) ws' + ps = [ (fromIntegral i + u0) / fromIntegral n | i <- [0 .. n - 1] ] + pick p = pickAt p cdf particles + pure (map pick ps) + where + pickAt p (c : cs) (x : xs) + | p <= c = x + | otherwise = pickAt p cs xs + pickAt _ _ (x : _) = x -- fallback (numeric edge) + pickAt _ _ [] = error "systematicResample: empty particle list" + +-- | [日本語]: K 回の Random Walk MH 移動。 target は @log π_t = logPrior + β · logLik@。 +-- [English]: K random-walk MH moves. The target is +-- @log π_t = logPrior + β · logLik@. +moveK + :: forall r m. PrimMonad m => ModelP r + -> Map.Map Text Double + -> Double -- ^ β + -> Int + -> [Params] + -> Gen (PrimState m) + -> m [Params] +moveK model steps beta k particles gen = + mapM (mhKSteps model steps beta k gen) particles + +mhKSteps + :: forall r m. PrimMonad m => ModelP r + -> Map.Map Text Double + -> Double + -> Int + -> Gen (PrimState m) + -> Params + -> m Params +mhKSteps model steps beta k gen p0 = go k p0 + where + target p = logPrior model p + beta * logLikelihood model p + go 0 p = pure p + go i p = do + let names = Map.keys p + proposed <- fmap Map.fromList $ forM names $ \n -> do + let s = Map.findWithDefault 1.0 n steps + cur = Map.findWithDefault 0.0 n p + eps <- normal 0 s gen + pure (n, cur + eps) + let logA = target proposed - target p + u <- uniform gen :: m Double + let !next = if log u < logA then proposed else p + go (i - 1) next + +-- --------------------------------------------------------------------------- +-- Helpers +-- --------------------------------------------------------------------------- + +jitterInit + :: forall m. PrimMonad m => [Text] -> Double -> Map.Map Text Double -> Params -> Gen (PrimState m) -> m Params +jitterInit names jitter steps init_ gen = + fmap Map.fromList $ forM names $ \n -> do + let s = jitter * Map.findWithDefault 1.0 n steps + cur = Map.findWithDefault 0.0 n init_ + eps <- normal 0 s gen + pure (n, cur + eps) + +-- | Numerically stable log-sum-exp. +logSumExp :: [Double] -> Double +logSumExp [] = -1 / 0 +logSumExp xs = + let m = maximum xs + in m + log (sum [ exp (x - m) | x <- xs ]) diff --git a/src/Hanalyze/MCMC/Slice.hs b/src/Hanalyze/MCMC/Slice.hs new file mode 100644 --- /dev/null +++ b/src/Hanalyze/MCMC/Slice.hs @@ -0,0 +1,170 @@ +-- | +-- Module : Hanalyze.MCMC.Slice +-- Description : Slice sampler (Neal 2003) — 受理率調整不要な単変量サンプリング法 +-- Copyright : (c) 2026 Aelysce Project (Toshiaki Honda) +-- License : BSD-3-Clause +-- +-- Slice sampler (Neal 2003) — a univariate method with no acceptance-rate +-- tuning. +-- +-- Each iteration: +-- +-- 1. Draw @log_y = log p(θ) − Exp(1)@ from the current log-density. +-- 2. Build a horizontal slice @[L, R]@ along each axis via stepping-out. +-- 3. Shrink: draw @θ_i'@ uniformly from @[L, R]@ and accept when +-- @log p > log_y@. +-- +-- One iteration is a Gibbs-style sweep over every coordinate. Like +-- HMC/NUTS no gradient is required, but each sweep involves many +-- log-density evaluations. +{-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE RankNTypes #-} +module Hanalyze.MCMC.Slice + ( SliceConfig (..) + , defaultSliceConfig + , slice + , sliceChains + , slicePure + , sliceChainsPure + ) where + +import Control.Concurrent.Async (mapConcurrently) +import Control.Monad (forM, replicateM, when) +import Control.Monad.Primitive (PrimMonad, PrimState) +import Control.Monad.ST (runST) +import Control.Parallel.Strategies (parList, rdeepseq, using) +import Data.Primitive.MutVar +import Data.Word (Word32) +import qualified Data.Map.Strict as Map +import qualified Data.Vector as V +import Data.Map.Strict (Map) +import Data.Text (Text) +import System.Random.MWC (Gen, GenIO, uniform, initialize) +import System.Random.MWC.Distributions (exponential) + +import Hanalyze.Model.HBM (ModelP, Params, logJoint, sampleNames) +import Hanalyze.MCMC.Core (Chain (..), spawnGen) + +-- | Slice-sampler configuration. +data SliceConfig = SliceConfig + { sliceIterations :: Int -- ^ Total iterations (burn-in included). + , sliceBurnIn :: Int -- ^ Burn-in iterations to discard. + , sliceWidths :: Map Text Double -- ^ Initial stepping-out width @w@ + -- per coordinate (default 1.0). + , sliceMaxSteps :: Int -- ^ Maximum number of stepping-out + -- steps (safety bound). + } deriving (Show) + +-- | Default configuration: 1000 iterations, 200 burn-in, width 1.0 per +-- parameter, max stepping-out 50. +defaultSliceConfig :: [Text] -> SliceConfig +defaultSliceConfig names = SliceConfig + { sliceIterations = 1000 + , sliceBurnIn = 200 + , sliceWidths = Map.fromList [(n, 1.0) | n <- names] + , sliceMaxSteps = 50 + } + +-- | Run the slice sampler. One iteration updates every coordinate in +-- turn (Gibbs-style sweep). +slice :: forall r m. PrimMonad m => ModelP r -> SliceConfig -> Params -> Gen (PrimState m) -> m Chain +slice model cfg init_ gen = do + let names = sampleNames model + total = sliceBurnIn cfg + sliceIterations cfg + widths = sliceWidths cfg + maxStep = sliceMaxSteps cfg + + logP :: Params -> Double + logP = logJoint model + + samplesRef <- newMutVar [] + acceptedRef <- newMutVar (0 :: Int) + + -- 1 coordinate 更新 (slice sampling on one axis) + let updateOne :: Text -> Params -> m Params + updateOne nm cur = do + let w = Map.findWithDefault 1.0 nm widths + x0 = Map.findWithDefault 0.0 nm cur + pAt v = logP (Map.insert nm v cur) + -- 水平スライス: log_y = log p(θ) - Exp(1) + e <- exponential 1.0 gen + let logY = pAt x0 - e + -- Stepping out + u <- uniform gen + let l0 = x0 - w * (u :: Double) + r0 = l0 + w + u2 <- uniform gen + let kL = floor (fromIntegral maxStep * (u2 :: Double)) :: Int + kR = maxStep - 1 - kL + expandLeft k l + | k <= 0 || pAt l <= logY = return l + | otherwise = expandLeft (k - 1) (l - w) + expandRight k r + | k <= 0 || pAt r <= logY = return r + | otherwise = expandRight (k - 1) (r + w) + l1 <- expandLeft kL l0 + r1 <- expandRight kR r0 + -- Shrinkage + let shrink l r = do + uS <- uniform gen + let xNew = l + (uS :: Double) * (r - l) + if pAt xNew > logY + then return xNew + else + if xNew < x0 + then shrink xNew r + else shrink l xNew + xNew <- shrink l1 r1 + modifyMutVar' acceptedRef (+1) + return (Map.insert nm xNew cur) + + let sweep current = foldr (\_ _ -> id) id [] `seq` + sweepGo names current + where + sweepGo [] c = return c + sweepGo (n:ns) c = do c' <- updateOne n c + sweepGo ns c' + + let loop 0 current = return current + loop i current = do + next <- sweep current + when (i <= sliceIterations cfg) $ + modifyMutVar' samplesRef (next :) + loop (i - 1) next + + _ <- loop total init_ + samples <- fmap reverse (readMutVar samplesRef) + accepted <- readMutVar acceptedRef + return Chain + { chainSamples = samples + , chainAccepted = accepted + , chainTotal = total * length names + , chainEnergy = [] + , chainDivergences = [] + , chainTreeDepths = [] + } + +-- | Run 'slice' on @numChains@ parallel chains. +sliceChains :: ModelP r -> SliceConfig -> Int -> Params -> GenIO -> IO [Chain] +sliceChains model cfg numChains initP baseGen = do + gens <- replicateM numChains (spawnGen baseGen) + mapConcurrently (\g -> slice model cfg initP g) gens + +-- | [日本語]: 純粋・決定的な slice sampler (seed → 確定 Chain)。 +-- [English]: A pure, deterministic slice sampler (seed → a fixed Chain). +slicePure :: ModelP r -> SliceConfig -> Params -> Word32 -> Chain +slicePure model cfg initP seed = + runST (initialize (V.singleton seed) >>= slice model cfg initP) + +-- | [日本語]: 純粋・決定的な multi-chain slice。 子 seed を純粋導出し @parList rdeepseq@ で並列。 +-- [English]: A pure, deterministic multi-chain slice sampler. Child +-- seeds are derived purely, and chains run in parallel via +-- @parList rdeepseq@. +sliceChainsPure :: ModelP r -> SliceConfig -> Int -> Params -> Word32 -> [Chain] +sliceChainsPure model cfg numChains initP seed = + let childSeeds :: [Word32] + childSeeds = runST $ do + g <- initialize (V.singleton seed) + replicateM numChains (uniform g) + chains = [ slicePure model cfg initP s | s <- childSeeds ] + in chains `using` parList rdeepseq diff --git a/src/Hanalyze/Model/HBM.hs b/src/Hanalyze/Model/HBM.hs new file mode 100644 --- /dev/null +++ b/src/Hanalyze/Model/HBM.hs @@ -0,0 +1,229 @@ +-- | +-- Module : Hanalyze.Model.HBM +-- Description : 多相階層ベイズモデル (Hierarchical Bayesian Model, HBM) DSL の facade +-- Copyright : (c) 2026 Aelysce Project (Toshiaki Honda) +-- License : BSD-3-Clause +-- +-- Polymorphic Hierarchical Bayesian Model (HBM) DSL. +-- +-- [日本語]: 責務別 submodule に分割済み。 本モジュールは __facade__: +-- 下位 8 module (Util/Distribution/Sampling/Model/Track/Eval/IR/Gradient) を +-- import し、 従来の公開 API を export list 経由でそのまま再公開する。 +-- 既存 importer (18 src module + test) は無改修で従来通り使える。 +-- [English]: Split into responsibility-scoped submodules; this module is +-- the __facade__: it imports the 8 lower-layer modules +-- (Util/Distribution/Sampling/Model/Track/Eval/IR/Gradient) and re-exports +-- the original public API unchanged via the export list. Existing +-- importers (18 src modules + tests) work without any modification. +-- +-- A free-monad embedded language for probabilistic programs. The +-- continuation type is left polymorphic so that a single model term can +-- be reinterpreted as: +-- +-- - a structural inspector (parameter / observation graph), +-- - a log-joint density, +-- - an automatically-differentiated log-joint +-- (via @Numeric.AD.Mode.Reverse.Double@ — +-- [日本語]: Double 特化の reverse モードゆえ勾配は latent 数 p に依らず +-- ~1 sweep。 forward 実装から切替: forward は勾配 1 本に p 回評価が要り +-- 階層モデルで線形悪化していた。 generic Reverse は tape boxing で +-- 低次元が遅く、 Reverse.Double が全 p で forward/generic を上回ると実測。 +-- [English]: specializing to @Double@ makes the reverse-mode gradient +-- cost ~1 sweep independent of the number of latent parameters p. The +-- previous forward-mode implementation needed p evaluations per +-- gradient and degraded linearly for hierarchical models; the generic +-- @Reverse@ mode is slow at low dimension due to tape boxing, and +-- measurements show @Reverse.Double@ beats both forward-mode and +-- generic @Reverse@ across all p.), +-- - a dependency tracker (the 'Track' interpretation, used by +-- @Hanalyze.Viz.ModelGraph@ to build a Mermaid DAG). +-- +-- See @docs/bayesian/02-probabilistic-model.md@ for an extended +-- introduction. +-- +-- @ +-- data ModelF a next +-- = Sample Text (Distribution a) (a -> next) +-- | Observe Text (Distribution a) [Double] next +-- deriving Functor +-- @ +-- +-- [日本語]: ユーザーは @forall a. (Floating a, Ord a) => Model a r@ という +-- 「型に多相なモデル」を一度だけ書き、 解釈時に @a@ を選ぶことで同じモデル +-- から複数の解釈 (サンプリング・log joint・AD 勾配・依存抽出) を取り出せる。 +-- [English]: Users write a "type-polymorphic model" — +-- @forall a. (Floating a, Ord a) => Model a r@ — just once, and choosing +-- @a@ at interpretation time lets the same model yield multiple readings +-- (sampling, log-joint, AD gradient, dependency extraction). +-- +-- == 使い方 / Usage +-- +-- @ +-- import Hanalyze.Model.HBM +-- +-- myModel :: ModelP () +-- myModel = do +-- mu <- sample "mu" (Normal 0 10) +-- sigma <- sample "sigma" (Exponential 1) +-- observe "y" (Normal mu sigma) [1.5, 2.0, 1.8] +-- +-- -- 異なる解釈: +-- logVal = logJoint myModel (Map.fromList [("mu",1),("sigma",2)]) -- 数値評価 +-- gVec = gradAD myModel ["mu","sigma"] [1, 2] -- AD 勾配 +-- deps = extractDeps myModel -- 依存関係 +-- @ +module Hanalyze.Model.HBM + ( -- * Polymorphic distributions + Distribution (..) + , distName + , logDensity + , logDensityObs + , sampleDist + , sampleMvDist + , distCDF + , logCDF + , logSF + -- * Polymorphic model DSL + , Free (..) + , liftF + , ModelF (..) + , Model + , ModelP + , sample + , observe + , observeMV + , observeColumns + , observeLM + , observeLMR + , observeNormalLM + , LMFamily (..) + , REff (..) + , REffect (..) + , reffNames + , reNormal + , at + , indexed + , (.#) + , potential + , deterministic + , runDeterministics + , deterministicNames + , augmentChainWithDeterministic + , nonCenteredNormal + , dirichlet + , orderedCuts + , dpStickBreaking + , hmmLatent + -- ** plate notation + , plate + , plateI + , plateI_ + , plateForM + , plateForM_ + , withPlate + , hmmForwardLogLik + , GlmmFamily (..) + , glmmRandomIntercept + , dataNamed + , dataNamedX + , dataNamedIx + , dataNamedObs + , Ix (..) + , TrackTag (..) + , (!!!) + , atIx + , withData + , withDataIx + , mvNormalLatent + , mvNormalLogDensity + , mvNormalCholLogDensity + , multinomialLogDensity + , mvStudentTLogDensity + , dirichletMultinomialLogDensity + , wishartLogDensity + , obsLogSum + , lkjCorrCholesky + , gpExpQuadCov + , gpLatent + , ar1Latent + -- * Structural inspection + , Node (..) + , NodeKind (..) + , collectNodes + , sampleNames + , dataSlots + , dataIxSlots + , extractDeps + -- * Type aliases + , Params + -- * Interpreters + , logJoint + , logPrior + , logLikelihood + , perObsLogLiks + , runObserveDists + , priorList + , describeModel + -- * Model graph (visualization) + , ModelGraph (..) + , buildModelGraph + , collapseIndexedPlateNodes + -- * AD gradient + , gradAD + , gradADU + , compileGradU + , compileGradUV + , compileGradValUV + , compileGradValUVM + , compileLogPU + , compileLogPUV + , synthGaussLMBlocks + , synthVecIR + , gradPathLabel + -- * Numeric utilities (test 用) + , lgammaApprox + , digamma + -- * Constraint transforms (for HMC) + , getTransforms + , logJointUnconstrained + , invTransformF + , logJacF + -- * Dependency-tracking interpretation + , Track (..) + , trackVar + , trackConst + ) where + +-- Phase 58.2: 純粋な数値・線形代数 leaf util を分離。 internal 利用に加え +-- 'lgammaApprox' / 'digamma' は export list 経由でそのまま再エクスポートされる。 +import Hanalyze.Model.HBM.Util +-- Phase 58.3/58.6a: 多相分布 ADT + 密度 + CDF を分離 (Util の上層)。 公開 API +-- (Distribution(..)/distName/logDensity/logDensityObs/obsLogSum/distCDF/logCDF/ +-- logSF/MV密度群) は export list 経由でそのまま再エクスポート。 ★58.6a で事前 +-- logDensity と観測 logDensityObs/obsLogSum を本体から Distribution へ集約 +-- (Eval の logJoint/logPrior が logDensity を参照する back-edge を解消・密度は +-- 本来 Distribution の責務。 INLINABLE は AD cross-module inlining 維持で保持)。 +import Hanalyze.Model.HBM.Distribution +-- Phase 58.4: 分布からのサンプリング (sampleDist/sampleMvDist) を分離。 +-- export list 経由でそのまま再エクスポート。 PrimMonad/mwc-random 依存・非ホット。 +import Hanalyze.Model.HBM.Sampling +-- Phase 58.5: 多相モデル DSL (Free monad + ModelF + plate + 構造検査) を分離。 +-- 公開 API (Free/liftF/ModelF/Model/ModelP/sample/observe/plate/collectNodes 等) +-- は export list 経由でそのまま再エクスポート。 +import Hanalyze.Model.HBM.Model +-- Phase 58.6b: 依存追跡型 Track (Track/trackVar/trackConst/extractDeps) を分離。 +-- Model/Distribution の上層・非ホット (DAG 抽出のみ・NUTS per-draw 非経路)。 +-- export list 経由でそのまま再エクスポート。 +import Hanalyze.Model.HBM.Track +-- Phase 58.6c: 評価層 (ObserveLM 評価 + logJoint/logPrior/logLikelihood interp + +-- 互換 API runDeterministics/buildModelGraph 等 + runTrack) を分離。 Track の上層。 +-- ★ホット (logJoint は AD 勾配経路)。 AD 勾配・IR (本体残置) は本モジュールを +-- forward import する。 公開 API は export list 経由でそのまま再エクスポート。 +import Hanalyze.Model.HBM.Eval +-- Phase 58.7: IR (中間表現) 層 (affine/非線形/密度 IR) を分離。 最ホット (gradVecIR)。 +import Hanalyze.Model.HBM.IR +-- Phase 58.8: AD 勾配コンパイラ層 (compileGradUV/hybridGradClosure/gaussLMBlocks/ +-- 定数 prior 解析勾配/制約変換) を分離。 IR の上層・最ホット (NUTS per-draw 本経路)。 +-- 公開 API (gradAD/gradADU/compileGradU/compileGradUV/compileLogPU/compileLogPUV/ +-- getTransforms/logJointUnconstrained/invTransformF/logJacF) は export list 経由で再公開。 +import Hanalyze.Model.HBM.Gradient diff --git a/src/Hanalyze/Model/HBM/Ast.hs b/src/Hanalyze/Model/HBM/Ast.hs new file mode 100644 --- /dev/null +++ b/src/Hanalyze/Model/HBM/Ast.hs @@ -0,0 +1,254 @@ +{-# LANGUAGE OverloadedStrings #-} + +-- | +-- Module : Hanalyze.Model.HBM.Ast +-- Description : HBM dialog DSL の AST 型と JSON decoder +-- Copyright : (c) 2026 Aelysce Project (Toshiaki Honda) +-- License : BSD-3-Clause +-- +-- HBM dialog DSL の AST 型と JSON decoder。 +-- +-- [日本語]: canvas-backend @CanvasApp.Analysis.HBM@ から移設。 frontend が +-- backend 統一 parser (@/api/v1/dsl/parse@) から得た @program_ast@ (JSON) を、 +-- streaming sidecar が直接 decode して実モデルを構築できるよう、 AST 型 + +-- 'parseAst' をライブラリ層 (hanalyze) に置く。 +-- [English]: Relocated from the canvas-backend +-- @CanvasApp.Analysis.HBM@. So the streaming sidecar can decode the +-- @program_ast@ (JSON) the frontend obtains from the backend's unified +-- parser (@/api/v1/dsl/parse@) directly into a real model, the AST type +-- and 'parseAst' live in this library layer (hanalyze). +-- +-- [日本語]: 本 module は __canvas wire 型・text parser (DSL frontend) いずれにも非依存__ +-- (= aeson のみ)。 text → AST 変換 (@parseHbmTextToExpr@ 等) は HT (DSL frontend) +-- に依存するため canvas-backend 側に残す。 +-- [English]: This module depends on __neither the canvas wire types nor the text parser (DSL frontend)__ +-- (aeson only). The text → AST conversion (e.g. @parseHbmTextToExpr@) +-- depends on HT (DSL frontend), so it stays on the canvas-backend side. +module Hanalyze.Model.HBM.Ast + ( -- * AST + Expr (..) + , Lit (..) + , Bind (..) + , DoStmt (..) + -- * JSON decode (= frontend program_ast → Expr) + , parseAst + , parseLit + , parseBind + , parseDoStmt + -- * JSON encode (= 'parseAst' の正確な逆。 backend が sidecar に + -- program_ast / top_binds を送る際に使う) + , exprToJSON + , litToJSON + , bindToJSON + , doStmtToJSON + -- * helpers + , Err + , collectApp + , getField + , getStr + , getNum + , getBool + , getArray + ) where + +import Data.Text (Text) +import qualified Data.Aeson as A +import Data.Aeson.Types (Pair) +import qualified Data.Aeson.Key as Key +import qualified Data.Aeson.KeyMap as KM +import qualified Data.Vector as V + +-- --------------------------------------------------------------------------- +-- AST (= frontend App.Hbm.Ast / DSL frontend hanalyze.HBM.Text.HbmExpr と同形、 +-- 11 ctor: ELit / ECol / EVar / EApp / ELam / EIf / ELet / ENeg / EOp / +-- EList / EDo) +-- --------------------------------------------------------------------------- + +data Expr + = ELit Lit + | ECol Text + | EVar Text + | EApp Expr Expr + | ELam Text Expr + | EIf Expr Expr Expr + | ELet [Bind] Expr + | ENeg Expr + | EOp Text Expr Expr + | EList [Expr] + | EDo [DoStmt] Expr + deriving (Show) + +data Lit = LNumber Double | LText Text | LBool Bool deriving (Show) + +data Bind = Bind { bindName :: Text, bindValue :: Expr } deriving (Show) + +data DoStmt + = DoBind Text Expr + | DoLet [Bind] + | DoExpr Expr + deriving (Show) + +-- | [日本語]: 評価系で多用する Either alias。 [English]: An @Either@ alias used heavily throughout the evaluator. +type Err a = Either Text a + +-- --------------------------------------------------------------------------- +-- JSON parser (= frontend が送る program_ast を Expr に decode) +-- --------------------------------------------------------------------------- + +parseAst :: A.Value -> Either Text Expr +parseAst v = case v of + A.Object o -> do + tag <- getStr o "tag" + case tag of + "ELit" -> ELit <$> (parseLit =<< getField o "lit") + "ECol" -> ECol <$> getStr o "name" + "EVar" -> EVar <$> getStr o "name" + "EApp" -> EApp <$> (parseAst =<< getField o "f") <*> (parseAst =<< getField o "x") + "ELam" -> ELam <$> getStr o "arg" <*> (parseAst =<< getField o "body") + "EIf" -> EIf <$> (parseAst =<< getField o "c") + <*> (parseAst =<< getField o "a") + <*> (parseAst =<< getField o "b") + "ELet" -> do + bs <- getArray o "binds" >>= mapM parseBind + body <- parseAst =<< getField o "body" + Right (ELet bs body) + "ENeg" -> ENeg <$> (parseAst =<< getField o "e") + "EOp" -> EOp <$> getStr o "op" + <*> (parseAst =<< getField o "a") + <*> (parseAst =<< getField o "b") + "EList" -> EList <$> (getArray o "items" >>= mapM parseAst) + "EDo" -> do + stmts <- getArray o "stmts" >>= mapM parseDoStmt + ret <- parseAst =<< getField o "ret" + Right (EDo stmts ret) + _ -> Left ("Unknown AST tag: " <> tag) + _ -> Left "AST root must be a JSON object" + +parseLit :: A.Value -> Either Text Lit +parseLit v = case v of + A.Object o -> do + tag <- getStr o "tag" + case tag of + "LNumber" -> do + n <- getNum o "value" + Right (LNumber n) + "LText" -> LText <$> getStr o "value" + "LBool" -> LBool <$> getBool o "value" + _ -> Left ("Unknown literal tag: " <> tag) + _ -> Left "Literal must be an object" + +parseBind :: A.Value -> Either Text Bind +parseBind v = case v of + A.Object o -> do + n <- getStr o "name" + e <- parseAst =<< getField o "value" + Right (Bind n e) + _ -> Left "Bind must be an object" + +parseDoStmt :: A.Value -> Either Text DoStmt +parseDoStmt v = case v of + A.Object o -> do + tag <- getStr o "tag" + case tag of + "DoBind" -> do + name <- getStr o "name" + rawValue <- parseAst =<< getField o "value" + -- Phase 9.1d-4 fix: frontend が `x <- sample "obsName" dist` を + -- DoBind の value に raw expression として渡してくる。 validateStmts + -- 以降は value が「純粋な distribution」 であることを期待するので、 + -- ここで sample wrapper を剥がす。 sample 形でなければそのまま通す + -- (互換: 直接 dist を入れた古い経路があった場合のため)。 + let distOnly = case collectApp rawValue of + Right ("sample", [ELit (LText _samplerName), d]) -> d + _ -> rawValue + pure (DoBind name distOnly) + "DoLet" -> DoLet <$> (getArray o "binds" >>= mapM parseBind) + "DoExpr" -> DoExpr <$> (parseAst =<< getField o "value") + _ -> Left ("Unknown DoStmt tag: " <> tag) + _ -> Left "DoStmt must be an object" + +-- --------------------------------------------------------------------------- +-- JSON encoder (= parseAst の正確な逆。 round-trip: parseAst . exprToJSON ≡ Right) +-- +-- Phase 27.5 step 3 (2026-06-01): topology B で backend が stream sidecar に +-- start.params を組む際、 resolveHbmModel が返す Expr / TopBind を worker +-- (= parseAst で decode) が読める JSON 文字列に直す必要がある。 decoder と +-- 同じ module に逆変換を置き、 tag / field 名のズレを構造的に防ぐ。 +-- +-- 注: DoBind の value は sample wrapper を剥がした dist-only を前提とする +-- (parseDoStmt は sample wrapper を剥がすが、 既に剥がれた式には作用しない = +-- idempotent。 resolveHbmModel 経由の Expr は剥がし済)。 +-- --------------------------------------------------------------------------- + +exprToJSON :: Expr -> A.Value +exprToJSON e = case e of + ELit l -> obj "ELit" ["lit" A..= litToJSON l] + ECol n -> obj "ECol" ["name" A..= n] + EVar n -> obj "EVar" ["name" A..= n] + EApp f x -> obj "EApp" ["f" A..= exprToJSON f, "x" A..= exprToJSON x] + ELam a b -> obj "ELam" ["arg" A..= a, "body" A..= exprToJSON b] + EIf c a b -> obj "EIf" ["c" A..= exprToJSON c, "a" A..= exprToJSON a, "b" A..= exprToJSON b] + ELet bs body -> obj "ELet" ["binds" A..= map bindToJSON bs, "body" A..= exprToJSON body] + ENeg x -> obj "ENeg" ["e" A..= exprToJSON x] + EOp op a b -> obj "EOp" ["op" A..= op, "a" A..= exprToJSON a, "b" A..= exprToJSON b] + EList xs -> obj "EList" ["items" A..= map exprToJSON xs] + EDo stmts r -> obj "EDo" ["stmts" A..= map doStmtToJSON stmts, "ret" A..= exprToJSON r] + where + obj :: Text -> [Pair] -> A.Value + obj tag fields = A.object (("tag" A..= tag) : fields) + +litToJSON :: Lit -> A.Value +litToJSON l = case l of + LNumber n -> A.object ["tag" A..= ("LNumber" :: Text), "value" A..= n] + LText t -> A.object ["tag" A..= ("LText" :: Text), "value" A..= t] + LBool b -> A.object ["tag" A..= ("LBool" :: Text), "value" A..= b] + +bindToJSON :: Bind -> A.Value +bindToJSON (Bind n v) = A.object ["name" A..= n, "value" A..= exprToJSON v] + +doStmtToJSON :: DoStmt -> A.Value +doStmtToJSON s = case s of + DoBind n v -> A.object ["tag" A..= ("DoBind" :: Text), "name" A..= n, "value" A..= exprToJSON v] + DoLet bs -> A.object ["tag" A..= ("DoLet" :: Text), "binds" A..= map bindToJSON bs] + DoExpr v -> A.object ["tag" A..= ("DoExpr" :: Text), "value" A..= exprToJSON v] + +-- | [日本語]: @EApp (EApp (EVar f) a) b@ → @(f, [a, b])@。 distribution / +-- 関数適用の head + 引数列を取り出す。 head が変数でなければ Left。 +-- [English]: @EApp (EApp (EVar f) a) b@ → @(f, [a, b])@. Extracts the +-- head and argument list of a distribution / function application. +-- Returns @Left@ if the head is not a variable. +collectApp :: Expr -> Err (Text, [Expr]) +collectApp e0 = go e0 [] + where + go (EVar n) acc = Right (n, acc) + go (EApp f x) acc = go f (x : acc) + go _ _ = Left "Distribution must be a function applied to scalar args" + +-- --------------------------------------------------------------------------- +-- helpers +-- --------------------------------------------------------------------------- + +getField :: A.Object -> Text -> Either Text A.Value +getField o k = case KM.lookup (Key.fromText k) o of + Just v -> Right v + Nothing -> Left ("Missing field: " <> k) + +getStr :: A.Object -> Text -> Either Text Text +getStr o k = case KM.lookup (Key.fromText k) o of + Just (A.String s) -> Right s + _ -> Left ("Field not string: " <> k) + +getNum :: A.Object -> Text -> Either Text Double +getNum o k = case KM.lookup (Key.fromText k) o of + Just (A.Number n) -> Right (realToFrac n) + _ -> Left ("Field not number: " <> k) + +getBool :: A.Object -> Text -> Either Text Bool +getBool o k = case KM.lookup (Key.fromText k) o of + Just (A.Bool b) -> Right b + _ -> Left ("Field not bool: " <> k) + +getArray :: A.Object -> Text -> Either Text [A.Value] +getArray o k = case KM.lookup (Key.fromText k) o of + Just (A.Array xs) -> Right (V.toList xs) + _ -> Left ("Field not array: " <> k) diff --git a/src/Hanalyze/Model/HBM/Distribution.hs b/src/Hanalyze/Model/HBM/Distribution.hs new file mode 100644 --- /dev/null +++ b/src/Hanalyze/Model/HBM/Distribution.hs @@ -0,0 +1,1626 @@ +{-# LANGUAGE RankNTypes #-} +{-# LANGUAGE DeriveFunctor #-} +{-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE ScopedTypeVariables #-} +-- | +-- Module : Hanalyze.Model.HBM.Distribution +-- Description : HBM の多相確率分布 ADT と密度・CDF +-- Copyright : (c) 2026 Aelysce Project (Toshiaki Honda) +-- License : BSD-3-Clause +-- +-- [日本語]: HBM の多相確率分布 ADT と密度・CDF。 +-- +-- 'Distribution' は値型 @a@ に多相な確率分布。 @a@ は @Double@ (サンプリング・ +-- 密度)、 @Reverse s Double@ (AD 勾配)、 @Track@ (依存追跡) を渡せる。 本モジュール +-- は型・名前・__事前密度__ @logDensity@・多変量密度・閉形式 CDF を提供し、 純粋 +-- leaf 'Hanalyze.Model.HBM.Util' のみに依存する。 +-- +-- ★観測尤度 'logDensityObs' / 'obsLogSum' は __含めない__ (Eval 層へ残置。 +-- Distribution→Eval の cycle を避けるため)。 +-- +-- 責務分離により 'Hanalyze.Model.HBM' から抽出。 数値は 1 bit 不変。 +-- +-- [English]: HBM's polymorphic probability-distribution ADT, densities, and +-- CDFs. +-- +-- 'Distribution' is a probability distribution polymorphic in its value +-- type @a@. @a@ can be @Double@ (sampling and density), @Reverse s Double@ +-- (AD gradient), or @Track@ (dependency tracking). This module provides the +-- type, names, the __prior density__ @logDensity@, multivariate densities, +-- and closed-form CDFs, and depends only on the pure leaf +-- 'Hanalyze.Model.HBM.Util'. +-- +-- ★The observation likelihood 'logDensityObs' \/ 'obsLogSum' is +-- __deliberately excluded__ (kept in the Eval layer to avoid a +-- Distribution→Eval cycle). +-- +-- Extracted from 'Hanalyze.Model.HBM' as part of a responsibility +-- split. Numerically bit-identical. +module Hanalyze.Model.HBM.Distribution + ( Distribution (..) + , distName + , nameToTransform + , distToTransform + , logDensity + , logDensityRD + , logDensityObs + , obsLogSum + , multinomialLogDensity + , mvNormalLogDensity + , mvNormalCholLogDensity + , mvStudentTLogDensity + , dirichletMultinomialLogDensity + , wishartLogDensity + , erfA + , phiCdfA + , distCDF + , logCDF + , logSF + , logCDFInterval + ) where + +import Data.List (mapAccumL, zip4) +import Data.Text (Text) +-- Phase 92 B3: 'logDensityRD' (AD 定数正規化項の畳み込み) 用。 多相 @logDensity@ +-- 本体は AD 非依存のまま。 +import Data.Reflection (Reifies) +import qualified Numeric.AD.Internal.Reverse.Double as ADRD +import Hanalyze.Model.HBM.Util +import Hanalyze.Stat.Distribution (Transform (..)) + +-- --------------------------------------------------------------------------- +-- 多相分布 +-- --------------------------------------------------------------------------- + +-- | A probability distribution polymorphic in its value type @a@. +-- +-- @a@ ranges over @Double@ (sampling and density), @Reverse s Double@ +-- (AD-based gradient), @Track@ (dependency tracking) and so on. +data Distribution a + = Normal a a -- ^ Normal(μ, σ) + | Exponential a -- ^ Exp(rate) + | Gamma a a -- ^ Gamma(shape, rate) + | Beta a a -- ^ Beta(α, β) + | Poisson a -- ^ Poisson(λ) + | Binomial Int a -- ^ Binomial(n, p) + | Uniform a a -- ^ Uniform(low, high) + | StudentT a a a -- ^ StudentT(ν degrees of freedom, μ location, σ scale) + | Cauchy a a -- ^ Cauchy(x₀ location, γ scale) + | HalfNormal a -- ^ HalfNormal(σ) — support: x ≥ 0 + | HalfCauchy a -- ^ HalfCauchy(γ scale) — support: x ≥ 0 + | LogNormal a a -- ^ LogNormal(μ log-mean, σ log-sd) — support: x > 0 + | Bernoulli a -- ^ Bernoulli(p) — observed: 0 or 1 + | Categorical [a] -- ^ Categorical(probs) — observed: 0..K-1 + | Mixture [a] [Distribution a] + -- ^ @Mixture(weights, components)@ — + -- @log p(x) = logSumExp(log w_k + log p_k(x))@. + -- Weights need only be positive; they are auto-normalized. + | Truncated (Distribution a) (Maybe a) (Maybe a) + -- ^ @Truncated(d, lo, hi)@: restrict the support of @d@ to + -- @[lo, hi]@. Out-of-range observations get @-∞@. + -- 'Nothing' bounds mean @-∞ / +∞@. Only base distributions with a + -- CDF (Normal / Exponential / LogNormal / Uniform) are supported. + | Censored (Distribution a) (Maybe a) (Maybe a) + -- ^ @Censored(d, lo, hi)@: censor @y ≤ lo@ on the left and + -- @y ≥ hi@ on the right. When @y_i@ equals a threshold the CDF/SF + -- is used. Useful for Tobit-style models. Only CDF-supporting + -- base distributions. + | MvNormal [a] [[a]] + -- ^ @MvNormal(μ, Σ)@: multivariate normal (observation-only). + -- @μ@ is a length-@k@ mean vector, @Σ@ is the @k×k@ + -- symmetric-positive-definite covariance. Pass @k@-vector + -- observations through @observeMV@. Density is computed via + -- Cholesky. /Not supported/ as a latent (@sample@ returns 0 + -- density). + | MvNormalChol [a] [a] [[a]] + -- ^ @MvNormalChol(μ, σ, L)@: multivariate normal parameterized by a + -- scale vector @σ@ (length @k@) and a /correlation/ Cholesky factor + -- @L@ (lower-triangular @k×k@, typically from @lkjCorrCholesky@). + -- The covariance is @Σ = (diag σ · L)(diag σ · L)ᵀ@. The density + -- uses the scaled Cholesky @M = diag σ · L@ directly (forward + -- substitution, no re-decomposition) — numerically the most stable + -- parameterization (Stan's @multi_normal_cholesky@ idiom). + -- Observation-only; pass @k@-vectors via @observeMV@. + | MvNormalGpRBF [a] a a a + -- ^ [日本語]: @MvNormalGpRBF(x, α, ρ, σ)@ — zero-mean GP 回帰尤度 + -- 専用の多変量正規 (observation-only)。 共分散は RBF (exp-quad) カーネル + -- @Σ_ij = α² exp(-0.5 (x_i-x_j)²/ρ²) + [i=j](1e-10 + σ)@ で内部構築する。 + -- 汎用 'MvNormal' と密度は同値だが、 カーネルの役割 (x/α/ρ/σ) を型で明示 + -- 保持することで、 勾配コンパイラ (@gpRBFAnalyticVG@) が + -- __Cholesky を AD tape に載せない閉形式随伴__ + -- (@∂Σ/∂α=2K'/α@・@∂Σ/∂ρ=K'∘d²/ρ³@・ + -- @∂Σ/∂σ=I@) を使える。 @x@ は共変量 data (定数)、 α/ρ/σ は latent。 + -- 観測は length-@k@ ベクトルを @observeMV@ で渡す (μ=0 固定)。 + -- [English]: @MvNormalGpRBF(x, α, ρ, σ)@ — a multivariate normal + -- dedicated to zero-mean GP regression likelihoods (observation-only). + -- The covariance is built internally from an RBF (exp-quad) kernel + -- @Σ_ij = α² exp(-0.5 (x_i-x_j)²/ρ²) + [i=j](1e-10 + σ)@. Its density + -- is equivalent to the generic 'MvNormal', but by keeping the + -- kernel's roles (x\/α\/ρ\/σ) explicit at the type level, the + -- gradient compiler (@gpRBFAnalyticVG@) can use a + -- __closed-form adjoint that never puts the Cholesky on the AD tape__ + -- (@∂Σ/∂α=2K'/α@; @∂Σ/∂ρ=K'∘d²/ρ³@; @∂Σ/∂σ=I@). @x@ is covariate + -- data (a constant); α\/ρ\/σ are latent. Observations are passed as + -- length-@k@ vectors via @observeMV@ (with @μ@ fixed at 0). + | HmmForwardNormal [a] [[a]] [a] a + -- ^ [日本語]: @HmmForwardNormal(π_0, trans, μs, σ)@ — Normal emission の + -- 隠れマルコフモデル周辺尤度 (observation-only)。 観測列 y_{1..T} 全体を + -- 1 つの多変量観測として @observeMV@ で渡す (@observeMV nm d [ys]@)。 + -- 密度は @'hmmForwardLogLik' π_0 trans emit@ (emit[t][k] = + -- Normal(μs[k], σ) の logpdf(y_t)) と同値。 状態役割 (π_0/遷移行/emission + -- 平均/σ) を型で明示保持することで、 勾配コンパイラ (@hmmAnalyticVG@) が + -- __forward-backward の閉形式随伴__ (∂logL/∂emit = γ_t・∂logL/∂T_ij = ξ + -- 集計・AD tape ゼロ) を使える。 π_0 は非正規化可 (log 空間で加算されるのみ)。 + -- [English]: @HmmForwardNormal(π_0, trans, μs, σ)@ — the marginal + -- likelihood of a hidden Markov model with Normal emissions + -- (observation-only). The whole observation sequence y_{1..T} is + -- passed as a single multivariate observation via @observeMV@ + -- (@observeMV nm d [ys]@). Its density is equivalent to + -- @'hmmForwardLogLik' π_0 trans emit@ (where emit[t][k] = the log + -- pdf of y_t under Normal(μs[k], σ)). By keeping the state roles + -- (π_0 \/ transition matrix \/ emission mean \/ σ) explicit at the + -- type level, the gradient compiler (@hmmAnalyticVG@) can use a + -- __closed-form forward-backward adjoint__ (∂logL/∂emit = γ_t; + -- ∂logL/∂T_ij = aggregated ξ; zero AD tape). @π_0@ may be + -- unnormalized (it is only ever summed in log space). + | ArmaNormal a a a a + -- ^ [日本語]: @ArmaNormal(μ, φ, θ, σ)@ — ARMA(1,1) の条件付き尤度 + -- (observation-only)。 観測列 y_{1..T} 全体を 1 つの多変量観測として + -- @observeMV@ で渡す (@observeMV nm d [ys]@)。 密度は Stan 原典 arma11 の + -- err 逐次再帰 (@err_1 = y_1 − (μ+φμ)@・@err_t = y_t − μ − φ·y_{t−1} − + -- θ·err_{t−1}@・@err_t ~ Normal(0, σ)@) と同値。 役割 (μ/φ/θ/σ) を型で + -- 明示保持することで、 勾配コンパイラ (@armaAnalyticVG@) が + -- __逆向き 1 パスの閉形式随伴__ + -- (@ē_t = −e_t/σ² − θ·ē_{t+1}@ の線形随伴再帰・ + -- AD tape ゼロ) を使える。 + -- [English]: @ArmaNormal(μ, φ, θ, σ)@ — the conditional likelihood of + -- an ARMA(1,1) process (observation-only). The whole observation + -- sequence y_{1..T} is passed as a single multivariate observation + -- via @observeMV@ (@observeMV nm d [ys]@). Its density is equivalent + -- to the original Stan arma11 err recursion (@err_1 = y_1 − + -- (μ+φμ)@; @err_t = y_t − μ − φ·y_{t−1} − θ·err_{t−1}@; @err_t ~ + -- Normal(0, σ)@). By keeping the roles (μ\/φ\/θ\/σ) explicit at the + -- type level, the gradient compiler (@armaAnalyticVG@) can use a + -- __closed-form adjoint computed in a single backward pass__ (the + -- linear adjoint recursion @ē_t = −e_t/σ² − θ·ē_{t+1}@; zero AD + -- tape). + | GradedResponseIrt [a] [Int] [Double] [[Double]] + -- ^ [日本語]: @GradedResponseIrt(θs, ncats, δs, γs)@ — graded response + -- IRT (順序ロジット・BUGS bones) の尤度 (observation-only)。 @θs@ = + -- 受験者能力 (latent・唯一の param 側)、 @ncats[j]@/@δs[j]@/@γs[j][k]@ = + -- 項目 j のカテゴリ数/識別力/カットポイント (__定数データ__)。 観測は + -- grade 行列 (nChild×nItem 行優先・1-based カテゴリ・欠測 = −1) を + -- @observeMV@ で 1 観測として渡す (@observeMV nm d [grades]@)。 + -- 密度は @Q_k = invlogit(δ(θ−γ_k))@ の隣接差 p のカテゴリ対数確率と + -- 同値。 θ_i (スカラ) 毎に独立なため、 勾配コンパイラ + -- (@gradedIrtAnalyticVG@) が __解析勾配__ (@dQ/dθ = δ·Q(1−Q)@ の差分・ + -- AD tape ゼロ) を使える。 + -- [English]: @GradedResponseIrt(θs, ncats, δs, γs)@ — the likelihood + -- of a graded response IRT model (ordered-logit \/ BUGS "bones", + -- observation-only). @θs@ = respondent ability (latent; the only + -- parameter side), @ncats[j]@\/@δs[j]@\/@γs[j][k]@ = item @j@'s + -- category count \/ discrimination \/ cutpoints (__constant data__). + -- Observations are passed as a single grade matrix (nChild×nItem, + -- row-major, 1-based categories, −1 for missing) via @observeMV@ + -- (@observeMV nm d [grades]@). Its density is equivalent to the + -- category log-probability from the adjacent difference of + -- @Q_k = invlogit(δ(θ−γ_k))@. Since it is independent per scalar + -- θ_i, the gradient compiler (@gradedIrtAnalyticVG@) can use an + -- __analytic gradient__ (the difference @dQ/dθ = δ·Q(1−Q)@; zero + -- AD tape). + | NegativeBinomial a a + -- ^ @NegativeBinomial(μ, α)@ (PyMC parameterization). + -- @mean = μ@, @var = μ + μ²/α@ (Poisson in the limit + -- @α → ∞@). Likelihood for over-dispersed count data; + -- observations are non-negative integers. + | Multinomial Int [a] + -- ^ @Multinomial(n, [p_0, …, p_{K-1}])@ (observation-only). + -- @n@ is the trial count and @p@ the probability vector. + -- Observations are @K@-dimensional count vectors summing to @n@, + -- passed via @observeMV@. + | ZeroInflatedPoisson a a + -- ^ @ZeroInflatedPoisson(ψ, λ)@: zero-inflated Poisson. + -- @ψ ∈ [0, 1]@ is the structural-zero probability. + -- @P(0) = ψ + (1-ψ) e^{-λ}@, + -- @P(k>0) = (1-ψ) λ^k e^{-λ} / k!@. + | ZeroInflatedBinomial Int a a + -- ^ @ZeroInflatedBinomial(n, ψ, p)@: zero-inflated binomial. + -- @P(0) = ψ + (1-ψ) (1-p)^n@, + -- @P(k>0) = (1-ψ) C(n,k) p^k (1-p)^{n-k}@. + | InverseGamma a a + -- ^ @InverseGamma(α, β)@. Support @x > 0@. If + -- @X ~ InverseGamma(α, β)@ then @1/X ~ Gamma(α, β)@ (rate + -- parameterization). Common conjugate prior on variance + -- (@mean = β/(α−1)@, finite when @α > 1@). + | Weibull a a + -- ^ @Weibull(k shape, λ scale)@: a standard survival distribution. + -- Support @x > 0@. @pdf = (k/λ) (x/λ)^{k-1} exp(-(x/λ)^k)@. + -- With @k = 1@ this is @Exponential(rate = 1/λ)@. + | Pareto a a + -- ^ @Pareto(α shape, x_m scale)@: heavy-tailed power law. + -- Support @x ≥ x_m > 0@. @pdf = α x_m^α / x^{α+1}@. + -- Mean @= α x_m / (α-1)@ when @α > 1@. + | BetaBinomial Int a a + -- ^ @BetaBinomial(n, α, β)@ overdispersed binomial + -- (observation-only). + -- @P(k) = C(n, k) B(k+α, n-k+β) / B(α, β)@. With @α = β = 1@ + -- this is uniform on @{0, …, n}@; large @α/β@ tends to a + -- binomial. + | VonMises a a + -- ^ @VonMises(μ location, κ concentration)@: distribution on the + -- circle @(-π, π]@. + -- @pdf = exp(κ cos(x − μ)) / (2π I_0(κ))@. + -- @κ → 0@ approaches uniform; @κ → ∞@ approaches + -- @Normal(μ, 1/√κ)@. + | SkewNormal a a a + -- ^ [日本語]: @SkewNormal(μ location, σ scale, α shape)@。 + -- @pdf = (2/σ) φ((x−μ)/σ) Φ(α(x−μ)/σ)@. + -- @α = 0@ で標準正規。 @α > 0@ で右側に歪み、 @α < 0@ で左側。 + -- Sample は Henze 1986: @δ = α/√(1+α²)@, + -- @X = μ + σ(δ |U₀| + √(1−δ²) U₁)@ with i.i.d. @U_i ~ N(0,1)@. + -- [English]: @SkewNormal(μ location, σ scale, α shape)@. + -- @pdf = (2/σ) φ((x−μ)/σ) Φ(α(x−μ)/σ)@. @α = 0@ gives the standard + -- normal; @α > 0@ skews right, @α < 0@ skews left. Sampling follows + -- Henze 1986: @δ = α/√(1+α²)@, + -- @X = μ + σ(δ |U₀| + √(1−δ²) U₁)@ with i.i.d. @U_i ~ N(0,1)@. + | Logistic a a + -- ^ [日本語]: @Logistic(μ location, s scale)@。 + -- @pdf = e^{−z} / (s(1+e^{−z})²)@ with @z = (x−μ)/s@. + -- 平均 @μ@、 分散 @s²π²/3@。 closed-form CDF あり。 + -- [English]: @Logistic(μ location, s scale)@. + -- @pdf = e^{−z} / (s(1+e^{−z})²)@ with @z = (x−μ)/s@. Mean @μ@, + -- variance @s²π²/3@. Has a closed-form CDF. + | Gumbel a a + -- ^ [日本語]: @Gumbel(μ location, β scale)@ (最大値型極値分布)。 + -- @pdf = (1/β) exp(−z − e^{−z})@ with @z = (x−μ)/β@. + -- 平均 @μ + βγ@ (γ ≈ 0.5772 オイラー定数)、 分散 @β²π²/6@。 + -- closed-form CDF: @F(x) = exp(−exp(−z))@. + -- [English]: @Gumbel(μ location, β scale)@ (a maximum-type extreme + -- value distribution). @pdf = (1/β) exp(−z − e^{−z})@ with + -- @z = (x−μ)/β@. Mean @μ + βγ@ (γ ≈ 0.5772, the Euler-Mascheroni + -- constant), variance @β²π²/6@. Closed-form CDF: + -- @F(x) = exp(−exp(−z))@. + | AsymmetricLaplace a a a + -- ^ [日本語]: @AsymmetricLaplace(b scale > 0, κ asymmetry > 0, μ location)@ + -- (PyMC parameterization、 分位点回帰の尤度)。 + -- @pdf = b/(κ+1/κ) · exp(−b·κ·(x−μ))@ for @x ≥ μ@、 + -- @pdf = b/(κ+1/κ) · exp(b/κ·(x−μ))@ for @x < μ@。 + -- @κ = 1@ で対称ラプラス、 @κ > 1@ で右側裾長。 + -- [English]: @AsymmetricLaplace(b scale > 0, κ asymmetry > 0, + -- μ location)@ (PyMC parameterization; the likelihood used for + -- quantile regression). @pdf = b/(κ+1/κ) · exp(−b·κ·(x−μ))@ for + -- @x ≥ μ@, @pdf = b/(κ+1/κ) · exp(b/κ·(x−μ))@ for @x < μ@. @κ = 1@ + -- gives the symmetric Laplace; @κ > 1@ gives a heavier right tail. + | OrderedLogistic a [a] + -- ^ [日本語]: @OrderedLogistic(η linear predictor, cuts = [c₁, …, c_{K-1}])@ + -- (順序ロジット回帰)。 + -- 観測 @y ∈ {0, …, K-1}@、 + -- @P(y=k) = σ(c_{k+1} − η) − σ(c_k − η)@ with + -- @σ(x) = 1/(1+e^{-x})@, @c_0 = −∞, c_K = +∞@. + -- cuts は __increasing__ 列、 入力側で確保すること。 + -- observation-only。 + -- [English]: @OrderedLogistic(η linear predictor, cuts = [c₁, …, + -- c_{K-1}])@ (ordered-logit regression). Observations + -- @y ∈ {0, …, K-1}@, + -- @P(y=k) = σ(c_{k+1} − η) − σ(c_k − η)@ with + -- @σ(x) = 1/(1+e^{-x})@, @c_0 = −∞, c_K = +∞@. @cuts@ must be an + -- __increasing__ sequence; ensure this on the caller's side. + -- Observation-only. + | DiscreteUniform Int Int + -- ^ [日本語]: @DiscreteUniform(lo, hi)@ (両端を含む)。 + -- @pmf = 1/(hi-lo+1)@ for @lo ≤ y ≤ hi@。 observation-only。 + -- [English]: @DiscreteUniform(lo, hi)@ (both endpoints inclusive). + -- @pmf = 1/(hi-lo+1)@ for @lo ≤ y ≤ hi@. Observation-only. + | Geometric a + -- ^ [日本語]: @Geometric(p)@ (PyMC 慣例 = 初回成功までの試行回数)。 + -- support @y = 1, 2, 3, …@、 @pmf = (1−p)^{y-1} p@。 + -- observation-only。 + -- [English]: @Geometric(p)@ (PyMC convention = number of trials up + -- to and including the first success). Support @y = 1, 2, 3, …@, + -- @pmf = (1−p)^{y-1} p@. Observation-only. + | HyperGeometric Int Int Int + -- ^ [日本語]: @HyperGeometric(N total, K successes, n draws)@ + -- (非復元抽出の成功数)。 + -- @pmf = C(K, y) C(N-K, n-y) / C(N, n)@、 + -- support @max(0, n+K-N) ≤ y ≤ min(n, K)@。 observation-only。 + -- [English]: @HyperGeometric(N total, K successes, n draws)@ (the + -- number of successes in sampling without replacement). + -- @pmf = C(K, y) C(N-K, n-y) / C(N, n)@, + -- support @max(0, n+K-N) ≤ y ≤ min(n, K)@. Observation-only. + | ZeroInflatedNegativeBinomial a a a + -- ^ [日本語]: @ZeroInflatedNegativeBinomial(ψ, μ, α)@ (過分散ゼロ過剰)。 + -- @P(0) = ψ + (1-ψ) (α/(α+μ))^α@、 + -- @P(k>0) = (1-ψ) · NegBin(k | μ, α)@。 + -- [English]: @ZeroInflatedNegativeBinomial(ψ, μ, α)@ (overdispersed, + -- zero-inflated). @P(0) = ψ + (1-ψ) (α/(α+μ))^α@, + -- @P(k>0) = (1-ψ) · NegBin(k | μ, α)@. + | MvStudentT a [a] [[a]] + -- ^ [日本語]: @MvStudentT(ν, μ, Σ)@ (ロバスト多変量)。 + -- @ν > 0@ 自由度、 @μ@ は @k@ 次元平均、 @Σ@ は @k×k@ SPD scale matrix。 + -- 観測 (observation-only)、 @y :: [Double]@ は flatten された + -- @k@ ベクトル列 (@observeMV@ で渡す)。 + -- @ν → ∞@ で MvNormal に収束。 + -- [English]: @MvStudentT(ν, μ, Σ)@ (a robust multivariate + -- distribution). @ν > 0@ degrees of freedom, @μ@ is a @k@-dimensional + -- mean, @Σ@ is a @k×k@ SPD scale matrix. Observation-only; + -- @y :: [Double]@ is a flattened sequence of @k@-vectors (passed via + -- @observeMV@). Converges to MvNormal as @ν → ∞@. + | DirichletMultinomial Int [a] + -- ^ [日本語]: @DirichletMultinomial(n trials, α concentration K-vector)@ + -- (過分散 multinomial)。 + -- 観測 y は @K@ 次元 counts、 @Σ yᵢ = n@。 + -- @logpmf = log Γ(α₀) − log Γ(α₀+n) + -- + Σ [log Γ(yᵢ+αᵢ) − log Γ(αᵢ)] + -- + log n! − Σ log yᵢ!@、 @α₀ = Σαᵢ@. + -- observation-only。 + -- [English]: @DirichletMultinomial(n trials, α concentration + -- K-vector)@ (an overdispersed multinomial). Observations @y@ are + -- @K@-dimensional counts with @Σ yᵢ = n@. + -- @logpmf = log Γ(α₀) − log Γ(α₀+n) + -- + Σ [log Γ(yᵢ+αᵢ) − log Γ(αᵢ)] + -- + log n! − Σ log yᵢ!@, @α₀ = Σαᵢ@. + -- Observation-only. + | Triangular a a a + -- ^ [日本語]: @Triangular(lower, c mode, upper)@ (弱情報事前)。 + -- Support @[lower, upper]@、 @lower ≤ c ≤ upper@。 + -- @pdf = 2(x-lower)/((upper-lower)(c-lower))@ for @lower ≤ x ≤ c@、 + -- @pdf = 2(upper-x)/((upper-lower)(upper-c))@ for @c < x ≤ upper@。 + -- closed-form CDF / 逆 CDF sample。 + -- [English]: @Triangular(lower, c mode, upper)@ (a weakly + -- informative prior). Support @[lower, upper]@, @lower ≤ c ≤ upper@. + -- @pdf = 2(x-lower)/((upper-lower)(c-lower))@ for @lower ≤ x ≤ c@, + -- @pdf = 2(upper-x)/((upper-lower)(upper-c))@ for @c < x ≤ upper@. + -- Has a closed-form CDF and inverse-CDF sampling. + | Kumaraswamy a a + -- ^ [日本語]: @Kumaraswamy(a, b)@ (Beta の代替、 closed-form CDF)。 + -- Support @(0, 1)@、 @pdf = a·b·x^{a-1}(1-x^a)^{b-1}@。 + -- CDF @= 1 - (1-x^a)^b@、 sample @x = (1-(1-u)^{1/b})^{1/a}@。 + -- [English]: @Kumaraswamy(a, b)@ (a Beta alternative with a + -- closed-form CDF). Support @(0, 1)@, + -- @pdf = a·b·x^{a-1}(1-x^a)^{b-1}@. CDF @= 1 - (1-x^a)^b@; sample + -- @x = (1-(1-u)^{1/b})^{1/a}@. + | Rice a a + -- ^ [日本語]: @Rice(ν, σ)@ (MRI / Rayleigh 拡張)。 + -- Support @x ≥ 0@、 @ν ≥ 0@、 @σ > 0@。 + -- @pdf = (x/σ²) exp(-(x²+ν²)/(2σ²)) I_0(xν/σ²)@、 + -- @ν = 0@ で Rayleigh(σ)。 @logBesselI0@ で評価。 + -- sample: @X = √(Y₁² + Y₂²)@ with @Y₁ ~ N(ν, σ²), Y₂ ~ N(0, σ²)@。 + -- [English]: @Rice(ν, σ)@ (an MRI \/ Rayleigh extension). Support + -- @x ≥ 0@, @ν ≥ 0@, @σ > 0@. + -- @pdf = (x/σ²) exp(-(x²+ν²)/(2σ²)) I_0(xν/σ²)@; reduces to + -- Rayleigh(σ) at @ν = 0@. Evaluated via @logBesselI0@. Sample: + -- @X = √(Y₁² + Y₂²)@ with @Y₁ ~ N(ν, σ²), Y₂ ~ N(0, σ²)@. + | DiscreteWeibull a a + -- ^ [日本語]: @DiscreteWeibull(q, β)@ (整数 Weibull)。 + -- Support @{0, 1, 2, …}@、 @0 < q < 1, β > 0@。 + -- @P(X ≤ k) = 1 - q^{(k+1)^β}@、 + -- @pmf(k) = q^{k^β} - q^{(k+1)^β}@。 observation-only。 + -- sample: @k = ⌈(log(1-u)/log q)^{1/β}⌉ - 1@。 + -- [English]: @DiscreteWeibull(q, β)@ (an integer-valued Weibull). + -- Support @{0, 1, 2, …}@, @0 < q < 1, β > 0@. + -- @P(X ≤ k) = 1 - q^{(k+1)^β}@, + -- @pmf(k) = q^{k^β} - q^{(k+1)^β}@. Observation-only. Sample: + -- @k = ⌈(log(1-u)/log q)^{1/β}⌉ - 1@. + | Wishart a [[a]] + -- ^ [日本語]: @Wishart(ν degrees, V scale matrix)@ (共分散プライアの直接表現)。 + -- @ν > k-1@、 @V@ は @k×k@ SPD scale matrix。 + -- 観測 (observation-only)、 @k×k@ 観測行列 W を flatten で渡す + -- (長さ @k²@、 row-major)。 @observeMV@ で渡す想定。 + -- @logpdf(W) = -(νk/2) log 2 - (ν/2) log|V| - log Γ_k(ν/2) + -- + ((ν-k-1)/2) log|W| - (1/2) tr(V⁻¹ W)@、 + -- @log Γ_k(z) = (k(k-1)/4) log π + Σ_{i=1}^k log Γ((z+1-i)/2)@。 + -- [English]: @Wishart(ν degrees, V scale matrix)@ (a direct + -- representation of a covariance prior). @ν > k-1@, @V@ is a + -- @k×k@ SPD scale matrix. Observation-only; the @k×k@ observation + -- matrix W is passed flattened (length @k²@, row-major), intended + -- to be passed via @observeMV@. + -- @logpdf(W) = -(νk/2) log 2 - (ν/2) log|V| - log Γ_k(ν/2) + -- + ((ν-k-1)/2) log|W| - (1/2) tr(V⁻¹ W)@, + -- @log Γ_k(z) = (k(k-1)/4) log π + Σ_{i=1}^k log Γ((z+1-i)/2)@. + | Bound (Distribution a) (Maybe a) (Maybe a) + -- ^ [日本語]: @Bound(d, lo, hi)@ (PyMC 互換)。 + -- @d@ の支持を @[lo, hi]@ に制限する。 'Truncated' とほぼ同義 + -- (実装も委譲)。 'Nothing' は @-∞ / +∞@。 + -- 違いは語用論のみ: PyMC では prior 寄りで Bound、 観測寄りで + -- Truncated を使う慣例があるため API として並べた。 + -- [English]: @Bound(d, lo, hi)@ (PyMC-compatible). Restricts the + -- support of @d@ to @[lo, hi]@. Nearly synonymous with 'Truncated' + -- (and delegates to its implementation). 'Nothing' bounds mean + -- @-∞ \/ +∞@. The difference is purely one of convention: PyMC + -- uses @Bound@ on the prior side and @Truncated@ on the + -- observation side, so both are exposed in the API to match. + | OrderedProbit a [a] + -- ^ [日本語]: @OrderedProbit(η linear predictor, cuts = [c₁, …, c_{K-1}])@ + -- (順序プロビット回帰)。 + -- @P(y=k) = Φ(c_{k+1} − η) − Φ(c_k − η)@ with + -- @c_0 = −∞, c_K = +∞@、 Φ は標準正規 CDF (@phiCdfA@)。 + -- cuts は increasing 列、 入力側で確保。 observation-only。 + -- [English]: @OrderedProbit(η linear predictor, cuts = [c₁, …, + -- c_{K-1}])@ (ordered-probit regression). + -- @P(y=k) = Φ(c_{k+1} − η) − Φ(c_k − η)@ with + -- @c_0 = −∞, c_K = +∞@, where Φ is the standard normal CDF + -- (@phiCdfA@). @cuts@ must be an increasing sequence, ensured on the + -- caller's side. Observation-only. + deriving (Show, Functor) + +-- | Display name of a distribution constructor (e.g. @\"Normal\"@). +distName :: Distribution a -> Text +distName Normal{} = "Normal" +distName Exponential{} = "Exponential" +distName Gamma{} = "Gamma" +distName Beta{} = "Beta" +distName Poisson{} = "Poisson" +distName Binomial{} = "Binomial" +distName Uniform{} = "Uniform" +distName StudentT{} = "StudentT" +distName Cauchy{} = "Cauchy" +distName HalfNormal{} = "HalfNormal" +distName HalfCauchy{} = "HalfCauchy" +distName LogNormal{} = "LogNormal" +distName Bernoulli{} = "Bernoulli" +distName Categorical{} = "Categorical" +distName Mixture{} = "Mixture" +distName Truncated{} = "Truncated" +distName Censored{} = "Censored" +distName MvNormal{} = "MvNormal" +distName MvNormalChol{} = "MvNormalChol" +distName MvNormalGpRBF{} = "MvNormalGpRBF" +distName HmmForwardNormal{} = "HmmForwardNormal" +distName ArmaNormal{} = "ArmaNormal" +distName GradedResponseIrt{} = "GradedResponseIrt" +distName NegativeBinomial{} = "NegativeBinomial" +distName Multinomial{} = "Multinomial" +distName ZeroInflatedPoisson{} = "ZeroInflatedPoisson" +distName ZeroInflatedBinomial{} = "ZeroInflatedBinomial" +distName InverseGamma{} = "InverseGamma" +distName Weibull{} = "Weibull" +distName Pareto{} = "Pareto" +distName BetaBinomial{} = "BetaBinomial" +distName VonMises{} = "VonMises" +distName SkewNormal{} = "SkewNormal" +distName Logistic{} = "Logistic" +distName Gumbel{} = "Gumbel" +distName AsymmetricLaplace{} = "AsymmetricLaplace" +distName OrderedLogistic{} = "OrderedLogistic" +distName DiscreteUniform{} = "DiscreteUniform" +distName Geometric{} = "Geometric" +distName HyperGeometric{} = "HyperGeometric" +distName ZeroInflatedNegativeBinomial{} = "ZeroInflatedNegativeBinomial" +distName MvStudentT{} = "MvStudentT" +distName DirichletMultinomial{} = "DirichletMultinomial" +distName Triangular{} = "Triangular" +distName Kumaraswamy{} = "Kumaraswamy" +distName Rice{} = "Rice" +distName DiscreteWeibull{} = "DiscreteWeibull" +distName Wishart{} = "Wishart" +distName Bound{} = "Bound" +distName OrderedProbit{} = "OrderedProbit" + +-- | [日本語]: 分布名 → NUTS が探索する __unconstrained 変換種別__。 latent の制約付き台 +-- (正値・単位区間) を実数空間へ写す種別を返す。 +-- +-- ★これが分布→変換の __唯一の表__。 @getTransforms@ +-- (@Gradient@・node walk 版) も本関数へ委譲する。 分布を latent 化して台が +-- 変わる場合はここを更新する (1 箇所)。 未列挙は保守的に 'UnconstrainedT'。 +-- +-- [English]: Distribution name → the __unconstrained transform kind__ +-- that NUTS explores. Returns the kind that maps a latent's constrained +-- support (positive \/ unit interval) to real space. +-- +-- ★This is the __single source of truth__ for the distribution→transform +-- mapping. @getTransforms@ (the @Gradient@ \/ node-walk version) also +-- delegates to this function. Update it here (one place) whenever a +-- distribution's support changes when latent-ized. Anything not +-- enumerated conservatively falls back to 'UnconstrainedT'. +nameToTransform :: Text -> Transform +nameToTransform "Exponential" = PositiveT +nameToTransform "Gamma" = PositiveT +nameToTransform "HalfNormal" = PositiveT +nameToTransform "HalfCauchy" = PositiveT +nameToTransform "LogNormal" = PositiveT -- support: x>0 (log は AD 安全) +nameToTransform "InverseGamma" = PositiveT +nameToTransform "Weibull" = PositiveT +nameToTransform "Pareto" = PositiveT +nameToTransform "Beta" = UnitIntervalT +nameToTransform "Bernoulli" = UnitIntervalT -- p ∈ (0,1) +nameToTransform "BetaBinomial" = UnitIntervalT +nameToTransform _ = UnconstrainedT -- Normal/StudentT/Cauchy/Uniform 等 +-- 注: Uniform の真の制約変換は logit-on-(lo,hi) だが現状未実装 (unconstrained 扱い)。 + +-- | [日本語]: 分布 (ADT) → unconstrained 変換種別。 'nameToTransform' の値レベル版。 +-- [English]: Distribution (ADT) → unconstrained transform kind. The +-- value-level counterpart of 'nameToTransform'. +distToTransform :: Distribution a -> Transform +distToTransform = nameToTransform . distName + +-- | Log probability of a single multinomial observation (a @K@-vector +-- of counts). +-- log P(k_1, …, k_K) = log n!/Π k_i! + Σ k_i log p_i +{-# INLINABLE multinomialLogDensity #-} +multinomialLogDensity :: forall a. (Floating a, Ord a) + => Int -> [a] -> [Double] -> a +multinomialLogDensity n probs counts + | length probs /= length counts = negInf + | sum (map round counts :: [Int]) /= n = negInf + | any (< 0) counts = negInf + | any (\p -> p <= 0) probs = negInf + | otherwise = + let logFactN = realToFrac (logFactorial n) :: a + logFactSum = sum [ realToFrac (logFactorial (round c :: Int)) :: a + | c <- counts ] + dotPart = sum (zipWith (\c p -> realToFrac c * log p) counts probs) + in logFactN - logFactSum + dotPart + +-- | [日本語]: 'MvNormal' の 1 観測 (k-vector) の log density。 +-- log p(y) = -k/2 log(2π) - 0.5 log|Σ| - 0.5 (y-μ)ᵀ Σ⁻¹ (y-μ) +-- Σ⁻¹ と log|Σ| は Cholesky 分解 Σ = L Lᵀ から計算。 +-- [English]: Log density of an 'MvNormal' at a single @k@-vector +-- observation. +-- log p(y) = -k/2 log(2π) - 0.5 log|Σ| - 0.5 (y-μ)ᵀ Σ⁻¹ (y-μ) +-- Σ⁻¹ and log|Σ| are computed from the Cholesky decomposition +-- Σ = L Lᵀ. +{-# INLINABLE mvNormalLogDensity #-} +mvNormalLogDensity :: forall a. (Floating a, Ord a) => [a] -> [[a]] -> [a] -> a +mvNormalLogDensity mu cov yObs + | length mu == 0 = 0 + | length yObs /= length mu = negInf + | otherwise = + case choleskyL cov of + Nothing -> negInf + Just l -> + let k = length mu + kA = fromIntegral k :: a + d = zipWith (-) yObs mu + z = forwardSub l d -- L z = d + quad = sum (map (\zi -> zi * zi) z) + logDet = 2 * sum [ log ((l !! i) !! i) | i <- [0 .. k - 1] ] + in -0.5 * kA * log (2 * pi) - 0.5 * logDet - 0.5 * quad + +-- | [日本語]: 'MvNormalChol' の 1 観測 (k-vector) の log density。 +-- scale vector @σ@ と /相関/ Cholesky 因子 @L@ から scaled Cholesky +-- @M = diag σ · L@ (= @M_ij = σ_i · L_ij@) を直接構成し、 共分散 +-- @Σ = M Mᵀ@ を /再分解せず/ 評価する: +-- @log p(y) = -k/2 log(2π) - Σ log M_ii - 0.5 |z|²@、 @M z = (y-μ)@ を +-- 前進代入で解く。 @log|Σ| = 2 Σ log M_ii@ なので密度の @-0.5 log|Σ|@ は +-- @-Σ log M_ii@。 'mvNormalLogDensity' (full Σ → choleskyL) と @Σ = M Mᵀ@ で +-- 数値一致する。 Stan の @multi_normal_cholesky@ と同じ idiom。 +-- [English]: Log density of an 'MvNormalChol' at a single @k@-vector +-- observation. Builds the scaled Cholesky factor +-- @M = diag σ · L@ (i.e. @M_ij = σ_i · L_ij@) directly from the scale +-- vector @σ@ and the /correlation/ Cholesky factor @L@, then evaluates +-- the covariance @Σ = M Mᵀ@ /without re-decomposing/ it: +-- @log p(y) = -k/2 log(2π) - Σ log M_ii - 0.5 |z|²@, where +-- @M z = (y-μ)@ is solved by forward substitution. Since +-- @log|Σ| = 2 Σ log M_ii@, the density's @-0.5 log|Σ|@ term becomes +-- @-Σ log M_ii@. This is numerically identical to 'mvNormalLogDensity' +-- (full Σ → choleskyL) under @Σ = M Mᵀ@. Same idiom as Stan's +-- @multi_normal_cholesky@. +{-# INLINABLE mvNormalCholLogDensity #-} +mvNormalCholLogDensity :: forall a. (Floating a, Ord a) => [a] -> [a] -> [[a]] -> [a] -> a +mvNormalCholLogDensity mu sigma l yObs + | k == 0 = 0 + | length yObs /= k || length sigma /= k = negInf + | length l /= k || any ((/= k) . length) l = negInf + | otherwise = + let m = [ [ (sigma !! i) * ((l !! i) !! j) | j <- [0 .. k - 1] ] + | i <- [0 .. k - 1] ] + kA = fromIntegral k :: a + d = zipWith (-) yObs mu + z = forwardSub m d -- M z = d (M 下三角) + quad = sum (map (\zi -> zi * zi) z) + logDet = sum [ log ((m !! i) !! i) | i <- [0 .. k - 1] ] -- = 0.5 log|Σ| + in -0.5 * kA * log (2 * pi) - logDet - 0.5 * quad + where k = length mu + +-- | [日本語]: MvStudentT(ν, μ, Σ) の 1 観測 (k-vector) の log density。 +-- @logpdf(y) = log Γ((ν+k)/2) − log Γ(ν/2) − (k/2) log(νπ) − (1/2) log|Σ| +-- − ((ν+k)/2) log(1 + m²/ν)@、 +-- @m² = (y−μ)ᵀ Σ⁻¹ (y−μ)@ を Cholesky で評価。 +-- [English]: The log density of a single observation (k-vector) from +-- MvStudentT(ν, μ, Σ). +-- @logpdf(y) = log Γ((ν+k)/2) − log Γ(ν/2) − (k/2) log(νπ) − (1/2) log|Σ| +-- − ((ν+k)/2) log(1 + m²/ν)@, +-- evaluating @m² = (y−μ)ᵀ Σ⁻¹ (y−μ)@ via Cholesky. +{-# INLINABLE mvStudentTLogDensity #-} +mvStudentTLogDensity :: forall a. (Floating a, Ord a) + => a -> [a] -> [[a]] -> [a] -> a +mvStudentTLogDensity nu mu cov yObs + | nu <= 0 = negInf + | length mu == 0 = 0 + | length yObs /= length mu = negInf + | otherwise = + case choleskyL cov of + Nothing -> negInf + Just l -> + let k = length mu + kA = fromIntegral k :: a + d = zipWith (-) yObs mu + z = forwardSub l d + quad = sum (map (\zi -> zi * zi) z) + logDet = 2 * sum [ log ((l !! i) !! i) | i <- [0 .. k - 1] ] + in lgammaApprox ((nu + kA) / 2) + - lgammaApprox (nu / 2) + - 0.5 * kA * log (nu * pi) + - 0.5 * logDet + - 0.5 * (nu + kA) * log (1 + quad / nu) + +-- | [日本語]: DirichletMultinomial(n, α) の 1 観測 (K-vector counts) の log pmf。 +-- @logpmf = log Γ(α₀) − log Γ(α₀+n) + Σ [log Γ(yᵢ+αᵢ) − log Γ(αᵢ)] +-- + log n! − Σ log yᵢ!@、 @α₀ = Σ αᵢ@. +-- [English]: The log pmf of a single observation (K-vector counts) from +-- DirichletMultinomial(n, α). +-- @logpmf = log Γ(α₀) − log Γ(α₀+n) + Σ [log Γ(yᵢ+αᵢ) − log Γ(αᵢ)] +-- + log n! − Σ log yᵢ!@, where @α₀ = Σ αᵢ@. +{-# INLINABLE dirichletMultinomialLogDensity #-} +dirichletMultinomialLogDensity :: forall a. (Floating a, Ord a) + => Int -> [a] -> [Double] -> a +dirichletMultinomialLogDensity n alpha counts + | length alpha /= length counts = negInf + | sum (map round counts :: [Int]) /= n = negInf + | any (< 0) counts = negInf + | any (\al -> al <= 0) alpha = negInf + | otherwise = + let nA = realToFrac (fromIntegral n :: Double) :: a + a0 = sum alpha + logFactN = realToFrac (logFactorial n) :: a + logFactSum = sum + [ realToFrac (logFactorial (round c :: Int)) :: a | c <- counts ] + term = sum + [ lgammaApprox (realToFrac c + ai) -- yᵢ + αᵢ + - lgammaApprox ai + | (c, ai) <- zip counts alpha + ] + in lgammaApprox a0 + - lgammaApprox (a0 + nA) + + term + + logFactN + - logFactSum + +-- | [日本語]: Wishart(ν, V) の 1 観測 (k×k 行列を flatten した長さ k² の列) の +-- log density。 +-- @logpdf(W) = -(νk/2) log 2 - (ν/2) log|V| - log Γ_k(ν/2) +-- + ((ν-k-1)/2) log|W| - (1/2) tr(V⁻¹ W)@、 +-- @log Γ_k(z) = (k(k-1)/4) log π + Σ_{i=1}^k log Γ((z+1-i)/2)@。 +-- V / W の Cholesky で log determinant と tr(V⁻¹ W) を評価。 +-- [English]: The log density of a single observation (a k×k matrix, +-- flattened to a length-k² sequence) from Wishart(ν, V). +-- @logpdf(W) = -(νk/2) log 2 - (ν/2) log|V| - log Γ_k(ν/2) +-- + ((ν-k-1)/2) log|W| - (1/2) tr(V⁻¹ W)@, +-- @log Γ_k(z) = (k(k-1)/4) log π + Σ_{i=1}^k log Γ((z+1-i)/2)@. +-- Evaluates the log determinant and tr(V⁻¹ W) via the Cholesky +-- factorization of V / W. +{-# INLINABLE wishartLogDensity #-} +wishartLogDensity :: forall a. (Floating a, Ord a) + => a -> [[a]] -> [a] -> a +wishartLogDensity nu vRows wFlat + | nu <= fromIntegral (k - 1) = negInf + | length wFlat /= k * k = negInf + | otherwise = + case (choleskyL vRows, choleskyL wRows) of + (Just lV, Just lW) -> + let logDetV = 2 * sum [ log ((lV !! i) !! i) | i <- [0 .. k - 1] ] + logDetW = 2 * sum [ log ((lW !! i) !! i) | i <- [0 .. k - 1] ] + -- tr(V⁻¹ W) を列ごとに solve V z_j = w_j で計算 + wCols = [ [ (wRows !! i) !! j | i <- [0 .. k - 1] ] + | j <- [0 .. k - 1] ] + solveV b = + let y = forwardSub lV b + x = backSubLT lV y -- Lᵀ x = y + in x + traceVW = sum [ solveV (wCols !! j) !! j + | j <- [0 .. k - 1] ] + kA = fromIntegral k :: a + -- log Γ_k(ν/2) + logMvGam = + (kA * (kA - 1) / 4) * log pi + + sum [ lgammaApprox ((nu + 1 - fromIntegral i) / 2) + | i <- [1 .. k] ] + in -(nu * kA / 2) * log 2 + - (nu / 2) * logDetV + - logMvGam + + ((nu - kA - 1) / 2) * logDetW + - 0.5 * traceVW + _ -> negInf + where + k = length vRows + wRows = chunksOf k wFlat + +-- --------------------------------------------------------------------------- +-- 多相 CDF / log-CDF (Truncated / Censored 用) +-- --------------------------------------------------------------------------- + +-- | [日本語]: 多相 erf 近似 (Abramowitz & Stegun 7.1.26)。誤差 < 1.5e-7。 +-- AD でも Track でも動く。 +-- [English]: A polymorphic erf approximation (Abramowitz & Stegun +-- 7.1.26). Error < 1.5e-7. Works with either AD or Track. +{-# INLINABLE erfA #-} +erfA :: (Floating a, Ord a) => a -> a +erfA x = + let p = 0.3275911 + a1 = 0.254829592 + a2 = -0.284496736 + a3 = 1.421413741 + a4 = -1.453152027 + a5 = 1.061405429 + sgn = if x < 0 then -1 else 1 + ax = abs x + t = 1 / (1 + p * ax) + poly = a1*t + a2*t*t + a3*t*t*t + a4*t*t*t*t + a5*t*t*t*t*t + in sgn * (1 - poly * exp (- ax * ax)) + +-- | [日本語]: 標準正規 CDF Φ(x)。 +-- [English]: The standard normal CDF Φ(x). +{-# INLINABLE phiCdfA #-} +phiCdfA :: (Floating a, Ord a) => a -> a +phiCdfA x = 0.5 * (1 + erfA (x / sqrt 2)) + +-- | CDF @F(x) = P(Y ≤ x)@ of a 'Distribution'. Returns 'Nothing' for +-- distributions that do not have a closed-form CDF in this library. +{-# INLINABLE distCDF #-} +distCDF :: (Floating a, Ord a) => Distribution a -> a -> Maybe a +distCDF (Normal mu sig) x + | sig <= 0 = Nothing + | otherwise = Just (phiCdfA ((x - mu) / sig)) +distCDF (Exponential rate) x + | rate <= 0 = Nothing + | x <= 0 = Just 0 + | otherwise = Just (1 - exp (-rate * x)) +distCDF (LogNormal mu sig) x + | sig <= 0 || x <= 0 = Nothing + | otherwise = Just (phiCdfA ((log x - mu) / sig)) +distCDF (Uniform lo hi) x + | hi <= lo = Nothing + | x <= lo = Just 0 + | x >= hi = Just 1 + | otherwise = Just ((x - lo) / (hi - lo)) +distCDF (HalfNormal sig) x + | sig <= 0 = Nothing + | x <= 0 = Just 0 + | otherwise = Just (erfA (x / (sig * sqrt 2))) +distCDF (HalfCauchy sc) x + | sc <= 0 = Nothing + | x <= 0 = Just 0 + | otherwise = Just (2 * atan (x / sc) / pi) +distCDF (Cauchy loc sc) x + | sc <= 0 = Nothing + | otherwise = Just (0.5 + atan ((x - loc) / sc) / pi) +distCDF (Gamma shape rate) x + | shape <= 0 || rate <= 0 = Nothing + | x <= 0 = Just 0 + | otherwise = Just (incGammaPA shape (rate * x)) +distCDF (Beta a b) x + | a <= 0 || b <= 0 = Nothing + | x <= 0 = Just 0 + | x >= 1 = Just 1 + | otherwise = Just (incBetaA x a b) +distCDF (StudentT df mu sig) x + | df <= 0 || sig <= 0 = Nothing + | otherwise = + let z = (x - mu) / sig + -- F_t(z; df) = 1 - 0.5 * I(df/(df+z²); df/2, 1/2) (z >= 0) + -- = 0.5 * I(df/(df+z²); df/2, 1/2) (z < 0) + ratio = df / (df + z * z) + ix = incBetaA ratio (df / 2) 0.5 + in Just (if z >= 0 then 1 - 0.5 * ix else 0.5 * ix) +distCDF (Logistic mu s) x + | s <= 0 = Nothing + | otherwise = Just (1 / (1 + exp (-((x - mu) / s)))) +distCDF (Gumbel mu beta) x + | beta <= 0 = Nothing + | otherwise = Just (exp (- exp (-((x - mu) / beta)))) +distCDF (AsymmetricLaplace b kappa mu) x + | b <= 0 || kappa <= 0 = Nothing + | otherwise = + let k2 = kappa * kappa + pc = k2 / (1 + k2) -- F(μ) + d = x - mu + in if d < 0 + then Just (pc * exp ((b / kappa) * d)) + else Just (1 - (1 - pc) * exp (- b * kappa * d)) +distCDF _ _ = Nothing -- SkewNormal / 離散・Mixture・Truncated 内 Truncated 等は未対応 + +-- | @log F(x)@. Computed as @log(F)@ directly to avoid loss of +-- precision near the tails where @F@ approaches 0 or 1. +{-# INLINABLE logCDF #-} +logCDF :: (Floating a, Ord a) => Distribution a -> a -> a +logCDF d x = case distCDF d x of + Nothing -> negInf + Just c | c <= 0 -> negInf + | c >= 1 -> 0 + | otherwise -> log c + +-- | Log of the right-tail survival function @log(1 − F(x))@. +{-# INLINABLE logSF #-} +logSF :: (Floating a, Ord a) => Distribution a -> a -> a +logSF d x = case distCDF d x of + Nothing -> negInf + Just c | c <= 0 -> 0 + | c >= 1 -> negInf + | otherwise -> log (1 - c) + +-- | [日本語]: log(F(hi) − F(lo)) — Truncated の正規化定数。 +-- [English]: log(F(hi) − F(lo)) — the normalization constant for +-- Truncated. +{-# INLINABLE logCDFInterval #-} +logCDFInterval :: (Floating a, Ord a) => Distribution a -> Maybe a -> Maybe a -> a +logCDFInterval d mLo mHi = case (mLo, mHi) of + (Nothing, Nothing) -> 0 -- log(1) + (Just lo, Nothing) -> logSF d lo + (Nothing, Just hi) -> logCDF d hi + (Just lo, Just hi) -> + case (distCDF d lo, distCDF d hi) of + (Just cl, Just ch) + | ch <= cl -> negInf + | otherwise -> log (ch - cl) + _ -> negInf + +-- --------------------------------------------------------------------------- +-- 多相 log 密度 (事前 logDensity + 観測 logDensityObs/obsLogSum) +-- --------------------------------------------------------------------------- +-- Phase 58.6: 元 HBM.hs の「事前 log 密度」節 (logDensity は 58.3 で AD 勾配と +-- 同居のため残置していたが、 logJoint/logPrior が参照するため Eval 抽出 (58.6c) で +-- back-edge になる。 密度は本来 Distribution の責務 (Phase 58 計画の module sketch) +-- ゆえここへ集約する。 INLINABLE は AD 経路の cross-module inlining 維持のため保持。 + +-- | Log prior density at a sample value of type @a@. +{-# INLINABLE logDensity #-} +logDensity :: (Floating a, Ord a) => Distribution a -> a -> a +logDensity (Normal mu sig) x + | sig <= 0 = negInf + | otherwise = -0.5 * log (2 * pi) - log sig + - 0.5 * ((x - mu) / sig) ^ (2::Int) +logDensity (Exponential rate) x + | x < 0 || rate <= 0 = negInf + | otherwise = log rate - rate * x +logDensity (Gamma shape rate) x + | x <= 0 || shape <= 0 || rate <= 0 = negInf + | otherwise = + (shape - 1) * log x - rate * x + + shape * log rate - lgammaApprox shape +logDensity (Beta alpha beta) x + | x <= 0 || x >= 1 || alpha <= 0 || beta <= 0 = negInf + | otherwise = + (alpha - 1) * log x + (beta - 1) * log (1 - x) + - (lgammaApprox alpha + lgammaApprox beta - lgammaApprox (alpha + beta)) +logDensity (Poisson lam) x + | lam <= 0 = negInf + | x < 0 = negInf + | otherwise = + -- x はサンプル値なので連続として扱う (整数化はしない) + x * log lam - lam +logDensity (Binomial _ p) _ + | p <= 0 || p >= 1 = negInf + | otherwise = 0 -- サンプル時は使わない (構造のみ) +logDensity (Uniform lo hi) x + | hi <= lo = negInf + | x < lo || x > hi = negInf + | otherwise = -log (hi - lo) +logDensity (StudentT df mu sig) x + | df <= 0 || sig <= 0 = negInf + | otherwise = + let z = (x - mu) / sig + in lgammaApprox ((df + 1) / 2) + - lgammaApprox (df / 2) + - 0.5 * log (df * pi) + - log sig + - ((df + 1) / 2) * log (1 + z * z / df) +logDensity (Cauchy loc sc) x + | sc <= 0 = negInf + | otherwise = + let z = (x - loc) / sc + in -log pi - log sc - log (1 + z * z) +logDensity (HalfNormal sig) x + | sig <= 0 = negInf + | x < 0 = negInf + | otherwise = + 0.5 * log 2 - 0.5 * log pi - log sig + - 0.5 * (x / sig) ^ (2::Int) +logDensity (HalfCauchy sc) x + | sc <= 0 = negInf + | x < 0 = negInf + | otherwise = + log 2 - log pi - log sc - log (1 + (x / sc) ^ (2::Int)) +logDensity (LogNormal mu sig) x + | sig <= 0 = negInf + | x <= 0 = negInf + | otherwise = + let lx = log x + in -0.5 * log (2 * pi) - log sig - lx + - 0.5 * ((lx - mu) / sig) ^ (2::Int) +logDensity (Bernoulli p) _ + | p <= 0 || p >= 1 = negInf + | otherwise = 0 -- 構造のみ (離散なので連続 prior 評価には使わない) +logDensity (Categorical _) _ = 0 -- 同上 +logDensity (Mixture ws comps) x + | null ws || length ws /= length comps = negInf + | otherwise = + let total = sum ws + logTotal = log total + -- log(w_k / Σw) + log p_k(x) + logTerms = zipWith (\w d -> log w - logTotal + logDensity d x) ws comps + in logSumExpA logTerms +logDensity (Truncated d mLo mHi) x = + -- 範囲外なら 0 (=> log で −∞) + let outOfRange = case (mLo, mHi) of + (Just lo, _ ) | x < lo -> True + (_, Just hi) | x > hi -> True + _ -> False + in if outOfRange + then negInf + else logDensity d x - logCDFInterval d mLo mHi +logDensity (Censored d _ _) x = + -- prior 評価では通常の密度を使う (打ち切りは観測時のみ意味を持つ) + logDensity d x +logDensity MvNormal{} _ = 0 -- observation-only: latent としては使わない +logDensity MvNormalChol{} _ = 0 -- observation-only +logDensity MvNormalGpRBF{} _ = 0 -- observation-only (Phase 95 B-dsl) +logDensity HmmForwardNormal{} _ = 0 -- observation-only (Phase 92 A2) +logDensity ArmaNormal{} _ = 0 -- observation-only (Phase 101 A2) +logDensity GradedResponseIrt{} _ = 0 -- observation-only (Phase 101 A3) +logDensity Multinomial{} _ = 0 -- observation-only +logDensity (InverseGamma alpha beta) x + | alpha <= 0 || beta <= 0 || x <= 0 = negInf + | otherwise = + alpha * log beta - lgammaApprox alpha + - (alpha + 1) * log x - beta / x +logDensity (Weibull kShape lam) x + | kShape <= 0 || lam <= 0 || x <= 0 = negInf + | otherwise = + log kShape - log lam + + (kShape - 1) * (log x - log lam) + - (x / lam) ** kShape +logDensity (Pareto alpha xm) x + | alpha <= 0 || xm <= 0 || x < xm = negInf + | otherwise = + log alpha + alpha * log xm - (alpha + 1) * log x +logDensity BetaBinomial{} _ = 0 -- 観測専用 (離散) +logDensity (VonMises mu kappa) x + | kappa <= 0 = negInf + | otherwise = + kappa * cos (x - mu) + - log (2 * pi) + - logBesselI0 kappa +logDensity (ZeroInflatedPoisson psi lam) x + | psi < 0 || psi > 1 || lam <= 0 || x < 0 = negInf + | x == 0 = + -- log(ψ + (1-ψ) e^{-λ}) + logSumExpA [log psi, log (1 - psi) - lam] + | otherwise = + -- log(1-ψ) + Poisson logpmf + log (1 - psi) + x * log lam - lam - lgammaApprox (x + 1) +logDensity (ZeroInflatedBinomial n psi p) x + | psi < 0 || psi > 1 || p <= 0 || p >= 1 || x < 0 = negInf + | otherwise = + let nA = realToFrac (fromIntegral n :: Double) + -- log(C(n,k)) = lgamma(n+1) - lgamma(k+1) - lgamma(n-k+1) (多相) + logC = lgammaApprox (nA + 1) + - lgammaApprox (x + 1) + - lgammaApprox (nA - x + 1) + in if x == 0 + then logSumExpA [log psi + , log (1 - psi) + nA * log (1 - p)] + else log (1 - psi) + + logC + x * log p + (nA - x) * log (1 - p) +logDensity (NegativeBinomial mu alpha) x + | mu <= 0 || alpha <= 0 || x < 0 = negInf + | otherwise = + let p = alpha / (alpha + mu) -- success prob + in lgammaApprox (x + alpha) + - lgammaApprox alpha + - lgammaApprox (x + 1) + + alpha * log p + + x * log (1 - p) +logDensity (SkewNormal mu sig alpha) x + | sig <= 0 = negInf + | otherwise = + let z = (x - mu) / sig + logPhi = -0.5 * log (2 * pi) - 0.5 * z * z + -- log Φ(αz) を phiCdfA 経由で。 引数が大きく負だと数値的に困るが、 + -- phiCdfA は erfA ベースなので clip して log を取る + cdfArg = phiCdfA (alpha * z) + -- 数値下限 1e-300 程度に防御 + logCdf = log (max cdfArg 1e-300) + in log 2 - log sig + logPhi + logCdf +logDensity (Logistic mu s) x + | s <= 0 = negInf + | otherwise = + let z = (x - mu) / s + in -z - log s - 2 * log (1 + exp (-z)) +logDensity (Gumbel mu beta) x + | beta <= 0 = negInf + | otherwise = + let z = (x - mu) / beta + in -log beta - z - exp (-z) +logDensity (AsymmetricLaplace b kappa mu) x + | b <= 0 || kappa <= 0 = negInf + | otherwise = + let logNorm = log b - log (kappa + 1 / kappa) + d = x - mu + in if d >= 0 + then logNorm - b * kappa * d + else logNorm + (b / kappa) * d +-- 離散分布は構造のみ (observation-only の意味で logDensity は使われない) +logDensity OrderedLogistic{} _ = 0 +logDensity DiscreteUniform{} _ = 0 +logDensity (Geometric p) _ + | p <= 0 || p >= 1 = negInf + | otherwise = 0 +logDensity HyperGeometric{} _ = 0 +logDensity (ZeroInflatedNegativeBinomial psi mu alpha) _ + | psi < 0 || psi > 1 || mu <= 0 || alpha <= 0 = negInf + | otherwise = 0 +logDensity MvStudentT{} _ = 0 -- observation-only +logDensity DirichletMultinomial{} _ = 0 -- observation-only +logDensity (Triangular lo c hi) x + | hi <= lo || c < lo || c > hi = negInf + | x < lo || x > hi = negInf + | x <= c = + log 2 + log (x - lo) + - log (hi - lo) - log (c - lo) + | otherwise = + log 2 + log (hi - x) + - log (hi - lo) - log (hi - c) +logDensity (Kumaraswamy a b) x + | a <= 0 || b <= 0 || x <= 0 || x >= 1 = negInf + | otherwise = + let xa = x ** a + in log a + log b + (a - 1) * log x + (b - 1) * log (1 - xa) +logDensity (Rice nu sig) x + | sig <= 0 || nu < 0 || x < 0 = negInf + | otherwise = + let s2 = sig * sig + z = x * nu / s2 + in log x - 2 * log sig - (x * x + nu * nu) / (2 * s2) + + logBesselI0 z +logDensity DiscreteWeibull{} _ = 0 -- 離散: structure only +logDensity Wishart{} _ = 0 -- observation-only (k×k 行列観測) +logDensity (Bound d mLo mHi) x = logDensity (Truncated d mLo mHi) x +logDensity OrderedProbit{} _ = 0 -- observation-only (離散) + +-- | [日本語]: @logDensity@ の AD ('ADRD.ReverseDouble') 特化版。 +-- hyperparameter が__定数__ ('ADRD.Zero' / 'ADRD.Lift' = tape 由来でない) の +-- lgamma 正規化項を Double で 1 発計算して 'ADRD.Lift' で戻す。 'ADRD.Lift' +-- 同士の AD 演算は @Lift (f b c)@ (同一の Double 演算列・tape 追記なし) なので +-- 結果は generic @logDensity@ と __bit-identical__、 定数に勾配は流れないので +-- 微分も不変。 hyperparameter が tape 変数 (階層 prior) なら generic へ +-- fallback し勾配は AD がそのまま構成する。 +-- +-- 動機 (hmm reduced prof): Dirichlet(1,…,1) = 棒折り Beta(1,1) の定数濃度 +-- lgamma が AD walk 上で毎 eval Stirling recurrence (z<12 の梯子 ~11 段 × +-- lgamma 3 呼び出し) を boxed 'ADRD.Lift' で歩いていた (550,480 entries = +-- 70 call/eval・time 6.6%/alloc 14.4%)。 対象は lgamma を持つ定数 prior 3 種 +-- (Beta / Gamma / StudentT の ν) のみ・折り畳み式の結合順は generic 実装と +-- 完全一致させてある (bit 一致の根拠)。 +-- ※ 'lgammaApprox' への RULES 書き換えは過負荷関数 + 辞書引数で発火せず断念 +-- (実測)、 呼び出し点注入 (@logPriorWith@) 方式にした。 +-- [English]: The AD ('ADRD.ReverseDouble')-specialized version of +-- @logDensity@. +-- +-- When a hyperparameter is __constant__ ('ADRD.Zero' / 'ADRD.Lift' — not +-- tape-derived), computes its lgamma normalization term once as a +-- Double and wraps it back with 'ADRD.Lift'. AD operations between +-- 'ADRD.Lift' values reduce to @Lift (f b c)@ (the same sequence of +-- Double operations, no tape entries appended), so the result is +-- __bit-identical__ to the generic @logDensity@, and since no gradient +-- flows through a constant, the derivative is unaffected either. When a +-- hyperparameter is a tape variable (hierarchical prior), this falls +-- back to the generic path and AD constructs the gradient as usual. +-- +-- Motivation (from a reduced-model hmm profile): the constant +-- concentration lgamma of Dirichlet(1,…,1) = stick-breaking Beta(1,1) +-- was walking the Stirling recurrence (a ladder of ~11 steps for z<12, +-- × 3 lgamma calls) through boxed 'ADRD.Lift' on every eval of the AD +-- walk (550,480 entries = 70 calls/eval, 6.6% time / 14.4% alloc). The +-- targets are only the 3 constant priors that have an lgamma (Beta / +-- Gamma / StudentT's ν); the associativity order of the folded +-- expression is made to exactly match the generic implementation (the +-- basis for bit-identical results). +-- Note: rewriting 'lgammaApprox' via RULES was abandoned since it did +-- not fire for the overloaded function with a dictionary argument +-- (measured), so injection at the call site (@logPriorWith@) was used +-- instead. +logDensityRD + :: forall s. Reifies s ADRD.Tape + => Distribution (ADRD.ReverseDouble s) -> ADRD.ReverseDouble s + -> ADRD.ReverseDouble s +logDensityRD d x = case d of + Beta a b + | Just a' <- constRD a, Just b' <- constRD b + , not (x <= 0 || x >= 1 || a' <= 0 || b' <= 0) -> + (a - 1) * log x + (b - 1) * log (1 - x) + - ADRD.Lift (lgammaApprox a' + lgammaApprox b' - lgammaApprox (a' + b')) + Gamma sh ra + | Just sh' <- constRD sh, Just ra' <- constRD ra + , not (x <= 0 || sh' <= 0 || ra' <= 0) -> + (sh - 1) * log x - ra * x + + ADRD.Lift (sh' * log ra') - ADRD.Lift (lgammaApprox sh') + StudentT df mu sig + | Just df' <- constRD df + , not (df' <= 0 || sig <= 0) -> + let z = (x - mu) / sig + in ADRD.Lift (lgammaApprox ((df' + 1) / 2) - lgammaApprox (df' / 2) + - 0.5 * log (df' * pi)) + - log sig + - ((df + 1) / 2) * log (1 + z * z / df) + _ -> logDensity d x + where + -- Zero/Lift = tape に乗らない定数 (Lift 同士の演算は Lift に閉じる) + constRD :: ADRD.ReverseDouble s -> Maybe Double + constRD ADRD.Zero = Just 0 + constRD (ADRD.Lift v) = Just v + constRD ADRD.ReverseDouble{} = Nothing + +-- | [日本語]: 観測点における log likelihood density (固定 @Double@)。 +-- 観測値は @[Double]@ で渡されるため、 ここでは @Floating a@ 制約のみを使う。 +-- +-- ObserveLM 評価 (lmObsLogLiks) と logJoint が AD で微分しながら呼ぶ +-- ホット経路。 本体から移したため cross-module になった。 INLINABLE で +-- 境界跨ぎ inline を維持 (M1/M2 の +25% 劣化を解消・bench 実測)。 +-- [English]: Log likelihood density at an observation (a fixed +-- @Double@). Observations are passed as @[Double]@, so this uses only +-- the @Floating a@ constraint. +-- +-- A hot path called by ObserveLM evaluation (lmObsLogLiks) and logJoint +-- while differentiating via AD. It became cross-module after being +-- moved out of the main body. INLINABLE keeps the cross-module inline +-- (resolving M1/M2's +25% degradation, measured by bench). +{-# INLINABLE logDensityObs #-} +logDensityObs :: forall a. (Floating a, Ord a) => Distribution a -> Double -> a +logDensityObs (Normal mu sig) y + | sig <= 0 = negInf + | otherwise = + let yA = realToFrac y :: a + in -0.5 * log (2 * pi) - log sig - 0.5 * ((yA - mu) / sig) ^ (2::Int) +logDensityObs (Exponential rate) y + | y < 0 = negInf + | rate <= 0 = negInf + | otherwise = log rate - rate * (realToFrac y :: a) +logDensityObs (Gamma shape rate) y + | y <= 0 = negInf + | shape <= 0 || rate <= 0 = negInf + | otherwise = + let yA = realToFrac y :: a + in (shape - 1) * log yA - rate * yA + + shape * log rate - lgammaApprox shape +logDensityObs (Beta alpha beta) y + | y <= 0 || y >= 1 || alpha <= 0 || beta <= 0 = negInf + | otherwise = + let yA = realToFrac y :: a + in (alpha - 1) * log yA + (beta - 1) * log (1 - yA) + - (lgammaApprox alpha + lgammaApprox beta - lgammaApprox (alpha + beta)) +logDensityObs (Poisson lam) y + | lam <= 0 = negInf + | y < 0 = negInf + | otherwise = + let kA = realToFrac y :: a + kInt = round y :: Int + logFactK = realToFrac (logFactorial kInt) :: a + in kA * log lam - lam - logFactK +logDensityObs (Binomial n p) y + | p <= 0 || p >= 1 = negInf + | otherwise = + let k = round y :: Int + kA = realToFrac y :: a + nA = realToFrac (fromIntegral n :: Double) :: a + logC = realToFrac (logBinomCoeff n k) :: a + in logC + kA * log p + (nA - kA) * log (1 - p) +logDensityObs (Uniform lo hi) y + | hi <= lo = negInf + | otherwise = + let yA = realToFrac y :: a + in if yA < lo || yA > hi then negInf else -log (hi - lo) +logDensityObs (StudentT df mu sig) y + | df <= 0 || sig <= 0 = negInf + | otherwise = + let yA = realToFrac y :: a + z = (yA - mu) / sig + in lgammaApprox ((df + 1) / 2) + - lgammaApprox (df / 2) + - 0.5 * log (df * pi) + - log sig + - ((df + 1) / 2) * log (1 + z * z / df) +logDensityObs (Cauchy loc sc) y + | sc <= 0 = negInf + | otherwise = + let yA = realToFrac y :: a + z = (yA - loc) / sc + in -log pi - log sc - log (1 + z * z) +logDensityObs (HalfNormal sig) y + | sig <= 0 = negInf + | y < 0 = negInf + | otherwise = + let yA = realToFrac y :: a + in 0.5 * log 2 - 0.5 * log pi - log sig + - 0.5 * (yA / sig) ^ (2::Int) +logDensityObs (HalfCauchy sc) y + | sc <= 0 = negInf + | y < 0 = negInf + | otherwise = + let yA = realToFrac y :: a + in log 2 - log pi - log sc - log (1 + (yA / sc) ^ (2::Int)) +logDensityObs (LogNormal mu sig) y + | sig <= 0 = negInf + | y <= 0 = negInf + | otherwise = + let yA = realToFrac y :: a + lx = log yA + in -0.5 * log (2 * pi) - log sig - lx + - 0.5 * ((lx - mu) / sig) ^ (2::Int) +logDensityObs (Bernoulli p) y + | p <= 0 || p >= 1 = negInf + | otherwise = + let k = round y :: Int + in case k of + 1 -> log p + 0 -> log (1 - p) + _ -> negInf +logDensityObs (Categorical probs) y = + let k = round y :: Int + n = length probs + in if k < 0 || k >= n + then negInf + else + -- log p_k - log(Σ p_i) (probs を正規化) + let pk = probs !! k + total = sum probs + in if pk <= 0 || total <= 0 + then negInf + else log pk - log total +logDensityObs (Mixture ws comps) y + | null ws || length ws /= length comps = negInf + | otherwise = + let total = sum ws + logTotal = log total + logTerms = zipWith (\w d -> log w - logTotal + logDensityObs d y) ws comps + in logSumExpA logTerms +logDensityObs (Truncated d mLo mHi) y = + let yA = realToFrac y :: a + outOfRange = case (mLo, mHi) of + (Just lo, _ ) | yA < lo -> True + (_, Just hi) | yA > hi -> True + _ -> False + in if outOfRange + then negInf + else logDensityObs d y - logCDFInterval d mLo mHi +logDensityObs (Censored d mLo mHi) y = + -- 観測値 y が境界 lo / hi に等しい場合は左/右打ち切り尤度 + let yA = realToFrac y :: a + eps = 1e-9 :: a + isAt v target = abs (v - target) < eps + in case (mLo, mHi) of + (Just lo, _) | yA <= lo || isAt yA lo -> logCDF d lo -- 左打ち切り + (_, Just hi) | yA >= hi || isAt yA hi -> logSF d hi -- 右打ち切り + _ -> logDensityObs d y -- 通常観測 +logDensityObs MvNormal{} _ = 0 +logDensityObs MvNormalChol{} _ = 0 +logDensityObs MvNormalGpRBF{} _ = 0 -- Phase 95 B-dsl: obsLogSum 経由 (下と同じ) +logDensityObs HmmForwardNormal{} _ = 0 -- Phase 92 A2: obsLogSum 経由 (下と同じ) +logDensityObs ArmaNormal{} _ = 0 -- Phase 101 A2: obsLogSum 経由 (下と同じ) +logDensityObs GradedResponseIrt{} _ = 0 -- Phase 101 A3: obsLogSum 経由 (下と同じ) + -- スカラー観測経路では使わない (chunk して 'mvNormalLogDensity' を呼ぶ obsLogSum 経由) +logDensityObs Multinomial{} _ = 0 + -- スカラー観測経路では使わない (k 次元 chunk で multinomialLogDensity を呼ぶ) +logDensityObs (InverseGamma alpha beta) y + | alpha <= 0 || beta <= 0 || y <= 0 = negInf + | otherwise = + let yA = realToFrac y :: a + in alpha * log beta - lgammaApprox alpha + - (alpha + 1) * log yA - beta / yA +logDensityObs (Weibull kShape lam) y + | kShape <= 0 || lam <= 0 || y <= 0 = negInf + | otherwise = + let yA = realToFrac y :: a + in log kShape - log lam + + (kShape - 1) * (log yA - log lam) + - (yA / lam) ** kShape +logDensityObs (Pareto alpha xm) y + | alpha <= 0 || xm <= 0 = negInf + | otherwise = + let yA = realToFrac y :: a + in if yA < xm + then negInf + else log alpha + alpha * log xm - (alpha + 1) * log yA +logDensityObs (BetaBinomial n alpha beta) y + | alpha <= 0 || beta <= 0 || y < 0 = negInf + | otherwise = + let yA = realToFrac y :: a + nA = realToFrac (fromIntegral n :: Double) :: a + k = round y :: Int + logC = realToFrac (logBinomCoeff n k) :: a + in logC + + lgammaApprox (yA + alpha) + + lgammaApprox (nA - yA + beta) + - lgammaApprox (nA + alpha + beta) + - (lgammaApprox alpha + lgammaApprox beta - lgammaApprox (alpha + beta)) +logDensityObs (VonMises mu kappa) y + | kappa <= 0 = negInf + | otherwise = + let yA = realToFrac y :: a + in kappa * cos (yA - mu) - log (2 * pi) - logBesselI0 kappa +logDensityObs (ZeroInflatedPoisson psi lam) y + | psi < 0 || psi > 1 || lam <= 0 || y < 0 = negInf + | y == 0 = + logSumExpA [log psi, log (1 - psi) - lam] + | otherwise = + let kA = realToFrac y :: a + kInt = round y :: Int + logFactK = realToFrac (logFactorial kInt) :: a + in log (1 - psi) + kA * log lam - lam - logFactK +logDensityObs (ZeroInflatedBinomial n psi p) y + | psi < 0 || psi > 1 || p <= 0 || p >= 1 || y < 0 = negInf + | otherwise = + let kA = realToFrac y :: a + k = round y :: Int + nA = realToFrac (fromIntegral n :: Double) :: a + logC = realToFrac (logBinomCoeff n k) :: a + in if y == 0 + then logSumExpA [log psi + , log (1 - psi) + nA * log (1 - p)] + else log (1 - psi) + + logC + kA * log p + (nA - kA) * log (1 - p) +logDensityObs (NegativeBinomial mu alpha) y + | mu <= 0 || alpha <= 0 || y < 0 = negInf + | otherwise = + let kA = realToFrac y :: a + p = alpha / (alpha + mu) + in lgammaApprox (kA + alpha) + - lgammaApprox alpha + - lgammaApprox (kA + 1) + + alpha * log p + + kA * log (1 - p) +logDensityObs (SkewNormal mu sig alpha) y + | sig <= 0 = negInf + | otherwise = + let yA = realToFrac y :: a + z = (yA - mu) / sig + logPhi = -0.5 * log (2 * pi) - 0.5 * z * z + cdfArg = phiCdfA (alpha * z) + logCdf = log (max cdfArg 1e-300) + in log 2 - log sig + logPhi + logCdf +logDensityObs (Logistic mu s) y + | s <= 0 = negInf + | otherwise = + let yA = realToFrac y :: a + z = (yA - mu) / s + in -z - log s - 2 * log (1 + exp (-z)) +logDensityObs (Gumbel mu beta) y + | beta <= 0 = negInf + | otherwise = + let yA = realToFrac y :: a + z = (yA - mu) / beta + in -log beta - z - exp (-z) +logDensityObs (AsymmetricLaplace b kappa mu) y + | b <= 0 || kappa <= 0 = negInf + | otherwise = + let yA = realToFrac y :: a + logNorm = log b - log (kappa + 1 / kappa) + d = yA - mu + in if d >= 0 + then logNorm - b * kappa * d + else logNorm + (b / kappa) * d +logDensityObs (OrderedLogistic eta cuts) y + | null cuts = negInf + | k < 0 || k > kMax = negInf + | otherwise = + -- σ(c_{k+1} − η) − σ(c_k − η)、 c_0 = −∞、 c_K = +∞ + let sigm x = 1 / (1 + exp (-x)) + kMax_a = kMax -- 上限カテゴリ index + probHi + | k == kMax_a = 1 + | otherwise = sigm (cuts !! k - eta) + probLo + | k == 0 = 0 + | otherwise = sigm (cuts !! (k - 1) - eta) + pK = probHi - probLo + in if pK <= 0 then negInf else log pK + where + k = round y :: Int + kMax = length cuts +logDensityObs (DiscreteUniform lo hi) y + | hi < lo = negInf + | yI < lo || yI > hi = negInf + | otherwise = -log (realToFrac (hi - lo + 1) :: a) + where + yI = round y :: Int +logDensityObs (Geometric p) y + | p <= 0 || p >= 1 = negInf + | yI < 1 = negInf + | otherwise = + let kA = realToFrac y :: a + in (kA - 1) * log (1 - p) + log p + where + yI = round y :: Int +logDensityObs (HyperGeometric nN kK nDraw) y + | nN <= 0 || kK < 0 || kK > nN || nDraw < 0 || nDraw > nN = negInf + | yI < max 0 (nDraw + kK - nN) || yI > min nDraw kK = negInf + | otherwise = + let lc = realToFrac (logBinomCoeff kK yI + + logBinomCoeff (nN - kK) (nDraw - yI) + - logBinomCoeff nN nDraw) :: a + in lc + where + yI = round y :: Int +logDensityObs (ZeroInflatedNegativeBinomial psi mu alpha) y + | psi < 0 || psi > 1 || mu <= 0 || alpha <= 0 || y < 0 = negInf + | y == 0 = + -- log(ψ + (1-ψ) (α/(α+μ))^α) + let p0NB = alpha * (log alpha - log (alpha + mu)) + in logSumExpA [log psi, log (1 - psi) + p0NB] + | otherwise = + let kA = realToFrac y :: a + p = alpha / (alpha + mu) + logNB = lgammaApprox (kA + alpha) + - lgammaApprox alpha + - lgammaApprox (kA + 1) + + alpha * log p + + kA * log (1 - p) + in log (1 - psi) + logNB +logDensityObs MvStudentT{} _ = 0 + -- スカラー観測経路では使わない (k chunk で mvStudentTLogDensity 経由) +logDensityObs DirichletMultinomial{} _ = 0 + -- スカラー観測経路では使わない (K chunk で dirichletMultinomialLogDensity 経由) +logDensityObs (Triangular lo c hi) y + | hi <= lo || c < lo || c > hi = negInf + | otherwise = + let yA = realToFrac y :: a + in if yA < lo || yA > hi + then negInf + else if yA <= c + then log 2 + log (yA - lo) + - log (hi - lo) - log (c - lo) + else log 2 + log (hi - yA) + - log (hi - lo) - log (hi - c) +logDensityObs (Kumaraswamy a b) y + | a <= 0 || b <= 0 = negInf + | otherwise = + let yA = realToFrac y :: a + in if yA <= 0 || yA >= 1 + then negInf + else let xa = yA ** a + in log a + log b + (a - 1) * log yA + (b - 1) * log (1 - xa) +logDensityObs (Rice nu sig) y + | sig <= 0 || nu < 0 || y < 0 = negInf + | otherwise = + let yA = realToFrac y :: a + s2 = sig * sig + z = yA * nu / s2 + in log yA - 2 * log sig - (yA * yA + nu * nu) / (2 * s2) + + logBesselI0 z +logDensityObs Wishart{} _ = 0 + -- スカラー観測経路では使わない (k² chunk で wishartLogDensity 経由) +logDensityObs (Bound d mLo mHi) y = logDensityObs (Truncated d mLo mHi) y +logDensityObs (OrderedProbit eta cuts) y + | null cuts = negInf + | k < 0 || k > kMax = negInf + | otherwise = + let probHi + | k == kMax = 1 + | otherwise = phiCdfA (cuts !! k - eta) + probLo + | k == 0 = 0 + | otherwise = phiCdfA (cuts !! (k - 1) - eta) + pK = probHi - probLo + in if pK <= 0 then negInf else log pK + where + k = round y :: Int + kMax = length cuts +logDensityObs (DiscreteWeibull q beta) y + | y < 0 = negInf + | otherwise = + -- q は (0,1)、 β > 0 + -- pmf(k) = q^(k^β) - q^((k+1)^β) + let qVal :: a + qVal = q + bVal :: a + bVal = beta + in if qVal <= 0 || qVal >= 1 || bVal <= 0 + then negInf + else + let kI = round y :: Int + kA = realToFrac (fromIntegral kI :: Double) :: a + logQ = log qVal + -- log(q^(k^β) - q^((k+1)^β)) + -- = log q^(k^β) + log(1 - q^((k+1)^β - k^β)) + -- 安定化: a1 = (k+1)^β - k^β > 0 (β>0) + pk = kA ** bVal + pk1 = (kA + 1) ** bVal + diffP = pk1 - pk + -- log(1 - q^diffP) = log(1 - exp(diffP * logQ)) + -- diffP * logQ <= 0 + expArg = diffP * logQ + log1mE = log (1 - exp expArg) + in pk * logQ + log1mE + +-- | [日本語]: 観測値のリストに対する log likelihood の和。 通常の分布では、 +-- 1 観測が 1 スカラー log-density を寄与する。 'MvNormal' (@k@-vector を +-- 期待する) では、 flatten された @[Double]@ を評価前に長さ @k@ のグループに +-- chunk する。 +-- +-- logJoint/logLikelihood の Observe 分岐が AD で呼ぶ。 cross-module +-- inline 維持のため INLINABLE。 +-- [English]: Sum of log likelihoods over a list of observations. For +-- ordinary distributions one observation contributes one scalar +-- log-density. For 'MvNormal' (which expects @k@-vectors), the +-- flattened @[Double]@ is chunked into length-@k@ groups before +-- evaluation. +-- +-- Called by the Observe branch of logJoint/logLikelihood while +-- differentiating via AD. INLINABLE to keep the cross-module inline. +{-# INLINABLE obsLogSum #-} +obsLogSum :: forall a. (Floating a, Ord a) => Distribution a -> [Double] -> a +obsLogSum (MvNormal mu cov) ys = + let k = length mu + chunks = chunksOf k ys + in sum [ mvNormalLogDensity mu cov (map realToFrac yv :: [a]) + | yv <- chunks ] +obsLogSum (MvNormalGpRBF xs alpha rho sigma) ys = + -- Phase 95 B-dsl: zero-mean・cov = RBF カーネル + (1e-10 + σ)·I。 値は汎用 + -- 'MvNormal' 経路と同値 (ホット勾配のみ @gpRBFAnalyticVG@ で閉形式化)。 + let k = length xs + cov = gpRBFCovList xs alpha rho sigma + mu = replicate k 0 + chunks = chunksOf k ys + in sum [ mvNormalLogDensity mu cov (map realToFrac yv :: [a]) + | yv <- chunks ] +obsLogSum (GradedResponseIrt thetas ncats deltas gammas) ys = + -- Phase 101 A3: grade 行列 (nChild×nItem 行優先・欠測 −1) 全体を 1 観測として + -- 評価。 値は従来の @logCatProb + potential@ 書きと同値 + -- (ホット勾配のみ @gradedIrtAnalyticVG@ で閉形式化)。 + let nItem = length ncats + rows = chunksOf nItem ys + logCatP th nc dl gm gr = + let kMax = nc - 1 + qs = [ 1 / (1 + exp (negate (realToFrac dl * (th - realToFrac (gm !! (kk - 1)))))) + | kk <- [1 .. kMax] ] + ps = [ if k == 1 then 1 - head qs + else if k == nc then qs !! (kMax - 1) + else (qs !! (k - 2)) - (qs !! (k - 1)) + | k <- [1 .. nc] ] + in log (ps !! (gr - 1)) + in sum [ logCatP th nc dl gm (round gr) + | (th, row) <- zip thetas rows + , (nc, dl, gm, gr) <- zip4 ncats deltas gammas row + , gr /= -1 ] +obsLogSum (ArmaNormal mu phi theta sg) ys = + -- Phase 101 A2: 観測列全体 (長さ T) を 1 観測として err 逐次再帰で評価。 + -- 値は従来の @mapAccumL + potential@ 書きと同値 + -- (ホット勾配のみ @armaAnalyticVG@ で閉形式化)。 + case ys of + [] -> 0 + (y1 : rest) -> + let e1 = realToFrac y1 - (mu + phi * mu) + step (prevY, prevErr) yt = + let err = realToFrac yt - (mu + phi * realToFrac prevY + theta * prevErr) + in ((yt, err), err) + errs = e1 : snd (mapAccumL step (y1, e1) rest) + in sum [ logDensity (Normal 0 sg) e | e <- errs ] +obsLogSum (HmmForwardNormal pi0 trans mus sg) ys = + -- Phase 92 A2: 観測列全体 (長さ T) を 1 観測として forward algorithm で周辺化。 + -- 値は従来の @potential nm (hmmForwardLogLik pi0 trans emit)@ 書きと同値 + -- (ホット勾配のみ @hmmAnalyticVG@ で閉形式化)。 + let emit = [ [ logDensity (Normal mu sg) (realToFrac y) | mu <- mus ] | y <- ys ] + in hmmForwardLogLik pi0 trans emit +obsLogSum (Multinomial n probs) ys = + let k = length probs + chunks = chunksOf k ys + in sum [ multinomialLogDensity n probs yv | yv <- chunks ] +obsLogSum (MvNormalChol mu sigma l) ys = + let k = length mu + chunks = chunksOf k ys + in sum [ mvNormalCholLogDensity mu sigma l (map realToFrac yv :: [a]) + | yv <- chunks ] +obsLogSum (MvStudentT nu mu cov) ys = + let k = length mu + chunks = chunksOf k ys + in sum [ mvStudentTLogDensity nu mu cov (map realToFrac yv :: [a]) + | yv <- chunks ] +obsLogSum (DirichletMultinomial n alpha) ys = + let k = length alpha + chunks = chunksOf k ys + in sum [ dirichletMultinomialLogDensity n alpha yv | yv <- chunks ] +obsLogSum (Wishart nu vRows) ys = + let k = length vRows + chunks = chunksOf (k * k) ys + in sum [ wishartLogDensity nu vRows (map realToFrac yv :: [a]) + | yv <- chunks ] +obsLogSum d ys = sum [ logDensityObs d y | y <- ys ] diff --git a/src/Hanalyze/Model/HBM/Eval.hs b/src/Hanalyze/Model/HBM/Eval.hs new file mode 100644 --- /dev/null +++ b/src/Hanalyze/Model/HBM/Eval.hs @@ -0,0 +1,697 @@ +{-# LANGUAGE RankNTypes #-} +{-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE ScopedTypeVariables #-} + +-- | +-- Module : Hanalyze.Model.HBM.Eval +-- Description : HBM のモデル評価層 (log-joint/尤度インタープリタ + DAG 構築) +-- Copyright : (c) 2026 Aelysce Project (Toshiaki Honda) +-- License : BSD-3-Clause +-- +-- [日本語]: モデル評価層を 'Hanalyze.Model.HBM' から分離したもの。 +-- +-- PPL の __評価層__ (記述層 'Hanalyze.Model.HBM.Model' の上): +-- +-- - 構造化線形予測子 observe ('ObserveLM') の評価 (lmObsLogSum 等) +-- - log-joint / log-prior / log-likelihood の多相インタープリタ +-- - Gibbs 共役検出向けの runObserveDists / priorList +-- - 派生量評価 (runDeterministics / augmentChainWithDeterministic) と +-- DAG 構築 (buildModelGraph / collapseIndexedPlateNodes) +-- +-- 依存は下層 Model / Distribution (密度) / Track (extractDeps) / Util / +-- MCMC.Core のみ。 AD 勾配・IR は __上層__ に置かれ本モジュールへ依存する +-- (一方向)。 +-- [English]: The model evaluation layer, split out from +-- 'Hanalyze.Model.HBM'. +-- +-- The PPL's __evaluation layer__ (built on top of the description layer +-- 'Hanalyze.Model.HBM.Model'): +-- +-- - Evaluation of the structured linear-predictor observe ('ObserveLM') +-- (lmObsLogSum, etc.) +-- - Polymorphic interpreters for log-joint / log-prior / log-likelihood +-- - runObserveDists / priorList for Gibbs conjugacy detection +-- - Derived-quantity evaluation (runDeterministics / +-- augmentChainWithDeterministic) and DAG construction +-- (buildModelGraph / collapseIndexedPlateNodes) +-- +-- Depends only on the lower layers Model / Distribution (densities) / +-- Track (extractDeps) / Util / MCMC.Core. The AD gradient / IR live in +-- the __upper__ layer and depend on this module (one direction only). +module Hanalyze.Model.HBM.Eval + ( -- * ObserveLM 評価 + lmObsLogSum + -- * Interpreters + , logJoint + , logPrior + , logPriorWith + , logLikelihood + , perObsLogLiks + , runObserveDists + , mvNormalObserveOf + , priorList + , describeModel + -- * Type aliases + , Params + -- * 派生量 + , runDeterministics + , deterministicNames + , augmentChainWithDeterministic + -- * Model graph (visualization) + , ModelGraph (..) + , buildModelGraph + , collapseIndexedPlateNodes + ) where + +import Data.List (nub) +import Data.Map.Strict (Map) +import qualified Data.Map.Strict as Map +import qualified Data.Set as Set +import Data.Text (Text) +import qualified Data.Text as T + +import Hanalyze.MCMC.Core (Chain (..)) +import Hanalyze.Model.HBM.Util (negInf, chunksOf) +import Hanalyze.Model.HBM.Distribution +import Hanalyze.Model.HBM.Model +import Hanalyze.Model.HBM.Track (Track, extractDeps) + +-- --------------------------------------------------------------------------- +-- ObserveLM (構造化線形予測子 observe) の評価 (Phase 54.1) +-- --------------------------------------------------------------------------- + +-- | [日本語]: 線形予測子 η_i = Σ_j β_j·X_ij。 +-- synthGaussLMBlocks (本体) / IR が AD で微分しながら呼ぶホット経路。 +-- monolith では同一モジュール inline されていた。 境界跨ぎで失われると M1/M2 が +-- 約 +25% 劣化する (bench で実測) ため INLINABLE で cross-module inline を維持。 +-- [English]: The linear predictor η_i = Σ_j β_j·X_ij. +-- +-- A hot path called by synthGaussLMBlocks (the main body) / IR while +-- differentiating via AD. It used to be inlined in the same module in +-- the monolith. Losing it across a module boundary degrades M1/M2 by +-- about +25% (measured by bench), so INLINABLE keeps the cross-module +-- inline. +{-# INLINABLE lmEta #-} +lmEta :: Fractional a => [a] -> [Double] -> a +lmEta betas xrow = sum (zipWith (\b x -> b * realToFrac x) betas xrow) + +-- | [日本語]: ランダム効果項の per-obs 寄与 @Σ_re w_i·u^{re}[gid_i]@ (長さ n)。 +-- 重み @Nothing@ = 全 1。 +-- [English]: The per-observation contribution of the random-effect term +-- @Σ_re w_i·u^{re}[gid_i]@ (length n). Weight @Nothing@ = all 1s. +{-# INLINABLE lmReffEta #-} +lmReffEta :: forall a. Fractional a => [REff] -> Int -> Map Text a -> [a] +lmReffEta reffs n params = + foldr (zipWith (+)) (replicate n 0) + [ let uvals = [ Map.findWithDefault 0 nm params | nm <- uNames ] + base = [ uvals !! g | g <- gids ] + in case mw of + Nothing -> base + Just ws -> zipWith (\v w -> v * realToFrac w) base ws + | REff uNames gids _ mw _ <- reffs ] + +-- | [日本語]: 'ObserveLM' ブロックの各観測の log-density (per-obs)。 param Map から +-- β / u / (Gaussian の) σ を名前で引く。 η_i = Σ_j β_j X_ij + Σ_re u^{re}[gid_i] +-- を scalar 経路と同じ式で評価する。 +-- [English]: The per-observation log-density of an 'ObserveLM' block. +-- Looks up β / u / (Gaussian's) σ by name from the param Map, and +-- evaluates η_i = Σ_j β_j X_ij + Σ_re u^{re}[gid_i] with the same +-- formula as the scalar path. +{-# INLINABLE lmObsLogLiks #-} +lmObsLogLiks :: forall a. (Floating a, Ord a) + => [Text] -> [[Double]] -> [REff] -> LMFamily -> [Double] -> Map Text a -> [a] +lmObsLogLiks betaNames designX reffs fam ys params = + let betas = [ Map.findWithDefault 0 n params | n <- betaNames ] + reEta = lmReffEta reffs (length ys) params + etas = zipWith (\xr re -> lmEta betas xr + re) designX reEta + rows = zip etas ys + in case fam of + LMGaussian sName -> + let sigma = Map.findWithDefault 0 sName params + in [ logDensityObs (Normal eta sigma) y | (eta, y) <- rows ] + LMPoisson -> + [ logDensityObs (Poisson (exp eta)) y | (eta, y) <- rows ] + LMBernoulli -> + [ logDensityObs (Bernoulli (1 / (1 + exp (negate eta)))) y + | (eta, y) <- rows ] + +-- | [日本語]: 'ObserveLM' ブロックの log-likelihood 和。 +-- [English]: The sum of the log-likelihood of an 'ObserveLM' block. +{-# INLINABLE lmObsLogSum #-} +lmObsLogSum :: (Floating a, Ord a) + => [Text] -> [[Double]] -> [REff] -> LMFamily -> [Double] -> Map Text a -> a +lmObsLogSum betaNames designX reffs fam ys params = + sum (lmObsLogLiks betaNames designX reffs fam ys params) + +-- --------------------------------------------------------------------------- +-- 評価インタープリタ +-- --------------------------------------------------------------------------- + +-- | [日本語]: log-joint @log p(θ, y)@ を計算する多相インタープリタ。 +-- 引数 @a@ を @Double@ にすると数値評価、@Reverse s Double@ にすると AD 評価が可能。 +-- [English]: Polymorphic interpreter that computes the log-joint +-- @log p(θ, y)@. Instantiating the argument @a@ with @Double@ gives +-- numeric evaluation; with @Reverse s Double@, AD evaluation. +logJoint :: (Floating a, Ord a) => Model a r -> Map Text a -> a +logJoint model params = go model 0 + where + go (Pure _) acc = acc + go (Free (Sample n d k)) acc = + case Map.lookup n params of + Nothing -> negInf + Just v -> + let lp = logDensity d v + in go (k v) (acc + lp) + go (Free (Observe _ d ys next)) acc = + let ll = obsLogSum d ys + in go next (acc + ll) + go (Free (ObserveLM _ bs xs re fam ys next)) acc = + go next (acc + lmObsLogSum bs xs re fam ys params) + go (Free (Potential _ v next)) acc = go next (acc + v) + go (Free (Deterministic _ v k)) acc = go (k v) acc + -- Phase 60.2: Data 継続は [a] を受ける (lazy list・消費 1 回で O(n)/eval) + go (Free (Data _ ys k)) acc = go (k (map realToFrac ys, ys)) acc + go (Free (DataIx _ is k)) acc = go (k is) acc + go (Free (PlateBegin _ _ next)) acc = go next acc + go (Free (PlateEnd next)) acc = go next acc + +-- | [日本語]: log p(θ) のみ (prior 部分)。 +-- [English]: Just log p(θ) (the prior part). +logPrior :: (Floating a, Ord a) => Model a r -> Map Text a -> a +logPrior = logPriorWith logDensity + +-- | [日本語]: 'logPrior' の密度関数注入版。 AD 経路が定数 hyperparameter の +-- lgamma 正規化項を Double へ畳み込む 'logDensityRD' を差し込むために使う +-- (@Gradient@ の fRest 参照)。 @logPriorWith logDensity@ = 従来の 'logPrior'。 +-- [English]: A density-function-injectable version of 'logPrior'. Used +-- to plug in 'logDensityRD', which folds the lgamma normalization term +-- of constant hyperparameters into a @Double@ on the AD path +-- (referenced by @Gradient@'s fRest). @logPriorWith logDensity@ is +-- equivalent to the plain 'logPrior'. +logPriorWith :: (Floating a, Ord a) + => (Distribution a -> a -> a) -> Model a r -> Map Text a -> a +logPriorWith density model params = go model 0 + where + go (Pure _) acc = acc + go (Free (Sample n d k)) acc = + case Map.lookup n params of + Nothing -> negInf + Just v -> go (k v) (acc + density d v) + go (Free (Observe _ _ _ next)) acc = go next acc + go (Free (ObserveLM _ _ _ _ _ _ next)) acc = go next acc -- prior 部分には寄与しない + go (Free (Potential _ v next)) acc = go next (acc + v) + go (Free (Deterministic _ v k)) acc = go (k v) acc + go (Free (Data _ ys k)) acc = go (k (map realToFrac ys, ys)) acc + go (Free (DataIx _ is k)) acc = go (k is) acc + go (Free (PlateBegin _ _ next)) acc = go next acc + go (Free (PlateEnd next)) acc = go next acc + +-- | [日本語]: log p(y | θ) のみ (likelihood 部分)。 +-- [English]: Just log p(y | θ) (the likelihood part). +logLikelihood :: (Floating a, Ord a) => Model a r -> Map Text a -> a +logLikelihood model params = go model 0 + where + go (Pure _) acc = acc + go (Free (Sample n _ k)) acc = + case Map.lookup n params of + Nothing -> go (k 0) acc + Just v -> go (k v) acc + go (Free (Observe _ d ys next)) acc = + let ll = obsLogSum d ys + in go next (acc + ll) + go (Free (ObserveLM _ bs xs re fam ys next)) acc = + go next (acc + lmObsLogSum bs xs re fam ys params) + go (Free (Potential _ _ next)) acc = go next acc -- Potential は事前項とみなす + go (Free (Deterministic _ v k)) acc = go (k v) acc + go (Free (Data _ ys k)) acc = go (k (map realToFrac ys, ys)) acc + go (Free (DataIx _ is k)) acc = go (k is) acc + go (Free (PlateBegin _ _ next)) acc = go next acc + go (Free (PlateEnd next)) acc = go next acc + +-- | [日本語]: 各 observe ノードについて、 現在のパラメータ値で評価した分布を +-- 観測データと共に返す。 Gibbs サンプラーが共役構造を検出する際に、 +-- 潜在変数の現在値に対する観測分布のパラメータを得るために使う +-- (Double 特殊化版)。 +-- +-- 例: @y ~ Normal(mu, sigma)@ で @ps = {mu=2, sigma=0.5}@ を渡すと +-- @[(\"y\", Normal 2 0.5, [...])]@ を返す。 +-- [English]: For each observe node, return its distribution evaluated at +-- the current parameter values together with the observed data. Used +-- by the Gibbs sampler when detecting conjugate structure, to obtain the +-- observation distribution's parameters at the latent variables' current +-- values (the @Double@-specialized form). +-- +-- Example: given @y ~ Normal(mu, sigma)@ and @ps = {mu=2, sigma=0.5}@, +-- returns @[(\"y\", Normal 2 0.5, [...])]@. +runObserveDists :: Model Double r + -> Map Text Double + -> [(Text, Distribution Double, [Double])] +runObserveDists (Pure _) _ = [] +runObserveDists (Free (Sample n _ k)) ps = + runObserveDists (k (Map.findWithDefault 0 n ps)) ps +runObserveDists (Free (Observe n d ys next)) ps = + (n, d, ys) : runObserveDists next ps +runObserveDists (Free (ObserveLM _ _ _ _ _ _ next)) ps = + -- ObserveLM は per-obs で μ が異なり単一 Distribution に収まらない。 + -- Gibbs 共役検出 (この関数の用途) の対象外ゆえスキップ。 + runObserveDists next ps +runObserveDists (Free (Potential _ _ next)) ps = + runObserveDists next ps +runObserveDists (Free (Deterministic _ v k)) ps = + runObserveDists (k v) ps +runObserveDists (Free (Data _ ys k)) ps = + runObserveDists (k (ys, ys)) ps +runObserveDists (Free (DataIx _ is k)) ps = + runObserveDists (k is) ps +runObserveDists (Free (PlateBegin _ _ next)) ps = runObserveDists next ps +runObserveDists (Free (PlateEnd next)) ps = runObserveDists next ps + +-- | [日本語]: 解析随伴 (detach) パスの適格判定 + 抽出。 +-- モデルの尤度項が __ちょうど 1 個の 'MvNormal' observe__ のみ (他 'Observe' / +-- 'ObserveLM' 無し) のとき、その @(μ, Σ, ys)@ を __現在の param 値で評価__ して +-- 返す。 それ以外は 'Nothing' (= 呼び出し側は従来の walk+ad / vecIR 経路へ)。 +-- +-- 多相 (@Floating a@) ゆえ Double でも AD 型でも走らせられる: Double 版で +-- LAPACK 用の Σ⁻¹/logdet を作り (G,h 定数化)、 AD 版で surrogate @@ を +-- 微分する (@Gradient.compileGradUV@ の解析枝)。 walk は 'logJoint' 等と同一 +-- (Sample 継続に @params Map.! name@ を流す)。 μ/Σ は Observe ノードの +-- 'Distribution' に格納された式ゆえ、 現在の param 値で lazy に具体化される。 +-- +-- 適格条件を __1 個の MvNormal に限定__するのは正しさのため: 尤度が MvNormal +-- 単独なら @grad(logPrior+logJac) + detach(observe)@ で厳密に総勾配を再構成できる +-- (@logJoint = logPrior + logLikelihood@・@logLikelihood = obsLogSum(MvNormal)@)。 +-- [English]: Eligibility check + extraction for the analytic adjoint +-- (detach) path. +-- +-- When the model's likelihood term consists of +-- __exactly one 'MvNormal' observe__ (no other 'Observe' / 'ObserveLM'), +-- returns its @(μ, Σ, ys)@ __evaluated at the current param values__. +-- Otherwise 'Nothing' (the caller then falls back to the usual +-- walk+ad / vecIR path). +-- +-- Being polymorphic (@Floating a@) lets it run with either Double or an +-- AD type: the Double version builds Σ⁻¹/logdet for LAPACK (fixing G, h +-- as constants), and the AD version differentiates the surrogate +-- @@ (the analytic branch of @Gradient.compileGradUV@). The walk +-- is the same as 'logJoint' etc. (feeding @params Map.! name@ into the +-- Sample continuation). μ/Σ come from the expression stored in the +-- Observe node's 'Distribution', so they are lazily materialized at the +-- current param values. +-- +-- The eligibility condition is __restricted to exactly 1 MvNormal__ for +-- correctness: only when the likelihood is a single MvNormal can the +-- total gradient be exactly reconstructed as +-- @grad(logPrior+logJac) + detach(observe)@ (since +-- @logJoint = logPrior + logLikelihood@ and +-- @logLikelihood = obsLogSum(MvNormal)@). +mvNormalObserveOf :: (Floating a, Ord a) + => Model a r -> Map Text a -> Maybe ([a], [[a]], [Double]) +mvNormalObserveOf model params = + case go model of + Just [(MvNormal mu cov, ys)] -> Just (mu, cov, ys) + _ -> Nothing + where + -- Observe ノードの (dist, ys) を集める。 ObserveLM が在れば失格 (Nothing)。 + go (Pure _) = Just [] + go (Free (Sample n _ k)) = + case Map.lookup n params of + Nothing -> Nothing -- param 欠落 = 失格 (通常起きない) + Just v -> go (k v) + go (Free (Observe _ d ys next)) = ((d, ys) :) <$> go next + go (Free (ObserveLM {})) = Nothing -- 構造化尤度は対象外 + go (Free (Potential _ _ next)) = go next -- prior 側 (logPrior が処理) + go (Free (Deterministic _ v k)) = go (k v) + go (Free (Data _ ys k)) = go (k (map realToFrac ys, ys)) + go (Free (DataIx _ is k)) = go (k is) + go (Free (PlateBegin _ _ next)) = go next + go (Free (PlateEnd next)) = go next + +-- | [日本語]: 各 sample ノードについて @(name, prior distribution)@ を +-- @Double@ 特殊化形式で返す。 +-- Gibbs サンプラーの共役検出で「この潜在変数の事前は Gamma か Beta か」を +-- 判定するために使う。継続値はプレースホルダ 0 を流す。 +-- [English]: For each sample node, return @(name, prior distribution)@ in +-- the @Double@-specialized form. Used by the Gibbs sampler's conjugacy +-- detection to determine "is this latent variable's prior a Gamma or a +-- Beta?". Feeds the placeholder 0 into the continuation. +priorList :: Model Double r -> [(Text, Distribution Double)] +priorList (Pure _) = [] +priorList (Free (Sample n d k)) = (n, d) : priorList (k 0) +priorList (Free (Observe _ _ _ next)) = priorList next +priorList (Free (ObserveLM _ _ _ _ _ _ next)) = priorList next +priorList (Free (Potential _ _ next)) = priorList next +priorList (Free (Deterministic _ v k)) = priorList (k v) +priorList (Free (Data _ ys k)) = priorList (k (ys, ys)) +priorList (Free (DataIx _ is k)) = priorList (k is) +priorList (Free (PlateBegin _ _ next)) = priorList next +priorList (Free (PlateEnd next)) = priorList next + +-- --------------------------------------------------------------------------- +-- 互換 API +-- --------------------------------------------------------------------------- + +-- | [日本語]: パラメータ名 → 値 のマップ (constrained 空間)。 +-- [English]: A map from parameter name to value (in the constrained space). +type Params = Map Text Double + +-- | [日本語]: Per-observation log-likelihood (WAIC / LOO-CV で使用)。 +-- 各 Observe ノードのすべての観測値の logDensity を平坦リストで返す。 +-- [English]: Per-observation log-likelihood (used by WAIC / LOO-CV). +-- Returns the logDensity of every observation of every Observe node as +-- a flat list. +perObsLogLiks :: forall r. ModelP r -> Params -> [Double] +perObsLogLiks m params = go m [] + where + go :: Model Double r -> [Double] -> [Double] + go (Pure _) acc = reverse acc + go (Free (Sample n _ k)) acc = + go (k (Map.findWithDefault 0 n params)) acc + go (Free (Observe _ d ys next)) acc = + let lls = case d of + MvNormal mu cov -> + let k = length mu + in [ mvNormalLogDensity mu cov (map realToFrac yv :: [Double]) + | yv <- chunksOf k ys ] + Multinomial nn pp -> + let k = length pp + in [ multinomialLogDensity nn pp yv | yv <- chunksOf k ys ] + _ -> [ logDensityObs d y | y <- ys ] + in go next (reverse lls ++ acc) + go (Free (ObserveLM _ bs xs re fam ys next)) acc = + let lls = lmObsLogLiks bs xs re fam ys params + in go next (reverse lls ++ acc) + go (Free (Potential _ _ next)) acc = go next acc + go (Free (Deterministic _ v k)) acc = go (k v) acc + go (Free (Data _ ys k)) acc = go (k (ys, ys)) acc + go (Free (DataIx _ is k)) acc = go (k is) acc + go (Free (PlateBegin _ _ next)) acc = go next acc + go (Free (PlateEnd next)) acc = go next acc + +-- | [日本語]: 全ての 'Deterministic' ノードを評価し、 導出量の @Map@ を返す。 +-- +-- @params@ は latent 変数 (sample) の値を表す Map。Deterministic は +-- それらから導出される量で、ここでは Double 特殊化で評価する。 +-- [English]: Evaluate every 'Deterministic' node and return the +-- resulting derived-quantity @Map@. +-- +-- @params@ is a Map representing the values of the latent variables +-- (sample). Deterministic quantities are derived from them, evaluated +-- here in the @Double@-specialized form. +runDeterministics :: forall r. ModelP r -> Params -> Map Text Double +runDeterministics m params = go m Map.empty + where + go :: Model Double r -> Map Text Double -> Map Text Double + go (Pure _) acc = acc + go (Free (Sample n _ k)) acc = + go (k (Map.findWithDefault 0 n params)) acc + go (Free (Observe _ _ _ next)) acc = go next acc + go (Free (ObserveLM _ _ _ _ _ _ next)) acc = go next acc + go (Free (Potential _ _ next)) acc = go next acc + go (Free (Deterministic n v k)) acc = + go (k v) (Map.insert n v acc) + go (Free (Data _ ys k)) acc = go (k (ys, ys)) acc + go (Free (DataIx _ is k)) acc = go (k is) acc + go (Free (PlateBegin _ _ next)) acc = go next acc + go (Free (PlateEnd next)) acc = go next acc + +-- | [日本語]: モデル中の 'Deterministic' 宣言名を宣言順で列挙する。 +-- 同名の重複宣言 (plate 内反復等) は初出のみ残す。'collectNodes' は +-- Deterministic を素通しして 'Node' 化しないため専用 walker で拾う。 +-- 'runDeterministics' の返す Map の key 集合と一致する (順序のみ異なる)。 +-- [English]: Enumerates the model's 'Deterministic' declaration names in +-- declaration order. Duplicate declarations of the same name (e.g. plate +-- iteration) keep only the first occurrence. Since 'collectNodes' passes +-- Deterministic through without turning it into a 'Node', a dedicated +-- walker collects it here. Matches the key set of the Map returned by +-- 'runDeterministics' (only the order differs). +deterministicNames :: forall r. ModelP r -> [Text] +deterministicNames m = nub (go m []) + where + go :: Model Double r -> [Text] -> [Text] + go (Pure _) acc = reverse acc + go (Free (Sample n _ k)) acc = go (k 0) acc -- placeholder 0 (collectNodes 同型) + go (Free (Observe _ _ _ next)) acc = go next acc + go (Free (ObserveLM _ _ _ _ _ _ next)) acc = go next acc + go (Free (Potential _ _ next)) acc = go next acc + go (Free (Deterministic n v k)) acc = go (k v) (n : acc) + go (Free (Data _ ys k)) acc = go (k (ys, ys)) acc + go (Free (DataIx _ is k)) acc = go (k is) acc + go (Free (PlateBegin _ _ next)) acc = go next acc + go (Free (PlateEnd next)) acc = go next acc + +-- | [日本語]: 全 posterior サンプルに対して 'runDeterministics' を評価し、 +-- 結果を 'chainSamples' の Map にマージした新しい Chain を返す。 +-- これにより @chainVals@ / @posteriorSummary@ などのヘルパで派生量を +-- そのまま参照できる。 +-- [English]: Evaluate 'runDeterministics' on every posterior sample and +-- return a new Chain with the results merged into 'chainSamples's Map. +-- This lets helpers such as @chainVals@ / @posteriorSummary@ reference +-- derived quantities directly. +augmentChainWithDeterministic :: ModelP r -> Chain -> Chain +augmentChainWithDeterministic m ch = + let aug ps = Map.union (runDeterministics m ps) ps + in ch { chainSamples = map aug (chainSamples ch) } + +-- | Human-readable summary of the model structure (no inference is run). +describeModel :: ModelP r -> Text +describeModel m = T.unlines (header : map fmtNode (collectNodes m)) + where + header = "Model nodes:" + fmtNode n = case nodeKind n of + LatentN -> " [latent] " <> nodeName n <> " ~ " <> nodeDist n + ObservedN k -> " [observed] " <> nodeName n <> " ~ " <> nodeDist n + <> " (n=" <> T.pack (show k) <> ")" + DeterministicN -> " [determ] " <> nodeName n <> " = " <> nodeDist n + DataN k -> " [data] " <> nodeName n + <> " (n=" <> T.pack (show k) <> ")" + +-- | DAG representation of the model. Edges are derived automatically by +-- 'extractDeps'. +data ModelGraph = ModelGraph + { mgNodes :: [Node] + , mgEdges :: [(Text, Text)] -- (parent, child) + , mgPlates :: Map Text Int -- Phase 40: plate 名 → サイズ N + } deriving (Show) + +-- | [日本語]: Plate 内の indexed RV (`eta_0, eta_1, …, eta_{n-1}`) を +-- __代表 1 ノードに集約__して、 PyMC `pm.model_to_graphviz` 流の true plate +-- 描画用に変換する。 +-- +-- 集約条件 (heuristic): +-- +-- - 同じ `nodePlates` (= plate スタック) に属する +-- - 名前が @\_\$@ パターン (末尾が _ + 数字) +-- - 同じ @prefix@ を持つノード群が 2 個以上 +-- - 同じ `nodeDist` (= 分布名が一致) +-- +-- 集約結果: +-- +-- - 代表ノード名は @prefix@ (例: @eta_0..eta_7@ → @eta@) +-- - `nodeKind`: 元の集合内で最初の出現を維持 (LatentN / ObservedN)。 +-- ObservedN の場合は観測数を全集約 (Σ) +-- - `nodeDeps`: 全集合の親集合の和 (ただし、 同じ集合内のメンバ間 deps は +-- 削除 — 自己集約のため) +-- - edges: 集約後の名前で dedupe +-- +-- plate 文脈外で起きる「同じ命名規則の名前衝突」 (e.g. @beta_0@ 固定効果 vs +-- @u_0@ 群効果) はこの heuristic で誤って集約されない (plate 制約)。 +-- +-- 元 graph をそのまま渡せば不変 (idempotent)。 plate に属さない / 単独 +-- のノードは触らない。 +-- [English]: Converts the indexed RVs within a plate (`eta_0, eta_1, …, +-- eta_{n-1}`) by __collapsing them into a single representative node__, +-- for PyMC `pm.model_to_graphviz`-style true plate rendering. +-- +-- Collapse conditions (heuristic): +-- +-- - Belong to the same `nodePlates` (= plate stack) +-- - Name matches the @\_\$@ pattern (ends in _ + digits) +-- - 2 or more nodes share the same @prefix@ +-- - Share the same `nodeDist` (= same distribution name) +-- +-- Collapse result: +-- +-- - The representative node's name is @prefix@ (e.g. @eta_0..eta_7@ → +-- @eta@) +-- - `nodeKind`: keeps the first occurrence within the original set +-- (LatentN / ObservedN). For ObservedN, the observation counts are +-- summed (Σ) +-- - `nodeDeps`: the union of the parent sets of the whole set (deps +-- between members of the same set are removed, since that would be +-- self-collapse) +-- - edges: deduped under the collapsed names +-- +-- A "name collision under the same naming convention" occurring outside +-- a plate context (e.g. @beta_0@ a fixed effect vs @u_0@ a group effect) +-- is not mistakenly collapsed by this heuristic (the plate constraint). +-- +-- Passing the original graph through unchanged is a no-op (idempotent). +-- Nodes that don't belong to a plate, or stand alone, are untouched. +collapseIndexedPlateNodes :: ModelGraph -> ModelGraph +collapseIndexedPlateNodes mg0 = + -- 不動点: 1 回の集約で取りこぼした多段 plate (e.g. y_0_0..y_2_1 → y_0..y_2 → + -- 残り index suffix を持つ → y) を順次潰す。 mgNodes 数が減らなくなれば終了。 + let step g = collapseIndexedPlateNodesOnce g + iter g = let g' = step g in if length (mgNodes g') == length (mgNodes g) + then g else iter g' + in iter mg0 + +-- | [日本語]: `collapseIndexedPlateNodes` の 1 段集約 (内部、 不動点を作る材料)。 +-- [English]: One step of `collapseIndexedPlateNodes`'s collapsing +-- (internal; the building block used to reach the fixed point). +collapseIndexedPlateNodesOnce :: ModelGraph -> ModelGraph +collapseIndexedPlateNodesOnce mg = + let ns = mgNodes mg + es = mgEdges mg + -- 1. 各ノードについて (plate path, prefix) または Nothing を計算 + keyOf n = case T.breakOnEnd "_" (nodeName n) of + (pre, digits) + | not (T.null pre) && not (T.null digits) + && T.all (`elem` ("0123456789" :: String)) digits -> + Just (nodePlates n, T.init pre) -- _ を除いた prefix + _ -> Nothing + -- 2. キー単位で groupings + keyed = [(keyOf n, n) | n <- ns] + -- 3. グループ化 (Just key) のみ、 Nothing は単独 + grouped :: Map.Map ([Text], Text) [Node] + grouped = Map.fromListWith (flip (++)) + [ (k, [n]) | (Just k, n) <- keyed ] + -- 4. 集約候補: size ≥ 2 かつ全 nodeDist 一致 + collapsible = Map.filter + (\g -> length g >= 2 + && all (\n -> nodeDist n == nodeDist (head g)) g) + grouped + -- 5. name → 代表名 のマップ + nameMap :: Map.Map Text Text + nameMap = Map.fromList + [ (nodeName n, prefix) + | ((_plates, prefix), grp) <- Map.toList collapsible + , n <- grp + ] + mapName n = Map.findWithDefault n n nameMap + -- 6. 集約後ノード作成 + mkRepresentative (_, prefix) grp = + let first = head grp + kind = case nodeKind first of + ObservedN _ -> + ObservedN (sum [k | n <- grp, + let ObservedN k = nodeKind n]) + LatentN -> LatentN + DeterministicN -> DeterministicN + dk@(DataN _) -> dk + -- 自己集約 (同じ集合のメンバへの deps) を除外 + memberNames = Set.fromList (map nodeName grp) + externalDeps = Set.unions (map nodeDeps grp) + `Set.difference` memberNames + -- 親側の名前も mapName で remap (e.g. mu_0..mu_K-1 集約済の場合) + remappedDeps = Set.map mapName externalDeps + in first { nodeName = prefix + , nodeKind = kind + , nodeDeps = remappedDeps + } + -- 7. ノードリスト再構築: 集約対象は代表 1 個、 非対象はそのまま + isInGroup n = case keyOf n of + Just k -> Map.member k collapsible + Nothing -> False + seenGroups :: [([Text], Text)] + seenGroups = [] + walk [] _ acc = reverse acc + walk (n:rest) seen acc + | isInGroup n = + let Just k = keyOf n + in if k `elem` seen + then walk rest seen acc + else let rep = mkRepresentative k (collapsible Map.! k) + in walk rest (k : seen) (rep : acc) + | otherwise = walk rest seen + (n { nodeDeps = Set.map mapName (nodeDeps n) } : acc) + newNodes = walk ns seenGroups [] + -- 8. edges を remap + dedupe + 自己ループ除去 + newEdges = Set.toList $ Set.fromList + [ (s', t') + | (s, t) <- es + , let s' = mapName s + , let t' = mapName t + , s' /= t' -- 自己ループ除外 + ] + in mg { mgNodes = newNodes, mgEdges = newEdges } + +-- | [日本語]: 多相モデルから DAG を自動構築する (Track 型による依存追跡)。 +-- +-- 同じ名前で複数登場する Observe ノード (例: 回帰モデルで観測点ごとに +-- @observe \"y\"@ を発行する場合) は 1 つに統合される。観測数の合計と +-- 親変数集合の和をマージし、エッジも重複排除する。 +-- [English]: Automatically builds a DAG from a polymorphic model +-- (dependency tracking via the Track type). +-- +-- Observe nodes that appear repeatedly under the same name (e.g. issuing +-- @observe \"y\"@ per observation in a regression model) are merged into +-- one. The observation counts are summed and the parent-variable sets +-- are unioned; edges are deduped as well. +buildModelGraph :: ModelP r -> ModelGraph +buildModelGraph m = + let (rawNodes, plates) = extractDeps m + merged = assignDataPlates plates (mergeByName rawNodes) + edges = Set.toList $ Set.fromList + [ (parent, nodeName n) + | n <- merged + , parent <- Set.toList (nodeDeps n) ] + in ModelGraph merged edges plates + where + -- Phase 60.6 追補: 宣言位置が plate 外 (nodePlates = []) の DataN を、 + -- PyMC の dims 同様「データ長 = plate サイズ」 の一意 match で plate に + -- 割り当てる (典型 = モデル冒頭で宣言した dataNamedX n=150 が obs(150) + -- cluster 内に描かれる)。 一致 plate が複数 / なし は据え置き (外に描く)。 + -- 入れ子 plate の full path は、 既にその plate に居る他ノードの + -- nodePlates から逆引きする (plate 内ノードが無い場合は単独 path)。 + assignDataPlates plates ns = + let paths = [ nodePlates n | n <- ns, not (null (nodePlates n)) ] + pathFor nm = case [ p | p <- paths, last p == nm ] of + (p : _) -> p + [] -> [nm] + assign n = case nodeKind n of + DataN k | null (nodePlates n) -> + case [ nm | (nm, sz) <- Map.toList plates, sz == k ] of + [nm] -> n { nodePlates = pathFor nm } + _ -> n + _ -> n + in map assign ns + -- 同名ノードを統合: ObservedN n1 + ObservedN n2 → ObservedN (n1+n2) + -- LatentN は最初の出現を残す。deps は和集合。 + -- nodePlates は最初の出現のものを維持 (同名は同 plate 前提)。 + mergeByName ns = mergeGo ns Map.empty [] + mergeGo [] _ acc = reverse acc + mergeGo (n:ns) seen acc = + let nm = nodeName n + in case Map.lookup nm seen of + Nothing -> mergeGo ns (Map.insert nm n seen) (n : acc) + Just prev -> + -- Phase 60.4: DataN は最弱 — 同名の非 DataN ノード (典型 = + -- dataNamedObs "y" + observe "y" の docs 慣例) があれば吸収される + -- (PyMC で observed RV が data 容器を内包して表示されるのと同型)。 + let (kind', dist', plates') = + case (nodeKind prev, nodeKind n) of + (ObservedN a, ObservedN b) -> + (ObservedN (a + b), nodeDist prev, nodePlates prev) + (DataN _, k2) -> (k2, nodeDist n, nodePlates n) + (k1, _) -> (k1, nodeDist prev, nodePlates prev) + merged' = Node + { nodeName = nm + , nodeKind = kind' + , nodeDist = dist' + , nodeDeps = nodeDeps prev <> nodeDeps n + , nodePlates = plates' + } + acc' = map (\x -> if nodeName x == nm then merged' else x) acc + in mergeGo ns (Map.insert nm merged' seen) acc' + + +-- --------------------------------------------------------------------------- +-- Track 評価 (logJoint の Track 特殊化) +-- --------------------------------------------------------------------------- + +-- | [日本語]: Track でモデルを評価する (log joint も依存集合付きで計算)。 +-- [English]: Evaluates the model with Track (computes the log joint +-- along with its dependency set as well). +runTrack :: forall r. ModelP r -> Map Text Track -> Track +runTrack m params = logJoint (m :: Model Track r) params diff --git a/src/Hanalyze/Model/HBM/Gradient.hs b/src/Hanalyze/Model/HBM/Gradient.hs new file mode 100644 --- /dev/null +++ b/src/Hanalyze/Model/HBM/Gradient.hs @@ -0,0 +1,2483 @@ +{-# LANGUAGE RankNTypes #-} +{-# LANGUAGE DeriveFunctor #-} +{-# LANGUAGE FlexibleInstances #-} +{-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE ImpredicativeTypes #-} + +-- | +-- Module : Hanalyze.Model.HBM.Gradient +-- Description : HBM の AD 勾配コンパイラ層 (NUTS per-draw のホット経路) +-- Copyright : (c) 2026 Aelysce Project (Toshiaki Honda) +-- License : BSD-3-Clause +-- +-- [日本語]: AD 勾配コンパイラ層を 'Hanalyze.Model.HBM' 本体から分離。 +-- IR (中間表現) 層の __上層__ であり、 NUTS per-draw の本経路 (compileGradUV → +-- gradVecIR / hybridGradClosure) を担う最ホット モジュール。 unconstrained 空間の +-- log-joint・解析閉形式勾配 (Gaussian LM ブロック)・ハイブリッド勾配クロージャ・ +-- 定数 prior 解析勾配・制約変換 (invTransformF/logJacF) を含む。 +-- +-- 全 top-level を export し ('module ... where' = 暗黙全公開)、 公開 API +-- (gradAD/gradADU/compileGradU/compileGradUV/compileLogPU/compileLogPUV/ +-- getTransforms/logJointUnconstrained/invTransformF/logJacF) は facade +-- 'Hanalyze.Model.HBM' の export list 経由で再エクスポートされる。 +-- +-- [English]: Splits the AD gradient compiler layer out from the +-- 'Hanalyze.Model.HBM' main body. It sits __above__ the IR +-- (intermediate representation) layer, and is the hottest module, +-- responsible for the main NUTS per-draw path (compileGradUV → +-- gradVecIR / hybridGradClosure). It contains the log-joint in +-- unconstrained space, analytic closed-form gradients (Gaussian LM +-- blocks), hybrid gradient closures, analytic gradients for constant +-- priors, and constraint transforms (invTransformF/logJacF). +-- +-- All top-level bindings are exported ('module ... where' means +-- everything is implicitly public); the public API +-- (gradAD/gradADU/compileGradU/compileGradUV/compileLogPU/compileLogPUV/ +-- getTransforms/logJointUnconstrained/invTransformF/logJacF) is +-- re-exported via the facade 'Hanalyze.Model.HBM''s export list. +module Hanalyze.Model.HBM.Gradient where + +import Control.DeepSeq (NFData (..), force) +import Control.Exception (SomeException, evaluate, try) +import Control.Monad (forM, forM_, replicateM, when) +import Data.List (foldl', zip4) +import System.IO.Unsafe (unsafePerformIO) +import qualified Data.Map.Strict as Map +import Data.Map.Strict (Map) +import qualified Data.Set as Set +import Data.Set (Set) +import Data.Text (Text) +import qualified Data.Text as T +import Data.Reflection (Reifies) +import Numeric.AD.Mode.Reverse.Double (grad, grad') +import qualified Numeric.AD.Internal.Reverse.Double as ADRD +import qualified System.Random.MWC as MWCBase +import qualified System.Random.MWC.Distributions as MWC +import System.Random.MWC (Gen) +import Control.Monad.Primitive (PrimMonad, PrimState, stToPrim) + +import Control.Monad.ST (ST, runST) +import qualified Data.Vector as BV +import qualified Data.Vector.Storable as VS +import qualified Data.Vector.Storable.Mutable as VSM +import qualified Data.Vector.Unboxed as VU + +-- Phase 95 A6: dense MvNormal observe の解析随伴 (detach) で Σ⁻¹/logdet を LAPACK +-- で 1 度だけ作る (cholesky を AD tape に載せない)。 +import qualified Numeric.LinearAlgebra as LA + +import Hanalyze.Stat.Distribution (Transform (..)) +import Hanalyze.MCMC.Core (Chain (..)) + +-- Phase 58.2: 純粋な数値・線形代数 leaf util を分離。 internal 利用に加え +-- 'lgammaApprox' / 'digamma' は export list 経由でそのまま再エクスポートされる。 +import Hanalyze.Model.HBM.Util +-- Phase 58.3/58.6a: 多相分布 ADT + 密度 + CDF を分離 (Util の上層)。 公開 API +-- (Distribution(..)/distName/logDensity/logDensityObs/obsLogSum/distCDF/logCDF/ +-- logSF/MV密度群) は export list 経由でそのまま再エクスポート。 ★58.6a で事前 +-- logDensity と観測 logDensityObs/obsLogSum を本体から Distribution へ集約 +-- (Eval の logJoint/logPrior が logDensity を参照する back-edge を解消・密度は +-- 本来 Distribution の責務。 INLINABLE は AD cross-module inlining 維持で保持)。 +import Hanalyze.Model.HBM.Distribution +-- Phase 58.4: 分布からのサンプリング (sampleDist/sampleMvDist) を分離。 +-- export list 経由でそのまま再エクスポート。 PrimMonad/mwc-random 依存・非ホット。 +import Hanalyze.Model.HBM.Sampling +-- Phase 58.5: 多相モデル DSL (Free monad + ModelF + plate + 構造検査) を分離。 +-- 公開 API (Free/liftF/ModelF/Model/ModelP/sample/observe/plate/collectNodes 等) +-- は export list 経由でそのまま再エクスポート。 +import Hanalyze.Model.HBM.Model + +-- Phase 58.6b: 依存追跡型 Track (Track/trackVar/trackConst/extractDeps) を分離。 +-- Model/Distribution の上層・非ホット (DAG 抽出のみ・NUTS per-draw 非経路)。 +-- export list 経由でそのまま再エクスポート。 +import Hanalyze.Model.HBM.Track +-- Phase 58.6c: 評価層 (ObserveLM 評価 + logJoint/logPrior/logLikelihood interp + +-- 互換 API runDeterministics/buildModelGraph 等 + runTrack) を分離。 Track の上層。 +-- ★ホット (logJoint は AD 勾配経路)。 AD 勾配・IR (本体残置) は本モジュールを +-- forward import する。 公開 API は export list 経由でそのまま再エクスポート。 +import Hanalyze.Model.HBM.Eval +-- Phase 58.7: IR (中間表現) 層 (affine/非線形/密度 IR) を分離。 最ホット (gradVecIR)。 +import Hanalyze.Model.HBM.IR + +-- Phase 60.7: '!!!' の依存タグは AD 勾配には無関係 (既定 id = サンプリング +-- ビット不変)。 ReverseDouble は ad の internal 型ゆえ orphan instance だが、 +-- AD 経路の instantiate はこのモジュールに閉じている。 +instance TrackTag (ADRD.ReverseDouble s) + +-- --------------------------------------------------------------------------- +-- AD 勾配 +-- --------------------------------------------------------------------------- + +-- | [日本語]: AD で勾配を計算する。@names@ の順で各パラメータに対する偏微分を返す。 +-- [English]: Computes the gradient with AD. Returns the partial +-- derivative for each parameter, in the order given by @names@. +gradAD :: ModelP r -> [Text] -> [Double] -> [Double] +gradAD m names xs0 = grad f xs0 + where + f xs = + let params = Map.fromList (zip names xs) + in logJoint m params + +-- | [日本語]: unconstrained 空間で AD 勾配を計算する (HMC 用)。 +-- 各パラメータに制約変換を適用し、Jacobian 補正項込みの log-joint を微分する。 +-- +-- モデルが __Gaussian-恒等リンクの 'ObserveLM' ブロック__ を +-- 含む場合、 そのブロックの観測尤度勾配 (= 規模 n に比例する支配項) を解析 +-- 閉形式 (∂β=Xᵀr/σ² 等) で計算し、 prior / jacobian / scalar observe / +-- 非 Gaussian LM は従来 @ad@ で計算してから成分加算する (ハイブリッド)。 +-- Gaussian LM を含まないモデルは従来通り全体を @ad@ で微分 (後方互換)。 +-- [English]: Computes the AD gradient in unconstrained space (for HMC). +-- Applies the constraint transform to each parameter and differentiates +-- the log-joint including the Jacobian correction term. +-- +-- If the model contains a __Gaussian identity-link 'ObserveLM' block__, +-- that block's observation-likelihood gradient (the dominant term, +-- proportional to sample size n) is computed as an analytic closed +-- form (@∂β=Xᵀr/σ²@ etc.), while the prior / jacobian / scalar observe / +-- non-Gaussian LM parts are still computed with @ad@ and then added +-- componentwise (a hybrid approach). A model without any Gaussian LM +-- block is differentiated in its entirety with @ad@ as before (for +-- backward compatibility). +gradADU :: ModelP r -> [Text] -> [Transform] -> [Double] -> [Double] +gradADU m names trans = compileGradU m names trans + +-- | [日本語]: 'compileGradUV' の list wrapper (後方互換 API)。 NUTS は vector-native の +-- 'compileGradUV' を直接使う。 +-- [English]: A list wrapper around 'compileGradUV' (for backward +-- compatibility). NUTS uses the vector-native 'compileGradUV' directly. +compileGradU :: forall r. ModelP r -> [Text] -> [Transform] -> ([Double] -> [Double]) +compileGradU m names trans = + let gv = compileGradUV m names trans + in VS.toList . gv . VS.fromList + +-- | [日本語]: 'compileGradUV' が実際に選ぶ勾配経路のラベル (診断表示用)。 +-- 'compileGradUV' 本体の分岐順 (gaussLMBlocksAuto → synthVecIR → 全体 ad) を +-- そのまま反映する唯一の分類子。 +-- +-- ★__束縛済モデル__ (@hbmModelSpec@ 経由) を渡すこと。 生の未束縛モデル +-- ('dataNamed*' の既定 @[]@ のまま) を渡すと data 行が空になり、 Gaussian LM +-- 合成も 'collectSymRows' も 0 行となって経路判定が狂う (実測: +-- 17-nes/12-ark は実際は Gaussian LM 閉形式経路なのに、 生モデルを渡した +-- 診断が @synthVecIR = Nothing@ と誤表示していた)。 +-- [English]: The label for the gradient path that 'compileGradUV' +-- actually selects (for diagnostic display). The single classifier +-- that directly mirrors 'compileGradUV''s own branch order +-- (gaussLMBlocksAuto → synthVecIR → full ad). +-- +-- ★Pass a __bound model__ (via @hbmModelSpec@). Passing a raw, +-- unbound model (still with the default @[]@ for @dataNamed*@) leaves +-- the data rows empty, so the Gaussian LM synthesis and +-- 'collectSymRows' both see 0 rows and the path decision goes wrong +-- (measured: 17-nes/12-ark actually take the Gaussian LM closed-form +-- path, but diagnostics on the raw model misreported +-- @synthVecIR = Nothing@). +gradPathLabel :: ModelP r -> String +gradPathLabel m = case gaussLMBlocksAuto m of + ([], _) -> case synthVecIR m of + Nothing | hasHmmObserve m -> "HMM forward-backward 閉形式随伴 (Phase 92)" + | hasArmaObserve m -> "ARMA(1,1) 逆向き随伴の閉形式 (Phase 101)" + | hasGradedIrtObserve m -> "graded response IRT 解析勾配 (Phase 101)" + | otherwise -> "legacy walk+ad (全体 ad)" + Just _ -> "vecIR (ベクトル式 IR 高速経路)" + _ -> "Gaussian LM 閉形式ブロック (解析勾配)" + +-- | [日本語]: 'gradPathLabel' 用の軽量構造判定 — 尤度が単一 'HmmForwardNormal' +-- observe か。 param 値は不要 (latent へ 0 を給餌・分布値は強制しない)。 +-- 実際の経路選択は 'gradValPlan' 内の @hmmAnalyticVG@ (probe 同型) が行う。 +-- [English]: A lightweight structural check for 'gradPathLabel' — +-- whether the likelihood is a single 'HmmForwardNormal' observe. +-- Parameter values are not needed (0 is fed to the latents; the +-- distribution values are not forced). The actual path selection is +-- done by @hmmAnalyticVG@ inside 'gradValPlan' (structurally identical +-- probe). +hasHmmObserve :: ModelP r -> Bool +hasHmmObserve m = case go m of + Just [HmmForwardNormal {}] -> True + _ -> False + where + go :: Model Double r' -> Maybe [Distribution Double] + go (Pure _) = Just [] + go (Free (Sample _ _ k)) = go (k 0) + go (Free (Observe _ d _ next)) = (d :) <$> go next + go (Free (ObserveLM {})) = Nothing + go (Free (Potential _ _ next)) = go next + go (Free (Deterministic _ v k)) = go (k v) + go (Free (Data _ ys k)) = go (k (ys, ys)) + go (Free (DataIx _ is k)) = go (k is) + go (Free (PlateBegin _ _ next)) = go next + go (Free (PlateEnd next)) = go next + +-- | [日本語]: 'gradPathLabel' 用の軽量構造判定 — 尤度が単一 'ArmaNormal' +-- observe か。 実際の経路選択は 'gradValPlan' 内の @armaAnalyticVG@ が行う。 +-- [English]: A lightweight structural check for 'gradPathLabel' — +-- whether the likelihood is a single 'ArmaNormal' observe. The actual +-- path selection is done by @armaAnalyticVG@ inside 'gradValPlan'. +hasArmaObserve :: ModelP r -> Bool +hasArmaObserve m = case go m of + Just [ArmaNormal {}] -> True + _ -> False + where + go :: Model Double r' -> Maybe [Distribution Double] + go (Pure _) = Just [] + go (Free (Sample _ _ k)) = go (k 0) + go (Free (Observe _ d _ next)) = (d :) <$> go next + go (Free (ObserveLM {})) = Nothing + go (Free (Potential _ _ next)) = go next + go (Free (Deterministic _ v k)) = go (k v) + go (Free (Data _ ys k)) = go (k (ys, ys)) + go (Free (DataIx _ is k)) = go (k is) + go (Free (PlateBegin _ _ next)) = go next + go (Free (PlateEnd next)) = go next + +-- | [日本語]: 'gradPathLabel' 用の軽量構造判定 — 尤度が単一 +-- 'GradedResponseIrt' observe か。 実際の経路選択は 'gradValPlan' 内の +-- @gradedIrtAnalyticVG@ が行う。 +-- [English]: A lightweight structural check for 'gradPathLabel' — +-- whether the likelihood is a single 'GradedResponseIrt' observe. The +-- actual path selection is done by @gradedIrtAnalyticVG@ inside +-- 'gradValPlan'. +hasGradedIrtObserve :: ModelP r -> Bool +hasGradedIrtObserve m = case go m of + Just [GradedResponseIrt {}] -> True + _ -> False + where + go :: Model Double r' -> Maybe [Distribution Double] + go (Pure _) = Just [] + go (Free (Sample _ _ k)) = go (k 0) + go (Free (Observe _ d _ next)) = (d :) <$> go next + go (Free (ObserveLM {})) = Nothing + go (Free (Potential _ _ next)) = go next + go (Free (Deterministic _ v k)) = go (k v) + go (Free (Data _ ys k)) = go (k (ys, ys)) + go (Free (DataIx _ is k)) = go (k is) + go (Free (PlateBegin _ _ next)) = go next + go (Free (PlateEnd next)) = go next + +-- --------------------------------------------------------------------------- +-- Phase 95 A6: dense MvNormal observe の解析随伴 (detach) 経路 +-- --------------------------------------------------------------------------- + +-- | [日本語]: 尤度が __単一 'MvNormal' observe__ のモデル (GP 回帰・ +-- dense-MvNormal) の value+grad を __解析随伴 (detach トリック)__ で計算する +-- クロージャを構築する。 'mvNormalObserveOf' が 'Just' のとき (= 適格) のみ +-- 'Just' を返し、 非適格なら 'Nothing' (呼び出し側は従来 walk+ad へ)。 +-- +-- === なぜ速いか +-- 現行 walk+ad は N×N Cholesky + solve + logdet を毎 leapfrog で +-- __reverse-AD tape 上に丸ごと展開__する (O(N³) のスカラー演算が各々 boxed AD +-- ノードを alloc)。 +-- 大 N で壊滅的 (§A4: N=50 で対 PyMC 62×)。 本経路は PyMC/Stan と同じく +-- __Cholesky を AD tape に載せず__ Σ⁻¹/logdet を LAPACK で 1 度だけ Double 計算し、 +-- @G = ∂logp/∂Σ@・@h = ∂logp/∂μ@ を定数化して surrogate @+@ の +-- 軽量 ad (O(N²)・cholesky 無し) だけを微分する。 +-- +-- === 数学 (detach) +-- 1 観測 (k-vector) の @logp = -k/2·log2π - 0.5·log|Σ| - 0.5·(y-μ)ᵀΣ⁻¹(y-μ)@ に対し +-- @α = Σ⁻¹(y-μ)@、 @∂logp/∂Σ = 0.5(ααᵀ - Σ⁻¹)@、 @∂logp/∂μ = α@。 複数 chunk では +-- @G = 0.5(Σ_m α_mα_mᵀ - nCh·Σ⁻¹)@、 @h = Σ_m α_m@。 surrogate @+@ の +-- θ 勾配は連鎖律で @∂logp/∂θ@ に厳密一致 (proto: 有限差分 2.8e-9・ad-through 2.2e-15)。 +-- θ=invTransform(u) を surrogate 内に組むので ∂/∂u が直接得られる。 +-- +-- 総勾配 = @grad(logPrior+logJac)@ (scalar・cheap) + @detach(observe)@。 +-- 値 = @(logPrior+logJac)@ + @logp_MvNormal@ (同じ LAPACK 分解から)。 +-- +-- === 常時 ON (次元しきい値なし) +-- 解析随伴は近似でなく厳密ゆえ正しさゲート不要。 PyMC/Stan も Cholesky Op の +-- 解析随伴を __常時__使う (次元しきい値を持たない)。 小 N (N≲40) では list 往復の +-- 定数倍で ad-through と拮抗〜僅遅だが (§A3-proto crossover ≈N40-50)、 該当する +-- shipping モデルは無い。 その帯の overhead 除去 (surrogate 配列化) は TODO。 +-- +-- [English]: For a model whose likelihood is a +-- __single 'MvNormal' observe__ (GP regression / dense MvNormal), +-- builds a closure that computes the value+grad via an +-- __analytic adjoint (the detach trick)__. Returns 'Just' only when +-- 'mvNormalObserveOf' returns 'Just' (i.e., the model qualifies); +-- otherwise 'Nothing' (the caller falls back to the usual walk+ad). +-- +-- === Why it's fast +-- The current walk+ad approach +-- __unrolls the N×N Cholesky + solve + logdet entirely onto the reverse-AD tape on every leapfrog step__ +-- (its O(N³) scalar operations each allocate a boxed AD node). This is +-- devastating for large N (§A4: at N=50, 62x slower than PyMC). This +-- path, like PyMC/Stan, __never puts the Cholesky on the AD tape__: +-- it computes Σ⁻¹/logdet as plain Doubles exactly once via LAPACK, +-- treats @G = ∂logp/∂Σ@ and @h = ∂logp/∂μ@ as constants, and +-- differentiates only the lightweight surrogate +-- @+@ (O(N²), no Cholesky). +-- +-- === Math (detach) +-- For one observation (a k-vector), given +-- @logp = -k/2·log2π - 0.5·log|Σ| - 0.5·(y-μ)ᵀΣ⁻¹(y-μ)@, let +-- @α = Σ⁻¹(y-μ)@, @∂logp/∂Σ = 0.5(ααᵀ - Σ⁻¹)@, @∂logp/∂μ = α@. For +-- multiple chunks, @G = 0.5(Σ_m α_mα_mᵀ - nCh·Σ⁻¹)@, +-- @h = Σ_m α_m@. By the chain rule, the θ-gradient of the surrogate +-- @+@ matches @∂logp/∂θ@ exactly (prototype: finite +-- difference 2.8e-9, ad-through 2.2e-15). Since @θ=invTransform(u)@ +-- is built into the surrogate, @∂/∂u@ is obtained directly. +-- +-- Total gradient = @grad(logPrior+logJac)@ (scalar, cheap) + +-- @detach(observe)@. Value = @(logPrior+logJac)@ + @logp_MvNormal@ +-- (from the same LAPACK factorization). +-- +-- === Always on (no dimension threshold) +-- The analytic adjoint is exact, not an approximation, so no +-- correctness gate is needed. PyMC/Stan also use the Cholesky op's +-- analytic adjoint __unconditionally__ (with no dimension threshold). +-- For small N (N≲40) the constant overhead of round-tripping through +-- lists makes it roughly on par with, or marginally slower than, +-- ad-through (§A3-proto crossover ≈N40-50), but no shipping model +-- falls in that range. Removing that overhead (by arraying the +-- surrogate) is a follow-up TODO. +mvNormalAnalyticVG + :: forall r. ModelP r -> [Text] -> [Transform] + -> Maybe (VS.Vector Double -> (Double, VS.Vector Double)) +mvNormalAnalyticVG m names trans = + -- 適格判定は構造のみ (param 値に依らない) — ダミー 0 で walk。 + case mvNormalObserveOf m probeParams of + Nothing -> Nothing + Just _ -> Just closure + where + probeParams :: Map Text Double + probeParams = Map.fromList [ (n, 0) | n <- names ] + + closure :: VS.Vector Double -> (Double, VS.Vector Double) + closure uv = + let us = VS.toList uv + thetaD = [ invTransformF t u | (t, u) <- zip trans us ] + paramD = Map.fromList (zip names thetaD) + in case mvNormalObserveOf m paramD of + Nothing -> (fFullA us, VS.fromList (gradFullA us)) -- 適格判定と矛盾: fallback + Just (muD, covD, ys) + | k == 0 || null chunks -> -- 退化: fallback + (fFullA us, VS.fromList (gradFullA us)) + | otherwise -> + let -- 片道 flat 化のみ (fromLists 不使用): [[Double]] を concat して + -- row-major で Matrix に積む。 戻りの toLists は一切しない。 + sig = LA.matrix k (concat covD) -- N×N + (inv, (lndet, _sgn)) = LA.invlndet sig + muV = LA.fromList muD + ds = [ LA.fromList (map realToFrac ym) - muV | ym <- chunks ] + alphas = [ inv LA.#> d | d <- ds ] -- α_m = Σ⁻¹(y_m-μ) + quadSum = sum [ d LA.<.> a | (d, a) <- zip ds alphas ] + kA = fromIntegral k :: Double + nChA = fromIntegral nCh :: Double + logpObs = nChA * (negate 0.5 * kA * log (2 * pi) - 0.5 * lndet) + - 0.5 * quadSum + -- G = 0.5(Σ_m α_mα_mᵀ - nCh·Σ⁻¹), h = Σ_m α_m + gMat = LA.scale 0.5 + (foldl1 (+) [ LA.outer a a | a <- alphas ] + - LA.scale nChA inv) + -- G/h は **flat Storable Vector のまま** (row-major)。 surrogate 側で + -- index して realToFrac lift = nested list 化 (toLists) を回避。 + gFlat = LA.flatten gMat :: VS.Vector Double -- length k*k, row-major + hVec = foldl1 (+) alphas :: VS.Vector Double -- length k + -- detach surrogate: ∂/∂u + 。 Σ(θ) は model が + -- [[a]] を吐くので concat で 1 回 flat 化し gFlat と flat×flat 内積。 + surrogate :: forall a. (Floating a, Ord a, TrackTag a) => [a] -> a + surrogate uu = + let thetaA = [ invTransformF t u | (t, u) <- zip trans uu ] + paramA = Map.fromList (zip names thetaA) + in case mvNormalObserveOf m paramA of + Just (muA, covA, _) -> + dotFlatL gFlat (concat covA) + dotFlatL hVec muA + Nothing -> 0 + gObs = grad surrogate us -- ∂obs/∂u + gRest = gradRestA us -- ∂(logPrior+logJac)/∂u + vRest = fRestA us -- (logPrior+logJac)(u) + in ( vRest + logpObs + , VS.fromList (zipWith (+) gRest gObs) ) + where k = length muD + chunks = chunksOf k ys + nCh = length chunks + + -- flat Double Vector · [a] リスト内積: Double 側 (G/h) を list 化せず index して + -- realToFrac で lift。 xs (concat covA / muA) だけを 1 回舐める (toLists 往復回避)。 + dotFlatL :: forall a. Floating a => VS.Vector Double -> [a] -> a + dotFlatL gv = go 0 0 + where go !i !acc (x : xs) = go (i + 1) (acc + realToFrac (VS.unsafeIndex gv i) * x) xs + go _ !acc [] = acc + + -- prior + logJac (尤度を除く) — scalar・dense 行列を含まない。 + fRestA :: (Floating a, Ord a, TrackTag a) => [a] -> a + fRestA us = + let paramsC = Map.fromList (zip names [ invTransformF t u | (t, u) <- zip trans us ]) + logJac = sum [ logJacF t u | (t, u) <- zip trans us ] + in logPrior m paramsC + logJac + -- AD 側は 'logDensityRD' 注入版 (Phase 92 B3・値/勾配とも fRestA と bit 一致) + fRestRD :: forall s. Reifies s ADRD.Tape + => [ADRD.ReverseDouble s] -> ADRD.ReverseDouble s + fRestRD us = + let paramsC = Map.fromList (zip names [ invTransformF t u | (t, u) <- zip trans us ]) + logJac = sum [ logJacF t u | (t, u) <- zip trans us ] + in logPriorWith logDensityRD m paramsC + logJac + gradRestA :: [Double] -> [Double] + gradRestA = grad fRestRD + + -- 適格判定が崩れた稀ケース用の完全 walk+ad fallback (従来経路と同一)。 + fFullA :: (Floating a, Ord a, TrackTag a) => [a] -> a + fFullA us = + let paramsC = Map.fromList (zip names [ invTransformF t u | (t, u) <- zip trans us ]) + logJac = sum [ logJacF t u | (t, u) <- zip trans us ] + in logJoint m paramsC + logJac + gradFullA :: [Double] -> [Double] + gradFullA = grad fFullA + +-- --------------------------------------------------------------------------- +-- Phase 95 B-dsl: GP (RBF) 尤度の閉形式随伴 (Cholesky を AD tape に載せない) +-- --------------------------------------------------------------------------- + +-- | [日本語]: 尤度が __単一 @MvNormalGpRBF@ observe__ のモデルから +-- @(x, α, ρ, σ, ys)@ を現在の param 値で抽出する。 それ以外 (他 Observe/ +-- ObserveLM 混在・非 GpRBF) は 'Nothing'。 walk は 'mvNormalObserveOf' と同型。 +-- [English]: For a model whose likelihood is a +-- __single @MvNormalGpRBF@ observe__, extracts @(x, α, ρ, σ, ys)@ +-- at the current parameter values. Otherwise (other Observe/ObserveLM +-- mixed in, or not GpRBF) returns 'Nothing'. The walk has the same +-- shape as 'mvNormalObserveOf'. +gpRBFObserveOf :: (Floating a, Ord a) + => Model a r -> Map Text a -> Maybe ([a], a, a, a, [Double]) +gpRBFObserveOf model params = + case go model of + Just [(MvNormalGpRBF xs al rh sg, ys)] -> Just (xs, al, rh, sg, ys) + _ -> Nothing + where + go (Pure _) = Just [] + go (Free (Sample n _ k)) = + case Map.lookup n params of + Nothing -> Nothing + Just v -> go (k v) + go (Free (Observe _ d ys next)) = ((d, ys) :) <$> go next + go (Free (ObserveLM {})) = Nothing + go (Free (Potential _ _ next)) = go next + go (Free (Deterministic _ v k)) = go (k v) + go (Free (Data _ ys k)) = go (k (map realToFrac ys, ys)) + go (Free (DataIx _ is k)) = go (k is) + go (Free (PlateBegin _ _ next)) = go next + go (Free (PlateEnd next)) = go next + +-- | [日本語]: @MvNormalGpRBF@ 尤度の value+grad を __閉形式随伴__で計算する +-- クロージャを構築する。 適格 (単一 GpRBF observe) のときのみ 'Just'。 +-- +-- === A案 (汎用 detach) との差 = 84% の除去 +-- A案 (§mvNormalAnalyticVG) は surrogate @@ を __AD で微分__するため、 +-- 毎 leaf で Σ(θ) を @[[a]]@ で組み直し (profile: gpExpQuadCov 55%) その全 N² +-- ノードを reverse-AD tape に載せていた (reifyTypeable+partials 29%)。 本経路は: +-- +-- 1. カーネル役割 (x/α/ρ/σ) が @MvNormalGpRBF@ の型で明示されているので、 +-- __∂Σ/∂θ を閉形式__ (@∂Σ/∂α=2K'/α@・@∂Σ/∂ρ=K'∘d²/ρ³@・@∂Σ/∂σ=I@) で書ける。 +-- 2. G=∂logp/∂Σ・K'・d² は __hmatrix Matrix__ (脱リスト) で計算し、 +-- @g_θ = @ を要素積 + trace で Double 算出 (__AD tape ゼロ__)。 +-- 3. u への連鎖律は __軽量 surrogate__ @g_α·α(u)+g_ρ·ρ(u)+g_σ·σ(u)@ を ad。 +-- ad は α/ρ/σ の __3 scalar 抽出のみ__ (cov は非展開・lazy)。 これで +-- 「どの u が α/ρ/σ か」の対応付けを AD が自動処理する (名前直書き不要)。 +-- 4. 値 logp と Σ⁻¹/logdet は同じ LAPACK 分解 (@invlndet@) から。 +-- +-- 距離行列 d² は x (data・定数) から __build 時に 1 度だけ__作り、全 leaf で再利用。 +-- +-- [English]: Builds a closure that computes the value+grad of the +-- @MvNormalGpRBF@ likelihood via a __closed-form adjoint__. Returns +-- 'Just' only when the model qualifies (a single GpRBF observe). +-- +-- === Difference from plan A (generic detach) = 84% removed +-- Plan A (§mvNormalAnalyticVG) +-- __differentiates the surrogate @@ with AD__, so at every +-- leaf it rebuilds Σ(θ) as +-- @[[a]]@ (profile: gpExpQuadCov 55%) and puts all of its N² nodes +-- on the reverse-AD tape (reifyTypeable+partials 29%). This path +-- instead: +-- +-- 1. Since the kernel's roles (x/α/ρ/σ) are made explicit by +-- 'MvNormalGpRBF''s type, __∂Σ/∂θ can be written in closed form__ +-- (@∂Σ/∂α=2K'/α@, @∂Σ/∂ρ=K'∘d²/ρ³@, @∂Σ/∂σ=I@). +-- 2. G=∂logp/∂Σ, K', and d² are computed as __hmatrix Matrix__ +-- values (delisted), and @g_θ = @ is obtained as a +-- Double via elementwise product + trace (__zero AD tape__). +-- 3. The chain rule into u is applied by ad-differentiating a +-- __lightweight surrogate__ @g_α·α(u)+g_ρ·ρ(u)+g_σ·σ(u)@; ad +-- only has to extract __3 scalars__ (α/ρ/σ; the covariance is +-- never expanded, staying lazy). This lets AD automatically +-- handle "which u corresponds to α/ρ/σ" (no hardcoded names +-- needed). +-- 4. The value logp and Σ⁻¹/logdet come from the same LAPACK +-- factorization (@invlndet@). +-- +-- The squared-distance matrix d² is built __once, at build time__, +-- from x (data, a constant), and reused across all leaves. +gpRBFAnalyticVG + :: forall r. ModelP r -> [Text] -> [Transform] + -> Maybe (VS.Vector Double -> (Double, VS.Vector Double)) +gpRBFAnalyticVG m names trans = + case gpRBFObserveOf m probeParams of + Just (xs0, _, _, _, _) | not (null xs0) -> Just (closure (buildD2 xs0)) + _ -> Nothing + where + probeParams :: Map Text Double + probeParams = Map.fromList [ (nm, 0) | nm <- names ] + + -- x (data・定数) から距離² 行列 D2_ij=(x_i-x_j)² を build 時 1 回だけ。 + buildD2 :: [Double] -> LA.Matrix Double + buildD2 xs = + let nn = length xs + xv = VS.fromList xs + in LA.matrix nn [ let d = VS.unsafeIndex xv i - VS.unsafeIndex xv j in d * d + | i <- [0 .. nn - 1], j <- [0 .. nn - 1] ] + + paramMapOf :: forall a. Floating a => [a] -> Map Text a + paramMapOf us = Map.fromList (zip names [ invTransformF t u | (t, u) <- zip trans us ]) + + closure :: LA.Matrix Double -> VS.Vector Double -> (Double, VS.Vector Double) + closure d2 uv = + let us = VS.toList uv + in case gpRBFObserveOf m (paramMapOf us) of + Nothing -> (fFullA us, VS.fromList (gradFullA us)) -- 適格崩れ: fallback + Just (xs, alphaD, rhoD, sigmaD, ys) + | k == 0 || null chunks -> (fFullA us, VS.fromList (gradFullA us)) + | otherwise -> + let -- 純カーネル K'_ij = α² exp(-0.5 D2/ρ²)。exp は hmatrix の + -- element-wise Floating instance (C ベクトル化) を使い、cmap の + -- Haskell ラムダ per-element boxing (2500 boxed Double/leaf) を回避。 + kMat = LA.scale (alphaD * alphaD) + (exp (LA.scale (negate 0.5 / (rhoD * rhoD)) d2)) + -- Σ = K' + (jitter+σ)I。対角に直接加算 (脱 ident: N×N を 3 パス→1)。 + covM = LA.accum kMat (+) [ ((i, i), 1e-10 + sigmaD) | i <- [0 .. k - 1] ] + -- covM = K'(PSD) + (1e-10+σ>0)I ゆえ本来 PD。Cholesky は LU 全逆行列より + -- O(N³) 定数が軽い (PyMC/Stan と同じ経路)。ただし σ→0⁺ の悪条件で LAPACK が + -- 非 PD 判定する稀ケースに備え mbChol で受け、崩れたら full-AD へ安全退避。 + in case LA.mbChol (LA.trustSym covM) of + Nothing -> (fFullA us, VS.fromList (gradFullA us)) -- 非 PD (稀): fallback + Just uchol -> + let inv = LA.cholSolve uchol (LA.ident k) -- Σ⁻¹ (SPD solve) + lndet = 2 * sum (map log (LA.toList (LA.takeDiag uchol))) -- log|Σ|=2Σlog U_ii + ds = [ LA.fromList (map realToFrac yv) | yv <- chunks ] -- (y - μ), μ=0 + alphas = [ inv LA.#> d | d <- ds ] -- Σ⁻¹(y-μ) + quadSum = sum [ d LA.<.> a | (d, a) <- zip ds alphas ] + kA = fromIntegral k :: Double + nChA = fromIntegral nCh :: Double + logpObs = nChA * (negate 0.5 * kA * log (2 * pi) - 0.5 * lndet) + - 0.5 * quadSum + -- 閉形式随伴 g_θ = , G = 0.5(Σ_m α_mα_mᵀ - nCh·Σ⁻¹)。 + -- Frobenius 恒等式で G を materialize せず算出 (N×N 一時行列を全廃): + -- <α_mα_mᵀ, M> = α_mᵀ M α_m (BLAS mat-vec + dot・N² alloc なし)、 + -- <Σ⁻¹, M> = (BLAS ddot・temp なし)。 + -- ∂Σ/∂α=2K'/α・∂Σ/∂σ=I・∂Σ/∂ρ=K'∘d²/ρ³。 + kd2 = kMat * d2 -- ∂Σ/∂ρ 用 Hadamard (1 回だけ) + invFlat = LA.flatten inv + quadK = sum [ a LA.<.> (kMat LA.#> a) | a <- alphas ] -- Σ α_mᵀK'α_m + quadKd2 = sum [ a LA.<.> (kd2 LA.#> a) | a <- alphas ] -- Σ α_mᵀ(K'∘d²)α_m + aaSum = sum [ a LA.<.> a | a <- alphas ] -- Σ α_mᵀα_m + frobIK = invFlat LA.<.> LA.flatten kMat -- <Σ⁻¹, K'> + frobIKd2 = invFlat LA.<.> LA.flatten kd2 -- <Σ⁻¹, K'∘d²> + trInv = LA.sumElements (LA.takeDiag inv) -- tr(Σ⁻¹) + gAlpha = (quadK - nChA * frobIK) / alphaD -- 2·/α + gSigma = 0.5 * (aaSum - nChA * trInv) -- = tr(G) + gRho = 0.5 * (quadKd2 - nChA * frobIKd2) / (rhoD ** 3) -- /ρ³ + -- 軽量 scatter: g_θ を u へ連鎖 (ad は α/ρ/σ の抽出のみ・cov 非展開) + surrogate :: forall a. (Floating a, Ord a, TrackTag a) => [a] -> a + surrogate uu = + case gpRBFObserveOf m (paramMapOf uu) of + Just (_, al, rh, sg, _) -> + realToFrac gAlpha * al + realToFrac gRho * rh + + realToFrac gSigma * sg + Nothing -> 0 + gObs = grad surrogate us -- ∂obs/∂u + gRest = gradRestA us -- ∂(logPrior+logJac)/∂u + vRest = fRestA us -- (logPrior+logJac)(u) + in ( vRest + logpObs + , VS.fromList (zipWith (+) gRest gObs) ) + where k = length xs -- GP 次元 + chunks = chunksOf k ys -- 通常 1 chunk + nCh = length chunks + + -- prior + logJac (尤度除く) — scalar・dense 行列なし。 + fRestA :: (Floating a, Ord a, TrackTag a) => [a] -> a + fRestA us = logPrior m (paramMapOf us) + + sum [ logJacF t u | (t, u) <- zip trans us ] + -- AD 側は 'logDensityRD' 注入版 (Phase 92 B3・値/勾配とも fRestA と bit 一致) + fRestRD :: forall s. Reifies s ADRD.Tape + => [ADRD.ReverseDouble s] -> ADRD.ReverseDouble s + fRestRD us = logPriorWith logDensityRD m (paramMapOf us) + + sum [ logJacF t u | (t, u) <- zip trans us ] + gradRestA :: [Double] -> [Double] + gradRestA = grad fRestRD + + -- 適格崩れ時の完全 walk+ad fallback。 + fFullA :: (Floating a, Ord a, TrackTag a) => [a] -> a + fFullA us = logJoint m (paramMapOf us) + + sum [ logJacF t u | (t, u) <- zip trans us ] + gradFullA :: [Double] -> [Double] + gradFullA = grad fFullA + +-- --------------------------------------------------------------------------- +-- Phase 92 A2: HMM forward 尤度の閉形式随伴 (forward-backward・AD tape ゼロ) +-- --------------------------------------------------------------------------- + +-- | [日本語]: 尤度が __単一 'HmmForwardNormal' observe__ のモデルから +-- @(π_0, trans, μs, σ, ys)@ を現在の param 値で抽出する。 それ以外 (他 Observe/ +-- ObserveLM 混在・非 HMM) は 'Nothing'。 walk は 'gpRBFObserveOf' と同型。 +-- [English]: For a model whose likelihood is a +-- __single 'HmmForwardNormal' observe__, extracts +-- @(π_0, trans, μs, σ, ys)@ at the current parameter values. +-- Otherwise (other Observe/ObserveLM mixed in, or not HMM) returns +-- 'Nothing'. The walk has the same shape as 'gpRBFObserveOf'. +hmmObserveOf :: (Floating a, Ord a) + => Model a r -> Map Text a -> Maybe ([a], [[a]], [a], a, [Double]) +hmmObserveOf model params = + case go model of + Just [(HmmForwardNormal pi0 tr mus sg, ys)] -> Just (pi0, tr, mus, sg, ys) + _ -> Nothing + where + go (Pure _) = Just [] + go (Free (Sample n _ k)) = + case Map.lookup n params of + Nothing -> Nothing + Just v -> go (k v) + go (Free (Observe _ d ys next)) = ((d, ys) :) <$> go next + go (Free (ObserveLM {})) = Nothing + go (Free (Potential _ _ next)) = go next + go (Free (Deterministic _ v k)) = go (k v) + go (Free (Data _ ys k)) = go (k (map realToFrac ys, ys)) + go (Free (DataIx _ is k)) = go (k is) + go (Free (PlateBegin _ _ next)) = go next + go (Free (PlateEnd next)) = go next + +-- | [日本語]: 'HmmForwardNormal' 尤度の value+grad を +-- __forward-backward の閉形式随伴__で計算するクロージャを構築する。 +-- 適格 (単一 HMM observe) のときのみ 'Just'。 構成は @gpRBFAnalyticVG@ と同じ 3 段: +-- +-- 1. forward α / backward β を __Double 空間__ (AD tape 外) で 1 回ずつ回し、 +-- 値 @logL = logSumExp_k α_T[k]@ と閉形式随伴 +-- @∂logL/∂μ_k = Σ_t γ_t[k]·(y_t-μ_k)/σ²@ (γ_t[k]=exp(α_t+β_t-logL))・ +-- @∂logL/∂T_ij = Σ_t exp(α_t[i]+emit_{t+1}[j]+β_{t+1}[j]-logL)@ (ξ 集計)・ +-- @∂logL/∂π_k = γ_0[k]/π_k@・σ も同様、 を Double で算出する。 +-- 2. u への連鎖律は __軽量 surrogate__ @Σ g_θ·θ(u)@ を ad。 ad は +-- π/T/μ/σ の __O(K²) scalar 抽出のみ__ (T 長の forward loop は非展開)。 +-- dirichlet の棒折り deterministic 等の合成は AD が自動処理する。 +-- 3. prior + logJac は 'logPrior' ベースの fRest (走査対象から尤度を除外)。 +-- +-- 従来 walk+ad は T×K² 個の logSumExp/logDensity を毎 leapfrog boxed AD で +-- 再評価していた (数値密度系 84% + AD 6%・alloc 23.7GB/5s)。 +-- 本経路は同じ O(TK²) を unboxed Double で 2 パス回すだけで tape に載せない。 +-- +-- [English]: Builds a closure that computes the value+grad of the +-- 'HmmForwardNormal' likelihood via a +-- __closed-form forward-backward adjoint__. Returns 'Just' only when +-- the model qualifies (a single HMM observe). The construction has +-- the same 3 stages as @gpRBFAnalyticVG@: +-- +-- 1. Run forward α / backward β once each in __Double space__ +-- (outside the AD tape), and compute the value +-- @logL = logSumExp_k α_T[k]@ along with the closed-form adjoints +-- @∂logL/∂μ_k = Σ_t γ_t[k]·(y_t-μ_k)/σ²@ (where +-- γ_t[k]=exp(α_t+β_t-logL)), +-- @∂logL/∂T_ij = Σ_t exp(α_t[i]+emit_{t+1}[j]+β_{t+1}[j]-logL)@ +-- (a ξ aggregate), and @∂logL/∂π_k = γ_0[k]/π_k@ (σ likewise), all +-- in Double. +-- 2. The chain rule into u is applied by ad-differentiating a +-- __lightweight surrogate__ @Σ g_θ·θ(u)@; ad only has to extract +-- __O(K²) scalars__ for π/T/μ/σ (the length-T forward loop is +-- never expanded). AD automatically handles compositions such as +-- the Dirichlet stick-breaking deterministic. +-- 3. The prior + logJac term is 'logPrior'-based fRest (excluding the +-- likelihood from the traversal). +-- +-- The previous walk+ad approach re-evaluated T×K² logSumExp/logDensity +-- calls with boxed AD on every leapfrog step (numerical density +-- family 84% + AD 6%, alloc 23.7GB/5s). This path instead runs the +-- same O(TK²) work as two passes over unboxed Doubles, never touching +-- the tape. +hmmAnalyticVG + :: forall r. ModelP r -> [Text] -> [Transform] + -> Maybe (VS.Vector Double -> (Double, VS.Vector Double)) +hmmAnalyticVG m names trans = + case hmmObserveOf m probeParams of + Just (pi0, tr, mus, _, ys) + | kDim > 0, length tr == kDim, all ((== kDim) . length) tr + , length mus == kDim, not (null ys) -> Just closure + where kDim = length pi0 + _ -> Nothing + where + probeParams :: Map Text Double + probeParams = Map.fromList [ (nm, 0) | nm <- names ] + + paramMapOf :: forall a. Floating a => [a] -> Map Text a + paramMapOf us = Map.fromList (zip names [ invTransformF t u | (t, u) <- zip trans us ]) + + closure :: VS.Vector Double -> (Double, VS.Vector Double) + closure uv = + let us = VS.toList uv + in case hmmObserveOf m (paramMapOf us) of + Nothing -> (fFullA us, VS.fromList (gradFullA us)) -- 適格崩れ: fallback + Just (pi0D, transD, musD, sgD, ys) + | sgD <= 0 || null ys -> (fFullA us, VS.fromList (gradFullA us)) + | otherwise -> + -- B2-① (2026-07-17): 脱リスト化 — α/β/emit を unboxed 行 vector で持ち、 + -- 内側 K ループは lseK (list 非 alloc の 2 パス logSumExp)。γ は非実体化。 + let kk = length pi0D + tT = length ys + ixs = [0 .. kk - 1] + ysV = VU.fromList ys + musV = VU.fromList musD + lPi0 = VU.fromList (map safeLog pi0D) + lTr = VU.fromList (map safeLog (concat transD)) -- 行優先 K×K flat + lsg = log sgD + c2pi = 0.5 * log (2 * pi) + emitAt t k' = let z = (VU.unsafeIndex ysV t - VU.unsafeIndex musV k') / sgD + in -0.5 * z * z - lsg - c2pi + emitRows = [ VU.generate kk (emitAt t) | t <- [0 .. tT - 1] ] + -- K 要素 logSumExp (max → sumexp の 2 パス・中間 list 無し) + lseK f = let mx = foldl' (\acc i -> max acc (f i)) negInf ixs + in if mx == negInf then negInf + else mx + log (foldl' (\acc i -> acc + exp (f i - mx)) 0 ixs) + -- forward: α_t 全行を保持 (随伴の γ/ξ に使う) + alpha0 = VU.zipWith (+) lPi0 (head emitRows) + stepF aPrev emT = VU.generate kk $ \j -> + lseK (\i -> VU.unsafeIndex aPrev i + VU.unsafeIndex lTr (i * kk + j)) + + VU.unsafeIndex emT j + alphaRows = scanl stepF alpha0 (tail emitRows) -- 長さ T + logL = lseK (VU.unsafeIndex (last alphaRows)) + -- backward: β_{T-1}=0・β_t[i] = lse_j (lT_ij + emit_{t+1}[j] + β_{t+1}[j]) + stepB emNext bNext = VU.generate kk $ \i -> + lseK (\j -> VU.unsafeIndex lTr (i * kk + j) + + VU.unsafeIndex emNext j + VU.unsafeIndex bNext j) + betaRows = scanr stepB (VU.replicate kk 0) (tail emitRows) -- 長さ T + -- 閉形式随伴 (全て Double・tape ゼロ・γ_t[k] = exp(α+β-logL) は都度計算) + gammaAt aR bR k' = exp (VU.unsafeIndex aR k' + VU.unsafeIndex bR k' - logL) + abRows = zip3 alphaRows betaRows [0 ..] + gMu = [ foldl' (\acc (aR, bR, t) -> + acc + gammaAt aR bR k' + * (VU.unsafeIndex ysV t - VU.unsafeIndex musV k') + / (sgD * sgD)) + 0 abRows + | k' <- ixs ] + gSg = foldl' (\acc (aR, bR, t) -> + foldl' (\a2 k' -> + let z = (VU.unsafeIndex ysV t + - VU.unsafeIndex musV k') / sgD + in a2 + gammaAt aR bR k' * (z * z - 1) / sgD) + acc ixs) + 0 abRows + gPi0 = [ if p > 0 then gammaAt (head alphaRows) (head betaRows) k' / p else 0 + | (k', p) <- zip ixs pi0D ] + -- ξ 集計: ∂logL/∂T_ij = Σ_{t + acc + exp (VU.unsafeIndex aR i + VU.unsafeIndex emN j + + VU.unsafeIndex bN j - logL)) + 0 xiRows + | j <- ixs ] + | i <- ixs ] + -- 軽量 scatter: g_θ を u へ連鎖 (ad は π/T/μ/σ の抽出のみ・T loop 非展開) + surrogate :: forall a. (Floating a, Ord a, TrackTag a) => [a] -> a + surrogate uu = + case hmmObserveOf m (paramMapOf uu) of + Just (p0, tr', ms, s, _) -> + sum (zipWith (\g v -> realToFrac g * v) gPi0 p0) + + sum (zipWith (\gr r -> sum (zipWith (\g v -> realToFrac g * v) gr r)) + gTr tr') + + sum (zipWith (\g v -> realToFrac g * v) gMu ms) + + realToFrac gSg * s + Nothing -> 0 + -- B2-② (2026-07-17): prior+logJac と surrogate を 1 本の AD tape に合流し + -- grad' で値+勾配を同時取得 (walk 4 回/eval → Double 1 + AD 1)。 + -- fRest の値は vComb から surrogate の Double 値を引いて復元する。 + -- B3: prior 密度は 'logDensityRD' 注入 (fRestRD) = 定数 + -- hyperparam の lgamma 正規化項を Double へ畳み込み (bit 一致)。 + fCombRD :: forall s. Reifies s ADRD.Tape + => [ADRD.ReverseDouble s] -> ADRD.ReverseDouble s + fCombRD uu = fRestRD uu + surrogate uu + (vComb, gComb) = grad' fCombRD us + surrAtUs = sum (zipWith (*) gPi0 pi0D) + + sum (zipWith (\gr r -> sum (zipWith (*) gr r)) gTr transD) + + sum (zipWith (*) gMu musD) + + gSg * sgD + in ( (vComb - surrAtUs) + logL + , VS.fromList gComb ) + + safeLog :: Double -> Double + safeLog x = if x <= 0 then negInf else log x + + -- prior + logJac (尤度除く) — @gpRBFAnalyticVG@ と同じ。 + fRestA :: (Floating a, Ord a, TrackTag a) => [a] -> a + fRestA us = logPrior m (paramMapOf us) + + sum [ logJacF t u | (t, u) <- zip trans us ] + gradRestA :: [Double] -> [Double] + gradRestA = grad fRestA + + -- fRestA の AD 特化 (Phase 92 B3): 'logDensityRD' 注入で定数 hyperparam の + -- lgamma 正規化項を Double へ畳み込む。 値・勾配とも fRestA と bit 一致 + -- ('logDensityRD' の注釈参照)。 + fRestRD :: forall s. Reifies s ADRD.Tape + => [ADRD.ReverseDouble s] -> ADRD.ReverseDouble s + fRestRD us = logPriorWith logDensityRD m (paramMapOf us) + + sum [ logJacF t u | (t, u) <- zip trans us ] + + -- 適格崩れ時の完全 walk+ad fallback。 + fFullA :: (Floating a, Ord a, TrackTag a) => [a] -> a + fFullA us = logJoint m (paramMapOf us) + + sum [ logJacF t u | (t, u) <- zip trans us ] + gradFullA :: [Double] -> [Double] + gradFullA = grad fFullA + +-- --------------------------------------------------------------------------- +-- Phase 101 A2: ARMA(1,1) 尤度の閉形式随伴 (逆向き随伴再帰・AD tape ゼロ) +-- --------------------------------------------------------------------------- + +-- | [日本語]: 尤度が __単一 'ArmaNormal' observe__ のモデルから +-- @(μ, φ, θ, σ, ys)@ を現在の param 値で抽出する。 それ以外 (他 Observe 混在・ +-- 非 ARMA) は 'Nothing'。 walk は 'hmmObserveOf' と同型。 +-- [English]: For a model whose likelihood is a +-- __single 'ArmaNormal' observe__, extracts @(μ, φ, θ, σ, ys)@ at +-- the current parameter values. Otherwise (other Observe mixed in, +-- or not ARMA) returns 'Nothing'. The walk has the same shape as +-- 'hmmObserveOf'. +armaObserveOf :: (Floating a, Ord a) + => Model a r -> Map Text a -> Maybe (a, a, a, a, [Double]) +armaObserveOf model params = + case go model of + Just [(ArmaNormal mu phi theta sg, ys)] -> Just (mu, phi, theta, sg, ys) + _ -> Nothing + where + go (Pure _) = Just [] + go (Free (Sample n _ k)) = + case Map.lookup n params of + Nothing -> Nothing + Just v -> go (k v) + go (Free (Observe _ d ys next)) = ((d, ys) :) <$> go next + go (Free (ObserveLM {})) = Nothing + go (Free (Potential _ _ next)) = go next + go (Free (Deterministic _ v k)) = go (k v) + go (Free (Data _ ys k)) = go (k (map realToFrac ys, ys)) + go (Free (DataIx _ is k)) = go (k is) + go (Free (PlateBegin _ _ next)) = go next + go (Free (PlateEnd next)) = go next + +-- | [日本語]: 'ArmaNormal' 尤度の value+grad を +-- __逆向き 1 パスの閉形式随伴__で計算するクロージャを構築する。 +-- 適格 (単一 ArmaNormal observe) のときのみ 'Just'。 構成は @hmmAnalyticVG@ と同じ 3 段: +-- +-- 1. err 前向き再帰 (@e_1 = y_1 − (μ+φμ)@・@e_t = y_t − μ − φ·y_{t−1} − +-- θ·e_{t−1}@) と随伴の逆向き再帰 (@ē_t = −e_t/σ² − θ·ē_{t+1}@・ +-- @ē_T = −e_T/σ²@) を __Double 空間__ (AD tape 外) で 1 回ずつ回し、 +-- 値 @logL = Σ_t log N(e_t; 0, σ)@ と閉形式随伴 +-- @∂logL/∂μ = ē_1·(−(1+φ)) − Σ_{t≥2} ē_t@・ +-- @∂logL/∂φ = ē_1·(−μ) − Σ_{t≥2} ē_t·y_{t−1}@・ +-- @∂logL/∂θ = −Σ_{t≥2} ē_t·e_{t−1}@・ +-- @∂logL/∂σ = −T/σ + (Σ_t e_t²)/σ³@ を算出する。 +-- 2. u への連鎖律は __軽量 surrogate__ @Σ g_θ·θ(u)@ を ad (μ/φ/θ/σ の +-- 4 scalar 抽出のみ・T 長の再帰は非展開)。 +-- 3. prior + logJac は 'logDensityRD' 注入の fRestRD (前述の @hmmAnalyticVG@ +-- と同じ手法)。 +-- +-- 従来 walk+ad は T 本の @logDensity@ + mapAccumL 再帰を毎 leapfrog boxed AD +-- で再評価していた (logDensity 31.2% + armaModel 20.9%・ +-- alloc 72%)。 本経路は同じ O(T) を unboxed Double で 2 パス回すだけ。 +-- +-- [English]: Builds a closure that computes the value+grad of the +-- 'ArmaNormal' likelihood via a +-- __closed-form, single reverse-pass adjoint__. Returns 'Just' only +-- when the model qualifies (a single ArmaNormal observe). The +-- construction has the same 3 stages as @hmmAnalyticVG@: +-- +-- 1. Run the err forward recursion (@e_1 = y_1 − (μ+φμ)@, +-- @e_t = y_t − μ − φ·y_{t−1} − θ·e_{t−1}@) and the adjoint's +-- backward recursion (@ē_t = −e_t/σ² − θ·ē_{t+1}@, +-- @ē_T = −e_T/σ²@) once each in __Double space__ (outside the AD +-- tape), and compute the value @logL = Σ_t log N(e_t; 0, σ)@ +-- along with the closed-form adjoints +-- @∂logL/∂μ = ē_1·(−(1+φ)) − Σ_{t≥2} ē_t@, +-- @∂logL/∂φ = ē_1·(−μ) − Σ_{t≥2} ē_t·y_{t−1}@, +-- @∂logL/∂θ = −Σ_{t≥2} ē_t·e_{t−1}@, and +-- @∂logL/∂σ = −T/σ + (Σ_t e_t²)/σ³@. +-- 2. The chain rule into u is applied by ad-differentiating a +-- __lightweight surrogate__ @Σ g_θ·θ(u)@ (extracting only the 4 +-- scalars μ/φ/θ/σ; the length-T recursion is never expanded). +-- 3. The prior + logJac term is the 'logDensityRD'-injected fRestRD +-- (the same technique used earlier in @hmmAnalyticVG@). +-- +-- The previous walk+ad approach re-evaluated T calls to @logDensity@ +-- plus a mapAccumL recursion with boxed AD on every leapfrog step +-- (logDensity 31.2% + armaModel 20.9%, alloc 72%). This path instead +-- runs the same O(T) work as two passes over unboxed Doubles. +armaAnalyticVG + :: forall r. ModelP r -> [Text] -> [Transform] + -> Maybe (VS.Vector Double -> (Double, VS.Vector Double)) +armaAnalyticVG m names trans = + case armaObserveOf m probeParams of + Just (_, _, _, _, ys) | not (null ys) -> Just closure + _ -> Nothing + where + probeParams :: Map Text Double + probeParams = Map.fromList [ (nm, 0) | nm <- names ] + + paramMapOf :: forall a. Floating a => [a] -> Map Text a + paramMapOf us = Map.fromList (zip names [ invTransformF t u | (t, u) <- zip trans us ]) + + closure :: VS.Vector Double -> (Double, VS.Vector Double) + closure uv = + let us = VS.toList uv + in case armaObserveOf m (paramMapOf us) of + Nothing -> (fFullA us, VS.fromList (gradFullA us)) -- 適格崩れ: fallback + Just (muD, phiD, thD, sgD, ys) + | sgD <= 0 || null ys -> (fFullA us, VS.fromList (gradFullA us)) + | otherwise -> + let ysV = VU.fromList ys + tT = VU.length ysV + s2 = sgD * sgD + -- forward: err 列 (unboxed・prefix 参照の constructN) + errsV = VU.constructN tT $ \pre -> + let i = VU.length pre + in if i == 0 + then VU.unsafeIndex ysV 0 - (muD + phiD * muD) + else VU.unsafeIndex ysV i + - (muD + phiD * VU.unsafeIndex ysV (i - 1) + + thD * VU.unsafeIndex pre (i - 1)) + sumE2 = VU.foldl' (\acc e -> acc + e * e) 0 errsV + logL = fromIntegral tT * (-0.5 * log (2 * pi) - log sgD) + - sumE2 / (2 * s2) + -- backward: 随伴 ē (suffix 参照の constructrN・ē_T = −e_T/σ²) + ebarV = VU.constructrN tT $ \suf -> + let t = tT - 1 - VU.length suf + direct = negate (VU.unsafeIndex errsV t) / s2 + in if VU.null suf + then direct + else direct - thD * VU.unsafeIndex suf 0 + -- 閉形式随伴 (全て Double・tape ゼロ) + gMu = VU.unsafeIndex ebarV 0 * negate (1 + phiD) + - VU.ifoldl' (\acc t eb -> if t == 0 then acc else acc + eb) + 0 ebarV + gPhi = VU.unsafeIndex ebarV 0 * negate muD + - VU.ifoldl' (\acc t eb -> + if t == 0 then acc + else acc + eb * VU.unsafeIndex ysV (t - 1)) + 0 ebarV + gTh = negate (VU.ifoldl' (\acc t eb -> + if t == 0 then acc + else acc + eb * VU.unsafeIndex errsV (t - 1)) + 0 ebarV) + gSg = negate (fromIntegral tT) / sgD + sumE2 / (s2 * sgD) + -- 軽量 scatter: g_θ を u へ連鎖 (ad は μ/φ/θ/σ の 4 scalar 抽出のみ) + surrogate :: forall a. (Floating a, Ord a, TrackTag a) => [a] -> a + surrogate uu = + case armaObserveOf m (paramMapOf uu) of + Just (mu', phi', th', sg', _) -> + realToFrac gMu * mu' + realToFrac gPhi * phi' + + realToFrac gTh * th' + realToFrac gSg * sg' + Nothing -> 0 + fCombRD :: forall s. Reifies s ADRD.Tape + => [ADRD.ReverseDouble s] -> ADRD.ReverseDouble s + fCombRD uu = fRestRD uu + surrogate uu + (vComb, gComb) = grad' fCombRD us + surrAtUs = gMu * muD + gPhi * phiD + gTh * thD + gSg * sgD + in ( (vComb - surrAtUs) + logL + , VS.fromList gComb ) + + -- prior + logJac (尤度除く)・'logDensityRD' 注入 — @hmmAnalyticVG@ と同じ。 + fRestRD :: forall s. Reifies s ADRD.Tape + => [ADRD.ReverseDouble s] -> ADRD.ReverseDouble s + fRestRD us = logPriorWith logDensityRD m (paramMapOf us) + + sum [ logJacF t u | (t, u) <- zip trans us ] + + -- 適格崩れ時の完全 walk+ad fallback。 + fFullA :: (Floating a, Ord a, TrackTag a) => [a] -> a + fFullA us = logJoint m (paramMapOf us) + + sum [ logJacF t u | (t, u) <- zip trans us ] + gradFullA :: [Double] -> [Double] + gradFullA = grad fFullA + +-- --------------------------------------------------------------------------- +-- Phase 101 A3: graded response IRT 尤度の解析勾配 (AD tape ゼロ) +-- --------------------------------------------------------------------------- + +-- | [日本語]: 尤度が __単一 'GradedResponseIrt' observe__ のモデルから +-- @(θs, ncats, δs, γs, ys)@ を現在の param 値で抽出する。 walk は +-- 'armaObserveOf' と同型。 +-- [English]: For a model whose likelihood is a +-- __single 'GradedResponseIrt' observe__, extracts +-- @(θs, ncats, δs, γs, ys)@ at the current parameter values. The +-- walk has the same shape as 'armaObserveOf'. +gradedIrtObserveOf :: (Floating a, Ord a) + => Model a r -> Map Text a + -> Maybe ([a], [Int], [Double], [[Double]], [Double]) +gradedIrtObserveOf model params = + case go model of + Just [(GradedResponseIrt ths ncats dls gms, ys)] -> Just (ths, ncats, dls, gms, ys) + _ -> Nothing + where + go (Pure _) = Just [] + go (Free (Sample n _ k)) = + case Map.lookup n params of + Nothing -> Nothing + Just v -> go (k v) + go (Free (Observe _ d ys next)) = ((d, ys) :) <$> go next + go (Free (ObserveLM {})) = Nothing + go (Free (Potential _ _ next)) = go next + go (Free (Deterministic _ v k)) = go (k v) + go (Free (Data _ ys k)) = go (k (map realToFrac ys, ys)) + go (Free (DataIx _ is k)) = go (k is) + go (Free (PlateBegin _ _ next)) = go next + go (Free (PlateEnd next)) = go next + +-- | [日本語]: 'GradedResponseIrt' 尤度の value+grad を __解析勾配__で +-- 計算するクロージャを構築する。 適格 (単一 GradedResponseIrt observe) の +-- ときのみ 'Just'。 構成は @armaAnalyticVG@ と同じ 3 段: +-- +-- 1. 各 (child i, item j, grade≠−1) の @Q_k = invlogit(δ_j(θ_i−γ_jk))@ と +-- カテゴリ確率 p (隣接差) を __Double 空間__で評価し、 値 +-- @logL = Σ log p@ と解析勾配 @∂logL/∂θ_i = Σ_j (dp/dθ)/p@ +-- (@dQ/dθ = δ·Q(1−Q)@ の隣接差) を算出する。 +-- 2. u への連鎖律は __軽量 surrogate__ @Σ g_i·θ_i(u)@ を ad +-- (θs の nChild scalar 抽出のみ)。 +-- 3. prior + logJac は 'logDensityRD' 注入の fRestRD。 +-- +-- 従来 walk+ad は nChild×nItem×ncat の Q/p リスト構築 (`!!` 索引込) を毎 +-- leapfrog boxed AD で再評価していた (logCatProb 64.8% time / +-- 73.2% alloc)。 本経路は同じ O(Σ ncat) を Double で 1 パス回すだけ。 +-- +-- [English]: Builds a closure that computes the value+grad of the +-- 'GradedResponseIrt' likelihood via an __analytic gradient__. Returns +-- 'Just' only when the model qualifies (a single GradedResponseIrt +-- observe). The construction has the same 3 stages as +-- @armaAnalyticVG@: +-- +-- 1. For each (child i, item j, grade≠−1), evaluate +-- @Q_k = invlogit(δ_j(θ_i−γ_jk))@ and the category probability p +-- (an adjacent difference) in __Double space__, and compute the +-- value @logL = Σ log p@ and the analytic gradient +-- @∂logL/∂θ_i = Σ_j (dp/dθ)/p@ (an adjacent difference of +-- @dQ/dθ = δ·Q(1−Q)@). +-- 2. The chain rule into u is applied by ad-differentiating a +-- __lightweight surrogate__ @Σ g_i·θ_i(u)@ (extracting only the +-- nChild scalars of θs). +-- 3. The prior + logJac term is the 'logDensityRD'-injected fRestRD. +-- +-- The previous walk+ad approach re-evaluated the nChild×nItem×ncat +-- Q/p list construction (including @!!@ indexing) with boxed AD on +-- every leapfrog step (logCatProb 64.8% time / 73.2% alloc). This +-- path instead runs the same O(Σ ncat) work as a single pass over +-- Doubles. +gradedIrtAnalyticVG + :: forall r. ModelP r -> [Text] -> [Transform] + -> Maybe (VS.Vector Double -> (Double, VS.Vector Double)) +gradedIrtAnalyticVG m names trans = + case gradedIrtObserveOf m probeParams of + Just (ths, ncats, dls, gms, ys) + | not (null ths), not (null ys) + , length ncats == length dls, length ncats == length gms -> Just closure + _ -> Nothing + where + probeParams :: Map Text Double + probeParams = Map.fromList [ (nm, 0) | nm <- names ] + + paramMapOf :: forall a. Floating a => [a] -> Map Text a + paramMapOf us = Map.fromList (zip names [ invTransformF t u | (t, u) <- zip trans us ]) + + closure :: VS.Vector Double -> (Double, VS.Vector Double) + closure uv = + let us = VS.toList uv + in case gradedIrtObserveOf m (paramMapOf us) of + Nothing -> (fFullA us, VS.fromList (gradFullA us)) -- 適格崩れ: fallback + Just (thsD, ncats, dls, gms, ys) + | null ys -> (fFullA us, VS.fromList (gradFullA us)) + | otherwise -> + let nItem = length ncats + rows = chunksOf nItem ys + -- (logL, g) を child i 毎に Double で 1 パス集計 + childLG th row = foldl' step (0, 0) (zip4 ncats dls gms row) + where + step (accL, accG) (nc, dl, gm, grD) + | grD == -1 = (accL, accG) + | otherwise = + let gr = round grD :: Int + kMax = nc - 1 + q kk = 1 / (1 + exp (negate (dl * (th - gm !! (kk - 1))))) + dq kk = let qv = q kk in dl * qv * (1 - qv) + (p, dp) + | gr == 1 = (1 - q 1, negate (dq 1)) + | gr == nc = (q kMax, dq kMax) + | otherwise = (q (gr - 1) - q gr, dq (gr - 1) - dq gr) + in (accL + log p, accG + dp / p) + lgs = [ childLG th row | (th, row) <- zip thsD rows ] + logL = sum (map fst lgs) + gThs = map snd lgs + -- 軽量 scatter: g_i を u へ連鎖 (ad は θs の scalar 抽出のみ) + surrogate :: forall a. (Floating a, Ord a, TrackTag a) => [a] -> a + surrogate uu = + case gradedIrtObserveOf m (paramMapOf uu) of + Just (ths', _, _, _, _) -> + sum (zipWith (\g v -> realToFrac g * v) gThs ths') + Nothing -> 0 + fCombRD :: forall s. Reifies s ADRD.Tape + => [ADRD.ReverseDouble s] -> ADRD.ReverseDouble s + fCombRD uu = fRestRD uu + surrogate uu + (vComb, gComb) = grad' fCombRD us + surrAtUs = sum (zipWith (*) gThs thsD) + in ( (vComb - surrAtUs) + logL + , VS.fromList gComb ) + + -- prior + logJac (尤度除く)・'logDensityRD' 注入 — @armaAnalyticVG@ と同じ。 + fRestRD :: forall s. Reifies s ADRD.Tape + => [ADRD.ReverseDouble s] -> ADRD.ReverseDouble s + fRestRD us = logPriorWith logDensityRD m (paramMapOf us) + + sum [ logJacF t u | (t, u) <- zip trans us ] + + -- 適格崩れ時の完全 walk+ad fallback。 + fFullA :: (Floating a, Ord a, TrackTag a) => [a] -> a + fFullA us = logJoint m (paramMapOf us) + + sum [ logJacF t u | (t, u) <- zip trans us ] + gradFullA :: [Double] -> [Double] + gradFullA = grad fFullA + +-- | [日本語]: 'gradADU' の __静的部分__ (Gaussian LM ブロック抽出・ +-- 設計列のベクトル化・名前→index 解決・@ad@ クロージャ構築) を 1 度だけ行い、 +-- unconstrained ベクトルを受けて勾配ベクトルを返すクロージャを構築する。 +-- NUTS / HMC は draw ループの __外__ で 1 度呼び、 全 leapfrog で再利用する。 +-- +-- per-op 計測 (prof-nuts-54.4e.prof) で per-call の Text-key +-- `Map.fromList` 組立 + `Map.fromListWith` 勾配集約 (compileGradU self 17.9%) と +-- vec-tape の演算毎ベクトル割当 (~52%) が残ボトルネックと確定 → 名前は compile +-- 時に index へ解決し、 勾配は ST mutable vector に解析閉形式で直接集約する +-- (Gaussian LM の勾配は ∂β_k=X_kᵀr/σ²・∂u_j=Σ_{i∈g_j}r_i/σ²・∂σ=-n/σ+sumR2/σ³ +-- の閉形式ゆえ汎用 tape 不要)。 +-- [English]: Builds, once only, the __static part__ of 'gradADU' +-- (Gaussian LM block extraction, design-column vectorization, +-- name→index resolution, @ad@ closure construction), producing a +-- closure that takes an unconstrained vector and returns a gradient +-- vector. NUTS / HMC calls this once, __outside__ the draw loop, and +-- reuses it across all leapfrog steps. +-- +-- Per-op measurement (prof-nuts-54.4e.prof) established that the +-- remaining bottleneck was the per-call Text-key `Map.fromList` +-- assembly + `Map.fromListWith` gradient aggregation (compileGradU +-- self 17.9%) and the per-operation vector allocation in the +-- vec-tape (~52%); so names are resolved to indices at compile time, +-- and gradients are aggregated directly, in closed analytic form, +-- into an ST mutable vector (the Gaussian LM gradient has the closed +-- form ∂β_k=X_kᵀr/σ², ∂u_j=Σ_{i∈g_j}r_i/σ², ∂σ=-n/σ+sumR2/σ³, so no +-- generic tape is needed). +compileGradUV :: forall r. ModelP r -> [Text] -> [Transform] + -> (VS.Vector Double -> VS.Vector Double) +compileGradUV m names trans = + case gaussLMBlocksAuto m of + ([], _) -> case synthVecIR m of + -- Phase 95: 尤度が単一 dense MvNormal observe なら解析随伴。 Gp-RBF (B-dsl・ + -- 閉形式随伴) を最優先、 次に汎用 MvNormal (A案・flat detach)、 いずれも + -- 非適格なら従来の全体 ad (後方互換)。 + Nothing -> case gpRBFAnalyticVG m names trans of + Just vg -> \uv -> snd (vg uv) + Nothing -> case mvNormalAnalyticVG m names trans of + Just vg -> \uv -> snd (vg uv) + Nothing -> \uv -> VS.fromList (gradFull (VS.toList uv)) + Just (gs, fams, sObs) -> -- 54.11: ベクトル式 IR (非線形 μ) + let ixOf = Map.fromList (zip names [0 ..]) + nP = length names + transB = BV.fromList trans + cvi = compileVecIR ixOf gs fams + famSet = Set.fromList (concat [ ms | (ms, _, _) <- fams ]) + cps = constPriorsOf m famSet + lnGroups = collectLogNormalGroups m -- Phase 98 A3: LogNormal 群 + lnUNames = concat [ us | (us, _, _) <- lnGroups ] + lnIx = map (resolveLogNormal ixOf) lnGroups + exclNames = sObs `Set.union` famSet + `Set.union` Set.fromList (map fst cps) + `Set.union` Set.fromList lnUNames + noResid = residualFreeOfDensity exclNames m + cpIx = [ (ixOf Map.! n, d) | (n, d) <- cps ] + mPriorGrad + | noResid = Nothing + | otherwise = Just (grad (fExcl (compileResidual exclNames m) exclNames)) + in \uv -> + let pc = VS.generate nP $ \i -> + invTransformF (transB BV.! i) (uv `VS.unsafeIndex` i) + mgc = runST $ do + mg <- VSM.replicate nP 0 + ok <- gradVecIR cvi pc mg + if ok + then do + mapM_ (\(i, d) -> + case constPriorGradD d (pc `VS.unsafeIndex` i) of + Just g -> VSM.modify mg (+ g) i + Nothing -> pure ()) cpIx + mapM_ (\ln -> gradLogNormalIx ln pc mg) lnIx -- A3 + Just <$> VS.unsafeFreeze mg + else pure Nothing + in case mgc of + -- guard 違反 (観測項が定数 -∞ の境界領域・例 invLogit の FP + -- 飽和 p==1): walk+ad に per-call fallback し従来経路と同一の + -- 勾配 (違反行 = 定数 → 勾配 0・他の行は有効) を返す。 + -- ★旧 tape (54.11-55.4) は unguarded で NaN が全勾配を汚染し + -- NUTS が max-depth 迷走する潜在バグだった (56.2 で修正・ + -- 'constPriorGradD' の「guard 違反 = 勾配 0 で ad と一致」 と + -- 同じ原則)。 境界点のみの稀ケースゆえ per-draw 影響なし。 + Nothing -> VS.fromList (gradFull (VS.toList uv)) + Just gC -> case mPriorGrad of + Nothing -> + VS.generate nP $ \i -> + let t = transB BV.! i + u = uv `VS.unsafeIndex` i + in gC `VS.unsafeIndex` i * dInvTransform t u + + dLogJacU t u + Just priorGrad -> + let pg = priorGrad (VS.toList uv) + in VS.fromList + [ pg_i + gC `VS.unsafeIndex` i * dInvTransform t u + | (i, (pg_i, (t, u))) <- zip [0 :: Int ..] + (zip pg (zip trans (VS.toList uv))) ] + (gbs, synthObs) -> -- ハイブリッド (静的 hoist) + let -- 54.4c: REff (Just scale) の u-prior は解析勾配・u_j を ad から除外。 + -- 54.4e: 定数パラメタ prior も解析勾配・ad から除外。 密度項が残らな + -- ければ ad クロージャを丸ごと省略 (logJac 勾配も解析式)。 + (priorREs, cps, exclNames, cblocks, hierGroups, noResid) = analyzeGaussModel m gbs synthObs + ixOf = Map.fromList (zip names [0 ..]) + nP = length names + transB = BV.fromList trans -- boxed (Storable 不可) + cbIx = map (resolveLMBlock ixOf) cblocks + reIx = [ ReffPriorIx (VU.fromList (map (ixOf Map.!) uNames)) + (ixOf Map.! scaleName) + | (uNames, scaleName) <- priorREs ] + hniIx = map (resolveHierNormal ixOf) hierGroups + cpIx = [ (ixOf Map.! n, d) | (n, d) <- cps ] + mPriorGrad -- 残 prior 等の ad (fallback) + | noResid = Nothing + | otherwise = Just (grad (fExcl (compileResidual exclNames m) exclNames)) + in hybridGradClosure nP transB trans + (\pc mg -> do + mapM_ (\cb -> gradLMBlockIx cb pc mg) cbIx + mapM_ (\ri -> gradReffPriorIx ri pc mg) reIx + mapM_ (\hn -> gradHierNormalIx hn pc mg) hniIx + mapM_ (\(i, d) -> + case constPriorGradD d (pc `VS.unsafeIndex` i) of + Just g -> VSM.modify mg (+ g) i + Nothing -> pure ()) cpIx) + mPriorGrad + where + gradFull = grad fFull + fFull us = + let paramsC = Map.fromList + [ (n, invTransformF t u) | (n, t, u) <- zip3 names trans us ] + logJac = sum [ logJacF t u | (t, u) <- zip trans us ] + in logJoint m paramsC + logJac + fExcl mcr excl us = + let paramsC = Map.fromList + [ (n, invTransformF t u) | (n, t, u) <- zip3 names trans us ] + logJac = sum [ logJacF t u | (t, u) <- zip trans us ] + in residualExcl mcr excl m paramsC + logJac + +-- | [日本語]: 'compileGradUV' の value-and-grad 融合版 (JAX @value_and_grad@ +-- 相当)。 返り値 = (logπ(u) (logJac 込・'compileLogPUV' と同値)・∇logπ(u) +-- ('compileGradUV' と同値))。 NUTS の葉は leapfrog 最終勾配とエネルギー (logπ) +-- を同一点で別々に評価していた (prof 実測で葉 logPi が全体の 19%) — 本閉包は +-- forward pass を 1 度だけ走らせて両方を返す。 経路分岐・fallback 意味論は +-- 'compileGradUV' / 'compileLogPUV' と対 (vecIR guard 違反 = 値 -∞ + +-- 勾配 walk+ad fallback)。 +-- [English]: A value-and-grad fused version of 'compileGradUV' +-- (equivalent to JAX's @value_and_grad@). Returns +-- (logπ(u) (including logJac; equal to 'compileLogPUV'), +-- ∇logπ(u) (equal to 'compileGradUV')). NUTS's leaf nodes used to +-- evaluate the final leapfrog gradient and the energy (logπ) +-- separately at the same point (profiling measured leaf logPi at +-- 19% of the total) — this closure instead runs the forward pass +-- once and returns both. The branch selection and fallback semantics +-- mirror 'compileGradUV' / 'compileLogPUV' (a vecIR guard violation +-- means value = -∞ plus a walk+ad gradient fallback). +compileGradValUV :: forall r. ModelP r -> [Text] -> [Transform] + -> (VS.Vector Double -> (Double, VS.Vector Double)) +compileGradValUV m names trans = + case gradValPlan m names trans of + GVPure f -> f + GVVecIR sz prep core finish -> \uv -> + let pc = prep uv + mvg = runST $ do + ar <- VSM.unsafeNew sz + adj <- VSM.unsafeNew sz + core ar adj pc + in finish uv pc mvg + +-- | [日本語]: 'compileGradValUV' の monadic 版。 vecIR 経路の作業 +-- バッファ (forward arena + 随伴 arena・13-traffic 実測で 34k セル × 2 = +-- 葉勾配 0.139ms 中の確保 0.031ms + GC churn) を __閉包生成時に 1 度だけ__ +-- 確保し、 全 leapfrog 呼出で再利用する。 閉包は chain ごとに生成される +-- (@nutsStream@ 内) ため、 chain 横断 spark 並列 (@nutsChainsPure@) とも +-- 干渉しない。 返る値/勾配は毎回 fresh に freeze されるので alias しない。 +-- 非 vecIR 経路 (walk+ad / ハイブリッド) は従来 pure 閉包をそのまま包む。 +-- [English]: A monadic version of 'compileGradValUV'. Allocates the +-- vecIR path's working buffers (the forward arena + adjoint arena — +-- measured on 13-traffic at 34k cells × 2, with 0.031ms of the +-- 0.139ms leaf-gradient time spent on allocation plus GC churn) +-- __only once, at closure-creation time__, and reuses them across +-- all leapfrog calls. Since the closure is created per chain (inside +-- @nutsStream@), this does not interfere with cross-chain spark +-- parallelism (@nutsChainsPure@) either. The returned value/gradient +-- are freshly frozen on every call, so there is no aliasing. The +-- non-vecIR paths (walk+ad / hybrid) simply wrap the usual pure +-- closure as before. +compileGradValUVM :: forall m r. PrimMonad m + => ModelP r -> [Text] -> [Transform] + -> m (VS.Vector Double -> m (Double, VS.Vector Double)) +compileGradValUVM m names trans = + case gradValPlan m names trans of + GVPure f -> pure (\uv -> pure (f uv)) + GVVecIR sz prep core finish -> do + ar <- VSM.unsafeNew sz + adj <- VSM.unsafeNew sz + pure $ \uv -> do + let pc = prep uv + mvg <- stToPrim (core ar adj pc) + pure (finish uv pc mvg) + +-- | [日本語]: 'compileGradValUV' / 'compileGradValUVM' が共有する +-- 静的解析結果。 vecIR 経路のみ per-call の arena/adj 確保をバッファ注入 +-- (prep / core / finish の 3 分割) に分離し、 pure 版 (毎回確保・従来意味論) +-- と monadic 版 (chain 閉包で 1 回確保) が同一 per-call コードを共有する。 +-- [English]: The static analysis result shared by 'compileGradValUV' +-- and 'compileGradValUVM'. Only the vecIR path separates out the +-- per-call arena/adj allocation as an injectable buffer (split into +-- prep / core / finish), so that the pure version (allocates every +-- call, the previous semantics) and the monadic version (allocates +-- once per chain closure) share the same per-call code. +data GradValPlan + = GVPure (VS.Vector Double -> (Double, VS.Vector Double)) + -- ^ [日本語]: walk+ad fallback / ハイブリッド経路 (arena 非使用・従来 pure 閉包)。 + -- [English]: The walk+ad fallback / hybrid path (no arena, the usual pure closure). + | GVVecIR + !Int -- ^ [日本語]: arena/adj サイズ ('vpSize')。 [English]: The arena/adj size ('vpSize'). + (VS.Vector Double -> VS.Vector Double) -- ^ prep: uv → pc (invTransform) + (forall s. VSM.MVector s Double -> VSM.MVector s Double + -> VS.Vector Double + -> ST s (Maybe (Double, VS.Vector Double))) + -- ^ [日本語]: core: ar adj pc → (値, constrained 勾配)。 guard 違反 = Nothing。 + -- [English]: core: ar adj pc → (value, constrained gradient). A guard violation is Nothing. + (VS.Vector Double -> VS.Vector Double + -> Maybe (Double, VS.Vector Double) + -> (Double, VS.Vector Double)) + -- ^ [日本語]: finish: uv pc mvg → 最終 (logπ, ∇logπ) (chain rule + fallback)。 + -- [English]: finish: uv pc mvg → the final (logπ, ∇logπ) (chain rule + fallback). + +gradValPlan :: forall r. ModelP r -> [Text] -> [Transform] -> GradValPlan +gradValPlan m names trans = + case gaussLMBlocksAuto m of + ([], _) -> case synthVecIR m of + -- Phase 95: 尤度が単一 dense MvNormal observe なら解析随伴 (pure 閉包 = GVPure)。 + -- Phase 92: 単一 HmmForwardNormal observe も同様 (forward-backward 閉形式)。 + -- Phase 101: 単一 ArmaNormal observe も同様 (逆向き随伴再帰の閉形式)。 + -- HMM → ARMA → Gp-RBF (B-dsl・閉形式) → 汎用 MvNormal (A案) → 従来の全体 walk+ad。 + Nothing -> case hmmAnalyticVG m names trans of + Just vg -> GVPure vg + Nothing -> case armaAnalyticVG m names trans of + Just vg -> GVPure vg + Nothing -> case gradedIrtAnalyticVG m names trans of + Just vg -> GVPure vg + Nothing -> case gpRBFAnalyticVG m names trans of + Just vg -> GVPure vg + Nothing -> case mvNormalAnalyticVG m names trans of + Just vg -> GVPure vg + Nothing -> GVPure $ \uv -> -- 後方互換: 全体を walk + ad (融合なし) + let us = VS.toList uv + in (fFull us, VS.fromList (gradFull us)) + Just (gs, fams, sObs) -> -- ベクトル式 IR (compileGradUV と同静的) + let ixOf = Map.fromList (zip names [0 ..]) + nP = length names + transB = BV.fromList trans + cvi = compileVecIR ixOf gs fams + famSet = Set.fromList (concat [ ms | (ms, _, _) <- fams ]) + cps = constPriorsOf m famSet + lnGroups = collectLogNormalGroups m -- Phase 98 A3: LogNormal 群 + lnUNames = concat [ us | (us, _, _) <- lnGroups ] + lnIx = map (resolveLogNormal ixOf) lnGroups + exclNames = sObs `Set.union` famSet + `Set.union` Set.fromList (map fst cps) + `Set.union` Set.fromList lnUNames + noResid = residualFreeOfDensity exclNames m + cpIx = [ (ixOf Map.! n, d) | (n, d) <- cps ] + mPrior -- (勾配, 値) の対 (fExcl は logJac 込) + | noResid = Nothing + | otherwise = let mcr = compileResidual exclNames m + in Just (grad (fExcl mcr exclNames), fExcl mcr exclNames) + prep uv = VS.generate nP $ \i -> + invTransformF (transB BV.! i) (uv `VS.unsafeIndex` i) + core :: forall s. VSM.MVector s Double -> VSM.MVector s Double + -> VS.Vector Double -> ST s (Maybe (Double, VS.Vector Double)) + core ar adj pc = do + mg <- VSM.replicate nP 0 + mv <- gradVecIRValWith cvi ar adj pc mg + case mv of + Nothing -> pure Nothing + Just v -> do + mapM_ (\(i, d) -> + case constPriorGradD d (pc `VS.unsafeIndex` i) of + Just g -> VSM.modify mg (+ g) i + Nothing -> pure ()) cpIx + mapM_ (\ln -> gradLogNormalIx ln pc mg) lnIx -- A3 + gv <- VS.unsafeFreeze mg + pure (Just (v, gv)) + finish uv pc mvg = case mvg of + -- guard 違反: 値 = -∞ ('vecIRValue' と同一)・勾配 = walk+ad + -- fallback ('compileGradUV' と同一)。 + Nothing -> ((-1) / 0, VS.fromList (gradFull (VS.toList uv))) + Just (vIR, gC) -> + let cpVal = sum [ logDensity d (pc `VS.unsafeIndex` i) + | (i, d) <- cpIx ] + + sum [ valueLogNormalIx ln pc | ln <- lnIx ] -- A3 + in case mPrior of + Nothing -> + let logJac = sum [ logJacF (transB BV.! i) + (uv `VS.unsafeIndex` i) + | i <- [0 .. nP - 1] ] + g = VS.generate nP $ \i -> + let t = transB BV.! i + u = uv `VS.unsafeIndex` i + in gC `VS.unsafeIndex` i * dInvTransform t u + + dLogJacU t u + in (vIR + cpVal + logJac, g) + Just (priorGrad, priorVal) -> + let us = VS.toList uv + pg = priorGrad us + g = VS.fromList + [ pg_i + gC `VS.unsafeIndex` i * dInvTransform t u + | (i, (pg_i, (t, u))) <- zip [0 :: Int ..] + (zip pg (zip trans us)) ] + in (vIR + cpVal + priorVal us, g) + in GVVecIR (vpSize (cvProg cvi)) prep core finish + (gbs, synthObs) -> GVPure $ -- ハイブリッド (compileGradUV と同静的) + let (priorREs, cps, exclNames, cblocks, hierGroups, noResid) = analyzeGaussModel m gbs synthObs + ixOf = Map.fromList (zip names [0 ..]) + nP = length names + transB = BV.fromList trans + cbIx = map (resolveLMBlock ixOf) cblocks + reIx = [ ReffPriorIx (VU.fromList (map (ixOf Map.!) uNames)) + (ixOf Map.! scaleName) + | (uNames, scaleName) <- priorREs ] + hniIx = map (resolveHierNormal ixOf) hierGroups + cpIx = [ (ixOf Map.! n, d) | (n, d) <- cps ] + mPrior + | noResid = Nothing + | otherwise = let mcr = compileResidual exclNames m + in Just (grad (fExcl mcr exclNames), fExcl mcr exclNames) + in hybridGradValClosure nP transB trans + (\pc mg -> do + mapM_ (\cb -> gradLMBlockIx cb pc mg) cbIx + mapM_ (\ri -> gradReffPriorIx ri pc mg) reIx + mapM_ (\hn -> gradHierNormalIx hn pc mg) hniIx + mapM_ (\(i, d) -> + case constPriorGradD d (pc `VS.unsafeIndex` i) of + Just g -> VSM.modify mg (+ g) i + Nothing -> pure ()) cpIx) + (\pc -> sum [ valueLMBlockIx cb pc | cb <- cbIx ] + + sum [ valueReffPriorIx ri pc | ri <- reIx ] + + sum [ valueHierNormalIx hn pc | hn <- hniIx ] + + sum [ logDensity d (pc `VS.unsafeIndex` i) + | (i, d) <- cpIx ]) + mPrior + where + gradFull = grad fFull + fFull us = + let paramsC = Map.fromList + [ (n, invTransformF t u) | (n, t, u) <- zip3 names trans us ] + logJac = sum [ logJacF t u | (t, u) <- zip trans us ] + in logJoint m paramsC + logJac + fExcl mcr excl us = + let paramsC = Map.fromList + [ (n, invTransformF t u) | (n, t, u) <- zip3 names trans us ] + logJac = sum [ logJacF t u | (t, u) <- zip trans us ] + in residualExcl mcr excl m paramsC + logJac + +-- | [日本語]: 'hybridGradClosure' の value-and-grad 融合版。 解析勾配 +-- (@gradC@) に加えて解析__値__ (@valC@ = 'compileLogPUV' ハイブリッド経路の +-- analytic と同一) を計算し、 (値 (logJac 込)・勾配) を返す。 +-- [English]: A value-and-grad fused version of 'hybridGradClosure'. +-- In addition to the analytic gradient (@gradC@), computes the +-- analytic __value__ (@valC@ = the same as the analytic term in +-- 'compileLogPUV''s hybrid path), and returns (value, including +-- logJac, and gradient). +hybridGradValClosure + :: Int -> BV.Vector Transform -> [Transform] + -> (forall s. VS.Vector Double -> VSM.MVector s Double -> ST s ()) + -> (VS.Vector Double -> Double) + -> Maybe ([Double] -> [Double], [Double] -> Double) + -> (VS.Vector Double -> (Double, VS.Vector Double)) +hybridGradValClosure nP transB trans gradC valC mPrior = \uv -> + let pc = VS.generate nP $ \i -> + invTransformF (transB BV.! i) (uv `VS.unsafeIndex` i) + gC = runST $ do + mg <- VSM.replicate nP 0 + gradC pc mg + VS.unsafeFreeze mg + aVal = valC pc + in case mPrior of + Nothing -> + let logJac = sum [ logJacF (transB BV.! i) (uv `VS.unsafeIndex` i) + | i <- [0 .. nP - 1] ] + g = VS.generate nP $ \i -> + let t = transB BV.! i + u = uv `VS.unsafeIndex` i + in gC `VS.unsafeIndex` i * dInvTransform t u + dLogJacU t u + in (aVal + logJac, g) + Just (priorGrad, priorVal) -> -- fExcl は logJac 込 + let us = VS.toList uv + pg = priorGrad us + g = VS.fromList + [ p + gC `VS.unsafeIndex` i * dInvTransform t u + | (i, (p, (t, u))) <- zip [0 :: Int ..] + (zip pg (zip trans us)) ] + in (aVal + priorVal us, g) + +-- | [日本語]: 'compileGradUV' の per-call 本体 (affine 経路と IR 経路の +-- 共有部を関数化): unconstrained ベクトル → constrained 値 → 解析/ベクトル +-- 経路の constrained 勾配 (@gradC@ が mutable ベクトルへ加算) → chain rule。 +-- @mPriorGrad@ = 残差 ad クロージャ ('Nothing' = 密度項が残らず logJac も解析)。 +-- [English]: The per-call body of 'compileGradUV' (factored out as +-- the part shared by the affine path and the IR path): unconstrained +-- vector → constrained value → the analytic/vectorized path's +-- constrained gradient (@gradC@ adds into the mutable vector) → +-- chain rule. @mPriorGrad@ is the residual ad closure ('Nothing' +-- means no density term remains, and logJac is analytic too). +hybridGradClosure + :: Int -> BV.Vector Transform -> [Transform] + -> (forall s. VS.Vector Double -> VSM.MVector s Double -> ST s ()) + -> Maybe ([Double] -> [Double]) + -> (VS.Vector Double -> VS.Vector Double) +hybridGradClosure nP transB trans gradC mPriorGrad = \uv -> + let pc = VS.generate nP $ \i -> + invTransformF (transB BV.! i) (uv `VS.unsafeIndex` i) + gC = runST $ do -- constrained 空間の解析勾配 + mg <- VSM.replicate nP 0 + gradC pc mg + VS.unsafeFreeze mg + in case mPriorGrad of + Nothing -> -- ad 完全省略 (logJac 解析) + VS.generate nP $ \i -> + let t = transB BV.! i + u = uv `VS.unsafeIndex` i + in gC `VS.unsafeIndex` i * dInvTransform t u + dLogJacU t u + Just priorGrad -> -- 残りは ad (chain は ad 内) + let pg = priorGrad (VS.toList uv) + in VS.fromList + [ p + gC `VS.unsafeIndex` i * dInvTransform t u + | (i, (p, (t, u))) <- zip [0 :: Int ..] + (zip pg (zip trans (VS.toList uv))) ] + +-- | [日本語]: __定数パラメタ prior__ の解析勾配 @d logDensity(d, θ)/dθ@ +-- (constrained 空間)。 @Nothing@ = 未対応分布 (従来 @ad@ に fallback)。 +-- +-- prior のパラメタが他 latent に依存しない (extractDeps で deps ∅) latent に +-- のみ使う前提 (パラメタを定数として θ でだけ微分する)。 各分岐は @logDensity@ +-- の実装・ガードと対にしてある: ガード違反域では @logDensity@ が定数 negInf を +-- 返し @ad@ の勾配は 0 になるので、 ここでも 0 を返して一致させる。 +-- [English]: The analytic gradient @d logDensity(d, θ)/dθ@ +-- (in constrained space) for a __constant-parameter prior__. +-- @Nothing@ means an unsupported distribution (falls back to @ad@ +-- as before). +-- +-- Assumes this is used only for latents whose prior parameters do +-- not depend on any other latent (deps ∅ via extractDeps), i.e. the +-- parameters are treated as constants and only θ is differentiated. +-- Each branch is paired with 'logDensity''s implementation and +-- guards: in a guard-violation region, @logDensity@ returns the +-- constant negInf and @ad@'s gradient becomes 0, so this returns 0 +-- there too, to match. +constPriorGradD :: Distribution Double -> Double -> Maybe Double +constPriorGradD d x = case d of + Normal mu sig + | sig <= 0 -> Just 0 + | otherwise -> Just (negate (x - mu) / (sig * sig)) + Exponential rate + | x < 0 || rate <= 0 -> Just 0 + | otherwise -> Just (negate rate) + Gamma shape rate + | x <= 0 || shape <= 0 || rate <= 0 -> Just 0 + | otherwise -> Just ((shape - 1) / x - rate) + Beta alpha beta + | x <= 0 || x >= 1 || alpha <= 0 || beta <= 0 -> Just 0 + | otherwise -> Just ((alpha - 1) / x - (beta - 1) / (1 - x)) + Uniform lo hi + | hi <= lo || x < lo || x > hi -> Just 0 + | otherwise -> Just 0 + StudentT df mu sig + | df <= 0 || sig <= 0 -> Just 0 + | otherwise -> + let z = x - mu + in Just (negate ((df + 1) * z) / (df * sig * sig + z * z)) + Cauchy loc sc + | sc <= 0 -> Just 0 + | otherwise -> + let z = x - loc + in Just (negate (2 * z) / (sc * sc + z * z)) + HalfNormal sig + | sig <= 0 || x < 0 -> Just 0 + | otherwise -> Just (negate x / (sig * sig)) + HalfCauchy sc + | sc <= 0 || x < 0 -> Just 0 + | otherwise -> Just (negate (2 * x) / (sc * sc + x * x)) + LogNormal mu sig + | sig <= 0 || x <= 0 -> Just 0 + | otherwise -> + Just (negate (1 + (log x - mu) / (sig * sig)) / x) + InverseGamma alpha beta + | alpha <= 0 || beta <= 0 || x <= 0 -> Just 0 + | otherwise -> Just (negate (alpha + 1) / x + beta / (x * x)) + Weibull kShape lam + | kShape <= 0 || lam <= 0 || x <= 0 -> Just 0 + | otherwise -> + Just ((kShape - 1) / x - (kShape / lam) * (x / lam) ** (kShape - 1)) + Pareto alpha xm + | alpha <= 0 || xm <= 0 || x < xm -> Just 0 + | otherwise -> Just (negate (alpha + 1) / x) + _ -> Nothing + +-- | [日本語]: 'logJacF' の u 微分 (ad 省略時に解析で加算)。 'logJacF' と対。 +-- [English]: The u-derivative of 'logJacF' (added analytically when +-- ad is skipped). Paired with 'logJacF'. +dLogJacU :: Transform -> Double -> Double +dLogJacU UnconstrainedT _ = 0 +dLogJacU PositiveT _ = 1 +dLogJacU UnitIntervalT u = let s = 1 / (1 + exp (-u)) in 1 - 2 * s + +-- | [日本語]: @excl@ 除外後の walk に log-density 寄与が残らないか。 +-- 残らなければ 'compileGradU' は @ad@ クロージャを丸ごと省略でき +-- (reflection tape 生成 = profile の 18.9% がゼロに)、 'compileLogPU' は +-- Free walk 自体を省略できる。 scalar 'Observe' は名前が @excl@ になければ +-- False (自動合成で吸収済みの Observe は除外扱い)。 +-- 'Potential' があれば常に False (従来 ad / walk 経路に fallback・正しさ担保)。 +-- [English]: Whether any log-density contribution remains in the +-- walk after excluding @excl@. If none remains, 'compileGradU' can +-- skip the @ad@ closure entirely (reflection tape generation, 18.9% +-- of profile, drops to zero), and 'compileLogPU' can skip the Free +-- walk itself. A scalar 'Observe' is False unless its name is in +-- @excl@ (an Observe already absorbed by automatic synthesis is +-- treated as excluded). If a 'Potential' is present, this is always +-- False (falls back to the usual ad / walk path, to guarantee +-- correctness). +residualFreeOfDensity :: Set Text -> Model Double r -> Bool +residualFreeOfDensity excl = go + where + go (Pure _) = True + go (Free (Sample n _ k)) = n `Set.member` excl && go (k 0) + go (Free (Observe n _ _ next)) = n `Set.member` excl && go next + go (Free (ObserveLM nm _ _ _ _ _ next)) = nm `Set.member` excl && go next + -- Phase 90 A10: vecIR ('VGPot') に吸収済みの potential は残差に数えない。 + go (Free (Potential n _ next)) = n `Set.member` excl && go next + go (Free (Deterministic _ v k)) = go (k v) + go (Free (Data _ ys k)) = go (k (ys, ys)) + go (Free (DataIx _ is k)) = go (k is) + go (Free (PlateBegin _ _ next)) = go next + go (Free (PlateEnd next)) = go next + +-- | [日本語]: 'compileGradU' / 'compileLogPU' 共通の静的解析。 +-- Gaussian LM ブロック群から (ブロック名, 解析 u-prior, 定数パラメタ prior, +-- 除外集合, 前処理済みブロック, residual 空フラグ) を 1 度だけ求める。 +-- @synthObs@ = 自動合成ブロックに吸収済みの scalar 'Observe' 名 +-- (除外集合に合流させ、 residual walk で二重加算しない)。 +-- [English]: The static analysis shared by 'compileGradU' and +-- 'compileLogPU'. From the Gaussian LM blocks, computes, once only, +-- (block names, analytic u-priors, constant-parameter priors, the +-- exclusion set, preprocessed blocks, and the residual-empty flag). +-- @synthObs@ is the set of scalar 'Observe' names already absorbed +-- by automatically synthesized blocks (merged into the exclusion +-- set so the residual walk doesn't double-count them). +analyzeGaussModel + :: ModelP r + -> [(Text, [Text], [[Double]], [REff], Text, [Double])] + -> Set Text -- synthObs (吸収済 scalar Observe 名) + -> ( [([Text], Text)] -- priorREs (uNames, scaleName) + , [(Text, Distribution Double)] -- constPriors + , Set Text -- exclNames + , [CompiledLMBlock] + , [([Text], Text, Text)] -- Phase 93: 階層 Normal 群 (uNames, μ名, τ名) + , Bool ) -- residual に密度項が残らないか +analyzeGaussModel m gbs synthObs = + let blockNames = [ bn | (bn, _, _, _, _, _) <- gbs ] + priorREs = [ (uNames, scaleName) + | (_, _, _, res, _, _) <- gbs + , REff uNames _ (Just scaleName) _ _ <- res ] + exclUNames = concat [ uNames | (uNames, _) <- priorREs ] + -- Phase 93: 非ゼロ latent 平均の階層 Normal prior 群 (mean-0 reff とは disjoint)。 + -- u_i の prior を解析勾配 ('gradHierNormalIx') で扱い残差 ad から外す。 + -- μ・τ は自身の prior を持つので cps 側に残す (u のみ除外)。 + hierGroups = collectHierNormalGroups m + hierUNames = concat [ us | (us, _, _) <- hierGroups ] + -- 定数パラメタ prior。 u_j (REff / 階層群 経由で解析済) は除く。 + cps = constPriorsOf m (Set.fromList (exclUNames ++ hierUNames)) + exclNames = Set.fromList (blockNames ++ exclUNames ++ hierUNames ++ map fst cps) + `Set.union` synthObs + cblocks = [ compileLMBlock (bs, xs, re, sn, ys) + | (_, bs, xs, re, sn, ys) <- gbs ] + noResid = residualFreeOfDensity exclNames m + in (priorREs, cps, exclNames, cblocks, hierGroups, noResid) + +-- | [日本語]: 定数パラメタ prior の抽出: extractDeps で親 latent 無し +-- (deps ∅) かつ解析勾配対応分布の latent。 @exclSet@ = 別経路 (REff 族 / +-- IR 族) で扱う latent は除く。 複数の呼び出し元で共有する。 +-- [English]: Extracts constant-parameter priors: latents with no +-- parent latent (deps ∅ via extractDeps) whose distribution has an +-- analytic gradient. @exclSet@ excludes latents handled by another +-- path (the REff family / the IR family). Shared across call sites. +constPriorsOf :: ModelP r -> Set Text -> [(Text, Distribution Double)] +constPriorsOf m exclSet = + let (depNodes, _) = extractDeps m + latentDeps = Map.fromList [ (nodeName nd, nodeDeps nd) + | nd <- depNodes, nodeKind nd == LatentN ] + in [ (n, dist) + | (n, dist) <- priorList m + , not (Set.member n exclSet) + , Just deps <- [Map.lookup n latentDeps], Set.null deps + , Just _ <- [constPriorGradD dist 0.5] ] + +-- | [日本語]: logp __値__ 評価のコンパイル ('compileGradU' の値版)。 +-- +-- NUTS は tree node ごとにエネルギー (logp の値) を評価する。 ある時点の +-- cost-centre profile で、 勾配は vec 化済みなのに値評価が Free walk + +-- per-obs スカラ 'logDensityObs' のままで per-draw の 46% を占めると判明 +-- (`prof-nuts-54.4c.prof`)。 本関数は 'compileGradU' と同じ静的前処理 +-- ('CompiledLMBlock') を 1 度だけ行い、 unconstrained ベクトルを受けて +-- log-joint + log-jacobian を返すクロージャを構築する: +-- +-- - Gaussian-恒等リンク 'ObserveLM' ブロックの観測尤度値 → 素な Double +-- ベクトル演算 (@valueCompiledLMBlock@・tape 不要) +-- - @REff (Just scale)@ の u-prior 値 → 解析式 (@reffPriorValue@) +-- - 残り (他 prior / scalar observe / 非 Gauss LM / jacobian) +-- → 'logJointExclBlocks' の Double walk +-- +-- Gaussian LM を含まないモデルは従来 'logJointUnconstrained' 相当に fallback +-- (後方互換)。 数値は 'logJointUnconstrained' と一致 (test で担保)。 +-- [English]: Compiles the __value__ evaluation of logp (the value +-- counterpart of 'compileGradU'). +-- +-- NUTS evaluates the energy (the value of logp) for every tree node. +-- A cost-centre profile found that, even though the gradient had +-- already been vectorized, the value evaluation still went through +-- a Free walk plus per-observation scalar 'logDensityObs' calls, +-- accounting for 46% of per-draw time (`prof-nuts-54.4c.prof`). This +-- function performs the same static preprocessing as 'compileGradU' +-- ('CompiledLMBlock') once, and builds a closure that takes an +-- unconstrained vector and returns log-joint + log-jacobian: +-- +-- - The observation-likelihood value of a Gaussian identity-link +-- 'ObserveLM' block → plain Double vector arithmetic +-- (@valueCompiledLMBlock@, no tape needed). +-- - The u-prior value of @REff (Just scale)@ → an analytic formula +-- (@reffPriorValue@). +-- - Everything else (other priors / scalar observe / non-Gaussian +-- LM / jacobian) → a Double walk via 'logJointExclBlocks'. +-- +-- A model without any Gaussian LM falls back to the equivalent of +-- 'logJointUnconstrained' as before (for backward compatibility). +-- The numeric value matches 'logJointUnconstrained' (guaranteed by +-- tests). +compileLogPU :: forall r. ModelP r -> [Text] -> [Transform] -> ([Double] -> Double) +compileLogPU m names trans = + let lv = compileLogPUV m names trans + in lv . VS.fromList + +-- | [日本語]: 'compileLogPU' の vector-native 版。 NUTS のエネルギー評価が +-- 直接使う。 名前は compile 時に index へ解決し、 per-call は Storable vector +-- 上の素な Double 演算のみ (Text-key Map 組立なし)。 +-- [English]: A vector-native version of 'compileLogPU', used +-- directly by NUTS's energy evaluation. Names are resolved to +-- indices at compile time, and each call does plain Double +-- arithmetic over a Storable vector only (no Text-key Map assembly). +compileLogPUV :: forall r. ModelP r -> [Text] -> [Transform] + -> (VS.Vector Double -> Double) +compileLogPUV m names trans = + case gaussLMBlocksAuto m of + ([], _) -> case synthVecIR m of + Nothing -> fFull . VS.toList -- 後方互換: 従来の walk 評価 + Just (gs, fams, sObs) -> -- 54.11: ベクトル式 IR (非線形 μ) + let ixOf = Map.fromList (zip names [0 ..]) + nP = length names + transB = BV.fromList trans + cvi = compileVecIR ixOf gs fams + famSet = Set.fromList (concat [ ms | (ms, _, _) <- fams ]) + cps = constPriorsOf m famSet + lnGroups = collectLogNormalGroups m -- Phase 98 A3: LogNormal 群 + lnUNames = concat [ us | (us, _, _) <- lnGroups ] + lnIx = map (resolveLogNormal ixOf) lnGroups + exclNames = sObs `Set.union` famSet + `Set.union` Set.fromList (map fst cps) + `Set.union` Set.fromList lnUNames + noResid = residualFreeOfDensity exclNames m + cpIx = [ (ixOf Map.! n, d) | (n, d) <- cps ] + mResid + | noResid = Nothing + | otherwise = Just (residualExcl (compileResidual exclNames m) exclNames m) + in hybridLogPClosure nP transB names + (\pc -> vecIRValue cvi pc + + sum [ logDensity d (pc `VS.unsafeIndex` i) + | (i, d) <- cpIx ] + + sum [ valueLogNormalIx ln pc | ln <- lnIx ]) -- A3 + mResid + (gbs, synthObs) -> + let -- 54.4c/54.4e と同じ静的解析: u-prior は解析値・定数パラメタ prior は + -- 直接 logDensity・残りだけ walk。 密度項が残らなければ Free walk 自体を + -- 省略する (モデル再構築 = reNormal の Text 名生成等も消える)。 + (priorREs, cps, exclNames, cblocks, hierGroups, noResid) = analyzeGaussModel m gbs synthObs + ixOf = Map.fromList (zip names [0 ..]) + nP = length names + transB = BV.fromList trans + cbIx = map (resolveLMBlock ixOf) cblocks + reIx = [ ReffPriorIx (VU.fromList (map (ixOf Map.!) uNames)) + (ixOf Map.! scaleName) + | (uNames, scaleName) <- priorREs ] + hniIx = map (resolveHierNormal ixOf) hierGroups + cpIx = [ (ixOf Map.! n, d) | (n, d) <- cps ] + mResid -- 残 walk (fallback のみ) + | noResid = Nothing + | otherwise = Just (residualExcl (compileResidual exclNames m) exclNames m) + in hybridLogPClosure nP transB names + (\pc -> sum [ valueLMBlockIx cb pc | cb <- cbIx ] + + sum [ valueReffPriorIx ri pc | ri <- reIx ] + + sum [ valueHierNormalIx hn pc | hn <- hniIx ] + + sum [ logDensity d (pc `VS.unsafeIndex` i) + | (i, d) <- cpIx ]) + mResid + where + fFull us = + let paramsC = Map.fromList + [ (n, invTransformF t u) | (n, t, u) <- zip3 names trans us ] + logJac = sum [ logJacF t u | (t, u) <- zip trans us ] + in logJoint m paramsC + logJac + +-- | [日本語]: 'compileLogPUV' の per-call 本体 (affine 経路と IR 経路の +-- 共有部を関数化): unconstrained ベクトル → constrained 値 → 解析/ベクトル +-- 経路の log-density 値 (@analytic@) + 残差 walk (@mResid@) + log-jacobian。 +-- [English]: The per-call body of 'compileLogPUV' (factored out as +-- the part shared by the affine path and the IR path): unconstrained +-- vector → constrained value → the analytic/vectorized path's +-- log-density value (@analytic@) + the residual walk (@mResid@) + +-- log-jacobian. +hybridLogPClosure + :: Int -> BV.Vector Transform -> [Text] + -> (VS.Vector Double -> Double) + -> Maybe (Map Text Double -> Double) + -> (VS.Vector Double -> Double) +hybridLogPClosure nP transB names analytic mResid = \uv -> + let pc = VS.generate nP $ \i -> + invTransformF (transB BV.! i) (uv `VS.unsafeIndex` i) + logJac = sum [ logJacF (transB BV.! i) (uv `VS.unsafeIndex` i) + | i <- [0 .. nP - 1] ] + residV = case mResid of + Nothing -> 0 + Just rv -> rv (Map.fromList (zip names (VS.toList pc))) + in residV + analytic pc + logJac + +-- | [日本語]: invTransform の導関数 dθ/du (chain rule 用)。 'invTransformF' と対。 +-- [English]: The derivative dθ/du of invTransform (for the chain +-- rule). Paired with 'invTransformF'. +dInvTransform :: Transform -> Double -> Double +dInvTransform UnconstrainedT _ = 1 +dInvTransform PositiveT u = exp u +dInvTransform UnitIntervalT u = let s = 1 / (1 + exp (-u)) in s * (1 - s) + +-- | [日本語]: モデル中の Gaussian-恒等リンク 'ObserveLM' ブロックを収集する +-- (ブロック名 / β 名 / 設計行列 / ランダム効果 / σ 名 / 観測 ys)。 非 Gaussian は除外。 +-- [English]: Collects the Gaussian identity-link 'ObserveLM' blocks +-- in the model (block name / β names / design matrix / random +-- effects / σ name / observations ys). Non-Gaussian blocks are +-- excluded. +gaussLMBlocks :: ModelP r -> [(Text, [Text], [[Double]], [REff], Text, [Double])] +gaussLMBlocks m = go m [] + where + go (Pure _) acc = reverse acc + go (Free f) acc = case f of + Sample _ _ k -> go (k 0) acc + Observe _ _ _ next -> go next acc + ObserveLM nm bs xs re fam ys next -> + case fam of + LMGaussian sn -> go next ((nm, bs, xs, re, sn, ys) : acc) + _ -> go next acc + Potential _ _ next -> go next acc + Deterministic _ v k -> go (k v) acc + Data _ ys k -> go (k (ys, ys)) acc + DataIx _ is k -> go (k is) acc + PlateBegin _ _ next -> go next acc + PlateEnd next -> go next acc + +-- | [日本語]: 'gaussLMBlocks' + 自動合成。 明示 'ObserveLM' ブロックに、 +-- per-obs scalar 'Observe' から自動合成したブロックを連結して返す +-- (合成に吸収した scalar Observe 名集合も返す → 'analyzeGaussModel' で除外)。 +-- [English]: 'gaussLMBlocks' plus automatic synthesis. Returns the +-- explicit 'ObserveLM' blocks concatenated with blocks automatically +-- synthesized from per-observation scalar 'Observe's (also returns +-- the set of scalar Observe names absorbed by the synthesis, so +-- 'analyzeGaussModel' can exclude them). +gaussLMBlocksAuto + :: ModelP r + -> ([(Text, [Text], [[Double]], [REff], Text, [Double])], Set Text) +gaussLMBlocksAuto m = + let (sblocks, sObs) = synthGaussLMBlocks m + in (gaussLMBlocks m ++ sblocks, sObs) + + +-- | [日本語]: 'logJoint' と同じだが、 名前が @excl@ に含まれる項を __加算しない__: +-- +-- - @excl@ に含まれる 'ObserveLM' ブロックの観測尤度 (vec-tape 経路で別計算) +-- - @excl@ に含まれる scalar 'Observe' の観測尤度 +-- ('synthGaussLMBlocks' が合成ブロックへ吸収済みのもの) +-- - @excl@ に含まれる 'Sample' ノードの prior log-density +-- (群効果 @u_j@ の prior を解析勾配経路で別計算するため)。 +-- 値は継続に必要なので 'Sample' 自体は walk するが log-density は足さない。 +-- [English]: The same as 'logJoint', except terms whose name is in +-- @excl@ are __not added__: +-- +-- - The observation likelihood of an 'ObserveLM' block in @excl@ +-- (computed separately by the vec-tape path). +-- - The observation likelihood of a scalar 'Observe' in @excl@ +-- (already absorbed into a synthesized block by +-- 'synthGaussLMBlocks'). +-- - The prior log-density of a 'Sample' node in @excl@ (because the +-- prior of a group effect @u_j@ is computed separately via the +-- analytic gradient path). The value is still needed for the +-- continuation, so 'Sample' itself is still walked, but its +-- log-density is not added. +logJointExclBlocks :: (Floating a, Ord a) + => Set Text -> Model a r -> Map Text a -> a +logJointExclBlocks excl model params = go model 0 + where + go (Pure _) acc = acc + go (Free (Sample n d k)) acc = + case Map.lookup n params of + Nothing -> negInf + Just v + | n `Set.member` excl -> go (k v) acc + | otherwise -> go (k v) (acc + logDensity d v) + go (Free (Observe n d ys next)) acc + | n `Set.member` excl = go next acc + | otherwise = go next (acc + obsLogSum d ys) + go (Free (ObserveLM nm bs xs re fam ys next)) acc + | nm `Set.member` excl = go next acc + | otherwise = go next (acc + lmObsLogSum bs xs re fam ys params) + -- Phase 90 A10: vecIR ('VGPot') に吸収済みの potential は二重加算しない。 + go (Free (Potential n v next)) acc + | n `Set.member` excl = go next acc + | otherwise = go next (acc + v) + go (Free (Deterministic _ v k)) acc = go (k v) acc + go (Free (Data _ ys k)) acc = go (k (map realToFrac ys, ys)) acc + go (Free (DataIx _ is k)) acc = go (k is) acc + go (Free (PlateBegin _ _ next)) acc = go next acc + go (Free (PlateEnd next)) acc = go next acc + +-- | [日本語]: excl 吸収後の残余 log-density。 'compileResidual' が成功すれば +-- flat 畳み込み ('residualValueA'・Free walk 無し)、 失敗すれば従来の +-- 'logJointExclBlocks' walk に fallback する。 呼び出し側は @mcr@ を 1 度だけ +-- ('compileResidual' で) 構築して値/勾配の両クロージャに渡す ('CompiledResidual' +-- は 'SExp' 保持の純データなので型非依存で共有できる)。 +-- [English]: The residual log-density remaining after excl +-- absorption. If 'compileResidual' succeeds, a flat fold +-- ('residualValueA', no Free walk); otherwise falls back to the +-- usual 'logJointExclBlocks' walk. The caller builds @mcr@ once +-- (via 'compileResidual') and passes it to both the value and +-- gradient closures ('CompiledResidual' is plain data holding an +-- 'SExp', so it can be shared regardless of type). +residualExcl :: (Floating a, Ord a) + => Maybe CompiledResidual -> Set Text -> Model a r -> Map Text a -> a +residualExcl (Just cr) _ _ params = residualValueA cr params +residualExcl Nothing excl m params = logJointExclBlocks excl m params + +-- | [日本語]: Gaussian-恒等リンク 'ObserveLM' ブロックの __静的部分__を 1 度 +-- だけ前処理した中間表現。 NUTS の draw ループの外で構築し全 leapfrog で再利用する +-- ことで、 設計列のベクトル化 (@row !! k@ = O(n·p²)) や群 id の unbox 変換・ys の +-- Storable 化といった「値に依らず draw 間で不変な仕事」 を毎勾配評価から外す。 +-- [English]: An intermediate representation that preprocesses the +-- __static part__ of a Gaussian identity-link 'ObserveLM' block +-- once. Built outside NUTS's draw loop and reused across all +-- leapfrog steps, so that "work that is invariant across draws and +-- independent of the value" — vectorizing the design columns +-- (@row !! k@ = O(n·p²)), unboxing group ids, and converting ys to +-- Storable — is removed from every gradient evaluation. +data CompiledLMBlock = CompiledLMBlock + { clbBetas :: ![Text] -- ^ [日本語]: β パラメタ名 (列順)。 [English]: The β parameter names (column order). + , clbCols :: ![VS.Vector Double] -- ^ [日本語]: 設計列 (p 本・各 length n)。 [English]: The design columns (p of them, each length n). + , clbReff :: ![([Text], Int, VU.Vector Int, Maybe (VS.Vector Double))] + -- ^ [日本語]: (u 名, nG, gids, per-row 重み) のランダム効果 (重み Nothing = 全 1)。 + -- [English]: A random effect as (u names, nG, gids, per-row weights); weights Nothing means all 1. + , clbSname :: !Text -- ^ [日本語]: σ パラメタ名。 [English]: The σ parameter name. + , clbYs :: !(VS.Vector Double) -- ^ [日本語]: 観測 (length n)。 [English]: The observations (length n). + , clbN :: !Int + , clbP :: !Int + } + +-- | [日本語]: 'gaussLMBlocks' の 1 ブロックを 'CompiledLMBlock' に前処理する (静的・1 回)。 +-- [English]: Preprocesses one block from 'gaussLMBlocks' into a +-- 'CompiledLMBlock' (static, done once). +compileLMBlock :: ([Text], [[Double]], [REff], Text, [Double]) -> CompiledLMBlock +compileLMBlock (betaNames, designX, reffs, sName, ys) = + let p = length betaNames + n = length ys + cols = [ VS.fromList [ row !! k | row <- designX ] | k <- [0 .. p - 1] ] + reff = [ (uNames, length uNames, VU.fromList gids, fmap VS.fromList mw) + | REff uNames gids _ mw _ <- reffs ] + in CompiledLMBlock betaNames cols reff sName (VS.fromList ys) n p + +-- | [日本語]: 'CompiledLMBlock' の名前参照を param index に解決した形。 +-- compile 時に 1 度だけ作り、 per-call は Storable vector への index 参照のみ +-- (Text-key Map lookup なし)。 +-- [English]: The form of 'CompiledLMBlock' with name references +-- resolved to param indices. Built once at compile time; each call +-- only does index lookups into a Storable vector (no Text-key Map +-- lookup). +data CompiledLMBlockIx = CompiledLMBlockIx + { cliBetaIx :: !(VU.Vector Int) -- ^ [日本語]: β の param index (列順)。 [English]: The param indices of β (column order). + , cliXMat :: !(VS.Vector Double) -- ^ [日本語]: 設計行列 row-major (n×p・X[i*p+k])。 [English]: The design matrix, row-major (n×p, X[i*p+k]). + , cliCols :: !(BV.Vector (VS.Vector Double)) -- ^ [日本語]: 設計列 (∂β dot 用・O(1) 添字)。 [English]: The design columns (for the ∂β dot product; O(1) indexing). + , cliReff :: ![(VU.Vector Int, Int, VU.Vector Int, Maybe (VS.Vector Double))] + -- ^ [日本語]: (u indices, nG, gids, per-row 重み)。 重み Nothing = 全 1。 + -- [English]: (u indices, nG, gids, per-row weights). Weights Nothing means all 1. + , cliSIx :: !Int -- ^ [日本語]: σ の param index。 [English]: The param index of σ. + , cliYs :: !(VS.Vector Double) -- ^ [日本語]: 観測 (length n)。 [English]: The observations (length n). + , cliN :: !Int + , cliP :: !Int + } + +-- | [日本語]: 'CompiledLMBlock' の名前を index に解決する (静的・1 回)。 +-- row-major 設計行列も前計算する (残差ループのキャッシュ局所性 + リスト走査排除)。 +-- [English]: Resolves 'CompiledLMBlock''s names to indices (static, +-- done once). Also precomputes the row-major design matrix (for +-- cache locality in the residual loop, and to eliminate list +-- traversal). +resolveLMBlock :: Map Text Int -> CompiledLMBlock -> CompiledLMBlockIx +resolveLMBlock ixOf clb = + let n = clbN clb + p = clbP clb + cols = clbCols clb + in CompiledLMBlockIx + { cliBetaIx = VU.fromList [ ixOf Map.! nm | nm <- clbBetas clb ] + , cliXMat = VS.generate (n * p) $ \ix -> + let (i, k) = ix `divMod` p + in (cols !! k) `VS.unsafeIndex` i + , cliCols = BV.fromList cols + , cliReff = [ (VU.fromList [ ixOf Map.! nm | nm <- uNames ], nG, gids, mw) + | (uNames, nG, gids, mw) <- clbReff clb ] + , cliSIx = ixOf Map.! clbSname clb + , cliYs = clbYs clb + , cliN = n + , cliP = p + } + +-- | [日本語]: 階層 Normal 群 (uNames, μ名, τ名) の名前を param index に解決する +-- ('resolveLMBlock' と同様に compile 時 1 回)。 +-- [English]: Resolves the names of a hierarchical Normal group +-- (uNames, mean name, scale name) to param indices (once at compile +-- time, same as 'resolveLMBlock'). +resolveHierNormal :: Map Text Int -> ([Text], Text, Text) -> HierNormalIx +resolveHierNormal ixOf (uNames, meanName, scaleName) = HierNormalIx + { hniUIx = VU.fromList [ ixOf Map.! nm | nm <- uNames ] + , hniMeanIx = ixOf Map.! meanName + , hniScaleIx = ixOf Map.! scaleName + } + +-- | [日本語]: 残差 @r_i = y_i - Σ_k β_k X_ik - Σ_re u^{re}[gid_i]@ と @Σr²@ を +-- __1 パスの手動ループ__ で計算する ((a)-0 実測で per-call +-- ~48-82KB の割当が本物と確定 — `VS.generate` 内のリスト fold・`zip`/@toList@ +-- の毎回再構築・dot/sumR2 の中間ベクトルが原因。 unboxed アキュムレータの +-- 明示ループ + row-major X で割当を r 1 本に削減)。 +-- [English]: Computes the residual +-- @r_i = y_i - Σ_k β_k X_ik - Σ_re u^{re}[gid_i]@ and @Σr²@ with a +-- __single hand-written pass__ (measurement (a)-0 established that +-- the ~48-82KB per-call allocation was real — caused by the list +-- fold inside `VS.generate`, the per-call rebuilding of `zip`/ +-- @toList@, and intermediate vectors for dot/sumR2. An explicit +-- loop with unboxed accumulators plus a row-major X reduces the +-- allocation to a single r vector). +lmResidualS :: CompiledLMBlockIx -> VS.Vector Double -> (VS.Vector Double, Double) +lmResidualS blk pc = runST $ do + let n = cliN blk + p = cliP blk + xm = cliXMat blk + ys = cliYs blk + res = cliReff blk + bv = VS.generate p (\k -> pc `VS.unsafeIndex` (cliBetaIx blk `VU.unsafeIndex` k)) + mr <- VSM.unsafeNew n + let goObs !i !acc + | i >= n = pure acc + | otherwise = do + let base = i * p + goK !k !s + | k >= p = s + | otherwise = goK (k + 1) + (s + bv `VS.unsafeIndex` k * (xm `VS.unsafeIndex` (base + k))) + reS = foldl' (\ !a (uix, _, gids, mw) -> + let u = pc `VS.unsafeIndex` (uix `VU.unsafeIndex` + (gids `VU.unsafeIndex` i)) + in a + case mw of + Nothing -> u + Just w -> w `VS.unsafeIndex` i * u) 0 res + ri = ys `VS.unsafeIndex` i - goK 0 0 - reS + VSM.unsafeWrite mr i ri + goObs (i + 1) (acc + ri * ri) + sumR2 <- goObs 0 0 + r <- VS.unsafeFreeze mr + pure (r, sumR2) + +-- | [日本語]: 前処理済みブロックの観測尤度 @Σ_i logDensityObs(Normal η_i σ) y_i@ の +-- __constrained 空間__での勾配を解析閉形式で mutable 勾配ベクトルに加算する +-- (Gaussian-恒等リンクは閉形式が書けるので汎用 tape 不要): +-- +-- > ∂/∂β_k = X_kᵀ r / σ² +-- > ∂/∂u_j = (Σ_{i: gid_i=j} w_i·r_i) / σ² (scatter・O(n)・重み無しは w_i=1) +-- > ∂/∂σ = -n/σ + (Σ r²)/σ³ +-- +-- dot / scatter とも unboxed アキュムレータの明示ループを使う +-- (中間ベクトル・`VU.convert`・@accumulate@ 割当なし)。 +-- [English]: Adds the __constrained-space__ gradient of the +-- preprocessed block's observation likelihood +-- @Σ_i logDensityObs(Normal η_i σ) y_i@ into the mutable gradient +-- vector, in closed analytic form (a Gaussian identity link has a +-- closed form, so no generic tape is needed): +-- +-- > ∂/∂β_k = X_kᵀ r / σ² +-- > ∂/∂u_j = (Σ_{i: gid_i=j} w_i·r_i) / σ² (scatter, O(n); w_i=1 when unweighted) +-- > ∂/∂σ = -n/σ + (Σ r²)/σ³ +-- +-- Both dot and scatter use an explicit loop with unboxed +-- accumulators (no intermediate vectors, `VU.convert`, or +-- @accumulate@ allocation). +gradLMBlockIx :: CompiledLMBlockIx -> VS.Vector Double + -> VSM.MVector s Double -> ST s () +gradLMBlockIx blk pc mg = do + let sigma = pc `VS.unsafeIndex` cliSIx blk + s2 = sigma * sigma + n = cliN blk + (r, sumR2) = lmResidualS blk pc + n' = fromIntegral n + forM_ [0 .. cliP blk - 1] $ \k -> do + let c = cliCols blk `BV.unsafeIndex` k + dot !i !acc + | i >= n = acc + | otherwise = dot (i + 1) + (acc + c `VS.unsafeIndex` i * r `VS.unsafeIndex` i) + VSM.modify mg (+ (dot 0 0 / s2)) (cliBetaIx blk `VU.unsafeIndex` k) + forM_ (cliReff blk) $ \(uix, nG, gids, mw) -> do + macc <- VSM.replicate nG 0 + let scat !i + | i >= n = pure () + | otherwise = do + let g = gids `VU.unsafeIndex` i + ri = r `VS.unsafeIndex` i + wr = case mw of + Nothing -> ri + Just w -> w `VS.unsafeIndex` i * ri + v <- VSM.unsafeRead macc g + VSM.unsafeWrite macc g (v + wr) + scat (i + 1) + scat 0 + forM_ [0 .. nG - 1] $ \j -> do + gj <- VSM.unsafeRead macc j + VSM.modify mg (+ (gj / s2)) (uix `VU.unsafeIndex` j) + VSM.modify mg (+ (negate n' / sigma + sumR2 / (s2 * sigma))) (cliSIx blk) + +-- | [日本語]: 前処理済みブロックの観測尤度の __値__ +-- @-n/2·log2π - n·logσ - Σr²/(2σ²)@。 r を materialize せず +-- sumR2 だけを 1 パスの明示ループで累積する (割当ゼロ)。 +-- guard (σ≤0 → -∞) は 'logDensityObs' の Normal 分岐と一致させる。 +-- [English]: The __value__ of the preprocessed block's observation +-- likelihood, @-n/2·log2π - n·logσ - Σr²/(2σ²)@. Does not +-- materialize r; accumulates only sumR2 in a single explicit-loop +-- pass (zero allocation). The guard (σ≤0 → -∞) matches +-- 'logDensityObs''s Normal branch. +valueLMBlockIx :: CompiledLMBlockIx -> VS.Vector Double -> Double +valueLMBlockIx blk pc + | sigma <= 0 = negInf + | otherwise = + negate (0.5 * n' * log (2 * pi)) - n' * log sigma + - sumR2 / (2 * sigma * sigma) + where + sigma = pc `VS.unsafeIndex` cliSIx blk + n = cliN blk + p = cliP blk + n' = fromIntegral n + xm = cliXMat blk + ys = cliYs blk + res = cliReff blk + bv = VS.generate p (\k -> pc `VS.unsafeIndex` (cliBetaIx blk `VU.unsafeIndex` k)) + sumR2 = goObs 0 0 + goObs !i !acc + | i >= n = acc + | otherwise = + let base = i * p + goK !k !s + | k >= p = s + | otherwise = goK (k + 1) + (s + bv `VS.unsafeIndex` k * (xm `VS.unsafeIndex` (base + k))) + reS = foldl' (\ !a (uix, _, gids, mw) -> + let u = pc `VS.unsafeIndex` (uix `VU.unsafeIndex` + (gids `VU.unsafeIndex` i)) + in a + case mw of + Nothing -> u + Just w -> w `VS.unsafeIndex` i * u) 0 res + ri = ys `VS.unsafeIndex` i - goK 0 0 - reS + in goObs (i + 1) (acc + ri * ri) + +-- | [日本語]: 群効果 prior @u_j ~ Normal(0, τ)@ の index 解決形。 +-- [English]: The index-resolved form of the group-effect prior +-- @u_j ~ Normal(0, τ)@. +data ReffPriorIx = ReffPriorIx + { rpiUIx :: !(VU.Vector Int) -- ^ [日本語]: u_j の param index (長さ nG)。 [English]: The param indices of u_j (length nG). + , rpiScaleIx :: !Int -- ^ [日本語]: τ の param index。 [English]: The param index of τ. + } + +-- | [日本語]: 群効果 prior の __constrained 空間__での解析勾配を mutable 勾配ベクトルに +-- 加算する (@ad@ のスカラ tape を回避): +-- +-- > log p(u | τ) = -nG/2·log(2π) - nG·log τ - (Σ u_j²)/(2τ²) +-- > ∂/∂u_j = -u_j / τ² +-- > ∂/∂τ = -nG/τ + (Σ u_j²)/τ³ +-- +-- τ 成分は τ 自身の prior (解析 or @ad@ 経路) と加算合流する。 unconstrained への +-- chain rule ('dInvTransform') は呼出側で適用する。 +-- [English]: Adds the analytic gradient of the group-effect prior, +-- in __constrained space__, into the mutable gradient vector +-- (avoiding a scalar @ad@ tape): +-- +-- > log p(u | τ) = -nG/2·log(2π) - nG·log τ - (Σ u_j²)/(2τ²) +-- > ∂/∂u_j = -u_j / τ² +-- > ∂/∂τ = -nG/τ + (Σ u_j²)/τ³ +-- +-- The τ component is merged additively with τ's own prior (analytic +-- or @ad@ path). The chain rule into unconstrained space +-- ('dInvTransform') is applied by the caller. +gradReffPriorIx :: ReffPriorIx -> VS.Vector Double -> VSM.MVector s Double -> ST s () +gradReffPriorIx (ReffPriorIx uix six) pc mg = do + let tau = pc `VS.unsafeIndex` six + tau2 = tau * tau + nG = VU.length uix + sumU2 <- VU.foldM' (\ !acc i -> do + let u = pc `VS.unsafeIndex` i + VSM.modify mg (+ (negate u / tau2)) i + pure (acc + u * u)) 0 uix + VSM.modify mg (+ (negate (fromIntegral nG) / tau + sumU2 / (tau2 * tau))) six + +-- | [日本語]: 群効果 prior の log-density 和の __値__ ('gradReffPriorIx' の値版)。 +-- guard (τ≤0 → -∞) は @logDensity@ の Normal 分岐と一致させる。 +-- [English]: The __value__ of the sum of the group-effect prior's +-- log-density (the value counterpart of 'gradReffPriorIx'). The +-- guard (τ≤0 → -∞) matches 'logDensity''s Normal branch. +valueReffPriorIx :: ReffPriorIx -> VS.Vector Double -> Double +valueReffPriorIx (ReffPriorIx uix six) pc + | tau <= 0 = negInf + | otherwise = + negate (0.5 * nG' * log (2 * pi)) - nG' * log tau + - sumU2 / (2 * tau * tau) + where + tau = pc `VS.unsafeIndex` six + nG' = fromIntegral (VU.length uix) + sumU2 = VU.foldl' (\ !acc i -> let u = pc `VS.unsafeIndex` i + in acc + u * u) 0 uix + +-- | [日本語]: __非ゼロ latent 平均__の階層 Normal prior の解析勾配経路。 +-- 'ReffPriorIx' (mean-0 専用) の一般化で、 平均 μ・スケール τ とも latent の +-- @u_i ~ Normal(μ, τ)@ 群を扱う (rats の @alpha[i]~Normal(muAlpha,sigmaAlpha)@ 等)。 +-- [English]: The analytic-gradient path for a hierarchical Normal +-- prior with a __non-zero latent mean__. A generalization of +-- 'ReffPriorIx' (which is mean-0 only): handles a group +-- @u_i ~ Normal(μ, τ)@ where both the mean μ and the scale τ are +-- latent (e.g. rats' @alpha[i]~Normal(muAlpha,sigmaAlpha)@). +data HierNormalIx = HierNormalIx + { hniUIx :: !(VU.Vector Int) -- ^ [日本語]: u_i の param index (長さ nG)。 [English]: The param indices of u_i (length nG). + , hniMeanIx :: !Int -- ^ [日本語]: μ の param index。 [English]: The param index of μ. + , hniScaleIx :: !Int -- ^ [日本語]: τ の param index。 [English]: The param index of τ. + } + +-- | [日本語]: 'HierNormalIx' の __constrained 空間__での解析勾配を mutable 勾配ベクトルに +-- 加算する (@ad@ のスカラ tape を回避): +-- +-- > log p(u | μ, τ) = -nG/2·log(2π) - nG·log τ - (Σ (u_i-μ)²)/(2τ²) +-- > ∂/∂u_i = -(u_i - μ) / τ² +-- > ∂/∂μ = (Σ (u_i - μ)) / τ² +-- > ∂/∂τ = -nG/τ + (Σ (u_i-μ)²)/τ³ +-- +-- μ・τ 成分は各自の prior (解析 or @ad@ 経路) と加算合流する。 unconstrained への +-- chain rule ('dInvTransform') は呼出側で適用する。 +-- [English]: Adds 'HierNormalIx''s analytic gradient, in +-- __constrained space__, into the mutable gradient vector (avoiding +-- a scalar @ad@ tape): +-- +-- > log p(u | μ, τ) = -nG/2·log(2π) - nG·log τ - (Σ (u_i-μ)²)/(2τ²) +-- > ∂/∂u_i = -(u_i - μ) / τ² +-- > ∂/∂μ = (Σ (u_i - μ)) / τ² +-- > ∂/∂τ = -nG/τ + (Σ (u_i-μ)²)/τ³ +-- +-- The μ and τ components are merged additively with their own +-- priors (analytic or @ad@ path). The chain rule into unconstrained +-- space ('dInvTransform') is applied by the caller. +gradHierNormalIx :: HierNormalIx -> VS.Vector Double -> VSM.MVector s Double -> ST s () +gradHierNormalIx (HierNormalIx uix mIx sIx) pc mg = do + let mu = pc `VS.unsafeIndex` mIx + tau = pc `VS.unsafeIndex` sIx + tau2 = tau * tau + nG = VU.length uix + (sumD, sumD2) <- + VU.foldM' (\ (!accD, !accD2) i -> do + let u = pc `VS.unsafeIndex` i + d = u - mu + VSM.modify mg (+ (negate d / tau2)) i + pure (accD + d, accD2 + d * d)) (0, 0) uix + VSM.modify mg (+ (sumD / tau2)) mIx + VSM.modify mg (+ (negate (fromIntegral nG) / tau + sumD2 / (tau2 * tau))) sIx + +-- | [日本語]: 'HierNormalIx' の log-density 和の __値__ ('gradHierNormalIx' の値版)。 +-- guard (τ≤0 → -∞) は @logDensity@ の Normal 分岐と一致させる。 +-- [English]: The __value__ of the sum of 'HierNormalIx''s +-- log-density (the value counterpart of 'gradHierNormalIx'). The +-- guard (τ≤0 → -∞) matches 'logDensity''s Normal branch. +valueHierNormalIx :: HierNormalIx -> VS.Vector Double -> Double +valueHierNormalIx (HierNormalIx uix mIx sIx) pc + | tau <= 0 = negInf + | otherwise = + negate (0.5 * nG' * log (2 * pi)) - nG' * log tau + - sumD2 / (2 * tau * tau) + where + mu = pc `VS.unsafeIndex` mIx + tau = pc `VS.unsafeIndex` sIx + nG' = fromIntegral (VU.length uix) + sumD2 = VU.foldl' (\ !acc i -> let d = pc `VS.unsafeIndex` i - mu + in acc + d * d) 0 uix + +-- --------------------------------------------------------------------------- +-- Phase 98 A3: LogNormal 群 prior の解析勾配 ('HierNormalIx' の LogNormal 版) +-- --------------------------------------------------------------------------- +-- @a_i ~ LogNormal(μ, σ)@ 群 (μ = 定数 or 単一 latent・σ = 単一 latent) の値/勾配を +-- 解析式で扱い、 vecIR 経路の残余 reverse-AD tape (irt-2pl で ~30%time/~85%alloc) を消す。 + +-- | [日本語]: 'collectLogNormalGroups' の結果を param index へ解決した中間表現。 +-- μ が定数なら @hlnMeanIx = Left c@、 latent なら @Right ix@。 +-- [English]: An intermediate representation with +-- 'collectLogNormalGroups''s result resolved to param indices. If μ +-- is constant, @hlnMeanIx = Left c@; if it is latent, @Right ix@. +data LogNormalIx = LogNormalIx + { hlnUIx :: !(VU.Vector Int) -- ^ [日本語]: a_i の param index (長さ nG)。 [English]: The param indices of a_i (length nG). + , hlnMeanIx :: !(Either Double Int) -- ^ [日本語]: μ (定数 or param index)。 [English]: μ (a constant, or a param index). + , hlnScaleIx :: !Int -- ^ [日本語]: σ の param index。 [English]: The param index of σ. + } + +resolveLogNormal :: Map Text Int -> ([Text], Either Double Text, Text) -> LogNormalIx +resolveLogNormal ixOf (uNames, mean, scaleName) = LogNormalIx + { hlnUIx = VU.fromList [ ixOf Map.! nm | nm <- uNames ] + , hlnMeanIx = either Left (Right . (ixOf Map.!)) mean + , hlnScaleIx = ixOf Map.! scaleName + } + +-- | [日本語]: 'LogNormalIx' の __constrained 空間__での解析勾配を mutable 勾配ベクトルに +-- 加算する (@ad@ のスカラ tape を回避)。 L_i = log a_i, d_i = L_i - μ として: +-- +-- > log p(a | μ, σ) = -nG/2·log(2π) - nG·log σ - Σ L_i - (Σ d_i²)/(2σ²) +-- > ∂/∂a_i = -(1 + d_i/σ²) / a_i +-- > ∂/∂μ = (Σ d_i) / σ² (μ が latent のときのみ) +-- > ∂/∂σ = -nG/σ + (Σ d_i²)/σ³ +-- +-- unconstrained への chain rule ('dInvTransform') は呼出側で適用する。 +-- [English]: Adds 'LogNormalIx''s analytic gradient, in +-- __constrained space__, into the mutable gradient vector (avoiding +-- a scalar @ad@ tape). Writing L_i = log a_i, d_i = L_i - μ: +-- +-- > log p(a | μ, σ) = -nG/2·log(2π) - nG·log σ - Σ L_i - (Σ d_i²)/(2σ²) +-- > ∂/∂a_i = -(1 + d_i/σ²) / a_i +-- > ∂/∂μ = (Σ d_i) / σ² (only when μ is latent) +-- > ∂/∂σ = -nG/σ + (Σ d_i²)/σ³ +-- +-- The chain rule into unconstrained space ('dInvTransform') is +-- applied by the caller. +gradLogNormalIx :: LogNormalIx -> VS.Vector Double -> VSM.MVector s Double -> ST s () +gradLogNormalIx (LogNormalIx uix meanIx sIx) pc mg = do + let mu = either id (pc `VS.unsafeIndex`) meanIx + sig = pc `VS.unsafeIndex` sIx + sig2 = sig * sig + nG = VU.length uix + (sumD, sumD2) <- + VU.foldM' (\ (!accD, !accD2) i -> do + let a = pc `VS.unsafeIndex` i + d = log a - mu + VSM.modify mg (+ (negate (1 + d / sig2) / a)) i + pure (accD + d, accD2 + d * d)) (0, 0) uix + case meanIx of + Right mIx -> VSM.modify mg (+ (sumD / sig2)) mIx + Left _ -> pure () + VSM.modify mg (+ (negate (fromIntegral nG) / sig + sumD2 / (sig2 * sig))) sIx + +-- | [日本語]: 'LogNormalIx' の log-density 和の __値__ ('gradLogNormalIx' の値版)。 +-- guard (σ≤0 / a_i≤0 → -∞) は @logDensity@ の LogNormal 分岐と一致させる +-- (a は PositiveT 変換で a>0 だが安全のため一致させる)。 +-- [English]: The __value__ of the sum of 'LogNormalIx''s +-- log-density (the value counterpart of 'gradLogNormalIx'). The +-- guard (σ≤0 / a_i≤0 → -∞) matches 'logDensity''s LogNormal branch +-- (a>0 is guaranteed by the PositiveT transform, but this matches +-- the guard anyway for safety). +valueLogNormalIx :: LogNormalIx -> VS.Vector Double -> Double +valueLogNormalIx (LogNormalIx uix meanIx sIx) pc + | sig <= 0 = negInf + | VU.any (\i -> pc `VS.unsafeIndex` i <= 0) uix = negInf + | otherwise = + negate (0.5 * nG' * log (2 * pi)) - nG' * log sig - sumL + - sumD2 / (2 * sig * sig) + where + mu = either id (pc `VS.unsafeIndex`) meanIx + sig = pc `VS.unsafeIndex` sIx + nG' = fromIntegral (VU.length uix) + (sumL, sumD2) = + VU.foldl' (\ (!aL, !aD2) i -> + let l = log (pc `VS.unsafeIndex` i) + d = l - mu + in (aL + l, aD2 + d * d)) (0, 0) uix + +-- --------------------------------------------------------------------------- +-- 制約変換 (Floating 多相版) +-- --------------------------------------------------------------------------- + +-- | [日本語]: unconstrained → constrained 変換 (Floating 多相)。 +-- +-- > UnconstrainedT: θ = u +-- > PositiveT: θ = exp(u) +-- > UnitIntervalT: θ = sigmoid(u) = 1/(1+exp(-u)) +-- [English]: The unconstrained → constrained transform (polymorphic +-- in Floating). +-- +-- > UnconstrainedT: θ = u +-- > PositiveT: θ = exp(u) +-- > UnitIntervalT: θ = sigmoid(u) = 1/(1+exp(-u)) +invTransformF :: Floating a => Transform -> a -> a +invTransformF UnconstrainedT u = u +invTransformF PositiveT u = exp u +invTransformF UnitIntervalT u = 1 / (1 + exp (-u)) + +-- | [日本語]: log |∂θ/∂u| — Jacobian 行列式の対数 (Floating 多相)。 +-- [English]: log |∂θ/∂u| — the log of the Jacobian determinant +-- (polymorphic in Floating). +logJacF :: Floating a => Transform -> a -> a +logJacF UnconstrainedT _ = 0 +logJacF PositiveT u = u -- log(exp u) = u +logJacF UnitIntervalT u = + let p = 1 / (1 + exp (-u)) + in log p + log (1 - p) -- log σ(u)(1-σ(u)) + +-- | [日本語]: 各 latent 変数の事前分布から制約変換を自動検出する。 分布名→変換の表は +-- 'nameToTransform' (@HBM.Distribution@) に一元化されている (probe 側 +-- 'vecIRProbeOK' と同一 source)。 +-- [English]: Auto-detects the constraint transform for each latent +-- variable from its prior distribution. The distribution-name → +-- transform table is centralized in 'nameToTransform' +-- (@HBM.Distribution@), the same source used by the probe-side +-- 'vecIRProbeOK'. +getTransforms :: ModelP r -> Map Text Transform +getTransforms m = Map.fromList + [ (nodeName n, nameToTransform (nodeDist n)) + | n <- collectNodes m + , nodeKind n == LatentN + ] + +-- | [日本語]: unconstrained 空間における log-joint (Jacobian 補正込み)。 +-- Jacobian 補正で確率密度の積分を保存する。 +-- [English]: The log-joint in unconstrained space (including the +-- Jacobian correction). The Jacobian correction preserves the +-- integral of the probability density. +logJointUnconstrained :: forall a r. (Floating a, Ord a) + => Model a r + -> [Text] -- ^ [日本語]: パラメータ順序。 [English]: The parameter order. + -> [Transform] -- ^ [日本語]: 各パラメータの変換種別。 [English]: The transform kind for each parameter. + -> Map Text a -- ^ [日本語]: unconstrained パラメータ値。 [English]: The unconstrained parameter values. + -> a +logJointUnconstrained m names trans paramsU = + let paramsC = Map.fromList + [ (n, invTransformF t (Map.findWithDefault 0 n paramsU)) + | (n, t) <- zip names trans ] + logJac = sum + [ logJacF t (Map.findWithDefault 0 n paramsU) + | (n, t) <- zip names trans ] + in logJoint m paramsC + logJac diff --git a/src/Hanalyze/Model/HBM/IR.hs b/src/Hanalyze/Model/HBM/IR.hs new file mode 100644 --- /dev/null +++ b/src/Hanalyze/Model/HBM/IR.hs @@ -0,0 +1,3666 @@ +{-# LANGUAGE RankNTypes #-} +{-# LANGUAGE DeriveFunctor #-} +{-# LANGUAGE FlexibleInstances #-} +{-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE ImpredicativeTypes #-} +{-# LANGUAGE BangPatterns #-} +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE TupleSections #-} + +-- | +-- Module : Hanalyze.Model.HBM.IR +-- Description : HBM の中間表現 (IR) 層 (affine 追跡・SExp/UExp コンパイル) +-- Copyright : (c) 2026 Aelysce Project (Toshiaki Honda) +-- License : BSD-3-Clause +-- +-- [日本語]: IR (中間表現) 層を 'Hanalyze.Model.HBM' から分離。 +-- +-- AD 勾配の高速経路で使う __中間表現__ (記述層 Model / 評価層 Eval の上層): +-- +-- - affine 追跡 ('AffV') による per-obs 手書き Gaussian モデルの自動 +-- ObserveLM 化 ('synthGaussLMBlocks') +-- - 非線形 μ の「スカラ式 IR」 ('SExp') → 「ベクトル式 IR」 ('UExp') 合成 +-- ('synthVecIR' / 'compileVecIR') +-- - 観測密度の IR 式化 ('VecObsIR' → 'CompiledVecIR') と arena 上の値/勾配 +-- 評価 ('vecIRValue' / 'gradVecIR') +-- +-- ★__最ホット__: NUTS per-draw の勾配本経路 ('gradVecIR')。 monolith では AD 勾配 +-- コンパイラ (@compileGradUV@ 本体残置) と同一モジュールで inline されていた。 +-- 境界跨ぎ inline 喪失を防ぐため定義と一緒に INLINABLE/SPECIALIZE を移送。 依存は +-- 下層 Model / Distribution / Eval (lmObsLogSum) / Util のみ (一方向)。 +-- +-- export list は省略 (内部実装層)。 公開 surface (synthGaussLMBlocks / synthVecIR) +-- は facade 'Hanalyze.Model.HBM' の export list が制御する。 +-- +-- [English]: The IR (intermediate representation) layer, separated out from +-- 'Hanalyze.Model.HBM'. +-- +-- The __intermediate representation__ used on the fast path for AD +-- gradients (a layer above the description layer Model / evaluation layer +-- Eval): +-- +-- - Automatic conversion of hand-written per-observation Gaussian models +-- into ObserveLM form via affine tracking ('AffV') +-- ('synthGaussLMBlocks') +-- - Composing a "scalar-expression IR" ('SExp') for nonlinear μ into a +-- "vector-expression IR" ('UExp') ('synthVecIR' \/ 'compileVecIR') +-- - Turning observation densities into IR expressions ('VecObsIR' -> +-- 'CompiledVecIR') and evaluating values\/gradients over the arena +-- ('vecIRValue' \/ 'gradVecIR') +-- +-- ★__The single hottest path__: the main gradient path for NUTS per-draw +-- ('gradVecIR'). In the monolith, this was inlined into the same module as +-- the AD gradient compiler (@compileGradUV@, whose body remains there). The +-- INLINABLE\/SPECIALIZE pragmas were moved together with the definitions to +-- avoid losing cross-module inlining. Dependencies flow one way only, on +-- the lower layers Model \/ Distribution \/ Eval (lmObsLogSum) \/ Util. +-- +-- The export list is omitted (this is an internal implementation layer). +-- The public surface (synthGaussLMBlocks \/ synthVecIR) is controlled by +-- the facade 'Hanalyze.Model.HBM' export list. +module Hanalyze.Model.HBM.IR where + +import Control.DeepSeq (NFData (..), force) +import Control.Exception (SomeAsyncException (..), SomeException, evaluate, + fromException, throwIO, try) +import Control.Monad (forM, forM_, replicateM, when) +import Data.IORef (IORef, modifyIORef', newIORef, readIORef, writeIORef) +import Data.List (foldl') +import System.IO.Unsafe (unsafePerformIO) +import System.Mem.StableName (StableName, hashStableName, makeStableName) +import qualified Data.IntMap.Strict as IM +import qualified Data.Map.Strict as Map +import Data.Map.Strict (Map) +import qualified Data.Set as Set +import Data.Set (Set) +import Data.Text (Text) +import qualified Data.Text as T +import Numeric.AD.Mode.Reverse.Double (grad) +import Control.Monad.Primitive (PrimMonad, PrimState) + +import Control.Monad.ST (ST, runST) +import qualified Data.Vector as BV +import qualified Data.Vector.Storable as VS +import qualified Data.Vector.Storable.Mutable as VSM +import qualified Data.Vector.Unboxed as VU + +import Hanalyze.Stat.Distribution (Transform (..), fromUnconstrained) +import Hanalyze.MCMC.Core (Chain (..)) + +import Hanalyze.Model.HBM.Util +import Hanalyze.Model.HBM.Distribution +import Hanalyze.Model.HBM.Sampling +import Hanalyze.Model.HBM.Model +import Hanalyze.Model.HBM.Track +import Hanalyze.Model.HBM.Eval + +-- --------------------------------------------------------------------------- +-- Phase 54.8: per-obs 手書きモデルの自動 ObserveLM 化 (M1 救済) +-- --------------------------------------------------------------------------- + +-- | [日本語]: affine 追跡値。 latent 値の場所に流して、 式が +-- @Σ coeff_i · latent_i + offset@ (係数は定数) の形に留まるかを追跡する。 +-- 非線形演算 (latent 同士の積・exp 等) が掛かった時点で @NA@ に落ちる。 +-- [English]: An affine-tracking value. Fed in place of a latent value, it +-- tracks whether the expression stays in the form @Σ coeff_i · latent_i + +-- offset@ (with constant coefficients). Once a nonlinear operation hits it +-- (product of latents, exp, etc.), it collapses to @NA@. +data AffV + = AffC !Double -- ^ [日本語]: 定数。 [English]: A constant. + | AffL !(Map Text Double) !Double -- ^ [日本語]: Σ coeff·latent + offset (affine)。 [English]: Σ coeff·latent + offset (affine). + | NA -- ^ [日本語]: 非 affine (追跡断念)。 [English]: Non-affine (tracking gave up). + +-- Phase 60.7: '!!!' の依存タグは IR 抽出には無関係 (既定 id)。 +instance TrackTag AffV + +-- | [日本語]: 非定数値の比較 = 値依存分岐。 構造抽出は分岐の片側しか見られないため +-- 誤抽出になる → error poison で walk 全体を失敗させ、 呼出側の +-- @try/evaluate/force@ で捕捉して fallback する (安全網①)。 +-- [English]: Comparing non-constant values means a value-dependent +-- branch. Structural extraction can only see one side of a branch, which +-- would cause a mis-extraction, so we error-poison the whole walk, and +-- the caller catches it with @try/evaluate/force@ and falls back (safety +-- net (1)). +affPoison :: a +affPoison = error "AffV: non-constant comparison (value-dependent branch)" + +instance Eq AffV where + AffC a == AffC b = a == b + _ == _ = affPoison + +instance Ord AffV where + compare (AffC a) (AffC b) = compare a b + compare _ _ = affPoison + +instance Num AffV where + AffC a + AffC b = AffC (a + b) + AffC a + AffL m c = AffL m (a + c) + AffL m c + AffC a = AffL m (c + a) + AffL m1 c1 + AffL m2 c2 = AffL (Map.unionWith (+) m1 m2) (c1 + c2) + _ + _ = NA + AffC a * AffC b = AffC (a * b) + AffC a * AffL m c = scaleAffV a m c + AffL m c * AffC a = scaleAffV a m c + _ * _ = NA + negate (AffC a) = AffC (negate a) + negate (AffL m c) = AffL (Map.map negate m) (negate c) + negate NA = NA + abs (AffC a) = AffC (abs a) + abs _ = NA + signum (AffC a) = AffC (signum a) + signum _ = NA + fromInteger = AffC . fromInteger + +-- | [日本語]: 定数倍。 0 倍は affine 情報ごと消えて定数 0 (係数 0 の死に列を作らない)。 +-- [English]: Scale by a constant. Multiplying by 0 wipes out the affine +-- info entirely, giving the constant 0 (so we never build a dead column +-- of zero coefficients). +scaleAffV :: Double -> Map Text Double -> Double -> AffV +scaleAffV a m c + | a == 0 = AffC 0 + | otherwise = AffL (Map.map (a *) m) (a * c) + +instance Fractional AffV where + AffC a / AffC b = AffC (a / b) + AffL m c / AffC b = scaleAffV (recip b) m c + _ / _ = NA + recip (AffC a) = AffC (recip a) + recip _ = NA + fromRational = AffC . fromRational + +instance Floating AffV where + pi = AffC pi + exp = affLift1 exp + log = affLift1 log + sqrt = affLift1 sqrt + sin = affLift1 sin + cos = affLift1 cos + tan = affLift1 tan + asin = affLift1 asin + acos = affLift1 acos + atan = affLift1 atan + sinh = affLift1 sinh + cosh = affLift1 cosh + tanh = affLift1 tanh + asinh = affLift1 asinh + acosh = affLift1 acosh + atanh = affLift1 atanh + +-- | [日本語]: 超越関数: 定数には適用、 latent が絡んだら非 affine。 +-- [English]: A transcendental function: applied to constants directly, +-- but treated as non-affine once a latent is involved. +affLift1 :: (Double -> Double) -> AffV -> AffV +affLift1 f (AffC a) = AffC (f a) +affLift1 _ _ = NA + +-- | [日本語]: per-obs 手書き scalar 'Observe' 群から Gaussian LM ブロックを +-- __自動合成__する。 返り値は (合成ブロック群, 吸収した Observe ノード名集合)。 +-- 検出できない / 安全網に掛かった場合は @([], ∅)@ (従来経路に fallback)。 +-- +-- 仕組み: 'Sample' の継続に @AffL {name:1} 0@ を給餌して model を walk し、 +-- @Observe nm (Normal μ σ) ys@ の μ が affine・σ が単一 latent (係数 1・offset 0) +-- の行を収集する。 定数 offset は ys 側に畳む (Normal は y−μ のみに依存)。 +-- σ 名ごとに 1 ブロックへまとめ、 prior が @Normal(0, τ)@ (τ 単一 latent) を +-- 共有し __各行にちょうど 1 つ__現れる latent 族を 'REff' gather に昇格する +-- (dense one-hot は O(nG·n) で階層に逆効果)。 係数は任意で、 +-- per-row 重みとして 'REff' に載せる (random slope @v_g·x_i@ も +-- gather 化。 全 1 なら重みスロットは @Nothing@ = 従来の random intercept)。 +-- 族抽出に失敗した latent は dense β 列のまま (正しいが遅い・安全方向)。 +-- +-- 安全網 2 段: ① 'AffV' の Eq/Ord は非定数比較で error poison → +-- 'unsafePerformIO' + 'try' + 'force' で捕捉し全体 fallback (値依存分岐モデルの +-- 誤抽出防止・Nonlinear 系の前例に同じ)。 ② 合成ブロックの観測尤度を probe +-- 2 点で walk 評価 ('obsOnlySum') と突合し、 不一致なら fallback。 +-- [English]: __Automatically synthesizes__ Gaussian LM blocks from a group +-- of hand-written per-observation scalar 'Observe' nodes. Returns +-- (synthesized blocks, the set of absorbed Observe node names). If +-- nothing can be detected, or a safety net trips, returns @([], ∅)@ (fall +-- back to the previous path). +-- +-- How it works: walks the model, feeding @AffL {name:1} 0@ into the +-- continuation of each 'Sample', and collects the rows where +-- @Observe nm (Normal μ σ) ys@ has an affine μ and a σ that is a single +-- latent (coefficient 1, offset 0). Constant offsets are folded into +-- @ys@ (Normal depends only on y−μ). Rows are grouped into one block per +-- σ name; when the prior is @Normal(0, τ)@ (a single latent τ) shared +-- across the group, and a latent family appears __exactly once per row__, +-- that family is promoted to an 'REff' gather (dense one-hot would be +-- O(nG·n), which hurts hierarchical models). Coefficients are arbitrary +-- and are carried in 'REff' as per-row weights (random slopes +-- @v_g·x_i@ are also gathered this way; if all weights are 1 the weight +-- slot is @Nothing@, i.e. an ordinary random intercept). Any latent whose +-- family extraction fails stays as a dense β column (correct but slower +-- — the safe direction). +-- +-- Two safety nets: (1) 'AffV'\'s Eq\/Ord poisons on non-constant +-- comparisons -> caught via 'unsafePerformIO' + 'try' + 'force' and the +-- whole thing falls back (this prevents mis-extraction on +-- value-dependent-branch models, following the precedent set for the +-- nonlinear family). (2) The synthesized blocks' observation likelihood +-- is cross-checked at two probe points against the walked evaluation +-- ('obsOnlySum'); a mismatch triggers fallback. +synthGaussLMBlocks + :: ModelP r + -> ([(Text, [Text], [[Double]], [REff], Text, [Double])], Set Text) +synthGaussLMBlocks m = unsafePerformIO $ do + r <- try (evaluate (force (synthGaussLMWalk m))) + pure $ case r :: Either SomeException + ([(Text, [Text], [[Double]], [REff], Text, [Double])], Set Text) of + Left _ -> ([], Set.empty) + Right v@(blocks, obsNames) + | null blocks -> ([], Set.empty) + | synthProbeOK m blocks obsNames -> v + | otherwise -> ([], Set.empty) +{-# NOINLINE synthGaussLMBlocks #-} + +-- | [日本語]: 'synthGaussLMBlocks' の純粋部 (walk + 族抽出)。 poison は遅延に潜むので +-- 呼出側が force してから使う。 +-- [English]: The pure part of 'synthGaussLMBlocks' (walking + family +-- extraction). The poison hides in laziness, so the caller must force +-- before use. +synthGaussLMWalk + :: ModelP r + -> ([(Text, [Text], [[Double]], [REff], Text, [Double])], Set Text) +synthGaussLMWalk m = + let (rows, priors) = collectAffRows m + sigmas = ordNubT [ sn | (_, _, _, sn, _) <- rows ] + blocks = [ synthBlock priors sn [ r | r@(_, _, _, sn', _) <- rows, sn' == sn ] + | sn <- sigmas ] + obsNames = Set.fromList [ nm | (nm, _, _, _, _) <- rows ] + in (blocks, obsNames) + where + ordNubT = go Set.empty + where go _ [] = [] + go seen (x:xs) + | x `Set.member` seen = go seen xs + | otherwise = x : go (Set.insert x seen) xs + +-- | [日本語]: model を 'AffV' で walk し、 合成可能な行 (Observe 名, μ 係数, μ offset, +-- σ 名, 観測値) と latent prior のスケール検出 (@Normal(0, τ)@ → @Just τ@) を集める。 +-- [English]: Walks the model with 'AffV' and collects the rows that can +-- be synthesized (Observe name, μ coefficients, μ offset, σ name, +-- observed value) along with scale detection for the latent prior +-- (@Normal(0, τ)@ -> @Just τ@). +collectAffRows + :: Model AffV r + -> ([(Text, Map Text Double, Double, Text, Double)], Map Text (Maybe Text)) +collectAffRows = go [] Map.empty + where + go rows priors (Pure _) = (reverse rows, priors) + go rows priors (Free f) = case f of + Sample n d k -> + let sc = case d of + Normal (AffC 0) (AffL tm 0) + | [(tn, 1)] <- Map.toList tm -> Just tn + _ -> Nothing + in go rows (Map.insert n sc priors) (k (AffL (Map.singleton n 1) 0)) + Observe nm (Normal mu sg) ys next + | AffL sm 0 <- sg, [(sn, 1)] <- Map.toList sm + , Just (cs, off) <- affParts mu -> + go ([ (nm, cs, off, sn, y) | y <- ys ] ++ rows) priors next + Observe _ _ _ next -> go rows priors next + ObserveLM _ _ _ _ _ _ next -> go rows priors next + Potential _ _ next -> go rows priors next + Deterministic _ v k -> go rows priors (k v) + Data _ ys k -> go rows priors (k (map realToFrac ys, ys)) + DataIx _ is k -> go rows priors (k is) + PlateBegin _ _ next -> go rows priors next + PlateEnd next -> go rows priors next + affParts (AffC c) = Just (Map.empty, c) + affParts (AffL m c) = Just (m, c) + affParts NA = Nothing + +-- | [日本語]: 非ゼロ __latent 平均__ の階層 Normal prior を検出する。 +-- @u_i ~ Normal(μ, τ)@ で μ・τ __ともに単一 latent__ (係数 1・offset 0) の +-- 'Sample' を集め、 (μ, τ) の組ごとに出現順を保って群化する。 返り値の各要素は +-- @(u 名の群, μ 名, τ 名)@。 +-- +-- 平均が定数 (@AffC 0@) の reff は既存の mean-0 解析経路 (@ReffPriorIx@) が +-- 扱うのでここでは検出しない (μ が 'AffL' でないと不一致)。 係数≠1・多項・ +-- offset≠0 の平均や非 affine な μ/τ も対象外 (残差 ad に残す安全側)。 +-- rats の @alpha[i]~Normal(muAlpha,sigmaAlpha)@ / @beta[i]~Normal(muBeta,sigmaBeta)@ +-- のような varying-intercept/slope の中心化階層 prior を解析勾配へ載せるための検出器。 +-- [English]: Detects hierarchical Normal priors with a __nonzero latent mean__. +-- Collects 'Sample' nodes of the form @u_i ~ Normal(μ, τ)@ where +-- __both μ and τ are a single latent__ (coefficient 1, offset 0), and +-- groups them by the (μ, τ) pair, preserving order of appearance. Each +-- returned element is @(the group of u names, μ name, τ name)@. +-- +-- Random effects whose mean is a constant (@AffC 0@) are already handled +-- by the existing mean-0 analytic path (@ReffPriorIx@), so they are not +-- detected here (they wouldn't match anyway, since μ isn't an 'AffL'). +-- Means with coefficient ≠ 1, polynomial terms, offset ≠ 0, or non-affine +-- μ\/τ are also excluded (left on the residual AD path, the safe +-- direction). This is the detector that puts centered hierarchical +-- priors for varying-intercept\/slope models — such as rats's +-- @alpha[i]~Normal(muAlpha,sigmaAlpha)@ \/ +-- @beta[i]~Normal(muBeta,sigmaBeta)@ — onto the analytic gradient. +collectHierNormalGroups :: Model AffV r -> [([Text], Text, Text)] +collectHierNormalGroups = regroup . go + where + go (Pure _) = [] + go (Free f) = case f of + Sample n d k -> + let hit = case d of + Normal (AffL mm 0) (AffL tm 0) + | [(mn, 1)] <- Map.toList mm + , [(tn, 1)] <- Map.toList tm -> Just (n, mn, tn) + _ -> Nothing + rest = go (k (AffL (Map.singleton n 1) 0)) + in maybe rest (: rest) hit + Observe _ _ _ next -> go next + ObserveLM _ _ _ _ _ _ next -> go next + Potential _ _ next -> go next + Deterministic _ v k -> go (k v) + Data _ ys k -> go (k (map realToFrac ys, ys)) + DataIx _ is k -> go (k is) + PlateBegin _ _ next -> go next + PlateEnd next -> go next + -- (μ,τ) ごとに、u の出現順を保って群化する。 + regroup hits = + [ ([ u | (u, mn', tn') <- hits, mn' == mn, tn' == tn ], mn, tn) + | (mn, tn) <- ordNub [ (mn, tn) | (_, mn, tn) <- hits ] ] + ordNub = goN Set.empty + where goN _ [] = [] + goN s (x:xs) | x `Set.member` s = goN s xs + | otherwise = x : goN (Set.insert x s) xs + +-- | [日本語]: @a_i ~ LogNormal(μ, σ)@ 群を検出する ('collectHierNormalGroups' の +-- LogNormal 版)。μ は定数 (@Left c@・例 irt-2pl の 0) か単一 latent (@Right mn@)、 +-- σ は単一 latent (@AffL {sn:1} 0@)。σ 定数は @constPriorsOf@ が拾うのでここでは対象外。 +-- 返り値 = [(u 名, μ, σ 名)]。vecIR 経路で解析勾配 (@gradLogNormalIx@) に載せ残差 ad から +-- 外す (irt-2pl の 20-項 LogNormal prior が reverse-AD tape を張っていたのを解消)。 +-- [English]: Detects @a_i ~ LogNormal(μ, σ)@ groups (the LogNormal +-- analogue of 'collectHierNormalGroups'). μ is either a constant +-- (@Left c@; e.g. 0 in irt-2pl) or a single latent (@Right mn@); σ is a +-- single latent (@AffL {sn:1} 0@). Constant σ is picked up by +-- @constPriorsOf@, so it's excluded here. Returns [(u names, μ, σ name)]. +-- Putting these onto the analytic gradient (@gradLogNormalIx@) on the +-- vecIR path removes them from the residual AD path (this fixed the +-- 20-term LogNormal prior in irt-2pl that was building a reverse-AD +-- tape). +collectLogNormalGroups :: Model AffV r -> [([Text], Either Double Text, Text)] +collectLogNormalGroups = regroup . go + where + go (Pure _) = [] + go (Free f) = case f of + Sample n d k -> + let hit = case d of + LogNormal muA (AffL tm 0) + | [(tn, 1)] <- Map.toList tm + , Just mean <- affMean muA -> Just (n, mean, tn) + _ -> Nothing + rest = go (k (AffL (Map.singleton n 1) 0)) + in maybe rest (: rest) hit + Observe _ _ _ next -> go next + ObserveLM _ _ _ _ _ _ next -> go next + Potential _ _ next -> go next + Deterministic _ v k -> go (k v) + Data _ ys k -> go (k (map realToFrac ys, ys)) + DataIx _ is k -> go (k is) + PlateBegin _ _ next -> go next + PlateEnd next -> go next + -- μ = 定数 (AffC c / offset のみの AffL) か単一 latent (係数 1・offset 0)。 + affMean (AffC c) = Just (Left c) + affMean (AffL mm c) + | Map.null mm = Just (Left c) + | [(mn, 1)] <- Map.toList mm, c == 0 = Just (Right mn) + affMean _ = Nothing + -- (μ, σ) ごとに u の出現順を保って群化する。 + regroup hits = + [ ([ u | (u, mean', sn') <- hits, mean' == mean, sn' == sn ], mean, sn) + | (mean, sn) <- ordNub [ (mean, sn) | (_, mean, sn) <- hits ] ] + ordNub = goN Set.empty + where goN _ [] = [] + goN s (x:xs) | x `Set.member` s = goN s xs + | otherwise = x : goN (Set.insert x s) xs + +-- | [日本語]: 1 つの σ 名グループから (ブロック名, β 名, X, REff 族, σ 名, ys') を合成する。 +-- [English]: Synthesizes (block name, β names, X, REff families, σ name, +-- ys') from a single σ-name group. +synthBlock + :: Map Text (Maybe Text) + -> Text + -> [(Text, Map Text Double, Double, Text, Double)] + -> (Text, [Text], [[Double]], [REff], Text, [Double]) +synthBlock priors sn rows = + let coeffs = [ cs | (_, cs, _, _, _) <- rows ] + latents = Set.toAscList (Set.unions (map Map.keysSet coeffs)) + -- 族候補 (Phase 54.10 で「係数常 1」 を撤廃): prior = Normal(0, τ) 検出済 + -- なら係数任意。 係数は per-row 重みとして 'REff' に載せる (random slope)。 + isCand l = maybe False (/= Nothing) (Map.lookup l priors) + -- スケール τ ごとに族を貪欲抽出: 全行にちょうど 1 つ現れる族のみ採用。 + famsByTau = Map.fromListWith (++) + [ (tn, [l]) | l <- latents, isCand l + , Just (Just tn) <- [Map.lookup l priors] ] + accepted = [ (tn, Set.toAscList (Set.fromList ls)) + | (tn, ls) <- Map.toList famsByTau + , let fam = Set.fromList ls + , all (\cs -> length (filter (`Map.member` cs) (Set.toList fam)) == 1) + coeffs ] + famSet = Set.fromList (concatMap snd accepted) + reffs = [ let gws = [ gwOf fam cs | cs <- coeffs ] + ws = map snd gws + mw = if all (== 1) ws then Nothing else Just ws + in REff fam (map fst gws) (Just tn) mw Nothing + | (tn, fam) <- accepted ] + -- 各行で族中ちょうど 1 つ現れる latent の (族内 index, 係数 = 重み)。 + gwOf fam cs = head [ (j, cs Map.! l) | (j, l) <- zip [0 ..] fam + , l `Map.member` cs ] + betas = [ l | l <- latents, not (l `Set.member` famSet) ] + xs = [ [ Map.findWithDefault 0 b cs | b <- betas ] | cs <- coeffs ] + ys' = [ y - off | (_, _, off, _, y) <- rows ] + in ("__synth_lm_" <> sn, betas, xs, reffs, sn, ys') + +-- | [日本語]: 安全網②: 合成ブロックの観測尤度を、 元 model の walk 評価 +-- ('obsOnlySum' = 吸収した scalar Observe だけ足す) と probe 2 点で突合する。 +-- prior は足さないので guard 起因の ±∞ で比較が壊れない。 probe 値は +-- per-param に変えて係数の取り違えも検出する (全 latent 正値 → σ guard 安全)。 +-- [English]: Safety net (2): cross-checks the synthesized blocks' +-- observation likelihood against the original model's walked evaluation +-- ('obsOnlySum', which sums only the absorbed scalar Observes) at two +-- probe points. Since priors aren't added, ±∞ from guards doesn't break +-- the comparison. The probe values vary per parameter so that swapped +-- coefficients are also detected (all latents positive, so σ guards are +-- safe). +synthProbeOK + :: ModelP r + -> [(Text, [Text], [[Double]], [REff], Text, [Double])] + -> Set Text -> Bool +synthProbeOK m blocks obsNames = all check [(0.5, 0.07), (1.3, 0.11)] + where + names = sampleNames m + check (base, step) = + let pm = Map.fromList [ (n, base + step * fromIntegral i) + | (n, i) <- zip names [0 :: Int ..] ] + ref = obsOnlySum obsNames m pm + syn = sum [ lmObsLogSum bs xs re (LMGaussian sn) ys pm + | (_, bs, xs, re, sn, ys) <- blocks ] + in abs (ref - syn) <= 1e-9 * (1 + abs ref) + +-- | [日本語]: 名前が @sel@ に含まれる scalar 'Observe' の log-likelihood __だけ__を足す +-- walk (probe 用)。 +-- [English]: A walk that sums __only__ the log-likelihood of scalar +-- 'Observe' nodes whose name is in @sel@ (used for probing). +obsOnlySum :: Set Text -> Model Double r -> Map Text Double -> Double +obsOnlySum sel model params = go model 0 + where + go (Pure _) acc = acc + go (Free (Sample n _ k)) acc = go (k (Map.findWithDefault 0 n params)) acc + go (Free (Observe n d ys next)) acc + | n `Set.member` sel = go next (acc + obsLogSum d ys) + | otherwise = go next acc + go (Free (ObserveLM _ _ _ _ _ _ next)) acc = go next acc + go (Free (Potential _ _ next)) acc = go next acc + go (Free (Deterministic _ v k)) acc = go (k v) acc + go (Free (Data _ ys k)) acc = go (k (ys, ys)) acc + go (Free (DataIx _ is k)) acc = go (k is) acc + go (Free (PlateBegin _ _ next)) acc = go next acc + go (Free (PlateEnd next)) acc = go next acc + +-- --------------------------------------------------------------------------- +-- Phase 54.11: 非線形 μ のベクトル式 IR (M5/M6 救済) +-- --------------------------------------------------------------------------- +-- +-- 54.8 の AffV (affine 限定) の代役として、 latent 値の場所に **スカラ式ノード** +-- ('SExp') を給餌して model を walk し、 per-obs scalar @Observe (Normal μ σ)@ の +-- μ 式 (非線形可) を行ごとに収集する。 行間で式の形が同型 (定数 leaf だけが行 +-- ごとに違う) なら定数列をベクトル leaf に束ねて「ベクトル式 IR」 ('UExp') へ +-- 持ち上げ (μ⃗ = f(θ, x⃗))、 評価は VecAD の vector-op tape で行う (勾配) / +-- 素な Double ベクトル演算で行う (値)。 階層 prior (a_g ~ Normal(m, τ) の族) の +-- スカラ密度も同 IR に乗せる (M6 要件・54.9)。 +-- +-- IR 持ち上げ + 静的解析は compile 時 1 回・draw 間で再利用する (54.4b 前例)。 +-- VecAD tape 自体は per-call 構築 (spike `bench-hbm-vecir` の実測はこの構築込み)。 + +-- | [日本語]: スカラ単項演算子 ('SExp' の節)。 導関数は 'sUnD' と対。 +-- [English]: A scalar unary operator ('SExp' node). Its derivative is +-- 'sUnD'. +data SUn + = SNegO | SAbsO | SSignumO | SExpO | SLogO | SSqrtO | SRecipO + | SSinO | SCosO | STanO | SAsinO | SAcosO | SAtanO + | SSinhO | SCoshO | STanhO | SAsinhO | SAcoshO | SAtanhO + | SLgammaO -- ^ [日本語]: log Γ (密度 IR 用。 'Floating' 経由では現れない)。 [English]: log Γ (used by the density IR; never appears via 'Floating'). + deriving (Eq, Ord, Show) + +-- | [日本語]: 'SUn' の評価関数を known-function として継続に渡す CPS dispatcher。 +-- arena 実行ループが @f = sUnF op@ で closure を束縛してから +-- 要素毎に間接呼出すると GHC が unbox できず per-element boxing が出る +-- (irt-2pl prof で sUnF/sBinF/sUnD 計 30.5% time・38.9% alloc)。 call site を +-- INLINE 展開して op の case をループの外に出し、 各分岐を known-function の +-- 特殊化 unboxed ループに落とす。 演算内容・FP 順序は不変 (= posterior bit 一致)。 +-- [English]: A CPS dispatcher that hands 'SUn'\'s evaluation function to +-- the continuation as a known function. If the arena execution loop +-- binds a closure via @f = sUnF op@ and calls it indirectly per element, +-- GHC can't unbox it, producing per-element boxing (in the irt-2pl +-- profile, sUnF\/sBinF\/sUnD accounted for 30.5% of time and 38.9% of +-- allocation combined). By INLINE-expanding the call site, the case on +-- @op@ is hoisted out of the loop, and each branch falls into a +-- specialized unboxed loop for its known function. The operation content +-- and floating-point order are unchanged (i.e. the posterior is +-- bit-identical). +withSUnF :: SUn -> ((Double -> Double) -> r) -> r +withSUnF o k = case o of + SNegO -> k negate + SAbsO -> k abs + SSignumO -> k signum + SExpO -> k exp + SLogO -> k log + SSqrtO -> k sqrt + SRecipO -> k recip + SSinO -> k sin + SCosO -> k cos + STanO -> k tan + SAsinO -> k asin + SAcosO -> k acos + SAtanO -> k atan + SSinhO -> k sinh + SCoshO -> k cosh + STanhO -> k tanh + SAsinhO -> k asinh + SAcoshO -> k acosh + SAtanhO -> k atanh + SLgammaO -> k lgammaApprox +{-# INLINE withSUnF #-} + +sUnF :: SUn -> Double -> Double +sUnF o = withSUnF o id + +-- | [日本語]: 'sUnF' の導関数の CPS dispatcher ('withSUnF' と同じ意図)。 +-- [English]: A CPS dispatcher for 'sUnF'\'s derivative (same intent as +-- 'withSUnF'). +withSUnD :: SUn -> ((Double -> Double) -> r) -> r +withSUnD o k = case o of + SNegO -> k (const (-1)) + SAbsO -> k signum + SSignumO -> k (const 0) + SExpO -> k exp + SLogO -> k recip + SSqrtO -> k (\x -> 0.5 / sqrt x) + SRecipO -> k (\x -> negate (recip (x * x))) + SSinO -> k cos + SCosO -> k (negate . sin) + STanO -> k (\x -> let t = tan x in 1 + t * t) + SAsinO -> k (\x -> 1 / sqrt (1 - x * x)) + SAcosO -> k (\x -> negate (1 / sqrt (1 - x * x))) + SAtanO -> k (\x -> 1 / (1 + x * x)) + SSinhO -> k cosh + SCoshO -> k sinh + STanhO -> k (\x -> let t = tanh x in 1 - t * t) + SAsinhO -> k (\x -> 1 / sqrt (x * x + 1)) + SAcoshO -> k (\x -> 1 / sqrt (x * x - 1)) + SAtanhO -> k (\x -> 1 / (1 - x * x)) + -- digamma でなく項別微分: 評価関数 lgammaApprox の AD 微分 (walk+ad fallback / + -- 参照勾配) とビット近傍一致させる (digamma だと z=12 境界で ~1.3e-9 ズレ・56.4) + SLgammaO -> k lgammaApproxDeriv +{-# INLINE withSUnD #-} + +-- | [日本語]: 'sUnF' の導関数。 +-- [English]: 'sUnF'\'s derivative. +sUnD :: SUn -> Double -> Double +sUnD o = withSUnD o id + +-- | [日本語]: スカラ二項演算子 ('SExp' の節)。 'SMaxO' は Mixture/ZeroInflatedBinomial の +-- log-sum-exp を数値安定に組むための elementwise max (勾配は winner-take-all の +-- subgradient・'gradVecIRGo' 参照)。 'SExp' の 'Num' インスタンス経由では +-- 構築しない (Num に max が無い) — 'logSumExp2' からのみ直接 'RU2 SMaxO' で +-- 使う。 +-- [English]: A scalar binary operator ('SExp' node). 'SMaxO' is the +-- elementwise max used to build a numerically stable log-sum-exp for +-- Mixture\/ZeroInflatedBinomial (its gradient is a winner-take-all +-- subgradient; see 'gradVecIRGo'). It is never built via 'SExp'\'s 'Num' +-- instance (Num has no max) — only 'logSumExp2' constructs it directly +-- as 'RU2 SMaxO'. +data SBin = SAddO | SSubO | SMulO | SDivO | SMaxO + deriving (Eq, Ord, Show) + +-- | [日本語]: 二項演算子の CPS dispatcher ('withSUnF' と同じ意図)。 +-- [English]: A CPS dispatcher for binary operators (same intent as +-- 'withSUnF'). +withSBinF :: SBin -> ((Double -> Double -> Double) -> r) -> r +withSBinF o k = case o of + SAddO -> k (+) + SSubO -> k (-) + SMulO -> k (*) + SDivO -> k (/) + SMaxO -> k max +{-# INLINE withSBinF #-} + +sBinF :: SBin -> Double -> Double -> Double +sBinF o = withSBinF o id + +-- | [日本語]: スカラ式 IR。 latent 値の場所に流して式の木を構築する (AffV と違い +-- 非線形演算も leaf に潜らず木に残る)。 定数同士は即畳み込む ('sc1'/'sc2') ので +-- データ由来の値は常に 'SC' leaf に正規化され、 行間の形状照合が成立する。 +-- [English]: The scalar-expression IR. Fed in place of a latent value, it +-- builds an expression tree (unlike AffV, nonlinear operations don't sink +-- into a leaf but stay in the tree). Constants fold together immediately +-- ('sc1'\/'sc2'), so data-derived values always normalize to an 'SC' +-- leaf, which makes shape matching across rows work. +data SExp + = SC !Double -- ^ [日本語]: 定数 (データ・リテラル)。 [English]: A constant (data or literal). + | SV !Text -- ^ [日本語]: latent 参照。 [English]: A latent reference. + | S1 !SUn SExp + | S2 !SBin SExp SExp + +instance NFData SExp where + rnf (SC x) = rnf x + rnf (SV n) = rnf n + rnf (S1 o e) = o `seq` rnf e + rnf (S2 o a b) = o `seq` rnf a `seq` rnf b + +-- Phase 60.7: '!!!' の依存タグは IR 抽出には無関係 (既定 id)。 +instance TrackTag SExp + +-- | [日本語]: 非定数値の比較 = 値依存分岐 → error poison (AffV と同じ安全網①)。 +-- [English]: Comparing non-constant values means a value-dependent branch +-- -> error poison (the same safety net (1) as for AffV). +symPoison :: a +symPoison = error "SExp: non-constant comparison (value-dependent branch)" + +instance Eq SExp where + SC a == SC b = a == b + _ == _ = symPoison + +instance Ord SExp where + compare (SC a) (SC b) = compare a b + compare _ _ = symPoison + +-- | [日本語]: 定数畳み込み付きノード構築。 +-- [English]: Node construction with constant folding. +sc1 :: SUn -> SExp -> SExp +sc1 o (SC a) = SC (sUnF o a) +sc1 o e = S1 o e + +sc2 :: SBin -> SExp -> SExp -> SExp +sc2 o (SC a) (SC b) = SC (sBinF o a b) +sc2 o a b = S2 o a b + +instance Num SExp where + (+) = sc2 SAddO + (-) = sc2 SSubO + (*) = sc2 SMulO + negate = sc1 SNegO + abs = sc1 SAbsO + signum = sc1 SSignumO + fromInteger = SC . fromInteger + +instance Fractional SExp where + (/) = sc2 SDivO + recip = sc1 SRecipO + fromRational = SC . fromRational + +instance Floating SExp where + pi = SC pi + exp = sc1 SExpO + log = sc1 SLogO + sqrt = sc1 SSqrtO + sin = sc1 SSinO + cos = sc1 SCosO + tan = sc1 STanO + asin = sc1 SAsinO + acos = sc1 SAcosO + atan = sc1 SAtanO + sinh = sc1 SSinhO + cosh = sc1 SCoshO + tanh = sc1 STanhO + asinh = sc1 SAsinhO + acosh = sc1 SAcoshO + atanh = sc1 SAtanhO + +-- | [日本語]: 構造一致 (total・poison しない。 族 prior の同型判定用)。 +-- [English]: Structural equality (total, never poisons; used to judge +-- whether family priors are isomorphic). +sexpEq :: SExp -> SExp -> Bool +sexpEq (SC a) (SC b) = a == b +sexpEq (SV a) (SV b) = a == b +sexpEq (S1 o a) (S1 p b) = o == p && sexpEq a b +sexpEq (S2 o a c) (S2 p b d) = o == p && sexpEq a b && sexpEq c d +sexpEq _ _ = False + +-- | [日本語]: 式中の latent 参照名。 +-- [English]: The names of latent references appearing in the expression. +sexpVars :: SExp -> Set Text +sexpVars (SC _) = Set.empty +sexpVars (SV n) = Set.singleton n +sexpVars (S1 _ e) = sexpVars e +sexpVars (S2 _ a b) = sexpVars a `Set.union` sexpVars b + +-- | [日本語]: μ 式の「形の指紋」。 演算子木の形と leaf の SC/SV 区別のみで、 +-- 値・名前は含めない。 同一 σ 下で式形が混在しても指紋ごとに独立のグループとして +-- 'unifyMany' に掛けるためのキー (形違いで σ グループ丸ごと drop しない)。 +-- 「全行同一 SV」 と「行で異なる SV (族 gather)」 の区別は従来どおり unify 側の仕事。 +-- [English]: The "shape fingerprint" of a μ expression. Includes only the +-- operator tree's shape and the SC\/SV distinction at leaves — no values +-- or names. Used as a key so that, even when expression shapes mix under +-- the same σ, each fingerprint forms an independent group fed into +-- 'unifyMany' (so a σ group with mixed shapes isn't dropped entirely). +-- Distinguishing "same SV across all rows" from "different SV per row (a +-- family gather)" remains the job of the unify step, as before. +sexpShape :: SExp -> String +sexpShape (SC _) = "c" +sexpShape (SV _) = "v" +sexpShape (S1 o e) = show o ++ '(' : sexpShape e ++ ")" +sexpShape (S2 o a b) = show o ++ '(' : sexpShape a ++ ',' : sexpShape b ++ ")" + +-- | [日本語]: σ 式の「名前付き指紋」。 'sexpShape' と違い SV は latent 名を +-- 含める: σ 側は名前が違えば別グループに分ける (σ leaf を行で混ぜて族 gather に +-- 持ち上げると、 族 prior 条件を満たさない σ 同士の合流でグループ全体が drop する +-- 退行が起き得るため、 σ は保守的に「同一式 (定数値のみ行依存可)」 でキーする)。 +-- heteroscedastic (例 @exp(g0 + g1·z_i)@) は名前が全行同一・データ定数だけ行で +-- 違う形なので、 このキーで 1 グループに揃い 'unifyMany' が UC 列に持ち上げる。 +-- [English]: The "named fingerprint" of a σ expression. Unlike +-- 'sexpShape', SV includes the latent name here: on the σ side, rows +-- with different names go into different groups (mixing σ leaves across +-- rows into a family gather could regress by merging σ's that don't +-- satisfy the family-prior condition and dropping the whole group, so σ +-- is keyed conservatively by "the same expression, only constant data +-- values may vary by row"). A heteroscedastic case (e.g. +-- @exp(g0 + g1·z_i)@) has the same name across all rows with only the +-- data constant differing by row, so it lines up into one group under +-- this key and 'unifyMany' lifts it to a UC column. +sexpKeyNamed :: SExp -> String +sexpKeyNamed (SC _) = "c" +sexpKeyNamed (SV n) = "v:" ++ T.unpack n +sexpKeyNamed (S1 o e) = show o ++ '(' : sexpKeyNamed e ++ ")" +sexpKeyNamed (S2 o a b) = + show o ++ '(' : sexpKeyNamed a ++ ',' : sexpKeyNamed b ++ ")" + +-- | [日本語]: scalar 'Observe' 行の分布部。 IR 化対象の分布のみ。 +-- +-- ★分布追加チェックリスト (1 分布 = 6 箇所・1 commit): +-- 1. 'collectSymRows' に Observe 分岐 (+観測値定義域チェック → 域外行を含む +-- グループは収集時に弾く = walk の -∞ 縮退を残す安全方向) +-- 2. @keyOf@ に family タグ (位置-尺度系は scale 側を 'sexpKeyNamed') +-- 3. @tryGroup@ の unify 分岐 +-- 4. 'VecGroupSrc' / 'VecObsIR' ctor (+NFData) + 観測値定数の compile 時前計算 +-- 5. 密度式 + 値 guard ('logDensityObs' の該当分岐と完全一致。 densityIR の +-- 式のみ・勾配は記号微分で自動) +-- 6. test: 吸収確認 + 値 1e-9 + 勾配 ad 1e-9 + 中心差分 1e-4 + fallback 確認。 +-- probe 点 (0.5/1.3) の定義域を分布別に確認 (link 経由は構造上域内・ +-- パラメタ latent 直で域外なら fallback = 既知制限) +-- [English]: The distribution part of a scalar 'Observe' row. Only +-- distributions targeted for IR-ification. +-- +-- ★Checklist for adding a distribution (one distribution = 6 spots, 1 +-- commit): +-- 1. An Observe branch in 'collectSymRows' (+ an observed-value domain +-- check -> groups containing out-of-domain rows are rejected at +-- collection time, the safe direction that preserves the walk's -∞ +-- degeneration) +-- 2. A family tag in @keyOf@ (for location-scale families, key the +-- scale side with 'sexpKeyNamed') +-- 3. A unify branch in @tryGroup@ +-- 4. A 'VecGroupSrc' \/ 'VecObsIR' constructor (+ NFData) and +-- precomputing the observed-value constants at compile time +-- 5. The density expression + value guard (must exactly match the +-- corresponding branch of 'logDensityObs'; only the densityIR +-- expression is needed — the gradient is automatic via symbolic +-- differentiation) +-- 6. A test: confirm absorption + value within 1e-9 + AD gradient +-- within 1e-9 + central-difference within 1e-4 + confirm fallback. +-- Check the probe points' (0.5\/1.3) domain per distribution (via a +-- link function they're structurally in-domain; with a parameter +-- directly on a latent, out-of-domain triggers fallback — a known +-- limitation) +data SymDist + = SDGauss SExp SExp -- ^ [日本語]: Normal μ σ (σ は任意式)。 [English]: Normal μ σ (σ may be any expression). + | SDPois SExp -- ^ [日本語]: Poisson λ (λ は任意式・GLM log link は exp が式に入る)。 [English]: Poisson λ (λ may be any expression; a GLM log link puts exp in the expression). + | SDBern SExp -- ^ [日本語]: Bernoulli p (同・invLogit が式に入る)。 [English]: Bernoulli p (likewise; invLogit appears in the expression). + | SDStudT !Double SExp SExp + -- ^ [日本語]: StudentT ν μ σ (ν は SC 定数のみ吸収 = lgamma 項が定数化。 + -- ν latent は fallback・計画の scope どおり)。 + -- [English]: StudentT ν μ σ (only an SC constant ν is absorbed, i.e. + -- the lgamma term becomes a constant; a latent ν falls back, as + -- planned). + | SDCauchy SExp SExp -- ^ [日本語]: Cauchy x₀ γ。 [English]: Cauchy x₀ γ. + | SDLogis SExp SExp -- ^ [日本語]: Logistic μ s。 [English]: Logistic μ s. + | SDGumbel SExp SExp -- ^ [日本語]: Gumbel μ β。 [English]: Gumbel μ β. + | SDExpo SExp -- ^ [日本語]: Exponential rate (y ≥ 0 は収集時に確認)。 [English]: Exponential rate (y ≥ 0 is checked at collection time). + | SDWeib SExp SExp -- ^ [日本語]: Weibull k λ (y > 0 は収集時に確認)。 [English]: Weibull k λ (y > 0 is checked at collection time). + | SDLogN SExp SExp -- ^ [日本語]: LogNormal μ σ (y > 0 は収集時に確認)。 [English]: LogNormal μ σ (y > 0 is checked at collection time). + | SDGamma SExp SExp -- ^ [日本語]: Gamma α rate (y > 0 は収集時に確認)。 [English]: Gamma α rate (y > 0 is checked at collection time). + | SDBeta SExp SExp -- ^ [日本語]: Beta α β (0 < y < 1 は収集時に確認)。 [English]: Beta α β (0 < y < 1 is checked at collection time). + | SDBinom !Int SExp -- ^ [日本語]: Binomial n p (n は ctor 定数・ + -- 0 ≤ round y ≤ n は収集時に確認)。 + -- [English]: Binomial n p (n is a constructor + -- constant; 0 ≤ round y ≤ n is checked at + -- collection time). + | SDGeom SExp -- ^ [日本語]: Geometric p (round y ≥ 1 は収集時に確認)。 [English]: Geometric p (round y ≥ 1 is checked at collection time). + | SDNegBin SExp SExp -- ^ [日本語]: NegativeBinomial μ α (y ≥ 0 は収集時に確認)。 [English]: NegativeBinomial μ α (y ≥ 0 is checked at collection time). + | SDMixNorm2 SExp SExp SExp SExp SExp SExp + -- ^ [日本語]: Mixture [w1,w2] [Normal μ1 σ1, Normal μ2 σ2] (2 成分 + -- Normal 混合限定 — 任意分布族・K 成分への一般化は対象外。 w1 w2 は + -- 'Distribution.hs' の Mixture 定義どおり Σw で自動正規化するので + -- w1+w2=1 を仮定しない (w1 w2 μ1 σ1 μ2 σ2)。 + -- [English]: Mixture [w1,w2] [Normal μ1 σ1, Normal μ2 σ2] (limited to + -- a 2-component Normal mixture — generalizing to arbitrary + -- distribution families or K components is out of scope. w1 and w2 + -- are automatically normalized by Σw, per the Mixture definition in + -- 'Distribution.hs', so w1+w2=1 is not assumed (w1 w2 μ1 σ1 μ2 σ2)). + | SDZIBinom !Int SExp SExp + -- ^ [日本語]: ZeroInflatedBinomial n ψ p (n は ctor 定数・ + -- 0 ≤ round y ≤ n は収集時に確認)。 + -- [English]: ZeroInflatedBinomial n ψ p (n is a constructor constant; + -- 0 ≤ round y ≤ n is checked at collection time). + +-- | [日本語]: model を 'SExp' で walk し、 scalar @Observe@ 行 (Observe 名, 分布部, 観測値) +-- と latent prior を集める ('collectAffRows' の後継版)。 σ を任意式にし、 +-- 対象分布を Normal 限定から Poisson / Bernoulli 等にも拡張している。 他のノードは +-- 素通し (residual walk に残す)。 +-- [English]: Walks the model with 'SExp' and collects scalar @Observe@ +-- rows (Observe name, distribution part, observed value) and latent +-- priors (the successor to 'collectAffRows'). σ may be any expression, +-- and the target distributions are extended from Normal-only to also +-- include Poisson \/ Bernoulli and others. Other nodes pass through +-- unchanged (left on the residual walk). +collectSymRows + :: Model SExp r + -> ([(Text, SymDist, Double)], Map Text (Distribution SExp)) +collectSymRows = go [] Map.empty + where + go rows priors (Pure _) = (reverse rows, priors) + go rows priors (Free f) = case f of + Sample n d k -> go rows (Map.insert n d priors) (k (SV n)) + Observe nm (Normal mu sg) ys next -> + go ([ (nm, SDGauss mu sg, y) | y <- ys ] ++ rows) priors next + Observe nm (Poisson lam) ys next -> + go ([ (nm, SDPois lam, y) | y <- ys ] ++ rows) priors next + Observe nm (Bernoulli p) ys next -> + go ([ (nm, SDBern p, y) | y <- ys ] ++ rows) priors next + -- 56.3 位置-尺度系 (support = ℝ → 観測値定義域チェック不要)。 + -- StudentT は ν=SC かつ ν>0 のみ (ν≤0 は walk の -∞ を残す安全方向)。 + Observe nm (StudentT (SC nu) mu sg) ys next | nu > 0 -> + go ([ (nm, SDStudT nu mu sg, y) | y <- ys ] ++ rows) priors next + Observe nm (Cauchy loc sc) ys next -> + go ([ (nm, SDCauchy loc sc, y) | y <- ys ] ++ rows) priors next + Observe nm (Logistic mu s) ys next -> + go ([ (nm, SDLogis mu s, y) | y <- ys ] ++ rows) priors next + Observe nm (Gumbel mu be) ys next -> + go ([ (nm, SDGumbel mu be, y) | y <- ys ] ++ rows) priors next + -- 56.4 正値・区間系 (観測値定義域チェックは tryGroup の ysV 検査で)。 + Observe nm (Exponential rate) ys next -> + go ([ (nm, SDExpo rate, y) | y <- ys ] ++ rows) priors next + Observe nm (Weibull k lam) ys next -> + go ([ (nm, SDWeib k lam, y) | y <- ys ] ++ rows) priors next + Observe nm (LogNormal mu sg) ys next -> + go ([ (nm, SDLogN mu sg, y) | y <- ys ] ++ rows) priors next + Observe nm (Gamma sh rt) ys next -> + go ([ (nm, SDGamma sh rt, y) | y <- ys ] ++ rows) priors next + Observe nm (Beta al be) ys next -> + go ([ (nm, SDBeta al be, y) | y <- ys ] ++ rows) priors next + -- 56.5 離散系。 + Observe nm (Binomial n p) ys next -> + go ([ (nm, SDBinom n p, y) | y <- ys ] ++ rows) priors next + Observe nm (Geometric p) ys next -> + go ([ (nm, SDGeom p, y) | y <- ys ] ++ rows) priors next + Observe nm (NegativeBinomial mu al) ys next -> + go ([ (nm, SDNegBin mu al, y) | y <- ys ] ++ rows) priors next + -- Phase 90 A3: 2成分 Normal 混合限定 (04-low-dim-gauss-mix の + -- log_mix(θ, normal_lpdf(μ1,σ1), normal_lpdf(μ2,σ2)) と同型)。 + -- 任意分布族・3成分以上は対象外 (素通し → residual walk+ad)。 + Observe nm (Mixture [w1, w2] [Normal mu1 sg1, Normal mu2 sg2]) ys next -> + go ([ (nm, SDMixNorm2 w1 w2 mu1 sg1 mu2 sg2, y) | y <- ys ] ++ rows) priors next + Observe nm (ZeroInflatedBinomial n psi p) ys next -> + go ([ (nm, SDZIBinom n psi p, y) | y <- ys ] ++ rows) priors next + Observe _ _ _ next -> go rows priors next + ObserveLM _ _ _ _ _ _ next -> go rows priors next + Potential _ _ next -> go rows priors next + Deterministic _ v k -> go rows priors (k v) + Data _ ys k -> go rows priors (k (map realToFrac ys, ys)) + DataIx _ is k -> go rows priors (k is) + PlateBegin _ _ next -> go rows priors next + PlateEnd next -> go rows priors next + +-- | [日本語]: model を 'SExp' で walk し、 raw 'Potential' の (名前, 式) を出現順に +-- 集める。 'collectSymRows' と同じ給餌 (latent = @SV n@)。 +-- [English]: Walks the model with 'SExp' and collects raw 'Potential' +-- (name, expression) pairs in order of appearance. Uses the same feed as +-- 'collectSymRows' (latent = @SV n@). +collectSymPots :: Model SExp r -> [(Text, SExp)] +collectSymPots = go [] + where + go acc (Pure _) = reverse acc + go acc (Free f) = case f of + Sample n _ k -> go acc (k (SV n)) + Observe _ _ _ next -> go acc next + ObserveLM _ _ _ _ _ _ next -> go acc next + Potential nm v next -> go ((nm, v) : acc) next + Deterministic _ v k -> go acc (k v) + Data _ ys k -> go acc (k (map realToFrac ys, ys)) + DataIx _ is k -> go acc (k is) + PlateBegin _ _ next -> go acc next + PlateEnd next -> go acc next + +-- =========================================================================== +-- Phase 98 A2: 残余 log-joint の flat compile (Free AST 再解釈の廃止) +-- =========================================================================== +-- @logJointExclBlocks@ (Gradient.hs) は excl 吸収後の残余 log-density を求める +-- ため 'Model a r' の Free 構造を毎勾配評価で頭から walk する。 vecIR arena に +-- 吸収し切れない項 (例 06-irt-2pl: `a` の LogNormal 事前分布) が残るモデルでは、 +-- 大量の吸収済み Observe plate まで「継続のため素通り walk」する純オーバーヘッド +-- が支配する (Phase 98 A1c prof: logJointExclBlocks = 31.5% time / 41.7% alloc・ +-- Free monad `>>=`/`fmap` が十数億 entry)。 +-- +-- 本 IR は残余を **1 度の symbolic walk で flat 化**し ('CompiledResidual')、 全 +-- leapfrog で「非吸収項の畳み込み」だけを行う (Free walk 廃止)。 @CompiledLMBlock@ +-- の残余版に相当。 'SExp' 保持の純データなので値 (Double) と勾配 (AD 型) の双方で +-- 共有できる ('residualValueA' が多相)。 + +-- | [日本語]: 残余 log-joint の非吸収項を出現順に flat 化した中間表現。 +-- [English]: An intermediate representation that flattens the +-- non-absorbed terms of the residual log-joint, in order of appearance. +data CompiledResidual = CompiledResidual + { crPriors :: ![(Text, Distribution SExp)] -- ^ [日本語]: 非吸収 Sample: logDensity d (params!n)。 [English]: A non-absorbed Sample: logDensity d (params!n). + , crObs :: ![(Distribution SExp, [Double])] -- ^ [日本語]: 非吸収 Observe: obsLogSum d ys。 [English]: A non-absorbed Observe: obsLogSum d ys. + , crPots :: ![SExp] -- ^ [日本語]: 非吸収 Potential: 式値。 [English]: A non-absorbed Potential: the expression value. + } + +-- | [日本語]: 'sUnF' の 'Floating' 一般化 (density IR 専用の @SLgammaO@ を除く — +-- SLgammaO は 'Floating SExp' インスタンス経由では現れず 'Distribution SExp' に +-- 入らない)。 +-- [English]: A 'Floating' generalization of 'sUnF' (excluding @SLgammaO@, +-- which is specific to the density IR — it never shows up via the +-- 'Floating SExp' instance, so it never ends up inside 'Distribution SExp'). +sUnG :: Floating a => SUn -> a -> a +sUnG SNegO = negate +sUnG SAbsO = abs +sUnG SSignumO = signum +sUnG SExpO = exp +sUnG SLogO = log +sUnG SSqrtO = sqrt +sUnG SRecipO = recip +sUnG SSinO = sin +sUnG SCosO = cos +sUnG STanO = tan +sUnG SAsinO = asin +sUnG SAcosO = acos +sUnG SAtanO = atan +sUnG SSinhO = sinh +sUnG SCoshO = cosh +sUnG STanhO = tanh +sUnG SAsinhO = asinh +sUnG SAcoshO = acosh +sUnG SAtanhO = atanh +sUnG SLgammaO = error "sUnG: SLgammaO は残余 SExp には現れない (compileResidual の不変条件)" + +-- | [日本語]: 'sBinF' の 'Floating'+'Ord' 一般化。 +-- [English]: A 'Floating'+'Ord' generalization of 'sBinF'. +sBinG :: (Floating a, Ord a) => SBin -> a -> a -> a +sBinG SAddO = (+) +sBinG SSubO = (-) +sBinG SMulO = (*) +sBinG SDivO = (/) +sBinG SMaxO = max + +-- | [日本語]: 'SExp' を任意の 'Floating' 型で評価する (latent 参照は @lookupVar@ +-- 経由)。 'CompiledResidual' の per-eval 評価に使う (SExp 木の畳み込み・Free +-- walk 無し)。 +-- [English]: Evaluates an 'SExp' at any 'Floating' type (latent references +-- go through @lookupVar@). Used for the per-eval evaluation of +-- 'CompiledResidual' (folding the SExp tree, with no Free-monad walk). +evalSExpA :: (Floating a, Ord a) => (Text -> a) -> SExp -> a +evalSExpA lookupVar = ev + where + ev (SC x) = realToFrac x + ev (SV n) = lookupVar n + ev (S1 o e) = sUnG o (ev e) + ev (S2 o a b) = sBinG o (ev a) (ev b) + +-- | [日本語]: 残余 (excl 吸収後) を 1 度の symbolic walk で 'CompiledResidual' に +-- flat 化する。 compiled 経路で忠実再現できない残余 (非吸収 'ObserveLM') が +-- あれば 'Nothing' を返し、 呼び出し側は従来の @logJointExclBlocks@ walk に +-- fallback する。 'Deterministic'\/'Data' は walk 時に 'SExp' へインライン +-- 展開されるので収集式の 'SV' は必ず sampled latent を指す (per-eval の params +-- に存在)。 +-- [English]: Flattens the residual (after excl absorption) into a +-- 'CompiledResidual' with a single symbolic walk. Returns 'Nothing' if the +-- residual contains something the compiled path cannot faithfully +-- reproduce (a non-absorbed 'ObserveLM'), and the caller falls back to the +-- previous @logJointExclBlocks@ walk. 'Deterministic'\/'Data' get inlined +-- into 'SExp' during the walk, so any 'SV' in the collected expression +-- always refers to a sampled latent (present in the per-eval params). +compileResidual :: Set Text -> Model SExp r -> Maybe CompiledResidual +compileResidual excl = go [] [] [] + where + go ps os pots (Pure _) = + Just (CompiledResidual (reverse ps) (reverse os) (reverse pots)) + go ps os pots (Free f) = case f of + Sample n d k + | n `Set.member` excl -> go ps os pots (k (SV n)) + | otherwise -> go ((n, d) : ps) os pots (k (SV n)) + Observe n d ys next + | n `Set.member` excl -> go ps os pots next + | otherwise -> go ps ((d, ys) : os) pots next + ObserveLM nm _ _ _ _ _ next + | nm `Set.member` excl -> go ps os pots next + | otherwise -> Nothing -- 非吸収 ObserveLM は flat 化不可 → fallback + Potential n v next + | n `Set.member` excl -> go ps os pots next + | otherwise -> go ps os (v : pots) next + Deterministic _ v k -> go ps os pots (k v) + Data _ ys k -> go ps os pots (k (map realToFrac ys, ys)) + DataIx _ is k -> go ps os pots (k is) + PlateBegin _ _ next -> go ps os pots next + PlateEnd next -> go ps os pots next + +-- | [日本語]: 'CompiledResidual' の per-eval 評価 (Free walk 無し・flat list の +-- 畳み込み)。 'logJointExclBlocks excl m params' と同値 (同じ +-- @logDensity@\/'obsLogSum'・同じ params)。 sampled latent が params に無い +-- 場合は @logJointExclBlocks@ と同じく -∞ (安全網)。 +-- [English]: Per-eval evaluation of a 'CompiledResidual' (no Free-monad +-- walk, just folding the flat list). Equivalent to +-- 'logJointExclBlocks excl m params' (same @logDensity@\/'obsLogSum', same +-- params). If a sampled latent is missing from params, returns -∞ just +-- like @logJointExclBlocks@ does (a safety net). +residualValueA :: (Floating a, Ord a) => CompiledResidual -> Map Text a -> a +residualValueA cr params = priorSum + obsSum + potSum + where + ev = evalSExpA (\n -> Map.findWithDefault 0 n params) + priorSum = sum [ case Map.lookup n params of + Nothing -> negInf + Just v -> logDensity (fmap ev d) v + | (n, d) <- crPriors cr ] + obsSum = sum [ obsLogSum (fmap ev d) ys | (d, ys) <- crObs cr ] + potSum = sum [ ev v | v <- crPots cr ] + +-- | [日本語]: ベクトル式 IR。 'unifyMany' が行ごとの 'SExp' を束ねた結果で、 +-- leaf はスカラ (行に依らない) かベクトル (行ごとに値が違う) のいずれか。 +-- [English]: The vector-expression IR. The result of 'unifyMany' bundling +-- together the per-row 'SExp' values; each leaf is either a scalar +-- (row-independent) or a vector (differs per row). +data UExp + = UK !Double -- ^ [日本語]: 全行同一の定数 (スカラ)。 [English]: A constant shared by all rows (scalar). + | UC !(VS.Vector Double) -- ^ [日本語]: 行ごとの定数列 (データ列・長さ n)。 [English]: A per-row constant column (a data column, length n). + | UV !Text -- ^ [日本語]: 全行同一の latent (スカラ・broadcast)。 [English]: A latent shared by all rows (scalar, broadcast). + | UG ![Text] !(VU.Vector Int) -- ^ [日本語]: 族 gather: 行 i は member[gids_i] (長さ n)。 [English]: A family gather: row i is member[gids_i] (length n). + | U1 !SUn UExp + | U2 !SBin UExp UExp + | USum UExp -- ^ [日本語]: Σ (ベクトル → スカラ)。 raw potential 内の + -- 同型 Σ チェーンのベクトル化に使う ('absorbPot')。 + -- 中身は行依存 ('uexpIsVec') であること。 + -- [English]: A Σ (vector -> scalar). Used to + -- vectorize isomorphic Σ chains inside a raw + -- potential ('absorbPot'). The contents must be + -- row-dependent ('uexpIsVec'). + +instance NFData UExp where + rnf (UK v) = rnf v + rnf (UC v) = v `seq` () + rnf (UV n) = rnf n + rnf (UG ms g) = rnf ms `seq` g `seq` () + rnf (U1 o e) = o `seq` rnf e + rnf (U2 o a b) = o `seq` rnf a `seq` rnf b + rnf (USum e) = rnf e + +-- | [日本語]: 行ごとのスカラ式を 1 本のベクトル式に持ち上げる (形状照合)。 +-- 演算子木が全行同型で、 leaf が「全行 SC」「全行同一 SV」「行ごとに違う SV +-- (→ 族 gather 候補)」 のいずれかに揃う場合のみ成功。 +-- [English]: Lifts per-row scalar expressions into a single vector +-- expression (shape matching). Succeeds only if the operator tree is +-- isomorphic across all rows and the leaves are uniformly one of: "SC in +-- every row", "the same SV in every row", or "a different SV per row" +-- (a candidate for a family gather). +unifyMany :: [SExp] -> Maybe UExp +unifyMany [] = Nothing +unifyMany es@(e0 : _) = case e0 of + SC _ -> do + vs <- mapM (\e -> case e of SC v -> Just v; _ -> Nothing) es + Just $ case vs of + (v : rest) | all (== v) rest -> UK v + _ -> UC (VS.fromList vs) + SV _ -> do + ns <- mapM (\e -> case e of SV n -> Just n; _ -> Nothing) es + Just $ case ns of + (n0 : rest) | all (== n0) rest -> UV n0 + _ -> + let mems = Set.toAscList (Set.fromList ns) + ixm = Map.fromList (zip mems [0 :: Int ..]) + in UG mems (VU.fromList [ ixm Map.! n | n <- ns ]) + S1 o _ -> do + cs <- mapM (\e -> case e of S1 o' c | o' == o -> Just c; _ -> Nothing) es + U1 o <$> unifyMany cs + S2 o _ _ -> do + ps <- mapM (\e -> case e of S2 o' a b | o' == o -> Just (a, b); _ -> Nothing) es + U2 o <$> unifyMany (map fst ps) <*> unifyMany (map snd ps) + +-- | [日本語]: IR 中のスカラ latent 参照 (出現順・重複あり)。 +-- [English]: Scalar latent references inside the IR (in order of +-- appearance, duplicates included). +uexpScalNames :: UExp -> [Text] +uexpScalNames (UV n) = [n] +uexpScalNames (U1 _ e) = uexpScalNames e +uexpScalNames (U2 _ a b) = uexpScalNames a ++ uexpScalNames b +uexpScalNames (USum e) = uexpScalNames e +uexpScalNames _ = [] + +-- | [日本語]: IR 中の族 gather の member リスト (出現順・重複あり)。 +-- [English]: The member lists of family gathers inside the IR (in order +-- of appearance, duplicates included). +uexpFamilies :: UExp -> [[Text]] +uexpFamilies (UG ms _) = [ms] +uexpFamilies (U1 _ e) = uexpFamilies e +uexpFamilies (U2 _ a b) = uexpFamilies a ++ uexpFamilies b +uexpFamilies (USum e) = uexpFamilies e +uexpFamilies _ = [] + +-- | [日本語]: 式が行依存 (ベクトル形) か ('ruIsVec' の 'UExp' 版)。 +-- 'USum' は Σ 済みなのでスカラ。 'ruIsVec' と同じ共有無視走査の残党 +-- (absorbPot 経路で同じ指数爆発があり得る) のため、 同時に StableName memo +-- walk 化 (詳細は 'ruIsVec')。 +-- [English]: Whether an expression is row-dependent (vector-shaped); the +-- 'UExp' counterpart of 'ruIsVec'. 'USum' has already summed, so it's a +-- scalar. Left over from the same share-ignoring traversal as 'ruIsVec' +-- (the same exponential blowup can happen on the absorbPot path), so it +-- is likewise turned into a StableName memo walk (see 'ruIsVec' for +-- details). +uexpIsVec :: UExp -> Bool +uexpIsVec e0 = unsafePerformIO $ do + memo <- newIORef IM.empty + let go x0 = do + x <- evaluate x0 + sn <- makeStableName x + let h = hashStableName sn + mm <- readIORef memo + case lookup sn =<< IM.lookup h mm of + Just r -> pure r + Nothing -> do + r <- case x of + UC _ -> pure True + UG _ _ -> pure True + U1 _ e -> go e + U2 _ a b -> (||) <$> go a <*> go b + _ -> pure False + modifyIORef' memo (IM.insertWith (++) h [(sn, r)]) + pure r + go e0 +{-# NOINLINE uexpIsVec #-} + +-- | [日本語]: 出現順を保つ重複排除。 +-- [English]: Deduplication that preserves order of appearance. +ordNubO :: Ord a => [a] -> [a] +ordNubO = go Set.empty + where + go _ [] = [] + go seen (x : xs) + | x `Set.member` seen = go seen xs + | otherwise = x : go (Set.insert x seen) xs + +-- | [日本語]: IR グループ (unify 後・compile 前)。 family ごとに観測密度の +-- 組み方が違う (Gaussian 限定から Poisson / Bernoulli を追加していった経緯あり)。 +-- [English]: An IR group (after unify, before compile). The way the +-- observation density is assembled differs per family (originally +-- Gaussian-only, with Poisson / Bernoulli added later). +data VecGroupSrc + = VGGauss !UExp !UExp !(VS.Vector Double) -- ^ [日本語]: μ IR, σ IR, ys。 [English]: μ IR, σ IR, ys. + | VGPois !UExp !(VS.Vector Double) -- ^ [日本語]: λ IR, ys (全行 y ≥ 0 を確認済)。 [English]: λ IR, ys (confirmed y ≥ 0 for every row). + | VGBern !UExp !(VS.Vector Double) -- ^ [日本語]: p IR, ys (全行 round y ∈ {0,1})。 [English]: p IR, ys (round y ∈ {0,1} for every row). + | VGStudT !Double !UExp !UExp !(VS.Vector Double) + -- ^ [日本語]: ν (SC 定数), μ IR, σ IR, ys。 [English]: ν (an SC constant), μ IR, σ IR, ys. + | VGCauchy !UExp !UExp !(VS.Vector Double) -- ^ [日本語]: x₀ IR, γ IR, ys。 [English]: x₀ IR, γ IR, ys. + | VGLogis !UExp !UExp !(VS.Vector Double) -- ^ [日本語]: μ IR, s IR, ys。 [English]: μ IR, s IR, ys. + | VGGumbel !UExp !UExp !(VS.Vector Double) -- ^ [日本語]: μ IR, β IR, ys。 [English]: μ IR, β IR, ys. + | VGExpo !UExp !(VS.Vector Double) -- ^ [日本語]: rate IR, ys (全行 y ≥ 0)。 [English]: rate IR, ys (y ≥ 0 for every row). + | VGWeib !UExp !UExp !(VS.Vector Double) -- ^ [日本語]: k IR, λ IR, ys (全行 y > 0)。 [English]: k IR, λ IR, ys (y > 0 for every row). + | VGLogN !UExp !UExp !(VS.Vector Double) -- ^ [日本語]: μ IR, σ IR, ys (全行 y > 0)。 [English]: μ IR, σ IR, ys (y > 0 for every row). + | VGGamma !UExp !UExp !(VS.Vector Double) -- ^ [日本語]: α IR, rate IR, ys (全行 y > 0)。 [English]: α IR, rate IR, ys (y > 0 for every row). + | VGBeta !UExp !UExp !(VS.Vector Double) -- ^ [日本語]: α IR, β IR, ys (全行 0 node. + , sdShapes :: !(Map ShapeKey Int) -- ^ [日本語]: 形状 intern。 [English]: Shape interning. + , sdShape :: !(IM.IntMap Int) -- ^ [日本語]: ID → 形状クラス ID。 [English]: ID -> shape-class ID. + , sdNamedKs :: !(Map NamedKey Int) -- ^ [日本語]: 名前付き指紋 intern。 [English]: Named-fingerprint interning. + , sdNamed :: !(IM.IntMap Int) -- ^ [日本語]: ID → 名前付き指紋 ID。 [English]: ID -> named-fingerprint ID. + , sdVars :: !(IM.IntMap (Set Text)) -- ^ [日本語]: ID → 自由 latent 集合。 [English]: ID -> the set of free latents. + , sdNext :: !Int + , sdUnify :: !(Map [Int] UExp) + -- ^ [日本語]: unify memo (ID 列 → 'UExp')。 __全 group 共有__ = 同一部分式列は + -- 同一 'UExp' heap オブジェクトに合流し、 出力も共有付き DAG になる + -- (garch11 のような group 跨ぎ共有が下流 'compileVecIR' の identity + -- memo で 1 回だけコンパイルされるために必須)。 + -- [English]: The unify memo (ID list -> 'UExp'). + -- __Shared across all groups__: identical subexpression sequences + -- merge into the same 'UExp' heap object, so the output is also a + -- shared DAG (this is required so that cross-group sharing, e.g. in + -- garch11, is compiled only once by the identity memo in the + -- downstream 'compileVecIR'). + } + +newSDag :: IO (IORef SDagSt) +newSDag = newIORef (SDagSt IM.empty Map.empty IM.empty Map.empty + IM.empty Map.empty IM.empty IM.empty 0 Map.empty) + +sdagNodeOf :: SDagSt -> Int -> SNode +sdagNodeOf st i = sdNodes st IM.! i + +sdagShapeOf :: SDagSt -> Int -> Int +sdagShapeOf st i = sdShape st IM.! i + +sdagNamedOf :: SDagSt -> Int -> Int +sdagNamedOf st i = sdNamed st IM.! i + +sdagVarsOf :: SDagSt -> Int -> Set Text +sdagVarsOf st i = sdVars st IM.! i + +-- | [日本語]: 'SExp' を DAG に intern して ID を返す。 各 heap ノードの訪問は +-- 1 回 (StableName memo)・poison ('symPoison' 等の error thunk) はここで +-- 顕在化する ('synthVecIR' の try が捕捉する範囲内で呼ぶこと)。 +-- [English]: Interns an 'SExp' into the DAG and returns its ID. Each heap +-- node is visited exactly once (via the StableName memo); a poison (an +-- error thunk such as 'symPoison') surfaces here (call this only within +-- the scope that 'synthVecIR''s try catches). +internS :: IORef SDagSt -> SExp -> IO Int +internS ref = go + where + go e0 = do + e <- evaluate e0 + sn <- makeStableName e + let h = hashStableName sn + st <- readIORef ref + case lookup sn =<< IM.lookup h (sdStable st) of + Just i -> pure i + Nothing -> do + nd <- case e of + SC v -> pure (NC v) + SV n -> pure (NV n) + S1 o a -> N1 o <$> go a + S2 o a b -> N2 o <$> go a <*> go b + st1 <- readIORef ref + i <- case Map.lookup nd (sdStruct st1) of + Just j -> pure j + Nothing -> do + let j = sdNext st1 + shKey = case nd of + NC _ -> KC + NV _ -> KV + N1 o a -> K1 o (sdShape st1 IM.! a) + N2 o a b -> K2 o (sdShape st1 IM.! a) (sdShape st1 IM.! b) + (shId, shapes') = case Map.lookup shKey (sdShapes st1) of + Just s -> (s, sdShapes st1) + Nothing -> let s = Map.size (sdShapes st1) + in (s, Map.insert shKey s (sdShapes st1)) + nmKey = case nd of + NC _ -> MC + NV n -> MV n + N1 o a -> M1 o (sdNamed st1 IM.! a) + N2 o a b -> M2 o (sdNamed st1 IM.! a) (sdNamed st1 IM.! b) + (nmId, nameds') = case Map.lookup nmKey (sdNamedKs st1) of + Just s -> (s, sdNamedKs st1) + Nothing -> let s = Map.size (sdNamedKs st1) + in (s, Map.insert nmKey s (sdNamedKs st1)) + vs = case nd of + NC _ -> Set.empty + NV n -> Set.singleton n + N1 _ a -> sdVars st1 IM.! a + N2 _ a b -> (sdVars st1 IM.! a) `Set.union` (sdVars st1 IM.! b) + writeIORef ref st1 + { sdStruct = Map.insert nd j (sdStruct st1) + , sdNodes = IM.insert j nd (sdNodes st1) + , sdShapes = shapes' + , sdShape = IM.insert j shId (sdShape st1) + , sdNamedKs = nameds' + , sdNamed = IM.insert j nmId (sdNamed st1) + , sdVars = IM.insert j vs (sdVars st1) + , sdNext = j + 1 } + pure j + modifyIORef' ref $ \s -> + s { sdStable = IM.insertWith (++) h [(sn, i)] (sdStable s) } + pure i + +-- | [日本語]: 'unifyMany' の DAG 版: 行ごとの ID で lockstep 再帰し、 位置 +-- (= ID 列) ごとに結果 'UExp' を memo する。 同一 ID 列は同一 'UExp' オブジェクトに +-- 合流するので出力も共有付き DAG (leaf 判定・失敗条件は 'unifyMany' と同一)。 +-- [English]: The DAG version of 'unifyMany': recurses in lockstep over +-- per-row IDs, memoizing the resulting 'UExp' by position (= the ID +-- list). Identical ID lists merge into the same 'UExp' object, so the +-- output is also a shared DAG (leaf detection and failure conditions are +-- the same as in 'unifyMany'). +unifyManyD :: IORef SDagSt -> [SExp] -> IO (Maybe UExp) +unifyManyD ref es = mapM (internS ref) es >>= goIds + where + goIds [] = pure Nothing + goIds is = do + st <- readIORef ref + case Map.lookup is (sdUnify st) of + Just u -> pure (Just u) + Nothing -> do + mu <- case map (sdagNodeOf st) is of + nds@(NC _ : _) -> pure $ do + vs <- mapM (\n -> case n of NC v -> Just v; _ -> Nothing) nds + Just $ case vs of + (v : rest) | all (== v) rest -> UK v + _ -> UC (VS.fromList vs) + nds@(NV _ : _) -> pure $ do + ns <- mapM (\n -> case n of NV nm -> Just nm; _ -> Nothing) nds + Just $ case ns of + (n0 : rest) | all (== n0) rest -> UV n0 + _ -> + let mems = Set.toAscList (Set.fromList ns) + ixm = Map.fromList (zip mems [0 :: Int ..]) + in UG mems (VU.fromList [ ixm Map.! n | n <- ns ]) + nds@(N1 o _ : _) -> + case mapM (\n -> case n of N1 o' c | o' == o -> Just c + _ -> Nothing) nds of + Nothing -> pure Nothing + Just cs -> fmap (U1 o) <$> goIds cs + nds@(N2 o _ _ : _) -> + case mapM (\n -> case n of N2 o' a b | o' == o -> Just (a, b) + _ -> Nothing) nds of + Nothing -> pure Nothing + Just ps -> do + ma <- goIds (map fst ps) + case ma of + Nothing -> pure Nothing + Just ua -> fmap (U2 o ua) <$> goIds (map snd ps) + [] -> pure Nothing + case mu of + Nothing -> pure Nothing + Just u -> do + u' <- evaluate u + modifyIORef' ref $ \s -> s { sdUnify = Map.insert is u' (sdUnify s) } + pure (Just u') + +-- | [日本語]: 'UExp' の scalar leaf 名と族 gather member リストを +-- __初出順__で収集する ('uexpScalNames'/'uexpFamilies' の共有保存版)。 +-- memo (visited 集合) を IORef で外から渡し、 複数式・複数 group を跨いで +-- 1 本の memo で走る = 共有部分式は 1 回だけ訪問。 収集結果を 'ordNubO' に +-- 掛ける用途ではスキップされた再訪問分は重複除去されるだけなので結果は +-- 木 walk と一致する。 +-- [English]: Collects an 'UExp''s scalar leaf names and family-gather +-- member lists __in order of first appearance__ (the share-preserving +-- version of 'uexpScalNames'\/'uexpFamilies'). The memo (the visited +-- set) is passed in from outside via an IORef, so it runs as a single +-- memo across multiple expressions and groups — shared subexpressions +-- are visited only once. When the collected result is fed into +-- 'ordNubO', skipped re-visits are simply deduplicated away, so the +-- result matches a plain tree walk. +uexpLeavesIO :: IORef (IM.IntMap [StableName UExp]) -> UExp + -> IO ([Text], [[Text]]) +uexpLeavesIO seenRef = go + where + go u0 = do + u <- evaluate u0 + sn <- makeStableName u + let h = hashStableName sn + seen <- readIORef seenRef + if maybe False (elem sn) (IM.lookup h seen) + then pure ([], []) + else do + modifyIORef' seenRef (IM.insertWith (++) h [sn]) + case u of + UK _ -> pure ([], []) + UC _ -> pure ([], []) + UV n -> pure ([n], []) + UG ms _ -> pure ([], [ms]) + U1 _ e -> go e + U2 _ a b -> do + (s1, f1) <- go a + (s2, f2) <- go b + pure (s1 ++ s2, f1 ++ f2) + USum e -> go e + +-- | [日本語]: 'absorbPot' が 'USum' 化を試みる加算チェーンの最小項数。 +-- これ未満の和はスカラ 'U2' 連鎖のまま持つ (コスト無視できる規模)。 +-- [English]: The minimum number of terms in an addition chain before +-- 'absorbPot' attempts to turn it into a 'USum'. Sums with fewer terms +-- than this stay as a scalar 'U2' chain (the cost is negligible at that +-- size). +potSumThreshold :: Int +potSumThreshold = 8 + +-- | [日本語]: raw `potential` 式を scalar 'UExp' へ吸収する。 大きな同型 +-- 加算チェーン (項数 ≥ 'potSumThreshold') は 'unifyManyD' でベクトル化 +-- して 'USum' へ落とす (チェーン中の定数項は畳んで加算)。 吸収できない +-- 構造 (unify 失敗・行依存にならない縮退 Σ 等) は Nothing = その +-- potential ごと残差 ad に残す (安全方向・値は walk と同値のまま)。 +-- 走査は StableName memo で共有保存 (式 DAG 化での教訓: 素朴な木 walk は +-- 共有式で指数爆発する)。 +-- [English]: Absorbs a raw `potential` expression into a scalar 'UExp'. +-- Large isomorphic addition chains (number of terms ≥ +-- 'potSumThreshold') are vectorized via 'unifyManyD' and folded down +-- into a 'USum' (constant terms in the chain are collapsed and added +-- in). Structures that can't be absorbed (unify failure, a degenerate +-- Σ that doesn't come out row-dependent, etc.) yield Nothing — that +-- potential is left on the residual AD path (the safe direction; the +-- value stays equivalent to the walk). The traversal preserves sharing +-- via a StableName memo (a lesson from the expression-DAG work: a naive +-- tree walk blows up exponentially on shared subexpressions). +absorbPot :: IORef SDagSt + -> IORef (IM.IntMap [(StableName SExp, Maybe UExp)]) + -> SExp -> IO (Maybe UExp) +absorbPot ref memoRef = go + where + go e0 = do + e <- evaluate e0 + sn <- makeStableName e + let h = hashStableName sn + mm <- readIORef memoRef + case lookup sn =<< IM.lookup h mm of + Just r -> pure r + Nothing -> do + r <- build e + modifyIORef' memoRef (IM.insertWith (++) h [(sn, r)]) + pure r + build e = case e of + SC v -> pure (Just (UK v)) + SV n -> pure (Just (UV n)) + S2 SAddO _ _ -> do + terms <- flat e [] + let (cs, ts) = foldr part (0, []) terms + part t (c, acc) = case t of + SC v -> (c + v, acc) + _ -> (c, t : acc) + if length ts >= potSumThreshold + then do + mu <- unifyManyD ref ts + case mu of + Just u | uexpIsVec u -> + pure (Just (if cs == 0 then USum u + else U2 SAddO (USum u) (UK cs))) + -- 巨大チェーンをスカラ連鎖のまま素通しすると compile 側が + -- 肥大するため、 unify 不能なら吸収ごと断念 (残差 ad へ)。 + _ -> pure Nothing + else bin e + S1 o a -> fmap (U1 o) <$> go a + S2 {} -> bin e + bin (S2 o a b) = do + ma <- go a + case ma of + Nothing -> pure Nothing + Just ua -> fmap (U2 o ua) <$> go b + bin _ = pure Nothing + -- 加算 spine の平坦化 (foldl 'sum' 由来の深い左スパイン・O(項数))。 + flat e0 acc = do + e <- evaluate e0 + case e of + S2 SAddO a b -> flat a =<< flat b acc + _ -> pure (e : acc) + +-- | [日本語]: per-obs 手書き scalar 'Observe' 群から「ベクトル式 IR」 を +-- __自動合成__する ('synthGaussLMBlocks' の非線形版)。 検出できない / 安全網に +-- 掛かった場合は 'Nothing' (従来経路に fallback)。 +-- +-- 安全網 2 段 ('synthGaussLMBlocks' と同じ): ① 'SExp' の Eq/Ord は非定数 +-- 比較で error poison → 'unsafePerformIO' + 'try' で捕捉し全体 fallback +-- (poison は 'internS' の走査中に顕在化する)。 async 例外 (timeout / +-- Ctrl-C 等) は fallback にせず__透過__する (飲み込むとハングの中断が +-- 「fallback」に誤報告されることを実測確認した)。 ② IR の値評価 (観測尤度 + +-- 族 prior) を probe 2 点で walk 評価 ('obsOnlySum' + 'priorOnlySum') と +-- 突合し、 不一致なら fallback。 +-- [English]: __Automatically synthesizes__ a "vector-expression IR" from +-- a group of hand-written per-observation scalar 'Observe' nodes (the +-- nonlinear counterpart of 'synthGaussLMBlocks'). Returns 'Nothing' if +-- nothing can be detected, or a safety net trips (falls back to the +-- previous path). +-- +-- Two safety nets (same as in 'synthGaussLMBlocks'): (1) 'SExp''s Eq\/Ord +-- poisons on non-constant comparisons -> caught via 'unsafePerformIO' + +-- 'try' and the whole thing falls back (the poison surfaces during +-- 'internS''s traversal). Async exceptions (timeout, Ctrl-C, etc.) are +-- __passed through__ rather than turned into a fallback (measurement +-- confirmed that swallowing them misreports a hang's interruption as a +-- "fallback"). (2) The IR's value evaluation (observation likelihood + +-- family prior) is cross-checked at two probe points against the walked +-- evaluation ('obsOnlySum' + 'priorOnlySum'); a mismatch triggers +-- fallback. +synthVecIR :: ModelP r -> Maybe VecIRSrc +synthVecIR m = unsafePerformIO $ do + r <- try (synthVecIRWalkIO m) + case r :: Either SomeException VecIRSrc of + Left e + | Just (SomeAsyncException _) <- fromException e -> throwIO e + | otherwise -> pure Nothing + Right v@(gs, _, _) + | null gs -> pure Nothing + | vecIRProbeOK m v -> pure (Just v) + | otherwise -> pure Nothing +{-# NOINLINE synthVecIR #-} + +-- | [日本語]: 互換 wrapper (旧 pure 版と同じ表面)。 内部は 'synthVecIRWalkIO'。 +-- [English]: A compatibility wrapper (same surface as the old pure +-- version). Internally delegates to 'synthVecIRWalkIO'. +synthVecIRWalk :: ModelP r -> VecIRSrc +synthVecIRWalk = unsafePerformIO . synthVecIRWalkIO +{-# NOINLINE synthVecIRWalk #-} + +-- | [日本語]: 'synthVecIR' の合成部 (walk + 形状照合 + 族抽出)。 共有保存 +-- DAG ('internS'/'unifyManyD') ベースで実装しており、 解析は全て ID 経由 +-- O(distinct ノード数) で、 RK4 のような深い自己参照式でも指数爆発しない。 +-- 照合に失敗した σ グループは丸ごと残す (residual ad に fallback・安全方向)。 +-- 結果の式部分は構築時に正格化済み (旧実装の「呼出側が force」 は不要 — +-- 共有 DAG に rnf を掛けると経路数比例で逆に爆発するため__禁止__)。 +-- [English]: The synthesis part of 'synthVecIR' (walking + shape +-- matching + family extraction). Built on the sharing-preserving DAG +-- ('internS'\/'unifyManyD'), so all analysis goes through IDs in +-- O(distinct nodes), and even deeply self-referential expressions like +-- RK4 don't blow up exponentially. σ groups that fail to match are left +-- whole (falls back to residual AD — the safe direction). The +-- expression part of the result is already strict by construction (the +-- old implementation's "caller must force" is unnecessary — and +-- __forbidden__, since running rnf over the shared DAG would itself blow +-- up in proportion to the path count). +synthVecIRWalkIO :: ModelP r -> IO VecIRSrc +synthVecIRWalkIO m = do + let (rows, priors) = collectSymRows m + ref <- newSDag + leafSeen <- newIORef IM.empty + -- 族条件: 全 member の prior が構造同一の Normal(m, τ) で、 m/τ が member + -- 自身を参照しない (ベクトル化密度 -nG·logτ - Σ(a_j-m)²/(2τ²) が成立する形)。 + -- 構造同一判定 ('sexpEq' 相当) は intern ID の等値。 + let famOf ms = case mapM (`Map.lookup` priors) ms of + Just ds@(Normal m0 t0 : _) -> do + i0 <- internS ref m0 + j0 <- internS ref t0 + oks <- forM ds $ \d -> case d of + Normal mm tt -> do + im <- internS ref mm + jt <- internS ref tt + pure (im == i0 && jt == j0) + _ -> pure False + st <- readIORef ref + pure $ if and oks + && Set.null ((sdagVarsOf st i0 `Set.union` sdagVarsOf st j0) + `Set.intersection` Set.fromList ms) + then Just (ms, m0, t0) else Nothing + _ -> pure Nothing + -- IO 上の Maybe 連結 (MaybeT 相当の局所定義・unify 失敗の短絡用)。 + mIO >>=? k = mIO >>= maybe (pure Nothing) k + okIf cond g = pure (if cond then Just g else Nothing) + -- family 別の unify + 観測値の妥当性 (値 guard を walk と一致させるため、 + -- 観測値側の guard に掛かる行を含むグループは吸収しない = walk が -∞ を + -- 返す縮退ケースをそのまま残す安全方向)。 + tryGroup grows = do + let ysV = VS.fromList [ y | (_, _, y) <- grows ] + mg <- case [ d | (_, d, _) <- grows ] of + ds@(SDGauss{} : _) -> + unifyManyD ref [ mu | SDGauss mu _ <- ds ] >>=? \u -> + unifyManyD ref [ sg | SDGauss _ sg <- ds ] >>=? \sgU -> + pure (Just (VGGauss u sgU ysV)) + ds@(SDPois{} : _) -> + unifyManyD ref [ lam | SDPois lam <- ds ] >>=? \u -> + okIf (VS.all (>= 0) ysV) (VGPois u ysV) + ds@(SDBern{} : _) -> + unifyManyD ref [ p | SDBern p <- ds ] >>=? \u -> + okIf (VS.all (\y -> let k = round y :: Int in k == 0 || k == 1) ysV) + (VGBern u ysV) + ds@(SDStudT nu _ _ : _) -> + unifyManyD ref [ mu | SDStudT _ mu _ <- ds ] >>=? \u -> + unifyManyD ref [ sg | SDStudT _ _ sg <- ds ] >>=? \sgU -> + pure (Just (VGStudT nu u sgU ysV)) + ds@(SDCauchy{} : _) -> + unifyManyD ref [ loc | SDCauchy loc _ <- ds ] >>=? \u -> + unifyManyD ref [ sc | SDCauchy _ sc <- ds ] >>=? \scU -> + pure (Just (VGCauchy u scU ysV)) + ds@(SDLogis{} : _) -> + unifyManyD ref [ mu | SDLogis mu _ <- ds ] >>=? \u -> + unifyManyD ref [ s | SDLogis _ s <- ds ] >>=? \sU -> + pure (Just (VGLogis u sU ysV)) + ds@(SDGumbel{} : _) -> + unifyManyD ref [ mu | SDGumbel mu _ <- ds ] >>=? \u -> + unifyManyD ref [ be | SDGumbel _ be <- ds ] >>=? \beU -> + pure (Just (VGGumbel u beU ysV)) + ds@(SDExpo{} : _) -> + unifyManyD ref [ rate | SDExpo rate <- ds ] >>=? \u -> + okIf (VS.all (>= 0) ysV) (VGExpo u ysV) + ds@(SDWeib{} : _) -> + unifyManyD ref [ k | SDWeib k _ <- ds ] >>=? \kU -> + unifyManyD ref [ lam | SDWeib _ lam <- ds ] >>=? \u -> + okIf (VS.all (> 0) ysV) (VGWeib kU u ysV) + ds@(SDLogN{} : _) -> + unifyManyD ref [ mu | SDLogN mu _ <- ds ] >>=? \u -> + unifyManyD ref [ sg | SDLogN _ sg <- ds ] >>=? \sgU -> + okIf (VS.all (> 0) ysV) (VGLogN u sgU ysV) + ds@(SDGamma{} : _) -> + unifyManyD ref [ sh | SDGamma sh _ <- ds ] >>=? \shU -> + unifyManyD ref [ rt | SDGamma _ rt <- ds ] >>=? \u -> + okIf (VS.all (> 0) ysV) (VGGamma shU u ysV) + ds@(SDBeta{} : _) -> + unifyManyD ref [ al | SDBeta al _ <- ds ] >>=? \alU -> + unifyManyD ref [ be | SDBeta _ be <- ds ] >>=? \u -> + okIf (VS.all (\y -> y > 0 && y < 1) ysV) (VGBeta alU u ysV) + ds@(SDBinom{} : _) -> + unifyManyD ref [ p | SDBinom _ p <- ds ] >>=? \u -> + let nsV = VS.fromList [ fromIntegral n | SDBinom n _ <- ds ] + -- Phase 94: n を行対応化したので、 各行を自分の n で域内判定 + -- (旧: 先頭行の n を全行に流用 = merge 前提が単一 n だった)。 + domOk = and [ let k = round y :: Int in k >= 0 && k <= round nn + | (nn, y) <- zip (VS.toList nsV) (VS.toList ysV) ] + in okIf domOk (VGBinom nsV u ysV) + ds@(SDGeom{} : _) -> + unifyManyD ref [ p | SDGeom p <- ds ] >>=? \u -> + okIf (VS.all (\y -> (round y :: Int) >= 1) ysV) (VGGeom u ysV) + ds@(SDNegBin{} : _) -> + unifyManyD ref [ mu | SDNegBin mu _ <- ds ] >>=? \u -> + unifyManyD ref [ al | SDNegBin _ al <- ds ] >>=? \alU -> + okIf (VS.all (>= 0) ysV) (VGNegBin u alU ysV) + ds@(SDMixNorm2{} : _) -> + unifyManyD ref [ w1 | SDMixNorm2 w1 _ _ _ _ _ <- ds ] >>=? \w1U -> + unifyManyD ref [ w2 | SDMixNorm2 _ w2 _ _ _ _ <- ds ] >>=? \w2U -> + unifyManyD ref [ m1 | SDMixNorm2 _ _ m1 _ _ _ <- ds ] >>=? \m1U -> + unifyManyD ref [ s1 | SDMixNorm2 _ _ _ s1 _ _ <- ds ] >>=? \s1U -> + unifyManyD ref [ m2 | SDMixNorm2 _ _ _ _ m2 _ <- ds ] >>=? \m2U -> + unifyManyD ref [ s2 | SDMixNorm2 _ _ _ _ _ s2 <- ds ] >>=? \s2U -> + pure (Just (VGMixNorm2 w1U w2U m1U s1U m2U s2U ysV)) + ds@(SDZIBinom{} : _) -> + unifyManyD ref [ psi | SDZIBinom _ psi _ <- ds ] >>=? \psiU -> + unifyManyD ref [ p | SDZIBinom _ _ p <- ds ] >>=? \pU -> + let nsV = VS.fromList [ fromIntegral n | SDZIBinom n _ _ <- ds ] + domOk = and [ let k = round y :: Int in k >= 0 && k <= round nn + | (nn, y) <- zip (VS.toList nsV) (VS.toList ysV) ] + in okIf domOk (VGZIBinom nsV psiU pU ysV) + [] -> pure Nothing + case mg of + Nothing -> pure Nothing + Just g -> do + -- Phase 90 A5: family absorb (prior のベクトル化) は likelihood 側の + -- vecIR 吸収と独立の最適化。 famOf に失敗した family は fams から + -- 単に除外し (absorb しない)、 その prior は既存の `constPriorsOf` + -- (`Gradient.hs`) 経由で扱わせる。 leaf 収集の memo (leafSeen) は + -- group 跨ぎ共有 — スキップされた再訪問分の family は前の group が + -- 同一 (ms, m0, τ0) を famsAll に登録済みなので結果は不変。 + famLs <- concatMap snd <$> mapM (uexpLeavesIO leafSeen) (vgExprAll g) + famRs <- mapM famOf (ordNubO famLs) + let fams = [ f | Just f <- famRs ] + pure (Just (g, fams, Set.fromList [ nm | (nm, _, _) <- grows ])) + -- Phase 55.2-56.3 のグループキー (family タグ + σ 名前付き指紋 + μ 形状) + -- を DAG の ID で表現: 名前付き指紋 = 構造 intern ID ('sexpKeyNamed' と + -- 同値)、 形状 = 形状クラス ID ('sexpShape' と同値)。 String 指紋は + -- 長さが式の展開サイズ (= 経路数) 比例で指数爆発するため廃止 (A6)。 + key tag named shaped = do + nids <- mapM (internS ref) named + si <- internS ref shaped + st <- readIORef ref + pure (tag :: String, map (sdagNamedOf st) nids, sdagShapeOf st si) + keyOf d = case d of + SDGauss mu sg -> key "g" [sg] mu + SDPois lam -> key "p" [] lam + SDBern p -> key "b" [] p + SDStudT nu mu sg -> key ("t:" ++ show nu) [sg] mu + SDCauchy loc sc -> key "cy" [sc] loc + SDLogis mu s -> key "lg" [s] mu + SDGumbel mu be -> key "gb" [be] mu + SDExpo rate -> key "e" [] rate + SDWeib k lam -> key "w" [k] lam + SDLogN mu sg -> key "ln" [sg] mu + SDGamma sh rt -> key "ga" [sh] rt + SDBeta al be -> key "be" [al] be + SDBinom _ p -> key "bi:" [] p -- Phase 94: n を key から除外 (行対応 Vector 化で 1 group に merge) + SDGeom p -> key "ge" [] p + SDNegBin mu al -> key "nb" [al] mu + SDMixNorm2 w1 w2 m1 s1 m2 s2 -> key "mx2" [w1, w2, s1, m2, s2] m1 + SDZIBinom _ psi p -> key "zb:" [psi] p -- Phase 94: n を key から除外 + rowsK <- forM rows $ \r@(_, d, _) -> (,) r <$> keyOf d + let gkeys = ordNubO (map snd rowsK) + cands <- concat <$> forM gkeys (\gk -> do + mc <- tryGroup [ r | (r, gk') <- rowsK, gk' == gk ] + pure (maybe [] (: []) mc)) + -- Phase 90 A10: raw potential の吸収。 吸収成功した potential は + -- 'VGPot' グループ + 吸収名集合 (第3成分) に合流する。 吸収できない + -- potential はここに現れない = 従来どおり残差 ad が担う。 + -- ★potential の gather は族の**部分集合** member list になり得る + -- (icar の node1/node2 等) ため、 leaf family を famOf に掛けると + -- Observe 群由来の全体族と member が重複し disjoint チェックで全体 + -- fallback してしまう。 potential 側では族 prior 吸収を行わない — + -- 族に吸収されなかった latent の prior は @constPriorsOf@ + -- (`Gradient.hs`) が per-scalar 解析勾配で拾うので残差ゼロは保たれる。 + potMemo <- newIORef IM.empty + potRs <- forM (collectSymPots m) $ \(nm, e) -> do + mu <- absorbPot ref potMemo e + pure (fmap ((,) nm) mu) + let pots = [ p | Just p <- potRs ] + let famsAll = Map.toList (Map.fromList + [ (ms, (mx, tx)) | (_, fs, _) <- cands, (ms, mx, tx) <- fs ]) + famNames = concatMap fst famsAll + disjoint = length famNames == Set.size (Set.fromList famNames) + evaluate $ if not disjoint + then ([], [], Set.empty) -- 族 member が重複 (二重計上の危険) → 全体 fallback + else ( [ g | (g, _, _) <- cands ] ++ [ VGPot u | (_, u) <- pots ] + , [ (ms, mx, tx) | (ms, (mx, tx)) <- famsAll ] + , Set.unions [ obs | (_, _, obs) <- cands ] + `Set.union` Set.fromList (map fst pots) ) + +-- | [日本語]: グループ中の全 'UExp' フィールド (出現順)。 従来の +-- @vgExpr1@/@vgExpr2@ (最大2フィールド限定) を、 Mixture (6フィールド) 等 +-- 任意個数のフィールドを持つ family にも対応できる形に一般化した +-- (呼び出し側 'compileVecIR' の scalNames/vecLists 収集は 1 パスに統合)。 +-- [English]: All the 'UExp' fields in a group (in order of appearance). +-- Generalizes the old @vgExpr1@\/@vgExpr2@ (limited to 2 fields) so it +-- can also handle families with an arbitrary number of fields, such as +-- Mixture (6 fields). (The caller-side 'compileVecIR' scalNames\/vecLists +-- collection is unified into a single pass.) +vgExprAll :: VecGroupSrc -> [UExp] +vgExprAll (VGGauss u sg _) = [u, sg] +vgExprAll (VGPois u _) = [u] +vgExprAll (VGBern u _) = [u] +vgExprAll (VGStudT _ u sg _) = [u, sg] +vgExprAll (VGCauchy u sc _) = [u, sc] +vgExprAll (VGLogis u s _) = [u, s] +vgExprAll (VGGumbel u be _) = [u, be] +vgExprAll (VGExpo u _) = [u] +vgExprAll (VGWeib k u _) = [u, k] +vgExprAll (VGLogN u sg _) = [u, sg] +vgExprAll (VGGamma sh u _) = [u, sh] +vgExprAll (VGBeta al u _) = [u, al] +vgExprAll (VGBinom _ u _) = [u] +vgExprAll (VGGeom u _) = [u] +vgExprAll (VGNegBin u al _) = [u, al] +vgExprAll (VGMixNorm2 w1 w2 m1 s1 m2 s2 _) = [w1, w2, m1, s1, m2, s2] +vgExprAll (VGZIBinom _ psi p _) = [psi, p] +vgExprAll (VGPot u) = [u] + +-- | [日本語]: グループ中の族 gather member リスト (全 'UExp' フィールドから)。 +-- [English]: The family-gather member lists in a group (gathered from +-- all its 'UExp' fields). +vgFamilies :: VecGroupSrc -> [[Text]] +vgFamilies = concatMap uexpFamilies . vgExprAll + +-- | [日本語]: 名前が @sel@ に含まれる raw 'Potential' の値__だけ__を足す walk +-- ('obsOnlySum' の potential 版・probe 用)。 +-- [English]: A walk that sums __only__ the values of raw 'Potential' +-- nodes whose name is in @sel@ (the potential counterpart of +-- 'obsOnlySum', used for probing). +potOnlySum :: Set Text -> Model Double r -> Map Text Double -> Double +potOnlySum sel model params = go model 0 + where + go (Pure _) acc = acc + go (Free (Sample n _ k)) acc = go (k (Map.findWithDefault 0 n params)) acc + go (Free (Observe _ _ _ next)) acc = go next acc + go (Free (ObserveLM _ _ _ _ _ _ next)) acc = go next acc + go (Free (Potential n v next)) acc + | n `Set.member` sel = go next (acc + v) + | otherwise = go next acc + go (Free (Deterministic _ v k)) acc = go (k v) acc + go (Free (Data _ ys k)) acc = go (k (ys, ys)) acc + go (Free (DataIx _ is k)) acc = go (k is) acc + go (Free (PlateBegin _ _ next)) acc = go next acc + go (Free (PlateEnd next)) acc = go next acc + +-- | [日本語]: 名前が @sel@ に含まれる 'Sample' の prior log-density +-- __だけ__を足す walk ('obsOnlySum' の prior 版・probe 用)。 +-- [English]: A walk that sums __only__ the prior log-density of +-- 'Sample' nodes whose name is in @sel@ (the prior counterpart of +-- 'obsOnlySum', used for probing). +priorOnlySum :: Set Text -> Model Double r -> Map Text Double -> Double +priorOnlySum sel model params = go model 0 + where + go (Pure _) acc = acc + go (Free (Sample n d k)) acc = + let v = Map.findWithDefault 0 n params + in go (k v) (if n `Set.member` sel then acc + logDensity d v else acc) + go (Free (Observe _ _ _ next)) acc = go next acc + go (Free (ObserveLM _ _ _ _ _ _ next)) acc = go next acc + go (Free (Potential _ _ next)) acc = go next acc + go (Free (Deterministic _ v k)) acc = go (k v) acc + go (Free (Data _ ys k)) acc = go (k (ys, ys)) acc + go (Free (DataIx _ is k)) acc = go (k is) acc + go (Free (PlateBegin _ _ next)) acc = go next acc + go (Free (PlateEnd next)) acc = go next acc + +-- | [日本語]: 安全網②: IR の値 (観測尤度 + 族 prior) を、 元 model の +-- walk 評価と probe 2 点で突合する ('synthProbeOK' と同じ流儀。 +-- probe 値は per-param に変えて係数取り違えも検出・全 latent 正値で guard 安全)。 +-- [English]: Safety net (2): cross-checks the IR's value (observation +-- likelihood + family prior) against the original model's walked +-- evaluation at two probe points (the same approach as +-- 'synthProbeOK'; probe values vary per-parameter so a mixed-up +-- coefficient can also be caught, and are kept positive for every +-- latent so any domain guard stays safe). +vecIRProbeOK :: ModelP r -> VecIRSrc -> Bool +vecIRProbeOK m (gs, fams, obsNames) = all check [(0.5, 0.07), (1.3, 0.11)] + where + names = sampleNames m + -- 各 latent の prior 分布 → 制約変換種別。 probe 点を **その latent の台** + -- (正値 / 単位区間) に写すため。 素の base+step は有界台の latent + -- (Beta 等) で域外になり sqrt(1-pc²) 等が NaN 化 → 誤 fallback していた + -- (Phase 80.2)。 'fromUnconstrained' で恒等 / exp / sigmoid を通し常に域内。 + -- + -- Phase 90 A3: @base + step·i@ の @i@ は全 latent 通し番号なので、 + -- 個体ごとの random effect 等で latent 数が多い階層モデル (M=385 等) では + -- 添字の大きい latent の probe 値が発散する (例: i=387 → 27.6)。 + -- unconstrained (Normal 等・変換なし) latent は exp/log の非線形演算を + -- 経由すると `exp(-eps)` アンダーフロー → `log(1-p) = -Infinity` → + -- `ref - syn = -Inf - (-Inf) = NaN` で誤って probe 不一致 (05-mh で実測 + -- 発覚)。 通し番号を法 16 で折り返し、 添字が多くても probe 値の広がりを + -- 一定に保つ (元の「係数取り違え検出のため異なる値を使う」意図は + -- 16 通りの相異なる値で十分に保たれる)。 + (_, priors) = collectSymRows m + trOf n = maybe UnconstrainedT distToTransform (Map.lookup n priors) + ixOf = Map.fromList (zip names [0 ..]) + cvi = compileVecIR ixOf gs fams + famSet = Set.fromList (concat [ ms | (ms, _, _) <- fams ]) + check (base, step) = + let pm = Map.fromList + [ (n, fromUnconstrained (trOf n) (base + step * fromIntegral (i `mod` 16))) + | (n, i) <- zip names [0 :: Int ..] ] + pc = VS.fromList [ pm Map.! n | n <- names ] + ref = obsOnlySum obsNames m pm + priorOnlySum famSet m pm + + potOnlySum obsNames m pm -- 吸収済み raw potential (A10) + syn = vecIRValue cvi pc + in abs (ref - syn) <= 1e-9 * (1 + abs ref) + +-- | [日本語]: index 解決済みのベクトル式 IR ノード。 latent 参照は leaf +-- __位置__ ('cvScalIx' \/ 'cvVecIxs' の添字) に解決済み (per-call の Text +-- lookup なし)。 +-- [English]: An index-resolved vector-expression IR node. Latent +-- references are already resolved to leaf __positions__ (indices into +-- 'cvScalIx' \/ 'cvVecIxs'), so there is no per-call Text lookup. +data RUExp + = RUK !Double + | RUC !(VS.Vector Double) + | RUV !Int -- ^ [日本語]: scalar leaf 位置。 [English]: A scalar-leaf position. + | RUG !Int !(VU.Vector Int) -- ^ [日本語]: vector leaf 位置 + gids (gather・長さ = 行数)。 [English]: A vector-leaf position plus gids (a gather, length = number of rows). + | RUVec !Int -- ^ [日本語]: vector leaf そのもの (族 prior 用)。 [English]: The vector leaf itself (used for family priors). + | RU1 !SUn RUExp + | RU2 !SBin RUExp RUExp + | RUSum RUExp -- ^ [日本語]: Σ (ベクトル → スカラ)。 [English]: A Σ (vector -> scalar). + deriving (Eq, Ord) -- ^ [日本語]: compile 時 hash-consing (CSE) 用。 [English]: Used for hash-consing (CSE) at compile time. + +-- | [日本語]: 式が行依存 (ベクトル形) か (compile 時に静的に決まる)。 +-- 素朴な構造再帰は共有 DAG 上で__経路数__に比例して走り、 garch11 の σ +-- 逐次再帰 (sPrev² = 2 参照 × T 段 → Σ2^t 経路) で指数ハングした +-- (prof 99.8% time・entries 2^30)。 StableName memo walk で +-- O(distinct)/呼出に是正 ('compileVecIR' と同流儀・引数決定的なので参照透過)。 +-- [English]: Whether an expression is row-dependent (vector-shaped); +-- this is decided statically at compile time. A naive structural +-- recursion runs in time proportional to the __path count__ over the +-- shared DAG, and hung exponentially on garch11's sequential σ +-- recursion (sPrev² referenced twice per step over T steps -> Σ2^t +-- paths; profiling showed 99.8% of time, 2^30 entries). Fixed to +-- O(distinct nodes) per call with a StableName memo walk (the same +-- approach as 'compileVecIR'; referentially transparent since the +-- arguments are deterministic). +ruIsVec :: RUExp -> Bool +ruIsVec e0 = unsafePerformIO $ do + memo <- newIORef IM.empty + let go x0 = do + x <- evaluate x0 + sn <- makeStableName x + let h = hashStableName sn + mm <- readIORef memo + case lookup sn =<< IM.lookup h mm of + Just r -> pure r + Nothing -> do + r <- case x of + RUC{} -> pure True + RUG{} -> pure True + RUVec{} -> pure True + RU1 _ e -> go e + RU2 _ a b -> (||) <$> go a <*> go b + _ -> pure False + modifyIORef' memo (IM.insertWith (++) h [(sn, r)]) + pure r + go e0 +{-# NOINLINE ruIsVec #-} + +infixl 6 .+#, .-# +infixl 7 .*#, ./# +-- | [日本語]: 密度 IR 構築用の局所演算子 (export しない)。 恒等演算 +-- (x·1 / x+0 / x-0 / x÷1) は構築時に畳む = 'ru2Smart'。 μ 合成 +-- (designHBMProgram) が汎用に作る @0 + coef·x@ 連鎖が radon で 919 セル級 +-- ベクトル命令 20 本中 6 本 (~29%) を占めると命令列 dump のプロファイルで +-- 実測されたため。 x·0→0 は IEEE 非保存 (x=Inf/NaN で NaN) ゆえ畳まない。 +-- [English]: Local operators for building the density IR (not +-- exported). Identity operations (x·1 \/ x+0 \/ x-0 \/ x÷1) are folded at +-- construction time via 'ru2Smart'. This is because a profiler +-- instruction-stream dump measured that the @0 + coef·x@ chains +-- generically produced by μ composition (designHBMProgram) accounted for +-- 6 of 20 (~29%) vector instructions at radon's ~919-cell scale. x·0→0 +-- is not folded, since it isn't IEEE-safe (x=Inf\/NaN gives NaN). +(.+#), (.-#), (.*#), (./#) :: RUExp -> RUExp -> RUExp +(.+#) = ru2Smart SAddO +(.-#) = ru2Smart SSubO +(.*#) = ru2Smart SMulO +(./#) = ru2Smart SDivO + +-- | [日本語]: 'RU2' の恒等演算畳み込み smart constructor。 定数同士は +-- 即値化 (SExp の 'sc2' と同じ流儀)。 +-- [English]: The identity-folding smart constructor for 'RU2'. Two +-- constants are folded into a literal value immediately (the same +-- approach as SExp's 'sc2'). +ru2Smart :: SBin -> RUExp -> RUExp -> RUExp +ru2Smart o (RUK a) (RUK b) = RUK (sBinF o a b) +ru2Smart SAddO (RUK 0) b = b +ru2Smart SAddO a (RUK 0) = a +ru2Smart SSubO a (RUK 0) = a +ru2Smart SMulO (RUK 1) b = b +ru2Smart SMulO a (RUK 1) = a +ru2Smart SDivO a (RUK 1) = a +ru2Smart o a b = RU2 o a b + +-- | [日本語]: 'RU1' の smart constructor (命令融合)。 +-- +-- - 定数は即値化 ('ru2Smart' と同流儀)。 +-- - __@log(exp x) → x@ の代数畳み込み__: GLM log-link で観測密度が +-- @Σ y·log(λ)@・@λ = exp(η)@ を組むと @log(exp η)@ の往復が 1 命令 +-- (観測長ぶんの `SLogO` pass + その backward) として残る。 これを恒等に +-- 畳んで η を直接使う (`Σ y·η − exp η` = Poisson の log-link 標準形)。 +-- 数学的に厳密な恒等 (exp は常に正・log(exp x)=x)。 FP では ulp 差が出る +-- ため __draws は変わる__ (回帰判定は PyMC 事後突合による別 gate)。 +-- ⚠ @exp(log x) → x@ は x>0 でしか成立せず (log(負)=NaN) 一般には不正 = +-- 畳まない。 log∘exp のみ。 +-- +-- [English]: The smart constructor for 'RU1' (instruction fusion). +-- +-- - Constants are folded into a literal value immediately (same approach +-- as 'ru2Smart'). +-- - __Algebraic folding of @log(exp x) → x@__: when a GLM log-link +-- observation density builds @Σ y·log(λ)@ with @λ = exp(η)@, the +-- round trip @log(exp η)@ remains as one instruction (an `SLogO` pass +-- over the observation length, plus its backward pass). This folds +-- that away as an identity, using η directly (giving +-- `Σ y·η − exp η`, the standard Poisson log-link form). Mathematically +-- an exact identity (exp is always positive, so log(exp x)=x), but +-- floating point produces ulp-level differences, so +-- __draws do change__ (regression checking is a separate gate, done +-- by comparing posteriors against PyMC). +-- ⚠ @exp(log x) → x@ only holds for x>0 (log of a negative is NaN), so +-- it is not valid in general and is not folded. Only log∘exp is. +ru1Smart :: SUn -> RUExp -> RUExp +ru1Smart o (RUK a) = RUK (sUnF o a) +ru1Smart SLogO (RU1 SExpO x) = x +ru1Smart o e = RU1 o e + +-- | [日本語]: 前処理済みの IR (@CompiledLMBlock@ の IR 版)。 compile 時に +-- 1 度だけ作り、 per-call は leaf 値の差し替え + tape/値評価のみ。 +-- index 解決 + 観測値由来の定数前計算済みのグループ。 +-- [English]: A pre-processed IR (the IR counterpart of +-- @CompiledLMBlock@). Built once at compile time; each call only swaps +-- in leaf values and evaluates the tape\/value. A group with indices +-- already resolved and observation-derived constants already +-- precomputed. +data VecObsIR + = VOGauss !RUExp !RUExp !(VS.Vector Double) + -- ^ [日本語]: (μ IR, σ IR, ys)。 σ IR が行依存 (RUC を含む) なら + -- heteroscedastic ベクトル密度。 + -- [English]: (μ IR, σ IR, ys). If σ IR is row-dependent (contains + -- RUC), this is a heteroscedastic vector density. + | VOPois !RUExp !(VS.Vector Double) !Double + -- ^ [日本語]: (λ IR, ys, Σ log y_i! 前計算)。 + -- logp = Σ(y_i·logλ_i - λ_i) - Σlog y_i! (y は定数なので factorial 項は + -- compile 時前計算・勾配に寄与しない)。 + -- [English]: (λ IR, ys, precomputed Σ log y_i!). logp = + -- Σ(y_i·logλ_i - λ_i) - Σlog y_i! (since y is a constant, the + -- factorial term is precomputed at compile time and contributes + -- nothing to the gradient). + | VOBern !RUExp !(VS.Vector Double) + -- ^ [日本語]: (p IR, yb)。 yb = round 済 0/1 列。 + -- logp = Σ(yb_i·log p_i + (1-yb_i)·log(1-p_i)) + -- ('logDensityObs' の round 分岐を係数化)。 + -- [English]: (p IR, yb). yb is the already-rounded 0\/1 column. + -- logp = Σ(yb_i·log p_i + (1-yb_i)·log(1-p_i)) (turns + -- 'logDensityObs''s round branch into coefficients). + | VOStudT !Double !RUExp !RUExp !(VS.Vector Double) + -- ^ [日本語]: (ν 定数, μ IR, σ IR, ys)。 lgamma 項は ν=SC なので + -- compile 時定数 ('lgammaApprox' で walk と完全一致)。 + -- [English]: (a constant ν, μ IR, σ IR, ys). The lgamma term is a + -- compile-time constant since ν=SC (matches the walk exactly via + -- 'lgammaApprox'). + | VOCauchy !RUExp !RUExp !(VS.Vector Double) + -- ^ [日本語]: (x₀ IR, γ IR, ys)。 logp = -n·logπ - Σlogγ - Σlog(1+z_i²)。 + -- [English]: (x₀ IR, γ IR, ys). logp = -n·logπ - Σlogγ - Σlog(1+z_i²). + | VOLogis !RUExp !RUExp !(VS.Vector Double) + -- ^ [日本語]: (μ IR, s IR, ys)。 logp = -Σz_i - Σlog s - 2·Σlog(1+exp(-z_i))。 + -- [English]: (μ IR, s IR, ys). logp = -Σz_i - Σlog s - 2·Σlog(1+exp(-z_i)). + | VOGumbel !RUExp !RUExp !(VS.Vector Double) + -- ^ [日本語]: (μ IR, β IR, ys)。 logp = -Σlog β - Σz_i - Σexp(-z_i)。 + -- [English]: (μ IR, β IR, ys). logp = -Σlog β - Σz_i - Σexp(-z_i). + | VOExpo !RUExp !(VS.Vector Double) + -- ^ [日本語]: (rate IR, ys)。 logp = Σlog rate_i - Σ rate_i·y_i。 + -- [English]: (rate IR, ys). logp = Σlog rate_i - Σ rate_i·y_i. + | VOWeib !RUExp !RUExp !(VS.Vector Double) + -- ^ [日本語]: (k IR, λ IR, log ys 前計算)。 (y/λ)^k は + -- exp(k·(log y - log λ)) で初等化 (@**@ と ulp 差のみ)。 + -- [English]: (k IR, λ IR, precomputed log ys). (y/λ)^k is reduced to + -- exp(k·(log y - log λ)) (differs from @**@ only at the ulp level). + | VOLogN !RUExp !RUExp !(VS.Vector Double) !Double + -- ^ [日本語]: (μ IR, σ IR, log ys 前計算, -Σlog y 定数)。 密度 = + -- Gaussian ノード ('VOGauss' の densityIR) 再利用 + 定数。 + -- [English]: (μ IR, σ IR, precomputed log ys, the constant -Σlog y). + -- The density reuses the Gaussian node ('VOGauss''s densityIR) plus + -- a constant. + | VOGamma !RUExp !RUExp !(VS.Vector Double) !(VS.Vector Double) + -- ^ [日本語]: (α IR, rate IR, ys, log ys 前計算)。 lgammaΓ(α) は + -- @SLgammaO@ (値 lgammaApprox / 導関数 lgammaApproxDeriv)。 + -- [English]: (α IR, rate IR, ys, precomputed log ys). lgammaΓ(α) is + -- @SLgammaO@ (value via lgammaApprox, derivative via + -- lgammaApproxDeriv). + | VOBeta !RUExp !RUExp !(VS.Vector Double) !(VS.Vector Double) + -- ^ [日本語]: (α IR, β IR, log ys, log (1-ys) 前計算)。 + -- [English]: (α IR, β IR, precomputed log ys, precomputed log (1-ys)). + | VOBinom !RUExp !(VS.Vector Double) !(VS.Vector Double) !Double + -- ^ [日本語]: (p IR, k 列 (raw y・walk の kA と一致), n-k 列, + -- Σ logC(n,round y) 定数)。 Bernoulli 式の係数一般化。 + -- [English]: (p IR, a k column (raw y, matching the walk's kA), + -- an n-k column, the constant Σ logC(n,round y)). A coefficient + -- generalization of the Bernoulli form. + | VOGeom !RUExp !(VS.Vector Double) + -- ^ [日本語]: (p IR, k 列 (raw y))。 + -- logp = Σ(k_i-1)·log(1-p_i) + Σlog p_i。 + -- [English]: (p IR, a k column (raw y)). + -- logp = Σ(k_i-1)·log(1-p_i) + Σlog p_i. + | VONegBin !RUExp !RUExp !(VS.Vector Double) !Double + -- ^ [日本語]: (μ IR, α IR, k 列 (raw y), Σ lgammaΓ(k_i+1) 定数)。 + -- lgammaΓ(k_i+α) は @SLgammaO@ の elementwise 適用。 + -- [English]: (μ IR, α IR, a k column (raw y), the constant + -- Σ lgammaΓ(k_i+1)). lgammaΓ(k_i+α) is an elementwise application of + -- @SLgammaO@. + | VOPot !RUExp + -- ^ [日本語]: raw `potential` 項。 scalar 形 (内部 Σ は 'RUSum')。 + -- logp 寄与 = 式の値そのもの・guard なし (walk の 'Potential' 加算と + -- 同値)。 + -- [English]: A raw `potential` term. Scalar-shaped (any internal Σ + -- is an 'RUSum'). Its contribution to logp is the expression's + -- value itself, with no guard (equivalent to the 'Potential' + -- addition performed by the walk). + | VOMixNorm2 !RUExp !RUExp !RUExp !RUExp !RUExp !RUExp !(VS.Vector Double) + -- ^ [日本語]: (w1 IR, w2 IR, μ1 IR, σ1 IR, μ2 IR, σ2 IR, ys)。 + -- 2成分 Normal 混合限定。 logp_i = logsumexp(logw1-logtotal+lpdf1_i, + -- logw2-logtotal+lpdf2_i) ('Distribution.hs' の Mixture と数式一致)。 + -- [English]: (w1 IR, w2 IR, μ1 IR, σ1 IR, μ2 IR, σ2 IR, ys). + -- Two-component Normal mixtures only. logp_i = + -- logsumexp(logw1-logtotal+lpdf1_i, logw2-logtotal+lpdf2_i) (matches + -- the formula for Mixture in 'Distribution.hs'). + | VOZIBinom !(VS.Vector Double) !RUExp !RUExp !(VS.Vector Double) !(VS.Vector Double) + !(VS.Vector Double) !(VS.Vector Double) + -- ^ [日本語]: (n 列 (行対応), ψ IR, p IR, mask0 列 (y=0なら1), y 列 (raw), + -- n-y 列, logC(n,y) 列)。 y==0/y>0 の分岐は compile 時に mask 列へ + -- 落とし、 両方の分岐式を全行 elementwise に計算してから mask で選択 + -- (group 分割はしない — family gather の disjoint 検査を壊さない + -- 安全設計)。 + -- [English]: (an n column (per row), ψ IR, p IR, a mask0 column (1 + -- if y=0), a y column (raw), an n-y column, a logC(n,y) column). + -- The y==0\/y>0 branch is lowered to a mask column at compile time: + -- both branch expressions are computed elementwise for every row, + -- then selected via the mask (no group splitting — a safety design + -- that avoids breaking the family-gather disjointness check). + +data CompiledVecIR = CompiledVecIR + { cvProg :: !VecProgram + -- ^ [日本語]: 値 + 勾配の静的命令列 (compile 時 1 回生成 = per-call の + -- tape 構築を撤去し「tape を compile 時に固定」)。 + -- [English]: The static value+gradient instruction stream (generated + -- once at compile time — removes per-call tape construction by + -- "fixing the tape at compile time"). + , cvScalIx :: !(VU.Vector Int) + -- ^ [日本語]: scalar leaf → param index。 [English]: scalar leaf -> param index. + , cvVecIxs :: ![VU.Vector Int] + -- ^ [日本語]: vector leaf → member param indices。 [English]: vector leaf -> member param indices. + } + + +-- | [日本語]: 'VecIRSrc' を param index に解決する (静的・1 回)。 +-- 'UExp'\/'SExp' → 'RUExp' の変換と leaf 収集を identity memo +-- (StableName) で共有保存し、 命令列生成は intern 済み DAG ('RUNode') 上で +-- 行う (旧実装は全て共有無視の木 walk + 'Map' 構造キーの CSE = 深い共有式 +-- で指数)。 表面は pure のまま ('Gradient.hs' の呼出互換・引数に対し +-- 決定的なので参照透過)。 +-- [English]: Resolves a 'VecIRSrc' into param indices (statically, once). +-- The 'UExp'\/'SExp' -> 'RUExp' conversion and leaf collection preserve +-- sharing via an identity memo (StableName), and instruction-stream +-- generation runs over the already-interned DAG ('RUNode') (the old +-- implementation used a share-ignoring tree walk plus 'Map'-structural-key +-- CSE throughout, which was exponential on deeply shared expressions). +-- The surface stays pure (compatible with how 'Gradient.hs' calls it; +-- referentially transparent since it's deterministic in its arguments). +compileVecIR + :: Map Text Int + -> [VecGroupSrc] -> [([Text], SExp, SExp)] + -> CompiledVecIR +compileVecIR ixOf gs fams = unsafePerformIO (compileVecIRIO ixOf gs fams) +{-# NOINLINE compileVecIR #-} + +compileVecIRIO + :: Map Text Int + -> [VecGroupSrc] -> [([Text], SExp, SExp)] + -> IO CompiledVecIR +compileVecIRIO ixOf gs fams = do + -- leaf 収集 (初出順・memo は全式共有 = ordNubO 後の結果は旧木 walk と同一) + seenRef <- newIORef IM.empty + leafPairs <- mapM (uexpLeavesIO seenRef) [ e | g <- gs, e <- vgExprAll g ] + sdag <- newSDag + famVars <- forM fams $ \(_, mx, tx) -> do + im <- internS sdag mx + it <- internS sdag tx + st <- readIORef sdag + pure (Set.toList (sdagVarsOf st im `Set.union` sdagVarsOf st it)) + let scalNames = ordNubO (concatMap fst leafPairs ++ concat famVars) + vecLists = ordNubO (concatMap snd leafPairs + ++ [ ms | (ms, _, _) <- fams ]) + sPos = Map.fromList (zip scalNames [0 :: Int ..]) + vPos = Map.fromList (zip vecLists [0 :: Int ..]) + -- UExp/SExp → RUExp (identity memo で共有保存・'ru2Smart' の畳みは従来どおり) + rUMemo <- newIORef IM.empty + rSMemo <- newIORef IM.empty + let rU u0 = do + u <- evaluate u0 + sn <- makeStableName u + let h = hashStableName sn + mm <- readIORef rUMemo + case lookup sn =<< IM.lookup h mm of + Just r -> pure r + Nothing -> do + r <- case u of + UK v -> pure (RUK v) + UC v -> pure (RUC v) + UV n -> pure (RUV (sPos Map.! n)) + UG ms gids -> pure (RUG (vPos Map.! ms) gids) + U1 o e -> RU1 o <$> rU e + U2 o a b -> ru2Smart o <$> rU a <*> rU b + USum e -> RUSum <$> rU e + r' <- evaluate r + modifyIORef' rUMemo (IM.insertWith (++) h [(sn, r')]) + pure r' + rS e0 = do + e <- evaluate e0 + sn <- makeStableName e + let h = hashStableName sn + mm <- readIORef rSMemo + case lookup sn =<< IM.lookup h mm of + Just r -> pure r + Nothing -> do + r <- case e of + SC v -> pure (RUK v) + SV n -> pure (RUV (sPos Map.! n)) + S1 o x -> RU1 o <$> rS x + S2 o a b -> ru2Smart o <$> rS a <*> rS b + r' <- evaluate r + modifyIORef' rSMemo (IM.insertWith (++) h [(sn, r')]) + pure r' + cgOf g = case g of + VGGauss u sg ys -> VOGauss <$> rU u <*> rU sg <*> pure ys + VGPois u ys -> + (\r -> VOPois r ys + (VS.sum (VS.map (logFactorial . (round :: Double -> Int)) ys))) + <$> rU u + VGBern u ys -> + (\r -> VOBern r (VS.map (\y -> fromIntegral (round y :: Int)) ys)) + <$> rU u + VGStudT nu u sg ys -> VOStudT nu <$> rU u <*> rU sg <*> pure ys + VGCauchy u sc ys -> VOCauchy <$> rU u <*> rU sc <*> pure ys + VGLogis u s ys -> VOLogis <$> rU u <*> rU s <*> pure ys + VGGumbel u be ys -> VOGumbel <$> rU u <*> rU be <*> pure ys + VGExpo u ys -> (\r -> VOExpo r ys) <$> rU u + VGWeib k u ys -> + (\rk r -> VOWeib rk r (VS.map log ys)) <$> rU k <*> rU u + VGLogN u sg ys -> + let lys = VS.map log ys + in (\r rsg -> VOLogN r rsg lys (negate (VS.sum lys))) + <$> rU u <*> rU sg + VGGamma sh u ys -> + (\rsh r -> VOGamma rsh r ys (VS.map log ys)) <$> rU sh <*> rU u + VGBeta al u ys -> + (\ral r -> VOBeta ral r (VS.map log ys) + (VS.map (\y -> log (1 - y)) ys)) + <$> rU al <*> rU u + VGBinom nv u ys -> + (\r -> VOBinom r ys (VS.zipWith (-) nv ys) + (VS.sum (VS.zipWith (\nn y -> logBinomCoeff (round nn) (round y)) nv ys))) + <$> rU u + VGGeom u ys -> (\r -> VOGeom r ys) <$> rU u + VGNegBin u al ys -> + (\r ral -> VONegBin r ral ys + (VS.sum (VS.map (\y -> lgammaApprox (y + 1)) ys))) + <$> rU u <*> rU al + VGMixNorm2 w1 w2 m1 s1 m2 s2 ys -> + (\a b c d e f -> VOMixNorm2 a b c d e f ys) + <$> rU w1 <*> rU w2 <*> rU m1 <*> rU s1 <*> rU m2 <*> rU s2 + VGZIBinom nv psi p ys -> + (\rpsi rp -> VOZIBinom nv rpsi rp + (VS.map (\y -> if round y == (0 :: Int) then 1 else 0) ys) + ys + (VS.zipWith (-) nv ys) + (VS.zipWith (\nn y -> logBinomCoeff (round nn) (round y)) nv ys)) + <$> rU psi <*> rU p + VGPot u -> VOPot <$> rU u + gd <- map densityIR <$> mapM cgOf gs + fd <- forM fams $ \(ms, mx, tx) -> + famDensityIR (vPos Map.! ms) (length ms) <$> rS mx <*> rS tx + let obj = foldl1 (RU2 SAddO) (map fst gd ++ map fst fd) + grds = concatMap snd gd ++ concatMap snd fd + -- RUExp → intern 済み DAG → 命令列 (構造 intern が旧 CSE cache と同じ + -- 重複排除を構造キー比較なしで与える) + rud <- newRUDag + k0 <- internRU rud (RUK 0) + objI <- internRU rud obj + gIs <- forM grds $ \(k, ge) -> (,) k <$> internRU rud ge + st <- readIORef rud + pure CompiledVecIR + { cvProg = compileVecProgramD (map length vecLists) (rudNodes st) + k0 objI gIs + , cvScalIx = VU.fromList [ ixOf Map.! n | n <- scalNames ] + , cvVecIxs = [ VU.fromList [ ixOf Map.! n | n <- ms ] | ms <- vecLists ] + } + +-- --------------------------------------------------------------------------- +-- Phase 56.2: 観測密度の IR 式化 + 静的命令列 (記号 reverse-mode・arena 実行) +-- --------------------------------------------------------------------------- + +-- | [日本語]: 値側 guard の種別 (勾配側は unguarded・従来の前例どおり)。 +-- [English]: The kind of guard applied on the value side (the gradient +-- side stays unguarded, following past precedent). +data GuardKind = GPos | GUnit + +-- | [日本語]: 数値安定な 2 項 log-sum-exp: +-- log(exp a + exp b) = max(a,b) + log(1+exp(-|a-b|))。 Mixture (log_mix)・ +-- ZeroInflatedBinomial の x=0 分岐で使う。 'SMaxO' (勾配は winner-take-all +-- subgradient) を経由するので、 常に有限差分 (|a-b| は必ず ≥0) のみ exp +-- する = オーバーフロー安全。 +-- [English]: A numerically stable pairwise log-sum-exp: +-- log(exp a + exp b) = max(a,b) + log(1+exp(-|a-b|)). Used for the +-- Mixture (log_mix) and ZeroInflatedBinomial x=0 branches. Since it goes +-- through 'SMaxO' (whose gradient is a winner-take-all subgradient), it +-- only ever exponentiates a finite difference (|a-b| is always ≥0), +-- making it overflow-safe. +logSumExp2 :: RUExp -> RUExp -> RUExp +logSumExp2 a b = + RU2 SMaxO a b .+# RU1 SLogO (RUK 1 .+# RU1 SExpO (RU1 SNegO (RU1 SAbsO (a .-# b)))) + +-- | [日本語]: family 別の観測密度を IR 式として組む。 式・guard とも +-- 'logDensityObs' の該当分岐と値一致 (test/probe で担保)。 旧 groupVal / +-- groupNode (手書き tape ノード) の置換 — 勾配は記号微分で自動。 +-- [English]: Assembles the observation density for each family as an IR +-- expression. Both the expression and the guard match the corresponding +-- branch of 'logDensityObs' (guaranteed by tests\/probes). Replaces the +-- old groupVal\/groupNode (hand-written tape nodes) — the gradient is +-- now automatic via symbolic differentiation. +densityIR :: VecObsIR -> (RUExp, [(GuardKind, RUExp)]) +densityIR g = case g of + -- raw potential (Phase 90 A10): 式値そのまま・guard なし。 + VOPot re -> (re, []) + -- Gaussian: -n/2·log2π - (n·logσ + Σr²/(2σ²)) (σ スカラ) / + -- -n/2·log2π - Σlogσ_i - Σ(r_i/σ_i)²/2 (σ 行依存・55.3) + VOGauss mu sge ys -> + let n' = fromIntegral (VS.length ys) :: Double + c0 = RUK (negate (0.5 * n' * log (2 * pi))) + r = RUC ys .-# mu + in if ruIsVec sge + then ( c0 .-# RUSum (RU1 SLogO sge) + .-# (RUSum (let t = r ./# sge in t .*# t) ./# RUK 2) + , [(GPos, sge)] ) + else ( c0 .-# (RUK n' .*# RU1 SLogO sge) + .-# (RUSum (r .*# r) ./# (RUK 2 .*# sge .*# sge)) + , [(GPos, sge)] ) + -- Poisson: Σ(y_i·logλ_i - λ_i) - Σlog y_i! (lfk 前計算・y は raw = kA 一致) + VOPois lam ys lfk -> + let n' = fromIntegral (VS.length ys) :: Double + -- Phase 90 A11-4② (F3b): log(exp η) を η に畳む (log-link 標準形)。 + logLam = ru1Smart SLogO lam + in if ruIsVec lam + then ( RUSum (RUC ys .*# logLam) .-# RUSum lam .-# RUK lfk + , [(GPos, lam)] ) + else ( (RUK (VS.sum ys) .*# logLam) + .-# (RUK n' .*# lam) .-# RUK lfk + , [(GPos, lam)] ) + -- Bernoulli: Σ(yb·log p + (1-yb)·log(1-p)) (yb = round 済 0/1 定数) + VOBern p yb -> + let n' = fromIntegral (VS.length yb) :: Double + c1 = VS.sum yb + omp = RUK 1 .-# p + in if ruIsVec p + then ( RUSum (RUC yb .*# RU1 SLogO p) + .+# RUSum (RUC (VS.map (1 -) yb) .*# RU1 SLogO omp) + , [(GUnit, p)] ) + else ( (RUK c1 .*# RU1 SLogO p) .+# (RUK (n' - c1) .*# RU1 SLogO omp) + , [(GUnit, p)] ) + -- StudentT (ν=SC・56.3): n·[lgamma((ν+1)/2) - lgamma(ν/2) - ½log(νπ)] + -- - Σlogσ - ((ν+1)/2)·Σ log(1 + z_i²/ν)。 ν≤0 は収集時に排除済。 + VOStudT nu mu sge ys -> + let n' = fromIntegral (VS.length ys) :: Double + c0 = RUK (n' * (lgammaApprox ((nu + 1) / 2) - lgammaApprox (nu / 2) + - 0.5 * log (nu * pi))) + z = zOf mu sge ys + in ( c0 .-# sumLogScale (VS.length ys) sge + .-# (RUK ((nu + 1) / 2) + .*# RUSum (RU1 SLogO (RUK 1 .+# (z .*# z ./# RUK nu)))) + , [(GPos, sge)] ) + -- Cauchy (56.3): -n·logπ - Σlogγ - Σ log(1 + z_i²) + VOCauchy loc sce ys -> + let n' = fromIntegral (VS.length ys) :: Double + z = zOf loc sce ys + in ( RUK (negate (n' * log pi)) .-# sumLogScale (VS.length ys) sce + .-# RUSum (RU1 SLogO (RUK 1 .+# z .*# z)) + , [(GPos, sce)] ) + -- Logistic (56.3): -Σz_i - Σlog s - 2·Σ log(1 + exp(-z_i)) + VOLogis mu se ys -> + let z = zOf mu se ys + in ( RU1 SNegO (RUSum z) .-# sumLogScale (VS.length ys) se + .-# (RUK 2 + .*# RUSum (RU1 SLogO (RUK 1 .+# RU1 SExpO (RU1 SNegO z)))) + , [(GPos, se)] ) + -- Gumbel (56.3): -Σlog β - Σz_i - Σ exp(-z_i) + VOGumbel mu bee ys -> + let z = zOf mu bee ys + in ( RU1 SNegO (sumLogScale (VS.length ys) bee) + .-# RUSum z .-# RUSum (RU1 SExpO (RU1 SNegO z)) + , [(GPos, bee)] ) + -- Exponential (56.4): Σ log rate_i - Σ rate_i·y_i (y ≥ 0 は収集時確認済) + VOExpo rate ys -> + ( sumLogScale (VS.length ys) rate + .-# (if ruIsVec rate then RUSum (rate .*# RUC ys) + else rate .*# RUK (VS.sum ys)) + , [(GPos, rate)] ) + -- Weibull (56.4): Σlog k - Σlog λ + Σ(k-1)·z_i - Σ exp(k·z_i) + -- (z_i = log y_i - log λ_i。 walk の (y/λ)**k と ulp 差のみ) + VOWeib k lam lys -> + let n = VS.length lys + z = RUC lys .-# RU1 SLogO lam + in ( sumLogScale n k .-# sumLogScale n lam + .+# RUSum ((k .-# RUK 1) .*# z) + .-# RUSum (RU1 SExpO (k .*# z)) + , [(GPos, k), (GPos, lam)] ) + -- LogNormal (56.4): N(log y⃗ | μ, σ) - Σlog y (Gaussian ノード再利用・guard も流用) + VOLogN mu sge lys c -> + let (e, gds) = densityIR (VOGauss mu sge lys) + in (RUK c .+# e, gds) + -- Gamma (56.4): Σ(α-1)·log y - Σ rate·y + Σ α·log rate - Σ lgammaΓ(α) + VOGamma al rate ys lys -> + let n = VS.length ys + in ( RUSum ((al .-# RUK 1) .*# RUC lys) + .-# (if ruIsVec rate then RUSum (rate .*# RUC ys) + else rate .*# RUK (VS.sum ys)) + .+# sumOf n (al .*# RU1 SLogO rate) + .-# sumOf n (RU1 SLgammaO al) + , [(GPos, al), (GPos, rate)] ) + -- Beta (56.4): Σ(α-1)·log y + Σ(β-1)·log(1-y) - Σ[lgΓα + lgΓβ - lgΓ(α+β)] + VOBeta al be lys l1ys -> + let n = VS.length lys + in ( RUSum ((al .-# RUK 1) .*# RUC lys) + .+# RUSum ((be .-# RUK 1) .*# RUC l1ys) + .-# sumOf n (RU1 SLgammaO al .+# RU1 SLgammaO be + .-# RU1 SLgammaO (al .+# be)) + , [(GPos, al), (GPos, be)] ) + -- Binomial (56.5): ΣlogC + Σk_i·log p_i + Σ(n-k_i)·log(1-p_i) + -- (Bernoulli の yb/1-yb 係数を k/n-k に一般化・logC は compile 時定数) + VOBinom p kv nkv lc -> + let omp = RUK 1 .-# p + in if ruIsVec p + then ( RUK lc .+# RUSum (RUC kv .*# RU1 SLogO p) + .+# RUSum (RUC nkv .*# RU1 SLogO omp) + , [(GUnit, p)] ) + else ( RUK lc .+# (RUK (VS.sum kv) .*# RU1 SLogO p) + .+# (RUK (VS.sum nkv) .*# RU1 SLogO omp) + , [(GUnit, p)] ) + -- Geometric (56.5): Σ(k_i-1)·log(1-p_i) + Σlog p_i (round y ≥ 1 は収集時確認済) + VOGeom p kv -> + let n' = fromIntegral (VS.length kv) :: Double + omp = RUK 1 .-# p + in if ruIsVec p + then ( RUSum ((RUC kv .-# RUK 1) .*# RU1 SLogO omp) + .+# RUSum (RU1 SLogO p) + , [(GUnit, p)] ) + else ( (RUK (VS.sum kv - n') .*# RU1 SLogO omp) + .+# (RUK n' .*# RU1 SLogO p) + , [(GUnit, p)] ) + -- NegativeBinomial (56.5 本命): p = α/(α+μ) として + -- Σ lgΓ(k_i+α) - Σ lgΓ(α) - Σ lgΓ(k_i+1) + Σ α·log p_i + Σ k_i·log(1-p_i) + -- (lgΓ(k_i+α) は SLgammaO の elementwise・lgΓ(k_i+1) は compile 時定数) + VONegBin mu al kv lgk1 -> + let n = VS.length kv + p = al ./# (al .+# mu) + in ( RUSum (RU1 SLgammaO (RUC kv .+# al)) + .-# sumOf n (RU1 SLgammaO al) + .-# RUK lgk1 + .+# sumOf n (al .*# RU1 SLogO p) + .+# RUSum (RUC kv .*# RU1 SLogO (RUK 1 .-# p)) + , [(GPos, mu), (GPos, al)] ) + -- Mixture (Phase 90 A3・2成分 Normal 限定): 'Distribution.hs' の + -- @logDensity (Mixture ws comps) x = logSumExpA [log(w_k/Σw)+logDensity d_k x]@ + -- と数式一致 (w1+w2=1 を仮定せず Σw で正規化)。 各成分の Gaussian 対数密度 + -- (per-row) を 'gaussLpdfElem' で作り、 'logSumExp2' で数値安定に合成。 + VOMixNorm2 w1 w2 m1 s1 m2 s2 ys -> + let total = w1 .+# w2 + logw1 = RU1 SLogO w1 .-# RU1 SLogO total + logw2 = RU1 SLogO w2 .-# RU1 SLogO total + lpdf1 = gaussLpdfElem m1 s1 ys + lpdf2 = gaussLpdfElem m2 s2 ys + perRow = logSumExp2 (logw1 .+# lpdf1) (logw2 .+# lpdf2) + in ( RUSum perRow + , [(GPos, w1), (GPos, w2), (GPos, s1), (GPos, s2)] ) + -- ZeroInflatedBinomial (Phase 90 A3): 'Distribution.hs' の + -- @logDensity (ZeroInflatedBinomial n psi p) x@ と数式一致。 y==0/y>0 の + -- データ分岐は group を分けず mask 列で elementwise に選択する (family + -- gather の disjoint 検査を壊さない安全設計・A1 調査で確定した方針)。 + -- branch0 は「もしこの行が y=0 だったら」の仮想値を **行ごとの n** (nv) で + -- 計算する (nmy=n-y_i は y=0 行以外で n と異なるため使えず、 専用の n 列 nv を + -- 使う。 Phase 94 で n を行対応化 = n 別 group 分裂を解消)。 + VOZIBinom nv psi p mask0 yv nmy logc -> + let omp = RUK 1 .-# p + ompsi = RUK 1 .-# psi + branch0 = logSumExp2 (RU1 SLogO psi) + (RU1 SLogO ompsi .+# (RUC nv .*# RU1 SLogO omp)) + branch1 = RU1 SLogO ompsi .+# RUC logc + .+# (RUC yv .*# RU1 SLogO p) .+# (RUC nmy .*# RU1 SLogO omp) + perRow = (RUC mask0 .*# branch0) .+# ((RUK 1 .-# RUC mask0) .*# branch1) + in ( RUSum perRow, [(GUnit, psi), (GUnit, p)] ) + where + -- 位置-尺度系の共通形 (56.3): z⃗ = (y⃗ - μ)/s (y⃗ は RUC 定数なので常にベクトル形) + zOf mu sge ys = (RUC ys .-# mu) ./# sge + -- Σ e: e がスカラ式なら n·e に畳む (走査回避・56.4 で一般化) + sumOf n e + | ruIsVec e = RUSum e + | otherwise = RUK (fromIntegral n) .*# e + -- Σ log s (スカラ時 n·log s・Gauss の 55.3 と同型) + sumLogScale n sge = sumOf n (RU1 SLogO sge) + -- Phase 90 A3 (Mixture 用): Gaussian 対数密度の **行ごとの値** + -- (-0.5·log2π - logσ - (y-μ)²/(2σ²))。 'VOGauss' の densityIR と違い + -- ここでは Σ を取らずベクトルのまま返す (logSumExp2 で行ごとに混合してから + -- 最後に 1 回だけ Σ する必要があるため)。 + gaussLpdfElem mu sge ys = + let r = RUC ys .-# mu + in RUK (negate (0.5 * log (2 * pi))) .-# RU1 SLogO sge + .-# ((r .*# r) ./# (RUK 2 .*# sge .*# sge)) + +-- | [日本語]: 族 prior 密度の IR 式: -nG/2·log2π - nG·logτ - Σ(a_j-m)²/(2τ²)。 +-- [English]: The IR expression for a family prior density: +-- -nG/2·log2π - nG·logτ - Σ(a_j-m)²/(2τ²). +famDensityIR :: Int -> Int -> RUExp -> RUExp -> (RUExp, [(GuardKind, RUExp)]) +famDensityIR vp nG mx tx = + let nG' = fromIntegral nG :: Double + c0 = RUK (negate (0.5 * nG' * log (2 * pi))) + ra = RUVec vp .-# mx + in ( c0 .-# (RUK nG' .*# RU1 SLogO tx) + .-# (RUSum (ra .*# ra) ./# (RUK 2 .*# tx .*# tx)) + , [(GPos, tx)] ) + +-- | [日本語]: 静的命令列の 1 命令。 slot i = 命令 i の結果 (SSA/ANF・共有保存)。 +-- 全 slot の形 (スカラ / 長さ n) は compile 時に確定し、 1 本の unboxed +-- arena にオフセット解決して敷き詰める (boxed 中間表現なし)。 +-- [English]: A single instruction in the static instruction stream. Slot +-- i holds the result of instruction i (SSA\/ANF, sharing preserved). The +-- shape of every slot (scalar, or length n) is fixed at compile time, +-- and all slots are laid out with resolved offsets into a single +-- unboxed arena (no boxed intermediate representation). +data VInstr + = VIK !Double -- ^ [日本語]: スカラ定数。 [English]: A scalar constant. + | VIKV !(VS.Vector Double) -- ^ [日本語]: ベクトル定数 (データ列)。 [English]: A vector constant (a data column). + | VILeafS !Int -- ^ [日本語]: scalar leaf p の値。 [English]: The value of scalar leaf p. + | VILeafV !Int -- ^ [日本語]: vector leaf p (member 列そのもの)。 [English]: Vector leaf p (the member column itself). + | VIGath !Int !(VU.Vector Int) !Int -- ^ [日本語]: gather (vector leaf p, gids, 行数)。 [English]: A gather (vector leaf p, gids, row count). + | VIUn !SUn !Int + | VIBin !SBin !Int !Int -- ^ [日本語]: broadcast は形 (静的) で解決。 [English]: Broadcasting is resolved via the (static) shape. + | VISum !Int -- ^ [日本語]: Σ (ベクトル → スカラ)。 [English]: A Σ (vector -> scalar). + + -- Phase 85.3-ii: superinstruction (radon 命令列 dump 由来の頻出パターンを + -- compile 時に融合。 pass 数と中間 slot を削減 = 85.3a spike の融合利得) + | VIAxpy !Int !Int !Int -- ^ [日本語]: out = a + s·v (a: slot・len 0 は + -- broadcast、 s: スカラ slot、 v: ベクトル slot)。 + -- [English]: out = a + s·v (a: a slot, + -- length 0 means broadcast; s: a scalar + -- slot; v: a vector slot). + | VIAxpyC !Int !Int !(VS.Vector Double) + -- ^ [日本語]: 同上・v がデータ列定数 + -- (VIKV copy 消滅)。 + -- [English]: Same as above, but v is a + -- data-column constant (eliminates a + -- VIKV copy). + | VISumSqD !Int !Int -- ^ [日本語]: out(スカラ) = Σ (x_j − m_j)² + -- (x/m: slot・スカラ側は broadcast)。 + -- [English]: out (scalar) = Σ (x_j − m_j)² + -- (x\/m: slots; the scalar side broadcasts). + | VISumSqC !(VS.Vector Double) !Int -- ^ [日本語]: 同上・x がデータ列定数。 + -- [English]: Same as above, but x is a + -- data-column constant. + + -- Phase 85.3-iv: RE 連鎖の gather 内蔵化 + 3 項融合 (radon 残 pass の削減) + | VIMulG !Int !Int !(VU.Vector Int) !Int + -- ^ [日本語]: out = s·gather(p) (s: スカラ + -- slot・gather は VIGath と同形で命令内蔵 + -- = gather の実体化 pass 消滅)。 + -- [English]: out = s·gather(p) (s: a + -- scalar slot; the gather has the same + -- shape as VIGath but is folded into the + -- instruction, eliminating the pass that + -- would materialize the gather). + | VIAxpyG !Int !Int !Int !(VU.Vector Int) !Int + -- ^ [日本語]: out = a + s·gather(p)。 + -- [English]: out = a + s·gather(p). + | VIMulVC !Int !Int !(VS.Vector Double) + -- ^ [日本語]: out = s·v⊙c + -- (スカラ×ベクトル×データ列定数)。 + -- [English]: out = s·v⊙c (scalar times + -- vector times data-column constant). + | VISumSqC2 !(VS.Vector Double) !Int !Int + -- ^ [日本語]: out(スカラ) = Σ (c_j − m1_j − m2_j)² + -- (m1/m2: ベクトル slot・和の実体化 + -- pass 消滅)。 + -- [English]: out (scalar) = + -- Σ (c_j − m1_j − m2_j)² (m1\/m2: vector + -- slots; eliminates the pass that would + -- materialize the sum). + | VISumSqDGG !Int !(VU.Vector Int) !Int !(VU.Vector Int) !Int + -- ^ [日本語]: out(スカラ) = + -- Σ (φ[px·gx_j] − φ[pm·gm_j])²。 gather 2 本を + -- SumSqD に内蔵 (ICAR ペア差分・5461 セルの + -- gather 実体化 2 本を消す)。 gather 値は pc + -- から直読み・随伴は param へ直 scatter。 + -- [English]: out (scalar) = + -- Σ (φ[px·gx_j] − φ[pm·gm_j])². Folds two + -- gathers into SumSqD (an ICAR pairwise + -- difference; eliminates the two gather + -- materialization passes at a 5461-cell + -- scale). Gather values are read + -- directly from pc; adjoints are + -- scattered directly into param. + +-- | [日本語]: compile 済みの値+勾配プログラム。 生成は 1 回・per-call は +-- forward (値) / forward+backward (勾配) の実行のみ = per-call の tape +-- 構築を撤去し「tape を compile 時に固定」。 arena は per-call 確保 +-- (共有 mutable なし = @nutsChainsPure@ の spark 並列と整合)。 +-- [English]: The compiled value+gradient program. Generated once; each +-- call only runs forward (for the value) or forward+backward (for the +-- gradient) — this removes per-call tape construction by "fixing the +-- tape at compile time". The arena is allocated per call (no shared +-- mutable state, consistent with 'nutsChainsPure''s spark-based +-- parallelism). +data VecProgram = VecProgram + { vpInstrs :: !(BV.Vector VInstr) + , vpOff :: !(VU.Vector Int) -- ^ [日本語]: slot → arena オフセット。 [English]: slot -> arena offset. + , vpLen :: !(VU.Vector Int) -- ^ [日本語]: slot → 0 (スカラ) / n (ベクトル)。 [English]: slot -> 0 (scalar) \/ n (vector). + , vpSize :: !Int -- ^ [日本語]: arena 総長。 [English]: The total arena length. + , vpObj :: !Int -- ^ [日本語]: 目的 (log-density 和) の slot。 [English]: The slot for the objective (the summed log-density). + , vpGuards :: ![(GuardKind, Int)] -- ^ [日本語]: 値側 guard (slot 参照)。 [English]: Value-side guards (slot references). + } + +-- =========================================================================== +-- Phase 90 A8: 'RUExp' の DAG intern + 共有保存の命令列生成 +-- =========================================================================== + +-- | [日本語]: 'RUExp' の DAG ノード (子は intern 済み ID)。 +-- [English]: A DAG node for 'RUExp' (children are already-interned IDs). +data RUNode + = RNK !Double + | RNC !(VS.Vector Double) + | RNV !Int + | RNG !Int !(VU.Vector Int) + | RNVec !Int + | RN1 !SUn !Int + | RN2 !SBin !Int !Int + | RNSum !Int + deriving (Eq, Ord) + +data RUDagSt = RUDagSt + { rudStable :: !(IM.IntMap [(StableName RUExp, Int)]) + , rudStruct :: !(Map RUNode Int) + , rudNodes :: !(IM.IntMap RUNode) + , rudNext :: !Int + } + +newRUDag :: IO (IORef RUDagSt) +newRUDag = newIORef (RUDagSt IM.empty Map.empty IM.empty 0) + +-- | [日本語]: 'RUExp' を DAG に intern して ID を返す ('internS' の RUExp +-- 版)。 構造 intern が旧 'compileVecProgram' の @Map RUExp Int@ CSE と同じ +-- 重複排除を与える (旧実装はキー比較が構造 walk = 共有木で経路数比例、 +-- こちらは子 ID 比較のみで O(ノード数 · log))。 +-- [English]: Interns an 'RUExp' into the DAG and returns its ID (the +-- 'RUExp' counterpart of 'internS'). Structural interning gives the same +-- deduplication as the old 'compileVecProgram''s @Map RUExp Int@ CSE +-- (the old implementation's key comparison was a structural walk, +-- proportional to path count over a shared tree; this one only compares +-- child IDs, giving O(number of nodes · log)). +internRU :: IORef RUDagSt -> RUExp -> IO Int +internRU ref = go + where + go e0 = do + e <- evaluate e0 + sn <- makeStableName e + let h = hashStableName sn + st <- readIORef ref + case lookup sn =<< IM.lookup h (rudStable st) of + Just i -> pure i + Nothing -> do + nd <- case e of + RUK v -> pure (RNK v) + RUC v -> pure (RNC v) + RUV p -> pure (RNV p) + RUG p gids -> pure (RNG p gids) + RUVec p -> pure (RNVec p) + RU1 o x -> RN1 o <$> go x + RU2 o a b -> RN2 o <$> go a <*> go b + RUSum x -> RNSum <$> go x + st1 <- readIORef ref + i <- case Map.lookup nd (rudStruct st1) of + Just j -> pure j + Nothing -> do + let j = rudNext st1 + writeIORef ref st1 + { rudStruct = Map.insert nd j (rudStruct st1) + , rudNodes = IM.insert j nd (rudNodes st1) + , rudNext = j + 1 } + pure j + modifyIORef' ref $ \s -> + s { rudStable = IM.insertWith (++) h [(sn, i)] (rudStable s) } + pure i + +-- | [日本語]: 'compileVecProgram' の DAG 版。 ノードは intern 済み ID で +-- 参照し、 CSE cache は ID → slot の 'IM.IntMap'。 superinstruction 融合 +-- (85.3-ii/iv) の構造判定は ID 経由の 1 段 lookup。 意味は旧実装と同一 — +-- 「構造等値 ⇔ ID 等値」 が intern で保証されるため、 Σ(x−m)² 融合の +-- @r1 == r2@ も ID 比較で厳密に旧構造比較と一致する。 +-- [English]: The DAG version of 'compileVecProgram'. Nodes are +-- referenced by their interned IDs, and the CSE cache is an +-- 'IM.IntMap' from ID to slot. The structural checks for +-- superinstruction fusion (85.3-ii\/iv) are a single ID-based lookup. +-- The semantics match the old implementation exactly — since interning +-- guarantees "structurally equal ⇔ same ID", the @r1 == r2@ check in the +-- Σ(x−m)² fusion, done via ID comparison, matches the old structural +-- comparison precisely. +compileVecProgramD + :: [Int] -- ^ [日本語]: vector leaf 長。 [English]: Vector-leaf lengths. + -> IM.IntMap RUNode -- ^ [日本語]: ID → ノード (子 ID < 親 ID)。 [English]: ID -> node (child IDs < parent ID). + -> Int -- ^ [日本語]: @RUK 0@ の ID (Σx² 融合で m 側が無い時の代用)。 [English]: The ID of @RUK 0@ (a stand-in for when the Σx² fusion has no m side). + -> Int -- ^ [日本語]: 目的 (log-density 和) root ID。 [English]: The root ID of the objective (the summed log-density). + -> [(GuardKind, Int)] -- ^ [日本語]: guard root ID。 [English]: Guard root IDs. + -> VecProgram +compileVecProgramD vecLens nodes k0 objI guardIs = + let nodeOf i = nodes IM.! i + isVecA = IM.foldlWithKey' + (\mp i nd -> IM.insert i + (case nd of + RNC{} -> True + RNG{} -> True + RNVec{} -> True + RN1 _ a -> mp IM.! a + RN2 _ a b -> (mp IM.! a) || (mp IM.! b) + _ -> False) mp) + IM.empty nodes + isVec i = isVecA IM.! i + emit ins l i (cache, acc, lens, n) = + (n, (IM.insert i n cache, ins : acc, BV.snoc lens l, n + 1 :: Int)) + -- Phase 85.3-ii: a + s·v (AXPY) の融合対象判定。 + mulSV i = case nodeOf i of + RN2 SMulO p q + | not (isVec p), isVec q -> Just (p, q) + | isVec p, not (isVec q) -> Just (q, p) + _ -> Nothing + axpyMatch a b = case mulSV b of + Just (se, ve) -> Just (a, se, ve) + Nothing -> case mulSV a of + Just (se, ve) -> Just (b, se, ve) + Nothing -> Nothing + -- Phase 85.3-iv: スカラ × gather (VIMulG 判定) + mulSG x y = case (nodeOf x, nodeOf y) of + (_, RNG p gids) | not (isVec x) -> Just (x, p, gids) + (RNG p gids, _) | not (isVec y) -> Just (y, p, gids) + _ -> Nothing + -- Phase 85.3-iv: (スカラ×ベクトル) ⊙ データ列定数 (VIMulVC 判定) + mulVC x y = case (nodeOf x, nodeOf y) of + (RNC c, _) -> goVC c y + (_, RNC c) -> goVC c x + _ -> Nothing + where + goVC c mi = case nodeOf mi of + RN2 SMulO p q + | not (isVec p), isVec q -> Just (p, q, c) + | isVec p, not (isVec q) -> Just (q, p, c) + _ -> Nothing + comp i st@(cache, _, _, _) = case IM.lookup i cache of + Just sl -> (sl, st) + Nothing -> case nodeOf i of + RNK v -> emit (VIK v) 0 i st + RNC v -> emit (VIKV v) (VS.length v) i st + RNV p -> emit (VILeafS p) 0 i st + RNVec p -> emit (VILeafV p) (vecLens !! p) i st + RNG p gids -> + emit (VIGath p gids (VU.length gids)) (VU.length gids) i st + RN1 o x -> + let (sx, st1@(_, _, lens1, _)) = comp x st + in emit (VIUn o sx) (lens1 BV.! sx) i st1 + -- Phase 85.3-ii: Σ(x−m)² / Σx² を 1 命令に融合。 + RNSum mI | RN2 SMulO r1 r2 <- nodeOf mI, r1 == r2, isVec r1 -> + let (xe, me) = case nodeOf r1 of + RN2 SSubO x mm -> (x, mm) + _ -> (r1, k0) + in case nodeOf xe of + RNC c | RN2 SAddO m1 m2 <- nodeOf me, isVec m1, isVec m2 -> + let (s1, st1) = comp m1 st + (s2, st2) = comp m2 st1 + in emit (VISumSqC2 c s1 s2) 0 i st2 + RNC c -> + let (sm, st1) = comp me st + in emit (VISumSqC c sm) 0 i st1 + -- Phase 90 A11-4② (F2): Σ(gather − gather)² は gather 2 本を + -- SumSqD 命令に内蔵し 5461 セルの arena 実体化を消す (ICAR)。 + _ | RNG px gx <- nodeOf xe, RNG pm gm <- nodeOf me + , VU.length gx == VU.length gm -> + emit (VISumSqDGG px gx pm gm (VU.length gx)) 0 i st + _ -> + let (sx, st1) = comp xe st + (sm, st2) = comp me st1 + in emit (VISumSqD sx sm) 0 i st2 + -- Phase 85.3-ii: a + s·v → VIAxpy。 85.3-iv: v が gather なら VIAxpyG。 + RN2 SAddO a b | Just (ae, se, ve) <- axpyMatch a b -> + let (sa, st1) = comp ae st + (ss, st2) = comp se st1 + in case nodeOf ve of + RNC c -> emit (VIAxpyC sa ss c) (VS.length c) i st2 + RNG p gids -> + emit (VIAxpyG sa ss p gids (VU.length gids)) + (VU.length gids) i st2 + _ -> + let (sv, st3@(_, _, lens3, _)) = comp ve st2 + in emit (VIAxpy sa ss sv) (lens3 BV.! sv) i st3 + -- Phase 85.3-iv: スカラ×gather → VIMulG。 + RN2 SMulO a b | Just (se, p, gids) <- mulSG a b -> + let (ss, st1) = comp se st + in emit (VIMulG ss p gids (VU.length gids)) (VU.length gids) i st1 + -- Phase 85.3-iv: (スカラ×ベクトル)⊙データ列定数 → VIMulVC。 + RN2 SMulO a b | Just (se, ve, c) <- mulVC a b -> + let (ss, st1) = comp se st + (sv, st2) = comp ve st1 + in emit (VIMulVC ss sv c) (VS.length c) i st2 + RN2 o a b -> + let (sa, st1) = comp a st + (sb, st2@(_, _, lens2, _)) = comp b st1 + in emit (VIBin o sa sb) (max (lens2 BV.! sa) (lens2 BV.! sb)) i st2 + RNSum x -> let (sx, st1) = comp x st in emit (VISum sx) 0 i st1 + st0 = (IM.empty :: IM.IntMap Int, [], BV.empty, 0) + (sObj, st1) = comp objI st0 + (gss, (_, accF, lensF, _)) = + foldl (\(gacc, st) (k, gi) -> + let (sl, st') = comp gi st in (gacc ++ [(k, sl)], st')) + ([], st1) guardIs + lensL = BV.toList lensF + offs = scanl (+) 0 (map (max 1) lensL) -- スカラ slot は 1 セル + in VecProgram + { vpInstrs = BV.fromList (reverse accF) + , vpOff = VU.fromList (init offs) + , vpLen = VU.fromList lensL + , vpSize = last offs + , vpObj = sObj + , vpGuards = gss + } + +-- | [日本語]: 'RUExp' (目的 + guard 式) を命令列へ。 leaf は重複排除・ +-- それ以外は木のまま (密度式は小さいので CSE なしで十分。 随伴は slot +-- 単位で共有されるため記号微分でも式膨張しない)。 +-- [English]: Lowers an 'RUExp' (the objective plus guard expressions) +-- into an instruction stream. Leaves are deduplicated; everything else +-- stays a tree (density expressions are small enough that CSE isn't +-- needed. Adjoints are shared per slot, so symbolic differentiation +-- doesn't blow up the expression either). +compileVecProgram :: [Int] -> RUExp -> [(GuardKind, RUExp)] -> VecProgram +compileVecProgram vecLens obj guards = + let emit ins l e (cache, acc, lens, n) = + (n, (Map.insert e n cache, ins : acc, BV.snoc lens l, n + 1 :: Int)) + -- Phase 85.3-ii: a + s·v (AXPY) の融合対象判定。 加数のどちらかが + -- (スカラ × ベクトル) 積なら Just (残りの加数, スカラ式, ベクトル式)。 + mulSV (RU2 SMulO p q) + | not (ruIsVec p), ruIsVec q = Just (p, q) + | ruIsVec p, not (ruIsVec q) = Just (q, p) + mulSV _ = Nothing + axpyMatch a b = case mulSV b of + Just (se, ve) -> Just (a, se, ve) + Nothing -> case mulSV a of + Just (se, ve) -> Just (b, se, ve) + Nothing -> Nothing + -- Phase 85.3-iv: スカラ × gather (VIMulG 判定) + mulSG x y = case (x, y) of + (s, RUG p gids) | not (ruIsVec s) -> Just (s, p, gids) + (RUG p gids, s) | not (ruIsVec s) -> Just (s, p, gids) + _ -> Nothing + -- Phase 85.3-iv: (スカラ×ベクトル) ⊙ データ列定数 (VIMulVC 判定) + mulVC x y = case (x, y) of + (RUC c, m) -> goVC c m + (m, RUC c) -> goVC c m + _ -> Nothing + where goVC c (RU2 SMulO p q) + | not (ruIsVec p), ruIsVec q = Just (p, q, c) + | ruIsVec p, not (ruIsVec q) = Just (q, p, c) + goVC _ _ = Nothing + comp e st@(cache, _, _, _) = case Map.lookup e cache of + Just sl -> (sl, st) + Nothing -> case e of + RUK v -> emit (VIK v) 0 e st + RUC v -> emit (VIKV v) (VS.length v) e st + RUV p -> emit (VILeafS p) 0 e st + RUVec p -> emit (VILeafV p) (vecLens !! p) e st + RUG p gids -> + emit (VIGath p gids (VU.length gids)) (VU.length gids) e st + RU1 o x -> + let (sx, st1@(_, _, lens1, _)) = comp x st + in emit (VIUn o sx) (lens1 BV.! sx) e st1 + -- Phase 85.3-ii: Σ(x−m)² / Σx² を 1 命令に融合 (residual→二乗→Σ の + -- 3 pass → 1 pass・中間 slot 消滅)。 x がデータ列定数なら VIKV copy + -- ごと消す。 ruIsVec 条件はスカラ RUSum の従来意味 (0) を保存。 + RUSum (RU2 SMulO r1 r2) | r1 == r2, ruIsVec r1 -> + let (xe, me) = case r1 of + RU2 SSubO x m -> (x, m) + x -> (x, RUK 0) + in case xe of + -- 85.3-iv: Σ(c − (m1+m2))² は和の実体化も畳む (radon の + -- 固定効果 μ + RE 項の和がここに来る) + RUC c | RU2 SAddO m1 m2 <- me, ruIsVec m1, ruIsVec m2 -> + let (s1, st1) = comp m1 st + (s2, st2) = comp m2 st1 + in emit (VISumSqC2 c s1 s2) 0 e st2 + RUC c -> + let (sm, st1) = comp me st + in emit (VISumSqC c sm) 0 e st1 + _ -> + let (sx, st1) = comp xe st + (sm, st2) = comp me st1 + in emit (VISumSqD sx sm) 0 e st2 + -- Phase 85.3-ii: a + s·v → VIAxpy (2 pass → 1 pass)。 + -- 85.3-iv: v が gather ならそれも内蔵 (VIAxpyG)。 + RU2 SAddO a b | Just (ae, se, ve) <- axpyMatch a b -> + let (sa, st1) = comp ae st + (ss, st2) = comp se st1 + in case ve of + RUC c -> emit (VIAxpyC sa ss c) (VS.length c) e st2 + RUG p gids -> + emit (VIAxpyG sa ss p gids (VU.length gids)) + (VU.length gids) e st2 + _ -> + let (sv, st3@(_, _, lens3, _)) = comp ve st2 + in emit (VIAxpy sa ss sv) (lens3 BV.! sv) e st3 + -- Phase 85.3-iv: スカラ×gather → VIMulG (gather 実体化の消滅)。 + RU2 SMulO a b | Just (se, p, gids) <- mulSG a b -> + let (ss, st1) = comp se st + in emit (VIMulG ss p gids (VU.length gids)) (VU.length gids) e st1 + -- Phase 85.3-iv: (スカラ×ベクトル)⊙データ列定数 → VIMulVC。 + RU2 SMulO a b | Just (se, ve, c) <- mulVC a b -> + let (ss, st1) = comp se st + (sv, st2) = comp ve st1 + in emit (VIMulVC ss sv c) (VS.length c) e st2 + RU2 o a b -> + let (sa, st1) = comp a st + (sb, st2@(_, _, lens2, _)) = comp b st1 + in emit (VIBin o sa sb) (max (lens2 BV.! sa) (lens2 BV.! sb)) e st2 + RUSum x -> let (sx, st1) = comp x st in emit (VISum sx) 0 e st1 + st0 = (Map.empty :: Map RUExp Int, [], BV.empty, 0) + (sObj, st1) = comp obj st0 + (gss, (_, accF, lensF, _)) = + foldl (\(gacc, st) (k, ge) -> + let (sl, st') = comp ge st in (gacc ++ [(k, sl)], st')) + ([], st1) guards + lensL = BV.toList lensF + offs = scanl (+) 0 (map (max 1) lensL) -- スカラ slot は 1 セル + in VecProgram + { vpInstrs = BV.fromList (reverse accF) + , vpOff = VU.fromList (init offs) + , vpLen = VU.fromList lensL + , vpSize = last offs + , vpObj = sObj + , vpGuards = gss + } + +-- | [日本語]: forward 実行: 全 slot の値を 1 本の arena に書く +-- (ST・per-call 確保)。 +-- [English]: The forward pass: writes every slot's value into a single +-- arena (in ST, allocated per call). +forwardArena + :: CompiledVecIR -> VS.Vector Double -> ST s (VSM.MVector s Double) +forwardArena cvi pc = do + ar <- VSM.unsafeNew (vpSize (cvProg cvi)) + forwardArenaInto cvi pc ar + pure ar + +-- | [日本語]: 'forwardArena' の呼出側バッファ版 (NUTS 葉勾配の per-call +-- arena 確保 (34k セル級) を chain 閉包での 1 回確保 + 再利用に変える)。 +-- 全 slot を毎回上書きするため zero-fill 不要。 +-- [English]: A caller-supplied-buffer version of 'forwardArena' (turns +-- NUTS leaf-gradient per-call arena allocation at the ~34k-cell scale +-- into a single allocation in the chain closure, reused thereafter). +-- Every slot is overwritten on each call, so no zero-fill is needed. +forwardArenaInto + :: CompiledVecIR -> VS.Vector Double -> VSM.MVector s Double -> ST s () +forwardArenaInto cvi pc ar = do + let prog = cvProg cvi + instrs = vpInstrs prog + offV = vpOff prog + lenV = vpLen prog + misB = BV.fromList (cvVecIxs cvi) + scal p = pc `VS.unsafeIndex` (cvScalIx cvi `VU.unsafeIndex` p) + let off i = offV `VU.unsafeIndex` i + len i = lenV `VU.unsafeIndex` i + rd = VSM.unsafeRead ar + wr = VSM.unsafeWrite ar + step i = do + let o = off i + case instrs BV.! i of + VIK v -> wr o v + VIKV v -> + let go !j | j >= VS.length v = pure () + | otherwise = do + wr (o + j) (v `VS.unsafeIndex` j) + go (j + 1) + in go 0 + VILeafS p -> wr o (scal p) + VILeafV p -> + let mis = misB BV.! p + go !j | j >= VU.length mis = pure () + | otherwise = do + wr (o + j) + (pc `VS.unsafeIndex` (mis `VU.unsafeIndex` j)) + go (j + 1) + in go 0 + VIGath p gids n -> + let mis = misB BV.! p + go !r | r >= n = pure () + | otherwise = do + wr (o + r) (pc `VS.unsafeIndex` + (mis `VU.unsafeIndex` (gids `VU.unsafeIndex` r))) + go (r + 1) + in go 0 + -- Phase 105 A3: withSUnF/withSBinF (INLINE CPS) で op の case を + -- ループ外に出し、 known-function の特殊化 unboxed ループに落とす + -- (closure 間接呼出の per-element boxing 排除。 FP 順序不変)。 + VIUn op x -> withSUnF op $ \f -> do + let xo = off x + case len i of + 0 -> rd xo >>= wr o . f + n -> + let go !j | j >= n = pure () + | otherwise = do + v <- rd (xo + j) + wr (o + j) (f v) + go (j + 1) + in go 0 + VIBin op x y -> withSBinF op $ \f -> do + let xo = off x + yo = off y + case (len x, len y) of + (0, 0) -> do + a <- rd xo + b <- rd yo + wr o (f a b) + (0, n) -> do + a <- rd xo + let go !j | j >= n = pure () + | otherwise = do + b <- rd (yo + j) + wr (o + j) (f a b) + go (j + 1) + go 0 + (n, 0) -> do + b <- rd yo + let go !j | j >= n = pure () + | otherwise = do + a <- rd (xo + j) + wr (o + j) (f a b) + go (j + 1) + go 0 + (n, _) -> + let go !j | j >= n = pure () + | otherwise = do + a <- rd (xo + j) + b <- rd (yo + j) + wr (o + j) (f a b) + go (j + 1) + in go 0 + VISum x -> do + let xo = off x + n = len x + go !acc !j | j >= n = wr o acc + | otherwise = do + v <- rd (xo + j) + go (acc + v) (j + 1) + go 0 0 + -- Phase 85.3-ii superinstruction + VIAxpy a s v -> do + sv <- rd (off s) + let vo = off v + n = len i + case len a of + 0 -> do + av <- rd (off a) + let go !j | j >= n = pure () + | otherwise = do + b <- rd (vo + j) + wr (o + j) (av + sv * b) + go (j + 1) + go 0 + _ -> do + let ao = off a + go !j | j >= n = pure () + | otherwise = do + av <- rd (ao + j) + b <- rd (vo + j) + wr (o + j) (av + sv * b) + go (j + 1) + go 0 + VIAxpyC a s c -> do + sv <- rd (off s) + let n = len i + case len a of + 0 -> do + av <- rd (off a) + let go !j | j >= n = pure () + | otherwise = do + wr (o + j) (av + sv * (c `VS.unsafeIndex` j)) + go (j + 1) + go 0 + _ -> do + let ao = off a + go !j | j >= n = pure () + | otherwise = do + av <- rd (ao + j) + wr (o + j) (av + sv * (c `VS.unsafeIndex` j)) + go (j + 1) + go 0 + VISumSqD x m -> do + let xo = off x + mo = off m + bx = len x /= 0 + bm = len m /= 0 + n = max (len x) (len m) + go !acc !j + | j >= n = wr o acc + | otherwise = do + a <- rd (if bx then xo + j else xo) + b <- rd (if bm then mo + j else mo) + let d = a - b + go (acc + d * d) (j + 1) + go 0 0 + VISumSqC c m -> do + let mo = off m + bm = len m /= 0 + n = VS.length c + go !acc !j + | j >= n = wr o acc + | otherwise = do + b <- rd (if bm then mo + j else mo) + let d = c `VS.unsafeIndex` j - b + go (acc + d * d) (j + 1) + go 0 0 + -- Phase 85.3-iv superinstruction + VIMulG s p gids n -> do + sv <- rd (off s) + let mis = misB BV.! p + go !j | j >= n = pure () + | otherwise = do + wr (o + j) (sv * (pc `VS.unsafeIndex` + (mis `VU.unsafeIndex` (gids `VU.unsafeIndex` j)))) + go (j + 1) + go 0 + VIAxpyG a s p gids n -> do + sv <- rd (off s) + let mis = misB BV.! p + gv j = pc `VS.unsafeIndex` + (mis `VU.unsafeIndex` (gids `VU.unsafeIndex` j)) + case len a of + 0 -> do + av <- rd (off a) + let go !j | j >= n = pure () + | otherwise = do + wr (o + j) (av + sv * gv j) + go (j + 1) + go 0 + _ -> do + let ao = off a + go !j | j >= n = pure () + | otherwise = do + av <- rd (ao + j) + wr (o + j) (av + sv * gv j) + go (j + 1) + go 0 + VIMulVC s v c -> do + sv <- rd (off s) + let vo = off v + n = VS.length c + go !j | j >= n = pure () + | otherwise = do + b <- rd (vo + j) + wr (o + j) (sv * b * (c `VS.unsafeIndex` j)) + go (j + 1) + go 0 + VISumSqC2 c m1 m2 -> do + let m1o = off m1 + m2o = off m2 + n = VS.length c + go !acc !j + | j >= n = wr o acc + | otherwise = do + b1 <- rd (m1o + j) + b2 <- rd (m2o + j) + let d = c `VS.unsafeIndex` j - b1 - b2 + go (acc + d * d) (j + 1) + go 0 0 + -- Phase 90 A11-4② (F2): gather 2 本内蔵の Σ(φ_a − φ_b)²。 gather 値は + -- pc から直読み (VIGath forward と同経路)・arena 実体化なし。 + VISumSqDGG px gx pm gm n -> do + let misx = misB BV.! px + mism = misB BV.! pm + go !acc !j + | j >= n = wr o acc + | otherwise = do + let a = pc `VS.unsafeIndex` + (misx `VU.unsafeIndex` (gx `VU.unsafeIndex` j)) + b = pc `VS.unsafeIndex` + (mism `VU.unsafeIndex` (gm `VU.unsafeIndex` j)) + d = a - b + go (acc + d * d) (j + 1) + go 0 0 + loop !i | i >= BV.length instrs = pure () + | otherwise = step i >> loop (i + 1) + loop 0 + +-- | [日本語]: IR の log-density __値__ (観測尤度 + 族 prior)。 guard +-- (σ/τ/λ ≤ 0・p ∉ (0,1) → -∞) は 'logDensityObs' / @logDensity@ の +-- 該当分岐と一致。 +-- [English]: The IR's log-density __value__ (observation likelihood + +-- family prior). Guards (σ\/τ\/λ ≤ 0, or p ∉ (0,1), give -∞) match the +-- corresponding branches of 'logDensityObs' \/ @logDensity@. +vecIRValue :: CompiledVecIR -> VS.Vector Double -> Double +vecIRValue cvi pc = runST $ do + let prog = cvProg cvi + ar <- forwardArena cvi pc + ok <- arenaGuardsOK prog ar + if ok then VSM.unsafeRead ar (vpOff prog `VU.unsafeIndex` vpObj prog) + else pure negInf + +-- | [日本語]: 'gradVecIR' の value-and-grad 融合版。 forward arena を 1 度 +-- だけ構築し、 log-density __値__ (objective slot・'vecIRValue' と同一) と +-- constrained 勾配 (mg へ加算・'gradVecIR' と同一) を同時に返す。 NUTS の葉が +-- leapfrog 最終勾配と同一点でエネルギー (logπ) を別途評価していた重複 +-- (プロファイル実測 19%) を除去するためのエントリポイント。 guard 違反 = +-- Nothing (呼出側が 値 -∞ / 勾配 walk+ad fallback で従来意味論と一致させる)。 +-- [English]: The value-and-gradient fused version of 'gradVecIR'. +-- Builds the forward arena only once, and returns both the log-density +-- __value__ (the objective slot, identical to 'vecIRValue') and the +-- constrained gradient (added into mg, identical to 'gradVecIR') at the +-- same time. This entry point removes the duplicated work where a NUTS +-- leaf separately evaluated the energy (logπ) at the same point as the +-- final leapfrog gradient (measured at 19% of profiled time). A guard +-- violation gives Nothing (the caller matches the previous semantics +-- by treating that as value -∞ \/ falling back to the walk+AD gradient). +gradVecIRVal :: CompiledVecIR -> VS.Vector Double -> VSM.MVector s Double + -> ST s (Maybe Double) +gradVecIRVal cvi pc mg = do + let sz = vpSize (cvProg cvi) + ar <- VSM.unsafeNew sz + adj <- VSM.unsafeNew sz + gradVecIRValWith cvi ar adj pc mg + +-- | [日本語]: 'gradVecIRVal' の呼出側バッファ版。 @ar@ / @adj@ は +-- 長さ 'vpSize' の作業バッファで、 呼出間で再利用してよい (初期化不要・ +-- 毎回全上書き / zero-fill される)。 NUTS の葉勾配 closure が chain ごとに +-- 1 度だけ確保して全 leapfrog で使い回すためのエントリポイント。 +-- [English]: A caller-supplied-buffer version of 'gradVecIRVal'. @ar@ \/ +-- @adj@ are working buffers of length 'vpSize' that may be reused +-- across calls (no initialization needed — everything is overwritten +-- or zero-filled on each call). This entry point lets a NUTS leaf's +-- gradient closure allocate once per chain and reuse the buffers for +-- every leapfrog step. +gradVecIRValWith :: CompiledVecIR + -> VSM.MVector s Double -> VSM.MVector s Double + -> VS.Vector Double -> VSM.MVector s Double + -> ST s (Maybe Double) +gradVecIRValWith cvi ar adj pc mg = do + let prog = cvProg cvi + forwardArenaInto cvi pc ar + ok <- arenaGuardsOK prog ar + if not ok + then pure Nothing + else do + v <- VSM.unsafeRead ar (vpOff prog `VU.unsafeIndex` vpObj prog) + gradVecIRGoWith cvi pc ar adj mg + pure (Just v) + +-- | [日本語]: forward arena 上で値側 guard を検査 (vecIRValue / gradVecIR 共有)。 +-- [English]: Checks the value-side guards over the forward arena +-- (shared by vecIRValue \/ gradVecIR). +arenaGuardsOK :: VecProgram -> VSM.MVector s Double -> ST s Bool +arenaGuardsOK prog ar = fmap and (mapM gOK (vpGuards prog)) + where + gOK (k, sl) = do + let o = vpOff prog `VU.unsafeIndex` sl + n = max 1 (vpLen prog `VU.unsafeIndex` sl) + chk = case k of + GPos -> (> 0) + GUnit -> \pv -> pv > 0 && pv < 1 + go !j | j >= n = pure True + | otherwise = do + v <- VSM.unsafeRead ar (o + j) + if chk v then go (j + 1) else pure False + go 0 + +-- | [日本語]: IR の constrained 勾配を mutable 勾配ベクトルへ__直接__加算 +-- する (記号 reverse-mode・arena backward)。 forward arena と同形の随伴 +-- arena に逆順伝播し、 leaf 随伴は param 位置へその場で scatter。 +-- 命令列・形・オフセットは compile 時に固定済み = per-call の tape +-- 構築なし。 勾配側は unguarded (従来の前例どおり・-∞ 状態は NUTS が +-- 値側で棄却する)。 unconstrained への chain rule は呼出側。 +-- [English]: Adds the IR's constrained gradient __directly__ into a +-- mutable gradient vector (symbolic reverse-mode, arena-based +-- backward pass). Propagates in reverse through an adjoint arena +-- shaped like the forward arena, scattering leaf adjoints straight +-- into their param positions as it goes. The instruction stream, +-- shapes, and offsets are already fixed at compile time, so there is +-- no per-call tape construction. The gradient side is unguarded +-- (following past precedent — -∞ states are rejected by NUTS on the +-- value side). The chain rule to unconstrained space is the caller's +-- responsibility. +gradVecIR :: CompiledVecIR -> VS.Vector Double -> VSM.MVector s Double + -> ST s Bool +gradVecIR cvi pc mg = do + let prog = cvProg cvi + ar <- forwardArena cvi pc + ok <- arenaGuardsOK prog ar + if not ok then pure False + else gradVecIRGo cvi pc ar mg >> pure True + +-- | [日本語]: 'gradVecIR' の backward 本体 (guard 通過後)。 gather 内蔵 +-- 命令 (VIMulG/VIAxpyG) が gather 値を読むため pc (constrained params) を取る。 +-- [English]: The backward-pass body of 'gradVecIR' (after guards pass). +-- Takes pc (the constrained params) because the gather-fused +-- instructions (VIMulG\/VIAxpyG) read gather values from it. +gradVecIRGo + :: CompiledVecIR -> VS.Vector Double -> VSM.MVector s Double + -> VSM.MVector s Double -> ST s () +gradVecIRGo cvi pc ar mg = do + adj <- VSM.unsafeNew (vpSize (cvProg cvi)) + gradVecIRGoWith cvi pc ar adj mg + +-- | [日本語]: 'gradVecIRGo' の呼出側 adj バッファ版。 zero-fill は +-- 本関数が行う (旧 @VSM.replicate (vpSize prog) 0@ と同値) ため、 +-- 呼出側は確保のみで初期化不要。 +-- [English]: A caller-supplied-adj-buffer version of 'gradVecIRGo'. +-- Zero-filling is done by this function itself (equivalent to the old +-- @VSM.replicate (vpSize prog) 0@), so the caller only needs to +-- allocate the buffer, not initialize it. +gradVecIRGoWith + :: CompiledVecIR -> VS.Vector Double -> VSM.MVector s Double + -> VSM.MVector s Double -> VSM.MVector s Double -> ST s () +gradVecIRGoWith cvi pc ar adj mg = do + let prog = cvProg cvi + instrs = vpInstrs prog + offV = vpOff prog + lenV = vpLen prog + misB = BV.fromList (cvVecIxs cvi) + nSlots = BV.length instrs + VSM.set adj 0 + let off i = offV `VU.unsafeIndex` i + len i = lenV `VU.unsafeIndex` i + rdV = VSM.unsafeRead ar + rdA = VSM.unsafeRead adj + addA o d = VSM.unsafeModify adj (+ d) o + addG ix d = VSM.unsafeModify mg (+ d) ix + step i = do + let o = off i + case instrs BV.! i of + VIK _ -> pure () + VIKV _ -> pure () + VILeafS p -> do + d <- rdA o + addG (cvScalIx cvi `VU.unsafeIndex` p) d + VILeafV p -> + let mis = misB BV.! p + go !j | j >= len i = pure () + | otherwise = do + d <- rdA (o + j) + addG (mis `VU.unsafeIndex` j) d + go (j + 1) + in go 0 + VIGath p gids n -> + let mis = misB BV.! p + go !r | r >= n = pure () + | otherwise = do + d <- rdA (o + r) + addG (mis `VU.unsafeIndex` + (gids `VU.unsafeIndex` r)) d + go (r + 1) + in go 0 + -- Phase 105 A3: withSUnD (INLINE CPS) で特殊化 (forward 側と同じ意図)。 + VIUn op x -> withSUnD op $ \df -> do + let xo = off x + case len i of + 0 -> do + a <- rdA o + v <- rdV xo + addA xo (df v * a) + n -> + let go !j | j >= n = pure () + | otherwise = do + a <- rdA (o + j) + v <- rdV (xo + j) + addA (xo + j) (df v * a) + go (j + 1) + in go 0 + VIBin op x y -> do + let xo = off x + yo = off y + n = max 1 (len i) + bx = len x /= 0 -- x がベクトルか + by = len y /= 0 + xi j = if bx then xo + j else xo + yi j = if by then yo + j else yo + case op of + SAddO -> + let go !j | j >= n = pure () + | otherwise = do + a <- rdA (o + j) + addA (xi j) a + addA (yi j) a + go (j + 1) + in go 0 + SSubO -> + let go !j | j >= n = pure () + | otherwise = do + a <- rdA (o + j) + addA (xi j) a + addA (yi j) (negate a) + go (j + 1) + in go 0 + SMulO -> + let go !j | j >= n = pure () + | otherwise = do + a <- rdA (o + j) + vx <- rdV (xi j) + vy <- rdV (yi j) + addA (xi j) (a * vy) + addA (yi j) (a * vx) + go (j + 1) + in go 0 + SDivO -> + let go !j | j >= n = pure () + | otherwise = do + a <- rdA (o + j) + vx <- rdV (xi j) + vy <- rdV (yi j) + addA (xi j) (a / vy) + addA (yi j) (negate (a * vx / (vy * vy))) + go (j + 1) + in go 0 + -- winner-take-all subgradient (tie は測度0・x側に付与で十分)。 + SMaxO -> + let go !j | j >= n = pure () + | otherwise = do + a <- rdA (o + j) + vx <- rdV (xi j) + vy <- rdV (yi j) + if vx >= vy + then addA (xi j) a + else addA (yi j) a + go (j + 1) + in go 0 + VISum x -> do + a <- rdA o + let xo = off x + n = len x + go !j | j >= n = pure () + | otherwise = addA (xo + j) a >> go (j + 1) + go 0 + -- Phase 85.3-ii superinstruction: out = a + s·v の随伴 = + -- adj a += g (スカラ a は Σg)・adj s += Σ g·v・adj v += g·s。 + VIAxpy a s v -> do + sv <- rdV (off s) + let vo = off v + ao = off a + n = len i + case len a of + 0 -> + let go !ga !gs !j + | j >= n = addA ao ga >> addA (off s) gs + | otherwise = do + g <- rdA (o + j) + bv <- rdV (vo + j) + addA (vo + j) (g * sv) + go (ga + g) (gs + g * bv) (j + 1) + in go 0 0 0 + _ -> + let go !gs !j + | j >= n = addA (off s) gs + | otherwise = do + g <- rdA (o + j) + bv <- rdV (vo + j) + addA (ao + j) g + addA (vo + j) (g * sv) + go (gs + g * bv) (j + 1) + in go 0 0 + VIAxpyC a s c -> do + let ao = off a + n = len i + case len a of + 0 -> + let go !ga !gs !j + | j >= n = addA ao ga >> addA (off s) gs + | otherwise = do + g <- rdA (o + j) + go (ga + g) (gs + g * (c `VS.unsafeIndex` j)) (j + 1) + in go 0 0 0 + _ -> + let go !gs !j + | j >= n = addA (off s) gs + | otherwise = do + g <- rdA (o + j) + addA (ao + j) g + go (gs + g * (c `VS.unsafeIndex` j)) (j + 1) + in go 0 0 + -- out = Σ(x−m)² の随伴 = adj x_j += 2(x_j−m_j)·g・adj m_j −= 同 + -- (スカラ側は Σ を単発加算)。 2(x−m)g は旧 (r·r 同一 slot 2 加算 + + -- SSubO 伝播) と IEEE 同値 (x+x ≡ 2x)。 + VISumSqD x m -> do + g <- rdA o + let xo = off x + mo = off m + bx = len x /= 0 + bm = len m /= 0 + n = max (len x) (len m) + go !sx !sm !j + | j >= n = do + if bx then pure () else addA xo sx + if bm then pure () else addA mo sm + | otherwise = do + a <- rdV (if bx then xo + j else xo) + b <- rdV (if bm then mo + j else mo) + let d = 2 * (a - b) * g + if bx then addA (xo + j) d else pure () + if bm then addA (mo + j) (negate d) else pure () + go (if bx then sx else sx + d) + (if bm then sm else sm - d) (j + 1) + go 0 0 0 + VISumSqC c m -> do + g <- rdA o + let mo = off m + bm = len m /= 0 + n = VS.length c + go !sm !j + | j >= n = if bm then pure () else addA mo sm + | otherwise = do + b <- rdV (if bm then mo + j else mo) + let d = 2 * (c `VS.unsafeIndex` j - b) * g + if bm then addA (mo + j) (negate d) >> go sm (j + 1) + else go (sm - d) (j + 1) + go 0 0 + -- Phase 85.3-iv superinstruction: gather 内蔵命令の随伴は leaf + -- (param) へ直接 scatter ('VIGath' backward と同じ) + gather 値は + -- pc から読む。 + VIMulG s p gids n -> do + sv <- rdV (off s) + let mis = misB BV.! p + go !gs !j + | j >= n = addA (off s) gs + | otherwise = do + g <- rdA (o + j) + let ix = mis `VU.unsafeIndex` (gids `VU.unsafeIndex` j) + addG ix (g * sv) + go (gs + g * (pc `VS.unsafeIndex` ix)) (j + 1) + go 0 0 + VIAxpyG a s p gids n -> do + sv <- rdV (off s) + let mis = misB BV.! p + ao = off a + case len a of + 0 -> + let go !ga !gs !j + | j >= n = addA ao ga >> addA (off s) gs + | otherwise = do + g <- rdA (o + j) + let ix = mis `VU.unsafeIndex` (gids `VU.unsafeIndex` j) + addG ix (g * sv) + go (ga + g) (gs + g * (pc `VS.unsafeIndex` ix)) (j + 1) + in go 0 0 0 + _ -> + let go !gs !j + | j >= n = addA (off s) gs + | otherwise = do + g <- rdA (o + j) + let ix = mis `VU.unsafeIndex` (gids `VU.unsafeIndex` j) + addA (ao + j) g + addG ix (g * sv) + go (gs + g * (pc `VS.unsafeIndex` ix)) (j + 1) + in go 0 0 + VIMulVC s v c -> do + sv <- rdV (off s) + let vo = off v + n = VS.length c + go !gs !j + | j >= n = addA (off s) gs + | otherwise = do + g <- rdA (o + j) + b <- rdV (vo + j) + let cj = c `VS.unsafeIndex` j + addA (vo + j) (g * sv * cj) + go (gs + g * b * cj) (j + 1) + go 0 0 + VISumSqC2 c m1 m2 -> do + g <- rdA o + let m1o = off m1 + m2o = off m2 + n = VS.length c + go !j | j >= n = pure () + | otherwise = do + b1 <- rdV (m1o + j) + b2 <- rdV (m2o + j) + let d = 2 * (c `VS.unsafeIndex` j - b1 - b2) * g + addA (m1o + j) (negate d) + addA (m2o + j) (negate d) + go (j + 1) + go 0 + -- Phase 90 A11-4② (F2): 随伴 = ∂/∂φ_a[Σ(φ_a−φ_b)²] = 2(φ_a−φ_b)·g を + -- param へ直 scatter (φ_b は −同)。 gather 値は pc から直読み。 + VISumSqDGG px gx pm gm n -> do + g <- rdA o + let misx = misB BV.! px + mism = misB BV.! pm + go !j | j >= n = pure () + | otherwise = do + let ixa = misx `VU.unsafeIndex` (gx `VU.unsafeIndex` j) + ixb = mism `VU.unsafeIndex` (gm `VU.unsafeIndex` j) + d = 2 * (pc `VS.unsafeIndex` ixa + - pc `VS.unsafeIndex` ixb) * g + addG ixa d + addG ixb (negate d) + go (j + 1) + go 0 + loop !i | i < 0 = pure () + | otherwise = step i >> loop (i - 1) + VSM.unsafeWrite adj (off (vpObj prog)) 1 + loop (nSlots - 1) diff --git a/src/Hanalyze/Model/HBM/Interp.hs b/src/Hanalyze/Model/HBM/Interp.hs new file mode 100644 --- /dev/null +++ b/src/Hanalyze/Model/HBM/Interp.hs @@ -0,0 +1,2058 @@ +{-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE LambdaCase #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE RankNTypes #-} +{-# LANGUAGE FlexibleContexts #-} + +-- | +-- Module : Hanalyze.Model.HBM.Interp +-- Description : HBM dialog DSL の評価系 (interpreter) と NUTS 設定・結果整形 +-- Copyright : (c) 2026 Aelysce Project (Toshiaki Honda) +-- License : BSD-3-Clause +-- +-- [日本語]: HBM dialog DSL の評価系 (interpreter) + NUTS 設定 reader + 結果整形。 +-- +-- (2026-05-31) step 1: canvas-backend @CanvasApp.Analysis.HBM@ +-- から eval/interp/curve コアを移設。 frontend が backend 統一 parser から得た +-- @program_ast@ を streaming sidecar が直接 interpret して実モデルを構築できる +-- よう、 DSL の評価系をライブラリ層 (hanalyze) に置く。 +-- +-- 本 module は __canvas wire 型 (AnalysisRequest 等) にも text parser (DSL frontend) にも依存しない__。 +-- 依存は 'Hanalyze.Model.HBM.Ast' (AST) + +-- @Hanalyze.Model.HBM@ (ModelP/Distribution) + @Hanalyze.Stat.*@ / +-- @Hanalyze.MCMC.*@ + aeson / hmatrix のみ。 +-- +-- text → AST 変換 (@parseHbmText@ 経路) と canvas 専用の @runHbm@ / +-- @buildDataMap@ / @ProgramInfo@ 解決は canvas-backend 側に残す。 +-- +-- [English]: HBM dialog DSL evaluator (interpreter) + NUTS config reader + +-- result formatting. +-- +-- (2026-05-31) step 1: relocated the eval\/interp\/curve core from the +-- canvas-backend @CanvasApp.Analysis.HBM@. So that a streaming sidecar +-- can directly interpret the @program_ast@ obtained by the frontend from the +-- backend's unified parser and build the actual model, the DSL's evaluator +-- lives in the library layer (hanalyze). +-- +-- This module depends on __neither the canvas wire types (AnalysisRequest etc.) nor the text parser (DSL frontend)__. +-- Its only dependencies are +-- 'Hanalyze.Model.HBM.Ast' (AST), @Hanalyze.Model.HBM@ +-- (ModelP\/Distribution), @Hanalyze.Stat.*@ \/ @Hanalyze.MCMC.*@, +-- and aeson \/ hmatrix. +-- +-- The text → AST conversion (the @parseHbmText@ path) and the +-- canvas-specific @runHbm@ \/ @buildDataMap@ \/ @ProgramInfo@ resolution +-- remain on the canvas-backend side. +module Hanalyze.Model.HBM.Interp + ( -- * core types + DataMap + , Column (..) + , colDoubles + , colLength + , colLevels + , lookupDoubles + , EnvA + , Value (..) + , PlateCtx (..) + , topCtx + , TopBind (..) + , ParamSummary (..) + , HbmMeanCurve (..) + -- * evaluation + , hasColRef + , liftD + , builtinTable + , asNum + , asBool + , asList + , asMatrix + , evalScalar + , evalValue + , evalDist + , buildTopEnv + -- * plate / groups (GLMM forEachGroup) + , matchForEachGroup + , lamBodyToStmts + , retToStmts + , groupValsIn + , rowsForGroup + , groupSuffix + , groupSuffixFor + -- * validation / interpretation + , inferTransforms + , validateAst + , preprocessAliases + , validateStmts + , interpStmts + , observeNodeMap + -- * NUTS config readers + , readChainCount + , readNutsConfig + -- * result shaping + , paramSummaryMulti + , fmtSummary + , round4 + , takeEvery + , summaryToJson + , hbmMeanCurveToJson + , extractObserveMeans + , collectCols + , percentileOf + , computeMeanCurves + -- * WAIC / LOO / posterior predictive + , ObsDistSet (..) + , computeObsDists + , pointwiseLogLik + , finitePointwiseLogLik + -- * multi-column observe (observeMV) WAIC / PPC + , MvObsDistSet (..) + , computeMvObsDists + , pointwiseLogLikMv + , reconstructMatrixComb + , MatrixCombSpec (..) + -- * model graph plate aggregation (GLMM forEachGroup) + , GraphPlate (..) + , plateRenameMap + , collectGraphPlates + , collapsePlateGraph + ) where + +import Control.Monad (forM_, when) +import Data.Char (isAlpha, isAlphaNum) +import Data.List (sort, nub, transpose) +import Data.Maybe (fromMaybe) +import qualified Data.Set as Set +import qualified Data.Aeson as A +import qualified Data.Aeson.Key as Key +import qualified Data.Aeson.KeyMap as KM +import qualified Data.Map.Strict as Map +import Data.Text (Text) +import qualified Data.Text as T +import qualified Numeric.LinearAlgebra as LA + +import qualified Hanalyze.MCMC.Core as MC +import qualified Hanalyze.MCMC.NUTS as NUTS +import qualified Hanalyze.Model.HBM as HBM +import qualified Hanalyze.Stat.Distribution as HD +import qualified Hanalyze.Stat.MCMC as SMC + +import Hanalyze.Model.HBM.Ast + ( Expr (..) + , Lit (..) + , Bind (..) + , DoStmt (..) + , collectApp + , Err + ) + +-- =========================================================================== +-- Interpreter コア型 +-- =========================================================================== + +-- | [日本語]: 列参照を含む式かどうか(observe の dist 引数に列が混じったら per-row 展開する)。 +-- [English]: Whether the expression contains a column reference (if a +-- column is mixed into an observe's dist argument, expand it per-row). +hasColRef :: Expr -> Bool +hasColRef = go + where + go (ECol _) = True + go (EApp f x) = go f || go x + go (ENeg e) = go e + go (EOp _ a b) = go a || go b + go (EList xs) = any go xs + go (ELet bs e) = any (go . bindValue) bs || go e + go (EIf c a b) = go c || go a || go b + go (ELam _ b) = go b + go (EDo _ _) = False -- nested do 不可 + go (ELit _) = False + go (EVar _) = False + +-- 'Err' (= Either Text) は Hanalyze.Model.HBM.Ast から import。 + +-- | [日本語]: スカラ式評価。col 参照は許可されるなら row idx を指定。Nothing なら無効でエラー。 +-- [English]: Scalar expression evaluation. If column references are +-- allowed, give the row index; 'Nothing' makes them invalid (an error). +liftD :: Floating a => Double -> a +liftD d = fromRational (toRational d) + +-- | [日本語]: 環境: sample / let / top-level で束縛された値 (数値・真偽・関数)。 +-- 拡張時に scalar から 'Value' に拡張 (ユーザ定義関数を持てる)。 +-- [English]: Environment: values bound by sample \/ let \/ top-level +-- (numbers, booleans, functions). Extended from scalar to 'Value' (so it +-- can hold user-defined functions). +type EnvA a = Map.Map Text (Value a) + +-- | [日本語]: データ列の値。 数値列 (連続/整数) と categorical 列 +-- (factor = level 辞書 + 整数 code) を区別する sum 型。 'Numeric' は従来の +-- @[Double]@ 相当で後方互換、 'Factor' は R の factor / PyMC coords 相当 +-- (level 出現順に 0,1,2,... の code を振る)。 +-- [English]: A data column's value. A sum type distinguishing numeric +-- columns (continuous\/integer) from categorical columns (factor = level +-- dictionary + integer code). 'Numeric' is backward-compatible with the +-- previous @[Double]@ representation; 'Factor' corresponds to R's factor +-- \/ PyMC coords (codes 0,1,2,... assigned in level appearance order). +data Column + = Numeric ![Double] -- ^ [日本語]: 連続 / 整数列 [English]: continuous \/ integer column + | Factor { facLevels :: ![Text], facCodes :: ![Int] } -- ^ categorical + deriving (Eq, Show) + +-- | [日本語]: データ列の Map。 列名 → 'Column'。 +-- [English]: A map of data columns. Column name → 'Column'. +type DataMap = Map.Map Text Column + +-- | [日本語]: 列を @[Double]@ として見る (群比較 / 数値 observe / mean-curve 用)。 +-- 'Numeric' はそのまま、 'Factor' は code を Double 化 (0,1,2,...)。 既存の +-- 数値ロジックはこの accessor 経由で Factor も透過に扱える。 +-- [English]: Views a column as @[Double]@ (used for group comparisons \/ +-- numeric observe \/ mean-curve computation). 'Numeric' passes through as +-- is; for 'Factor', the codes are converted to Double (0,1,2,...). Existing +-- numeric logic can transparently handle 'Factor' too via this accessor. +colDoubles :: Column -> [Double] +colDoubles (Numeric xs) = xs +colDoubles (Factor _ cs) = map fromIntegral cs + +-- | [日本語]: 列長 (行数)。 +-- [English]: Column length (row count). +colLength :: Column -> Int +colLength (Numeric xs) = length xs +colLength (Factor _ cs) = length cs + +-- | [日本語]: 'Factor' なら level 辞書 (出現順)、 'Numeric' なら Nothing。 +-- [English]: For 'Factor', the level dictionary (in appearance order); for +-- 'Numeric', 'Nothing'. +colLevels :: Column -> Maybe [Text] +colLevels (Factor ls _) = Just ls +colLevels (Numeric _) = Nothing + +-- | [日本語]: 列を @[Double]@ として引く (無ければ空)。 旧 @Map.findWithDefault [] k dm@ の +-- 'Column' 対応版。 +-- [English]: Looks up a column as @[Double]@ (empty if absent). The +-- 'Column'-aware counterpart of the old @Map.findWithDefault [] k dm@. +lookupDoubles :: Text -> DataMap -> [Double] +lookupDoubles k = maybe [] colDoubles . Map.lookup k + +-- | [日本語]: モデル本体評価中の値。 Double 閉の数値 + 真偽 + +-- 一階クロージャ (ユーザ定義関数 / lambda) + 組込関数マーカ + 遅延エラー +-- (top-level 値束縛の評価失敗を lookup まで遅延運搬する)。 +-- [English]: A value during model-body evaluation. Numbers closed over +-- Double, booleans, first-order closures (user-defined functions \/ +-- lambdas), a builtin-function marker, and a deferred error (carries a +-- top-level value binding's evaluation failure forward until it is looked +-- up). +data Value a + = VNum a + | VBool Bool + | VList [Value a] -- Phase 42: list リテラル ([e₁, …, eₙ])。 多値分布 + -- (Categorical [probs] / OrderedLogistic eta [cuts]) + -- の list 引数評価の土台。 要素はスカラ閉前提。 + | VClosure (EnvA a) [Text] Expr -- 捕捉環境, 残り仮引数, 本体 (一階・カリー化) + | VBuiltin Text Int -- 組込関数 (名前, arity)。 適用時に builtinTable で解決 + | VErr Text -- 遅延エラー (top-level 値 thunk の評価失敗) + +-- | [日本語]: 組込数学関数ホワイトリスト (name → (arity, impl))。 +-- すべて Double 上で閉じる純関数 (IO/import なし)。 GLM リンク +-- (invLogit/logistic) もここに含む。 必要に応じ追加。 +-- [English]: Whitelist of builtin math functions (name → (arity, impl)). +-- All are pure functions closed over Double (no IO\/import). GLM link +-- functions (invLogit\/logistic) are included here too. Add more as +-- needed. +builtinTable :: forall a. (Floating a, Ord a) => Map.Map Text (Int, [a] -> a) +builtinTable = Map.fromList + [ ("exp", (1, \xs -> exp (head xs))) + , ("log", (1, \xs -> log (head xs))) + , ("log1p", (1, \xs -> log (1 + head xs))) + , ("sqrt", (1, \xs -> sqrt (head xs))) + , ("abs", (1, \xs -> abs (head xs))) + , ("signum", (1, \xs -> signum (head xs))) + , ("recip", (1, \xs -> recip (head xs))) + , ("negate", (1, \xs -> negate (head xs))) + , ("tanh", (1, \xs -> tanh (head xs))) + , ("sin", (1, \xs -> sin (head xs))) + , ("cos", (1, \xs -> cos (head xs))) + , ("logistic", (1, \xs -> 1 / (1 + exp (negate (head xs))))) + , ("invLogit", (1, \xs -> 1 / (1 + exp (negate (head xs))))) + , ("min", (2, \xs -> min (xs !! 0) (xs !! 1))) + , ("max", (2, \xs -> max (xs !! 0) (xs !! 1))) + ] + +-- | [日本語]: list builtin (VList→VList) の名前集合。 スカラ 'builtinTable' +-- (= @[a] -> a@) には乗らないので別管理。 softmax 多項ロジット +-- (@Categorical (softmax [η₀, η₁, …])@) で多クラス線形予測子を確率に変換する。 +-- [English]: The set of list-builtin names (VList→VList). Managed +-- separately since they don't fit the scalar 'builtinTable' (= +-- @[a] -> a@). Used by the softmax multinomial logit +-- (@Categorical (softmax [η₀, η₁, …])@) to convert multi-class linear +-- predictors into probabilities. +listBuiltins :: Set.Set Text +listBuiltins = Set.fromList ["softmax"] + +-- | [日本語]: 安定 softmax (= @exp(xₖ − max x) / Σ@)。 識別性のため基準クラスは η=0 を +-- 明示的に並べる前提 (例 @softmax [0, b₁·x, b₂·x]@)。 空リストはエラー。 +-- [English]: Numerically stable softmax (= @exp(xₖ − max x) / Σ@). For +-- identifiability, assumes the reference class is explicitly listed with +-- η=0 (e.g. @softmax [0, b₁·x, b₂·x]@). An empty list is an error. +softmaxList :: forall a. (Floating a, Ord a) => [a] -> Err [a] +softmaxList [] = Left "softmax: 空リストには適用できません (クラス数 ≥ 1 の [..] が必要です)" +softmaxList xs = + let m = maximum xs + es = map (\x -> exp (x - m)) xs + s = sum es + in Right (map (/ s) es) + +-- | [日本語]: Value を数値に落とす (= 親切な日本語エラー)。 +-- [English]: Coerces a Value down to a number (with a friendly Japanese +-- error message). +asNum :: Value a -> Err a +asNum (VNum x) = Right x +asNum (VBool _) = Left "真偽値が数値の位置に現れました (比較式を算術に混ぜていませんか)" +asNum (VClosure{}) = Left "関数が数値の位置に現れました (引数不足か、 適用し忘れ)" +asNum (VBuiltin n _) = Left ("組込関数 " <> n <> " が数値の位置に現れました (引数を渡してください)") +asNum (VList _) = Left "リストが数値の位置に現れました (list 引数はスカラとして使えません)" +asNum (VErr msg) = Left msg + +-- | [日本語]: Value を真偽に落とす。 +-- [English]: Coerces a Value down to a boolean. +asBool :: Value a -> Err Bool +asBool (VBool b) = Right b +asBool (VNum _) = Left "数値が真偽の位置に現れました (if の条件は比較式である必要があります)" +asBool (VClosure{}) = Left "関数が真偽の位置に現れました" +asBool (VBuiltin n _) = Left ("組込関数 " <> n <> " が真偽の位置に現れました") +asBool (VList _) = Left "リストが真偽の位置に現れました" +asBool (VErr msg) = Left msg + +-- | [日本語]: Value を list に落とす (多値分布の list 引数評価)。 各要素は +-- 'asNum' でスカラに落とせる前提。 +-- [English]: Coerces a Value down to a list (for evaluating multi-valued +-- distributions' list arguments). Assumes each element can be coerced to +-- a scalar via 'asNum'. +asList :: Value a -> Err [Value a] +asList (VList xs) = Right xs +asList (VNum _) = Left "数値がリストの位置に現れました (list 引数には [..] を渡してください)" +asList (VBool _) = Left "真偽値がリストの位置に現れました" +asList (VClosure{}) = Left "関数がリストの位置に現れました" +asList (VBuiltin n _) = Left ("組込関数 " <> n <> " がリストの位置に現れました") +asList (VErr msg) = Left msg + +-- | [日本語]: Value を行列 ([[a]]) に落とす (MvNormal の cov / lkjCorrCholesky の +-- 行列値を評価する用途)。 VList-of-VList を期待し、 各内側 VList の長さ一致と +-- 数値性を検査する。 list 操作で書くのは DSL スカラ値 'Value' 上の小さい構造 +-- 変換のためで、 hmatrix Matrix 経路ではない (density 計算側の choleskyL / +-- forwardSub は既存実装を流用)。 +-- [English]: Coerces a Value down to a matrix ([[a]]) (for evaluating +-- MvNormal's cov \/ lkjCorrCholesky's matrix values). Expects a +-- VList-of-VList and checks each inner VList's length matches and is +-- numeric. Written as list operations because it's a small structural +-- conversion over the DSL's scalar 'Value', not the hmatrix Matrix path +-- (the density computation's choleskyL \/ forwardSub reuse the existing +-- implementation). +asMatrix :: Value a -> Err [[a]] +asMatrix v = do + rows <- asList v + mat <- mapM (\r -> asList r >>= mapM asNum) rows + case mat of + [] -> Left "行列が空です (cov / 行列引数には [[..],..] が必要です)" + (r0 : rs) + | all ((== length r0) . length) rs -> Right mat + | otherwise -> Left "行列の各行の長さが揃っていません (cov は正方行列である必要があります)" + +-- | [日本語]: スカラ式評価 (= 数値を返す)。 'evalValue' の薄いラッパ。 col 参照は +-- mi=Just i なら行 i、 Nothing なら不可。 +-- [English]: Scalar expression evaluation (returns a number). A thin +-- wrapper around 'evalValue'. For column references, mi=Just i means row +-- i, and 'Nothing' means they are disallowed. +evalScalar + :: forall a. (Floating a, Ord a) + => EnvA a -> DataMap -> Maybe Int -> Expr -> Err a +evalScalar env dataMap mi e = evalValue env dataMap mi e >>= asNum + +-- | [日本語]: 式を 'Value' に評価する interpreter 中核。 +-- 適用 (EApp) / if (EIf) / 比較・論理 (EOp) / lambda (ELam) / let を解釈。 +-- ユーザ定義関数 (一階・カリー化) と組込数学関数を呼べる。 +-- 再帰なし (total) ・ IO/ADT/型クラスなし。 +-- [English]: The interpreter's core, evaluating an expression to a +-- 'Value'. Interprets application (EApp) \/ if (EIf) \/ +-- comparison-and-logic (EOp) \/ lambda (ELam) \/ let. Can call +-- user-defined functions (first-order, curried) and builtin math +-- functions. No recursion (total); no IO\/ADT\/type classes. +evalValue + :: forall a. (Floating a, Ord a) + => EnvA a -> DataMap -> Maybe Int -> Expr -> Err (Value a) +evalValue env dataMap mi = go + where + go :: Expr -> Err (Value a) + go (ELit (LNumber d)) = Right (VNum (liftD d)) + go (ELit (LBool b)) = Right (VBool b) + go (ELit (LText t)) = Left ("文字列リテラル \"" <> t <> "\" は数式の中では使えません") + go (EVar n) = case Map.lookup n env of + Just v -> Right v + Nothing -> case Map.lookup n (builtinTable :: Map.Map Text (Int, [a] -> a)) of + Just (ar, _) -> Right (VBuiltin n ar) + Nothing + -- Phase 43: list builtin (softmax 等、 VList→VList) はスカラ + -- builtinTable に乗らないので別途 VBuiltin として解決する。 + | n `Set.member` listBuiltins -> Right (VBuiltin n 1) + | otherwise -> Left ("未定義の変数です: " <> n) + go (ECol c) = case mi of + Nothing -> Left ("列参照 '" <> c <> "' は行ごとの評価文脈 (observe の per-row) でのみ使えます") + Just i -> case Map.lookup c dataMap of + Just col -> case drop i (colDoubles col) of + (x : _) -> Right (VNum (liftD x)) + [] -> Left ("列のインデックスが範囲外です: " <> c <> "[" <> T.pack (show i) <> "]") + Nothing -> Left ("未知の列です: " <> c) + go (ENeg e) = do v <- go e; x <- asNum v; Right (VNum (negate x)) + go (EOp op a b) = evalOp op a b + go (EIf c a b) = do + cv <- go c + cond <- asBool cv + if cond then go a else go b + go (ELet bs body) = do + env' <- foldEnv env bs + evalValue env' dataMap mi body + go (ELam x body) = Right (VClosure env [x] body) + go appE@(EApp _ _) = + let (h, args) = spine appE [] + in do + hv <- go h + argVs <- mapM go args + applyValue hv argVs + -- Phase 42: list リテラルを VList に評価。 多値分布 (Categorical / + -- OrderedLogistic) の list 引数で使う。 要素は順次スカラ評価される。 + go (EList xs) = VList <$> mapM go xs + go (EDo _ _) = Left "入れ子の do ブロックは未対応です" + + -- 適用の脊柱を平坦化: f x y → (f, [x, y])。 + spine :: Expr -> [Expr] -> (Expr, [Expr]) + spine (EApp f x) acc = spine f (x : acc) + spine e acc = (e, acc) + + -- 値を引数列に適用 (一階・カリー化)。 + applyValue :: Value a -> [Value a] -> Err (Value a) + applyValue v [] = Right v + -- Phase 43: list builtin (softmax)。 VList を取り VList を返す (多項ロジット: + -- Categorical (softmax [η₀, η₁, …]))。 スカラ builtinTable とは別経路。 + applyValue (VBuiltin "softmax" _) args = case args of + [VList xs] -> do ns <- mapM asNum xs; VList . map VNum <$> softmaxList ns + [_] -> Left "softmax はリスト引数 ([..]) が必要です" + _ -> Left ("softmax はリスト引数 1 個が必要ですが " <> T.pack (show (length args)) <> " 個でした") + applyValue (VBuiltin name ar) args = + case Map.lookup name (builtinTable :: Map.Map Text (Int, [a] -> a)) of + Nothing -> Left ("内部エラー: 未知の組込関数 " <> name) + Just (_, impl) + | length args == ar -> do ns <- mapM asNum args; Right (VNum (impl ns)) + | length args < ar -> Left ("組込関数 " <> name <> " は引数 " + <> T.pack (show ar) <> " 個が必要ですが " <> T.pack (show (length args)) + <> " 個でした (部分適用は未対応)") + | otherwise -> Left ("組込関数 " <> name <> " に引数が多すぎます (必要 " + <> T.pack (show ar) <> " 個)") + applyValue (VClosure cenv params body) args = applyClosure cenv params body args + applyValue (VNum _) _ = Left "数値を関数として適用しています (関数ではない値に引数を渡しています)" + applyValue (VBool _) _ = Left "真偽値を関数として適用しています" + applyValue (VList _) _ = Left "リストを関数として適用しています" + applyValue (VErr m) _ = Left m + + -- クロージャ適用: 引数を仮引数に順次束縛。 仮引数が尽きたら (= 完全 + -- 適用) 本体を評価して残り引数をさらに適用、 引数が尽きたが仮引数が + -- 残るなら部分適用 (VClosure を返す)。 節順が重要: 完全適用 + -- (ps=[], args=[]) は本体評価を先に判定する。 + applyClosure :: EnvA a -> [Text] -> Expr -> [Value a] -> Err (Value a) + applyClosure cenv [] body args = do + r <- evalValue cenv dataMap mi body + applyValue r args + applyClosure cenv ps body [] = Right (VClosure cenv ps body) + applyClosure cenv (p : ps) body (a : as) = + applyClosure (Map.insert p a cenv) ps body as + + evalOp :: Text -> Expr -> Expr -> Err (Value a) + evalOp op a b = case op of + "+" -> num2 (+) + "-" -> num2 (-) + "*" -> num2 (*) + "/" -> num2 (/) + "**" -> num2 (**) + "^" -> num2 (**) -- DSL では Double 冪 (整数冪に限定しない) + "==" -> cmp (==) + "/=" -> cmp (/=) + "<" -> cmp (<) + "<=" -> cmp (<=) + ">" -> cmp (>) + ">=" -> cmp (>=) + "&&" -> bool2 (&&) + "||" -> bool2 (||) + _ -> Left ("モデル中で未対応の演算子です: " <> op) + where + num2 f = do av <- asNum =<< go a; bv <- asNum =<< go b; Right (VNum (f av bv)) + cmp f = do av <- asNum =<< go a; bv <- asNum =<< go b; Right (VBool (f av bv)) + bool2 f = do av <- asBool =<< go a; bv <- asBool =<< go b; Right (VBool (f av bv)) + + foldEnv :: EnvA a -> [Bind] -> Err (EnvA a) + foldEnv e [] = Right e + foldEnv e (Bind n v : rest) = do + vv <- evalValue e dataMap mi v + foldEnv (Map.insert n vv e) rest + +-- =========================================================================== +-- Phase 27 §F-3a: top-level 束縛環境 +-- =========================================================================== + +-- | [日本語]: 1 つの top-level 値/関数束縛 (= ユーザが model と並べて書く +-- `tmpvar = 1` / `linkfunc x = log x`)。 model (= do-block 束縛) は含まない。 +-- [English]: A single top-level value \/ function binding (i.e. what a +-- user writes alongside the model, like `tmpvar = 1` \/ +-- `linkfunc x = log x`). Does not include the model (do-block bindings). +data TopBind = TopBind + { tbName :: Text + , tbParams :: [Text] + , tbBody :: Expr + } deriving (Show) + +-- | [日本語]: top-level 束縛から評価環境を組む。 値束縛 (引数なし) は env の中で +-- 遅延評価して 'Value' に (相互参照は laziness で解決、 再帰は無い前提)。 +-- 関数束縛 (引数あり) は env を捕捉した 'VClosure' に。 評価失敗は 'VErr' +-- として運び、 実際に参照された時にエラーを出す。 +-- [English]: Builds the evaluation environment from top-level bindings. +-- Value bindings (no arguments) are lazily evaluated to a 'Value' within +-- the env (mutual references are resolved via laziness; no recursion is +-- assumed). Function bindings (with arguments) become a 'VClosure' +-- capturing the env. Evaluation failures are carried as 'VErr' and raised +-- only when actually referenced. +buildTopEnv :: forall a. (Floating a, Ord a) => DataMap -> [TopBind] -> EnvA a +buildTopEnv dataMap binds = env + where + env :: EnvA a + env = Map.fromList [ (tbName b, toVal b) | b <- binds ] + toVal b + | null (tbParams b) = case evalValue env dataMap Nothing (tbBody b) of + Right v -> v + Left e -> VErr e + | otherwise = VClosure env (tbParams b) (tbBody b) + +-- =========================================================================== +-- Phase 27 §F-3c: GLMM plate (forEachGroup) +-- =========================================================================== +-- +-- `forEachGroup "gcol" $ \g -> do { … }` は群列 gcol の distinct 値ごとに +-- 内部 do-block を展開する専用構文。 native ModelP の +-- `forM_ groups $ \j -> do { theta <- sample ("theta_"++show j) …; observe … }` +-- (= Phase37 demo randomSlope/multiLevel) を AST 経由で再現する。 +-- * sample / observe 名は群値で suffix を付け、 群ごとに別 latent にする。 +-- * observe は当該群の行のみを対象にする (= 行サブセット)。 +-- * lambda 引数 g は群値 (Double) に束縛 (式で使える)。 +-- * nest 可能 (forEachGroup の中に forEachGroup): suffix 連結 + 行積集合。 + +-- | [日本語]: plate 評価コンテキスト。 top-level は 'topCtx' = ("", 全行) で従来挙動を保つ。 +-- [English]: The plate evaluation context. At the top level, 'topCtx' = +-- ("", all rows) preserves the previous behavior. +data PlateCtx = PlateCtx + { pcSuffix :: Text -- sample/observe 名に付ける suffix (例 "_1_2") + , pcRows :: Maybe [Int] -- observe 対象行 (Nothing = 全行) + } + +topCtx :: PlateCtx +topCtx = PlateCtx "" Nothing + +-- | [日本語]: `forEachGroup "gcol" (\g -> do { … })` を検出し (群列, 引数名, 内部 stmts)。 +-- [English]: Detects `forEachGroup "gcol" (\g -> do { … })` and returns +-- (group column, parameter name, inner stmts). +matchForEachGroup :: Expr -> Maybe (Text, Text, [DoStmt]) +matchForEachGroup e = case collectApp e of + Right ("forEachGroup", [ELit (LText gcol), ELam param body]) -> + Just (gcol, param, lamBodyToStmts body) + _ -> Nothing + +-- | [日本語]: lambda 本体を DoStmt 列に。 do なら stmts + 末尾式、 それ以外は単一 DoExpr。 +-- [English]: Converts a lambda body into a DoStmt list. For a do-block, +-- stmts plus the trailing expression; otherwise, a single DoExpr. +lamBodyToStmts :: Expr -> [DoStmt] +lamBodyToStmts (EDo stmts ret) = stmts ++ retToStmts ret +lamBodyToStmts other = [DoExpr other] + +-- | [日本語]: do-block 末尾式: pure/return は捨て、 それ以外 (observe 等) は DoExpr に。 +-- [English]: The do-block's trailing expression: pure\/return is dropped; +-- anything else (e.g. observe) becomes a DoExpr. +retToStmts :: Expr -> [DoStmt] +retToStmts r = case r of + EApp (EVar "pure") _ -> [] + EApp (EVar "return") _ -> [] + ELit (LBool _) -> [] + _ -> [DoExpr r] + +-- | [日本語]: ctx の対象行に限定した群列の distinct 値 (昇順)。 +-- [English]: The group column's distinct values, restricted to ctx's +-- target rows (ascending order). +groupValsIn :: DataMap -> Text -> Maybe [Int] -> [Double] +groupValsIn dm gcol mrows = + let col = lookupDoubles gcol dm + idxs = fromMaybe [0 .. length col - 1] mrows + in sort (nub [ col !! i | i <- idxs, i >= 0, i < length col ]) + +-- | [日本語]: ctx の対象行のうち 群列 == gval の行 index。 +-- [English]: The row indices, among ctx's target rows, where the group +-- column equals gval. +rowsForGroup :: DataMap -> Text -> Double -> Maybe [Int] -> [Int] +rowsForGroup dm gcol gval mrows = + let col = lookupDoubles gcol dm + idxs = fromMaybe [0 .. length col - 1] mrows + in [ i | i <- idxs, i >= 0, i < length col, col !! i == gval ] + +-- | [日本語]: 群値を name suffix に (整数なら "_3"、 非整数なら小数表記)。 +-- 群列は整数コード前提なので通常は整数 suffix。 +-- [English]: Turns a group value into a name suffix (an integer value +-- like "_3"; otherwise, decimal notation). Group columns are assumed to +-- be integer codes, so the suffix is normally an integer. +groupSuffix :: Double -> Text +groupSuffix g + | g == fromIntegral (round g :: Integer) = "_" <> T.pack (show (round g :: Integer)) + | otherwise = "_" <> T.pack (show g) + +-- | [日本語]: 群 suffix。 群列が 'Factor' で code が level を指し、 その +-- level が安全な識別子 (先頭英字/下線 + 英数字/下線のみ) なら可読 suffix +-- "_" (例 "_setosa")、 それ以外は 'groupSuffix' (数値 code suffix) に +-- フォールバック。 charset / 衝突安全のため不安全な level は code に落とす。 +-- interpStmts / collectObsInstances / plateRenameMap の 3 経路で同一規律を使う +-- 必要がある (node 名が一致しないと観測値/グラフが噛み合わない)。 +-- [English]: The group suffix. If the group column is 'Factor' and the +-- code points to a level, and that level is a safe identifier (leading +-- alpha\/underscore, then alphanumerics\/underscore only), the readable +-- suffix "_" (e.g. "_setosa") is used; otherwise it falls back to +-- 'groupSuffix' (a numeric-code suffix). For charset\/collision safety, +-- unsafe levels fall back to the code. The same rule must be shared across +-- the three code paths interpStmts \/ collectObsInstances \/ +-- plateRenameMap (if the node names don't match, observed values and the +-- graph fall out of sync). +groupSuffixFor :: Maybe Column -> Double -> Text +groupSuffixFor (Just (Factor levels _)) g + | i >= 0, i < length levels, isSafeIdent (levels !! i) = "_" <> (levels !! i) + where i = round g :: Int +groupSuffixFor _ g = groupSuffix g + +-- | [日本語]: node 名 suffix に使える安全な識別子か (先頭英字/下線、 以降英数字/下線)。 +-- [English]: Whether this is a safe identifier usable as a node-name +-- suffix (leading alpha\/underscore, then alphanumerics\/underscore). +isSafeIdent :: Text -> Bool +isSafeIdent t = case T.uncons t of + Nothing -> False + Just (c0, rest) -> (isAlpha c0 || c0 == '_') + && T.all (\c -> isAlphaNum c || c == '_') rest + +-- | [日本語]: Distribution AST を Hanalyze.Model.HBM.Distribution に変換。 +-- mi が Nothing なら列参照不可。 mi=Just i なら行 i で評価。 +-- [English]: Converts the Distribution AST to +-- Hanalyze.Model.HBM.Distribution. If mi is Nothing, column +-- references are disallowed; if mi=Just i, evaluation happens at row i. +evalDist + :: forall a. (Floating a, Ord a) + => EnvA a -> DataMap -> Maybe Int -> Expr -> Err (HBM.Distribution a) +evalDist env dataMap mi expr = do + (name, args) <- collectApp expr + case (name, args) of + ("Normal", [m, s]) -> mk2 HBM.Normal m s + ("HalfNormal", [s]) -> mk1 HBM.HalfNormal s + ("Beta", [a, b]) -> mk2 HBM.Beta a b + ("Gamma", [s, r]) -> mk2 HBM.Gamma s r + ("Exponential", [r]) -> mk1 HBM.Exponential r + ("Poisson", [l]) -> mk1 HBM.Poisson l + ("Bernoulli", [p]) -> mk1 HBM.Bernoulli p + ("Uniform", [l, h]) -> mk2 HBM.Uniform l h + ("StudentT", [df, m, s]) -> mk3 HBM.StudentT df m s + ("Cauchy", [l, s]) -> mk2 HBM.Cauchy l s + ("HalfCauchy", [s]) -> mk1 HBM.HalfCauchy s + ("LogNormal", [m, s]) -> mk2 HBM.LogNormal m s + -- Phase 42: 多値 categorical 応答。 list 引数は VList 経由でスカラ列に + -- 評価する (observe は factor code 0..K-1)。 2 値応答は Bernoulli + + -- factor code 0/1 で Phase 41.5 対応済。 + ("Categorical", [probs]) -> HBM.Categorical <$> evalList probs + ("OrderedLogistic", [eta, cuts]) -> HBM.OrderedLogistic <$> eval eta <*> evalList cuts + -- Phase 44: 多変量正規 (観測専用)。 mu = 平均ベクトル ([a])、 cov = full Σ + -- (VList-of-VList → [[a]])。 observeMV 経由で k-vector を観測する。 cov の + -- 正定値性は density 評価時 (choleskyL→-∞) 任せ、 ここでは正方性のみ検査。 + ("MvNormal", [mu, cov]) -> do + muV <- evalList mu + covV <- evalMat cov + let k = length muV + when (length covV /= k) + (Left ("MvNormal: 平均ベクトル長 " <> T.pack (show k) + <> " と共分散の行数 " <> T.pack (show (length covV)) <> " が一致しません")) + when (any ((/= k) . length) covV) + (Left ("MvNormal: 共分散は " <> T.pack (show k) <> "×" <> T.pack (show k) + <> " 正方行列である必要があります")) + Right (HBM.MvNormal muV covV) + -- Phase 44: scale vector σ + 相関 Cholesky L パラメタ化 (観測専用)。 L は + -- lkjCorrCholesky bind 由来の VList-of-VList。 covariance = (diag σ·L)(diag σ·L)ᵀ。 + ("MvNormalChol", [mu, sigma, lExpr]) -> do + muV <- evalList mu + sigmaV <- evalList sigma + lV <- evalMat lExpr + let k = length muV + when (length sigmaV /= k) + (Left ("MvNormalChol: 平均ベクトル長 " <> T.pack (show k) + <> " と scale ベクトル長 " <> T.pack (show (length sigmaV)) <> " が一致しません")) + when (length lV /= k || any ((/= k) . length) lV) + (Left ("MvNormalChol: 相関 Cholesky L は " <> T.pack (show k) <> "×" + <> T.pack (show k) <> " 行列である必要があります")) + Right (HBM.MvNormalChol muV sigmaV lV) + -- Phase 45: 混合分布 (スカラ単一列、 観測は scalar observe)。 第 1 引数 = + -- 重みベクトル ([a]、 literal `[0.3,0.7]` or dirichlet 由来 VList ref、 既存 + -- evalList で評価)、 第 2 引数 = 成分分布リスト (EList の各要素を **再帰 + -- evalDist** = `[Distribution a]`)。 component 数 K は EList 要素数で静的決定。 + -- MvNormal (Phase 44) と異なり multi-column ではない (logDensity はスカラ x)。 + ("Mixture", [weights, EList distExprs]) -> do + ws <- evalList weights + comps <- mapM (evalDist env dataMap mi) distExprs + when (null comps) + (Left "Mixture: 成分分布が空です (第 2 引数に少なくとも 1 つの分布が必要)") + when (length ws /= length comps) + (Left ("Mixture: 重み数 " <> T.pack (show (length ws)) + <> " と成分分布数 " <> T.pack (show (length comps)) <> " が一致しません")) + Right (HBM.Mixture ws comps) + _ -> Left ("Unsupported distribution: " <> name <> " with " <> T.pack (show (length args)) <> " args") + where + eval = evalScalar env dataMap mi + -- list 引数 ([a]): EList を VList に評価 → 各要素をスカラに。 + evalList e = evalValue env dataMap mi e >>= asList >>= mapM asNum + -- 行列引数 ([[a]]): VList-of-VList に評価 → 'asMatrix' で正方性検査。 + evalMat e = evalValue env dataMap mi e >>= asMatrix + mk1 f a = f <$> eval a + mk2 f a b = f <$> eval a <*> eval b + mk3 f a b c = f <$> eval a <*> eval b <*> eval c + +-- | [日本語]: k-vector 観測を取る多変量分布か。 @observeMV@ はこれらのみ +-- 受理し、 scalar 分布が渡されたら親切エラーにする。 obsLogSum +-- (HBM.hs:987) が chunk 処理する分布と対応する。 +-- [English]: Whether this is a multivariate distribution that takes a +-- k-vector observation. @observeMV@ accepts only these, and gives a +-- friendly error if a scalar distribution is passed. Corresponds to the +-- distributions that obsLogSum (HBM.hs:987) chunk-processes. +isMultivariateDist :: HBM.Distribution a -> Bool +isMultivariateDist d = HBM.distName d `elem` + [ "MvNormal", "MvNormalChol", "MvStudentT" + , "Multinomial", "DirichletMultinomial", "Wishart" ] + +-- 'collectApp' は Hanalyze.Model.HBM.Ast から import。 + +-- =========================================================================== +-- Phase 43: list 値 Model combinator (latent vector を返す DoBind) +-- =========================================================================== +-- +-- 現 'DoBind' は scalar @sample@ 専用 (= @x <- Dist …@ で 1 値)。 だが +-- @cuts <- orderedCuts "cut" 2 (-2) 1@ や @probs <- dirichlet "pi" [1,1,1]@ の +-- ように **latent vector を返す** Model combinator (HBM.orderedCuts / +-- HBM.dirichlet、 いずれも @Model a [a]@) は scalar に乗らない。 これらは +-- 'evalDist' (= Distribution を返す) とは別経路で、 DoBind の RHS を +-- 'matchListComb' で検出し、 Model モナドで実行して 'VList' に束縛する。 +-- +-- 消費側 (@OrderedLogistic eta cuts@ / @Categorical probs@) は env 内で cuts / +-- probs が VList に束縛されるので Phase 42 の evalList (evalValue >>= asList >>= +-- mapM asNum) がそのまま解決する (本機構の追加は **bind 側のみ**)。 + +-- | [日本語]: 検出した list 値 combinator 呼び出し (引数は評価済 = base 名 + 構築情報)。 +-- [English]: A detected list-valued combinator call (arguments already +-- evaluated = base name + construction info). +-- 結果ベクトルの長さは引数から静的に決まる ('listCombLen')。 +data ListComb a + = OrderedCutsComb Text Int a a -- ^ name, nCuts (= K-1 ≥ 1), cMin, HalfNormal scale + | DirichletComb Text [a] -- ^ [日本語]: name, α 集中度ベクトル (長さ K ≥ 2) [English]: name, the α concentration vector (length K ≥ 2) + +-- | [日本語]: combinator が返すベクトルの長さ (validateStmts の placeholder VList 長 / +-- interpStmts は実値から決まるので参照不要)。 +-- [English]: The length of the vector the combinator returns (used for +-- validateStmts's placeholder VList length; interpStmts derives it from +-- the actual value, so it doesn't need to consult this). +listCombLen :: ListComb a -> Int +listCombLen (OrderedCutsComb _ n _ _) = n +listCombLen (DirichletComb _ as) = length as + +-- | [日本語]: base 名に plate suffix を付ける (forEachGroup 内で群ごとに別 latent にする)。 +-- [English]: Appends the plate suffix to the base name (so that +-- forEachGroup gives each group a distinct latent). +listCombSuffix :: Text -> ListComb a -> ListComb a +listCombSuffix suf (OrderedCutsComb nm n cm sc) = OrderedCutsComb (nm <> suf) n cm sc +listCombSuffix suf (DirichletComb nm as) = DirichletComb (nm <> suf) as + +-- | [日本語]: DoBind の RHS が list 値 combinator (orderedCuts / dirichlet) なら +-- 引数を評価して 'ListComb' に。 combinator でなければ 'Nothing' (= scalar +-- sample 経路へ)。 名前は文字列リテラル必須、 nCuts は数値リテラル必須 +-- (静的に長さを決めるため。 現状の制約、 doc 想定リスク参照)。 +-- [English]: If the DoBind's RHS is a list-valued combinator (orderedCuts +-- \/ dirichlet), evaluates the arguments into a 'ListComb'; if it is not a +-- combinator, returns 'Nothing' (falling through to the scalar sample +-- path). The name must be a string literal, and nCuts must be a numeric +-- literal (to determine the length statically; this is a current +-- limitation — see the doc's noted risk). +matchListComb + :: forall a. (Floating a, Ord a) + => EnvA a -> DataMap -> Expr -> Maybe (Err (ListComb a)) +matchListComb env dataMap expr = case collectApp expr of + Right ("orderedCuts", [nameE, nCutsE, cMinE, scaleE]) -> Just $ do + nm <- textLit "orderedCuts" nameE + n <- intLit "orderedCuts" nCutsE + when (n < 1) (Left "orderedCuts: カット数 (第 2 引数) は 1 以上である必要があります") + cm <- evalScalar env dataMap Nothing cMinE + sc <- evalScalar env dataMap Nothing scaleE + Right (OrderedCutsComb nm n cm sc) + Right ("dirichlet", [nameE, alphasE]) -> Just $ do + nm <- textLit "dirichlet" nameE + as <- evalValue env dataMap Nothing alphasE >>= asList >>= mapM asNum + when (length as < 2) (Left "dirichlet: α ベクトル (第 2 引数) は長さ 2 以上の [..] である必要があります") + Right (DirichletComb nm as) + _ -> Nothing + where + textLit _ (ELit (LText t)) = Right t + textLit fn _ = Left (fn <> " の名前引数 (第 1 引数) は文字列リテラルである必要があります") + intLit :: Text -> Expr -> Err Int + intLit _ (ELit (LNumber d)) = Right (round d) + intLit fn _ = Left (fn <> " のカット数引数 (第 2 引数) は数値リテラルである必要があります (変数経由は未対応)") + +-- | [日本語]: 'ListComb' を実際の Model アクション (latent vector を sample) に。 +-- [English]: Turns a 'ListComb' into an actual Model action (sampling a +-- latent vector). +runListComb :: forall a. (Floating a, Ord a) => ListComb a -> HBM.Model a [a] +runListComb (OrderedCutsComb nm n cm sc) = HBM.orderedCuts nm n cm sc +runListComb (DirichletComb nm as) = HBM.dirichlet nm as + +-- =========================================================================== +-- Phase 44: 行列値 Model combinator (latent 相関行列を返す DoBind) +-- =========================================================================== +-- +-- 'ListComb' (Phase 43、 @Model a [a]@) の行列版。 @lkjCorrCholesky@ は +-- @Model a [[a]]@ で k×k 下三角の相関 Cholesky 因子 L を返す latent +-- combinator。 @L <- lkjCorrCholesky "L" 2 2.0@ を 'VList'-of-'VList' に束縛し、 +-- 消費側 ('MvNormalChol' の第 3 引数) は env 内の VList-of-VList を 'asMatrix' +-- で解決する。 内部 latent (@L_pc*@ / @L_L*@ 等) は Model が自動登録する +-- (DSL は latent を再実装しない)。 + +-- | [日本語]: 検出した行列値 combinator 呼び出し (引数評価済)。 +-- [English]: A detected matrix-valued combinator call (arguments already +-- evaluated). +data MatrixComb a + = LkjCholComb Text Int a -- ^ [日本語]: name, dim k (≥ 2), eta (LKJ 集中度) [English]: name, dim k (≥ 2), eta (LKJ concentration) + +-- | [日本語]: combinator が返す行列の次元 k (validateStmts の placeholder 用)。 +-- [English]: The dimension k of the matrix the combinator returns (for +-- validateStmts's placeholder). +matrixCombDim :: MatrixComb a -> Int +matrixCombDim (LkjCholComb _ k _) = k + +-- | [日本語]: base 名に plate suffix を付ける (群ごとに別 latent にする)。 +-- [English]: Appends the plate suffix to the base name (giving each group +-- a distinct latent). +matrixCombSuffix :: Text -> MatrixComb a -> MatrixComb a +matrixCombSuffix suf (LkjCholComb nm k eta) = LkjCholComb (nm <> suf) k eta + +-- | [日本語]: DoBind の RHS が行列値 combinator (lkjCorrCholesky) なら引数を評価して +-- 'MatrixComb' に。 combinator でなければ 'Nothing'。 名前は文字列リテラル、 +-- 次元 k は数値リテラル必須 (静的に行列サイズを決めるため)。 +-- [English]: If the DoBind's RHS is a matrix-valued combinator +-- (lkjCorrCholesky), evaluates the arguments into a 'MatrixComb'; if it is +-- not a combinator, returns 'Nothing'. The name must be a string literal, +-- and the dimension k must be a numeric literal (to determine the matrix +-- size statically). +matchMatrixComb + :: forall a. (Floating a, Ord a) + => EnvA a -> DataMap -> Expr -> Maybe (Err (MatrixComb a)) +matchMatrixComb env dataMap expr = case collectApp expr of + Right ("lkjCorrCholesky", [nameE, kE, etaE]) -> Just $ do + nm <- textLit nameE + k <- intLit kE + when (k < 2) (Left "lkjCorrCholesky: 次元 (第 2 引数) は 2 以上である必要があります") + eta <- evalScalar env dataMap Nothing etaE + Right (LkjCholComb nm k eta) + _ -> Nothing + where + textLit (ELit (LText t)) = Right t + textLit _ = Left "lkjCorrCholesky の名前引数 (第 1 引数) は文字列リテラルである必要があります" + intLit :: Expr -> Err Int + intLit (ELit (LNumber d)) = Right (round d) + intLit _ = Left "lkjCorrCholesky の次元引数 (第 2 引数) は数値リテラルである必要があります (変数経由は未対応)" + +-- | [日本語]: 'MatrixComb' を実際の Model アクション (latent 相関行列を sample) に。 +-- [English]: Turns a 'MatrixComb' into an actual Model action (sampling a +-- latent correlation matrix). +runMatrixComb :: forall a. (Floating a, Ord a) => MatrixComb a -> HBM.Model a [[a]] +runMatrixComb (LkjCholComb nm k eta) = HBM.lkjCorrCholesky nm k eta + +-- | [日本語]: stmts を walk して各 latent 変数(DoBind の左辺)に対する +-- Transform を返す。 constrained 空間で 0 初期化は PositiveT で log 0 = -∞ +-- 発散するため、 streaming endpoint で transform 別に初期値を選ぶのに使う。 +-- [English]: Walks the stmts and returns the Transform for each latent +-- variable (a DoBind's left-hand side). Initializing at 0 in constrained +-- space diverges for PositiveT (log 0 = -∞), so this is used by the +-- streaming endpoint to pick an initial value per transform. +inferTransforms :: [DoStmt] -> Map.Map Text HD.Transform +inferTransforms = Map.fromList . concatMap extract + where + extract (DoBind name distExpr) = case collectApp distExpr of + Right (dname, _) -> [(name, distNameToTransform dname)] + _ -> [(name, HD.UnconstrainedT)] + extract _ = [] + + distNameToTransform "Normal" = HD.UnconstrainedT + distNameToTransform "StudentT" = HD.UnconstrainedT + distNameToTransform "Cauchy" = HD.UnconstrainedT + distNameToTransform "Uniform" = HD.UnconstrainedT + distNameToTransform "HalfNormal" = HD.PositiveT + distNameToTransform "HalfCauchy" = HD.PositiveT + distNameToTransform "Gamma" = HD.PositiveT + distNameToTransform "Exponential" = HD.PositiveT + distNameToTransform "LogNormal" = HD.PositiveT + distNameToTransform "InverseGamma" = HD.PositiveT + distNameToTransform "Weibull" = HD.PositiveT + distNameToTransform "Beta" = HD.UnitIntervalT + distNameToTransform "Bernoulli" = HD.UnitIntervalT + distNameToTransform _ = HD.UnconstrainedT + +-- =========================================================================== +-- AST → Model モナド構築 +-- =========================================================================== + +-- | [日本語]: EDo 内の各 stmt を Model モナドに翻訳。実装は 'forall a' のもとに +-- 動作する必要があるが、Err は Haskell の純粋値なので外側で先に検査して +-- Model 構築は失敗しない前提にする。エラーは事前検証で全部捕まえる方針。 +-- [English]: Translates each stmt inside an EDo into the Model monad. The +-- implementation must work under 'forall a', but since Err is a pure +-- Haskell value, we check it up front on the outside and assume Model +-- construction never fails from here on. The policy is to catch all +-- errors during pre-validation. +-- +-- ここでは「事前検証ありで、検証通過後に Model を直接組み上げる」設計。 +-- ModelP は forall を含む rank-1 polymorphic 型なので Either に直接乗らない +-- (ImpredicativeTypes を避けるため)。validate と build を分離する。 +validateAst :: [TopBind] -> Expr -> DataMap -> Either Text [DoStmt] +validateAst topBinds body0 dataMap = do + rawStmts <- case body0 of + -- Phase 13 §9.3c-2: frontend parser は最終 DoExpr を ret として分離する + -- (hanalyze 慣行で「observe が最後 + pure 省略」 が許される)。 ret が + -- pure / return のときは捨て、 それ以外(例: observe)は body 末尾に + -- DoExpr として戻して扱う。 + EDo s ret -> + let isDiscard = case ret of + EApp (EVar "pure") _ -> True + EApp (EVar "return") _ -> True + ELit (LBool _) -> True -- frontend の implicit ret fallback + _ -> False + full = if isDiscard then s else s ++ [DoExpr ret] + in Right full + _ -> Left "Model body must be a do-block (`do { ... }`)" + -- Phase 26.1 §A-2 (2026-05-27): `x <- Data "label" expr` を pre-process で + -- substitute (= 列 alias 経路)。 詳細は streaming bridge/.../HbmAst.hs の同名 + -- 関数 doc 参照。 hanalyze の pm.Data 厳密対応 (withData 経由) は将来 phase。 + let stmts = preprocessAliases rawStmts + validateStmts topBinds dataMap stmts + pure stmts + +-- | [日本語]: alias 経路 (2026-05-27 王道方針に切替): Haskell の `let` +-- を syntactic alias として扱う pre-processing。 spec §4.2 と整合 +-- (= ∀LIC∃Code は Haskell サブセット、 `let x = col "..."` で Expression +-- Language alias)。 詳細は streaming bridge/.../HbmAst.hs の同名関数 doc 参照。 +-- [English]: The alias path (switched to the "orthodox" approach on +-- 2026-05-27): pre-processing that treats Haskell's `let` as a syntactic +-- alias. Consistent with spec §4.2 (= ∀LIC∃Code is a Haskell subset; +-- `let x = col "..."` is an Expression Language alias). See the +-- like-named function's doc in streaming bridge/.../HbmAst.hs for details. +preprocessAliases :: [DoStmt] -> [DoStmt] +preprocessAliases = go Map.empty + where + go _ [] = [] + go aliases (s : rest) = case extractLetAliases aliases s of + Just (newAliases, mStmt) -> case mStmt of + Nothing -> go newAliases rest + Just stmt -> stmt : go newAliases rest + Nothing -> + substituteStmt aliases s : go aliases rest + + extractLetAliases + :: Map.Map Text Expr -> DoStmt + -> Maybe (Map.Map Text Expr, Maybe DoStmt) + extractLetAliases aliases (DoLet bs) = + let (newAliases, kept) = foldl step (aliases, []) bs + step (al, acc) (Bind n v) = + let substV = substitute al v + in if hasColRef substV + then (Map.insert n substV al, acc) + else (al, acc ++ [Bind n substV]) + in Just (newAliases, if null kept then Nothing else Just (DoLet kept)) + extractLetAliases _ _ = Nothing + + substituteStmt :: Map.Map Text Expr -> DoStmt -> DoStmt + substituteStmt aliases (DoBind n v) = DoBind n (substitute aliases v) + substituteStmt aliases (DoLet bs) = DoLet (map (substBind aliases) bs) + substituteStmt aliases (DoExpr e) = DoExpr (substitute aliases e) + + substBind :: Map.Map Text Expr -> Bind -> Bind + substBind aliases (Bind n v) = Bind n (substitute aliases v) + + substitute :: Map.Map Text Expr -> Expr -> Expr + substitute env = go' + where + go' (EVar n) = case Map.lookup n env of + Just e -> e + Nothing -> EVar n + go' (EOp op a b) = EOp op (go' a) (go' b) + go' (EApp f x) = EApp (go' f) (go' x) + go' (ENeg e) = ENeg (go' e) + go' (ELet bs body) = ELet (map (substBind env) bs) (go' body) + go' (EList xs) = EList (map go' xs) + go' (EIf c a b) = EIf (go' c) (go' a) (go' b) + go' (ELam x b) = ELam x (go' b) + go' (EDo s r) = EDo s r -- 入れ子 do は触らない + go' x = x + +-- | [日本語]: 静的検証(変数名スコープ / 列存在)。型は (Double, Double) 環境で一度評価して +-- 実行時エラーが起きないかを確認する。 +-- [English]: Static validation (variable scoping \/ column existence). +-- Evaluates once in a (Double, Double) environment to confirm no runtime +-- error occurs. +validateStmts :: [TopBind] -> DataMap -> [DoStmt] -> Err () +validateStmts topBinds dataMap stmts0 = go (buildTopEnv dataMap topBinds) stmts0 + where + go :: EnvA Double -> [DoStmt] -> Err () + go _env [] = Right () + -- Phase 43: RHS が list 値 combinator (orderedCuts / dirichlet) なら、 + -- 構造検証 + 結果長 K の placeholder VList を束縛する (Model モナドが無い + -- 検証経路では実行できないため。 後続 Categorical/OrderedLogistic の検証が + -- 通るように長さだけ合わせる)。 + go env (DoBind name distExpr : rest) + | Just ecomb <- matchListComb env dataMap distExpr :: Maybe (Err (ListComb Double)) = do + comb <- ecomb + let k = listCombLen comb + go (Map.insert name (VList (replicate k (VNum 0.0))) env) rest + -- Phase 44: 行列値 combinator は k×k placeholder VList-of-VList を束縛する + -- (Model モナドが無い検証経路では実行できないため、 次元だけ合わせる)。 + go env (DoBind name distExpr : rest) + | Just ecomb <- matchMatrixComb env dataMap distExpr :: Maybe (Err (MatrixComb Double)) = do + comb <- ecomb + let k = matrixCombDim comb + go (Map.insert name (VList (replicate k (VList (replicate k (VNum 0.0))))) env) rest + go env (DoBind name distExpr : rest) = do + _ <- if hasColRef distExpr + then Left ("sample distribution cannot reference data column: " <> name) + else evalDist env dataMap Nothing distExpr :: Err (HBM.Distribution Double) + go (Map.insert name (VNum 0.0) env) rest + go env (DoLet bs : rest) = do + env' <- foldEnv env bs + go env' rest + go env (DoExpr e : rest) + -- Phase 27 §F-3c: forEachGroup は群列の存在を確認し、 内部 stmts を + -- 代表コンテキスト (param = 0) で検証する (行サブセットは検証不要)。 + | Just (gcol, param, inner) <- matchForEachGroup e = do + when (Map.notMember gcol dataMap) + (Left ("forEachGroup の群列が見つかりません: " <> gcol)) + go (Map.insert param (VNum 0) env) inner + go env rest + | otherwise = do + -- observe の構文形を検査 + validateObserve env e + go env rest + + foldEnv e [] = Right e + foldEnv e (Bind n v : rest) = do + vv <- evalValue e dataMap Nothing v + foldEnv (Map.insert n vv e) rest + + validateObserve env e = do + (fname, args) <- collectApp e + case (fname, args) of + ("observe", [ELit (LText _obsName), distExpr, dataRef]) -> do + colName <- requireColRef dataRef + when (Map.notMember colName dataMap) (Left ("Unknown column in observe: " <> colName)) + if hasColRef distExpr + then do + -- 列参照あり: 各行で評価できることを確認(row 0 でテスト) + _ <- evalDist env dataMap (Just 0) distExpr :: Err (HBM.Distribution Double) + pure () + else do + _ <- evalDist env dataMap Nothing distExpr :: Err (HBM.Distribution Double) + pure () + -- Phase 44: multi-column observe。 第 3 引数は観測列リスト ['y1','y2',..]。 + -- 全列の存在 + 列数 ≥ 2 + 列長一致 + dist が多変量かを検査する。 dist の + -- μ/cov は行不変前提なので row 参照不可 (Nothing) で評価する。 + ("observeMV", [ELit (LText _obsName), distExpr, EList colRefs]) -> do + cols <- mapM requireColRef colRefs + when (length cols < 2) + (Left "observeMV: 観測列は 2 列以上必要です (1 列なら scalar observe を使ってください)") + forM_ cols $ \c -> + when (Map.notMember c dataMap) (Left ("Unknown column in observeMV: " <> c)) + let lens = map (length . (`lookupDoubles` dataMap)) cols + when (any (/= head lens) (tail lens)) + (Left "observeMV: 観測列の長さが揃っていません (k-vector 組成には全列同長が必要です)") + d <- evalDist env dataMap Nothing distExpr :: Err (HBM.Distribution Double) + when (not (isMultivariateDist d)) + (Left ("observeMV: 第 2 引数は多変量分布 (MvNormal 等) が必要ですが、 scalar 分布 " + <> HBM.distName d <> " が渡されました")) + ("observeMV", _) -> + Left "observeMV: 第 3 引数は観測列リスト ['y1','y2',..] が必要です" + ("pure", _) -> Right () -- pure x は最終行用 + ("return", _) -> Right () + _ -> Left ("Unsupported statement: " <> fname) + + requireColRef (ECol n) = Right n + requireColRef _ = Left "observe's third argument must be a column reference 'colname'" + +-- | [日本語]: 検証通過後の Model 構築(polymorphic in a)。エラーは想定外なので +-- error で落とす(validation で漏れたバグは fail-fast)。 +-- [English]: Builds the Model after validation passes (polymorphic in a). +-- Errors are unexpected at this point, so it fails via error (a bug that +-- slipped past validation should fail fast). +interpStmts :: [TopBind] -> DataMap -> [DoStmt] -> HBM.ModelP () +interpStmts topBinds dataMap stmts = goM topCtx (buildTopEnv dataMap topBinds) stmts + where + goM :: forall a. (Floating a, Ord a) => PlateCtx -> EnvA a -> [DoStmt] -> HBM.Model a () + goM _ _env [] = pure () + -- Phase 43: list 値 combinator (orderedCuts / dirichlet) は scalar sample + -- ではなく Model アクションを実行して latent vector を VList 束縛する。 + -- base 名に plate suffix を付け、 forEachGroup 内で群ごとに別 latent にする。 + goM ctx env (DoBind name distExpr : rest) + | Just ecomb <- matchListComb env dataMap distExpr = do + let comb = case ecomb of + Right c -> listCombSuffix (pcSuffix ctx) c + Left e -> error (T.unpack e) + xs <- runListComb comb + goM ctx (Map.insert name (VList (map VNum xs)) env) rest + -- Phase 44: 行列値 combinator (lkjCorrCholesky) は latent 相関 Cholesky 因子を + -- Model で実行し、 VList-of-VList に束縛する (MvNormalChol の L 引数で消費)。 + goM ctx env (DoBind name distExpr : rest) + | Just ecomb <- matchMatrixComb env dataMap distExpr = do + let comb = case ecomb of + Right c -> matrixCombSuffix (pcSuffix ctx) c + Left e -> error (T.unpack e) + m <- runMatrixComb comb + goM ctx (Map.insert name (VList (map (VList . map VNum) m)) env) rest + goM ctx env (DoBind name distExpr : rest) = do + let dist = case evalDist env dataMap Nothing distExpr of + Right d -> d + Left e -> error (T.unpack e) + x <- HBM.sample (name <> pcSuffix ctx) dist + goM ctx (Map.insert name (VNum x) env) rest + goM ctx env (DoLet bs : rest) = do + let env' = foldlBinds env bs + goM ctx env' rest + goM ctx env (DoExpr e : rest) + -- Phase 27 §F-3c: forEachGroup は群ごとに内部 do-block を展開する。 + | Just (gcol, param, inner) <- matchForEachGroup e = do + let gvals = groupValsIn dataMap gcol (pcRows ctx) + forM_ gvals $ \gval -> do + let ctx' = ctx + { pcSuffix = pcSuffix ctx <> groupSuffixFor (Map.lookup gcol dataMap) gval + , pcRows = Just (rowsForGroup dataMap gcol gval (pcRows ctx)) + } + env' = Map.insert param (VNum (liftD gval)) env + goM ctx' env' inner + goM ctx env rest + | otherwise = do + execObserve ctx env e + goM ctx env rest + + foldlBinds :: forall a. (Floating a, Ord a) => EnvA a -> [Bind] -> EnvA a + foldlBinds e [] = e + foldlBinds e (Bind n v : rs) = + let vv = case evalValue e dataMap Nothing v of + Right x -> x + Left err -> VErr err -- 検証通過しているはずなので通常来ない + in foldlBinds (Map.insert n vv e) rs + + execObserve :: forall a. (Floating a, Ord a) => PlateCtx -> EnvA a -> Expr -> HBM.Model a () + execObserve ctx env e = case collectApp e of + Right ("observe", [ELit (LText obsName), distExpr, ECol colName]) -> + let ys = lookupDoubles colName dataMap + allRows = [0 .. length ys - 1] + rows = fromMaybe allRows (pcRows ctx) -- 群コンテキストなら当該群の行 + nm = obsName <> pcSuffix ctx + in if hasColRef distExpr + then do + -- per-row distribution。 対象行のみ observeColumns でまとめる。 + let pairs = [ (case evalDist env dataMap (Just i) distExpr of + Right d -> d + Left _ -> error "validation should have caught this" + , [ys !! i]) + | i <- rows, i >= 0, i < length ys + ] + HBM.observeColumns nm pairs + else do + let dist = case evalDist env dataMap Nothing distExpr of + Right d -> d + Left _ -> error "validation should have caught this" + HBM.observe nm dist [ ys !! i | i <- rows, i >= 0, i < length ys ] + -- Phase 44: multi-column observe。 dist (μ/cov) は行不変なので 1 回だけ + -- 評価し、 観測列だけを行ごとに k-vector に組んで HBM.observeMV に流す。 + Right ("observeMV", [ELit (LText obsName), distExpr, EList colRefs]) -> + let cols = [ c | ECol c <- colRefs ] + colVecs = map (`lookupDoubles` dataMap) cols + n = if null colVecs then 0 else minimum (map length colVecs) + allRows = [0 .. n - 1] + rows = fromMaybe allRows (pcRows ctx) + nm = obsName <> pcSuffix ctx + dist = case evalDist env dataMap Nothing distExpr of + Right d -> d + Left _ -> error "validation should have caught this" + obss = [ [ cv !! i | cv <- colVecs ] | i <- rows, i >= 0, i < n ] + in HBM.observeMV nm dist obss + Right ("pure", _) -> pure () + Right ("return", _) -> pure () + _ -> pure () -- validateObserve で弾く想定 + +-- =========================================================================== +-- NUTS 設定 reader +-- =========================================================================== + +-- | [日本語]: extra から NUTS 設定を読む(欠落時はデフォルト)。 +-- [English]: Reads the NUTS config from extra (defaults when absent). +readChainCount :: A.Object -> Int +readChainCount o = case KM.lookup (Key.fromText "hbmChains") o of + Just (A.Number n) -> let v = floor (realToFrac n :: Double) in max 1 (min 16 v) + _ -> 4 + +readNutsConfig :: A.Object -> NUTS.NUTSConfig +readNutsConfig o = + let + def = NUTS.defaultNUTSConfig + getInt :: Text -> Int + getInt k = case KM.lookup (Key.fromText k) o of + Just (A.Number n) -> floor (realToFrac n :: Double) + _ -> 0 + getDbl :: Text -> Double + getDbl k = case KM.lookup (Key.fromText k) o of + Just (A.Number n) -> realToFrac n + _ -> 0 + getBool' k = case KM.lookup (Key.fromText k) o of + Just (A.Bool b) -> b + _ -> False + in def + { NUTS.nutsIterations = if getInt "hbmIterations" > 0 then getInt "hbmIterations" else NUTS.nutsIterations def + , NUTS.nutsBurnIn = max 0 (getInt "hbmBurnIn") + , NUTS.nutsStepSize = if getDbl "hbmStepSize" > 0 then getDbl "hbmStepSize" else NUTS.nutsStepSize def + , NUTS.nutsMaxDepth = if getInt "hbmMaxDepth" > 0 then getInt "hbmMaxDepth" else NUTS.nutsMaxDepth def + , NUTS.nutsAdaptStepSize = getBool' "hbmAdaptStepSize" + , NUTS.nutsTargetAccept = let v = getDbl "hbmTargetAccept" in if v > 0 && v < 1 then v else NUTS.nutsTargetAccept def + , NUTS.nutsAdaptMass = getBool' "hbmAdaptMass" + } + +-- | [日本語]: observe ノード名 → 観測列名 のマッピングを取り出す。 +-- DSL 構文 @observe "NAME" DIST @COL@@ から (NAME, COL) を抽出。 +-- frontend で「observe ノード "y" の観測値はどの列か」 を解決する用途。 +-- [English]: Extracts the observe-node-name → observed-column-name +-- mapping. Extracts (NAME, COL) from the DSL syntax +-- @observe "NAME" DIST @COL@@. Used by the frontend to resolve "which +-- column is the observe node \"y\"'s observed value?". +observeNodeMap :: [DoStmt] -> [(Text, Text)] +observeNodeMap = concatMap step + where + step (DoExpr e) = case collectApp e of + Right ("observe", [ELit (LText obsName), _distExpr, ECol colName]) -> + [(obsName, colName)] + _ -> [] + step _ = [] + +-- =========================================================================== +-- 結果整形 (param summary / posterior mean curves) +-- =========================================================================== + +data ParamSummary = ParamSummary + { psName :: !Text + , psMean :: !Double + , psSd :: !Double + , psLow :: !Double + , psHigh :: !Double + , psRhat :: !(Maybe Double) + , psEss :: !Double + } deriving (Show) + +-- | [日本語]: SC29: 複数 chain から事後統計を計算。R̂ は split-R̂ (hanalyze)、 +-- ESS は Geyer initial monotone(全チェーン pool)。 +-- [English]: SC29: Computes posterior statistics from multiple chains. R̂ +-- is split-R̂ (hanalyze); ESS is Geyer's initial monotone sequence +-- (pooled across all chains). +paramSummaryMulti :: [MC.Chain] -> Text -> ParamSummary +paramSummaryMulti chains name = + let perChain = map (MC.chainVals name) chains + pooled = concat perChain + n = length pooled + m = if n == 0 then 0 else sum pooled / fromIntegral n + sd2 = if n < 2 then 0 + else sum (map (\v -> (v - m) ** 2) pooled) / fromIntegral (n - 1) + sd = sqrt sd2 + sorted = LA.toList (LA.sortVector (LA.fromList pooled)) + pct :: Double -> Double + pct q = if n == 0 then 0 + else let idx = max 0 (min (n - 1) (floor (q * fromIntegral (n - 1)) :: Int)) + in case drop idx sorted of (x:_) -> x; [] -> 0 + rh = SMC.rhat perChain + e = SMC.ess pooled + in ParamSummary name m sd (pct 0.025) (pct 0.975) rh e + +fmtSummary :: ParamSummary -> Text +fmtSummary p = psName p <> "=" + <> T.pack (show (round4 (psMean p))) + <> "±" <> T.pack (show (round4 (psSd p))) + +round4 :: Double -> Double +round4 x = fromIntegral (round (x * 10000) :: Int) / 10000 + +-- | [日本語]: thinning: stride 飛ばしに要素を取る。 +-- [English]: Thinning: takes elements while skipping by stride. +takeEvery :: Int -> [a] -> [a] +takeEvery _ [] = [] +takeEvery n (x:xs) = x : takeEvery n (drop (max 0 (n - 1)) xs) + +summaryToJson :: ParamSummary -> A.Value +summaryToJson p = A.object + [ Key.fromText "name" A..= psName p + , Key.fromText "mean" A..= psMean p + , Key.fromText "sd" A..= psSd p + , Key.fromText "ci2_5" A..= psLow p + , Key.fromText "ci97_5" A..= psHigh p + , Key.fromText "rhat" A..= psRhat p + , Key.fromText "ess" A..= psEss p + ] + +-- --------------------------------------------------------------------------- +-- Phase NN §A (2026-05-27): HBM posterior predictive mean curves +-- --------------------------------------------------------------------------- +-- DSL の observe 文の mean expression (= 例: `alpha + beta * @x@`) を +-- 直接評価して posterior predictive curve を計算する。 frontend の +-- buildHbmOverlay 内 heuristic (= "alpha" / "beta_" 等の名前推定) を +-- 撤去するための backend AST driven 経路。 + +data HbmMeanCurve = HbmMeanCurve + { hmcObsName :: !Text -- observe 文の名前 (top="y"、 per-group="y_1") + , hmcPredictor :: !Text -- 説明変数列名 (= 例: "x") + , hmcX :: ![Double] -- 64 grid points + , hmcMedian :: ![Double] -- posterior median of mu(x) + , hmcLower :: ![Double] -- 2.5% + , hmcUpper :: ![Double] -- 97.5% + -- Phase 27 GLMM overlay: per-group curve の群ラベル (top-level は Nothing)。 + -- frontend が群ごと色分け + 凡例に使う。 + , hmcGroupCol :: !(Maybe Text) -- 群列名 (= forEachGroup "gcol") + , hmcGroupVal :: !(Maybe Double) -- 群値 + } deriving (Show) + +hbmMeanCurveToJson :: HbmMeanCurve -> A.Value +hbmMeanCurveToJson c = A.object + [ Key.fromText "obsName" A..= hmcObsName c + , Key.fromText "predictor" A..= hmcPredictor c + , Key.fromText "x" A..= hmcX c + , Key.fromText "median" A..= hmcMedian c + , Key.fromText "lower" A..= hmcLower c + , Key.fromText "upper" A..= hmcUpper c + , Key.fromText "groupCol" A..= hmcGroupCol c + , Key.fromText "groupVal" A..= hmcGroupVal c + ] + +-- | [日本語]: stmts から各 observe の mean 式 (= Distribution の第 1 引数) を抽出。 +-- [English]: Extracts each observe's mean expression (= the +-- Distribution's first argument) from the stmts. +extractObserveMeans :: [DoStmt] -> [(Text, Expr)] +extractObserveMeans = concatMap go + where + go (DoExpr e) = case collectApp e of + Right ("observe", [ELit (LText nm), distExpr, _]) -> + case collectApp distExpr of + Right (_, m : _) -> [(nm, m)] + _ -> [] + _ -> [] + go _ = [] + +-- | [日本語]: 式中の全 ECol 列名を集める (重複排除)。 +-- [English]: Collects all ECol column names appearing in the expression +-- (deduplicated). +collectCols :: Expr -> [Text] +collectCols = nub . go + where + go (ECol c) = [c] + go (EOp _ a b) = go a <> go b + go (EApp f x) = go f <> go x + go (ENeg e) = go e + go (ELet bs body) = concatMap (\(Bind _ v) -> go v) bs <> go body + go _ = [] + +-- | [日本語]: percentile (= 0..1) を sorted リストから線形補間で取る。 +-- [English]: Reads a percentile (= 0..1) from a sorted list via linear +-- interpolation. +percentileOf :: Double -> [Double] -> Double +percentileOf q xs0 = + let xs = sort xs0 + n = length xs + in if n == 0 then 0 + else + let idxF = q * fromIntegral (n - 1) + idx = max 0 (min (n - 1) (floor idxF)) + in case drop idx xs of + (v:_) -> v + [] -> 0 + +-- | [日本語]: observe 1 個分の mean-curve 計算文脈 (top-level または per-group)。 +-- GLMM overlay: @forEachGroup@ 内の observe も拾えるよう、 plate +-- 展開を replay しながら集める ('collectObsInstances')。 +-- [English]: The mean-curve computation context for one observe +-- (top-level or per-group). For the GLMM overlay: to also pick up +-- observes inside @forEachGroup@, this replays the plate expansion while +-- collecting them ('collectObsInstances'). +-- | [日本語]: WAIC/PPC 再評価用、 list 値 combinator 由来の latent vector を +-- posterior サンプルから再構築する仕様。 stream worker は @sampleNames@ +-- (= sample された latent のみ) を samples に載せ、 combinator の deterministic +-- (cut_c_*/pi_*) は載せないため、 再評価 env で sampled latent (cut_d_*/pi_b*) から +-- VList を組み直す ('reconstructComb')。 +-- [English]: The spec for reconstructing a list-valued combinator's latent +-- vector from posterior samples, for WAIC\/PPC re-evaluation. The stream +-- worker puts only @sampleNames@ (= sampled latents) into samples and +-- omits the combinator's deterministic quantities (cut_c_*\/pi_*), so the +-- re-evaluation env rebuilds the VList from the sampled latents +-- (cut_d_*\/pi_b*) via 'reconstructComb'. +data ListCombSpec + = OCutsSpec Text Int Expr -- ^ [日本語]: suffix 付き base 名, nCuts, cMin 式 (定数想定) [English]: suffixed base name, nCuts, the cMin expression (assumed constant) + | DirSpec Text Int -- ^ [日本語]: suffix 付き base 名, K (= α ベクトル長) [English]: suffixed base name, K (= the α vector length) + +data ObsInstance = ObsInstance + { oiObsName :: Text -- 表示名 (top="y"、 group="y_1") + , oiMeanExpr :: Expr -- distribution の第 1 引数 (mean 式) + , oiDistExpr :: Expr -- distribution 式全体 (例: `Normal mu sigma`) + , oiObsCol :: Text -- 観測列名 (observe の第 3 引数 `col "y"`) + , oiNameKey :: Map.Map Text Text -- local latent 名 → samples のキー (theta→"theta_1") + , oiParamEnv :: Map.Map Text Double -- 群変数 (forEachGroup の \g) → 群値 + , oiRows :: Maybe [Int] -- 対象行 (Nothing=全行) + , oiGroup :: Maybe (Text, Double) -- 最内 forEachGroup の (群列, 群値) + , oiListCombs :: [(Text, ListCombSpec)] -- Phase 43: list latent bind (cuts/probs) の再構築仕様 + } + +-- | [日本語]: list 値 combinator の latent vector を 1 posterior サンプルから再構築する。 +-- @base@ は cMin 定数評価用の env (sample 非依存)、 @sm@ は sampled latent の Map。 +-- sampled latent (cut_d_*/pi_b*) が欠けていれば 'Nothing'。 +-- [English]: Reconstructs a list-valued combinator's latent vector from +-- one posterior sample. @base@ is the env for evaluating the cMin +-- constant (sample-independent); @sm@ is the Map of sampled latents. +-- Returns 'Nothing' if a sampled latent (cut_d_*\/pi_b*) is missing. +reconstructComb + :: EnvA Double -> DataMap -> Map.Map Text Double -> ListCombSpec -> Maybe [Double] +reconstructComb base dataMap sm spec = case spec of + -- orderedCuts: c_1 = cMin、 c_i = c_{i-1} + d_i (d_i = name_d_i, i=2..n)。 + OCutsSpec nm n cMinE -> do + cMin <- either (const Nothing) Just (evalScalar base dataMap Nothing cMinE) + ds <- mapM (\j -> Map.lookup (nm <> "_d_" <> T.pack (show j)) sm) [2 .. n] + pure (scanl (+) cMin ds) -- 長さ n、 単調増加 + -- dirichlet: stick-breaking。 betas = name_b0..name_b{K-2}、 π を復元。 + DirSpec nm k -> do + betas <- mapM (\j -> Map.lookup (nm <> "_b" <> T.pack (show j)) sm) [0 .. k - 2] + let prods = scanl (\acc b -> acc * (1 - b)) 1 betas + pure [ if j < length betas then (betas !! j) * (prods !! j) else prods !! j + | j <- [0 .. k - 1] ] + +-- | [日本語]: DoBind の RHS が list 値 combinator なら 'ListCombSpec' を返す (suffix 付き +-- base 名で。 'matchListComb' / 'listCombSuffix' と同じ命名規律)。 +-- [English]: If the DoBind's RHS is a list-valued combinator, returns a +-- 'ListCombSpec' (with a suffixed base name, following the same naming +-- rule as 'matchListComb' \/ 'listCombSuffix'). +matchListCombSpec :: Text -> Expr -> Maybe ListCombSpec +matchListCombSpec suf rhs = case collectApp rhs of + Right ("orderedCuts", [ELit (LText bn), ELit (LNumber d), cMinE, _scaleE]) -> + Just (OCutsSpec (bn <> suf) (round d) cMinE) + Right ("dirichlet", [ELit (LText bn), EList alphas]) -> + Just (DirSpec (bn <> suf) (length alphas)) + _ -> Nothing + +-- | [日本語]: observe 式から (名前, mean 式, distribution 式全体, 観測列名) を抽出。 +-- [English]: Extracts (name, mean expression, the whole distribution +-- expression, observed column name) from an observe expression. +-- mean 式 = Distribution の第 1 引数、 dist 式全体は WAIC/PPC で logDensity / +-- sampleDist を回すために必要。 観測列名は第 3 引数 `col "y"` の列。 +observeFull :: Expr -> Maybe (Text, Expr, Expr, Text) +observeFull e = case collectApp e of + Right ("observe", [ELit (LText nm), distExpr, ECol colName]) -> + case collectApp distExpr of + Right (_, m : _) -> Just (nm, m, distExpr, colName) + _ -> Nothing + _ -> Nothing + +-- | [日本語]: stmts を plate 展開しながら observe instance を集める。 'interpStmts' と +-- 同じ suffix (groupSuffix) / 行 (rowsForGroup) / 群変数束縛の規律を replay し、 +-- top-level observe は suffix=""・全行、 forEachGroup 内 observe は群ごとに +-- suffix 付き・群行で展開する。 latent 名は sample された scope の suffix で +-- samples のキーに対応づける (top latent="mu"、 群 latent="theta_1" 等)。 +-- [English]: Collects observe instances while plate-expanding the stmts. +-- Replays the same suffix (groupSuffix) \/ rows (rowsForGroup) \/ +-- group-variable-binding rule as 'interpStmts': a top-level observe +-- expands with suffix="" over all rows, while an observe inside +-- forEachGroup expands per group with a suffix and that group's rows. +-- Latent names are mapped to their samples key by the suffix of the scope +-- in which they were sampled (top latent="mu", group latent="theta_1", +-- etc.). +collectObsInstances :: DataMap -> [DoStmt] -> [ObsInstance] +collectObsInstances dm = go "" Nothing Map.empty Map.empty Nothing [] + where + go :: Text -> Maybe [Int] -> Map.Map Text Text -> Map.Map Text Double + -> Maybe (Text, Double) -> [(Text, ListCombSpec)] -> [DoStmt] -> [ObsInstance] + go _ _ _ _ _ _ [] = [] + go suf rows nameKey penv grp combs (st : rest) = case st of + -- latent: 現 suffix 付きで samples に載る (= キー対応を記録)。 Phase 43: + -- list 値 combinator なら再構築仕様も記録 (= WAIC/PPC 再評価で VList 復元)。 + DoBind name rhs -> + let combs' = case matchListCombSpec suf rhs of + Just spec -> (name, spec) : combs + Nothing -> combs + in go suf rows (Map.insert name (name <> suf) nameKey) penv grp combs' rest + DoLet _ -> go suf rows nameKey penv grp combs rest + DoExpr e + | Just (gcol, param, inner) <- matchForEachGroup e -> + let gvals = groupValsIn dm gcol rows + here = concatMap + (\gv -> + go (suf <> groupSuffixFor (Map.lookup gcol dm) gv) + (Just (rowsForGroup dm gcol gv rows)) + nameKey + (Map.insert param gv penv) + (Just (gcol, gv)) + combs + inner) + gvals + in here ++ go suf rows nameKey penv grp combs rest + | Just (obsName, meanExpr, distExpr, obsCol) <- observeFull e -> + ObsInstance (obsName <> suf) meanExpr distExpr obsCol nameKey penv rows grp combs + : go suf rows nameKey penv grp combs rest + | otherwise -> go suf rows nameKey penv grp combs rest + +-- =========================================================================== +-- model graph plate aggregation (Phase 27.5 後続 TODO3、 2026-06-02) +-- =========================================================================== +-- +-- forEachGroup は群ごとに内部 do-block を展開するため、 @buildModelGraph@ が +-- 見る realized ModelP には alpha_1 / alpha_2 / alpha_3 … と群数ぶんの latent / +-- observe ノードが並ぶ (3 群 × 数 latent で 40 ノード級に肥大)。 PyMC 流の +-- plate 表記では「群コピーを 1 つの代表ノードに畳み、 箱のラベルに群数」 を +-- 出すので、 ここでは AST (= forEachGroup 構造が残る層) から +-- +-- * realized 名 → base 名 の rename map ('plateRenameMap') +-- * forEachGroup ごとの plate (ラベル "<群列> (<群数>)" + 直下 base 名) +-- ('collectGraphPlates') +-- +-- を導き、 realized ModelGraph を base 名へ collapse する +-- ('collapsePlateGraph')。 rename は interpStmts と同じ groupSuffix 規律を +-- replay して作る total な対応なので、 文字列推測ではない。 + +-- | [日本語]: model graph 上の plate (= forEachGroup 1 サイト)。 frontend +-- hgg DAGPlate (label + member id 群) に対応。 +-- [English]: A plate in the model graph (= one forEachGroup site). +-- Corresponds to the frontend hgg DAGPlate (label + set of +-- member ids). +data GraphPlate = GraphPlate + { gpLabel :: Text -- ^ [日本語]: "<群列> (<群数>)" [English]: " ()" + , gpMembers :: [Text] -- ^ [日本語]: この plate 直下の base 名 (latent + observe) [English]: the base names directly under this plate (latent + observe) + } deriving (Show, Eq) + +-- | [日本語]: realized node 名 (alpha_1 等) → base 名 (alpha) の rename map。 +-- 'collectObsInstances' / 'interpStmts' と同じ suffix (groupSuffix) / +-- 行 (rowsForGroup) 規律を replay し、 各 DoBind latent / observe を +-- その出現 suffix 付き名 → base 名 で登録する。 top-level は suffix="" なので +-- 恒等 (alpha → alpha)。 +-- [English]: A rename map from realized node names (e.g. alpha_1) to base +-- names (alpha). Replays the same suffix (groupSuffix) \/ rows +-- (rowsForGroup) rule as 'collectObsInstances' \/ 'interpStmts', and +-- registers each DoBind latent \/ observe under its occurrence's +-- suffixed name → base name. At the top level, suffix="" so it's the +-- identity (alpha → alpha). +plateRenameMap :: DataMap -> [DoStmt] -> Map.Map Text Text +plateRenameMap dm = go "" Nothing + where + go :: Text -> Maybe [Int] -> [DoStmt] -> Map.Map Text Text + go _ _ [] = Map.empty + go suf rows (st : rest) = case st of + DoBind name _ -> + Map.insert (name <> suf) name (go suf rows rest) + DoLet _ -> go suf rows rest + DoExpr e + | Just (gcol, _param, inner) <- matchForEachGroup e -> + let gvals = groupValsIn dm gcol rows + inners = Map.unions + [ go (suf <> groupSuffixFor (Map.lookup gcol dm) gv) + (Just (rowsForGroup dm gcol gv rows)) inner + | gv <- gvals ] + in inners `Map.union` go suf rows rest + | Just (obsName, _meanE, distExpr, obsCol) <- observeFull e -> + -- observe は col 参照を含むと @observeColumns@ で per-row 展開され、 + -- 実ノードは "_" (j = 当該 plate 対象行の 0 始まり + -- 連番、 = execObserve の規律) になる。 col 参照無しなら単一 + -- ""。 どちらも base 名 obsName に畳む。 + let colLen = length (lookupDoubles obsCol dm) + baseRows = fromMaybe [0 .. colLen - 1] rows + validCount = length [ i | i <- baseRows, i >= 0, i < colLen ] + names | hasColRef distExpr = + [ obsName <> suf <> "_" <> T.pack (show j) + | j <- [0 .. validCount - 1] ] + | otherwise = [ obsName <> suf ] + in Map.union (Map.fromList [ (nm, obsName) | nm <- names ]) + (go suf rows rest) + | otherwise -> go suf rows rest + +-- | [日本語]: forEachGroup ごとに 1 plate を集める (群値ぶんは展開しない)。 ラベルは +-- "<群列> (<群数>)"、 member は当該 forEachGroup 直下の base 名 (latent + +-- observe、 ネストした forEachGroup の中身は含めない = ネストは別 plate)。 +-- ネスト plate は代表 1 群の行で再帰的に拾う。 +-- [English]: Collects one plate per forEachGroup (not expanded per group +-- value). The label is " ()", and the members +-- are the base names directly under that forEachGroup (latent + observe; +-- the contents of a nested forEachGroup are not included — a nest is a +-- separate plate). Nested plates are collected recursively using one +-- representative group's rows. +collectGraphPlates :: DataMap -> [DoStmt] -> [GraphPlate] +collectGraphPlates dm = go Nothing + where + go :: Maybe [Int] -> [DoStmt] -> [GraphPlate] + go _ [] = [] + go rows (st : rest) = case st of + DoExpr e + | Just (gcol, _param, inner) <- matchForEachGroup e -> + let gvals = groupValsIn dm gcol rows + count = length gvals + label = gcol <> " (" <> T.pack (show count) <> ")" + members = directMembers inner + nested = case gvals of + (gv : _) -> go (Just (rowsForGroup dm gcol gv rows)) inner + [] -> [] + in GraphPlate label members : nested ++ go rows rest + | otherwise -> go rows rest + _ -> go rows rest + -- 直下の DoBind latent + observe 名 (ネスト forEachGroup は DoExpr なので除外)。 + directMembers :: [DoStmt] -> [Text] + directMembers stmts = + [ name | DoBind name _ <- stmts ] + ++ [ obsName | DoExpr e <- stmts, Just (obsName, _, _, _) <- [observeFull e] ] + +-- | [日本語]: realized 'HBM.ModelGraph' を plate 単位に collapse。 群展開ノードを base 名に +-- 畳み (重複ノードは初出を残す)、 辺は両端を rename して自己ループ除去 + 重複 +-- 除去。 併せて plate 一覧を返す。 +-- [English]: Collapses the realized 'HBM.ModelGraph' down to plate units. +-- Folds group-expanded nodes into base names (keeping the first +-- occurrence for duplicate nodes); renames both endpoints of each edge, +-- removing self-loops and duplicates. Also returns the plate list. +collapsePlateGraph + :: DataMap -> [DoStmt] -> HBM.ModelGraph -> (HBM.ModelGraph, [GraphPlate]) +collapsePlateGraph dm stmts mg = + let rn = plateRenameMap dm stmts + ren x = Map.findWithDefault x x rn + nodes' = dedupNodes [ renameNode ren n | n <- HBM.mgNodes mg ] + edges' = nub [ (ren a, ren b) + | (a, b) <- HBM.mgEdges mg, ren a /= ren b ] + -- Phase 40 merge: ModelGraph に mgPlates (plate→size) フィールドが追加された。 + -- DSL forEachGroup の collapse は独自の GraphPlate 列 (collectGraphPlates) を + -- 別途返すので、 ここでは入力 graph の mgPlates をそのまま引き継ぐ。 + in (HBM.ModelGraph nodes' edges' (HBM.mgPlates mg), collectGraphPlates dm stmts) + where + renameNode ren n = n + { HBM.nodeName = ren (HBM.nodeName n) + , HBM.nodeDeps = Set.map ren (HBM.nodeDeps n) + } + dedupNodes = goD Set.empty + where + goD _ [] = [] + goD seen (n : ns) + | HBM.nodeName n `Set.member` seen = goD seen ns + | otherwise = n : goD (Set.insert (HBM.nodeName n) seen) ns + +-- | [日本語]: observe ごと × 列ごとに 64 点 curve を計算。 GLMM overlay: +-- top-level observe に加え forEachGroup 内の per-group observe も対象 +-- ('collectObsInstances' が plate 展開)。 +-- [English]: Computes a 64-point curve per observe × per column. For the +-- GLMM overlay: targets per-group observes inside forEachGroup in +-- addition to top-level observes ('collectObsInstances' does the plate +-- expansion). +-- | - [日本語]: 主 predictor = 当該列、 grid は (per-group なら群の) data min..max +-- [English]: The primary predictor = this column; the grid spans the +-- data's (per-group: that group's) min..max. +-- | - [日本語]: 他 predictor は (per-group なら群の) data median で固定 +-- [English]: Other predictors are held fixed at the data's (per-group: +-- that group's) median. +-- | - [日本語]: 群 latent は suffix 付きキー (theta_1 等)、 群変数は定数として env に注入 +-- [English]: Group latents use their suffixed key (e.g. theta_1); +-- group variables are injected into the env as constants. +-- | - [日本語]: 各 sample × 各 grid 点で mean 式を Double 評価 +-- [English]: The mean expression is evaluated as a Double at each +-- sample × each grid point. +-- | - [日本語]: 各 grid 点で全 sample から median + 2.5% / 97.5% percentile +-- [English]: At each grid point, the median + 2.5% \/ 97.5% +-- percentiles are taken across all samples. +computeMeanCurves + :: [TopBind] -- top-level 値/関数束縛 (ユーザ定義リンク等) + -> [DoStmt] + -> DataMap -- data: col → values + -> [Map.Map Text Double] -- posterior samples + -> [HbmMeanCurve] +computeMeanCurves topBinds stmts dataMap samples = + [ HbmMeanCurve + { hmcObsName = oiObsName inst + , hmcPredictor = col + , hmcX = xGrid + , hmcMedian = map (percentileOf 0.5) valuesPerX + , hmcLower = map (percentileOf 0.025) valuesPerX + , hmcUpper = map (percentileOf 0.975) valuesPerX + , hmcGroupCol = fst <$> oiGroup inst + , hmcGroupVal = snd <$> oiGroup inst + } + | inst <- collectObsInstances dataMap stmts + , let meanExpr = oiMeanExpr inst + mrows = oiRows inst + -- per-group なら群の行に限定して列値を取り出す。 + colValsFor c = + let ca = lookupDoubles c dataMap + in case mrows of + Nothing -> ca + Just rs -> [ ca !! i | i <- rs, i >= 0, i < length ca ] + , col <- collectCols meanExpr + , let xs = colValsFor col + , not (null xs) + , let nGrid = 64 :: Int + xLo = minimum xs + xHi = maximum xs + step = if xHi == xLo then 1.0 else (xHi - xLo) / fromIntegral (nGrid - 1) + xGrid = [ xLo + step * fromIntegral i | i <- [0 .. nGrid - 1] ] + otherCols = filter (/= col) (collectCols meanExpr) + medianOf vs = case sort vs of + [] -> 0 + ss -> ss !! (length ss `div` 2) + fixedOther = + Map.fromList [ (c, Numeric [medianOf (colValsFor c)]) | c <- otherCols ] + -- 群変数 (forEachGroup の \g) を定数として env に注入。 + paramVNums = Map.map VNum (oiParamEnv inst) + -- sample を nameKey で remap: local latent 名 → samples の suffix 付きキー。 + -- (Map.union は left-biased なので renamed が元キーより優先) + remapSample sample = + let renamed = Map.fromList + [ (localNm, v) + | (localNm, key) <- Map.toList (oiNameKey inst) + , Just v <- [Map.lookup key sample] ] + in Map.union renamed sample + evalAt sample x = + let synthetic = Map.insert col (Numeric [x]) fixedOther + sm = remapSample sample + -- posterior サンプル (alpha/beta 等) を env に、 top-level + -- 値/関数 (ユーザ定義リンク等) + 群変数も併せて見えるようにする。 + senv = Map.union paramVNums + (Map.union (Map.map VNum sm) (buildTopEnv dataMap topBinds)) + in case evalScalar @Double senv synthetic (Just 0) meanExpr of + Right v -> v + Left _ -> 0 / 0 -- NaN + valuesPerX = [ [ evalAt sample x | sample <- samples ] | x <- xGrid ] + ] + +-- --------------------------------------------------------------------------- +-- Phase 27.5 (2026-06-02): WAIC / LOO / posterior predictive 用の observe +-- distribution 評価。 computeMeanCurves と同じ plate 展開 (collectObsInstances) +-- を使い、 mean 式ではなく distribution 式全体を各 sample × 各対象行で評価する。 +-- これにより pointwise log-likelihood (WAIC/LOO) と posterior predictive draw +-- (PPC、 worker 側で sampleDist) の共通基盤を 1 度で作る。 +-- --------------------------------------------------------------------------- + +-- | [日本語]: observe instance 1 個分の、 全 posterior sample × 対象行で評価した結果。 +-- [English]: The result of evaluating one observe instance across all +-- posterior samples × target rows. +data ObsDistSet = ObsDistSet + { odsName :: !Text -- observe ノード名 (suffix 付き) + , odsObserved :: ![Double] -- 対象行の観測値 (= col の当該行) + , odsDists :: ![[HBM.Distribution Double]] -- [sample][row] の Distribution + } + +-- | [日本語]: 各 observe instance について、 各 posterior sample × 各対象行で +-- distribution を評価する。 mean に列参照がある GLM 形 (例: Normal (a+b*@x@) s) +-- も evalDist の per-row 評価 ((Just i)) で正しく行ごとに展開される。 +-- [English]: For each observe instance, evaluates the distribution at +-- each posterior sample × each target row. GLM forms whose mean contains +-- a column reference (e.g. Normal (a+b*@x@) s) are also correctly +-- expanded per row via evalDist's per-row evaluation ((Just i)). +computeObsDists + :: [TopBind] + -> [DoStmt] + -> DataMap + -> [Map.Map Text Double] + -> [ObsDistSet] +computeObsDists topBinds stmts dataMap samples = + [ ObsDistSet + { odsName = oiObsName inst + , odsObserved = ys + , odsDists = [ [ distAt sample i | i <- rows ] | sample <- samples ] + } + | inst <- collectObsInstances dataMap stmts + , let distExpr = oiDistExpr inst + yCol = lookupDoubles (oiObsCol inst) dataMap + allRows = [0 .. length yCol - 1] + rows = filter (\i -> i >= 0 && i < length yCol) + (fromMaybe allRows (oiRows inst)) + ys = [ yCol !! i | i <- rows ] + paramVNums = Map.map VNum (oiParamEnv inst) + remapSample sample = + let renamed = Map.fromList + [ (localNm, v) + | (localNm, key) <- Map.toList (oiNameKey inst) + , Just v <- [Map.lookup key sample] ] + in Map.union renamed sample + distAt sample i = + let sm = remapSample sample + base = Map.union paramVNums + (Map.union (Map.map VNum sm) (buildTopEnv dataMap topBinds)) + -- Phase 43: combinator 由来 latent vector (cuts/probs) を sampled + -- latent から再構築して VList 束縛 (= 消費分布の list 引数を解決)。 + combBinds = [ (bindNm, VList (map VNum vals)) + | (bindNm, spec) <- oiListCombs inst + , Just vals <- [reconstructComb base dataMap sample spec] ] + senv = Map.union (Map.fromList combBinds) base + in case evalDist senv dataMap (Just i) distExpr :: Err (HBM.Distribution Double) of + Right d -> d + Left _ -> HBM.Normal (0 / 0) 1 -- 評価不能は NaN 化 (logDensity→NaN) + , not (null rows) + ] + +-- | [日本語]: 'ObsDistSet' 群から WAIC/LOO 用の log-likelihood 行列 (S × N) を作る。 +-- 行 = posterior sample、 列 = 全 observe instance の全対象行を連結。 +-- @Hanalyze.Stat.ModelSelect.waic@ / @loo@ がこの shape を期待する。 +-- [English]: Builds the log-likelihood matrix (S × N) for WAIC\/LOO from +-- a set of 'ObsDistSet's. Rows = posterior samples; columns = all target +-- rows of all observe instances, concatenated. +-- @Hanalyze.Stat.ModelSelect.waic@ \/ @loo@ expect this shape. +pointwiseLogLik :: [ObsDistSet] -> [[Double]] +pointwiseLogLik sets = + [ concatMap (\set -> zipWith HBM.logDensity (sampleRow set s) (odsObserved set)) sets + | s <- [0 .. nSamples - 1] ] + where + nSamples = case sets of + (set : _) -> length (odsDists set) + [] -> 0 + sampleRow set s = case drop s (odsDists set) of + (row : _) -> row + [] -> [] + +-- | [日本語]: log-lik 行列 (S×N) から非有限 (NaN / ±Inf) を含む観測列を除外する。 +-- @distAt@ の eval 失敗 (= Normal NaN) や +-- 退化パラメータで @logDensity@ が NaN / Inf になった観測点は、 そのまま waic/loo に +-- 渡すと per-observation 集計 (lppd / pwaic) を汚染して全体が NaN → JSON null に +-- なる。 該当する観測列 (= 全 sample で同一観測点) を丸ごと落として残りで waic/loo を +-- 計算できるようにする。 返り値は (除外後行列, 落とした列数)。 全列が非有限なら +-- ([], N) を返し、 呼び出し側は waic/loo を Nothing にできる。 +-- [English]: Excludes observation columns containing non-finite values +-- (NaN \/ ±Inf) from the log-lik matrix (S×N). Observation points where +-- @distAt@ fails to evaluate (= Normal NaN) or @logDensity@ becomes NaN \/ +-- Inf due to degenerate parameters would, if passed straight into +-- waic\/loo, contaminate the per-observation aggregates (lppd \/ pwaic) +-- and turn the whole result into NaN → JSON null. This drops the affected +-- observation columns (= the same observation point across all samples) +-- entirely, so waic\/loo can be computed from what remains. The return +-- value is (the matrix after exclusion, the number of dropped columns). +-- If every column is non-finite, returns ([], N), letting the caller turn +-- waic\/loo into Nothing. +finitePointwiseLogLik :: [[Double]] -> ([[Double]], Int) +finitePointwiseLogLik mat = + let cols = transpose mat -- [obs点][sample] + keptCols = filter (all isFiniteD) cols + dropped = length cols - length keptCols + in (transpose keptCols, dropped) + where + isFiniteD x = not (isNaN x || isInfinite x) + +-- =========================================================================== +-- Phase 44: multi-column observe (observeMV) の WAIC / PPC 経路 +-- =========================================================================== +-- +-- scalar observe (1 列) の WAIC/PPC ('collectObsInstances' / 'computeObsDists' / +-- 'pointwiseLogLik') は per-row スカラ logDensity 前提で、 multi-column の +-- k-vector joint density (= 'HBM.obsLogSum') を扱えない。 そこで Phase 44 の +-- 設計方針 (observeMV は scalar observe と別 builtin の並行経路) を WAIC/PPC まで +-- 貫き、 **MV 専用の並行経路** を足す (scalar 経路は無傷)。 +-- +-- latent Σ (lkjCorrCholesky 由来の相関 Cholesky L) は posterior sample に +-- 内部 latent (@L_u_@ = partial-correlation の Beta latent) として載るので、 +-- bind 名 @L@ を 'reconstructMatrixComb' で再構築して MvNormalChol に渡す +-- (Phase 43 'reconstructComb' の行列版)。 + +-- | [日本語]: 行列値 combinator の再構築仕様 (suffix 付き base 名 + 次元 k)。 +-- [English]: The reconstruction spec for a matrix-valued combinator +-- (suffixed base name + dimension k). +data MatrixCombSpec + = LkjCholSpec Text Int -- ^ [日本語]: suffix 付き base 名, k (= 行列次元) [English]: suffixed base name, k (= the matrix dimension) + deriving (Show) + +-- | [日本語]: DoBind の RHS が行列値 combinator (lkjCorrCholesky) なら 'MatrixCombSpec' を +-- 返す ('matchMatrixComb' / 'matrixCombSuffix' と同じ命名規律)。 +-- [English]: If the DoBind's RHS is a matrix-valued combinator +-- (lkjCorrCholesky), returns a 'MatrixCombSpec' (following the same +-- naming rule as 'matchMatrixComb' \/ 'matrixCombSuffix'). +matchMatrixCombSpec :: Text -> Expr -> Maybe MatrixCombSpec +matchMatrixCombSpec suf rhs = case collectApp rhs of + Right ("lkjCorrCholesky", [ELit (LText bn), ELit (LNumber d), _etaE]) -> + Just (LkjCholSpec (bn <> suf) (round d)) + _ -> Nothing + +-- | [日本語]: lkjCorrCholesky の相関 Cholesky L を 1 posterior サンプルから再構築する。 +-- sampled latent は @_u_@ (Beta in (0,1))、 partial correlation は +-- @z_ij = 2u - 1@。 L は @HBM.lkjCorrCholesky@ の deterministic 構築を replay: +-- L_00 = 1、 対角 L_ii = √(1 - Σ_{k_u_@ (Beta in (0,1)); the partial correlation is +-- @z_ij = 2u - 1@. L replays 'HBM.lkjCorrCholesky''s deterministic +-- construction: L_00 = 1, the diagonal L_ii = √(1 - Σ_{k MatrixCombSpec -> Maybe [[Double]] +reconstructMatrixComb sm (LkjCholSpec nm k) = do + let uKey i j = nm <> "_u" <> T.pack (show i) <> "_" <> T.pack (show j) + pcPairs <- mapM + (\(i, j) -> do u <- Map.lookup (uKey i j) sm; pure ((i, j), 2 * u - 1)) + [(i, j) | i <- [1 .. k - 1], j <- [0 .. i - 1]] + let pcMap = Map.fromList pcPairs + pc i j = Map.findWithDefault 0 (i, j) pcMap + sq z = z * z + lRow i = + [ if j > i then 0 + else if i == 0 && j == 0 then 1 + else if j == i + then sqrt (max 0 (1 - sum [ sq (pc i kk) | kk <- [0 .. i - 1] ])) + else pc i j * sqrt (max 0 (product [ 1 - sq (pc i kk) | kk <- [0 .. j - 1] ])) + | j <- [0 .. k - 1] ] + pure [ lRow i | i <- [0 .. k - 1] ] + +-- | [日本語]: observeMV 式から (名前, distribution 式全体, 観測列名リスト) を抽出。 +-- [English]: Extracts (name, the whole distribution expression, the list +-- of observed column names) from an observeMV expression. +observeMVFull :: Expr -> Maybe (Text, Expr, [Text]) +observeMVFull e = case collectApp e of + Right ("observeMV", [ELit (LText nm), distExpr, EList colRefs]) -> + let cols = [ c | ECol c <- colRefs ] + in if length cols == length colRefs && length cols >= 2 + then Just (nm, distExpr, cols) else Nothing + _ -> Nothing + +-- | [日本語]: observeMV instance (plate 展開済)。 'collectObsInstances' の MV 版で、 +-- 単一 'oiObsCol' でなく __列リスト__ を持ち、 list/matrix combinator の +-- 再構築仕様も保持する。 +-- [English]: An observeMV instance (plate-expanded). The MV counterpart of +-- 'collectObsInstances': instead of a single 'oiObsCol', it holds +-- __a list of columns__, and also retains the reconstruction specs for +-- list \/ matrix combinators. +data MvObsInstance = MvObsInstance + { mviObsName :: Text + , mviDistExpr :: Expr + , mviObsCols :: [Text] + , mviNameKey :: Map.Map Text Text + , mviParamEnv :: Map.Map Text Double + , mviRows :: Maybe [Int] + , mviListCombs :: [(Text, ListCombSpec)] + , mviMatrixCombs :: [(Text, MatrixCombSpec)] + } + +-- | [日本語]: stmts を plate 展開しながら observeMV instance を集める +-- ('collectObsInstances' と同じ suffix / 行 / 群変数 / latent 名規律を replay)。 +-- [English]: Collects observeMV instances while plate-expanding the stmts +-- (replaying the same suffix \/ row \/ group-variable \/ latent-naming +-- rule as 'collectObsInstances'). +collectMvObsInstances :: DataMap -> [DoStmt] -> [MvObsInstance] +collectMvObsInstances dm = go "" Nothing Map.empty Map.empty [] [] + where + go :: Text -> Maybe [Int] -> Map.Map Text Text -> Map.Map Text Double + -> [(Text, ListCombSpec)] -> [(Text, MatrixCombSpec)] -> [DoStmt] + -> [MvObsInstance] + go _ _ _ _ _ _ [] = [] + go suf rows nameKey penv lcombs mcombs (st : rest) = case st of + DoBind name rhs -> + let lcombs' = case matchListCombSpec suf rhs of + Just spec -> (name, spec) : lcombs + Nothing -> lcombs + mcombs' = case matchMatrixCombSpec suf rhs of + Just spec -> (name, spec) : mcombs + Nothing -> mcombs + in go suf rows (Map.insert name (name <> suf) nameKey) penv lcombs' mcombs' rest + DoLet _ -> go suf rows nameKey penv lcombs mcombs rest + DoExpr e + | Just (gcol, param, inner) <- matchForEachGroup e -> + let gvals = groupValsIn dm gcol rows + here = concatMap + (\gv -> + go (suf <> groupSuffixFor (Map.lookup gcol dm) gv) + (Just (rowsForGroup dm gcol gv rows)) + nameKey (Map.insert param gv penv) lcombs mcombs inner) + gvals + in here ++ go suf rows nameKey penv lcombs mcombs rest + | Just (obsName, distExpr, cols) <- observeMVFull e -> + MvObsInstance (obsName <> suf) distExpr cols nameKey penv rows lcombs mcombs + : go suf rows nameKey penv lcombs mcombs rest + | otherwise -> go suf rows nameKey penv lcombs mcombs rest + +-- | [日本語]: observeMV 1 個分の WAIC/PPC 評価結果。 'ObsDistSet' の MV 版。 +-- [English]: The WAIC\/PPC evaluation result for one observeMV. The MV +-- counterpart of 'ObsDistSet'. +data MvObsDistSet = MvObsDistSet + { mvodsName :: !Text -- ^ [日本語]: 表示名 (top="y"、 group="y_1") [English]: the display name (top="y", group="y_1") + , mvodsCols :: ![Text] -- ^ [日本語]: 観測列名リスト (長さ k) [English]: the list of observed column names (length k) + , mvodsObserved :: ![[Double]] -- ^ [日本語]: [row][component] = 各行の k-vector [English]: [row][component] = each row's k-vector + , mvodsDists :: ![[HBM.Distribution Double]] -- ^ [日本語]: [sample][row] の Distribution [English]: the Distribution at [sample][row] + } + +-- | [日本語]: observeMV instance × sample × 行で多変量 Distribution を評価する +-- ('computeObsDists' の MV 版)。 dist は行不変 (μ/Σ は latent) なので各行で +-- 同一だが、 既存経路に合わせ row 評価する。 latent Σ は 'reconstructMatrixComb' +-- で L を、 list 引数は 'reconstructComb' で復元して env に束縛する。 +-- [English]: Evaluates the multivariate Distribution at each observeMV +-- instance × sample × row (the MV counterpart of 'computeObsDists'). The +-- dist is row-invariant (μ\/Σ are latents), so it's the same for every +-- row, but evaluation is done per row to match the existing path. The +-- latent Σ's L is restored via 'reconstructMatrixComb', and list +-- arguments are restored via 'reconstructComb', then bound into the env. +computeMvObsDists + :: [TopBind] + -> [DoStmt] + -> DataMap + -> [Map.Map Text Double] + -> [MvObsDistSet] +computeMvObsDists topBinds stmts dataMap samples = + [ MvObsDistSet + { mvodsName = mviObsName inst + , mvodsCols = cols + , mvodsObserved = [ [ lookupDoubles c dataMap !! i | c <- cols ] | i <- rows ] + , mvodsDists = [ [ distAt sample i | i <- rows ] | sample <- samples ] + } + | inst <- collectMvObsInstances dataMap stmts + , let distExpr = mviDistExpr inst + cols = mviObsCols inst + colLens = map (length . (`lookupDoubles` dataMap)) cols + colLen = if null colLens then 0 else minimum colLens + allRows = [0 .. colLen - 1] + rows = filter (\i -> i >= 0 && i < colLen) + (fromMaybe allRows (mviRows inst)) + paramVNums = Map.map VNum (mviParamEnv inst) + remapSample sample = + let renamed = Map.fromList + [ (localNm, v) + | (localNm, key) <- Map.toList (mviNameKey inst) + , Just v <- [Map.lookup key sample] ] + in Map.union renamed sample + distAt sample i = + let sm = remapSample sample + base = Map.union paramVNums + (Map.union (Map.map VNum sm) (buildTopEnv dataMap topBinds)) + listBinds = [ (bn, VList (map VNum vals)) + | (bn, spec) <- mviListCombs inst + , Just vals <- [reconstructComb base dataMap sample spec] ] + matBinds = [ (bn, VList (map (VList . map VNum) m)) + | (bn, spec) <- mviMatrixCombs inst + , Just m <- [reconstructMatrixComb sample spec] ] + senv = Map.union (Map.fromList (listBinds ++ matBinds)) base + in case evalDist senv dataMap (Just i) distExpr :: Err (HBM.Distribution Double) of + Right d -> d + Left _ -> HBM.MvNormal [0 / 0] [[1]] -- eval 不能は NaN 化 + , not (null rows) + ] + +-- | [日本語]: MV observe の pointwise log-lik 行列 (S×N)。 各行 (= 1 観測点) の寄与は +-- k-vector joint density 'HBM.obsLogSum'。 scalar 経路の 'pointwiseLogLik' と +-- 列方向に連結して使う (worker 側)。 +-- [English]: The pointwise log-lik matrix (S×N) for MV observe. Each row's +-- (= one observation point's) contribution is the k-vector joint density +-- 'HBM.obsLogSum'. Used by concatenating column-wise with the scalar +-- path's 'pointwiseLogLik' (on the worker side). +pointwiseLogLikMv :: [MvObsDistSet] -> [[Double]] +pointwiseLogLikMv sets = + [ concatMap (\set -> [ HBM.obsLogSum (distRow set s !! r) (mvodsObserved set !! r) + | r <- [0 .. nRows set - 1] ]) sets + | s <- [0 .. nSamples - 1] ] + where + nSamples = case sets of + (set : _) -> length (mvodsDists set) + [] -> 0 + nRows set = length (mvodsObserved set) + distRow set s = case drop s (mvodsDists set) of + (row : _) -> row + [] -> [] diff --git a/src/Hanalyze/Model/HBM/Model.hs b/src/Hanalyze/Model/HBM/Model.hs new file mode 100644 --- /dev/null +++ b/src/Hanalyze/Model/HBM/Model.hs @@ -0,0 +1,1611 @@ +{-# LANGUAGE RankNTypes #-} +{-# LANGUAGE DeriveFunctor #-} +{-# LANGUAGE FlexibleInstances #-} +{-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE ImpredicativeTypes #-} + +-- | +-- Module : Hanalyze.Model.HBM.Model +-- Description : HBM の多相モデル DSL (Free monad) 記述層 +-- Copyright : (c) 2026 Aelysce Project (Toshiaki Honda) +-- License : BSD-3-Clause +-- +-- [日本語]: 多相モデル DSL (Free monad) を 'Hanalyze.Model.HBM' から分離した。 +-- +-- 本モジュールは PPL の __記述層__ を担う: +-- +-- - @Free@ monad 再実装 (型は 'Hanalyze.Model.HBM' 公開のものと別個) +-- - 'ModelF' プリミティブ (sample / observe / observeLM / deterministic / +-- plate / Data / Potential) と 'Model' / 'ModelP' 型エイリアス +-- - 第一級ランダム効果値 'REffect' / 'REff' と階層モデル helper 群 +-- (reNormal / mvNormalLatent / lkjCorrCholesky / ar1Latent / dirichlet / +-- orderedCuts / dpStickBreaking / hmmLatent / glmmRandomIntercept 等) +-- - Plate notation と構造検査 ('collectNodes' / 'sampleNames') +-- +-- 評価 (logJoint 等)・AD 勾配・IR は __上層__ に置かれ、 本モジュールは +-- それらに依存しない (leaf-first・facade 非 import の規律。 計画参照)。 +-- 依存は下層 'Hanalyze.Model.HBM.Util' / '...Distribution' のみ。 +-- +-- [English]: The polymorphic model DSL (Free monad) was split out of +-- 'Hanalyze.Model.HBM'. +-- +-- This module owns the PPL's __description layer__: +-- +-- - The reimplemented @Free@ monad (a type distinct from the one +-- publicly exposed by 'Hanalyze.Model.HBM') +-- - The 'ModelF' primitives (sample \/ observe \/ observeLM \/ +-- deterministic \/ plate \/ Data \/ Potential) and the 'Model' \/ +-- 'ModelP' type aliases +-- - First-class random-effect values 'REffect' \/ 'REff' and the +-- hierarchical-model helper family (reNormal \/ mvNormalLatent \/ +-- lkjCorrCholesky \/ ar1Latent \/ dirichlet \/ orderedCuts \/ +-- dpStickBreaking \/ hmmLatent \/ glmmRandomIntercept, etc.) +-- - Plate notation and structural inspection ('collectNodes' \/ +-- 'sampleNames') +-- +-- Evaluation (logJoint etc.), AD gradients, and the IR live in the +-- __upper layer__; this module does not depend on them (leaf-first, +-- no-facade-import discipline; see the plan). Its only dependencies +-- are the lower-layer 'Hanalyze.Model.HBM.Util' \/ +-- '...Distribution'. +module Hanalyze.Model.HBM.Model + ( -- * Free monad + Free (..) + , liftF + -- * Polymorphic model DSL + , ModelF (..) + , Model + , ModelP + , sample + , observe + , observeMV + , observeColumns + , observeLM + , observeLMR + , observeNormalLM + , LMFamily (..) + , lmFamilyName + , lmParents + , REff (..) + , REffect (..) + , reffNames + , reNormal + , at + , indexed + , (.#) + , potential + , deterministic + , nonCenteredNormal + , dirichlet + , orderedCuts + , dpStickBreaking + , hmmLatent + , hmmForwardLogLik + , GlmmFamily (..) + , glmmRandomIntercept + , dataNamed + , dataNamedX + , dataNamedIx + , dataNamedObs + , Ix (..) + , TrackTag (..) + , (!!!) + , atIx + , withData + , withDataIx + , mvNormalLatent + , lkjCorrCholesky + , gpExpQuadCov + , gpLatent + , ar1Latent + -- ** plate notation + , plate + , plateI + , plateI_ + , plateForM + , plateForM_ + , withPlate + -- * Structural inspection + , Node (..) + , NodeKind (..) + , collectNodes + , sampleNames + , dataSlots + , dataIxSlots + ) where + +import Control.DeepSeq (NFData (..)) +import Control.Monad (forM, forM_) +import Data.List (foldl', nub) +import qualified Data.Set as Set +import Data.Set (Set) +import Data.Text (Text) +import qualified Data.Text as T + +import Hanalyze.Model.HBM.Util (negInf, logSumExpA, choleskyL, hmmForwardLogLik) +import Hanalyze.Model.HBM.Distribution + +-- --------------------------------------------------------------------------- +-- @Free@ monad (再実装。Hanalyze.Model.HBM のものとは型が違うので別途定義) +-- --------------------------------------------------------------------------- + +data Free f a = Pure a | Free (f (Free f a)) + +instance Functor f => Functor (Free f) where + fmap g (Pure a) = Pure (g a) + fmap g (Free x) = Free (fmap (fmap g) x) + +instance Functor f => Applicative (Free f) where + pure = Pure + Pure g <*> x = fmap g x + Free fg <*> x = Free (fmap (<*> x) fg) + +instance Functor f => Monad (Free f) where + return = pure + Pure a >>= g = g a + Free x >>= g = Free (fmap (>>= g) x) + +liftF :: Functor f => f a -> Free f a +liftF fa = Free (fmap Pure fa) + +-- --------------------------------------------------------------------------- +-- 多相モデル (@Free@ monad) +-- --------------------------------------------------------------------------- + +-- | [日本語]: DSL のプリミティブ。継続が @a -> next@ なので任意の @a@ を流せる。 +-- +-- 'Potential' は PyMC の @pm.Potential@ 相当で、任意の log-prob 項を +-- log-joint に加える。ソフト制約・カスタム尤度・正則化項などに使える。 +-- +-- [English]: The DSL's primitives. Since the continuation is +-- @a -> next@, any @a@ can be threaded through. +-- +-- 'Potential' is the equivalent of PyMC's @pm.Potential@: it adds an +-- arbitrary log-prob term to the log-joint. Useful for soft +-- constraints, custom likelihoods, regularization terms, and the like. +-- | [日本語]: 構造化線形予測子 observe の family / link。 +-- +-- 通常の 'Observe' は平均が不透明な AD 値ゆえ「β に線形」 という構造を +-- 保持できない。 'ObserveLM' は設計行列 X (Double) と β パラメタ名を __分離__ +-- して持つことで線形構造をライブラリが知り、 Gaussian-恒等リンクの +-- 十分統計量 collapse (観測和を tape O(p²) に畳む) を可能にする。 +-- +-- [English]: The family \/ link for the structured linear-predictor +-- observe. +-- +-- The regular 'Observe' cannot preserve the "linear in β" structure, +-- since its mean is an opaque AD value. 'ObserveLM' keeps the design +-- matrix X (Double) and the β parameter names __separate__, so the +-- library knows the linear structure, enabling the collapse of the +-- sufficient statistic for the Gaussian-identity link (folding the +-- observation sum into an O(p²) tape). +data LMFamily + = LMGaussian Text -- ^ [日本語]: identity link。 引数 = σ (誤差 SD) パラメタ名。 [English]: identity link; the argument is the σ (residual SD) parameter name. + | LMPoisson -- ^ [日本語]: log link (μ = exp η)。 [English]: log link (μ = exp η). + | LMBernoulli -- ^ [日本語]: logit link (p = 1/(1+e^{-η}))。 [English]: logit link (p = 1/(1+e^{-η})). + deriving (Show, Eq) + +-- | [日本語]: 'ObserveLM' のランダム効果項。 線形予測子に +-- @η_i += u^{re}[gid_i]@ を gather で加える。 設計行列の one-hot 指示列として +-- 密に展開する代わりに、 群 id ベクトルで疎に保持することで vec-tape の +-- 観測尤度勾配が群効果に対しても O(n) で済む (密展開は O(nG·n) で階層モデルで +-- 逆効果になる・計測で確認済)。 +-- +-- フィールド: u パラメタ名 (長さ nG・既に 'sample' 済の latent を参照) / +-- 各観測の群 id (長さ n・0..nG-1) / +-- prior スケール名: @Just τName@ なら各 u_j が +-- @u_j ~ Normal(0, τ)@ という標準的な階層 prior を持つことを宣言する。 +-- これがあると @compileGradU@ は u-prior 勾配を __解析的に__ (ベクトル化して) +-- 計算し、 対応する @u_j@ 'Sample' ノードを @ad@ walk から除外できる +-- (per-grad の支配項だった O(nG) スカラ @ad@ を排除)。 @Nothing@ なら +-- prior は従来通り @ad@ 経路で扱う (後方互換)。 通常は 'reNormal'/'at' で +-- 自動的に @Just@ が載るので、 ユーザがこの構築子を直接書く必要はない。 +-- +-- per-row 重み: @Just ws@ (長さ n) なら @η_i += w_i·u^{re}[gid_i]@ +-- (random slope = 群別係数 × 共変量)。 @Nothing@ = 全 1 (random intercept・ +-- 後方互換)。 prior 解析勾配 (@u_j ~ Normal(0,τ)@) は重みと無関係に同形。 +-- 由来 slot 名: 5 番目 field は gids がどのデータ slot +-- ('dataNamedIx') 由来かの静的属性。 @Just slot@ なら 'lmParents' が slot 名を +-- 親集合に加え、 DAG に slot (DataN)→観測ノードのエッジが出る (PyMC +-- @b0[gid]@ 同型)。 'atIx' が自動で載せる。 'at' / IR 合成経路は @Nothing@ +-- (従来挙動)。 hot closure (@CompiledLMBlock@) には乗らない = per-draw 無影響。 +-- +-- [English]: The random-effect term of 'ObserveLM'. Adds +-- @η_i += u^{re}[gid_i]@ to the linear predictor via a gather. Instead +-- of densely expanding the design matrix into one-hot indicator +-- columns, keeping the group-id vector sparse means the vec-tape's +-- observation-likelihood gradient stays O(n) even for group effects +-- (dense expansion is O(nG·n), which backfires for hierarchical +-- models, as confirmed by measurement). +-- +-- Fields: the u parameter names (length nG; refer to latents already +-- declared via 'sample') \/ each observation's group id (length n, +-- 0..nG-1) \/ the prior scale name: @Just τName@ declares that each +-- u_j has the standard hierarchical prior @u_j ~ Normal(0, τ)@. When +-- present, @compileGradU@ computes the u-prior gradient __analytically__ +-- (vectorized) and can exclude the corresponding @u_j@ 'Sample' nodes +-- from the @ad@ walk (removing the O(nG) scalar @ad@ calls that used +-- to dominate the per-gradient cost). @Nothing@ means the prior is +-- still handled via the @ad@ path as before (backward compatible). +-- 'reNormal'\/'at' normally attach @Just@ automatically, so users +-- should not need to construct this directly. +-- +-- Per-row weights: @Just ws@ (length n) means +-- @η_i += w_i·u^{re}[gid_i]@ (random slope = per-group coefficient × +-- covariate). @Nothing@ = all 1s (random intercept, backward +-- compatible). The analytic prior gradient (@u_j ~ Normal(0,τ)@) has +-- the same shape regardless of the weights. +-- +-- Origin slot name: the 5th field is a static attribute recording +-- which data slot ('dataNamedIx') the gids came from. @Just slot@ +-- makes 'lmParents' add the slot name to the parent set, so the DAG +-- gets an edge from the slot (DataN) to the observation node (the +-- same shape as PyMC's @b0[gid]@). 'atIx' attaches this automatically. +-- 'at' \/ the IR-composition path leave it @Nothing@ (legacy behavior). +-- It is not carried by the hot closure (@CompiledLMBlock@), so it has +-- no per-draw impact. +data REff = REff [Text] [Int] (Maybe Text) (Maybe [Double]) !(Maybe Text) + deriving (Show, Eq) + +-- Phase 54.8: synthGaussLMBlocks の安全網 (force で全評価し poison を捕捉) 用。 +instance NFData REff where + rnf (REff us gids sc mw ms) = + rnf us `seq` rnf gids `seq` rnf sc `seq` rnf mw `seq` rnf ms + +data ModelF a next + = Sample Text (Distribution a) (a -> next) + | Observe Text (Distribution a) [Double] next + | ObserveLM Text [Text] [[Double]] [REff] LMFamily [Double] next + -- ^ [日本語]: 構造化線形予測子 observe。 + -- フィールド: ブロック名 / β パラメタ名 (順序 = X の列) / + -- 設計行列 X (n 行 × p 列、 Double) / ランダム効果項 (gather) / + -- family-link / 観測 ys (長さ n)。 + -- 各 i について η_i = Σ_j β_j·X_ij + Σ_re u^{re}[gid^{re}_i]、 + -- μ_i = link⁻¹(η_i)、 log-lik = Σ_i logDensityObs(family μ_i) y_i。 + -- β / u / 分散パラメタは別途 'sample' で宣言された latent を + -- __名前参照__する (prior は持たない)。 + -- DAG 上は 1 観測ノード (親 = β + u + 分散パラメタ名)。 + -- + -- [English]: The structured linear-predictor observe. + -- Fields: block name \/ β parameter names (order = X's columns) \/ + -- design matrix X (n rows × p columns, Double) \/ random-effect + -- terms (gather) \/ family-link \/ observations ys (length n). + -- For each i, η_i = Σ_j β_j·X_ij + Σ_re u^{re}[gid^{re}_i], + -- μ_i = link⁻¹(η_i), log-lik = Σ_i logDensityObs(family μ_i) y_i. + -- β \/ u \/ the variance parameter are __referenced by name__ from + -- latents declared separately via 'sample' (they carry no prior + -- here). On the DAG this is a single observation node (parents = + -- β + u + the variance parameter name). + | Potential Text a next + -- ^ [日本語]: 名前付きの ad-hoc な log-prob 項。値 @a@ がそのまま log-joint に加算される。 + -- [English]: A named ad-hoc log-prob term. The value @a@ is added + -- directly to the log-joint. + | Deterministic Text a (a -> next) + -- ^ [日本語]: 名前付きの派生量 (PyMC `pm.Deterministic`)。log-joint には寄与せず、 + -- サンプルごとに値を保存する。継続には値そのものを通すので、その後の + -- モデル中でも参照可能。 + -- + -- [English]: A named derived quantity (PyMC's `pm.Deterministic`). + -- Does not contribute to the log-joint; its value is saved for + -- each sample. Since the value itself is threaded through the + -- continuation, it can also be referenced later in the model. + | Data Text [Double] (([a], [Double]) -> next) + -- ^ [日本語]: 名前付き観測データプレースホルダ (PyMC `pm.Data`)。 + -- モデル内でデータを保持し、`withData` で外部から差し替え可能。 + -- 観測値を直接 `observe` に渡す代わりに、`dataNamed` で受け取って + -- `observe` に渡すと、後でデータ差し替えができる。 + -- ★破壊的変更 (旧バージョン比): 継続は ([a], [Double]) の 2 view を受ける + -- (格納は [Double] のまま・各 interpreter が lift)。 fst = モデル数値型 + -- ('dataNamed'、 covariate 用・realToFrac 不要)、 snd = 生 [Double] + -- ('dataNamedObs'、 'observe' の観測値用)。 tuple は lazy なので + -- 未使用側の lift コストは掛からない。 + -- + -- [English]: A named observed-data placeholder (PyMC's `pm.Data`). + -- Holds data inside the model, replaceable from the outside via + -- `withData`. Instead of passing observed values directly to + -- `observe`, receiving them via `dataNamed` and passing that to + -- `observe` allows the data to be swapped out later. + -- ★Breaking change (vs. an earlier version): the continuation now + -- receives two views, ([a], [Double]) (storage stays [Double]; + -- each interpreter lifts as needed). fst = the model's numeric + -- type ('dataNamed', for covariates, no realToFrac needed), snd = + -- raw [Double] ('dataNamedObs', for 'observe''s observed values). + -- The tuple is lazy, so the unused side incurs no lifting cost. + | DataIx Text [Int] ([Int] -> next) + -- ^ [日本語]: 離散 index 専用のデータプレースホルダ。 群 index 等の + -- 名義尺度を [Int] のまま運ぶ (= AD 型に持ち上げない・round 罠の根治)。 + -- 継続型は @a@ に依らず [Int] なので interpreter の lift も不要。 + -- + -- [English]: A data placeholder dedicated to discrete indices. + -- Carries nominal-scale values such as group indices as-is, as + -- [Int] (i.e. never lifted to the AD type — the fix for the + -- rounding trap at its root). Since the continuation type is + -- [Int] regardless of @a@, no interpreter lifting is needed either. + | PlateBegin Text Int next + -- ^ [日本語]: Plate 開始マーカー (Pyro/NumPyro 流の plate-block 糖衣)。 + -- 名前 + サイズ N を持つ plate スコープの開始。 直後から 'PlateEnd' + -- までに登録される 'Sample' / 'Observe' / 'Deterministic' は + -- buildModelGraph で「plate メンバ」 として描画される。 + -- nested plate は LIFO スタックで対応。 log eval interpreter (logJoint + -- 等) は __透過__ に処理する (何もしない)。 + -- + -- [English]: A plate-begin marker (sugar for a Pyro\/NumPyro-style + -- plate block). Opens a plate scope with a name and size N. Every + -- 'Sample' \/ 'Observe' \/ 'Deterministic' registered from here up + -- to the matching 'PlateEnd' is drawn as a "plate member" by + -- buildModelGraph. Nested plates are handled with a LIFO stack. + -- Log-eval interpreters (logJoint etc.) treat this __transparently__ + -- (i.e. do nothing). + | PlateEnd next + -- ^ [日本語]: Plate 終了マーカー。 最新の PlateBegin スコープを閉じる。 + -- [English]: A plate-end marker. Closes the most recent + -- 'PlateBegin' scope. + deriving Functor + +type Model a = Free (ModelF a) + +-- | [日本語]: 多相モデル DSL の型エイリアス。 +-- @ModelP r = forall a. (Floating a, Ord a, TrackTag a) => Model a r@ +-- ('TrackTag' は '!!!' の依存タグ注入用。 数値解釈は既定 id)。 +-- +-- [English]: Type alias for the polymorphic model DSL. +-- @ModelP r = forall a. (Floating a, Ord a, TrackTag a) => Model a r@ +-- ('TrackTag' is used to inject the dependency tag for '!!!'; the +-- default numeric interpretation is the identity). +type ModelP r = forall a. (Floating a, Ord a, TrackTag a) => Model a r + +sample :: Text -> Distribution a -> Model a a +sample n d = liftF (Sample n d id) + +observe :: Text -> Distribution a -> [Double] -> Model a () +observe n d ys = liftF (Observe n d ys ()) + +-- | [日本語]: 構造化線形予測子 observe。 +-- +-- @observeLM name betaNames designX family ys@ は、 設計行列 @designX@ +-- (n 行 × p 列) と β パラメタ名 @betaNames@ (長さ p・既に 'sample' で宣言済の +-- latent を参照) を __分離して__保持する観測ブロック。 各観測 i について +-- η_i = Σ_j β_j·X_ij を作り、 @family@ のリンク逆関数で μ_i に写して +-- 観測 @ys !! i@ の log-density を加算する。 +-- +-- 通常の per-obs @observe@ を N 回呼ぶのと数値的に等価だが、 線形構造を +-- 保持するので Gaussian-恒等リンクの十分統計量 collapse に乗せられる。 +-- +-- [English]: The structured linear-predictor observe. +-- +-- @observeLM name betaNames designX family ys@ is an observation block +-- that keeps the design matrix @designX@ (n rows × p columns) and the +-- β parameter names @betaNames@ (length p; referencing latents already +-- declared via 'sample') __separate__. For each observation i, it forms +-- η_i = Σ_j β_j·X_ij, maps it to μ_i via @family@'s inverse link, and +-- adds the log-density of the observation @ys !! i@. +-- +-- Numerically equivalent to calling per-obs @observe@ N times, but +-- since it preserves the linear structure it can be put through the +-- sufficient-statistic collapse for the Gaussian-identity link. +observeLM :: Text -> [Text] -> [[Double]] -> LMFamily -> [Double] -> Model a () +observeLM n betas designX fam ys = liftF (ObserveLM n betas designX [] fam ys ()) + +-- | [日本語]: ランダム効果付き 'observeLM'。 +-- +-- @observeLMR name betaNames designX reffs family ys@ は 'observeLM' に +-- ランダム効果項 @reffs@ を加えたもの。 各 'REff' は (u パラメタ名, 群 id) で +-- @η_i += u^{re}[gid_i]@ を __gather__ で寄与する。 群効果を設計行列の one-hot +-- 指示列に密展開すると vec-tape 勾配が O(nG·n) になり階層モデルで逆効果になる +-- (計測済) ため、 群構造は疎に保持して gather で O(n) に保つ。 +-- +-- [English]: 'observeLM' with random effects. +-- +-- @observeLMR name betaNames designX reffs family ys@ is 'observeLM' +-- with the random-effect terms @reffs@ added. Each 'REff' contributes +-- @η_i += u^{re}[gid_i]@ via a __gather__, from a (u parameter name, +-- group id) pair. Densely expanding group effects into one-hot +-- indicator columns of the design matrix would make the vec-tape +-- gradient O(nG·n), which backfires for hierarchical models (as +-- measured), so the group structure is kept sparse and applied via +-- gather to stay O(n). +observeLMR :: Text -> [Text] -> [[Double]] -> [REff] -> LMFamily -> [Double] + -> Model a () +observeLMR n betas designX reffs fam ys = + liftF (ObserveLM n betas designX reffs fam ys ()) + +-- --------------------------------------------------------------------------- +-- 第一級ランダム効果値 (Phase 54.4c) +-- --------------------------------------------------------------------------- + +-- | [日本語]: 第一級ランダム効果値。 'reNormal' で宣言した nG 個の +-- iid @Normal(0, τ)@ latent を、 構造 (基底名・群数・スケール名・値) ごと +-- ひとつの値に載せて持ち運ぶ。 これにより観測の線形予測子に効果を載せるとき +-- 文字列添字 (@"u_" <> show j@) も @us !! g@ も書かずに 'at' で gather でき +-- (Haskell 王道の「構造を値に載せて流す」)、 さらにスケール名が構造として +-- 保持されるので @compileGradU@ が u-prior 勾配を解析的にベクトル化できる。 +-- +-- [English]: A first-class random-effect value. Carries the nG iid +-- @Normal(0, τ)@ latents declared via 'reNormal' as a single value, +-- together with their structure (base name, group count, scale name, +-- values). This lets an observation's linear predictor gather the +-- effect via 'at' without writing string subscripts (@"u_" <> show j@) +-- or @us !! g@ (the idiomatic Haskell approach of "threading structure +-- through as a value"), and since the scale name is preserved as part +-- of the structure, @compileGradU@ can vectorize the u-prior gradient +-- analytically. +data REffect a = REffect + { reffBase :: !Text -- ^ [日本語]: 基底名 (例 @"u"@)。 latent 名は @base_@。 [English]: The base name (e.g. @"u"@); latent names are @base_@. + , reffNG :: !Int -- ^ [日本語]: 群数 nG [English]: The number of groups, nG. + , reffScale :: !Text -- ^ [日本語]: スケール latent の名前 (@u_j ~ Normal(0, scale)@) [English]: The name of the scale latent (@u_j ~ Normal(0, scale)@). + , reffValues :: [a] -- ^ [日本語]: サンプル済 nG 個の値 (forward 評価・deterministic 用) [English]: The nG already-sampled values (for forward evaluation \/ deterministic). + } + +-- | [日本語]: 'REffect' の latent 名 (@base_0 .. base_{nG-1}@)。 +-- [English]: 'REffect'\'s latent names (@base_0 .. base_{nG-1}@). +reffNames :: REffect a -> [Text] +reffNames re = [ indexed (reffBase re) j | j <- [0 .. reffNG re - 1] ] + +-- | [日本語]: 群別ランダム効果を第一級値として宣言する。 +-- +-- @reNormal base nG scaleName scaleVal@ は @base_0 .. base_{nG-1}@ という +-- nG 個の latent を各々 @Normal(0, scaleVal)@ として 'sample' し、 その構造 +-- (基底名 / nG / スケール名 / 値) を 'REffect' にまとめて返す。 @scaleName@ は +-- @scaleVal@ を生んだスケール latent の名前 (例 @"tau_u"@) で、 解析 prior 勾配 +-- 経路 (@compileGradU@) がスケール変数を引くために構造として保持する +-- (値は名前を覚えていないため明示的に渡す)。 +-- +-- @ +-- tau <- sample "tau_u" (HalfNormal 5) +-- u <- reNormal "u" nG "tau_u" tau +-- observeNormalLM "y" xRows betaNames [u \`at\` gids] "sigma" ys +-- @ +-- +-- [English]: Declares per-group random effects as a first-class value. +-- +-- @reNormal base nG scaleName scaleVal@ samples nG latents, +-- @base_0 .. base_{nG-1}@, each as @Normal(0, scaleVal)@ via 'sample', +-- and returns their structure (base name \/ nG \/ scale name \/ values) +-- bundled into an 'REffect'. @scaleName@ is the name of the scale +-- latent that produced @scaleVal@ (e.g. @"tau_u"@), kept as part of +-- the structure so that the analytic prior-gradient path +-- (@compileGradU@) can look up the scale variable (the value alone +-- doesn't remember its name, so it must be passed explicitly). +-- +-- @ +-- tau <- sample "tau_u" (HalfNormal 5) +-- u <- reNormal "u" nG "tau_u" tau +-- observeNormalLM "y" xRows betaNames [u \`at\` gids] "sigma" ys +-- @ +reNormal :: Num a => Text -> Int -> Text -> a -> Model a (REffect a) +reNormal base nG scaleName scaleVal = do + vals <- forM [0 .. nG - 1] $ \j -> + sample (indexed base j) (Normal 0 scaleVal) + pure (REffect base nG scaleName vals) + +-- | [日本語]: 'REffect' を観測の群 id 列に対して gather し 'REff' (観測ブロック用) に変換する。 +-- @η_i += u^{re}[gid_i]@。 スケール名を 'REff' に載せるので、 これ経由で観測に +-- 入った効果は @compileGradU@ の解析 prior 勾配経路に乗る。 +-- +-- [English]: Gathers an 'REffect' over an observation's group-id list +-- and converts it to an 'REff' (for use in observation blocks): +-- @η_i += u^{re}[gid_i]@. Since the scale name is carried on the +-- 'REff', an effect entered this way is put through @compileGradU@\'s +-- analytic prior-gradient path. +at :: REffect a -> [Int] -> REff +at re gids = REff (reffNames re) gids (Just (reffScale re)) Nothing Nothing + +-- | [日本語]: Gaussian-恒等リンク版の構造化 observe。 'observeLMR' の +-- @LMGaussian@ 特化で、 'at' で作った 'REff' をそのまま渡せる薄いラッパ。 +-- +-- @observeNormalLM name designX betaNames reffs sigmaName ys@。 +-- +-- [English]: The structured observe specialized to the +-- Gaussian-identity link. A thin wrapper around 'observeLMR'\'s +-- @LMGaussian@ case that lets you pass an 'REff' built with 'at' +-- directly. +-- +-- @observeNormalLM name designX betaNames reffs sigmaName ys@. +observeNormalLM :: Text -> [[Double]] -> [Text] -> [REff] -> Text -> [Double] + -> Model a () +observeNormalLM name designX betaNames reffs sName ys = + observeLMR name betaNames designX reffs (LMGaussian sName) ys + +-- | [日本語]: 多変量観測 ('MvNormal' 用)。 各観測は長さ @k@ のベクトルで、 +-- リストとして @[[Double]]@ で渡す。 内部的には @concat@ で flatten され、 +-- 評価時に Distribution の次元 k で chunk される。 +-- +-- [English]: Multivariate observation (for 'MvNormal'). Each +-- observation is a length-@k@ vector; pass them as a list +-- @[[Double]]@. Internally it is flattened via @concat@, then +-- re-chunked into groups of the Distribution's dimension k at +-- evaluation time. +observeMV :: Text -> Distribution a -> [[Double]] -> Model a () +observeMV n d obss = liftF (Observe n d (concat obss) ()) + +-- | [日本語]: 多出力観測 helper。 @q@ 組の +-- @observe (prefix <> \"_\" <> j) dist_j ys_j@ を順に発行する。 +-- +-- 多出力回帰の尤度を 1 行で書きたいときに使う: +-- +-- @ +-- observeColumns \"y\" [(Normal mu_j sigma_j, ysCol j) | j <- [0 .. q - 1]] +-- @ +-- +-- [English]: Multi-output observation helper. Emits @q@ pairs of +-- @observe (prefix <> \"_\" <> j) dist_j ys_j@ in order. +-- +-- Useful when you want to write a multi-output regression's +-- likelihood in one line: +-- +-- @ +-- observeColumns \"y\" [(Normal mu_j sigma_j, ysCol j) | j <- [0 .. q - 1]] +-- @ +observeColumns :: Text -> [(Distribution a, [Double])] -> Model a () +observeColumns prefix pairs = + mapM_ (\(j, (d, ys)) -> + observe (prefix <> "_" <> T.pack (show (j :: Int))) d ys) + (zip [0..] pairs) + +-- | [日本語]: インデックス付きノード名を作る: @indexed "theta" 1 == "theta_1"@。 +-- +-- 階層モデルで群ごとの 'sample' / 'observe' 名を作るときに頻出する +-- @T.pack ("theta_" ++ show j)@ ボイラープレートを畳む。 アンダースコアは +-- 自動付与 (= 'observeColumns' / 'nonCenteredNormal' 等の命名規約に一致)。 +-- +-- [English]: Builds an indexed node name: +-- @indexed "theta" 1 == "theta_1"@. +-- +-- Folds the common boilerplate @T.pack ("theta_" ++ show j)@ used to +-- name per-group 'sample' / 'observe' calls in hierarchical models. The +-- underscore is added automatically (matching the naming convention of +-- 'observeColumns' / 'nonCenteredNormal', etc.). +-- +-- > forM_ (zip [1..] groupData) $ \(j, ys) -> do +-- > theta <- sample (indexed "theta" j) (Normal mu tau) -- "theta_1" … +-- > observe (indexed "y" j) (Normal theta 1) ys +indexed :: Text -> Int -> Text +indexed pre i = pre <> "_" <> T.pack (show i) + +-- | [日本語]: 'indexed' の中置演算子版: @"theta" .# j == "theta_1"@。 +-- (Haskell の演算子記号に @_@ は使えないため @.#@ を採用。) +-- [English]: The infix-operator form of 'indexed': +-- @"theta" .# j == "theta_1"@. (@.#@ is used because Haskell operator +-- symbols cannot contain @_@.) +infixl 9 .# +(.#) :: Text -> Int -> Text +(.#) = indexed + +-- | Add an arbitrary log-probability term to the model (analogous to +-- PyMC's @pm.Potential@). +-- +-- [日本語]: 通常のサンプリング/観測では表せない log-density 寄与を入れるのに +-- 使う。 典型用途: +-- +-- - __ソフト制約__: @potential \"order\" (if mu1 < mu2 then 0 else (-1e10))@ +-- - __カスタム尤度__: 既存 'Distribution' で表せない尤度項 +-- - __正則化__: ベイズ的な正則化 (e.g. ridge: @-0.5 * lambda * sum (map (^2) betas)@) +-- +-- @Potential@ の値は @logJoint@ と @logPrior@ に加算される +-- (@logLikelihood@ には含まれない — これらは @observe@ 専用)。 +-- +-- [English]: Used to add log-density contributions that ordinary +-- sampling/observation cannot express. Typical uses: +-- +-- - __Soft constraints__: @potential \"order\" (if mu1 < mu2 then 0 else (-1e10))@ +-- - __Custom likelihoods__: likelihood terms not expressible with an existing 'Distribution' +-- - __Regularization__: Bayesian regularization (e.g. ridge: @-0.5 * lambda * sum (map (^2) betas)@) +-- +-- @Potential@'s value is added to @logJoint@ and @logPrior@ (it is not +-- included in @logLikelihood@ — those are @observe@-only). +potential :: Text -> a -> Model a () +potential nm v = liftF (Potential nm v ()) + +-- | [日本語]: 派生量を名前付きで保存する (PyMC `pm.Deterministic` 相当)。 +-- log-joint には寄与しないが、 各 posterior サンプルごとに値が記録され +-- @augmentChainWithDeterministic@ で Chain に注入できる。 +-- [English]: Saves a derived quantity under a name (equivalent to +-- PyMC's @pm.Deterministic@). It does not contribute to the log-joint, +-- but its value is recorded for each posterior draw and can be injected +-- into the Chain via @augmentChainWithDeterministic@. +-- +-- 例 / Example: +-- +-- > tau <- deterministic "tau" (1 / (sigma * sigma)) +deterministic :: Text -> a -> Model a a +deterministic nm v = liftF (Deterministic nm v id) + +-- | [日本語]: DAG / Node 表示用の分布名 (リンク逆関数を適用した観測分布の名前)。 [English]: The distribution name for DAG / Node display (the observation distribution's name after applying the inverse link). +lmFamilyName :: LMFamily -> Text +lmFamilyName (LMGaussian _) = "Normal" +lmFamilyName LMPoisson = "Poisson" +lmFamilyName LMBernoulli = "Bernoulli" + +-- | [日本語]: 'ObserveLM' が参照する latent パラメタ名の集合 (DAG の親)。 +-- β + ランダム効果 u + (Gaussian の) σ。 +-- [English]: The set of latent parameter names 'ObserveLM' references +-- (the DAG's parents): β + random effects u + (for Gaussian) σ. +lmParents :: [Text] -> [REff] -> LMFamily -> Set Text +lmParents betaNames reffs fam = + Set.fromList betaNames + <> Set.fromList (concat [ uNames | REff uNames _ _ _ _ <- reffs ]) + -- Phase 62: gids の由来 slot 名 ('atIx' 経由) も親に = slot→観測ノードのエッジ + <> Set.fromList [ s | REff _ _ _ _ (Just s) <- reffs ] + <> case fam of + LMGaussian sName -> Set.singleton sName + LMPoisson -> Set.empty + LMBernoulli -> Set.empty + +-- --------------------------------------------------------------------------- +-- Phase 40-A1: Plate notation +-- --------------------------------------------------------------------------- + +-- | [日本語]: Pyro / NumPyro 流の plate-block。 +-- [English]: A Pyro-/NumPyro-style plate block. +-- +-- [日本語]: @plate name n body@ は、 do-block 内で繰り返し作られる indexed RV 群 +-- (e.g. @eta_0, eta_1, …, eta_{n-1}@) を __同じ plate に属する__ と +-- マークする bracket。 @buildModelGraph@ で plate 集約描画される。 +-- [English]: @plate name n body@ is a bracket that marks the indexed RVs +-- repeatedly created inside a do-block (e.g. @eta_0, eta_1, …, +-- eta_{n-1}@) as __belonging to the same plate__. @buildModelGraph@ +-- renders plates aggregated. +-- +-- 例 (8-schools) / Example (8-schools): +-- +-- > mu <- sample "mu" (Normal 0 5) +-- > tau <- sample "tau" (HalfCauchy 5) +-- > etas <- plate "school" 8 $ forM [0..7] $ \j -> +-- > sample ("eta_" <> T.pack (show j)) (Normal 0 1) +-- > _ <- plate "school" 8 $ forM_ [0..7] $ \j -> +-- > observe ("y_" <> T.pack (show j)) +-- > (Normal (mu + tau * (etas !! j)) 1) [ys !! j] +-- +-- [日本語]: 内部: 'PlateBegin' / 'PlateEnd' マーカーで囲む。 log eval (logJoint +-- / logPrior 等) は __透過__ に動作し、 plate は描画レイヤーでのみ +-- 意味を持つ。 NUTS / Gibbs / VI への影響なし。 +-- [English]: Internally, this wraps the body with 'PlateBegin' / +-- 'PlateEnd' markers. Log evaluation (logJoint / logPrior, etc.) works +-- __transparently__ through it — plates only carry meaning at the +-- rendering layer, and have no effect on NUTS / Gibbs / VI. +plate :: Text -> Int -> Model a r -> Model a r +plate name n body = do + liftF (PlateBegin name n ()) + r <- body + liftF (PlateEnd ()) + return r + +-- | [日本語]: 'plate' の利便 helper: @plateI name n f@ = +-- @plate name n (forM [0..n-1] f)@。 「N 個の indexed RV を作る」 という +-- 最頻パターン向け糖衣。 +-- [English]: A convenience helper over 'plate': @plateI name n f@ = +-- @plate name n (forM [0..n-1] f)@. Sugar for the most common pattern, +-- "create N indexed RVs." +-- +-- 例 / Example: +-- +-- > etas <- plateI "school" 8 $ \j -> +-- > sample ("eta_" <> T.pack (show j)) (Normal 0 1) +plateI :: Text -> Int -> (Int -> Model a r) -> Model a [r] +plateI name n action = plate name n (forM [0 .. n - 1] action) + +-- | [日本語]: 'plateI' の返り値を捨てる版 (@forM_@ の plate 版・index 反復)。 +-- @plateI_ name n f = plate name n (forM_ [0..n-1] f)@。 観測のみの index +-- ループ向け (@plateForM_ name [0..n-1] f@ と同義だが index 反復の意図が明示的・ +-- 'plateForM' / 'plateForM_' の対称に合わせ index 版にも破棄形を用意)。 +-- [English]: The value-discarding version of 'plateI' (the plate +-- version of @forM_@, index-driven). @plateI_ name n f = plate name n +-- (forM_ [0..n-1] f)@. For observation-only index loops (equivalent to +-- @plateForM_ name [0..n-1] f@, but makes index-driven iteration +-- explicit; provided so the index-based variant has a discarding form +-- symmetric with 'plateForM' / 'plateForM_'). +-- +-- 例 (8-schools の観測) / Example (8-schools observations): +-- +-- > plateI_ "school" 8 $ \j -> +-- > observe ("y" .# j) (Normal (mu + tau * etas !! j) 1) [ys !! j] +plateI_ :: Text -> Int -> (Int -> Model a r) -> Model a () +plateI_ name n action = plate name n (forM_ [0 .. n - 1] action) + +-- | [日本語]: データ行リストを plate で囲んで反復する糖衣 (@forM@ の plate +-- 版・引数順も @forM@ 形)。 @plateForM name rows f = plate name (length rows) +-- (forM rows f)@。 plate サイズは行数から自動。 観測ループの定番 +-- @plate name (length rows) $ forM_ … rows@ を畳む。 +-- [English]: Sugar for iterating over a list of data rows wrapped in a +-- plate (the plate version of @forM@, argument order matches @forM@ +-- too). @plateForM name rows f = plate name (length rows) (forM rows +-- f)@. The plate size is derived automatically from the row count, +-- folding the common observation-loop pattern +-- @plate name (length rows) $ forM_ … rows@. +-- +-- 例 (ベイズ線形回帰の観測) / Example (Bayesian linear regression observations): +-- +-- > plateForM_ "obs" (zip x y) $ \(xi, yi) -> do +-- > mu <- deterministic "mu" (a + b * realToFrac xi) +-- > observe "obs" (Normal mu s) [yi] +plateForM :: Text -> [b] -> (b -> Model a r) -> Model a [r] +plateForM name rows f = plate name (length rows) (forM rows f) + +-- | [日本語]: 返り値を捨てる版 (@forM_@ の plate 版)。 観測のみのループに。 [English]: The value-discarding version (the plate version of @forM_@), for observation-only loops. +plateForM_ :: Text -> [b] -> (b -> Model a r) -> Model a () +plateForM_ name rows f = plate name (length rows) (forM_ rows f) + +-- | [日本語]: 低レベル plate API: 任意の Model action を plate スコープで包む。 +-- 'plate' は @withPlate name n@ + body の組合せに分解される。 nested +-- plate を独自構築する際の primitive。 +-- [English]: The low-level plate API: wraps an arbitrary Model action in +-- a plate scope. 'plate' decomposes into @withPlate name n@ + body; this +-- is the primitive for building custom nested plates. +withPlate :: Text -> Int -> Model a r -> Model a r +withPlate = plate + +-- | [日本語]: 名前付きデータプレースホルダを宣言する (PyMC `pm.Data` 相当)。 +-- 既定値 @ys@ を持ち、 後で 'withData' により差し替え可能。 +-- [English]: Declares a named data placeholder (equivalent to PyMC's +-- @pm.Data@). Has a default value @ys@, later swappable via 'withData'. +-- +-- 典型的な使い方 / Typical usage: +-- +-- > model = do +-- > y <- dataNamed "y" trainData +-- > mu <- sample "mu" (Normal 0 5) +-- > observe "y" (Normal mu 1) y +-- +-- [日本語]: そして @withData \"y\" testData model@ で同じ構造で別データを使う。 +-- [English]: Then @withData \"y\" testData model@ reuses the same +-- structure with different data. +-- +-- [日本語]: ★破壊的変更: 戻り値は @[a]@ (モデルの数値型)。 受け取った値は +-- そのまま式に入る (@realToFrac@ 不要)。 @a@ には @Real@ 制約が無いので、 +-- 旧コードの @realToFrac xi@ は型エラーになる (= 無言の挙動変化が起きない +-- 壊れ方)。 機械的に @realToFrac@ を消せば移行完了。 +-- 観測値として 'observe' に渡す側 (@[Double]@ が要る) は 'dataNamedObs' を使う。 +-- [English]: ★A breaking change: the return type is @[a]@ (the model's +-- numeric type). The received value flows directly into expressions (no +-- @realToFrac@ needed). Since @a@ carries no @Real@ constraint, old code +-- with @realToFrac xi@ now fails to type-check (a loud failure, not a +-- silent behavior change). Migration is complete once @realToFrac@ is +-- mechanically removed. Use 'dataNamedObs' for the side that passes +-- observed values to 'observe' (which needs @[Double]@). +dataNamed :: Text -> [Double] -> Model a [a] +dataNamed n ys = liftF (Data n ys fst) + +-- | [日本語]: 'dataNamed' の同義。 役割 suffix 三点セットの正書き: +-- [English]: A synonym for 'dataNamed'. The canonical spelling of the +-- three role-suffixed variants: +-- +-- > x <- dataNamedX "x" [] -- 説明変数 / covariate: モデル数値型 [a] +-- > ys <- dataNamedObs "y" [] -- 目的変数 / response: 生 [Double] ('observe' へ) +-- > gs <- dataNamedIx "g" [] -- 群 index / group index: [Int] +-- +-- [日本語]: 既存コードの 'dataNamed' もそのまま使える (削除予定なし)。 +-- [English]: Existing code's 'dataNamed' also still works as-is (no plan +-- to remove it). +dataNamedX :: Text -> [Double] -> Model a [a] +dataNamedX = dataNamed + +-- | [日本語]: 'dataNamed' と同じ slot の __観測値 view__ (生 @[Double]@)。 +-- 'observe' / 'observeLM' の観測値引数は AD に持ち上げない @[Double]@ 固定 +-- なので、 y 側のデータ slot はこちらで受ける: +-- [English]: The __observed-value view__ (raw @[Double]@) of the same +-- slot as 'dataNamed'. Since 'observe' / 'observeLM''s observed-value +-- argument is always @[Double]@ and never lifted to AD, y-side data +-- slots should be received this way: +-- +-- > x <- dataNamed "x" [] -- covariate: モデル数値型 [a] +-- > ys <- dataNamedObs "y" [] -- 観測値 / observed value: 生 [Double] +-- > ... +-- > observe "y" (Normal mu s) ys +-- +-- [日本語]: 同名 slot を 'dataNamed' と 'dataNamedObs' の両 view で読んでもよい +-- (差し替えは 'withData' / 列 bind が slot 名単位で行うため一貫する)。 +-- [English]: The same-named slot may be read through both the +-- 'dataNamed' and 'dataNamedObs' views — this stays consistent because +-- swapping (via 'withData' / column binding) operates per slot name. +dataNamedObs :: Text -> [Double] -> Model a [Double] +dataNamedObs n ys = liftF (Data n ys snd) + +-- | [日本語]: 離散 index 専用のデータプレースホルダ (後に 'Ix' 戻りへ刷新)。 +-- 群 index 等を slot 名タグ付き index 'Ix' で運ぶ。 @bs '!!!' g@ で引くと +-- DAG に slot→利用先のエッジが自動で出る (PyMC の @b0[gid]@ 同型)。 +-- 'Ix' は Num でないので誤って算術に混ぜると型エラーで止まる +-- (= 連続値経路の round 罠を根治し、 以後も維持)。 +-- [English]: A data placeholder dedicated to discrete indices (later +-- revised to return 'Ix'). Carries group indices, etc. as a slot-name- +-- tagged index 'Ix'. Indexing with @bs '!!!' g@ automatically emits a +-- slot→use-site edge in the DAG (matching PyMC's @b0[gid]@). Since 'Ix' +-- is not a @Num@, accidentally mixing it into arithmetic fails at +-- compile time (this permanently closes the continuous-value rounding +-- trap). +-- +-- > gs <- dataNamedIx "g" [0,0,1,1,2] +-- > let mu_i = b0s !!! g -- round 不要 / no round needed・DAG に g→mu エッジ +dataNamedIx :: Text -> [Int] -> Model a [Ix] +dataNamedIx n is = liftF (DataIx n is (map (\i -> Ix i (Just n)))) + +-- | [日本語]: slot 名タグ付き離散 index。 'dataNamedIx' が返し、 '!!!' で +-- 使う。 由来 slot 名 ('ixSlot') は DAG 抽出 (Track 解釈) のエッジ生成にだけ +-- 使われ、 数値評価では 'ixVal' のみが意味を持つ。 +-- [English]: A slot-name-tagged discrete index, returned by +-- 'dataNamedIx' and consumed by '!!!'. The originating slot name +-- ('ixSlot') is used only for edge generation in DAG extraction (the +-- @Track@ interpretation); numeric evaluation looks only at 'ixVal'. +data Ix = Ix + { ixVal :: !Int -- ^ [日本語]: index 本体 (0..nG-1)。 [English]: the index value itself (0..nG-1). + , ixSlot :: !(Maybe Text) -- ^ [日本語]: 由来 slot 名 ('dataNamedIx' なら Just)。 [English]: the originating slot name (@Just@ when from 'dataNamedIx'). + } deriving (Show, Eq) + +-- | [日本語]: 解釈ごとの依存タグ注入。 既定 = 何もしない (数値解釈は +-- ゼロコスト・サンプリングはビット不変)。 @Track@ 解釈だけが override して +-- 依存集合に slot 名を足し、 DAG にエッジを出す。 +-- [English]: Injects a dependency tag, per interpretation. The default +-- does nothing (numeric evaluation pays zero cost; sampling is bit- +-- identical). Only the @Track@ interpretation overrides it, adding the +-- slot name to the dependency set and emitting a DAG edge. +class TrackTag a where + tagDep :: Text -> a -> a + tagDep _ = id + {-# INLINE tagDep #-} + +instance TrackTag Double + +-- dogfood 典型 (群別係数のタプル) 用: 成分ごとに伝播 +instance (TrackTag a, TrackTag b) => TrackTag (a, b) where + tagDep nm (a, b) = (tagDep nm a, tagDep nm b) +instance (TrackTag a, TrackTag b, TrackTag c) => TrackTag (a, b, c) where + tagDep nm (a, b, c) = (tagDep nm a, tagDep nm b, tagDep nm c) +instance (TrackTag a, TrackTag b, TrackTag c, TrackTag d) + => TrackTag (a, b, c, d) where + tagDep nm (a, b, c, d) = (tagDep nm a, tagDep nm b, tagDep nm c, tagDep nm d) + +-- | [日本語]: slot 名タグ付き索引。 @bs '!!!' g@ = @bs !! ixVal g@ に、 +-- Track 解釈でのみ g の由来 slot 名を依存タグとして注入する +-- (= DAG に slot→利用先エッジ。 数値解釈は '!!' と同コスト)。 +-- [English]: Slot-name-tagged indexing. @bs '!!!' g@ is @bs !! ixVal g@, +-- with g's originating slot name injected as a dependency tag only +-- under the @Track@ interpretation (emitting a slot→use-site edge in +-- the DAG; numeric evaluation costs the same as '!!'). +(!!!) :: TrackTag b => [b] -> Ix -> b +xs !!! Ix i ms = maybe id tagDep ms (xs !! i) +infixl 9 !!! +{-# INLINE (!!!) #-} + +-- | [日本語]: 'at' の 'Ix' 版。 'dataNamedIx' の gids を random effect の +-- gather に渡す。 先頭 'Ix' の由来 slot 名 ('ixSlot') を 'REff' に +-- 載せるので、 DAG に slot→観測ノードのエッジが出る (gather の gids は単一 +-- slot 由来が通常形ゆえ先頭で代表)。 '!!!' (deterministic μ 経路) と並ぶ +-- PyMC @b0[gid]@ 同型の両経路対応。 +-- [English]: The 'Ix' version of 'at'. Passes 'dataNamedIx''s gids to a +-- random effect's gather. Carries the first 'Ix''s originating slot name +-- ('ixSlot') onto 'REff', so the DAG gets a slot→observation-node edge +-- (since a gather's gids typically come from a single slot, the first +-- one is taken as representative). Covers both paths matching PyMC's +-- @b0[gid]@, alongside '!!!' (the deterministic-μ path). +atIx :: REffect a -> [Ix] -> REff +atIx re gids = + REff (reffNames re) (map ixVal gids) (Just (reffScale re)) Nothing + (case gids of { Ix _ ms : _ -> ms; [] -> Nothing }) + +-- | Replace a named data block in the model. If no match exists the +-- model is returned unchanged. +-- +-- [日本語]: 同じ名前が複数回出現する場合は全箇所で差し替わる。 +-- 型シグネチャは @Model a r@ なので、 ユーザーが @ModelP r@ から呼ぶ場合 +-- そのまま多相的に使える (各 @a@ で個別に適用される)。 +-- [English]: If the same name occurs multiple times, all occurrences are +-- replaced. Since the type signature is @Model a r@, users calling from +-- @ModelP r@ can use it polymorphically as-is (applied separately for +-- each @a@). +withData :: forall r. Text -> [Double] -> ModelP r -> ModelP r +withData n new m = mPoly + where + -- 戻り値を多相モデルとして再構築。各 @a@ 個別に元の m を走査する。 + mPoly :: forall a. (Floating a, Ord a, TrackTag a) => Model a r + mPoly = go m + where + go :: Model a r -> Model a r + go (Pure r) = Pure r + go (Free f) = Free (case f of + Data n' ys k + | n == n' -> Data n' new (\d -> go (k d)) + | otherwise -> Data n' ys (\d -> go (k d)) + DataIx n' is k -> DataIx n' is (\d -> go (k d)) + Sample nm d k -> Sample nm d (\v -> go (k v)) + Observe nm d ys nx -> Observe nm d ys (go nx) + ObserveLM nm bs xs re fam ys nx -> ObserveLM nm bs xs re fam ys (go nx) + Potential nm v nx -> Potential nm v (go nx) + Deterministic nm v k -> Deterministic nm v (\v' -> go (k v')) + PlateBegin nm sz nx -> PlateBegin nm sz (go nx) + PlateEnd nx -> PlateEnd (go nx)) + +-- | [日本語]: 'withData' の離散 index 版: 名前付き @DataIx@ ブロックを +-- 外部から差し替える。 一致しなければモデルは不変。 +-- [English]: The discrete-index version of 'withData': externally +-- replaces a named @DataIx@ block. If nothing matches, the model is +-- returned unchanged. +withDataIx :: forall r. Text -> [Int] -> ModelP r -> ModelP r +withDataIx n new m = mPoly + where + mPoly :: forall a. (Floating a, Ord a, TrackTag a) => Model a r + mPoly = go m + where + go :: Model a r -> Model a r + go (Pure r) = Pure r + go (Free f) = Free (case f of + DataIx n' is k + | n == n' -> DataIx n' new (\d -> go (k d)) + | otherwise -> DataIx n' is (\d -> go (k d)) + Data n' ys k -> Data n' ys (\d -> go (k d)) + Sample nm d k -> Sample nm d (\v -> go (k v)) + Observe nm d ys nx -> Observe nm d ys (go nx) + ObserveLM nm bs xs re fam ys nx -> ObserveLM nm bs xs re fam ys (go nx) + Potential nm v nx -> Potential nm v (go nx) + Deterministic nm v k -> Deterministic nm v (\v' -> go (k v')) + PlateBegin nm sz nx -> PlateBegin nm sz (go nx) + PlateEnd nx -> PlateEnd (go nx)) + +-- | Latent multivariate-normal vector (analogous to PyMC's +-- @pm.MvNormal@ used as a latent). +-- +-- [日本語]: 非中心化パラメタ化 + Cholesky 分解で実装: +-- [English]: Implemented via non-centered parameterization + Cholesky +-- decomposition: +-- +-- z_i ~ Normal(0, 1) (i = 0..K-1, 独立な latent / independent latents) +-- x = μ + L z (L = Cholesky(Σ)) +-- +-- [日本語]: 各 z_i は通常の latent として NUTS が探索し、 x は派生量として +-- Chain に記録される。 共分散行列が他の latent に依存する形でも +-- 動作する (choleskyL は @(Floating a, Ord a)@ 多相)。 +-- [English]: Each z_i is explored by NUTS as an ordinary latent, and x +-- is recorded in the Chain as a derived quantity. This also works when +-- the covariance matrix depends on other latents (@choleskyL@ is +-- polymorphic over @(Floating a, Ord a)@). +-- +-- [日本語]: 共分散が非正定値のときは μ をそのまま返す (NUTS 探索中の不正領域 +-- に対する graceful fallback)。 +-- [English]: When the covariance is not positive-definite, μ is returned +-- as-is (a graceful fallback for invalid regions visited during NUTS +-- exploration). +-- +-- [日本語]: 戻り値: K 次元 latent ベクトル @[a]@ (μ + L z)。 Chain には +-- @_z@ (raw latent) と @_@ (派生量) を保存。 +-- [English]: Returns: the K-dimensional latent vector @[a]@ (μ + L z). +-- The Chain stores @_z@ (the raw latent) and @_@ (the +-- derived quantity). +mvNormalLatent :: forall a. (Floating a, Ord a) + => Text -> [a] -> [[a]] -> Model a [a] +mvNormalLatent name muVec covMatrix = do + let k = length muVec + zs <- mapM (\i -> sample (name <> "_z" <> T.pack (show i)) (Normal 0 1)) + [0 .. k - 1] + let xs = case choleskyL covMatrix of + Just l -> [ (muVec !! i) + + sum [ ((l !! i) !! j) * (zs !! j) + | j <- [0 .. i] ] + | i <- [0 .. k - 1] ] + Nothing -> muVec -- non-PD のフォールバック + mapM + (\(i, x) -> deterministic (name <> "_" <> T.pack (show i)) x) + (zip [0 :: Int ..] xs) + +-- | [日本語]: LKJ 相関行列の Cholesky factor (PyMC @LKJCholeskyCov@ 相当)。 +-- [English]: The Cholesky factor of an LKJ correlation matrix +-- (equivalent to PyMC's @LKJCholeskyCov@). +-- +-- [日本語]: LKJ(η) 事前分布: p(R) ∝ |R|^(η-1)。 η = 1 で uniform、 +-- η > 1 で I に集中。 +-- [English]: LKJ(η) prior: p(R) ∝ |R|^(η-1). Uniform at η = 1, +-- concentrating toward I as η > 1. +-- +-- [日本語]: 実装は canonical partial correlations (CPC) 法: +-- [English]: Implemented via the canonical partial correlations (CPC) +-- method: +-- +-- z_ij ~ scaled Beta(α_i, α_i) on (-1, 1), α_i = η + (K - i - 1) / 2 +-- (i = 1..K-1, j = 0..i-1) +-- +-- [日本語]: 各 z_ij は @_pc_@ (Beta latent in (0,1)、内部で +-- 2u-1 に変換) として保存。 Cholesky factor の各要素は派生量 +-- @_L_@。 +-- [English]: Each z_ij is stored as @_pc_@ (a Beta latent in +-- (0,1), converted internally to 2u-1). Each Cholesky factor element is +-- the derived quantity @_L_@. +-- +-- [日本語]: 戻り値: K×K 下三角行列 L (R = L Lᵀ となる相関の Cholesky)。 +-- 対角は √(1 - Σ z_{i,k}²)、対角下は z_ij × √(Π_{k Text -> Int -> a -> Model a [[a]] +lkjCorrCholesky name k eta + | k < 2 = error "lkjCorrCholesky: dimension must be >= 2" + | otherwise = do + -- 各 (i, j) で 1 <= j < i <= K-1 の partial correlation を sample + let pcIndices = [(i, j) | i <- [1 .. k - 1], j <- [0 .. i - 1]] + pcs <- mapM + (\(i, j) -> do + let alpha = eta + fromIntegral (k - i - 1) / 2 + tag = T.pack (show i) <> "_" <> T.pack (show j) + u <- sample (name <> "_u" <> tag) (Beta alpha alpha) + deterministic (name <> "_pc" <> tag) (2 * u - 1)) + pcIndices + -- (i,j) → z_ij マップ + let pcMap = zip pcIndices pcs + lookupPC i j = head [v | ((ii, jj), v) <- pcMap, ii == i, jj == j] + -- Cholesky factor を構築 (下三角) + let lRow i = + [ if j > i then 0 + else if i == 0 && j == 0 then 1 + else if j == i -- 対角 + then sqrt (1 - sum [ let z = lookupPC i kk + in z * z | kk <- [0 .. i - 1] ]) + else -- 対角下 j < i + let z = lookupPC i j + factor2 = product [ let z' = lookupPC i kk + in 1 - z' * z' | kk <- [0 .. j - 1] ] + in z * sqrt factor2 + | j <- [0 .. k - 1] ] + lMat = [lRow i | i <- [0 .. k - 1]] + -- L 各要素を deterministic として保存 + _ <- mapM + (\(i, j) -> + deterministic (name <> "_L" <> T.pack (show i) <> "_" <> T.pack (show j)) + ((lMat !! i) !! j)) + [(i, j) | i <- [0 .. k - 1], j <- [0 .. i]] + return lMat + +-- | [日本語]: RBF (exponentiated quadratic) カーネルによる GP 共分散行列 +-- (Stan @gp_exp_quad_cov(x, alpha, rho)@ 相当)。 +-- [English]: A GP covariance matrix via the RBF (exponentiated +-- quadratic) kernel (equivalent to Stan's +-- @gp_exp_quad_cov(x, alpha, rho)@). +-- +-- [日本語]: @K[i][j] = alpha^2 * exp(-0.5 * (x_i - x_j)^2 / rho^2)@、対角には +-- 数値安定化の jitter (1e-10) を加える (Stan 原典の +-- @+ diag_matrix(rep_vector(1e-10, N))@ に対応)。 @x@ は 'dataNamedX' で +-- 束縛した @[a]@ をそのまま渡す (data とハイパーパラメータ alpha/rho は +-- 共に @a@ 型なので realToFrac 不要)。 +-- [English]: @K[i][j] = alpha^2 * exp(-0.5 * (x_i - x_j)^2 / rho^2)@, +-- with a numerical-stability jitter (1e-10) added on the diagonal +-- (corresponding to Stan's original +-- @+ diag_matrix(rep_vector(1e-10, N))@). @x@ can be passed straight +-- through as the @[a]@ bound by 'dataNamedX' (no @realToFrac@ needed +-- since both the data and the alpha/rho hyperparameters share type @a@). +-- +-- [日本語]: vecIR (per-row 独立項の和が前提) には密行列が構造的に載らない +-- ため、 legacy walk+ad 経路 (@grad fFull@) で使う想定の孤立関数。 +-- [English]: A dense matrix cannot structurally fit vecIR (which assumes +-- a sum of per-row-independent terms), so this is an isolated function +-- intended for use via the legacy walk+ad path (@grad fFull@). +gpExpQuadCov :: forall a. Floating a => [a] -> a -> a -> [[a]] +gpExpQuadCov xs alpha rho = + [ [ let d = xi - xj + in alpha * alpha * exp (negate 0.5 * d * d / (rho * rho)) + + (if i == j then 1e-10 else 0) + | (j, xj) <- zip [0 :: Int ..] xs ] + | (i, xi) <- zip [0 :: Int ..] xs ] + +-- | [日本語]: Gaussian Process 潜在関数 (Stan の non-centered GP +-- パラメタ化相当): +-- [English]: A Gaussian Process latent function (equivalent to Stan's +-- non-centered GP parameterization): +-- +-- > f_tilde ~ Normal(0, 1) (各点独立 / independent per point) +-- > L_cov = cholesky_decompose(gp_exp_quad_cov(x, alpha, rho)) +-- > f = L_cov * f_tilde +-- +-- [日本語]: 既存の 'choleskyL' ('mvNormalLatent' と同じ AD 対応 Cholesky +-- 分解) をそのまま流用する。 共分散が非正定値のときは全ゼロに +-- フォールバックする ('mvNormalLatent' と同型の graceful fallback)。 +-- [English]: Reuses the existing 'choleskyL' (the same AD-compatible +-- Cholesky decomposition as 'mvNormalLatent'). When the covariance is +-- not positive-definite, falls back to all zeros (a graceful fallback +-- of the same form as 'mvNormalLatent'). +-- +-- [日本語]: 戻り値: N 次元 latent ベクトル @[a]@ (GP 事後関数値 f)。 各要素は +-- @_f@ として deterministic 保存される。 +-- [English]: Returns: the N-dimensional latent vector @[a]@ (the GP +-- posterior function value f). Each element is stored as a +-- deterministic under @_f@. +gpLatent :: forall a. (Floating a, Ord a) + => Text -> [a] -> a -> a -> Model a [a] +gpLatent name xs alpha rho = do + let n = length xs + ftilde <- mapM (\i -> sample (name <> "_ftilde" <> T.pack (show i)) (Normal 0 1)) + [0 .. n - 1] + let cov = gpExpQuadCov xs alpha rho + fs = case choleskyL cov of + Just l -> [ sum [ (l !! i !! j) * (ftilde !! j) | j <- [0 .. i] ] + | i <- [0 .. n - 1] ] + Nothing -> replicate n 0 -- non-PD のフォールバック + mapM + (\(i, f) -> deterministic (name <> "_f" <> T.pack (show i)) f) + (zip [0 :: Int ..] fs) + +-- | [日本語]: AR(1) latent 時系列 (PyMC `pm.AR1` 相当)。 +-- [English]: An AR(1) latent time series (equivalent to PyMC's +-- @pm.AR1@). +-- +-- [日本語]: 状態方程式: x_t = ϕ x_{t−1} + ε_t, ε_t ~ Normal(0, σ) +-- 初期分布: x_0 ~ Normal(0, σ / √(1 − ϕ²)) (定常分布、 |ϕ| < 1 なら有限) +-- [English]: State equation: x_t = ϕ x_{t−1} + ε_t, ε_t ~ Normal(0, σ). +-- Initial distribution: x_0 ~ Normal(0, σ / √(1 − ϕ²)) (the stationary +-- distribution, finite when |ϕ| < 1). +-- +-- [日本語]: 引数 @phi@ は AR 係数、 @sigma@ は innovation の sd。 N 個の +-- latent 状態 x_0 .. x_{N-1} を非中心化パラメタ化で sample する: +-- [English]: The @phi@ argument is the AR coefficient, and @sigma@ is +-- the innovation's sd. Samples N latent states x_0 .. x_{N-1} using +-- non-centered parameterization: +-- +-- raw_t ~ Normal(0, 1) +-- x_t = phi * x_{t-1} + sigma * raw_t (t > 0) +-- x_0 = (sigma / √(1 - ϕ²)) * raw_0 +-- +-- [日本語]: 戻り値: x_0 .. x_{N-1} の latent 値リスト ([a])。 各 raw_t は +-- @_raw@、 x_t 自体は派生量 @_@ として保存。 +-- [English]: Returns: the list of latent values x_0 .. x_{N-1} ([a]). +-- Each raw_t is stored as @_raw@, and x_t itself as the derived +-- quantity @_@. +-- +-- [日本語]: |ϕ| ≥ 1 のフォールバック: 初期 sd を sigma に置き換える。 +-- [English]: Fallback for |ϕ| ≥ 1: replaces the initial sd with sigma. +ar1Latent :: forall a. (Floating a, Ord a) + => Text -> Int -> a -> a -> Model a [a] +ar1Latent name nT phi sigma + | nT < 1 = error "ar1Latent: length must be >= 1" + | otherwise = do + raws <- mapM + (\t -> sample (name <> "_raw" <> T.pack (show t)) (Normal 0 1)) + [0 .. nT - 1] + let phi2 = phi * phi + stat = if phi2 < 1 + then sigma / sqrt (1 - phi2) + else sigma -- フォールバック + -- Phase 38: scanl で xs を先に組み立てると、 各 x_t の Track が + -- {x_raw0, …, x_raw_t} という遠い親集合を保持してしまい、 後で + -- deterministic 登録しても下流の親が plate-style にならない。 + -- 各 step で deterministic の戻り値 (det 名で再ラベルされた Track) + -- を次の step に渡す monadic recursion で組む。 + x0 <- deterministic (name <> "_0") (stat * head raws) + let chain _ [] = return [] + chain xPrev ((t, rt):rest) = do + xt <- deterministic + (name <> "_" <> T.pack (show t)) + (phi * xPrev + sigma * rt) + xs' <- chain xt rest + return (xt : xs') + xs' <- chain x0 (zip [(1 :: Int) .. ] (tail raws)) + return (x0 : xs') + +-- | [日本語]: 非中心化 (non-centered) 正規分布。 +-- [English]: A non-centered normal distribution. +-- +-- [日本語]: @x ~ Normal(loc, scale)@ を直接サンプリングする代わりに、 +-- [English]: Instead of sampling @x ~ Normal(loc, scale)@ directly, it +-- expands to: +-- +-- > raw <- sample (name <> "_raw") (Normal 0 1) +-- > deterministic name (loc + scale * raw) +-- +-- [日本語]: loc / scale が他の latent に依存するとき、 centered +-- パラメタ化は HMC の posterior が病的になりやすいので、 それを +-- 緩和するヘルパ。 Neal's funnel が代表例。 +-- [English]: When loc / scale depend on other latents, centered +-- parameterization tends to make HMC's posterior pathological; this +-- helper mitigates that. Neal's funnel is the canonical example. +-- +-- [日本語]: 戻り値は constrained な値 @loc + scale * raw@。 Chain には +-- @_raw@ (latent) と @@ (derived) の両方が保存される。 +-- [English]: Returns the constrained value @loc + scale * raw@. Both +-- @_raw@ (the latent) and @@ (the derived quantity) are +-- stored in the Chain. +nonCenteredNormal :: Num a => Text -> a -> a -> Model a a +nonCenteredNormal name loc scale = do + raw <- sample (name <> "_raw") (Normal 0 1) + deterministic name (loc + scale * raw) + +-- | [日本語]: 'glmmRandomIntercept' の GLMM family。 [English]: The GLMM family for 'glmmRandomIntercept'. +data GlmmFamily + = GlmmGaussian -- ^ [日本語]: 連続 y、 残差 SD @sigma@ も sample される。 [English]: continuous y; the residual SD @sigma@ is also sampled. + | GlmmBinomial -- ^ [日本語]: 0/1 y、 Bernoulli(σ(η))。 [English]: 0/1 y, Bernoulli(σ(η)). + | GlmmPoisson -- ^ [日本語]: 非負整数 y、 Poisson(exp η)。 [English]: non-negative integer y, Poisson(exp η). + deriving (Show, Eq) + +-- | [日本語]: Random intercept GLMM helper。 +-- [English]: A random-intercept GLMM helper. +-- +-- [日本語]: `y ~ X β + u_{group(i)} + (error)` を 1 関数で組み立てる: +-- [English]: Assembles `y ~ X β + u_{group(i)} + (error)` in a single +-- function: +-- +-- - 固定効果 @β_k ~ Normal(0, 5)@ (p 個) / fixed effects @β_k ~ Normal(0, 5)@ (p of them) +-- - 群レベル SD @τ_u ~ HalfNormal(5)@ / group-level SD @τ_u ~ HalfNormal(5)@ +-- - 群効果 @u_j ~ Normal(0, τ_u)@ (nG 個、 centered パラメタ化。 +-- 群数大 / 群内 N 小なら別途 'nonCenteredNormal' を直接使う) / +-- group effects @u_j ~ Normal(0, τ_u)@ (nG of them, centered +-- parameterization; for many groups / small within-group N, use +-- 'nonCenteredNormal' directly instead) +-- - family に応じた観測 / observation depending on the family: +-- - Gaussian: 残差 @σ ~ Exp(1)@ を sample → @y ~ Normal(X β + u_j, σ)@ / +-- sample the residual @σ ~ Exp(1)@ → @y ~ Normal(X β + u_j, σ)@ +-- - Binomial: @y ~ Bernoulli(σ(X β + u_j))@、 y は 0/1 / +-- @y ~ Bernoulli(σ(X β + u_j))@, y is 0/1 +-- - Poisson: @y ~ Poisson(exp(X β + u_j))@、 y は非負整数 / +-- @y ~ Poisson(exp(X β + u_j))@, y is a non-negative integer +-- +-- [日本語]: 観測は単一の構造化ブロック @observeLMR \"y\"@ として発行される +-- (PyMC/Stan と同じく 1 ベクトル化観測ノード。 旧実装は per-obs @y_i@ を +-- n 個展開)。 固定効果は密設計行列・群効果は gather で表現するので +-- vec-tape ハイブリッド gradADU の高速経路に乗る。 chain 上の latent 名: +-- @beta_0, …, beta_{p-1}, tau_u, u_0, …, u_{nG-1}, sigma?@. +-- [English]: The observation is emitted as a single structured block +-- @observeLMR \"y\"@ (one vectorized observation node, matching +-- PyMC/Stan; the earlier implementation expanded per-obs @y_i@ into n +-- nodes). Fixed effects are represented with a dense design matrix and +-- group effects with a gather, so this rides the fast path of the +-- vec-tape hybrid gradADU. Latent names in the chain: +-- @beta_0, …, beta_{p-1}, tau_u, u_0, …, u_{nG-1}, sigma?@. +-- +-- [日本語]: 個別 (random slope や non-centered) が必要ならパターン 5 +-- (random slope) / 形式 C (non-centered) を直接書く方が柔軟。 本 helper +-- は最頻ユースケース 「固定効果 + 群別切片」 専用の shorthand。 +-- [English]: For custom needs (random slopes, non-centered), writing +-- Pattern 5 (random slope) / Form C (non-centered) directly is more +-- flexible. This helper is shorthand dedicated to the most common use +-- case, "fixed effects + per-group intercept." +glmmRandomIntercept + :: forall a. (Floating a, Ord a) + => GlmmFamily -- ^ [日本語]: 尤度の family。 [English]: the likelihood family. + -> [[Double]] -- ^ [日本語]: 固定効果 design X (n × p)、 切片は手で 1 列追加すること。 [English]: the fixed-effect design X (n × p); add an intercept column by hand. + -> [Int] -- ^ [日本語]: 各観測の group id (0..nG-1)。 [English]: each observation's group id (0..nG-1). + -> [Double] -- ^ [日本語]: 観測 y (length n)。 [English]: the observed y (length n). + -> Model a () +glmmRandomIntercept fam xRows gids ys = do + let n = length ys + p = if null xRows then 0 else length (head xRows) + nG = if null gids then 0 else maximum gids + 1 + -- 固定効果 + betas <- forM [0 .. p - 1] $ \k -> + sample (T.pack ("beta_" ++ show k)) (Normal 0 5) + -- 群レベル SD + tauU <- sample "tau_u" (HalfNormal 5) + -- 群別切片を第一級ランダム効果値として宣言 (Phase 54.4c)。 reNormal が + -- u_0..u_{nG-1} ~ Normal(0, tauU) を sample しつつスケール名 "tau_u" を構造に + -- 載せるので、 観測に `at` で gather すると compileGradU の **解析 prior 勾配** + -- 経路に乗り、 prior の O(nG) スカラ ad が排除される。 + u <- reNormal "u" nG "tau_u" tauU + -- Gaussian のみ残差 SD + _mSig <- case fam of + GlmmGaussian -> Just <$> sample "sigma" (Exponential 1) + _ -> return Nothing + -- 観測は単一の構造化ブロック (observeLMR) として発行する (Phase 54.4a)。 + -- η_i = Σ_k β_k X_ik + u_{g(i)} を固定効果 (密設計行列) + 群効果 (gather) で + -- 表現するので、 vec-tape ハイブリッド gradADU の高速経路に乗る。 PyMC/Stan と + -- 同じく観測は 1 ベクトル化ノード "y" (旧: per-obs y_i を n 個展開)。 + let betaNames = [ T.pack ("beta_" ++ show k) | k <- [0 .. p - 1] ] + reffs = [ u `at` gids ] + lmFam = case fam of + GlmmGaussian -> LMGaussian "sigma" + GlmmBinomial -> LMBernoulli + GlmmPoisson -> LMPoisson + -- betas/n は名前参照ゆえ値は使わないが、 latent 宣言として必要。 + _ <- pure (betas, n) + observeLMR "y" betaNames xRows reffs lmFam ys + +-- | Dirichlet distribution (analogous to PyMC's @pm.Dirichlet@), expanded +-- via stick-breaking into a +-- [日本語]: latent ベクトル。 [English]: latent vector. +-- +-- [日本語]: 引数: +-- [English]: Arguments: +-- +-- - @name@ : ベース名。 展開後は @_b@ (i=0..K-2) が Beta 由来の +-- 棒折り変数、 @_@ (i=0..K-1) が deterministic で +-- 記録された π 成分。 / +-- the base name. After expansion, @_b@ (i=0..K-2) +-- are the Beta-derived stick-breaking variables, and +-- @_@ (i=0..K-1) are the π components recorded as +-- deterministics. +-- - @alphas@ : 集中度ベクトル α = (α_1,...,α_K)。 長さ K ≥ 2。 / +-- the concentration vector α = (α_1,...,α_K), length K ≥ 2. +-- +-- [日本語]: アルゴリズム: +-- k = 1..K-1 で β_k ~ Beta(α_k, Σ_{j>k} α_j) を sample する。 +-- π_1 = β_1, π_k = β_k Π_{jk} α_j). Then +-- π_1 = β_1, π_k = β_k Π_{j Text -> [a] -> Model a [a] +dirichlet name alphas = do + let k = length alphas + if k < 2 + then error "dirichlet: 長さ 2 未満のベクトルは未対応" + else do + let -- α_k+1..K の累積和 (右から)。長さ K (最後の要素は 0) + tailSums = scanr (+) 0 alphas + -- β_0..β_{K-2} を sample + betas <- mapM + (\i -> sample (name <> "_b" <> T.pack (show i)) + (Beta (alphas !! i) (tailSums !! (i + 1)))) + [0 .. k - 2] + -- 残り棒の累積積 prods[i] = Π_{j acc * (1 - b)) (1 :: a) betas + -- π_i = β_i * prods[i] for i < K-1, π_{K-1} = prods[K-1] + pis = [ if i < length betas + then (betas !! i) * (prods !! i) + else prods !! i + | i <- [0 .. k - 1] ] + -- 各 π_i を deterministic として保存し戻り値にも返す + mapM (\(i, p) -> + deterministic (name <> "_" <> T.pack (show i)) p) + (zip [0 :: Int ..] pis) + +-- | Increasing cuts helper for 'OrderedLogistic' / 'OrderedProbit'. +-- [日本語]: @c_1 = c_min@、 @c_k = c_{k-1} + d_k@ with +-- @d_k ~ HalfNormal(scale)@ により自動的に increasing 列を保証する。 +-- [English]: @c_1 = c_min@, @c_k = c_{k-1} + d_k@ with +-- @d_k ~ HalfNormal(scale)@ automatically guarantees an increasing +-- sequence. +-- +-- [日本語]: 戻り値は長さ @nCuts@ の Track が通る deterministic 値の列 +-- (@name_c_0@, …, @name_c_{nCuts-1}@)。 各 @d_k@ は @name_d_k@ で +-- latent として登録される。 cuts は OrderedLogistic / OrderedProbit に +-- そのまま渡せる。 +-- [English]: Returns a length-@nCuts@ list of deterministic values that +-- Track passes through (@name_c_0@, …, @name_c_{nCuts-1}@). Each @d_k@ +-- is registered as a latent under @name_d_k@. The cuts can be passed +-- directly to OrderedLogistic / OrderedProbit. +-- +-- [日本語]: DAG-safe pattern: monadic recursion で @deterministic@ の +-- 戻り値 (det 名で relabel された Track) を次 step に渡すことで +-- plate-style の親集合を保つ。 +-- [English]: A DAG-safe pattern: threads @deterministic@'s return value +-- (a Track relabeled under the deterministic's name) into the next step +-- via monadic recursion, preserving a plate-style parent set. +orderedCuts :: forall a. (Floating a, Ord a) + => Text -- ^ [日本語]: ベース名。 [English]: the base name. + -> Int -- ^ [日本語]: カット数 K-1 (≥ 1)。 [English]: the number of cuts K-1 (≥ 1). + -> a -- ^ [日本語]: 最小値 c_min。 [English]: the minimum value c_min. + -> a -- ^ [日本語]: 増分の HalfNormal スケール。 [English]: the HalfNormal scale for the increments. + -> Model a [a] +orderedCuts name nCuts cMin scale + | nCuts < 1 = error "orderedCuts: nCuts < 1 は未対応" + | otherwise = do + -- c_1 = c_min (定数を deterministic で登録、 Track 透過のため) + c1 <- deterministic (name <> "_c_1") cMin + -- c_2, ..., c_nCuts を monadic recursion で順に作る + -- chain prev i: 現在の前 cut Track が prev、 次に作るのは index i (1-based) + let chain prev i acc + | i > nCuts = return (reverse acc) + | otherwise = do + d <- sample (name <> "_d_" <> T.pack (show i)) + (HalfNormal scale) + ci <- deterministic (name <> "_c_" <> T.pack (show i)) + (prev + d) + chain ci (i + 1) (ci : acc) + rest <- chain c1 2 [] + return (c1 : rest) + +-- | [日本語]: Dirichlet Process の有限近似 stick-breaking。 +-- [English]: A finite stick-breaking approximation of a Dirichlet +-- Process. +-- +-- [日本語]: @β_k ~ Beta(1, α)@ for @k = 1, …, T-1@、 重み +-- @π_k = β_k Π_{j Text -- ^ [日本語]: ベース名。 [English]: the base name. + -> Int -- ^ [日本語]: truncation level T (≥ 2)。 [English]: the truncation level T (≥ 2). + -> a -- ^ [日本語]: concentration α (> 0)。 [English]: the concentration α (> 0). + -> Model a [a] +dpStickBreaking name truncT alpha + | truncT < 2 = error "dpStickBreaking: truncation level < 2 は未対応" + | otherwise = do + -- β_1, …, β_{T-1} を sample + betas <- mapM + (\i -> sample (name <> "_b_" <> T.pack (show i)) + (Beta 1 alpha)) + [1 .. truncT - 1] + -- 累積積 stick_k = Π_{j "_stick_1") (1 :: a) + let stickChain prev i acc + | i > truncT = return (reverse acc) + | otherwise = do + let bIdx = i - 1 + beta = betas !! (bIdx - 1) -- 1-based β_{i-1} + sNext <- deterministic + (name <> "_stick_" <> T.pack (show i)) + (prev * (1 - beta)) + stickChain sNext (i + 1) (sNext : acc) + restSticks <- stickChain stick1 2 [] + let sticks = stick1 : restSticks -- 長さ T + -- π_k = β_k * stick_k for k < T、 π_T = stick_T + pis <- mapM + (\i -> + let stickI = sticks !! (i - 1) + piVal = if i < truncT + then (betas !! (i - 1)) * stickI + else stickI + in deterministic (name <> "_pi_" <> T.pack (show i)) piVal) + [1 .. truncT] + return pis + +-- | [日本語]: Hidden Markov Model 用の遷移行列 + 初期分布 prior helper。 +-- K 状態の HMM について、 初期分布 π_0 と K×K 遷移行列の各行に +-- Dirichlet(α, …, α) prior を置く。 +-- [English]: A helper for Hidden Markov Model transition-matrix + initial- +-- distribution priors. For a K-state HMM, places a Dirichlet(α, …, α) +-- prior on the initial distribution π_0 and on each row of the K×K +-- transition matrix. +-- +-- [日本語]: 戻り値は @(π_0, transitions)@: +-- [English]: Returns @(π_0, transitions)@: +-- +-- - @π_0@: 長さ K の確率列 (Σ = 1)、 @name_pi0_@ で deterministic 登録 / +-- a length-K probability vector (Σ = 1), registered as a deterministic +-- under @name_pi0_@ +-- - @transitions@: 長さ K のリスト、 i 番目は遷移行列 i 行目 +-- (@name_trans_i_@ で deterministic) / +-- a length-K list whose i-th element is row i of the transition matrix +-- (@name_trans_i_@ as a deterministic) +-- +-- [日本語]: 離散状態列は __直接 latent としない__ (NUTS は離散変数を扱えない)。 +-- 代わりに、 ユーザは観測列 @y@ の emission log-prob 行列を計算し、 +-- 'hmmForwardLogLik' で状態列をマージナル化した周辺対数尤度を求め、 +-- 'potential' で組み込む形を取る。 +-- [English]: The discrete state sequence is __never a direct latent__ +-- (NUTS cannot handle discrete variables). Instead, users compute the +-- emission log-prob matrix for the observation sequence @y@, obtain the +-- marginal log-likelihood by marginalizing out the state sequence via +-- 'hmmForwardLogLik', and incorporate it with 'potential'. +-- +-- [日本語]: 内部実装は既存 'dirichlet' helper を K+1 回呼ぶだけ。 すべて +-- deterministic chain で DAG-safe。 +-- [English]: Internally this simply calls the existing 'dirichlet' +-- helper K+1 times; everything is DAG-safe via a deterministic chain. +hmmLatent :: forall a. (Floating a, Ord a) + => Text -- ^ [日本語]: ベース名。 [English]: the base name. + -> Int -- ^ [日本語]: 状態数 K (≥ 2)。 [English]: the number of states K (≥ 2). + -> a -- ^ [日本語]: Dirichlet concentration α (> 0、 1 で uniform prior)。 [English]: the Dirichlet concentration α (> 0; 1 gives a uniform prior). + -> Model a ([a], [[a]]) +hmmLatent name k alpha + | k < 2 = error "hmmLatent: K < 2 は未対応" + | otherwise = do + pi0 <- dirichlet (name <> "_pi0") (replicate k alpha) + trans <- mapM + (\i -> dirichlet (name <> "_trans_" <> T.pack (show i)) + (replicate k alpha)) + [0 .. k - 1] + return (pi0, trans) + +-- | [日本語]: HMM forward algorithm marginal log-likelihood。 +-- 'Hanalyze.Model.HBM.Util' へ純粋移設済 (ここは re-export +-- のみ・API 不変)。 用法は従来の @'potential' nm (hmmForwardLogLik ...)@ に加え、 +-- Normal emission の場合は 'HmmForwardNormal' + 'observeMV' が推奨 +-- (勾配コンパイラが forward-backward の閉形式随伴を使えるため大幅に速い)。 +-- [English]: The HMM forward algorithm's marginal log-likelihood. Purely +-- relocated to 'Hanalyze.Model.HBM.Util' (this is a re-export +-- only; the API is unchanged). Besides the traditional usage +-- @'potential' nm (hmmForwardLogLik ...)@, for Normal emissions, +-- 'HmmForwardNormal' + 'observeMV' is recommended (much faster, since +-- the gradient compiler can use the closed-form forward-backward +-- adjoint). + +-- --------------------------------------------------------------------------- +-- 構造検査 +-- --------------------------------------------------------------------------- + +data NodeKind = LatentN | ObservedN Int | DeterministicN + | DataN Int -- ^ [日本語]: データ slot ('dataNamed' / 'dataNamedIx')。 + -- Int = 長さ。 PyMC の pm.Data (ConstantData) 相当。 + -- [English]: A data slot ('dataNamed' / + -- 'dataNamedIx'); the @Int@ is the length, + -- equivalent to PyMC's @pm.Data@ + -- (@ConstantData@). + deriving (Show, Eq) + +data Node = Node + { nodeName :: Text + , nodeKind :: NodeKind + , nodeDist :: Text -- 分布名 (e.g. "Normal") + , nodeDeps :: Set Text -- 直接の親 (依存変数) + , nodePlates :: [Text] -- Phase 40: plate スタック (外側から内側、 空 = 任意の plate に属さない) + } deriving (Show) + +-- | Walk the model with placeholder zeros and collect 'Node' metadata. +-- [日本語]: 依存関係 ('nodeDeps') は @extractDeps@ を使うこと (placeholder +-- 走査では取れない)。 +-- [English]: For dependencies ('nodeDeps'), use @extractDeps@ instead — a +-- placeholder walk cannot recover them. +collectNodes :: forall r. ModelP r -> [Node] +collectNodes m = go m [] + where + go :: Model Double r -> [Node] -> [Node] + go (Pure _) acc = reverse acc + go (Free (Sample n d k)) acc = + go (k 0) (Node n LatentN (distName d) Set.empty [] : acc) + go (Free (Observe n d ys next)) acc = + go next (Node n (ObservedN (length ys)) (distName d) Set.empty [] : acc) + go (Free (ObserveLM n _ _ _ fam ys next)) acc = + go next (Node n (ObservedN (length ys)) (lmFamilyName fam) Set.empty [] : acc) + go (Free (Potential _ _ next)) acc = go next acc -- Node 表示には含めない + go (Free (Deterministic _ v k)) acc = go (k v) acc + go (Free (Data n ys k)) acc = + go (k (ys, ys)) (Node n (DataN (length ys)) "Data" Set.empty [] : acc) + go (Free (DataIx n is k)) acc = + go (k is) (Node n (DataN (length is)) "DataIx" Set.empty [] : acc) + go (Free (PlateBegin _ _ next)) acc = go next acc -- Phase 40: 透過 + go (Free (PlateEnd next)) acc = go next acc + +sampleNames :: ModelP r -> [Text] +sampleNames m = [nodeName n | n <- collectNodes m, nodeKind n == LatentN] + +-- | [日本語]: モデル中の 'Data' slot を (名前, placeholder が空か) で列挙する。 +-- 同名 slot が複数回現れる場合は 1 entry に集約し、 __いずれかが空なら空扱い__ +-- (束縛層の loud error 判定は保守側に倒す)。 @DataIx@ slot は 'dataIxSlots'。 +-- [English]: Enumerates the model's 'Data' slots as (name, is the +-- placeholder empty). If the same-named slot occurs multiple times, they +-- are collapsed to one entry — __treated as empty if any occurrence is empty__ +-- (erring conservative for the binding layer's loud-error +-- check). Use 'dataIxSlots' for @DataIx@ slots. +dataSlots :: forall r. ModelP r -> [(Text, Bool)] +dataSlots m = dedupSlots (go m []) + where + go :: Model Double r -> [(Text, Bool)] -> [(Text, Bool)] + go (Pure _) acc = reverse acc + go (Free (Sample _ _ k)) acc = go (k 0) acc + go (Free (Observe _ _ _ next)) acc = go next acc + go (Free (ObserveLM _ _ _ _ _ _ next)) acc = go next acc + go (Free (Potential _ _ next)) acc = go next acc + go (Free (Deterministic _ v k)) acc = go (k v) acc + go (Free (Data n ys k)) acc = go (k (ys, ys)) ((n, null ys) : acc) + go (Free (DataIx _ is k)) acc = go (k is) acc + go (Free (PlateBegin _ _ next)) acc = go next acc + go (Free (PlateEnd next)) acc = go next acc + +-- | [日本語]: モデル中の @DataIx@ slot を (名前, placeholder が空か) で列挙する。 [English]: Enumerates the model's @DataIx@ slots as (name, is the placeholder empty). +dataIxSlots :: forall r. ModelP r -> [(Text, Bool)] +dataIxSlots m = dedupSlots (go m []) + where + go :: Model Double r -> [(Text, Bool)] -> [(Text, Bool)] + go (Pure _) acc = reverse acc + go (Free (Sample _ _ k)) acc = go (k 0) acc + go (Free (Observe _ _ _ next)) acc = go next acc + go (Free (ObserveLM _ _ _ _ _ _ next)) acc = go next acc + go (Free (Potential _ _ next)) acc = go next acc + go (Free (Deterministic _ v k)) acc = go (k v) acc + go (Free (Data _ ys k)) acc = go (k (ys, ys)) acc + go (Free (DataIx n is k)) acc = go (k is) ((n, null is) : acc) + go (Free (PlateBegin _ _ next)) acc = go next acc + go (Free (PlateEnd next)) acc = go next acc + +-- | [日本語]: slot 列挙の重複集約 (先頭出現順を保ち、 空 flag は OR)。 [English]: Deduplicates a slot enumeration, preserving first-occurrence order and OR-ing the empty flag. +dedupSlots :: [(Text, Bool)] -> [(Text, Bool)] +dedupSlots xs = + [ (n, or [ e | (n', e) <- xs, n' == n ]) + | n <- nub (map fst xs) ] + diff --git a/src/Hanalyze/Model/HBM/Sampling.hs b/src/Hanalyze/Model/HBM/Sampling.hs new file mode 100644 --- /dev/null +++ b/src/Hanalyze/Model/HBM/Sampling.hs @@ -0,0 +1,438 @@ +{-# LANGUAGE ScopedTypeVariables #-} + +-- | +-- Module : Hanalyze.Model.HBM.Sampling +-- Description : HBM の分布サンプリング (事前/事後予測用) +-- Copyright : (c) 2026 Aelysce Project (Toshiaki Honda) +-- License : BSD-3-Clause +-- +-- [日本語]: 分布からのサンプリング (事前/事後予測用) を分離。 +-- +-- 'Distribution' / 'HBM.Util' の上層。 PrimMonad + mwc-random に依存し、 +-- mwc-random が直接提供しない分布 (Cauchy, HalfCauchy, Weibull, …) は +-- 逆 CDF 法 / rejection でここに実装する。 NUTS の per-draw 経路には乗らず +-- (事前/事後予測のみ)、 性能ホットではない。 +-- [English]: Sampling from distributions (for prior/posterior predictive +-- checks), factored out as its own module. +-- +-- Sits above 'Distribution' / 'HBM.Util'. Depends on PrimMonad + +-- mwc-random; distributions mwc-random doesn't provide directly (Cauchy, +-- HalfCauchy, Weibull, …) are implemented here via inverse-CDF / +-- rejection sampling. This does not sit on NUTS's per-draw path (used +-- only for prior/posterior predictive checks), so it is not +-- performance-hot. +module Hanalyze.Model.HBM.Sampling + ( sampleDist + , sampleMvDist + , sampleObsRep + ) where + +import Control.Monad (replicateM) +import Data.List (zip4) +import Control.Monad.Primitive (PrimMonad, PrimState) +import qualified System.Random.MWC as MWCBase +import qualified System.Random.MWC.Distributions as MWC +import System.Random.MWC (Gen) + +import Hanalyze.Model.HBM.Util (choleskyL, chunksOf, gpRBFCovList) +import Hanalyze.Model.HBM.Distribution (Distribution (..), phiCdfA) + +-- --------------------------------------------------------------------------- +-- 分布からのサンプリング (事前/事後予測用) +-- --------------------------------------------------------------------------- + +-- | Draw a single sample from a 'Distribution Double'. +-- [日本語]: 事前予測サンプリング、事後予測サンプリング、観測値の生成に使う。 +-- +-- mwc-random が直接提供しない分布はここで実装する (Cauchy, HalfCauchy, etc.)。 +-- [English]: Used for prior predictive sampling, posterior predictive +-- sampling, and generating observed values. +-- +-- Distributions mwc-random doesn't provide directly are implemented here +-- (Cauchy, HalfCauchy, etc.). +sampleDist :: forall m. PrimMonad m => Distribution Double -> Gen (PrimState m) -> m Double +sampleDist (Normal mu sig) gen = MWC.normal mu sig gen +sampleDist (Exponential rate) gen = do + u <- MWCBase.uniform gen :: m Double + return (-log u / rate) +sampleDist (Gamma shape rate) gen = + -- mwc-random の gamma は scale パラメタ化なので 1/rate を渡す + MWC.gamma shape (1 / rate) gen +sampleDist (Beta a b) gen = do + x <- MWC.gamma a 1 gen + y <- MWC.gamma b 1 gen + return (x / (x + y)) +sampleDist (Poisson lam) gen = samplePoissonKnuth lam gen +sampleDist (Binomial n p) gen = do + -- n 回のベルヌーイ試行 + let go 0 acc = return acc + go k acc = do + u <- MWCBase.uniform gen :: m Double + go (k - 1) (if u < p then acc + 1 else acc) + fmap fromIntegral (go n (0 :: Int)) +sampleDist (Uniform lo hi) gen = do + u <- MWCBase.uniform gen :: m Double + return (lo + u * (hi - lo)) +sampleDist (StudentT df mu sig) gen = do + -- t = mu + sig * Normal(0,1) / sqrt(Chi2(df) / df) + z <- MWC.standard gen + chi2 <- MWC.gamma (df / 2) 2 gen -- Chi2(df) = Gamma(df/2, scale=2) + return (mu + sig * z / sqrt (chi2 / df)) +sampleDist (Cauchy loc sc) gen = do + u <- MWCBase.uniform gen :: m Double + return (loc + sc * tan (pi * (u - 0.5))) +sampleDist (HalfNormal sig) gen = do + z <- MWC.standard gen + return (abs (sig * z)) +sampleDist (HalfCauchy sc) gen = do + u <- MWCBase.uniform gen :: m Double + return (sc * abs (tan (pi * (u - 0.5)))) +sampleDist (LogNormal mu sig) gen = do + z <- MWC.standard gen + return (exp (mu + sig * z)) +sampleDist (Bernoulli p) gen = do + u <- MWCBase.uniform gen :: m Double + return (if u < p then 1.0 else 0.0) +sampleDist (Categorical probs) gen = do + u <- MWCBase.uniform gen :: m Double + let total = sum probs + go _ [] = fromIntegral (length probs - 1) + go acc (p:ps) = + let acc' = acc + p / total + in if u < acc' then 0 else 1 + go acc' ps + return (go 0 probs) +sampleDist (Mixture ws comps) gen + | null ws || length ws /= length comps = return (0/0) -- NaN: 不正 + | otherwise = do + -- 1) 重みに比例して成分 k を選ぶ + u <- MWCBase.uniform gen :: m Double + let total = sum ws + pickIdx _ [] = length ws - 1 + pickIdx acc (w:rest) = + let acc' = acc + w / total + in if u < acc' then 0 else 1 + pickIdx acc' rest + k = pickIdx 0 ws + -- 2) 選んだ成分からサンプリング + sampleDist (comps !! k) gen +sampleDist (Truncated d mLo mHi) gen = + -- 単純なリジェクション・サンプリング (範囲が極めて狭いと収束遅い) + let inRange y = case (mLo, mHi) of + (Just lo, _ ) | y < lo -> False + (_, Just hi) | y > hi -> False + _ -> True + tryOnce maxAttempts + | maxAttempts <= 0 = return (0/0) -- 諦め + | otherwise = do + y <- sampleDist d gen + if inRange y then return y else tryOnce (maxAttempts - 1) + in tryOnce (10000 :: Int) +sampleDist MvNormal{} _ = + error "MvNormal: observation-only — 'sample' 経由でのドローは未対応" +sampleDist MvNormalChol{} _ = + error "MvNormalChol: observation-only — 'sample' 経由でのドローは未対応" +sampleDist MvNormalGpRBF{} _ = + error "MvNormalGpRBF: observation-only — 'sample' 経由でのドローは未対応" +sampleDist HmmForwardNormal{} _ = + error "HmmForwardNormal: observation-only — 'sample' 経由でのドローは未対応" +sampleDist ArmaNormal{} _ = + error "ArmaNormal: observation-only — 'sample' 経由でのドローは未対応" +sampleDist GradedResponseIrt{} _ = + error "GradedResponseIrt: observation-only — 'sample' 経由でのドローは未対応" +sampleDist Multinomial{} _ = + error "Multinomial: observation-only — 'sample' 経由でのドローは未対応" +sampleDist (InverseGamma alpha beta) gen = do + -- 1 / Gamma(α, rate=β) = 1 / Gamma(α, scale=1/β) + y <- MWC.gamma alpha (1 / beta) gen + return (1 / y) +sampleDist (Weibull kShape lam) gen = do + -- 逆 CDF 法: x = λ (-log(1-u))^(1/k) + u <- MWCBase.uniform gen :: m Double + return (lam * ((-log (1 - u)) ** (1 / kShape))) +sampleDist (Pareto alpha xm) gen = do + -- 逆 CDF 法: x = x_m / u^(1/α) + u <- MWCBase.uniform gen :: m Double + return (xm / (u ** (1 / alpha))) +sampleDist (BetaBinomial n alpha beta) gen = do + -- p ~ Beta(α, β); k ~ Binomial(n, p) + p <- sampleDist (Beta alpha beta) gen + sampleDist (Binomial n p) gen +sampleDist (VonMises mu kappa) gen = do + -- Best-Fisher の rejection sampler + let a = 1 + sqrt (1 + 4 * kappa * kappa) + b = (a - sqrt (2 * a)) / (2 * kappa) + r = (1 + b * b) / (2 * b) + tryOnce = do + u1 <- MWCBase.uniform gen :: m Double + let z = cos (pi * u1) + f = (1 + r * z) / (r + z) + c = kappa * (r - f) + u2 <- MWCBase.uniform gen :: m Double + if c * (2 - c) - u2 > 0 || log (c / u2) + 1 - c >= 0 + then do + u3 <- MWCBase.uniform gen :: m Double + let sign = if u3 - 0.5 < 0 then (-1.0) else 1.0 + return (mu + sign * acos f) + else tryOnce + tryOnce +sampleDist (ZeroInflatedPoisson psi lam) gen = do + u <- MWCBase.uniform gen :: m Double + if u < psi + then return 0 + else samplePoissonKnuth lam gen +sampleDist (ZeroInflatedBinomial n psi p) gen = do + u <- MWCBase.uniform gen :: m Double + if u < psi + then return 0 + else sampleDist (Binomial n p) gen +sampleDist (SkewNormal mu sig alpha) gen = do + -- Henze 1986: δ = α/√(1+α²), X = μ + σ(δ|U₀| + √(1-δ²)U₁) + let delta = alpha / sqrt (1 + alpha * alpha) + u0 <- MWC.standard gen + u1 <- MWC.standard gen + return (mu + sig * (delta * abs u0 + sqrt (1 - delta * delta) * u1)) +sampleDist (Logistic mu s) gen = do + -- 逆 CDF: X = μ + s · log(u/(1-u)) + u <- MWCBase.uniform gen :: m Double + return (mu + s * log (u / (1 - u))) +sampleDist (Gumbel mu beta) gen = do + -- 逆 CDF: X = μ - β · log(-log u) + u <- MWCBase.uniform gen :: m Double + return (mu - beta * log (- log u)) +sampleDist (AsymmetricLaplace b kappa mu) gen = do + -- 逆 CDF。 pc = κ²/(1+κ²)、 U < pc なら左尾、 そうでなければ右尾 + u <- MWCBase.uniform gen :: m Double + let k2 = kappa * kappa + pc = k2 / (1 + k2) + if u < pc + then return (mu + (kappa / b) * log (u / pc)) + else return (mu - (1 / (b * kappa)) * log ((1 - u) / (1 - pc))) +sampleDist (OrderedLogistic eta cuts) gen = do + -- η と cuts から各カテゴリの確率を計算して Categorical で sample + let sigm x = 1 / (1 + exp (-x)) + probs = catProbs cuts + catProbs [] = [1] + catProbs (c:rest) = sigm (c - eta) : restProbs (sigm (c - eta)) rest + restProbs _ [] = [1 - sigm (last cuts - eta)] + restProbs prev (c:rest) = + let cur = sigm (c - eta) + in (cur - prev) : restProbs cur rest + u <- MWCBase.uniform gen :: m Double + let go acc k [] = realToFrac (k - 1 :: Int) + go acc k (p:ps) = + let acc' = acc + p + in if u < acc' then realToFrac k else go acc' (k + 1) ps + return (go 0 (0 :: Int) probs) +sampleDist (DiscreteUniform lo hi) gen = do + u <- MWCBase.uniform gen :: m Double + let span_ = hi - lo + 1 + k = lo + floor (u * realToFrac span_) + kClip = min hi k -- u が 1 のとき span_ になるのを防ぐ + return (realToFrac kClip) +sampleDist (Geometric p) gen = do + -- 逆 CDF: X = ceil(log U / log(1-p))、 PyMC convention で support から 1 + u <- MWCBase.uniform gen :: m Double + let lq = log (1 - p) + x = ceiling (log u / lq) :: Int + return (realToFrac (max 1 x)) +sampleDist (HyperGeometric nN kK nDraw) gen = do + -- 単純な urn sampling: 各引きで残り成功/失敗の割合から二項 + let loop remN remK remDraw acc + | remDraw <= 0 = return (realToFrac (acc :: Int)) + | otherwise = do + u <- MWCBase.uniform gen :: m Double + let pSucc = realToFrac remK / realToFrac remN :: Double + pick = if u < pSucc then 1 else 0 + loop (remN - 1) (remK - pick) (remDraw - 1) (acc + pick) + loop nN kK nDraw 0 +sampleDist (ZeroInflatedNegativeBinomial psi mu alpha) gen = do + u <- MWCBase.uniform gen :: m Double + if u < psi + then return 0 + else sampleDist (NegativeBinomial mu alpha) gen +sampleDist MvStudentT{} _ = + error "MvStudentT: observation-only — 'sample' 経由でのドローは未対応 (latent helper を別途用意予定)" +sampleDist DirichletMultinomial{} _ = + error "DirichletMultinomial: observation-only — 'sample' 経由でのドローは未対応" +sampleDist (NegativeBinomial mu alpha) gen = do + -- Gamma-Poisson mixture: λ ~ Gamma(α, β=α/μ); X ~ Poisson(λ) + lam <- MWC.gamma alpha (mu / alpha) gen + samplePoissonKnuth lam gen +sampleDist (Censored d _ _) gen = + -- 元分布から普通にサンプリング (打ち切りは「観測過程」の話で生成側ではない) + sampleDist d gen +sampleDist (Triangular lo c hi) gen = do + -- 逆 CDF。 Fc = (c-lo)/(hi-lo) 未満なら左側、 以上なら右側 + u <- MWCBase.uniform gen :: m Double + let fc = (c - lo) / (hi - lo) + if u < fc + then return (lo + sqrt (u * (hi - lo) * (c - lo))) + else return (hi - sqrt ((1 - u) * (hi - lo) * (hi - c))) +sampleDist (Kumaraswamy a b) gen = do + -- 逆 CDF: x = (1 - (1-u)^{1/b})^{1/a} + u <- MWCBase.uniform gen :: m Double + return ((1 - (1 - u) ** (1 / b)) ** (1 / a)) +sampleDist (Rice nu sig) gen = do + -- Y1 ~ N(ν, σ), Y2 ~ N(0, σ); X = sqrt(Y1²+Y2²) + y1 <- MWC.normal nu sig gen + y2 <- MWC.normal 0 sig gen + return (sqrt (y1 * y1 + y2 * y2)) +sampleDist Wishart{} _ = + error "Wishart: observation-only — 'sample' 経由でのドローは未対応 (Bartlett decomp の latent helper を別途用意予定)" +sampleDist (Bound d mLo mHi) gen = + -- Bound は Truncated とほぼ同じ。 sample は rejection。 + sampleDist (Truncated d mLo mHi) gen +sampleDist (OrderedProbit eta cuts) gen = do + -- η と cuts から各カテゴリの確率を計算して Categorical で sample + let probs = catProbs cuts + catProbs [] = [1] + catProbs (c:rest) = phiCdfA (c - eta) : restProbs (phiCdfA (c - eta)) rest + restProbs _ [] = [1 - phiCdfA (last cuts - eta)] + restProbs prev (c:rest) = + let cur = phiCdfA (c - eta) + in (cur - prev) : restProbs cur rest + u <- MWCBase.uniform gen :: m Double + let go acc k [] = realToFrac (k - 1 :: Int) + go acc k (p:ps) = + let acc' = acc + p + in if u < acc' then realToFrac k else go acc' (k + 1) ps + return (go 0 (0 :: Int) probs) +sampleDist (DiscreteWeibull q beta) gen = do + -- 逆 CDF: k = ceil((log(1-u)/log q)^{1/β}) - 1 + u <- MWCBase.uniform gen :: m Double + let r = log (1 - u) / log q + k0 = ceiling (r ** (1 / beta)) - 1 :: Int + k = max 0 k0 + return (fromIntegral k) + +-- | [日本語]: 多変量分布から 1 観測 (k-vector) を draw する (PPC 用)。 +-- @sampleDist@ (スカラ観測専用 'error') と別経路。 @y = μ + C z@ で z ~ N(0,1)、 +-- C は MvNormal なら @choleskyL Σ@、 MvNormalChol なら scaled Cholesky +-- @M = diag σ · L@ (再分解不要)。 対応外の多変量分布は空リスト (= worker 側で +-- graceful にスキップ。 MvNormal/MvNormalChol に集中)。 +-- [English]: Draws one observation (a k-vector) from a multivariate +-- distribution (for posterior predictive checks). A separate path from +-- @sampleDist@ (which 'error's on scalar-only observations). Computes +-- @y = μ + C z@ with z ~ N(0,1), where C is @choleskyL Σ@ for MvNormal, +-- or the scaled Cholesky @M = diag σ · L@ for MvNormalChol (no +-- re-decomposition needed). Unsupported multivariate distributions +-- return an empty list (gracefully skipped on the worker side; coverage +-- is focused on MvNormal/MvNormalChol). +sampleMvDist :: forall m. PrimMonad m => Distribution Double -> Gen (PrimState m) -> m [Double] +sampleMvDist (MvNormal mu cov) gen = do + let k = length mu + zs <- replicateM k (MWC.standard gen) + pure $ case choleskyL cov of + Just c -> [ (mu !! i) + sum [ ((c !! i) !! j) * (zs !! j) | j <- [0 .. i] ] + | i <- [0 .. k - 1] ] + Nothing -> mu +sampleMvDist (MvNormalChol mu sigma l) gen = do + let k = length mu + m = [ [ (sigma !! i) * ((l !! i) !! j) | j <- [0 .. k - 1] ] | i <- [0 .. k - 1] ] + zs <- replicateM k (MWC.standard gen) + pure [ (mu !! i) + sum [ ((m !! i) !! j) * (zs !! j) | j <- [0 .. i] ] + | i <- [0 .. k - 1] ] +sampleMvDist (MvNormalGpRBF xs alpha rho sigma) gen = -- Phase 95 B-dsl: cov 展開して MvNormal と同じ + sampleMvDist (MvNormal (replicate (length xs) 0) (gpRBFCovList xs alpha rho sigma)) gen +sampleMvDist _ _ = pure [] + +-- | [日本語]: 1 observe ノード分の複製データ (y_rep) をまとめて draw する。 +-- PPC (@sampleYRep@/@epredPIAtHeld@ 等) はこれまで観測分布を問わず ys の要素 +-- ごとに @sampleDist@ (スカラ専用) を呼んでいたため、 多変量分布 (MvNormal/ +-- MvNormalChol、 ys = 1 つの k-vector 観測がフラット化された形) では即座に +-- @error@ していた (07-gp-regr で実測発覚)。 ここで次元 k ごとに ys をチャンク +-- し 'sampleMvDist' に委譲する。 Multinomial は 'sampleMvDist' が未対応の +-- ままなので、 従来どおり @sampleDist@ に委ねる。 +-- [English]: Draws replicated data (y_rep) for one observe node, all at +-- once. PPC helpers (@sampleYRep@/@epredPIAtHeld@, etc.) used to call +-- @sampleDist@ (scalar-only) per element of ys regardless of the +-- observation distribution, which immediately @error@'d for multivariate +-- distributions (MvNormal/MvNormalChol, where ys is a single flattened +-- k-vector observation) — discovered via 07-gp-regr. Here, ys is chunked +-- by dimension k and delegated to 'sampleMvDist'. Multinomial is still +-- unsupported by 'sampleMvDist', so it continues to be delegated to +-- @sampleDist@ as before. +sampleObsRep :: forall m. PrimMonad m + => Gen (PrimState m) -> Distribution Double -> [Double] -> m [Double] +sampleObsRep gen d@(MvNormal mu _) ys = + concat <$> mapM (const (sampleMvDist d gen)) (chunksOf (length mu) ys) +sampleObsRep gen d@(MvNormalChol mu _ _) ys = + concat <$> mapM (const (sampleMvDist d gen)) (chunksOf (length mu) ys) +sampleObsRep gen d@(MvNormalGpRBF xs _ _ _) ys = -- Phase 95 B-dsl + concat <$> mapM (const (sampleMvDist d gen)) (chunksOf (length xs) ys) +sampleObsRep gen (HmmForwardNormal pi0 trans mus sg) ys = do + -- Phase 92 A2 (PPC): 状態列を π_0/遷移行列から draw → Normal(μ_s, σ) で emission。 + -- 観測列全体 = 1 観測なので T = length ys の系列を 1 本生成する。 + let kk = length pi0 + pick ws = do -- 重み ws (非正規化可) からカテゴリを 1 つ draw + let s = sum ws + u <- (* s) <$> MWCBase.uniform gen + let go i acc (w:rest) | null rest || u <= acc + w = pure i + | otherwise = go (i + 1) (acc + w) rest + go i _ [] = pure (max 0 (i - 1)) + go 0 0 ws + stepState s = pick (if s < length trans then trans !! s else replicate kk 1) + emitAt s = do + z <- MWC.standard gen + pure ((if s < length mus then mus !! s else 0) + sg * z) + go' _ 0 acc = pure (reverse acc) + go' s n acc = do + y <- emitAt s + s' <- stepState s + go' s' (n - 1 :: Int) (y : acc) + s0 <- pick pi0 + go' s0 (length ys) [] +sampleObsRep gen (ArmaNormal mu phi theta sg) ys = do + -- Phase 101 A2 (PPC): err_t ~ Normal(0, σ) を draw し、y を前向き再帰で生成 + -- (y_1 = μ+φμ+e_1・y_t = μ + φ·y_{t−1} + θ·e_{t−1} + e_t)。 + -- 観測列全体 = 1 観測なので T = length ys の系列を 1 本生成する。 + let drawE = (sg *) <$> MWC.standard gen + go' _ _ 0 acc = pure (reverse acc) + go' prevY prevE n acc = do + e <- drawE + let y = mu + phi * prevY + theta * prevE + e + go' y e (n - 1 :: Int) (y : acc) + case length ys of + 0 -> pure [] + t -> do + e1 <- drawE + let y1 = mu + phi * mu + e1 + go' y1 e1 (t - 1) [y1] +sampleObsRep gen (GradedResponseIrt thetas ncats deltas gammas) ys = do + -- Phase 101 A3 (PPC): 各 (child, item) の p ベクトルからカテゴリを draw。 + -- 欠測 (−1) 位置は −1 のまま返す (観測の欠測パターンを保存)。 + let nItem = length ncats + rows = chunksOf nItem ys + catPs th nc dl gm = + let kMax = nc - 1 + qs = [ 1 / (1 + exp (negate (dl * (th - gm !! (kk - 1))))) + | kk <- [1 .. kMax] ] + in [ if k == 1 then 1 - head qs + else if k == nc then qs !! (kMax - 1) + else (qs !! (k - 2)) - (qs !! (k - 1)) + | k <- [1 .. nc] ] + pickCat ps = do + u <- MWCBase.uniform gen + let go k acc (w:rest) | null rest || u <= acc + w = pure k + | otherwise = go (k + 1) (acc + w) rest + go k _ [] = pure (max 1 (k - 1)) + go 1 0 ps + drawRow th row = + sequence [ if gr == -1 then pure (-1) + else fromIntegral <$> pickCat (catPs th nc dl gm) + | (nc, dl, gm, gr) <- zip4 ncats deltas gammas row ] + concat <$> sequence [ drawRow th row | (th, row) <- zip thetas rows ] +sampleObsRep gen d ys = mapM (const (sampleDist d gen)) ys + +-- | [日本語]: Knuth のアルゴリズムで Poisson(λ) サンプル。λ < 30 程度なら十分高速。 +-- [English]: Samples Poisson(λ) via Knuth's algorithm. Fast enough for +-- λ around 30 or less. +samplePoissonKnuth :: forall m. PrimMonad m => Double -> Gen (PrimState m) -> m Double +samplePoissonKnuth lam gen = do + let l = exp (-lam) + go k p = do + u <- MWCBase.uniform gen :: m Double + let p' = p * u + if p' < l + then return (fromIntegral k) + else go (k + 1) p' + go 0 (1.0 :: Double) diff --git a/src/Hanalyze/Model/HBM/Track.hs b/src/Hanalyze/Model/HBM/Track.hs new file mode 100644 --- /dev/null +++ b/src/Hanalyze/Model/HBM/Track.hs @@ -0,0 +1,331 @@ +{-# LANGUAGE RankNTypes #-} +{-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE ScopedTypeVariables #-} + +-- | +-- Module : Hanalyze.Model.HBM.Track +-- Description : HBM の依存追跡型 Track (latent 変数への依存伝播) +-- Copyright : (c) 2026 Aelysce Project (Toshiaki Honda) +-- License : BSD-3-Clause +-- +-- [日本語]: 依存追跡型 'Track' を 'Hanalyze.Model.HBM' から分離。 +-- +-- 'Track' は @Floating@ 演算を通して「この値はどの latent 変数に依存するか」を +-- 伝播する型。 'ModelP' をこの型で特殊化することで各 Observe / Deterministic +-- ノードの親集合を自動抽出する ('extractDeps')。 DAG 可視化 (buildModelGraph) +-- の基盤。 +-- +-- 依存は下層 'Hanalyze.Model.HBM.Model' (Node / ModelF / lmParents 等) と +-- '...Distribution' (Distribution / distName) のみ。 評価層 (logJoint 等) には +-- 依存しない (runTrack = logJoint の Track 特殊化は Eval 層に置く)。 +-- [English]: The dependency-tracking type 'Track', factored out of +-- 'Hanalyze.Model.HBM'. +-- +-- 'Track' is a type that propagates "which latent variable does this +-- value depend on" through @Floating@ operations. Specializing 'ModelP' +-- to this type automatically extracts each Observe / Deterministic +-- node's set of parents ('extractDeps'). This underlies DAG +-- visualization (buildModelGraph). +-- +-- Dependencies only reach into the lower layers +-- 'Hanalyze.Model.HBM.Model' (Node / ModelF / lmParents, etc.) and +-- '...Distribution' (Distribution / distName). There is no dependency on +-- the evaluation layer (logJoint, etc.) — runTrack, the Track +-- specialization of logJoint, lives in the Eval layer instead. +module Hanalyze.Model.HBM.Track + ( Track (..) + , trackVar + , trackConst + , extractDeps + ) where + +import Data.Set (Set) +import qualified Data.Set as Set +import Data.Map.Strict (Map) +import qualified Data.Map.Strict as Map +import Data.Text (Text) + +import Hanalyze.Model.HBM.Model +import Hanalyze.Model.HBM.Distribution + +-- --------------------------------------------------------------------------- +-- 依存追跡型 Track +-- --------------------------------------------------------------------------- + +-- | [日本語]: Floating 演算を通して「この値はどの変数に依存するか」を伝播する型。 +-- +-- @ModelP@ をこの型で特殊化することで、各 Observe ノードが +-- どの latent 変数に依存しているか自動抽出できる。 +-- [English]: A type that propagates "which variable does this value +-- depend on" through @Floating@ operations. +-- +-- Specializing @ModelP@ to this type lets us automatically extract which +-- latent variables each Observe node depends on. +data Track = Track + { trackVal :: !Double + , trackDeps :: !(Set Text) + } deriving (Show, Eq) + +-- | [日本語]: 変数として登場する Track (deps に自分の名前を入れる)。 +-- [English]: A Track that appears as a variable (its own name is added +-- to deps). +trackVar :: Text -> Double -> Track +trackVar n v = Track v (Set.singleton n) + +-- | [日本語]: 定数として扱う Track (deps なし)。 +-- [English]: A Track treated as a constant (empty deps). +trackConst :: Double -> Track +trackConst v = Track v Set.empty + +-- Phase 60.7: '!!!' の依存タグ注入。 Track 解釈だけが slot 名を依存集合に +-- 足し、 DAG に slot→利用先のエッジを出す (数値解釈は既定 id)。 +instance TrackTag Track where + tagDep nm (Track v ds) = Track v (Set.insert nm ds) + +-- 自然な順序関係 (Double の比較を使う) +instance Ord Track where + compare a b = compare (trackVal a) (trackVal b) + +-- Floating の階段 +instance Num Track where + fromInteger n = trackConst (fromInteger n) + Track a sa + Track b sb = Track (a + b) (sa <> sb) + Track a sa - Track b sb = Track (a - b) (sa <> sb) + Track a sa * Track b sb = Track (a * b) (sa <> sb) + abs (Track a sa) = Track (abs a) sa + signum (Track a sa) = Track (signum a) sa + negate (Track a sa) = Track (negate a) sa + +instance Fractional Track where + fromRational r = trackConst (fromRational r) + Track a sa / Track b sb = Track (a / b) (sa <> sb) + +instance Floating Track where + pi = trackConst pi + exp (Track a sa) = Track (exp a) sa + log (Track a sa) = Track (log a) sa + sin (Track a sa) = Track (sin a) sa + cos (Track a sa) = Track (cos a) sa + tan (Track a sa) = Track (tan a) sa + asin (Track a sa) = Track (asin a) sa + acos (Track a sa) = Track (acos a) sa + atan (Track a sa) = Track (atan a) sa + sinh (Track a sa) = Track (sinh a) sa + cosh (Track a sa) = Track (cosh a) sa + tanh (Track a sa) = Track (tanh a) sa + asinh (Track a sa) = Track (asinh a) sa + acosh (Track a sa) = Track (acosh a) sa + atanh (Track a sa) = Track (atanh a) sa + sqrt (Track a sa) = Track (sqrt a) sa + Track a sa ** Track b sb = Track (a ** b) (sa <> sb) + logBase (Track a sa) (Track b sb) = Track (logBase a b) (sa <> sb) + +instance Real Track where + toRational = toRational . trackVal + +instance RealFrac Track where + properFraction (Track a sa) = let (i, f) = properFraction a in (i, Track f sa) + +-- | [日本語]: モデルを Track 型で実行し、各ノードの依存関係を抽出する。 +-- +-- Sample n: その変数自体は @{n}@ に依存する (自己依存)。 +-- Observe n: 分布のパラメータに含まれる latent 変数の集合を deps とする。 +-- +-- plate スタックを保持し、 各 Node に 'nodePlates' を埋める。 +-- 同時に出現した plate (name, size) を 'Map Text Int' で返す。 +-- [English]: Runs the model with the Track type and extracts each node's +-- dependencies. +-- +-- Sample n: the variable itself depends on @{n}@ (self-dependency). +-- Observe n: deps is the set of latent variables appearing in the +-- distribution's parameters. +-- +-- Maintains a plate stack, filling in each Node's 'nodePlates'. Returns +-- the plates encountered along the way as (name, size) in a +-- 'Map Text Int'. +extractDeps :: forall r. ModelP r -> ([Node], Map Text Int) +extractDeps m = + let (ns, plates) = go m [] [] Map.empty Map.empty Map.empty in (ns, plates) + where + -- 引数 stack は **inner-most が head** の plate 名スタック。 + -- slots / obsAcc は Phase 63.1 の side map: slots = データ slot の生値 + -- (slot 名 → ys)、 obsAcc = observe の生 ys を obs 名ごとに chunk 蓄積 + -- (per-point loop の observe \"y\" … [y] も連結すれば slot 全列と一致する)。 + -- walk 終端 (Pure) で値一致逆引きし obs→slot エッジを張る ('linkObsSlots')。 + go :: Model Track r -> [Text] -> [Node] -> Map Text Int + -> Map Text [Double] -> Map Text [[Double]] -> ([Node], Map Text Int) + go (Pure _) _ acc plates slots obsAcc = + (reverse (linkObsSlots slots obsAcc acc), plates) + go (Free (Sample n d k)) stack acc plates slots obsAcc = + let parentDeps = distDepsT d + node = Node n LatentN (distName d) parentDeps (reverse stack) + v = trackVar n 1.0 -- 1 にすると log/exp が安全 + in go (k v) stack (node : acc) plates slots obsAcc + go (Free (Observe n d ys next)) stack acc plates slots obsAcc = + let parentDeps = distDepsT d + node = Node n (ObservedN (length ys)) (distName d) parentDeps (reverse stack) + in go next stack (node : acc) plates slots (obsChunk n ys obsAcc) + go (Free (ObserveLM n bs _ re fam ys next)) stack acc plates slots obsAcc = + -- 親 = β + u + 分散パラメタ名 (lmParents)。 観測ブロックは 1 ノード。 + let parentDeps = lmParents bs re fam + node = Node n (ObservedN (length ys)) (lmFamilyName fam) parentDeps (reverse stack) + in go next stack (node : acc) plates slots (obsChunk n ys obsAcc) + go (Free (Potential nm v next)) stack acc plates slots obsAcc = + -- Potential も DAG 上は「依存を持つ無形ノード」として可視化 + let parentDeps = trackDeps v + node = Node nm LatentN "Potential" parentDeps (reverse stack) + in go next stack (node : acc) plates slots obsAcc + go (Free (Deterministic nm v k)) stack acc plates slots obsAcc = + -- Deterministic ノードの親は @v@ が触れた latent 集合。 + -- 継続には deps を @{nm}@ に「再ラベル」 した Track を渡し、 下流が + -- @v@ の遠い親 (mu, tau 等) ではなく **det 名 nm そのもの** を + -- 親として認識するようにする (Phase 38 で plate-style DAG に修正)。 + -- 数値値は元の @trackVal v@ を保持 (下流の log/exp 等が安全)。 + let parentDeps = trackDeps v + node = Node nm DeterministicN "Deterministic" parentDeps (reverse stack) + v' = Track (trackVal v) (Set.singleton nm) + in go (k v') stack (node : acc) plates slots obsAcc + go (Free (Data n ys k)) stack acc plates slots obsAcc = + -- Phase 60.4: pm.Data 相当のデータノード。 値 (fst view) には slot 名の + -- dep タグを載せ、 下流 (deterministic / observe の dist パラメタ) が + -- x→mu のエッジを自動で張れるようにする (Phase 38 deterministic + -- re-label と同手法)。 snd view (dataNamedObs の生 [Double]) には + -- deps を載せられないため、 slots に生値を控えて walk 終端で + -- 値一致逆引きの obs→slot エッジを張る (Phase 63.1)。 + let node = Node n (DataN (length ys)) "Data" Set.empty (reverse stack) + vals = map (\v -> Track v (Set.singleton n)) ys + in go (k (vals, ys)) stack (node : acc) plates (Map.insert n ys slots) obsAcc + go (Free (DataIx n is k)) stack acc plates slots obsAcc = + -- DataIx は [Int] のまま継続に渡すため dep タグは載らない (ノードのみ)。 + -- observe の ys ([Double]) と一致し得ないので slots にも入れない。 + let node = Node n (DataN (length is)) "DataIx" Set.empty (reverse stack) + in go (k is) stack (node : acc) plates slots obsAcc + go (Free (PlateBegin nm sz next)) stack acc plates slots obsAcc = + -- plate を開始 = stack に push、 サイズも記録 (重複時は新値で上書き + -- = 同名 plate は同サイズ前提) + go next (nm : stack) acc (Map.insert nm sz plates) slots obsAcc + go (Free (PlateEnd next)) stack acc plates slots obsAcc = + -- plate を終了 = stack から pop。 空 stack は誤用 (PlateBegin 抜き + -- で PlateEnd が来た等) — 黙って無視する + let stack' = case stack of { _ : t -> t; [] -> [] } + in go next stack' acc plates slots obsAcc + + -- obs 名ごとの ys chunk 蓄積 (新 chunk を先頭 prepend = 逆順保持。 + -- per-observe の list append による O(n²) を避ける)。 + obsChunk :: Text -> [Double] -> Map Text [[Double]] -> Map Text [[Double]] + obsChunk n ys = Map.insertWith (++) n [ys] + + -- Phase 63.1: observe の連結 ys と値一致するデータ slot へ obs→slot エッジ + -- (= 該当 DataN Node の nodeDeps に obs 名を追加。 nodeDeps は「直接の親」 + -- ゆえ slot は obs の子 = PyMC `make_compute_graph` の obs→y と同型)。 + -- + -- - 値一致は plate 長さ match (60.6) と同種の表示専用ヒューリスティック: + -- 偶然同値の slot にも張られる (既知 caveat・doc 明記)、 同値 slot 複数は + -- 全部に張る。 空 slot (未 bind placeholder) は対象外。 + -- - 同名 (dataNamedObs \"y\" + observe \"y\" の docs 慣例) は対象外: + -- mergeByName で 1 ノードに統合されるため自己ループになる。 + -- - 引数 acc は逆順のまま受けて逆順のまま返す (呼び元 Pure 節で reverse)。 + linkObsSlots :: Map Text [Double] -> Map Text [[Double]] -> [Node] -> [Node] + linkObsSlots slots obsAcc acc + | Map.null links = acc + | otherwise = map upd acc + where + -- obs 名 → 連結 ys (chunk は新しい順 prepend 蓄積ゆえ reverse) + obsYs = Map.map (concat . reverse) obsAcc + -- slot 名 → 親として足す obs 名集合 + links = Map.fromListWith Set.union + [ (slotName, Set.singleton obsName) + | (slotName, sv) <- Map.toList slots + , not (null sv) + , (obsName, ys) <- Map.toList obsYs + , obsName /= slotName + , sv == ys ] + upd nd = case nodeKind nd of + DataN _ | Just parents <- Map.lookup (nodeName nd) links -> + nd { nodeDeps = nodeDeps nd <> parents } + _ -> nd + +-- | [日本語]: Distribution Track に含まれる依存変数集合を取り出す。 +-- [English]: Extracts the set of dependency variables contained in a +-- Distribution Track. +distDepsT :: Distribution Track -> Set Text +distDepsT (Normal mu sig) = trackDeps mu <> trackDeps sig +distDepsT (Exponential r) = trackDeps r +distDepsT (Gamma s r) = trackDeps s <> trackDeps r +distDepsT (Beta a b) = trackDeps a <> trackDeps b +distDepsT (Poisson lam) = trackDeps lam +distDepsT (Binomial _ p) = trackDeps p +distDepsT (Uniform lo hi) = trackDeps lo <> trackDeps hi +distDepsT (StudentT df mu s) = trackDeps df <> trackDeps mu <> trackDeps s +distDepsT (Cauchy loc s) = trackDeps loc <> trackDeps s +distDepsT (HalfNormal s) = trackDeps s +distDepsT (HalfCauchy s) = trackDeps s +distDepsT (LogNormal mu s) = trackDeps mu <> trackDeps s +distDepsT (Bernoulli p) = trackDeps p +distDepsT (Categorical ps) = mconcat (map trackDeps ps) +distDepsT (Mixture ws ds) = mconcat (map trackDeps ws) <> mconcat (map distDepsT ds) +distDepsT (Truncated d mLo mHi) = + distDepsT d <> maybe mempty trackDeps mLo <> maybe mempty trackDeps mHi +distDepsT (Censored d mLo mHi) = + distDepsT d <> maybe mempty trackDeps mLo <> maybe mempty trackDeps mHi +distDepsT (MvNormal mus covRows) = + mconcat (map trackDeps mus) + <> mconcat (concatMap (map trackDeps) covRows) +distDepsT (MvNormalChol mus sigmas lRows) = + mconcat (map trackDeps mus) + <> mconcat (map trackDeps sigmas) + <> mconcat (concatMap (map trackDeps) lRows) +distDepsT (MvNormalGpRBF xs alpha rho sigma) = -- Phase 95 B-dsl: x は data・α/ρ/σ が param + mconcat (map trackDeps xs) + <> trackDeps alpha <> trackDeps rho <> trackDeps sigma +distDepsT (HmmForwardNormal pi0 trans mus sg) = -- Phase 92 A2: 全て param 側 (data は Observe に載る) + mconcat (map trackDeps pi0) + <> mconcat (concatMap (map trackDeps) trans) + <> mconcat (map trackDeps mus) <> trackDeps sg +distDepsT (ArmaNormal mu phi theta sg) = -- Phase 101 A2: 全て param 側 (data は Observe に載る) + trackDeps mu <> trackDeps phi <> trackDeps theta <> trackDeps sg +distDepsT (GradedResponseIrt thetas _ _ _) = -- Phase 101 A3: θs のみ param 側 (他は定数 data) + mconcat (map trackDeps thetas) +distDepsT (NegativeBinomial mu alpha) = trackDeps mu <> trackDeps alpha +distDepsT (Multinomial _ ps) = mconcat (map trackDeps ps) +distDepsT (ZeroInflatedPoisson psi lam) = trackDeps psi <> trackDeps lam +distDepsT (ZeroInflatedBinomial _ psi p) = trackDeps psi <> trackDeps p +distDepsT (InverseGamma a b) = trackDeps a <> trackDeps b +distDepsT (Weibull k l) = trackDeps k <> trackDeps l +distDepsT (Pareto a xm) = trackDeps a <> trackDeps xm +distDepsT (BetaBinomial _ a b) = trackDeps a <> trackDeps b +distDepsT (VonMises mu k) = trackDeps mu <> trackDeps k +-- Phase 37 で追加した分布 (Phase 38 補修で網羅追加) +distDepsT (SkewNormal mu sig alpha) = + trackDeps mu <> trackDeps sig <> trackDeps alpha +distDepsT (Logistic mu s) = trackDeps mu <> trackDeps s +distDepsT (Gumbel mu beta) = trackDeps mu <> trackDeps beta +distDepsT (AsymmetricLaplace b kappa mu) = + trackDeps b <> trackDeps kappa <> trackDeps mu +distDepsT (OrderedLogistic eta cuts) = + trackDeps eta <> mconcat (map trackDeps cuts) +distDepsT DiscreteUniform{} = mempty -- Int 引数のみ +distDepsT (Geometric p) = trackDeps p +distDepsT HyperGeometric{} = mempty -- Int 引数のみ +distDepsT (ZeroInflatedNegativeBinomial psi mu alpha) = + trackDeps psi <> trackDeps mu <> trackDeps alpha +distDepsT (MvStudentT nu mus covRows) = + trackDeps nu + <> mconcat (map trackDeps mus) + <> mconcat (concatMap (map trackDeps) covRows) +distDepsT (DirichletMultinomial _ alphas) = + mconcat (map trackDeps alphas) +distDepsT (Triangular lo c hi) = + trackDeps lo <> trackDeps c <> trackDeps hi +distDepsT (Kumaraswamy a b) = trackDeps a <> trackDeps b +distDepsT (Rice nu sig) = trackDeps nu <> trackDeps sig +distDepsT (DiscreteWeibull q beta) = trackDeps q <> trackDeps beta +distDepsT (Wishart nu vRows) = + trackDeps nu <> mconcat (concatMap (map trackDeps) vRows) +distDepsT (Bound d mLo mHi) = + distDepsT d + <> maybe mempty trackDeps mLo + <> maybe mempty trackDeps mHi +distDepsT (OrderedProbit eta cuts) = + trackDeps eta <> mconcat (map trackDeps cuts) + diff --git a/src/Hanalyze/Model/HBM/Util.hs b/src/Hanalyze/Model/HBM/Util.hs new file mode 100644 --- /dev/null +++ b/src/Hanalyze/Model/HBM/Util.hs @@ -0,0 +1,464 @@ +{-# LANGUAGE ScopedTypeVariables #-} +-- | +-- Module : Hanalyze.Model.HBM.Util +-- Description : HBM の純粋な数値・線形代数 leaf ユーティリティ +-- Copyright : (c) 2026 Aelysce Project (Toshiaki Honda) +-- License : BSD-3-Clause +-- +-- HBM の純粋な数値・線形代数 leaf ユーティリティ。 +-- +-- [日本語]: ここに集めた定義は HBM のいずれの型 (Distribution / Model / +-- Track 等) にも依存しない葉 (leaf) であり、 Floating / Ord のみで多相に +-- 書かれている。 AD (Reverse.Double) でも Track でも評価できるよう型クラス +-- 制約を最小に保つ。 'Hanalyze.Model.HBM' は本モジュールを import +-- して内部利用 + 一部を re-export する (公開シンボル: 'lgammaApprox' / +-- 'digamma')。 +-- [English]: The definitions gathered here are leaves that depend on none +-- of HBM's types (Distribution / Model / Track, etc.), written +-- polymorphically with only @Floating@ / @Ord@ constraints, kept minimal +-- so they can be evaluated via AD (@Reverse.Double@) or @Track@ alike. +-- 'Hanalyze.Model.HBM' imports this module for internal use and +-- re-exports part of it (public symbols: 'lgammaApprox' / 'digamma'). +-- +-- [日本語]: 'Hanalyze.Model.HBM' (5,519 行) から責務分離して抽出。 +-- 数値は 1 bit も変えていない (純粋な移設)。 +-- [English]: Extracted from 'Hanalyze.Model.HBM' (5,519 lines) as +-- a responsibility split; the numerics are unchanged to the bit (a pure +-- relocation). +module Hanalyze.Model.HBM.Util + ( -- * 線形代数 (下三角ソルバ / Cholesky / リスト整形) + backSubLT + , chunksOf + , choleskyL + , forwardSub + , gpRBFCovList + -- * log-sum-exp / HMM forward + , negInf + , logSumExpA + , hmmForwardLogLik + -- * 不完全ガンマ関数 P(a, x) + , incGammaPA + , igammSer + , igammCF + -- * 正則化不完全ベータ関数 I_x(a, b) + , incBetaA + , betaCFA + -- * 数値ユーティリティ (Γ / digamma / 階乗 / Bessel) + , lgammaApprox + , digamma + , lgammaApproxDeriv + , logFactorial + , logBinomCoeff + , logBesselI0 + ) where + +import Data.List (foldl') +import qualified Data.Vector as V + +-- =========================================================================== +-- 線形代数 (下三角ソルバ / Cholesky / リスト整形) +-- =========================================================================== +-- Phase 95 A2 (2026-07-13): choleskyL/forwardSub/backSubLT の内部を nested-list +-- ([[a]] + !! O(n)索引 + ++ O(n)追記) から Data.Vector (O(1) 索引 + snoc) へ +-- 脱リスト化。公開シグネチャ ([[a]]) は不変 = 呼び出し元は無改修。数値は回帰 +-- テスト内で一致 (posterior bit 一致を実測)。★N=11 の gp-regr では効果ゼロ +-- (真因は AD tape ノード alloc・§A2 参照) だが、大 N の密行列では list !!/++ が +-- O(N⁴) 化して支配的になるため user 判断で先行 infra として採用 (2026-07-13)。 +-- ※さらなる高速化には interface 自体の Vector 化 (呼出側の per-call 変換除去) が +-- 要・大 N 密行列モデル出現時の TODO。 + +-- | [日本語]: RBF (exponentiated-quadratic) GP カーネルの共分散行列を +-- nested list で構築する。 @Σ_ij = α² exp(-0.5 (x_i-x_j)²/ρ²) + [i=j](1e-10 + σ)@。 +-- @Hanalyze.Model.HBM.gpExpQuadCov@ (jitter 1e-10 込) + 対角 σ と一致する +-- = @MvNormalGpRBF@ 密度が呼ぶ (値は既存 gp-regr モデルと bit 一致)。 下層 (Util) +-- に置くことで @Distribution@ の @obsLogSum@ から参照できる (Model 層の +-- @gpExpQuadCov@ は上層ゆえ密度からは呼べない)。 ★ホット経路 (Gradient の +-- @gpRBFAnalyticVG@) は本 list 版を使わず hmatrix Matrix で直接組む (脱リスト)。 +-- [English]: Builds the RBF (exponentiated-quadratic) GP kernel's +-- covariance matrix as a nested list: +-- @Σ_ij = α² exp(-0.5 (x_i-x_j)²/ρ²) + [i=j](1e-10 + σ)@. Matches +-- @Hanalyze.Model.HBM.gpExpQuadCov@ (including the 1e-10 jitter) +-- plus the diagonal σ — called by the @MvNormalGpRBF@ density (bit-identical +-- to existing gp-regr models). Placed in this lower layer (Util) so +-- 'Distribution''s @obsLogSum@ can reference it (the Model-layer +-- @gpExpQuadCov@ sits above and cannot be called from a density). ★The hot +-- path (Gradient's @gpRBFAnalyticVG@) does not use this list version but +-- builds directly with an hmatrix @Matrix@ instead (delisted). +{-# INLINABLE gpRBFCovList #-} +gpRBFCovList :: forall a. Floating a => [a] -> a -> a -> a -> [[a]] +gpRBFCovList xs alpha rho sigma = + [ [ let d = xi - xj + k = alpha * alpha * exp (negate 0.5 * d * d / (rho * rho)) + in k + (if i == j then 1e-10 + sigma else 0) + | (j, xj) <- zip [0 :: Int ..] xs ] + | (i, xi) <- zip [0 :: Int ..] xs ] + +-- | [日本語]: 下三角 L から Lᵀ x = b を後退代入で解く (L は @choleskyL@ 形式)。 [English]: Solves Lᵀ x = b by back-substitution given lower-triangular L (in @choleskyL@ form). +{-# INLINABLE backSubLT #-} +backSubLT :: forall a. Floating a => [[a]] -> [a] -> [a] +backSubLT l b = + let n = length b + lV = V.fromList [ V.fromList r | r <- l ] + arr = V.fromListN n (b ++ repeat 0) + go :: Int -> V.Vector a -> V.Vector a + go i acc -- acc = x[i+1..n-1] + | i < 0 = acc + | otherwise = + -- (acc は index i+1..n-1 の解、 i 番目を解く) + -- Lᵀ x = b → 行 i: Σ_{j>=i} L[j][i] x_j = b_i + -- → x_i = (b_i - Σ_{j>i} L[j][i] x_j) / L[i][i] + let lii = (lV V.! i) V.! i + bi = arr V.! i + s = V.sum (V.imap (\t xj -> (lV V.! (i + 1 + t)) V.! i * xj) acc) + xi = (bi - s) / lii + in go (i - 1) (V.cons xi acc) + in V.toList (go (n - 1) V.empty) + +-- | [日本語]: リストを長さ @n@ ごとに分割。 最後が短ければそのまま (本実装では使わない想定)。 [English]: Splits a list into chunks of length @n@; the last chunk stays short if it doesn't divide evenly (unused in this implementation). +chunksOf :: Int -> [a] -> [[a]] +chunksOf _ [] = [] +chunksOf n xs = let (h, t) = splitAt n xs in h : chunksOf n t + +-- | [日本語]: 対称正定値行列 Σ の Cholesky 下三角分解 L (Σ = L Lᵀ)。 +-- 行列は行リスト @[[a]]@ で、 l[i] は長さ @i+1@ の下三角行 +-- ([L[i][0]..L[i][i]])。 対角が非正になれば @Nothing@。 +-- [English]: Cholesky lower-triangular decomposition L of a symmetric +-- positive-definite matrix Σ (Σ = L Lᵀ). The matrix is a list of rows +-- @[[a]]@, where l[i] is the length-@i+1@ lower-triangular row +-- ([L[i][0]..L[i][i]]). Returns @Nothing@ if a diagonal entry becomes +-- non-positive. +{-# INLINABLE choleskyL #-} +choleskyL :: forall a. (Floating a, Ord a) => [[a]] -> Maybe [[a]] +choleskyL a0 = + let n = length a0 + aV = V.fromList [ V.fromList r | r <- a0 ] -- 入力行 (各行長 >= i+1) + step :: Int -> V.Vector (V.Vector a) -> Maybe (V.Vector (V.Vector a)) + step i prev -- prev = 確定済 L[0..i-1] + | i == n = Just prev + | otherwise = + let row = aV V.! i + buildCol :: Int -> V.Vector a -> Maybe (V.Vector a) + buildCol j cur -- cur = L[i][0..j-1] + | j > i = Just cur + | j == i = + let s = V.sum (V.map (\v -> v * v) cur) + d2 = (row V.! i) - s + in if d2 <= 0 + then Nothing + else buildCol (j + 1) (V.snoc cur (sqrt d2)) + | otherwise = + let lj = prev V.! j -- 長さ j+1 + s = V.sum (V.zipWith (*) cur lj) + ljj = lj V.! j + in if ljj == 0 + then Nothing + else buildCol (j + 1) (V.snoc cur ((row V.! j - s) / ljj)) + in case buildCol 0 V.empty of + Nothing -> Nothing + Just nr -> step (i + 1) (V.snoc prev nr) + in fmap (\v -> [ V.toList r | r <- V.toList v ]) (step 0 V.empty) + +-- | [日本語]: 下三角系 L z = b の前進代入 (L は @choleskyL@ 形式、 長さ各 i+1)。 [English]: Solves the lower-triangular system L z = b by forward substitution (L in @choleskyL@ form, each row length i+1). +{-# INLINABLE forwardSub #-} +forwardSub :: forall a. Floating a => [[a]] -> [a] -> [a] +forwardSub l b = + let n = length b + lV = V.fromList [ V.fromList r | r <- l ] + bV = V.fromList b + go :: Int -> V.Vector a -> V.Vector a + go i acc -- acc = z[0..i-1] + | i == n = acc + | otherwise = + let lrow = lV V.! i -- 長さ i+1 + lii = lrow V.! i + lpre = V.take i lrow -- L[i][0..i-1] + bi = bV V.! i + s = V.sum (V.zipWith (*) lpre acc) + zi = (bi - s) / lii + in go (i + 1) (V.snoc acc zi) + in V.toList (go 0 V.empty) + +-- =========================================================================== +-- log-sum-exp +-- =========================================================================== + +negInf :: Floating a => a +negInf = -1/0 + +-- | [日本語]: 多相 log-sum-exp。 AD でも Track でも使えるよう Floating + Ord +-- で書く。 @logSumExpA xs = log (Σ exp x)@ を最大値シフトで安定化。 +-- [English]: Polymorphic log-sum-exp, written with @Floating@ + @Ord@ so +-- it works under both AD and @Track@. Stabilizes +-- @logSumExpA xs = log (Σ exp x)@ via a max-value shift. +{-# INLINABLE logSumExpA #-} +logSumExpA :: (Floating a, Ord a) => [a] -> a +logSumExpA [] = negInf +logSumExpA [x] = x +logSumExpA xs = + let m = maximum xs + -- 全要素が -∞ なら m - m = NaN になるので早期 return + in if m == negInf + then negInf + else m + log (sum (map (\x -> exp (x - m)) xs)) + +-- =========================================================================== +-- HMM forward algorithm (状態列の周辺化) +-- =========================================================================== +-- Phase 92 A2 (2026-07-17): Model.hs:1071 から純粋移設 (数値は 1 bit も不変)。 +-- @Distribution@ の 'HmmForwardNormal' 密度 ('obsLogSum') が呼ぶため、 +-- Model 非依存の leaf である本モジュールへ降ろした。 +-- 'Hanalyze.Model.HBM.Model' が従来どおり re-export する。 + +-- | [日本語]: 隠れマルコフモデルの周辺対数尤度 (forward algorithm)。 +-- [English]: The hidden Markov model's marginal log-likelihood (forward +-- algorithm). +-- +-- Recursion in log-space (underflow 防止 / to prevent underflow): +-- +-- - @α_1[k] = log π_0[k] + emit[0][k]@ +-- - @α_{t+1}[k'] = logSumExp_j (α_t[j] + log T[j][k']) + emit[t+1][k']@ +-- - @log P(y_{1..T}) = logSumExp_k α_T[k]@ +-- +-- [日本語]: 多相 (@Floating a, Ord a@) のため Track / AD 経由でも動く。 +-- 計算量 @O(T K²)@。 大 T では list-based なので O(K²) の内部ループは +-- そのまま、 step は foldl' で過去 α を破棄しメモリは @O(K)@。 +-- [English]: Polymorphic (@Floating a, Ord a@), so it also works via +-- @Track@ / AD. Complexity is @O(T K²)@; for large T the inner O(K²) loop +-- stays list-based, but @step@ uses @foldl'@ to discard past α, keeping +-- memory at @O(K)@. +hmmForwardLogLik :: forall a. (Floating a, Ord a) + => [a] -- ^ [日本語]: 初期分布 π_0 (length K)。 [English]: Initial distribution π_0 (length K). + -> [[a]] -- ^ [日本語]: 遷移行列 (K×K rows of length K)。 [English]: Transition matrix (K×K rows of length K). + -> [[a]] -- ^ log emission [T][K] + -> a +hmmForwardLogLik pi0 trans emit + | null emit = 0 -- T=0: 観測なし + | null pi0 = negInf + | length pi0 /= length trans = negInf + | any ((/= k) . length) trans = negInf + | otherwise = + let -- α_1[s] = log π_0[s] + emit[0][s] + alpha0 = zipWith (\p e -> log p + e) pi0 (head emit) + -- 1 step: α_{t+1}[s'] = logSumExp_s (α_t[s] + log T[s][s']) + emit_{t+1}[s'] + step :: [a] -> [a] -> [a] + step alphaT emT = + [ logSumExpA + [ (alphaT !! s) + log ((trans !! s) !! s') + | s <- [0 .. k - 1] ] + + (emT !! s') + | s' <- [0 .. k - 1] ] + alphaFinal = foldl' step alpha0 (tail emit) + in logSumExpA alphaFinal + where + k = length pi0 + +-- =========================================================================== +-- 不完全ガンマ関数 P(a, x) = γ(a, x) / Γ(a) (Numerical Recipes 6.2) +-- =========================================================================== + +-- | [日本語]: 正則化された下側不完全ガンマ関数 P(a, x) = γ(a, x) / Γ(a) ∈ [0, 1]。 +-- これは Gamma(shape=a, rate=1) の CDF F(x)。 +-- [English]: The regularized lower incomplete gamma function +-- P(a, x) = γ(a, x) / Γ(a) ∈ [0, 1] — this is the CDF F(x) of +-- Gamma(shape=a, rate=1). +{-# INLINABLE incGammaPA #-} +incGammaPA :: (Floating a, Ord a) => a -> a -> a +incGammaPA a x + | x <= 0 || a <= 0 = 0 + | x < a + 1 = igammSer a x -- 級数展開で P(a,x) + | otherwise = 1 - igammCF a x -- 連分数で Q(a,x)、P = 1 - Q + +-- 級数展開: P(a, x) = e^{-x} x^a / Γ(a) * Σ x^n / (a(a+1)...(a+n)) +{-# INLINABLE igammSer #-} +igammSer :: forall a. (Floating a, Ord a) => a -> a -> a +igammSer a x = sumSer * exp (-x + a * log x - lgammaApprox a) + where + -- 反復: term_{n+1} = term_n * x / (a + n + 1) + sumSer = go (0 :: Int) (1 / a) (1 / a) + eps :: a + eps = 1e-13 + maxIt = 200 :: Int + go n term acc + | n >= maxIt = acc + | abs term < abs acc * eps = acc + | otherwise = + let n' = n + 1 + term' = term * x / (a + fromIntegral n') + acc' = acc + term' + in go n' term' acc' + +-- 連分数 (Lentz 法): Q(a, x) = e^{-x} x^a / Γ(a) * CF +-- CF = 1/(x+1-a - 1·(1-a)/(x+3-a - 2·(2-a)/(...)) +{-# INLINABLE igammCF #-} +igammCF :: forall a. (Floating a, Ord a) => a -> a -> a +igammCF a x = exp (-x + a * log x - lgammaApprox a) * h + where + fpmin, eps :: a + fpmin = 1e-300 + eps = 1e-13 + maxIt = 200 :: Int + -- modified Lentz's method + b0 = x + 1 - a + c0 = 1 / fpmin + d0 = 1 / b0 + h = goCF (1 :: Int) b0 c0 d0 d0 + goCF i b c d hh + | i > maxIt = hh + | abs (del - 1) < eps = hh' + | otherwise = goCF (i + 1) b' c'' d''' hh' + where + an = -fromIntegral i * (fromIntegral i - a) + b' = b + 2 + d' = b' + an * d + d'' = if abs d' < fpmin then fpmin else d' + c' = b' + an / c + c'' = if abs c' < fpmin then fpmin else c' + d''' = 1 / d'' + del = d''' * c'' + hh' = hh * del + _ = c0 -- 未使用ダミー (修正された Lentz 法の起動値: 別経路) + +-- =========================================================================== +-- 正則化された不完全ベータ関数 I_x(a, b) = B(x; a, b) / B(a, b) +-- =========================================================================== + +-- | [日本語]: 正則化された不完全ベータ関数 I_x(a, b) ∈ [0, 1]。 +-- これは Beta(a, b) の CDF F(x)。 StudentT の CDF にも内部で使用。 +-- [English]: The regularized incomplete beta function I_x(a, b) ∈ [0, 1] +-- — this is the CDF F(x) of Beta(a, b). Also used internally by +-- StudentT's CDF. +{-# INLINABLE incBetaA #-} +incBetaA :: (Floating a, Ord a) => a -> a -> a -> a +incBetaA x a b + | x <= 0 = 0 + | x >= 1 = 1 + | otherwise = + -- 対数ベータ正規化定数 + let bt = exp ( lgammaApprox (a + b) + - lgammaApprox a + - lgammaApprox b + + a * log x + + b * log (1 - x)) + in if x < (a + 1) / (a + b + 2) + then bt * betaCFA x a b / a + else 1 - bt * betaCFA (1 - x) b a / b + +-- 連分数 (modified Lentz, Numerical Recipes §6.4) +{-# INLINABLE betaCFA #-} +betaCFA :: forall a. (Floating a, Ord a) => a -> a -> a -> a +betaCFA x a b = iterate' (1 :: Int) 1 d0 h0 + where + fpmin, eps :: a + fpmin = 1e-300 + eps = 1e-13 + maxIt = 200 :: Int + qab = a + b + qap = a + 1 + qam = a - 1 + capLent v = if abs v < fpmin then fpmin else v + d0 = 1 / capLent (1 - qab * x / qap) + h0 = d0 + + iterate' m c d h + | m > maxIt = h + | abs (del - 1) < eps = hO + | otherwise = iterate' (m + 1) cO dO hO + where + mD = fromIntegral m :: a + -- 偶数項: aa_2m = m(b-m)x / ((qam+2m)(a+2m)) + aaE = mD * (b - mD) * x / ((qam + 2 * mD) * (a + 2 * mD)) + dE = 1 / capLent (1 + aaE * d) + cE = capLent (1 + aaE / c) + hE = h * dE * cE + -- 奇数項: aa_2m+1 = -(a+m)(qab+m)x / ((a+2m)(qap+2m)) + aaO = -(a + mD) * (qab + mD) * x / ((a + 2 * mD) * (qap + 2 * mD)) + dO = 1 / capLent (1 + aaO * dE) + cO = capLent (1 + aaO / cE) + del = dO * cO + hO = hE * del + +-- =========================================================================== +-- 数値ユーティリティ (Γ / digamma / 階乗 / Bessel) +-- =========================================================================== + +-- | [日本語]: log Γ(z) の Stirling 近似 (z > 0)。 AD でも Track でも使える多相版。 [English]: Stirling's approximation of log Γ(z) (z > 0); a polymorphic version usable under both AD and @Track@. +{-# INLINABLE lgammaApprox #-} +lgammaApprox :: (Floating a, Ord a) => a -> a +lgammaApprox z + | z < 12 = lgammaApprox (z + 1) - log z + | otherwise = (z - 0.5) * log z - z + 0.5 * log (2 * pi) + + 1 / (12 * z) - 1 / (360 * z ^ (3::Int)) + +-- | [日本語]: ψ(z) = d/dz log Γ(z) (z > 0)。 記号微分 IR の lgamma 単項 op +-- (@SLgammaO@ 予定) の導関数用。 'lgammaApprox' と同一の recurrence +-- (z < 12 を押し上げ) + 漸近級数を lgammaApprox の Stirling 微分より 1 項深く +-- (-1/(252 z⁶) まで) 打切り: 真の ψ との差は z=12 で ~1e-11、 +-- lgammaApprox の数値微分との差は lgammaApprox 側の打切り由来 ~1.3e-9 +-- (試験許容 1e-8 内)。 z ≤ 0 は未対応 (利用箇所は正値前提)。 +-- [English]: ψ(z) = d/dz log Γ(z) (z > 0). For the derivative of the +-- symbolic-differentiation IR's lgamma unary op (planned @SLgammaO@). +-- Uses the same recurrence as 'lgammaApprox' (push up z < 12), but +-- truncates the asymptotic series one term deeper than lgammaApprox's +-- Stirling derivative (down to -1/(252 z⁶)): the difference from the true +-- ψ is ~1e-11 at z=12, and the difference from lgammaApprox's numerical +-- derivative is ~1.3e-9, coming from lgammaApprox's own truncation +-- (within the 1e-8 test tolerance). z ≤ 0 is unsupported (call sites +-- assume positive values). +digamma :: Double -> Double +digamma z + | z < 12 = digamma (z + 1) - 1 / z + | otherwise = log z - 1 / (2 * z) - 1 / (12 * z * z) + + 1 / (120 * z ^ (4 :: Int)) - 1 / (252 * z ^ (6 :: Int)) + +-- | [日本語]: 'lgammaApprox' の __厳密な項別導関数__。 'digamma' とは最終項 +-- 1/(252z⁶) の有無だけ違う (digamma は真の ψ に 1 項深い分この差 ~1.3e-9 が +-- z=12 境界で出る・実測)。 記号微分 IR (@SLgammaO@) の導関数は、 評価関数 +-- (lgammaApprox) の AD 微分 = walk+ad fallback / 参照勾配と一致させる必要が +-- あるためこちらを使う。 +-- [English]: The __exact term-by-term derivative__ of 'lgammaApprox'. +-- Differs from 'digamma' only in whether the final term 1/(252z⁶) is +-- included (digamma is one term deeper toward the true ψ, so this ~1.3e-9 +-- gap shows up at the z=12 boundary, measured). The symbolic-differentiation +-- IR's (@SLgammaO@) derivative must match the AD derivative of the +-- evaluation function (lgammaApprox) — i.e. the walk+ad fallback / +-- reference gradient — so this one is used for that purpose. +lgammaApproxDeriv :: Double -> Double +lgammaApproxDeriv z + | z < 12 = lgammaApproxDeriv (z + 1) - 1 / z + | otherwise = log z - 1 / (2 * z) - 1 / (12 * z * z) + + 1 / (120 * z ^ (4 :: Int)) + +logFactorial :: Int -> Double +logFactorial n + | n <= 1 = 0 + | otherwise = sum (map log [2 .. fromIntegral n]) + +logBinomCoeff :: Int -> Int -> Double +logBinomCoeff n k = logFactorial n - logFactorial k - logFactorial (n - k) + +-- | [日本語]: log I_0(x) — 修正 Bessel 関数 (第一種・order 0) の対数。 +-- VonMises 用。 小 x: 級数 I_0(x) = Σ (x/2)^(2k) / (k!)² (k = 0..)。 +-- 大 x: 漸近展開 I_0(x) ≈ exp(x) / √(2πx) × [1 + 1/(8x) + 9/(128x²) + …]。 +-- AD/Track 互換のため (Floating a, Ord a) 多相。 +-- [English]: log I_0(x) — the log of the modified Bessel function (first +-- kind, order 0), used for VonMises. Small x: series +-- I_0(x) = Σ (x/2)^(2k) / (k!)² (k = 0..). Large x: asymptotic expansion +-- I_0(x) ≈ exp(x) / √(2πx) × [1 + 1/(8x) + 9/(128x²) + …]. Polymorphic +-- over (@Floating a, Ord a@) for AD/@Track@ compatibility. +{-# INLINABLE logBesselI0 #-} +logBesselI0 :: (Floating a, Ord a) => a -> a +logBesselI0 x + | x < 0 = logBesselI0 (-x) -- 偶関数 + | x < 3.75 = + -- Abramowitz & Stegun 9.8.1: 多項式近似 (誤差 < 1.6e-7) + let t = (x / 3.75) ^ (2::Int) + i0 = 1 + t * (3.5156229 + t * (3.0899424 + t * (1.2067492 + + t * (0.2659732 + t * (0.0360768 + t * 0.0045813))))) + in log i0 + | otherwise = + -- Abramowitz & Stegun 9.8.2: 漸近 (誤差 < 1.9e-7) + let t = 3.75 / x + poly = 0.39894228 + t * (0.01328592 + t * (0.00225319 + + t * (-0.00157565 + t * (0.00916281 + t * (-0.02057706 + + t * (0.02635537 + t * (-0.01647633 + t * 0.00392377))))))) + in x - 0.5 * log x + log poly diff --git a/src/Hanalyze/Model/HBM/VecAD.hs b/src/Hanalyze/Model/HBM/VecAD.hs new file mode 100644 --- /dev/null +++ b/src/Hanalyze/Model/HBM/VecAD.hs @@ -0,0 +1,397 @@ +{-# LANGUAGE BangPatterns #-} + +-- | +-- Module : Hanalyze.Model.HBM.VecAD +-- Description : 自作の最小 reverse-mode AD (vector-op tape) +-- Copyright : (c) 2026 Aelysce Project (Toshiaki Honda) +-- License : BSD-3-Clause +-- +-- [日本語]: 自作・最小 reverse-mode AD (vector-op tape)。 「採用 = 案B +-- (自前 vector-op tape)」 と判断したエンジンを本実装用に移植したもの +-- (`bench/haskell/BenchHBMVecADSpike.hs` の gradHandroll 系)。 +-- +-- 設計: forward で「ベクトル演算ごとにノードを発番」 し、 各ノードの随伴更新 +-- クロージャを逆順リストに積む (= 自前 Wengert tape)。 backward で出力に 1 を +-- seed し、 逆位相順 (= 発番の逆順 = prepend したリストの先頭) にクロージャを +-- replay して入力 (leaf) の随伴を得る。 tape は「ベクトル演算 1 個 = 1 ノード」 +-- ゆえ @ad@ のスカラ tape (per-scalar-op で O(n) ノード) より桁で小さい。 +-- +-- スカラは長さ 1 の Storable Vector として随伴を持ち、 ノード随伴は単一の +-- mutable 配列に統一格納する。 +-- +-- 注意: 値依存制御フロー (分布の台チェック等) は tape に乗らない。 本エンジンは +-- 構造が値に依らず静的な部分 (Gaussian-恒等リンクの線形予測子 + 二乗和) 専用。 +-- 非対応の構造は呼出側で scalar (@ad@) 経路に fallback する。 +-- [English]: A hand-rolled minimal reverse-mode AD (vector-op tape). This is +-- the engine chosen ("adopt = plan B, a custom vector-op tape") and ported +-- into the production implementation from the spike +-- (`bench/haskell/BenchHBMVecADSpike.hs`'s gradHandroll family). +-- +-- Design: on the forward pass, each vector operation is numbered as a node, +-- and each node's adjoint-update closure is pushed onto a reverse-order list +-- (= a custom Wengert tape). On the backward pass, the output is seeded with +-- 1, and closures are replayed in reverse topological order (= reverse +-- numbering order = the head of the prepended list) to obtain the adjoints +-- of the input leaves. Because the tape records "one vector operation = one +-- node," it is orders of magnitude smaller than @ad@'s scalar tape (which +-- creates O(n) nodes per scalar op). +-- +-- Scalars carry their adjoint as a length-1 Storable Vector, and all node +-- adjoints are stored uniformly in a single mutable array. +-- +-- Caveat: value-dependent control flow (e.g. distribution support checks) +-- does not survive on the tape. This engine only covers the part of the +-- structure that is static regardless of value (Gaussian-identity-link +-- linear predictor + sum of squares). Unsupported structures fall back to +-- the scalar (@ad@) path at the call site. +module Hanalyze.Model.HBM.VecAD + ( -- * 値ハンドルと文脈 + Rval (..) + , Ctx + , ridOf + -- * tape の実行 + , runTape + -- * leaf + , inputVec + , inputScal + , constVec + -- * ベクトル演算 (随伴付き) + , idxHR + , sliceHR + , scaleHR + , vaddHR + , vsubHR + , dotHR + , gatherHR + , vexpHR + , bcastAddHR + , hadamardHR + , vmap1HR + -- * スカラ演算 (随伴付き) + , map1S + , cstS + , addS + , subS + , mulS + , divByS + , expS + , logS + , mulConstS + , addConstS + , foldVadd + ) where + +import Control.Monad (when) +import Control.Monad.ST +import Data.Array.ST (STArray, newArray, readArray, writeArray) +import Data.STRef +import qualified Data.Vector.Storable as VS +import qualified Data.Vector.Unboxed as VU + +-- =========================================================================== +-- 値ハンドルと tape 文脈 +-- =========================================================================== + +-- | [日本語]: reverse-mode の値ハンドル: ノード id + primal (scalar / vector)。 +-- [English]: A reverse-mode value handle: node id + primal (scalar / vector). +data Rval = RScal !Int !Double | RVec !Int !(VS.Vector Double) + +ridOf :: Rval -> Int +ridOf (RScal i _) = i +ridOf (RVec i _) = i + +type Adj s = STArray s Int (VS.Vector Double) + +-- | [日本語]: 発番カウンタ + backward クロージャ列 (prepend = 発番の逆順)。 +-- [English]: A numbering counter plus the list of backward closures +-- (prepend order = reverse numbering order). +data Ctx s = Ctx !(STRef s Int) !(STRef s [Adj s -> ST s ()]) + +fresh :: Ctx s -> ST s Int +fresh (Ctx cnt _) = do + n <- readSTRef cnt + writeSTRef cnt (n + 1) + pure n + +record :: Ctx s -> (Adj s -> ST s ()) -> ST s () +record (Ctx _ bw) f = modifySTRef' bw (f :) + +-- | [日本語]: 随伴の加算 (空 = ゼロ扱い)。 +-- [English]: Accumulate into an adjoint (an empty vector is treated as zero). +bumpA :: Adj s -> Int -> VS.Vector Double -> ST s () +bumpA adj i contrib = do + cur <- readArray adj i + writeArray adj i (if VS.null cur then contrib else VS.zipWith (+) cur contrib) + +readAdjS :: Adj s -> Int -> ST s Double +readAdjS adj i = do + v <- readArray adj i + pure (if VS.null v then 0 else v VS.! 0) + +-- =========================================================================== +-- tape の実行 (forward build → seed → backward replay) +-- =========================================================================== + +-- | [日本語]: tape を構築するアクション (出力ノード + 勾配を読みたい leaf 群を +-- 返す) を受け取り、 forward 評価 → 出力に 1 を seed → backward replay の +-- 上で、 各 leaf の随伴 (= 出力の各 leaf に対する勾配ベクトル) を返す。 +-- +-- @build@ は @(出力 Rval, [leaf Rval])@ を返す。 結果は leaf ごとの随伴 +-- ベクトル (RScal leaf は長さ1、 RVec leaf は元の長さ)。 +-- [English]: Takes an action that builds the tape (returning the output +-- node plus the leaves whose gradient we want to read), then performs a +-- forward evaluation, seeds the output with 1, replays backward, and +-- returns the adjoint of each leaf (= the gradient vector with respect to +-- that leaf). +-- +-- @build@ returns @(output Rval, [leaf Rval])@. The result is one adjoint +-- vector per leaf (length 1 for an RScal leaf, original length for an +-- RVec leaf). +runTape :: (forall s. Ctx s -> ST s (Rval, [Rval])) -> [VS.Vector Double] +runTape build = runST $ do + cnt <- newSTRef 0 + bw <- newSTRef [] + let ctx = Ctx cnt bw + (out, leaves) <- build ctx + total <- readSTRef cnt + adj <- newArray (0, max 0 (total - 1)) VS.empty + writeArray adj (ridOf out) (VS.singleton 1) + closures <- readSTRef bw + mapM_ ($ adj) closures + mapM (\lf -> readArray adj (ridOf lf)) leaves + +-- =========================================================================== +-- leaf +-- =========================================================================== + +-- | [日本語]: ベクトル leaf (勾配を読む入力)。 +-- [English]: A vector leaf (an input whose gradient we read). +inputVec :: Ctx s -> VS.Vector Double -> ST s Rval +inputVec ctx v = do + i <- fresh ctx + pure (RVec i v) + +-- | [日本語]: スカラ leaf (勾配を読む入力)。 +-- [English]: A scalar leaf (an input whose gradient we read). +inputScal :: Ctx s -> Double -> ST s Rval +inputScal ctx x = do + i <- fresh ctx + pure (RScal i x) + +-- | [日本語]: 定数ベクトルノード (backward 無し)。 +-- [English]: A constant vector node (no backward closure). +constVec :: Ctx s -> VS.Vector Double -> ST s Rval +constVec ctx v = do { i <- fresh ctx; pure (RVec i v) } + +-- =========================================================================== +-- ベクトル演算 (随伴付き) +-- =========================================================================== + +-- | [日本語]: 全長 @l@ の vec から要素 @i@ を取り出す (scalar 化)。 随伴 = e_i·dy。 +-- [English]: Extracts element @i@ from a length-@l@ vec (turning it into a +-- scalar). Adjoint = e_i·dy. +idxHR :: Ctx s -> Int -> Int -> Rval -> ST s Rval +idxHR ctx l i (RVec vid v) = do + o <- fresh ctx + record ctx $ \adj -> do + g <- readAdjS adj o + when (g /= 0) $ bumpA adj vid (VS.generate l (\j -> if j == i then g else 0)) + pure (RScal o (v VS.! i)) +idxHR _ _ _ _ = error "idxHR: scalar input" + +-- | [日本語]: 全長 @l@ の vec から @[off, off+len)@ を切り出す。 随伴は +-- zeros l に散布。 +-- [English]: Slices out @[off, off+len)@ from a length-@l@ vec. The +-- adjoint is scattered into zeros of length l. +sliceHR :: Ctx s -> Int -> Int -> Int -> Rval -> ST s Rval +sliceHR ctx l off len (RVec vid v) = do + o <- fresh ctx + record ctx $ \adj -> do + dy <- readArray adj o + when (not (VS.null dy)) $ + bumpA adj vid (VS.generate l (\j -> if j >= off && j < off + len then dy VS.! (j - off) else 0)) + pure (RVec o (VS.slice off len v)) +sliceHR _ _ _ _ _ = error "sliceHR: scalar input" + +-- | [日本語]: scalar * vector。 ∂scalar = dy·v、 ∂v = scalar·dy。 +-- [English]: scalar * vector. ∂scalar = dy·v, ∂v = scalar·dy. +scaleHR :: Ctx s -> Rval -> Rval -> ST s Rval +scaleHR ctx (RScal kid k) (RVec vid v) = do + o <- fresh ctx + record ctx $ \adj -> do + dy <- readArray adj o + when (not (VS.null dy)) $ do + bumpA adj kid (VS.singleton (VS.sum (VS.zipWith (*) dy v))) + bumpA adj vid (VS.map (* k) dy) + pure (RVec o (VS.map (* k) v)) +scaleHR _ _ _ = error "scaleHR: shape" + +-- | [日本語]: vector + vector。 +-- [English]: vector + vector. +vaddHR :: Ctx s -> Rval -> Rval -> ST s Rval +vaddHR ctx (RVec aid a) (RVec bid b) = do + o <- fresh ctx + record ctx $ \adj -> do + dy <- readArray adj o + when (not (VS.null dy)) $ do + bumpA adj aid dy + bumpA adj bid dy + pure (RVec o (VS.zipWith (+) a b)) +vaddHR _ _ _ = error "vaddHR: shape" + +-- | [日本語]: vector - vector。 +-- [English]: vector - vector. +vsubHR :: Ctx s -> Rval -> Rval -> ST s Rval +vsubHR ctx (RVec aid a) (RVec bid b) = do + o <- fresh ctx + record ctx $ \adj -> do + dy <- readArray adj o + when (not (VS.null dy)) $ do + bumpA adj aid dy + bumpA adj bid (VS.map negate dy) + pure (RVec o (VS.zipWith (-) a b)) +vsubHR _ _ _ = error "vsubHR: shape" + +-- | [日本語]: 内積。 ∂a = dy·b、 ∂b = dy·a。 +-- [English]: Dot product. ∂a = dy·b, ∂b = dy·a. +dotHR :: Ctx s -> Rval -> Rval -> ST s Rval +dotHR ctx (RVec aid a) (RVec bid b) = do + o <- fresh ctx + record ctx $ \adj -> do + g <- readAdjS adj o + when (g /= 0) $ do + bumpA adj aid (VS.map (* g) b) + bumpA adj bid (VS.map (* g) a) + pure (RScal o (VS.sum (VS.zipWith (*) a b))) +dotHR _ _ _ = error "dotHR: shape" + +-- | [日本語]: @u[gids]@ gather (gids/nG は定数)。 随伴は scatter-add で O(n)。 +-- [English]: A @u[gids]@ gather (gids/nG are constants). The adjoint is +-- an O(n) scatter-add. +gatherHR :: Ctx s -> VU.Vector Int -> Int -> Rval -> ST s Rval +gatherHR ctx gids nG (RVec uid u) = do + let n = VU.length gids + o <- fresh ctx + record ctx $ \adj -> do + dy <- readArray adj o + when (not (VS.null dy)) $ + bumpA adj uid (VS.convert $ + VU.accumulate (+) (VU.replicate nG 0) (VU.zip gids (VU.convert dy :: VU.Vector Double))) + pure (RVec o (VS.generate n (\i -> u VS.! (gids VU.! i)))) +gatherHR _ _ _ _ = error "gatherHR: shape" + +-- | [日本語]: elementwise exp (非線形 μ 用)。 ∂v = dy ⊙ exp(v)。 +-- [English]: Elementwise exp (for a non-linear μ). ∂v = dy ⊙ exp(v). +vexpHR :: Ctx s -> Rval -> ST s Rval +vexpHR ctx (RVec vid v) = do + let ev = VS.map exp v + o <- fresh ctx + record ctx $ \adj -> do + dy <- readArray adj o + when (not (VS.null dy)) $ + bumpA adj vid (VS.zipWith (*) dy ev) + pure (RVec o ev) +vexpHR _ _ = error "vexpHR: scalar input" + +-- | [日本語]: scalar + vector の broadcast 加算。 ∂scalar = Σ dy。 +-- [English]: A broadcast addition of scalar + vector. ∂scalar = Σ dy. +bcastAddHR :: Ctx s -> Rval -> Rval -> ST s Rval +bcastAddHR ctx (RScal kid k) (RVec vid v) = do + o <- fresh ctx + record ctx $ \adj -> do + dy <- readArray adj o + when (not (VS.null dy)) $ do + bumpA adj kid (VS.singleton (VS.sum dy)) + bumpA adj vid dy + pure (RVec o (VS.map (+ k) v)) +bcastAddHR _ _ _ = error "bcastAddHR: shape" + +-- | [日本語]: elementwise 積 v ⊙ w (gather(a)[i]·exp(-b·x_i) 用)。 +-- ∂v = dy ⊙ w、 ∂w = dy ⊙ v。 +-- [English]: An elementwise product v ⊙ w (for gather(a)[i]·exp(-b·x_i)). +-- ∂v = dy ⊙ w, ∂w = dy ⊙ v. +hadamardHR :: Ctx s -> Rval -> Rval -> ST s Rval +hadamardHR ctx (RVec aid a) (RVec bid b) = do + o <- fresh ctx + record ctx $ \adj -> do + dy <- readArray adj o + when (not (VS.null dy)) $ do + bumpA adj aid (VS.zipWith (*) dy b) + bumpA adj bid (VS.zipWith (*) dy a) + pure (RVec o (VS.zipWith (*) a b)) +hadamardHR _ _ _ = error "hadamardHR: shape" + +-- | [日本語]: 汎用 elementwise 単項 (ベクトル式 IR の log/recip/sqrt/tanh 等)。 +-- @f@ とその導関数 @f'@ を受け、 ∂v = dy ⊙ f'(v) (v は入力 primal)。 +-- [English]: A generic elementwise unary op (for the vector-expression IR's +-- log/recip/sqrt/tanh, etc.). Takes @f@ and its derivative @f'@; +-- ∂v = dy ⊙ f'(v) (v is the input primal). +vmap1HR :: Ctx s -> (Double -> Double) -> (Double -> Double) -> Rval -> ST s Rval +vmap1HR ctx f df (RVec vid v) = do + o <- fresh ctx + record ctx $ \adj -> do + dy <- readArray adj o + when (not (VS.null dy)) $ + bumpA adj vid (VS.zipWith (\g x -> g * df x) dy v) + pure (RVec o (VS.map f v)) +vmap1HR _ _ _ _ = error "vmap1HR: scalar input" + +-- | [日本語]: 非空ベクトルノード列を vadd で畳む。 +-- [English]: Folds a non-empty list of vector nodes with vadd. +foldVadd :: Ctx s -> [Rval] -> ST s Rval +foldVadd _ [] = error "foldVadd: empty" +foldVadd _ [x] = pure x +foldVadd ctx (x:y:xs) = vaddHR ctx x y >>= \z -> foldVadd ctx (z : xs) + +-- =========================================================================== +-- スカラ演算 (随伴付き) +-- =========================================================================== + +cstS :: Ctx s -> Double -> ST s Rval +cstS ctx x = do { i <- fresh ctx; pure (RScal i x) } + +binS :: Ctx s -> (Double -> Double -> Double) -> (Double -> Double -> (Double, Double)) + -> Rval -> Rval -> ST s Rval +binS ctx f df (RScal aid a) (RScal bid b) = do + o <- fresh ctx + record ctx $ \adj -> do + g <- readAdjS adj o + when (g /= 0) $ do + let (da, db) = df a b + bumpA adj aid (VS.singleton (g * da)) + bumpA adj bid (VS.singleton (g * db)) + pure (RScal o (f a b)) +binS _ _ _ _ _ = error "binS: scalar expected" + +addS, subS, mulS :: Ctx s -> Rval -> Rval -> ST s Rval +addS ctx = binS ctx (+) (\_ _ -> (1, 1)) +subS ctx = binS ctx (-) (\_ _ -> (1, -1)) +mulS ctx = binS ctx (*) (\a b -> (b, a)) + +-- | [日本語]: scalar 除算 (a/b)。 +-- [English]: Scalar division (a/b). +divByS :: Ctx s -> Rval -> Rval -> ST s Rval +divByS ctx = binS ctx (/) (\a b -> (1 / b, negate a / (b * b))) + +unS :: Ctx s -> (Double -> Double) -> (Double -> Double) -> Rval -> ST s Rval +unS ctx f df (RScal aid a) = do + o <- fresh ctx + record ctx $ \adj -> do + g <- readAdjS adj o + when (g /= 0) $ bumpA adj aid (VS.singleton (g * df a)) + pure (RScal o (f a)) +unS _ _ _ _ = error "unS: scalar expected" + +expS, logS :: Ctx s -> Rval -> ST s Rval +expS ctx = unS ctx exp exp +logS ctx = unS ctx log (\a -> 1 / a) + +-- | [日本語]: 汎用スカラ単項。 @f@ と導関数 @f'@ を受ける ('unS' の公開形)。 +-- [English]: A generic scalar unary op. Takes @f@ and its derivative @f'@ +-- (the public form of 'unS'). +map1S :: Ctx s -> (Double -> Double) -> (Double -> Double) -> Rval -> ST s Rval +map1S = unS + +mulConstS, addConstS :: Ctx s -> Double -> Rval -> ST s Rval +mulConstS ctx c = unS ctx (* c) (const c) +addConstS ctx c = unS ctx (+ c) (const 1) diff --git a/src/Hanalyze/Stat/AD.hs b/src/Hanalyze/Stat/AD.hs new file mode 100644 --- /dev/null +++ b/src/Hanalyze/Stat/AD.hs @@ -0,0 +1,294 @@ +{-# LANGUAGE RankNTypes #-} +-- | +-- Module : Hanalyze.Stat.AD +-- Description : automatic differentiation (AD) による正確な勾配計算 (HMC 連携) +-- Copyright : (c) 2026 Aelysce Project (Toshiaki Honda) +-- License : BSD-3-Clause +-- +-- Exact gradient computation via automatic differentiation (AD), with HMC +-- integration. +-- +-- Uses reverse-mode AD from @Numeric.AD@ (ekmett/ad) to compute gradients. +-- More accurate than central-difference numerical differentiation, and runs +-- at comparable speed when the parameter count is small (< 100). +-- +-- == Usage +-- +-- The user writes @log p(θ, y)@ as a /Floating-polymorphic/ function. Fixed +-- observation values are lifted via @realToFrac@: +-- +-- @ +-- import Hanalyze.Stat.AD +-- import Hanalyze.Stat.Distribution (Transform (..)) +-- +-- -- θ = [mu, sigma] +-- myLogJoint :: [Double] -> LogJointF +-- myLogJoint obs [mu, sigma] = +-- logNormalF 0 10 mu -- prior: μ ~ N(0,10) +-- + logExpF 1 sigma -- prior: σ ~ Exp(1) +-- + sum [ logNormalObsF y mu sigma | y <- obs ] -- lik +-- +-- chain <- hmcAD (myLogJoint myData) +-- [UnconstrainedT, PositiveT] +-- defaultHMCConfig +-- ["mu","sigma"] +-- (Map.fromList [("mu",0),("sigma",1)]) +-- gen +-- @ +module Hanalyze.Stat.AD + ( -- * 多相対数密度関数 (log-joint 記述用) + LogJointF + , Params + , logNormalF + , logNormalObsF + , logExpF + , logGammaF + , logBetaF + , logPoissonObsF + , logBernoulliObsF + -- * AD-gradient computation + , gradAD + , gradADU + -- * HMC (AD variant) + , hmcAD + , hmcADChains + ) where + +import Control.Concurrent.Async (mapConcurrently) +import Control.Monad (forM, replicateM) +import Data.IORef +import qualified Data.Map.Strict as Map +import Data.Text (Text) +import Numeric.AD.Mode.Reverse.Double (grad) +import System.Random.MWC (GenIO, uniform) +import System.Random.MWC.Distributions (standard) + +import Hanalyze.MCMC.Core (Chain (..), spawnGen) +import Hanalyze.MCMC.HMC (HMCConfig (..), leapfrogWith, kinetic) +import Hanalyze.Stat.Distribution (Transform (..), toUnconstrained, fromUnconstrained) + +-- | Named parameter map (parameter name → constrained-space value). +type Params = Map.Map Text Double + +-- | Type alias for a 'Floating'-polymorphic log-joint function. The +-- argument @[a]@ is the constrained-space parameter vector. +type LogJointF = forall a. Floating a => [a] -> a + +-- --------------------------------------------------------------------------- +-- 多相対数密度関数 +-- --------------------------------------------------------------------------- + +-- | @log N(x; μ₀, σ₀)@ where @μ₀@ and @σ₀@ are fixed @Double@ +-- hyperparameters and @x@ is differentiable. +logNormalF :: Floating a => Double -> Double -> a -> a +logNormalF mu0 sig0 x = + let mu = realToFrac mu0 + sig = realToFrac sig0 + in negate (0.5 * log (2 * pi)) - log sig - 0.5 * ((x - mu) / sig) ^ (2::Int) +{-# INLINE logNormalF #-} + +-- | @log N(y_obs; μ, σ)@ where @y_obs@ is a fixed observation and @μ@, +-- @σ@ are differentiable. +logNormalObsF :: Floating a => Double -> a -> a -> a +logNormalObsF y_obs mu sig = + let y = realToFrac y_obs + in negate (0.5 * log (2 * pi)) - log sig - 0.5 * ((y - mu) / sig) ^ (2::Int) +{-# INLINE logNormalObsF #-} + +-- | @log Exp(x; rate)@ with fixed rate. +logExpF :: Floating a => Double -> a -> a +logExpF rate0 x = + let r = realToFrac rate0 + in log r - r * x +{-# INLINE logExpF #-} + +-- | @log Gamma(x; shape, rate)@ with fixed shape and rate. +-- +-- @log p(x) = (α-1) log x − β x + α log β − log Γ(α)@. +-- @log Γ(α)@ is Stirling's approximation (treated as a fixed constant). +logGammaF :: Floating a => Double -> Double -> a -> a +logGammaF shape0 rate0 x = + let a = realToFrac shape0 + b = realToFrac rate0 + lgA = realToFrac (stirlingLogGamma shape0) + in (a - 1) * log x - b * x + a * log b - lgA +{-# INLINE logGammaF #-} + +-- | @log Beta(x; α, β)@ with fixed shape parameters. +-- @log p(x) = (α-1) log x + (β-1) log(1-x) − log B(α,β)@. +logBetaF :: Floating a => Double -> Double -> a -> a +logBetaF alpha0 beta0 x = + let a = realToFrac alpha0 + b = realToFrac beta0 + lbB = realToFrac (stirlingLogGamma alpha0 + stirlingLogGamma beta0 + - stirlingLogGamma (alpha0 + beta0)) + in (a - 1) * log x + (b - 1) * log (1 - x) - lbB +{-# INLINE logBetaF #-} + +-- | @log Poisson(k | λ)@ with @k@ a fixed (rounded) observation and @λ@ +-- differentiable. +logPoissonObsF :: Floating a => Double -> a -> a +logPoissonObsF y_obs lam = + let k = fromIntegral (round y_obs :: Int) :: Double + lf = realToFrac (logFactorial (round y_obs :: Int)) + in realToFrac k * log lam - lam - lf +{-# INLINE logPoissonObsF #-} + +-- | @log Bernoulli(y | p)@ with @y ∈ {0, 1}@ a fixed observation and @p@ +-- differentiable. +logBernoulliObsF :: Floating a => Double -> a -> a +logBernoulliObsF y_obs p + | y_obs > 0.5 = log p + | otherwise = log (1 - p) +{-# INLINE logBernoulliObsF #-} + +-- --------------------------------------------------------------------------- +-- AD 勾配計算 +-- --------------------------------------------------------------------------- + +-- | Compute the gradient of a constrained-space log-joint via AD. +-- +-- @ +-- gradAD logJoint [1.0, 0.5] -- [∂/∂θ₁, ∂/∂θ₂] +-- @ +gradAD :: LogJointF -> [Double] -> [Double] +gradAD f xs = grad f xs + +-- | AD gradient of the log-joint in unconstrained space (with constraint +-- transforms and Jacobian correction applied automatically). +gradADU :: LogJointF -> [Transform] -> [Double] -> [Double] +gradADU logJointC transforms us = + grad (logJointUF transforms logJointC) us + +-- --------------------------------------------------------------------------- +-- 制約変換 (Floating 多相版) +-- --------------------------------------------------------------------------- + +-- | Map an unconstrained value to its constrained image +-- (Floating-polymorphic). +invTransformF :: Floating a => Transform -> a -> a +invTransformF UnconstrainedT u = u +invTransformF PositiveT u = exp u +invTransformF UnitIntervalT u = 1 / (1 + exp (-u)) -- sigmoid +{-# INLINE invTransformF #-} + +-- | Log-Jacobian @log |∂θ/∂u|@ for one parameter (Floating-polymorphic). +logJacF :: Floating a => Transform -> a -> a +logJacF UnconstrainedT _ = 0 +logJacF PositiveT u = u -- log(exp u) = u +logJacF UnitIntervalT u = + let p = 1 / (1 + exp (-u)) + in log p + log (1 - p) -- log σ(u)(1−σ(u)) +{-# INLINE logJacF #-} + +-- | Log-joint in unconstrained space, including constraint transforms +-- and the Jacobian correction. +logJointUF :: Floating a => [Transform] -> LogJointF -> [a] -> a +logJointUF transforms logJointC us = + let thetas = zipWith invTransformF transforms us + logJac = sum (zipWith logJacF transforms us) + in logJointC thetas + logJac + +-- --------------------------------------------------------------------------- +-- HMC AD 版サンプラー +-- --------------------------------------------------------------------------- + +-- | HMC sampler using AD gradients. +-- +-- Same algorithm as 'Hanalyze.MCMC.HMC.hmc', but gradients are computed exactly +-- with 'Numeric.AD.grad'. The user writes the log-joint in 'LogJointF' +-- form (i.e. @Floating@-polymorphic). +hmcAD + :: LogJointF -- ^ @log p(θ, y)@ as a 'LogJointF' (constrained space). + -> [Transform] -- ^ Per-parameter constraint kind (same order as the + -- parameter-name list). + -> HMCConfig + -> [Text] -- ^ Parameter names (matches the initial-value @Params@ keys). + -> Params -- ^ Initial values (constrained space). + -> GenIO + -> IO Chain +hmcAD logJointC transforms cfg names initC gen = do + let total = hmcBurnIn cfg + hmcIterations cfg + -- Unconstrained log-joint + logJU u = logJointUF transforms logJointC + [Map.findWithDefault 0 n u | n <- names] + -- AD gradient function: leapfrogWith の規約は ∇U = -∇logπ なので符号を反転 + gradFn ns paramsU = + let xs = [Map.findWithDefault 0 n paramsU | n <- ns] + in map negate (grad (logJointUF transforms logJointC) xs) + -- Initial unconstrained params + initU = Map.fromList + [ (n, toUnconstrained t v) + | (n, t) <- zip names transforms + , Just v <- [Map.lookup n initC] + ] + + samplesRef <- newIORef [] + acceptedRef <- newIORef (0 :: Int) + + let step currentU = do + r <- forM names (\_ -> standard gen) + let (proposedU, rFinal) = + leapfrogWith gradFn names + (hmcStepSize cfg) (hmcLeapfrogSteps cfg) + currentU r + logAlpha = (logJU proposedU - kinetic rFinal) + - (logJU currentU - kinetic r) + u <- uniform gen + if log (u :: Double) < logAlpha + then do modifyIORef' acceptedRef (+1); return proposedU + else return currentU + + let loop 0 currentU = return currentU + loop i currentU = do + nextU <- step currentU + when (i <= hmcIterations cfg) $ + modifyIORef' samplesRef + (Map.fromList + [ (n, fromUnconstrained t (Map.findWithDefault 0 n nextU)) + | (n, t) <- zip names transforms + ] :) + loop (i - 1) nextU + + _ <- loop total initU + samples <- fmap reverse (readIORef samplesRef) + accepted <- readIORef acceptedRef + return Chain + { chainSamples = samples + , chainAccepted = accepted + , chainTotal = total + , chainEnergy = [] + , chainDivergences = [] + } + where + when True action = action + when False _ = return () + +-- | Run 'hmcAD' on @numChains@ parallel chains. +hmcADChains + :: LogJointF + -> [Transform] + -> HMCConfig + -> Int + -> [Text] + -> Params + -> GenIO + -> IO [Chain] +hmcADChains logJointC transforms cfg numChains names initC baseGen = do + gens <- replicateM numChains (spawnGen baseGen) + mapConcurrently (\g -> hmcAD logJointC transforms cfg names initC g) gens + +-- --------------------------------------------------------------------------- +-- 数値ユーティリティ +-- --------------------------------------------------------------------------- + +-- Stirling 近似による log Γ(z) — z は固定 Double ハイパーパラメータ用 +stirlingLogGamma :: Double -> Double +stirlingLogGamma z + | z < 0.5 = log pi - log (sin (pi * z)) - stirlingLogGamma (1 - z) + | z < 12 = stirlingLogGamma (z + 1) - log z + | otherwise = (z - 0.5) * log z - z + 0.5 * log (2 * pi) + + 1/(12*z) - 1/(360*z^(3::Int)) + +logFactorial :: Int -> Double +logFactorial n = sum (map log [2 .. fromIntegral n]) diff --git a/src/Hanalyze/Stat/BayesFactor.hs b/src/Hanalyze/Stat/BayesFactor.hs new file mode 100644 --- /dev/null +++ b/src/Hanalyze/Stat/BayesFactor.hs @@ -0,0 +1,131 @@ +{-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE RankNTypes #-} +-- | +-- Module : Hanalyze.Stat.BayesFactor +-- Description : Bridge Sampling による Bayes Factor (Kass & Raftery 1995) 計算 +-- Copyright : (c) 2026 Aelysce Project (Toshiaki Honda) +-- License : BSD-3-Clause +-- +-- [日本語]: Bridge Sampling による Bayes Factor (Kass & Raftery 1995) の計算。 +-- +-- @ +-- BF_{10} = p(y | M_1) / p(y | M_0) +-- @ +-- +-- 計算は 2 モデルそれぞれに対し +-- 'Hanalyze.Stat.BridgeSampling.bridgeSampling' を呼び、 log marginal +-- 同士の差を取る。 解釈表は Kass-Raftery (1995) Table 1。 +-- +-- Reference: Kass & Raftery (1995) "Bayes factors". JASA 90:773-795. +-- [English]: Bayes Factor (Kass & Raftery 1995) via Bridge Sampling. +-- +-- @ +-- BF_{10} = p(y | M_1) / p(y | M_0) +-- @ +-- +-- The computation calls +-- 'Hanalyze.Stat.BridgeSampling.bridgeSampling' for each of the 2 +-- models and takes the difference between the log marginals. The +-- interpretation table is Kass-Raftery (1995) Table 1. +-- +-- Reference: Kass & Raftery (1995) "Bayes factors". JASA 90:773-795. +module Hanalyze.Stat.BayesFactor + ( BayesFactorResult (..) + , bayesFactor + , BFInterpretation (..) + , interpretBF + ) where + +import System.Random.MWC (GenIO) + +import Hanalyze.Model.HBM (ModelP) +import Hanalyze.MCMC.Core (Chain) +import Hanalyze.Stat.BridgeSampling + (BridgeConfig, BridgeResult (..), bridgeSampling) + +-- --------------------------------------------------------------------------- +-- Bayes Factor +-- --------------------------------------------------------------------------- + +data BayesFactorResult = BayesFactorResult + { bfLog10 :: !Double -- ^ log_10 BF_{10} + , bfLogE :: !Double -- ^ log_e BF_{10} = log p(y|M_1) - log p(y|M_0) + , bfLogMarginal0 :: !Double -- ^ log p(y | M_0) + , bfLogMarginal1 :: !Double -- ^ log p(y | M_1) + , bfConverged0 :: !Bool + , bfConverged1 :: !Bool + } deriving (Show) + +-- | [日本語]: 2 モデル間の Bayes Factor BF_{10} = p(y|M_1) / p(y|M_0)。 +-- 各モデルに対し Bridge Sampling で log marginal を推定、 差を取る。 +-- [English]: The Bayes Factor between 2 models, BF_{10} = p(y|M_1) / +-- p(y|M_0). Estimates the log marginal for each model via Bridge +-- Sampling and takes the difference. +bayesFactor + :: forall r0 r1. + ModelP r0 -> Chain -- ^ M_0 + posterior chain + -> ModelP r1 -> Chain -- ^ M_1 + posterior chain + -> BridgeConfig + -> GenIO + -> IO BayesFactorResult +bayesFactor m0 ch0 m1 ch1 cfg gen = do + r0 <- bridgeSampling m0 cfg ch0 gen + r1 <- bridgeSampling m1 cfg ch1 gen + let logE = brLogMarginal r1 - brLogMarginal r0 + log10v = logE / log 10 + pure BayesFactorResult + { bfLog10 = log10v + , bfLogE = logE + , bfLogMarginal0 = brLogMarginal r0 + , bfLogMarginal1 = brLogMarginal r1 + , bfConverged0 = brConverged r0 + , bfConverged1 = brConverged r1 + } + +-- --------------------------------------------------------------------------- +-- Kass-Raftery 解釈表 +-- --------------------------------------------------------------------------- + +-- | [日本語]: Bayes Factor の強度区分 (Kass & Raftery 1995 Table 1)。 +-- 区分の境界は @log_e BF@ で定義 (= log_10 ≈ /2.303): +-- +-- @ +-- 0 < log_e BF < 1 (1 < BF < 2.7) : Negligible +-- 1 ≤ log_e BF < 3 (2.7 ≤ BF < 20) : Positive (substantial) +-- 3 ≤ log_e BF < 5 (20 ≤ BF < 150) : Strong +-- 5 ≤ log_e BF (BF ≥ 150) : Very strong (decisive) +-- @ +-- +-- 負側は対称 (M_0 寄り)。 +-- [English]: The strength classification of a Bayes Factor (Kass & +-- Raftery 1995 Table 1). The boundaries between classes are defined in +-- terms of @log_e BF@ (= log_10 ≈ /2.303): +-- +-- @ +-- 0 < log_e BF < 1 (1 < BF < 2.7) : Negligible +-- 1 ≤ log_e BF < 3 (2.7 ≤ BF < 20) : Positive (substantial) +-- 3 ≤ log_e BF < 5 (20 ≤ BF < 150) : Strong +-- 5 ≤ log_e BF (BF ≥ 150) : Very strong (decisive) +-- @ +-- +-- The negative side is symmetric (favoring M_0). +data BFInterpretation + = BFNegligible + | BFPositive -- substantial evidence + | BFStrong + | BFVeryStrong + deriving (Show, Eq) + +-- | [日本語]: log_e BF 値から強度区分を返す。 符号で方向 (M_0 / M_1 どちらに寄与) は +-- 呼び出し側が判定する想定 (@abs logE@ を渡しても OK)。 +-- [English]: Returns the strength classification from a log_e BF value. +-- The sign (which direction, M_0 or M_1) is assumed to be judged by the +-- caller (passing @abs logE@ is also fine). +interpretBF :: Double -> BFInterpretation +interpretBF logE + | a < 1 = BFNegligible + | a < 3 = BFPositive + | a < 5 = BFStrong + | otherwise = BFVeryStrong + where + a = abs logE diff --git a/src/Hanalyze/Stat/BayesianModelAveraging.hs b/src/Hanalyze/Stat/BayesianModelAveraging.hs new file mode 100644 --- /dev/null +++ b/src/Hanalyze/Stat/BayesianModelAveraging.hs @@ -0,0 +1,121 @@ +{-# LANGUAGE OverloadedStrings #-} +-- | +-- Module : Hanalyze.Stat.BayesianModelAveraging +-- Description : Bridge Sampling の log marginal を用いた真の Bayesian Model Averaging (BMA) +-- Copyright : (c) 2026 Aelysce Project (Toshiaki Honda) +-- License : BSD-3-Clause +-- +-- True Bayesian Model Averaging (BMA) via Bridge Sampling log marginals. +-- +-- @ +-- p(θ | y) = Σ_k p(θ | y, M_k) · p(M_k | y) +-- p(M_k | y) ∝ p(y | M_k) · p(M_k) +-- @ +-- +-- [日本語]: +-- +-- 入力: 各モデルの __Bridge Sampling 推定 log marginal likelihood__ + +-- prior model weights (省略時 uniform 1/K)。 +-- +-- 出力: posterior model weights + 重み付き予測の helper。 +-- +-- ## 既存 pseudo-BMA との位置付け +-- +-- 既存 'Hanalyze.Stat.ModelSelect' の pseudo-BMA (= PSIS-LOO ベース近似) は +-- 軽量だが marginal likelihood を正しく計算していない (= LOO予測精度を代理 +-- 指標として使う)。 本 module は Bridge Sampling 経由で __真の log marginal__ +-- を使う BMA で、 解釈が一貫している (= Bayes Factor / 仮説検定と同じ基盤)。 +-- +-- [English]: +-- +-- Input: each model's __Bridge-Sampling-estimated log marginal likelihood__, +-- plus prior model weights (default: uniform 1/K). +-- +-- Output: posterior model weights + a helper for weighted prediction. +-- +-- ## Positioning relative to the existing pseudo-BMA +-- +-- The existing pseudo-BMA in 'Hanalyze.Stat.ModelSelect' (= a +-- PSIS-LOO-based approximation) is lightweight but does not correctly +-- compute the marginal likelihood (it uses LOO predictive accuracy as a +-- proxy metric instead). This module performs BMA using the +-- __true log marginal__ via Bridge Sampling, giving a consistent +-- interpretation (= the same foundation as a Bayes Factor / hypothesis +-- testing). +-- +-- Reference: Hoeting, Madigan, Raftery, Volinsky (1999) "Bayesian Model +-- Averaging: A Tutorial". Statistical Science 14(4):382-417. +module Hanalyze.Stat.BayesianModelAveraging + ( BMAResult (..) + , bayesianModelAveraging + , averagePredictions + ) where + +import qualified Numeric.LinearAlgebra as LA + +-- --------------------------------------------------------------------------- +-- BMA +-- --------------------------------------------------------------------------- + +data BMAResult = BMAResult + { bmaWeights :: ![Double] -- ^ [日本語]: posterior model weights @p(M_k|y)@、 Σ = 1 [English]: The posterior model weights @p(M_k|y)@, summing to 1. + , bmaLogMarginals :: ![Double] -- ^ [日本語]: 入力された per-model @log p(y|M_k)@ (引き継ぎ) [English]: The input per-model @log p(y|M_k)@ (passed through). + , bmaLogPriors :: ![Double] -- ^ [日本語]: 入力された per-model @log p(M_k)@ (引き継ぎ) [English]: The input per-model @log p(M_k)@ (passed through). + } deriving (Show) + +-- | [日本語]: log marginal + log prior weights (省略時 uniform) から +-- posterior model weights を計算 (softmax 安定化)。 +-- +-- @ +-- p(M_k | y) ∝ exp(log p(y|M_k) + log p(M_k)) +-- @ +-- +-- 引数の長さは同じである必要 (異なる場合は短い方に合わせる)。 全 log +-- marginal が -∞ なら uniform fallback。 +-- [English]: Computes posterior model weights from the log marginals + +-- log prior weights (default: uniform), via a numerically stabilized +-- softmax. +-- +-- @ +-- p(M_k | y) ∝ exp(log p(y|M_k) + log p(M_k)) +-- @ +-- +-- The argument lists must have equal length (if not, truncated to the +-- shorter one). Falls back to uniform weights if all log marginals are +-- -∞. +bayesianModelAveraging + :: [Double] -- ^ [日本語]: log marginals (Bridge Sampling 推定値 等) [English]: The log marginals (e.g. Bridge Sampling estimates) + -> Maybe [Double] -- ^ optional log prior weights (Nothing = uniform) + -> BMAResult +bayesianModelAveraging logMs mPriors = + let k = length logMs + logPriors = case mPriors of + Just ps | length ps == k -> ps + _ -> replicate k (- log (fromIntegral k)) + logUnnorm = zipWith (+) logMs logPriors + ws = if all isInfinite logUnnorm + then replicate k (1 / fromIntegral k) -- fallback uniform + else + let m = maximum logUnnorm + es = map (\x -> exp (x - m)) logUnnorm + s = sum es + in if s == 0 then replicate k (1 / fromIntegral k) + else map (/ s) es + in BMAResult + { bmaWeights = ws + , bmaLogMarginals = logMs + , bmaLogPriors = logPriors + } + +-- | [日本語]: per-model 予測ベクトル (= 各モデルから出した y* の posterior mean 等) を +-- BMA weights で加重平均。 全ベクトルは同じ長さである必要。 +-- [English]: Computes a weighted average of per-model prediction vectors +-- (e.g. each model's posterior mean of y*), using the BMA weights. All +-- vectors must have the same length. +averagePredictions :: BMAResult -> [LA.Vector Double] -> LA.Vector Double +averagePredictions bma preds + | null preds = LA.fromList [] + | length preds /= length (bmaWeights bma) = + error "averagePredictions: number of predictions ≠ number of weights" + | otherwise = + foldr1 (+) [ LA.scale w v | (w, v) <- zip (bmaWeights bma) preds ] diff --git a/src/Hanalyze/Stat/BridgeSampling.hs b/src/Hanalyze/Stat/BridgeSampling.hs new file mode 100644 --- /dev/null +++ b/src/Hanalyze/Stat/BridgeSampling.hs @@ -0,0 +1,301 @@ +{-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE RankNTypes #-} +{-# LANGUAGE BangPatterns #-} +-- | +-- Module : Hanalyze.Stat.BridgeSampling +-- Description : Bridge Sampling による周辺尤度 log p(y) 推定 (Meng & Wong 1996) +-- Copyright : (c) 2026 Aelysce Project (Toshiaki Honda) +-- License : BSD-3-Clause +-- +-- [日本語]: Bridge Sampling による周辺尤度 @log p(y)@ 推定 +-- (Meng & Wong 1996)。 +-- +-- Reference: +-- +-- * Meng & Wong (1996) "Simulating ratios of normalising constants +-- via a simple identity: a theoretical exploration". Statistica +-- Sinica 6:831-860. +-- * Gronau, Sarafoglou, Matzke, Ly, Boehm, Marsman, Leslie, Forster, +-- Wagenmakers, Steingroever (2017) "A tutorial on bridge sampling". +-- Journal of Mathematical Psychology 81:80-97. +-- +-- ## アルゴリズム +-- +-- 目的: 周辺尤度 @log p(y) = log ∫ p(y|θ) p(θ) dθ@ を、 既存 MCMC chain +-- (posterior samples) と diagonal Gaussian proposal @g(θ)@ から推定する。 +-- +-- Bridge identity (Meng-Wong): +-- +-- @ +-- p(y) = E_g[α(θ) q(θ)] / E_p[α(θ) g(θ)] +-- @ +-- +-- 最適 bridge function @α*(θ) = 1 / (s_1 q(θ) + s_2 r g(θ))@ を使った +-- iterative scheme で @r̂@ を求める: +-- +-- @ +-- r̂_{t+1} = [(1/N_2) Σ_i q(θ̃_2,i) / (s_1 q(θ̃_2,i) + s_2 r̂_t g(θ̃_2,i))] +-- / [(1/N_1) Σ_j g(θ̃_1,j) / (s_1 q(θ̃_1,j) + s_2 r̂_t g(θ̃_1,j))] +-- @ +-- +-- ここで: +-- * @θ̃_1@ は proposal @g@ から (本実装では Gaussian fit-to-chain) +-- * @θ̃_2@ は posterior chain サンプル +-- * @s_1 = N_1/(N_1+N_2)@、 @s_2 = N_2/(N_1+N_2)@ +-- * @q(θ) = p(y|θ)·p(θ)@ = @logJoint@ の exp 化 +-- +-- 全計算は __log space__ で行い (log-sum-exp 安定化)、 浮動小数 underflow を回避。 +-- +-- ## SMC との関係 +-- +-- SMC は副産物として log marginal を推定する (= temperature schedule の +-- incremental log-mean-weight 累積)。 Bridge Sampling は MCMC chain + proposal +-- から __独立な推定経路__ で求めるので、 両者が 5% 以内で一致すれば妥当性が裏付け。 +-- 不一致なら chain の収束不足 / SMC schedule 粗さ / proposal 不適切のサイン。 +-- +-- [English]: Bridge Sampling estimator of the marginal likelihood +-- @log p(y)@ (Meng & Wong 1996). +-- +-- Reference: +-- +-- * Meng & Wong (1996) "Simulating ratios of normalising constants +-- via a simple identity: a theoretical exploration". Statistica +-- Sinica 6:831-860. +-- * Gronau, Sarafoglou, Matzke, Ly, Boehm, Marsman, Leslie, Forster, +-- Wagenmakers, Steingroever (2017) "A tutorial on bridge sampling". +-- Journal of Mathematical Psychology 81:80-97. +-- +-- ## Algorithm +-- +-- Goal: estimate the marginal likelihood @log p(y) = log ∫ p(y|θ) p(θ) dθ@ +-- from an existing MCMC chain (posterior samples) and a diagonal Gaussian +-- proposal @g(θ)@. +-- +-- Bridge identity (Meng-Wong): +-- +-- @ +-- p(y) = E_g[α(θ) q(θ)] / E_p[α(θ) g(θ)] +-- @ +-- +-- Using the optimal bridge function @α*(θ) = 1 / (s_1 q(θ) + s_2 r g(θ))@, +-- @r̂@ is found via the following iterative scheme: +-- +-- @ +-- r̂_{t+1} = [(1/N_2) Σ_i q(θ̃_2,i) / (s_1 q(θ̃_2,i) + s_2 r̂_t g(θ̃_2,i))] +-- / [(1/N_1) Σ_j g(θ̃_1,j) / (s_1 q(θ̃_1,j) + s_2 r̂_t g(θ̃_1,j))] +-- @ +-- +-- where: +-- * @θ̃_1@ comes from the proposal @g@ (a Gaussian fit-to-chain in this +-- implementation) +-- * @θ̃_2@ are posterior chain samples +-- * @s_1 = N_1\/(N_1+N_2)@, @s_2 = N_2\/(N_1+N_2)@ +-- * @q(θ) = p(y|θ)·p(θ)@, i.e. the exponentiated @logJoint@ +-- +-- All computation is done in __log space__ (log-sum-exp stabilized) to +-- avoid floating-point underflow. +-- +-- ## Relationship to SMC +-- +-- SMC estimates the log marginal as a byproduct (= the cumulative +-- incremental log-mean-weight over the temperature schedule). Bridge +-- Sampling derives its estimate via an __independent path__ from the MCMC +-- chain and proposal, so agreement between the two within 5% corroborates +-- validity. Disagreement signals insufficient chain convergence, a coarse +-- SMC schedule, or an unsuitable proposal. +module Hanalyze.Stat.BridgeSampling + ( BridgeConfig (..) + , defaultBridgeConfig + , BridgeResult (..) + , bridgeSampling + ) where + +import Control.Monad (replicateM, forM) +import qualified Data.Map.Strict as Map +import Data.Text (Text) +import System.Random.MWC (GenIO) +import System.Random.MWC.Distributions (normal) + +import Hanalyze.Model.HBM (ModelP, Params, logJoint, sampleNames) +import Hanalyze.MCMC.Core (Chain (..), chainVals) + +-- --------------------------------------------------------------------------- +-- Configuration +-- --------------------------------------------------------------------------- + +-- | [日本語]: Bridge Sampling 設定。 +-- [English]: Bridge Sampling configuration. +data BridgeConfig = BridgeConfig + { bcNProposal :: !Int -- ^ [日本語]: N_1: proposal samples 数 (典型 chain サンプル数と同等)。 [English]: N_1: the number of proposal samples (typically comparable to the chain's sample count). + , bcMaxIter :: !Int -- ^ [日本語]: 反復解の最大回数 (典型 100、 通常 < 20 で収束)。 [English]: Maximum number of solver iterations (typically 100; usually converges in < 20). + , bcTolerance :: !Double -- ^ [日本語]: 反復収束判定 |Δ log r̂| < tol (典型 1e-6)。 [English]: Iterative convergence threshold |Δ log r̂| < tol (typically 1e-6). + } deriving (Show) + +defaultBridgeConfig :: BridgeConfig +defaultBridgeConfig = BridgeConfig + { bcNProposal = 500 + , bcMaxIter = 100 + , bcTolerance = 1e-6 + } + +-- | [日本語]: Bridge Sampling 結果。 +-- [English]: Bridge Sampling result. +data BridgeResult = BridgeResult + { brLogMarginal :: !Double -- ^ [日本語]: 推定 @log p(y)@。 [English]: Estimated @log p(y)@. + , brIterations :: !Int -- ^ [日本語]: 収束に要した反復数。 [English]: Number of iterations needed to converge. + , brConverged :: !Bool -- ^ [日本語]: tol 以内で収束したか。 [English]: Whether it converged within the tolerance. + } deriving (Show) + +-- --------------------------------------------------------------------------- +-- 公開 API +-- --------------------------------------------------------------------------- + +-- | [日本語]: Bridge Sampling で @log p(y)@ を推定。 +-- +-- 入力: +-- * モデル (logJoint = log q(θ) = log p(y|θ) + log p(θ)) +-- * posterior chain (既存 NUTS / MH / SMC 等の結果) +-- * proposal は __diagonal Gaussian fit__ to chain (各パラメータの sample +-- mean / SD から構築) +-- +-- 出力: log marginal likelihood 推定値 + 収束情報。 +-- [English]: Estimate @log p(y)@ via Bridge Sampling. +-- +-- Inputs: +-- * the model (logJoint = log q(θ) = log p(y|θ) + log p(θ)) +-- * a posterior chain (the result of an existing NUTS \/ MH \/ SMC run, +-- etc.) +-- * the proposal is a __diagonal Gaussian fit__ to the chain (built +-- from each parameter's sample mean \/ SD) +-- +-- Output: the estimated log marginal likelihood + convergence info. +bridgeSampling + :: forall r. ModelP r + -> BridgeConfig + -> Chain -- ^ posterior chain + -> GenIO + -> IO BridgeResult +bridgeSampling model cfg chain gen = do + let names = sampleNames model + posterior = chainSamples chain + n2 = length posterior + (mus, sds) = fitDiagGaussian names chain + -- 1. Sample N_1 from proposal g (diagonal Gaussian) + proposal <- replicateM (bcNProposal cfg) (sampleProposal names mus sds gen) + let n1 = length proposal + s1 = fromIntegral n1 / fromIntegral (n1 + n2) + s2 = fromIntegral n2 / fromIntegral (n1 + n2) + -- 2. Precompute log q (logJoint) and log g (proposal log-density) + logq2 = map (logJoint model) posterior + logq1 = map (logJoint model) proposal + logg2 = map (logProposal names mus sds) posterior + logg1 = map (logProposal names mus sds) proposal + -- 3. Iterative solve for log r̂ + let (logR, niter, converged) = + iterateBridge cfg logq1 logg1 logq2 logg2 s1 s2 0.0 + pure BridgeResult + { brLogMarginal = logR + , brIterations = niter + , brConverged = converged + } + +-- | Meng-Wong iterative formula in log space. +iterateBridge + :: BridgeConfig + -> [Double] -> [Double] -- ^ logq1, logg1 (proposal samples) + -> [Double] -> [Double] -- ^ logq2, logg2 (posterior samples) + -> Double -- ^ s_1 + -> Double -- ^ s_2 + -> Double -- ^ [日本語]: 初期 log r̂。 [English]: Initial log r̂. + -> (Double, Int, Bool) +iterateBridge cfg logq1 logg1 logq2 logg2 s1 s2 logR0 = go 0 logR0 + where + ls1 = log s1 + ls2 = log s2 + go !it !logR + | it >= bcMaxIter cfg = (logR, it, False) + | otherwise = + let -- Numerator: posterior 側の logq2 - logSumExp(s1·q2, s2·r·g2) + numTerms = + [ lq - logSumExp2 (ls1 + lq) (ls2 + logR + lg) + | (lq, lg) <- zip logq2 logg2 ] + -- Denominator: proposal 側の logg1 - logSumExp(s1·q1, s2·r·g1) + denTerms = + [ lg - logSumExp2 (ls1 + lq) (ls2 + logR + lg) + | (lq, lg) <- zip logq1 logg1 ] + num = logMeanExp numTerms + den = logMeanExp denTerms + logR' = num - den + diff = abs (logR' - logR) + in if diff < bcTolerance cfg + then (logR', it + 1, True) + else go (it + 1) logR' + +-- --------------------------------------------------------------------------- +-- Diagonal Gaussian proposal (fit-to-chain) +-- --------------------------------------------------------------------------- + +-- | [日本語]: chain から各パラメータの sample mean / SD を抽出。 SD = 0 になりうる +-- (定数推定) 場合は 1e-6 で下駄を履かせる (g(θ) 評価で除算 0 を避ける safety)。 +-- [English]: Extract each parameter's sample mean \/ SD from the chain. +-- When the SD could be 0 (a constant estimate), it is floored at 1e-6 +-- as a safety measure to avoid division by zero when evaluating g(θ). +fitDiagGaussian + :: [Text] -> Chain -> (Map.Map Text Double, Map.Map Text Double) +fitDiagGaussian names chain = + let mus = Map.fromList + [ (n, mean (chainVals n chain)) | n <- names ] + sds = Map.fromList + [ (n, max 1e-6 (stddev (chainVals n chain))) | n <- names ] + in (mus, sds) + where + mean xs = sum xs / fromIntegral (length xs) + stddev xs = + let mu = mean xs + n = fromIntegral (length xs) :: Double + in if n <= 1 then 0 + else sqrt (sum [(x - mu) ^ (2 :: Int) | x <- xs] / (n - 1)) + +-- | [日本語]: Diagonal Gaussian proposal からサンプル抽出。 +-- [English]: Draw a sample from the diagonal Gaussian proposal. +sampleProposal + :: [Text] -> Map.Map Text Double -> Map.Map Text Double -> GenIO + -> IO Params +sampleProposal names mus sds gen = + fmap Map.fromList $ forM names $ \n -> do + let mu = Map.findWithDefault 0 n mus + sd = Map.findWithDefault 1 n sds + x <- normal mu sd gen + pure (n, x) + +-- | [日本語]: θ における diagonal Gaussian proposal の log density。 +-- [English]: Log density of the diagonal Gaussian proposal at θ. +logProposal + :: [Text] -> Map.Map Text Double -> Map.Map Text Double -> Params + -> Double +logProposal names mus sds theta = + sum + [ let mu = Map.findWithDefault 0 n mus + sd = Map.findWithDefault 1 n sds + x = Map.findWithDefault 0 n theta + z = (x - mu) / sd + in -0.5 * log (2 * pi) - log sd - 0.5 * z * z + | n <- names ] + +-- --------------------------------------------------------------------------- +-- log-sum-exp helpers +-- --------------------------------------------------------------------------- + +logSumExp2 :: Double -> Double -> Double +logSumExp2 a b + | a > b = a + log (1 + exp (b - a)) + | otherwise = b + log (1 + exp (a - b)) + +logMeanExp :: [Double] -> Double +logMeanExp xs + | null xs = -1 / 0 + | otherwise = + let m = maximum xs + s = sum [ exp (x - m) | x <- xs ] + n = fromIntegral (length xs) :: Double + in m + log (s / n) diff --git a/src/Hanalyze/Stat/PosteriorPredictive.hs b/src/Hanalyze/Stat/PosteriorPredictive.hs new file mode 100644 --- /dev/null +++ b/src/Hanalyze/Stat/PosteriorPredictive.hs @@ -0,0 +1,153 @@ +{-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE RankNTypes #-} +-- | +-- Module : Hanalyze.Stat.PosteriorPredictive +-- Description : 事前/事後予測サンプリング (PyMC の sample_prior/posterior_predictive 相当) +-- Copyright : (c) 2026 Aelysce Project (Toshiaki Honda) +-- License : BSD-3-Clause +-- +-- Prior- and posterior-predictive sampling (analogous to PyMC's +-- @sample_prior_predictive@ / @sample_posterior_predictive@). +-- +-- @ +-- import Hanalyze.Stat.PosteriorPredictive +-- +-- chain <- nuts model cfg initP gen +-- ppc <- posteriorPredictive model chain gen +-- -- ppc :: [Map Text [Double]] -- predicted observations per sample +-- @ +module Hanalyze.Stat.PosteriorPredictive + ( -- * 事後予測サンプリング (chain ベース) + posteriorPredictive + , posteriorPredictiveSummary + -- * Prior predictive sampling (chain not required) + , priorPredictive + -- * Prior sampling (including latents) + , samplePrior + ) where + +import Control.Monad (replicateM) +import qualified Data.Map.Strict as Map +import Data.Map.Strict (Map) +import Data.Text (Text) +import Data.List (sort) +import System.Random.MWC (GenIO) + +import Hanalyze.MCMC.Core (Chain (..)) +import Hanalyze.Model.HBM + ( ModelP, sampleDist, runObserveDists, priorList ) + +-- --------------------------------------------------------------------------- +-- 事後予測サンプリング +-- --------------------------------------------------------------------------- + +-- | Posterior-predictive samples for every observe node in the model. +-- +-- Algorithm: +-- +-- 1. Walk the chain's latent samples. +-- 2. At each sample, evaluate 'runObserveDists' to obtain the +-- conditional distribution at every observe node. +-- 3. Draw as many fresh @y@ values from that distribution as the +-- original observation count. +-- +-- The returned list has the same length as @chainSamples@; each element +-- is a @Map@ from observe-node name to a fresh predicted-value list of +-- the original length. +posteriorPredictive + :: forall r. ModelP r + -> Chain + -> GenIO + -> IO [Map Text [Double]] +posteriorPredictive m chain gen = + mapM (\ps -> genFromObserves m ps gen) (chainSamples chain) + +-- | Per-observation posterior-predictive summary statistics +-- (mean and 95 % credible interval). +-- +-- Returns: observation name ↦ a list of @(mean, 2.5%, 97.5%)@ triples, +-- one per original observation index. +posteriorPredictiveSummary + :: [Map Text [Double]] -- posteriorPredictive の出力 + -> Map Text [(Double, Double, Double)] +posteriorPredictiveSummary preds = + let names = case preds of + [] -> [] + (m:_) -> Map.keys m + in Map.fromList + [ (n, summarizePerObs (perSamplePerObs n preds)) | n <- names ] + where + -- 観測 n: 各サンプルの観測 i 番目を集めて [[Double]] (列ごと) + perSamplePerObs :: Text -> [Map Text [Double]] -> [[Double]] + perSamplePerObs nm samples = + transpose (map (Map.findWithDefault [] nm) samples) + + summarizePerObs :: [[Double]] -> [(Double, Double, Double)] + summarizePerObs cols = map oneObs cols + where + oneObs xs = + let s = sort xs + n = length s + mu = if n == 0 then 0 else sum xs / fromIntegral n + q p = if n == 0 then 0 + else s !! min (n - 1) (max 0 (floor (p * fromIntegral n) :: Int)) + in (mu, q 0.025, q 0.975) + + transpose :: [[a]] -> [[a]] + transpose [] = [] + transpose xss + | all null xss = [] + | otherwise = + let heads = [h | (h:_) <- xss] + tails = [t | (_:t) <- xss] + in heads : transpose tails + +-- --------------------------------------------------------------------------- +-- 事前予測サンプリング (チェーン不要) +-- --------------------------------------------------------------------------- + +-- | Generate @N@ predictive samples from the prior alone (without any +-- observed data). Useful for sanity-checking what the model predicts +-- /before/ conditioning on observations. +priorPredictive + :: forall r. ModelP r + -> Int -- ^ Number of samples @N@. + -> GenIO + -> IO [Map Text [Double]] +priorPredictive m n gen = replicateM n $ do + ps <- samplePrior m gen + genFromObserves m ps gen + +-- | Draw one sample of every latent variable from its prior. +-- +-- Note: 'priorList' walks the model with placeholder zeros to extract its +-- structure. This function then samples each latent independently from +-- its individual prior. For hierarchical models this does not match +-- PyMC's @sample_prior_predictive@ (which threads downstream dependencies), +-- but it is enough for quick prior sanity checks. +samplePrior :: forall r. ModelP r -> GenIO -> IO (Map Text Double) +samplePrior m gen = do + let priors = priorList m -- [(name, Distribution Double)] (placeholder=0 走査) + vals <- mapM (\(_, d) -> sampleDist d gen) priors + return (Map.fromList (zip (map fst priors) vals)) + +-- --------------------------------------------------------------------------- +-- 内部: 与えられた latent 値で観測を生成 +-- --------------------------------------------------------------------------- + +-- 各 observe ノードについて、元データの個数だけ新しいサンプルを生成。 +genFromObserves + :: forall r. ModelP r + -> Map Text Double + -> GenIO + -> IO (Map Text [Double]) +genFromObserves m ps gen = do + let observes = runObserveDists m ps -- [(name, Distribution Double, [Double])] + newGroups <- mapM + (\(nm, d, ys) -> do + let nObs = length ys + newYs <- replicateM nObs (sampleDist d gen) + return (nm, newYs)) + observes + -- 同名 observe が複数ある場合はリスト連結 + return $ Map.fromListWith (++) newGroups diff --git a/src/Hanalyze/Stat/VI.hs b/src/Hanalyze/Stat/VI.hs new file mode 100644 --- /dev/null +++ b/src/Hanalyze/Stat/VI.hs @@ -0,0 +1,459 @@ +{-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE RankNTypes #-} +-- | +-- Module : Hanalyze.Stat.VI +-- Description : 変分推論 (ADVI: Automatic Differentiation Variational Inference) +-- Copyright : (c) 2026 Aelysce Project (Toshiaki Honda) +-- License : BSD-3-Clause +-- +-- Variational inference (ADVI — Automatic Differentiation Variational +-- Inference). +-- +-- Implements the mean-field normal VI of Kucukelbir et al. (2017). Uses +-- the same unconstrained transform as HMC/NUTS and maximizes the ELBO +-- with Adam. +-- +-- Approximating family: @q(u; φ) = Π_i Normal(u_i; μ_i, σ_i)@ +-- +-- @ +-- ELBO = E_q[log p(θ,y) + log|J|] + Σ_i H[Normal(μ_i, σ_i)] +-- = E_q[logJointU(u)] + Σ_i ω_i + N/2 × (1 + log 2π) +-- @ +-- +-- Gradient (reparameterization trick): +-- +-- @ +-- u^s = μ + σ ⊙ ε^s, ε^s ~ N(0, I) +-- ∂ELBO/∂μ_i ≈ (1/S) Σ_s ∂logJointU/∂u_i |_{u^s} +-- ∂ELBO/∂ω_i ≈ (1/S) Σ_s ε_i^s × σ_i × ∂logJointU/∂u_i |_{u^s} + 1 +-- @ +-- +-- @ +-- let cfg = defaultVIConfig { viIterations = 1000 } +-- result <- advi model cfg initParams gen +-- print (viPostMeans result) +-- @ +module Hanalyze.Stat.VI + ( VIConfig (..) + , defaultVIConfig + , VIResult (..) + , VIMethod (..) + , advi + , fullRankAdvi + ) where + +import Control.DeepSeq (force) +import Control.Monad (forM, forM_, replicateM, when) +import Data.IORef +import qualified Data.Map.Strict as Map +import System.Random.MWC (GenIO) +import System.Random.MWC.Distributions (standard) + +import Hanalyze.Model.HBM (ModelP, Params, sampleNames, getTransforms) +import Hanalyze.Optim.Adam (adamStep) +import Hanalyze.MCMC.HMC ( logJointU, paramsToVec, vecToParams + , toUnconstrainedParams, fromUnconstrainedParams ) + +-- --------------------------------------------------------------------------- +-- 設定 +-- --------------------------------------------------------------------------- + +-- | ADVI configuration. +data VIConfig = VIConfig + { viIterations :: Int -- ^ Number of Adam iterations. + , viSamples :: Int -- ^ Monte Carlo samples per ELBO gradient (5–10 typical). + , viLearningRate :: Double -- ^ Adam learning rate @α@. + , viBeta1 :: Double -- ^ Adam @β₁@ (default 0.9). + , viBeta2 :: Double -- ^ Adam @β₂@ (default 0.999). + , viEpsilon :: Double -- ^ Adam @ε@ (default 1e-8). + , viNumDraws :: Int -- ^ Number of post-fit draws from @q@. + , viGradStep :: Double -- ^ Finite-difference step for numeric gradients. + } deriving (Show) + +-- | Sensible defaults for ADVI: 1000 iterations, 5 MC samples, Adam at +-- @α = 0.1@. +defaultVIConfig :: VIConfig +defaultVIConfig = VIConfig + { viIterations = 1000 + , viSamples = 5 + , viLearningRate = 0.1 + , viBeta1 = 0.9 + , viBeta2 = 0.999 + , viEpsilon = 1e-8 + , viNumDraws = 2000 + , viGradStep = 1e-5 + } + +-- --------------------------------------------------------------------------- +-- 結果 +-- --------------------------------------------------------------------------- + +-- | [日本語]: VI 近似法。 mean-field (`advi`) と full-rank (`fullRankAdvi`) を区別する。 +-- [English]: The VI approximation method. Distinguishes mean-field +-- (`advi`) from full-rank (`fullRankAdvi`). +data VIMethod = MeanField | FullRank + deriving (Show, Eq) + +-- | [日本語]: ADVI result。 mean-field と full-rank の両方が返す。 full-rank では +-- @viCovU@ に @n×n@ 下三角 Cholesky 因子 @L@ (unconstrained 空間) が入る。 +-- [English]: The ADVI result. Returned by both mean-field and full-rank. +-- For full-rank, @viCovU@ holds the @n×n@ lower-triangular Cholesky +-- factor @L@ (in the unconstrained space). +data VIResult = VIResult + { viPostMeans :: Params -- ^ Posterior means (constrained space, sample mean). + , viPostSDs :: Params -- ^ Posterior SDs (constrained space). + , viMuU :: [Double] -- ^ Variational mean @μ@ (unconstrained). + , viSigmaU :: [Double] -- ^ [日本語]: Variational SD @σ@ (unconstrained、 mean-field の場合は対角要素、 full-rank なら L_ii)。 [English]: The variational SD @σ@ (unconstrained; the diagonal elements for mean-field, or L_ii for full-rank). + , viCovU :: Maybe [[Double]] -- ^ [日本語]: Full-rank ADVI: 下三角 Cholesky 因子 @L@ ([row][col])、 @LLᵀ = Σ@。 mean-field では @Nothing@。 [English]: Full-rank ADVI: the lower-triangular Cholesky factor @L@ ([row][col]), where @LLᵀ = Σ@. @Nothing@ for mean-field. + , viMethod :: VIMethod -- ^ [日本語]: どちらの近似法か。 [English]: Which approximation method was used. + , viElboHistory :: [Double] -- ^ ELBO trajectory (for convergence inspection). + , viDraws :: [Params] -- ^ Posterior draws in the constrained space (length 'viNumDraws'). + } deriving (Show) + +-- --------------------------------------------------------------------------- +-- ADVI +-- --------------------------------------------------------------------------- + +-- | Run mean-field normal ADVI. +-- +-- Optimization happens in unconstrained space; samples are mapped back +-- to the constrained space on the way out. Constrained parameters +-- (e.g. @Exponential → PositiveT@) are transformed automatically. +advi :: ModelP r -> VIConfig -> Params -> GenIO -> IO VIResult +advi model cfg initP gen = do + let names = sampleNames model + transforms = getTransforms model + n = length names + initU = paramsToVec names (toUnconstrainedParams transforms initP) + + -- unconstrained 空間での log p(θ,y) + log|J| (Jacobian 補正済み) + logJ :: [Double] -> Double + logJ uVec = logJointU model transforms (vecToParams names uVec) + + -- 有限差分勾配 ∂logJ/∂u + h = viGradStep cfg + numGrad :: [Double] -> [Double] + numGrad uVec = + [ let ui = uVec !! i + lp = logJ (replaceAt i (ui + h) uVec) + lm = logJ (replaceAt i (ui - h) uVec) + raw = (lp - lm) / (2 * h) + in if isNaN raw || isInfinite raw then 0 else raw + | i <- [0 .. n-1] + ] + + -- 変分パラメータ: μ (unconstrained 平均), ω = log(σ) (log 標準偏差) + muRef <- newIORef initU + omegaRef <- newIORef (replicate n 0.0) -- σ = exp(0) = 1 で初期化 + + -- Adam の 1次/2次モーメント + m1MuRef <- newIORef (replicate n 0.0) + m2MuRef <- newIORef (replicate n 0.0) + m1OmRef <- newIORef (replicate n 0.0) + m2OmRef <- newIORef (replicate n 0.0) + + elboRef <- newIORef [] + + let b1 = viBeta1 cfg + b2 = viBeta2 cfg + eps_ = viEpsilon cfg + alpha = viLearningRate cfg + sNum = viSamples cfg + + -- Adam ループ + forM_ [1 .. viIterations cfg] $ \t -> do + mu <- readIORef muRef + omega <- readIORef omegaRef + let sigma = map exp omega + + -- MC 勾配推定 + mcResults <- forM [1 .. sNum] $ \_ -> do + epsilons <- replicateM n (standard gen) + let -- u^s = μ + σ ⊙ ε (reparameterization) + uVec = zipWith3 (\m s e -> m + s * e) mu sigma epsilons + lj = logJ uVec + g = numGrad uVec + -- ∂ELBO/∂μ_i = ∂logJ/∂u_i + dMu = g + -- ∂ELBO/∂ω_i = ε_i × σ_i × ∂logJ/∂u_i + 1 (+1 はエントロピー項) + dOm = zipWith3 (\e s gi -> e * s * gi + 1) epsilons sigma g + return (lj, dMu, dOm) + + let sD = fromIntegral sNum :: Double + !ljMC = sum (map (\(l,_,_) -> l) mcResults) / sD + -- ELBO = E[logJointU] + Σω + N/2×(1+log2π) + !elboV = ljMC + sum omega + fromIntegral n * 0.5 * (1 + log (2*pi)) + !gMu = force (map (/ sD) $ foldr1 (zipWith (+)) (map (\(_,g,_) -> g) mcResults)) + !gOm = force (map (/ sD) $ foldr1 (zipWith (+)) (map (\(_,_,g) -> g) mcResults)) + + modifyIORef' elboRef (elboV :) + + -- Adam で μ を更新 + m1Mu <- readIORef m1MuRef + m2Mu <- readIORef m2MuRef + let (m1Mu', m2Mu', dxMu) = adamStep b1 b2 eps_ alpha t m1Mu m2Mu gMu + -- Phase Q3 (2026-05-14): 'zipWith (+)' / Adam の各リストは lazy で、 + -- IORef に書き戻すとそのまま thunk のまま積まれ、次イテレーションで + -- 読み出されると `zipWith (+) thunk_{t-1} ...` が再帰的に重なる。 + -- iter=10000 K=20 で max residency 85 MB / 総 alloc 222 GB を観測。 + -- 'force' で spine + 各要素を NF にし、t 階層の thunk チェーンを断つ。 + writeIORef m1MuRef (force m1Mu') + writeIORef m2MuRef (force m2Mu') + writeIORef muRef (force (zipWith (+) mu dxMu)) + + -- Adam で ω を更新 + m1Om <- readIORef m1OmRef + m2Om <- readIORef m2OmRef + let (m1Om', m2Om', dxOm) = adamStep b1 b2 eps_ alpha t m1Om m2Om gOm + writeIORef m1OmRef (force m1Om') + writeIORef m2OmRef (force m2Om') + writeIORef omegaRef (force (zipWith (+) omega dxOm)) + + -- 収束後: q(u; φ*) からサンプリングして constrained 空間に変換 + muFinal <- readIORef muRef + omegaFinal <- readIORef omegaRef + let sigmaFinal = map exp omegaFinal + + draws <- forM [1 .. viNumDraws cfg] $ \_ -> do + epsilons <- replicateM n (standard gen) + let uVec = zipWith3 (\m s e -> m + s * e) muFinal sigmaFinal epsilons + return (fromUnconstrainedParams transforms (vecToParams names uVec)) + + -- サンプルから事後平均・SD を計算 + let nD = fromIntegral (viNumDraws cfg) :: Double + getVals p = map (Map.findWithDefault 0 p) draws + muP p = let vs = getVals p in sum vs / nD + sdP p = let vs = getVals p + mu = muP p + in sqrt (sum (map (\v -> (v - mu) ^ (2::Int)) vs) / nD) + postMeans = Map.fromList [(nm, muP nm) | nm <- names] + postSDs = Map.fromList [(nm, sdP nm) | nm <- names] + + elboHistory <- fmap reverse (readIORef elboRef) + + return VIResult + { viPostMeans = postMeans + , viPostSDs = postSDs + , viMuU = muFinal + , viSigmaU = sigmaFinal + , viCovU = Nothing + , viMethod = MeanField + , viElboHistory = elboHistory + , viDraws = draws + } + +-- --------------------------------------------------------------------------- +-- Full-rank ADVI (Phase 37-A5) +-- --------------------------------------------------------------------------- + +-- | [日本語]: Full-rank ADVI: 共分散を含めた変分近似 @q(u) = N(μ, LLᵀ)@ を最適化する。 +-- +-- 平均場 'advi' との違い: +-- +-- - 変分パラメータは @μ@ (n-vector) と @L@ (下三角 n×n、 対角は log で +-- parameterize して正値保証) +-- - @u = μ + L·ε@ の reparameterization で勾配を取り、 ELBO の補正項は +-- @log|L| = Σ log L_ii = Σ ω_i@ +-- - 推定共分散 @Σ = LLᵀ@ は @viCovU@ に入る (下三角 @L@ そのもの) +-- +-- 平均場と比べて posterior の相関を捉えられるが、 パラメタ数 @O(n²)@、 +-- 計算量も @O(n² S)@ per iteration なので n が大きいモデルでは重い。 +-- 平均場が「SD を過小評価」 する hierarchical model で特に有用。 +-- [English]: Full-rank ADVI: optimizes the variational approximation +-- @q(u) = N(μ, LLᵀ)@, including the covariance. +-- +-- Differences from mean-field 'advi': +-- +-- - The variational parameters are @μ@ (n-vector) and @L@ (lower +-- triangular n×n; the diagonal is parameterized via log to guarantee +-- positivity) +-- - Gradients are taken via the reparameterization @u = μ + L·ε@; the +-- ELBO's correction term is @log|L| = Σ log L_ii = Σ ω_i@ +-- - The estimated covariance @Σ = LLᵀ@ is stored in @viCovU@ (the +-- lower-triangular @L@ itself) +-- +-- Compared to mean-field, this captures posterior correlations, but has +-- @O(n²)@ parameters and @O(n² S)@ compute per iteration, so it becomes +-- heavy for models with large n. Especially useful for hierarchical +-- models where mean-field "underestimates the SD". +fullRankAdvi :: ModelP r -> VIConfig -> Params -> GenIO -> IO VIResult +fullRankAdvi model cfg initP gen = do + let names = sampleNames model + transforms = getTransforms model + n = length names + initU = paramsToVec names (toUnconstrainedParams transforms initP) + + logJ :: [Double] -> Double + logJ uVec = logJointU model transforms (vecToParams names uVec) + + h = viGradStep cfg + numGrad :: [Double] -> [Double] + numGrad uVec = + [ let ui = uVec !! i + lp = logJ (replaceAt i (ui + h) uVec) + lm = logJ (replaceAt i (ui - h) uVec) + raw = (lp - lm) / (2 * h) + in if isNaN raw || isInfinite raw then 0 else raw + | i <- [0 .. n-1] + ] + + -- 変分パラメータ: μ (n-vector)、 ω (n-vector、 ω_i = log L_ii)、 + -- offdiag (下三角の i > j 要素を行優先で並べた長さ n(n-1)/2 のリスト) + muRef <- newIORef initU + omegaRef <- newIORef (replicate n 0.0) -- L_ii = exp(0) = 1 + let nOff = n * (n - 1) `div` 2 + offRef <- newIORef (replicate nOff 0.0) -- off-diag は 0 で初期化 + + -- Adam モーメント (μ / ω / offdiag それぞれ) + m1MuRef <- newIORef (replicate n 0.0) + m2MuRef <- newIORef (replicate n 0.0) + m1OmRef <- newIORef (replicate n 0.0) + m2OmRef <- newIORef (replicate n 0.0) + m1OffRef <- newIORef (replicate nOff 0.0) + m2OffRef <- newIORef (replicate nOff 0.0) + + elboRef <- newIORef [] + + let b1 = viBeta1 cfg + b2 = viBeta2 cfg + eps_ = viEpsilon cfg + alpha = viLearningRate cfg + sNum = viSamples cfg + + forM_ [1 .. viIterations cfg] $ \t -> do + mu <- readIORef muRef + omega <- readIORef omegaRef + offdg <- readIORef offRef + let lMat = buildL n omega offdg -- 下三角 L + + -- MC 勾配 + mcResults <- forM [1 .. sNum] $ \_ -> do + epsilons <- replicateM n (standard gen) + let uVec = vecAdd mu (matVec lMat epsilons) + lj = logJ uVec + gU = numGrad uVec -- ∂lp/∂u_i, length n + dMu = gU -- ∂ELBO/∂μ_i = gU_i + -- ∂ELBO/∂ω_i = ε_i × L_ii × gU_i + 1 (entropy +1) + dOm = [ epsilons !! i + * (lMat !! i !! i) + * (gU !! i) + 1 + | i <- [0 .. n-1] ] + -- ∂ELBO/∂L_ij (i > j) = ε_j × gU_i (no entropy contribution) + dOff = [ (epsilons !! j) * (gU !! i) + | i <- [1 .. n-1], j <- [0 .. i-1] ] + return (lj, dMu, dOm, dOff) + + let sD = fromIntegral sNum :: Double + !ljMC = sum (map (\(l,_,_,_) -> l) mcResults) / sD + -- ELBO = E[logJointU] + log|L| + n/2 (1 + log 2π) + !elboV = ljMC + sum omega + fromIntegral n * 0.5 * (1 + log (2*pi)) + !gMu = force (map (/ sD) $ foldr1 (zipWith (+)) + (map (\(_,g,_,_) -> g) mcResults)) + !gOm = force (map (/ sD) $ foldr1 (zipWith (+)) + (map (\(_,_,g,_) -> g) mcResults)) + !gOff = if nOff == 0 + then [] + else force (map (/ sD) $ foldr1 (zipWith (+)) + (map (\(_,_,_,g) -> g) mcResults)) + + modifyIORef' elboRef (elboV :) + + -- Adam で μ + m1Mu <- readIORef m1MuRef + m2Mu <- readIORef m2MuRef + let (m1Mu', m2Mu', dxMu) = adamStep b1 b2 eps_ alpha t m1Mu m2Mu gMu + writeIORef m1MuRef (force m1Mu') + writeIORef m2MuRef (force m2Mu') + writeIORef muRef (force (zipWith (+) mu dxMu)) + + -- Adam で ω + m1Om <- readIORef m1OmRef + m2Om <- readIORef m2OmRef + let (m1Om', m2Om', dxOm) = adamStep b1 b2 eps_ alpha t m1Om m2Om gOm + writeIORef m1OmRef (force m1Om') + writeIORef m2OmRef (force m2Om') + writeIORef omegaRef (force (zipWith (+) omega dxOm)) + + -- Adam で off-diagonal (n=1 のときは空) + when (nOff > 0) $ do + m1Off <- readIORef m1OffRef + m2Off <- readIORef m2OffRef + let (m1Off', m2Off', dxOff) = adamStep b1 b2 eps_ alpha t m1Off m2Off gOff + writeIORef m1OffRef (force m1Off') + writeIORef m2OffRef (force m2Off') + writeIORef offRef (force (zipWith (+) offdg dxOff)) + + -- 収束後 + muFinal <- readIORef muRef + omegaFinal <- readIORef omegaRef + offFinal <- readIORef offRef + let lFinal = buildL n omegaFinal offFinal + lDiag = [ lFinal !! i !! i | i <- [0 .. n-1] ] + + draws <- forM [1 .. viNumDraws cfg] $ \_ -> do + epsilons <- replicateM n (standard gen) + let uVec = vecAdd muFinal (matVec lFinal epsilons) + return (fromUnconstrainedParams transforms (vecToParams names uVec)) + + let nD = fromIntegral (viNumDraws cfg) :: Double + getVals p = map (Map.findWithDefault 0 p) draws + muP p = let vs = getVals p in sum vs / nD + sdP p = let vs = getVals p + mu = muP p + in sqrt (sum (map (\v -> (v - mu) ^ (2::Int)) vs) / nD) + postMeans = Map.fromList [(nm, muP nm) | nm <- names] + postSDs = Map.fromList [(nm, sdP nm) | nm <- names] + + elboHistory <- fmap reverse (readIORef elboRef) + + return VIResult + { viPostMeans = postMeans + , viPostSDs = postSDs + , viMuU = muFinal + , viSigmaU = lDiag + , viCovU = Just lFinal + , viMethod = FullRank + , viElboHistory = elboHistory + , viDraws = draws + } + +-- | [日本語]: 下三角 L を構築。 @omega@ は対角 (L_ii = exp ω_i)、 +-- @offdg@ は (i, j) for i > j を行優先 (i 昇順、 同 i 内で j 昇順) で +-- 並べたリスト。 結果は @n × n@ 行列、 上三角は 0。 +-- [English]: Builds the lower-triangular @L@. @omega@ is the diagonal +-- (L_ii = exp ω_i); @offdg@ is the list of (i, j) for i > j in row-major +-- order (ascending i, then ascending j within the same i). The result is +-- an @n × n@ matrix with the upper triangle set to 0. +buildL :: Int -> [Double] -> [Double] -> [[Double]] +buildL n omega offdg = + let -- offdg をインデックス map に変換 + offMap = Map.fromList (zip pairs offdg) + pairs = [ (i, j) | i <- [1 .. n-1], j <- [0 .. i-1] ] + diag i = exp (omega !! i) + row i = [ if j < i then Map.findWithDefault 0 (i, j) offMap + else if j == i then diag i + else 0 + | j <- [0 .. n-1] ] + in [ row i | i <- [0 .. n-1] ] + +-- | [日本語]: 行列・ベクトル積 @y = M·x@。 +-- [English]: Matrix-vector product @y = M·x@. +matVec :: [[Double]] -> [Double] -> [Double] +matVec mat x = [ sum (zipWith (*) row x) | row <- mat ] + +-- | [日本語]: ベクトル足し算。 +-- [English]: Vector addition. +vecAdd :: [Double] -> [Double] -> [Double] +vecAdd = zipWith (+) + +-- --------------------------------------------------------------------------- +-- 補助関数 +-- --------------------------------------------------------------------------- + +-- adamStep は Hanalyze.Optim.Adam に集約 (Phase R0)。 +-- 再 export することで既存の利用箇所はそのまま動く。 + +-- | [日本語]: リストの i 番目要素を x で置換する。 +-- [English]: Replaces the i-th element of a list with x. +replaceAt :: Int -> Double -> [Double] -> [Double] +replaceAt i x xs = take i xs ++ [x] ++ drop (i + 1) xs