import chroma modules from subpackages with import chroma.module_name

19 files changed, 74 insertions, 110 deletions

diff --git a/camera.py b/camera.py
index fa38060..98ea1fc 100755
--- a/camera.py
+++ b/camera.py
@@ -4,12 +4,12 @@ from itertools import product, count
 from threading import Thread, Lock
 import os
 import sys
-from color import map_wavelength, map_to_color
+from chroma.color import map_wavelength, map_to_color
 import matplotlib.pyplot as plt
 
-from gpu import GPU, CUDAFuncs
-from tools import timeit
-from transform import rotate
+from chroma.gpu import GPU, CUDAFuncs
+from chroma.tools import timeit
+from chroma.transform import rotate
 
 import pygame
 from pygame.locals import *
@@ -483,11 +483,11 @@ if __name__ == '__main__':
     import optparse
     import inspect
 
-    import solids
-    import detectors
-    import scenes
-    from stl import mesh_from_stl
-    from view import build
+    import chroma.solids
+    import chroma.detectors
+    import chroma.scenes
+    from chroma.stl import mesh_from_stl
+    from chroma.view import build
 
     parser = optparse.OptionParser('%prog filename.stl')
     parser.add_option('-b', '--bits', type='int', dest='bits',
@@ -511,7 +511,7 @@ if __name__ == '__main__':
             obj = mesh_from_stl(args[0])
 
     else:
-        members = dict(inspect.getmembers(detectors) + inspect.getmembers(solids) + inspect.getmembers(scenes))
+        members = dict(inspect.getmembers(chroma.detectors) + inspect.getmembers(chroma.solids) + inspect.getmembers(chroma.scenes))
 
         buildable_lookup = {}
         for member in members.values():
diff --git a/color/chromaticity.py b/color/chromaticity.py
index 859cc14..0076957 100644
--- a/color/chromaticity.py
+++ b/color/chromaticity.py
@@ -1,14 +1,8 @@
 import numpy as np
-import os
-import sys
+from os.path import realpath, dirname
+from chroma.tools import read_csv
 
-dir = os.path.split(os.path.realpath(__file__))[0]
-
-sys.path.append(dir + '/..')
-
-from tools import read_csv
-
-color_map = read_csv(dir + '/ciexyz64_1.csv')
+color_map = read_csv(dirname(realpath(__file__)) + '/ciexyz64_1.csv')
 
 def map_wavelength(wavelength):
     r = np.interp(wavelength, color_map[:,0], color_map[:,1])
diff --git a/detectors/__init__.py b/detectors/__init__.py
index 8550751..84e853b 100644
--- a/detectors/__init__.py
+++ b/detectors/__init__.py
@@ -1,15 +1,8 @@
 from lbne import build_lbne
 from sno import build_sno as build_sno_detector
-
-import os
-import sys
+from chroma.view import buildable
 import inspect
 
-dir = os.path.split(os.path.realpath(__file__))[0]
-sys.path.append(dir + '/..')
-
-from view import buildable
-
 # from LBNE document #3951
 radius = 63.30/2
 height = 76.60
diff --git a/detectors/lbne.py b/detectors/lbne.py
index 8942cb2..ca13c2d 100644
--- a/detectors/lbne.py
+++ b/detectors/lbne.py
@@ -1,27 +1,21 @@
-import os
-import sys
 import numpy as np
-
-dir = os.path.split(os.path.realpath(__file__))[0]
-sys.path.append(dir + '/..')
-
-from geometry import *
-from solids import build_12inch_pmt, build_12inch_pmt_with_lc, build_12inch_pmt_shell
-from optics import *
-from transform import rotate, make_rotation_matrix
+from chroma.geometry import Solid, Geometry
+import chroma.solids as solids
+from chroma.optics import vacuum, water, water_wcsim, black_surface
+from chroma.transform import rotate, make_rotation_matrix
 from itertools import product
-from make import cylinder, segmented_cylinder
+import chroma.make as make
 
 def build_lbne(radius, height, nstrings, pmts_per_string, endcap_spacing, physical_model=True):
     if physical_model:
-        pmt = build_12inch_pmt_with_lc(outer_material=water_wcsim)
+        pmt = solids.build_12inch_pmt_with_lc(outer_material=water_wcsim)
     else:
-        pmt = build_12inch_pmt_shell(outer_material=water_wcsim)
+        pmt = solids.build_12inch_pmt_shell(outer_material=water_wcsim)
 
     lbne = Geometry()
 
     # outer cylinder
-    cylinder_mesh = segmented_cylinder(radius, height+height/(pmts_per_string-1), nsteps=16*nstrings, nsegments=1200)
+    cylinder_mesh = make.segmented_cylinder(radius, height+height/(pmts_per_string-1), nsteps=16*nstrings, nsegments=1200)
     cylinder_mesh.vertices = rotate(cylinder_mesh.vertices, np.pi/2, (-1,0,0))
     lbne.add_solid(Solid(cylinder_mesh, water_wcsim, vacuum, black_surface, 0xff0000ff))
 
diff --git a/detectors/sno.py b/detectors/sno.py
index 3e748b4..d0334d3 100644
--- a/detectors/sno.py
+++ b/detectors/sno.py
@@ -1,12 +1,12 @@
 import numpy as np
 import numpy.linalg
 import math
-from make import rotate_extrude
-from stl import mesh_from_stl
-from geometry import *
-from optics import *
-from solids import build_8inch_pmt_with_lc
-from transform import rotate, make_rotation_matrix
+from chroma.make import rotate_extrude
+from chroma.stl import mesh_from_stl
+from chroma.geometry import *
+from chroma.optics import *
+from chroma.solids import build_8inch_pmt_with_lc
+from chroma.transform import rotate, make_rotation_matrix
 import os
 
 def sno_vessel(sphere_radius, neck_radius, neck_top, nsteps=40):
diff --git a/fileio/root.py b/fileio/root.py
index 77c23ff..ae9e741 100644
--- a/fileio/root.py
+++ b/fileio/root.py
@@ -4,7 +4,7 @@ import os.path
 ROOT.gROOT.ProcessLine('.L '+os.path.join(os.path.dirname(__file__), 'root.C+g'))
 
 import ROOT
-import event
+import chroma.event as event
 
 class RootWriter(object):
     def __init__(self, filename):
diff --git a/generator/g4gen.py b/generator/g4gen.py
index cec53f7..b43ce92 100644
--- a/generator/g4gen.py
+++ b/generator/g4gen.py
@@ -4,7 +4,7 @@ import g4py.NISTmaterials
 import g4py.ParticleGun
 import pyublas
 import numpy as np
-import event
+import chroma.event as event
 
 try:
     import G4chroma
diff --git a/generator/vertex.py b/generator/vertex.py
index 1fa61cb..3998e6e 100644
--- a/generator/vertex.py
+++ b/generator/vertex.py
@@ -3,8 +3,8 @@ import numpy as np
 from math import sin, cos, sqrt
 import copy
 
-import pi0
-from event import Event, Subtrack
+import chroma.pi0 as pi0
+from chroma.event import Event, Subtrack
 
 # generator parts for use with gun()
 
diff --git a/geometry.py b/geometry.py
index 856ded3..6ea8994 100644
--- a/geometry.py
+++ b/geometry.py
@@ -1,8 +1,8 @@
 import sys
 import os
 import numpy as np
-from itertoolset import *
-from tools import timeit, profile_if_possible
+from chroma.itertoolset import *
+from chroma.tools import timeit, profile_if_possible
 from hashlib import md5
 import cPickle as pickle
 import gzip
@@ -274,7 +274,7 @@ class Geometry(object):
 
         try:
             self.surface_index[self.surface_index == surface_lookup[None]] = -1
-        except ValueError:
+        except KeyError:
             pass
 
         checksum = md5(str(bits))
diff --git a/gpu.py b/gpu.py
index 327e692..372bc7b 100644
--- a/gpu.py
+++ b/gpu.py
@@ -7,11 +7,12 @@ import pycuda.driver as cuda
 from pycuda import gpuarray
 from copy import copy
 from itertools import izip
-from tools import timeit
+from chroma.tools import timeit
+from os.path import dirname
 
-import src
-from geometry import standard_wavelengths
-from color import map_to_color
+import chroma.src
+from chroma.geometry import standard_wavelengths
+from chroma.color import map_to_color
 import sys
 import event
 
@@ -97,15 +98,15 @@ class GPU(object):
 
         self.context.set_cache_config(cuda.func_cache.PREFER_L1)
 
-        cuda_options = ['-I' + src.dir, '--use_fast_math', '--ptxas-options=-v']
+        cuda_options = ['-I' + dirname(chroma.src.__file__), '--use_fast_math', '--ptxas-options=-v']
 
-        self.module = SourceModule(src.kernel, options=cuda_options, no_extern_c=True)
+        self.module = SourceModule(chroma.src.kernel, options=cuda_options, no_extern_c=True)
         self.geo_funcs = CUDAFuncs(self.module, ['set_wavelength_range', 'set_material', 'set_surface', 'set_global_mesh_variables', 'color_solids'])
 
         self.prop_funcs = CUDAFuncs(self.module, ['init_rng', 'propagate', 'swap'])
         self.nthread_per_block = 64
         self.max_blocks = 1024
-        self.daq_module = SourceModule(src.daq, options=cuda_options, no_extern_c=True)
+        self.daq_module = SourceModule(chroma.src.daq, options=cuda_options, no_extern_c=True)
         self.daq_funcs = CUDAFuncs(self.daq_module, 
                               ['reset_earliest_time_int', 'run_daq',
                                'convert_sortable_int_to_float'])
diff --git a/make.py b/make.py
index 1da5d0e..0375084 100644
--- a/make.py
+++ b/make.py
@@ -1,7 +1,7 @@
 import numpy as np
-from geometry import Mesh
-from transform import rotate
-from itertoolset import *
+from chroma.geometry import Mesh
+from chroma.transform import rotate
+from chroma.itertoolset import *
 
 def mesh_grid(grid):
     return np.vstack(zip(grid[:-1].flatten(),grid[1:].flatten(),np.roll(grid[1:],-1,1).flatten()) + zip(grid[:-1].flatten(),np.roll(grid[1:],-1,1).flatten(),np.roll(grid[:-1],-1,1).flatten()))
@@ -65,8 +65,11 @@ def rotate_extrude(x, y, nsteps=64):
 
     return Mesh(vertices, triangles, remove_duplicate_vertices=True)
 
-def cube(size=1):
+def cube(size=1, height=None):
     "Return a cube mesh whose sides have length `size`."
+    if height is None:
+        height = size
+
     return linear_extrude([-size/2.0,size/2.0,size/2.0,-size/2.0],[-size/2.0,-size/2.0,size/2.0,size/2.0], height=size)
 
 def cylinder(radius=1, height=2, radius2=None, nsteps=64):
diff --git a/models/__init__.py b/models/__init__.py
deleted file mode 100644
index f369c9b..0000000
--- a/models/__init__.py
+++ /dev/null
@@ -1,3 +0,0 @@
-import os
-
-dir = os.path.split(os.path.realpath(__file__))[0]
diff --git a/optics.py b/optics.py
index 4e7530e..02b8e0d 100644
--- a/optics.py
+++ b/optics.py
@@ -1,5 +1,5 @@
 import numpy as np
-from geometry import Material, Surface
+from chroma.geometry import Material, Surface
 
 vacuum = Material('vacuum')
 vacuum.set('refractive_index', 1.0)
diff --git a/sample.py b/sample.py
index 0b2ffc6..1909370 100644
--- a/sample.py
+++ b/sample.py
@@ -1,5 +1,5 @@
 import numpy as np
-from transform import rotate
+from chroma.transform import rotate
 
 def uniform_sphere(size=None, dtype=np.double):
     """
diff --git a/scenes/checkerboard.py b/scenes/checkerboard.py
index d2a9ec8..4ad1120 100644
--- a/scenes/checkerboard.py
+++ b/scenes/checkerboard.py
@@ -1,15 +1,9 @@
-import os
-import sys
-
-dir = os.path.split(os.path.realpath(__file__))[0]
-sys.path.append(dir + '/..')
-
 import numpy as np
-from itertoolset import *
-from geometry import Mesh, Solid, Geometry
-from optics import *
-from make import sphere
-from view import buildable
+from chroma.itertoolset import *
+from chroma.geometry import Mesh, Solid, Geometry
+from chroma.optics import *
+from chroma.make import sphere
+from chroma.view import buildable
 
 @buildable('checkerboard_scene')
 def build_checkerboard_scene(checkers_per_side=10, squares_per_checker=50):
diff --git a/solids/__init__.py b/solids/__init__.py
index a24ce9c..c23be98 100644
--- a/solids/__init__.py
+++ b/solids/__init__.py
@@ -1,23 +1,16 @@
 import numpy as np
-
 from pmts import build_pmt, build_light_collector, build_light_collector_from_file, build_pmt_shell
-
-import os
-import sys
-
-dir = os.path.split(os.path.realpath(__file__))[0]
-sys.path.append(dir + '/..')
-
-from optics import *
-from view import buildable
+from chroma.optics import *
+from chroma.view import buildable
+from os.path import dirname
 
 @buildable('12inch_pmt')
 def build_12inch_pmt(outer_material=water, nsteps=16):
-    return build_pmt(dir + '/hamamatsu_12inch.txt', 0.003, outer_material, nsteps)
+    return build_pmt(dirname(__file__) + '/hamamatsu_12inch.txt', 0.003, outer_material, nsteps)
 
 @buildable('12inch_pmt_shell')
 def build_12inch_pmt_shell(outer_material=water, nsteps=16):
-    return build_pmt_shell(dir + '/hamamatsu_12inch.txt')
+    return build_pmt_shell(dirname(__file__) + '/hamamatsu_12inch.txt')
 
 # from Jelena Maricic
 lc_12inch_a = 0.16597
@@ -39,10 +32,10 @@ def build_12inch_pmt_with_lc_hd(outer_material=water, nsteps=128):
 
 @buildable('8inch_pmt')
 def build_8inch_pmt(outer_material=water, nsteps=24):
-    return build_pmt(dir + '/sno_pmt_reduced.txt', 0.003, outer_material, nsteps)
+    return build_pmt(dirname(__file__) + '/sno_pmt_reduced.txt', 0.003, outer_material, nsteps)
 
 @buildable('8inch_pmt_with_lc')
 def build_8inch_pmt_with_lc(outer_material=water, nsteps=24):
     pmt = build_8inch_pmt(outer_material, nsteps)
-    lc = build_light_collector_from_file(dir + '/sno_cone.txt', outer_material, nsteps)
+    lc = build_light_collector_from_file(dirname(__file__) + '/sno_cone.txt', outer_material, nsteps)
     return pmt + lc
diff --git a/solids/pmts.py b/solids/pmts.py
index fcbc8fc..43aed3f 100644
--- a/solids/pmts.py
+++ b/solids/pmts.py
@@ -1,14 +1,8 @@
-import os
-import sys
-
-dir = os.path.split(os.path.realpath(__file__))[0]
-sys.path.append(dir + '/..')
-
 import numpy as np
-from geometry import Solid
-from make import rotate_extrude
-from optics import *
-from tools import read_csv, offset
+from chroma.geometry import Solid
+from chroma.make import rotate_extrude
+from chroma.optics import *
+from chroma.tools import read_csv, offset
 
 def get_lc_profile(radii, a, b, d, rmin, rmax):
     c = -b*np.sqrt(1 - (rmin-d)**2/a**2)
diff --git a/src/__init__.py b/src/__init__.py
index 865d612..f2f5215 100644
--- a/src/__init__.py
+++ b/src/__init__.py
@@ -1,6 +1,7 @@
-import os
+from os.path import dirname
 
-dir = os.path.split(os.path.realpath(__file__))[0]
+with open(dirname(__file__) + '/kernel.cu') as f:
+    kernel = f.read()
 
-kernel = open(dir + '/kernel.cu').read()
-daq = open(dir + '/daq.cu').read()
+with open(dirname(__file__) + '/daq.cu') as f:
+    daq = f.read()
diff --git a/stl.py b/stl.py
index 9bc14de..0806c29 100644
--- a/stl.py
+++ b/stl.py
@@ -1,7 +1,7 @@
 import numpy as np
 import string
 import struct
-from geometry import Mesh
+from chroma.geometry import Mesh
 import bz2
 
 def mesh_from_stl(filename):