saddm 0.1.0__tar.gz

saddm-0.1.0/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2026 Kianté Fernandez, Blair R K Shevlin, Roger Ratcliff, Ian Krajbich
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

saddm-0.1.0/PKG-INFO

@@ -0,0 +1,103 @@
+Metadata-Version: 2.4
+Name: saddm
+Version: 0.1.0
+Summary: Drift diffusion model with across-trial variability in boundary separation (DDM-SA): differentiable PyTensor likelihood for PyMC and HSSM
+Author: Kianté Fernandez, Blair R K Shevlin, Roger Ratcliff, Ian Krajbich
+License: MIT
+Project-URL: Repository, https://github.com/kiante-fernandez/saddm
+Classifier: Programming Language :: Python :: 3
+Classifier: License :: OSI Approved :: MIT License
+Classifier: Topic :: Scientific/Engineering
+Requires-Python: >=3.11
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Requires-Dist: numpy
+Requires-Dist: scipy
+Requires-Dist: pytensor
+Provides-Extra: sampling
+Requires-Dist: pymc>=5.20; extra == "sampling"
+Requires-Dist: arviz; extra == "sampling"
+Requires-Dist: jax==0.5.3; extra == "sampling"
+Requires-Dist: jaxlib==0.5.3; extra == "sampling"
+Requires-Dist: numpyro==0.19.0; extra == "sampling"
+Provides-Extra: hssm
+Requires-Dist: saddm[sampling]; extra == "hssm"
+Requires-Dist: hssm>=0.3.0; extra == "hssm"
+Requires-Dist: matplotlib<3.11; extra == "hssm"
+Provides-Extra: reference
+Requires-Dist: numba; extra == "reference"
+Provides-Extra: test
+Requires-Dist: pytest; extra == "test"
+Requires-Dist: numba; extra == "test"
+Requires-Dist: pymc>=5.20; extra == "test"
+Requires-Dist: arviz; extra == "test"
+Dynamic: license-file
+
+# Implementation of Diffusion decision model with across-trial variability in boundary separation
+
+Diffusion decision model with across-trial variability in boundary separation
+(`sa`), drift (`sv`), and non-decision time (`st`). This codebase presents the 
+DDM-SA as a fully differentiable PyTensor likelihood for gradient-based Bayesian estimation.
+
+The likelihood is analytic (Navarro–Fuss density; drift variability integrated
+in closed form, uniform variability by Gauss–Legendre quadrature), so NUTS gets
+exact gradients. It is validated against the Fortran implementation the model
+was originally developed in; the Fortran programs, the data, and reference
+results ship in this repository so every validation is reproducible from a
+clone.
+
+## Install
+
+```bash
+pip install -e .                # core: numpy, scipy, pytensor
+pip install -e ".[sampling]"    # + pymc, arviz, numpyro (pinned jax)
+pip install -e ".[hssm]"        # + hssm
+pip install -e ".[test]"        # + pytest, numba (reference backend)
+```
+
+## Quickstart
+
+```python
+from saddm import make_ddmsa_model, sample_ddmsa, sample_ddmsa_exact
+
+data  = sample_ddmsa_exact(a=1.1, z=0.5, v=1.5, t=0.25,
+                           sv=0.8, sa=0.5, st=0.08, n_trials=2000)
+idata = sample_ddmsa(make_ddmsa_model(data), backend="numpyro")
+```
+
+`saddm.ddmsa_logp(rt, response, a, z, v, t, sv, sa, st, sz)` is the per-trial
+log-likelihood; every parameter may be a scalar or a per-trial vector. It has
+no lapse mixture: HSSM applies its own `p_outlier` on top of any analytical
+likelihood, and in plain PyMC one is a `pm.logaddexp` away.
+
+With HSSM, register `saddm.ddmsa_logp` as a `loglik_kind="analytical"`
+likelihood; `examples/estimate_HSSM_saddm.py` is the minimal adapter.
+
+## Layout
+
+| path | contents |
+|---|---|
+| `saddm/` | `ddmsa.py`: the likelihood and PyMC glue. `core.py`/`integrator.py`/`model.py`: the Numba reference implementation (`reference` extra). |
+| `tests/` | `test_ddmsa.py`: verification suite — s = 1 closed forms, agreement with the Numba/Fortran reference, finite-difference gradients, corner finiteness, per-trial broadcasting, backend agreement, static-zero collapse (run directly with `--sample` for an end-to-end NUTS check). Remaining `test_*.py` cover the Numba reference. |
+| `verification/` | `parameter_recovery.py`: 100-config NUTS recovery study. `recovery_figure.py`, `compare_to_fortran.py`, `likelihood_figure.py`: analysis and figures (read `results/reference/` by default; set `RESULTS` for a fresh run). |
+| `examples/` | HSSM applications: flat fit on cavanagh_theta, the per-subject + k-sweep replication of the Fortran intertemporal-choice analysis, hierarchical variants, and the random-effects figure. |
+| `fortran/` | The Fortran programs that produced the benchmarks, with build notes. |
+| `data/itc_amasino/` | Amasino et al. (2019) trials, the Fortran benchmarks, and the exact k-sweep permutation files. |
+| `results/reference/` | Reference outputs: recovery, ITC, hierarchical, and cavanagh results with figures. Everything else under `results/` is gitignored, and every script writes there by default. |
+
+## Citation
+
+```bibtex
+@unpublished{shevlin2026little,
+  author = {Shevlin, Blair R. K. and Fernandez, Kiant{\'e} and Ratcliff, Roger and Krajbich, Ian},
+  title  = {A little goes a long way: Fitting one-shot decisions with cognitive models},
+  note   = {Manuscript in preparation},
+  year   = {2026},
+}
+```
+
+Blair R. K. Shevlin\* and Kianté Fernandez\* contributed equally.
+
+:Author: Kianté Fernandez, Blair R K Shevlin, Roger Ratcliff, Ian Krajbich
+
+:Contact: kiante@ucla.edu, blair.shevlin@mssm.edu, ratcliff.22@osu.edu, krajbich@ucla.edu

