Merge pull request #91 from brittonsmith/raycuts

Add ability to make spectrum from data container.

doc/source/advanced_spectra.rst

@@ -142,3 +142,35 @@ Giving us:
 To understand how to further customize your spectra, look at the documentation 
 for the :class:`~trident.SpectrumGenerator` and :class:`~trident.LineDatabase`
 classes and other :ref:`API <api-reference>` documentation.
+
+Making Spectra from a Subset of a Ray
+-------------------------------------
+
+The situation may arise where you want to see the spectrum that is generated
+by only a portion of the gas along a line of sight. For example, you may want to
+see the spectrum of only the cold gas. This can be done by creating a
+:class:`~yt.data_objects.selection_data_containers.YTCutRegion` from a loaded ray
+dataset::
+
+    import trident
+    import yt
+
+    ds = yt.load('ray.h5')
+    all_data = ds.all_data()
+    cold_gas = ds.cut_region(all_data, 'obj["gas", "temperature"] < 10000')
+
+    sg = trident.SpectrumGenerator(lambda_min=1200, lambda_max=1225,
+                                   dlambda=0.01)
+
+    # spectrum of entire ray
+    sg.make_spectrum(all_data, lines=['H I 1216'])
+    all_spectrum = sg.flux_field[:]
+
+    # spectrum of cold gas
+    sg.make_spectrum(cold_gas, lines=['H I 1216'])
+    cold_spectrum = sg.flux_field[:]
+
+    trident.plot_spectrum(sg.lambda_field, [all_spectrum, cold_spectrum],
+                          label=['all gas', 'cold gas'], stagger=None)
+
+.. image:: https://raw.githubusercontent.com/trident-project/trident-docs-images/master/spec_cutregion.png

doc/source/conf.py

@@ -34,11 +34,15 @@
 extensions = [
     'sphinx.ext.autodoc',
     'sphinx.ext.autosummary',
+    'sphinx.ext.intersphinx',
     'sphinx.ext.napoleon',
     'sphinx.ext.viewcode',
     'sphinx.ext.imgmath',
 ]
 
+intersphinx_mapping = \
+  {'yt': ('http://yt-project.org/docs/dev/', None),}
+
 autosummary_generate = glob.glob("reference.rst")
 
 # Add any paths that contain templates here, relative to this directory.

tests/test_spectrum_generator.py

@@ -12,6 +12,7 @@
 #-----------------------------------------------------------------------------
 
 from __future__ import absolute_import
+from trident.plotting import plot_spectrum
 from trident.utilities import make_onezone_ray
 from trident.spectrum_generator import \
     SpectrumGenerator, \
@@ -35,7 +36,6 @@ def test_save_load_spectrum_ascii():
     sg = load_spectrum(os.path.join(dirpath, 'spec.txt'))
     sg.plot_spectrum(filename=os.path.join(dirpath, 'spec.png'))
     shutil.rmtree(dirpath)
-    assert True
 
 def test_save_load_spectrum_hdf5():
     """
@@ -52,7 +52,6 @@ def test_save_load_spectrum_hdf5():
     sg = load_spectrum(os.path.join(dirpath, 'spec.h5'))
     sg.plot_spectrum(filename=os.path.join(dirpath, 'spec.png'))
     shutil.rmtree(dirpath)
-    assert True
 
 def test_save_load_spectrum_fits():
     """
@@ -69,7 +68,6 @@ def test_save_load_spectrum_fits():
     sg = load_spectrum(os.path.join(dirpath, 'spec.fits'))
     sg.plot_spectrum(filename=os.path.join(dirpath, 'spec.png'))
     shutil.rmtree(dirpath)
-    assert True
 
 def test_create_spectrum_all_lines():
     """
@@ -82,7 +80,6 @@ def test_create_spectrum_all_lines():
     sg.make_spectrum(ray, lines='all')
     sg.plot_spectrum(os.path.join(dirpath, 'spec.png'))
     shutil.rmtree(dirpath)
-    assert True
 
 def test_create_spectrum_H_O_lines():
     """
