Add new plotting tools and do some code cleanup

scimes/orion2scimes.py

@@ -4,6 +4,9 @@
 import numpy as np
 import math
 import aplpy
+import random
+
+from matplotlib import pyplot as plt
 
 from astrodendro import Dendrogram, ppv_catalog
 from astropy import units as u
@@ -64,13 +67,13 @@ def make_asgn(dendro, data_file, cores_idx = [], tag = '_', collapse = True):
 
     # Making the assignment cube
     if len(data.shape) == 3:
-        asgn = np.zeros(data.shape, dtype = int32)
+        asgn = np.zeros(data.shape, dtype = np.int32)
 
     if len(data.shape) == 4:
-        asgn = np.zeros((data.shape[1],data.shape[2],data.shape[3]), dtype = int32)    
+        asgn = np.zeros((data.shape[1],data.shape[2],data.shape[3]), dtype = np.int32)    
 
     for i in cores_idx:
-        asgn[where(d[i].get_mask(shape = asgn.shape))] = i
+        asgn[np.where(d[i].get_mask(shape = asgn.shape))] = i
 
 
     # Write the fits file
@@ -85,8 +88,8 @@ def make_asgn(dendro, data_file, cores_idx = [], tag = '_', collapse = True):
 
         asgn_map = np.amax(asgn.data, axis = 0) 
 
-        matshow(asgn_map, origin = "lower")
-        colorbar()
+        plt.matshow(asgn_map, origin = "lower")
+        plt.colorbar()
 
 
     return
@@ -95,13 +98,14 @@ def make_asgn(dendro, data_file, cores_idx = [], tag = '_', collapse = True):
 
 
 
-path = '/Volumes/Zeruel_data/ORION/'
+path = './'
+#path = '/Volumes/Zeruel_data/ORION/'
 #path = '/Users/Dario/Documents/dendrograms/'
-filename = path+'orion'
+filename = os.path.join(path, 'orion')
 
-do_make = False
-do_catalog = False
-do_load = False
+do_make = True
+do_catalog = True
+do_load = True
 
 
 if do_make:
@@ -115,7 +119,7 @@ def make_asgn(dendro, data_file, cores_idx = [], tag = '_', collapse = True):
     data = hdu.data
     hd = hdu.header
 
-    if size(shape(data))==4:
+    if data.ndim==4:
         data = data[0,:,:,:]
 
 

scimes/plotting.py

@@ -0,0 +1,82 @@
+import warnings
+import os.path
+
+import numpy as np
+import math
+import aplpy
+import random
+import itertools
+
+from matplotlib import pyplot as plt
+
+from astrodendro import Dendrogram, ppv_catalog
+from astropy import units as u
+from astropy.io import fits
+from astropy import wcs
+from astropy.table.table import Column
+from astropy.table import Table
+
+from sklearn.cluster import spectral_clustering
+from skimage.measure import regionprops
+
+from scimes import SpectralCloudstering
+
+from datetime import datetime
+
+from pdb import set_trace as stop
+
+
+
+def dendroplot_clusters(clusters,
+                        dend,
+                        cat,
+                        axname1='flux',
+                        axname2='v_rms',
+                        axscale1=1.,
+                        axscale2=1.,
+                        colors=itertools.cycle('rgbcmyk'),
+                        marker='s',
+                        marker2=None,
+                        linestyle='-', **kwargs):
+    """
+    Plot all of the clusters with (partially transparent) lines connecting
+    leaves to their parents all the way up to the trunk
+
+    Examples
+    --------
+    >>> SC = SpectralCloudstering(dend, cat)
+    >>> colors = showdendro(SC.dendrogram, SC.clusters)
+    >>> dendroplot_clusters(SC.clusters, SC.dendrogram, SC.catalog,
+    ...                     colors=colors, axname1='radius')
+    """
+
+    axis = plt.gca()
+
+    for cluster, color in zip(clusters, colors):
+        leaves = dend[cluster].sorted_leaves()
+        last_parent = cluster
+
+        for leaf in leaves:
+            xax,yax = ([cat[leaf.idx][axname1]*axscale1],
+                       [cat[leaf.idx][axname2]*axscale2])
+            #if axname1 in ('v_rms','reff'):
+            #    xax *= gcorfactor[leaf.idx]
+            #if axname2 in ('v_rms','reff'):
+            #    yax *= gcorfactor[leaf.idx]
+            axis.plot(xax, yax, marker, color=color, markeredgecolor='none', alpha=0.5)
+            obj = leaf.parent
+            while obj.parent:
+                xax.append(cat[obj.idx][axname1]*axscale1)
+                yax.append(cat[obj.idx][axname2]*axscale2)
+                if obj.idx == last_parent:
+                    break
+                obj = obj.parent
+            if np.any(np.isnan(yax)):
+                ok = ~np.isnan(yax)
+                axis.plot(np.array(xax)[ok], np.array(yax)[ok], alpha=0.5,
+                          label=leaf.idx, color='b', zorder=5,
+                          linestyle=linestyle, marker=marker2, **kwargs)
+            else:
+                axis.plot(xax, yax, alpha=0.1, label=leaf.idx, color=color,
+                          zorder=5, linestyle=linestyle, marker=marker2,
+                          **kwargs)

