Merge pull request #17 from chummels/add_tests

Add some additional tests for particle-based codes

examples/gizmo_script.py

@@ -0,0 +1,53 @@
+# Example of a currently working script using trident to generate a COS
+# spectrum from a Gizmo dataset.  Because the sightline passes through the
+# center of the galaxy, it should produce a DLA with lots of absorption lines.
+# You can get this dataset at: # http://yt-project.org/data/
+#
+# wget http://yt-project.org/data/FIRE_M12i_ref11.tar.gz
+
+import yt
+import trident
+
+# Set the dataset filename, load it into yt and define the trajectory
+# of the LightRay.  This uses the maximum density location as the one end of
+# the ray.  Define desired spectral features to include # all H, C, N, O, and 
+# Mg lines.
+fn = 'FIRE_M12i_ref11/snapshot_600.hdf5'
+ds = yt.load(fn)
+_, c = ds.find_max(('gas', 'density'))
+ray_start = c
+ray_end = ds.domain_right_edge
+line_list = ['H', 'C', 'N', 'O', 'Mg']
+
+# Make a LightRay object including all necessary fields so you can add
+# all H, C, N, O, and Mg fields to the resulting spectrum from your dataset.
+# Save LightRay to ray.h5 and use it locally as ray object.
+# Note: We use PartType0, the gas particle field type as our ftype!
+ray = trident.make_simple_ray(ds, start_position=ray_start,
+                              end_position=ray_end, data_filename='ray.h5',
+                              lines=line_list, ftype='PartType0')
+
+# Create a projection of the dataset in density along the x axis,
+# overplot the trajectory of the ray, and save it.
+# Note that this is using the 'gas' field type, which is the field deposited
+# to the grid.
+p = yt.ProjectionPlot(ds, 'x', ('gas', 'density'))
+p.annotate_ray(ray, arrow=True)
+p.save('projection.png')
+
+# Now use the ray object to actually generate an absorption spectrum
+# Use the settings (spectral range, LSF, and spectral resolution) for COS
+# And save it as an output text file and plot it to an image.
+sg = trident.SpectrumGenerator('COS')
+sg.make_spectrum(ray, lines=line_list)
+sg.save_spectrum('spec_raw.txt')
+sg.plot_spectrum('spec_raw.png')
+
+# "Final" spectrum with added quasar, MW background, applied line-spread 
+# function, and added gaussian noise (SNR=30)
+sg.add_qso_spectrum()
+sg.add_milky_way_foreground()
+sg.apply_lsf()
+sg.add_gaussian_noise(30)
+sg.save_spectrum('spec_final.txt')
+sg.plot_spectrum('spec_final.png')

tests/test_ion_balance.py

@@ -27,9 +27,6 @@
 import shutil
 
 import numpy as np
-# TODO
-# Make particle fields added work; cannot see 'gas' counterparts
-# and grid*particle error problems when looking at particle fields
 
 def test_add_ion_fraction_field_to_grid_ds():
     """
@@ -259,54 +256,36 @@ def test_add_all_ion_fields_to_amr_ds():
         yt.SlicePlot(ds, 'x', field).save(dirpath)
     shutil.rmtree(dirpath)
 
-#def test_add_all_ion_fields_to_particle_ds():
-#    """
-#    Test to add various ion fields
-#    """
-#    ds = fake_particle_ds(fields=('particle_mass',
-#                                  'particle_position_x',
-#                                  'particle_position_y',
-#                                  'particle_position_z',
-#                                  "particle_velocity_x",
-#                                  "particle_velocity_y",
-#                                  "particle_velocity_z",
-#                                  "density",
-#                                  "temperature",
-#                                  "metallicity",
-#                                  "smoothing_length"),
-#                           units=('g',
-#                                  'cm',
-#                                  'cm',
-#                                  'cm',
-#                                  'cm/s',
-#                                  'cm/s',
-#                                  'cm/s',
-#                                  'g/cm**3',
-#                                  'K',
-#                                  '',
-#                                  'cm'),
-#                           negative=(False,
-#                                     False,
-#                                     False,
-#                                     False,
-#                                     True,
-#                                     True,
-#                                     True,
-#                                     False,
-#                                     False,
-#                                     False,
-#                                     False))
-#    ftype = 'io',
-#    ad = ds.all_data()
-#    tri.add_ion_fields(ds, ['H', 'N IV'], ftype='io')
-#    #len(ad[('gas', 'S_p3_ion_fraction')])
-#    #len(ad[('gas', 'H_mass')])
-#    #import pdb; pdb.set_trace()
-#    fields = ['H_ion_fraction', 'H_p0_number_density', 'O_p5_mass', 'N_p4_density']
-#    #dirpath = tempfile.mkdtemp()
-#    #for field in fields:
-#    #    field = (ftype, field)
-#    #    assert field in ds.derived_field_list
-#    #    assert isinstance(ad[field], np.ndarray)
-#    #    yt.SlicePlot(ds, 'x', field).save(dirpath)
-#    #shutil.rmtree(dirpath)
+def test_add_ion_fields_to_enzo():
+    """
+    Test to add various ion fields to Enzo dataset and slice on them
+    """
+    ds = yt.load('IsolatedGalaxy/galaxy0030/galaxy0030')
+    tri.add_ion_fields(ds, ['H', 'O VI'], ftype='gas')
+    ad = ds.all_data()
+    fields = ['H_p0_number_density', 'O_p5_density']
+    # Assure that a sampling of fields are added and can be sliced
+    dirpath = tempfile.mkdtemp()
+    for field in fields:
+        field = ('gas', field)
+        assert field in ds.derived_field_list
+        assert isinstance(ad[field], np.ndarray)
+        yt.SlicePlot(ds, 'x', field).save(dirpath)
+    shutil.rmtree(dirpath)
+
+def test_add_ion_fields_to_gizmo():
+    """
+    Test to add various ion fields to gizmo dataset and slice on them
+    """
+    ds = yt.load('FIRE_M12i_ref11/snapshot_600.hdf5')
+    tri.add_ion_fields(ds, ['H', 'O VI'], ftype='PartType0')
+    ad = ds.all_data()
+    fields = ['H_ion_fraction', 'O_p5_mass']
+    # Assure that a sampling of fields are added and can be sliced
+    dirpath = tempfile.mkdtemp()
+    for field in fields:
+        field = ('gas', field)
+        assert field in ds.derived_field_list
+        assert isinstance(ad[field], np.ndarray)
+        yt.SlicePlot(ds, 'x', field).save(dirpath)
+    shutil.rmtree(dirpath)

