Merge pull request #2 from brenjohn/velocity

Velocity

Author: brenjohn

dmsr/data_tools/__init__.py

@@ -0,0 +1,2 @@
+from .dataset import *
+from .utils import *

dmsr/data_tools/dataset.py

@@ -0,0 +1,81 @@
+#!/usr/bin/env python3
+# -*- coding: utf-8 -*-
+"""
+Created on Fri Sep 13 13:45:00 2024
+
+@author: brennan
+
+This file defines a DMSR-Dataset class for handling training data for a 
+DMSR-WGAN model.
+"""
+
+import torch
+
+from torch.utils.data import Dataset
+from ..field_operations.augmentation import permute_tensor
+
+
+class DMSRDataset(Dataset):
+    """A Dataset class for holding training data for the DMSR-WGAN.
+    """
+    
+    def __init__(
+            self, 
+            lr_position, 
+            hr_position,
+            lr_velocity = None,
+            hr_velocity = None,
+            augment     = True
+        ):
+        self.lr_position = lr_position
+        self.hr_position = hr_position
+        self.lr_velocity = lr_velocity
+        self.hr_velocity = hr_velocity
+        self.velocities_included = not lr_velocity is None
+        self.augment = augment
+        
+    
+    def __len__(self):
+        return self.lr_position.size(0)
+    
+    
+    def __getitem__(self, idx):
+        lr_data = self.lr_position[idx]
+        hr_data = self.hr_position[idx]
+
+        # Apply augmentation (random flip/permutation) if specified
+        if self.augment:
+            random_perm = torch.randperm(3)
+            lr_data = permute_tensor(lr_data, random_perm)
+            hr_data = permute_tensor(hr_data, random_perm)
+                
+        if self.velocities_included:
+            lr_velocity = self.lr_velocity[idx]
+            hr_velocity = self.hr_velocity[idx]
+            
+            if self.augment:
+                lr_velocity = permute_tensor(lr_velocity, random_perm)
+                hr_velocity = permute_tensor(hr_velocity, random_perm)
+            
+            lr_data = torch.concat((lr_data, lr_velocity))
+            hr_data = torch.concat((hr_data, hr_velocity))
+        
+        return lr_data, hr_data
+    
+    
+    def normalise_dataset(self):
+        """Scales position and velocity data by dividing by their respective
+        standard deviations. The standard deviations are also returned as a
+        dictionary.
+        """
+        params = {}
+        field_names = ["lr_position", "hr_position"]
+        if self.velocities_included:
+            field_names += ["lr_velocity", "hr_velocity"]
+        
+        for field in field_names:
+            standard_deviation = vars(self)[field].std()
+            params[field + "_std"] = standard_deviation.item()
+            vars(self)[field] /= standard_deviation
+        
+        return params

dmsr/data_tools/utils.py

@@ -0,0 +1,57 @@
+#!/usr/bin/env python3
+# -*- coding: utf-8 -*-
+"""
+Created on Thu Feb 27 11:56:32 2025
+
+@author: brennan
+"""
+
+import torch
+import numpy as np
+
+from os.path import exists
+
+
+def load_numpy_dataset(data_directory):
+    """Returns LR and HR data contained in numpy files saved in the given 
+    directory.
+    """
+    LR_data = np.load(data_directory + 'LR_fields.npy')
+    LR_data = torch.from_numpy(LR_data)
+    
+    HR_data = np.load(data_directory + 'HR_fields.npy')
+    HR_data = torch.from_numpy(HR_data)
+    
+    meta_file = data_directory + 'metadata.npy'
+    meta_data = np.load(meta_file)
+    box_size, HR_patch_size, LR_size, HR_size, LR_mass, HR_mass = meta_data
+    
+    return LR_data, HR_data, HR_patch_size, LR_size, HR_size
+
+
+def load_normalisation_parameters(param_file):
+    """Reads the standard deviations from the given .npy file used to noramlise
+    dmsr training data.
+    """
+    lr_pos_std = hr_pos_std = lr_vel_std = hr_vel_std = 1
+    
+    if exists(param_file):
+        scale_params = np.load(param_file, allow_pickle=True).item()
+        scale_params = {k : v.item() for k, v in scale_params.items()}
+        lr_pos_std = scale_params.get('lr_position_std', 1)
+        hr_pos_std = scale_params.get('hr_position_std', 1)
+        lr_vel_std = scale_params.get('lr_velocity_std', 1)
+        hr_vel_std = scale_params.get('hr_velocity_std', 1)
+    
+    return lr_pos_std, hr_pos_std, lr_vel_std, hr_vel_std
+
+
+def generate_mock_data(lr_grid_size, hr_grid_size, channels, samples):
+    """Create a mock training data set for testing.
+    """
+    box_size = 1
+    shape = (samples, channels, lr_grid_size, lr_grid_size, lr_grid_size)
+    LR_data = torch.rand(*shape)
+    shape = (samples, channels, hr_grid_size, hr_grid_size, hr_grid_size)
+    HR_data = torch.rand(*shape)
+    return LR_data, HR_data, box_size, lr_grid_size, hr_grid_size

