Merge pull request #1798 from passagemath/matrix_keys_fix-2

Improved handling of `matrix(row_keys=..., column_keys=...)` for callables

Author: mkoeppe

src/sage/matrix/args.pyx

@@ -1139,6 +1139,42 @@ cdef class MatrixArgs:
                to Free module generated by {0, 1, 2, 3, 4} over Integer Ring
                in Category of finite dimensional modules with basis over Integer Ring;
              typ=SEQ_FLAT; entries=[0, 1, 1, 2, 2, 3, 3, 4, 4, 5]>
+            sage: MatrixArgs(lambda i, j: i + j,
+            ....:            ncols=3, row_keys=[10, 20, 30, 40]).finalized()
+            <MatrixArgs for Set of Morphisms
+             from Ambient free module of rank 3 over the principal ideal domain Integer Ring
+               to Free module generated by {10, 20, 30, 40} over Integer Ring
+               in Category of finite dimensional modules with basis over
+                  (Dedekind domains and euclidean domains and noetherian rings
+                   and infinite enumerated sets and metric spaces);
+             typ=SEQ_FLAT; entries=[10, 11, 12, 20, 21, 22, 30, 31, 32, 40, 41, 42]>
+            sage: MatrixArgs(lambda i, j: i + j,
+            ....:            column_keys=[0, 2, 5], row_keys=[10, 20, 30, 40]).finalized()
+            <MatrixArgs for Set of Morphisms
+             from Free module generated by {0, 2, 5} over Integer Ring
+               to Free module generated by {10, 20, 30, 40} over Integer Ring
+               in Category of finite dimensional modules with basis over Integer Ring;
+             typ=SEQ_FLAT; entries=[10, 12, 15, 20, 22, 25, 30, 32, 35, 40, 42, 45]>
+            sage: MatrixArgs(lambda i, j: i + j,
+            ....:            column_keys=[0, 2, 5], nrows=4).finalized()
+            <MatrixArgs for Set of Morphisms
+             from Free module generated by {0, 2, 5} over Integer Ring
+               to Ambient free module of rank 4 over the principal ideal domain Integer Ring
+               in Category of finite dimensional modules with basis over
+                  (Dedekind domains and euclidean domains and noetherian rings
+                   and infinite enumerated sets and metric spaces);
+             typ=SEQ_FLAT; entries=[0, 2, 5, 1, 3, 6, 2, 4, 7, 3, 5, 8]>
+            sage: MatrixArgs(lambda i, j: i + j,
+            ....:            ncols=4, nrows=3).finalized()
+            <MatrixArgs for Full MatrixSpace of 3 by 4 dense matrices over Integer Ring;
+             typ=SEQ_FLAT; entries=[0, 1, 2, 3, 1, 2, 3, 4, 2, 3, 4, 5]>
+            sage: MatrixArgs(ZZ, lambda i, j: i + j,
+            ....:            column_keys=[0, 2, 5], row_keys=[10, 20, 30, 40]).finalized()
+            <MatrixArgs for Set of Morphisms
+             from Free module generated by {0, 2, 5} over Integer Ring
+               to Free module generated by {10, 20, 30, 40} over Integer Ring
+               in Category of finite dimensional modules with basis over Integer Ring;
+             typ=SEQ_FLAT; entries=[10, 12, 15, 20, 22, 25, 30, 32, 35, 40, 42, 45]>
 
         Check :issue:`36065`::
 
@@ -1241,6 +1277,9 @@ cdef class MatrixArgs:
             else:
                 raise AssertionError(f"nrows={self.nrows}  ncols={self.ncols}  base={self.base}  type={self.typ}")
 
+        if self.typ == MA_ENTRIES_CALLABLE and (self.row_keys is not None or self.column_keys is not None):
+            self.finalize_callable()
+
         # Non-zero scalar matrices must be square
         # also ensure type is MA_ENTRIES_ZERO for scalar zero matrices
         if self.typ == MA_ENTRIES_SCALAR:
@@ -1518,14 +1557,33 @@ cdef class MatrixArgs:
         self._ensure_nrows_ncols()
         # Dimensions are required, so we must determine the base ring.
         # We do this by converting to the SEQ_FLAT type
-        f = self.entries
         cdef list entries = []
         cdef long i, j
-        for i in range(self.nrows):
-            row = <object>i
-            for j in range(self.ncols):
-                sig_check()
-                entries.append(f(row, j))
+
+        f = self.entries
+
+        if self.row_keys is None:
+            if self.column_keys is None:
+                for i in range(self.nrows):
+                    for j in range(self.ncols):
+                        sig_check()
+                        entries.append(f(i, j))
+            else:
+                for i in range(self.nrows):
+                    for col in self.column_keys:
+                        sig_check()
+                        entries.append(f(i, col))
+        else:
+            if self.column_keys is None:
+                for row in self.row_keys:
+                    for j in range(self.ncols):
+                        sig_check()
+                        entries.append(f(row, j))
+            else:
+                for row in self.row_keys:
+                    for col in self.column_keys:
+                        sig_check()
+                        entries.append(f(row, col))
 
         self.set_seq_flat(entries)