Refactor SOM and metrics modules:

- Removed unused imports in `som.py` for improved clarity and organization.
- Commented out unnecessary device transfer in the `SOM` class to prevent potential issues.
- Added back essential imports in `metrics.py` for axial distance calculations.
- Enhanced inline comments for better understanding of logic in distance normalization.

Author: LouisTier

torchsom/core/som.py

@@ -11,12 +11,7 @@
 
 from ..utils.decay import DECAY_FUNCTIONS
 from ..utils.distances import DISTANCE_FUNCTIONS
-from ..utils.grid import (
-    adjust_meshgrid_topology,
-    axial_distance,
-    convert_to_axial_coords,
-    create_mesh_grid,
-)
+from ..utils.grid import adjust_meshgrid_topology, create_mesh_grid
 from ..utils.initialization import initialize_weights
 from ..utils.metrics import calculate_quantization_error, calculate_topographic_error
 from ..utils.neighborhood import NEIGHBORHOOD_FUNCTIONS
@@ -338,7 +333,7 @@ def fit(
             Tuple[List[float], List[float]]: Quantization and topographic errors [epoch]
         """
 
-        data = data.to(self.device)
+        # data = data.to(self.device)
         dataset = TensorDataset(data)
         dataloader = DataLoader(
             dataset, batch_size=self.batch_size, shuffle=True, pin_memory=False
@@ -682,7 +677,7 @@ def build_distance_map(
         # Normalize the distance map
         max_distance = torch.max(
             distance_matrix.masked_fill(torch.isnan(distance_matrix), float("-inf"))
-        )
+        )  # Replace NaNs with -inf to be ignored by max()
         return distance_matrix / max_distance if max_distance > 0 else distance_matrix
 
     def build_bmus_data_map(

torchsom/utils/metrics.py

@@ -3,6 +3,8 @@
 
 import torch
 
+from ..utils.grid import axial_distance, convert_to_axial_coords
+
 
 def calculate_quantization_error(
     data: torch.Tensor,
@@ -95,9 +97,6 @@ def calculate_topographic_error(
             bmu2_row = int(torch.div(indices[i, 1], y_dim, rounding_mode="floor"))
             bmu2_col = int(indices[i, 1] % y_dim)
 
-            # Convert to axial coordinates
-            from ..utils.grid import axial_distance, convert_to_axial_coords
-
             q1, r1 = convert_to_axial_coords(bmu1_row, bmu1_col)
             q2, r2 = convert_to_axial_coords(bmu2_row, bmu2_col)