Update analysis scripts

Author: brenjohn

scripts/data_processing/dataset_preperation.py

@@ -14,25 +14,30 @@
 import numpy as np
 
 from swift_tools.data import read_snapshots
-from dmsr.field_operations.resize import cut_field
+from swift_tools.fields import cut_field
 
 
 #%%
 data_directory = '../../data/dmsr_runs/'
 
-LR_snapshots = np.sort(glob.glob(data_directory + '*/064/snap_0002.hdf5'))[:16]
-HR_snapshots = np.sort(glob.glob(data_directory + '*/128/snap_0002.hdf5'))[:16]
+# LR_snapshots = np.sort(glob.glob(data_directory + '*/064/snap_0002.hdf5'))[:16]
+# HR_snapshots = np.sort(glob.glob(data_directory + '*/256/snap_0002.hdf5'))[:16]
 
-# LR_snapshots = np.sort(glob.glob(data_directory + 'run17/064/snap_0002.hdf5'))
-# HR_snapshots = np.sort(glob.glob(data_directory + 'run17/128/snap_0002.hdf5'))
+LR_snapshots = np.sort(glob.glob(data_directory + 'run17/064/snap_0002.hdf5'))
+HR_snapshots = np.sort(glob.glob(data_directory + 'run17/256/snap_0002.hdf5'))
 
 
 #%%
 LR_disp, LR_vel, box_size, LR_grid_size, LR_mass = read_snapshots(LR_snapshots)
 HR_disp, HR_vel, box_size, HR_grid_size, HR_mass = read_snapshots(HR_snapshots)
 
 LR_fields = np.concatenate((LR_disp, LR_vel), axis=1)
+del LR_disp
+del LR_vel
+
 HR_fields = np.concatenate((HR_disp, HR_vel), axis=1)
+del HR_disp
+del HR_vel
 
 # # Normalise values so that box size is 1
 # LR_fields /= box_size
@@ -43,23 +48,23 @@
 #%%
 padding = 2
 LR_patch_size = 16
-HR_patch_size = 32
+HR_patch_size = 64
 
 LR_fields = cut_field(LR_fields, LR_patch_size, LR_patch_size, pad=padding)
 HR_fields = cut_field(HR_fields, HR_patch_size, HR_patch_size)
 
 
 #%%
-LR_file = '../../data/dmsr_training/LR_fields.npy'
-# LR_file = '../../data/dmsr_validation/LR_fields.npy'
+# LR_file = '../../data/dmsr_training/LR_fields.npy'
+LR_file = '../../data/dmsr_validation/LR_fields.npy'
 np.save(LR_file, LR_fields)
 
-HR_file = '../../data/dmsr_training/HR_fields.npy'
-# HR_file = '../../data/dmsr_validation/HR_fields.npy'
+# HR_file = '../../data/dmsr_training/HR_fields.npy'
+HR_file = '../../data/dmsr_validation/HR_fields.npy'
 np.save(HR_file, HR_fields)
 
