From 242021dac27a9d8f6bcf84a8fea9f0298f2d4c6d Mon Sep 17 00:00:00 2001
From: Felix Geyer <felix.geyer@tu-dortmund.de>
Date: Tue, 17 Jan 2023 11:59:50 +0000
Subject: [PATCH 01/31] Add swp files to gitignore

---
 .gitignore | 1 +
 1 file changed, 1 insertion(+)

diff --git a/.gitignore b/.gitignore
index 1be06ed5..9968eecb 100644
--- a/.gitignore
+++ b/.gitignore
@@ -44,3 +44,4 @@ __pycache__/
 dist/
 radionets/dl_framework/values.yaml
 *.csv
+*.swp

From 0cb36d2adaa2688c8286cd4792ec1deed3d318a9 Mon Sep 17 00:00:00 2001
From: Felix Geyer <felix.geyer@tu-dortmund.de>
Date: Tue, 17 Jan 2023 12:00:06 +0000
Subject: [PATCH 02/31] Change framework for using half images

---
 .../dl_framework/architectures/res_exp.py     |  4 ++--
 .../dl_framework/architectures/unc_archs.py   | 20 +++++++++----------
 radionets/dl_framework/callbacks.py           |  2 +-
 radionets/dl_framework/data.py                |  2 ++
 radionets/dl_framework/model.py               | 14 ++++++++++++-
 5 files changed, 28 insertions(+), 14 deletions(-)

diff --git a/radionets/dl_framework/architectures/res_exp.py b/radionets/dl_framework/architectures/res_exp.py
index 7fb91cc7..4a58e22d 100644
--- a/radionets/dl_framework/architectures/res_exp.py
+++ b/radionets/dl_framework/architectures/res_exp.py
@@ -167,9 +167,9 @@ def forward(self, x):
 
         x = self.final(x)
 
