diff --git a/src/graph2mat/bindings/torch/data/formats.py b/src/graph2mat/bindings/torch/data/formats.py index fba58f1..a477437 100644 --- a/src/graph2mat/bindings/torch/data/formats.py +++ b/src/graph2mat/bindings/torch/data/formats.py @@ -37,6 +37,8 @@ def _csr_to_dense(csr: torch.sparse_csr_tensor) -> torch.Tensor: @converter def _torch_to_numpy(tensor: torch.Tensor) -> np.ndarray: + # if isinstance(tensor, np.ndarray): + # return tensor return tensor.numpy(force=True) @@ -53,7 +55,8 @@ def nodes_and_edges_to_coo( node_vals: torch.Tensor, edge_vals: torch.Tensor, edge_index: torch.Tensor, - orbitals: torch.Tensor, + orbitals_row: torch.Tensor, + orbitals_col: torch.Tensor, n_supercells: int = 1, edge_neigh_isc: Optional[torch.Tensor] = None, threshold: Optional[float] = None, @@ -100,7 +103,8 @@ def _init_coo(data, rows, cols, shape): node_vals=node_vals, edge_vals=edge_vals, edge_index=edge_index, - orbitals=orbitals, + orbitals_row=orbitals_row, + orbitals_col=orbitals_col, n_supercells=n_supercells, edge_neigh_isc=edge_neigh_isc, threshold=threshold, diff --git a/src/graph2mat/core/data/basis.py b/src/graph2mat/core/data/basis.py index 560100e..41b77e2 100644 --- a/src/graph2mat/core/data/basis.py +++ b/src/graph2mat/core/data/basis.py @@ -150,7 +150,6 @@ class PointBasis: You can also use the aliases that we provide, such as ``"cartesian"`` or ``"spherical"``. - Examples ---------- diff --git a/src/graph2mat/core/data/processing.py b/src/graph2mat/core/data/processing.py index 3e37b77..e611914 100644 --- a/src/graph2mat/core/data/processing.py +++ b/src/graph2mat/core/data/processing.py @@ -510,6 +510,9 @@ def sort_edge_index( inplace: bool, optional Whether the output should be placed in the input arrays, otherwise new arrays are created. + is_square: bool, optional + Whether the matrix is square or not. If it is not square, the edge types + are not signed, and therefore ther will not be negative edge types. Return --------- @@ -804,8 +807,10 @@ def labels_to( else: kwargs["n_supercells"] = nsc.prod() - n_orbitals = [point.basis_size for point in self.basis_table.basis] - kwargs["orbitals"] = [n_orbitals[at_type] for at_type in point_types] + n_orbitals_row = [point.basis_size for point in self.basis_table.row_basis] + n_orbitals_col = [point.basis_size for point in self.basis_table.col_basis] + kwargs["orbitals_row"] = [n_orbitals_row[at_type] for at_type in point_types] + kwargs["orbitals_col"] = [n_orbitals_col[at_type] for at_type in point_types] # Add back atomic contributions to the node blocks in case they were removed if self.sub_point_matrix: @@ -1354,14 +1359,13 @@ def from_config( # array that converts from sc shifts (3D) to a single supercell index. This is isc_off. supercell = sisl.Lattice(config.cell, nsc=nsc) - # Get the edge types edge_types = data_processor.basis_table.point_type_to_edge_type( indices[edge_index] ) # Sort the edges to make it easier for the reading routines data_processor.sort_edge_index( - edge_index, sc_shifts, shifts.T, edge_types, supercell.isc_off, inplace=True + edge_index, sc_shifts, shifts.T, edge_types, supercell.isc_off, inplace=True, ) # Then, get the supercell index of each interaction. diff --git a/src/graph2mat/core/data/sparse.py b/src/graph2mat/core/data/sparse.py index 057c613..79b8459 100644 --- a/src/graph2mat/core/data/sparse.py +++ b/src/graph2mat/core/data/sparse.py @@ -155,7 +155,8 @@ def block_dict_to_coo( def _blockmatrix_coo_coords( - orbitals: ArrayLike, + orbitals_row: ArrayLike, + orbitals_col: ArrayLike, edge_index: ArrayLike, n_supercells: int = 1, edge_neigh_isc: Optional[ArrayLike] = None, @@ -182,8 +183,11 @@ def _blockmatrix_coo_coords( Parameters ---------- - orbitals: - for each atom, the amount of orbitals it has. + orbitals_row / orbitals_col: + for each atom, the amount of orbitals it has. In case of contracted basis sets, + rows represent the original basis, and columns the contracted basis. In case of + non-contracted basis sets, both are equal. + edge_index: shape (2, n_edges), for each edge the indices of the atoms that participate. If `symmetrize_edges` is `True`, this must @@ -202,17 +206,22 @@ def _blockmatrix_coo_coords( # Initialize the arrays to store the coordinates. rows = [] cols = [] + is_square = np.array_equal(orbitals_row, orbitals_col) # Store index of first orbital for each atom, as well as total number of orbitals. - first_orb = np.cumsum([0, *orbitals]) - no = first_orb[-1] + first_orb_row = np.cumsum([0, *orbitals_row]) + first_orb_col = np.cumsum([0, *orbitals_col]) + no_row = first_orb_row[-1] # Total number of orbitals in the row. + no_col = first_orb_col[-1] # Total number of orbitals in the column. # First, compute the coordinates for the node blocks. - for i_at, dim in enumerate(orbitals): - i_start = first_orb[i_at] - i_end = i_start + dim + for dim in range(len(orbitals_row)): + i_start_row = first_orb_row[dim] + i_end_row = i_start_row + orbitals_row[dim] + j_start_col = first_orb_col[dim] + j_end_col = j_start_col + orbitals_col[dim] - block_rows, block_cols = np.mgrid[i_start:i_end, i_start:i_end].reshape(2, -1) + block_rows, block_cols = np.mgrid[i_start_row:i_end_row, j_start_col:j_end_col].reshape(2, -1) rows.extend(block_rows) cols.extend(block_cols) @@ -232,19 +241,23 @@ def _blockmatrix_coo_coords( for i_edge, ((i_at, j_at), neigh_isc) in enumerate( zip(edge_index.T, edge_neigh_isc) - ): - i_start = first_orb[i_at] - i_end = i_start + orbitals[i_at] - j_start = first_orb[j_at] - j_end = j_start + orbitals[j_at] + ): + if neigh_isc != 0 and not is_square: + raise NotImplementedError( + "Supercell interactions are not yet implemented for non-square matrices (contractions)." + ) + i_start = first_orb_row[i_at] + i_end = i_start + orbitals_row[i_at] + j_start = first_orb_col[j_at] + j_end = j_start + orbitals_col[j_at] block_rows, block_cols = np.mgrid[i_start:i_end, j_start:j_end].reshape(2, -1) - sc_block_cols = block_cols + no * neigh_isc + sc_block_cols = block_cols + no_col * neigh_isc rows.extend(block_rows) cols.extend(sc_block_cols) - if symmetrize_edges: + if symmetrize_edges and is_square: # Columns and rows are easy to determine if the connection is in the unit # cell, as the opposite block is in the transposed location. opp_block_cols = block_rows @@ -253,16 +266,18 @@ def _blockmatrix_coo_coords( if neigh_isc != 0: # For supercell connections we need to find out what is the the supercell # index of the neighbor in the opposite connection. - opp_block_cols += no * (n_supercells - neigh_isc) + opp_block_cols += no_col * (n_supercells - neigh_isc) rows_symm.extend(opp_block_rows) cols_symm.extend(opp_block_cols) - + elif symmetrize_edges and not is_square: + raise ValueError( + "We cannot use symmetrized edges for non-square matrices (contractions).") # Add coordinates of symmetrized edges to the list of coordinates. rows.extend(rows_symm) cols.extend(cols_symm) - return np.array(rows), np.array(cols), (no, no * n_supercells) + return np.array(rows), np.array(cols), (no_row, no_col * n_supercells) def _nodes_and_edges_to_coo( @@ -271,7 +286,8 @@ def _nodes_and_edges_to_coo( node_vals: ArrayLike, edge_vals: ArrayLike, edge_index: ArrayLike, - orbitals: ArrayLike, + orbitals_row: ArrayLike, + orbitals_col: ArrayLike, n_supercells: int = 1, edge_neigh_isc: Optional[ArrayLike] = None, threshold: Optional[float] = None, @@ -306,8 +322,10 @@ def _nodes_and_edges_to_coo( edge_index Array of shape (2, n_edges) containing the indices of the atoms that participate in each edge. - orbitals + orbitals_row / orbitals_col Array of shape (n_nodes, ) containing the number of orbitals for each atom. + For contracted basis sets, rows represent the original basis, and columns the + contracted basis. For non-contracted basis sets, both are equal. n_supercells Number of auxiliary supercells. edge_neigh_isc @@ -320,15 +338,14 @@ def _nodes_and_edges_to_coo( whether for each edge only one direction is provided. The edge block for the opposite direction is then created as the transpose. """ - rows, cols, shape = _blockmatrix_coo_coords( - orbitals=orbitals, + orbitals_row=orbitals_row, + orbitals_col=orbitals_col, edge_index=edge_index, n_supercells=n_supercells, edge_neigh_isc=edge_neigh_isc, symmetrize_edges=symmetrize_edges, ) - if symmetrize_edges: sparse_data = concatenate([node_vals, edge_vals, edge_vals]) else: @@ -352,7 +369,8 @@ def nodes_and_edges_to_coo( node_vals: np.ndarray, edge_vals: np.ndarray, edge_index: np.ndarray, - orbitals: np.ndarray, + orbitals_row: np.ndarray, + orbitals_col: np.ndarray, n_supercells: int = 1, edge_neigh_isc: Optional[np.ndarray] = None, threshold: Optional[float] = None, @@ -373,8 +391,10 @@ def nodes_and_edges_to_coo( edge_index Array of shape (2, n_edges) containing the indices of the atoms that participate in each edge. - orbitals + orbitals_row / orbitals_col Array of shape (n_nodes, ) containing the number of orbitals for each atom. + For contracted basis sets, rows represent the original basis, and columns the + contracted basis. For non-contracted basis sets, both are equal. n_supercells Number of auxiliary supercells. edge_neigh_isc @@ -397,7 +417,8 @@ def _init_coo(data, rows, cols, shape): node_vals=node_vals, edge_vals=edge_vals, edge_index=edge_index, - orbitals=orbitals, + orbitals_row=orbitals_row, + orbitals_col=orbitals_col, n_supercells=n_supercells, edge_neigh_isc=edge_neigh_isc, threshold=threshold, diff --git a/src/graph2mat/core/data/table.py b/src/graph2mat/core/data/table.py index 4a839b0..5d89735 100644 --- a/src/graph2mat/core/data/table.py +++ b/src/graph2mat/core/data/table.py @@ -15,7 +15,8 @@ import itertools from io import StringIO from pathlib import Path -from typing import Callable, Generator, List, Literal, Optional, Sequence, Union +from typing import Callable, Generator, List, Literal, Optional, Sequence, Union, Tuple +from itertools import chain import numpy as np import sisl @@ -31,11 +32,16 @@ class BasisTableWithEdges: Its function is to assist in pre and post processing data by providing a centralized source of truth for the basis that a model should be able to deal with. + It allows non-square matrices, in which case the row and column basis are different. In that case, + the basis should be passed as a dict with keys "row" and "col" and the corresponding basis for each, + passed as a list of `PointBasis` objects with matching types. The order of the types in the row and + column basis should be the same (although it handles reordering), and the types should be unique. + Parameters ---------- basis: List of `PointBasis` objects for types that are (possibly) present in the systems - of interest. + of interest. If dict, it should have keys "row" and "col" with the corresponding basis for each. get_point_matrix: A function that takes a `PointBasis` object and returns the matrix that is a constant for that type of point. @@ -52,7 +58,8 @@ class BasisTableWithEdges: """ #: List of ``PointBasis`` objects that this table knows about. - basis: List[PointBasis] + row_basis: List[PointBasis] + col_basis: List[PointBasis] #: The spherical harmonics convention used for the basis #: (same for all ``PointBasis``). basis_convention: BasisConvention @@ -104,20 +111,40 @@ class BasisTableWithEdges: #: If the basis was read from files, this might store the contents of the files. #: For saving/loading purposes. file_contents: Optional[List[str]] - - def __init__( - self, basis: Sequence[PointBasis], get_point_matrix: Optional[Callable] = None - ): + def __init__(self, basis: Sequence[PointBasis], get_point_matrix: Optional[Callable] = None): self._init_args = {"atoms": basis, "get_point_matrix": get_point_matrix} - self.basis = list(basis) - self.types = [point_basis.type for point_basis in self.basis] + # Non-square matrices are allowed, but only if the basis is passed as a dict with keys "row" and "col". + self.is_square = not isinstance(basis, dict) + + # Define row and col basis + # The funtion also orders the basis by types, taking the order of the rows basis + # Also verifyes that len(row)==len(col), and types_row == types_col + row_basis, col_basis = self._process_basis_rows_cols(basis) + self.row_basis = row_basis + self.col_basis = col_basis + + # Get types + self.types = [point_basis.type for point_basis in self.row_basis] assert len(set(self.types)) == len( - self.basis + self.row_basis ), f"The tag of each basis must be unique. Got {self.types}." + n_types = len(self.types) + + self.R = self.get_R() + + + # Define the point matrix + if get_point_matrix is None: + self.point_matrix = None + else: + self.point_matrix = [ + get_point_matrix(point_basis_row, point_basis_col) + for point_basis_row, point_basis_col in zip(self.row_basis, self.col_basis) + ] # Define the basis convention and make sure that all the point basis adhere to that convention. - for point_basis in self.basis: + for point_basis in chain(self.row_basis, self.col_basis): if len(point_basis.basis) > 0: basis_convention = point_basis.basis_convention break @@ -126,7 +153,7 @@ def __init__( all_conventions = [ point_basis.basis_convention - for point_basis in self.basis + for point_basis in chain(self.row_basis, self.col_basis) if len(point_basis.basis) > 0 ] if len(all_conventions) > 0: @@ -142,7 +169,6 @@ def __init__( "cartesian", self.basis_convention ) - n_types = len(self.types) # Array to get the edge type from point types. point_types_to_edge_types = np.empty((n_types, n_types), dtype=np.int32) edge_type = 0 @@ -150,7 +176,6 @@ def __init__( # The diagonal edge type, always positive point_types_to_edge_types[i, i] = edge_type edge_type += 1 - # The non diagonal edge types, which are negative for the lower triangular part, # to account for the fact that the direction is different. for j in range(i + 1, n_types): @@ -158,91 +183,115 @@ def __init__( point_types_to_edge_types[j, i] = -edge_type edge_type += 1 - 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. - # Get the point matrix for each type. This is the matrix that a point would - # have if it was the only one in the system, and it depends only on the type. - if get_point_matrix is None: - self.point_matrix = None - else: - self.point_matrix = [ - get_point_matrix(point_basis) for point_basis in self.basis - ] - - # Get also the cutoff radii for each point. - self.R = np.array([point_basis.maxR() for point_basis in self.basis]) + self.edge_type = point_types_to_edge_types # Store the sizes of each point's basis. - self.basis_size = np.array( - [basis.basis_size for basis in self.basis], dtype=np.int32 + row_basis_size = np.array( + [basis.basis_size for basis in self.row_basis], dtype=np.int32 + ) + self.row_basis_size = row_basis_size + col_basis_size = np.array( + [basis.basis_size for basis in self.col_basis], dtype=np.int32 ) + self.col_basis_size = col_basis_size # And also the sizes of the blocks. - self.point_block_shape = np.array([self.basis_size, self.basis_size]) - self.point_block_size = self.basis_size**2 + self.point_block_shape = np.array([row_basis_size, col_basis_size]) + self.point_block_size = row_basis_size * col_basis_size point_types_combinations = np.array( - list(itertools.combinations_with_replacement(range(n_types), 2)) - ).T - self.edge_type_to_point_types = point_types_combinations.T - self.edge_block_shape = self.basis_size[point_types_combinations] - self.edge_block_size = self.edge_block_shape.prod(axis=0) - - def __repr__(self): - return f"{self.__class__.__name__}({self.basis_convention}, basis={self.basis})" + list(itertools.combinations_with_replacement(range(n_types), 2))).T + + # Store the sizes of each point's basis + # Even if its not square, we take each direction once: + # we define the inverse shape, where we pass from (Ar, Bc) to (Ac, Br) - def _repr_html_(self): - table = "" - table += f"" - - def _basis_string(basis): - s = "" - - for basis_set in basis: - mul, l, parity = basis_set - s += f"{mul}x{l}{'e' if parity == 1 else 'o'} + " - s = s[:-3] - - return s + self.edge_type_to_point_types = point_types_combinations.T - for i, point_basis in enumerate(self.basis): - table += f"" + row_type_indices = point_types_combinations[0] # which row-type each combo refers to + col_type_indices = point_types_combinations[1] # which col-type each combo refers to - table += "
IndexTypeIrrepsMax R
{i}{point_basis.type}{_basis_string(point_basis.basis)}{point_basis.maxR()}
" + # 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) - def __str__(self): - return "\n".join([f"\t- {point_basis}" for point_basis in self.basis]) + self.edge_block_size = self.edge_block_shape.prod(axis=0) # shape: (n_combos,) + self.edge_block_size_inv = self.edge_block_shape_inv.prod(axis=0) # shape: (n_combos,) - def __len__(self): - return len(self.basis) + @staticmethod + def _process_basis_rows_cols(basis: Union[List[PointBasis], dict[str, List[PointBasis]]]) \ + -> Tuple[List[PointBasis], List[PointBasis]]: + """Processes the basis to determine which basis is used for rows and which for columns. - def __eq__(self, other): - if not isinstance(other, self.__class__): - return False + If the basis has a `matrix_role` attribute, it will be used to determine the role of each basis. + If not, the first basis will be used for rows and the second for columns. - same = all(x == y for x, y in itertools.zip_longest(self.basis, other.basis)) - same &= all(x == y for x, y in itertools.zip_longest(self.types, other.types)) + Parameters + ---------- + basis: List[PointBasis] or dict[str, List[PointBasis]] + In the case of the list, it is a square matrix with the same + basis for rows and columns. In the case of the dict, it should + have keys "row" and "col" with the corresponding basis. - if self.point_matrix is None: - same &= other.point_matrix is None + Returns + ------- + i_basis: List[PointBasis] + The list of point bases for rows. + j_basis: List[PointBasis] + The list of point bases for columns. + """ + i_basis = [] + j_basis = [] + if isinstance(basis, dict): + if "row" not in basis or "col" not in basis: + raise ValueError("If basis is a dict, it must have keys 'row' and 'col'.") + i_basis = basis["row"] + j_basis = basis["col"] else: - if other.point_matrix is None: - return False - same &= all( - np.allclose(x, y) - for x, y in itertools.zip_longest(self.point_matrix, other.point_matrix) - ) + i_basis = j_basis = basis + assert len(i_basis) == len(j_basis), "The number of row and column bases must be the same." + # Make sure order if correct + # If it is correct, return the basis as is. If not, reorder the basis to match the order of the other basis. + # 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.") + + type_to_index = {} + for idx, elem in enumerate(j_basis): + if elem.type in type_to_index: + raise ValueError(f"Duplicate type '{elem.type}' found in j_basis") + type_to_index[elem.type] = idx + # Reorder basis: check that the point type is the same for i_basis[k] and j_basis[k] + ordered_j_basis = [] + for elem in i_basis: + idx = type_to_index.get(elem.type) + if idx is None: + raise ValueError(f"Type '{elem.type}' not found in j_basis") + ordered_j_basis.append(j_basis[idx]) + return i_basis, ordered_j_basis + + def get_R(self) -> np.ndarray: + """Returns the cutoff radii in the correct format.""" + # Get also the cutoff radii for each point. + r_rows = np.array([point_basis.maxR() for point_basis in self.row_basis]) + r_cols = np.array([point_basis.maxR() for point_basis in self.col_basis]) + if not np.allclose(r_rows, r_cols): + print(f"Warning: Row and column point types have different cutoff radii. Taking the maximum of the two.") + return np.max([r_rows, r_cols], axis=0) - same &= np.allclose(self.edge_type, other.edge_type) - same &= np.allclose(self.R, other.R) - same &= np.allclose(self.basis_size, other.basis_size) - same &= np.allclose(self.point_block_shape, other.point_block_shape) - same &= np.allclose(self.point_block_size, other.point_block_size) - same &= np.allclose(self.edge_block_shape, other.edge_block_shape) - same &= np.allclose(self.edge_block_size, other.edge_block_size) - return same def group( self, grouping: Literal["basis_shape", "point_type", "max"] @@ -306,37 +355,39 @@ def group( values = ... # some computation with the grouped basis (dim_new_basis, ) type0_values = values[filters[0]] - """ + + "TODO: define this for nonsquare in this scheme" + # SN: fix I did: for non square, raises error if its not point type. + # Hence, the square grouping takes only row_basis, and does the same as before. + + if not self.is_square and grouping != "point_type": + raise NotImplementedError(f"Grouping for non-square matrices with grouping {grouping} is not implemented yet.") filters = None if grouping == "point_type": new_table = self - class A: - def __getitem__(self, key): - return key - - point_type_conversion = A() - edge_type_conversion = A() + point_type_conversion = IdentityConversion() + edge_type_conversion = IdentityConversion() elif grouping == "basis_shape": # Get all basis sizes: - basis_sizes = np.zeros((len(self.basis), 5), dtype=int) - for i, point_type_basis in enumerate(self.basis): + basis_sizes = np.zeros((len(self.row_basis), 5), dtype=int) + for i, point_type_basis in enumerate(self.row_basis): for n, l, _ in point_type_basis.basis: basis_sizes[i, l] += n - + # Get the unique basis sizes unique_sizes, unique_indices, pseudo_types = np.unique( basis_sizes, axis=0, return_index=True, return_inverse=True ) - + # Create the new table with the unique basis sizes. Here we just # take the first point type that has that basis size, but perhaps # it would be better to create a new point type that represents # that basis size (e.g. so that the name of the type is not misleading). - new_table = self.__class__([self.basis[i] for i in unique_indices]) + new_table = self.__class__([self.row_basis[i] for i in unique_indices]) point_type_conversion = pseudo_types - + # Get conversions old_edgetypes_to_new_point_types = point_type_conversion[ self.edge_type_to_point_types @@ -352,14 +403,14 @@ def __getitem__(self, key): ) elif grouping == "max": # Get all basis sizes: - basis_sizes = np.zeros((len(self.basis), 5), dtype=int) - for i, point_type_basis in enumerate(self.basis): + basis_sizes = np.zeros((len(self.row_basis), 5), dtype=int) + for i, point_type_basis in enumerate(self.row_basis): for n, l, _ in point_type_basis.basis: basis_sizes[i, l] += n - + # Maximum sizes: max_sizes = basis_sizes.max(axis=0) - + # Build the new point basis max_basis = PointBasis( "all", @@ -369,31 +420,71 @@ def __getitem__(self, key): ) # Create the new table with only one type. new_table = self.__class__([max_basis]) - + # For each original point type, compute a mask that allows us to # select the values of the new basis that correspond to that point type. # (i.e. discard the values that are not present in that original point type). missing_ls = max_sizes - basis_sizes - filters = np.zeros((len(self.basis), max_basis.basis_size), dtype=bool) + filters = np.zeros((len(self.row_basis), max_basis.basis_size), dtype=bool) i = 0 for l, n in enumerate(max_sizes): if n == 0: continue - + for i_point, point_missing_ls in enumerate(missing_ls[:, l]): filters[i_point, i : i + (n - point_missing_ls) * (2 * l + 1)] = 1 - + i += (2 * l + 1) * n - + # Point and edge type conversions, just map any type to 0. - point_type_conversion = np.zeros(len(self.basis), dtype=int) + point_type_conversion = np.zeros(len(self.row_basis), dtype=int) edge_type_conversion = np.zeros( len(self.edge_type_to_point_types), dtype=int ) else: raise NotImplementedError(f"Grouping by {grouping} is not implemented.") + + return new_table, point_type_conversion, edge_type_conversion, filters + + + + def _repr_html_(self): + def _repr_html_1basis(my_basis): + table = "" + table += f"" + + def _basis_string(basis1): + s = "" + + for basis_set in basis1: + mul, l, parity = basis_set + s += f"{mul}x{l}{'e' if parity == 1 else 'o'} + " + s = s[:-3] + + return s + + for i, point_basis in enumerate(my_basis): + table += f"" + + table += "
IndexTypeIrrepsMax R
{i}{point_basis.type}{_basis_string(point_basis.basis)}{point_basis.maxR()}
" + return table + + if self.is_square: + return _repr_html_1basis(self.row_basis) + else: + table = "" + table += f"" + for i in range(len(self.row_basis)): + table += f"" + table += "
Row BasisColumn Basis
{self.row_basis[i]}{self.col_basis[i]}
" + return table + + def __repr__(self): + if self.is_square: + return f"{self.__class__.__name__}({self.basis_convention}, basis={self.row_basis})" + else: + return f"{self.__class__.__name__}({self.basis_convention}, row_basis={self.row_basis}, col_basis={self.col_basis})" - return new_table, point_type_conversion, edge_type_conversion, filters def index_to_type(self, index: int) -> Union[str, int]: """Converts from the index of the point type to the type ID. @@ -405,6 +496,68 @@ def index_to_type(self, index: int) -> Union[str, int]: """ return self.types[index] + + + def get_sisl_atoms(self) -> List[sisl.Atom]: + """Returns a list of sisl atoms corresponding to the basis. + + If the basis does not contain atoms, `PointBasis` objects are + converted to atoms. + + WARNING: to sislfunctions only working well for SQUARE matrices. + """ + if hasattr(self, "atoms"): + return self.atoms + else: + return [point.to_sisl_atom() for point in self.row_basis] + + + + def __str__(self): + if self.is_square: + # Both are the same, print one + return "\n".join([f"\t- {point_basis}" for point_basis in self.row_basis]) + else: + str_print = "Row basis:\n" + str_print += "\n".join([f"\t- {point_basis}" for point_basis in self.row_basis]) + str_print += "\nColumn basis:\n" + str_print += "\n".join([f"\t- {point_basis}" for point_basis in self.col_basis]) + return str_print + + def __len__(self): + if len(self.col_basis) != len(self.row_basis): + raise ValueError("Row and column basis have different lengths.") + return len(self.row_basis) # Both have same length + + + def __eq__(self, other): + if not isinstance(other, self.__class__): + return False + same = self.is_square == other.is_square + + same = all(x == y for x, y in itertools.zip_longest(self.row_basis, other.row_basis)) + same &= all(x == y for x, y in itertools.zip_longest(self.col_basis, other.col_basis)) + same &= all(x == y for x, y in itertools.zip_longest(self.types, other.types)) + + if self.point_matrix is None: + same &= other.point_matrix is None + else: + if other.point_matrix is None: + return False + same &= all( + np.allclose(x, y) + for x, y in itertools.zip_longest(self.point_matrix, other.point_matrix) + ) + + same &= np.allclose(self.edge_type, other.edge_type) + same &= np.allclose(self.R, other.R) + same &= np.allclose(self.basis_size, other.basis_size) + same &= np.allclose(self.point_block_shape, other.point_block_shape) + same &= np.allclose(self.point_block_size, other.point_block_size) + same &= np.allclose(self.edge_block_shape, other.edge_block_shape) + same &= np.allclose(self.edge_block_size, other.edge_block_size) + return same + def type_to_index(self, point_type: Union[str, int]) -> int: """Converts from the type ID to the index of the point type in the table. @@ -449,10 +602,6 @@ def point_type_to_edge_type(self, point_type: np.ndarray) -> Union[int, np.ndarr """ return self.edge_type[point_type[0], point_type[1]] - def maxR(self) -> float: - """Maximum cutoff radius in the basis.""" - return self.R.max() - def point_block_pointer(self, point_types: Sequence[int]) -> np.ndarray: """Pointers to the beggining of node blocks in a flattened matrix. @@ -489,19 +638,23 @@ def edge_block_pointer(self, edge_types: Sequence[int]): which they appear in the flattened matrix. """ pointers = np.zeros(len(edge_types) + 1, dtype=np.int32) - np.cumsum(self.edge_block_size[np.abs(edge_types)], out=pointers[1:]) + if self.is_square: + np.cumsum(self.edge_block_size[np.abs(edge_types)], out=pointers[1:]) + else: + # In non-square case, the edge block size is not symmetric, so for + # negative values of + # New conditional selection: + sizes = np.where(edge_types > 0, + self.edge_block_size[edge_types], # positive → forward size + self.edge_block_size_inv[-edge_types]) # non‑positive → backward size (use absolute index) + + np.cumsum(sizes, out=pointers[1:]) return pointers - def get_sisl_atoms(self) -> List[sisl.Atom]: - """Returns a list of sisl atoms corresponding to the basis. - - If the basis does not contain atoms, `PointBasis` objects are - converted to atoms. - """ - if hasattr(self, "atoms"): - return self.atoms - else: - return [point.to_sisl_atom() for point in self.basis] + def maxR(self) -> float: + """Maximum cutoff radius in the basis.""" + return self.R.max() + class AtomicTableWithEdges(BasisTableWithEdges): @@ -666,3 +819,104 @@ def __setstate__(self, d): self._set_state_by_filecontents(file_names, file_contents) else: self._set_state_by_atoms(d["atoms"]) + +# SN: defined this outside to be able to compare the poijnt type conversions +# for nonsquare correclty: it should be fulfilled by default, but porsiacaso. +class IdentityConversion: + def __getitem__(self, key): + return key + + def __eq__(self, other): + return isinstance(other, IdentityConversion) + + +# OLD CODE: + +# SN: I just copied how to do it for a single basis, so that the changes are minimal, +# and I do not mess up the grouping +def group(self, grouping: Literal["basis_shape", "point_type", "max"]) \ + -> tuple["BasisTableWithEdges", + np.ndarray, np.ndarray, Optional[np.ndarray]]: + filters = None + if grouping == "point_type": + new_table = self + + point_type_conversion = IdentityConversion() + edge_type_conversion = IdentityConversion() + + elif grouping == "basis_shape": + # Get all basis sizes: + basis_sizes = np.zeros((len(self.basis), 5), dtype=int) + for i, point_type_basis in enumerate(self.basis): + for n, l, _ in point_type_basis.basis: + basis_sizes[i, l] += n + + # Get the unique basis sizes + unique_sizes, unique_indices, pseudo_types = np.unique( + basis_sizes, axis=0, return_index=True, return_inverse=True + ) + + # Create the new table with the unique basis sizes. Here we just + # take the first point type that has that basis size, but perhaps + # it would be better to create a new point type that represents + # that basis size (e.g. so that the name of the type is not misleading). + new_table = self.__class__([self.basis[i] for i in unique_indices]) + point_type_conversion = pseudo_types + + # Get conversions + old_edgetypes_to_new_point_types = point_type_conversion[ + self.edge_type_to_point_types + ] + edge_type_conversion = new_table.edge_type[ + old_edgetypes_to_new_point_types[:, 0], + old_edgetypes_to_new_point_types[:, 1], + ] + # For edge type conversions handle the case in which the edge type + # is negative. + edge_type_conversion = np.concatenate( + [edge_type_conversion, -1 * np.flip(edge_type_conversion[1:])] + ) + elif grouping == "max": + # Get all basis sizes: + basis_sizes = np.zeros((len(self.basis), 5), dtype=int) + for i, point_type_basis in enumerate(self.basis): + for n, l, _ in point_type_basis.basis: + basis_sizes[i, l] += n + + # Maximum sizes: + max_sizes = basis_sizes.max(axis=0) + + # Build the new point basis + max_basis = PointBasis( + "all", + R=self.maxR(), + basis=[(int(n), l, (-1) ** l) for l, n in enumerate(max_sizes)], + basis_convention=self.basis_convention, + ) + # Create the new table with only one type. + new_table = self.__class__([max_basis]) + + # For each original point type, compute a mask that allows us to + # select the values of the new basis that correspond to that point type. + # (i.e. discard the values that are not present in that original point type). + missing_ls = max_sizes - basis_sizes + filters = np.zeros((len(self.basis), max_basis.basis_size), dtype=bool) + i = 0 + for l, n in enumerate(max_sizes): + if n == 0: + continue + + for i_point, point_missing_ls in enumerate(missing_ls[:, l]): + filters[i_point, i : i + (n - point_missing_ls) * (2 * l + 1)] = 1 + + i += (2 * l + 1) * n + + # Point and edge type conversions, just map any type to 0. + point_type_conversion = np.zeros(len(self.basis), dtype=int) + edge_type_conversion = np.zeros( + len(self.edge_type_to_point_types), dtype=int + ) + else: + raise NotImplementedError(f"Grouping by {grouping} is not implemented.") + + return new_table, point_type_conversion, edge_type_conversion, filters diff --git a/src/graph2mat/core/modules/_labels_resort.py b/src/graph2mat/core/modules/_labels_resort.py index bcf945e..bccca70 100644 --- a/src/graph2mat/core/modules/_labels_resort.py +++ b/src/graph2mat/core/modules/_labels_resort.py @@ -4,10 +4,11 @@ @cython.boundscheck(False) -@cython.wraparound(False) +# @cython.wraparound(False) def get_labels_resorting_array( types: cython.integral[:], shapes: cython.integral[:, :], + shapes_inv: cython.integral[:, :] = None, # SN: for cases where is not square is needed: the shape of edge i ! = -i transpose_neg: cython.bint = False, ): """ @@ -55,7 +56,8 @@ def get_labels_resorting_array( one direction is predicted. """ n_entries = types.shape[0] - n_types: cython.int = shapes.shape[1] + n_types: cython.int = shapes.shape[1]*2-1 # positive and negative : with shape_inv + ntypes_int: cython.int = shapes.shape[1] # this is the number of integers >=0 we have. type: cython.int rows: cython.int @@ -66,26 +68,44 @@ def get_labels_resorting_array( type_nlabels: cython.long[:] = np.zeros(n_types, dtype=int) offset: cython.long[:] = np.zeros(n_types, dtype=int) + if shapes_inv is None: + shapes_inv = shapes + + # Check node 0 shape is the same - it should be ALWAYS. + i: cython.int + for i in range(shapes.shape[0]): + if shapes[i, 0] != shapes_inv[i, 0]: + raise ValueError( + "For nodes, they must have the same shape, so shapes[:,0] == shapes_inv[:,0]" + ) + # Compute the sizes for each type sizes: cython.int[:] = np.zeros(n_types, dtype=np.int32) - for type in range(n_types): - sizes[type] = shapes[0, type] * shapes[1, type] + for type in range(ntypes_int): + sizes[type + ntypes_int - 1] = shapes[0, type] * shapes[1, type] + sizes[-type + ntypes_int - 1] = shapes_inv[0, type] * shapes_inv[1, type] # Count the number of entries of each type for i_edge in range(n_entries): - type: cython.int = abs(types[i_edge]) - - type_nlabels[type] += sizes[type] + type: cython.int = types[i_edge] + type_nlabels[type + ntypes_int - 1] += sizes[type + ntypes_int - 1] # Cumsum of type_nlabels to understand where do the labels for # each type start. - for type in range(1, n_types): - offset[type] = offset[type - 1] + type_nlabels[type - 1] + + prev_type: cython.int = ntypes_int - 1 # Start from the position that corresponds to type 0 + for type in range(1, ntypes_int): # Here we just range in n_types, but this takes the negatives + offset[type + ntypes_int - 1] = offset[prev_type] + type_nlabels[prev_type] # >0 + offset[-type + ntypes_int - 1] = offset[type + ntypes_int - 1] + type_nlabels[type + ntypes_int - 1] # < 0 + + # We have to continue from the negative type, because the next positive type will be after it. + prev_type = -type + ntypes_int - 1 # Initialize the indices array. # (for each label value, index of the unsorted array where it is located) + # The +1 is to avoid that ntypes_int=n means that we have n-1 integers types, avoiding 0. indices: cython.long[:] = np.empty( - offset[n_types - 1] + type_nlabels[n_types - 1], dtype=int + offset[0] + type_nlabels[0], dtype=int ) type_i: cython.long[:] = np.zeros_like(sizes, dtype=int) @@ -95,8 +115,8 @@ def get_labels_resorting_array( type = types[i_edge] abs_type: cython.int = abs(type) - block_size: cython.int = sizes[abs_type] - start: cython.int = offset[abs_type] + type_i[abs_type] + block_size: cython.int = sizes[type + ntypes_int - 1] + start: cython.int = offset[type + ntypes_int - 1] + type_i[type + ntypes_int - 1] if transpose_neg and type < 0: # Get the transposed shape @@ -105,13 +125,12 @@ def get_labels_resorting_array( for jcol in range(cols): indices[i] = start + jcol * rows + jrow i += 1 - else: for j in range(start, start + block_size): indices[i] = j i += 1 - type_i[abs_type] += block_size + type_i[type + ntypes_int - 1] += block_size return np.asarray(indices) diff --git a/src/graph2mat/core/modules/graph2mat.py b/src/graph2mat/core/modules/graph2mat.py index 808ff46..ea9aa55 100644 --- a/src/graph2mat/core/modules/graph2mat.py +++ b/src/graph2mat/core/modules/graph2mat.py @@ -17,6 +17,7 @@ ) import numpy as np +import torch from ..data import BasisMatrixData, BasisTableWithEdges from ..data.basis import PointBasis @@ -272,11 +273,14 @@ def __call__(self, **kwargs): interactions: Dict[Tuple[int, int], MatrixBlock] #: The basis table used internally by graph2mat - graph2mat_table: List[PointBasis] + graph2mat_table_row: List[PointBasis] + graph2mat_table_col: List[PointBasis] #: The mapping of types from the original basis to the graph2mat basis. types_to_graph2mat: ArrayType #: The mapping of edge types from the original basis to the graph2mat basis. edge_types_to_graph2mat: ArrayType + # If the matrix is square (same basis for rows and columns) + is_square: bool #: If the ``basis_grouping`` is "max", this is a mask that is used to select #: the values for the original basis from the new grouped basis. This has #: shape (n_point_types, dim_new_basis). @@ -319,11 +323,21 @@ def __init__( self_blocks_symmetry = blocks_symmetry self.symmetric = symmetric - self.basis_table = ( - unique_basis - if isinstance(unique_basis, BasisTableWithEdges) - else BasisTableWithEdges(unique_basis) - ) + + # Asses symmetry: if symmetric, it means that all the unique basis have matrix_role=None, + # and the matrix is square. If not symmetric, we can have different basis for rows and columns. + + if not isinstance(unique_basis, BasisTableWithEdges): + unique_basis = BasisTableWithEdges(unique_basis) + + if not unique_basis.is_square and symmetric: + raise ValueError( + f"Asked for symmetric={symmetric}, but the matrix is not square. The basis has matrix_role={[(b.type, b.matrix_role) for b in (unique_basis.row.basis)]}\ +(for rows) and matrix_role={[(b.type, b.matrix_role) for b in (unique_basis.col.basis)]} (for columns).\ +This is inconsistent. If symmetric, all basis must have matrix_role=None. If not symmetric, basis must have matrix_role='row' or 'col'." + ) + self.is_square = unique_basis.is_square + self.basis_table = unique_basis self._matrix_block_cls = matrix_block_cls self.numpy = numpy if numpy is not None else np self._self_interactions_list = self_interactions_list @@ -369,28 +383,35 @@ def __init__( ) self.interactions = self._interactions_dict(interactions) + def _init_center_types(self, basis_grouping): self.basis_grouping = basis_grouping + # TODO: cambiar -- hacerlo para rows y cols # Do the grouping + # SN: Done separately for rows and columns because the basis + # can be different for rows and columns. ( - self.graph2mat_table, + graph2mat_tables, self.types_to_graph2mat, self.edge_types_to_graph2mat, self.basis_filters, ) = self.basis_table.group(self.basis_grouping) + self.graph2mat_table_rowb = graph2mat_tables.row_basis + self.graph2mat_table_colb = graph2mat_tables.col_basis + # Prepare the filters to mask the output of the operations # Currently self.basis_filters is only not None when # basis_grouping is "max". Otherwise, we don't need to apply - # any mask because the + # any mask if self.basis_filters is not None: original_edgetypes = self.basis_table.edge_type_to_point_types - self.node_filters = np.einsum( + self.node_filters_row = np.einsum( "ia, ib ->iab", self.basis_filters, self.basis_filters ) - self.edge_filters = np.einsum( + self.edge_filters_row = np.einsum( "ia, ib ->iab", self.basis_filters[original_edgetypes[:, 0]], self.basis_filters[original_edgetypes[:, 1]], @@ -402,15 +423,17 @@ def _init_center_types(self, basis_grouping): def _init_self_interactions(self, **kwargs) -> List[MatrixBlock]: self_interactions = [] - for point_type_basis in self.graph2mat_table.basis: - if len(point_type_basis.basis) == 0: + for i_point_type_basis in range(len(self.graph2mat_table_rowb)): + row_basis = self.graph2mat_table_rowb[i_point_type_basis] + col_basis = self.graph2mat_table_colb[i_point_type_basis] + if len(row_basis.basis) == 0 or len(col_basis.basis) == 0: # The point type has no basis functions self_interactions.append(None) else: self_interactions.append( self._matrix_block_cls( - i_basis=point_type_basis, - j_basis=point_type_basis, + i_basis=row_basis, + j_basis=col_basis, **kwargs, ) ) @@ -419,7 +442,7 @@ def _init_self_interactions(self, **kwargs) -> List[MatrixBlock]: def _init_interactions(self, **kwargs) -> Dict[Tuple[int, int], MatrixBlock]: point_type_combinations = itertools.combinations_with_replacement( - range(len(self.graph2mat_table.basis)), 2 + range(len(self.graph2mat_table_rowb)), 2 # row and col basis have the same length, so we can use either one ) interactions = {} @@ -433,8 +456,8 @@ def _init_interactions(self, **kwargs) -> Dict[Tuple[int, int], MatrixBlock]: perms.append((-edge_type, neigh_type, point_type)) for signed_edge_type, point_i, point_j in perms: - i_basis = self.graph2mat_table.basis[point_i] - j_basis = self.graph2mat_table.basis[point_j] + i_basis = self.graph2mat_table_rowb[point_i] + j_basis = self.graph2mat_table_colb[point_j] if len(i_basis.basis) == 0 or len(j_basis.basis) == 0: # One of the involved point types has no basis functions @@ -448,7 +471,6 @@ def _init_interactions(self, **kwargs) -> Dict[Tuple[int, int], MatrixBlock]: symm_transpose=(self.symmetric and neigh_type == point_type), **kwargs, ) - return {str(k): v for k, v in interactions.items()} def _get_preprocessing_nodes_summary(self) -> str: @@ -487,47 +509,66 @@ def summary(self) -> str: It is better than the pytorch repr to understand the high level architecture of the module, but it is not as detailed. """ + try: + s = "" - s = "" - - s += f"Preprocessing nodes: {self._get_preprocessing_nodes_summary()}\n" + s += f"Preprocessing nodes: {self._get_preprocessing_nodes_summary()}\n" - s += f"Preprocessing edges: {self._get_preprocessing_edges_summary()}\n" + s += f"Preprocessing edges: {self._get_preprocessing_edges_summary()}\n" - s += "Node operations:" - for i, x in enumerate(self.self_interactions): - point = self.graph2mat_table.basis[i] - - if x is None: - s += f"\n ({point.type}) No basis functions." - continue - - s += f"\n ({point.type}) " - - if x.symm_transpose: - s += " [XY = YX.T]" - - s += f" {self._get_node_operation_summary(x)}" + s += "Node operations:" + for i, x in enumerate(self.self_interactions): + if x.symm_transpose and not self.is_square: + raise ValueError( + f"Node operation {i} is symmetric transpose, but the matrix is NOT square. This is not allowed." + ) + point_r = self.graph2mat_table_rowb[i] + point_c = self.graph2mat_table_colb[i] + + if x is None: + s += f"\n ({point_r.type}, {point_c.type}) No basis functions." + continue + if self.is_square: + s += f"\n ({point_r.type}) " + else: + s += f"\n ({point_r.type}r, {point_c.type}c) " - s += "\nEdge operations:" - for k, x in self.interactions.items(): - point_type, neigh_type, edge_type = map(int, k[1:-1].split(",")) + if x.symm_transpose: + s += " [XY = YX.T]" - point = self.graph2mat_table.basis[point_type] - neigh = self.graph2mat_table.basis[neigh_type] + s += f" {self._get_node_operation_summary(x)}" - if x is None: - s += f"\n ({point.type}, {neigh.type}) No basis functions." - continue + s += "\nEdge operations:" + + for k, x in self.interactions.items(): + if x.symm_transpose and not self.is_square: + raise ValueError( + f"Edge operation {k} is symmetric transpose, but the matrix is NOT square. This is not allowed." + ) + + point_type, neigh_type, edge_type = map(int, k[1:-1].split(",")) - s += f"\n ({point.type}, {neigh.type})" + point = self.graph2mat_table_rowb[point_type] + neigh = self.graph2mat_table_colb[neigh_type] - if x.symm_transpose: - s += " [XY = YX.T]" + if x is None: + s += f"\n ({point.type}, {neigh.type}) No basis functions." + continue + if self.is_square: + s += f"\n ({point.type}, {neigh.type})" + else: + # SN: the pairs should contain all the combinations + # of the edge types, so we just have to add the specification of row/col + # TODO: review this in nonsym case + s += f"\n ({point.type}r, {neigh.type}c)" - s += f" {self._get_edge_operation_summary(x)}." + if x.symm_transpose: + s += " [XY = YX.T]" - return s + s += f" {self._get_edge_operation_summary(x)}." + return s + except Exception as e: + return f"Error generating summary: {e}" def forward( self, @@ -828,10 +869,10 @@ def _forward_interactions( # Get the indices that will resort the edge outputs to produce # the target. (i.e. go back to the order the edges came in). + sort_indices = self._get_edgelabels_resort_index( graph2mat_edge_types, original_types=edge_types ) - # Do the resorting and return the result. return unsorted_edge_labels[sort_indices] @@ -858,7 +899,8 @@ def _get_nodelabels_resort_index( return self._get_labels_resort_index( types=types, original_types=original_types, - shapes=self.graph2mat_table.point_block_shape, + shapes=self.basis_table.point_block_shape, # SN: changed: the full info is in basis_table, graph2mat table has row and col. + shapes_inv=self.basis_table.point_block_shape, filters=self.node_filters, # original_sizes=self.basis_table.point_block_size, transpose_neg=False, @@ -892,7 +934,8 @@ def _get_edgelabels_resort_index( return self._get_labels_resort_index( types=types, original_types=original_types, - shapes=self.graph2mat_table.edge_block_shape, + shapes=self.basis_table.edge_block_shape, + shapes_inv=self.basis_table.edge_block_shape_inv, filters=self.edge_filters, transpose_neg=self.symmetric and self.basis_grouping == "basis_shape", **kwargs, @@ -902,6 +945,7 @@ def _get_labels_resort_index( self, types: np.ndarray, shapes: np.ndarray, + shapes_inv: np.ndarray, original_types: ArrayType, filters: ArrayType, transpose_neg: bool = False, @@ -948,6 +992,7 @@ def _get_labels_resort_index( indices = get_labels_resorting_array( types, shapes=shapes.astype(types.dtype), + shapes_inv=shapes_inv.astype(types.dtype), transpose_neg=transpose_neg, **kwargs, )