@@ -95,7 +92,6 @@ def test_create_spectrum_H_O_lines():
     sg.make_spectrum(ray, lines=['H', 'O'])
     sg.plot_spectrum(os.path.join(dirpath, 'spec.png'))
     shutil.rmtree(dirpath)
-    assert True
 
 def test_create_spectrum_H_lines_no_continuum():
     """
@@ -109,4 +105,34 @@ def test_create_spectrum_H_lines_no_continuum():
     sg.make_spectrum(ray, lines=['H I'], ly_continuum=False)
     sg.plot_spectrum(os.path.join(dirpath, 'spec.png'))
     shutil.rmtree(dirpath)
-    assert True
+
+def test_input_types():
+    """
+    Test that spectra can be generated from ray file, dataset, or
+    data container.
+    """
+
+    dirpath = tempfile.mkdtemp()
+    filename = os.path.join(dirpath, 'ray.h5')
+    ray = make_onezone_ray(filename=filename)
+
+    sg = SpectrumGenerator(lambda_min=1200, lambda_max=1300, dlambda=0.5)
+    spectra = []
+
+    # from file
+    sg.make_spectrum(filename, lines=['H I'], ly_continuum=False)
+    spectra.append(sg.flux_field[:])
+
+    # from dataset
+    sg.make_spectrum(ray, lines=['H I'], ly_continuum=False)
+    spectra.append(sg.flux_field[:])
+    assert (spectra[0] == spectra[1]).all()
+
+    # from data container
+    sg.make_spectrum(ray.all_data(), lines=['H I'], ly_continuum=False)
+    spectra.append(sg.flux_field[:])
+    assert (spectra[0] == spectra[2]).all()
+
+    plot_spectrum(sg.lambda_field, spectra,
+                  filename=os.path.join(dirpath, 'spec.png'))
+    shutil.rmtree(dirpath)

trident/absorption_spectrum/absorption_spectrum.py

@@ -22,6 +22,10 @@
 from trident.absorption_spectrum.absorption_line import \
     tau_profile
 
+from yt.data_objects.data_containers import \
+    YTDataContainer
+from yt.data_objects.static_output import \
+    Dataset
 from yt.extern.six import string_types
 from yt.convenience import load
 from yt.funcs import get_pbar, mylog
@@ -153,7 +157,7 @@ def add_continuum(self, label, field_name, wavelength,
                                     'normalization': normalization,
                                     'index': index})
 
-    def make_spectrum(self, input_file, output_file=None,
+    def make_spectrum(self, input_object, output_file=None,
                       line_list_file=None, output_absorbers_file=None,
                       use_peculiar_velocity=True,
                       store_observables=False,
@@ -164,9 +168,12 @@ def make_spectrum(self, input_file, output_file=None,
 
         **Parameters**
 
-        :input_file: string or dataset
+        :input_object: string, dataset, or data container
 
-           path to input ray data or a loaded ray dataset
+           If a string, the path to the ray dataset. As a dataset,
+           this is the ray dataset loaded by yt. As a data container,
+           this is a data object created from a ray dataset, such as
+           a cut region.
 
         :output_file: optional, string
 
@@ -266,18 +273,22 @@ def make_spectrum(self, input_file, output_file=None,
                 input_fields.append(feature['field_name'])
                 field_units[feature["field_name"]] = "cm**-3"
 
-        if isinstance(input_file, string_types):
-            input_ds = load(input_file)
-        else:
-            input_ds = input_file
-        field_data = input_ds.all_data()
+        if isinstance(input_object, string_types):
+            input_ds = load(input_object)
+            field_data = input_ds.all_data()
+        elif isinstance(input_object, Dataset):
+            input_ds = input_object
+            field_data = input_ds.all_data()
+        elif isinstance(input_object, YTDataContainer):
+            input_ds = input_object.ds
+            field_data = input_object
 
         # temperature field required to calculate voigt profile widths
         if ('temperature' not in input_ds.derived_field_list) and \
            (('gas', 'temperature') not in input_ds.derived_field_list):
             raise RuntimeError(
                 "('gas', 'temperature') field required to be present in %s "
-                "for AbsorptionSpectrum to function." % input_file)
+                "for AbsorptionSpectrum to function." % str(input_object))
 
         self.tau_field = np.zeros(self.lambda_field.size)
         self.absorbers_list = []

trident/plotting.py

@@ -184,28 +184,30 @@ def plot_spectrum(wavelength, flux, filename="spectrum.png",
     my_axes = figure.add_axes((left_side, bottom_side, panel_width, panel_height))
 
     # Are we overplotting several spectra?  or just one?
-    if not (isinstance(wavelength, list) and isinstance(flux, list)):
-        wavelengths = [wavelength]
-        fluxs = [flux]
-        labels = [label]
-        steps = [step]
+    if isinstance(flux, list):
+        fluxs = flux
     else:
+        fluxs = [flux]
+
+    if isinstance(wavelength, list):
         wavelengths = wavelength
-        fluxs = flux
-        if label is not None:
-            labels = label
-        else:
-            labels = [None]*len(fluxs)
-        if step is not None:
-            steps = step
-        else:
-            steps = [None]*len(fluxs)
+    else:
+        wavelengths = [wavelength]*len(fluxs)
+
+    if isinstance(step, list):
+        steps = step
+    else:
+        steps = [step]*len(fluxs)
+
+    if isinstance(label, list):
+        labels = label
+    else:
+        labels = [label]*len(fluxs)
 
     # A running maximum of flux for use in ylim scaling in final plot
     max_flux = 0.
 
     for i, (wavelength, flux) in enumerate(zip(wavelengths, fluxs)):
-
         # Do we stagger the fluxes?
         if stagger is not None:
             flux -= stagger * i

trident/spectrum_generator.py

@@ -18,6 +18,10 @@
     AbsorptionSpectrum
 from yt.convenience import \
     load
+from yt.data_objects.data_containers import \
+    YTDataContainer
+from yt.data_objects.static_output import \
+    Dataset
 from yt.funcs import \
     mylog, \
     YTArray
@@ -232,9 +236,12 @@ def make_spectrum(self, ray, lines='all',
 
         **Parameters**
 
-        :ray: string or dataset
+        :ray: string, dataset, or data container
 
-            Ray dataset filename or a loaded ray dataset
+            If a string, the path to the ray dataset. As a dataset,
+            this is the ray dataset loaded by yt. As a data container,
+            this is a data object created from a ray dataset, such as
+            a cut region.
 
         :lines: list of strings
 
@@ -334,7 +341,13 @@ def make_spectrum(self, ray, lines='all',
 
         if isinstance(ray, str):
             ray = load(ray)
-        ad = ray.all_data()
+        if isinstance(ray, Dataset):
+            ad = ray.all_data()
+        elif isinstance(ray, YTDataContainer):
+            ad = ray
+            ray = ad.ds
+        else:
+            raise RuntimeError("Unrecognized ray type.")
 
         # Clear out any previous spectrum that existed first
         self.clear_spectrum()
@@ -396,7 +409,7 @@ def make_spectrum(self, ray, lines='all',
         if len(H_lines) > 0 and ly_continuum:
             self.add_continuum('Ly C', H_lines[0].field, 912.32336, 1.6e17, 3.0)
 
-        AbsorptionSpectrum.make_spectrum(self, ray,
+        AbsorptionSpectrum.make_spectrum(self, ad,
                                          output_file=None,
                                          line_list_file=None,
                                          output_absorbers_file=output_absorbers_file,