code quality

example.py

@@ -1,5 +1,5 @@
 from swarms_torch import ParticleSwarmOptimization
 
-#test
+# test
 pso = ParticleSwarmOptimization(goal="Attention is all you need", n_particles=100)
-pso.optimize(iterations=1000)
+pso.optimize(iterations=1000)

nnt.py

@@ -4,22 +4,18 @@
 x = torch.randn(1, 10)
 
 network = NNTransformer(
-    #transformer cells
-    neuron_count = 5, 
-
-    #num states
-    num_states = 10,
-
-    #input dim
-    input_dim = 10,
-
-    #output dim
-    output_dim = 10,
-
-    #nhead
-    nhead = 2,
+    # transformer cells
+    neuron_count=5,
+    # num states
+    num_states=10,
+    # input dim
+    input_dim=10,
+    # output dim
+    output_dim=10,
+    # nhead
+    nhead=2,
 )
 
 
 output = network(x)
-print(output)
+print(output)

swarms_torch/__init__.py

@@ -10,5 +10,4 @@
     "ParticleSwarmOptimization",
     "AntColonyOptimization",
     "QueenBeeGa",
-    "SPO"
-]
+    "SPO"]

swarms_torch/ant_colony_swarm.py

@@ -1,5 +1,6 @@
 import torch
 
+
 class AntColonyOptimization:
     """
     Ant Colony Optimization
@@ -27,9 +28,9 @@ class AntColonyOptimization:
     goal_string = "Hello ACO"
     aco = AntColonyOptimization(goal_string, num_iterations=1000)
     best_solution = aco.optimize()
-    
+
     print("Best Matched String:", best_solution)
-    
+
     Features to implement
     --------
     1. Add a stopping criterion
@@ -39,22 +40,24 @@ class AntColonyOptimization:
     5. Add a function to plot the best solution
 
     """
+
     def __init__(
-        self, 
-        goal: str = None, 
-        num_ants: int = 10000, 
-        evaporation_rate: float = 0.1, 
-        alpha: int = 1, 
-        beta: int = 1, 
-        num_iterations: int = 10010
+        self,
+        goal: str = None,
+        num_ants: int = 10000,
+        evaporation_rate: float = 0.1,
+        alpha: int = 1,
+        beta: int = 1,
+        num_iterations: int = 10010,
     ):
         self.goal = torch.tensor([ord(c) for c in goal], dtype=torch.float32)
         self.num_ants = num_ants
         self.evaporation_rate = evaporation_rate
         self.alpha = alpha
         self.beta = beta
         self.num_iterations = num_iterations
-        # Pheromone levels can be initialized for different paths (architectures)
+        # Pheromone levels can be initialized for different paths
+        # (architectures)
         self.pheromones = torch.ones(num_ants)
         self.solutions = []
 
@@ -65,15 +68,15 @@ def fitness(self, solution):
     def update_pheromones(self):
         """Update pheromone levels"""
         for i, solution in enumerate(self.solutions):
-            self.pheromones[i] = (
-                1 - self.evaporation_rate
-            ) * self.pheromones[i] + self.fitness(solution)
+            self.pheromones[i] = (1 - self.evaporation_rate) * self.pheromones[
+                i
+            ] + self.fitness(solution)
 
     def choose_next_path(self):
         """Choose the next path based on the pheromone levels"""
-        probabilities = (
-            self.pheromones ** self.alpha
-        ) * ((1.0 / (1 + self.pheromones)) ** self.beta)
+        probabilities = (self.pheromones**self.alpha) * (
+            (1.0 / (1 + self.pheromones)) ** self.beta
+        )
 
         probabilities /= probabilities.sum()
 
@@ -84,13 +87,11 @@ def optimize(self):
         for iteration in range(self.num_iterations):
             self.solutions = []
             for _ in range(self.num_ants):
-                # This is a placeholder. Actual implementation will define how ants traverse the search space.
+                # This is a placeholder. Actual implementation will define how
+                # ants traverse the search space.
                 solution = torch.randint(
-                    32, 
-                    127, 
-                    (
-                        len(self.goal),
-                    ), dtype=torch.float32)  # Random characters.
+                    32, 127, (len(self.goal),), dtype=torch.float32
+                )  # Random characters.
                 self.solutions.append(solution)
             self.update_pheromones()
 

swarms_torch/cellular_transformer.py

@@ -1,6 +1,7 @@
-import torch 
+import torch
 from torch import nn
 
