def setUp(self):
self.cache_dir = random_tempdir('chroma_cache_test')
self.cache = Cache(self.cache_dir)
self.a = Geometry()
self.a.add_solid(Solid(box(1,1,1)))
self.a.add_solid(Solid(box(1,1,1)), displacement=(10,10,10))
self.a.flatten()
self.b = Geometry()
self.b.add_solid(Solid(box(2,2,2)))
self.b.add_solid(Solid(box(2,2,2)), displacement=(10,10,10))
self.b.add_solid(Solid(box(2,2,2)), displacement=(-10,-10,-10))
self.b.flatten()
# c is not in cache
self.c = Geometry()
self.c.add_solid(Solid(box(2,2,2)))
self.c.flatten()
self.a_hash = self.a.mesh.md5()
self.b_hash = self.b.mesh.md5()
self.c_hash = self.c.mesh.md5()
self.cache.save_geometry('a', self.a)
self.cache.save_geometry('b', self.b)
def render_particle_track(self):
x = 10.0
h = x * np.sqrt(3) / 2
pyramid = make.linear_extrude([-x / 2, 0, x / 2],
[-h / 2, h / 2, -h / 2], h, [0] * 3,
[0] * 3)
marker = Solid(pyramid, vacuum, vacuum)
if self.photon_display_mode == 'beg':
photons = self.ev.photons_beg
else:
photons = self.ev.photons_end
geometry = Geometry()
sample_factor = max(1, len(photons.pos) / 10000)
for pos in photons.pos[::sample_factor]:
geometry.add_solid(marker,
displacement=pos,
rotation=make_rotation_matrix(
np.random.uniform(0, 2 * np.pi),
uniform_sphere()))
geometry = create_geometry_from_obj(geometry)
gpu_geometry = gpu.GPUGeometry(geometry)
self.gpu_geometries = [self.gpu_geometry, gpu_geometry]
def build_detector():
"""Returns a cubic detector made of cubic photodetectors."""
g = Geometry(water)
for pos, dir in iter_box(nx,ny,nz,spacing):
# convert to arrays
pos, dir = np.array(pos), np.array(dir)
# we need to figure out what rotation matrix to apply to each
# photodetector so that the photosensitive surface will be facing
# `dir`.
if tuple(dir) == (0,0,1):
rotation = None
elif tuple(dir) == (0,0,-1):
rotation = make_rotation_matrix(np.pi,(1,0,0))
else:
rotation = make_rotation_matrix(np.arccos(dir.dot((0,0,1))),
np.cross(dir,(0,0,1)))
# add the photodetector
g.add_solid(build_pd(size,glass_thickness),rotation=rotation,
displacement=pos)
world = Solid(make.box(spacing*nx,spacing*ny,spacing*nz),water,vacuum,
color=0x33ffffff)
g.add_solid(world)
return g
def flatten(self):
self.solid_id_to_channel_index = np.asarray(
self.solid_id_to_channel_index, dtype=np.int32)
self.channel_index_to_solid_id = np.asarray(
self.channel_index_to_solid_id, dtype=np.int32)
self.channel_index_to_channel_id = np.asarray(
self.channel_index_to_channel_id, dtype=np.int32)
Geometry.flatten(self)
def testAbort(self):
'''Photons that hit a triangle at normal incidence should not abort.
Photons that hit a triangle at exactly normal incidence can sometimes
produce a dot product that is outside the range allowed by acos().
Trigger these with axis aligned photons in a box.
'''
# Setup geometry
cube = Geometry(vacuum)
cube.add_solid(Solid(box(100, 100, 100), vacuum, vacuum))
geo = create_geometry_from_obj(cube, update_bvh_cache=False)
# Initialize simulation (without geant4)
sim = Simulation(geo, geant4_processes=0)
# Create initial photons
nphotons = 5
pos = np.tile([0, 0, 0], (nphotons, 1)).astype(np.float32)
dir = np.tile([0, 0, 1], (nphotons, 1)).astype(np.float32)
pol = np.zeros_like(pos)
phi = np.random.uniform(0, 2 * np.pi, nphotons).astype(np.float32)
pol[:, 0] = np.cos(phi)
pol[:, 1] = np.sin(phi)
t = np.zeros(nphotons, dtype=np.float32)
wavelengths = np.empty(nphotons, np.float32)
wavelengths.fill(400.0)
photons = Photons(pos=pos,
dir=dir,
pol=pol,
t=t,
wavelengths=wavelengths)
# First make one step to check for strangeness
photons_end = sim.simulate([photons],
keep_photons_end=True,
max_steps=1).next().photons_end
print "FIRST STEP"
print photons_end.pos[0:10]
self.assertFalse(np.isnan(photons_end.pos).any())
self.assertFalse(np.isnan(photons_end.dir).any())
self.assertFalse(np.isnan(photons_end.pol).any())
self.assertFalse(np.isnan(photons_end.t).any())
self.assertFalse(np.isnan(photons_end.wavelengths).any())
# Now let it run the usual ten steps
photons_end = sim.simulate([photons],
keep_photons_end=True,
max_steps=10).next().photons_end
aborted = (photons_end.flags & (1 << 31)) > 0
print 'aborted photon fracttion: %1.1f' % (
float(count_nonzero(aborted)) / nphotons)
self.assertFalse(aborted.any())
print "LAST STEPS"
print photons_end.pos[0:10]
def create_geometry_from_obj(obj, bvh_name="default",
auto_build_bvh=True, read_bvh_cache=True,
update_bvh_cache=True, cache_dir=None,
cuda_device=None):
if callable(obj):
obj = obj()
if isinstance(obj, Detector):
geometry = obj
if isinstance(obj, Geometry):
geometry = obj
elif isinstance(obj, Solid):
geometry = Geometry()
geometry.add_solid(obj)
elif isinstance(obj, Mesh):
geometry = Geometry()
geometry.add_solid(Solid(obj, vacuum, vacuum, color=0x33ffffff))
else:
raise TypeError('cannot build type %s' % type(obj))
geometry.flatten()
if geometry.bvh is None:
geometry.bvh = load_bvh(geometry, auto_build_bvh=auto_build_bvh,
read_bvh_cache=read_bvh_cache,
update_bvh_cache=update_bvh_cache,
cache_dir=cache_dir,
cuda_device=cuda_device)
return geometry
def render_mc_info(self,max_photons=1000,cher_only=True):
#need to render photon tracking info if available
self.gpu_geometries = [self.gpu_geometry]
if self.sum_mode or self.ev is None:
return
if self.photon_display_mode == 'beg':
photons = self.ev.photons_beg
elif self.photon_display_mode == 'end':
photons = self.ev.photons_end
else:
photons = None
if photons is not None:
geometry = Geometry()
self.render_photons(geometry,photons)
geometry = create_geometry_from_obj(geometry)
gpu_geometry = gpu.GPUGeometry(geometry)
self.gpu_geometries.append(gpu_geometry)
if self.track_display_mode in ['geant4', 'both'] and self.ev.vertices is not None:
geometry = Geometry()
any = False
for vertex in self.ev.vertices:
if vertex.steps:
any = True
self.render_vertex(geometry,vertex,children=True)
if any:
geometry = create_geometry_from_obj(geometry)
gpu_geometry = gpu.GPUGeometry(geometry)
self.gpu_geometries.append(gpu_geometry)
if self.track_display_mode in ['chroma', 'both'] and self.ev.photon_tracks is not None:
geometry = Geometry()
print('Total Photons',len(self.ev.photon_tracks))
cherenkov = np.asarray([track.flags[0] & event.CHERENKOV == event.CHERENKOV for track in self.ev.photon_tracks])
ncherenkov = np.count_nonzero(cherenkov)
nphotons = ncherenkov if cher_only else len(self.ev.photon_tracks)
prob = max_photons/nphotons
selector = np.random.random(len(self.ev.photon_tracks)) < prob
nphotons = 0
for track in (t for s,t in zip(selector,self.ev.photon_tracks) if s):
if cher_only and track.flags[0] & event.CHERENKOV == event.CHERENKOV:
self.render_photon_track(geometry,track[:min(len(track),5)],sz=0.05)
nphotons = nphotons + 1
elif not cher_only:
self.render_photon_track(geometry,track[:min(len(track),5)])
nphotons = nphotons + 1
if nphotons > 0:
print('Rendered Photons',nphotons)
geometry = create_geometry_from_obj(geometry)
gpu_geometry = gpu.GPUGeometry(geometry)
self.gpu_geometries.append(gpu_geometry)
def flatten(self):
# Using numpy arrays here to allow for fancy indexing
self.solid_id_to_channel_index = np.asarray(
self.solid_id_to_channel_index, dtype=np.int32)
self.channel_index_to_solid_id = np.asarray(
self.channel_index_to_solid_id, dtype=np.int32)
self.channel_index_to_channel_type = np.asarray(
self.channel_index_to_channel_type, dtype=np.int32)
self.channel_index_to_position = np.asarray(
self.channel_index_to_position, dtype=np.int32)
Geometry.flatten(self)
def __init__(self, detector_material=None):
Geometry.__init__(self, detector_material=detector_material)
self.solid_id_to_channel_index = []
self.channel_index_to_solid_id = []
self.channel_index_to_channel_type = []
self.channel_index_to_position = []
# zero time and unit charge distributions
self.time_cdf = (np.array([-0.00000001,
0.00000001]), np.array([0.0, 1.0]))
self.charge_cdf = (np.array([0.999999999,
1.00000000]), np.array([0.0, 1.0]))
def setUp(self):
self.cache_dir = random_tempdir('chroma_cache_test')
self.cache = Cache(self.cache_dir)
self.a = Geometry()
self.a.add_solid(Solid(box(1, 1, 1)))
self.a.add_solid(Solid(box(1, 1, 1)), displacement=(10, 10, 10))
self.a.flatten()
self.b = Geometry()
self.b.add_solid(Solid(box(2, 2, 2)))
self.b.add_solid(Solid(box(2, 2, 2)), displacement=(10, 10, 10))
self.b.add_solid(Solid(box(2, 2, 2)), displacement=(-10, -10, -10))
self.b.flatten()
def setUp(self):
self.cache_dir = random_tempdir('chroma_cache_test')
self.cache = Cache(self.cache_dir)
self.a = Geometry()
self.a.add_solid(Solid(box(1,1,1)))
self.a.add_solid(Solid(box(1,1,1)), displacement=(10,10,10))
self.a.flatten()
self.b = Geometry()
self.b.add_solid(Solid(box(2,2,2)))
self.b.add_solid(Solid(box(2,2,2)), displacement=(10,10,10))
self.b.add_solid(Solid(box(2,2,2)), displacement=(-10,-10,-10))
self.b.flatten()
class TestRayleigh(unittest.TestCase):
def setUp(self):
self.cube = Geometry(water)
self.cube.add_solid(Solid(box(100, 100, 100), water, water))
self.geo = create_geometry_from_obj(self.cube, update_bvh_cache=False)
self.sim = Simulation(self.geo, geant4_processes=0)
nphotons = 100000
pos = np.tile([0, 0, 0], (nphotons, 1)).astype(np.float32)
dir = np.tile([0, 0, 1], (nphotons, 1)).astype(np.float32)
pol = np.zeros_like(pos)
phi = np.random.uniform(0, 2 * np.pi, nphotons).astype(np.float32)
pol[:, 0] = np.cos(phi)
pol[:, 1] = np.sin(phi)
t = np.zeros(nphotons, dtype=np.float32)
wavelengths = np.empty(nphotons, np.float32)
wavelengths.fill(400.0)
self.photons = Photons(pos=pos,
dir=dir,
pol=pol,
t=t,
wavelengths=wavelengths)
def testAngularDistributionPolarized(self):
# Fully polarized photons
self.photons.pol[:] = [1.0, 0.0, 0.0]
photons_end = self.sim.simulate([self.photons],
keep_photons_end=True,
max_steps=1).next().photons_end
aborted = (photons_end.flags & (1 << 31)) > 0
self.assertFalse(aborted.any())
# Compute the dot product between initial and final dir
rayleigh_scatters = (photons_end.flags & (1 << 4)) > 0
cos_scatter = (self.photons.dir[rayleigh_scatters] *
photons_end.dir[rayleigh_scatters]).sum(axis=1)
theta_scatter = np.arccos(cos_scatter)
h = Histogram(bins=100, range=(0, np.pi))
h.fill(theta_scatter)
h = rootify(h)
# The functional form for polarized light should be
# (1 + \cos^2 \theta)\sin \theta according to GEANT4 physics
# reference manual.
f = ROOT.TF1("pol_func", "[0]*(1+cos(x)**2)*sin(x)", 0, np.pi)
h.Fit(f, 'NQ')
self.assertGreater(f.GetProb(), 1e-3)
def __init__(self, detector_material=None):
Geometry.__init__(self, detector_material=detector_material)
# Using numpy arrays here to allow for fancy indexing
self.solid_id_to_channel_index = np.zeros(0, dtype=np.int32)
self.channel_index_to_solid_id = np.zeros(0, dtype=np.int32)
self.channel_index_to_channel_id = np.zeros(0, dtype=np.int32)
# If the ID numbers are arbitrary, we can't treat them
# as array indices, so have to use a dictionary
self.channel_id_to_channel_index = {}
# zero time and unit charge distributions
self.time_cdf = (np.array([-0.00000001, 0.00000001]), np.array([0.0, 1.0]))
self.charge_cdf = (np.array([0.999999999, 1.00000000]), np.array([0.0, 1.0]))
def __init__(self, detector_material=None):
Geometry.__init__(self, detector_material=detector_material)
# Using numpy arrays here to allow for fancy indexing
self.solid_id_to_channel_index = []
self.channel_index_to_solid_id = []
self.channel_index_to_channel_id = []
# If the ID numbers are arbitrary, we can't treat them
# as array indices, so have to use a dictionary
self.channel_id_to_channel_index = {}
# zero time and unit charge distributions
self.time_cdf = (np.array([-0.00000001,
0.00000001]), np.array([0.0, 1.0]))
self.charge_cdf = (np.array([0.999999999,
1.00000000]), np.array([0.0, 1.0]))
def testAbort(self):
'''Photons that hit a triangle at normal incidence should not abort.
Photons that hit a triangle at exactly normal incidence can sometimes
produce a dot product that is outside the range allowed by acos().
Trigger these with axis aligned photons in a box.
'''
# Setup geometry
cube = Geometry(vacuum)
cube.add_solid(Solid(box(100,100,100), vacuum, vacuum))
geo = create_geometry_from_obj(cube, update_bvh_cache=False)
sim = Simulation(geo, geant4_processes=0)
# Create initial photons
nphotons = 10000
pos = np.tile([0,0,0], (nphotons,1)).astype(np.float32)
dir = np.tile([0,0,1], (nphotons,1)).astype(np.float32)
pol = np.zeros_like(pos)
phi = np.random.uniform(0, 2*np.pi, nphotons).astype(np.float32)
pol[:,0] = np.cos(phi)
pol[:,1] = np.sin(phi)
t = np.zeros(nphotons, dtype=np.float32)
wavelengths = np.empty(nphotons, np.float32)
wavelengths.fill(400.0)
photons = Photons(pos=pos, dir=dir, pol=pol, t=t,
wavelengths=wavelengths)
# First make one step to check for strangeness
photons_end = sim.simulate([photons], keep_photons_end=True,
max_steps=1).next().photons_end
self.assertFalse(np.isnan(photons_end.pos).any())
self.assertFalse(np.isnan(photons_end.dir).any())
self.assertFalse(np.isnan(photons_end.pol).any())
self.assertFalse(np.isnan(photons_end.t).any())
self.assertFalse(np.isnan(photons_end.wavelengths).any())
# Now let it run the usual ten steps
photons_end = sim.simulate([photons], keep_photons_end=True,
max_steps=10).next().photons_end
aborted = (photons_end.flags & (1 << 31)) > 0
print 'aborted photons: %1.1f' % \
(float(count_nonzero(aborted)) / nphotons)
self.assertFalse(aborted.any())
def setUp(self):
self.cube = Geometry(water)
self.cube.add_solid(Solid(box(100,100,100), water, water))
self.geo = create_geometry_from_obj(self.cube, update_bvh_cache=False)
self.sim = Simulation(self.geo, geant4_processes=0)
nphotons = 100000
pos = np.tile([0,0,0], (nphotons,1)).astype(np.float32)
dir = np.tile([0,0,1], (nphotons,1)).astype(np.float32)
pol = np.zeros_like(pos)
phi = np.random.uniform(0, 2*np.pi, nphotons).astype(np.float32)
pol[:,0] = np.cos(phi)
pol[:,1] = np.sin(phi)
t = np.zeros(nphotons, dtype=np.float32)
wavelengths = np.empty(nphotons, np.float32)
wavelengths.fill(400.0)
self.photons = Photons(pos=pos, dir=dir, pol=pol, t=t, wavelengths=wavelengths)
class TestRayleigh(unittest.TestCase):
def setUp(self):
self.cube = Geometry(water)
self.cube.add_solid(Solid(box(100,100,100), water, water))
self.geo = create_geometry_from_obj(self.cube, update_bvh_cache=False)
self.sim = Simulation(self.geo, geant4_processes=0)
nphotons = 100000
pos = np.tile([0,0,0], (nphotons,1)).astype(np.float32)
dir = np.tile([0,0,1], (nphotons,1)).astype(np.float32)
pol = np.zeros_like(pos)
phi = np.random.uniform(0, 2*np.pi, nphotons).astype(np.float32)
pol[:,0] = np.cos(phi)
pol[:,1] = np.sin(phi)
t = np.zeros(nphotons, dtype=np.float32)
wavelengths = np.empty(nphotons, np.float32)
wavelengths.fill(400.0)
self.photons = Photons(pos=pos, dir=dir, pol=pol, t=t, wavelengths=wavelengths)
def testAngularDistributionPolarized(self):
# Fully polarized photons
self.photons.pol[:] = [1.0, 0.0, 0.0]
photons_end = self.sim.simulate([self.photons], keep_photons_end=True, max_steps=1).next().photons_end
aborted = (photons_end.flags & (1 << 31)) > 0
self.assertFalse(aborted.any())
# Compute the dot product between initial and final dir
rayleigh_scatters = (photons_end.flags & (1 << 4)) > 0
cos_scatter = (self.photons.dir[rayleigh_scatters] * photons_end.dir[rayleigh_scatters]).sum(axis=1)
theta_scatter = np.arccos(cos_scatter)
h = Histogram(bins=100, range=(0, np.pi))
h.fill(theta_scatter)
h = rootify(h)
# The functional form for polarized light should be
# (1 + \cos^2 \theta)\sin \theta according to GEANT4 physics
# reference manual.
f = ROOT.TF1("pol_func", "[0]*(1+cos(x)**2)*sin(x)", 0, np.pi)
h.Fit(f, 'NQ')
self.assertGreater(f.GetProb(), 1e-3)
def add_solid(self, solid, rotation=None, displacement=None):
"""
Add the solid `solid` to the geometry. When building the final triangle
mesh, `solid` will be placed by rotating it with the rotation matrix
`rotation` and displacing it by the vector `displacement`.
"""
solid_id = Geometry.add_solid(self, solid=solid, rotation=rotation,
displacement=displacement)
self.solid_id_to_channel_index.resize(solid_id+1)
self.solid_id_to_channel_index[solid_id] = -1 # solid maps to no channel
return solid_id
def render_particle_track(self):
x = 10.0
h = x*np.sqrt(3)/2
pyramid = make.linear_extrude([-x/2,0,x/2], [-h/2,h/2,-h/2], h,
[0]*3, [0]*3)
marker = Solid(pyramid, vacuum, vacuum)
if self.photon_display_mode == 'beg':
photons = self.ev.photons_beg
else:
photons = self.ev.photons_end
geometry = Geometry()
sample_factor = max(1, len(photons.pos) / 10000)
for pos in photons.pos[::sample_factor]:
geometry.add_solid(marker, displacement=pos, rotation=make_rotation_matrix(np.random.uniform(0,2*np.pi), uniform_sphere()))
geometry = create_geometry_from_obj(geometry)
gpu_geometry = gpu.GPUGeometry(geometry)
self.gpu_geometries = [self.gpu_geometry, gpu_geometry]
def add_solid(self, solid, rotation=None, displacement=None):
"""
Add the solid `solid` to the geometry. When building the final triangle
mesh, `solid` will be placed by rotating it with the rotation matrix
`rotation` and displacing it by the vector `displacement`.
"""
solid_id = Geometry.add_solid(self,
solid=solid,
rotation=rotation,
displacement=displacement)
self.solid_id_to_channel_index.append(-1) # solid maps to no channel
return solid_id
def setUp(self):
self.cache_dir = random_tempdir('chroma_cache_test')
self.cache = Cache(self.cache_dir)
self.a = Geometry()
self.a.add_solid(Solid(box(1, 1, 1)))
self.a.add_solid(Solid(box(1, 1, 1)), displacement=(10, 10, 10))
self.a.flatten()
self.b = Geometry()
self.b.add_solid(Solid(box(2, 2, 2)))
self.b.add_solid(Solid(box(2, 2, 2)), displacement=(10, 10, 10))
self.b.add_solid(Solid(box(2, 2, 2)), displacement=(-10, -10, -10))
self.b.flatten()
# c is not in cache
self.c = Geometry()
self.c.add_solid(Solid(box(2, 2, 2)))
self.c.flatten()
self.a_hash = self.a.mesh.md5()
self.b_hash = self.b.mesh.md5()
self.c_hash = self.c.mesh.md5()
self.cache.save_geometry('a', self.a)
self.cache.save_geometry('b', self.b)
def build_detector():
"""Returns a cubic detector made of cubic photodetectors."""
g = Geometry(water)
for pos, dir in iter_box(nx, ny, nz, spacing):
# convert to arrays
pos, dir = np.array(pos), np.array(dir)
# we need to figure out what rotation matrix to apply to each
# photodetector so that the photosensitive surface will be facing
# `dir`.
if tuple(dir) == (0, 0, 1):
rotation = None
elif tuple(dir) == (0, 0, -1):
rotation = make_rotation_matrix(np.pi, (1, 0, 0))
else:
rotation = make_rotation_matrix(np.arccos(dir.dot((0, 0, 1))),
np.cross(dir, (0, 0, 1)))
# add the photodetector
g.add_solid(build_pd(size, glass_thickness),
rotation=rotation,
displacement=pos)
world = Solid(make.box(spacing * nx, spacing * ny, spacing * nz),
water,
vacuum,
color=0x33ffffff)
g.add_solid(world)
return g
def bvh_mesh(geometry, layer):
lower_bounds, upper_bounds = geometry.bvh.get_layer(layer).get_bounds()
if len(lower_bounds) == 0 or len(upper_bounds) == 0:
raise Exception('no nodes at layer %i' % layer)
dx, dy, dz = upper_bounds[0] - lower_bounds[0]
center = np.mean([upper_bounds[0],lower_bounds[0]], axis=0)
geometry = Geometry()
geometry.add_solid(Solid(make.box(dx,dy,dz,center), vacuum, vacuum, color=0x33ffffff))
for center, dx, dy, dz in list(zip(np.mean([lower_bounds,upper_bounds],axis=0),
*list(zip(*upper_bounds-lower_bounds))))[1:]:
geometry.add_solid(Solid(make.box(dx,dy,dz,center), vacuum, vacuum, color=0x33ffffff))
return create_geometry_from_obj(geometry)
def build_checkerboard_scene(checkers_per_side=10, squares_per_checker=50):
x = np.linspace(-5000.0, 5000.0, checkers_per_side*squares_per_checker+1)
y = np.linspace(-5000.0, 5000.0, checkers_per_side*squares_per_checker+1)
vertices = np.array(tuple(product(x,y,[0])))
triangles = []
for j in range(y.size-1):
for i in range(x.size-1):
triangles.append([j*len(x)+i, (j+1)*len(x)+i,(j+1)*len(x)+i+1])
triangles.append([j*len(x)+i, j*len(x)+i+1,(j+1)*len(x)+i+1])
checkerboard_mesh = Mesh(vertices, triangles, remove_duplicate_vertices=True)
checkerboard_color_line1 = take(checkers_per_side*squares_per_checker*2, cycle([0]*2*squares_per_checker + [0xffffff]*2*squares_per_checker))*squares_per_checker
checkerboard_color_line2 = take(checkers_per_side*squares_per_checker*2, cycle([0xffffff]*2*squares_per_checker + [0]*2*squares_per_checker))*squares_per_checker
checkerboard_color = take(len(checkerboard_mesh.triangles), cycle(checkerboard_color_line1 + checkerboard_color_line2))
checkerboard_surface_line1 = take(checkers_per_side*squares_per_checker*2, cycle([black_surface]*2*squares_per_checker + [lambertian_surface]*2*squares_per_checker))*squares_per_checker
checkerboard_surface_line2 = take(checkers_per_side*squares_per_checker*2, cycle([lambertian_surface]*2*squares_per_checker + [black_surface]*2*squares_per_checker))*squares_per_checker
checkerboard_surface = take(len(checkerboard_mesh.triangles), cycle(checkerboard_surface_line1 + checkerboard_surface_line2))
checkerboard = Solid(checkerboard_mesh, vacuum, vacuum, surface=checkerboard_surface, color=checkerboard_color)
sphere1 = Solid(sphere(1000.0, nsteps=512), water, vacuum)
sphere2 = Solid(sphere(1000.0, nsteps=512), vacuum, vacuum,
surface=shiny_surface)
sphere3 = Solid(sphere(1000.0, nsteps=512), vacuum, vacuum, surface=lambertian_surface)
checkerboard_scene = Geometry()
checkerboard_scene.add_solid(checkerboard, displacement=(0,0,-1500.0))
checkerboard_scene.add_solid(sphere1, displacement=(2000.0,-2000.0,0))
checkerboard_scene.add_solid(sphere2, displacement=(-2000.0,-2000.0,0))
checkerboard_scene.add_solid(sphere3, displacement=(0.0,2000.0,0))
return checkerboard_scene
class TestCacheBVH(unittest.TestCase):
def setUp(self):
self.cache_dir = random_tempdir('chroma_cache_test')
self.cache = Cache(self.cache_dir)
self.a = Geometry()
self.a.add_solid(Solid(box(1,1,1)))
self.a.add_solid(Solid(box(1,1,1)), displacement=(10,10,10))
self.a.flatten()
self.b = Geometry()
self.b.add_solid(Solid(box(2,2,2)))
self.b.add_solid(Solid(box(2,2,2)), displacement=(10,10,10))
self.b.add_solid(Solid(box(2,2,2)), displacement=(-10,-10,-10))
self.b.flatten()
# c is not in cache
self.c = Geometry()
self.c.add_solid(Solid(box(2,2,2)))
self.c.flatten()
self.a_hash = self.a.mesh.md5()
self.b_hash = self.b.mesh.md5()
self.c_hash = self.c.mesh.md5()
self.cache.save_geometry('a', self.a)
self.cache.save_geometry('b', self.b)
def test_list_bvh(self):
self.assertEqual(len(self.cache.list_bvh(self.a_hash)), 0)
self.cache.save_bvh([], self.a_hash)
self.assertIn('default', self.cache.list_bvh(self.a_hash))
self.cache.save_bvh([], self.a_hash, 'foo')
self.assertIn('foo', self.cache.list_bvh(self.a_hash))
self.assertEqual(len(self.cache.list_bvh(self.a_hash)), 2)
def test_exist_bvh(self):
self.cache.save_bvh([], self.a_hash)
assert self.cache.exist_bvh(self.a_hash)
self.cache.save_bvh([], self.a_hash, 'foo')
assert self.cache.exist_bvh(self.a_hash, 'foo')
def test_load_bvh_not_found(self):
with self.assertRaises(BVHNotFoundError):
self.cache.load_bvh(self.c_hash)
with self.assertRaises(BVHNotFoundError):
self.cache.load_bvh(self.a_hash, 'foo')
def test_save_load_new_bvh(self):
self.cache.save_bvh([], self.a_hash)
self.cache.load_bvh(self.a_hash)
self.cache.save_bvh([], self.a_hash, 'foo')
self.cache.load_bvh(self.a_hash, 'foo')
def test_remove_bvh(self):
self.cache.remove_bvh(self.a_hash, 'does_not_exist')
self.cache.save_bvh([], self.a_hash)
self.cache.save_bvh([], self.a_hash, 'foo')
assert self.cache.exist_bvh(self.a_hash)
assert self.cache.exist_bvh(self.a_hash, 'foo')
self.cache.remove_bvh(self.a_hash)
assert not self.cache.exist_bvh(self.a_hash)
assert self.cache.exist_bvh(self.a_hash, 'foo')
self.cache.remove_bvh(self.a_hash, 'foo')
assert not self.cache.exist_bvh(self.a_hash)
assert not self.cache.exist_bvh(self.a_hash, 'foo')
def tearDown(self):
remove_path(self.cache_dir)
def build_checkerboard_scene(checkers_per_side=10, squares_per_checker=50):
x = np.linspace(-5000.0, 5000.0,
checkers_per_side * squares_per_checker + 1)
y = np.linspace(-5000.0, 5000.0,
checkers_per_side * squares_per_checker + 1)
vertices = np.array(tuple(product(x, y, [0])))
triangles = []
for j in range(y.size - 1):
for i in range(x.size - 1):
triangles.append([
j * len(x) + i, (j + 1) * len(x) + i, (j + 1) * len(x) + i + 1
])
triangles.append(
[j * len(x) + i, j * len(x) + i + 1, (j + 1) * len(x) + i + 1])
checkerboard_mesh = Mesh(vertices,
triangles,
remove_duplicate_vertices=True)
checkerboard_color_line1 = take(
checkers_per_side * squares_per_checker * 2,
cycle([0] * 2 * squares_per_checker +
[0xffffff] * 2 * squares_per_checker)) * squares_per_checker
checkerboard_color_line2 = take(
checkers_per_side * squares_per_checker * 2,
cycle([0xffffff] * 2 * squares_per_checker +
[0] * 2 * squares_per_checker)) * squares_per_checker
checkerboard_color = take(
len(checkerboard_mesh.triangles),
cycle(checkerboard_color_line1 + checkerboard_color_line2))
checkerboard_surface_line1 = take(
checkers_per_side * squares_per_checker * 2,
cycle([black_surface] * 2 * squares_per_checker +
[lambertian_surface] * 2 *
squares_per_checker)) * squares_per_checker
checkerboard_surface_line2 = take(
checkers_per_side * squares_per_checker * 2,
cycle([lambertian_surface] * 2 * squares_per_checker +
[black_surface] * 2 * squares_per_checker)) * squares_per_checker
checkerboard_surface = take(
len(checkerboard_mesh.triangles),
cycle(checkerboard_surface_line1 + checkerboard_surface_line2))
checkerboard = Solid(checkerboard_mesh,
vacuum,
vacuum,
surface=checkerboard_surface,
color=checkerboard_color)
sphere1 = Solid(sphere(1000.0, nsteps=512), water, vacuum)
sphere2 = Solid(sphere(1000.0, nsteps=512),
vacuum,
vacuum,
surface=shiny_surface)
sphere3 = Solid(sphere(1000.0, nsteps=512),
vacuum,
vacuum,
surface=lambertian_surface)
checkerboard_scene = Geometry()
checkerboard_scene.add_solid(checkerboard, displacement=(0, 0, -1500.0))
checkerboard_scene.add_solid(sphere1, displacement=(2000.0, -2000.0, 0))
checkerboard_scene.add_solid(sphere2, displacement=(-2000.0, -2000.0, 0))
checkerboard_scene.add_solid(sphere3, displacement=(0.0, 2000.0, 0))
return checkerboard_scene
def render_mc_info_all_events(self):
#function added to visualize all events in a root file at the same time
self.gpu_geometries = [self.gpu_geometry]
for i, ev in enumerate(self.rr):
print('Evaluating event %i' % i)
if self.sum_mode or ev is None:
return
if self.track_display_mode in ['chroma', 'both'
] and ev.photon_tracks is not None:
geometry = Geometry()
print('Total Photons', len(ev.photon_tracks))
def has(flags, test):
return flags & test == test
tracks = ev.photon_tracks
if self.photons_detected_only:
detected = np.asarray([
has(track.flags[-1], event.SURFACE_DETECT)
for track in tracks
])
tracks = [t for t, m in zip(tracks, detected) if m]
cherenkov = np.asarray([
has(track.flags[0], event.CHERENKOV)
and not has(track.flags[-1], event.BULK_REEMIT)
for track in tracks
])
scintillation = np.asarray([
has(track.flags[0], event.SCINTILLATION)
and not has(track.flags[-1], event.BULK_REEMIT)
for track in tracks
])
reemission = np.asarray([
has(track.flags[-1], event.BULK_REEMIT) for track in tracks
])
if self.photons_only_type is not None:
if self.photons_only_type == 'cher':
selector = cherenkov
elif self.photons_only_type == 'scint':
selector = scintillation
elif self.photons_only_type == 'reemit':
selector = reemission
else:
raise Exception('Unknown only type: %s' % only)
tracks = [t for t, m in zip(tracks, selector) if m]
cherenkov = cherenkov[selector]
scintillation = scintillation[selector]
reemission = reemission[selector]
nphotons = len(tracks)
prob = self.photons_max / nphotons if self.photons_max is not None and nphotons is not 0 else 1.0
selector = np.random.random(len(tracks)) < prob
nphotons = np.count_nonzero(selector)
for i, track in ((i, t)
for i, (s,
t) in enumerate(zip(selector, tracks))
if s):
if cherenkov[i]:
color = [255, 0, 0]
elif scintillation[i]:
color = [0, 0, 255]
elif reemission[i]:
color = [0, 255, 0]
else:
color = [255, 255, 255]
steps = min(
len(track), self.photons_max_steps
) if self.photons_max_steps is not None else len(track)
self.render_photon_track(geometry,
track[:steps],
sz=self.photons_track_size,
color=color)
if nphotons > 0:
print('Rendered Photons', nphotons)
geometry = create_geometry_from_obj(geometry)
gpu_geometry = gpu.GPUGeometry(geometry)
self.gpu_geometries.append(gpu_geometry)
class TestCacheGeometry(unittest.TestCase):
def setUp(self):
self.cache_dir = random_tempdir('chroma_cache_test')
self.cache = Cache(self.cache_dir)
self.a = Geometry()
self.a.add_solid(Solid(box(1, 1, 1)))
self.a.add_solid(Solid(box(1, 1, 1)), displacement=(10, 10, 10))
self.a.flatten()
self.b = Geometry()
self.b.add_solid(Solid(box(2, 2, 2)))
self.b.add_solid(Solid(box(2, 2, 2)), displacement=(10, 10, 10))
self.b.add_solid(Solid(box(2, 2, 2)), displacement=(-10, -10, -10))
self.b.flatten()
def test_list_geometry(self):
self.assertEqual(len(self.cache.list_geometry()), 0)
self.cache.save_geometry('a', self.a)
l = self.cache.list_geometry()
self.assertEqual(len(l), 1)
self.assertIn('a', l)
self.cache.save_geometry('b', self.b)
l = self.cache.list_geometry()
self.assertEqual(len(l), 2)
self.assertIn('a', l)
self.assertIn('b', l)
self.cache.save_geometry('a', self.a)
l = self.cache.list_geometry()
self.assertEqual(len(l), 2)
self.assertIn('a', l)
self.assertIn('b', l)
def test_load_geometry_not_found(self):
with self.assertRaises(GeometryNotFoundError):
self.cache.load_geometry('a')
def test_save_load_new_geometry(self):
self.cache.save_geometry('b', self.b)
b = self.cache.load_geometry('b')
def test_replace_geometry(self):
self.cache.save_geometry('b', self.b)
b = self.cache.load_geometry('b')
self.assertEqual(b.mesh.md5(), self.b.mesh.md5())
self.cache.save_geometry('b', self.b)
b = self.cache.load_geometry('b')
self.assertEqual(b.mesh.md5(), self.b.mesh.md5())
def test_remove_geometry(self):
self.cache.save_geometry('b', self.b)
self.assertIn('b', self.cache.list_geometry())
self.cache.remove_geometry('b')
self.assertNotIn('b', self.cache.list_geometry())
def test_get_geometry_hash(self):
self.cache.save_geometry('b', self.b)
self.assertEqual(self.cache.get_geometry_hash('b'), self.b.mesh.md5())
def test_get_geometry_hash_not_found(self):
with self.assertRaises(GeometryNotFoundError):
self.cache.get_geometry_hash('a')
def test_default_geometry(self):
self.cache.save_geometry('a', self.a)
self.cache.save_geometry('b', self.b)
with self.assertRaises(GeometryNotFoundError):
self.cache.set_default_geometry('c')
self.cache.set_default_geometry('b')
b = self.cache.load_default_geometry()
self.cache.set_default_geometry('a')
a = self.cache.load_default_geometry()
def test_default_geometry_corruption(self):
self.cache.save_geometry('a', self.a)
self.cache.save_geometry('b', self.b)
# Put a file where a symlink should be
default_symlink_path = self.cache.get_geometry_filename('.default')
with open(default_symlink_path, 'w') as f:
f.write('foo')
with self.assertRaises(IOError):
self.cache.set_default_geometry('b')
# Verify file not modified
assert os.path.isfile(default_symlink_path)
with open(default_symlink_path) as f:
self.assertEqual(f.read(), 'foo')
def tearDown(self):
remove_path(self.cache_dir)
class TestCacheBVH(unittest.TestCase):
def setUp(self):
self.cache_dir = random_tempdir('chroma_cache_test')
self.cache = Cache(self.cache_dir)
self.a = Geometry()
self.a.add_solid(Solid(box(1, 1, 1)))
self.a.add_solid(Solid(box(1, 1, 1)), displacement=(10, 10, 10))
self.a.flatten()
self.b = Geometry()
self.b.add_solid(Solid(box(2, 2, 2)))
self.b.add_solid(Solid(box(2, 2, 2)), displacement=(10, 10, 10))
self.b.add_solid(Solid(box(2, 2, 2)), displacement=(-10, -10, -10))
self.b.flatten()
# c is not in cache
self.c = Geometry()
self.c.add_solid(Solid(box(2, 2, 2)))
self.c.flatten()
self.a_hash = self.a.mesh.md5()
self.b_hash = self.b.mesh.md5()
self.c_hash = self.c.mesh.md5()
self.cache.save_geometry('a', self.a)
self.cache.save_geometry('b', self.b)
def test_list_bvh(self):
self.assertEqual(len(self.cache.list_bvh(self.a_hash)), 0)
self.cache.save_bvh([], self.a_hash)
self.assertIn('default', self.cache.list_bvh(self.a_hash))
self.cache.save_bvh([], self.a_hash, 'foo')
self.assertIn('foo', self.cache.list_bvh(self.a_hash))
self.assertEqual(len(self.cache.list_bvh(self.a_hash)), 2)
def test_exist_bvh(self):
self.cache.save_bvh([], self.a_hash)
assert self.cache.exist_bvh(self.a_hash)
self.cache.save_bvh([], self.a_hash, 'foo')
assert self.cache.exist_bvh(self.a_hash, 'foo')
def test_load_bvh_not_found(self):
with self.assertRaises(BVHNotFoundError):
self.cache.load_bvh(self.c_hash)
with self.assertRaises(BVHNotFoundError):
self.cache.load_bvh(self.a_hash, 'foo')
def test_save_load_new_bvh(self):
self.cache.save_bvh([], self.a_hash)
self.cache.load_bvh(self.a_hash)
self.cache.save_bvh([], self.a_hash, 'foo')
self.cache.load_bvh(self.a_hash, 'foo')
def test_remove_bvh(self):
self.cache.remove_bvh(self.a_hash, 'does_not_exist')
self.cache.save_bvh([], self.a_hash)
self.cache.save_bvh([], self.a_hash, 'foo')
assert self.cache.exist_bvh(self.a_hash)
assert self.cache.exist_bvh(self.a_hash, 'foo')
self.cache.remove_bvh(self.a_hash)
assert not self.cache.exist_bvh(self.a_hash)
assert self.cache.exist_bvh(self.a_hash, 'foo')
self.cache.remove_bvh(self.a_hash, 'foo')
assert not self.cache.exist_bvh(self.a_hash)
assert not self.cache.exist_bvh(self.a_hash, 'foo')
def tearDown(self):
remove_path(self.cache_dir)
def testBulkReemission(self):
'''Test bulk reemission
Start a bunch of monoenergetic photons at the center of a wavelength-
shifting sphere, forcing reemission, and check that the final
wavelength distribution matches the wls spectrum.
'''
import scipy.stats
nphotons = 1e5
# set up detector -- a sphere of 'scintillator' surrounded by a
# detecting sphere
scint = Material('scint')
scint.set('refractive_index', 1)
scint.set('absorption_length', 1.0)
scint.set('scattering_length', 1e7)
scint.set('reemission_prob', 1)
x = np.arange(0,1000,10)
norm = scipy.stats.norm(scale=50, loc=600)
pdf = 10 * norm.pdf(x)
cdf = norm.cdf(x)
scint.reemission_cdf = np.array(zip(x, cdf))
detector = Surface('detector')
detector.set('detect', 1)
world = Geometry(vacuum)
world.add_solid(Solid(sphere(1000), vacuum, vacuum, surface=detector))
world.add_solid(Solid(sphere(500), scint, vacuum))
w = create_geometry_from_obj(world, update_bvh_cache=False)
sim = Simulation(w, geant4_processes=0)
# initial photons -- isotropic 250 nm at the origin
pos = np.tile([0,0,0], (nphotons,1)).astype(np.float32)
dir = np.random.rand(nphotons, 3).astype(np.float32) * 2 - 1
dir /= np.sqrt(dir[:,0]**2 + dir[:,1]**2 + dir[:,2]**2)[:,np.newaxis]
pol = np.zeros_like(pos)
t = np.zeros(nphotons, dtype=np.float32)
wavelengths = np.ones(nphotons).astype(np.float32) * 250
photons = Photons(pos=pos, dir=dir, pol=pol, t=t, wavelengths=wavelengths)
# run simulation and extract final wavelengths
event = sim.simulate([photons], keep_photons_end=True).next()
mask = (event.photons_end.flags & SURFACE_DETECT) > 0
final_wavelengths = event.photons_end.wavelengths[mask]
# compare wavelength distribution to scintillator's reemission pdf
hist, edges = np.histogram(final_wavelengths, bins=x)
print 'detected', hist.sum(), 'of', nphotons, 'photons'
hist_norm = 1.0 * hist / (1.0 * hist.sum() / 1000)
pdf /= (1.0 * pdf.sum() / 1000)
chi2 = scipy.stats.chisquare(hist_norm, pdf[:-1])[1]
print 'chi2 =', chi2
# show histogram comparison
#plt.figure(1)
#width = edges[1] - edges[0]
#plt.bar(left=edges, height=pdf, width=width, color='red')
#plt.bar(left=edges[:-1], height=hist_norm, width=width)
#plt.show()
self.assertTrue(chi2 > 0.75)
def build(slit_width):
directory = 'stls/' #directory containing individual detector .stls
##### Generating World #####
print 'Generating world...'
size = (40,40,40) #Bounds for world box
world = Geometry(vacuum)
bounds = Solid(make.box(*size),vacuum,vacuum,color=0x33ffffff)
world.add_solid(bounds,displacement=(18,-18,10))
##### Generating Non-Special Solids #####
print 'Generating non-special solids...'
blacklist = mfunctions.get_blacklist() #Mesh read-in blacklist
tol = 0.005 #Tolerance for painting meshes using conditionals
setup_solid = Solid(make.cube(1),vacuum,vacuum,color=0xf0fc03) #origin
for filename in os.listdir(directory):
path = os.path.join(directory,filename)
tmp_mesh = mesh_from_stl(path)
mesh_no = re.findall('\d+',path)[0] #Select mesh number
if mesh_no in ['141','142']: #Angled mirrors
setup_surf, setup_col = mfunctions.mirror(mesh_no)
"""
elif mesh_no in ['161','206']: #Slit housing
setup_surf = shiny_surface
setup_col = 0x7d6b56 #grey
"""
else: #Non-optical components
setup_surf = black_surface
setup_col = 0x7d6b56 #grey
#print 'Adding solid '+mesh_no+'/249'
if mesh_no not in blacklist: #Remove blacklisted meshes from setup
setup_solid += Solid(tmp_mesh,vacuum,vacuum,surface=setup_surf,color=setup_col)
world.add_solid(setup_solid)
##### Generating Special Solids #####
print 'Generating special solids...'
"""
#Laser Absorb
laser_absorb = Solid(make.segmented_cylinder(1,0.1),vacuum,vacuum,surface=black_surface,color=0xffff00)
cap_center = np.array((18.470600,-26.771999,19.469999))
world.add_solid(laser_absorb,displacement=cap_center)
"""
#SiPM Plate
sipm_plate = Solid(mesh_from_stl(directory+'vuv_setup - Mesh 188.stl'),vacuum,vacuum,surface=black_surface,color=0xfc7f03)
rotation_matrix = mfunctions.Ry(np.arctan(0.0755832)) #Correction for SiPM plate rotation in the STL
sipm_center = np.array((18.48695,-22.89630,9.08933)) #Centre of SiPM plate
correction = sipm_center-np.matmul(rotation_matrix,sipm_center) #Off-centre rotation induces displacement, this corrects for this
world.add_solid(sipm_plate,rotation=rotation_matrix,displacement=correction+np.array((0,1,0)))
#Detector Plate/SiPM
detector = Solid(make.box(2,3,0.01),vacuum,vacuum,surface=mfunctions.get_sipm_surface(),color=0x0dff00)
world.add_solid(detector,displacement=sipm_center+np.array((0,0.8,0)))
#Closed Mirror
#mirror_surf,mirror_col = mfunctions.mirror('250')
mirror_surf = mfunctions.get_mirror_surface()
mirror_col = 0xe6ad12
closed_mirror = Solid(mesh_from_stl(directory+'vuv_setup - Mesh 250.stl'),vacuum,vacuum,surface=mirror_surf,color=mirror_col)
world.add_solid(closed_mirror,rotation=mfunctions.Ry(np.pi),displacement=np.array((18.470901,-21.916050,19.602501))+np.array((0,0,0.613293)))
#Post-Monochromator Slit Doors
basex = np.array((0.0,0.0,0.074744,0.6,0.6))#x-points for door base
basey = np.array((0.0,0.00488564,0.0800017,0.0800017,0.0))#y-points
height = 0.65
#Max slit width: 1000micron
door1 = make.linear_extrude(basex,basey,height)
door2 = make.linear_extrude(-basex,basey,height)
slit_center = mfunctions.get_center('161')
door1_solid = Solid(door1,vacuum,vacuum,surface=shiny_surface,color=0xffff00)
door2_solid = Solid(door2,vacuum,vacuum,surface=shiny_surface,color=0xffff00)
world.add_solid(door1_solid,displacement=slit_center+np.array((0.5*slit_width,-0.144076,-0.05)))
world.add_solid(door2_solid,displacement=slit_center+np.array((-0.5*slit_width,-0.144076,-0.05)))
#PMT plate
pmt_center = mfunctions.get_center('210')
pmt_center += np.array((0.4,0,0))
pmt_solid = Solid(make.segmented_cylinder(1.5,0.1),vacuum,vacuum,surface=mfunctions.get_pmt_surface(),color=0x0dff00)
world.add_solid(pmt_solid,rotation=mfunctions.Rz(np.pi/2),displacement=pmt_center)
##### Saving Geometry to File #####
config_path = open('config/geometry.pickle','wb')
cPickle.dump(world,config_path)
config_path.close()
config_path = open('config/vars.txt','wb')
config_path.write(str(slit_width)+'\n')
config_path.write('This is the config file for geometry.pickle.\nAltering this file will cause geometry.pickle to regenerate!')
config_path.close()
print 'Done.'
return world
def load_geometry_from_string(geometry_str,
auto_build_bvh=True, read_bvh_cache=True,
update_bvh_cache=True, cache_dir=None,
cuda_device=None):
'''Create or load a geometry and optionally load/build a BVH for it.
This is a convenience interface to the geometry and BVH construction code,
as well as the Chroma caching layer. Most applications should use
this function rather than manually building a Geometry and BVH.
The geometry string passed to this function has several forms:
"" (empty string) - Load the default geometry from the cache and
the default BVH for that geometry.
"filename.stl" or "filename.stl.bz2" - Create a geometry from a
3D mesh on disk. This model will not be cached, but the
BVH can be, depending on whether update_bvh_cache is True.
"geometry_name" - Load a geometry from the cache with this name
and the default BVH for that geometry.
"geometry_name:bvh_name" - Load a geometry from the cache and
the requested BVH by name.
"@chroma.models.lionsolid" - Run this function inside a Python
module, found in the current $PYTHONPATH, to create the
geometry, and load the default BVH. For convenience, the
current directory is also added to the $PYTHONPATH.
"@chroma.models.lionsolid:bvh_name" - Run this function to
create the Geometry and load a BVH by name.
By default, the Chroma cache in the user's home directory is
consulted for both the geometry and the BVH. A different cache
directory can be selected by passing the path in via the
``cache_dir`` parameter.
If ``read_bvh_cache`` is set to False, then the BVH cache will not
be inspected for BVH objects.
If the requested BVH (default, or named) does not exist for this
geometry (checked by MD5 hashing the geometry mesh) and
``auto_build_bvh`` is true, then a BVH will be automatically
generated using the "simple" BVH algorithm. The simple algorithm
is very fast, but produces a poor quality BVH.
Any newly created BVH will be saved in the Chroma cache if the
``update_cache_bvh`` parameter is True.
BVH construction requires a GPU, so the CUDA device number can be
specified with the ``cuda_device`` parameter.
Returns: a Geometry object (or subclass) with the ``bvh`` property
set if the options allow.
'''
# Find BVH id if given
bvh_name = 'default'
if ':' in geometry_str:
geometry_id, bvh_name = geometry_str.split(':')
else:
geometry_id = geometry_str
if cache_dir is None:
cache = Cache()
else:
cache = Cache(cache_dir)
# Where is the geometry coming from?
if os.path.exists(geometry_id) and \
geometry_id.lower().endswith(('.stl', '.bz2')):
# Load from file
mesh = mesh_from_stl(geometry_id)
geometry = Geometry()
geometry.add_solid(Solid(mesh, vacuum, vacuum, color=0x33ffffff))
geometry.flatten()
elif geometry_id.startswith('@'):
# Load from function
function_path = geometry_id[1:]
module_name, obj_name = function_path.rsplit('.', 1)
orig_sys_path = list(sys.path)
try:
sys.path.append('.')
module = __import__(module_name, fromlist=[obj_name])
sys.path = orig_sys_path
except ImportError:
sys.path = orig_sys_path
raise
obj = getattr(module, obj_name)
geometry = create_geometry_from_obj(obj, bvh_name=bvh_name,
auto_build_bvh=auto_build_bvh,
read_bvh_cache=read_bvh_cache,
update_bvh_cache=update_bvh_cache,
cache_dir=cache_dir,
cuda_device=cuda_device)
return geometry # RETURN EARLY HERE! ALREADY GOT BVH
else:
# Load from cache
if geometry_id == '':
geometry = cache.load_default_geometry()
else:
geometry = cache.load_geometry(geometry_id)
# Cached geometries are flattened already
geometry.bvh = load_bvh(geometry, bvh_name=bvh_name,
auto_build_bvh=auto_build_bvh,
read_bvh_cache=read_bvh_cache,
update_bvh_cache=update_bvh_cache,
cache_dir=cache_dir,
cuda_device=cuda_device)
return geometry
def load_geometry_from_string(geometry_str,
auto_build_bvh=True,
read_bvh_cache=True,
update_bvh_cache=True,
cache_dir=None,
cuda_device=None):
'''Create or load a geometry and optionally load/build a BVH for it.
This is a convenience interface to the geometry and BVH construction code,
as well as the Chroma caching layer. Most applications should use
this function rather than manually building a Geometry and BVH.
The geometry string passed to this function has several forms:
"" (empty string) - Load the default geometry from the cache and
the default BVH for that geometry.
"filename.stl" or "filename.stl.bz2" - Create a geometry from a
3D mesh on disk. This model will not be cached, but the
BVH can be, depending on whether update_bvh_cache is True.
"geometry_name" - Load a geometry from the cache with this name
and the default BVH for that geometry.
"geometry_name:bvh_name" - Load a geometry from the cache and
the requested BVH by name.
"@chroma.models.lionsolid" - Run this function inside a Python
module, found in the current $PYTHONPATH, to create the
geometry, and load the default BVH. For convenience, the
current directory is also added to the $PYTHONPATH.
"@chroma.models.lionsolid:bvh_name" - Run this function to
create the Geometry and load a BVH by name.
By default, the Chroma cache in the user's home directory is
consulted for both the geometry and the BVH. A different cache
directory can be selected by passing the path in via the
``cache_dir`` parameter.
If ``read_bvh_cache`` is set to False, then the BVH cache will not
be inspected for BVH objects.
If the requested BVH (default, or named) does not exist for this
geometry (checked by MD5 hashing the geometry mesh) and
``auto_build_bvh`` is true, then a BVH will be automatically
generated using the "simple" BVH algorithm. The simple algorithm
is very fast, but produces a poor quality BVH.
Any newly created BVH will be saved in the Chroma cache if the
``update_cache_bvh`` parameter is True.
BVH construction requires a GPU, so the CUDA device number can be
specified with the ``cuda_device`` parameter.
Returns: a Geometry object (or subclass) with the ``bvh`` property
set if the options allow.
'''
# Find BVH id if given
bvh_name = 'default'
if ':' in geometry_str:
geometry_id, bvh_name = geometry_str.split(':')
else:
geometry_id = geometry_str
if cache_dir is None:
cache = Cache()
else:
cache = Cache(cache_dir)
# Where is the geometry coming from?
if os.path.exists(geometry_id) and \
geometry_id.lower().endswith(('.stl', '.bz2')):
# Load from file
mesh = mesh_from_stl(geometry_id)
geometry = Geometry()
geometry.add_solid(Solid(mesh, vacuum, vacuum, color=0x33ffffff))
geometry.flatten()
elif geometry_id.startswith('@'):
# Load from function
function_path = geometry_id[1:]
module_name, obj_name = function_path.rsplit('.', 1)
orig_sys_path = list(sys.path)
try:
sys.path.append('.')
module = __import__(module_name, fromlist=[obj_name])
sys.path = orig_sys_path
except ImportError:
sys.path = orig_sys_path
raise
obj = getattr(module, obj_name)
geometry = create_geometry_from_obj(obj,
bvh_name=bvh_name,
auto_build_bvh=auto_build_bvh,
read_bvh_cache=read_bvh_cache,
update_bvh_cache=update_bvh_cache,
cache_dir=cache_dir,
cuda_device=cuda_device)
return geometry # RETURN EARLY HERE! ALREADY GOT BVH
else:
# Load from cache
if geometry_id == '':
geometry = cache.load_default_geometry()
else:
geometry = cache.load_geometry(geometry_id)
# Cached geometries are flattened already
geometry.bvh = load_bvh(geometry,
bvh_name=bvh_name,
auto_build_bvh=auto_build_bvh,
read_bvh_cache=read_bvh_cache,
update_bvh_cache=update_bvh_cache,
cache_dir=cache_dir,
cuda_device=cuda_device)
return geometry
def create_geometry_from_obj(obj,
bvh_name="default",
auto_build_bvh=True,
read_bvh_cache=True,
update_bvh_cache=True,
cache_dir=None,
cuda_device=None):
if callable(obj):
obj = obj()
if isinstance(obj, Detector):
geometry = obj
if isinstance(obj, Geometry):
geometry = obj
elif isinstance(obj, Solid):
geometry = Geometry()
geometry.add_solid(obj)
elif isinstance(obj, Mesh):
geometry = Geometry()
geometry.add_solid(Solid(obj, vacuum, vacuum, color=0x33ffffff))
else:
raise TypeError('cannot build type %s' % type(obj))
geometry.flatten()
if geometry.bvh is None:
geometry.bvh = load_bvh(geometry,
auto_build_bvh=auto_build_bvh,
read_bvh_cache=read_bvh_cache,
update_bvh_cache=update_bvh_cache,
cache_dir=cache_dir,
cuda_device=cuda_device)
return geometry
def render_mc_info(self):
#need to render photon tracking info if available
self.gpu_geometries = [self.gpu_geometry]
if self.sum_mode or self.ev is None:
return
if self.photon_display_mode == 'beg':
photons = self.ev.photons_beg
elif self.photon_display_mode == 'end':
photons = self.ev.photons_end
else:
photons = None
if photons is not None:
geometry = Geometry()
self.render_photons(geometry, photons)
geometry = create_geometry_from_obj(geometry)
gpu_geometry = gpu.GPUGeometry(geometry)
self.gpu_geometries.append(gpu_geometry)
if self.track_display_mode in ['geant4', 'both'
] and self.ev.vertices is not None:
geometry = Geometry()
any = False
for vertex in self.ev.vertices:
if vertex.steps:
any = True
self.render_vertex(geometry, vertex, children=True)
if any:
geometry = create_geometry_from_obj(geometry)
gpu_geometry = gpu.GPUGeometry(geometry)
self.gpu_geometries.append(gpu_geometry)
if self.track_display_mode in ['chroma', 'both'
] and self.ev.photon_tracks is not None:
geometry = Geometry()
print('Total Photons', len(self.ev.photon_tracks))
def has(flags, test):
return flags & test == test
tracks = self.ev.photon_tracks
if self.photons_detected_only:
detected = np.asarray([
has(track.flags[-1], event.SURFACE_DETECT)
for track in tracks
])
tracks = [t for t, m in zip(tracks, detected) if m]
cherenkov = np.asarray([
has(track.flags[0], event.CHERENKOV)
and not has(track.flags[-1], event.BULK_REEMIT)
for track in tracks
])
scintillation = np.asarray([
has(track.flags[0], event.SCINTILLATION)
and not has(track.flags[-1], event.BULK_REEMIT)
for track in tracks
])
reemission = np.asarray(
[has(track.flags[-1], event.BULK_REEMIT) for track in tracks])
if self.photons_only_type is not None:
if self.photons_only_type == 'cher':
selector = cherenkov
elif self.photons_only_type == 'scint':
selector = scintillation
elif self.photons_only_type == 'reemit':
selector = reemission
else:
raise Exception('Unknown only type: %s' % only)
tracks = [t for t, m in zip(tracks, selector) if m]
cherenkov = cherenkov[selector]
scintillation = scintillation[selector]
reemission = reemission[selector]
nphotons = len(tracks)
prob = self.photons_max / nphotons if self.photons_max is not None else 1.0
selector = np.random.random(len(tracks)) < prob
nphotons = np.count_nonzero(selector)
for i, track in ((i, t)
for i, (s, t) in enumerate(zip(selector, tracks))
if s):
if cherenkov[i]:
color = [255, 0, 0]
elif scintillation[i]:
color = [0, 0, 255]
elif reemission[i]:
color = [0, 255, 0]
else:
color = [255, 255, 255]
steps = min(
len(track), self.photons_max_steps
) if self.photons_max_steps is not None else len(track)
self.render_photon_track(geometry,
track[:steps],
sz=self.photons_track_size,
color=color)
if nphotons > 0:
print('Rendered Photons', nphotons)
geometry = create_geometry_from_obj(geometry)
gpu_geometry = gpu.GPUGeometry(geometry)
self.gpu_geometries.append(gpu_geometry)
class TestCacheGeometry(unittest.TestCase):
def setUp(self):
self.cache_dir = random_tempdir('chroma_cache_test')
self.cache = Cache(self.cache_dir)
self.a = Geometry()
self.a.add_solid(Solid(box(1,1,1)))
self.a.add_solid(Solid(box(1,1,1)), displacement=(10,10,10))
self.a.flatten()
self.b = Geometry()
self.b.add_solid(Solid(box(2,2,2)))
self.b.add_solid(Solid(box(2,2,2)), displacement=(10,10,10))
self.b.add_solid(Solid(box(2,2,2)), displacement=(-10,-10,-10))
self.b.flatten()
def test_list_geometry(self):
self.assertEqual(len(self.cache.list_geometry()), 0)
self.cache.save_geometry('a', self.a)
l = self.cache.list_geometry()
self.assertEqual(len(l), 1)
self.assertIn('a', l)
self.cache.save_geometry('b', self.b)
l = self.cache.list_geometry()
self.assertEquals(len(l), 2)
self.assertIn('a', l)
self.assertIn('b', l)
self.cache.save_geometry('a', self.a)
l = self.cache.list_geometry()
self.assertEquals(len(l), 2)
self.assertIn('a', l)
self.assertIn('b', l)
def test_load_geometry_not_found(self):
with self.assertRaises(GeometryNotFoundError):
self.cache.load_geometry('a')
def test_save_load_new_geometry(self):
self.cache.save_geometry('b', self.b)
b = self.cache.load_geometry('b')
def test_replace_geometry(self):
self.cache.save_geometry('b', self.b)
b = self.cache.load_geometry('b')
self.assertEqual(b.mesh.md5(), self.b.mesh.md5())
self.cache.save_geometry('b', self.b)
b = self.cache.load_geometry('b')
self.assertEqual(b.mesh.md5(), self.b.mesh.md5())
def test_remove_geometry(self):
self.cache.save_geometry('b', self.b)
self.assertIn('b', self.cache.list_geometry())
self.cache.remove_geometry('b')
self.assertNotIn('b', self.cache.list_geometry())
def test_get_geometry_hash(self):
self.cache.save_geometry('b', self.b)
self.assertEqual(self.cache.get_geometry_hash('b'), self.b.mesh.md5())
def test_get_geometry_hash_not_found(self):
with self.assertRaises(GeometryNotFoundError):
self.cache.get_geometry_hash('a')
def test_default_geometry(self):
self.cache.save_geometry('a', self.a)
self.cache.save_geometry('b', self.b)
with self.assertRaises(GeometryNotFoundError):
self.cache.set_default_geometry('c')
self.cache.set_default_geometry('b')
b = self.cache.load_default_geometry()
self.cache.set_default_geometry('a')
a = self.cache.load_default_geometry()
def test_default_geometry_corruption(self):
self.cache.save_geometry('a', self.a)
self.cache.save_geometry('b', self.b)
# Put a file where a symlink should be
default_symlink_path = self.cache.get_geometry_filename('.default')
with open(default_symlink_path, 'w') as f:
f.write('foo')
with self.assertRaises(IOError):
self.cache.set_default_geometry('b')
# Verify file not modified
assert os.path.isfile(default_symlink_path)
with open(default_symlink_path) as f:
self.assertEqual(f.read(), 'foo')
def tearDown(self):
remove_path(self.cache_dir)