Merge pull request #61 from Kevin2/511_pixel

Add functionalities for 511 pixel pictures

Author: FeGeyer

Diff for: dl_framework/architectures.py

@@ -692,7 +692,7 @@ def forward(self, x):
 
 
 class filter_deep_amp(nn.Module):
-    def __init__(self):
+    def __init__(self, img_size):
         super().__init__()
         self.conv1_amp = nn.Sequential(
             *conv_amp(1, 4, (23, 23), 1, 11, 1)
@@ -704,7 +704,7 @@ def __init__(self):
              *conv_amp(8, 12, (17, 17), 1, 8, 1)
         )
         self.conv_con1_amp = nn.Sequential(
-            LocallyConnected2d(12, 1, 63, 1, stride=1, bias=False),
+            LocallyConnected2d(12, 1, img_size, 1, stride=1, bias=False),
             nn.BatchNorm2d(1),
             nn.ReLU(),
         )
@@ -722,7 +722,7 @@ def __init__(self):
              *conv_amp(12, 16, (3, 3), 1, 1, 1)
         )
         self.conv_con2_amp = nn.Sequential(
-            LocallyConnected2d(16, 1, 63, 1, stride=1, bias=False),
+            LocallyConnected2d(16, 1, img_size, 1, stride=1, bias=False),
             nn.BatchNorm2d(1),
             nn.ReLU(),
         )
@@ -737,7 +737,7 @@ def __init__(self):
              *conv_amp(8, 12, (3, 3), 1, 2, 2)
         )
         self.conv_con3_amp = nn.Sequential(
-            LocallyConnected2d(12, 1, 63, 1, stride=1, bias=False),
+            LocallyConnected2d(12, 1, img_size, 1, stride=1, bias=False),
             nn.BatchNorm2d(1),
             nn.ReLU(),
         )
@@ -776,13 +776,15 @@ def forward(self, x):
 
         amp = self.conv_con3_amp(amp)
 
-        amp = amp + inp[:, 0].unsqueeze(1)
-        x0 = self.symmetry_real(amp).reshape(-1, 1, 63, 63)
+        inp_amp = inp[:, 0].unsqueeze(1)
+        x0 = self.symmetry_real(amp).reshape(-1, 1, amp.shape[2], amp.shape[2])
+        x0[inp_amp != 0] = inp_amp[inp_amp != 0]
+
         return x0
 
 
 class filter_deep_phase(nn.Module):
-    def __init__(self):
+    def __init__(self, img_size):
         super().__init__()
         self.conv1_phase = nn.Sequential(
             *conv_phase(1, 4, (23, 23), 1, 11, 1, add=-2.1415)
@@ -794,7 +796,7 @@ def __init__(self):
              *conv_phase(8, 12, (17, 17), 1, 8, 1, add=-2.1415)
         )
         self.conv_con1_phase = nn.Sequential(
-            LocallyConnected2d(12, 1, 63, 1, stride=1, bias=False),
+            LocallyConnected2d(12, 1, img_size, 1, stride=1, bias=False),
             nn.BatchNorm2d(1),
             GeneralELU(-2.1415),
         )
@@ -812,7 +814,7 @@ def __init__(self):
              *conv_phase(12, 16, (3, 3), 1, 1, 1, add=-2.1415)
         )
         self.conv_con2_phase = nn.Sequential(
-            LocallyConnected2d(16, 1, 63, 1, stride=1, bias=False),
+            LocallyConnected2d(16, 1, img_size, 1, stride=1, bias=False),
             nn.BatchNorm2d(1),
             GeneralELU(-2.1415),
         )
@@ -827,7 +829,7 @@ def __init__(self):
              *conv_phase(8, 12, (3, 3), 1, 2, 2, add=-2.1415)
         )
         self.conv_con3_phase = nn.Sequential(
-            LocallyConnected2d(12, 1, 63, 1, stride=1, bias=False),
+            LocallyConnected2d(12, 1, img_size, 1, stride=1, bias=False),
             nn.BatchNorm2d(1),
             GeneralELU(-2.1415),
         )
