From 8341b1fae4454c1305a615b8069fcda42042a0cb Mon Sep 17 00:00:00 2001 From: Michael Schneeberger Date: Wed, 17 Aug 2022 17:27:48 +0200 Subject: 'parametrize' method accepts any size matrix, not only vectors --- polymatrix/expression/mixins/parametrizeexprmixin.py | 9 ++++----- 1 file 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, -- cgit v1.2.1