saddm-0.1.0/README.md

@@ -0,0 +1,68 @@
+# Implementation of Diffusion decision model with across-trial variability in boundary separation
+
+Diffusion decision model with across-trial variability in boundary separation
+(`sa`), drift (`sv`), and non-decision time (`st`). This codebase presents the 
+DDM-SA as a fully differentiable PyTensor likelihood for gradient-based Bayesian estimation.
+
+The likelihood is analytic (Navarro–Fuss density; drift variability integrated
+in closed form, uniform variability by Gauss–Legendre quadrature), so NUTS gets
+exact gradients. It is validated against the Fortran implementation the model
+was originally developed in; the Fortran programs, the data, and reference
+results ship in this repository so every validation is reproducible from a
+clone.
+
+## Install
+
+```bash
+pip install -e .                # core: numpy, scipy, pytensor
+pip install -e ".[sampling]"    # + pymc, arviz, numpyro (pinned jax)
+pip install -e ".[hssm]"        # + hssm
+pip install -e ".[test]"        # + pytest, numba (reference backend)
+```
+
+## Quickstart
+
+```python
+from saddm import make_ddmsa_model, sample_ddmsa, sample_ddmsa_exact
+
+data  = sample_ddmsa_exact(a=1.1, z=0.5, v=1.5, t=0.25,
+                           sv=0.8, sa=0.5, st=0.08, n_trials=2000)
+idata = sample_ddmsa(make_ddmsa_model(data), backend="numpyro")
+```
+
+`saddm.ddmsa_logp(rt, response, a, z, v, t, sv, sa, st, sz)` is the per-trial
+log-likelihood; every parameter may be a scalar or a per-trial vector. It has
+no lapse mixture: HSSM applies its own `p_outlier` on top of any analytical
+likelihood, and in plain PyMC one is a `pm.logaddexp` away.
+
+With HSSM, register `saddm.ddmsa_logp` as a `loglik_kind="analytical"`
+likelihood; `examples/estimate_HSSM_saddm.py` is the minimal adapter.
+
+## Layout
+
+| path | contents |
+|---|---|
+| `saddm/` | `ddmsa.py`: the likelihood and PyMC glue. `core.py`/`integrator.py`/`model.py`: the Numba reference implementation (`reference` extra). |
+| `tests/` | `test_ddmsa.py`: verification suite — s = 1 closed forms, agreement with the Numba/Fortran reference, finite-difference gradients, corner finiteness, per-trial broadcasting, backend agreement, static-zero collapse (run directly with `--sample` for an end-to-end NUTS check). Remaining `test_*.py` cover the Numba reference. |
+| `verification/` | `parameter_recovery.py`: 100-config NUTS recovery study. `recovery_figure.py`, `compare_to_fortran.py`, `likelihood_figure.py`: analysis and figures (read `results/reference/` by default; set `RESULTS` for a fresh run). |
+| `examples/` | HSSM applications: flat fit on cavanagh_theta, the per-subject + k-sweep replication of the Fortran intertemporal-choice analysis, hierarchical variants, and the random-effects figure. |
+| `fortran/` | The Fortran programs that produced the benchmarks, with build notes. |
+| `data/itc_amasino/` | Amasino et al. (2019) trials, the Fortran benchmarks, and the exact k-sweep permutation files. |
+| `results/reference/` | Reference outputs: recovery, ITC, hierarchical, and cavanagh results with figures. Everything else under `results/` is gitignored, and every script writes there by default. |
+
+## Citation
+
+```bibtex
+@unpublished{shevlin2026little,
+  author = {Shevlin, Blair R. K. and Fernandez, Kiant{\'e} and Ratcliff, Roger and Krajbich, Ian},
+  title  = {A little goes a long way: Fitting one-shot decisions with cognitive models},
+  note   = {Manuscript in preparation},
+  year   = {2026},
+}
+```
+
+Blair R. K. Shevlin\* and Kianté Fernandez\* contributed equally.
+
+:Author: Kianté Fernandez, Blair R K Shevlin, Roger Ratcliff, Ian Krajbich
+
+:Contact: kiante@ucla.edu, blair.shevlin@mssm.edu, ratcliff.22@osu.edu, krajbich@ucla.edu