@@ -866,6 +868,8 @@ def forward(self, x):
 
         phase = self.conv_con3_phase(phase)
 
-        phase = phase + inp[:, 1].unsqueeze(1)
-        x1 = self.symmetry_imag(phase).reshape(-1, 1, 63, 63)
+        inp_phase = inp[:, 1].unsqueeze(1)
+
+        x1 = self.symmetry_imag(phase).reshape(-1, 1, phase.shape[2], phase.shape[2])
+        x1[inp_phase != 0] = inp_phase[inp_phase != 0]
         return x1

Diff for: dl_framework/data.py

@@ -64,7 +64,11 @@ def __getitem__(self, i):
         return x, y
 
     def open_bundle(self, bundle_path, var):
-        bundle = h5py.File(bundle_path, "r")
+        # distinguish between compressed (npz) or not compressed (h5)
+        if re.search(".npz", str(bundle_path)):
+            bundle = np.load(bundle_path, mmap_mode="r")
+        else:
+            bundle = h5py.File(bundle_path, "r")
         data = bundle[var]
         return data
 
@@ -77,14 +81,21 @@ def open_image(self, var, i):
         bundle = indices // self.num_img
         image = indices - bundle * self.num_img
         bundle_unique = torch.unique(bundle)
-        bundle_paths = [
-            h5py.File(self.bundles[bundle], "r") for bundle in bundle_unique
-        ]
+        # distinguish between compressed (npz) or not compressed (h5)
+        if re.search(".npz", str(self.bundles[bundle[0]])):
+            bundle_paths = [
+                np.load(self.bundles[bundle], mmap_mode='r') for bundle in bundle_unique
+            ]
+        else:
+            bundle_paths = [
+                h5py.File(self.bundles[bundle], 'r') for bundle in bundle_unique
+            ]
+        bundle_paths_str = list(map(str, bundle_paths))
         data = torch.tensor(
             [
                 bund[var][img]
-                for bund in bundle_paths
-                for img in image[bundle == bundle_unique[bundle_paths.index(bund)]]
+                for bund, bund_str in zip(bundle_paths, bundle_paths_str)
+                for img in image[bundle == bundle_unique[bundle_paths_str.index(bund_str)]]
             ]
         )
         if var == "x" or self.tar_fourier is True:
@@ -193,6 +204,16 @@ def open_fft_pair(path):
     return bundle_x, bundle_y
 
 
+def open_fft_pair_npz(path):
+    """
+    open fft_pairs for files saved in .npz format
+    """
+    f = np.load(path)
+    bundle_x = np.array(f["x"])
+    bundle_y = np.array(f["y"])
+    return bundle_x, bundle_y
+
+
 def mean_and_std(array):
     return array.mean(), array.std()
 

Diff for: gaussian_sources/Makefile

@@ -26,7 +26,8 @@ ${data}/fft_samp_train0.h5:
 		-amp_phase ${amp_phase} \
 		-noise ${noise} \
 		-preview ${preview} \
-		-specific_mask True \
+		-specific_mask False \
+		-compressed True\
 		-lon -80 \
 		-lat 50 \
 		-steps 50

Diff for: gaussian_sources/calculate_normalization.py

@@ -1,5 +1,10 @@
 import click
-from dl_framework.data import open_fft_pair, get_bundles, mean_and_std
+from dl_framework.data import (
+    open_fft_pair_npz,
+    open_fft_pair,
+    get_bundles,
+    mean_and_std,
+)
 import pandas as pd
 import numpy as np
 import re
@@ -21,7 +26,11 @@ def main(data_path, out_path):
     stds_imag = np.array([])
 
     for path in tqdm(bundle_paths):
-        x, _ = open_fft_pair(path)
+        # distinguish between compressed (.npz) and not compressed (.h5) files
+        if re.search(".npz", str(path)):
+            x, _ = open_fft_pair_npz(path)
+        else:
+            x, _ = open_fft_pair(path)
         x_amp, x_imag = np.double(x[:, 0]), np.double(x[:, 1])
         mean_amp, std_amp = mean_and_std(x_amp)
         mean_imag, std_imag = mean_and_std(x_imag)
