Introduce PolyMatrix.variables() to retrieve all variables

1 files changed, 25 insertions, 1 deletions

diff --git a/polymatrix/polymatrix/mixins.py b/polymatrix/polymatrix/mixins.py
index bb2fc17..b35490d 100644
--- a/polymatrix/polymatrix/mixins.py
+++ b/polymatrix/polymatrix/mixins.py
@@ -63,6 +63,20 @@ class PolyMatrixMixin(ABC):
                 if poly := self.at(row, col):
                     yield MatrixIndex(row, col), poly
 
+    def variables(self) -> set[VariableIndex]:
+        """
+        Get the indices of all variable indices that are in the polymatrix
+        object. Only scalar VariableIndices are returned, i.e. they all have a
+        power of one.
+        """
+        variables = set()
+        for _, poly in self.entries():
+            for monomial, _ in poly.terms():
+                for var in monomial:
+                    variables.add(VariableIndex(var.index, power=1))
+
+        return variables
+
     # -- Old API ---
 
     @deprecated("Replaced by PolyMatrixMixin.entries()")
@@ -124,7 +138,17 @@ class BroadcastPolyMatrixMixin(PolyMatrixMixin, ABC):
     @override
     def at(self, row: int, col: int) -> PolyDict:
         """ See :py:meth:`PolyMatrixMixin.at`. """
-        return self.data
+        return self.polynomial
+
+    @override
+    def variables(self) -> set[VariableIndex]:
+        # Override to be more efficient
+        variables = set()
+        for monomial, _ in self.poly.terms():
+            for var in monomial:
+                variables.add(VariableIndex(var.index, power=1))
+
+        return variables
 
     # --- Old API ---