diff --git a/.devcontainer/devcontainer.json b/.devcontainer/devcontainer.json
index 975b9d8..0908520 100644
--- a/.devcontainer/devcontainer.json
+++ b/.devcontainer/devcontainer.json
@@ -7,22 +7,16 @@
 		"dockerfile": "Dockerfile"
 	},
 	"features": {
-		"ghcr.io/devcontainers/features/anaconda:1": {
-			"version": "latest"
-		},
-		"ghcr.io/devcontainers/features/nvidia-cuda:2": {
-			"installCudnn": true,
-			"installCudnnDev": true,
-			"installNvtx": true,
-			"installToolkit": true,
-			"cudaVersion": "11.8",
-			"cudnnVersion": "automatic"
-		},
-		"ghcr.io/raucha/devcontainer-features/pytorch:1": {}
+		"ghcr.io/devcontainers/features/anaconda:1": {},
+		"ghcr.io/devcontainers/features/nvidia-cuda:2": {},
+		"ghcr.io/rocker-org/devcontainer-features/miniforge:2": {}
 	},
 	"runArgs": [
-		"--gpus=all"
-	]
+		"--gpus", "all"
+	],
+	// allow gpu
+
+
 	// Features to add to the dev container. More info: https://containers.dev/features.
 	// "features": {},
 
diff --git a/Adversarial_Observation/Attacks.py b/Adversarial_Observation/Attacks.py
index 40dc714..8b24667 100755
--- a/Adversarial_Observation/Attacks.py
+++ b/Adversarial_Observation/Attacks.py
@@ -1,74 +1,92 @@
+import numpy as np
 import torch
+import torch.nn.functional as F
+import matplotlib.pyplot as plt
 import logging
-import os
-from datetime import datetime
-from torch.nn import Softmax
-from .utils import fgsm_attack, compute_success_rate, log_metrics, visualize_adversarial_examples
-from .utils import seed_everything
-
-class AdversarialTester:
-    def __init__(self, model: torch.nn.Module, epsilon: float = 0.1, attack_method: str = 'fgsm', alpha: float = 0.01, 
-                 num_steps: int = 40, device=None, save_dir: str = './results', seed: int = 42):
-        seed_everything(seed)
-        self.model = model
+
+# Set up logging
+logging.basicConfig(level=logging.INFO)
+
+def fgsm_attack(input_batch_data: torch.Tensor, model: torch.nn.Module, input_shape: tuple, epsilon: float) -> torch.Tensor:
+    """
+    Apply the FGSM attack to input images given a pre-trained PyTorch model.
+    """
+    model.eval()
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    input_batch_data = input_batch_data.to(device)
+
+    adversarial_batch_data = []
+
+    for img in input_batch_data:
+        # Make a copy of the image and enable gradient tracking
+        img = img.clone().detach().unsqueeze(0).to(device)
+        img.requires_grad = True
+
+        # Forward pass
+        preds = model(img)
+        target = torch.argmax(preds, dim=1)
+        loss = F.cross_entropy(preds, target)
+
+        # Backward pass
+        model.zero_grad()
+        loss.backward()
+
+        # Generate perturbation
+        grad = img.grad.data
+        adversarial_img = img + epsilon * grad.sign()
+        adversarial_img = torch.clamp(adversarial_img, 0, 1)
+
+        adversarial_batch_data.append(adversarial_img.squeeze(0).detach())
+
+    return torch.stack(adversarial_batch_data)
+
+def compute_gradients(model, img, target_class):
+    preds = model(img)
+    target_score = preds[0, target_class]
+    return torch.autograd.grad(target_score, img)[0]
+
+def generate_adversarial_examples(input_batch_data, model, method='fgsm', **kwargs):
+    if method == 'fgsm':
+        return fgsm_attack(input_batch_data, model, **kwargs)
+    # Implement other attack methods as needed
+
+def gradient_ascent(input_image, model, input_shape, target_class, num_iterations=100, step_size=0.01):
+    model.eval()
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    input_image = input_image.to(device).detach().requires_grad_(True)
+
+    for _ in range(num_iterations):
+        gradients = compute_gradients(model, input_image.reshape(input_shape), target_class)
+        input_image = input_image + step_size * gradients.sign()
+        input_image = torch.clamp(input_image, 0, 1)
+        input_image = input_image.detach().requires_grad_(True)
+
+    return input_image.cpu().detach().numpy()
+
+def gradient_map(input_image, model, input_shape):
+    model.eval()
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    input_image = torch.tensor(input_image).to(device).detach().requires_grad_(True)
+
+    preds = model(input_image.reshape(input_shape))
+    target_class = torch.argmax(preds)
+    loss = F.cross_entropy(preds, target_class.unsqueeze(0))
+
+    model.zero_grad()
+    loss.backward()
+
+    gradient = input_image.grad.data.cpu().numpy()
+    gradient = np.abs(gradient).mean(axis=1)  # Average over channels if needed
+    return gradient
+    
+def visualize_adversarial_examples(original, adversarial):
+    # Code to visualize original vs adversarial images
+    pass
+
+def log_metrics(success_rate, average_perturbation):
+    logging.info(f'Success Rate: {success_rate}, Average Perturbation: {average_perturbation}')
+
+class Config:
+    def __init__(self, epsilon=0.1, attack_method='fgsm'):
         self.epsilon = epsilon
-        self.attack_method = attack_method
-        self.alpha = alpha
-        self.num_steps = num_steps
-        self.device = device or ("cuda" if torch.cuda.is_available() else "cpu")
-        self.save_dir = save_dir
-
-        # Create save directory if it doesn't exist
-        os.makedirs(self.save_dir, exist_ok=True)
-        self.model.to(self.device)
-        self.model.eval()
-
-        self._setup_logging()
-
-    def _setup_logging(self):
-        log_file = os.path.join(self.save_dir, f"attack_log_{datetime.now().strftime('%Y%m%d_%H%M%S')}.log")
-        logging.basicConfig(filename=log_file, level=logging.DEBUG)
-        logging.info(f"Started adversarial testing at {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}")
-        logging.info(f"Using model: {self.model.__class__.__name__}")
-        logging.info(f"Attack Method: {self.attack_method}, Epsilon: {self.epsilon}, Alpha: {self.alpha}, Steps: {self.num_steps}")
-
-    def test_attack(self, input_batch_data: torch.Tensor):
-        input_batch_data = input_batch_data.to(self.device)
-        adversarial_images = self._generate_adversarial_images(input_batch_data)
-
-        # Save and log images
-        self._save_images(input_batch_data, adversarial_images)
-        self._compute_and_log_metrics(input_batch_data, adversarial_images)
-
-    def _generate_adversarial_images(self, input_batch_data: torch.Tensor):
-        logging.info(f"Starting attack with method: {self.attack_method}")
-        if self.attack_method == 'fgsm':
-            return fgsm_attack(input_batch_data, self.model, self.epsilon, self.device)
-        else:
-            raise ValueError(f"Unsupported attack method: {self.attack_method}")
-
-    def _save_images(self, original_images: torch.Tensor, adversarial_images: torch.Tensor):
-        for i in range(original_images.size(0)):
-            original_image_path = os.path.join(self.save_dir, f"original_{i}.png")
-            adversarial_image_path = os.path.join(self.save_dir, f"adversarial_{i}.png")
-            visualize_adversarial_examples(original_images, adversarial_images, original_image_path, adversarial_image_path)
-
-    def _compute_and_log_metrics(self, original_images: torch.Tensor, adversarial_images: torch.Tensor):
-        original_predictions = torch.argmax(self.model(original_images), dim=1)
-        adversarial_predictions = torch.argmax(self.model(adversarial_images), dim=1)
-
-        success_rate = compute_success_rate(original_predictions, adversarial_predictions)
-        average_perturbation = torch.mean(torch.abs(adversarial_images - original_images)).item()
-
-        log_metrics(success_rate, average_perturbation)
-        self._save_metrics(success_rate, average_perturbation)
-
-        logging.info(f"Success Rate: {success_rate:.4f}, Average Perturbation: {average_perturbation:.4f}")
-
-    def _save_metrics(self, success_rate: float, avg_perturbation: float):
-        """
-        Save the metrics (success rate and average perturbation) to a file.
-        """
-        metrics_file = os.path.join(self.save_dir, "attack_metrics.txt")
-        with open(metrics_file, 'a') as f:
-            f.write(f"Success Rate: {success_rate:.4f}, Average Perturbation: {avg_perturbation:.4f}\n")
\ No newline at end of file
+        self.attack_method = attack_method
\ No newline at end of file
diff --git a/Adversarial_Observation/BirdParticle.py b/Adversarial_Observation/BirdParticle.py
index a9ebcce..ccaf63d 100755
--- a/Adversarial_Observation/BirdParticle.py
+++ b/Adversarial_Observation/BirdParticle.py
@@ -1,42 +1,44 @@
-import tensorflow as tf
+import torch
+import torch.nn.functional as F
 import numpy as np
 
 class BirdParticle:
     """
-    Represents a particle in the Particle Swarm Optimization (PSO) algorithm for adversarial attacks.
-    
-    The BirdParticle class encapsulates the state of each particle, including its position, velocity,
-    fitness evaluation, and the updates to its velocity and position based on the PSO algorithm.
+    Represents a particle in the Particle Swarm Optimization (PSO) algorithm for adversarial attacks (PyTorch version).
     """
 
-    def __init__(self, model: tf.keras.Model, input_data: tf.Tensor, target_class: int, num_iterations: int = 20,
-                 velocity: tf.Tensor = None, inertia_weight: float = 0.5, 
+    def __init__(self, model: torch.nn.Module, input_data: torch.Tensor, target_class: int, num_iterations: int = 20,
+                 velocity: torch.Tensor = None, inertia_weight: float = 0.5, 
                  cognitive_weight: float = 1.0, social_weight: float = 1.0, 
-                 momentum: float = 0.9, clip_value_position: float = 1.0):
+                 momentum: float = 0.9, clip_value_position: float = 1.0, device='cpu'):
         """
         Initialize a particle in the PSO algorithm.
         
         Args:
-            model (tf.keras.Model): The model to attack.
-            input_data (tf.Tensor): The input data (image) to attack.
+            model (torch.nn.Module): The model to attack.
+            input_data (torch.Tensor): The input data (image) to attack.
             target_class (int): The target class for misclassification.
-            velocity (tf.Tensor, optional): The initial velocity for the particle's movement. Defaults to zero velocity if not provided.
-            inertia_weight (float): The inertia weight for the velocity update. Default is 0.5.
-            cognitive_weight (float): The cognitive weight for the velocity update. Default is 1.0.
-            social_weight (float): The social weight for the velocity update. Default is 1.0.
-            momentum (float): The momentum for the velocity update. Default is 0.9.
+            velocity (torch.Tensor, optional): Initial velocity; defaults to zero.
+            inertia_weight (float): Inertia weight for velocity update.
+            cognitive_weight (float): Cognitive weight for velocity update.
+            social_weight (float): Social weight for velocity update.
+            momentum (float): Momentum for velocity update.
+            clip_value_position (float): Max absolute value to clip position.
+            device (str): Device to run on.
         """
-        self.model = model
+        self.device = device
+        self.model = model.to(device)
         self.num_iterations = num_iterations
-        self.original_data = tf.identity(input_data)  # Clone the input data
+        self.original_data = input_data.clone().detach().to(device)
+        self.position = input_data.clone().detach().to(device)
         self.target_class = target_class
-        self.best_position = tf.identity(input_data)  # Clone the input data
+        self.velocity = velocity.clone().detach().to(device) if velocity is not None else torch.zeros_like(input_data).to(device)
+        self.best_position = self.position.clone().detach()
         self.best_score = -np.inf
-        self.position = tf.identity(input_data)  # Clone the input data
-        self.velocity = velocity if velocity is not None else tf.zeros_like(input_data)
-        self.history = [self.position]
+        self.history = [self.position.clone().detach()]
         self.clip_value_position = clip_value_position
-        # Class attributes
+
+        # PSO hyperparameters
         self.inertia_weight = inertia_weight
         self.cognitive_weight = cognitive_weight
         self.social_weight = social_weight
