parametrize operator acts only on column vector

1 files changed, 5 insertions, 5 deletions

diff --git a/polymatrix/expression/mixins/parametrizeexprmixin.py b/polymatrix/expression/mixins/parametrizeexprmixin.py
index b8054ef..c33ca5b 100644
--- a/polymatrix/expression/mixins/parametrizeexprmixin.py
+++ b/polymatrix/expression/mixins/parametrizeexprmixin.py
@@ -31,13 +31,13 @@ class ParametrizeExprMixin(ExpressionBaseMixin):
 
         state, underlying = self.underlying.apply(state)
 
-        start_index = state.n_param
+        assert underlying.shape[1] == 1
+
         terms = {}
 
-        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}
+        for row in range(underlying.shape[0]):
+            var_index = state.n_param + row
+            terms[row, 0] = {((var_index, 1),): 1.0}
 
         state = state.register(
             key=self,