From a7699f4dd8dbbb8635994ce9e10d60fc4c755cd0 Mon Sep 17 00:00:00 2001 From: Saru local ICN2 Date: Mon, 20 Jul 2026 10:50:20 +0200 Subject: [PATCH 1/3] Added all necesary to handle non-square matriced from branch contractions --- .../e3nn/modules/tests/test_e3nngraph2mat.py | 1 + src/graph2mat/bindings/torch/data/formats.py | 8 +- src/graph2mat/core/data/__init__.py | 2 +- src/graph2mat/core/data/basis.py | 12 + .../core/data/matrices/basis_matrix.py | 2 +- src/graph2mat/core/data/neighborhood.py | 3 + src/graph2mat/core/data/processing.py | 74 +++- src/graph2mat/core/data/sparse.py | 114 +++-- src/graph2mat/core/data/table.py | 397 +++++++++++++++--- src/graph2mat/core/modules/_labels_resort.py | 60 ++- src/graph2mat/core/modules/graph2mat.py | 220 +++++++--- 11 files changed, 740 insertions(+), 153 deletions(-) diff --git a/src/graph2mat/bindings/e3nn/modules/tests/test_e3nngraph2mat.py b/src/graph2mat/bindings/e3nn/modules/tests/test_e3nngraph2mat.py index e7ee7d6..6c1ce30 100644 --- a/src/graph2mat/bindings/e3nn/modules/tests/test_e3nngraph2mat.py +++ b/src/graph2mat/bindings/e3nn/modules/tests/test_e3nngraph2mat.py @@ -8,6 +8,7 @@ from graph2mat import ( BasisConfiguration, BasisTableWithEdges, + BasisTableWithEdges, MatrixDataProcessor, PointBasis, conversions, diff --git a/src/graph2mat/bindings/torch/data/formats.py b/src/graph2mat/bindings/torch/data/formats.py index fba58f1..4135cb7 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/__init__.py b/src/graph2mat/core/data/__init__.py index af17e1a..b812325 100644 --- a/src/graph2mat/core/data/__init__.py +++ b/src/graph2mat/core/data/__init__.py @@ -27,5 +27,5 @@ from .metrics import OrbitalMatrixMetric from . import metrics from .processing import MatrixDataProcessor, BasisMatrixData, BasisMatrixDataBase -from .table import BasisTableWithEdges, AtomicTableWithEdges +from .table import BasisTableWithEdges, AtomicTableWithEdges, BasisTableWithEdges_rowcol # BORRAR el ultimo from .formats import Formats, conversions, ConversionManager diff --git a/src/graph2mat/core/data/basis.py b/src/graph2mat/core/data/basis.py index 560100e..caf8e85 100644 --- a/src/graph2mat/core/data/basis.py +++ b/src/graph2mat/core/data/basis.py @@ -149,6 +149,10 @@ class PointBasis: with optional minus signs, e.g. ``"xyz"``, ``"-yz-x"``, ``"z-x-y"``. You can also use the aliases that we provide, such as ``"cartesian"`` or ``"spherical"``. + matrix_role: + role that the specified basis plays in the output matrix. It can be 'row' + or 'col'. If None, assumes the matrix is square and the basis is used for + both rows and columns. Examples @@ -175,12 +179,20 @@ class PointBasis: R: Union[float, np.ndarray] basis: Union[str, Sequence[Union[int, Tuple[int, int, int]]]] = () basis_convention: BasisConvention = "spherical" + matrix_role: Union[Literal["row", "col"], None] = None def __post_init__(self): basis = self._sanitize_basis(self.basis) object.__setattr__(self, "basis", basis) + if self.matrix_role is not None: + assert self.matrix_role in ["row", "col"], "matrix_role must be 'row', 'col' or None." + print( + f"Defined {self.matrix_role} for type {self.type} with basis {self.basis} and reach {self.R}." + ) + object.__setattr__(self, "matrix_role", self.matrix_role) + assert isinstance(self.R, Number) or ( isinstance(self.R, np.ndarray) and len(self.R) == self.basis_size ), f"R must be a float or an array of length {self.basis_size} (the number of functions)." diff --git a/src/graph2mat/core/data/matrices/basis_matrix.py b/src/graph2mat/core/data/matrices/basis_matrix.py index c71f82e..96c91b6 100644 --- a/src/graph2mat/core/data/matrices/basis_matrix.py +++ b/src/graph2mat/core/data/matrices/basis_matrix.py @@ -5,7 +5,7 @@ from dataclasses import dataclass import numpy as np -from ..table import BasisTableWithEdges +from ..table import BasisTableWithEdges, BasisTableWithEdges BasisCount = np.ndarray # [num_points] diff --git a/src/graph2mat/core/data/neighborhood.py b/src/graph2mat/core/data/neighborhood.py index 185c754..df1634b 100644 --- a/src/graph2mat/core/data/neighborhood.py +++ b/src/graph2mat/core/data/neighborhood.py @@ -29,6 +29,9 @@ def get_neighborhood( assert len(pbc) == 3 and all(isinstance(i, (bool, np.bool_)) for i in pbc) assert cell.shape == (3, 3) + # BORRAR + # print('Cutoff:', cutoff) + sender, receiver, unit_shifts = ase.neighborlist.primitive_neighbor_list( quantities="ijS", pbc=pbc, diff --git a/src/graph2mat/core/data/processing.py b/src/graph2mat/core/data/processing.py index 3e37b77..7097882 100644 --- a/src/graph2mat/core/data/processing.py +++ b/src/graph2mat/core/data/processing.py @@ -198,6 +198,10 @@ def matrix_from_data( if is_batch is None: is_batch = isinstance(data, Batch) + # BORRAR + # print("In processing.py matrix_from_data:") + # print("data:", data) + # print("is_batch:", is_batch) if is_batch: return tuple( self.yield_from_batch( @@ -268,6 +272,12 @@ def yield_from_batch( # Types for both atoms and edges. point_types = arrays.point_types edge_types = arrays.edge_types + + # BORRAR + # print("In MatrixDataProcessor.yield_from_batch:") + # print("arrays=data.numpy_arrays():", arrays) + # print("atom_ptr:", atom_ptr) + # print("edge_ptr:", edge_ptr) # Get the values for the node blocks and the pointer to the start of each block. node_labels_ptr = self.basis_table.point_block_pointer(point_types) @@ -298,6 +308,17 @@ def yield_from_batch( edge_labels_ptr[edge_start] : edge_labels_ptr[edge_end] ] + # BORRAR + # print(f"example {i} (batch):") + # print(" atom_start:", atom_start) + # print(" atom_end:", atom_end) + # print(" edge_start:", edge_start) + # print(" edge_end:", edge_end) + # print(" new_edge_label = edge_labels[edge_labels_ptr[edge_start]: edge_labels_ptr[edge_end]]:") + # print(f" this takes the edge labels from edge_labels_ptr[edge_start] = edge_labels_ptr[{edge_start}] = {edge_labels_ptr[edge_start]}") + # print(f" to edge_labels_ptr[edge_end] = edge_labels_ptr[{edge_end}] = {edge_labels_ptr[edge_end]}") + # print(f" len(new_edge_label) = {len(new_edge_label)}") + if getattr(example, "point_labels", None) is not None: assert len(new_atom_label) == len(example.point_labels) if getattr(example, "edge_labels", None) is not None: @@ -510,6 +531,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 --------- @@ -520,6 +544,12 @@ def sort_edge_index( # Get the supercell index of the neighbor in each interaction isc = isc_off[sc_shifts[0], sc_shifts[1], sc_shifts[2]] + # BORRAR + # print("In sort_edge_index:") + # print("isc_off:", isc_off) + # print("sc_shifts:", sc_shifts) + # print("isc:", isc) + # Find unique edges: # - For edges that are between different supercells: We get just connections from # the unit cell to half of the supercells. One can then reproduce all the connections @@ -651,8 +681,15 @@ def get_cutoff(self, point_types: np.ndarray) -> Union[float, np.ndarray]: """ if isinstance(self.basis_table.R, float): + # BORRAR + # print('self.basis_table.R is a float:', self.basis_table.R) + return self.basis_table.R * 2 else: + # BORRAR + # print('self.basis_table.R is an array:', self.basis_table.R) + # print('point_types:', point_types) + # print('self.basis_table.R[point_types]:', self.basis_table.R[point_types]) return self.basis_table.R[point_types] def get_nlabels_per_edge_type(self, edge_types: np.ndarray) -> np.ndarray: @@ -804,8 +841,11 @@ 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] + # SN: changed this -- the n_orbitals can be diff in cols and rows if its non square + 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: @@ -822,6 +862,16 @@ def labels_to( edge_index = edge_index[:, ::2] edge_types = edge_types[::2] neigh_isc = neigh_isc[::2] + # BORRAR + # print("In labels_to: of MatrixDataProcessor") + # print("data_format:", data_format) + # print("out_format:", out_format) + # print("len node_labels:", len(node_labels)) + # print("len edge_labels:", len(edge_labels)) + # print("edge_index:", edge_index) + # print("neigh_isc:", neigh_isc) + # print("self.symmetric_matrix:", self.symmetric_matrix) + # print("calling conversions.get_converter with data_format:", data_format, "and out_format:", out_format) # Construct the matrix. matrix = conversions.get_converter(data_format, out_format)( @@ -1354,15 +1404,28 @@ 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) + # BORRAR + # print('IN BasisMatrixData.from_config:') + # print('all indices:', indices) + # print('edge_index:', edge_index) # Get the edge types edge_types = data_processor.basis_table.point_type_to_edge_type( indices[edge_index] ) + # BORRAR + # print('In BasisMatrixData.from_config 1:') + # print('edge_index:', edge_index) + # print('edge_types:', edge_types) + # 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, ) + # BORRAR + # print('In BasisMatrixData.from_config 2:') + # print('edge_index:', edge_index) + # print('edge_types:', edge_types) # Then, get the supercell index of each interaction. neigh_isc = supercell.isc_off[sc_shifts[0], sc_shifts[1], sc_shifts[2]] @@ -1372,6 +1435,11 @@ def from_config( config, point_types=indices, edge_index=edge_index, neigh_isc=neigh_isc ) + # BORRAR + # print('In BasisMatrixData.from_config 3:') + # print('edge_index:', edge_index) + # print('edge_types:', edge_types) + return cls( edge_index=edge_index, neigh_isc=neigh_isc, diff --git a/src/graph2mat/core/data/sparse.py b/src/graph2mat/core/data/sparse.py index 057c613..aa411bd 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,12 +183,16 @@ 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 ONLY contain the edges in one of the directions. + # TODO: verify this work with non sym in the non-square case. n_supercells: number of supercells in the matrix. edge_neigh_isc: @@ -203,16 +208,29 @@ def _blockmatrix_coo_coords( rows = [] cols = [] + # BORRAR + # print(f"In sparse.py _blockmatrix_coo_coords:") + # print(f"orbitals_row: {orbitals_row}") + # print(f"orbitals_col: {orbitals_col}") + # print(f"edge_index: {edge_index}") + + 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 +250,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 +275,32 @@ 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( + "SN: TODO - I think thsi not work as I intended to... maybe it cannot be") + # We assume that the edge_index contains only the edges in one direction, + # but as it it not square, we cannot assume that the opposite direction is just the transpose of the + # original edge. + i_start_symm = first_orb_row[j_at] + i_end_symm = i_start_symm + orbitals_row[j_at] + j_start_symm = first_orb_col[i_at] + j_end_symm = j_start_symm + orbitals_col[i_at] + block_rows_symm, block_cols_symm = np.mgrid[i_start_symm:i_end_symm, j_start_symm:j_end_symm].reshape(2, -1) + rows_symm.extend(block_rows_symm) + cols_symm.extend(block_cols_symm) # 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) + # BORRAR + # print(f"Len rows in _blockmatrix_coo_coords: {len(rows)}") + # print(f"Len cols in _blockmatrix_coo_coords: {len(cols)}") + + return np.array(rows), np.array(cols), (no_row, no_col * n_supercells) def _nodes_and_edges_to_coo( @@ -271,7 +309,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 +345,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 @@ -321,18 +362,37 @@ def _nodes_and_edges_to_coo( opposite direction is then created as the transpose. """ + # BORRAR + # print(f"In sparse.py _nodes_and_edges_to_coo:") + # print(f"calling _blockmatrix_coo_coords with:") + # print(f"orbitals_row: {orbitals_row}") + # print(f"orbitals_col: {orbitals_col}") + # print(f"edge_index: {edge_index}") + # print(f"n_supercells: {n_supercells}") + # print(f"edge_neigh_isc: {edge_neigh_isc}") + # print(f"symmetrize_edges: {symmetrize_edges}") + 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, ) + # BORRAR + # print(f"Len rows in _nodes_and_edges_to_coo: {len(rows)}") + # print(f"Len cols in _nodes_and_edges_to_coo: {len(cols)}") + # print(f"Shape in _nodes_and_edges_to_coo: {shape}") + # print(f"Len node_vals in _nodes_and_edges_to_coo: {len(node_vals)}") + # print(f"Len edge_vals in _nodes_and_edges_to_coo: {len(edge_vals)}") if symmetrize_edges: sparse_data = concatenate([node_vals, edge_vals, edge_vals]) else: sparse_data = concatenate([node_vals, edge_vals]) + # BORRAR + # print(f"Len sparse_data in _nodes_and_edges_to_coo: {len(sparse_data)}") if threshold is not None: mask = abs(sparse_data) > threshold @@ -352,7 +412,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 +434,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 +460,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..0031ea5 100644 --- a/src/graph2mat/core/data/table.py +++ b/src/graph2mat/core/data/table.py @@ -15,7 +15,7 @@ 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 import numpy as np import sisl @@ -23,7 +23,7 @@ from .basis import BasisConvention, PointBasis, get_change_of_basis -class BasisTableWithEdges: +class BasisTableWithEdges_rowcol: """Stores the unique types of points in the system, with their basis and the possible edges. It also knows the size of the blocks, and other type dependent variables. @@ -137,31 +137,9 @@ def __init__( self.basis_convention = basis_convention - # For the basis convention, get the matrices to change from cartesian to our convention. - self.change_of_basis, self.change_of_basis_inv = get_change_of_basis( - "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 - for i in range(n_types): - # 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): - point_types_to_edge_types[i, j] = edge_type - point_types_to_edge_types[j, i] = -edge_type - edge_type += 1 - - self.edge_type = point_types_to_edge_types - # 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. + # TODO: pass this to the non-square(to Sara_) if get_point_matrix is None: self.point_matrix = None else: @@ -177,16 +155,16 @@ def __init__( [basis.basis_size for basis in self.basis], dtype=np.int32 ) - # 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 + # BORRAR + # print(f"Basis sizes: {self.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) + # SN: MOVED THIS TO BasisTableWithEdges BUT I AM NOT SURE OF HOW IT WORKS + # 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})" @@ -246,7 +224,7 @@ def __eq__(self, other): def group( self, grouping: Literal["basis_shape", "point_type", "max"] - ) -> tuple["BasisTableWithEdges", np.ndarray, np.ndarray, Optional[np.ndarray]]: + ) -> tuple["BasisTableWithEdges_rowcol", np.ndarray, np.ndarray, Optional[np.ndarray]]: r"""Groups the basis in this table and creates a new table. It also returns useful objects to convert between the ungrouped @@ -312,12 +290,17 @@ def group( if grouping == "point_type": new_table = self - class A: - def __getitem__(self, key): - return key + # class A: + # def __getitem__(self, key): + # return key + # def __eq__(self, other): + # return isinstance(other, A) + + # point_type_conversion = A() + # edge_type_conversion = A() + point_type_conversion = IdentityConversion() + edge_type_conversion = IdentityConversion() - point_type_conversion = A() - edge_type_conversion = A() elif grouping == "basis_shape": # Get all basis sizes: basis_sizes = np.zeros((len(self.basis), 5), dtype=int) @@ -405,6 +388,284 @@ def index_to_type(self, index: int) -> Union[str, int]: """ return self.types[index] + def maxR(self) -> float: + """Maximum cutoff radius in the basis.""" + return self.R.max() + + 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] + + +class BasisTableWithEdges: + """Storing point information accounting for different row and column point types.""" + def __init__(self, basis: Sequence[PointBasis], get_point_matrix: Optional[Callable] = None): + self.is_square = all(b.matrix_role is None for b in basis) + + # BORRAR + # print(f"BasisTableWithEdges: is_square = {self.is_square}") + # print(f"BasisTableWithEdges: all matrix roles = {[b.matrix_role for b in basis]}") + + row_basis, col_basis = self._process_basis_rows_cols(basis) + self.row = (BasisTableWithEdges_rowcol(row_basis, get_point_matrix) + if len(row_basis) > 0 else None) + self.col = (BasisTableWithEdges_rowcol(col_basis, get_point_matrix) + if len(col_basis) > 0 else None) + + if (self.row is None)^(self.col is None): + raise ValueError("If one of row or column basis is None, the other must be None too.") + + # Check, in the case of non square matrix, properties match as they should + if not self.is_square: + self.check_nonsquare_properties() + + self.basis_convention = self.row.basis_convention if self.row is not None else None + + # For the basis convention, get the matrices to change from cartesian to our convention. + self.change_of_basis, self.change_of_basis_inv = get_change_of_basis( + "cartesian", self.basis_convention + ) + + self.types = self.row.types if self.row is not None else None + + n_types = len(self.types) + + if self.is_square: + self.R = self.row.R + else: + # Each row and column point type can have a different cutoff radius. + # Warn if they are different, and take the maximum of the two. + if not np.allclose(self.row.R, self.col.R): + print(f"Warning: Row and column point types have different cutoff radii. Taking the maximum of the two.") + # BORRAR + # print(f"Row cutoff radii: {self.row.R}") + # print(f"Col cutoff radii: {self.col.R}") + # print(f"Max cutoff radius: {np.max([self.row.R, self.col.R], axis=0)}") + + self.R = np.max([self.row.R, self.col.R], axis=0) + + # 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 + for i in range(n_types): + # 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): + point_types_to_edge_types[i, j] = edge_type + point_types_to_edge_types[j, i] = -edge_type + edge_type += 1 + + # 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. + + self.edge_type = point_types_to_edge_types + + + # And also the sizes of the blocks. + # SN: Here we have all possible blocks + self.point_block_shape = np.array([self.row.basis_size, self.col.basis_size]) + self.point_block_size = self.row.basis_size * self.col.basis_size + + # BORRAR + # print("In BasisTableWithEdges: ") + # print(f"self.edge_type == point_types_to_edge_types:\n{self.edge_type}") + # print(f"self.point_block_shape:\n{self.point_block_shape}") + # print(f"self.point_block_size:\n{self.point_block_size}") + + point_types_combinations = np.array( + list(itertools.combinations_with_replacement(range(n_types), 2))).T + # Even if its not sqare, we take each direction once : we define the inverse shape, where we pass from (Ar, Bc) to (Ac, Br) + if self.is_square: + + + self.edge_type_to_point_types = point_types_combinations.T + self.edge_block_shape = self.row.basis_size[point_types_combinations] + self.edge_block_shape_inv = self.row.basis_size[point_types_combinations] + self.edge_block_size = self.edge_block_shape.prod(axis=0) + self.edge_block_size_inv = self.edge_block_shape_inv.prod(axis=0) + + # BORRAR + # print(f"Point type to edge type:\n{self.edge_type}") + # print(f"Edge type to point types:\n{self.edge_type_to_point_types}") + # print(f"Edge block shape:\n{self.edge_block_shape}") + # print(f"Edge block size:\n{self.edge_block_size}") + + else: + # Store the sizes of each point's basis — now separate for rows and cols. + row_basis_size = self.row.basis_size + col_basis_size = self.col.basis_size + + # BORRAR + # print(f"Row basis sizes: {row_basis_size}") + # print(f"Col basis sizes: {col_basis_size}") + + self.edge_type_to_point_types = point_types_combinations.T + + # TODO: comprobe if this is okay for non-square matrices + # BORRAR + # print(f"Point type to edge type:\n{self.edge_type}") + # print(f"Edge type to point types:\n{self.edge_type_to_point_types}") + + 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 + + # 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) + self.edge_block_shape_inv = np.array([ + row_basis_size[col_type_indices], + col_basis_size[row_type_indices], + ]) # shape: (2, n_combos) + + # BORRAR + # print(f"Edge block shape:\n{self.edge_block_shape}") + # print(f"Edge block shape inv:\n{self.edge_block_shape_inv}") + + 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,) + + @staticmethod + def _process_basis_rows_cols(basis: Union[List[PointBasis], BasisTableWithEdges_rowcol]) \ + -> Tuple[List[PointBasis], List[PointBasis]]: + """Processes the basis to determine which basis is used for rows and which for columns. + + 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. + + Parameters + ---------- + basis: List[PointBasis] + The list of point bases. + + Returns + ------- + i_basis: List[PointBasis] + The list of point bases for rows. + j_basis: List[PointBasis] + The list of point bases for columns. + """ + if isinstance(basis, BasisTableWithEdges_rowcol): + basis = basis.basis + i_basis = [] + j_basis = [] + for b in basis: + if b.matrix_role == "row": + i_basis.append(b) + elif b.matrix_role == "col": + j_basis.append(b) + else: + # If None, it is a square matrix: same basis for rows and columns + i_basis.append(b) + j_basis.append(b) + 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 check_nonsquare_properties(self): + + basis_convention_row = self.row.basis_convention if self.row is not None else None + basis_convention_col = self.col.basis_convention if self.col is not None else None + + assert basis_convention_row == basis_convention_col, \ + f"Row and column basis conventions must be the same. \ +Got {basis_convention_row} and {basis_convention_col}." + + types_row = self.row.types if self.row is not None else None + types_col = self.col.types if self.col is not None else None + + assert types_row == types_col, \ + f"Row and column basis types must be the same. \ +Got {types_row} and {types_col}." + + def group( + self, grouping: Literal["basis_shape", "point_type", "max"] + ) -> tuple["BasisTableWithEdges_rowcol", np.ndarray, np.ndarray, Optional[np.ndarray]]: + + "TODO: define this for nonsquare in this scheme" + + # SN: changed the output: returns a tuple of two tables, one for rows and one for columns + if self.is_square: + new_table, point_type_conversion, edge_type_conversion, filters = self.row.group(grouping) + return (new_table, new_table), point_type_conversion, edge_type_conversion, filters + else: + if grouping == "point_type": + new_table_row, point_type_conversion_row, edge_type_conversion_row, filters_row = self.row.group(grouping) + new_table_col, point_type_conversion_col, edge_type_conversion_col, filters_col = self.col.group(grouping) + if point_type_conversion_row != point_type_conversion_col: + raise ValueError(f"Row and column point type conversions must be the same for non-square matrices. Got {point_type_conversion_row} and {point_type_conversion_col}.") + if edge_type_conversion_row != edge_type_conversion_col: + raise ValueError(f"Row and column edge type conversions must be the same for non-square matrices. Got {edge_type_conversion_row} and {edge_type_conversion_col}.") + # TODO: for the filters ? + return (new_table_row, new_table_col), point_type_conversion_row, edge_type_conversion_row, filters_row + else: + raise NotImplementedError(f"Grouping for non-square matrices \ +with grouping {grouping} is not implemented yet.") + + def _repr_html_(self): + if self.is_square: + return self.row._repr_html_() + else: + # TODO: comprobar + 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 __str__(self): + if self.is_square: + return str(self.row) + else: + return f"Row basis:\n{self.row}\nColumn basis:\n{self.col}" + + def __len__(self): + if self.is_square: + return len(self.row) + else: + if len(self.col) != len(self.row): + raise ValueError("Row and column basis have different lengths.") + return len(self.row) + + def __eq__(self, other): + if not isinstance(other, self.__class__): + return False + same = self.is_square == other.is_square + same &= self.row == other.row + same &= self.col == other.col + return same + + # SN: moved from the old BasisTableWithEdges_rowcol class to here 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. @@ -415,6 +676,7 @@ def type_to_index(self, point_type: Union[str, int]) -> int: """ return self.types.index(point_type) + # SN: moved from the old BasisTableWithEdges_rowcol class to here def types_to_indices(self, types: Sequence) -> np.ndarray: """Converts from an array of types IDs to their indices in the basis table. @@ -438,6 +700,7 @@ def types_to_indices(self, types: Sequence) -> np.ndarray: # And reconstruct the original array, which is now an array of indices instead of types return unique_indices[inverse_indices] + # SN: moved from the old BasisTableWithEdges_rowcol class to here def point_type_to_edge_type(self, point_type: np.ndarray) -> Union[int, np.ndarray]: """Converts pairs of point types to edge types. @@ -449,10 +712,7 @@ 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() - + # SN: moved from the old BasisTableWithEdges_rowcol class to here def point_block_pointer(self, point_types: Sequence[int]) -> np.ndarray: """Pointers to the beggining of node blocks in a flattened matrix. @@ -470,8 +730,15 @@ def point_block_pointer(self, point_types: Sequence[int]) -> np.ndarray: """ pointers = np.zeros(len(point_types) + 1, dtype=np.int32) np.cumsum(self.point_block_size[point_types], out=pointers[1:]) + # BORRAR + # print(f"In BasisTableWithEdges.point_block_pointer:") + # print(f" point_types = {point_types}") + # print(f" point_block_size = {self.point_block_size}") + # print(f" pointers = {pointers}") + return pointers + # SN: moved from the old BasisTableWithEdges_rowcol class to here def edge_block_pointer(self, edge_types: Sequence[int]): """Pointers to the beggining of edge blocks in a flattened matrix. @@ -489,19 +756,30 @@ 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:]) - return pointers + 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) - def get_sisl_atoms(self) -> List[sisl.Atom]: - """Returns a list of sisl atoms corresponding to the basis. + np.cumsum(sizes, out=pointers[1:]) - 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] + # BORRAR + # print(f"In BasisTableWithEdges.edge_block_pointer:") + # print(f" edge_types = {edge_types}") + # print(f" edge_block_size = {self.edge_block_size}") + # print(f" pointers = {pointers}") + + return pointers + + # SN: the max R must be the max among rows and cols + def maxR(self) -> float: + """Maximum cutoff radius in the basis.""" + return max(self.row.maxR(), self.col.maxR()) class AtomicTableWithEdges(BasisTableWithEdges): @@ -666,3 +944,12 @@ 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) \ No newline at end of file diff --git a/src/graph2mat/core/modules/_labels_resort.py b/src/graph2mat/core/modules/_labels_resort.py index bcf945e..a28589b 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,28 +68,59 @@ 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] # 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 + # BORRAR + # print(f"Type {type}: offset[prev_type] = {offset[prev_type]}, type_nlabels[prev_type] = {type_nlabels[prev_type]}") + + offset[type + ntypes_int - 1] = offset[prev_type] + type_nlabels[prev_type] # >0 + # BORRAR + # print(f"Type {-1*type}: offset[-type] = {offset[type]}, type_nlabels[type] = {type_nlabels[type]}") + offset[-type + ntypes_int - 1] = offset[type + ntypes_int - 1] + type_nlabels[type + ntypes_int - 1] # < 0 + + # BORRAR + # print(f"calculated offset[{type}] = {offset[type]}, offset[{-type}] = {offset[-type]}") + prev_type = -type + ntypes_int - 1 # we have to continue from the negative type, because the next positive type will be after it. # 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 ) + # BORRAR + # print("In get_labels_resorting_array:") + # print("sizes: ", np.asarray(sizes)) + # print("type_nlabels: ", np.asarray(type_nlabels)) + # print("offset: ", np.asarray(offset)) + # print("indices.shape: ", indices.shape) + type_i: cython.long[:] = np.zeros_like(sizes, dtype=int) i: cython.int = 0 @@ -95,8 +128,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 +138,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..c3954e1 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_row = graph2mat_tables[0] + self.graph2mat_table_col = graph2mat_tables[1] + # 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,24 +423,28 @@ 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_row.basis)): + row_basis = self.graph2mat_table_row.basis[i_point_type_basis] + col_basis = self.graph2mat_table_col.basis[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, ) ) + # BORRAR + # print('self_interactions: \n', self_interactions) return self_interactions 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_row.basis)), 2 # row and col basis have the same length, so we can use either one ) interactions = {} @@ -433,8 +458,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_row.basis[point_i] + j_basis = self.graph2mat_table_col.basis[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 +473,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 +511,90 @@ 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 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}) " + s += f"Preprocessing nodes: {self._get_preprocessing_nodes_summary()}\n" - if x.symm_transpose: - s += " [XY = YX.T]" + s += f"Preprocessing edges: {self._get_preprocessing_edges_summary()}\n" - 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_row.basis[i] + point_c = self.graph2mat_table_col.basis[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:" - s += f"\n ({point.type}, {neigh.type})" - - if x.symm_transpose: - s += " [XY = YX.T]" + # # BORRAR + # print(f"self.interactions: {self.interactions}") + + 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" {self._get_edge_operation_summary(x)}." + point = self.graph2mat_table_row.basis[point_type] + neigh = self.graph2mat_table_col.basis[neigh_type] - return s + 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)" + + if x.symm_transpose: + s += " [XY = YX.T]" + + s += f" {self._get_edge_operation_summary(x)}." + + # # BORRAR + # print(s) + + # point = self.graph2mat_table_row.basis[point_type] + # neigh = self.graph2mat_table_col.basis[neigh_type] + + # if x is None: + # print(" x is None ") + # print(f" ({point.type}, {neigh.type}) No basis functions.") + # continue + # print(" ({point.type}, {neigh.type})") + # print(f"({point.type}r, {neigh.type})") + + # if x.symm_transpose: + # print(" x.symm_transpose is True ") + # print(" [XY = YX.T]") + # print(" self._get_edge_operation_summary(x) , being x ", x) + # print(f" {self._get_edge_operation_summary(x)}.") + # # END BORRAR + # raise NotImplementedError("Printing edge operations for non-square matrices are not implemented yet.") + return s + except Exception as e: + return f"Error generating summary: {e}" def forward( self, @@ -752,6 +819,12 @@ def _forward_interactions( outputs = [] graph2mat_edge_types = self.edge_types_to_graph2mat[edge_types] + # BORRAR + # print("In Graph2Mat _forward_interactions: ") + # print("graph2mat_edge_types: ", graph2mat_edge_types) + # print("edge_types: ", edge_types) + # print("Order of interactions (edge type is last): ", list(self.interactions.keys())) + # print('This is the order of indexes: then all the indexes of same type se cogen seguidos') # Call each unique interaction function with only the features # of edges that correspond to that type. @@ -828,10 +901,26 @@ 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). + + # BORRAR + # print("In Graph2Mat _forward_interactions (before get_edgelabels_resort_index): ") + # print("graph2mat_edge_types: ", graph2mat_edge_types) + # print("edge_types (original_types): ", edge_types) + sort_indices = self._get_edgelabels_resort_index( graph2mat_edge_types, original_types=edge_types ) + # BORRAR + # print("In Graph2Mat _forward_interactions (after get_edgelabels_resort_index): ") + # print("graph2mat_edge_types: ", graph2mat_edge_types) + # print("edge_types (original_types): ", edge_types) + # print("unsorted_edge_labels: ", unsorted_edge_labels) + # print("sort_indices: ", sort_indices) + # print("unsorted_edge_labels[sort_indices]: ", unsorted_edge_labels[sort_indices]) + # print("unsorted_edge_labels.shape: ", unsorted_edge_labels.shape) + # print("new shape: ", unsorted_edge_labels[sort_indices].shape) + # Do the resorting and return the result. return unsorted_edge_labels[sort_indices] @@ -858,7 +947,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, @@ -889,10 +979,21 @@ def _get_edgelabels_resort_index( types = types[::2] original_types = original_types[::2] + # # BORRAR + # print("In Graph2Mat _get_edgelabels_resort_index: ") + # print("types: ", types) + # print("types.type: ", types.dtype) + # types = reorder_array(types) + + # # BORRAR + # print("after reorder_array(types): ", types) + # print("types.type: ", types.dtype) + 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 +1003,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, @@ -945,11 +1047,25 @@ def _get_labels_resort_index( return np.where(mask)[0] else: + # BORRAR + # print("In Graph2Mat _get_labels_resort_index: ") + # print('BEFORE CALLING get_labels_resorting_array') + # print("types: ", types) + # print("shapes: ", shapes) + # print("shapes_inv: ", shapes_inv) + # print("transpose_neg: ", transpose_neg) + # print("kwargs: ", kwargs) + indices = get_labels_resorting_array( types, shapes=shapes.astype(types.dtype), + shapes_inv=shapes_inv.astype(types.dtype), transpose_neg=transpose_neg, **kwargs, ) + # BORRAR + # print('AFTER CALLING get_labels_resorting_array') + # print("len(indices): ", len(indices)) + # print("indices: ", indices) return indices From 4795d45a661c4fa78ee7093a5ffd86b97beb12b8 Mon Sep 17 00:00:00 2001 From: Saru local ICN2 Date: Tue, 21 Jul 2026 11:49:14 +0200 Subject: [PATCH 2/3] Deleted all the commented prints for debugging --- src/graph2mat/core/data/__init__.py | 2 +- src/graph2mat/core/data/neighborhood.py | 3 - src/graph2mat/core/data/processing.py | 64 +----------------- src/graph2mat/core/data/sparse.py | 31 --------- src/graph2mat/core/data/table.py | 64 ++---------------- src/graph2mat/core/modules/_labels_resort.py | 18 +---- src/graph2mat/core/modules/graph2mat.py | 71 -------------------- 7 files changed, 8 insertions(+), 245 deletions(-) diff --git a/src/graph2mat/core/data/__init__.py b/src/graph2mat/core/data/__init__.py index b812325..8f4a0b2 100644 --- a/src/graph2mat/core/data/__init__.py +++ b/src/graph2mat/core/data/__init__.py @@ -27,5 +27,5 @@ from .metrics import OrbitalMatrixMetric from . import metrics from .processing import MatrixDataProcessor, BasisMatrixData, BasisMatrixDataBase -from .table import BasisTableWithEdges, AtomicTableWithEdges, BasisTableWithEdges_rowcol # BORRAR el ultimo +from .table import BasisTableWithEdges, AtomicTableWithEdges, BasisTableWithEdges_rowcol from .formats import Formats, conversions, ConversionManager diff --git a/src/graph2mat/core/data/neighborhood.py b/src/graph2mat/core/data/neighborhood.py index df1634b..185c754 100644 --- a/src/graph2mat/core/data/neighborhood.py +++ b/src/graph2mat/core/data/neighborhood.py @@ -29,9 +29,6 @@ def get_neighborhood( assert len(pbc) == 3 and all(isinstance(i, (bool, np.bool_)) for i in pbc) assert cell.shape == (3, 3) - # BORRAR - # print('Cutoff:', cutoff) - sender, receiver, unit_shifts = ase.neighborlist.primitive_neighbor_list( quantities="ijS", pbc=pbc, diff --git a/src/graph2mat/core/data/processing.py b/src/graph2mat/core/data/processing.py index 7097882..2661ac6 100644 --- a/src/graph2mat/core/data/processing.py +++ b/src/graph2mat/core/data/processing.py @@ -198,10 +198,7 @@ def matrix_from_data( if is_batch is None: is_batch = isinstance(data, Batch) - # BORRAR - # print("In processing.py matrix_from_data:") - # print("data:", data) - # print("is_batch:", is_batch) + if is_batch: return tuple( self.yield_from_batch( @@ -272,12 +269,6 @@ def yield_from_batch( # Types for both atoms and edges. point_types = arrays.point_types edge_types = arrays.edge_types - - # BORRAR - # print("In MatrixDataProcessor.yield_from_batch:") - # print("arrays=data.numpy_arrays():", arrays) - # print("atom_ptr:", atom_ptr) - # print("edge_ptr:", edge_ptr) # Get the values for the node blocks and the pointer to the start of each block. node_labels_ptr = self.basis_table.point_block_pointer(point_types) @@ -308,17 +299,6 @@ def yield_from_batch( edge_labels_ptr[edge_start] : edge_labels_ptr[edge_end] ] - # BORRAR - # print(f"example {i} (batch):") - # print(" atom_start:", atom_start) - # print(" atom_end:", atom_end) - # print(" edge_start:", edge_start) - # print(" edge_end:", edge_end) - # print(" new_edge_label = edge_labels[edge_labels_ptr[edge_start]: edge_labels_ptr[edge_end]]:") - # print(f" this takes the edge labels from edge_labels_ptr[edge_start] = edge_labels_ptr[{edge_start}] = {edge_labels_ptr[edge_start]}") - # print(f" to edge_labels_ptr[edge_end] = edge_labels_ptr[{edge_end}] = {edge_labels_ptr[edge_end]}") - # print(f" len(new_edge_label) = {len(new_edge_label)}") - if getattr(example, "point_labels", None) is not None: assert len(new_atom_label) == len(example.point_labels) if getattr(example, "edge_labels", None) is not None: @@ -544,12 +524,6 @@ def sort_edge_index( # Get the supercell index of the neighbor in each interaction isc = isc_off[sc_shifts[0], sc_shifts[1], sc_shifts[2]] - # BORRAR - # print("In sort_edge_index:") - # print("isc_off:", isc_off) - # print("sc_shifts:", sc_shifts) - # print("isc:", isc) - # Find unique edges: # - For edges that are between different supercells: We get just connections from # the unit cell to half of the supercells. One can then reproduce all the connections @@ -681,15 +655,8 @@ def get_cutoff(self, point_types: np.ndarray) -> Union[float, np.ndarray]: """ if isinstance(self.basis_table.R, float): - # BORRAR - # print('self.basis_table.R is a float:', self.basis_table.R) - return self.basis_table.R * 2 else: - # BORRAR - # print('self.basis_table.R is an array:', self.basis_table.R) - # print('point_types:', point_types) - # print('self.basis_table.R[point_types]:', self.basis_table.R[point_types]) return self.basis_table.R[point_types] def get_nlabels_per_edge_type(self, edge_types: np.ndarray) -> np.ndarray: @@ -862,16 +829,6 @@ def labels_to( edge_index = edge_index[:, ::2] edge_types = edge_types[::2] neigh_isc = neigh_isc[::2] - # BORRAR - # print("In labels_to: of MatrixDataProcessor") - # print("data_format:", data_format) - # print("out_format:", out_format) - # print("len node_labels:", len(node_labels)) - # print("len edge_labels:", len(edge_labels)) - # print("edge_index:", edge_index) - # print("neigh_isc:", neigh_isc) - # print("self.symmetric_matrix:", self.symmetric_matrix) - # print("calling conversions.get_converter with data_format:", data_format, "and out_format:", out_format) # Construct the matrix. matrix = conversions.get_converter(data_format, out_format)( @@ -1404,28 +1361,14 @@ 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) - # BORRAR - # print('IN BasisMatrixData.from_config:') - # print('all indices:', indices) - # print('edge_index:', edge_index) - # Get the edge types edge_types = data_processor.basis_table.point_type_to_edge_type( indices[edge_index] ) - # BORRAR - # print('In BasisMatrixData.from_config 1:') - # print('edge_index:', edge_index) - # print('edge_types:', edge_types) - # 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, ) - # BORRAR - # print('In BasisMatrixData.from_config 2:') - # print('edge_index:', edge_index) - # print('edge_types:', edge_types) # Then, get the supercell index of each interaction. neigh_isc = supercell.isc_off[sc_shifts[0], sc_shifts[1], sc_shifts[2]] @@ -1435,11 +1378,6 @@ def from_config( config, point_types=indices, edge_index=edge_index, neigh_isc=neigh_isc ) - # BORRAR - # print('In BasisMatrixData.from_config 3:') - # print('edge_index:', edge_index) - # print('edge_types:', edge_types) - return cls( edge_index=edge_index, neigh_isc=neigh_isc, diff --git a/src/graph2mat/core/data/sparse.py b/src/graph2mat/core/data/sparse.py index aa411bd..1548eca 100644 --- a/src/graph2mat/core/data/sparse.py +++ b/src/graph2mat/core/data/sparse.py @@ -207,16 +207,8 @@ def _blockmatrix_coo_coords( # Initialize the arrays to store the coordinates. rows = [] cols = [] - - # BORRAR - # print(f"In sparse.py _blockmatrix_coo_coords:") - # print(f"orbitals_row: {orbitals_row}") - # print(f"orbitals_col: {orbitals_col}") - # print(f"edge_index: {edge_index}") - 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_row = np.cumsum([0, *orbitals_row]) first_orb_col = np.cumsum([0, *orbitals_col]) @@ -296,10 +288,6 @@ def _blockmatrix_coo_coords( rows.extend(rows_symm) cols.extend(cols_symm) - # BORRAR - # print(f"Len rows in _blockmatrix_coo_coords: {len(rows)}") - # print(f"Len cols in _blockmatrix_coo_coords: {len(cols)}") - return np.array(rows), np.array(cols), (no_row, no_col * n_supercells) @@ -361,17 +349,6 @@ 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. """ - - # BORRAR - # print(f"In sparse.py _nodes_and_edges_to_coo:") - # print(f"calling _blockmatrix_coo_coords with:") - # print(f"orbitals_row: {orbitals_row}") - # print(f"orbitals_col: {orbitals_col}") - # print(f"edge_index: {edge_index}") - # print(f"n_supercells: {n_supercells}") - # print(f"edge_neigh_isc: {edge_neigh_isc}") - # print(f"symmetrize_edges: {symmetrize_edges}") - rows, cols, shape = _blockmatrix_coo_coords( orbitals_row=orbitals_row, orbitals_col=orbitals_col, @@ -380,19 +357,11 @@ def _nodes_and_edges_to_coo( edge_neigh_isc=edge_neigh_isc, symmetrize_edges=symmetrize_edges, ) - # BORRAR - # print(f"Len rows in _nodes_and_edges_to_coo: {len(rows)}") - # print(f"Len cols in _nodes_and_edges_to_coo: {len(cols)}") - # print(f"Shape in _nodes_and_edges_to_coo: {shape}") - # print(f"Len node_vals in _nodes_and_edges_to_coo: {len(node_vals)}") - # print(f"Len edge_vals in _nodes_and_edges_to_coo: {len(edge_vals)}") if symmetrize_edges: sparse_data = concatenate([node_vals, edge_vals, edge_vals]) else: sparse_data = concatenate([node_vals, edge_vals]) - # BORRAR - # print(f"Len sparse_data in _nodes_and_edges_to_coo: {len(sparse_data)}") if threshold is not None: mask = abs(sparse_data) > threshold diff --git a/src/graph2mat/core/data/table.py b/src/graph2mat/core/data/table.py index 0031ea5..2ca7e03 100644 --- a/src/graph2mat/core/data/table.py +++ b/src/graph2mat/core/data/table.py @@ -155,17 +155,6 @@ def __init__( [basis.basis_size for basis in self.basis], dtype=np.int32 ) - # BORRAR - # print(f"Basis sizes: {self.basis_size}") - - # SN: MOVED THIS TO BasisTableWithEdges BUT I AM NOT SURE OF HOW IT WORKS - # 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})" @@ -408,10 +397,6 @@ class BasisTableWithEdges: """Storing point information accounting for different row and column point types.""" def __init__(self, basis: Sequence[PointBasis], get_point_matrix: Optional[Callable] = None): self.is_square = all(b.matrix_role is None for b in basis) - - # BORRAR - # print(f"BasisTableWithEdges: is_square = {self.is_square}") - # print(f"BasisTableWithEdges: all matrix roles = {[b.matrix_role for b in basis]}") row_basis, col_basis = self._process_basis_rows_cols(basis) self.row = (BasisTableWithEdges_rowcol(row_basis, get_point_matrix) @@ -444,10 +429,6 @@ def __init__(self, basis: Sequence[PointBasis], get_point_matrix: Optional[Calla # Warn if they are different, and take the maximum of the two. if not np.allclose(self.row.R, self.col.R): print(f"Warning: Row and column point types have different cutoff radii. Taking the maximum of the two.") - # BORRAR - # print(f"Row cutoff radii: {self.row.R}") - # print(f"Col cutoff radii: {self.col.R}") - # print(f"Max cutoff radius: {np.max([self.row.R, self.col.R], axis=0)}") self.R = np.max([self.row.R, self.col.R], axis=0) @@ -472,16 +453,9 @@ def __init__(self, basis: Sequence[PointBasis], get_point_matrix: Optional[Calla # And also the sizes of the blocks. - # SN: Here we have all possible blocks self.point_block_shape = np.array([self.row.basis_size, self.col.basis_size]) self.point_block_size = self.row.basis_size * self.col.basis_size - - # BORRAR - # print("In BasisTableWithEdges: ") - # print(f"self.edge_type == point_types_to_edge_types:\n{self.edge_type}") - # print(f"self.point_block_shape:\n{self.point_block_shape}") - # print(f"self.point_block_size:\n{self.point_block_size}") - + point_types_combinations = np.array( list(itertools.combinations_with_replacement(range(n_types), 2))).T # Even if its not sqare, we take each direction once : we define the inverse shape, where we pass from (Ar, Bc) to (Ac, Br) @@ -493,29 +467,13 @@ def __init__(self, basis: Sequence[PointBasis], get_point_matrix: Optional[Calla self.edge_block_shape_inv = self.row.basis_size[point_types_combinations] self.edge_block_size = self.edge_block_shape.prod(axis=0) self.edge_block_size_inv = self.edge_block_shape_inv.prod(axis=0) - - # BORRAR - # print(f"Point type to edge type:\n{self.edge_type}") - # print(f"Edge type to point types:\n{self.edge_type_to_point_types}") - # print(f"Edge block shape:\n{self.edge_block_shape}") - # print(f"Edge block size:\n{self.edge_block_size}") - else: # Store the sizes of each point's basis — now separate for rows and cols. row_basis_size = self.row.basis_size col_basis_size = self.col.basis_size - # BORRAR - # print(f"Row basis sizes: {row_basis_size}") - # print(f"Col basis sizes: {col_basis_size}") - self.edge_type_to_point_types = point_types_combinations.T - # TODO: comprobe if this is okay for non-square matrices - # BORRAR - # print(f"Point type to edge type:\n{self.edge_type}") - # print(f"Edge type to point types:\n{self.edge_type_to_point_types}") - 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 @@ -524,15 +482,14 @@ def __init__(self, basis: Sequence[PointBasis], get_point_matrix: Optional[Calla row_basis_size[row_type_indices], col_basis_size[col_type_indices], ]) # shape: (2, n_combos) + + # 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) - # BORRAR - # print(f"Edge block shape:\n{self.edge_block_shape}") - # print(f"Edge block shape inv:\n{self.edge_block_shape_inv}") - 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,) @@ -730,12 +687,6 @@ def point_block_pointer(self, point_types: Sequence[int]) -> np.ndarray: """ pointers = np.zeros(len(point_types) + 1, dtype=np.int32) np.cumsum(self.point_block_size[point_types], out=pointers[1:]) - # BORRAR - # print(f"In BasisTableWithEdges.point_block_pointer:") - # print(f" point_types = {point_types}") - # print(f" point_block_size = {self.point_block_size}") - # print(f" pointers = {pointers}") - return pointers # SN: moved from the old BasisTableWithEdges_rowcol class to here @@ -767,13 +718,6 @@ def edge_block_pointer(self, edge_types: Sequence[int]): self.edge_block_size_inv[-edge_types]) # non‑positive → backward size (use absolute index) np.cumsum(sizes, out=pointers[1:]) - - # BORRAR - # print(f"In BasisTableWithEdges.edge_block_pointer:") - # print(f" edge_types = {edge_types}") - # print(f" edge_block_size = {self.edge_block_size}") - # print(f" pointers = {pointers}") - return pointers # SN: the max R must be the max among rows and cols diff --git a/src/graph2mat/core/modules/_labels_resort.py b/src/graph2mat/core/modules/_labels_resort.py index a28589b..2296943 100644 --- a/src/graph2mat/core/modules/_labels_resort.py +++ b/src/graph2mat/core/modules/_labels_resort.py @@ -95,17 +95,10 @@ def get_labels_resorting_array( 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 - # BORRAR - # print(f"Type {type}: offset[prev_type] = {offset[prev_type]}, type_nlabels[prev_type] = {type_nlabels[prev_type]}") - offset[type + ntypes_int - 1] = offset[prev_type] + type_nlabels[prev_type] # >0 - # BORRAR - # print(f"Type {-1*type}: offset[-type] = {offset[type]}, type_nlabels[type] = {type_nlabels[type]}") offset[-type + ntypes_int - 1] = offset[type + ntypes_int - 1] + type_nlabels[type + ntypes_int - 1] # < 0 - - # BORRAR - # print(f"calculated offset[{type}] = {offset[type]}, offset[{-type}] = {offset[-type]}") - prev_type = -type + ntypes_int - 1 # we have to continue from the negative type, because the next positive type will be after it. + # 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) @@ -114,13 +107,6 @@ def get_labels_resorting_array( offset[0] + type_nlabels[0], dtype=int ) - # BORRAR - # print("In get_labels_resorting_array:") - # print("sizes: ", np.asarray(sizes)) - # print("type_nlabels: ", np.asarray(type_nlabels)) - # print("offset: ", np.asarray(offset)) - # print("indices.shape: ", indices.shape) - type_i: cython.long[:] = np.zeros_like(sizes, dtype=int) i: cython.int = 0 diff --git a/src/graph2mat/core/modules/graph2mat.py b/src/graph2mat/core/modules/graph2mat.py index c3954e1..83dc328 100644 --- a/src/graph2mat/core/modules/graph2mat.py +++ b/src/graph2mat/core/modules/graph2mat.py @@ -437,8 +437,6 @@ def _init_self_interactions(self, **kwargs) -> List[MatrixBlock]: **kwargs, ) ) - # BORRAR - # print('self_interactions: \n', self_interactions) return self_interactions @@ -541,9 +539,6 @@ def summary(self) -> str: s += f" {self._get_node_operation_summary(x)}" s += "\nEdge operations:" - - # # BORRAR - # print(f"self.interactions: {self.interactions}") for k, x in self.interactions.items(): if x.symm_transpose and not self.is_square: @@ -571,27 +566,6 @@ def summary(self) -> str: s += " [XY = YX.T]" s += f" {self._get_edge_operation_summary(x)}." - - # # BORRAR - # print(s) - - # point = self.graph2mat_table_row.basis[point_type] - # neigh = self.graph2mat_table_col.basis[neigh_type] - - # if x is None: - # print(" x is None ") - # print(f" ({point.type}, {neigh.type}) No basis functions.") - # continue - # print(" ({point.type}, {neigh.type})") - # print(f"({point.type}r, {neigh.type})") - - # if x.symm_transpose: - # print(" x.symm_transpose is True ") - # print(" [XY = YX.T]") - # print(" self._get_edge_operation_summary(x) , being x ", x) - # print(f" {self._get_edge_operation_summary(x)}.") - # # END BORRAR - # raise NotImplementedError("Printing edge operations for non-square matrices are not implemented yet.") return s except Exception as e: return f"Error generating summary: {e}" @@ -819,12 +793,6 @@ def _forward_interactions( outputs = [] graph2mat_edge_types = self.edge_types_to_graph2mat[edge_types] - # BORRAR - # print("In Graph2Mat _forward_interactions: ") - # print("graph2mat_edge_types: ", graph2mat_edge_types) - # print("edge_types: ", edge_types) - # print("Order of interactions (edge type is last): ", list(self.interactions.keys())) - # print('This is the order of indexes: then all the indexes of same type se cogen seguidos') # Call each unique interaction function with only the features # of edges that correspond to that type. @@ -902,25 +870,9 @@ 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). - # BORRAR - # print("In Graph2Mat _forward_interactions (before get_edgelabels_resort_index): ") - # print("graph2mat_edge_types: ", graph2mat_edge_types) - # print("edge_types (original_types): ", edge_types) - sort_indices = self._get_edgelabels_resort_index( graph2mat_edge_types, original_types=edge_types ) - - # BORRAR - # print("In Graph2Mat _forward_interactions (after get_edgelabels_resort_index): ") - # print("graph2mat_edge_types: ", graph2mat_edge_types) - # print("edge_types (original_types): ", edge_types) - # print("unsorted_edge_labels: ", unsorted_edge_labels) - # print("sort_indices: ", sort_indices) - # print("unsorted_edge_labels[sort_indices]: ", unsorted_edge_labels[sort_indices]) - # print("unsorted_edge_labels.shape: ", unsorted_edge_labels.shape) - # print("new shape: ", unsorted_edge_labels[sort_indices].shape) - # Do the resorting and return the result. return unsorted_edge_labels[sort_indices] @@ -979,16 +931,6 @@ def _get_edgelabels_resort_index( types = types[::2] original_types = original_types[::2] - # # BORRAR - # print("In Graph2Mat _get_edgelabels_resort_index: ") - # print("types: ", types) - # print("types.type: ", types.dtype) - # types = reorder_array(types) - - # # BORRAR - # print("after reorder_array(types): ", types) - # print("types.type: ", types.dtype) - return self._get_labels_resort_index( types=types, original_types=original_types, @@ -1047,15 +989,6 @@ def _get_labels_resort_index( return np.where(mask)[0] else: - # BORRAR - # print("In Graph2Mat _get_labels_resort_index: ") - # print('BEFORE CALLING get_labels_resorting_array') - # print("types: ", types) - # print("shapes: ", shapes) - # print("shapes_inv: ", shapes_inv) - # print("transpose_neg: ", transpose_neg) - # print("kwargs: ", kwargs) - indices = get_labels_resorting_array( types, shapes=shapes.astype(types.dtype), @@ -1063,9 +996,5 @@ def _get_labels_resort_index( transpose_neg=transpose_neg, **kwargs, ) - # BORRAR - # print('AFTER CALLING get_labels_resorting_array') - # print("len(indices): ", len(indices)) - # print("indices: ", indices) return indices From e8f39b064eadc14e84b8bd147ebf0baf9fd93058 Mon Sep 17 00:00:00 2001 From: Saru local ICN2 Date: Fri, 31 Jul 2026 13:13:36 +0200 Subject: [PATCH 3/3] Handle BasisTablewithEdges without a subclass: pass PointBasis lists as dictionarieswith keys "row" and "col" --- .../e3nn/modules/tests/test_e3nngraph2mat.py | 1 - src/graph2mat/bindings/torch/data/formats.py | 4 +- src/graph2mat/core/data/__init__.py | 2 +- src/graph2mat/core/data/basis.py | 13 - .../core/data/matrices/basis_matrix.py | 2 +- src/graph2mat/core/data/processing.py | 6 +- src/graph2mat/core/data/sparse.py | 14 +- src/graph2mat/core/data/table.py | 673 +++++++++--------- src/graph2mat/core/modules/_labels_resort.py | 5 +- src/graph2mat/core/modules/graph2mat.py | 24 +- 10 files changed, 370 insertions(+), 374 deletions(-) diff --git a/src/graph2mat/bindings/e3nn/modules/tests/test_e3nngraph2mat.py b/src/graph2mat/bindings/e3nn/modules/tests/test_e3nngraph2mat.py index 6c1ce30..e7ee7d6 100644 --- a/src/graph2mat/bindings/e3nn/modules/tests/test_e3nngraph2mat.py +++ b/src/graph2mat/bindings/e3nn/modules/tests/test_e3nngraph2mat.py @@ -8,7 +8,6 @@ from graph2mat import ( BasisConfiguration, BasisTableWithEdges, - BasisTableWithEdges, MatrixDataProcessor, PointBasis, conversions, diff --git a/src/graph2mat/bindings/torch/data/formats.py b/src/graph2mat/bindings/torch/data/formats.py index 4135cb7..a477437 100644 --- a/src/graph2mat/bindings/torch/data/formats.py +++ b/src/graph2mat/bindings/torch/data/formats.py @@ -37,8 +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 + # if isinstance(tensor, np.ndarray): + # return tensor return tensor.numpy(force=True) diff --git a/src/graph2mat/core/data/__init__.py b/src/graph2mat/core/data/__init__.py index 8f4a0b2..af17e1a 100644 --- a/src/graph2mat/core/data/__init__.py +++ b/src/graph2mat/core/data/__init__.py @@ -27,5 +27,5 @@ from .metrics import OrbitalMatrixMetric from . import metrics from .processing import MatrixDataProcessor, BasisMatrixData, BasisMatrixDataBase -from .table import BasisTableWithEdges, AtomicTableWithEdges, BasisTableWithEdges_rowcol +from .table import BasisTableWithEdges, AtomicTableWithEdges from .formats import Formats, conversions, ConversionManager diff --git a/src/graph2mat/core/data/basis.py b/src/graph2mat/core/data/basis.py index caf8e85..41b77e2 100644 --- a/src/graph2mat/core/data/basis.py +++ b/src/graph2mat/core/data/basis.py @@ -149,11 +149,6 @@ class PointBasis: with optional minus signs, e.g. ``"xyz"``, ``"-yz-x"``, ``"z-x-y"``. You can also use the aliases that we provide, such as ``"cartesian"`` or ``"spherical"``. - matrix_role: - role that the specified basis plays in the output matrix. It can be 'row' - or 'col'. If None, assumes the matrix is square and the basis is used for - both rows and columns. - Examples ---------- @@ -179,20 +174,12 @@ class PointBasis: R: Union[float, np.ndarray] basis: Union[str, Sequence[Union[int, Tuple[int, int, int]]]] = () basis_convention: BasisConvention = "spherical" - matrix_role: Union[Literal["row", "col"], None] = None def __post_init__(self): basis = self._sanitize_basis(self.basis) object.__setattr__(self, "basis", basis) - if self.matrix_role is not None: - assert self.matrix_role in ["row", "col"], "matrix_role must be 'row', 'col' or None." - print( - f"Defined {self.matrix_role} for type {self.type} with basis {self.basis} and reach {self.R}." - ) - object.__setattr__(self, "matrix_role", self.matrix_role) - assert isinstance(self.R, Number) or ( isinstance(self.R, np.ndarray) and len(self.R) == self.basis_size ), f"R must be a float or an array of length {self.basis_size} (the number of functions)." diff --git a/src/graph2mat/core/data/matrices/basis_matrix.py b/src/graph2mat/core/data/matrices/basis_matrix.py index 96c91b6..c71f82e 100644 --- a/src/graph2mat/core/data/matrices/basis_matrix.py +++ b/src/graph2mat/core/data/matrices/basis_matrix.py @@ -5,7 +5,7 @@ from dataclasses import dataclass import numpy as np -from ..table import BasisTableWithEdges, BasisTableWithEdges +from ..table import BasisTableWithEdges BasisCount = np.ndarray # [num_points] diff --git a/src/graph2mat/core/data/processing.py b/src/graph2mat/core/data/processing.py index 2661ac6..e611914 100644 --- a/src/graph2mat/core/data/processing.py +++ b/src/graph2mat/core/data/processing.py @@ -198,7 +198,6 @@ def matrix_from_data( if is_batch is None: is_batch = isinstance(data, Batch) - if is_batch: return tuple( self.yield_from_batch( @@ -808,9 +807,8 @@ def labels_to( else: kwargs["n_supercells"] = nsc.prod() - # SN: changed this -- the n_orbitals can be diff in cols and rows if its non square - 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] + 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] diff --git a/src/graph2mat/core/data/sparse.py b/src/graph2mat/core/data/sparse.py index 1548eca..79b8459 100644 --- a/src/graph2mat/core/data/sparse.py +++ b/src/graph2mat/core/data/sparse.py @@ -192,7 +192,6 @@ def _blockmatrix_coo_coords( 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. n_supercells: number of supercells in the matrix. edge_neigh_isc: @@ -273,17 +272,7 @@ def _blockmatrix_coo_coords( cols_symm.extend(opp_block_cols) elif symmetrize_edges and not is_square: raise ValueError( - "SN: TODO - I think thsi not work as I intended to... maybe it cannot be") - # We assume that the edge_index contains only the edges in one direction, - # but as it it not square, we cannot assume that the opposite direction is just the transpose of the - # original edge. - i_start_symm = first_orb_row[j_at] - i_end_symm = i_start_symm + orbitals_row[j_at] - j_start_symm = first_orb_col[i_at] - j_end_symm = j_start_symm + orbitals_col[i_at] - block_rows_symm, block_cols_symm = np.mgrid[i_start_symm:i_end_symm, j_start_symm:j_end_symm].reshape(2, -1) - rows_symm.extend(block_rows_symm) - cols_symm.extend(block_cols_symm) + "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) @@ -357,7 +346,6 @@ def _nodes_and_edges_to_coo( edge_neigh_isc=edge_neigh_isc, symmetrize_edges=symmetrize_edges, ) - if symmetrize_edges: sparse_data = concatenate([node_vals, edge_vals, edge_vals]) else: diff --git a/src/graph2mat/core/data/table.py b/src/graph2mat/core/data/table.py index 2ca7e03..5d89735 100644 --- a/src/graph2mat/core/data/table.py +++ b/src/graph2mat/core/data/table.py @@ -16,6 +16,7 @@ from io import StringIO from pathlib import Path from typing import Callable, Generator, List, Literal, Optional, Sequence, Union, Tuple +from itertools import chain import numpy as np import sisl @@ -23,7 +24,7 @@ from .basis import BasisConvention, PointBasis, get_change_of_basis -class BasisTableWithEdges_rowcol: +class BasisTableWithEdges: """Stores the unique types of points in the system, with their basis and the possible edges. It also knows the size of the blocks, and other type dependent variables. @@ -31,11 +32,16 @@ class BasisTableWithEdges_rowcol: 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_rowcol: """ #: 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_rowcol: #: 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: @@ -137,83 +164,138 @@ def __init__( self.basis_convention = basis_convention - # 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. - # TODO: pass this to the non-square(to Sara_) - 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 - ] + # For the basis convention, get the matrices to change from cartesian to our convention. + self.change_of_basis, self.change_of_basis_inv = get_change_of_basis( + "cartesian", self.basis_convention + ) - # Get also the cutoff radii for each point. - self.R = np.array([point_basis.maxR() for point_basis in self.basis]) + # 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 + for i in range(n_types): + # 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): + point_types_to_edge_types[i, j] = edge_type + point_types_to_edge_types[j, i] = -edge_type + edge_type += 1 + + # 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. + + 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 - def __repr__(self): - return f"{self.__class__.__name__}({self.basis_convention}, basis={self.basis})" + # And also the sizes of the blocks. + self.point_block_shape = np.array([row_basis_size, col_basis_size]) + self.point_block_size = row_basis_size * col_basis_size - def _repr_html_(self): - table = "" - table += f"" + point_types_combinations = np.array( + 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 _basis_string(basis): - s = "" + self.edge_type_to_point_types = point_types_combinations.T - for basis_set in basis: - mul, l, parity = basis_set - s += f"{mul}x{l}{'e' if parity == 1 else 'o'} + " - s = s[:-3] + 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 - return s + # 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) - for i, point_basis in enumerate(self.basis): - table += f"" + # 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) - table += "
IndexTypeIrrepsMax R
{i}{point_basis.type}{_basis_string(point_basis.basis)}{point_basis.maxR()}
" + 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,) - return table + @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 __str__(self): - return "\n".join([f"\t- {point_basis}" for point_basis in self.basis]) + 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. - def __len__(self): - return len(self.basis) + 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. - def __eq__(self, other): - if not isinstance(other, self.__class__): - return False + 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: + 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.") - 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)) + 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 - 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) - ) + 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"] - ) -> tuple["BasisTableWithEdges_rowcol", np.ndarray, np.ndarray, Optional[np.ndarray]]: + ) -> tuple["BasisTableWithEdges", np.ndarray, np.ndarray, Optional[np.ndarray]]: r"""Groups the basis in this table and creates a new table. It also returns useful objects to convert between the ungrouped @@ -273,42 +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 - # def __eq__(self, other): - # return isinstance(other, A) - - # 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 @@ -324,14 +403,14 @@ def group( ) 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", @@ -341,31 +420,71 @@ def group( ) # 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. @@ -377,252 +496,68 @@ def index_to_type(self, index: int) -> Union[str, int]: """ return self.types[index] - def maxR(self) -> float: - """Maximum cutoff radius in the basis.""" - return self.R.max() + 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.basis] + return [point.to_sisl_atom() for point in self.row_basis] - -class BasisTableWithEdges: - """Storing point information accounting for different row and column point types.""" - def __init__(self, basis: Sequence[PointBasis], get_point_matrix: Optional[Callable] = None): - self.is_square = all(b.matrix_role is None for b in basis) - row_basis, col_basis = self._process_basis_rows_cols(basis) - self.row = (BasisTableWithEdges_rowcol(row_basis, get_point_matrix) - if len(row_basis) > 0 else None) - self.col = (BasisTableWithEdges_rowcol(col_basis, get_point_matrix) - if len(col_basis) > 0 else None) - - if (self.row is None)^(self.col is None): - raise ValueError("If one of row or column basis is None, the other must be None too.") - - # Check, in the case of non square matrix, properties match as they should - if not self.is_square: - self.check_nonsquare_properties() - - self.basis_convention = self.row.basis_convention if self.row is not None else None - - # For the basis convention, get the matrices to change from cartesian to our convention. - self.change_of_basis, self.change_of_basis_inv = get_change_of_basis( - "cartesian", self.basis_convention - ) - - self.types = self.row.types if self.row is not None else None - - n_types = len(self.types) - - if self.is_square: - self.R = self.row.R - else: - # Each row and column point type can have a different cutoff radius. - # Warn if they are different, and take the maximum of the two. - if not np.allclose(self.row.R, self.col.R): - print(f"Warning: Row and column point types have different cutoff radii. Taking the maximum of the two.") - - self.R = np.max([self.row.R, self.col.R], axis=0) - - # 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 - for i in range(n_types): - # 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): - point_types_to_edge_types[i, j] = edge_type - point_types_to_edge_types[j, i] = -edge_type - edge_type += 1 - - # 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. - - self.edge_type = point_types_to_edge_types - - - # And also the sizes of the blocks. - self.point_block_shape = np.array([self.row.basis_size, self.col.basis_size]) - self.point_block_size = self.row.basis_size * self.col.basis_size - - point_types_combinations = np.array( - list(itertools.combinations_with_replacement(range(n_types), 2))).T - # Even if its not sqare, we take each direction once : we define the inverse shape, where we pass from (Ar, Bc) to (Ac, Br) - if self.is_square: - - - self.edge_type_to_point_types = point_types_combinations.T - self.edge_block_shape = self.row.basis_size[point_types_combinations] - self.edge_block_shape_inv = self.row.basis_size[point_types_combinations] - self.edge_block_size = self.edge_block_shape.prod(axis=0) - self.edge_block_size_inv = self.edge_block_shape_inv.prod(axis=0) - else: - # Store the sizes of each point's basis — now separate for rows and cols. - row_basis_size = self.row.basis_size - col_basis_size = self.col.basis_size - - self.edge_type_to_point_types = point_types_combinations.T - - 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 - - # 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) - - # 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) - - 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,) - - @staticmethod - def _process_basis_rows_cols(basis: Union[List[PointBasis], BasisTableWithEdges_rowcol]) \ - -> Tuple[List[PointBasis], List[PointBasis]]: - """Processes the basis to determine which basis is used for rows and which for columns. - - 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. - - Parameters - ---------- - basis: List[PointBasis] - The list of point bases. - - Returns - ------- - i_basis: List[PointBasis] - The list of point bases for rows. - j_basis: List[PointBasis] - The list of point bases for columns. - """ - if isinstance(basis, BasisTableWithEdges_rowcol): - basis = basis.basis - i_basis = [] - j_basis = [] - for b in basis: - if b.matrix_role == "row": - i_basis.append(b) - elif b.matrix_role == "col": - j_basis.append(b) - else: - # If None, it is a square matrix: same basis for rows and columns - i_basis.append(b) - j_basis.append(b) - 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 check_nonsquare_properties(self): - - basis_convention_row = self.row.basis_convention if self.row is not None else None - basis_convention_col = self.col.basis_convention if self.col is not None else None - - assert basis_convention_row == basis_convention_col, \ - f"Row and column basis conventions must be the same. \ -Got {basis_convention_row} and {basis_convention_col}." - - types_row = self.row.types if self.row is not None else None - types_col = self.col.types if self.col is not None else None - - assert types_row == types_col, \ - f"Row and column basis types must be the same. \ -Got {types_row} and {types_col}." - - def group( - self, grouping: Literal["basis_shape", "point_type", "max"] - ) -> tuple["BasisTableWithEdges_rowcol", np.ndarray, np.ndarray, Optional[np.ndarray]]: - - "TODO: define this for nonsquare in this scheme" - - # SN: changed the output: returns a tuple of two tables, one for rows and one for columns - if self.is_square: - new_table, point_type_conversion, edge_type_conversion, filters = self.row.group(grouping) - return (new_table, new_table), point_type_conversion, edge_type_conversion, filters - else: - if grouping == "point_type": - new_table_row, point_type_conversion_row, edge_type_conversion_row, filters_row = self.row.group(grouping) - new_table_col, point_type_conversion_col, edge_type_conversion_col, filters_col = self.col.group(grouping) - if point_type_conversion_row != point_type_conversion_col: - raise ValueError(f"Row and column point type conversions must be the same for non-square matrices. Got {point_type_conversion_row} and {point_type_conversion_col}.") - if edge_type_conversion_row != edge_type_conversion_col: - raise ValueError(f"Row and column edge type conversions must be the same for non-square matrices. Got {edge_type_conversion_row} and {edge_type_conversion_col}.") - # TODO: for the filters ? - return (new_table_row, new_table_col), point_type_conversion_row, edge_type_conversion_row, filters_row - else: - raise NotImplementedError(f"Grouping for non-square matrices \ -with grouping {grouping} is not implemented yet.") - - def _repr_html_(self): - if self.is_square: - return self.row._repr_html_() - else: - # TODO: comprobar - 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 __str__(self): if self.is_square: - return str(self.row) + # Both are the same, print one + return "\n".join([f"\t- {point_basis}" for point_basis in self.row_basis]) else: - return f"Row basis:\n{self.row}\nColumn basis:\n{self.col}" + 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 self.is_square: - return len(self.row) - else: - if len(self.col) != len(self.row): - raise ValueError("Row and column basis have different lengths.") - return len(self.row) + 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 &= self.row == other.row - same &= self.col == other.col + + 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 - - # SN: moved from the old BasisTableWithEdges_rowcol class to here + 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. @@ -633,7 +568,6 @@ def type_to_index(self, point_type: Union[str, int]) -> int: """ return self.types.index(point_type) - # SN: moved from the old BasisTableWithEdges_rowcol class to here def types_to_indices(self, types: Sequence) -> np.ndarray: """Converts from an array of types IDs to their indices in the basis table. @@ -657,7 +591,6 @@ def types_to_indices(self, types: Sequence) -> np.ndarray: # And reconstruct the original array, which is now an array of indices instead of types return unique_indices[inverse_indices] - # SN: moved from the old BasisTableWithEdges_rowcol class to here def point_type_to_edge_type(self, point_type: np.ndarray) -> Union[int, np.ndarray]: """Converts pairs of point types to edge types. @@ -669,7 +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]] - # SN: moved from the old BasisTableWithEdges_rowcol class to here def point_block_pointer(self, point_types: Sequence[int]) -> np.ndarray: """Pointers to the beggining of node blocks in a flattened matrix. @@ -689,7 +621,6 @@ def point_block_pointer(self, point_types: Sequence[int]) -> np.ndarray: np.cumsum(self.point_block_size[point_types], out=pointers[1:]) return pointers - # SN: moved from the old BasisTableWithEdges_rowcol class to here def edge_block_pointer(self, edge_types: Sequence[int]): """Pointers to the beggining of edge blocks in a flattened matrix. @@ -719,11 +650,11 @@ def edge_block_pointer(self, edge_types: Sequence[int]): np.cumsum(sizes, out=pointers[1:]) return pointers - - # SN: the max R must be the max among rows and cols + def maxR(self) -> float: """Maximum cutoff radius in the basis.""" - return max(self.row.maxR(), self.col.maxR()) + return self.R.max() + class AtomicTableWithEdges(BasisTableWithEdges): @@ -896,4 +827,96 @@ def __getitem__(self, key): return key def __eq__(self, other): - return isinstance(other, IdentityConversion) \ No newline at end of file + 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 2296943..bccca70 100644 --- a/src/graph2mat/core/modules/_labels_resort.py +++ b/src/graph2mat/core/modules/_labels_resort.py @@ -88,7 +88,7 @@ def get_labels_resorting_array( # Count the number of entries of each type for i_edge in range(n_entries): type: cython.int = types[i_edge] - type_nlabels[type + ntypes_int - 1] += sizes[type] + 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. @@ -97,7 +97,8 @@ def get_labels_resorting_array( 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. + + # 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. diff --git a/src/graph2mat/core/modules/graph2mat.py b/src/graph2mat/core/modules/graph2mat.py index 83dc328..ea9aa55 100644 --- a/src/graph2mat/core/modules/graph2mat.py +++ b/src/graph2mat/core/modules/graph2mat.py @@ -398,8 +398,8 @@ def _init_center_types(self, basis_grouping): self.basis_filters, ) = self.basis_table.group(self.basis_grouping) - self.graph2mat_table_row = graph2mat_tables[0] - self.graph2mat_table_col = graph2mat_tables[1] + 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 @@ -423,9 +423,9 @@ def _init_center_types(self, basis_grouping): def _init_self_interactions(self, **kwargs) -> List[MatrixBlock]: self_interactions = [] - for i_point_type_basis in range(len(self.graph2mat_table_row.basis)): - row_basis = self.graph2mat_table_row.basis[i_point_type_basis] - col_basis = self.graph2mat_table_col.basis[i_point_type_basis] + 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) @@ -442,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_row.basis)), 2 # row and col basis have the same length, so we can use either one + range(len(self.graph2mat_table_rowb)), 2 # row and col basis have the same length, so we can use either one ) interactions = {} @@ -456,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_row.basis[point_i] - j_basis = self.graph2mat_table_col.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 @@ -522,8 +522,8 @@ def summary(self) -> str: raise ValueError( f"Node operation {i} is symmetric transpose, but the matrix is NOT square. This is not allowed." ) - point_r = self.graph2mat_table_row.basis[i] - point_c = self.graph2mat_table_col.basis[i] + 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." @@ -548,8 +548,8 @@ def summary(self) -> str: point_type, neigh_type, edge_type = map(int, k[1:-1].split(",")) - point = self.graph2mat_table_row.basis[point_type] - neigh = self.graph2mat_table_col.basis[neigh_type] + point = self.graph2mat_table_rowb[point_type] + neigh = self.graph2mat_table_colb[neigh_type] if x is None: s += f"\n ({point.type}, {neigh.type}) No basis functions."