HSSM regression -dealing with missing data go/nogo task

[EPinzuti]

Hi all,

It would useful for many applications with missing data in one condition like go/nogo task to be able to use HSSM like the old HDDM_regression.py

Now in the old package with the HDDMnn_regression.py I get:

Checking for missing data (LANs do not support this out of the box)

TODO: Add choice probability networks to allow missing values to be accounted for.

if data["rt"].abs().max() >= 998:
raise NotImplementedError(
"The HDDMnn classes can not yet deal with missing values. Support coming soon!"
)

Do you know if this will be possible soon ?
Thanks a lot
Edoardo

[digicosmos86]

Hi Edoardo,

We are currently working on adding go-nogo and deadline as the next major feature of HSSM. These should be available in a few weeks. Please stay tuned :).

Thanks!
Paul

[EPinzuti]

Perfect !!

thanks a lot

[EPinzuti]

Hi,

is there any update on this ? Thanks a lot

[digicosmos86]

Can you provide more details on what errors that you are getting?

From how your data looks like, you might need to specify condition as a categorical variable by using C(condition) in the formula. If that doesn't work, might be a good idea to code condition as (0, 1) so that all fields are numerical.

[EPinzuti]

I coded as numerical but I get the same error:
Auto-assigning NUTS sampler...
Initializing NUTS using adapt_diag...

TypeError Traceback (most recent call last)
File ~/hssm2/lib/python3.11/site-packages/pytensor/compile/function/types.py:970, in Function.call(self, *args, **kwargs)
968 try:
969 outputs = (
--> 970 self.vm()
971 if output_subset is None
972 else self.vm(output_subset=output_subset)
973 )
974 except Exception:

File ~/hssm2/lib/python3.11/site-packages/pytensor/graph/op.py:515, in Op.make_py_thunk..rval(p, i, o, n)
513 @is_thunk_type
514 def rval(p=p, i=node_input_storage, o=node_output_storage, n=node):
--> 515 r = p(n, [x[0] for x in i], o)
516 for o in node.outputs:

File ~/hssm2/lib/python3.11/site-packages/hssm/distribution_utils/onnx/onnx.py:195, in make_jax_logp_ops..LANLogpOp.perform(self, node, inputs, output_storage)
181 """Perform the Apply node.
182
183 Parameters
(...)
193 the Op.
194 """
--> 195 result = logp(*inputs)
196 output_storage[0][0] = np.asarray(result, dtype=node.outputs[0].dtype)

[... skipping hidden 14 frame]

File ~/hssm2/lib/python3.11/site-packages/hssm/distribution_utils/onnx/onnx.py:86, in make_jax_logp_funcs_from_onnx..logp(*inputs)
84 input_vector = jnp.concatenate((jnp.array(dist_params), data))
---> 86 return interpret_onnx(loaded_model.graph, input_vector)[0].squeeze()

File ~/hssm2/lib/python3.11/site-packages/hssm/distribution_utils/onnx/onnx2xla.py:151, in interpret_onnx(graph, *args)
150 attrs = {a.name: attribute_handlersa.type for a in node.attribute}
--> 151 outputs = onnx_ops[node.op_type](*args, **attrs)
152 for name, output in zip(node.output, outputs):

File ~/hssm2/lib/python3.11/site-packages/hssm/distribution_utils/onnx/onnx2xla.py:110, in onnx_gemm(a, b, c, alpha, beta, transA, transB)
108 b = jnp.transpose(b) if transB else b
--> 110 return [alpha * jnp.matmul(a, b) + beta * c]

[... skipping hidden 11 frame]

File ~/hssm2/lib/python3.11/site-packages/jax/_src/numpy/lax_numpy.py:4038, in matmul(a, b, precision, preferred_element_type)
4037 b = lax.squeeze(b, tuple(b_squeeze))
-> 4038 out = lax.dot_general(
4039 a, b, (((ndim(a) - 1,), (ndim(b) - 1 - b_is_mat,)), (a_batch, b_batch)),
4040 precision=precision, preferred_element_type=preferred_element_type)
4041 result = lax.transpose(out, perm)

[... skipping hidden 7 frame]

File ~/hssm2/lib/python3.11/site-packages/jax/_src/lax/lax.py:2665, in _dot_general_shape_rule(lhs, rhs, dimension_numbers, precision, preferred_element_type)
2663 msg = ("dot_general requires contracting dimensions to have the same "
2664 "shape, got {} and {}.")
-> 2665 raise TypeError(msg.format(lhs_contracting_shape, rhs_contracting_shape))
2667 return _dot_general_shape_computation(lhs.shape, rhs.shape, dimension_numbers)

TypeError: dot_general requires contracting dimensions to have the same shape, got (4,) and (5,).

[digicosmos86]

Your data only has missing values but not deadlines right? The correct network to download is ddm_cpn from here. Just modify your code to point to this network instead and that should work. the opn assumes that you also have a deadline column in your data, thus the incompatible dimensions.

Also feel free to try sample="nuts_numpyro" in model.sample and see if there's any improvement in sampling performance

[EPinzuti]

Thanks !! It runs now. I will start working with it and let you know. If the official release will come in the next two/three months maybe I can use Hssm for the paper. Let s see how far I go with it.

[digicosmos86]

Great! Let me know if you have any other questions. We'll work to bring the final versions of the networks available asap