dmsr/dmsr_gan/__init__.py

@@ -5,11 +5,9 @@
 
 @author: brennan
 
-This file defines functions for resizing tensors using various methods.
+This file defines functions for resizing tensors.
 """
 
-import numpy as np
-
 
 def crop(field, crop_size):
     """Crops the spatial dimensions of the given tensor by size crop_size.
@@ -20,65 +18,33 @@ def crop(field, crop_size):
     return field[ind]
 
 
-def cut_field(fields, cut_size, stride=0, pad=0):
-    """Cuts the given field tensor into blocks of size `cut_size`.
-    
-    Arguments:
-        - fields   : A numpy tensor of shape (batch_size, channels, N, N, N)
-                     where N is the grid size of the fields.
-        - cut_size : The base size of the blocks to cut the given fields into.
-        
-        - stride   : The number of cells to move in each direction before 
-                     extracting the next block.
-        - pad      : The number of cells to pad the base blocks on each side.
-        
-    Returns:
-        A numpy tensor containing the blocks/subfields cut from the given
-        fields tensor. The shape of the returned tensor is:
-                 (number_of_cuts * batch_size, channels, n, n, n),
-        where number_of_cuts is the number of subfields extracted from each 
-        field and n is the grid size of each subfield (ie cut_size + 2 * pad).
+def pixel_unshuffle(tensor, scale):
     """
