Create synthetic data cube for testing astrodendro

gistfile1.py

#!/usr/bin/env python
import numpy as np
import pyfits
import math
import random

random.seed(5678) # consistent random seed
ri = random.randint

def export_data(data, desc, filename):
	hdu = pyfits.PrimaryHDU(data)
	header = hdu.header
	header.update('OBJECT', "Synthetic Data")
	header.update('LINENAME', desc)
	print("Exporting {desc} to {f}".format(desc=desc, f=filename))
	hdu.writeto(filename)

# Set up desired data cube here
base_shape = (100,100,500) # Size of the data cube
#spheres - list of tuples of (x,y,z,radius,peak,min)
spheres = [
	( 50, 50, 50, 48,3,3), # Sphere 0
		#random fuzzy spheres get added inside this big sphere
]

# Add some strong random spheres inside main sphere:
for _ in range(0,20):
	x,y,z = 50,50,50
	radius = ri(1,40)
	spheres.append(( (x+ri(-radius,radius)), (y+ri(-radius,radius)), (z+ri(-radius,radius)), radius, 1.5,0))

# Add some weak random spheres inside sphere 7 at 75,75,75:
for _ in range(0,20):
	x,y,z = 75,75,75
	radius = ri(1,20)
	spheres.append(( (x+ri(-radius,radius)), (y+ri(-radius,radius)), (z+ri(-radius,radius)), radius, 1,0))

if __name__ == '__main__':
	print("Creating base data cube")
	base = np.ndarray(shape=base_shape,dtype=float, order='F')
	
	prog = -1 # for progress marker
	
	# Iterate over the array:
	it = np.nditer(base, flags=['multi_index'], op_flags=['writeonly'])
	while not it.finished:
		prog_now = it.iterindex*100//it.itersize
		if prog_now > prog:
			prog = prog_now
			if prog %10 == 0:
				print "{0}%".format(prog)
		coords = it.multi_index
		val = 0 # Initial state of this pixel:
		# Go through all the spheres that this pixel might touch:
		for (x,y,z,radius,peak,edgeval) in spheres:
			if abs(coords[0]-x) > radius or abs(coords[2]-z) > radius:
				continue # Simple faster check on X and Z to speed up computation
			dist = math.sqrt( (coords[0]-x)**2.0 + (coords[1]-y)**2.0 + (coords[2]-z)**2.0 )
			if dist < radius:
				val += peak - (peak-edgeval)*(dist/radius)
		it[0] = val
		it.iternext()
	
	print("100%")
	export_data(base, "base", "base.fits")
	print("All done.")