From 0bac19cf34887623449b16dc7e8b92df36984398 Mon Sep 17 00:00:00 2001 From: Fengler Date: Wed, 1 May 2024 22:04:38 +0200 Subject: [PATCH 01/11] change order between sz and sv in list_params --- src/hssm/defaults.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/hssm/defaults.py b/src/hssm/defaults.py index 40ed155f..7561c4ab 100644 --- a/src/hssm/defaults.py +++ b/src/hssm/defaults.py @@ -184,7 +184,7 @@ class DefaultConfig(TypedDict): }, "full_ddm": { "response": ["rt", "response"], - "list_params": ["v", "a", "z", "t", "sv", "sz", "st"], + "list_params": ["v", "a", "z", "t", "sz", "sv", "st"], "description": "The full Drift Diffusion Model (DDM)", "likelihoods": { "blackbox": { From 3404690da3a4c4b7e890c0ecb6b923355b2098a1 Mon Sep 17 00:00:00 2001 From: Fengler Date: Sun, 5 May 2024 21:38:20 +0200 Subject: [PATCH 02/11] get rid of some commented code --- src/hssm/prior.py | 8 -------- 1 file changed, 8 deletions(-) diff --git a/src/hssm/prior.py b/src/hssm/prior.py index 70d3bfab..fcab0d80 100644 --- a/src/hssm/prior.py +++ b/src/hssm/prior.py @@ -351,11 +351,3 @@ def get_hddm_default_prior( "sz": {"dist": "Gamma", "mu": HDDM_MU["sz"], "sigma": HDDM_SIGMA["sz"]}, "st": {"dist": "Gamma", "mu": HDDM_MU["st"], "sigma": HDDM_SIGMA["st"]}, } - -# INITVAL_SETTINGS_LOGIT: dict[Any, Any] = { -# "t" : {"initval": -4.0}, -# } - -# INITVAL_SETTINGS_NOLINK: dict[Any, Any] = { -# "t" : {"initval": 0.05}, -# } From 95922535f5f2173ccab00e418ffc8aba9b0b14e0 Mon Sep 17 00:00:00 2001 From: Fengler Date: Sun, 5 May 2024 22:54:00 +0200 Subject: [PATCH 03/11] add v to posterior directly --- .github/workflows/run_tests.yml | 2 +- src/hssm/hssm.py | 9 +++++++-- 2 files changed, 8 insertions(+), 3 deletions(-) diff --git a/.github/workflows/run_tests.yml b/.github/workflows/run_tests.yml index a143656e..3061862d 100644 --- a/.github/workflows/run_tests.yml +++ b/.github/workflows/run_tests.yml @@ -30,7 +30,7 @@ jobs: - name: install pymc and poetry run: | mamba info - mamba install -c conda-forge pymc=5.12 poetry + mamba install -c conda-forge pymc=5.14 poetry - name: install hssm run: | diff --git a/src/hssm/hssm.py b/src/hssm/hssm.py index 4a0fff17..0226f92f 100644 --- a/src/hssm/hssm.py +++ b/src/hssm/hssm.py @@ -497,6 +497,11 @@ def sample( inference_method=sampler, init=init, **kwargs ) + # 'v' was previously not part of deterministics --> compute it via posterior_predictive + # (works because it acts as the 'mu' parameter in the GLM as far as bambi is concerned) + self.sample_posterior_predictive(self._inference_obj, kind="mean") + # rename 'rt,response_mean' to 'v' so in the traces everything looks the way it should + self._inference_obj.rename_vars({"rt,response_mean": "v"}, inplace=True) return self.traces def sample_posterior_predictive( @@ -526,13 +531,13 @@ def sample_posterior_predictive( If `True` will make predictions including the group specific effects. Otherwise, predictions are made with common effects only (i.e. group- specific are set to zero), by default True. - kind + kind: optional Indicates the type of prediction required. Can be `"mean"` or `"pps"`. The first returns draws from the posterior distribution of the mean, while the latter returns the draws from the posterior predictive distribution (i.e. the posterior probability distribution for a new observation). Defaults to `"pps"`. - n_samples + n_samples: optional The number of samples to draw from the posterior predictive distribution from each chain. When it's an integer >= 1, the number of samples to be extracted from the From e3a4fc0be1dcbc7e9f9d6d24965ab2fd5656e371 Mon Sep 17 00:00:00 2001 From: Fengler Date: Sun, 5 May 2024 23:00:59 +0200 Subject: [PATCH 04/11] attempt at avoiding mypy error --- src/hssm/hssm.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/src/hssm/hssm.py b/src/hssm/hssm.py index 0226f92f..441828db 100644 --- a/src/hssm/hssm.py +++ b/src/hssm/hssm.py @@ -501,7 +501,8 @@ def sample( # (works because it acts as the 'mu' parameter in the GLM as far as bambi is concerned) self.sample_posterior_predictive(self._inference_obj, kind="mean") # rename 'rt,response_mean' to 'v' so in the traces everything looks the way it should - self._inference_obj.rename_vars({"rt,response_mean": "v"}, inplace=True) + if self._inference_obj is not None: + self._inference_obj.rename_vars({"rt,response_mean": "v"}, inplace=True) return self.traces def sample_posterior_predictive( From f14e8ed945a4764334bd19875cdee8e0c66a9500 Mon Sep 17 00:00:00 2001 From: Fengler Date: Sun, 5 May 2024 23:06:10 +0200 Subject: [PATCH 05/11] ruff --- src/hssm/hssm.py | 9 ++++++--- 1 file changed, 6 insertions(+), 3 deletions(-) diff --git a/src/hssm/hssm.py b/src/hssm/hssm.py index 441828db..d6dd7f71 100644 --- a/src/hssm/hssm.py +++ b/src/hssm/hssm.py @@ -497,10 +497,13 @@ def sample( inference_method=sampler, init=init, **kwargs ) - # 'v' was previously not part of deterministics --> compute it via posterior_predictive - # (works because it acts as the 'mu' parameter in the GLM as far as bambi is concerned) + # 'v' was previously not part of deterministics --> compute it via + # posterior_predictive (works because it acts as the 'mu' parameter + # in the GLM as far as bambi is concerned) self.sample_posterior_predictive(self._inference_obj, kind="mean") - # rename 'rt,response_mean' to 'v' so in the traces everything looks the way it should + + # rename 'rt,response_mean' to 'v' so in the traces everything + # looks the way it should if self._inference_obj is not None: self._inference_obj.rename_vars({"rt,response_mean": "v"}, inplace=True) return self.traces From 64e356e2e60820e6573dfee1e71b401f11d00e90 Mon Sep 17 00:00:00 2001 From: Fengler Date: Mon, 6 May 2024 16:34:25 +0200 Subject: [PATCH 06/11] fix sample function to not create naming conflict for v --- src/hssm/hssm.py | 9 ++++----- tests/test_hssm.py | 2 +- 2 files changed, 5 insertions(+), 6 deletions(-) diff --git a/src/hssm/hssm.py b/src/hssm/hssm.py index d6dd7f71..7ed548d4 100644 --- a/src/hssm/hssm.py +++ b/src/hssm/hssm.py @@ -500,11 +500,10 @@ def sample( # 'v' was previously not part of deterministics --> compute it via # posterior_predictive (works because it acts as the 'mu' parameter # in the GLM as far as bambi is concerned) - self.sample_posterior_predictive(self._inference_obj, kind="mean") - - # rename 'rt,response_mean' to 'v' so in the traces everything - # looks the way it should - if self._inference_obj is not None: + if "v" not in self._inference_obj.posterior.data_vars.keys(): + self.sample_posterior_predictive(self._inference_obj, kind="mean") + # rename 'rt,response_mean' to 'v' so in the traces everything + # looks the way it should self._inference_obj.rename_vars({"rt,response_mean": "v"}, inplace=True) return self.traces diff --git a/tests/test_hssm.py b/tests/test_hssm.py index 0a660303..76c0e3d7 100644 --- a/tests/test_hssm.py +++ b/tests/test_hssm.py @@ -7,7 +7,7 @@ from hssm.utils import download_hf from hssm.likelihoods import DDM, logp_ddm -hssm.set_floatX("float32") +hssm.set_floatX("float32", jax=True) param_v = { "name": "v", From 21ff60c00e9c6a0a720f44801aad8b3da2b12089 Mon Sep 17 00:00:00 2001 From: Fengler Date: Tue, 7 May 2024 21:00:34 +0200 Subject: [PATCH 07/11] attempt at fixing assertion errors from tests --- src/hssm/hssm.py | 20 ++++++++++++++++---- src/hssm/likelihoods/blackbox.py | 2 +- tests/slow/test_mcmc.py | 18 ++++++++++++++++-- 3 files changed, 33 insertions(+), 7 deletions(-) diff --git a/src/hssm/hssm.py b/src/hssm/hssm.py index 7ed548d4..2cd81fdc 100644 --- a/src/hssm/hssm.py +++ b/src/hssm/hssm.py @@ -497,14 +497,16 @@ def sample( inference_method=sampler, init=init, **kwargs ) - # 'v' was previously not part of deterministics --> compute it via + # The parent was previously not part of deterministics --> compute it via # posterior_predictive (works because it acts as the 'mu' parameter # in the GLM as far as bambi is concerned) - if "v" not in self._inference_obj.posterior.data_vars.keys(): + if self._parent not in self._inference_obj.posterior.data_vars.keys(): self.sample_posterior_predictive(self._inference_obj, kind="mean") # rename 'rt,response_mean' to 'v' so in the traces everything # looks the way it should - self._inference_obj.rename_vars({"rt,response_mean": "v"}, inplace=True) + self._inference_obj.rename_vars( + {"rt,response_mean": self._parent}, inplace=True + ) return self.traces def sample_posterior_predictive( @@ -1316,11 +1318,21 @@ def _get_deterministic_var_names(self, idata) -> list[str]: var_names = [ f"~{param_name}" for param_name, param in self.params.items() - if param.is_regression and not param.is_parent + if param.is_regression ] + # Handle specific case where parent is not explictly in traces + if ("~" + self._parent in var_names) and ( + self._parent not in idata.posterior.data_vars + ): + var_names.remove("~" + self._parent) + + # ] and not param.is_parent --> this can be skipped now, because `v` will be added to traces by default + # ] + if f"{self.response_str}_mean" in idata["posterior"].data_vars: var_names.append(f"~{self.response_str}_mean") + return var_names def _handle_missing_data_and_deadline(self): diff --git a/src/hssm/likelihoods/blackbox.py b/src/hssm/likelihoods/blackbox.py index c094c056..4aa91ca4 100644 --- a/src/hssm/likelihoods/blackbox.py +++ b/src/hssm/likelihoods/blackbox.py @@ -61,7 +61,7 @@ def logp_ddm_sdv_bbox(data: np.ndarray, v, a, z, t, sv) -> np.ndarray: @hddm_to_hssm -def logp_full_ddm(data: np.ndarray, v, a, z, t, sv, sz, st): +def logp_full_ddm(data: np.ndarray, v, a, z, t, sz, sv, st): """Compute blackbox log-likelihoods for full_ddm models.""" return wfpt.wiener_logp_array( x=data, diff --git a/tests/slow/test_mcmc.py b/tests/slow/test_mcmc.py index f0f8728a..1d99b291 100644 --- a/tests/slow/test_mcmc.py +++ b/tests/slow/test_mcmc.py @@ -124,6 +124,8 @@ def run_sample(model, sampler, step, expected): @pytest.mark.parametrize(parameter_names, parameter_grid) def test_simple_models(data_ddm, loglik_kind, backend, sampler, step, expected): + print("PYMC VERSION: ") + print(pm.__version__) print("TEST INPUTS WERE: ") print("REPORTING FROM SIMPLE MODELS TEST") print(loglik_kind, backend, sampler, step, expected) @@ -147,6 +149,8 @@ def test_simple_models(data_ddm, loglik_kind, backend, sampler, step, expected): @pytest.mark.parametrize(parameter_names, parameter_grid) def test_reg_models(data_ddm_reg, loglik_kind, backend, sampler, step, expected): + print("PYMC VERSION: ") + print(pm.__version__) print("TEST INPUTS WERE: ") print("REPORTING FROM REG MODELS TEST") print(loglik_kind, backend, sampler, step, expected) @@ -169,7 +173,7 @@ def test_reg_models(data_ddm_reg, loglik_kind, backend, sampler, step, expected) # Only runs once if loglik_kind == "analytical" and sampler is None: - assert not model._get_deterministic_var_names(model.traces) + assert model._get_deterministic_var_names(model.traces) == ["~v"] # test summary: summary = model.summary() assert summary.shape[0] == 6 @@ -181,6 +185,8 @@ def test_reg_models(data_ddm_reg, loglik_kind, backend, sampler, step, expected) @pytest.mark.parametrize(parameter_names, parameter_grid) def test_reg_models_v_a(data_ddm_reg, loglik_kind, backend, sampler, step, expected): + print("PYMC VERSION: ") + print(pm.__version__) print("TEST INPUTS WERE: ") print("REPORTING FROM REG MODELS V_A TEST") print(loglik_kind, backend, sampler, step, expected) @@ -218,7 +224,7 @@ def test_reg_models_v_a(data_ddm_reg, loglik_kind, backend, sampler, step, expec # Only runs once if loglik_kind == "analytical" and sampler is None: - assert model._get_deterministic_var_names(model.traces) == ["~a"] + assert model._get_deterministic_var_names(model.traces) == ["~a", "~v"] # test summary: summary = model.summary() assert summary.shape[0] == 8 @@ -253,6 +259,8 @@ def test_reg_models_v_a(data_ddm_reg, loglik_kind, backend, sampler, step, expec def test_simple_models_missing_data( data_ddm_missing, loglik_kind, backend, sampler, step, expected, cpn ): + print("PYMC VERSION: ") + print(pm.__version__) print("TEST INPUTS WERE: ") print("REPORTING FROM SIMPLE MODELS MISSING DATA TEST") print(loglik_kind, backend, sampler, step, expected) @@ -271,6 +279,8 @@ def test_simple_models_missing_data( def test_reg_models_missing_data( data_ddm_reg_missing, loglik_kind, backend, sampler, step, expected, cpn ): + print("PYMC VERSION: ") + print(pm.__version__) print("TEST INPUTS WERE: ") print("REPORTING FROM REG MODELS MISSING DATA TEST") print(loglik_kind, backend, sampler, step, expected) @@ -298,6 +308,8 @@ def test_reg_models_missing_data( def test_simple_models_deadline( data_ddm_deadline, loglik_kind, backend, sampler, step, expected, opn ): + print("PYMC VERSION: ") + print(pm.__version__) print("TEST INPUTS WERE: ") print("REPORTING FROM SIMPLE MODELS DEADLINE TEST") print(loglik_kind, backend, sampler, step, expected) @@ -315,6 +327,8 @@ def test_simple_models_deadline( def test_reg_models_deadline( data_ddm_reg_deadline, loglik_kind, backend, sampler, step, expected, opn ): + print("PYMC VERSION: ") + print(pm.__version__) print("TEST INPUTS WERE: ") print("REPORTING FROM REG MODELS DEADLINE TEST") print(loglik_kind, backend, sampler, step, expected) From 6342c72b1070879b400db590418d32840843173c Mon Sep 17 00:00:00 2001 From: Fengler Date: Tue, 7 May 2024 21:39:28 +0200 Subject: [PATCH 08/11] make mypy happy --- src/hssm/hssm.py | 15 ++++++++------- tests/slow/test_mcmc.py | 7 ++++++- 2 files changed, 14 insertions(+), 8 deletions(-) diff --git a/src/hssm/hssm.py b/src/hssm/hssm.py index 2cd81fdc..2f2b96a3 100644 --- a/src/hssm/hssm.py +++ b/src/hssm/hssm.py @@ -500,13 +500,14 @@ def sample( # The parent was previously not part of deterministics --> compute it via # posterior_predictive (works because it acts as the 'mu' parameter # in the GLM as far as bambi is concerned) - if self._parent not in self._inference_obj.posterior.data_vars.keys(): - self.sample_posterior_predictive(self._inference_obj, kind="mean") - # rename 'rt,response_mean' to 'v' so in the traces everything - # looks the way it should - self._inference_obj.rename_vars( - {"rt,response_mean": self._parent}, inplace=True - ) + if self._inference_obj is not None: + if self._parent not in self._inference_obj.posterior.data_vars.keys(): + self.sample_posterior_predictive(self._inference_obj, kind="mean") + # rename 'rt,response_mean' to 'v' so in the traces everything + # looks the way it should + self._inference_obj.rename_vars( + {"rt,response_mean": self._parent}, inplace=True + ) return self.traces def sample_posterior_predictive( diff --git a/tests/slow/test_mcmc.py b/tests/slow/test_mcmc.py index 1d99b291..96f86a0d 100644 --- a/tests/slow/test_mcmc.py +++ b/tests/slow/test_mcmc.py @@ -224,7 +224,12 @@ def test_reg_models_v_a(data_ddm_reg, loglik_kind, backend, sampler, step, expec # Only runs once if loglik_kind == "analytical" and sampler is None: - assert model._get_deterministic_var_names(model.traces) == ["~a", "~v"] + assert len(model._get_deterministic_var_names(model.traces)) == len( + ["~a", "~v"] + ) + assert set(model._get_deterministic_var_names(model.traces)) == set( + ["~a", "~v"] + ) # test summary: summary = model.summary() assert summary.shape[0] == 8 From e586ec035f28f9e2252b2ad1aeacc065312484a7 Mon Sep 17 00:00:00 2001 From: Fengler Date: Tue, 7 May 2024 21:56:38 +0200 Subject: [PATCH 09/11] make ruff happy --- src/hssm/hssm.py | 5 +---- 1 file changed, 1 insertion(+), 4 deletions(-) diff --git a/src/hssm/hssm.py b/src/hssm/hssm.py index 2f2b96a3..6eb14ce4 100644 --- a/src/hssm/hssm.py +++ b/src/hssm/hssm.py @@ -60,7 +60,7 @@ class HSSM: - """The Hierarchical Sequential Sampling Model (HSSM) class. + """The basic Hierarchical Sequential Sampling Model (HSSM) class. Parameters ---------- @@ -1328,9 +1328,6 @@ def _get_deterministic_var_names(self, idata) -> list[str]: ): var_names.remove("~" + self._parent) - # ] and not param.is_parent --> this can be skipped now, because `v` will be added to traces by default - # ] - if f"{self.response_str}_mean" in idata["posterior"].data_vars: var_names.append(f"~{self.response_str}_mean") From 3391ae484723322f7b2c026af7e8d2e660e625ab Mon Sep 17 00:00:00 2001 From: Fengler Date: Tue, 7 May 2024 22:28:19 +0200 Subject: [PATCH 10/11] explictly include jax=True --- tests/slow/test_mcmc.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/slow/test_mcmc.py b/tests/slow/test_mcmc.py index 96f86a0d..83b24347 100644 --- a/tests/slow/test_mcmc.py +++ b/tests/slow/test_mcmc.py @@ -10,7 +10,7 @@ from hssm.utils import _rearrange_data -hssm.set_floatX("float32") +hssm.set_floatX("float32", jax=True) # AF-TODO: Include more tests that use different link functions! From 613490526b802da261af34809a104990aeb353f0 Mon Sep 17 00:00:00 2001 From: Fengler Date: Wed, 8 May 2024 20:10:00 +0200 Subject: [PATCH 11/11] work on suggestions --- .github/workflows/run_tests.yml | 2 +- src/hssm/hssm.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/.github/workflows/run_tests.yml b/.github/workflows/run_tests.yml index 3061862d..5bd038f0 100644 --- a/.github/workflows/run_tests.yml +++ b/.github/workflows/run_tests.yml @@ -30,7 +30,7 @@ jobs: - name: install pymc and poetry run: | mamba info - mamba install -c conda-forge pymc=5.14 poetry + mamba install -c conda-forge pymc poetry - name: install hssm run: | diff --git a/src/hssm/hssm.py b/src/hssm/hssm.py index 6eb14ce4..97c520cf 100644 --- a/src/hssm/hssm.py +++ b/src/hssm/hssm.py @@ -502,7 +502,7 @@ def sample( # in the GLM as far as bambi is concerned) if self._inference_obj is not None: if self._parent not in self._inference_obj.posterior.data_vars.keys(): - self.sample_posterior_predictive(self._inference_obj, kind="mean") + self.model.predict(self._inference_obj, kind="mean", inplace=True) # rename 'rt,response_mean' to 'v' so in the traces everything # looks the way it should self._inference_obj.rename_vars(