diff options
| -rw-r--r-- | polymatrix/expression/mixins/parametrizeexprmixin.py | 9 |
1 files changed, 4 insertions, 5 deletions
diff --git a/polymatrix/expression/mixins/parametrizeexprmixin.py b/polymatrix/expression/mixins/parametrizeexprmixin.py index b9673e4..b8054ef 100644 --- a/polymatrix/expression/mixins/parametrizeexprmixin.py +++ b/polymatrix/expression/mixins/parametrizeexprmixin.py @@ -31,14 +31,13 @@ class ParametrizeExprMixin(ExpressionBaseMixin): state, underlying = self.underlying.apply(state) - assert underlying.shape[1] == 1 - start_index = state.n_param terms = {} - for row in range(underlying.shape[0]): - var_index = start_index + row - terms[row, 0] = {((var_index, 1),): 1.0} + for col in range(underlying.shape[1]): + for row in range(underlying.shape[0]): + var_index = start_index + col * underlying.shape[1] + row + terms[row, col] = {((var_index, 1),): 1.0} state = state.register( key=self, |