Astro_Dataton.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8286ee24-b703-48d8-ae25-db33e22a42b4",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import numpy as np\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.optim as optim\n",
    "\n",
    "from torch.utils.data import Dataset, DataLoader\n",
    "from transformers import ViTMAEModel, ViTMAEConfig"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "49280d41-3ef6-43db-be94-d6980b281ea5",
   "metadata": {},
   "outputs": [],
   "source": [
    "class GalaxyDataset(Dataset):\n",
    "    def __init__(self, input_dir, target_dir = None, transform = None):\n",
    "        self.input_files = sorted([os.path.join(input_dir, f) for f in os.listdir(input_dir) if f.endswith('.npy')])\n",
    "        self.target_files = sorted([os.path.join(target_dir, f) for f in os.listdir(target_dir) if f.endswith('.npy')]) if target_dir else None\n",
    "        self.transform = transform\n",
    "\n",
    "    def __len__(self):\n",
    "        return len(self.input_files)\n",
    "\n",
    "    def __getitem__(self, idx):\n",
    "        input_data = np.load(self.input_files[idx])\n",
    "        input_tensor = torch.Tensor(input_data)\n",
    "\n",
    "        # If target_dir is provided, load the target, otherwise return only the input\n",
    "        if self.target_files:\n",
    "            target_data = np.load(self.target_files[idx])\n",
    "            target_tensor = torch.Tensor(target_data)\n",
    "            return input_tensor, target_tensor  # Return input and target\n",
    "        else:\n",
    "            return input_tensor  # Only return input when no target_dir is provided"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4e8336de-0c60-48c6-ba3f-633a8816d782",
   "metadata": {},
   "outputs": [],
   "source": [
    "def preprocess_target_patches_batch(batch_tensor, patch_size = 16):\n",
    "    \"\"\"\n",
    "    Converts a batch of 2D target tensors into patches suitable for ViT.\n",
    "    Args:\n",
    "        batch_tensor (Tensor): Input tensor of shape (batch_size, H, W).\n",
    "        patch_size (int): Size of each patch.\n",
    "    Returns:\n",
    "        patches (Tensor): Tensor of shape (batch_size, num_patches, patch_size, patch_size).\n",
    "    \"\"\"\n",
    "    batch_size, h, w = batch_tensor.shape\n",
    "    assert h == w == 128, \"Target image must be 128x128\"\n",
    "\n",
    "    # Unfold (slice) the height and width dimensions to create patches\n",
    "    patches = batch_tensor.unfold(1, patch_size, patch_size).unfold(2, patch_size, patch_size)  # (batch_size, num_patches_h, num_patches_w, patch_size, patch_size)\n",
    "\n",
    "    # Rearrange the patches into the required shape\n",
    "    patches = patches.permute(0, 1, 2, 3, 4).reshape(batch_size, -1, patch_size, patch_size)  # Shape: (batch_size, num_patches, patch_size, patch_size)\n",
    "\n",
    "    return patches"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "73b35272-183b-4690-bdf6-e1dd0d8f8c06",
   "metadata": {},
   "outputs": [],
   "source": [
    "class WeakLensingViT(nn.Module):\n",
    "    def __init__(self, patch_size = 16):\n",
    "        super(WeakLensingViT, self).__init__()\n",
    "\n",
    "        # Load a pretrained ViT model\n",
    "        config = ViTMAEConfig(image_size = 128, patch_size = patch_size, num_channels = 3, hidden_size = 768)\n",
    "        self.vit = ViTMAEModel(config = config)\n",
    "\n",
    "        # Modify the output layer to output 256 values (16x16) for each patch\n",
    "        self.reconstruction_head = nn.Linear(in_features = 768, out_features = patch_size*patch_size)\n",
    "\n",
    "    def forward(self, patches):\n",
    "        # Pass the patches through the ViT model\n",
    "        vit_output = self.vit(patches).last_hidden_state  # Shape: (num_patches + 1, 768)\n",
    "\n",
    "        # Reconstruct the 2D map (16x16 = 256 values per patch)\n",
    "        map_output = self.reconstruction_head(vit_output)  # Output shape: (num_patches, 256)\n",
    "        map_output = map_output.view(-1, 64, 16, 16)  # Reshape to (num_patches, 16, 16) for each patch\n",
    "\n",
    "        return map_output"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "885c9f36-c2a8-4331-94c0-3b3816d3edee",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Instantiate the dataset and data loader\n",
    "\n",
    "input_dir = '1/EPSILON'\n",
    "target_dir = '1/KAPPA'\n",
    "\n",
    "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
    "\n",
    "# Use the dataset class\n",
    "train_dataset = GalaxyDataset(input_dir = input_dir, target_dir = target_dir)\n",
    "\n",
    "# Use more workers for faster parallel data loading\n",
    "train_loader = DataLoader(train_dataset, batch_size = 128, shuffle = True, num_workers = 8, pin_memory = True)\n",
    "\n",
    "# Initialize the model, optimizer\n",
    "model = WeakLensingViT(patch_size = 16)\n",
    "model = model.to(device)\n",
    "\n",
    "optimizer = optim.Adam(model.parameters(), lr = 1e-4)\n",
    "loss = nn.HubberLoss()\n",
    "\n",
    "# Training loop\n",
    "num_epochs = 10\n",
    "\n",
    "for epoch in range(num_epochs):\n",
    "    model.train()\n",
    "    running_loss = 0.0\n",
    "\n",
    "    for inputs, targets in train_loader:\n",
    "        inputs = inputs.to(device)\n",
    "        targets = targets.to(device)\n",
    "\n",
    "        targets_patches = preprocess_target_patches_batch(targets.squeeze())  # Shape: (num_patches, 16, 16)\n",
    "\n",
    "        # Zero gradients\n",
    "        optimizer.zero_grad()\n",
    "\n",
    "        outputs = model(inputs)\n",
    "\n",
    "        loss = loss(outputs, targets_patches)\n",
    "\n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "\n",
    "        running_loss += loss.item()\n",
    "\n",
    "    print(f\"Epoch [{epoch + 1}/{num_epochs}], Loss: {running_loss/len(train_loader)}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "431d05a9-c32b-4f89-96e7-af472eaaf68a",
   "metadata": {},
   "outputs": [],
   "source": [
    "torch.save(model.state_dict(), 'model.pth')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "230e44bc-e111-4c46-bea5-78c9eab992ad",
   "metadata": {},
   "outputs": [],
   "source": [
    "def assemble_patches(patches, original_size = (128, 128), patch_size = 16):\n",
    "    \"\"\"\n",
    "    Reassemble patches into the original full-sized image.\n",
    "    Args:\n",
    "        patches (Tensor): Patch tensor of shape (num_patches, patch_size, patch_size).\n",
    "        original_size (tuple): The original size of the image (H, W).\n",
    "        patch_size (int): The size of each patch.\n",
    "    Returns:\n",
    "        full_map (Tensor): The reassembled image tensor.\n",
    "    \"\"\"\n",
    "    num_patches = patches.shape[1]\n",
    "    h, w = original_size\n",
    "    assert num_patches == (h//patch_size)*(w//patch_size), \"Mismatch between patches and original size\"\n",
    "\n",
    "    full_map = patches.view(h//patch_size, w//patch_size, patch_size, patch_size)\n",
    "    full_map = full_map.permute(0, 2, 1, 3).reshape(h, w)  # Reassemble into (H, W)\n",
    "\n",
    "    return full_map"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8029e89e-9014-4e6c-9507-898904b728c4",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load the test dataset\n",
    "test_input_dir = 'EPSILON'\n",
    "output_dir = 'tierraplana'\n",
    "\n",
    "if not os.path.exists(output_dir):\n",
    "    os.makedirs(output_dir)\n",
    "\n",
    "test_dataset = GalaxyDataset(input_dir = test_input_dir)\n",
    "test_loader = DataLoader(test_dataset, batch_size = 1, shuffle = False)\n",
    "\n",
    "# Evaluate the model\n",
    "model.eval()\n",
    "with torch.no_grad():\n",
    "    for i, inputs in enumerate(test_loader):  # Only retrieve inputs since no target\n",
    "        inputs = inputs.to(device)\n",
    "\n",
    "        # Forward pass\n",
    "        outputs = model(inputs)\n",
    "\n",
    "        # Reassemble the patches back into a full image\n",
    "        full_output = assemble_patches(outputs.cpu(), original_size = (128, 128), patch_size = 16)\n",
    "\n",
    "        # Convert the output to float16 before saving\n",
    "        full_output_float16 = full_output.numpy().astype(np.float16)\n",
    "\n",
    "        # Get the original filename from the test folder\n",
    "        input_filename = test_dataset.input_files[i]\n",
    "        base_filename = os.path.basename(input_filename)  # Extract the filename from the full path\n",
    "\n",
    "        # Save the output as .npy with the same filename in the output directory in float16 format\n",
    "        output_filepath = os.path.join(output_dir, base_filename)\n",
    "        np.save(output_filepath, full_output_float16)  # Save the full output as .npy in float16"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "base",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.18"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}