Refactored code completely. Proposed modifications to test.

Author: LouisTier

torchsom/core/som.py

@@ -156,6 +156,10 @@ def _update_weights(
                     for row, col in bmus
                 ]
             )  # [batch_size, row_neurons, col_neurons]
+            # ! Modification to test
+            # # Vectorised: build a tensor of BMU coordinates and compute in one shot
+            # coords = torch.stack([bmus[:, 0], bmus[:, 1]], dim=1).to(torch.long)
+            # neighborhoods = self.neighborhood_fn(coords, sigma)  # update neighborhood_fn to accept batched coords  # [batch_size, row_neurons, col_neurons]
 
             # Reshape for broadcasting
             neighborhoods = neighborhoods.view(batch_size, self.x, self.y, 1)
@@ -595,6 +599,21 @@ def build_distance_map(
                 max_neighbors += len(neighbor_offsets)
                 all_offsets.append(neighbor_offsets)
 
+        # # ! Modification to test
+        # def _offsets_for_row(r: int) -> List[Tuple[int, int]]:
+        #     if self.topology == "hexagonal":
+        #         merged = []
+        #         for order in range(1, neighborhood_order + 1):
+        #             merged.extend(
+        #                 get_hexagonal_offsets(order)["even" if r % 2 == 0 else "odd"]
+        #             )
+        #         return merged
+        #     else:
+        #         merged = []
+        #         for order in range(1, neighborhood_order + 1):
+        #             merged.extend(get_rectangular_offsets(order))
+        #         return merged
+
         # Initialize distance map
         distance_matrix = torch.full(
             (self.weights.shape[0], self.weights.shape[1], max_neighbors),
@@ -924,6 +943,21 @@ def build_classification_map(
             for order in range(1, neighborhood_order + 1):
                 neighborhood_offsets.extend(get_rectangular_offsets(order))
 
+        # ! Method to test
+        # def _offsets_for_row(r: int) -> List[Tuple[int, int]]:
+        #     if self.topology == "hexagonal":
+        #         merged = []
+        #         for order in range(1, neighborhood_order + 1):
+        #             merged.extend(
+        #                 get_hexagonal_offsets(order)["even" if r % 2 == 0 else "odd"]
+        #             )
+        #         return merged
+        #     else:
+        #         merged = []
+        #         for order in range(1, neighborhood_order + 1):
+        #             merged.extend(get_rectangular_offsets(order))
+        #         return merged
+
         # Iterate through each activated neuron
         for bmu_pos, sample_indices in bmus_map.items():
             if len(sample_indices) > 0:

torchsom/utils/distances.py

@@ -45,7 +45,7 @@ def _euclidean_distance(
         torch.Tensor: euclidean distance between input and weights [row_neurons, col_neurons]
     """
 
-    return torch.norm(torch.subtract(data, weights), dim=-1)
+    return torch.max(torch.abs(data - weights), dim=-1).values
 
 
 def _manhattan_distance(
@@ -79,7 +79,8 @@ def _chebyshev_distance(
         torch.Tensor: chebyshev distance between input and weights [row_neurons, col_neurons]
     """
 
-    return torch.max(torch.subtract(data, weights), dim=-1).values
+    # return torch.max(torch.subtract(data, weights), dim=-1).values
+    return torch.max(torch.abs(data - weights), dim=-1).values
 
 
 # TODO Check if this method works and if it is more efficient (also ensure it is compatible with batch framework)

torchsom/utils/grid.py

@@ -86,6 +86,17 @@ def convert_to_axial_coords(
     r = row  # y axis
     return q, r
 
+    # ! Modification to test
+    # # Match the even-r layout used in `adjust_meshgrid_topology`
+    # #   even rows -> shifted left by 0.5
+    # #   odd  rows -> no horizontal shift
+    # if row % 2 == 0:  # even-r
+    #     q = col - 0.5 - (row // 2)
+    # else:  # odd-r
+    #     q = col - (row // 2)
+    # r = row
+    # return q, r
+
 
 def axial_distance(
     q1: float,

torchsom/utils/initialization.py

@@ -35,6 +35,11 @@ def random_init(
 
         return sampled_weights
 
+        # ! Modification to test
+        # # Return value ignores the original weights tensor
+        # weights.copy_(data[indices])
+        # return weights
+
     except RuntimeError as e:
         raise RuntimeError(f"Random initialization failed: {str(e)}")
 

torchsom/utils/metrics.py

@@ -79,6 +79,7 @@ def calculate_topographic_error(
     # Calculate distances between each data point and all neurons
     distances = distance_fn(data_expanded, weights_expanded)
 
+    # ! Modification to test: all the lines below could be vectorized
     # Get top 2 BMU indices for each sample
     batch_size = distances.shape[0]
     _, indices = torch.topk(distances.view(batch_size, -1), k=2, largest=False, dim=1)

torchsom/utils/neighborhood.py

@@ -6,7 +6,7 @@
 def _gaussian(
     xx: torch.Tensor,
     yy: torch.Tensor,
-    c: Tuple[torch.Tensor, torch.Tensor],
+    c: Tuple[torch.Tensor, torch.Tensor],  # ! Tuple[int, int]
     sigma: float,
 ) -> torch.Tensor:
     """Gaussian neighborhood function to update weights.
@@ -55,7 +55,7 @@ def _gaussian(
 def _mexican_hat(
     xx: torch.Tensor,
     yy: torch.Tensor,
-    c: Tuple[torch.Tensor, torch.Tensor],
+    c: Tuple[torch.Tensor, torch.Tensor],  # ! Tuple[int, int]
     sigma: float,
 ) -> torch.Tensor:
     """
@@ -91,7 +91,6 @@ def _mexican_hat(
     Returns:
         torch.Tensor: Mexican hat neighborhood weights. Element-wise product standing for the combined influence of mexican neighborhood around center c with a spread sigma [row_neurons, col_neurons].
     """
-
     denum = 2 * sigma * sigma
     cst = 1 / (torch.pi * torch.pow(torch.tensor(sigma), 4))
     squared_distances = torch.pow(xx - c[0], 2) + torch.pow(
@@ -100,6 +99,16 @@ def _mexican_hat(
     exp_distances = torch.exp(-squared_distances / denum)
     mexican_hat = cst * (1 - (1 / 2) * squared_distances / (2 * denum)) * exp_distances
 
+    # ! Modification to test
+    # denum = 2 * sigma * sigma
+    # sigma_t = torch.tensor(sigma, device=xx.device, dtype=xx.dtype)
+    # cst = 1 / (torch.pi * sigma_t.pow(4))
+    # squared_distances = torch.pow(xx - c[0], 2) + torch.pow(
+    #     yy - c[1], 2
+    # )  # Squared distances from center [row_neurons, col_neurons]
+    # exp_distances = torch.exp(-squared_distances / denum)
+    # mexican_hat = cst * (1 - (1 / 2) * squared_distances / (2 * denum)) * exp_distances
+
     # Ensure the central peak is exactly 1.0
     max_value = mexican_hat[c[0], c[1]]
     if max_value > 0:
@@ -110,7 +119,7 @@ def _mexican_hat(
 def _bubble(
     xx: torch.Tensor,
     yy: torch.Tensor,
-    c: Tuple[torch.Tensor, torch.Tensor],
+    c: Tuple[torch.Tensor, torch.Tensor],  # ! Tuple[int, int]
     sigma: float,
 ) -> torch.Tensor:
     """
@@ -159,7 +168,7 @@ def _bubble(
 def _triangle(
     xx: torch.Tensor,
     yy: torch.Tensor,
-    c: Tuple[torch.Tensor, torch.Tensor],
+    c: Tuple[torch.Tensor, torch.Tensor],  # ! Tuple[int, int]
     sigma: float,
 ) -> torch.Tensor:
     """