ECxEstimator for mixtures - ID selection
ValueError while using ECx calculator for substance mixtures in estimate()-function:
File c:...\Lib\site-packages\pymob\solvers\diffrax.py:129, in JaxSolver.solve(self, parameters, y0, x_in)
112 initialized_eval_func = partial(
113 self.odesolve_splitargs,
114 odestates = tuple(y0.keys()),
(...) 117 n_xin=len(x_in_flat)
118 )
120 loop_eval = jax.vmap(
121 initialized_eval_func,
122 in_axes=(
(...) 127 )
128 )
--> 129 result = loop_eval(*Y_0, *ode_args, *pp_args, *x_in_flat)
131 # if self.batch_dimension not in self.coordinates:
132 # this is not yet stable, because it may remove extra dimensions
133 # if there is a batch dimension of explicitly one specified
(...) 136 # this is added at the 0-axis
137 # if parameters are scalars, the returned shape is
138 for v, val in result.items():
[... skipping hidden 2 frame]File c:...\Lib\site-packages\jax_src\api.py:1340, in _mapped_axis_size(fn, tree, vals, dims, name) 1338 else: 1339 msg.append(f" * some axes ({ct} of them) had size {sz}, e.g. axis {ax} of {ex};\n") -> 1340 raise ValueError(''.join(msg)[:-2])
ValueError: vmap got inconsistent sizes for array axes to be mapped:
- most axes (7 of them) had size 1, e.g. axis 0 of argument args[0] of type float64[1,3];
- one axis had size 4: axis 0 of argument args[7] of type float64[4,9,3]
This error occurs because the selection of ID in ECxEstimator is not applied to the input exposure profile (x_in). The error can be solved by manually selecting the ID with x_in.sel(id='xyz'). The selection in the ECxEstimator could be adapted. ValueError message could be specified so that the error refers directly to the input data rather than first to the parameter set. Perhaps a distinction can be made in the error message between x_in and the parameter set.