@@ -45,52 +47,47 @@ def __init__(self, model: tf.keras.Model, input_data: tf.Tensor, target_class: i
     def fitness(self) -> float:
         """
         Compute the fitness score for the particle, which is the softmax probability of the target class.
-        
-        Higher fitness scores correspond to better success in the attack (misclassifying the image into the target class).
-        
         Returns:
-            float: Fitness score for this particle (higher is better).
+            float: Target class softmax probability.
         """
-        output = self.model(self.position)  # Add batch dimension and pass through the model
-        probabilities = tf.nn.softmax(output, axis=1)  # Get probabilities for each class
-        target_prob = probabilities[:, self.target_class]  # Target class probability
-        
-        return target_prob.numpy().item()  # Return the target class probability as fitness score
+        self.model.eval()
+        with torch.no_grad():
+            input_tensor = self.position 
+            output = self.model(input_tensor.to(self.device))
+            probabilities = F.softmax(output, dim=1)
+            target_prob = probabilities[:, self.target_class]
+            return target_prob.item()
 
-    def update_velocity(self, global_best_position: tf.Tensor) -> None:
+    def update_velocity(self, global_best_position: torch.Tensor) -> None:
         """
-        Update the velocity of the particle based on the PSO update rule.
+        Update the particle's velocity using the PSO rule.
         
         Args:
-            global_best_position (tf.Tensor): The global best position found by the swarm.
+            global_best_position (torch.Tensor): Global best position in the swarm.
         """
+        r1 = torch.rand_like(self.position).to(self.device)
+        r2 = torch.rand_like(self.position).to(self.device)
+        
         inertia = self.inertia_weight * self.velocity
-        cognitive = self.cognitive_weight * tf.random.uniform(self.position.shape) * (self.best_position - self.position)
-        social = self.social_weight * tf.random.uniform(self.position.shape) * (global_best_position - self.position)
-
-        # Apply momentum to velocity update:
-        self.velocity = self.momentum * self.velocity + inertia + cognitive + social  # Apply momentum
+        cognitive = self.cognitive_weight * r1 * (self.best_position - self.position)
+        social = self.social_weight * r2 * (global_best_position.to(self.device) - self.position)
+        
+        self.velocity = self.momentum * self.velocity + inertia + cognitive + social
 
     def update_position(self) -> None:
         """
-        Update the position of the particle based on the updated velocity.
-        
-        The position is updated directly without any bounds checking.
+        Update the particle's position based on the new velocity.
         """
-        self.position = self.position + self.velocity  # Update position directly without clipping
-        self.position = tf.clip_by_value(self.position + self.velocity, 0, 1) # Ensure position stays within bounds
-        self.history.append(tf.identity(self.position))  # Store the position history
-        self.position = tf.clip_by_value(self.position, -self.clip_value_position, self.clip_value_position)
-        
+        self.position = self.position + self.velocity
+        self.position = torch.clamp(self.position, 0.0, 1.0)  # Keep values in [0, 1]
+        self.position = torch.clamp(self.position, -self.clip_value_position, self.clip_value_position)
+        self.history.append(self.position.clone().detach())
+
     def evaluate(self) -> None:
         """
-        Evaluate the fitness of the current particle and update its personal best.
-        
-        The fitness score is calculated using the target class probability. If the current fitness score
-        is better than the personal best, update the best position and score.
+        Evaluate current fitness and update personal best if needed.
         """
-        score = self.fitness()  # Get the current fitness score based on the perturbation
-        if score > self.best_score:  # If score is better than the personal best, update the best position
+        score = self.fitness()
+        if score > self.best_score:
             self.best_score = score
-            self.best_position = tf.identity(self.position)  # Clone the current position
-
+            self.best_position = self.position.clone().detach()
diff --git a/Adversarial_Observation/Swarm.py b/Adversarial_Observation/Swarm.py
index eca34fe..fa016de 100755
--- a/Adversarial_Observation/Swarm.py
+++ b/Adversarial_Observation/Swarm.py
@@ -1,222 +1,159 @@
 import os
 import logging
 from typing import List
-from Adversarial_Observation.BirdParticle import BirdParticle
-from tqdm import tqdm
-import tensorflow as tf
+import torch
 import numpy as np
+from tqdm import tqdm
 import matplotlib.pyplot as plt
+from Adversarial_Observation.BirdParticle import BirdParticle
+
 
 class ParticleSwarm:
-    """
-    Represents the Particle Swarm Optimization (PSO) algorithm applied to adversarial attacks.
-    
-    The ParticleSwarm class manages the swarm of particles and optimizes the perturbations on the input data (image)
-    to misclassify it into the target class.
-    """
-
-    def __init__(self, model: tf.keras.Model, input_set: np.ndarray, starting_class: int, target_class: int,
+    def __init__(self, model: torch.nn.Module, input_set: np.ndarray, starting_class: int, target_class: int,
                  num_iterations: int = 20, save_dir: str = 'results', inertia_weight: float = 0.5,
-                 cognitive_weight: float = .5, social_weight: float = .5, momentum: float = 0.9,
+                 cognitive_weight: float = 0.5, social_weight: float = 0.5, momentum: float = 0.9,
                  clip_value_position: float = 0.2, enable_logging: bool = False, device: str = 'cpu'):
-        """
-        Initialize the Particle Swarm Optimization (PSO) for adversarial attacks.
-        
-        Args:
-            model (tf.keras.Model): The model to attack.
-            input_set (np.ndarray): The batch of input images to attack as a NumPy array.
-            target_class (int): The target class for misclassification.
-            num_iterations (int): The number of optimization iterations.
-            save_dir (str): The directory to save output images and logs.
-            inertia_weight (float): The inertia weight for the velocity update.
-            cognitive_weight (float): The cognitive weight for the velocity update.
-            social_weight (float): The social weight for the velocity update.
-            momentum (float): The momentum for the velocity update.
-            clip_value_position (float): The velocity clamp to limit the velocity.
-            device (str): The device for computation ('cpu' or 'gpu'). Default is 'cpu'.
-        """
-        self.model = model
-        self.input_set = tf.convert_to_tensor(input_set, dtype=tf.float32)  # Convert NumPy array to TensorFlow tensor
-        self.start_class = starting_class # The starting class index
-        self.target_class = target_class  # The target class index
+        self.model = model.to(device).eval()
+        self.device = torch.device(device)
+
+        self.input_set = input_set.float().view(-1, 1, 28, 28).to(self.device)
+
+        self.start_class = starting_class
+        self.target_class = target_class
         self.num_iterations = num_iterations
-        self.save_dir = save_dir  # Directory to save perturbed images
+        self.save_dir = save_dir
         self.enable_logging = enable_logging
-        self.device = device  # Device ('cpu' or 'gpu')
 
         self.particles: List[BirdParticle] = [
             BirdParticle(model, self.input_set[i:i + 1], target_class,
-                         inertia_weight=inertia_weight, cognitive_weight=cognitive_weight, social_weight=social_weight,
-                         momentum=momentum, clip_value_position=clip_value_position) 
+                         inertia_weight=inertia_weight, cognitive_weight=cognitive_weight,
+                         social_weight=social_weight, momentum=momentum,
+                         clip_value_position=clip_value_position, device=self.device)
             for i in range(len(input_set))
         ]
-        
-        self.global_best_position = tf.zeros_like(self.input_set[0]) # Global best position
-        self.global_best_score = -float('inf')  # Initialize with a very low score
-        
-        self.fitness_history: List[float] = []  # History of fitness scores to track progress
-
-        # Make output folder
-        iteration_dir = self.save_dir
-        os.makedirs(iteration_dir, exist_ok=True)
+
+        self.global_best_position = torch.zeros_like(self.input_set[0])
+        self.global_best_score = -float('inf')
+        self.fitness_history: List[float] = []
+
+        os.makedirs(self.save_dir, exist_ok=True)
         if self.enable_logging:
             self.setup_logging()
             self.log_progress(-1)
 
     def setup_logging(self):
-        """
-        Set up logging for each iteration. Each iteration will have a separate log file.
-        Also prints logs to the terminal.
-        """
-        log_file = os.path.join(self.save_dir, f'iteration_log.log')
+        log_file = os.path.join(self.save_dir, 'iteration_log.log')
         self.logger = logging.getLogger()
-        
-        # Create a file handler to save logs to a file
+        self.logger.setLevel(logging.INFO)
+
         file_handler = logging.FileHandler(log_file)
-        file_handler.setFormatter(logging.Formatter('%(asctime)s - %(message)s'))
-        
-        # Create a stream handler to output logs to the console (terminal)
         stream_handler = logging.StreamHandler()
-        stream_handler.setFormatter(logging.Formatter('%(message)s'))  # Keep it simple for console
-        
+
+        formatter = logging.Formatter('%(asctime)s - %(message)s')
+        file_handler.setFormatter(formatter)
+        stream_handler.setFormatter(logging.Formatter('%(message)s'))
+
         self.logger.addHandler(file_handler)
-        self.logger.addHandler(stream_handler)  # Add the stream handler to print to terminal
-        self.logger.setLevel(logging.INFO)
+        self.logger.addHandler(stream_handler)
 
-        # Log class initialization details
         self.logger.info(f"\n{'*' * 60}")
         self.logger.info(f"ParticleSwarm Optimization (PSO) for Adversarial Attack")
         self.logger.info(f"{'-' * 60}")
         self.logger.info(f"Model: {self.model.__class__.__name__}")
-        self.logger.info(f"Target Class: {self.target_class} (This is the class we want to misclassify the image into)")
-        self.logger.info(f"Number of Iterations: {self.num_iterations} (Optimization steps)")
+        self.logger.info(f"Target Class: {self.target_class}")
+        self.logger.info(f"Number of Iterations: {self.num_iterations}")
         self.logger.info(f"Save Directory: {self.save_dir}")
         self.logger.info(f"{'*' * 60}")
 
     def log_progress(self, iteration: int):
-        """
-        Log detailed information for the current iteration in a manually formatted table.
-        
-        Args:
-            iteration (int): The current iteration of the optimization process.
-        """
         if not self.enable_logging:
             return
 
-        # Log the header for the iteration
         self.logger.info(f"\n{'-'*60}")
         self.logger.info(f"Iteration {iteration + 1}/{self.num_iterations}")
         self.logger.info(f"{'='*60}")
-    
-        # Table header
-        header = f"{'Particle':<10}{'Original Pred':<15}{'Perturbed Pred':<18}{'Orig Start Prob':<20}{'Pert Start Prob':<20}{'Orig Target Prob':<20}" \
-                 f"{'Pert Target Prob':<20}{'Personal Best':<20}{'Global Best':<20}"
+
+        header = f"{'Particle':<10}{'Original Pred':<15}{'Perturbed Pred':<18}{'Orig Start Prob':<20}{'Pert Start Prob':<20}{'Orig Target Prob':<20}{'Pert Target Prob':<20}{'Personal Best':<20}{'Global Best':<20}"
         self.logger.info(header)
         self.logger.info(f"{'-'*60}")
-    
-        # Log particle information
+
         for i, particle in enumerate(self.particles):
-            # Get original and perturbed outputs
-            original_output = self.model(particle.original_data)  # Pass through the model
-            perturbed_output = self.model(particle.position)  # Pass through the model
-    
-            # Get predicted classes
-            original_pred = tf.argmax(original_output, axis=1).numpy().item()
-            perturbed_pred = tf.argmax(perturbed_output, axis=1).numpy().item()
-    
-            # Get softmax probabilities
-            original_probs = tf.nn.softmax(original_output, axis=1)
-            perturbed_probs = tf.nn.softmax(perturbed_output, axis=1)
-
-            # Get starting class probabilities (how far away we've moved)
-            original_prob_start = original_probs[0, self.start_class].numpy().item()
-            perturbed_prob_start = perturbed_probs[0, self.start_class].numpy().item()
-
-            # Get target class probabilities
-            original_prob_target = original_probs[0, self.target_class].numpy().item()
-            perturbed_prob_target = perturbed_probs[0, self.target_class].numpy().item()
-    
-            # Log each particle's data in a formatted row
-            self.logger.info(f"{i+1:<10}{original_pred:<15}{perturbed_pred:<18}{original_prob_start:<20.4f}{perturbed_prob_start:<20.4f}"
-                            f"{original_prob_target:<20.4f}{perturbed_prob_target:<20.4f}{particle.best_score:<20.4f}{self.global_best_score:<20.4f}")
-    
+            with torch.no_grad():
+                original_output = self.model(particle.original_data)
+                perturbed_output = self.model(particle.position)
+
+                original_probs = torch.softmax(original_output, dim=1)
+                perturbed_probs = torch.softmax(perturbed_output, dim=1)
+
+                original_pred = original_output.argmax(dim=1).item()
+                perturbed_pred = perturbed_output.argmax(dim=1).item()
+
+                orig_start_prob = original_probs[0, self.start_class].item()
+                pert_start_prob = perturbed_probs[0, self.start_class].item()
+                orig_target_prob = original_probs[0, self.target_class].item()
+                pert_target_prob = perturbed_probs[0, self.target_class].item()
+
+            self.logger.info(f"{i+1:<10}{original_pred:<15}{perturbed_pred:<18}"
+                             f"{orig_start_prob:<20.4f}{pert_start_prob:<20.4f}"
+                             f"{orig_target_prob:<20.4f}{pert_target_prob:<20.4f}"
+                             f"{particle.best_score:<20.4f}{self.global_best_score:<20.4f}")
+
         self.logger.info(f"{'='*60}")
-    
+
     def optimize(self):
-        """
-        Run the Particle Swarm Optimization process to optimize the perturbations.
-        """
-        with tf.device(f"/{self.device}:0"):  # Use the GPU/CPU based on the flag
-            for iteration in tqdm(range(self.num_iterations), desc="Running Swarm"):
-                # Update particles and velocities, evaluate them, and track global best
-                for particle in self.particles:
-                    particle.evaluate()
-                    particle.update_velocity(self.global_best_position)  # No need to pass inertia_weight explicitly
-                    particle.update_position()
-                
-                # Update the global best based on the personal best scores of particles
-                best_particle = max(self.particles, key=lambda p: p.best_score)
-                if best_particle.best_score > self.global_best_score:
-                    self.global_best_score = best_particle.best_score
-                    self.global_best_position = tf.identity(best_particle.best_position)
-                
-                self.log_progress(iteration)
-
-    def reduce_excess_perturbations(self, original_img: np.ndarray, target_label: int, model_shape: tuple = (1, 28, 28, 1)) -> np.ndarray:
-        """
-        Reduces excess perturbations in adversarial images while ensuring the misclassification remains.
-
-        Args:
-            original_img (np.ndarray): The original (clean) image.
-            target_label (int): The label we want the adversarial image to produce.
-            model_shape (tuple): The expected shape of the model input.
-
-        Returns:
-            list[np.ndarray]: A list of denoised adversarial images.
-        """
-        denoised_adv = []
-        total_pixels = np.prod(original_img.shape)  # Total pixels in the image
+        for iteration in tqdm(range(self.num_iterations), desc="Running Swarm"):
+            for particle in self.particles:
+                particle.evaluate()
+                particle.update_velocity(self.global_best_position)
+                particle.update_position()
 
-        for adv_particle in tqdm(self.particles, desc="Processing Particles", unit="particle"):
-            adv_img = np.copy(adv_particle.position).reshape(original_img.shape)  # Copy to avoid modifying the original
+            best_particle = max(self.particles, key=lambda p: p.best_score)
+            if best_particle.best_score > self.global_best_score:
+                self.global_best_score = best_particle.best_score
+                self.global_best_position = best_particle.best_position.clone()
 
-            if original_img.shape != adv_img.shape:
-                raise ValueError("original_img and adv_img must have the same shape.")
+            self.log_progress(iteration)
 
-            # Iterate over every pixel coordinate in the image with tqdm progress bar
-            with tqdm(total=total_pixels, desc="Processing Pixels", unit="pixel", leave=False) as pbar:
+    def reduce_excess_perturbations(self, original_img: np.ndarray, target_label: int, model_shape: tuple = (1, 1, 28, 28)) -> List[np.ndarray]:
+        denoised_adv = []
+        total_pixels = np.prod(original_img.shape)
+
+        for adv_particle in tqdm(self.particles, desc="Processing Particles"):
+            adv_img = adv_particle.position.clone().detach().cpu().numpy().reshape(original_img.shape)
+            orig = original_img.copy()
+
+            with tqdm(total=total_pixels, desc="Processing Pixels", leave=False) as pbar:
                 for idx in np.ndindex(original_img.shape):
-                    if original_img[idx] == adv_img[idx]:  # Ignore unchanged pixels
+                    if orig[idx] == adv_img[idx]:
                         pbar.update(1)
                         continue
 
-                    # Store old adversarial value
                     old_val = adv_img[idx]
+                    adv_img[idx] = orig[idx]
 
-                    # Try restoring the pixel to original
-                    adv_img[idx] = original_img[idx]
-
-                    # Check if the label is still the target label
-                    output = self.model(adv_img.reshape(model_shape))
-                    softmax_output = tf.nn.softmax(tf.squeeze(output), axis=0).numpy()
-                    current_label = np.argmax(softmax_output)
+                    test_img = torch.from_numpy(adv_img.reshape(model_shape)).float().to(self.device)
+                    with torch.no_grad():
+                        output = self.model(test_img)
+                        pred = output.softmax(dim=1).argmax(dim=1).item()
 
-                    if current_label != target_label:
-                        # If misclassification is lost, try halfway adjustment
-                        adv_img[idx] = old_val
-                        adv_img[idx] += (original_img[idx] - adv_img[idx]) * 0.5
+                    if pred != target_label:
+                        adv_img[idx] = old_val + (orig[idx] - old_val) * 0.5
+                        test_img = torch.from_numpy(adv_img.reshape(model_shape)).float().to(self.device)
+                        with torch.no_grad():
+                            output = self.model(test_img)
+                            pred = output.softmax(dim=1).argmax(dim=1).item()
 
-                        # Recheck if the label is still the target
-                        output = self.model(adv_img.reshape(model_shape))
-                        softmax_output = tf.nn.softmax(tf.squeeze(output), axis=0).numpy()
-                        current_label = np.argmax(softmax_output)
-
-                        if current_label != target_label:
-                            # If misclassification is still lost, revert back
+                        if pred != target_label:
                             adv_img[idx] = old_val
 
-                    pbar.update(1)  # Update pixel progress
+                    pbar.update(1)
 
             denoised_adv.append(adv_img)
 
         return denoised_adv
+
+    def getBest(self) -> np.ndarray:
+        return self.global_best_position.detach().cpu().numpy()
+
+    def getPoints(self) -> List[np.ndarray]:
+        return [particle.position.detach().cpu().numpy() for particle in self.particles]
diff --git a/Adversarial_Observation/utils.py b/Adversarial_Observation/utils.py
index 701a068..0ed2cbd 100644
--- a/Adversarial_Observation/utils.py
+++ b/Adversarial_Observation/utils.py
@@ -154,7 +154,7 @@ def load_MNIST_model():
 
     return model
 
-def load_data(batch_size=32):
+def load_MNIST_data(batch_size=32):
     """
     Loads MNIST train and test data and prepares it for evaluation.
 
diff --git a/manuscripts/PEARC24/MNIST/1_build_train.py b/manuscripts/PEARC24/MNIST/1_build_train.py
index c7d8c49..100df32 100644
--- a/manuscripts/PEARC24/MNIST/1_build_train.py
+++ b/manuscripts/PEARC24/MNIST/1_build_train.py
@@ -62,7 +62,7 @@ def testModel(model, test_loader, filename):
 
 def main():
     # Seed everything
-    AO.utils.seedEverything(42)
+    AO.utils.seed_everything(42)
 
     train_loader, test_loader = AO.utils.load_MNIST_data()
     model = AO.utils.load_MNIST_model()
diff --git a/manuscripts/PEARC24/MNIST/2_adversarial_attack.py b/manuscripts/PEARC24/MNIST/2_adversarial_attack.py
index 98c1973..7bd5666 100644
--- a/manuscripts/PEARC24/MNIST/2_adversarial_attack.py
+++ b/manuscripts/PEARC24/MNIST/2_adversarial_attack.py
@@ -8,7 +8,7 @@
 
 def main():
     # Seed everything
-    AO.utils.seedEverything(6991)
+    AO.utils.seed_everything(6991)
 
     train_loader, test_loader = AO.utils.load_MNIST_data()
     model = AO.utils.load_MNIST_model()
diff --git a/manuscripts/PEARC24/MNIST/3_create_APSO.py b/manuscripts/PEARC24/MNIST/3_create_APSO.py
index ccdae42..72d1825 100644
--- a/manuscripts/PEARC24/MNIST/3_create_APSO.py
+++ b/manuscripts/PEARC24/MNIST/3_create_APSO.py
@@ -1,4 +1,4 @@
-from Swarm_Observer import Swarm 
+from Adversarial_Observation import Swarm as Swarm
 import Adversarial_Observation as AO
 import os
 import torch
@@ -105,10 +105,13 @@ def plotSwarm(swarm, umap_model, epoch, otherpoints):
             y = [i[1] for i in otherpoints[key]]
             ax.scatter(x, y, label=key)
     # Get all the points
-    points = swarm.getPoints()
+    points = np.array(swarm.getPoints())
 
     # Transform the points
-    points = umap_model.transform(points)
+    points_np = points.cpu().detach().numpy() if hasattr(points, 'cpu') else points
+    points_reshaped = points_np.reshape(points_np.shape[0], -1)
+    points = umap_model.transform(points_reshaped)
+
     ax.scatter(points[:, 0], points[:, 1], c='black', label='Swarm', marker='x', s=100)                
     ax.legend()
     ax.set_title(f'Epoch: {epoch}')
@@ -120,7 +123,7 @@ def plotSwarm(swarm, umap_model, epoch, otherpoints):
 
 def plotImages(swarm, epoch):
     points = swarm.getPoints()
-    points = points.reshape(-1, 1, 28, 28)
+    points =np.array(points).reshape(-1, 1, 28, 28)
     #  plot the best
     best = swarm.getBest()
     best = best.reshape(28, 28)
@@ -130,9 +133,9 @@ def plotImages(swarm, epoch):
     ax.imshow(best, cmap='gray')
     ax.axis('off')
     global label
-    ax.set_title(f'Confidence of {label}: {round(swarm.model(best.reshape(1,1,28,28).to(torch.float32).to(device))[0][label].item(),2)}')
+    ax.set_title(f'Confidence of {label}: {round(swarm.model(torch.tensor(best).reshape(1,1,28,28).to(torch.float32).to(device))[0][label].item(),2)}')
 
-    grad = AO.Attacks.gradient_map(best.reshape(1,1,28,28), swarm.model, (1,1, 28, 28))[0].reshape(28,28)
+    grad = AO.Attacks.gradient_map(torch.tensor(best).reshape(1,1,28,28), swarm.model, (1,1, 28, 28))[0].reshape(28,28)
     grad = np.abs(grad)  # Make the gradients absolute for better visualization
     grad_normalized = (grad - np.min(grad)) / (np.max(grad) - np.min(grad))  # Normalize gradients to 0-1
 
@@ -152,9 +155,9 @@ def plotImages(swarm, epoch):
         fig, ax = plt.subplots(figsize=(6, 6))
         ax.imshow(point, cmap='gray')
         ax.axis('off')
-        ax.set_title(f'Confidence of {label}: {round(swarm.model(point.reshape(1,1,28,28).to(torch.float32).to(device))[0][label].item(),3)}')
+        ax.set_title(f'Confidence of {label}: {round(swarm.model(torch.tensor(point).reshape(1,1,28,28).to(torch.float32).to(device))[0][label].item(),3)}')
 
-        grad = AO.Attacks.gradient_map(point.reshape(1,1,28,28), swarm.model, (1,1, 28, 28))[0].reshape(28,28)
+        grad = AO.Attacks.gradient_map(torch.tensor(point).reshape(1,1,28,28), swarm.model, (1,1, 28, 28))[0].reshape(28,28)
         grad = np.abs(grad)  # Make the gradients absolute for better visualization
         grad_normalized = (grad - np.min(grad)) / (np.max(grad) - np.min(grad))  # Normalize gradients to 0-1
 
@@ -170,10 +173,26 @@ def plotImages(swarm, epoch):
         plt.close()
 
 def runSwarm(inital_points, model, device, umap_model, epochs, otherpoints):
-    APSO = Swarm.PSO(inital_points, cost_func, model, 1, .5, .5)
+    APSO = Swarm.ParticleSwarm(
+        model=model,
+        input_set=inital_points,
+        starting_class=label,
+        target_class=label,
+        num_iterations=1,  
+        save_dir='results',
+        inertia_weight=0.5,
+        cognitive_weight=0.5,
+        social_weight=0.5,
+        momentum=0.9,
+        clip_value_position=0.2,
+        enable_logging=True,
+        device='cuda'  # or 'cpu' if you want to run on CPU
+    )
+
+
     plotSwarm(APSO, umap_model, 0, otherpoints)
     for i in tqdm.tqdm(range(1, 1+epochs), desc='Running Swarm', total=epochs):
-        APSO.step()
+        APSO.optimize()
         plotSwarm(APSO, umap_model, i, otherpoints)
     plotSwarm(APSO, umap_model, epochs+1, otherpoints)
 
diff --git a/manuscripts/PEARC24/MNIST/3p1_APSO_W_Anchor.py b/manuscripts/PEARC24/MNIST/3p1_APSO_W_Anchor.py
index f36c13d..c18f46a 100644
--- a/manuscripts/PEARC24/MNIST/3p1_APSO_W_Anchor.py
+++ b/manuscripts/PEARC24/MNIST/3p1_APSO_W_Anchor.py
@@ -1,167 +1,177 @@
-from Swarm_Observer import Swarm 
-import Adversarial_Observation as AO
+# --- Imports ---
 import os
 import torch
 import umap
 import numpy as np
 import tqdm
 import matplotlib.pyplot as plt
-from scipy.sparse import csr_matrix
+
+import Adversarial_Observation as AO
+from Adversarial_Observation.Swarm import ParticleSwarm  # ← This is your custom swarm class
+
+# --- Global Config ---
+optimize = 3  # Target class for the attack
 
 
+# --- Main Function ---
 def main():
-    #  load the model
+    # Load model and data
     model = AO.utils.load_MNIST_model()
-    #  load the data
     train_loader, test_loader = AO.utils.load_MNIST_data()
-    #  update weights 
+
     if os.path.isfile('MNIST_cnn.pt'):
         model.load_state_dict(torch.load('MNIST_cnn.pt'))
     else:
-        raise Exception('MNIST_cnn.pt not found. Please run 1_build_train.py first')
-    
-    #  set the model to evaluation mode
-    model.eval()
+        raise FileNotFoundError("MNIST_cnn.pt not found. Please run 1_build_train.py first.")
 
+    model.eval()
 
-    otherpoints  = {}
+    # Setup UMAP and collect test points
     umap_model = umap.UMAP()
-
-    accumulated_data = []  # Accumulate data samples
-    targets = []  # Accumulate targets
-    global optimize
+    accumulated_data = []
+    targets = []
+    otherpoints = {}
     anchor = None
-    for idx, (data, target) in tqdm.tqdm(enumerate(test_loader), total=len(test_loader), desc='Training UMAP'):
+
+    for idx, (data, target) in tqdm.tqdm(enumerate(test_loader), total=len(test_loader), desc='Preparing UMAP'):
         if idx > 10:
             break
         for img, label in zip(data, target):
-            if label == optimize:
+            if label.item() == optimize and anchor is None:
                 anchor = img
-            img = img.reshape(1, 28*28)
-            img = img.detach().numpy()
-            accumulated_data.append(img)  # Accumulate data samples
-            targets.append(label)  # Accumulate targets
-            
-    # Convert accumulated data to a single NumPy array
-    accumulated_data = np.concatenate(accumulated_data, axis=0)
+            img_np = img.reshape(1, 28 * 28).detach().cpu().numpy()
+            accumulated_data.append(img_np)
+            targets.append(label)
 
+    accumulated_data = np.concatenate(accumulated_data, axis=0)
     reduced = umap_model.fit_transform(accumulated_data)
 
-    # Store the reduced data points
-    for idx, label in tqdm.tqdm(enumerate(targets), total=len(targets), desc='Storing reduced data'):
-        if label.item() not in otherpoints.keys():
-            otherpoints[label.item()] = []
-        otherpoints[label.item()].append(reduced[idx])
+    for idx, label in tqdm.tqdm(enumerate(targets), total=len(targets), desc='Storing UMAP data'):
+        label_value = label.item()
+        if label_value not in otherpoints:
+            otherpoints[label_value] = []
+        otherpoints[label_value].append(reduced[idx])
 
-    #  get the device
+    # Device
     device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
     model = model.to(device)
-    # parameters 
-    points = 300
-    shape = 28*28
+
+    # Swarm parameters
+    num_particles = 300
+    shape = 28 * 28
     epochs = 20
 
-    #  get initial 
-    initial_points = np.random.rand(points, shape)
+    initial_points = np.random.rand(num_particles, shape)
     mask = np.random.choice([0, 1], size=initial_points.shape, p=[0.99, 0.01])
-    initial_points = initial_points * mask
-    initial_points = torch.tensor(initial_points).to(torch.float32)
-    #  add anchor to initial points
-    initial_points = torch.cat((initial_points, anchor.reshape(1, 28*28)), dim=0)
+    initial_points *= mask
+    initial_points = torch.tensor(initial_points, dtype=torch.float32)
+
+    # Add anchor to the swarm
+    if anchor is None:
+        raise ValueError(f"No anchor found for class {optimize}.")
+    initial_points = torch.cat((initial_points, anchor.reshape(1, shape)), dim=0)
 
-    #  create the swarm
+    # Run the particle swarm
     runSwarm(initial_points, model, device, umap_model, epochs, otherpoints)
 
-optimize = 3
-def cost_func(model: torch.nn.Sequential, point: torch.tensor):
-    global optimize
-    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
-    model = model.to(device)
-    point = point.to(device)
-    point = point.reshape(1, 1, 28, 28)
-    output = model(point)
 
-    grad = np.abs(AO.Attacks.gradient_map(point, model, (1,1,28,28))[0].reshape(-1))
-    maxgrad = np.max(grad)
-    return .7*(output[0][optimize].item()) + .3*(np.sum(grad)/(maxgrad*len(grad)))
+# --- Run Swarm with Custom ParticleSwarm ---
+def runSwarm(initial_points, model, device, umap_model, epochs, otherpoints):
+    swarm = ParticleSwarm(
+        model=model,
+        input_set=initial_points,
+        starting_class=optimize,  # assuming starting and target are same
+        target_class=optimize,
+        num_iterations=epochs,
+        save_dir='./APSO_A',
+        enable_logging=False,
+        device=str(device)
+    )
+
+    plotSwarm(swarm, umap_model, 0, otherpoints)
+    swarm.optimize()
+
+    # After optimization, save final results
+    plotSwarm(swarm, umap_model, epochs, otherpoints)
+
 
+# --- Plotting Functions ---
 def plotSwarm(swarm, umap_model, epoch, otherpoints):
-    # Plot the points
     fig, ax = plt.subplots(figsize=(12, 12))
 
-    # Plot otherpoints 
-    for key in otherpoints.keys():
+    # Plot static test points
+    for key in otherpoints:
         x = [i[0] for i in otherpoints[key]]
         y = [i[1] for i in otherpoints[key]]
         ax.scatter(x, y, label=key)
-    # Get all the points
-    points = swarm.getPoints()
 
-    # Transform the points
-    points = umap_model.transform(points)
-    ax.scatter(points[:, 0], points[:, 1], c='black', label='Swarm')                
+    # Swarm points
+    points = np.array(swarm.getPoints())
+    if points.ndim > 2:
+        points = points.reshape(points.shape[0], -1)
+
+    transformed = umap_model.transform(points)
+    ax.scatter(transformed[:, 0], transformed[:, 1], c='black', label='Swarm')
+
     ax.legend()
     ax.set_title(f'Epoch: {epoch}')
     os.makedirs('./APSO_A/points', exist_ok=True)
     plt.savefig(f'./APSO_A/points/epoch_{epoch}.png')
     plt.close()
+
     plotImages(swarm, epoch)
 
+
 def plotImages(swarm, epoch):
-    points = swarm.getPoints()
+    points = np.array(swarm.getPoints())
     points = points.reshape(-1, 1, 28, 28)
-    #  plot the best
-    best = swarm.getBest()
-    best = best.reshape(28, 28)
+
+    best = swarm.getBest().reshape(28, 28)
     device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
 
+    # Plot best particle
     fig, ax = plt.subplots(figsize=(6, 6))
     ax.imshow(best, cmap='gray')
     ax.axis('off')
-    global optimize
-    ax.set_title(f'Confidence of {optimize}: {swarm.model(best.reshape(1,1,28,28).to(torch.float32).to(device))[0][optimize].item()}')
+    confidence = swarm.model(torch.tensor(best).reshape(1, 1, 28, 28).float().to(device))[0][optimize].item()
+    ax.set_title(f'Confidence of {optimize}: {confidence:.4f}')
 
-    grad = AO.Attacks.gradient_map(best.reshape(1,1,28,28), swarm.model, (1,1, 28, 28))[0].reshape(28,28)
-    grad = np.abs(grad)  # Make the gradients absolute for better visualization
-    grad_normalized = (grad - np.min(grad)) / (np.max(grad) - np.min(grad))  # Normalize gradients to 0-1
+    grad = AO.Attacks.gradient_map(
+        torch.tensor(best).reshape(1, 1, 28, 28).float().to(device), swarm.model,
+        (1, 1, 28, 28)
+    )[0].reshape(28, 28)
 
-    grad_cmap = plt.get_cmap('jet')
-    grad_rgba = grad_cmap(grad_normalized)
-    grad_rgba[..., 3] = 0.7  # Set alpha channel to control transparency
+    grad = np.abs(grad)
+    grad_norm = (grad - np.min(grad)) / (np.max(grad) - np.min(grad) + 1e-8)
+    ax.imshow(plt.get_cmap('jet')(grad_norm), alpha=0.7)
 
-    ax.imshow(grad_rgba, cmap='jet', interpolation='bilinear')
     os.makedirs(f'./APSO_A/images/epoch_{epoch}', exist_ok=True)
     plt.savefig(f'./APSO_A/images/epoch_{epoch}/best.png')
     plt.close()
-    # plot the rest
+
+    # Plot remaining particles
     for idx, point in enumerate(points):
-        point = point.reshape(28, 28)
+        point_img = point.reshape(28, 28)
         fig, ax = plt.subplots(figsize=(6, 6))
-        ax.imshow(point, cmap='gray')
+        ax.imshow(point_img, cmap='gray')
         ax.axis('off')
-        ax.set_title(f'Confidence of {optimize}: {swarm.model(point.reshape(1,1,28,28).to(torch.float32).to(device))[0][optimize].item()}')
-
-        grad = AO.Attacks.gradient_map(point.reshape(1,1,28,28), swarm.model, (1,1, 28, 28))[0].reshape(28,28)
-        grad = np.abs(grad)  # Make the gradients absolute for better visualization
-        grad_normalized = (grad - np.min(grad)) / (np.max(grad) - np.min(grad))  # Normalize gradients to 0-1
-
-        grad_cmap = plt.get_cmap('jet')
-        grad_rgba = grad_cmap(grad_normalized)
-        grad_rgba[..., 3] = 0.7  # Set alpha channel to control transparency
-
-        ax.imshow(grad_rgba, cmap='jet', interpolation='bilinear')
+        confidence = swarm.model(torch.tensor(point_img).reshape(1, 1, 28, 28).float().to(device))[0][optimize].item()
+        ax.set_title(f'Confidence of {optimize}: {confidence:.4f}')
+
+        grad = AO.Attacks.gradient_map(
+            torch.tensor(point_img).reshape(1, 1, 28, 28).float().to(device),
+            swarm.model,
+            (1, 1, 28, 28)
+        )[0].reshape(28, 28)
+        grad = np.abs(grad)
+        grad_norm = (grad - np.min(grad)) / (np.max(grad) - np.min(grad) + 1e-8)
+        ax.imshow(plt.get_cmap('jet')(grad_norm), alpha=0.7)
 
         plt.tight_layout()
         plt.savefig(f'./APSO_A/images/epoch_{epoch}/point_{idx}.png')
         plt.close()
 
-def runSwarm(inital_points, model, device, umap_model, epochs, otherpoints):
-    APSO = Swarm.PSO(inital_points, cost_func, model)
-    plotSwarm(APSO, umap_model, 0, otherpoints)
-    for i in tqdm.tqdm(range(1, 1+epochs), desc='Running Swarm', total=epochs):
-        APSO.step()
-        plotSwarm(APSO, umap_model, i, otherpoints)
-    plotSwarm(APSO, umap_model, epochs+1, otherpoints)
 
+# --- Entry ---
 if __name__ == '__main__':
-    main()
\ No newline at end of file
+    main()
diff --git a/manuscripts/PEARC24/MNIST/5_create_SHAP.py b/manuscripts/PEARC24/MNIST/5_create_SHAP.py
index b83d179..60cc43a 100644
--- a/manuscripts/PEARC24/MNIST/5_create_SHAP.py
+++ b/manuscripts/PEARC24/MNIST/5_create_SHAP.py
@@ -35,67 +35,70 @@ def save_and_plot_shap_values(dataloader, model):
     data = data.to(device)
     target = target.to(device)
 
-    # Move model to device
     model = model.to(device)
 
-    # Generate SHAP values
     explainer = shap.DeepExplainer(model, data)
-    shap_values = explainer.shap_values(data)
+    shap_values = explainer.shap_values(data)  # List of [class][samples, features]
 
     save_dir = 'SHAP'
     os.makedirs(save_dir, exist_ok=True)
 
-    # Create a 10x10 grid of subplots
+    # Create a 10x11 grid: 1 original + 10 SHAP values
     fig, axes = plt.subplots(10, 11, figsize=(20, 22))
+    last_img = None  # For colorbar
 
-    # Iterate over the SHAP values and plot on the subplots
     for i in range(len(data)):
-        shap_i = shap_values[i]
         label = target[i].item()
+        shap_i = [class_shap[i] for class_shap in shap_values]  # SHAP per class, for this image
 
-        # Save the original image as a numpy array
+        # Save original image
         np.save(f'{save_dir}/{i}_original.npy', data[i].cpu().numpy())
-
-        # Plot the original image
         axes[i, 0].imshow(data[i].cpu().reshape(28, 28), cmap='gray')
         axes[i, 0].set_title(f'Label: {label}')
-
-        # Plot the SHAP values
-        num_shap_values = min(10, len(shap_i))  # Adjust the number of SHAP values to fit within the grid
-        for j in range(num_shap_values):
-            # Save the SHAP value as a numpy array
-            np.save(f'{save_dir}/{i}_shap_{j}.npy', shap_i[j])
-            img = axes[i, j+1].imshow(shap_i[j].reshape(28, 28), cmap='jet')
+        axes[i, 0].axis('off')
+
+        for j in range(min(10, len(shap_i))):
+            shap_array = shap_i[j]
+            try:
+                reshaped = shap_array.reshape(10, 28, 28)[j]  # extract correct class
+            except Exception as e:
+                print(f"[ERROR] SHAP reshape failed for sample {i}, class {j}: {e}")
+                continue
+
+            np.save(f'{save_dir}/{i}_shap_{j}.npy', shap_array)
+            last_img = axes[i, j+1].imshow(reshaped, cmap='jet')
             axes[i, j+1].axis('off')
-            # axes[i, j+1].set_title(f'SHAP value {j+1}')
 
-        # Remove empty cells in the row
-        for j in range(num_shap_values + 1, 11):
+
+        # Fill remaining columns
+        for j in range(len(shap_i) + 1, 11):
             axes[i, j].axis('off')
 
-        # Save the row individually and remove the white space
-        row_fig = plt.figure(figsize=(10, 1))
-        row_axes = row_fig.subplots(1, num_shap_values + 1)
+        # Save row as standalone image
+        row_fig, row_axes = plt.subplots(1, 11, figsize=(20, 2))
         row_axes[0].imshow(data[i].cpu().reshape(28, 28), cmap='gray')
         row_axes[0].set_title(f'Label: {label}')
-        for j in range(num_shap_values):
-            row_axes[j+1].imshow(shap_i[j].reshape(28, 28), cmap='jet')
+        row_axes[0].axis('off')
+        for j in range(min(10, len(shap_i))):
+            row_axes[j+1].imshow(shap_i[j][:784].reshape(28, 28), cmap='jet')
             row_axes[j+1].axis('off')
+        for j in range(len(shap_i) + 1, 11):
+            row_axes[j].axis('off')
         plt.tight_layout()
         row_fig.savefig(f'{save_dir}/row_{i}.png')
         plt.close(row_fig)
 
-    # Remove empty rows
+    # Fill empty rows if less than 10 samples
     for i in range(len(data), 10):
-        for j in range(10):
+        for j in range(11):
             axes[i, j].axis('off')
 
-    # Add colorbar
-    cbar_ax = fig.add_axes([.93, 0.15, 0.02, 0.7])  # Adjust the position of the colorbar
-    fig.colorbar(img, cax=cbar_ax)
+    # Add colorbar only if a SHAP plot was rendered
+    if last_img is not None:
+        cbar_ax = fig.add_axes([.93, 0.15, 0.02, 0.7])
+        fig.colorbar(last_img, cax=cbar_ax)
 
-    # Save the figure
-    # plt.tight_layout()
+    plt.tight_layout()
     plt.savefig(f'{save_dir}/shap_values.png')
     plt.close()
 
diff --git a/manuscripts/PEARC24/MNIST/docs/1_build_train.md b/manuscripts/PEARC24/MNIST/docs/1_build_train.md
index 004a008..6df8c48 100644
--- a/manuscripts/PEARC24/MNIST/docs/1_build_train.md
+++ b/manuscripts/PEARC24/MNIST/docs/1_build_train.md
@@ -1,6 +1,8 @@
 ## MNIST Training and Evaluation
 
-This code, implemented in `1_build_train.py`, demonstrates the training and evaluation process for a Convolutional Neural Network (CNN) model on the MNIST dataset. The MNIST dataset consists of grayscale images of handwritten digits from 0 to 9.
+This code demonstrates the training and evaluation process for a Convolutional Neural Network (CNN) model on the MNIST dataset, which contains grayscale images of handwritten digits (0-9).
+
+The script, `1_build_train.py`, trains a CNN using the MNIST dataset, evaluates its performance, and saves the results and the trained model.
 
 ### Functions
 
@@ -10,46 +12,91 @@ The code includes the following functions:
 
 This function trains the model for one epoch using the provided data loader, optimizer, and loss function. It saves the training loss to the specified file.
 
-- `model`: The model to be trained.
-- `train_loader`: The data loader for the training data.
-- `optimizer`: The optimizer used for training.
-- `loss`: The loss function.
-- `epoch`: The current epoch number.
-- `filename`: The name of the file to save the training loss.
+* **`model`**: The CNN model to be trained.
+* **`train_loader`**: The data loader for the training data.
+* **`optimizer`**: The optimizer used for training (e.g., Adam).
+* **`loss`**: The loss function (e.g., CrossEntropyLoss).
+* **`epoch`**: The current epoch number.
+* **`filename`**: The file path to save the training loss.
 
 #### 2. `testModel(model, test_loader, filename)`
 
-This function evaluates the model using the provided data loader and calculates the test loss and accuracy. It saves the test loss to the specified file.
+This function evaluates the model on the test dataset, calculating the test loss and accuracy. It saves the test loss to the specified file.
 
-- `model`: The model to be evaluated.
-- `test_loader`: The data loader for the test data.
-- `filename`: The name of the file to save the test loss.
+* **`model`**: The CNN model to be evaluated.
+* **`test_loader`**: The data loader for the test data.
+* **`filename`**: The file path to save the test loss.
 
 #### 3. `seedEverything(seed)`
 
-This function seeds all the random number generators to ensure reproducibility.
+This function ensures reproducibility by seeding all random number generators across the framework.
 
-- `seed`: The seed value for random number generators.
+* **`seed`**: The seed value for random number generators.
 
 #### 4. `main()`
 
 This is the main function that orchestrates the training and evaluation process. It performs the following steps:
 
-- Seeds the random number generators for reproducibility.
-- Loads the MNIST dataset using the `load_MNIST_data()` function from the `util` module.
-- Builds the CNN model using the `build_MNIST_Model()` function from the `util` module.
-- Sets up the optimizer and loss function.
-- Trains the model for a specified number of epochs, calling the `trainModel()` and `testModel()` functions.
-- Saves the trained model to a file.
+* Seeds the random number generators for reproducibility using `seedEverything()`.
+* Loads the MNIST dataset with `load_MNIST_data()` from the `AO.utils` module.
+* Builds the CNN model using `load_MNIST_model()` from the `AO.utils` module.
+* Sets up the Adam optimizer and CrossEntropy loss function.
+* Trains the model for a specified number of epochs, calling the `trainModel()` and `testModel()` functions for each epoch.
+* Saves the trained model’s parameters to the file `MNIST_cnn.pt`.
 
 ### Usage
 
 To use this code:
 
-1. Make sure you have the necessary dependencies installed.
-2. Run the Python script named `1_build_train.py`.
-3. Execute the script to train the model and evaluate its performance on the MNIST dataset.
-4. The training loss will be saved in the `log.csv` file.
-5. The trained model will be saved in the `MNIST_cnn.pt` file.
+1. **Install Dependencies**
+   Ensure that you have the following Python libraries installed:
+
+   * `torch`
+   * `torchvision`
+   * `numpy`
+   * `tqdm`
+
+   You can install them using `pip`:
+
+   ```bash
+   pip install torch torchvision numpy tqdm
+   ```
+
+2. **Run the Script**
+   Execute the Python script `1_build_train.py`:
+
+   ```bash
+   python 1_build_train.py
+   ```
+
+3. **Monitor Training and Testing**
+   The script will output the training loss and test loss for each epoch in the terminal. The progress bar (from `tqdm`) will also show the current status of training.
+
+4. **Review the Log File**
+   The training loss and test loss will be appended to a file called `log.csv`. Each line in the file will record the training and test loss for each epoch.
+
+5. **Saved Model**
+   After training completes, the final model’s parameters (weights) will be saved to the file `MNIST_cnn.pt`. You can load this file later for inference or further training.
+
+6. **Modify Configuration**
+   If necessary, modify the paths, filenames, or number of epochs in the script to suit your needs.
+
+### Example Output
+
+In the `log.csv` file, you will see entries like:
+
+```
+Train Epoch: 1 Loss: 0.350100
+Test set: Average loss: 0.0512, Accuracy: 97%
+Train Epoch: 2 Loss: 0.090200
+Test set: Average loss: 0.0451, Accuracy: 98%
+```
+
+### Notes
+
+* **Device**: The model will be trained on a GPU if available, otherwise it will fall back to CPU. This is automatically detected in the code.
+
+* **Reproducibility**: The random seed is fixed (`42`), so the results should be consistent across multiple runs.
+
+* **Custom `AO` Module**: This script depends on a custom module `Adversarial_Observation` (imported as `AO`) for utilities like `load_MNIST_data()` and `load_MNIST_model()`. Make sure this module is available and correctly implemented.
 
-Note: Before running the code, you may need to modify the paths and filenames to match your desired configuration.
diff --git a/manuscripts/PEARC24/MNIST/docs/2_adversarial_attack.md b/manuscripts/PEARC24/MNIST/docs/2_adversarial_attack.md
index b099535..9382fb4 100644
--- a/manuscripts/PEARC24/MNIST/docs/2_adversarial_attack.md
+++ b/manuscripts/PEARC24/MNIST/docs/2_adversarial_attack.md
@@ -1,39 +1,103 @@
 ## Adversarial Image Generation and Evaluation
 
-This code, implemented in `2_adversarial_attack.py`, demonstrates the generation of adversarial images using the Fast Gradient Sign Method (FGSM) and evaluates their impact on the model's predictions. The code is designed to work with a pre-trained CNN model on the MNIST dataset.
-
+This project demonstrates the generation of adversarial images using the **Fast Gradient Sign Method (FGSM)** and evaluates their impact on a pre-trained CNN model trained on the **MNIST dataset**. The script, `2_adversarial_attack.py`, generates perturbed images and visualizes the differences between the original and adversarial images.
 
 ### Functions
 
-The code includes the following functions:
+The code includes several functions that facilitate adversarial image generation, evaluation, and visualization.
 
 #### 1. `seedEverything(seed)`
 
-This function seeds all the random number generators to ensure reproducibility.
+Seeds all random number generators to ensure reproducibility of results.
+
+* **Parameters**:
 
-- `seed`: The seed value for random number generators.
+  * `seed`: An integer seed value for the random number generators.
 
 #### 2. `main()`
 
-This is the main function that orchestrates the adversarial image generation and evaluation process. It performs the following steps:
+The main function that orchestrates the adversarial image generation process. It performs the following tasks:
+
+* Seeds the random number generators for reproducibility.
+* Loads the MNIST dataset and pre-trained model.
+* Sets the model to evaluation mode.
+* Generates adversarial images using the FGSM attack with varying epsilon values.
+* Saves the adversarial images, the perturbations, activation maps, and original images for analysis.
+
+#### 3. `grad_ascent(model, device)`
+
+This function performs **gradient ascent** to create adversarial examples for each class by perturbing an image to maximize the model’s confidence for the target class. It saves generated images as PNG files in the `./gradient_Ascent` directory.
+
+* **Parameters**:
+
+  * `model`: The pre-trained CNN model.
+  * `device`: The computing device (CPU or GPU).
+
+#### 4. `fgsm(imgs, labels, model, device)`
+
+Generates adversarial images using the **Fast Gradient Sign Method (FGSM)** with a list of epsilon values. The adversarial images, their gradients, and the activation maps are saved to the `./attack_results` directory.
+
+* **Parameters**:
+
+  * `imgs`: The input images for generating adversarial examples.
+  * `labels`: The true labels of the input images.
+  * `model`: The pre-trained CNN model.
+  * `device`: The computing device (CPU or GPU).
+
+#### 5. `plot_perterbed(perterbed, eps, imgs, labels, model)`
+
+Visualizes the adversarial images, their corresponding perturbations, and activation maps. The original, perturbed, and noise images are saved as PNG files, and their numpy arrays are stored in the `./attack_results` directory.
+
+* **Parameters**:
 
-- Seeds the random number generators for reproducibility.
-- Loads the MNIST dataset using the `load_MNIST_data()` function from the `util` module.
-- Builds the MNIST CNN model using the `build_MNIST_Model()` function from the `util` module.
-- Loads the pre-trained model weights from the `MNIST_cnn.pt` file.
-- Sets the model to evaluation mode.
-- Generates adversarial images using the FGSM attack on misclassified images.
-- Saves the generated adversarial images, their corresponding perturbations, activation maps, and original images for reference.
+  * `perterbed`: The perturbed images.
+  * `eps`: The epsilon value used in the FGSM attack.
+  * `imgs`: The original input images.
+  * `labels`: The true labels of the input images.
+  * `model`: The pre-trained CNN model.
 
 ### Usage
 
-To use this code:
+To use this code, follow these steps:
+
+1. **Install Dependencies**: Ensure that all the necessary dependencies (PyTorch, torchvision, numpy, matplotlib, etc.) are installed.
+
+2. **Run the Script**: Execute the Python script `2_adversarial_attack.py` to generate adversarial images and evaluate their impact.
+
+   ```bash
+   python 2_adversarial_attack.py
+   ```
+
+3. **Generated Results**:
+
+   * The generated adversarial images and the corresponding perturbations will be saved in the `./attack_results` directory.
+   * The visualizations of the original, perturbed, and noise images, along with the activation maps, will be saved as PNG files.
+   * Numpy arrays of the perturbed images are also saved in the same directory for further analysis.
+
+4. **Adjusting Attack Parameters**:
+
+   * The `eps` values (perturbation strength) used in the FGSM attack can be adjusted in the code.
+   * You can modify the saving directories and file names as needed.
+
+### Prerequisites
+
+* **Pre-trained Model**: Ensure that the pre-trained model weights `MNIST_cnn.pt` are available in the working directory. If not, you need to run the model training script (e.g., `1_build_train.py`) to obtain the `MNIST_cnn.pt` file.
+
+* **Paths**: Make sure the file paths for saving images and results match your desired directory structure.
+
+### Notes
+
+* This code supports both **CPU** and **GPU** execution. The device is automatically detected.
+* **Gradient ascent** and **FGSM** can be used to create adversarial examples, with results saved for each class and epsilon value.
+* The perturbation strength (`eps`) can be customized to explore the effect of different attack intensities.
+
+### Output Structure
+
+* `./attack_results/`: Directory containing:
 
-1. Make sure you have the necessary dependencies installed.
-2. Run the Python script named `2_adversarial_attack.py`.
-3. Execute the script to generate adversarial images using the FGSM attack and evaluate their impact.
-4. The generated adversarial images will be saved in the `Noise/fgsm` directory.
-5. The corresponding perturbations, activation maps, and original images will also be saved for each adversarial image.
-6. Adjust the attack parameters (`eps` values) and saving directories as needed.
+  * `fgsm_{eps}_{label}.png`: Original, perturbed, and noise images for each attack.
+  * `fgsm_{eps}_{label}.npy`: Numpy array of the perturbed image.
+  * `fgsm_{eps}_{label}_grad.png`: Activation maps of original and perturbed images.
+* `./gradient_Ascent/`: Directory containing:
 
-Note: Before running the code, make sure to have the pre-trained model weights in the `MNIST_cnn.pt` file and modify the paths and filenames to match your desired configuration.
+  * `ga_{i}.png`: Gradient ascent images for each class.
diff --git a/manuscripts/PEARC24/MNIST/docs/3_create_APSO.md b/manuscripts/PEARC24/MNIST/docs/3_create_APSO.md
index e2d7bd3..aaa0efb 100644
--- a/manuscripts/PEARC24/MNIST/docs/3_create_APSO.md
+++ b/manuscripts/PEARC24/MNIST/docs/3_create_APSO.md
@@ -1,40 +1,123 @@
 # Adversarial Particle Swarm Optimization (APSO) Description
 
-The given code is a Python script that implements the Adversarial Particle Swarm Optimization (APSO) algorithm. APSO aims to generate adversarial examples by optimizing a cost function using a particle swarm optimization algorithm.
-
+The provided Python script implements the **Adversarial Particle Swarm Optimization (APSO)** algorithm to generate adversarial examples for a machine learning model. APSO optimizes a cost function using Particle Swarm Optimization (PSO) to perturb input data (MNIST images) in a way that maximizes the model's misclassification probability. The generated adversarial examples are visualized using UMAP for dimensionality reduction and gradient-based techniques to highlight regions of vulnerability.
 
 ## Global Variables
-The code sets the following global variables:
-- `label` to determine the target label for adversarial attacks.
-- `initial` as the initial label of the dataset.
-- `epochs` to specify the number of optimization epochs.
-- `points` to define the number of attack points.
+
+* **`label`**: The target label for the adversarial attack (set to `3` by default). This is the class the algorithm tries to generate adversarial examples for.
+* **`initial`**: The initial label of the dataset (not used directly in the code but could be a placeholder for tracking the original label).
+* **`epochs`**: The number of iterations the PSO algorithm will run to optimize the adversarial examples (set to `20`).
+* **`points`**: The number of particles in the swarm (set to `50`).
 
 ## Functions
 
-### `cost_func(model, x)`
-This function calculates the cost (prediction score) for a given input data point using a pre-trained PyTorch model. It takes the model and input data as arguments and returns the cost as a float.
+### `cost_func(model, point)`
+
+This function calculates the cost (prediction score) for a given input data point using the pre-trained PyTorch model. It computes the product of the model’s confidence in predicting the target `label` and the sum of the gradients over the image's pixels. The function returns the cost, which represents how well the current point (image) serves as an adversarial example.
+
+**Parameters:**
+
+* `model`: The pre-trained PyTorch model used to classify the data.
+* `point`: The input data point (image) to evaluate.
+
+**Returns:**
+
+* A float value representing the adversarial cost.
+
+### `plotSwarm(swarm, umap_model, epoch, otherpoints)`
+
+This function visualizes the state of the particle swarm during optimization. It uses UMAP for dimensionality reduction and generates a scatter plot of the swarm's position in 2D space, along with the positions of the original class data points (MNIST images).
 
-### `umap_data(dataloader)`
-This function generates UMAP embeddings for the data in a given dataloader. It takes a dataloader as input and returns the UMAP object fitted on the flattened data and a dictionary mapping target labels to their corresponding data points.
+**Parameters:**
 
-### `plotData(data, umap, title, saveName)`
-This function plots the data points using UMAP embeddings. It takes the data, UMAP object, plot title, and save name as arguments.
+* `swarm`: The current state of the swarm (ParticleSwarm object).
+* `umap_model`: The trained UMAP model to reduce the dimensionality of data for plotting.
+* `epoch`: The current epoch number to label the plot.
+* `otherpoints`: A dictionary mapping the target labels to their corresponding UMAP-reduced data points.
 
-### `plotPSO(points, step, model, runName)`
-This function plots the Particle Swarm Optimization (PSO) results. It takes the list of points generated by the PSO algorithm, the current step or iteration number, the pre-trained PyTorch model, and the name of the PSO run as arguments.
+### `plotImages(swarm, epoch)`
 
-### `getPositions(APSO)`
-This function returns the positions of particles from an APSO object.
+This function visualizes the generated adversarial images and their corresponding gradient maps (showing the influence of each pixel on the model’s classification). It saves the visualizations of the best particles and others in the swarm at each epoch.
 
-### `plotBest(APSO, model, runName, epoch)`
-This function plots the best point and its activation map from APSO results. It takes the APSO object, pre-trained PyTorch model, name of the APSO run, and current epoch or iteration number as arguments.
+**Parameters:**
 
-### `runAPSO(points, epochs, model, cost_func, dataDic, umap, runName)`
-This function runs the Adversarial Particle Swarm Optimization (APSO) algorithm. It takes the initial set of points, number of optimization epochs, pre-trained PyTorch model, cost function, dictionary of data points, UMAP object, and name of the APSO run as arguments. It returns the final positions of the particles after optimization.
+* `swarm`: The current state of the swarm (ParticleSwarm object).
+* `epoch`: The current epoch number.
+
+### `runSwarm(initial_points, model, device, umap_model, epochs, otherpoints)`
+
+This function initializes and runs the particle swarm optimization process. It uses the `Swarm.ParticleSwarm` class to optimize a set of particles (initial adversarial candidates) for a specified number of epochs. After optimization, the best adversarial examples are returned.
+
+**Parameters:**
+
+* `initial_points`: The initial swarm points (random adversarial candidates).
+* `model`: The pre-trained PyTorch model used to classify the data.
+* `device`: The device (`'cuda'` or `'cpu'`) on which to run the model.
+* `umap_model`: The UMAP model to reduce the dimensionality for visualization.
+* `epochs`: The number of iterations for optimization.
+* `otherpoints`: A dictionary mapping target labels to their corresponding data points.
 
 ### `main()`
-This function is the main entry point of the script. It initializes global variables, loads the dataset, builds the model, and performs APSO for different attack scenarios. It calls the necessary functions to run APSO, plot clusters, and save results.
+
+This function is the main entry point of the script. It initializes global variables, loads the dataset (MNIST), builds the model, and runs APSO for generating adversarial examples. It also handles saving and visualizing the results during the optimization process.
+
+**Execution flow:**
+
+1. Load the MNIST dataset and pre-trained model.
+2. Initialize UMAP and other required variables.
+3. Run the `runSwarm` function to perform optimization.
+4. Plot and save the results (including UMAP visualizations and adversarial images).
 
 ## Execution
-The script starts by defining global variables and importing required libraries and modules. Then, it defines various functions for cost calculation, UMAP generation, plotting, and APSO optimization. Finally, the `main()` function is called to execute the script.
+
+### Prerequisites
+
+To run this script, you need:
+
+* **PyTorch**: For building and running the neural network model.
+* **UMAP**: For dimensionality reduction and visualization.
+* **tqdm**: For progress bars.
+* **matplotlib**: For plotting the swarm and images.
+* **SciPy**: For sparse matrix operations (if needed).
+
+### Running the Script
+
+To execute the script, follow these steps:
+
+1. **Install required dependencies**:
+
+   ```bash
+   pip install torch umap-learn matplotlib scipy tqdm
+   ```
+
+2. **Download or train the MNIST model**:
+
+   * If you have the pre-trained `MNIST_cnn.pt` model, place it in the working directory.
+   * If not, you need to train a model first using the appropriate training script (e.g., `1_build_train.py`).
+
+3. **Run the script**:
+
+   ```bash
+   python adversarial_pso.py
+   ```
+
+4. **Output**:
+
+   * The script will generate plots of the adversarial examples at each epoch and save them to the `./APSO/images/` directory.
+   * The swarm optimization results (including UMAP visualizations and gradient maps) will be saved under `./APSO/points/` and `./APSO/images/`.
+
+### Visualization
+
+The generated images and plots will allow you to observe how the particles (adversarial examples) evolve over time to deceive the model into misclassifying them as the target class.
+
+## Notes
+
+* **Swarm Configuration**: The PSO parameters like inertia weight, cognitive weight, and social weight are set to standard values but can be tuned for different optimization behavior.
+* **Visualization**: The code generates and saves multiple visualizations, including UMAP scatter plots and adversarial images with their corresponding gradient maps. These help to understand how the particles are evolving and which parts of the image influence the model's decision.
+* **File Storage**: The script saves the results in `./APSO/images/` and `./APSO/points/`. Ensure enough disk space for storing images.
+
+## References
+
+* [PyTorch](https://pytorch.org/): Deep learning framework used for the MNIST model.
+* [UMAP](https://umap-learn.readthedocs.io/): Used for dimensionality reduction and visualizing high-dimensional data.
+* [Particle Swarm Optimization (PSO)](https://en.wikipedia.org/wiki/Particle_swarm_optimization): Optimization algorithm inspired by the social behavior of birds flocking or fish schooling.
diff --git a/manuscripts/PEARC24/MNIST/docs/3p1_APSO_W_Anchor.md b/manuscripts/PEARC24/MNIST/docs/3p1_APSO_W_Anchor.md
new file mode 100644
index 0000000..b70620d
--- /dev/null
+++ b/manuscripts/PEARC24/MNIST/docs/3p1_APSO_W_Anchor.md
@@ -0,0 +1,106 @@
+
+# Adversarial Particle Swarm Optimization (APSO) with Custom Swarm Class
+
+This Python script implements an **Adversarial Particle Swarm Optimization (APSO)** algorithm to generate adversarial examples for a **MNIST** model. The algorithm uses **Particle Swarm Optimization (PSO)** to perturb input images in a way that maximizes the misclassification probability. The generated adversarial examples are visualized using **UMAP** for dimensionality reduction and gradient-based techniques to highlight areas of high vulnerability in the model.
+
+The script uses a custom `ParticleSwarm` class defined in the `Adversarial_Observation` package to perform the optimization.
+
+## Global Variables
+
+* **`optimize`**: The target class label for the adversarial attack (default is `3`).
+* **`epochs`**: The number of optimization iterations for the swarm (set to `20`).
+* **`num_particles`**: The number of particles in the swarm (set to `300`).
+* **`shape`**: The dimensionality of each particle (flattened MNIST images, i.e., `28 * 28 = 784`).
+
+## Functions
+
+### `main()`
+
+The main entry point of the script. This function:
+
+1. Loads the pre-trained MNIST model.
+2. Loads the MNIST data using PyTorch's DataLoader.
+3. Initializes **UMAP** for dimensionality reduction and prepares test points for visualization.
+4. Sets up the particle swarm by initializing random particles and adding an anchor point for the target class.
+5. Runs the **particle swarm optimization** by calling the `runSwarm()` function.
+
+### `runSwarm(initial_points, model, device, umap_model, epochs, otherpoints)`
+
+This function runs the **Particle Swarm Optimization (PSO)** algorithm. It initializes a **ParticleSwarm** object and optimizes the particles to generate adversarial examples that target the specified class (`optimize`). The function:
+
+1. Initializes the swarm with random starting points (particles).
+2. Calls the `optimize()` method of the **ParticleSwarm** class to perform the optimization.
+3. Saves and plots the swarm’s progress during each epoch.
+
+### `plotSwarm(swarm, umap_model, epoch, otherpoints)`
+
+This function visualizes the positions of the particles in the swarm at each optimization epoch. It:
+
+1. Visualizes the particles using **UMAP** to reduce the dimensionality of the swarm's position in 2D.
+2. Saves the scatter plot as an image to the `./APSO_A/points/` directory.
+3. Calls `plotImages()` to visualize the adversarial images.
+
+### `plotImages(swarm, epoch)`
+
+This function visualizes the adversarial images and their gradient maps (indicating which parts of the image most influence the model’s decision). It:
+
+1. Plots the **best particle** (adversarial example) and its confidence score for the target class.
+2. Visualizes the **gradient map** for the best particle, showing which pixels influence the model's decision the most.
+3. Saves the images of the best particle and all particles in the swarm for each epoch to the `./APSO_A/images/` directory.
+
+## Execution
+
+### Prerequisites
+
+To run this script, you will need the following dependencies:
+
+* **PyTorch**: For model training and inference.
+* **UMAP**: For dimensionality reduction and visualization.
+* **Matplotlib**: For plotting images and swarm positions.
+* **tqdm**: For progress bars during the execution.
+* **Adversarial\_Observation**: This is a custom package that contains the `ParticleSwarm` class and related functionality for adversarial attacks.
+
+Install the necessary dependencies:
+
+```bash
+pip install torch umap-learn matplotlib tqdm
+```
+
+Additionally, the **Adversarial\_Observation** package must be available. If you don't have this package, ensure that it's installed or accessible in your project directory.
+
+### Model Preparation
+
+The script assumes that a pre-trained MNIST model (`MNIST_cnn.pt`) is available. If you don’t have this model file, you need to train it using the provided training script (`1_build_train.py`) before running the adversarial attack script.
+
+### Running the Script
+
+To execute the script:
+
+1. Ensure you have the pre-trained model `MNIST_cnn.pt` in the current working directory.
+2. Run the script as follows:
+
+   ```bash
+   python adversarial_pso.py
+   ```
+
+### Output
+
+* The generated adversarial examples and swarm positions will be saved to the `./APSO_A/points/` and `./APSO_A/images/` directories.
+* The script saves **UMAP visualizations** and **gradient maps** of the best particle at each optimization epoch.
+
+### Visualization
+
+The **swarm visualizations** show how the adversarial particles evolve during the optimization process. The **image visualizations** show the generated adversarial examples and highlight the pixels that most influence the model's classification using a gradient map.
+
+## Notes
+
+* **Swarm Configuration**: The swarm's configuration, including the number of particles, epochs, and particle initialization, can be adjusted by modifying the relevant parameters in the script.
+* **Anchor**: The algorithm selects one test image with the target label (`optimize`) as the anchor point for the attack. If no anchor image is found in the dataset, a `ValueError` will be raised.
+* **Visualization**: The code generates and saves multiple visualizations, including UMAP scatter plots, adversarial images, and their gradient maps. These help to understand the impact of the optimization process on the model.
+
+## References
+
+* **PyTorch**: [https://pytorch.org/](https://pytorch.org/) – Deep learning framework used to build and run the MNIST model.
+* **UMAP**: [https://umap-learn.readthedocs.io/](https://umap-learn.readthedocs.io/) – Dimensionality reduction technique used for visualizing high-dimensional data.
+* **Particle Swarm Optimization (PSO)**: [https://en.wikipedia.org/wiki/Particle\_swarm\_optimization](https://en.wikipedia.org/wiki/Particle_swarm_optimization) – Optimization algorithm used for generating adversarial examples.
+
diff --git a/manuscripts/PEARC24/MNIST/docs/4_analyze_APSO.md b/manuscripts/PEARC24/MNIST/docs/4_analyze_APSO.md
index 8b35c45..dd1ac27 100644
--- a/manuscripts/PEARC24/MNIST/docs/4_analyze_APSO.md
+++ b/manuscripts/PEARC24/MNIST/docs/4_analyze_APSO.md
@@ -1,17 +1,75 @@
-# Image Smoothing with Convolution Description
 
-The given code is a Python script that performs image smoothing using convolution. It applies a Gaussian smoothing filter to a set of images and saves the smoothed images.
+# Image Smoothing with Convolution
+
+## Description
+
+This Python script applies Gaussian smoothing (blur) to a set of images using convolution. It loads `.npy` image files, applies a smoothing filter to each image, and then saves the smoothed images as `.png` files. The script uses parallel processing to handle large sets of images efficiently.
+
 
 ## Functions
 
 ### `main()`
-This function is the main entry point of the script. It starts by defining variables for image file paths, filter size, and filter sigma. Then, it creates a pool of processes for parallel execution. The `smooth_and_save` function is applied to each image file path using the `imap` method of the pool. The progress is tracked using a progress bar from `tqdm`. Once the smoothing is completed for all images, a message is printed.
+
+The entry point of the script. It:
+
+1. Gathers paths to all `.npy` image files from the directory `./APSO/images/*/*npy`.
+2. Sets default values for the Gaussian filter size and sigma.
+3. Creates a pool of processes to apply the `smooth_and_save` function to each image in parallel.
+4. Tracks progress with a `tqdm` progress bar.
+5. Once all images are processed, it prints a completion message.
 
 ### `smooth_and_save(img_path)`
-This function takes an image file path as input. It loads the image from the file using `np.load` and applies the `smooth_array` function to smooth the image. The smoothed image is then plotted using `matplotlib.pyplot` and saved to a file with the same name but appended with "_convl.png". The figure is closed to free up resources.
+
+This function processes each image:
+
+1. Loads the image from the file specified by `img_path` using `np.load()`.
+2. Applies a Gaussian smoothing filter via the `smooth_array()` function.
+3. Saves the smoothed image as a `.png` file with the same base name but appended with `_convl.png`.
+4. The figure is saved with no axis labels or ticks for a clean result.
+5. The figure is closed after saving to free resources.
 
 ### `smooth_array(img, filter_size=5, filter_sigma=1)`
-This function performs image smoothing using convolution. It takes an input image and optional parameters for the filter size and filter sigma. It defines a Gaussian smoothing filter based on the filter size and sigma values. The filter is then applied to the image using `scipy.signal.convolve2d` with mode set to 'same', which ensures the output has the same shape as the input. The smoothed array is returned.
+
+This function performs the actual image smoothing:
+
+1. Defines a Gaussian kernel based on the given `filter_size` and `filter_sigma`.
+2. Applies the kernel to the image using 2D convolution via `scipy.signal.convolve2d()`.
+3. The convolution is performed with the `mode='same'` option, ensuring the output image has the same size as the input image.
+4. The smoothed image array is returned.
 
 ## Execution
-The script starts by importing the required libraries and modules. Then, it defines the main entry point function `main()` and the supporting functions `smooth_and_save()` and `smooth_array()`. Finally, the `main()` function is called to execute the script.
+
+1. Clone or download the script to your local machine.
+2. Place your `.npy` image files in the `./APSO/images/` directory or modify the script to point to your own directory containing `.npy` files.
+3. Run the script:
+
+   ```bash
+   python <script_name>.py
+   ```
+
+The script will process each image in parallel, smooth it with a Gaussian filter, and save the results as `.png` images in the same directory as the original images.
+
+### Example Run:
+
+```bash
+['./APSO/images/folder1/image1.npy', './APSO/images/folder2/image2.npy', ... ]
+Smoothing completed.
+Saved ./APSO/images/folder1/image1_convl.png
+Saved ./APSO/images/folder2/image2_convl.png
+...
+```
+
+## Customization
+
+* **Filter Size**: The default filter size is set to 5. You can adjust this in the `smooth_array()` function to control the level of smoothing. Larger values result in a stronger blur effect.
+
+* **Sigma**: The default sigma is set to 1. Increasing the sigma will make the Gaussian filter more spread out, which can result in a more pronounced smoothing effect.
+
+* **Image File Paths**: The script assumes images are stored in `.npy` format. Ensure your images are in this format or modify the script to handle other formats if needed.
+
+## Notes
+
+* The script uses parallel processing (`multiprocessing.Pool`) to speed up the smoothing of large sets of images.
+* The images are displayed using `matplotlib` and saved without axis labels for a cleaner output.
+* By default, the smoothed images are saved in the same directory as the original images with a `_convl.png` suffix. Modify the file-saving path in the `smooth_and_save()` function if needed.
+
diff --git a/manuscripts/PEARC24/MNIST/docs/5_create_SHAP.md b/manuscripts/PEARC24/MNIST/docs/5_create_SHAP.md
index 52e7e75..9321943 100644
--- a/manuscripts/PEARC24/MNIST/docs/5_create_SHAP.md
+++ b/manuscripts/PEARC24/MNIST/docs/5_create_SHAP.md
@@ -1,21 +1,109 @@
-# SHAP Values Generation and Visualization Description
+# SHAP Values Generation and Visualization
 
-The provided code is a Python script that generates and saves SHAP (SHapley Additive exPlanations) values for a given model and dataset. It visualizes the original images, along with their corresponding SHAP values, and saves the results.
+This Python script generates and visualizes SHAP (SHapley Additive exPlanations) values for a given model and dataset. SHAP values help explain the predictions made by a model, showing which parts of an input image were most important for the model's prediction. The script visualizes the original images alongside their SHAP value explanations and saves both the images and the visualizations for further analysis.
 
 ## Functions
 
 ### `getData(dataloader)`
-This function retrieves one sample of each class from the given dataloader. It iterates over the dataloader and collects unique samples based on their target labels. The data and target tensors are sorted based on the target labels. The data tensor is stacked along the 0th dimension, and the target tensor is converted to a torch tensor. The function returns the data and target tensors.
+
+This function retrieves one sample of each class from the provided dataloader. It:
+
+* Iterates over the dataloader to collect unique samples based on their target labels.
+* Sorts the data and target tensors by the target labels.
+* Stacks the data samples along the 0th dimension and converts the target list into a PyTorch tensor.
+
+**Returns:**
+
+* A tuple `(data, target)`, where:
+
+  * `data`: A tensor containing the samples, stacked along the 0th dimension.
+  * `target`: A tensor containing the corresponding class labels.
 
 ### `save_and_plot_shap_values(dataloader, model)`
-This function generates and saves the SHAP values for the given model and dataloader. It first checks the availability of a GPU device and clears the GPU memory. It retrieves the data and target tensors using the `getData` function. The tensors are moved to the GPU if available. The model is also moved to the GPU. A `shap.DeepExplainer` object is created with the model and data tensors. SHAP values are generated using the explainer. 
 
-The function creates a directory named 'SHAP' to store the results. It then creates a figure with a 10x11 grid of subplots. The function iterates over the data samples and their corresponding SHAP values. For each sample, the original image is saved as a numpy array. The original image and its label are plotted in the first column of the subplot grid. The SHAP values are plotted in the subsequent columns, up to a maximum of 10 values. Each SHAP value is saved as a numpy array. The function also creates individual row figures to remove empty cells and white space. These row figures are saved separately.
+This function generates SHAP values for the provided model and dataloader, then saves and visualizes the results. It:
+
+1. Checks the availability of a GPU device and moves the model and data to the appropriate device.
+2. Retrieves the data and target tensors using the `getData` function.
+3. Initializes a `shap.DeepExplainer` to compute SHAP values for each image in the batch.
+4. Creates a directory named `SHAP` to store the results (if it doesn't already exist).
+5. Creates a 10x11 grid of subplots:
+
+   * The first column shows the original image.
+   * The following columns show SHAP values for each class, with a maximum of 10 SHAP values per image.
+6. Saves each image and SHAP visualization as `.npy` (for raw data) and `.png` (for visualizations) files.
+7. Adds a colorbar to the figure and ensures the figure is properly saved and closed to free up resources.
+8. If there are fewer than 10 samples in the dataloader, the empty subplot cells are removed.
 
-Empty rows in the subplot grid are removed. A colorbar is added to the figure. The final figure, row figures, and colorbar are saved. The figures and subplots are closed to free up resources.
+**Parameters:**
+
+* `dataloader`: The PyTorch dataloader providing the dataset (e.g., the test set).
+* `model`: The trained model to generate SHAP values for.
 
 ### `main()`
-This function is the main entry point of the script. It loads the MNIST dataset using `load_MNIST_data` function and builds the MNIST model using `build_MNIST_Model`. The pre-trained model weights are loaded from the 'MNIST_cnn.pt' file. The `save_and_plot_shap_values` function is called with the test dataloader and the loaded model.
+
+The main entry point of the script:
+
+1. Loads the MNIST dataset using the `load_MNIST_data` function.
+2. Loads the pre-trained MNIST model and its weights from the file `'MNIST_cnn.pt'`.
+3. Calls the `save_and_plot_shap_values` function to generate and save SHAP values for the test set.
 
 ## Execution
-The script starts by importing the required libraries and modules. Then, it defines the `getData`, `save_and_plot_shap_values`, and `main` functions. Finally, the `main` function is called to execute the script.
+
+To run the script, follow these steps:
+
+1. **Install dependencies**: Ensure that you have the following Python packages installed:
+
+   * `torch`
+   * `numpy`
+   * `matplotlib`
+   * `shap`
+   * `Adversarial_Observation` (custom module)
+
+2. **Run the script**:
+
+   * The script automatically loads the MNIST dataset and the pre-trained model, then generates and saves SHAP values visualizations.
+   * You can run the script using the command:
+
+     ```bash
+     python shap_values_generation.py
+     ```
+
+3. **Output**:
+
+   * SHAP visualizations are saved in the `SHAP` directory.
+   * Each sample's image and its corresponding SHAP values for each class are saved as `.npy` files.
+   * Visualizations are saved as `.png` images, with a colorbar indicating the magnitude of the SHAP values.
+
+### Example Output:
+
+* A directory structure like:
+
+  ```
+  SHAP/
+    0_original.npy
+    0_shap_0.npy
+    0_shap_1.npy
+    ...
+    shap_values.png
+    row_0.png
+    row_1.png
+    ...
+  ```
+
+## Notes
+
+* The script assumes the MNIST images are 28x28 pixels, and the model is a CNN (Convolutional Neural Network) trained on MNIST data. If using a different dataset or model, adjustments to the code (such as image reshaping) may be required.
+* The SHAP values for each image are reshaped and visualized for each class in a 28x28 grid. Ensure that the model outputs correspond to a classification task where each class has a separate SHAP explanation.
+
+## Requirements
+
+* **Python**: Version 3.x
+* **Libraries**:
+
+  * `torch` (for PyTorch)
+  * `shap` (for SHAP explanations)
+  * `matplotlib` (for visualization)
+  * `numpy` (for numerical operations)
+  * `Adversarial_Observation` (custom module, used for loading MNIST data and model)
+
diff --git a/manuscripts/PEARC24/README.md b/manuscripts/PEARC24/README.md
deleted file mode 100644
index 1339114..0000000
--- a/manuscripts/PEARC24/README.md
+++ /dev/null
@@ -1,25 +0,0 @@
-
-# Examples
----
-
-## 1_build_train_save_artifacts.py
----
-
-This code defines a deep convolutional neural network (CNN) that is trained to classify handwritten digits using the MNIST dataset. The code uses the PyTorch framework to build and train the model.
-
-The loadData() function loads the MNIST dataset using the torchvision.datasets module and returns two data loaders - one for training data and the other for testing data. The train() function trains the model using the optimizer, loss function, and training data, while the test() function evaluates the model's accuracy on the test data. The reduce_data() function preprocesses the data by reducing its dimensionality using principal component analysis (PCA) and saves the transformed data to disk for later use. Finally, the main() function calls all the other functions in the correct sequence, builds the CNN, trains it, and saves the model parameters to disk.
-
-Overall, this code demonstrates how to build, train, and evaluate a deep learning model for image classification using PyTorch and the MNIST dataset. The code also illustrates how to preprocess the data using PCA and save the transformed data to disk.
-
-## 2_generate_attacks.py
----
-
-This code performs adversarial attacks and visualizations on the MNIST dataset. First, it loads the MNIST dataset using the loadData function, which applies some transformations to the data, and creates data loaders for the training and testing sets. Then, it loads a pre-trained CNN model from a saved state dict using the buildCNN function and the load_state_dict method.
-
-Next, the code applies the Fast Gradient Sign Method (FGSM) adversarial attack on a randomly selected image from the test set with varying levels of epsilon (the attack strength) and saves the perturbed images and their confidence scores in separate folders using the fgsm_attack and saliency_map functions. Finally, it generates a saliency map for the original image and saves it to a separate folder. The code uses various visualization techniques and settings to display the original and perturbed images, confidence scores, and saliency maps using Matplotlib.
-
-## 3_swarm_attack.py
----
-This code implements the Particle Swarm Optimization (PSO) algorithm to optimize a column of a pre-trained Convolutional Neural Network (CNN). The SwarmPSO() function initializes the swarm, consisting of multiple particles, with random initial positions. The swarm then updates its position by considering its previous position, its previous best position, and the global best position. The global best position is updated every time a particle finds a better position than the previous global best. In each epoch, the function evaluates the cost function on each particle's position, finds the best position, and updates the global best position. Finally, it visualizes the best position after the last epoch using imshow().
-
-The SwarmPSOVisualize() function is similar to SwarmPSO() but includes visualizations of the swarm in each epoch using visualizeSwarm() and plotInfo() functions. visualizeSwarm() plots the swarm's positions in the 50-dimensional PCA space reduced from the original 784-dimensional space of MNIST images. plotInfo() plots the average and the best position of the swarm, which are saved in the ./artifacts folder along with the GIF of the swarm's movement throughout the epochs. The swarm's data is also saved in a CSV file named swarm.csv for later analysis. 
\ No newline at end of file
diff --git a/manuscripts/Posion25/.gitignore b/manuscripts/Posion25/.gitignore
index bcce73b..c8a2052 100644
--- a/manuscripts/Posion25/.gitignore
+++ b/manuscripts/Posion25/.gitignore
@@ -3,4 +3,7 @@ results/
 *.pyc
 __pycache__/
 *.pkl
-*.keras
\ No newline at end of file
+*.keras
+slurm_jobs/
+false_data
+MNIST_train_*
diff --git a/manuscripts/Posion25/0_trainModel.py b/manuscripts/Posion25/0_trainModel.py
new file mode 100644
index 0000000..36697d1
--- /dev/null
+++ b/manuscripts/Posion25/0_trainModel.py
@@ -0,0 +1,33 @@
+import argparse
+import os
+from train import load_data, train_model, evaluate_model, train_model_and_save
+
+def main():
+    parser = argparse.ArgumentParser(description="Train and evaluate MNIST or AudioMNIST models")
+
+    # Dataset and model type
+    parser.add_argument('--data', type=str, choices=['MNIST', 'MNIST_Audio'], required=True,
+                        help='Dataset to use: MNIST or MNIST_Audio')
+    parser.add_argument('--model_type', type=str,
+                        choices=['normal', 'complex', 'complex_augmented', 'complex_adversarial'],
+                        required=True, help='Model architecture type')
+
+    # Training parameters
+    parser.add_argument('--batch_size', type=int, default=1024, help='Batch size for training')
+    parser.add_argument('--epochs', type=int, default=5, help='Number of epochs')
+
+    # Adversarial training option
+    parser.add_argument('--adversarial', type=str, default='none',
+                        choices=['none', 'pgd', 'trades'],
+                        help='Adversarial training method to use (if any)')
+
+    # Output directory
+    parser.add_argument('--save_dir', type=str, default='results', help='Directory to save trained models and results')
+
+    args = parser.parse_args()
+
+    # Run training + evaluation
+    train_model_and_save(args)
+
+if __name__ == '__main__':
+    main()
diff --git a/manuscripts/Posion25/0_train_models.sh b/manuscripts/Posion25/0_train_models.sh
new file mode 100644
index 0000000..171b562
--- /dev/null
+++ b/manuscripts/Posion25/0_train_models.sh
@@ -0,0 +1,33 @@
+#!/bin/bash
+
+module load anaconda3_gpu
+
+mkdir -p models
+
+DATASET="MNIST"
+BATCH_SIZE=32
+EPOCHS=5
+SAVE_DIR="models"
+
+# Define model types and adversarial modes
+MODEL_TYPES=("normal" "complex" "complex_augmented" "complex_adversarial")
+# For adversarial modes, only complex_adversarial uses them, others use "none"
+ADVERSARIAL_MODES=("none" "pgd" "trades")
+
+for model_type in "${MODEL_TYPES[@]}"; do
+    if [[ "$model_type" == "complex_adversarial" ]]; then
+        for adv in "${ADVERSARIAL_MODES[@]:1}"; do
+            echo "Training $DATASET with model_type=$model_type and adversarial=$adv"
+            echo "Command: python 0_trainModel.py --data $DATASET --model_type $model_type --adversarial $adv --batch_size $BATCH_SIZE --epochs $EPOCHS --save_dir $SAVE_DIR"
+            python 0_trainModel.py --data $DATASET --model_type $model_type --adversarial $adv --batch_size $BATCH_SIZE --epochs $EPOCHS --save_dir $SAVE_DIR
+        done
+    else
+        echo "Training $DATASET with model_type=$model_type (no adversarial)"
+        echo "Command: python 0_trainModel.py --data $DATASET --model_type $model_type --batch_size $BATCH_SIZE --epochs $EPOCHS --save_dir $SAVE_DIR"
+        python 0_trainModel.py --data $DATASET --model_type $model_type --batch_size $BATCH_SIZE --epochs $EPOCHS --save_dir $SAVE_DIR
+    fi
+done
+
+echo "Training completed. Models saved under $SAVE_DIR/ directory."
+
+exit
diff --git a/manuscripts/Posion25/1_dataset.sh b/manuscripts/Posion25/1_dataset.sh
new file mode 100644
index 0000000..6112c0d
--- /dev/null
+++ b/manuscripts/Posion25/1_dataset.sh
@@ -0,0 +1,28 @@
+#!/bin/bash
+
+# Configuration
+SEED=42
+BATCH_SIZE=32
+DATASETS=("mnist")
+SPLITS=("train" "test")
+
+# Step 1: Generate label TSVs
+echo "Generating label TSVs..."
+for dataset in "${DATASETS[@]}"; do
+    echo "Processing $dataset..."
+    python 1_dataset_label_tool.py --generate --dataset $dataset --batch_size $BATCH_SIZE
+done
+
+# Step 2: Generate false labels
+echo "Generating false labels..."
+for dataset in "${DATASETS[@]}"; do
+    for split in "${SPLITS[@]}"; do
+        input_file="${dataset^^}_${split}_labels.tsv"         # e.g., MNIST_train_labels.tsv
+        output_file="${dataset^^}_${split}_false_labels.tsv"  # e.g., MNIST_train_false_labels.tsv
+
+        echo "Generating false labels for $input_file..."
+        python 1_dataset_label_tool.py --input "$input_file" --output "$output_file" --seed $SEED
+    done
+done
+
+echo "All label and false label files generated."
diff --git a/manuscripts/Posion25/1_dataset_label_tool.py b/manuscripts/Posion25/1_dataset_label_tool.py
new file mode 100644
index 0000000..6ef38fd
--- /dev/null
+++ b/manuscripts/Posion25/1_dataset_label_tool.py
@@ -0,0 +1,111 @@
+import argparse
+import os
+import pandas as pd
+import numpy as np
+import tensorflow as tf
+from tensorflow.keras.datasets import cifar10, mnist
+from tensorflow.keras.utils import to_categorical
+
+# -------------------- Utility Functions --------------------
+
+def generate_false_labels(df, seed=None):
+    if seed is not None:
+        np.random.seed(seed)
+
+    num_classes = 10
+    false_labels = []
+
+    for true_label in df['Label']:
+        choices = [i for i in range(num_classes) if i != true_label]
+        false_label = np.random.choice(choices)
+        false_labels.append(false_label)
+
+    df_out = pd.DataFrame({
+        'index': df['Index'],
+        'trueLabel': df['Label'],
+        'falseLabel': false_labels
+    })
+    return df_out
+
+def standardize_data(x):
+    """Normalize data to [0, 1] range."""
+    return x / 255.0
+
+def write_labels_to_tsv(dataset, output_file):
+    with open(output_file, 'w') as f:
+        f.write("Index\tLabel\n")
+        for idx, (_, label) in enumerate(dataset.unbatch()):
+            true_label = int(np.argmax(label.numpy()))
+            f.write(f"{idx}\t{true_label}\n")
+    print(f"Wrote labels to {output_file}")
+
+# -------------------- Dataset Loaders --------------------
+
+def load_mnist(batch_size):
+    (x_train, y_train), (x_test, y_test) = mnist.load_data()
+    x_train = standardize_data(x_train.reshape(-1, 28, 28, 1).astype('float32'))
+    x_test = standardize_data(x_test.reshape(-1, 28, 28, 1).astype('float32'))
+    y_train = to_categorical(y_train, 10)
+    y_test = to_categorical(y_test, 10)
+    train_dataset = tf.data.Dataset.from_tensor_slices((x_train, y_train)).batch(batch_size)
+    test_dataset = tf.data.Dataset.from_tensor_slices((x_test, y_test)).batch(batch_size)
+    return train_dataset, test_dataset
+
+def load_cifar10(batch_size):
+    (x_train, y_train), (x_test, y_test) = cifar10.load_data()
+    x_train = standardize_data(x_train.astype('float32'))
+    x_test = standardize_data(x_test.astype('float32'))
+    y_train = to_categorical(y_train, 10)
+    y_test = to_categorical(y_test, 10)
+    train_dataset = tf.data.Dataset.from_tensor_slices((x_train, y_train)).batch(batch_size)
+    test_dataset = tf.data.Dataset.from_tensor_slices((x_test, y_test)).batch(batch_size)
+    return train_dataset, test_dataset
+
+# -------------------- Main Script --------------------
+
+def main():
+    parser = argparse.ArgumentParser(description="Dataset label handler.")
+    parser.add_argument('--dataset', choices=['mnist', 'cifar10'], help="Dataset to process.")
+    parser.add_argument('--generate', action='store_true', help="Generate TSV files from dataset.")
+    parser.add_argument('--input', help="Input TSV file (for false label generation).")
+    parser.add_argument('--output', help="Output TSV file name (saved in false_data/).")
+    parser.add_argument('--seed', type=int, default=None, help="Random seed for reproducibility.")
+    parser.add_argument('--batch_size', type=int, default=32, help="Batch size for data loading.")
+
+    args = parser.parse_args()
+
+    # Create folder for false labels
+    false_data_dir = "false_data"
+    os.makedirs(false_data_dir, exist_ok=True)
+
+    # Generate TSV files from the dataset
+    if args.generate:
+        if args.dataset == 'mnist':
+            train_ds, test_ds = load_mnist(args.batch_size)
+            write_labels_to_tsv(train_ds, os.path.join(false_data_dir, "MNIST_train_labels.tsv"))
+            write_labels_to_tsv(test_ds, os.path.join(false_data_dir, "MNIST_test_labels.tsv"))
+        elif args.dataset == 'cifar10':
+            train_ds, test_ds = load_cifar10(args.batch_size)
+            write_labels_to_tsv(train_ds, os.path.join(false_data_dir, "CIFAR10_train_labels.tsv"))
+            write_labels_to_tsv(test_ds, os.path.join(false_data_dir, "CIFAR10_test_labels.tsv"))
+        else:
+            print("Please specify a valid dataset with --dataset.")
+
+    # Generate false labels from input TSV
+    if args.input and args.output:
+        input_path = args.input
+        if not os.path.exists(input_path):
+            input_path = os.path.join(false_data_dir, args.input)
+
+        if not os.path.exists(input_path):
+            raise FileNotFoundError(f"Input file not found: {input_path}")
+
+        df = pd.read_csv(input_path, sep='\t')
+        df_out = generate_false_labels(df, seed=args.seed)
+
+        output_path = os.path.join(false_data_dir, args.output)
+        df_out.to_csv(output_path, sep='\t', index=False)
+        print(f"False labels saved to {output_path}")
+
+if __name__ == '__main__':
+    main()
diff --git a/manuscripts/Posion25/1_trainModel.py b/manuscripts/Posion25/1_trainModel.py
deleted file mode 100644
index 955dd8e..0000000
--- a/manuscripts/Posion25/1_trainModel.py
+++ /dev/null
@@ -1,28 +0,0 @@
-import argparse
-import os
-from train import load_data, train_model, evaluate_model, train_model_and_save
-
-
-def main():
-    # Command-line arguments
-    parser = argparse.ArgumentParser()
-
-    # Data and model type arguments
-    parser.add_argument('--data', type=str, choices=['MNIST', 'MNIST_Audio'], required=True, help='Dataset to use')
-    parser.add_argument('--model_type', type=str, choices=['normal', 'complex', 'complex_augmented'], required=True, help='Model type to use')
-
-    # Training information arguments
-    parser.add_argument('--batch_size', type=int, default=1024, help='Batch size for training')
-    parser.add_argument('--epochs', type=int, default=5, help='Number of epochs for training')
-
-    # Folder saving argument
-    parser.add_argument('--save_dir', type=str, default='results', help='Directory to save model and results')
-
-    # Parse arguments
-    args = parser.parse_args()
-
-    # Train the model and evaluate it
-    train_model_and_save(args)
-
-if __name__ == '__main__':
-    main()
diff --git a/manuscripts/Posion25/2_attackModel.py b/manuscripts/Posion25/2_attackModel.py
index d720a48..fb75664 100644
--- a/manuscripts/Posion25/2_attackModel.py
+++ b/manuscripts/Posion25/2_attackModel.py
@@ -1,102 +1,67 @@
 import argparse
 import os
 import pickle
-from taint import adversarial_attack_blackbox
-from analysis import *
-from train import train_model_and_save
-import torch
 import tensorflow as tf
+import torch
+from taint import adversarial_attack_blackbox
+
+
+def load_model(model_path):
+    # Assumes it's a Keras model (update if using PyTorch)
+    return tf.keras.models.load_model(model_path)
+
+
+def get_test_dataset(data_name):
+    # Import here to avoid unnecessary dependencies if unused
+    from train import load_data  # Ensure get_data returns (train_ds, test_ds)
+
+    train_ds, test_ds = load_data(dataset_type=data_name)
+    return test_ds
 
-def attack_model(args, model, test_ds, save_dir, num_data=10):
-    # Get the labels by iterating through a batch from the test_ds
-    first_batch = next(iter(test_ds))  # Get the first batch
-    images, labels = first_batch  # Unpack the images and labels from the first batch
-    
-    # Check if labels are a TensorFlow tensor or PyTorch tensor
-    if isinstance(labels, tf.Tensor):
-        # If using TensorFlow, convert labels to class indices (from one-hot encoded)
-        labels = tf.argmax(labels, axis=1).numpy()  # Get class indices from one-hot encoded labels
-    elif isinstance(labels, torch.Tensor):
-        # If using PyTorch, convert labels to class indices (from one-hot encoded)
-        labels = torch.argmax(labels, dim=1).cpu().numpy()  # Get class indices from one-hot encoded labels
-
-    # Convert labels to a set of unique outputs
-    unique_outputs = set(labels)  # Convert to a Python set for unique labels
-
-    # Continue with the rest of the attack logic
-    for output in unique_outputs:
-        instances = [i for i, label in enumerate(labels) if label == output][:num_data]  # Select `num_data` instances with the current output label
-        
-        for image_index in instances:
-            # Create a subdirectory for each image_index and its original output label
-            sub_dir = os.path.join(save_dir, f'image_{image_index}_label_{output}')
-            
-            # Ensure the directory exists
-            os.makedirs(sub_dir, exist_ok=True)
-
-            # Correct dynamic pickle filename to include the original and target class
-            pickle_filename = f'attacker_{image_index}_{output}.pkl'
-            pickle_path = os.path.join(sub_dir, pickle_filename)
-            
-            # Check if the attacker pickle already exists for this image_index and output
-            if os.path.exists(pickle_path):
-                with open(pickle_path, 'rb') as f:
-                    attacker = pickle.load(f)
-                print(f"Loaded attacker for image {image_index} with label {output} from {pickle_path}")
-            else:
-                print(f"Running adversarial attack for image {image_index} with label {output}...")
-                
-                # For the current `output`, target all other classes
-                for target_output in unique_outputs:
-                    if target_output != output:  # We want to target all other outputs
-                        for _ in range(num_data):  # Attack the target output `num_data` times
-                            target_sub_dir = os.path.join(sub_dir, f'target_{target_output}')
-                            os.makedirs(target_sub_dir, exist_ok=True)  # Create a subdir for each target class
-
-                            # Correct dynamic pickle filename to include the original and target class
-                            target_pickle_filename = f'attacker_{image_index}_{output}_to_{target_output}.pkl'
-                            target_pickle_path = os.path.join(target_sub_dir, target_pickle_filename)
-
-                            # Perform the adversarial attack targeting `target_output`
-                            attacker = adversarial_attack_blackbox(
-                                model=model,
-                                dataset=test_ds,
-                                image_index=image_index,
-                                output_dir=target_sub_dir,
-                                num_iterations=args.iterations,
-                                num_particles=args.particles,
-                                target_class=target_output  # Specify the target class for the attack
-                            )
-                            print(f"Adversarial attack completed for image {image_index} targeting class {target_output}")
-
-                            # After performing the attack, save the attacker object to a pickle file
-                            with open(target_pickle_path, 'wb') as f:
-                                pickle.dump(attacker, f)
-                            print(f"Saved attacker for image {image_index} with label {output} targeting {target_output} to {target_pickle_path}")
 
 def main():
-    # Command-line arguments
     parser = argparse.ArgumentParser()
 
-    # Data and model type arguments (to align with the ones used in the training script)
-    parser.add_argument('--data', type=str, choices=['MNIST', 'MNIST_Audio'], required=True, help='Dataset to use')
-    parser.add_argument('--model_type', type=str, choices=['normal', 'complex', 'complex_augmented'], required=True, help='Model type to use')
+    # Required args
+    parser.add_argument('--model_path', type=str, required=True, help='Path to saved model (.keras)')
+    parser.add_argument('--save_dir', type=str, required=True, help='Directory to save attack results')
+    parser.add_argument('--source_index', type=int, required=True, help='Index of image to attack')
+    parser.add_argument('--target', type=int, required=True, help='Target class for adversarial attack')
 
-    # Attack parameters
-    parser.add_argument('--iterations', type=int, default=10, help='Number of iterations for attack')
-    parser.add_argument('--particles', type=int, default=100, help='Number of particles for attack')
+    # Dataset config
+    parser.add_argument('--data', type=str, choices=['MNIST', 'MNIST_Audio'], required=True, help='Dataset name')
 
-    # Folder saving argument
-    parser.add_argument('--save_dir', type=str, default='results', help='Directory to save model and results')
+    # Attack config
+    parser.add_argument('--iterations', type=int, default=30, help='Number of attack iterations')
+    parser.add_argument('--particles', type=int, default=100, help='Number of swarm particles')
 
-    # Parse arguments
     args = parser.parse_args()
 
-    # First, train the model and get the necessary details for attack
-    model, test_ds, save_dir, model_path = train_model_and_save(args)
+    # Load model and dataset
+    model = load_model(args.model_path)
+    test_ds = get_test_dataset(args.data)
+
+    # Create output directory
+    os.makedirs(args.save_dir, exist_ok=True)
+
+    # Run the blackbox adversarial attack
+    attacker = adversarial_attack_blackbox(
+        model=model,
+        dataset=test_ds,
+        image_index=args.source_index,
+        output_dir=args.save_dir,
+        num_iterations=args.iterations,
+        num_particles=args.particles,
+        target_class=args.target
+    )
+
+    # Save attacker object
+    output_path = os.path.join(args.save_dir, f'attacker_{args.source_index}_to_{args.target}.pkl')
+    with open(output_path, 'wb') as f:
+        pickle.dump(attacker, f)
+
+    print(f"Attack complete. Saved attacker to: {output_path}")
 
-    # Perform the adversarial attack
-    attack_model(args, model, test_ds, save_dir)
 
 if __name__ == '__main__':
     main()
diff --git a/manuscripts/Posion25/2_attackModel.sh b/manuscripts/Posion25/2_attackModel.sh
new file mode 100644
index 0000000..18122d6
--- /dev/null
+++ b/manuscripts/Posion25/2_attackModel.sh
@@ -0,0 +1,67 @@
+#!/bin/bash
+
+# Base working directory
+WORKINGDIR=/workspaces/Adversarial_Observation/manuscripts/Posion25
+
+# Label and model paths
+LABELDIR=$WORKINGDIR/false_data
+MODELDIR=$WORKINGDIR/models/MNIST_normal
+MODEL=$MODELDIR/MNIST_normal.keras
+DATA=MNIST
+
+# Attack script (adjust path if needed)
+POISON=$WORKINGDIR/2_attackModel.py
+
+# Output directory for SLURM scripts
+OUTPUT=$WORKINGDIR/slurm_jobs
+mkdir -p $OUTPUT
+cd $OUTPUT
+
+# SLURM script header
+HEADER="#!/bin/bash\n#SBATCH -A bbse-delta-cpu\n#SBATCH --partition=cpu\n#SBATCH --nodes=1\n#SBATCH --tasks=1\n#SBATCH --cpus-per-task=4\n#SBATCH --mem=24g\n#SBATCH --time=8:00:00\n"
+
+# Label file to use
+LABELS=$LABELDIR/MNIST_train_false_labels.tsv
+
+# Attack parameters
+ITERATION=30
+PARTICLE=500
+
+# Initialize cohort tracking
+COHORT_ID=0
+COHORT_INDEX=0
+
+# Start first SLURM script
+echo "Preparing cohort: $COHORT_ID"
+echo -e $HEADER > $OUTPUT/attack_$COHORT_ID.slurm
+echo "module load anaconda3_gpu" >> $OUTPUT/attack_$COHORT_ID.slurm
+echo "cd $OUTPUT" >> $OUTPUT/attack_$COHORT_ID.slurm
+
+# Read through label file
+while read line; do
+    # Skip header
+    if [[ "$line" == index* ]]; then
+        continue
+    fi
+
+    # Parse line
+    index=$(echo "$line" | awk '{print $1}')
+    trueLabel=$(echo "$line" | awk '{print $2}')
+    falseLabel=$(echo "$line" | awk '{print $3}')
+
+    # Add attack command
+    echo "time python $POISON --data $DATA --model_path $MODEL --iterations $ITERATION --particles $PARTICLE --save_dir MNIST_train_$index --target $falseLabel --source_index $index" >> $OUTPUT/attack_$COHORT_ID.slurm
+
+    ((COHORT_INDEX++))
+
+    if [ $COHORT_INDEX -gt 100 ]; then
+        COHORT_INDEX=0
+        ((COHORT_ID++))
+
+        echo "Preparing cohort: $COHORT_ID"
+        echo -e $HEADER > $OUTPUT/attack_$COHORT_ID.slurm
+        echo "module load anaconda3_gpu" >> $OUTPUT/attack_$COHORT_ID.slurm
+        echo "cd $OUTPUT" >> $OUTPUT/attack_$COHORT_ID.slurm
+    fi
+
+done < "$LABELS"
diff --git a/manuscripts/Posion25/models.py b/manuscripts/Posion25/models.py
index 40584b2..8c13f35 100644
--- a/manuscripts/Posion25/models.py
+++ b/manuscripts/Posion25/models.py
@@ -1,21 +1,22 @@
-
 def load_MNIST_model(model_path=None):
     from tensorflow.keras.models import Sequential
     from tensorflow.keras.layers import Conv2D, MaxPooling2D, Flatten, Dense, Dropout, BatchNormalization
     from tensorflow.keras.optimizers import Adam
     
+     # Create a new Keras model
     model = Sequential([
-        Conv2D(32, (3, 3), padding='same', activation='relu', input_shape=(28, 28, 1)),
-        MaxPooling2D((2, 2)),
-        Conv2D(64, (3, 3), padding='same', activation='relu'),
-        MaxPooling2D((2, 2)),
+        Conv2D(32, kernel_size=(3, 3), padding='same', activation='relu', input_shape=(28, 28, 1)),
+        MaxPooling2D(pool_size=(2, 2)),
+        Conv2D(64, kernel_size=(3, 3), padding='same', activation='relu'),
+        MaxPooling2D(pool_size=(2, 2)),
         Flatten(),
         Dense(128, activation='relu'),
-        Dropout(0.5),
         Dense(10, activation='softmax')
     ])
-
-    model.compile(optimizer=Adam(), loss='categorical_crossentropy', metrics=['accuracy'])
+    # Compile the model
+    model.compile(optimizer=Adam(learning_rate=0.001), loss='categorical_crossentropy', metrics=['accuracy'])
+    # Summarize the model
+    model.summary()
     return model
 
 def load_complex_MNIST_model(model_path=None):
@@ -23,49 +24,96 @@ def load_complex_MNIST_model(model_path=None):
     from tensorflow.keras.layers import Conv2D, MaxPooling2D, Flatten, Dense, Dropout, BatchNormalization, GlobalAveragePooling2D
     from tensorflow.keras.optimizers import Adam
 
+    # Create a new Keras model
     model = Sequential([
-        # First block: Convolution + Pooling
-        Conv2D(32, (3, 3), padding='same', activation='relu', input_shape=(28, 28, 1)),
+        Conv2D(32, kernel_size = 3, activation='relu', input_shape = (28, 28, 1)),
         BatchNormalization(),
-        MaxPooling2D((2, 2)),
-
-        # Second block: Deeper convolution + Pooling
-        Conv2D(64, (3, 3), padding='same', activation='relu'),
+        Conv2D(32, kernel_size = 3, activation='relu'),
         BatchNormalization(),
-        MaxPooling2D((2, 2)),
-
-        # Third block: Deeper convolution + Pooling
-        Conv2D(128, (3, 3), padding='same', activation='relu'),
+        Conv2D(32, kernel_size = 5, strides=2, padding='same', activation='relu'),
         BatchNormalization(),
-        MaxPooling2D((2, 2)),
-
-        # Fourth block: Additional convolutions to add complexity
-        Conv2D(256, (3, 3), padding='same', activation='relu'),
+        Dropout(0.4),    
+        Conv2D(64, kernel_size = 3, activation='relu'),
         BatchNormalization(),
-        MaxPooling2D((2, 2)),
-
-        # Global Average Pooling instead of Flatten for better generalization
-        GlobalAveragePooling2D(),
-
-        # Fully connected layers
-        Dense(512, activation='relu'),
-        Dropout(0.5),
-
-        Dense(256, activation='relu'),
+        Conv2D(64, kernel_size = 3, activation='relu'),
+        BatchNormalization(),
+        Conv2D(64, kernel_size = 5, strides=2, padding='same', activation='relu'),
+        BatchNormalization(),
+        Dropout(0.4),
+        Conv2D(128, kernel_size = 4, activation='relu'),
+        BatchNormalization(),
+        Flatten(),
         Dropout(0.4),
+        Dense(10, activation='softmax')
+    ])
+    # Compile the model
+    model.compile(optimizer=Adam(learning_rate=0.001), loss='categorical_crossentropy', metrics=['accuracy'])
+    # Summarize the model
+    model.summary()
+    return model
 
-        Dense(128, activation='relu'),
-        Dropout(0.3),
+def load_CIFAR10_model(model_path=None):
+    from tensorflow.keras.models import Sequential
+    from tensorflow.keras.layers import Conv2D, MaxPooling2D, Flatten, Dense, Dropout, BatchNormalization, GlobalAveragePooling2D
+    from tensorflow.keras.optimizers import Adam
 
-        # Output layer
+    """
+    Creates a CNN model for CIFAR-10 dataset.
+    """
+    model = Sequential([
+        Conv2D(32, (3, 3), activation='relu', input_shape=(32, 32, 3)),
+        MaxPooling2D((2, 2)),
+        Conv2D(64, (3, 3), activation='relu'),
+        MaxPooling2D((2, 2)),
+        Conv2D(64, (3, 3), activation='relu'),
+        Flatten(),
+        Dense(64, activation='relu'),
         Dense(10, activation='softmax')
     ])
+    # Compile the model
+    model.compile(optimizer=Adam(learning_rate=0.001), loss='categorical_crossentropy', metrics=['accuracy'])
+    # Summarize the model
+    model.summary()
+    return model
 
-    # Compile the model with Adam optimizer and categorical crossentropy loss
-    model.compile(optimizer=Adam(), loss='categorical_crossentropy', metrics=['accuracy'])
+def load_complex_CIFAR10_model(model_path=None):
+    from tensorflow.keras.models import Sequential
+    from tensorflow.keras.layers import Conv2D, MaxPooling2D, Flatten, Dense, Dropout, BatchNormalization, GlobalAveragePooling2D, MaxPool2D
+    from tensorflow.keras.optimizers import Adam
     
-    return model
 
+    """
+    Creates a CNN model for CIFAR-10 dataset.
+    """
+    model = Sequential([
+        Conv2D(filters=32, kernel_size=(3, 3), activation='relu', input_shape=(32, 32, 3)),
+        BatchNormalization(),
+        Conv2D(filters=32, kernel_size=(3, 3), activation='relu'),
+        BatchNormalization(),
+        MaxPool2D(pool_size=(2, 2)),
+        Dropout(0.25),
+        Conv2D(filters=64, kernel_size=(3, 3), activation='relu', padding='same'),
+        BatchNormalization(),
+        Conv2D(filters=64, kernel_size=(3, 3), activation='relu'),
+        BatchNormalization(),
+        MaxPool2D(pool_size=(2, 2)),
+        Dropout(0.25),
+        Conv2D(filters=128, kernel_size=(3, 3), activation='relu', padding='same'),
+        BatchNormalization(),
+        Conv2D(filters=128, kernel_size=(3, 3), activation='relu'),
+        BatchNormalization(),
+        MaxPool2D(pool_size=(2, 2)),
+        Dropout(0.25),
+        Flatten(),
+        Dense(128, activation='relu'),
+        Dropout(0.25),
+        Dense(10, activation='softmax')
+    ])
+    # Compile the model
+    model.compile(optimizer=Adam(learning_rate=0.001), loss='categorical_crossentropy', metrics=['accuracy'])
+    # Summarize the model
+    model.summary()
+    return model
 
 def load_AudioMNIST_model(model_path=None):
     from tensorflow.keras.models import Sequential
diff --git a/manuscripts/Posion25/taint.py b/manuscripts/Posion25/taint.py
index 1740bac..3d64f8f 100644
--- a/manuscripts/Posion25/taint.py
+++ b/manuscripts/Posion25/taint.py
@@ -1,13 +1,15 @@
 import os
+import json
 import numpy as np
 import tensorflow as tf
 import pickle
+from tqdm import tqdm
 from manuscripts.Posion25.analysis import *
 from Adversarial_Observation.Swarm import ParticleSwarm
 from analysis import *
 
+
 def adversarial_attack_blackbox(model, dataset, image_index, output_dir='results', num_iterations=30, num_particles=100, target_class=None):
-    
     pickle_path = os.path.join(output_dir, 'attacker.pkl')
 
     dataset_list = list(dataset.as_numpy_iterator())
@@ -23,8 +25,14 @@ def adversarial_attack_blackbox(model, dataset, image_index, output_dir='results
     if target_class is None:
         target_class = (single_target + 1) % 10
 
+    if single_target == target_class:
+        raise ValueError("Target class must be different from original class")
+
+    print(f"Original class: {single_target}")
+    print(f"Target misclassification class: {target_class}")
+
     input_set = np.stack([
-        single_input + (np.random.uniform(0, 1, single_input.shape) * (np.random.rand(*single_input.shape) < 0.9))
+        single_input + (np.random.uniform(0, 1, single_input.shape) * (np.random.rand(*single_input.shape) < 0.7))
         for _ in range(num_particles)
     ])
 
@@ -33,32 +41,35 @@ def adversarial_attack_blackbox(model, dataset, image_index, output_dir='results
             attacker = pickle.load(f)
         print(f"Loaded attacker from {pickle_path}")
     else:
-
         attacker = ParticleSwarm(
             model=model, input_set=input_set, starting_class=single_target,
             target_class=target_class, num_iterations=num_iterations,
             save_dir=output_dir, inertia_weight=0.01
         )
         attacker.optimize()
-        # save the attacker as a pickle
         with open(pickle_path, 'wb') as f:
             pickle.dump(attacker, f)
         print(f"Saved attacker to {pickle_path}")
+
     print("Adversarial attack completed. Analyzing results...")
     analyze_attack(attacker, single_input, target_class)
+
+    final_pred = predict_class(model, attacker.global_best_position.numpy())
+    print(f"Attack {'succeeded' if final_pred == target_class else 'failed'}. Final prediction: {final_pred}")
     return attacker
 
+
 def best_analysis(attacker, original_data, target):
     adv = attacker.global_best_position.numpy()
     save_dir = attacker.save_dir
     ensure_dir(save_dir)
 
-    # save the original data
-    save_array_csv(os.path.join(save_dir, "original_data.csv"), original_data)
-    save_ndarray_visualization(os.path.join(save_dir, "original_data.png"), original_data)
+    # Save original input
+    save_array_csv(os.path.join(save_dir, "original_image.csv"), original_data)
+    save_ndarray_visualization(os.path.join(save_dir, "original_image.png"), original_data)
     save_softmax_stats(os.path.join(save_dir, "original_data_stats.tsv"),
                        *get_softmax_stats(attacker.model, original_data), target)
-    
+
     # Save best particle
     save_array_csv(os.path.join(save_dir, "best_particle.csv"), adv)
     save_ndarray_visualization(os.path.join(save_dir, "best_particle.png"), adv)
@@ -73,9 +84,6 @@ def best_analysis(attacker, original_data, target):
         diff, mode="auto", cmap="seismic", vmin=-1, vmax=1
     )
 
-    # Save stats
-    softmax_output, max_val, max_class = get_softmax_stats(attacker.model, adv)
-    save_softmax_stats(os.path.join(save_dir, "best_particle_stats.tsv"), softmax_output, max_class, max_val, target)
 
 def denoise_analysis(attacker, original_data, denoised_data, target):
     save_dir = attacker.save_dir
@@ -92,13 +100,15 @@ def denoise_analysis(attacker, original_data, denoised_data, target):
     )
 
     softmax_output, max_val, max_class = get_softmax_stats(attacker.model, denoised_data)
-    save_softmax_stats(os.path.join(save_dir, "best_particle-clean_stats.tsv"), softmax_output, max_class, max_val, target)
+    save_softmax_stats(os.path.join(save_dir, "best_particle-clean_stats.tsv"),
+                       softmax_output, max_class, max_val, target)
+
+    print(f"Denoised image predicted class: {max_class} with confidence: {max_val}")
+    if max_class != target:
+        print("Warning: Denoised image no longer classified as target!")
+
 
 def reduce_excess_perturbations(attacker, original_data, adv_data, target_label):
-    """
-    Reduce unnecessary perturbations in adversarial data while maintaining misclassification.
-    Works for data of any shape.
-    """
     original_data = np.squeeze(original_data)
     adv_data = np.squeeze(adv_data)
 
@@ -106,13 +116,13 @@ def reduce_excess_perturbations(attacker, original_data, adv_data, target_label)
         raise ValueError("Original and adversarial data must have the same shape after squeezing.")
 
     adv_data = adv_data.copy()
-    changed = True
+    max_iterations = 5
+    iteration = 0
 
-    # Wrap the iteration with tqdm to monitor progress
-    while changed:
+    while iteration < max_iterations:
         changed = False
         indices = list(np.ndindex(original_data.shape))
-        for idx in tqdm(indices, desc="Reducing perturbations"):
+        for idx in tqdm(indices, desc=f"Reducing perturbations (iter {iteration+1})"):
             if np.isclose(original_data[idx], adv_data[idx]):
                 continue
 
@@ -134,8 +144,13 @@ def reduce_excess_perturbations(attacker, original_data, adv_data, target_label)
             else:
                 changed = True
 
+        if not changed:
+            break
+        iteration += 1
+
     return adv_data
 
+
 def reduce_excess_perturbations_scale(attacker, original_data, adv_data, target_label, tol=1e-4, max_iter=20):
     original_data = np.squeeze(original_data)
     adv_data = np.squeeze(adv_data)
@@ -163,10 +178,8 @@ def reduce_excess_perturbations_scale(attacker, original_data, adv_data, target_
 
     return best_scaled
 
+
 def full_analysis(attacker, input_data, target):
-    """
-    Save full analysis of all particles' histories and confidences.
-    """
     analysis = {
         "original_misclassification_input": input_data.tolist(),
         "original_misclassification_target": int(target),
@@ -183,7 +196,7 @@ def full_analysis(attacker, input_data, target):
             "differences_from_original": []
         }
 
-        for pos in tqdm(particle.history, desc=f"Particle {i} history", leave=False):
+        for step_idx, pos in enumerate(tqdm(particle.history, desc=f"Particle {i} history", leave=False)):
             pos_np = pos.numpy() if isinstance(pos, tf.Tensor) else np.array(pos)
             softmax, max_val, max_class = get_softmax_stats(attacker.model, pos_np)
             diff = float(np.linalg.norm(pos_np - input_data))
@@ -194,6 +207,12 @@ def full_analysis(attacker, input_data, target):
             pdata["max_output_classes"].append(max_class)
             pdata["differences_from_original"].append(diff)
 
+            # Optional: Save intermediate images
+            iter_dir = os.path.join(attacker.save_dir, f"particle_{i}_step_{step_idx}")
+            ensure_dir(iter_dir)
+            save_ndarray_visualization(os.path.join(iter_dir, "perturbed.png"), pos_np, cmap="gray")
+            save_ndarray_visualization(os.path.join(iter_dir, "perturbation.png"), pos_np - input_data, cmap="seismic", vmin=-1, vmax=1)
+
         analysis["particles"].append(pdata)
 
     path = os.path.join(attacker.save_dir, "attack_analysis.json")
@@ -202,11 +221,29 @@ def full_analysis(attacker, input_data, target):
 
     print(f"Full analysis saved to {path}")
 
+
 def analyze_attack(attacker, original_img, target):
     print("Starting analysis of the adversarial attack...")
     best_analysis(attacker, original_img, target)
+
     print("Reducing excess perturbations...")
     reduced_img = reduce_excess_perturbations(attacker, original_img, attacker.global_best_position.numpy(), target)
+
     denoise_analysis(attacker, original_img, reduced_img, target)
+
     print("Performing full analysis of the attack...")
     full_analysis(attacker, original_img, target)
+
+
+def pgd_attack(model, images, labels, eps=0.3, alpha=0.01, steps=40):
+    adv_images = tf.identity(images)
+    for _ in range(steps):
+        with tf.GradientTape() as tape:
+            tape.watch(adv_images)
+            preds = model(adv_images)
+            loss = tf.keras.losses.categorical_crossentropy(labels, preds)
+        gradients = tape.gradient(loss, adv_images)
+        adv_images = adv_images + alpha * tf.sign(gradients)
+        adv_images = tf.clip_by_value(adv_images, images - eps, images + eps)
+        adv_images = tf.clip_by_value(adv_images, 0.0, 1.0)
+    return adv_images
diff --git a/manuscripts/Posion25/temp/01_generate_attack_labels.sh b/manuscripts/Posion25/temp/01_generate_attack_labels.sh
new file mode 100644
index 0000000..c76680d
--- /dev/null
+++ b/manuscripts/Posion25/temp/01_generate_attack_labels.sh
@@ -0,0 +1,28 @@
+#!/bin/bash
+
+# Load environment
+module load anaconda3_cpu
+
+# Create output directory
+mkdir -p labels
+
+# Paths to scripts
+LABEL_SCRIPT="bin/utils/generate_FalseLabels.py"
+OUTPUT_SCRIPTS=("bin/utils/output_MNIST_labels.py" "bin/utils/output_CIFAR10_labels.py")
+
+# Loop through each dataset output script
+for OUTPUT_SCRIPT in "${OUTPUT_SCRIPTS[@]}"; do
+    python "$OUTPUT_SCRIPT"
+    
+    # Extract dataset name (e.g., MNIST or CIFAR10)
+    DATASET=$(basename "$OUTPUT_SCRIPT" | cut -d'_' -f2)
+
+    # Move generated label files
+    mv ${DATASET}*labels.tsv labels/
+
+    # Generate false labels
+    python "$LABEL_SCRIPT" \
+        --input "labels/${DATASET}_test_labels.tsv" \
+        --output "labels/${DATASET}_test_labels-misclassify.tsv" \
+        --seed 1
+done
diff --git a/manuscripts/Posion25/temp/3_generate_attack_label.py b/manuscripts/Posion25/temp/3_generate_attack_label.py
new file mode 100644
index 0000000..32e431f
--- /dev/null
+++ b/manuscripts/Posion25/temp/3_generate_attack_label.py
@@ -0,0 +1,87 @@
+import argparse
+import pandas as pd
+import numpy as np
+import tensorflow as tf
+from tensorflow.keras.utils import to_categorical
+from tensorflow.keras.datasets import mnist, cifar10
+
+def standardize_mnist(x):
+    return x / 255.0
+
+def standardize_cifar(x):
+    return x / 255.0
+
+def load_data(dataset_name, batch_size=32):
+    if dataset_name == 'mnist':
+        (x_train, y_train), (x_test, y_test) = mnist.load_data()
+        x_train = standardize_mnist(x_train.reshape(-1, 28, 28, 1).astype('float32'))
+        x_test = standardize_mnist(x_test.reshape(-1, 28, 28, 1).astype('float32'))
+    elif dataset_name == 'cifar10':
+        (x_train, y_train), (x_test, y_test) = cifar10.load_data()
+        x_train = standardize_cifar(x_train.astype('float32'))
+        x_test = standardize_cifar(x_test.astype('float32'))
+    else:
+        raise ValueError("Unsupported dataset. Choose 'mnist' or 'cifar10'.")
+
+    y_train = to_categorical(y_train, 10)
+    y_test = to_categorical(y_test, 10)
+
+    train_dataset = tf.data.Dataset.from_tensor_slices((x_train, y_train)).batch(batch_size)
+    test_dataset = tf.data.Dataset.from_tensor_slices((x_test, y_test)).batch(batch_size)
+
+    return train_dataset, test_dataset
+
+def write_labels_to_tsv(dataset, filename):
+    with open(filename, 'w') as f:
+        f.write("Index\tLabel\n")
+        for idx, (_, label) in enumerate(dataset.unbatch()):
+            true_label = int(np.argmax(label.numpy()))
+            f.write(f"{idx}\t{true_label}\n")
+    print(f"Wrote labels to {filename}")
+
+def generate_false_labels(df, num_classes=10, seed=None):
+    if seed is not None:
+        np.random.seed(seed)
+
+    false_labels = []
+    for true_label in df['Label']:
+        choices = [i for i in range(num_classes) if i != true_label]
+        false_label = np.random.choice(choices)
+        false_labels.append(false_label)
+
+    df_out = pd.DataFrame({
+        'Index': df['Index'],
+        'trueLabel': df['Label'],
+        'falseLabel': false_labels
+    })
+    return df_out
+
+def main():
+    parser = argparse.ArgumentParser(description="Load dataset, write labels, and optionally generate false labels TSV.")
+    parser.add_argument('--dataset', type=str, required=True, choices=['mnist', 'cifar10'], help='Dataset to load.')
+    parser.add_argument('--batch_size', type=int, default=32, help='Batch size for dataset.')
+    parser.add_argument('--output_prefix', type=str, default='output', help='Prefix/path for output TSV files.')
+    parser.add_argument('--seed', type=int, default=None, help='Random seed for false label generation.')
+    parser.add_argument('--generate_false_labels', action='store_true', help='Generate false labels TSV after label files are created.')
+    args = parser.parse_args()
+
+    train_dataset, test_dataset = load_data(args.dataset, batch_size=args.batch_size)
+
+    train_labels_file = f"{args.output_prefix}_{args.dataset}_train_labels.tsv"
+    test_labels_file = f"{args.output_prefix}_{args.dataset}_test_labels.tsv"
+
+    write_labels_to_tsv(train_dataset, train_labels_file)
+    write_labels_to_tsv(test_dataset, test_labels_file)
+
+    if args.generate_false_labels:
+        # Read back the label files and generate false labels
+        for split in ['train', 'test']:
+            label_file = f"{args.output_prefix}_{args.dataset}_{split}_labels.tsv"
+            df_labels = pd.read_csv(label_file, sep='\t')
+            df_false = generate_false_labels(df_labels, seed=args.seed)
+            false_label_file = f"{args.output_prefix}_{args.dataset}_{split}_false_labels.tsv"
+            df_false.to_csv(false_label_file, sep='\t', index=False)
+            print(f"Generated false labels TSV at {false_label_file}")
+
+if __name__ == '__main__':
+    main()
diff --git a/manuscripts/Posion25/3_stats.py b/manuscripts/Posion25/temp/4_stats.py
similarity index 99%
rename from manuscripts/Posion25/3_stats.py
rename to manuscripts/Posion25/temp/4_stats.py
index e45660f..527e0d7 100644
--- a/manuscripts/Posion25/3_stats.py
+++ b/manuscripts/Posion25/temp/4_stats.py
@@ -115,7 +115,7 @@ def run_statistics(args):
     # Save the collected statistics as a JSON file for later analysis
     stats_file = os.path.join(results_dir, f"{args.data}_attack_stats.json")
     with open(stats_file, 'w') as f:
-        json.dump(all_stats, f, indent=4)
+        json.dump(all_stats, f, indent=4)   
 
     print(f"Statistics saved to {stats_file}")
 
diff --git a/manuscripts/Posion25/train.py b/manuscripts/Posion25/train.py
index 09c829a..541e4d7 100644
--- a/manuscripts/Posion25/train.py
+++ b/manuscripts/Posion25/train.py
@@ -15,6 +15,8 @@
 from tensorflow.keras.preprocessing.image import ImageDataGenerator
 from tensorflow.keras.utils import to_categorical
 from models import *
+from taint import pgd_attack
+
 SAMPLING_RATE = 16000
 NUM_CLASSES = 10
 RANDOM_SEED = 42
@@ -44,26 +46,66 @@ def load_dataset(data_path):
     data = [pad_audio(audio, max_len) for audio in data]
     return np.array(data), np.array(labels), max_len
 
-def prepare_datasets(data, labels, max_len, test_size=0.2):
+def augment_audio(audio, label):
+    def add_noise(audio):
+        noise = tf.random.normal(shape=tf.shape(audio), mean=0.0, stddev=0.005)
+        return audio + noise
+
+    def time_shift(audio):
+        shift = tf.random.uniform([], minval=-1600, maxval=1600, dtype=tf.int32)
+        return tf.roll(audio, shift=shift, axis=0)
+
+    audio = add_noise(audio)
+    audio = time_shift(audio)
+    return audio, label
+
+def generate_adversarial_audio(audio, label):
+    # Placeholder for actual adversarial attack
+    perturbation = tf.random.normal(tf.shape(audio), mean=0.0, stddev=0.01)
+    adversarial_audio = tf.clip_by_value(audio + perturbation, -1.0, 1.0)
+    return adversarial_audio, label
+
+def prepare_datasets(data, labels, max_len, test_size=0.2, use_augmentation=False, adversarial="none"):
     data = np.array([pad_audio(audio, max_len) for audio in data])[..., np.newaxis]
     x_train, x_test, y_train, y_test = train_test_split(data, labels, test_size=test_size, random_state=RANDOM_SEED)
     x_train = x_train.astype(np.float32) / np.max(np.abs(x_train))
     x_test = x_test.astype(np.float32) / np.max(np.abs(x_test))
     y_train = to_categorical(y_train, NUM_CLASSES)
     y_test = to_categorical(y_test, NUM_CLASSES)
-    train_ds = tf.data.Dataset.from_tensor_slices((x_train, y_train)).batch(BATCH_SIZE)
-    test_ds = tf.data.Dataset.from_tensor_slices((x_test, y_test)).batch(BATCH_SIZE)
+
+    train_ds = tf.data.Dataset.from_tensor_slices((x_train, y_train))
+    test_ds = tf.data.Dataset.from_tensor_slices((x_test, y_test))
+
+    if use_augmentation:
+        train_ds = train_ds.map(augment_audio, num_parallel_calls=tf.data.AUTOTUNE)
+
+    if adversarial in ["train", "both"]:
+        train_ds = train_ds.map(generate_adversarial_audio, num_parallel_calls=tf.data.AUTOTUNE)
+
+    if adversarial in ["test", "both"]:
+        test_ds = test_ds.map(generate_adversarial_audio, num_parallel_calls=tf.data.AUTOTUNE)
+
+    train_ds = train_ds.batch(BATCH_SIZE).prefetch(tf.data.AUTOTUNE)
+    test_ds = test_ds.batch(BATCH_SIZE).prefetch(tf.data.AUTOTUNE)
+
     return train_ds, test_ds, (x_test, y_test)
 
-def load_audio_mnist_data(data_path):
+def load_audio_mnist_data(data_path, use_augmentation=False, adversarial="none"):
     if not os.path.exists(data_path):
         raise FileNotFoundError(f"Data path {data_path} does not exist.")
     data, labels, max_len = load_dataset(data_path)
-    train_ds, test_ds, _ = prepare_datasets(data, labels, max_len)
+    train_ds, test_ds, _ = prepare_datasets(data, labels, max_len, use_augmentation=use_augmentation, adversarial=adversarial)
     return train_ds, test_ds, max_len
 
+from tensorflow.keras.datasets import cifar10
 
-def load_data(batch_size=32, dataset_type="MNIST", use_augmentation=False):
+def normalize_cifar10(x):
+    return x.astype('float32') / 255.0
+
+def load_data(batch_size=32, dataset_type="MNIST", use_augmentation=False, adversarial="none"):
+    """
+    Loads the dataset (MNIST, CIFAR10, or MNIST_Audio) with optional augmentation and adversarial setting.
+    """
     if dataset_type == "MNIST":
         (x_train, y_train), (x_test, y_test) = mnist.load_data()
         x_train = normalize_mnist(x_train.reshape(-1, 28, 28, 1).astype('float32'))
@@ -72,7 +114,6 @@ def load_data(batch_size=32, dataset_type="MNIST", use_augmentation=False):
         y_test = to_categorical(y_test, 10)
 
         if use_augmentation:
-            # Data Augmentation with ImageDataGenerator
             datagen = ImageDataGenerator(
                 rotation_range=10,
                 zoom_range=0.10,
@@ -80,12 +121,41 @@ def load_data(batch_size=32, dataset_type="MNIST", use_augmentation=False):
                 height_shift_range=0.1
             )
             datagen.fit(x_train)
-            train_datagen = datagen.flow(x_train, y_train, batch_size=batch_size)
+            train_generator = datagen.flow(x_train, y_train, batch_size=batch_size)
+            train_dataset = tf.data.Dataset.from_generator(
+                lambda: train_generator,
+                output_signature=(
+                    tf.TensorSpec(shape=(None, 28, 28, 1), dtype=tf.float32),
+                    tf.TensorSpec(shape=(None, 10), dtype=tf.float32)
+                )
+            )
+        else:
+            train_dataset = tf.data.Dataset.from_tensor_slices((x_train, y_train)).batch(batch_size)
+
+        test_dataset = tf.data.Dataset.from_tensor_slices((x_test, y_test)).batch(batch_size)
+        return train_dataset, test_dataset
+
+    elif dataset_type == "CIFAR10":
+        (x_train, y_train), (x_test, y_test) = cifar10.load_data()
+        x_train = normalize_cifar10(x_train)
+        x_test = normalize_cifar10(x_test)
+        y_train = to_categorical(y_train, 10)
+        y_test = to_categorical(y_test, 10)
+
+        if use_augmentation:
+            datagen = ImageDataGenerator(
+                rotation_range=15,
+                width_shift_range=0.1,
+                height_shift_range=0.1,
+                horizontal_flip=True
+            )
+            datagen.fit(x_train)
+            train_generator = datagen.flow(x_train, y_train, batch_size=batch_size)
             train_dataset = tf.data.Dataset.from_generator(
-                lambda: train_datagen,
+                lambda: train_generator,
                 output_signature=(
-                    tf.TensorSpec(shape=(batch_size, 28, 28, 1), dtype=tf.float32),
-                    tf.TensorSpec(shape=(batch_size, 10), dtype=tf.float32)
+                    tf.TensorSpec(shape=(None, 32, 32, 3), dtype=tf.float32),
+                    tf.TensorSpec(shape=(None, 10), dtype=tf.float32)
                 )
             )
         else:
@@ -96,12 +166,12 @@ def load_data(batch_size=32, dataset_type="MNIST", use_augmentation=False):
 
     elif dataset_type == "MNIST_Audio":
         data_path = "./AudioMNIST/data"
-        train_ds, test_ds, max_len = load_audio_mnist_data(data_path)
+        train_ds, test_ds, _ = load_audio_mnist_data(data_path, use_augmentation=use_augmentation, adversarial=adversarial)
         return train_ds, test_ds
 
     else:
         raise ValueError(f"Unsupported dataset type: {dataset_type}")
-
+ 
 def train_model(model, train_dataset, epochs=10):
     model.compile(optimizer=Adam(0.001), loss='categorical_crossentropy', metrics=['accuracy'])
     for epoch in range(epochs):
@@ -127,49 +197,111 @@ def evaluate_model(model, test_dataset):
     auprc = average_precision_score(to_categorical(y_true, NUM_CLASSES), y_pred)
     print(f"Test Loss: {loss:.4f}, Accuracy: {acc:.4f}, AUROC: {auroc:.4f}, AUPRC: {auprc:.4f}")
 
+def trades_loss(model, x_natural, y, eps=0.3, alpha=0.01, steps=10, beta=6.0):
+    x_adv = tf.identity(x_natural) + 0.001 * tf.random.normal(tf.shape(x_natural))
+    for _ in range(steps):
+        with tf.GradientTape() as tape:
+            tape.watch(x_adv)
+            kl_loss = tf.keras.losses.KLDivergence()(
+                tf.nn.softmax(model(x_natural)), tf.nn.softmax(model(x_adv)))
+        grad = tape.gradient(kl_loss, x_adv)
+        x_adv = x_adv + alpha * tf.sign(grad)
+        x_adv = tf.clip_by_value(x_adv, x_natural - eps, x_natural + eps)
+        x_adv = tf.clip_by_value(x_adv, 0.0, 1.0)
+
+    loss_nat = tf.keras.losses.categorical_crossentropy(y, model(x_natural))
+    loss_rob = tf.keras.losses.KLDivergence()(tf.nn.softmax(model(x_natural)), tf.nn.softmax(model(x_adv)))
+    return tf.reduce_mean(loss_nat + beta * loss_rob)
+
+def train_trades(model, train_dataset, epochs=5):
+    optimizer = tf.keras.optimizers.Adam(learning_rate=0.001)
+    for epoch in range(epochs):
+        print(f"\nEpoch {epoch + 1}/{epochs} [TRADES]")
+        for images, labels in train_dataset:
+            with tf.GradientTape() as tape:
+                loss = trades_loss(model, images, labels)
+            grads = tape.gradient(loss, model.trainable_variables)
+            optimizer.apply_gradients(zip(grads, model.trainable_variables))
+    return model
+
+def train_pgd(model, train_dataset, eps=0.3, alpha=0.01, steps=40, epochs=5):
+    optimizer = tf.keras.optimizers.Adam(learning_rate=0.001)
+    loss_fn = tf.keras.losses.CategoricalCrossentropy()
+    for epoch in range(epochs):
+        print(f"\nEpoch {epoch + 1}/{epochs} [PGD]")
+        for x, y in train_dataset:
+            x_adv = pgd_attack(model, x, y, eps=eps, alpha=alpha, steps=steps)
+            with tf.GradientTape() as tape:
+                logits = model(x_adv, training=True)
+                loss = loss_fn(y, logits)
+            grads = tape.gradient(loss, model.trainable_variables)
+            optimizer.apply_gradients(zip(grads, model.trainable_variables))
+    return model
+
 def train_model_and_save(args):
-    # Create folder name based on dataset and model type
     folder_name = f"{args.data}_{args.model_type}"
-
     save_dir = os.path.join(args.save_dir, folder_name)
-
-    # Ensure the save directory exists
     os.makedirs(save_dir, exist_ok=True)
     model_path = os.path.join(save_dir, f'{folder_name}.keras')
 
-    # Load or train the model based on the data and model type
+    # Select model and load data
     if args.data == 'MNIST':
-        if args.model_type == 'normal':
-            model = load_MNIST_model(model_path)
-            train_ds, test_ds = load_data(dataset_type="MNIST", use_augmentation=False)
-        elif args.model_type == 'complex':
-            model = load_complex_MNIST_model(model_path)
-            train_ds, test_ds = load_data(dataset_type="MNIST", use_augmentation=False)
-        elif args.model_type == 'complex_augmented':
+        if args.model_type in ['normal', 'complex']:
+            model = load_complex_MNIST_model(model_path) if args.model_type == 'complex' else load_MNIST_model(model_path)
+            train_ds, test_ds = load_data(
+                dataset_type="MNIST",
+                use_augmentation=False,
+                adversarial=args.adversarial
+            )
+        elif args.model_type in ['complex_augmented', 'complex_adversarial']:
             model = load_complex_MNIST_model(model_path)
-            train_ds, test_ds = load_data(dataset_type="MNIST", use_augmentation=True)
+            train_ds, test_ds = load_data(
+                dataset_type="MNIST",
+                use_augmentation=True,
+                adversarial=args.adversarial
+            )
+
+    elif args.data == 'CIFAR10':
+        if args.model_type == 'normal':
+            model = load_CIFAR10_model(model_path)
+        else:
+            model = load_complex_CIFAR10_model(model_path)
+        train_ds, test_ds = load_data(
+            dataset_type="CIFAR10",
+            use_augmentation=args.model_type != 'normal',
+            adversarial=args.adversarial
+        )
+
     elif args.data == 'MNIST_Audio':
         if args.model_type == 'normal':
             model = load_AudioMNIST_model(model_path)
-            train_ds, test_ds = load_data(dataset_type="MNIST_Audio", use_augmentation=False)
-        elif args.model_type == 'complex':
-            model = load_complex_AudioMNIST_model(model_path)
-            train_ds, test_ds = load_data(dataset_type="MNIST_Audio", use_augmentation=False)
-        elif args.model_type == 'complex_augmented':
+        else:
             model = load_complex_AudioMNIST_model(model_path)
-            train_ds, test_ds = load_data(dataset_type="MNIST_Audio", use_augmentation=True)
+        train_ds, test_ds = load_data(
+            dataset_type="MNIST_Audio",
+            use_augmentation=args.model_type != 'normal',
+            adversarial=args.adversarial
+        )
 
-    # Check if model weights exist, if not, train and save
+    else:
+        raise ValueError(f"Unsupported dataset: {args.data}")
+
+    # Train or load
     if not os.path.exists(model_path):
         print("Training the model...")
-        model = train_model(model, train_ds, epochs=args.epochs)
+        if args.adversarial == "trades":
+            model = train_trades(model, train_ds, epochs=args.epochs)
+        elif args.adversarial == "pgd":
+            model = train_pgd(model, train_ds, epochs=args.epochs)
+        else:
+            model = train_model(model, train_ds, epochs=args.epochs)
+
         model.save(model_path)
         print(f"Model saved to {model_path}")
     else:
         print(f"Model found. Loading weights from {model_path}")
         model.load_weights(model_path)
 
-    # Evaluate the model
     print("Evaluating the model...")
     evaluate_model(model, test_ds)
 
diff --git a/manuscripts/Posion25/utils.py b/manuscripts/Posion25/utils.py
new file mode 100644
index 0000000..e69de29