Fits export

astrodendro/analysis.py

@@ -307,7 +307,8 @@ class PPVStatistic(SpatialBase):
     """
 
     velocity_scale = MetaData('velocity_scale', 'Velocity channel width')
-    vaxis = MetaData('vaxis', 'Index of velocity axis (numpy convention)')
+    vaxis = MetaData('vaxis', 'Index of velocity axis (numpy convention)',
+                     default=0)
 
     def __init__(self, stat, metadata=None):
         if isinstance(stat, Structure):

astrodendro/dendrogram.py

@@ -304,14 +304,19 @@ def next_idx():
         s = tuple(slice(0, s, 1) for s in data.shape)
         self.index_map = self.index_map[s]
 
+        self._index()
+
+        # Return the newly-created dendrogram:
+        return self
+
+
+    def _index(self):
         # add dendrogram index
         ti = TreeIndex(self)
 
         for s in self.structures_dict.itervalues():
             s._tree_index = ti
 
-        # Return the newly-created dendrogram:
-        return self
 
     @property
     def trunk(self):

astrodendro/io/fits.py

@@ -21,18 +21,30 @@
 
 import numpy as np
 
+from .util import parse_dendrogram
 # Import and export
 
 
 def dendro_export_fits(d, filename):
     """Export the dendrogram 'd' to the FITS file 'filename'"""
-    import pyfits
-    raise NotImplementedError("FITS export has not yet been implemented.")
+    from astropy.io import fits
+
+    hdus = [fits.PrimaryHDU(),
+            fits.ImageHDU(d.data),
+            fits.ImageHDU(d.index_map),
+            fits.ImageHDU(np.array([ord(x) for x in  d.to_newick()]))]
+    hdulist = fits.HDUList(hdus)
+
+    hdulist.writeto(filename, clobber=True)
 
 
 def dendro_import_fits(filename):
     """Import 'filename' and construct a dendrogram from it"""
-    import pyfits
-    from ..dendrogram import Dendrogram
-    from ..structure import Structure
-    raise NotImplementedError("FITS import has not yet been implemented.")
+    from astropy.io import fits
+
+    with fits.open(filename) as hdus:
+        data = hdus[1].data
+        index_map = hdus[2].data
+        newick = ''.join(chr(x) for x in hdus[3].data.flat)
+
+    return parse_dendrogram(newick, data, index_map)

astrodendro/io/hdf5.py

@@ -21,72 +21,7 @@
 
 import numpy as np
 
-# Helper functions:
-
-
-def _parse_newick(string):
-
-    items = {}
-
-    # Find maximum level
-    current_level = 0
-    max_level = 0
-    for i, c in enumerate(string):
-        if c == '(':
-            current_level += 1
-        if c == ')':
-            current_level -= 1
-        max_level = max(max_level, current_level)
-
-    # Loop through levels and construct tree
-    for level in range(max_level, 0, -1):
-
-        pairs = []
-
-        current_level = 0
-        for i, c in enumerate(string):
-            if c == '(':
-                current_level += 1
-                if current_level == level:
-                    start = i
-            if c == ')':
-                if current_level == level:
-                    pairs.append((start, i))
-                current_level -= 1
-
-        for pair in pairs[::-1]:
-
-            # Extract start and end of branch definition
-            start, end = pair
-
-            # Find the ID of the branch
-            colon = string.find(":", end)
-            branch_id = string[end + 1:colon]
-            if branch_id == '':
-                branch_id = 'trunk'
-            else:
-                branch_id = int(branch_id)
-
-            # Add branch definition to overall definition
-            items[branch_id] = eval("{%s}" % string[start + 1:end])
-
-            # Remove branch definition from string
-            string = string[:start] + string[end + 1:]
-
-    new_items = {}
-
-    def collect(d):
-        for item in d:
-            if item in items:
-                collect(items[item])
-                d[item] = (items[item], d[item])
-        return
-
-    collect(items['trunk'])
-
-    return items['trunk']
-
-# Import and export
+from .util import parse_dendrogram
 
 
 def dendro_export_hdf5(d, filename):
@@ -114,66 +49,10 @@ def dendro_export_hdf5(d, filename):
 def dendro_import_hdf5(filename):
     """Import 'filename' and construct a dendrogram from it"""
     import h5py
-    from ..dendrogram import Dendrogram
-    from ..structure import Structure
-    h5f = h5py.File(filename, 'r')
-    d = Dendrogram()
-    d.n_dim = h5f.attrs['n_dim']
-    d.data = h5f['data'].value
-    d.index_map = h5f['index_map'].value
-    d.structures_dict = {}
-
-    flux_by_structure = {}
-    indices_by_structure = {}
-
-    def _construct_tree(repr):
-        structures = []
-        for idx in repr:
-            structure_indices = indices_by_structure[idx]
-            f = flux_by_structure[idx]
-            if type(repr[idx]) == tuple:
-                sub_structures_repr = repr[idx][0]  # Parsed representation of sub structures
-                sub_structures = _construct_tree(sub_structures_repr)
-                for i in sub_structures:
-                    d.structures_dict[i.idx] = i
-                b = Structure(structure_indices, f, children=sub_structures, idx=idx)
-                # Correct merge levels - complicated because of the
-                # order in which we are building the tree.
-                # What we do is look at the heights of this branch's
-                # 1st child as stored in the newick representation, and then
-                # work backwards to compute the merge level of this branch
-                first_child_repr = sub_structures_repr.itervalues().next()
-                if type(first_child_repr) == tuple:
-                    height = first_child_repr[1]
-                else:
-                    height = first_child_repr
-                d.structures_dict[idx] = b
-                structures.append(b)
-            else:
-                l = Structure(structure_indices, f, idx=idx)
-                structures.append(l)
-                d.structures_dict[idx] = l
-        return structures
-
-    # Do a fast iteration through d.data, adding the indices and data values
-    # to the two dictionaries declared above:
-    indices = np.indices(d.data.shape).reshape(d.data.ndim, np.prod(d.data.shape)).transpose()
-
-    for coord in indices:
-        coord = tuple(coord)
-        idx = d.index_map[coord]
-        if idx:
-            try:
-                flux_by_structure[idx].append(d.data[coord])
-                indices_by_structure[idx].append(coord)
-            except KeyError:
-                flux_by_structure[idx] = [d.data[coord]]
-                indices_by_structure[idx] = [coord]
 
-    d.trunk = _construct_tree(_parse_newick(h5f['newick'].value))
-    # To make the structure.level property fast, we ensure all the items in the
-    # trunk have their level cached as "0"
-    for structure in d.trunk:
-        structure._level = 0  # See the @property level() definition in structure.py
+    with h5py.File(filename, 'r') as h5f:
+        newick = h5f['newick'].value
+        data = h5f['data'].value
+        index_map = h5f['index_map'].value
 
-    return d
+    return parse_dendrogram(newick, data, index_map)

astrodendro/io/util.py

@@ -0,0 +1,133 @@
+import numpy as np
+
+
+def parse_newick(string):
+
+    items = {}
+
+    # Find maximum level
+    current_level = 0
+    max_level = 0
+    for i, c in enumerate(string):
+        if c == '(':
+            current_level += 1
+        if c == ')':
+            current_level -= 1
+        max_level = max(max_level, current_level)
+
+    # Loop through levels and construct tree
+    for level in range(max_level, 0, -1):
+
+        pairs = []
+
+        current_level = 0
+        for i, c in enumerate(string):
+            if c == '(':
+                current_level += 1
+                if current_level == level:
+                    start = i
+            if c == ')':
+                if current_level == level:
+                    pairs.append((start, i))
+                current_level -= 1
+
+        for pair in pairs[::-1]:
+
+            # Extract start and end of branch definition
+            start, end = pair
+
+            # Find the ID of the branch
+            colon = string.find(":", end)
+            branch_id = string[end + 1:colon]
+            if branch_id == '':
+                branch_id = 'trunk'
+            else:
+                branch_id = int(branch_id)
+
+            # Add branch definition to overall definition
+            items[branch_id] = eval("{%s}" % string[start + 1:end])
+
+            # Remove branch definition from string
+            string = string[:start] + string[end + 1:]
+
+    new_items = {}
+
+    def collect(d):
+        for item in d:
+            if item in items:
+                collect(items[item])
+                d[item] = (items[item], d[item])
+        return
+
+    collect(items['trunk'])
+
+    return items['trunk']
+
+
+def parse_dendrogram(newick, data, index_map):
+    from ..dendrogram import Dendrogram
+    from ..structure import Structure
+
+    d = Dendrogram()
+    d.ndim = len(data.shape)
+
+    d.structures_dict = {}
+    d.data = data
+    d.index_map = index_map
+
+    flux_by_structure = {}
+    indices_by_structure = {}
+
+    def _construct_tree(repr):
+        structures = []
+        for idx in repr:
+            idx = int(idx)
+            structure_indices = indices_by_structure[idx]
+            f = flux_by_structure[idx]
+            if type(repr[idx]) == tuple:
+                sub_structures_repr = repr[idx][0]  # Parsed representation of sub structures
+                sub_structures = _construct_tree(sub_structures_repr)
+                for i in sub_structures:
+                    d.structures_dict[i.idx] = i
+                b = Structure(structure_indices, f, children=sub_structures, idx=idx)
+                # Correct merge levels - complicated because of the
+                # order in which we are building the tree.
+                # What we do is look at the heights of this branch's
+                # 1st child as stored in the newick representation, and then
+                # work backwards to compute the merge level of this branch
+                first_child_repr = sub_structures_repr.itervalues().next()
+                if type(first_child_repr) == tuple:
+                    height = first_child_repr[1]
+                else:
+                    height = first_child_repr
+                d.structures_dict[idx] = b
+                structures.append(b)
+            else:
+                l = Structure(structure_indices, f, idx=idx)
+                structures.append(l)
+                d.structures_dict[idx] = l
+        return structures
+
+    # Do a fast iteration through d.data, adding the indices and data values
+    # to the two dictionaries declared above:
+    indices = np.indices(d.data.shape).reshape(d.data.ndim, np.prod(d.data.shape)).transpose()
+
+    for coord in indices:
+        coord = tuple(coord)
+        idx = d.index_map[coord]
+        if idx:
+            try:
+                flux_by_structure[idx].append(d.data[coord])
+                indices_by_structure[idx].append(coord)
+            except KeyError:
+                flux_by_structure[idx] = [d.data[coord]]
+                indices_by_structure[idx] = [coord]
+
+    d.trunk = _construct_tree(parse_newick(newick))
+    # To make the structure.level property fast, we ensure all the items in the
+    # trunk have their level cached as "0"
+    for structure in d.trunk:
+        structure._level = 0  # See the @property level() definition in structure.py
+
+    d._index()
+    return d

astrodendro/test/test_io.py

@@ -59,12 +59,17 @@ def compare_dendrograms(self, d1, d2):
         np.testing.assert_array_equal(d1.data, d2.data)
         # Now check that the structures are the same:
         for idx in d2.structures_dict:
-            structure1, structure2 = d1.structures_dict[idx], d2.structures_dict[idx]
-            assert_permuted_fancyindex(structure1.indices(subtree=False), structure2.indices(subtree=False))
-            assert np.all(np.sort(structure1.values(subtree=False)) == np.sort(structure2.values(subtree=False)))
+            structure1, structure2 = d1[idx], d2[idx]
+            assert_permuted_fancyindex(structure1.indices(subtree=False),
+                                       structure2.indices(subtree=False))
+            assert np.all(np.sort(structure1.values(subtree=False)) ==
+                          np.sort(structure2.values(subtree=False)))
             assert type(structure1) == type(structure2)
             # Compare the coordinates and data values of all peak pixels:
-            assert structure1.get_peak(subtree=True) == structure2.get_peak(subtree=True)
+            assert structure1.get_peak(subtree=True) == \
+              structure2.get_peak(subtree=True)
+
+            assert structure2._tree_index is not None
 
     # Below are the actual tests for each import/export format:
 
@@ -74,3 +79,10 @@ def test_hdf5(self):
         d1.save_to(self.test_filename)
         d2 = Dendrogram.load_from(self.test_filename)
         self.compare_dendrograms(d1, d2)
+
+    def test_fits(self):
+        self.test_filename = 'astrodendro-test.fits'
+        d1 = Dendrogram.compute(self.data, verbose=False)
+        d1.save_to(self.test_filename)
+        d2 = Dendrogram.load_from(self.test_filename)
+        self.compare_dendrograms(d1, d2)