class for generation

simple_eta/__init__.py

@@ -1,2 +1,2 @@
-from .generate import gaussian_2d
+from .generate import gaussian_2d, Generator
 from .plotting import plot_gaussian

simple_eta/generate.py

@@ -22,7 +22,6 @@ def sum_pixels(t, 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
 
@@ -32,15 +31,151 @@ def sum_pixels(t, 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)
 
         return pixels
 
+
+class Generator:
+    def __init__(self, sigma, pixel_size, grid_size, resolution, device = 'cpu'):
+        self.sigma = sigma
+        self.pixel_size = pixel_size
+        self.grid_size = grid_size
+        self.resolution = resolution
+        self.device = device
+
+    def hit(self, mx, my):
+        x = torch.linspace(0, self.pixel_size*self.grid_size, self.resolution)
+        x,y = torch.meshgrid(x,x, indexing="ij")
+        t  = 1 / (2*math.pi*self.sigma**2) * \
+          torch.exp(-((x - my)**2 / (2*self.sigma**2) + (y - mx)**2 / (2*self.sigma**2)))
+
+        p = sum_pixels(t, self.grid_size)
+        return t, p
+
+    def uniform_hits(self, n_hits):
+        x = torch.linspace(0, self.pixel_size*self.grid_size, self.resolution, device = self.device)
+        x,y = torch.meshgrid(x,x, indexing="ij")
+        xs = x.unsqueeze(0).repeat(n_hits,1,1) 
+        ys = y.unsqueeze(0).repeat(n_hits,1,1)
+    
+    
+        
+        if self.grid_size % 2 == 0:
+            #For even grids spread hits over the inner quadrants around the center
+            low = (self.grid_size-1)//2*self.pixel_size+self.pixel_size/2
+            high =self.pixel_size+low
+        else:
+            #For odd grids spread hits over the central pixel
+            low = self.pixel_size*(self.grid_size//2)
+            high = low+self.pixel_size
+    
+        mx = torch.rand(n_hits,1,1, device = self.device) * (high-low)+low
+        my = torch.rand(n_hits,1,1, device = self.device) * (high-low) +low
+    
+        ts = 1 / (2*math.pi*self.sigma**2) * \
+          torch.exp(-((xs - my)**2 / (2*self.sigma**2) + (ys - mx)**2 / (2*self.sigma**2)))
+        
+        
+        #Sum signal in pixels for all N depositions
+        step = self.resolution//self.grid_size
+        pixels = torch.zeros((n_hits,self.grid_size,self.grid_size))
+        for i in range(self.grid_size):
+            for j in range(self.grid_size):
+                pixels[:,i,j] = ts[:,i*step:(i+1)*step, j*step:(j+1)*step].sum(axis = 1).sum(axis = 1)
+        
+        return mx, my, pixels
+
+    def triangle_hits(self, n_hits):
+        x = torch.linspace(0, self.pixel_size*self.grid_size, self.resolution, device = self.device)
+        x,y = torch.meshgrid(x,x, indexing="ij")
+        xs = x.unsqueeze(0).repeat(n_hits,1,1) 
+        ys = y.unsqueeze(0).repeat(n_hits,1,1)
+    
+    
+        
+        if self.grid_size % 2 == 0:
+            #For even grids spread hits over the inner quadrants around the center
+            low = (self.grid_size-1)//2*self.pixel_size+self.pixel_size/2
+            high =self.pixel_size+low
+        else:
+            #For odd grids spread hits over the central pixel
+            low = self.pixel_size*(self.grid_size//2)
+            high = low+self.pixel_size
+    
+        # mx = torch.rand(n_hits,1,1, device = self.device) * (high-low)+low
+        # my = torch.rand(n_hits,1,1, device = self.device) * (high-low) +low
+
+        mx = torch.rand(n_hits,1,1, device = self.device)
+        my = torch.rand(n_hits,1,1, device = self.device)
+        mask = (mx + my > 1)
+        mx[mask] = 1 - mx[mask]
+        my[mask] = 1 - my[mask]
+        mx = mx * (high-low) + low
+        my = my * (high-low) + low
+    
+        ts = 1 / (2*math.pi*self.sigma**2) * \
+          torch.exp(-((xs - my)**2 / (2*self.sigma**2) + (ys - mx)**2 / (2*self.sigma**2)))
+        
+        
+        #Sum signal in pixels for all N depositions
+        step = self.resolution//self.grid_size
+        pixels = torch.zeros((n_hits,self.grid_size,self.grid_size))
+        for i in range(self.grid_size):
+            for j in range(self.grid_size):
+                pixels[:,i,j] = ts[:,i*step:(i+1)*step, j*step:(j+1)*step].sum(axis = 1).sum(axis = 1)
+        
+        return mx, my, pixels
+def generate_triangle_hits(sigma, pixel_size, grid_size, resolution, N=100, device = 'cpu'):
+    """
+    Generate N gaussians. For even grid size use the four inner corners of the central pixel
+    For odd corners the whole central pixel
+    """
+    
+    x = torch.linspace(0, pixel_size*grid_size, resolution, device = device)
+    x,y = torch.meshgrid(x,x, indexing="ij")
+    xs = x.unsqueeze(0).repeat(N,1,1) 
+    ys = y.unsqueeze(0).repeat(N,1,1)
+
+
+    
+    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
+
+    mx = torch.rand(N,1,1, device = device)
+    my = torch.rand(N,1,1, device = device)
+    mask = (mx + my > 1)
+    mx[mask] = 1 - mx[mask]
+    my[mask] = 1 - my[mask]
+    mx = mx * (high-low) + low
+    my = my * (high-low) + low
+    
+    # mx = torch.rand(N,1,1, device = device) * (high-low)+low
+    # my = torch.rand(N,1,1, device = device) * (high-low) +low
+
+    ts = 1 / (2*math.pi*sigma**2) * \
+      torch.exp(-((xs - my)**2 / (2*sigma**2) + (ys - mx)**2 / (2*sigma**2)))
+    
+    
+    #Sum signal in pixels for all N depositions
+    step = resolution//grid_size
+    pixels = torch.zeros((N,grid_size,grid_size))
+    for i in range(grid_size):
+        for j in range(grid_size):
+            pixels[:,i,j] = ts[:,i*step:(i+1)*step, j*step:(j+1)*step].sum(axis = 1).sum(axis = 1)
+    
+    return mx, my, pixels
+
 def generate_uniform_hits(sigma, pixel_size, grid_size, resolution, N=100, device = 'cpu'):
     """
-    Generate N gaussians with sigma. Uniformly distributed over
-    the 2x2 pixel grid. 
+    Generate N gaussians. For even grid size use the four inner corners of the central pixel
+    For odd corners the whole central pixel
     """
     
     x = torch.linspace(0, pixel_size*grid_size, resolution, device = device)

simple_eta/plotting.py

@@ -14,7 +14,7 @@ def plot_gaussian(t, pixel_size, grid_size, ax = None):
         fig = None
 
         
-    ax.pcolormesh(xa,xa, t)
+    mesh = ax.pcolormesh(xa,xa, t)
     ax.set_xticks(ticks)
     ax.set_yticks(ticks)
     ax.set_xlim(0,grid_size*pixel_size)
@@ -25,4 +25,4 @@ def plot_gaussian(t, pixel_size, grid_size, ax = None):
 
     ax.set_xlabel(r'Position x [$\mu$m]')
     ax.set_ylabel(r'Position y [$\mu$m]')
-    return fig, ax 
+    return fig, ax, mesh