Source code for yt.analysis_modules.cosmological_observation.light_cone.light_cone

"""
LightCone class and member functions.



"""

#-----------------------------------------------------------------------------
# Copyright (c) 2013, yt Development Team.
#
# Distributed under the terms of the Modified BSD License.
#
# The full license is in the file COPYING.txt, distributed with this software.
#-----------------------------------------------------------------------------

from yt.utilities.on_demand_imports import _h5py as h5py
import numpy as np
import os

from yt.config import \
    ytcfg
from yt.funcs import \
    ensure_dir, \
    mylog, \
    only_on_root
from yt.analysis_modules.cosmological_observation.cosmology_splice import \
    CosmologySplice
from yt.convenience import \
    load
from yt.utilities.parallel_tools.parallel_analysis_interface import \
    parallel_objects, \
    parallel_root_only
from yt.visualization.image_writer import \
    write_image
from yt.units.yt_array import \
    YTArray
from .light_cone_projection import \
    _light_cone_projection

[docs]class LightCone(CosmologySplice): """ Initialize a LightCone object. Parameters ---------- near_redshift : float The near (lowest) redshift for the light cone. far_redshift : float The far (highest) redshift for the light cone. observer_redshift : float The redshift of the observer. Default: 0.0. use_minimum_datasets : bool If True, the minimum number of datasets is used to connect the initial and final redshift. If false, the light cone solution will contain as many entries as possible within the redshift interval. Default: True. deltaz_min : float Specifies the minimum :math:`\Delta z` between consecutive datasets in the returned list. Default: 0.0. minimum_coherent_box_fraction : float Used with use_minimum_datasets set to False, this parameter specifies the fraction of the total box size to be traversed before rerandomizing the projection axis and center. This was invented to allow light cones with thin slices to sample coherent large scale structure, but in practice does not work so well. Try setting this parameter to 1 and see what happens. Default: 0.0. time_data : bool Whether or not to include time outputs when gathering datasets for time series. Default: True. redshift_data : bool Whether or not to include redshift outputs when gathering datasets for time series. Default: True. find_outputs : bool Whether or not to search for datasets in the current directory. Default: False. set_parameters : dict Dictionary of parameters to attach to ds.parameters. Default: None. output_dir : string The directory in which images and data files will be written. Default: "LC". output_prefix : string The prefix of all images and data files. Default: "LightCone". """ def __init__(self, parameter_filename, simulation_type, near_redshift, far_redshift, observer_redshift=0.0, use_minimum_datasets=True, deltaz_min=0.0, minimum_coherent_box_fraction=0.0, time_data=True, redshift_data=True, find_outputs=False, set_parameters=None, output_dir="LC", output_prefix="LightCone"): self.near_redshift = near_redshift self.far_redshift = far_redshift self.observer_redshift = observer_redshift self.use_minimum_datasets = use_minimum_datasets self.deltaz_min = deltaz_min self.minimum_coherent_box_fraction = minimum_coherent_box_fraction if set_parameters is None: self.set_parameters = {} else: self.set_parameters = set_parameters self.output_dir = output_dir self.output_prefix = output_prefix # Create output directory. ensure_dir(self.output_dir) # Calculate light cone solution. CosmologySplice.__init__(self, parameter_filename, simulation_type, find_outputs=find_outputs) self.light_cone_solution = \ self.create_cosmology_splice(self.near_redshift, self.far_redshift, minimal=self.use_minimum_datasets, deltaz_min=self.deltaz_min, time_data=time_data, redshift_data=redshift_data)
[docs] def calculate_light_cone_solution(self, seed=None, filename=None): r"""Create list of projections to be added together to make the light cone. Several sentences providing an extended description. Refer to variables using back-ticks, e.g. `var`. Parameters ---------- seed : int The seed for the random number generator. Any light cone solution can be reproduced by giving the same random seed. Default: None (each solution will be distinct). filename : string If given, a text file detailing the solution will be written out. Default: None. """ # Don"t use box coherence with maximum projection depths. if self.use_minimum_datasets and \ self.minimum_coherent_box_fraction > 0: mylog.info("Setting minimum_coherent_box_fraction to 0 with " + "minimal light cone.") self.minimum_coherent_box_fraction = 0 # Calculate projection sizes, and get # random projection axes and centers. seed = int(seed) np.random.seed(seed) # For box coherence, keep track of effective depth travelled. box_fraction_used = 0.0 for q in range(len(self.light_cone_solution)): if "previous" in self.light_cone_solution[q]: del self.light_cone_solution[q]["previous"] if "next" in self.light_cone_solution[q]: del self.light_cone_solution[q]["next"] if q == len(self.light_cone_solution) - 1: z_next = self.near_redshift else: z_next = self.light_cone_solution[q+1]["redshift"] # Calculate fraction of box required for a depth of delta z self.light_cone_solution[q]["box_depth_fraction"] = \ (self.cosmology.comoving_radial_distance(z_next, \ self.light_cone_solution[q]["redshift"]) / \ self.simulation.box_size).in_units("") # Calculate fraction of box required for width corresponding to # requested image size. proper_box_size = self.simulation.box_size / \ (1.0 + self.light_cone_solution[q]["redshift"]) self.light_cone_solution[q]["box_width_per_angle"] = \ (self.cosmology.angular_scale(self.observer_redshift, self.light_cone_solution[q]["redshift"]) / proper_box_size).in_units("1 / degree") # Simple error check to make sure more than 100% of box depth # is never required. if self.light_cone_solution[q]["box_depth_fraction"] > 1.0: mylog.error(("Warning: box fraction required to go from " + "z = %f to %f is %f") % (self.light_cone_solution[q]["redshift"], z_next, self.light_cone_solution[q]["box_depth_fraction"])) mylog.error(("Full box delta z is %f, but it is %f to the " + "next data dump.") % (self.light_cone_solution[q]["dz_max"], self.light_cone_solution[q]["redshift"]-z_next)) # Get projection axis and center. # If using box coherence, only get random axis and center if enough # of the box has been used, or if box_fraction_used will be greater # than 1 after this slice. if (q == 0) or (self.minimum_coherent_box_fraction == 0) or \ (box_fraction_used > self.minimum_coherent_box_fraction) or \ (box_fraction_used + self.light_cone_solution[q]["box_depth_fraction"] > 1.0): # Random axis and center. self.light_cone_solution[q]["projection_axis"] = \ np.random.randint(0, 3) self.light_cone_solution[q]["projection_center"] = \ np.random.random(3) box_fraction_used = 0.0 else: # Same axis and center as previous slice, # but with depth center shifted. self.light_cone_solution[q]["projection_axis"] = \ self.light_cone_solution[q-1]["projection_axis"] self.light_cone_solution[q]["projection_center"] = \ self.light_cone_solution[q-1]["projection_center"].copy() self.light_cone_solution[q]["projection_center"]\ [self.light_cone_solution[q]["projection_axis"]] += \ 0.5 * (self.light_cone_solution[q]["box_depth_fraction"] + self.light_cone_solution[q-1]["box_depth_fraction"]) if self.light_cone_solution[q]["projection_center"]\ [self.light_cone_solution[q]["projection_axis"]] >= 1.0: self.light_cone_solution[q]["projection_center"]\ [self.light_cone_solution[q]["projection_axis"]] -= 1.0 box_fraction_used += self.light_cone_solution[q]["box_depth_fraction"] # Write solution to a file. if filename is not None: self._save_light_cone_solution(filename=filename)
[docs] def project_light_cone(self, field_of_view, image_resolution, field, weight_field=None, photon_field=False, save_stack=True, save_final_image=True, save_slice_images=False, cmap_name=None, njobs=1, dynamic=False): r"""Create projections for light cone, then add them together. Parameters ---------- field_of_view : YTQuantity or tuple of (float, str) The field of view of the image and the units. image_resolution : YTQuantity or tuple of (float, str) The size of each image pixel and the units. field : string The projected field. weight_field : string the weight field of the projection. This has the same meaning as in standard projections. Default: None. photon_field : bool if True, the projection data for each slice is decremented by 4 Pi R^2`, where R is the luminosity distance between the observer and the slice redshift. Default: False. save_stack : bool if True, the light cone data including each individual slice is written to an hdf5 file. Default: True. save_final_image : bool if True, save an image of the final light cone projection. Default: True. save_slice_images : bool save images for each individual projection slice. Default: False. cmap_name : string color map for images. Default: your default colormap. njobs : int The number of parallel jobs over which the light cone projection will be split. Choose -1 for one processor per individual projection and 1 to have all processors work together on each projection. Default: 1. dynamic : bool If True, use dynamic load balancing to create the projections. Default: False. """ if cmap_name is None: cmap_name = ytcfg.get("yt", "default_colormap") if isinstance(field_of_view, tuple) and len(field_of_view) == 2: field_of_view = self.simulation.quan(field_of_view[0], field_of_view[1]) elif not isinstance(field_of_view, YTArray): raise RuntimeError("field_of_view argument must be either a YTQauntity " + "or a tuple of type (float, str).") if isinstance(image_resolution, tuple) and len(image_resolution) == 2: image_resolution = self.simulation.quan(image_resolution[0], image_resolution[1]) elif not isinstance(image_resolution, YTArray): raise RuntimeError("image_resolution argument must be either a YTQauntity " + "or a tuple of type (float, str).") # Calculate number of pixels on a side. pixels = int((field_of_view / image_resolution).in_units("")) # Clear projection stack. projection_stack = [] projection_weight_stack = [] if "object" in self.light_cone_solution[-1]: del self.light_cone_solution[-1]["object"] # for q, output in enumerate(self.light_cone_solution): all_storage = {} for my_storage, output in parallel_objects(self.light_cone_solution, storage=all_storage, dynamic=dynamic): output["object"] = load(output["filename"]) output["object"].parameters.update(self.set_parameters) # Calculate fraction of box required for width corresponding to # requested image size. proper_box_size = self.simulation.box_size / \ (1.0 + output["redshift"]) output["box_width_fraction"] = (output["box_width_per_angle"] * field_of_view).in_units("") frb = _light_cone_projection(output, field, pixels, weight_field=weight_field) if photon_field: # Decrement the flux by the luminosity distance. # Assume field in frb is in erg/s/cm^2/Hz dL = self.cosmology.luminosity_distance(self.observer_redshift, output["redshift"]) proper_box_size = self.simulation.box_size / \ (1.0 + output["redshift"]) pixel_area = (proper_box_size.in_cgs() / pixels)**2 #in proper cm^2 factor = pixel_area / (4.0 * np.pi * dL.in_cgs()**2) mylog.info("Distance to slice = %s" % dL) frb[field] *= factor #in erg/s/cm^2/Hz on observer"s image plane. if weight_field is None: my_storage.result = {"field": frb[field]} else: my_storage.result = {"field": (frb[field] * frb["weight_field"]), "weight_field": frb["weight_field"]} del output["object"] # Combine results from each slice. all_slices = list(all_storage.keys()) all_slices.sort() for my_slice in all_slices: if save_slice_images: name = os.path.join(self.output_dir, "%s_%04d_%04d" % (self.output_prefix, my_slice, len(self.light_cone_solution))) if weight_field is None: my_image = all_storage[my_slice]["field"] else: my_image = all_storage[my_slice]["field"] / \ all_storage[my_slice]["weight_field"] only_on_root(write_image, np.log10(my_image), "%s_%s.png" % (name, field), cmap_name=cmap_name) projection_stack.append(all_storage[my_slice]["field"]) if weight_field is not None: projection_weight_stack.append(all_storage[my_slice]["field"]) projection_stack = self.simulation.arr(projection_stack) projection_weight_stack = self.simulation.arr(projection_weight_stack) # Add up slices to make light cone projection. if (weight_field is None): light_cone_projection = projection_stack.sum(axis=0) else: light_cone_projection = \ projection_stack.sum(axis=0) / \ self.simulation.arr(projection_weight_stack).sum(axis=0) filename = os.path.join(self.output_dir, self.output_prefix) # Write image. if save_final_image: only_on_root(write_image, np.log10(light_cone_projection), "%s_%s.png" % (filename, field), cmap_name=cmap_name) # Write stack to hdf5 file. if save_stack: self._save_light_cone_stack(field, weight_field, projection_stack, projection_weight_stack, filename=filename, attrs={"field_of_view": str(field_of_view), "image_resolution": str(image_resolution)})
@parallel_root_only def _save_light_cone_solution(self, filename="light_cone.dat"): "Write out a text file with information on light cone solution." mylog.info("Saving light cone solution to %s." % filename) f = open(filename, "w") f.write("# parameter_filename = %s\n" % self.parameter_filename) f.write("\n") f.write("# Slice Dataset Redshift depth/box " + \ "width/degree axis center\n") for q, output in enumerate(self.light_cone_solution): f.write(("%04d %s %f %f %f %d %f %f %f\n") % (q, output["filename"], output["redshift"], output["box_depth_fraction"], output["box_width_per_angle"], output["projection_axis"], output["projection_center"][0], output["projection_center"][1], output["projection_center"][2])) f.close() @parallel_root_only def _save_light_cone_stack(self, field, weight_field, pstack, wstack, filename=None, attrs=None): "Save the light cone projection stack as a 3d array in and hdf5 file." if attrs is None: attrs = {} # Make list of redshifts to include as a dataset attribute. redshift_list = np.array([my_slice["redshift"] \ for my_slice in self.light_cone_solution]) field_node = "%s_%s" % (field, weight_field) weight_field_node = "weight_field_%s" % weight_field if (filename is None): filename = os.path.join(self.output_dir, "%s_data" % self.output_prefix) if not(filename.endswith(".h5")): filename += ".h5" if pstack.size == 0: mylog.info("save_light_cone_stack: light cone projection is empty.") return mylog.info("Writing light cone data to %s." % filename) fh = h5py.File(filename, "a") if field_node in fh: del fh[field_node] mylog.info("Saving %s to %s." % (field_node, filename)) dataset = fh.create_dataset(field_node, data=pstack) dataset.attrs["units"] = str(pstack.units) dataset.attrs["redshifts"] = redshift_list dataset.attrs["observer_redshift"] = np.float(self.observer_redshift) for key, value in attrs.items(): dataset.attrs[key] = value if wstack.size > 0: if weight_field_node in fh: del fh[weight_field_node] mylog.info("Saving %s to %s." % (weight_field_node, filename)) dataset = fh.create_dataset(weight_field_node, data=wstack) dataset.attrs["units"] = str(wstack.units) dataset.attrs["redshifts"] = redshift_list dataset.attrs["observer_redshift"] = np.float(self.observer_redshift) for key, value in attrs.items(): dataset.attrs[key] = value fh.close()