diff --git a/src/hssm/_types.py b/src/hssm/_types.py index 82711c38..e02c07bf 100644 --- a/src/hssm/_types.py +++ b/src/hssm/_types.py @@ -25,6 +25,8 @@ "lba2", "racing_diffusion_3", "poisson_race", + "softmax_inv_temperature_2", + "softmax_inv_temperature_3", ] diff --git a/src/hssm/likelihoods/analytical.py b/src/hssm/likelihoods/analytical.py index 01659278..f4010303 100644 --- a/src/hssm/likelihoods/analytical.py +++ b/src/hssm/likelihoods/analytical.py @@ -5,8 +5,6 @@ https://gist.github.com/sammosummo/c1be633a74937efaca5215da776f194b. """ -from typing import Type - import jax.numpy as jnp import numpy as np import pymc as pm @@ -399,14 +397,14 @@ def logp_ddm_sdv( ddm_params = ["v", "a", "z", "t"] ddm_sdv_params = ddm_params + ["sv"] -DDM: Type[pm.Distribution] = make_distribution( +DDM: type[pm.Distribution] = make_distribution( rv="ddm", loglik=logp_ddm, list_params=ddm_params, bounds=ddm_bounds, ) -DDM_SDV: Type[pm.Distribution] = make_distribution( +DDM_SDV: type[pm.Distribution] = make_distribution( rv="ddm_sdv", loglik=logp_ddm_sdv, list_params=ddm_sdv_params, @@ -541,7 +539,7 @@ def _print_message(_): "t": (0.0, inf), } -RDM3: Type[pm.Distribution] = make_distribution( +RDM3: type[pm.Distribution] = make_distribution( rv="racing_diffusion_3", loglik=logp_rdm3, list_params=rdm3_params, @@ -685,14 +683,14 @@ def logp_lba3( "v2": (0.0, inf), } -LBA2: Type[pm.Distribution] = make_distribution( +LBA2: type[pm.Distribution] = make_distribution( rv="lba2", loglik=logp_lba2, list_params=lba2_params, bounds=lba2_bounds, ) -LBA3: Type[pm.Distribution] = make_distribution( +LBA3: type[pm.Distribution] = make_distribution( rv="lba3", loglik=logp_lba3, list_params=lba3_params, @@ -806,3 +804,39 @@ def logp_poisson_race( list_params=poisson_race_params, bounds=poisson_race_bounds, ) + + +def softmax_inv_temperature(data: np.ndarray, beta: np.ndarray, *logits): + """Compute the log-likelihood of the Inverse Softmax Temperature Model. + + Parameters + ---------- + data + 1D array of responses (choices). + beta + A scalar for the softmax temperature (0, inf). + *logits + Logits for each choice excluding logit0. + + Returns + ------- + pt.TensorVariable + The log likelihood of the Inverse Softmax Temperature Model. + """ + choices = pt.where(data < 1, 0.0, data).astype("int32") + zeros = pt.zeros_like(data, dtype=pytensor.config.floatX) + + logits_stacked = pt.stack( + [ + zeros, # logit0 is always 0 + *(logit + zeros for logit in logits), + ], + ) + + logits_scaled = logits_stacked * beta + choice_logits = logits_scaled[choices, pt.arange(data.shape[0])] + log_prob_choices = choice_logits - pt.logsumexp( + logits_scaled, axis=0, keepdims=False + ) + + return log_prob_choices diff --git a/src/hssm/modelconfig/_softmax_inv_temperature_config.py b/src/hssm/modelconfig/_softmax_inv_temperature_config.py new file mode 100644 index 00000000..0e118b0a --- /dev/null +++ b/src/hssm/modelconfig/_softmax_inv_temperature_config.py @@ -0,0 +1,59 @@ +"""The default configuration for the Inverse Softmax Temperature Model with 2 logits.""" + +import numpy as np + +from .._types import DefaultConfig, ParamSpec +from ..likelihoods.analytical import softmax_inv_temperature + + +def softmax_inv_temperature_config(n_choices: int = 2) -> DefaultConfig: + """ + Get the default config for the Softmax Inv. Temperature Model. + + Parameters + ---------- + n_choices : optional + The number of choices in the model. Must be at least 2. The number of logits + will be n_choices - 1. Default is 2 (i.e., 1 logit). + + Returns + ------- + DefaultConfig + A dictionary containing the default configuration settings for the + Inverse Softmax Temperature Model, including response variables, model + parameters, choices, description, and likelihood specifications. + """ + if n_choices < 2: + raise ValueError("n_choices must be at least 2.") + + bounds = {"beta": (0.0, np.inf)} + bounds.update({f"logit{i}": (-np.inf, np.inf) for i in range(1, n_choices)}) + default_priors: dict[str, ParamSpec] = { + "beta": { + "name": "Gamma", + "alpha": 2.0, + "beta": 0.5, + }, + } + default_priors.update( + { + f"logit{i}": {"name": "Normal", "mu": 0.0, "sigma": 1.0} + for i in range(1, n_choices) + } + ) + + return { + "response": ["response"], + "list_params": ["beta"] + [f"logit{i}" for i in range(1, n_choices)], + "choices": [-1, 1] if n_choices == 2 else list(range(n_choices)), + "description": f"The Softmax Inv. Temperature Model with {n_choices} choices", + "likelihoods": { + "analytical": { + "loglik": softmax_inv_temperature, + "backend": None, + "bounds": bounds, + "default_priors": default_priors, + "extra_fields": None, + }, + }, + } diff --git a/src/hssm/modelconfig/softmax_inv_temperature_2_config.py b/src/hssm/modelconfig/softmax_inv_temperature_2_config.py new file mode 100644 index 00000000..83c4f1ea --- /dev/null +++ b/src/hssm/modelconfig/softmax_inv_temperature_2_config.py @@ -0,0 +1,18 @@ +"""Softmax Inverse Temperature Model with 2 logits configuration.""" + +from .._types import DefaultConfig +from ._softmax_inv_temperature_config import softmax_inv_temperature_config + + +def get_softmax_inv_temperature_2_config() -> DefaultConfig: + """ + Get the default config for the Inverse Softmax Temperature Model with 2 logits. + + Returns + ------- + DefaultConfig + A dictionary containing the default configuration settings for the + Inverse Softmax Temperature Model with 2 logits, including response variables, + model parameters, choices, description, and likelihood specifications. + """ + return softmax_inv_temperature_config(n_choices=2) diff --git a/src/hssm/modelconfig/softmax_inv_temperature_3_config.py b/src/hssm/modelconfig/softmax_inv_temperature_3_config.py new file mode 100644 index 00000000..e4a6d7c3 --- /dev/null +++ b/src/hssm/modelconfig/softmax_inv_temperature_3_config.py @@ -0,0 +1,18 @@ +"""Softmax Inverse Temperature Model with 3 logits configuration.""" + +from .._types import DefaultConfig +from ._softmax_inv_temperature_config import softmax_inv_temperature_config + + +def get_softmax_inv_temperature_3_config() -> DefaultConfig: + """ + Get the default config for the Inverse Softmax Temperature Model with 3 logits. + + Returns + ------- + DefaultConfig + A dictionary containing the default configuration settings for the + Inverse Softmax Temperature Model with 3 logits, including response variables, + model parameters, choices, description, and likelihood specifications. + """ + return softmax_inv_temperature_config(n_choices=3) diff --git a/tests/test_likelihoods_choice_only.py b/tests/test_likelihoods_choice_only.py new file mode 100644 index 00000000..f3e8f061 --- /dev/null +++ b/tests/test_likelihoods_choice_only.py @@ -0,0 +1,45 @@ +import pytest + +import numpy as np + +from hssm.likelihoods.analytical import softmax_inv_temperature + + +_N = 10 +_rng = np.random.default_rng(42) + +_DATA_BINARY = _rng.choice([-1, 1], size=_N).astype(np.float32) +_DATA_TERNARY = _rng.choice([0, 1, 2], size=_N).astype(np.float32) + +_SCALAR_BETA = np.float32(1.5) +_VECTOR_BETA = np.full(_N, 1.5, dtype=np.float32) + +_SCALAR_LOGIT = np.float32(0.5) +_VECTOR_LOGIT = np.full(_N, 0.5, dtype=np.float32) + + +@pytest.mark.parametrize( + "beta", [_SCALAR_BETA, _VECTOR_BETA], ids=["scalar_beta", "vector_beta"] +) +@pytest.mark.parametrize( + "logit", [_SCALAR_LOGIT, _VECTOR_LOGIT], ids=["scalar_logit", "vector_logit"] +) +def test_softmax_inv_temperature_shape_2choice(beta, logit): + result = softmax_inv_temperature(_DATA_BINARY, beta, logit) + evaluated = result.eval() + assert evaluated.shape == (_N,) + + +@pytest.mark.parametrize( + "beta", [_SCALAR_BETA, _VECTOR_BETA], ids=["scalar_beta", "vector_beta"] +) +@pytest.mark.parametrize( + "logit1", [_SCALAR_LOGIT, _VECTOR_LOGIT], ids=["scalar_logit1", "vector_logit1"] +) +@pytest.mark.parametrize( + "logit2", [_SCALAR_LOGIT, _VECTOR_LOGIT], ids=["scalar_logit2", "vector_logit2"] +) +def test_softmax_inv_temperature_shape_3choice(beta, logit1, logit2): + result = softmax_inv_temperature(_DATA_TERNARY, beta, logit1, logit2) + evaluated = result.eval() + assert evaluated.shape == (_N,) diff --git a/tests/test_modelconfig.py b/tests/test_modelconfig.py index f4ed1f7d..520f843b 100644 --- a/tests/test_modelconfig.py +++ b/tests/test_modelconfig.py @@ -1,6 +1,10 @@ +import numpy as np import pytest from hssm.modelconfig import get_default_model_config +from hssm.modelconfig._softmax_inv_temperature_config import ( + softmax_inv_temperature_config, +) import hssm @@ -108,3 +112,70 @@ def test_load_all_supported_model_configs(model): def test_get_default_model_config_invalid(): with pytest.raises(ValueError): get_default_model_config("invalid_model") + + +def test_softmax_inv_temperature_default(): + """Test softmax_inv_temperature with default n_choices=2.""" + config = softmax_inv_temperature_config() + + assert config["response"] == ["response"] + assert config["choices"] == [-1, 1] + assert config["list_params"] == ["beta", "logit1"] + assert config["description"] == "The Softmax Inv. Temperature Model with 2 choices" + + likelihoods = config["likelihoods"] + lk_analytical = likelihoods["analytical"] + + assert lk_analytical["backend"] is None + assert lk_analytical["extra_fields"] is None + + # Test bounds + assert lk_analytical["bounds"]["beta"] == (0.0, np.inf) + assert lk_analytical["bounds"]["logit1"] == (-np.inf, np.inf) + + # Test default priors + assert lk_analytical["default_priors"]["beta"] == { + "name": "Gamma", + "alpha": 2.0, + "beta": 0.5, + } + assert lk_analytical["default_priors"]["logit1"] == { + "name": "Normal", + "mu": 0.0, + "sigma": 1.0, + } + + +def test_softmax_inv_temperature_3_choices(): + """Test softmax_inv_temperature with n_choices=3.""" + config = softmax_inv_temperature_config(n_choices=3) + + assert config["response"] == ["response"] + assert config["choices"] == [0, 1, 2] + assert config["list_params"] == ["beta", "logit1", "logit2"] + assert config["description"] == "The Softmax Inv. Temperature Model with 3 choices" + + likelihoods = config["likelihoods"] + lk_analytical = likelihoods["analytical"] + + # Test bounds + assert lk_analytical["bounds"]["beta"] == (0.0, np.inf) + assert lk_analytical["bounds"]["logit1"] == (-np.inf, np.inf) + assert lk_analytical["bounds"]["logit2"] == (-np.inf, np.inf) + + # Test default priors + assert lk_analytical["default_priors"]["beta"] == { + "name": "Gamma", + "alpha": 2.0, + "beta": 0.5, + } + assert lk_analytical["default_priors"]["logit1"] == { + "name": "Normal", + "mu": 0.0, + "sigma": 1.0, + } + assert lk_analytical["default_priors"]["logit2"] == { + "name": "Normal", + "mu": 0.0, + "sigma": 1.0, + }