-meta_file = '../../data/dmsr_training/metadata.npy'
-# meta_file = '../../data/dmsr_validation/metadata.npy'
+# meta_file = '../../data/dmsr_training/metadata.npy'
+meta_file = '../../data/dmsr_validation/metadata.npy'
 LR_size = LR_patch_size + 2 * padding
 HR_size = HR_patch_size
 np.save(meta_file, [

scripts/training/plot_loss.py

@@ -9,7 +9,7 @@
 import numpy as np
 import matplotlib.pyplot as plt
 
-losses = np.load('./losses.npz')
+losses = np.load('./velocity_run/losses.npz')
 
 critic_loss = losses['critic_loss']
 critic_batches = losses['critic_batches']
@@ -54,5 +54,6 @@
     generator_batches[window_size//2-1:-window_size//2], moving_average, 
     # color='black'
 )
+# plt.savefig('loss_curve.png', dpi=210)
 plt.show()
 plt.close()

scripts/training/plot_power_spectrum.py

@@ -17,6 +17,7 @@
 import numpy as np
 import matplotlib.pyplot as plt
 
+from swift_tools.data import load_normalisation_parameters
 from dmsr.analysis import displacement_power_spectrum
 
 
@@ -57,31 +58,35 @@ def plot_spectra(
 
 
 #%%
-plots_dir = 'plots/training_spectra/'
-data_dir = './data/samples/'
-sr_samples = glob.glob(data_dir + 'sr_sample_*.npy')
+data_dir = './velocity_run/'
+plots_dir = data_dir + 'plots/training_spectra/'
+samples_dir = data_dir + 'samples/'
+sr_samples = glob.glob(samples_dir + 'sr_sample_*.npy')
 sr_samples = np.sort(sr_samples)
 
 existing_plots = glob.glob(plots_dir + 'power_sprectrum_epoch_*.png')
 existing_plots = [re.split(r'[._]+', plot)[-2] for plot in existing_plots]
 sr_samples = [sample for sample in sr_samples 
               if not re.split(r'[._]+', sample)[-2] in existing_plots]
 
-lr_sample = np.load(data_dir + 'lr_sample.npy')
-lr_sample = torch.from_numpy(lr_sample)
-hr_sample = np.load(data_dir + 'hr_sample.npy')
-hr_sample = torch.from_numpy(hr_sample)
+scale_path = data_dir + 'normalisation.npy'
+lr_std, hr_std, _, _ = load_normalisation_parameters(scale_path)
+
+lr_sample = np.load(samples_dir + 'lr_sample.npy')
+lr_sample = torch.from_numpy(lr_sample)[:, :3, ...] * lr_std
+hr_sample = np.load(samples_dir + 'hr_sample.npy')
+hr_sample = torch.from_numpy(hr_sample)[:, :3, ...] * hr_std
 
 # TODO: read this from metadata
 box_size = 35.56187768431281
 
 for sr_sample in sr_samples:
     epoch = int(re.split(r'[._]+', sr_sample)[-2])
     sr_sample = np.load(sr_sample)
-    sr_sample = torch.from_numpy(sr_sample)
+    sr_sample = torch.from_numpy(sr_sample)[:, :3, ...] * hr_std
 
     plot_spectra(
         lr_sample, sr_sample, hr_sample,
-        64, 1, 20*box_size/14, box_size, 20, 56,
+        64, 1, 20*box_size/16, box_size, 20, 64,
         epoch, plots_dir
     )

scripts/training/plot_samples.py

@@ -17,6 +17,7 @@
 import numpy as np
 import matplotlib.pyplot as plt
 
+from swift_tools.data import load_normalisation_parameters
 from dmsr.field_operations.conversion import displacements_to_positions
 
 def plot_samples(
@@ -77,25 +78,29 @@ def get_xys(positions):
 
 
 #%%
-plots_dir = 'plots/training_samples/'
-data_dir = './data/samples/'
-sr_samples = glob.glob(data_dir + 'sr_sample_*.npy')
+data_dir = './velocity_run/'
+plots_dir = data_dir + 'plots/training_samples/'
+samples_dir = data_dir + 'samples/'
+sr_samples = glob.glob(samples_dir + 'sr_sample_*.npy')
 sr_samples = np.sort(sr_samples)
 
 existing_plots = glob.glob(plots_dir + 'particle_plot_epoch_*.png')
 existing_plots = [re.split(r'[._]+', plot)[-2] for plot in existing_plots]
 sr_samples = [sample for sample in sr_samples 
               if not re.split(r'[._]+', sample)[-2] in existing_plots]
 
-lr_sample = np.load(data_dir + 'lr_sample.npy')
-hr_sample = np.load(data_dir + 'hr_sample.npy')
+scale_path = data_dir + 'normalisation.npy'    
+lr_std, hr_std, _, _ = load_normalisation_parameters(scale_path)
+
+lr_sample = np.load(samples_dir + 'lr_sample.npy')[:, :3, ...] * lr_std
+hr_sample = np.load(samples_dir + 'hr_sample.npy')[:, :3, ...] * hr_std
 
 # TODO: read this from metadata
 box_size = 35.56187768431281
 
 for sr_sample in sr_samples:
     epoch = int(re.split(r'[._]+', sr_sample)[-2])
-    sr_sample = np.load(sr_sample)
+    sr_sample = np.load(sr_sample)[:, :3, ...] * hr_std
 
     plot_samples(
         lr_sample, sr_sample, hr_sample,