-        x0 = x[:, 0].reshape(-1, 1, s, s)
+        x0 = x[:, 0].reshape(-1, 1, s // 2 + 1, s)
         x0 = self.relu(x0)
-        x1 = self.hardtanh(x[:, 1]).reshape(-1, 1, s, s)
+        x1 = self.hardtanh(x[:, 1]).reshape(-1, 1, s // 2 + 1, s)
 
         return torch.cat([x0, x1], dim=1)
 
diff --git a/radionets/dl_framework/architectures/unc_archs.py b/radionets/dl_framework/architectures/unc_archs.py
index 8c201e20..b066a64e 100644
--- a/radionets/dl_framework/architectures/unc_archs.py
+++ b/radionets/dl_framework/architectures/unc_archs.py
@@ -4,7 +4,7 @@
     GeneralELU,
     LocallyConnected2d,
 )
-from radionets.dl_framework.architectures.res_exp import SRResNet
+from radionets.dl_framework.architectures.res_exp import SRResNet_16
 
 
 class Uncertainty(nn.Module):
@@ -12,28 +12,28 @@ def __init__(self, img_size):
         super().__init__()
 
         self.conv1 = nn.Sequential(
-            nn.Conv2d(4, 8, 9, stride=1, padding=4, groups=2),
-            nn.BatchNorm2d(8),
+            nn.Conv2d(4, 16, 9, stride=1, padding=4, groups=2),
+            nn.BatchNorm2d(16),
             nn.PReLU(),
         )
 
         self.conv2 = nn.Sequential(
-            nn.Conv2d(8, 16, 3, stride=1, padding=1),
-            nn.BatchNorm2d(16),
+            nn.Conv2d(16, 32, 3, stride=1, padding=1),
+            nn.BatchNorm2d(32),
             nn.PReLU(),
         )
 
         self.conv3 = nn.Sequential(
-            nn.Conv2d(16, 32, 9, stride=1, padding=4, groups=2),
-            nn.BatchNorm2d(32),
+            nn.Conv2d(32, 64, 9, stride=1, padding=4, groups=2),
+            nn.BatchNorm2d(64),
             nn.PReLU(),
         )
 
         self.final = nn.Sequential(
             LocallyConnected2d(
-                32,
+                64,
                 2,
-                img_size,
+                [img_size // 2 + 1, img_size],
                 1,
                 stride=1,
                 bias=False,
@@ -55,7 +55,7 @@ def forward(self, x):
 class UncertaintyWrapper(nn.Module):
     def __init__(self, img_size):
         super().__init__()
-        self.pred = SRResNet()
+        self.pred = SRResNet_16()
 
         self.uncertainty = Uncertainty(img_size)
 
diff --git a/radionets/dl_framework/callbacks.py b/radionets/dl_framework/callbacks.py
index 8b5e02dd..bfa2248b 100644
--- a/radionets/dl_framework/callbacks.py
+++ b/radionets/dl_framework/callbacks.py
@@ -183,7 +183,7 @@ class DataAug(Callback):
     def before_batch(self):
         x = self.xb[0].clone()
         y = self.yb[0].clone()
-        randint = np.random.randint(0, 4, x.shape[0])
+        randint = np.random.randint(0, 1, x.shape[0]) * 2
         last_axis = len(x.shape) - 1
         for i in range(x.shape[0]):
             x[i] = torch.rot90(x[i], int(randint[i]), [last_axis - 2, last_axis - 1])
diff --git a/radionets/dl_framework/data.py b/radionets/dl_framework/data.py
index afa11f04..bca94f6a 100644
--- a/radionets/dl_framework/data.py
+++ b/radionets/dl_framework/data.py
@@ -71,6 +71,8 @@ def open_image(self, var, i):
         if data.shape[0] == 1:
             data = data.squeeze(0)
 
+        data = data[:, :65, :]
+
         return data.float()
 
 
diff --git a/radionets/dl_framework/model.py b/radionets/dl_framework/model.py
index db103812..ffaa70aa 100644
--- a/radionets/dl_framework/model.py
+++ b/radionets/dl_framework/model.py
@@ -191,7 +191,7 @@ def __init__(
         self, in_channels, out_channels, output_size, kernel_size, stride, bias=False
     ):
         super(LocallyConnected2d, self).__init__()
-        output_size = _pair(output_size)
+        # output_size = _pair(output_size)
         self.weight = nn.Parameter(
             torch.randn(
                 1,
@@ -304,3 +304,15 @@ def symmetry(x):
     x[:, 0, x.shape[2] // 2 :, :] = upper_half[:, 0]
     x[:, 1, x.shape[2] // 2 :, :] = -upper_half[:, 1]
     return x
+
+
+def sym_new(x):
+    upper_half = x[:, :, :64, :].clone()
+    a = torch.rot90(upper_half, 2, dims=[-2, -1])
+    # amp
+    x[:, 0, 65:, 1:] = a[:, 0, :-1, :-1]
+    x[:, 0, 65:, 0] = a[:, 0, :-1, -1]
+    # phase
+    x[:, 1, 65:, 1:] = -a[:, 1, :-1, :-1]
+    x[:, 1, 65:, 0] = -a[:, 1, :-1, -1]
+    return x

From 2dd5bb8e349cef5a4630bf08daee7a87a6463471 Mon Sep 17 00:00:00 2001
From: Felix Geyer <felix.geyer@tu-dortmund.de>
Date: Tue, 17 Jan 2023 12:57:24 +0000
Subject: [PATCH 03/31] Plotting changes

---
 radionets/evaluation/plotting.py | 47 ++++++++++++++------------------
 radionets/evaluation/utils.py    | 10 +++----
 2 files changed, 25 insertions(+), 32 deletions(-)

diff --git a/radionets/evaluation/plotting.py b/radionets/evaluation/plotting.py
index c4b1a758..3dc3bbea 100644
--- a/radionets/evaluation/plotting.py
+++ b/radionets/evaluation/plotting.py
@@ -6,10 +6,10 @@
 import torch
 import matplotlib as mpl
 from matplotlib.colors import ListedColormap, LogNorm
-from matplotlib.lines import Line2D
-from matplotlib.patches import Arc, Rectangle
+from matplotlib.patches import Rectangle
 from mpl_toolkits.axes_grid1 import make_axes_locatable
 from pytorch_msssim import ms_ssim
+
 from radionets.evaluation.blob_detection import calc_blobs
 from radionets.evaluation.contour import compute_area_ratio
 from radionets.evaluation.dynamic_range import calc_dr, get_boxsize
@@ -291,9 +291,9 @@ def visualize_with_fourier_diff(
     real_pred, imag_pred = img_pred[0], img_pred[1]
     real_truth, imag_truth = img_truth[0], img_truth[1]
 
-    if amp_phase:
-        real_pred = 10 ** (10 * real_pred - 10) - 1e-10
-        real_truth = 10 ** (10 * real_truth - 10) - 1e-10
+    # if amp_phase:
+    #    real_pred = 10 ** (10 * real_pred - 10) - 1e-10
+    #    real_truth = 10 ** (10 * real_truth - 10) - 1e-10
 
     # plotting
     # plt.style.use('./paper_large_3_2.rc')
@@ -302,10 +302,10 @@ def visualize_with_fourier_diff(
     )
 
     if amp_phase:
-        im1 = ax1.imshow(real_pred, cmap="inferno", norm=LogNorm())
+        im1 = ax1.imshow(real_pred, cmap="inferno")
         make_axes_nice(fig, ax1, im1, r"Amplitude Prediction")
 
-        im2 = ax2.imshow(real_truth, cmap="inferno", norm=LogNorm())
+        im2 = ax2.imshow(real_truth, cmap="inferno")
         make_axes_nice(fig, ax2, im2, r"Amplitude Truth")
 
         a = check_vmin_vmax(real_pred - real_truth)
@@ -350,30 +350,30 @@ def visualize_source_reconstruction(
 ):
     m_truth, n_truth, alpha_truth = calc_jet_angle(ifft_truth)
     m_pred, n_pred, alpha_pred = calc_jet_angle(ifft_pred)
-    x_space = torch.arange(0, 63, 1)
+    # x_space = torch.arange(0, 63, 1)
 
     # plt.style.use("./paper_large_3.rc")
     fig, (ax1, ax2, ax3) = plt.subplots(1, 3, figsize=(16, 10), sharey=True)
 
     # Plot prediction
-    ax1.plot(x_space, m_pred * x_space + n_pred, "w--", alpha=0.5)
-    ax1.axvline(32, 0, 1, linestyle="--", color="white", alpha=0.5)
+    # ax1.plot(x_space, m_pred * x_space + n_pred, "w--", alpha=0.5)
+    # ax1.axvline(32, 0, 1, linestyle="--", color="white", alpha=0.5)
 
     # create angle visualization
-    theta1 = min(0, -alpha_pred.numpy()[0])
-    theta2 = max(0, -alpha_pred.numpy()[0])
-    ax1.add_patch(Arc([32, 32], 50, 50, 90, theta1, theta2, color="white"))
+    # theta1 = min(0, -alpha_pred.numpy()[0])
+    # theta2 = max(0, -alpha_pred.numpy()[0])
+    # ax1.add_patch(Arc([32, 32], 50, 50, 90, theta1, theta2, color="white"))
 
     im1 = ax1.imshow(ifft_pred, vmax=ifft_truth.max(), cmap="inferno")
 
     # Plot truth
-    ax2.plot(x_space, m_truth * x_space + n_truth, "w--", alpha=0.5)
-    ax2.axvline(32, 0, 1, linestyle="--", color="white", alpha=0.5)
+    # ax2.plot(x_space, m_truth * x_space + n_truth, "w--", alpha=0.5)
+    # ax2.axvline(32, 0, 1, linestyle="--", color="white", alpha=0.5)
 
     # create angle visualization
-    theta1 = min(0, -alpha_truth.numpy()[0])
-    theta2 = max(0, -alpha_truth.numpy()[0])
-    ax2.add_patch(Arc([32, 32], 50, 50, 90, theta1, theta2, color="white"))
+    # theta1 = min(0, -alpha_truth.numpy()[0])
+    # theta2 = max(0, -alpha_truth.numpy()[0])
+    # ax2.add_patch(Arc([32, 32], 50, 50, 90, theta1, theta2, color="white"))
 
     im2 = ax2.imshow(ifft_truth, cmap="inferno")
 
@@ -417,15 +417,8 @@ def visualize_source_reconstruction(
 
     outpath = str(out_path) + f"/fft_pred_{i}.{plot_format}"
 
-    line = Line2D(
-        [], [], linestyle="-", color="w", label=rf"$\alpha = {alpha_pred[0]:.2f}\,$deg"
-    )
-    line_truth = Line2D(
-        [], [], linestyle="-", color="w", label=rf"$\alpha = {alpha_truth[0]:.2f}\,$deg"
-    )
-
-    ax1.legend(loc="best", handles=[line])
-    ax2.legend(loc="best", handles=[line_truth])
+    # ax1.legend(loc="best", handles=[line])
+    # ax2.legend(loc="best", handles=[line_truth])
     fig.tight_layout(pad=1)
     plt.savefig(outpath, bbox_inches="tight", pad_inches=0.05)
     plt.close("all")
diff --git a/radionets/evaluation/utils.py b/radionets/evaluation/utils.py
index a9cac1d3..b23f3a16 100644
--- a/radionets/evaluation/utils.py
+++ b/radionets/evaluation/utils.py
@@ -307,10 +307,10 @@ def get_ifft(array, amp_phase=True):
     if len(array.shape) == 3:
         array = array.unsqueeze(0)
     if amp_phase:
-        amp = 10 ** (10 * array[:, 0] - 10) - 1e-10
+        # amp = 10 ** (10 * array[:, 0] - 10) - 1e-10
 
-        a = amp * np.cos(array[:, 1])
-        b = amp * np.sin(array[:, 1])
+        a = array[:, 0] * np.cos(array[:, 1])
+        b = array[:, 0] * np.sin(array[:, 1])
         compl = a + b * 1j
     else:
         compl = array[:, 0] + array[:, 1] * 1j
@@ -365,10 +365,10 @@ def fft_pred(pred, truth, amp_phase=True):
     b_true = truth[1, :, :]
 
     if amp_phase:
-        amp_pred_rescaled = (10 ** (10 * a) - 1) / 10 ** 10
+        amp_pred_rescaled = (10 ** (10 * a) - 1) / 10**10
         phase_pred = b
 
-        amp_true_rescaled = (10 ** (10 * a_true) - 1) / 10 ** 10
+        amp_true_rescaled = (10 ** (10 * a_true) - 1) / 10**10
         phase_true = b_true
 
         compl_pred = amp_pred_rescaled * np.exp(1j * phase_pred)

From af5c6e71834d73fb9c227cf952e2a4366dd3fff0 Mon Sep 17 00:00:00 2001
From: Felix Geyer <felix.geyer@tu-dortmund.de>
Date: Wed, 18 Jan 2023 11:00:50 +0000
Subject: [PATCH 04/31] Add methods for evaluation

---
 radionets/evaluation/train_inspection.py | 10 ++++++++++
 radionets/evaluation/utils.py            | 18 ++++++++++++++++++
 2 files changed, 28 insertions(+)

diff --git a/radionets/evaluation/train_inspection.py b/radionets/evaluation/train_inspection.py
index 5842d46f..afd72e6f 100644
--- a/radionets/evaluation/train_inspection.py
+++ b/radionets/evaluation/train_inspection.py
@@ -28,6 +28,7 @@
     pad_unsqueeze,
     save_pred,
     read_pred,
+    sym_new,
 )
 from radionets.evaluation.jet_angle import calc_jet_angle
 from radionets.evaluation.dynamic_range import calc_dr
@@ -36,6 +37,7 @@
 from radionets.evaluation.pointsources import flux_comparison
 from pytorch_msssim import ms_ssim
 from tqdm import tqdm
+import torch.nn.functional as F
 
 
 def create_predictions(conf):
@@ -82,6 +84,14 @@ def get_prediction(conf, mode="test"):
         pred_2 = pred[:, 2, :].unsqueeze(1)
         pred = torch.cat((pred_1, pred_2), dim=1)
 
+    pred = F.pad(input=pred, pad=(0, 0, 0, 63), mode="constant", value=0)
+    img_test = F.pad(input=img_test, pad=(0, 0, 0, 63), mode="constant", value=0)
+    img_true = F.pad(input=img_true, pad=(0, 0, 0, 63), mode="constant", value=0)
+
+    pred = sym_new(pred)
+    img_test = sym_new(img_test)
+    img_true = sym_new(img_true)
+
     return pred, img_test, img_true
 
 
diff --git a/radionets/evaluation/utils.py b/radionets/evaluation/utils.py
index b23f3a16..dc9934a6 100644
--- a/radionets/evaluation/utils.py
+++ b/radionets/evaluation/utils.py
@@ -257,6 +257,8 @@ def get_images(test_ds, num_images, rand=False):
         indices = torch.randint(0, len(test_ds), size=(num_images,))
     img_test = test_ds[indices][0]
     img_true = test_ds[indices][1]
+    img_test = img_test[:, :, :65, :]
+    img_true = img_true[:, :, :65, :]
     return img_test, img_true
 
 
@@ -429,3 +431,19 @@ def check_outpath(model_path):
     path = Path(model_path)
     exists = path.exists()
     return exists
+
+
+def sym_new(x):
+    if x.shape[1] == 4:
+        channel = 2
+    else:
+        channel = 1
+    upper_half = x[:, :, :64, :].clone()
+    a = torch.rot90(upper_half, 2, dims=[-2, -1])
+    # amp
+    x[:, 0, 65:, 1:] = a[:, 0, :-1, :-1]
+    x[:, 0, 65:, 0] = a[:, 0, :-1, -1]
+    # phase
+    x[:, channel, 65:, 1:] = -a[:, channel, :-1, :-1]
+    x[:, channel, 65:, 0] = -a[:, channel, :-1, -1]
+    return x

From 7102ba23ff468eba3c4001a93986aa6523249075 Mon Sep 17 00:00:00 2001
From: Felix Geyer <felix.geyer@tu-dortmund.de>
Date: Wed, 18 Jan 2023 11:02:31 +0000
Subject: [PATCH 05/31] Comment out

---
 radionets/evaluation/plotting.py | 7 ++++---
 1 file changed, 4 insertions(+), 3 deletions(-)

diff --git a/radionets/evaluation/plotting.py b/radionets/evaluation/plotting.py
index 3dc3bbea..76480691 100644
--- a/radionets/evaluation/plotting.py
+++ b/radionets/evaluation/plotting.py
@@ -13,7 +13,8 @@
 from radionets.evaluation.blob_detection import calc_blobs
 from radionets.evaluation.contour import compute_area_ratio
 from radionets.evaluation.dynamic_range import calc_dr, get_boxsize
-from radionets.evaluation.jet_angle import calc_jet_angle
+
+# from radionets.evaluation.jet_angle import calc_jet_angle
 from radionets.evaluation.utils import (
     check_vmin_vmax,
     make_axes_nice,
@@ -348,8 +349,8 @@ def visualize_source_reconstruction(
     msssim=False,
     plot_format="png",
 ):
-    m_truth, n_truth, alpha_truth = calc_jet_angle(ifft_truth)
-    m_pred, n_pred, alpha_pred = calc_jet_angle(ifft_pred)
+    # m_truth, n_truth, alpha_truth = calc_jet_angle(ifft_truth)
+    # m_pred, n_pred, alpha_pred = calc_jet_angle(ifft_pred)
     # x_space = torch.arange(0, 63, 1)
 
     # plt.style.use("./paper_large_3.rc")

From b928745d00aa285e842e8e577977a5d762079753 Mon Sep 17 00:00:00 2001
From: Felix Geyer <felix.geyer@tu-dortmund.de>
Date: Fri, 20 Jan 2023 10:49:03 +0000
Subject: [PATCH 06/31] Add Locally connected layer to SRResNet

---
 radionets/dl_framework/architectures/res_exp.py   | 12 ++++++++++--
 radionets/dl_framework/architectures/unc_archs.py |  2 +-
 2 files changed, 11 insertions(+), 3 deletions(-)

diff --git a/radionets/dl_framework/architectures/res_exp.py b/radionets/dl_framework/architectures/res_exp.py
index 4a58e22d..64d7b6e2 100644
--- a/radionets/dl_framework/architectures/res_exp.py
+++ b/radionets/dl_framework/architectures/res_exp.py
@@ -4,6 +4,7 @@
     SRBlock,
     Lambda,
     symmetry,
+    LocallyConnected2d,
 )
 from functools import partial
 from math import pi
@@ -120,7 +121,7 @@ def forward(self, x):
 
 
 class SRResNet_16(nn.Module):
-    def __init__(self):
+    def __init__(self, img_size):
         super().__init__()
 
         self.preBlock = nn.Sequential(
@@ -152,7 +153,14 @@ def __init__(self):
         )
 
         self.final = nn.Sequential(
-            nn.Conv2d(64, 2, 9, stride=1, padding=4, groups=2),
+            LocallyConnected2d(
+                64,
+                2,
+                [img_size // 2 + 1, img_size],
+                1,
+                stride=1,
+                bias=False,
+            )
         )
 
         self.hardtanh = nn.Hardtanh(-pi, pi)
diff --git a/radionets/dl_framework/architectures/unc_archs.py b/radionets/dl_framework/architectures/unc_archs.py
index b066a64e..7940eed4 100644
--- a/radionets/dl_framework/architectures/unc_archs.py
+++ b/radionets/dl_framework/architectures/unc_archs.py
@@ -55,7 +55,7 @@ def forward(self, x):
 class UncertaintyWrapper(nn.Module):
     def __init__(self, img_size):
         super().__init__()
-        self.pred = SRResNet_16()
+        self.pred = SRResNet_16(img_size)
 
         self.uncertainty = Uncertainty(img_size)
 

From 68bf35a0b27a66c04b1ddc482665e5e5e27a7203 Mon Sep 17 00:00:00 2001
From: Felix Geyer <felix.geyer@tu-dortmund.de>
Date: Fri, 20 Jan 2023 10:49:47 +0000
Subject: [PATCH 07/31] Padding_mode change to reflect

---
 radionets/dl_framework/model.py | 8 ++++++--
 1 file changed, 6 insertions(+), 2 deletions(-)

diff --git a/radionets/dl_framework/model.py b/radionets/dl_framework/model.py
index ffaa70aa..9293b31a 100644
--- a/radionets/dl_framework/model.py
+++ b/radionets/dl_framework/model.py
@@ -281,10 +281,14 @@ def forward(self, x):
 
     def _conv_block(self, ni, nf, stride):
         return nn.Sequential(
-            nn.Conv2d(ni, nf, 3, stride=stride, padding=1, bias=False),
+            nn.Conv2d(
+                ni, nf, 3, stride=stride, padding=1, bias=False, padding_mode="reflect"
+            ),
             nn.BatchNorm2d(nf),
             nn.PReLU(),
-            nn.Conv2d(nf, nf, 3, stride=1, padding=1, bias=False),
+            nn.Conv2d(
+                nf, nf, 3, stride=1, padding=1, bias=False, padding_mode="reflect"
+            ),
             nn.BatchNorm2d(nf),
         )
 

From 688d5d4c43d8274d5e6066f85ec999d52f210109 Mon Sep 17 00:00:00 2001
From: Felix Geyer <felix.geyer@tu-dortmund.de>
Date: Mon, 23 Jan 2023 14:02:50 +0100
Subject: [PATCH 08/31] Utilize symmetry only for 128 pixel images

---
 radionets/evaluation/train_inspection.py | 15 ++++++++-------
 1 file changed, 8 insertions(+), 7 deletions(-)

diff --git a/radionets/evaluation/train_inspection.py b/radionets/evaluation/train_inspection.py
index afd72e6f..1a2deec4 100644
--- a/radionets/evaluation/train_inspection.py
+++ b/radionets/evaluation/train_inspection.py
@@ -84,13 +84,14 @@ def get_prediction(conf, mode="test"):
         pred_2 = pred[:, 2, :].unsqueeze(1)
         pred = torch.cat((pred_1, pred_2), dim=1)
 
-    pred = F.pad(input=pred, pad=(0, 0, 0, 63), mode="constant", value=0)
-    img_test = F.pad(input=img_test, pad=(0, 0, 0, 63), mode="constant", value=0)
-    img_true = F.pad(input=img_true, pad=(0, 0, 0, 63), mode="constant", value=0)
-
-    pred = sym_new(pred)
-    img_test = sym_new(img_test)
-    img_true = sym_new(img_true)
+    if pred.shape[-1] == 128:
+        pred = F.pad(input=pred, pad=(0, 0, 0, 63), mode="constant", value=0)
+        img_test = F.pad(input=img_test, pad=(0, 0, 0, 63), mode="constant", value=0)
+        img_true = F.pad(input=img_true, pad=(0, 0, 0, 63), mode="constant", value=0)
+
+        pred = sym_new(pred)
+        img_test = sym_new(img_test)
+        img_true = sym_new(img_true)
 
     return pred, img_test, img_true
 

From aee1899d3bf716c5b864c366c6814190e3bb6979 Mon Sep 17 00:00:00 2001
From: Felix Geyer <felix.geyer@tu-dortmund.de>
Date: Mon, 23 Jan 2023 13:33:30 +0000
Subject: [PATCH 09/31] Delete Locally connected layer

---
 radionets/dl_framework/architectures/res_exp.py   | 14 ++------------
 radionets/dl_framework/architectures/unc_archs.py |  2 +-
 2 files changed, 3 insertions(+), 13 deletions(-)

diff --git a/radionets/dl_framework/architectures/res_exp.py b/radionets/dl_framework/architectures/res_exp.py
index 64d7b6e2..96770b8c 100644
--- a/radionets/dl_framework/architectures/res_exp.py
+++ b/radionets/dl_framework/architectures/res_exp.py
@@ -4,7 +4,6 @@
     SRBlock,
     Lambda,
     symmetry,
-    LocallyConnected2d,
 )
 from functools import partial
 from math import pi
@@ -121,7 +120,7 @@ def forward(self, x):
 
 
 class SRResNet_16(nn.Module):
-    def __init__(self, img_size):
+    def __init__(self):
         super().__init__()
 
         self.preBlock = nn.Sequential(
@@ -152,16 +151,7 @@ def __init__(self, img_size):
             nn.Conv2d(64, 64, 3, stride=1, padding=1, bias=False), nn.BatchNorm2d(64)
         )
 
-        self.final = nn.Sequential(
-            LocallyConnected2d(
-                64,
-                2,
-                [img_size // 2 + 1, img_size],
-                1,
-                stride=1,
-                bias=False,
-            )
-        )
+        self.final = nn.Sequential(nn.Conv2d(64, 2, 9, stride=1, padding=4, groups=2))
 
         self.hardtanh = nn.Hardtanh(-pi, pi)
         self.relu = nn.ReLU()
diff --git a/radionets/dl_framework/architectures/unc_archs.py b/radionets/dl_framework/architectures/unc_archs.py
index 7940eed4..b066a64e 100644
--- a/radionets/dl_framework/architectures/unc_archs.py
+++ b/radionets/dl_framework/architectures/unc_archs.py
@@ -55,7 +55,7 @@ def forward(self, x):
 class UncertaintyWrapper(nn.Module):
     def __init__(self, img_size):
         super().__init__()
-        self.pred = SRResNet_16(img_size)
+        self.pred = SRResNet_16()
 
         self.uncertainty = Uncertainty(img_size)
 

From 0e11ed0a269b7e0af20964300ca8b8fdc150954a Mon Sep 17 00:00:00 2001
From: Felix Geyer <felix.geyer@tu-dortmund.de>
Date: Mon, 23 Jan 2023 14:58:10 +0000
Subject: [PATCH 10/31] Delete double occurance of sym_new

---
 radionets/dl_framework/model.py | 12 ------------
 1 file changed, 12 deletions(-)

diff --git a/radionets/dl_framework/model.py b/radionets/dl_framework/model.py
index 9293b31a..5c0599d1 100644
--- a/radionets/dl_framework/model.py
+++ b/radionets/dl_framework/model.py
@@ -308,15 +308,3 @@ def symmetry(x):
     x[:, 0, x.shape[2] // 2 :, :] = upper_half[:, 0]
     x[:, 1, x.shape[2] // 2 :, :] = -upper_half[:, 1]
     return x
-
-
-def sym_new(x):
-    upper_half = x[:, :, :64, :].clone()
-    a = torch.rot90(upper_half, 2, dims=[-2, -1])
-    # amp
-    x[:, 0, 65:, 1:] = a[:, 0, :-1, :-1]
-    x[:, 0, 65:, 0] = a[:, 0, :-1, -1]
-    # phase
-    x[:, 1, 65:, 1:] = -a[:, 1, :-1, :-1]
-    x[:, 1, 65:, 0] = -a[:, 1, :-1, -1]
-    return x

From 51c34e00386fc5dff214bf44c372a2d6c7af4a08 Mon Sep 17 00:00:00 2001
From: Felix Geyer <felix.geyer@tu-dortmund.de>
Date: Mon, 23 Jan 2023 15:07:43 +0000
Subject: [PATCH 11/31] Add docstring for sym_new function

---
 radionets/evaluation/utils.py | 34 ++++++++++++++++++++++++++--------
 1 file changed, 26 insertions(+), 8 deletions(-)

diff --git a/radionets/evaluation/utils.py b/radionets/evaluation/utils.py
index dc9934a6..c269a826 100644
--- a/radionets/evaluation/utils.py
+++ b/radionets/evaluation/utils.py
@@ -433,17 +433,35 @@ def check_outpath(model_path):
     return exists
 
 
-def sym_new(x):
-    if x.shape[1] == 4:
+def sym_new(image):
+    """
+    Symmetry function to complete the images
+
+    Parameters
+    ----------
+    image : torch.Tensor
+        (stack of) half images
+
+    Returns
+    -------
+    torch.Tensor
+        quadratic images after utilizing symmetry
+    """
+    # set channel for phase dependent of uncertainty training
+    if image.shape[1] == 4:
         channel = 2
     else:
         channel = 1
-    upper_half = x[:, :, :64, :].clone()
+
+    upper_half = image[:, :, :64, :].clone()
     a = torch.rot90(upper_half, 2, dims=[-2, -1])
+
     # amp
-    x[:, 0, 65:, 1:] = a[:, 0, :-1, :-1]
-    x[:, 0, 65:, 0] = a[:, 0, :-1, -1]
+    image[:, 0, 65:, 1:] = a[:, 0, :-1, :-1]
+    image[:, 0, 65:, 0] = a[:, 0, :-1, -1]
+
     # phase
-    x[:, channel, 65:, 1:] = -a[:, channel, :-1, :-1]
-    x[:, channel, 65:, 0] = -a[:, channel, :-1, -1]
-    return x
+    image[:, channel, 65:, 1:] = -a[:, channel, :-1, :-1]
+    image[:, channel, 65:, 0] = -a[:, channel, :-1, -1]
+
+    return image

From 8d2b56bb5c463530f85a3e1181f482be9e01ccb2 Mon Sep 17 00:00:00 2001
From: Felix Geyer <felix.geyer@tu-dortmund.de>
Date: Mon, 30 Jan 2023 12:10:10 +0100
Subject: [PATCH 12/31] Fix sampling for half images

---
 radionets/evaluation/train_inspection.py | 22 ++++++--------------
 radionets/evaluation/utils.py            | 26 +++++++++++++++---------
 2 files changed, 22 insertions(+), 26 deletions(-)

diff --git a/radionets/evaluation/train_inspection.py b/radionets/evaluation/train_inspection.py
index fcf74150..6ba93f41 100644
--- a/radionets/evaluation/train_inspection.py
+++ b/radionets/evaluation/train_inspection.py
@@ -31,7 +31,7 @@
     pad_unsqueeze,
     save_pred,
     read_pred,
-    sym_new,
+    apply_symmetry,
     sample_images,
     mergeDictionary,
     sampled_dataset,
@@ -43,7 +43,6 @@
 from radionets.evaluation.pointsources import flux_comparison
 from pytorch_msssim import ms_ssim
 from tqdm import tqdm
-import torch.nn.functional as F
 
 
 def create_predictions(conf):
@@ -96,20 +95,8 @@ def get_prediction(conf, mode="test"):
         images["indices"] = indices
 
     if images["pred"].shape[-1] == 128:
-        pred = F.pad(input=images["pred"], pad=(0, 0, 0, 63), mode="constant", value=0)
-        img_test = F.pad(
-            input=images["test"], pad=(0, 0, 0, 63), mode="constant", value=0
-        )
-        img_true = F.pad(
-            input=images["true"], pad=(0, 0, 0, 63), mode="constant", value=0
-        )
-
-        pred = sym_new(pred)
-        img_test = sym_new(img_test)
-        img_true = sym_new(img_true)
-        images["pred"] = pred
-        images["test"] = img_test
-        images["true"] = img_true
+        images = apply_symmetry(images)
+    print(images["unc"].shape)
 
     return images
 
@@ -602,6 +589,9 @@ def save_sampled(conf):
             img["unc"] = unc
             img["pred"] = pred
 
+        if img["pred"].shape[-1] == 128:
+            img = apply_symmetry(img)
+
         result = sample_images(img["pred"], img["unc"], 100)
         ifft_truth = get_ifft(img["true"], amp_phase=conf["amp_phase"])
         dict_img_true = {"true": ifft_truth}
diff --git a/radionets/evaluation/utils.py b/radionets/evaluation/utils.py
index d439a6bf..8e367686 100644
--- a/radionets/evaluation/utils.py
+++ b/radionets/evaluation/utils.py
@@ -471,7 +471,7 @@ def check_outpath(model_path):
     return exists
 
 
-def sym_new(image):
+def sym_new(image, key):
     """
     Symmetry function to complete the images
 
@@ -486,25 +486,31 @@ def sym_new(image):
         quadratic images after utilizing symmetry
     """
     # set channel for phase dependent of uncertainty training
-    if image.shape[1] == 4:
-        channel = 2
-    else:
-        channel = 1
-
     upper_half = image[:, :, :64, :].clone()
     a = torch.rot90(upper_half, 2, dims=[-2, -1])
 
-    # amp
     image[:, 0, 65:, 1:] = a[:, 0, :-1, :-1]
     image[:, 0, 65:, 0] = a[:, 0, :-1, -1]
 
-    # phase
-    image[:, channel, 65:, 1:] = -a[:, channel, :-1, :-1]
-    image[:, channel, 65:, 0] = -a[:, channel, :-1, -1]
+    if key == "unc":
+        image[:, 1, 65:, 1:] = a[:, 1, :-1, :-1]
+        image[:, 1, 65:, 0] = a[:, 1, :-1, -1]
+    else:
+        image[:, 1, 65:, 1:] = -a[:, 1, :-1, :-1]
+        image[:, 1, 65:, 0] = -a[:, 1, :-1, -1]
 
     return image
 
 
+def apply_symmetry(img_dict):
+    for key in img_dict:
+        output = F.pad(input=img_dict[key], pad=(0, 0, 0, 63), mode="constant", value=0)
+        output = sym_new(output, key)
+        img_dict[key] = output
+
+    return img_dict
+
+
 def even_better_symmetry(x):
     upper_half = x[:, :, 0 : x.shape[2] // 2, :].copy()
     upper_left = upper_half[:, :, :, 0 : upper_half.shape[3] // 2].copy()

From f55723f6c76e5fc2b24ae22dbc63730c82f6f0c4 Mon Sep 17 00:00:00 2001
From: Felix Geyer <felix.geyer@tu-dortmund.de>
Date: Mon, 30 Jan 2023 13:07:57 +0000
Subject: [PATCH 13/31] Fix double occurance indices

---
 radionets/evaluation/utils.py | 14 +++++++++++---
 1 file changed, 11 insertions(+), 3 deletions(-)

diff --git a/radionets/evaluation/utils.py b/radionets/evaluation/utils.py
index 8e367686..281ff825 100644
--- a/radionets/evaluation/utils.py
+++ b/radionets/evaluation/utils.py
@@ -285,6 +285,11 @@ def get_images(test_ds, num_images, rand=False, indices=None):
         indices = torch.arange(num_images)
         if rand:
             indices = torch.randint(0, len(test_ds), size=(num_images,))
+            while len(torch.unique(indices)) < len(indices):
+                new_indices = torch.randint(
+                    0, len(test_ds), size=(num_images - len(torch.unique(indices)),)
+                )
+                indices = torch.cat((torch.unique(indices), new_indices))
             indices, _ = torch.sort(indices)
         img_test = test_ds[indices][0]
         img_true = test_ds[indices][1]
@@ -504,9 +509,12 @@ def sym_new(image, key):
 
 def apply_symmetry(img_dict):
     for key in img_dict:
-        output = F.pad(input=img_dict[key], pad=(0, 0, 0, 63), mode="constant", value=0)
-        output = sym_new(output, key)
-        img_dict[key] = output
+        if key != "indices":
+            output = F.pad(
+                input=img_dict[key], pad=(0, 0, 0, 63), mode="constant", value=0
+            )
+            output = sym_new(output, key)
+            img_dict[key] = output
 
     return img_dict
 

From b07dbd3a0c8ccf6af5ddc99fbb5e09556e1116b8 Mon Sep 17 00:00:00 2001
From: Felix Geyer <felix.geyer@tu-dortmund.de>
Date: Mon, 30 Jan 2023 13:08:50 +0000
Subject: [PATCH 14/31] Delete print

---
 radionets/evaluation/train_inspection.py | 1 -
 1 file changed, 1 deletion(-)

diff --git a/radionets/evaluation/train_inspection.py b/radionets/evaluation/train_inspection.py
index 6ba93f41..e7ef1e57 100644
--- a/radionets/evaluation/train_inspection.py
+++ b/radionets/evaluation/train_inspection.py
@@ -96,7 +96,6 @@ def get_prediction(conf, mode="test"):
 
     if images["pred"].shape[-1] == 128:
         images = apply_symmetry(images)
-    print(images["unc"].shape)
 
     return images
 

From a4f2d8e474d3ce37e2d55f078498ad09ffd45ffd Mon Sep 17 00:00:00 2001
From: Felix Geyer <felix.geyer@tu-dortmund.de>
Date: Mon, 30 Jan 2023 14:08:15 +0000
Subject: [PATCH 15/31] Nice plotting for uncertainty images

---
 radionets/evaluation/plotting.py | 10 +++++-----
 1 file changed, 5 insertions(+), 5 deletions(-)

diff --git a/radionets/evaluation/plotting.py b/radionets/evaluation/plotting.py
index 0d470c21..39e895e5 100644
--- a/radionets/evaluation/plotting.py
+++ b/radionets/evaluation/plotting.py
@@ -466,19 +466,19 @@ def visualize_uncertainty(
         2, 2, sharey=True, sharex=True, figsize=(12, 10)
     )
 
-    im1 = ax1.imshow(true_phase)
+    im1 = ax1.imshow(true_phase, cmap=OrBu, vmin=-np.pi, vmax=np.pi)
 
-    im2 = ax2.imshow(pred_phase)
+    im2 = ax2.imshow(pred_phase, cmap=OrBu, vmin=-np.pi, vmax=np.pi)
 
     im3 = ax3.imshow(unc_phase)
 
     a = check_vmin_vmax(true_phase - pred_phase)
     im4 = ax4.imshow(true_phase - pred_phase, cmap=OrBu, vmin=-a, vmax=a)
 
-    make_axes_nice(fig, ax1, im1, r"Simulation")
-    make_axes_nice(fig, ax2, im2, r"Predicted $\mu$")
+    make_axes_nice(fig, ax1, im1, r"Simulation", phase=True)
+    make_axes_nice(fig, ax2, im2, r"Predicted $\mu$", phase=True)
     make_axes_nice(fig, ax3, im3, r"Predicted $\sigma^2$", unc=True)
-    make_axes_nice(fig, ax4, im4, r"Difference")
+    make_axes_nice(fig, ax4, im4, r"Difference", phase_diff=True)
 
     ax1.set_ylabel(r"pixels")
     ax3.set_ylabel(r"pixels")

From a9c7ed561d25c44ec2b91ab0fd326966536817c2 Mon Sep 17 00:00:00 2001
From: Felix Geyer <felix.geyer@tu-dortmund.de>
Date: Mon, 30 Jan 2023 15:47:22 +0100
Subject: [PATCH 16/31] Working changes

---
 radionets/evaluation/train_inspection.py | 4 ++--
 radionets/evaluation/utils.py            | 9 ++++++---
 2 files changed, 8 insertions(+), 5 deletions(-)

diff --git a/radionets/evaluation/train_inspection.py b/radionets/evaluation/train_inspection.py
index e7ef1e57..a43fa29d 100644
--- a/radionets/evaluation/train_inspection.py
+++ b/radionets/evaluation/train_inspection.py
@@ -588,8 +588,8 @@ def save_sampled(conf):
             img["unc"] = unc
             img["pred"] = pred
 
-        if img["pred"].shape[-1] == 128:
-            img = apply_symmetry(img)
+        # if img["pred"].shape[-1] == 128:
+        #     img = apply_symmetry(img)
 
         result = sample_images(img["pred"], img["unc"], 100)
         ifft_truth = get_ifft(img["true"], amp_phase=conf["amp_phase"])
diff --git a/radionets/evaluation/utils.py b/radionets/evaluation/utils.py
index 281ff825..05b5dceb 100644
--- a/radionets/evaluation/utils.py
+++ b/radionets/evaluation/utils.py
@@ -59,7 +59,7 @@ def create_databunch(data_path, fourier, source_list, batch_size):
             test_ds, batch_size=batch_size, shuffle=True, collate_fn=source_list_collate
         )
     else:
-        data = DataLoader(test_ds, batch_size=batch_size, shuffle=True)
+        data = DataLoader(test_ds, batch_size=batch_size, shuffle=False)
     return data
 
 
@@ -491,6 +491,8 @@ def sym_new(image, key):
         quadratic images after utilizing symmetry
     """
     # set channel for phase dependent of uncertainty training
+    if isinstance(image, np.ndarray):
+        image = torch.tensor(image)
     upper_half = image[:, :, :64, :].clone()
     a = torch.rot90(upper_half, 2, dims=[-2, -1])
 
@@ -549,7 +551,7 @@ def trunc_rvs(loc, scale, mode, num_samples, num_img):
         raise ValueError("Wrong mode!")
     a, b = (myclip_a - loc) / scale, (myclip_b - loc) / scale
     sampled_gauss = truncnorm.rvs(
-        a, b, loc=loc, scale=scale, size=(num_samples, num_img, 128, 128)
+        a, b, loc=loc, scale=scale, size=(num_samples, num_img, 65, 128)
     )
 
     return sampled_gauss.swapaxes(0, 1)
@@ -578,7 +580,8 @@ def sample_images(mean, std, num_samples):
     assert mask_invalid_phase.sum() == 0
 
     sampled_gauss = np.stack([sampled_gauss_amp, sampled_gauss_phase], axis=1)
-    sampled_gauss_symmetry = even_better_symmetry(sampled_gauss)
+    # sampled_gauss_symmetry = even_better_symmetry(sampled_gauss)
+    sampled_gauss_symmetry = sym_new(sampled_gauss, None)
 
     fft_sampled_symmetry = get_ifft(sampled_gauss_symmetry, amp_phase=True, scale=False)
     results = {

From d7ba2d067136ea3e2a155e5bbcd0495ce5350b68 Mon Sep 17 00:00:00 2001
From: Felix Geyer <felix.geyer@tu-dortmund.de>
Date: Tue, 31 Jan 2023 07:14:29 +0000
Subject: [PATCH 17/31] Fix sampling

---
 radionets/evaluation/train_inspection.py | 37 ++++++++++++++----------
 radionets/evaluation/utils.py            | 19 +++++++++---
 2 files changed, 37 insertions(+), 19 deletions(-)

diff --git a/radionets/evaluation/train_inspection.py b/radionets/evaluation/train_inspection.py
index a43fa29d..4cc3174f 100644
--- a/radionets/evaluation/train_inspection.py
+++ b/radionets/evaluation/train_inspection.py
@@ -43,6 +43,8 @@
 from radionets.evaluation.pointsources import flux_comparison
 from pytorch_msssim import ms_ssim
 from tqdm import tqdm
+import torch.nn.functional as F
+from radionets.evaluation.utils import sym_new
 
 
 def create_predictions(conf):
@@ -578,27 +580,32 @@ def save_sampled(conf):
             pred = torch.cat((pred, pred_2), dim=1)
 
         img = {"pred": pred, "inp": img_test, "true": img_true}
-        if pred.shape[1] == 4:
-            unc_amp = pred[:, 1, :]
-            unc_phase = pred[:, 3, :]
-            unc = torch.stack([unc_amp, unc_phase], dim=1)
-            pred_1 = pred[:, 0, :]
-            pred_2 = pred[:, 2, :]
-            pred = torch.stack((pred_1, pred_2), dim=1)
-            img["unc"] = unc
-            img["pred"] = pred
-
-        # if img["pred"].shape[-1] == 128:
-        #     img = apply_symmetry(img)
-
-        result = sample_images(img["pred"], img["unc"], 100)
+        # separate prediction and uncertainty
+        unc_amp = pred[:, 1, :]
+        unc_phase = pred[:, 3, :]
+        unc = torch.stack([unc_amp, unc_phase], dim=1)
+        pred_1 = pred[:, 0, :]
+        pred_2 = pred[:, 2, :]
+        pred = torch.stack((pred_1, pred_2), dim=1)
+        img["unc"] = unc
+        img["pred"] = pred
+
+        result = sample_images(img["pred"], img["unc"], 10)
+
+        # pad true image
+        output = F.pad(input=img["true"], pad=(0, 0, 0, 63), mode="constant", value=0)
+        img["true"] = sym_new(output, None)
         ifft_truth = get_ifft(img["true"], amp_phase=conf["amp_phase"])
+
+        # add images to dict
         dict_img_true = {"true": ifft_truth}
-        # print(dict_img_true["true"].shape)
         results = mergeDictionary(results, result)
         results = mergeDictionary(results, dict_img_true)
+
+    # reshaping
     for key in results.keys():
         results[key] = results[key].reshape(num_img, img_size, img_size)
+
     save_pred(str(out_path) + "/sampled_imgs.h5", results)
 
 
diff --git a/radionets/evaluation/utils.py b/radionets/evaluation/utils.py
index 05b5dceb..40803530 100644
--- a/radionets/evaluation/utils.py
+++ b/radionets/evaluation/utils.py
@@ -512,6 +512,8 @@ def sym_new(image, key):
 def apply_symmetry(img_dict):
     for key in img_dict:
         if key != "indices":
+            if isinstance(img_dict[key], np.ndarray):
+                img_dict[key] = torch.tensor(img_dict[key])
             output = F.pad(
                 input=img_dict[key], pad=(0, 0, 0, 63), mode="constant", value=0
             )
@@ -562,12 +564,14 @@ def sample_images(mean, std, num_samples):
     std_amp, std_phase = std[:, 0], std[:, 1]
     num_img = mean_amp.shape[0]
     # amplitude
-    sampled_gauss_amp = trunc_rvs(mean_amp, std_amp, "amp", num_samples, num_img)
+    sampled_gauss_amp = trunc_rvs(
+        mean_amp, std_amp, "amp", num_samples, num_img
+    ).reshape(num_img * num_samples, 65, 128)
 
     # phase
     sampled_gauss_phase = trunc_rvs(
         mean_phase, std_phase, "phase", num_samples, num_img
-    )
+    ).reshape(num_img * num_samples, 65, 128)
 
     # masks
     mask_invalid_amp = sampled_gauss_amp <= (0 - 1e-4)
@@ -580,10 +584,17 @@ def sample_images(mean, std, num_samples):
     assert mask_invalid_phase.sum() == 0
 
     sampled_gauss = np.stack([sampled_gauss_amp, sampled_gauss_phase], axis=1)
-    # sampled_gauss_symmetry = even_better_symmetry(sampled_gauss)
+
+    # pad resulting images and utilize symmetry
+    sampled_gauss = F.pad(
+        input=torch.tensor(sampled_gauss), pad=(0, 0, 0, 63), mode="constant", value=0
+    )
     sampled_gauss_symmetry = sym_new(sampled_gauss, None)
 
-    fft_sampled_symmetry = get_ifft(sampled_gauss_symmetry, amp_phase=True, scale=False)
+    fft_sampled_symmetry = get_ifft(
+        sampled_gauss_symmetry, amp_phase=True, scale=False
+    ).reshape(num_img, num_samples, 128, 128)
+
     results = {
         "mean": fft_sampled_symmetry.mean(axis=1),
         "std": fft_sampled_symmetry.std(axis=1),

From e02536d0aa4b1255d19dab5655dcf54064569dfd Mon Sep 17 00:00:00 2001
From: Kevin Schmidt <Kevin2@users.noreply.github.com>
Date: Tue, 31 Jan 2023 09:06:55 +0100
Subject: [PATCH 18/31] bump version

---
 setup.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/setup.py b/setup.py
index 7e088763..3bf829f8 100644
--- a/setup.py
+++ b/setup.py
@@ -2,7 +2,7 @@
 
 setup(
     name="radionets",
-    version="0.1.18",
+    version="0.2.0",
     description="Imaging radio interferometric data with neural networks",
     url="https://github.com/radionets-project/radionets",
     author="Kevin Schmidt, Felix Geyer",

From 183f436f6b05036eb53ef30b21847b9ee338b0d6 Mon Sep 17 00:00:00 2001
From: Felix Geyer <felix.geyer@tu-dortmund.de>
Date: Tue, 31 Jan 2023 08:13:01 +0000
Subject: [PATCH 19/31] Delete calc_blobs option

---
 radionets/dl_framework/model.py               |  1 -
 radionets/evaluation/plotting.py              | 33 -------------------
 .../evaluation/scripts/start_evaluation.py    |  3 --
 radionets/evaluation/train_inspection.py      |  1 -
 radionets/evaluation/utils.py                 |  1 -
 tests/evaluate.toml                           |  1 -
 6 files changed, 40 deletions(-)

diff --git a/radionets/dl_framework/model.py b/radionets/dl_framework/model.py
index 5c0599d1..6fa1e29e 100644
--- a/radionets/dl_framework/model.py
+++ b/radionets/dl_framework/model.py
@@ -191,7 +191,6 @@ def __init__(
         self, in_channels, out_channels, output_size, kernel_size, stride, bias=False
     ):
         super(LocallyConnected2d, self).__init__()
-        # output_size = _pair(output_size)
         self.weight = nn.Parameter(
             torch.randn(
                 1,
diff --git a/radionets/evaluation/plotting.py b/radionets/evaluation/plotting.py
index 39e895e5..353a8199 100644
--- a/radionets/evaluation/plotting.py
+++ b/radionets/evaluation/plotting.py
@@ -10,7 +10,6 @@
 from mpl_toolkits.axes_grid1 import make_axes_locatable
 from pytorch_msssim import ms_ssim
 
-from radionets.evaluation.blob_detection import calc_blobs
 from radionets.evaluation.contour import compute_area_ratio
 from radionets.evaluation.dynamic_range import calc_dr, get_boxsize
 
@@ -345,37 +344,16 @@ def visualize_source_reconstruction(
     out_path,
     i,
     dr=False,
-    blobs=False,
     msssim=False,
     plot_format="png",
 ):
-    # m_truth, n_truth, alpha_truth = calc_jet_angle(ifft_truth)
-    # m_pred, n_pred, alpha_pred = calc_jet_angle(ifft_pred)
-    # x_space = torch.arange(0, 63, 1)
-
     # plt.style.use("./paper_large_3.rc")
     fig, (ax1, ax2, ax3) = plt.subplots(1, 3, figsize=(16, 10), sharey=True)
 
     # Plot prediction
-    # ax1.plot(x_space, m_pred * x_space + n_pred, "w--", alpha=0.5)
-    # ax1.axvline(32, 0, 1, linestyle="--", color="white", alpha=0.5)
-
-    # create angle visualization
-    # theta1 = min(0, -alpha_pred.numpy()[0])
-    # theta2 = max(0, -alpha_pred.numpy()[0])
-    # ax1.add_patch(Arc([32, 32], 50, 50, 90, theta1, theta2, color="white"))
-
     im1 = ax1.imshow(ifft_pred, vmax=ifft_truth.max(), cmap="inferno")
 
     # Plot truth
-    # ax2.plot(x_space, m_truth * x_space + n_truth, "w--", alpha=0.5)
-    # ax2.axvline(32, 0, 1, linestyle="--", color="white", alpha=0.5)
-
-    # create angle visualization
-    # theta1 = min(0, -alpha_truth.numpy()[0])
-    # theta2 = max(0, -alpha_truth.numpy()[0])
-    # ax2.add_patch(Arc([32, 32], 50, 50, 90, theta1, theta2, color="white"))
-
     im2 = ax2.imshow(ifft_truth, cmap="inferno")
 
     a = check_vmin_vmax(ifft_pred - ifft_truth)
@@ -390,15 +368,6 @@ def visualize_source_reconstruction(
     ax2.set_xlabel(r"Pixels")
     ax3.set_xlabel(r"Pixels")
 
-    # ax1.tick_params(axis="both", labelsize=20)
-    # ax2.tick_params(axis="both", labelsize=20)
-    # ax3.tick_params(axis="both", labelsize=20)
-
-    if blobs:
-        blobs_pred, blobs_truth = calc_blobs(ifft_pred, ifft_truth)
-        plot_blobs(blobs_pred, ax1)
-        plot_blobs(blobs_truth, ax2)
-
     if dr:
         dr_truth, dr_pred, num_boxes, corners = calc_dr(
             ifft_truth[None, ...], ifft_pred[None, ...]
@@ -418,8 +387,6 @@ def visualize_source_reconstruction(
 
     outpath = str(out_path) + f"/fft_pred_{i}.{plot_format}"
 
-    # ax1.legend(loc="best", handles=[line])
-    # ax2.legend(loc="best", handles=[line_truth])
     fig.tight_layout(pad=1)
     plt.savefig(outpath, bbox_inches="tight", pad_inches=0.05)
     plt.close("all")
diff --git a/radionets/evaluation/scripts/start_evaluation.py b/radionets/evaluation/scripts/start_evaluation.py
index b7ecefb9..24ec90d9 100644
--- a/radionets/evaluation/scripts/start_evaluation.py
+++ b/radionets/evaluation/scripts/start_evaluation.py
@@ -66,9 +66,6 @@ def main(configuration_path):
 
         click.echo(f"\nCreated {eval_conf['num_images']} test predictions.\n")
 
-    if eval_conf["vis_blobs"]:
-        click.echo("\nBlob visualization is enabled for source plots.\n")
-
     if eval_conf["vis_ms_ssim"]:
         click.echo("\nVisualization of ms ssim is enabled for source plots.\n")
 
diff --git a/radionets/evaluation/train_inspection.py b/radionets/evaluation/train_inspection.py
index 4cc3174f..6c361185 100644
--- a/radionets/evaluation/train_inspection.py
+++ b/radionets/evaluation/train_inspection.py
@@ -283,7 +283,6 @@ def create_source_plots(conf, num_images=3, rand=False):
             out_path,
             i,
             dr=conf["vis_dr"],
-            blobs=conf["vis_blobs"],
             msssim=conf["vis_ms_ssim"],
             plot_format=conf["format"],
         )
diff --git a/radionets/evaluation/utils.py b/radionets/evaluation/utils.py
index 40803530..c83f88c7 100644
--- a/radionets/evaluation/utils.py
+++ b/radionets/evaluation/utils.py
@@ -123,7 +123,6 @@ def read_config(config):
     eval_conf["unc"] = config["inspection"]["visualize_uncertainty"]
     eval_conf["plot_contour"] = config["inspection"]["visualize_contour"]
     eval_conf["vis_dr"] = config["inspection"]["visualize_dynamic_range"]
-    eval_conf["vis_blobs"] = config["inspection"]["visualize_blobs"]
     eval_conf["vis_ms_ssim"] = config["inspection"]["visualize_ms_ssim"]
     eval_conf["num_images"] = config["inspection"]["num_images"]
     eval_conf["random"] = config["inspection"]["random"]
diff --git a/tests/evaluate.toml b/tests/evaluate.toml
index 928b75db..bd5e73f0 100644
--- a/tests/evaluate.toml
+++ b/tests/evaluate.toml
@@ -26,7 +26,6 @@ sample_uncertainty = false
 visualize_uncertainty = false
 visualize_contour = true
 visualize_dynamic_range = false
-visualize_blobs = false
 visualize_ms_ssim = false
 random = true
 num_images = 10

From 7c865bbcaf118cd658b9cc03bc9e73b4124c4e6e Mon Sep 17 00:00:00 2001
From: Kevin Schmidt <Kevin2@users.noreply.github.com>
Date: Tue, 31 Jan 2023 09:26:21 +0100
Subject: [PATCH 20/31] bump version

---
 pyproject.toml | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/pyproject.toml b/pyproject.toml
index 733cbff0..881555c1 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "radionets"
-version = "0.1.16"
+version = "0.2.0"
 authors = [
   { name="Kevin Schmidt", email="kevin3.schmidt@tu-dortmund.de" },
 ]

From ab462475f58605a898d89b5aaf19a64c10df53d9 Mon Sep 17 00:00:00 2001
From: Kevin Schmidt <Kevin2@users.noreply.github.com>
Date: Tue, 31 Jan 2023 09:32:27 +0100
Subject: [PATCH 21/31] remove blob detection

---
 examples/default_eval_config.toml | 1 -
 1 file changed, 1 deletion(-)

diff --git a/examples/default_eval_config.toml b/examples/default_eval_config.toml
index f541d0c1..901a043c 100644
--- a/examples/default_eval_config.toml
+++ b/examples/default_eval_config.toml
@@ -24,7 +24,6 @@ visualize_prediction = true
 visualize_source_reconstruction = true
 visualize_contour = true
 visualize_dynamic_range = false
-visualize_blobs = true
 visualize_ms_ssim = false
 random = false
 num_images = 5

From 36eb629e1b1ed54152352d7073f7c674bb527e7a Mon Sep 17 00:00:00 2001
From: Felix Geyer <felix.geyer@tu-dortmund.de>
Date: Tue, 31 Jan 2023 08:32:57 +0000
Subject: [PATCH 22/31] Delete and comment

---
 radionets/evaluation/plotting.py | 5 -----
 radionets/evaluation/utils.py    | 5 +++++
 2 files changed, 5 insertions(+), 5 deletions(-)

diff --git a/radionets/evaluation/plotting.py b/radionets/evaluation/plotting.py
index 353a8199..5de2a578 100644
--- a/radionets/evaluation/plotting.py
+++ b/radionets/evaluation/plotting.py
@@ -13,7 +13,6 @@
 from radionets.evaluation.contour import compute_area_ratio
 from radionets.evaluation.dynamic_range import calc_dr, get_boxsize
 
-# from radionets.evaluation.jet_angle import calc_jet_angle
 from radionets.evaluation.utils import (
     check_vmin_vmax,
     make_axes_nice,
@@ -291,10 +290,6 @@ def visualize_with_fourier_diff(
     real_pred, imag_pred = img_pred[0], img_pred[1]
     real_truth, imag_truth = img_truth[0], img_truth[1]
 
-    # if amp_phase:
-    #    real_pred = 10 ** (10 * real_pred - 10) - 1e-10
-    #    real_truth = 10 ** (10 * real_truth - 10) - 1e-10
-
     # plotting
     # plt.style.use('./paper_large_3_2.rc')
     fig, ((ax1, ax2, ax3), (ax4, ax5, ax6)) = plt.subplots(
diff --git a/radionets/evaluation/utils.py b/radionets/evaluation/utils.py
index c83f88c7..950ea176 100644
--- a/radionets/evaluation/utils.py
+++ b/radionets/evaluation/utils.py
@@ -284,12 +284,17 @@ def get_images(test_ds, num_images, rand=False, indices=None):
         indices = torch.arange(num_images)
         if rand:
             indices = torch.randint(0, len(test_ds), size=(num_images,))
+
+            # remove dublicate indices
             while len(torch.unique(indices)) < len(indices):
                 new_indices = torch.randint(
                     0, len(test_ds), size=(num_images - len(torch.unique(indices)),)
                 )
                 indices = torch.cat((torch.unique(indices), new_indices))
+
+            # sort after getting indices
             indices, _ = torch.sort(indices)
+
         img_test = test_ds[indices][0]
         img_true = test_ds[indices][1]
         img_test = img_test[:, :, :65, :]

From 3597edbc5412e358885c3966f860a426e78f8fc6 Mon Sep 17 00:00:00 2001
From: Felix Geyer <felix.geyer@tu-dortmund.de>
Date: Tue, 31 Jan 2023 08:34:58 +0000
Subject: [PATCH 23/31] Remove more scaling

---
 radionets/evaluation/plotting.py | 5 -----
 1 file changed, 5 deletions(-)

diff --git a/radionets/evaluation/plotting.py b/radionets/evaluation/plotting.py
index 5de2a578..6a295cf2 100644
--- a/radionets/evaluation/plotting.py
+++ b/radionets/evaluation/plotting.py
@@ -208,11 +208,6 @@ def visualize_with_fourier(
     real_pred, imag_pred = img_pred[0], img_pred[1]
     real_truth, imag_truth = img_truth[0], img_truth[1]
 
-    if amp_phase:
-        inp_real = 10 ** (10 * inp_real - 10) - 1e-10
-        real_pred = 10 ** (10 * real_pred - 10) - 1e-10
-        real_truth = 10 ** (10 * real_truth - 10) - 1e-10
-
     # plotting
     fig, ((ax1, ax2, ax3), (ax4, ax5, ax6)) = plt.subplots(
         2, 3, figsize=(16, 10), sharex=True, sharey=True

From f62b7e99ae3f3aeafcc76e1e0b3b273dfefc99e8 Mon Sep 17 00:00:00 2001
From: Felix Geyer <felix.geyer@tu-dortmund.de>
Date: Tue, 31 Jan 2023 08:44:20 +0000
Subject: [PATCH 24/31] Write changes

---
 docs/changes/140.bugfix.rst      | 3 +++
 docs/changes/140.feature.rst     | 1 +
 docs/changes/141.maintenance.rst | 2 ++
 3 files changed, 6 insertions(+)
 create mode 100644 docs/changes/140.bugfix.rst
 create mode 100644 docs/changes/140.feature.rst
 create mode 100644 docs/changes/141.maintenance.rst

diff --git a/docs/changes/140.bugfix.rst b/docs/changes/140.bugfix.rst
new file mode 100644
index 00000000..39e39cae
--- /dev/null
+++ b/docs/changes/140.bugfix.rst
@@ -0,0 +1,3 @@
+fixed sampling of test data set
+fixed same indices for plots
+
diff --git a/docs/changes/140.feature.rst b/docs/changes/140.feature.rst
new file mode 100644
index 00000000..d77f5699
--- /dev/null
+++ b/docs/changes/140.feature.rst
@@ -0,0 +1 @@
+enabled training and evaluation of half sized images (for 128 pixel images)
diff --git a/docs/changes/141.maintenance.rst b/docs/changes/141.maintenance.rst
new file mode 100644
index 00000000..493a417e
--- /dev/null
+++ b/docs/changes/141.maintenance.rst
@@ -0,0 +1,2 @@
+Deleted unusable functions for new source types
+Deleted unused hardcoded scaling

From b6185f72ae08558c2c489359f91f348402beb778 Mon Sep 17 00:00:00 2001
From: Kevin Schmidt <Kevin2@users.noreply.github.com>
Date: Tue, 31 Jan 2023 09:44:46 +0100
Subject: [PATCH 25/31] Create 140.api.rst

---
 docs/changes/140.api.rst | 2 ++
 1 file changed, 2 insertions(+)
 create mode 100644 docs/changes/140.api.rst

diff --git a/docs/changes/140.api.rst b/docs/changes/140.api.rst
new file mode 100644
index 00000000..3774527d
--- /dev/null
+++ b/docs/changes/140.api.rst
@@ -0,0 +1,2 @@
+train on half-sized iamges and applying symmetry afterward is a backward incompatible change
+models trained with early versions of `radionets` are not supported anymore

From e947b3319ad348a3383ee04c0548482d9819bc05 Mon Sep 17 00:00:00 2001
From: Felix Geyer <felix.geyer@tu-dortmund.de>
Date: Tue, 31 Jan 2023 08:46:53 +0000
Subject: [PATCH 26/31] Fix typos

---
 docs/changes/140.bugfix.rst                               | 1 -
 docs/changes/{141.maintenance.rst => 140.maintenance.rst} | 0
 2 files changed, 1 deletion(-)
 rename docs/changes/{141.maintenance.rst => 140.maintenance.rst} (100%)

diff --git a/docs/changes/140.bugfix.rst b/docs/changes/140.bugfix.rst
index 39e39cae..ee1a0e36 100644
--- a/docs/changes/140.bugfix.rst
+++ b/docs/changes/140.bugfix.rst
@@ -1,3 +1,2 @@
 fixed sampling of test data set
 fixed same indices for plots
-
diff --git a/docs/changes/141.maintenance.rst b/docs/changes/140.maintenance.rst
similarity index 100%
rename from docs/changes/141.maintenance.rst
rename to docs/changes/140.maintenance.rst

From e5b3fcdf89e512623ca183492242dc4194ccbf78 Mon Sep 17 00:00:00 2001
From: Kevin Schmidt <kevin3.schmidt@tu-dortmund.de>
Date: Tue, 31 Jan 2023 09:52:40 +0100
Subject: [PATCH 27/31] fix changes overview

---
 CHANGES.rst                      | 35 ++++++++++++++++++++++++++++++++
 docs/changes/140.api.rst         |  2 --
 docs/changes/140.bugfix.rst      |  2 --
 docs/changes/140.feature.rst     |  1 -
 docs/changes/140.maintenance.rst |  2 --
 5 files changed, 35 insertions(+), 7 deletions(-)
 delete mode 100644 docs/changes/140.api.rst
 delete mode 100644 docs/changes/140.bugfix.rst
 delete mode 100644 docs/changes/140.feature.rst
 delete mode 100644 docs/changes/140.maintenance.rst

diff --git a/CHANGES.rst b/CHANGES.rst
index 73b363ab..f2fc2ba6 100644
--- a/CHANGES.rst
+++ b/CHANGES.rst
@@ -1,3 +1,38 @@
+Radionets 0.2.0 (2023-01-31)
+============================
+
+
+API Changes
+-----------
+
+- train on half-sized iamges and applying symmetry afterward is a backward incompatible change
+  models trained with early versions of `radionets` are not supported anymore [`#140 <https://github.com/radionets-project/radionets/pull/140>`__]
+
+
+Bug Fixes
+---------
+
+- fixed sampling of test data set
+  fixed same indices for plots [`#140 <https://github.com/radionets-project/radionets/pull/140>`__]
+
+
+New Features
+------------
+
+- enabled training and evaluation of half sized images (for 128 pixel images) [`#140 <https://github.com/radionets-project/radionets/pull/140>`__]
+
+
+Maintenance
+-----------
+
+- Deleted unusable functions for new source types
+  Deleted unused hardcoded scaling [`#140 <https://github.com/radionets-project/radionets/pull/140>`__]
+
+
+Refactoring and Optimization
+----------------------------
+
+
 Radionets 0.1.18 (2023-01-30)
 =============================
 
diff --git a/docs/changes/140.api.rst b/docs/changes/140.api.rst
deleted file mode 100644
index 3774527d..00000000
--- a/docs/changes/140.api.rst
+++ /dev/null
@@ -1,2 +0,0 @@
-train on half-sized iamges and applying symmetry afterward is a backward incompatible change
-models trained with early versions of `radionets` are not supported anymore
diff --git a/docs/changes/140.bugfix.rst b/docs/changes/140.bugfix.rst
deleted file mode 100644
index ee1a0e36..00000000
--- a/docs/changes/140.bugfix.rst
+++ /dev/null
@@ -1,2 +0,0 @@
-fixed sampling of test data set
-fixed same indices for plots
diff --git a/docs/changes/140.feature.rst b/docs/changes/140.feature.rst
deleted file mode 100644
index d77f5699..00000000
--- a/docs/changes/140.feature.rst
+++ /dev/null
@@ -1 +0,0 @@
-enabled training and evaluation of half sized images (for 128 pixel images)
diff --git a/docs/changes/140.maintenance.rst b/docs/changes/140.maintenance.rst
deleted file mode 100644
index 493a417e..00000000
--- a/docs/changes/140.maintenance.rst
+++ /dev/null
@@ -1,2 +0,0 @@
-Deleted unusable functions for new source types
-Deleted unused hardcoded scaling

From 6aa9b0a55dd3bebdb75706bb00f23f40163356da Mon Sep 17 00:00:00 2001
From: Kevin Schmidt <kevin3.schmidt@tu-dortmund.de>
Date: Tue, 31 Jan 2023 09:52:58 +0100
Subject: [PATCH 28/31] delete previous change files

---
 docs/changes/129.feature.rst     | 7 -------
 docs/changes/130.maintenance.rst | 1 -
 docs/changes/134.maintenance.rst | 1 -
 docs/changes/136.maintenance.rst | 1 -
 4 files changed, 10 deletions(-)
 delete mode 100644 docs/changes/129.feature.rst
 delete mode 100644 docs/changes/130.maintenance.rst
 delete mode 100644 docs/changes/134.maintenance.rst
 delete mode 100644 docs/changes/136.maintenance.rst

diff --git a/docs/changes/129.feature.rst b/docs/changes/129.feature.rst
deleted file mode 100644
index a469db05..00000000
--- a/docs/changes/129.feature.rst
+++ /dev/null
@@ -1,7 +0,0 @@
-added creation of uncertainty plots
-changed creation and saving/reading of predictions to ``dicts``
-prediction ``dicts`` have 3 or 4 entries depending on uncertainty
-added scaled option to ``get_ifft``
-created new dataset class for sampled images
-created option for sampling and saving the whole test dataset
-updated and wrote new tests
\ No newline at end of file
diff --git a/docs/changes/130.maintenance.rst b/docs/changes/130.maintenance.rst
deleted file mode 100644
index 77ffe76f..00000000
--- a/docs/changes/130.maintenance.rst
+++ /dev/null
@@ -1 +0,0 @@
-Add and enable ``towncrier`` in CI.
diff --git a/docs/changes/134.maintenance.rst b/docs/changes/134.maintenance.rst
deleted file mode 100644
index 9e561a7f..00000000
--- a/docs/changes/134.maintenance.rst
+++ /dev/null
@@ -1 +0,0 @@
-publish radionets on pypi
diff --git a/docs/changes/136.maintenance.rst b/docs/changes/136.maintenance.rst
deleted file mode 100644
index 4e73f7e3..00000000
--- a/docs/changes/136.maintenance.rst
+++ /dev/null
@@ -1 +0,0 @@
-Update README, use figures from the paper, minor text adjustments
\ No newline at end of file

From 1bcc95d18a583ab0c8f477bc96a2c89d06da0796 Mon Sep 17 00:00:00 2001
From: Kevin Schmidt <kevin3.schmidt@tu-dortmund.de>
Date: Tue, 31 Jan 2023 09:56:43 +0100
Subject: [PATCH 29/31] undo delete

---
 docs/changes/140.bugfixes.rst    | 2 ++
 docs/changes/140.feature.rst     | 1 +
 docs/changes/140.maintenance.rst | 2 ++
 3 files changed, 5 insertions(+)
 create mode 100644 docs/changes/140.bugfixes.rst
 create mode 100644 docs/changes/140.feature.rst
 create mode 100644 docs/changes/140.maintenance.rst

diff --git a/docs/changes/140.bugfixes.rst b/docs/changes/140.bugfixes.rst
new file mode 100644
index 00000000..ee1a0e36
--- /dev/null
+++ b/docs/changes/140.bugfixes.rst
@@ -0,0 +1,2 @@
+fixed sampling of test data set
+fixed same indices for plots
diff --git a/docs/changes/140.feature.rst b/docs/changes/140.feature.rst
new file mode 100644
index 00000000..d77f5699
--- /dev/null
+++ b/docs/changes/140.feature.rst
@@ -0,0 +1 @@
+enabled training and evaluation of half sized images (for 128 pixel images)
diff --git a/docs/changes/140.maintenance.rst b/docs/changes/140.maintenance.rst
new file mode 100644
index 00000000..493a417e
--- /dev/null
+++ b/docs/changes/140.maintenance.rst
@@ -0,0 +1,2 @@
+Deleted unusable functions for new source types
+Deleted unused hardcoded scaling

From 5d533141799b52d2dca555f46f227afcee27d332 Mon Sep 17 00:00:00 2001
From: Felix Geyer <felix.geyer@tu-dortmund.de>
Date: Wed, 1 Feb 2023 11:17:39 +0100
Subject: [PATCH 30/31] Documentation and click echo

---
 .../evaluation/scripts/start_evaluation.py    |  1 +
 radionets/evaluation/utils.py                 | 25 +++++++++++++++----
 2 files changed, 21 insertions(+), 5 deletions(-)

diff --git a/radionets/evaluation/scripts/start_evaluation.py b/radionets/evaluation/scripts/start_evaluation.py
index 24ec90d9..39b6a925 100644
--- a/radionets/evaluation/scripts/start_evaluation.py
+++ b/radionets/evaluation/scripts/start_evaluation.py
@@ -38,6 +38,7 @@ def main(configuration_path):
     print(eval_conf, "\n")
 
     if eval_conf["sample_unc"]:
+        click.echo("Sampling test data set.\n")
         save_sampled(eval_conf)
 
     for entry in conf["inspection"]:
diff --git a/radionets/evaluation/utils.py b/radionets/evaluation/utils.py
index 950ea176..d718dbae 100644
--- a/radionets/evaluation/utils.py
+++ b/radionets/evaluation/utils.py
@@ -564,9 +564,29 @@ def trunc_rvs(loc, scale, mode, num_samples, num_img):
 
 
 def sample_images(mean, std, num_samples):
+    """Samples for every pixel in Fourier space from a truncated Gaussian distribution
+    based on the output of the network.
+
+    Parameters
+    ----------
+    mean : torch.tensor
+        mean values of the pixels with shape (number of images, number of samples,
+        image size // 2 + 1, image_size)
+    std : torch.tensor
+        uncertainty values of the pixels with shape (number of images,
+        number of samples, image size // 2 + 1, image_size)
+    num_samples : int
+        number of samples in Fourier space
+
+    Returns
+    -------
+    dict
+        resulting mean and standard deviation
+    """
     mean_amp, mean_phase = mean[:, 0], mean[:, 1]
     std_amp, std_phase = std[:, 0], std[:, 1]
     num_img = mean_amp.shape[0]
+
     # amplitude
     sampled_gauss_amp = trunc_rvs(
         mean_amp, std_amp, "amp", num_samples, num_img
@@ -611,11 +631,6 @@ def mergeDictionary(dict_1, dict_2):
     for key, value in dict_3.items():
         if key in dict_1 and key in dict_2:
             dict_3[key] = np.append(dict_1[key], value)
-            # try:
-            #     dict_3[key] = np.stack((dict_1[key], value))
-            # except ValueError:
-            #     value = value.reshape(1, *value.shape)
-            #     dict_3[key] = np.concatenate((dict_1[key], value))
     return dict_3
 
 

From 888d43b2b9f7f460c58a1b5225ba703879ab3e61 Mon Sep 17 00:00:00 2001
From: Felix Geyer <felix.geyer@tu-dortmund.de>
Date: Wed, 1 Feb 2023 11:18:29 +0100
Subject: [PATCH 31/31] Update tests

---
 tests/test_evaluation.py | 72 ++++++++++++++++++++++++++++++----------
 1 file changed, 54 insertions(+), 18 deletions(-)

diff --git a/tests/test_evaluation.py b/tests/test_evaluation.py
index 67bd8986..b0606a80 100644
--- a/tests/test_evaluation.py
+++ b/tests/test_evaluation.py
@@ -335,34 +335,58 @@ def test_symmetry(self):
 
     def test_sample_images(self):
         import numpy as np
+        import torch
+        import torch.nn.functional as F
         from radionets.evaluation.utils import (
             read_pred,
             trunc_rvs,
-            even_better_symmetry,
             get_ifft,
+            sym_new,
         )
 
         num_samples = 100
+        num_img = 2
         img = read_pred("./tests/model/predictions_unc.h5")
-        mean = img["pred"][0]
-        std = img["unc"][0]
-        mean_amp, mean_phase = mean[0], mean[1]
-        std_amp, std_phase = std[0], std[1]
+        mean_amp, mean_phase = (
+            img["pred"][:num_img, 0, :65, :],
+            img["pred"][:num_img, 1, :65, :],
+        )
+        std_amp, std_phase = (
+            img["unc"][:num_img, 0, :65, :],
+            img["unc"][:num_img, 1, :65, :],
+        )
         img_size = mean_amp.shape[-1]
 
         # amplitude
-        sampled_gauss_amp = trunc_rvs(mean_amp, std_amp, "amp", num_samples, num_img=1)
+        sampled_gauss_amp = trunc_rvs(
+            mean_amp, std_amp, "amp", num_samples, num_img=num_img
+        )
 
         # phase
         sampled_gauss_phase = trunc_rvs(
-            mean_phase, std_phase, "phase", num_samples, num_img=1
+            mean_phase, std_phase, "phase", num_samples, num_img=num_img
         )
 
-        assert sampled_gauss_amp.shape == (1, num_samples, img_size, img_size)
-        assert sampled_gauss_phase.shape == (1, num_samples, img_size, img_size)
+        assert sampled_gauss_amp.shape == (
+            num_img,
+            num_samples,
+            img_size // 2 + 1,
+            img_size,
+        )
+        assert sampled_gauss_phase.shape == (
+            num_img,
+            num_samples,
+            img_size // 2 + 1,
+            img_size,
+        )
+
+        sampled_gauss_amp = sampled_gauss_amp.reshape(num_img * num_samples, 65, 128)
+        sampled_gauss_phase = sampled_gauss_phase.reshape(
+            num_img * num_samples, 65, 128
+        )
 
         with pytest.raises(ValueError):
-            trunc_rvs(mean_phase, std_phase, "pase", num_samples, num_img=1)
+            trunc_rvs(mean_phase, std_phase, "pase", num_samples, num_img=num_img)
 
         # masks
         mask_invalid_amp = sampled_gauss_amp < 0
@@ -373,24 +397,36 @@ def test_sample_images(self):
         assert mask_invalid_phase.sum() == 0
 
         sampled_gauss = np.stack([sampled_gauss_amp, sampled_gauss_phase], axis=1)
-        sampled_gauss_symmetry = even_better_symmetry(sampled_gauss)
+
+        # pad resulting images and utilize symmetry
+        sampled_gauss = F.pad(
+            input=torch.tensor(sampled_gauss),
+            pad=(0, 0, 0, 63),
+            mode="constant",
+            value=0,
+        )
+        sampled_gauss_symmetry = sym_new(sampled_gauss, None)
 
         fft_sampled_symmetry = get_ifft(
             sampled_gauss_symmetry, amp_phase=True, scale=False
-        )
+        ).reshape(num_img, num_samples, 128, 128)
+
         results = {
-            "mean": fft_sampled_symmetry.mean(axis=0),
-            "std": fft_sampled_symmetry.std(axis=0),
+            "mean": fft_sampled_symmetry.mean(axis=1),
+            "std": fft_sampled_symmetry.std(axis=1),
         }
-        assert results["mean"].shape == (img_size, img_size)
-        assert results["std"].shape == (img_size, img_size)
+        assert results["mean"].shape == (num_img, img_size, img_size)
+        assert results["std"].shape == (num_img, img_size, img_size)
 
     def test_uncertainty_plots(self):
         from radionets.evaluation.utils import read_pred, sample_images
 
         img = read_pred("./tests/model/predictions_unc.h5")
-        results = sample_images(img["pred"][0:2], img["unc"][0:2], 100)
-        assert results is not None
+        results = sample_images(
+            img["pred"][:2, :, :65, :], img["unc"][:2, :, :65, :], 100
+        )
+        assert results["mean"].shape == (2, 128, 128)
+        assert results["std"].shape == (2, 128, 128)
 
     def test_evaluation(self):
         import shutil