diff options
| -rwxr-xr-x | benchmark.py | 7 | ||||
| -rw-r--r-- | transform.py | 18 |
2 files changed, 21 insertions, 4 deletions
diff --git a/benchmark.py b/benchmark.py index 31c145c..329986c 100755 --- a/benchmark.py +++ b/benchmark.py @@ -14,6 +14,7 @@ from chroma import sample from chroma import generator from chroma import tools from chroma import optics +from chroma.transform import normalize # Generator processes need to fork BEFORE the GPU context is setup g4generator = generator.photon.G4ParallelGenerator(4, optics.water_wcsim) @@ -46,8 +47,7 @@ def load_photons(number=100, nphotons=500000): per second.""" pos = np.zeros((nphotons,3)) dir = sample.uniform_sphere(nphotons) - pol = np.cross(sample.uniform_sphere(nphotons), dir) - pol /= np.apply_along_axis(np.linalg.norm,1,pol)[:,np.newaxis] + pol = normalize(np.cross(sample.uniform_sphere(nphotons), dir)) wavelengths = np.random.uniform(400,800,size=nphotons) photons = event.Photons(pos, dir, pol, wavelengths) @@ -73,8 +73,7 @@ def propagate(gpu_geometry, number=10, nphotons=500000, nthreads_per_block=64, for i in tools.progress(range(number)): pos = np.zeros((nphotons,3)) dir = sample.uniform_sphere(nphotons) - pol = np.cross(sample.uniform_sphere(nphotons), dir) - pol /= np.apply_along_axis(np.linalg.norm,1,pol)[:,np.newaxis] + pol = normalize(np.cross(sample.uniform_sphere(nphotons), dir)) wavelengths = np.random.uniform(400,800,size=nphotons) photons = event.Photons(pos, dir, pol, wavelengths) gpu_photons = gpu.GPUPhotons(photons) diff --git a/transform.py b/transform.py index 4793c0f..299da46 100644 --- a/transform.py +++ b/transform.py @@ -18,3 +18,21 @@ def rotate(x, phi, n): around the axis `n` (when looking towards +infinity). """ return np.inner(np.asarray(x),make_rotation_matrix(phi, n)) + +def normalize(x): + "Returns unit vectors in the direction of `x`." + x = np.asarray(x) + + if x.shape[-1] != 3: + raise ValueError('dimension of last axis must be 3.') + + d = len(x.shape) + + if d == 1: + norm = np.sqrt(x.dot(x)) + elif d == 2: + norm = np.sqrt(np.sum(x*x, axis=1))[:,np.newaxis] + else: + raise ValueError('len(`x`.shape) must be zero or one.') + + return x/norm |