WIP

simple_eta/generate.py

@@ -15,20 +15,27 @@ def gaussian_2d(mx, my, sigma = 1, res=100, pixel_size = 25, grid_size = 2):
 
 
 def sum_pixels(t, grid):
-    resolution = t.shape[0]
-    pixels = np.zeros((grid,grid))
+    if t.ndim == 2:
+        resolution = t.shape[0]
+        pixels = np.zeros((grid,grid))
+    
+        step = resolution//grid
+        for i in range(grid):
+            for j in range(grid):
+                # print(f'[{i*step}:{(i+1)*step}, {j*step}:{(j+1)*step}]')
+                pixels[i,j] = t[i*step:(i+1)*step, j*step:(j+1)*step].sum()
+        return pixels
 
-    step = resolution//grid
-    for i in range(grid):
-        for j in range(grid):
-            print(f'[{i*step}:{(i+1)*step}, {j*step}:{(j+1)*step}]')
-            pixels[i,j] = t[i*step:(i+1)*step, j*step:(j+1)*step].sum()
+    elif t.ndim == 3:
+        resolution = t.shape[1]
+        step = resolution//grid
+        pixels = np.zeros((t.shape[0], grid, grid))
+        for i in range(grid):
+            for j in range(grid):
+                # print(f'[{i*step}:{(i+1)*step}, {j*step}:{(j+1)*step}]')
+                pixels[:,i,j] = t[:,i*step:(i+1)*step, j*step:(j+1)*step].sum(axis = 1).sum(axis = 1)
 
-    # pixels[0,0] = t[0:resolution//2, 0:resolution//2].sum()
-    # pixels[0,1] = t[0:resolution//2, resolution//2:].sum()
-    # pixels[1,0] = t[resolution//2:, 0:resolution//2].sum()
-    # pixels[1,1] = t[resolution//2:, resolution//2:].sum()
-    return pixels
+        return pixels
 
 def generate_uniform_hits(sigma, pixel_size, grid_size, resolution, N=100, device = 'cpu'):
     """
@@ -42,16 +49,16 @@ def generate_uniform_hits(sigma, pixel_size, grid_size, resolution, N=100, devic
     ys = y.unsqueeze(0).repeat(N,1,1)
 
 
-    #TODO! The genreal cases are actually odd and even
+    
     if grid_size % 2 == 0:
+        #For even grids spread hits over the inner quadrants around the center
         low = (grid_size-1)//2*pixel_size+pixel_size/2
         high =pixel_size+low
     else:
+        #For odd grids spread hits over the central pixel
         low = pixel_size*(grid_size//2)
         high = low+pixel_size
 
-    print(low, high)
-
     mx = torch.rand(N,1,1, device = device) * (high-low)+low
     my = torch.rand(N,1,1, device = device) * (high-low) +low
 

simple_eta/plotting.py

@@ -1,19 +1,25 @@
 import matplotlib.pyplot as plt
 import numpy as np
 
-def plot_gaussian(t, pixel_size, grid_size):
+def plot_gaussian(t, pixel_size, grid_size, ax = None):
     print(f'{t.shape=}')
 
     resolution = t.shape[0]
     xa = np.linspace(0,grid_size*pixel_size,resolution)
     ticks = [tick for tick in range(0,pixel_size*grid_size+1, pixel_size)]
 
-    fig, ax = plt.subplots(figsize = (7,7))
+    if ax is None:
+        fig, ax = plt.subplots(figsize = (7,7))
+    else:
+        fig = None
+
+        
     ax.pcolormesh(xa,xa, t)
     ax.set_xticks(ticks)
     ax.set_yticks(ticks)
     ax.set_xlim(0,grid_size*pixel_size)
     ax.set_ylim(0,grid_size*pixel_size)
+    ax.set_aspect('equal')
     
     ax.grid()