-    grid_size = fields.shape[-1]
-    if not stride:
-        stride = cut_size
+    Reshapes the given a tensor of shape (B, C, D, H, W) to shape 
+    (B, C * scale**3, D // scale, H // scale, W // scale).
     
-    cuts = []
-    for i in range(0, grid_size, stride):
-        for j in range(0, grid_size, stride):
-            for k in range(0, grid_size, stride):
-                
-                slice_x = [n % grid_size for n in range(i-pad, i+cut_size+pad)]
-                slice_y = [n % grid_size for n in range(j-pad, j+cut_size+pad)]
-                slice_z = [n % grid_size for n in range(k-pad, k+cut_size+pad)]
-                
-                patch = np.take(fields, slice_x, axis=2)
-                patch = np.take(patch, slice_y, axis=3)
-                patch = np.take(patch, slice_z, axis=4)
-                
-                cuts.append(patch)
-    
-    return np.concatenate(cuts)
-
-
-def stitch_fields(patches, patches_per_dim):
-    """Combines or stitches the given collection of patches into a single
-    tensor.
-    
-    This function can be thought of as performing the reverse operation
-    performed by `cut_field`.
+    The reshaping procedure uses the pixel shuffle method of Shi et al 2016 -
+    "Real-Time Single Image and Video Super-Resolution Using an Efficient 
+    Sub-Pixel Convolutional Neural Network"
     """
+    # Ensure tensor has the right shape
+    batch_size, channels, depth, height, width = tensor.shape
+
+    new_channels = channels * scale**3
+    new_depth    = depth  // scale
+    new_height   = height // scale
+    new_width    = width  // scale
+
+    # Reshape and permute to rearrange data
+    tensor = tensor.contiguous().view(
+        batch_size, channels,
+        new_depth, scale,
+        new_height, scale,
+        new_width, scale
+    )
+    tensor = tensor.permute(0, 1, 3, 5, 7, 2, 4, 6)
+    tensor = tensor.contiguous().view(
+        batch_size, new_channels, new_depth, new_height, new_width
+    )
 
-    patch_size = patches[0].shape[-1]
-    field_size = patch_size * patches_per_dim
-    field = np.zeros((3, field_size, field_size, field_size))
-    
-    for n, patch in enumerate(patches):
-        i = n // patches_per_dim**2
-        j = (n % patches_per_dim**2) // patches_per_dim
-        k = n % patches_per_dim
-        
-        N = patch_size
-        field[:, i*N:(i+1)*N, j*N:(j+1)*N, k*N:(k+1)*N] = patch
-        
-    return field
+    return tensor

dmsr/dmsr_gan/dmsr_dataset.py

@@ -0,0 +1,2 @@
+from .monitor import *
+from .manager import *

dmsr/field_operations/resize.py

@@ -0,0 +1,95 @@
+#!/usr/bin/env python3
+# -*- coding: utf-8 -*-
+"""
+Created on Thu Feb 27 12:25:54 2025
+
+@author: brennan
+"""
+
+import time
+
+
+class MonitorManager():
+    """A class to manage monitor objects.
+    
+    The Monitor Manager class stores and calls Monitor objects during DMSR-WGAN
+    training at appropriate times.
+    
+    Monitor objects are stored in a monitors dictionary. During DMSR training,
+    at the end of a batch update the `post_batch_processing` method of each 
+    monitor object is called. Similarly, at the end of each epoch, the
+    `post_epoch_processing` method of each monitor is called by the monitor
+    manager.
+    
+    Any messages returned by the `post_batch_processing` calls are passed to a
+    batch report method which prints them along with some information regarding
+    batch/epoch number and timings. At the end of each epoch, the monitor
+    manager also prints some timing information regarding the epoch and epoch
+    post processing.
+    """
+    
+    def __init__(self, report_rate, device):
+        self.device = device
+        self.report_rate = report_rate
+        
+    
+    def set_monitors(self, monitors):
+        self.monitors = monitors
+        
+    
+    def init_monitoring(self, num_epochs, num_batches):
+        """Initializes values for variables used for timing batches and epochs.
+        """
+        self.num_epochs = num_epochs
+        self.num_batches = num_batches
+        self.batch_start_time = time.time()
+        self.epoch_start_time = time.time()
+    
+    
+    def end_of_epoch(self, epoch):
+        """Calls the `post_epoch_processing` method of each monitor.
+        """
+        epoch_time = time.time() - self.epoch_start_time
+        print(f"[Epoch {epoch} took: {epoch_time:.4f} sec]")
+        post_processing_start_time = time.time()
+        
+        for monitor in self.monitors.values():
+            monitor.post_epoch_processing(epoch)
+        
+        self.epoch_start_time = time.time()
+        self.batch_start_time = time.time()
+        post_processing_time = time.time() - post_processing_start_time
+        print(
+            f"[Epoch post-processing took: {post_processing_time:.4f} sec]",
+            flush=True
+        )
+    
+        
+    def end_of_batch(self, epoch, batch, batch_counter, losses):
+        """Calls the `post_batch_processing` method of each monitor.
+        """
+        monitor_report = ''
+        
+        for monitor in self.monitors.values():
+            monitor_report += monitor.post_batch_processing(
+                epoch, batch, batch_counter, losses
+            )
+        
+        self.batch_report(epoch, batch, monitor_report)
+    
+    
+    def batch_report(self, epoch, batch, monitor_report):
+        """Report some satistics for the last few batch updates.
+        """
+        if (batch > 0 and batch % self.report_rate == 0):
+            time_curr = time.time()
+            time_prev = self.batch_start_time
+            average_batch_time = (time_curr - time_prev) / self.report_rate
+            
+            report  = f"[Epoch {epoch:04}/{self.num_epochs}]"
+            report += f"[Batch {batch:03}/{self.num_batches}]"
+            report += f"[time per batch: {average_batch_time*1000:.4f} ms]"
+            report += monitor_report
+            
+            print(report)
+            self.batch_start_time = time.time()