saddm-0.1.0/pyproject.toml

@@ -0,0 +1,43 @@
+[build-system]
+requires = ["setuptools>=64"]
+build-backend = "setuptools.build_meta"
+
+[project]
+name = "saddm"
+version = "0.1.0"
+description = "Drift diffusion model with across-trial variability in boundary separation (DDM-SA): differentiable PyTensor likelihood for PyMC and HSSM"
+readme = "README.md"
+license = { text = "MIT" }
+authors = [
+    { name = "Kianté Fernandez" },
+    { name = "Blair R K Shevlin" },
+    { name = "Roger Ratcliff" },
+    { name = "Ian Krajbich" },
+]
+requires-python = ">=3.11"
+classifiers = [
+    "Programming Language :: Python :: 3",
+    "License :: OSI Approved :: MIT License",
+    "Topic :: Scientific/Engineering",
+]
+dependencies = [
+    "numpy",
+    "scipy",
+    "pytensor",
+]
+
+[project.urls]
+Repository = "https://github.com/kiante-fernandez/saddm"
+
+[project.optional-dependencies]
+# jax is pinned: jax 0.10 silently freezes numpyro NUTS at its initial point.
+sampling = ["pymc>=5.20", "arviz", "jax==0.5.3", "jaxlib==0.5.3", "numpyro==0.19.0"]
+hssm = ["saddm[sampling]", "hssm>=0.3.0", "matplotlib<3.11"]
+reference = ["numba"]
+test = ["pytest", "numba", "pymc>=5.20", "arviz"]
+
+[tool.setuptools]
+packages = ["saddm"]
+
+[tool.pytest.ini_options]
+testpaths = ["tests"]

saddm-0.1.0/saddm/__init__.py

