allow both preflattened and non-flattened calls

raytrace.py

@@ -111,7 +111,7 @@ def sphere_raytrace_count_between(xx, yy, zz, RR, rho, a, b, c):
 	ndpointer(dtype=numpy.float64, ndim=1, flags='C_CONTIGUOUS'), 
 	]
 
-def grid_raytrace(rho, x, y, z, a, b, c):
+def grid_raytrace_flat(rho_flat, lenrho, x, y, z, a, b, c):
 	"""
 	ray tracing on a grid
 	
@@ -128,16 +128,34 @@ def grid_raytrace(rho, x, y, z, a, b, c):
 	NHout   double array: output; of size m
 	"""
 
-	lenrho = len(rho)
-	#rho_flat = numpy.array(rho.flatten())
 	lena = len(a)
 	NHout = numpy.zeros(shape=lena) - 1
-	args = [rho, lenrho, x, y, z, a, b, c, lena, NHout]
-	r = lib.grid_raytrace(*args)
+	r = lib.grid_raytrace(rho_flat, lenrho, x, y, z, a, b, c, lena, NHout)
 	if r != 0:
 		raise Exception("Calculation failed")
 	return NHout
 
+def grid_raytrace(rho, x, y, z, a, b, c):
+	"""
+	ray tracing on a grid
+	
+	Parameters regarding the spheres:
+	rho:    double array: density for conversion from length to column density
+	 * n:      length of rho
+	x:      double array: start vector
+	y:      double array: start vector
+	z:      double array: start vector
+	a:      double array: direction vector
+	b:      double array: direction vector
+	c:      double array: direction vector
+	 * m:      length of a, b, c, x, y, z
+	NHout   double array: output; of size m
+	"""
+
+	lenrho = len(rho)
+	rho_flat = numpy.array(rho.flatten())
+	return grid_raytrace_flat(rho_flat, lenrho, x, y, z, a, b, c)
+
 
 lib.voronoi_raytrace.argtypes = [
 	ndpointer(dtype=numpy.float64, ndim=1, flags='C_CONTIGUOUS'), 
@@ -364,13 +382,13 @@ def cone_raytrace_finite(thetas, rhos, x, y, z, a, b, c, d):
 	ndpointer(dtype=numpy.float64, ndim=1, flags='C_CONTIGUOUS'), 
 	]
 
-def grid_raytrace_finite(rho, x, y, z, a, b, c, d):
+def grid_raytrace_finite_flat(rho_flat, lenrho, x, y, z, a, b, c, d):
 	"""
 	ray tracing on a grid
 	
 	Parameters regarding the spheres:
 	rho:    double array: density for conversion from length to column density
-	 * n:      length of rho
+	lenrho: length of rho
 	x:      double array: start vector
 	y:      double array: start vector
 	z:      double array: start vector
@@ -384,8 +402,6 @@ def grid_raytrace_finite(rho, x, y, z, a, b, c, d):
 	t       double array: end position along direction vector. -1 if infinite
 	"""
 
-	lenrho = len(rho)
-	#rho_flat = numpy.array(rho.flatten())
 	lena = len(a)
 	assert len(b) == lena
 	assert len(c) == lena
@@ -394,10 +410,33 @@ def grid_raytrace_finite(rho, x, y, z, a, b, c, d):
 	assert len(z) == lena
 	assert len(d) == lena
 	t = numpy.zeros(shape=lena)
-	args = [rho, lenrho, x, y, z, a, b, c, lena, d, t]
-	r = lib.grid_raytrace_finite(*args)
+	r = lib.grid_raytrace_finite(rho_flat, lenrho, x, y, z, a, b, c, lena, d, t)
 	if r != 0:
 		raise Exception("Calculation failed")
 	return t
 
+def grid_raytrace_finite(rho, x, y, z, a, b, c, d):
+	"""
+	ray tracing on a grid
+	
+	Parameters regarding the spheres:
+	rho:    double array: density for conversion from length to column density
+	 * n:      length of rho
+	x:      double array: start vector
+	y:      double array: start vector
+	z:      double array: start vector
+	a:      double array: direction vector
+	b:      double array: direction vector
+	c:      double array: direction vector
+	 * m:      length of a, b, c, x, y, z
+	NHmax   double array: stop at this NH
+	
+	Returns:
+	t       double array: end position along direction vector. -1 if infinite
+	"""
+
+	lenrho = len(rho)
+	rho_flat = numpy.array(rho.flatten())
+	return grid_raytrace_finite_flat(rho_flat, lenrho, x, y, z, a, b, c, d)
+
 

