first implementation of direct convergence integration

lenstools/scripts/raytracing.py

@@ -56,13 +56,13 @@ def convergence_callback(jacobian,tracer,k,realization,angle,map_batch,settings)
 #######Single redshift ray tracing##############
 ################################################
 
-def singleRedshift(pool,batch,settings,id):
+def singleRedshift(pool,batch,settings,batch_id):
 
 	#Safety check
 	assert isinstance(pool,MPIWhirlPool) or (pool is None)
 	assert isinstance(batch,SimulationBatch)
 
-	parts = id.split("|")
+	parts = batch_id.split("|")
 
 	if len(parts)==2:
 
@@ -94,7 +94,7 @@ def singleRedshift(pool,batch,settings,id):
 	else:
 
 		if (pool is None) or (pool.is_master()):
-			logdriver.error("Format error in {0}: too many '|'".format(id))
+			logdriver.error("Format error in {0}: too many '|'".format(batch_id))
 		sys.exit(1)
 
 
@@ -353,6 +353,300 @@ def singleRedshift(pool,batch,settings,id):
 		now = time.time()
 		logdriver.info("Total runtime {0:.3f}s".format(now-begin))
 
+############################################################################################################################################################################
+
+##########################################################
+#######Direct integration of the convergence##############
+##########################################################
+
+def convergenceDirect(pool,batch,settings,batch_id):
+
+	#Safety check
+	assert isinstance(pool,MPIWhirlPool) or (pool is None)
+	assert isinstance(batch,SimulationBatch)
+
+	parts = batch_id.split("|")
+
+	if len(parts)==2:
+
+		assert isinstance(settings,MapSettings)
+
+		#Separate the id into cosmo_id and geometry_id
+		cosmo_id,geometry_id = parts
+
+		#Get a handle on the model
+		model = batch.getModel(cosmo_id)
+
+		#Get the corresponding simulation collection and map batch handlers
+		collection = [model.getCollection(geometry_id)]
+		map_batch = collection[0].getMapSet(settings.directory_name)
+		cut_redshifts = np.array([0.0])
+
+	elif len(parts)==1:
+
+		assert isinstance(settings,TelescopicMapSettings)
+
+		#Get a handle on the model
+		model = batch.getModel(parts[0])
+
+		#Get the corresponding simulation collection and map batch handlers
+		map_batch = model.getTelescopicMapSet(settings.directory_name)
+		collection = map_batch.mapcollections
+		cut_redshifts = map_batch.redshifts
+
+	else:
+
+		if (pool is None) or (pool.is_master()):
+			logdriver.error("Format error in {0}: too many '|'".format(batch_id))
+		sys.exit(1)
+
+
+	#Override the settings with the previously pickled ones, if prompted by user
+	if settings.override_with_local:
+
+		local_settings_file = os.path.join(map_batch.home_subdir,"settings.p")
+		settings = MapSettings.read(local_settings_file)
+		assert isinstance(settings,MapSettings)
+
+		if (pool is None) or (pool.is_master()):
+			logdriver.warning("Overriding settings with the previously pickled ones at {0}".format(local_settings_file))
+
+	##################################################################
+	##################Settings read###################################
+	##################################################################
+
+	#Set random seed to generate the realizations
+	if pool is not None:
+		np.random.seed(settings.seed + pool.rank)
+	else:
+		np.random.seed(settings.seed)
+
+	#Read map angle,redshift and resolution from the settings
+	map_angle = settings.map_angle
+	source_redshift = settings.source_redshift
+	resolution = settings.map_resolution
+
+	if len(parts)==2:
+
+		#########################
+		#Use a single collection#
+		#########################
+
+		#Read the plane set we should use
+		plane_set = (settings.plane_set,)
+
+		#Randomization
+		nbody_realizations = (settings.mix_nbody_realizations,)
+		cut_points = (settings.mix_cut_points,)
+		normals = (settings.mix_normals,)
+		map_realizations = settings.lens_map_realizations
+
+	elif len(parts)==1:
+
+		#######################
+		#####Telescopic########
+		#######################
+
+		#Check that we have enough info
+		for attr_name in ["plane_set","mix_nbody_realizations","mix_cut_points","mix_normals"]:
+			if len(getattr(settings,attr_name))!=len(collection):
+				if (pool is None) or (pool.is_master()):
+					logdriver.error("You need to specify a setting {0} for each collection!".format(attr_name))
+				sys.exit(1)
+
+		#Read the plane set we should use
+		plane_set = settings.plane_set
+
+		#Randomization
+		nbody_realizations = settings.mix_nbody_realizations
+		cut_points = settings.mix_cut_points
+		normals = settings.mix_normals
+		map_realizations = settings.lens_map_realizations
+
+
+
+	#Decide which map realizations this MPI task will take care of (if pool is None, all of them)
+	try:
+		realization_offset = settings.first_realization - 1
+	except AttributeError:
+		realization_offset = 0
+
+	if pool is None:
+		first_map_realization = 0 + realization_offset
+		last_map_realization = map_realizations + realization_offset
+		realizations_per_task = map_realizations
+		logdriver.debug("Generating lensing map realizations from {0} to {1}".format(first_map_realization+1,last_map_realization))
+	else:
+		assert map_realizations%(pool.size+1)==0,"Perfect load-balancing enforced, map_realizations must be a multiple of the number of MPI tasks!"
+		realizations_per_task = map_realizations//(pool.size+1)
+		first_map_realization = realizations_per_task*pool.rank + realization_offset
+		last_map_realization = realizations_per_task*(pool.rank+1) + realization_offset
+		logdriver.debug("Task {0} will generate lensing map realizations from {1} to {2}".format(pool.rank,first_map_realization+1,last_map_realization))
+
+	#Planes will be read from this path
+	plane_path = os.path.join("{0}","ic{1}","{2}")
+
+	if (pool is None) or (pool.is_master()):
+		for c,coll in enumerate(collection):
+			logdriver.info("Reading planes from {0}".format(plane_path.format(coll.storage_subdir,"-".join([str(n) for n in nbody_realizations[c]]),plane_set[c])))
+
+	#Plane info file is the same for all collections
+	if (not hasattr(settings,"plane_info_file")) or (settings.plane_info_file is None):
+		info_filename = batch.syshandler.map(os.path.join(plane_path.format(collection[0].storage_subdir,nbody_realizations[0][0],plane_set[0]),"info.txt"))
+	else:
+		info_filename = settings.plane_info_file
+
+	if (pool is None) or (pool.is_master()):
+		logdriver.info("Reading lens plane summary information from {0}".format(info_filename))
+
+	#Read how many snapshots are available
+	with open(info_filename,"r") as infofile:
+		num_snapshots = len(infofile.readlines())
+
+	#Save path for the maps
+	save_path = map_batch.storage_subdir
+
+	if (pool is None) or (pool.is_master()):
+		logdriver.info("Lensing maps will be saved to {0}".format(save_path))
+
+	begin = time.time()
+
+	#Log initial memory load
+	peak_memory_task,peak_memory_all = peakMemory(),peakMemoryAll(pool)
+	if (pool is None) or (pool.is_master()):
+		logstderr.info("Initial memory usage: {0:.3f} (task), {1[0]:.3f} (all {1[1]} tasks)".format(peak_memory_task,peak_memory_all))
+
+	#We need one of these for cycles for each map random realization
+	for rloc,r in enumerate(range(first_map_realization,last_map_realization)):
+
+		#Instantiate the RayTracer
+		tracer = RayTracer()
+
+		#Force garbage collection
+		gc.collect()
+
+		#Start timestep
+		start = time.time()
+		last_timestamp = start
+
+		#############################################################
+		###############Add the lenses to the system##################
+		#############################################################
+
+		#Open the info file to read the lens specifications (assume the info file is the same for all nbody realizations)
+		infofile = open(info_filename,"r")
+
+		#Read the info file line by line, and decide if we should add the particular lens corresponding to that line or not
+		for s in range(num_snapshots):
+
+			#Read the line
+			line = infofile.readline().strip("\n")
+
+			#Stop if there is nothing more to read
+			if line=="":
+				break
+
+			#Split the line in snapshot,distance,redshift
+			line = line.split(",")
+
+			snapshot_number = int(line[0].split("=")[1])
+
+			distance,unit = line[1].split("=")[1].split(" ")
+			if unit=="Mpc/h":
+				distance = float(distance)*model.Mpc_over_h
+			else:
+				distance = float(distance)*getattr(u,"unit")
+
+			lens_redshift = float(line[2].split("=")[1])
+
+			#Select the right collection
+			for n,z in enumerate(cut_redshifts):
+				if lens_redshift>=z:
+					c = n
+
+			#Randomization of planes
+			nbody = np.random.randint(low=0,high=len(nbody_realizations[c]))
+			cut = np.random.randint(low=0,high=len(cut_points[c]))
+			normal = np.random.randint(low=0,high=len(normals[c]))
+
+			#Log to user
+			logdriver.debug("Realization,snapshot=({0},{1}) --> NbodyIC,cut_point,normal=({2},{3},{4})".format(r,s,nbody_realizations[c][nbody],cut_points[c][cut],normals[c][normal]))
+
+			#Add the lens to the system
+			logdriver.info("Adding lens at redshift {0}".format(lens_redshift))
+			plane_name = batch.syshandler.map(os.path.join(plane_path.format(collection[c].storage_subdir,nbody_realizations[c][nbody],plane_set[c]),settings.plane_name_format.format(snapshot_number,cut_points[c][cut],normals[c][normal],settings.plane_format)))
+			tracer.addLens((plane_name,distance,lens_redshift))
+
+		#Close the infofile
+		infofile.close()
+
+		now = time.time()
+		logdriver.info("Plane specification reading completed in {0:.3f}s".format(now-start))
+		last_timestamp = now
+
+		#Rearrange the lenses according to redshift and roll them randomly along the axes
+		tracer.reorderLenses()
+
+		now = time.time()
+		logdriver.info("Reordering completed in {0:.3f}s".format(now-last_timestamp))
+		last_timestamp = now
+
+		#Start a bucket of light rays from a regular grid of initial positions
+		b = np.linspace(0.0,map_angle.value,resolution)
+		xx,yy = np.meshgrid(b,b)
+		pos = np.array([xx,yy]) * map_angle.unit
+
+		if settings.tomographic_convergence:
+
+			#Trace the ray deflections and save the convergence at every step
+			raise NotImplementedError("Tomographic convergence direct not implemented!")
+
+		else:
+
+			#Integrate the density on the line of sight
+			convergence = tracer.convergenceDirect(pos,z=source_redshift,save_intermediate=False)
+
+			now = time.time()
+			logdriver.info("Direct density integration for realization {0} completed in {1:.3f}s".format(r+1,now-last_timestamp))
+			last_timestamp = now
+
+			#Compute shear,convergence and omega from the jacobians
+			if settings.convergence:
+
+				convMap = ConvergenceMap(data=convergence,angle=map_angle)
+				savename = batch.syshandler.map(os.path.join(save_path,"WLconv_z{0:.2f}_{1:04d}r.{2}".format(source_redshift,r+1,settings.format)))
+				logdriver.info("Saving convergence map to {0}".format(savename)) 
+				convMap.save(savename)
+				logdriver.debug("Saved convergence map to {0}".format(savename)) 
+
+			##############################################################################################################################
+
+			if settings.shear:
+				raise NotImplementedError("Shear calculation from E mode not implemented")
+
+			##############################################################################################################################
+
+			if settings.omega:
+				pass
+
+		now = time.time()
+
+		#Log peak memory usage to stdout
+		peak_memory_task,peak_memory_all = peakMemory(),peakMemoryAll(pool)
+		logdriver.info("Weak lensing calculations for realization {0} completed in {1:.3f}s".format(r+1,now-last_timestamp))
+		logdriver.info("Peak memory usage: {0:.3f} (task), {1[0]:.3f} (all {1[1]} tasks)".format(peak_memory_task,peak_memory_all))
+
+		#Log progress and peak memory usage to stderr
+		if (pool is None) or (pool.is_master()):
+			logstderr.info("Progress: {0:.2f}%, peak memory usage: {1:.3f} (task), {2[0]:.3f} (all {2[1]} tasks)".format(100*(rloc+1.)/realizations_per_task,peak_memory_task,peak_memory_all))
+
+	#Safety sync barrier
+	if pool is not None:
+		pool.comm.Barrier()
+
+	if (pool is None) or (pool.is_master()):	
+		now = time.time()
+		logdriver.info("Total runtime {0:.3f}s".format(now-begin))
 
 ############################################################################################################################################################################
 
@@ -361,15 +655,15 @@ def singleRedshift(pool,batch,settings,id):
 #######Galaxy catalog ray tracing##############
 ###############################################
 
-def simulatedCatalog(pool,batch,settings,id):
+def simulatedCatalog(pool,batch,settings,batch_id):
 
 	#Safety check
 	assert isinstance(pool,MPIWhirlPool) or (pool is None)
 	assert isinstance(batch,SimulationBatch)
 	assert isinstance(settings,CatalogSettings)
 
 	#Separate the id into cosmo_id and geometry_id
-	cosmo_id,geometry_id = id.split("|")
+	cosmo_id,geometry_id = batch_id.split("|")
 
 	#Get a handle on the model
 	model = batch.getModel(cosmo_id)