@@ -0,0 +1,33 @@
+"""
+saddm - Drift Diffusion Model with across-trial variability in boundary separation.
+
+Parameters (s = 1): a boundary separation, z relative start point in (0, 1),
+v drift rate, t non-decision time (seconds); sv Gaussian drift variability;
+sa, st, sz uniform full-width variability of boundary, non-decision time, and
+start point.
+
+The Numba reference implementation (saddm.core, saddm.integrator, saddm.model)
+needs the `reference` extra and is imported explicitly.
+"""
+
+__version__ = "0.1.0"
+
+from .ddmsa import (
+    DDMSA,
+    ddmsa_logp,
+    ddmsa_potential,
+    make_ddmsa_model,
+    sample_ddmsa,
+    sample_ddmsa_exact,
+    simulate_ddmsa,
+)
+
+__all__ = [
+    "ddmsa_logp",
+    "ddmsa_potential",
+    "DDMSA",
+    "make_ddmsa_model",
+    "sample_ddmsa",
+    "sample_ddmsa_exact",
+    "simulate_ddmsa",
+]

saddm-0.1.0/saddm/core.py

@@ -0,0 +1,71 @@
+from numba import njit, float64
+from math import pi, sqrt, exp, log, sin, ceil, floor
+
+@njit(float64(float64, float64, float64), cache=True)
+def ftt_01w(tt, w, err):
+    """Compute f(t|0,1,w) following Navarro & Fuss, 2009."""
+    # calculate number of terms needed for large t
+    if pi * tt * err < 1:
+        kl = sqrt(-2 * log(pi * tt * err) / (pi * pi * tt))
+        kl = max(kl, 1. / (pi * sqrt(tt)))
+    else:
+        kl = 1. / (pi * sqrt(tt))
+
+    # calculate number of terms needed for small t
+    if 2 * sqrt(2 * pi * tt) * err < 1:
+        ks = 2 + sqrt(-2 * tt * log(2 * sqrt(2 * pi * tt) * err))
+        ks = max(ks, sqrt(tt) + 1)
+    else:
+        ks = 2
+
+    # compute f(tt|0,1,w)
+    p = 0.0
+    if ks < kl:  # small t approximation
+        K = int(ceil(ks))
+        lower = -int(floor((K-1)/2))
+        upper = int(ceil((K-1)/2))
+        for k in range(lower, upper + 1):
+            p += (w + 2 * k) * exp(-(pow(w + 2 * k, 2)) / 2 / tt)
+        p /= sqrt(2 * pi * pow(tt, 3))
+    else:  # large t approximation
+        K = int(ceil(kl))
+        for k in range(1, K + 1):
+            p += k * exp(-(k * k) * (pi * pi) * tt / 2) * sin(k * pi * w)
+        p *= pi
+
+    return p
+
+@njit(float64(float64, float64, float64, float64, float64, float64),
+      cache=True)
+def ddm_pdf_core(rt, a, z, v, ter, sv):
+    """Core DDM PDF calculation including drift rate variability.
+
+    Parameters use diffusion coefficient s=1:
+        rt: response time (seconds)
+        a:  boundary separation (~0.65-2.4)
+        z:  relative starting point (0 to 1, e.g. 0.5 for unbiased)
+        v:  drift rate (~0.1-0.5)
+        ter: non-decision time (seconds)
+        sv: drift rate variability (Gaussian SD)
+    """
+    if rt <= 0 or rt <= ter or a <= 0:
+        return 0.0
+
+    tt = (rt - ter) / (a * a)
+    err = 1e-4
+    p = ftt_01w(tt, z, err)
+
+    if sv <= 1e-10:
+        return p * exp(-v * a * z - (v * v * (rt - ter)) / 2.) / (a * a)
+    else:
+        sv_squared_rt = sv * sv * (rt - ter)
+
+        if sv_squared_rt > 100:
+            return 0.0
+
+        exp_term = ((a * z * sv)**2 - 2 * a * v * z - (v * v) * (rt - ter)) / (2 * (sv_squared_rt + 1))
+        scaling = 1 / sqrt(sv_squared_rt + 1)
+        if exp_term < -500 or exp_term > 500:
+            return 0.0
+
+        return p * exp(exp_term) * scaling / (a * a)