test_grid.py

@@ -9,7 +9,7 @@ def raytrace_grid_finite(x, v, d):
 	x, y, z = numpy.array([x0*1. + 128]), numpy.array([y0*1.+128]), numpy.array([z0*1.+128])
 	a, b, c = numpy.array([dx*1.]), numpy.array([dy*1.]), numpy.array([dz*1.])
 	NHmax = numpy.array([d*1.])
-	r = grid_raytrace_finite(numpy.array(rho.flatten()), x, y, z, a, b, c, NHmax)
+	r = grid_raytrace_finite(rho, x, y, z, a, b, c, NHmax)
 	t = r[0]
 	print('distance:', t, (x, y, z), NHmax)
 	return x0 + dx * t, y0 + dy * t, z0 + dz * t
@@ -232,7 +232,7 @@ def test_random_single():
 	print('direction:', a**2 + b**2 + c**2)
 	NHmax = numpy.array([63.1791943441])
 	rho = numpy.ones((256, 256, 256))
-	t = grid_raytrace_finite(numpy.array(rho.flatten()), x, y, z, a, b, c, NHmax)
+	t = grid_raytrace_finite(rho, x, y, z, a, b, c, NHmax)
 	print('distance:', t)
 	assert (t > 0).all(), t
 	assert numpy.isclose(t, NHmax), (t, NHmax)
@@ -246,11 +246,11 @@ def test_single_small1d_pos():
 	print('direction:', a**2 + b**2 + c**2)
 	NHmax = numpy.array([1000.])
 	rho = numpy.ones((256, 256, 256))
-	t = grid_raytrace_finite(numpy.array(rho.flatten()), x, y, z, a, b, c, NHmax)
+	t = grid_raytrace_finite(rho, x, y, z, a, b, c, NHmax)
 	print('distance:', t)
 	assert t == -1, t
 	NHmax = numpy.array([10.])
-	t = grid_raytrace_finite(numpy.array(rho.flatten()), x, y, z, a, b, c, NHmax)
+	t = grid_raytrace_finite(rho, x, y, z, a, b, c, NHmax)
 	assert (t > 0).all(), t
 	assert numpy.isclose(t, NHmax), (t, NHmax)
 
@@ -260,11 +260,11 @@ def test_single_small1d_neg():
 	print('direction:', a**2 + b**2 + c**2)
 	NHmax = numpy.array([1000.])
 	rho = numpy.ones((256, 256, 256))
-	t = grid_raytrace_finite(numpy.array(rho.flatten()), x, y, z, a, b, c, NHmax)
+	t = grid_raytrace_finite(rho, x, y, z, a, b, c, NHmax)
 	print('distance:', t)
 	assert t == -1, t
 	NHmax = numpy.array([10.])
-	t = grid_raytrace_finite(numpy.array(rho.flatten()), x, y, z, a, b, c, NHmax)
+	t = grid_raytrace_finite(rho, x, y, z, a, b, c, NHmax)
 	assert (t > 0).all(), t
 	assert numpy.isclose(t, NHmax), (t, NHmax)
 
@@ -277,7 +277,7 @@ def test_inhomogeneous_single():
 	a, b, c = numpy.array([-0.321689]), numpy.array([-0.901964]), numpy.array([-0.288056])
 	print('direction:', a**2 + b**2 + c**2)
 	NHmax = numpy.array([(10 + 1 + 1) * 3**0.5])
-	t = grid_raytrace_finite(numpy.array(rho.flatten()), x, y, z, a, b, c, NHmax)
+	t = grid_raytrace_finite(rho, x, y, z, a, b, c, NHmax)
 	print('distance:', t)
 	assert (t == -1).all(), t
 
@@ -290,37 +290,37 @@ def test_density_coord():
 	b = numpy.array([0., 1., 0., 0., -1., 0.])
 	c = numpy.array([0., 0., 1., 0., 0., -1.])
 	NHmax = numpy.array([3.]*6)
