Handle non-square matrices - #31
Conversation
pfebrer
left a comment
There was a problem hiding this comment.
Very nice thanksss!! I will try to move things around for the basis tables so that one can do:
BasisTableWithEdges(basis)
# or
BasisTableWithEdges({"rows": basis_rows, "cols": basis_cols})and then remove the matrix_role from PointBasis. Do you see any problem with this?
| from graph2mat import ( | ||
| BasisConfiguration, | ||
| BasisTableWithEdges, | ||
| BasisTableWithEdges, |
There was a problem hiding this comment.
BasisTableWithEdges is important, but no need to import it twice :)
| if isinstance(tensor, np.ndarray): | ||
| return tensor |
There was a problem hiding this comment.
This shouldn't be here, if you are using the converter from torch to numpy the input should be a torch tensor. Wherever it was needed you need to avoid calling this with a torch tensor
| import numpy as np | ||
|
|
||
| from ..table import BasisTableWithEdges | ||
| from ..table import BasisTableWithEdges, BasisTableWithEdges |
There was a problem hiding this comment.
You really like importing it twice 😅
|
|
||
| n_orbitals = [point.basis_size for point in self.basis_table.basis] | ||
| kwargs["orbitals"] = [n_orbitals[at_type] for at_type in point_types] | ||
| # SN: changed this -- the n_orbitals can be diff in cols and rows if its non square |
| shape (2, n_edges), for each edge the indices of the atoms | ||
| that participate. If `symmetrize_edges` is `True`, this must | ||
| ONLY contain the edges in one of the directions. | ||
| # TODO: verify this work with non sym in the non-square case. |
There was a problem hiding this comment.
TODO: Remove comment
Because this is already verified no?
…as dictionarieswith keys "row" and "col"
pfebrer
left a comment
There was a problem hiding this comment.
It is looking good, can you clean it so that is ready for merging? I.e. remove extra comments (e.g. those that start with SN), remove prints, etc... Then I will do a final review
| @converter | ||
| def _torch_to_numpy(tensor: torch.Tensor) -> np.ndarray: | ||
| # if isinstance(tensor, np.ndarray): | ||
| # return tensor |
There was a problem hiding this comment.
Remove these commented lines
| def __init__( | ||
| self, basis: Sequence[PointBasis], get_point_matrix: Optional[Callable] = None | ||
| ): | ||
| def __init__(self, basis: Sequence[PointBasis], get_point_matrix: Optional[Callable] = None): |
There was a problem hiding this comment.
basis: Sequence[PointBasis] | dict[Literal["row", "col"], Sequence[PointBasis]]
| self.edge_type = point_types_to_edge_types | ||
| # For non square matrices, we don't have a symmetric lower triangular part, | ||
| # but the change in sign will still be used to indicate the direction of the edge. | ||
|
|
There was a problem hiding this comment.
We are talking about the matrix of edge types, where there is a lower triangular part no matter what happens, so this comment is not needed
There was a problem hiding this comment.
True, I accidentally left it from the modification where all edges were possitive
| # Block shape: (row_basis_size[i], col_basis_size[j]) for each edge type (i, j). | ||
| self.edge_block_shape = np.array([ | ||
| row_basis_size[row_type_indices], | ||
| col_basis_size[col_type_indices], | ||
| ]) # shape: (2, n_combos) | ||
|
|
||
| return table | ||
| # Define the inverse block shape: (row_basis_size[j], col_basis_size[i]) for each edge type (i, j). | ||
| # This is befause for non-square, (Ac, Br) != (Bc, Ar) in general. | ||
| self.edge_block_shape_inv = np.array([ | ||
| row_basis_size[col_type_indices], | ||
| col_basis_size[row_type_indices], | ||
| ]) # shape: (2, n_combos) |
There was a problem hiding this comment.
Let's just define a self.edge_block_shape that is twice as long so that you can use it with positive and with negative indices. Same for all the other arrays like this. Then when you use them you don't need to check if self.is_square.
| # Quick check if already correctly ordered | ||
| if all(a.type == b.type for a, b in zip(i_basis, j_basis)): | ||
| return i_basis, j_basis | ||
| print("Reordering basis to match rows and columns.") |
Edited graph2mat to be able to handle non-square matrices. This is done by adding the option to specify the
matrix_rolewhen defining aPointBasisobject and modifying the indexing in matrices to take into account both row and column sizes.Example code:
point_1 = PointBasis("A", R=2, basis="0e", basis_convention="spherical", matrix_role='row')
point_2 = PointBasis("A", R=2, basis="2x0e", basis_convention="spherical", matrix_role='col')
point_3 = PointBasis("B", R=5, basis="0e + 1o", basis_convention="spherical", matrix_role='row')
point_4 = PointBasis("B", R=5, basis="2x0e + 1o", basis_convention="spherical", matrix_role='col')
For now, still some handlers, like the ones for plotting the limits correctly, are not implemented.