scimes/scimes.py

@@ -3,11 +3,20 @@
 import math
 
 import numpy as np
+import itertools
+from itertools import combinations
 
 from matplotlib import pyplot as plt
 
-from sklearn.cluster import spectral_clustering
-from itertools import combinations
+from astrodendro import Dendrogram, ppv_catalog
+from astropy import units as u
+from astropy.stats import median_absolute_deviation as mad
+from astropy.table import Column
+import aplpy
+
+from sklearn import metrics
+from spectral import spectral_clustering
+from skimage.measure import regionprops
 
 
 def mat_smooth(Mat, scalpar = 0, lscal = False):
@@ -450,14 +459,14 @@ def cloudstering(dendrogram, catalog, criteria, user_k, user_ams, user_scalpars,
 
         clusts = clusts + unclust_leaves
 
-        print "-- Unclustered leaves added. Final cluster number", len(clusts)        
+        print "-- Unclustered leaves added. Final cluster number", len(clusts)
 
     return clusts, AMs, escalpars, silhoutte 
 
 
 
 
-class SpectralCloudstering:
+class SpectralCloudstering(object):
     """
     Apply the spectral clustering to find the best 
     cloud segmentation out from a dendrogram.
@@ -466,11 +475,11 @@ class SpectralCloudstering:
     -----------
 
     dendrogram: 'astrodendro.dendrogram.Dendrogram' instance
-    The dendrogram to clusterize
+        The dendrogram to clusterize
 
     catalog: 'astropy.table.table.Table' instance
-    A catalog containing all properties of the dendrogram
-    structures. Generally generated with ppv_catalog module
+        A catalog containing all properties of the dendrogram
+        structures. Generally generated with ppv_catalog module
 
     cl_volume: bool
     Clusterize the dendrogram using the 'volume' criterium 
@@ -523,12 +532,19 @@ class SpectralCloudstering:
 
     """
 
-    def __init__(self, dendrogram, catalog, cl_volume = True, cl_luminosity=True, \
-                 user_k = None, user_ams = None, user_scalpars = None, \
+    def __init__(self, dendrogram, catalog, cl_volume = True, cl_luminosity=True,
+                 user_k = None, user_ams = None, user_scalpars = None,
                  savesingles = False, locscaling = False, blind = False):
 
         self.dendrogram = dendrogram
         self.catalog = catalog
+        if 'luminosity' not in catalog.colnames:
+            print("WARNING: adding luminosity = flux to the catalog.")
+            catalog.add_column(Column(catalog['flux'], 'luminosity'))
+        if 'volume' not in catalog.colnames:
+            print("WARNING: adding volume = pi * radius^2 * v_rms to the catalog.")
+            catalog.add_column(Column(catalog['radius']**2*np.pi *
+                                      catalog['v_rms'], 'volume'))
         self.cl_volume = cl_volume
         self.cl_luminosity = cl_luminosity
         self.user_k = user_k or 0
@@ -545,7 +561,95 @@ def __init__(self, dendrogram, catalog, cl_volume = True, cl_luminosity=True, \
         if self.cl_luminosity:
             self.criteria.append(1)
 
+        # default colors in case plot_connected_colors is called before showdendro
+        self.colors = itertools.cycle('rgbcmyk')
 
-        self.clusters, self.affmats, self.escalpars, self.silhouette = cloudstering(self.dendrogram, self.catalog, self.criteria, \
-                                                   self.user_k, self.user_ams, self.user_scalpars, \
+        self.clusters, self.affmats, self.escalpars, self.silhouette = cloudstering(self.dendrogram,
+                                                                                    self.catalog,
+                                                                                    self.criteria,
+                                                   self.user_k, self.user_ams, self.user_scalpars,
                                                    self.savesingles, self.locscaling, self.blind)
+
+    def showdendro(self):
+
+        dendro = self.dendro
+        cores_idx = self.clusters
+
+
+        # For the random colors
+        r = lambda: random.randint(0,255)
+
+        p = dendro.plotter()
+
+        fig = plt.figure(figsize=(20, 12))
+        ax = fig.add_subplot(111)
+
+        ax.set_yscale('log')
+
+        cols = []
+
+
+        # Plot the whole tree
+
+        p.plot_tree(ax, color='black')
+
+        for i in range(len(cores_idx)):
+
+            col = '#%02X%02X%02X' % (r(),r(),r())
+            cols.append(col)
+            p.plot_tree(ax, structure=dendro[cores_idx[i]], color=cols[i], lw=3)
+
+        ax.set_title("Final clustering configuration")
+
+        ax.set_xlabel("Structure")
+        ax.set_ylabel("Flux")
+
+        self.colors = cols
+
+
+    def plot_connected_clusters(self, **kwargs):
+        from plotting import dendroplot_clusters
+
+        return dendroplot_clusters(self.clusters, self.dendrogram, self.catalog,
+                                   colors=self.colors,
+                                   **kwargs)
+
+    def make_assignment_cube(self, outfileprefix, header, tag = '_',
+                             collapse = True):
+        """
+        Create a label cube with only the cluster (cloudster) IDs included, and
+        write to disk.
+
+        Parameters
+        ----------
+        outfileprefix : str
+            The prefix for the output filename.  The file has a format
+            `outfileprefix+'_asgn_'+tag+'.fits'`
+        header : `fits.Header`
+            The header of the output assignment cube.  Should be the same
+            header that the dendrogram was generated from
+        tag : str
+        collapse : bool
+        """
+
+        data = self.dendrogram.data.squeeze()
+
+        # Making the assignment cube
+        asgn = np.zeros(data.shape, dtype=np.int32)
+
+        for i in self.clusters:
+            asgn[np.where(dendro[i].get_mask(shape = asgn.shape))] = i+1
+
+
+        # Write the fits file
+        self.asgn = fits.PrimaryHDU(asgn.astype('short'), header)
+
+        self.asgn.writeto(outfileprefix+'_asgn_'+tag+'.fits', clobber=True)
+
+        # Collapsed version of the asgn cube
+        if collapse:
+
+            asgn_map = np.amax(asgn.data, axis = 0) 
+
+            plt.matshow(asgn_map, origin = "lower")
+            plt.colorbar()

scimes/spectral.py

@@ -9,13 +9,13 @@
 
 import numpy as np
 
-from ..base import BaseEstimator, ClusterMixin
-from ..utils import check_random_state, as_float_array, deprecated
-from ..utils.extmath import norm
-from ..metrics.pairwise import pairwise_kernels
-from ..neighbors import kneighbors_graph
-from ..manifold import spectral_embedding
-from .k_means_ import k_means
+from sklearn.base import BaseEstimator, ClusterMixin
+from sklearn.utils import check_random_state, as_float_array, deprecated
+from sklearn.utils.extmath import norm
+from sklearn.metrics.pairwise import pairwise_kernels
+from sklearn.neighbors import kneighbors_graph
+from sklearn.manifold import spectral_embedding
+from sklearn.cluster.k_means_ import k_means
 
 
 def discretize(vectors, copy=True, max_svd_restarts=30, n_iter_max=20,