tests/test_pipelines.py

@@ -31,6 +31,7 @@
 
 # Datasets required for running some answer tests
 enzo_small = os.path.join(answer_test_data_dir, 'enzo_small')
+FIRE = os.path.join(answer_test_data_dir, 'FIRE_M12I_ref11/snapshot_600.hdf5')
 err_precision = 8
 
 def test_verify():
@@ -159,7 +160,7 @@ def test_enzo_small_compound(self):
         sg.add_gaussian_noise(30, seed=1)
         final_file = 'enzo_small_compound_spec_final.h5'
         sg.save_spectrum(final_file)
-        final_file_compare = os.path.join(final_file)
+        final_file_compare = os.path.join(test_results_dir, final_file)
         sg.plot_spectrum('enzo_small_compound_spec_final.png')
 
         if generate_results:
@@ -186,3 +187,76 @@ def test_enzo_small_compound(self):
                                     'for enzo_small_compound answer test')
             old_spec.close()
             new_spec.close()
+
+    def test_gizmo_small_simple(self):
+        """
+        This is an answer test, which compares the results of this test
+        against answers generated from a previous version of the code.
+
+        This test generates a COS spectrum from a single Gizmo dataset 
+        using a simple ray and compare the ray and spectral output data
+        against a known answer.
+
+        """
+
+        # Set the dataset filename, load it into yt and define the trajectory
+        # of the LightRay to cross the box from one corner to the other.
+        ds = load(FIRE)
+        _, c = ds.find_max(('gas', 'density'))
+        ray_start = c
+        ray_end = ds.domain_right_edge
+        line_list = ['H', 'C', 'O']
+
+        # Make a LightRay object including all necessary fields so you can add
+        # all H, C, N, O, and Mg fields to the resulting spectrum from your dataset.
+        # Save LightRay to ray.h5 and use it locally as ray object.
+        # Note: We use PartType0, the gas particle field type as our ftype!
+        ray = make_simple_ray(ds, start_position=ray_start,
+                              end_position=ray_end, data_filename='ray.h5',
+                              lines=line_list, ftype='PartType0')
+
+        # Now use the ray object to actually generate an absorption spectrum
+        # Use the settings (spectral range, LSF, and spectral resolution) for COS
+        # And save it as an output text file and plot it to an image.
+        sg = SpectrumGenerator('COS')
+        sg.make_spectrum(ray, lines=line_list)
+        raw_file = 'gizmo_small_simple_spec_raw.h5'
+        raw_file_compare = os.path.join(test_results_dir, raw_file)
+        sg.save_spectrum(raw_file)
+        sg.plot_spectrum('gizmo_small_simple_spec_raw.png')
+
+        # "Final" spectrum with added quasar, MW background, applied line-spread
+        # function, and added gaussian noise (SNR=30)
+        sg.add_qso_spectrum()
+        sg.add_milky_way_foreground()
+        sg.apply_lsf()
+        sg.add_gaussian_noise(30, seed=1)
+        final_file = 'gizmo_small_simple_spec_final.h5'
+        sg.save_spectrum(final_file)
+        final_file_compare = os.path.join(test_results_dir, final_file)
+        sg.plot_spectrum('gizmo_small_simple_spec_final.png')
+
+        if generate_results:
+            os.rename(raw_file, raw_file_compare)
+            os.rename(final_file, final_file_compare)
+
+        else:
+            old_spec = h5py.File(raw_file_compare, 'r')
+            new_spec = h5py.File(raw_file, 'r')
+            for key in old_spec.keys():
+                assert_almost_equal(new_spec[key].value, old_spec[key].value, \
+                                    decimal=err_precision,
+                                    err_msg='Raw spectrum array does not match '+\
+                                    'for gizmo_small_simple answer test')
+            old_spec.close()
+            new_spec.close()
+
+            old_spec = h5py.File(final_file_compare, 'r')
+            new_spec = h5py.File(final_file, 'r')
+            for key in old_spec.keys():
+                assert_almost_equal(new_spec[key].value, old_spec[key].value, \
+                                    decimal=err_precision,
+                                    err_msg='Final spectrum array does not match '+\
+                                    'for gizmo_small_simple answer test')
+            old_spec.close()
+            new_spec.close()