Improve conversion to new index types in init_polymatrix

This is here for backwards compatibility

1 files changed, 22 insertions, 4 deletions

diff --git a/polymatrix/polymatrix/init.py b/polymatrix/polymatrix/init.py
index 07ceb4c..9a85a0d 100644
--- a/polymatrix/polymatrix/init.py
+++ b/polymatrix/polymatrix/init.py
@@ -1,6 +1,15 @@
-from polymatrix.polymatrix.mixins import BroadcastPolyMatrixMixin
+from __future__ import annotations
+
+import math
+import numpy as np
+
+from typing import TYPE_CHECKING
+
 from polymatrix.polymatrix.impl import BroadcastPolyMatrixImpl, PolyMatrixImpl
-from polymatrix.polymatrix.typing import PolyMatrixDict, PolyDict
+from polymatrix.polymatrix.typing import PolyMatrixDict, PolyDict, MatrixIndex, MonomialIndex, VariableIndex
+
+if TYPE_CHECKING:
+    from polymatrix.polymatrix.mixins import BroadcastPolyMatrixMixin, PolyMatrixMixin
 
 
 # FIXME: use polymatrix.typing
@@ -10,7 +19,13 @@ def init_poly_matrix(
 ):
     # This is here because the old code in expression.mixins usually gives raw dictionaries
     if not isinstance(data, PolyMatrixDict):
-        data = PolyMatrixDict(data)
+        data = PolyMatrixDict({
+            MatrixIndex(*entry): PolyDict({
+                MonomialIndex(VariableIndex(*v) for v in monomial): coeff
+            })
+            for entry, poly in data.items()
+                for monomial, coeff in poly.items()
+        })
     
     return PolyMatrixImpl(
         data=data,
@@ -24,7 +39,10 @@ def init_broadcast_poly_matrix(
 ) -> BroadcastPolyMatrixMixin:
     # This is here because the old code in expression.mixins usually gives raw dictionaries
     if not isinstance(data, PolyDict):
-        data = PolyDict(data)
+        data = PolyDict({
+            MonomialIndex(VariableIndex(*v) for v in monomial): coeff
+            for monomial, coeff in data.items()
+        })
 
     return BroadcastPolyMatrixImpl(
         data=data,