@@ -38,8 +47,8 @@ def main(data_path, out_path):
     d = {
         "train_mean_c0": [mean_amp],
         "train_std_c0": [std_amp],
-        'train_mean_c1': [mean_imag],
-        'train_std_c1': [std_imag]
+        "train_mean_c1": [mean_imag],
+        "train_std_c1": [std_imag],
     }
 
     df = pd.DataFrame(data=d)

Diff for: gaussian_sources/create_fft_pairs.py

@@ -1,4 +1,5 @@
 import click
+import os
 import numpy as np
 from tqdm import tqdm
 from dl_framework.data import (
@@ -11,6 +12,7 @@
 from simulations.uv_simulations import sample_freqs
 from simulations.gaussian_simulations import add_noise
 import re
+from numpy import savez_compressed
 
 
 @click.command()
@@ -20,6 +22,7 @@
 @click.option("-amp_phase", type=bool, required=True)
 @click.option("-fourier", type=bool, required=True)
 @click.option("-specific_mask", type=bool)
+@click.option("-compressed", type=bool)
 @click.option("-lon", type=float, required=False)
 @click.option("-lat", type=float, required=False)
 @click.option("-steps", type=float, required=False)
@@ -31,6 +34,7 @@ def main(
     antenna_config_path,
     amp_phase,
     fourier,
+    compressed=False,
     specific_mask=False,
     lon=None,
     lat=None,
@@ -69,6 +73,11 @@ def main(
             if amp_phase:
                 amp, phase = split_amp_phase(bundle_fft)
                 amp = (np.log10(amp + 1e-10) / 10) + 1
+
+                # Test new masking for 511 Pixel pictures
+                if amp.shape[1] == 511:
+                    mask = amp > 0.1
+                    phase[~mask] = 0
                 bundle_fft = np.stack((amp, phase), axis=1)
             else:
                 real, imag = split_real_imag(bundle_fft)
@@ -100,7 +109,11 @@ def main(
                 )
             out = out_path + path.name.split("_")[-1]
             if fourier:
-                save_fft_pair(out, bundle_samp, bundle_fft)
+                if compressed:
+                    savez_compressed(out, x=bundle_samp, y=bundle_fft)
+                    os.remove(path)
+                else:
+                    save_fft_pair(out, bundle_samp, bundle_fft)
             else:
                 save_fft_pair(out, bundle_samp, images)
 

Diff for: gaussian_sources/find_lr.py

@@ -88,7 +88,7 @@ def main(
     img_size = train_ds[0][0][0].shape[1]
     # Define model
     arch_name = arch
-    if arch == "filter_deep":
+    if arch == "filter_deep" or arch == "filter_deep_amp" or arch == "filter_deep_phase":
         arch = getattr(architecture, arch)(img_size)
     else:
         arch = getattr(architecture, arch)()

Diff for: gaussian_sources/inspection.py

@@ -317,6 +317,11 @@ def plot_difference(i, img_pred, img_truth, sensitivity, out_path):
         num_three = 60
         num_two = 50
         num_four = 40
+    elif img_size == 511:
+        # work in progress, these are dummy values for compilating
+        num_three = 60
+        num_two = 50
+        num_four = 40
     num = [num_three, num_two, num_four]
     dr_truth, mode = compute_dr(i, img_truth, sensitivity, num)
     dr_pred = compute_dr_pred(img_pred, mode, num)

Diff for: gaussian_sources/train_cnn.py

@@ -83,7 +83,7 @@ def main(
 
     img_size = train_ds[0][0][0].shape[1]
     # Define model
-    if arch == "filter_deep":
+    if arch == "filter_deep" or arch == "filter_deep_amp" or arch == "filter_deep_phase":
         arch = getattr(architecture, arch)(img_size)
     else:
         arch = getattr(architecture, arch)()