-	t = grid_raytrace_finite(numpy.array(rho.flatten()), x, y, z, a, b, c, NHmax)
+	t = grid_raytrace_finite(rho, x, y, z, a, b, c, NHmax)
 	t_expect = numpy.array([-1,-1,3., -1,-1,-1])
 	assert numpy.allclose(t, t_expect), (t, t_expect)
 
 	rho = numpy.zeros((256, 256, 256))
 	rho[128,128:140,128] = 1
-	t = grid_raytrace_finite(numpy.array(rho.flatten()), x, y, z, a, b, c, NHmax)
+	t = grid_raytrace_finite(rho, x, y, z, a, b, c, NHmax)
 	t_expect = numpy.array([-1,3.,-1, -1,-1,-1])
 	assert numpy.allclose(t, t_expect), (t, t_expect)
 
 	rho = numpy.zeros((256, 256, 256))
 	rho[128:140,128,128] = 1
-	t = grid_raytrace_finite(numpy.array(rho.flatten()), x, y, z, a, b, c, NHmax)
+	t = grid_raytrace_finite(rho, x, y, z, a, b, c, NHmax)
 	t_expect = numpy.array([3.,-1,-1, -1,-1,-1])
 	assert numpy.allclose(t, t_expect), (t, t_expect)
 
 	rho = numpy.zeros((256, 256, 256))
 	rho[128,128,120:129] = 1
-	t = grid_raytrace_finite(numpy.array(rho.flatten()), x, y, z, a, b, c, NHmax)
+	t = grid_raytrace_finite(rho, x, y, z, a, b, c, NHmax)
 	t_expect = numpy.array([-1,-1,-1, -1,-1,3.])
 	assert numpy.allclose(t, t_expect), (t, t_expect)
 
 	rho = numpy.zeros((256, 256, 256))
 	rho[128,120:129,128] = 1
-	t = grid_raytrace_finite(numpy.array(rho.flatten()), x, y, z, a, b, c, NHmax)
+	t = grid_raytrace_finite(rho, x, y, z, a, b, c, NHmax)
 	t_expect = numpy.array([-1,-1,-1, -1,3.,-1])
 	assert numpy.allclose(t, t_expect), (t, t_expect)
 
 	rho = numpy.zeros((256, 256, 256))
 	rho[120:129,128,128] = 1
-	t = grid_raytrace_finite(numpy.array(rho.flatten()), x, y, z, a, b, c, NHmax)
+	t = grid_raytrace_finite(rho, x, y, z, a, b, c, NHmax)
 	t_expect = numpy.array([-1,-1,-1, 3.,-1,-1])
 	assert numpy.allclose(t, t_expect), (t, t_expect)
 
@@ -338,15 +338,15 @@ def test_inhomogeneous():
 	a, b, c = numpy.array(a), numpy.array(b), numpy.array(c)
 	# call raytrace_grid_finite_c()
 	NHmax = 4.99 + numpy.zeros((N))
-	t = grid_raytrace_finite(numpy.array(rho.flatten()), x, y, z, a, b, c, NHmax)
+	t = grid_raytrace_finite(rho, x, y, z, a, b, c, NHmax)
 	print('distance:', t)
 	assert numpy.allclose(t, 0.499), t
 
 	# go one further
 	rho[124:132,124:132,124:132] = 1
 	rho[128,128,128] = 10
 	NHmax = 5.99 + numpy.zeros((N))
-	t = grid_raytrace_finite(numpy.array(rho.flatten()), x, y, z, a, b, c, NHmax)
+	t = grid_raytrace_finite(rho, x, y, z, a, b, c, NHmax)
 	print('distance:', t)
 	assert (t <= 1.499).all(), t
 	assert (t >= 0.599).all(), t
@@ -357,7 +357,7 @@ def test_inhomogeneous():
 	rho[128,128,128] = 10
 	NHmax = (10 + 1 + 1) * 3**0.5 + numpy.zeros((N))
 	#NHmax = 400 + numpy.zeros((20))
-	t = grid_raytrace_finite(numpy.array(rho.flatten()), x, y, z, a, b, c, NHmax)
+	t = grid_raytrace_finite(rho, x, y, z, a, b, c, NHmax)
 	print('distance:', t)
 	assert (t == -1).all(), t