+
 class TransformerCell(nn.Module):
     def __init__(
         self,
@@ -11,20 +12,17 @@ def __init__(
     ):
         super(TransformerCell, self).__init__()
         self.transformer = nn.Transformer(
-            input_dim,
-            nhead=nhead,
-            num_encoder_layers=num_layers
+            input_dim, nhead=nhead, num_encoder_layers=num_layers
         )
         self.neighborhood_size = neighborhood_size
-
-    def forward(self, x, neigbors):
 
+    def forward(self, x, neigbors):
         x = self.transformer(x, x)
 
         out = torch.cat([x] + neigbors, dim=0)
-        
+
         return out
-    
+
 
 class CellularSwarm(nn.Module):
     """
@@ -42,51 +40,39 @@ class CellularSwarm(nn.Module):
         input_dim (int): Input dimension
         nhead (int): Number of heads in the transformer cell
         time_steps (int): Number of time steps to run the network
-    
+
     Returns:
         torch.Tensor: Output tensor
-    
+
     Usage:
         >>> x = torch.randn(10, 32, 512)
         >>> model = CellularSwarm(cell_count=5, input_dim=512, nhead=8)
         >>> output = model(x)
         >>> print(output)
-    
-    
+
+
     """
-    def __init__(
-        self,
-        cell_count,
-        input_dim,
-        nhead,
-        time_steps=4
-    ):
+
+    def __init__(self, cell_count, input_dim, nhead, time_steps=4):
         super(CellularSwarm, self).__init__()
-        self.cells = nn.ModuleList([
-            TransformerCell(input_dim, nhead) for _ in range(cell_count)
-        ])
+        self.cells = nn.ModuleList(
+            [TransformerCell(input_dim, nhead) for _ in range(cell_count)]
+        )
         self.time_steps = time_steps
-    
+
     def forward(self, x):
         for _ in range(self.time_steps):
             for i, cell in enumerate(self.cells):
-                #get neighboring cells states
-                start_idx = max(
-                    0,
-                    i - cell.neighborhood_size
-                )
+                # get neighboring cells states
+                start_idx = max(0, i - cell.neighborhood_size)
 
-                end_idx = min(
-                    len(self.cells), i + cell.neighborhood_size + 1
-                )
+                end_idx = min(len(self.cells), i + cell.neighborhood_size + 1)
 
                 neighbors = [
-                    self.cells[j].transformer(x, x) for j in range(
-                        start_idx, 
-                        end_idx
-                    ) if j != i
+                    self.cells[j].transformer(x, x)
+                    for j in range(start_idx, end_idx)
+                    if j != i
                 ]
 
                 x = cell(x, neighbors)
         return x
-

swarms_torch/graph_cellular_automa.py

@@ -9,9 +9,9 @@ def __init__(self, input_dim, hidden_dim, output_dim):
         self.mlp = nn.Sequential(
             nn.Linear(input_dim, hidden_dim),
             nn.ReLU(),
-            nn.Linear(hidden_dim, output_dim)
+            nn.Linear(hidden_dim, output_dim),
         )
-        
+
     def forward(self, x):
         return self.mlp(x)
 
@@ -24,60 +24,70 @@ def __init__(self, input_dim, hidden_dim):
             nn.Linear(input_dim, hidden_dim),
             nn.ReLU(),
             nn.Linear(hidden_dim, 1),
-            nn.Sigmoid()  # for binary classification
+            nn.Sigmoid(),  # for binary classification
         )
-        
+
     def forward(self, x):
         return self.mlp(x)
 
+
 class WeightUpdateModel(nn.Module):
     def __init__(self, input_dim, hidden_dim):
         super(WeightUpdateModel, self).__init__()
 
         self.mlp = nn.Sequential(
-            nn.Linear(input_dim, hidden_dim),
-            nn.ReLU(),
-            nn.Linear(hidden_dim, 1)
-        )
-
+            nn.Linear(
+                input_dim, hidden_dim), nn.ReLU(), nn.Linear(
+                hidden_dim, 1))
+
     def forward(self, x):
         return self.mlp(x)
 
+
 class NDP(nn.Module):
     def __init__(self, embedding_dim, hidden_dim):
         super(NDP, self).__init__()
 
-        self.gc_automata = GraphCellularAutomata(embedding_dim, hidden_dim, embedding_dim)
+        self.gc_automata = GraphCellularAutomata(
+            embedding_dim, hidden_dim, embedding_dim
+        )
         self.replication_model = ReplicationModel(embedding_dim, hidden_dim)
-        self.weight_update_model = WeightUpdateModel(2 * embedding_dim, hidden_dim)
-
+        self.weight_update_model = WeightUpdateModel(
+            2 * embedding_dim, hidden_dim)
+
     def forward(self, node_embeddings, adjacency_matrix):
         # Update node embeddings using Graph Cellular Automata
         updated_embeddings = self.gc_automata(node_embeddings)
-        
+
         # Check which nodes need to replicate
         replication_decisions = self.replication_model(updated_embeddings)
-        
+
         # Weight update (assuming weighted network)
         num_nodes = node_embeddings.shape[0]
         edge_weights = torch.zeros((num_nodes, num_nodes))
 
         for i in range(num_nodes):
             for j in range(num_nodes):
-                combined_embedding = torch.cat((updated_embeddings[i], updated_embeddings[j]))
-
-                edge_weights[i, j] = self.weight_update_model(combined_embedding)
-
+                combined_embedding = torch.cat(
+                    (updated_embeddings[i], updated_embeddings[j])
+                )
+
+                edge_weights[i, j] = self.weight_update_model(
+                    combined_embedding)
+
         return updated_embeddings, replication_decisions, edge_weights
 
+
 # Usage examples
 embedding_dim = 16
 hidden_dim = 32
 node_embeddings = torch.rand((10, embedding_dim))  # For 10 nodes
 adjacency_matrix = torch.rand((10, 10))  # Dummy adjacency matrix for 10 nodes
 
 model = NDP(embedding_dim, hidden_dim)
-updated_embeddings, replication_decisions, edge_weights = model(node_embeddings, adjacency_matrix)
+updated_embeddings, replication_decisions, edge_weights = model(
+    node_embeddings, adjacency_matrix
+)
 
 print(updated_embeddings.shape)
 print(replication_decisions.shape)