MLXID/simple_eta
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

saving ground truth

Browse Source
This commit is contained in:
Erik Fröjdh
2025-11-20 22:42:04 +01:00
parent d11cd36335
commit 3f73549d96
3 changed files with 174 additions and 118 deletions
Download Patch File Download Diff File Expand all files Collapse all files

178
Vis.ipynb
View File

File diff suppressed because one or more lines are too long

2
simple_eta/__init__.py
View File

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

112
simple_eta/generate.py
View File

@@ -2,19 +2,11 @@ import torch
import math
import numpy as np
def gaussian_2d(mx, my, sigma = 1, res=100, pixel_size = 25, grid_size = 2):
"""
Generate a 2D gaussian as position mx, my, with sigma=sigma.
The gaussian is placed on a grid_size x grid_size pixel matrix with resolution
res in one dimesion.
"""
x = torch.linspace(0, pixel_size*grid_size, res)
x,y = torch.meshgrid(x,x, indexing="ij")
return 1 / (2*math.pi*sigma**2) * \
torch.exp(-((x - my)**2 / (2*sigma**2) + (y - mx)**2 / (2*sigma**2)))
def sum_pixels(t, grid):
"""
Given a charge density as a torch array sum it to pixels
"""
if t.ndim == 2:
resolution = t.shape[0]
pixels = np.zeros((grid,grid))
@@ -45,14 +37,25 @@ class Generator:
self.device = device
def hit(self, mx, my):
x = torch.linspace(0, self.pixel_size*self.grid_size, self.resolution)
x = torch.linspace(0, self.pixel_size*self.grid_size,
self.resolution, device = self.device)
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)
if self.device != 'cpu':
t = t.cpu()
return t, p
def dx(self):
"""
Return the step for normalizing of the gaussian
"""
x = torch.linspace(0, self.pixel_size*self.grid_size,
self.resolution, device = self.device)
return x[1]-x[0]
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")
@@ -126,87 +129,6 @@ class Generator:
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. 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) * (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