source: trunk/src/ATrous/ReconSearch.cc @ 845

// -----------------------------------------------------------------------
// ReconSearch.cc: Searching a wavelet-reconstructed cube.
// -----------------------------------------------------------------------
// Copyright (C) 2006, Matthew Whiting, ATNF
//
// This program is free software; you can redistribute it and/or modify it
// under the terms of the GNU General Public License as published by the
// Free Software Foundation; either version 2 of the License, or (at your
// option) any later version.
//
// Duchamp is distributed in the hope that it will be useful, but WITHOUT
// ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or
// FITNESS FOR A PARTICULAR PURPOSE.  See the GNU General Public License
// for more details.
//
// You should have received a copy of the GNU General Public License
// along with Duchamp; if not, write to the Free Software Foundation,
// Inc., 59 Temple Place, Suite 330, Boston, MA 02111-1307, USA
//
// Correspondence concerning Duchamp may be directed to:
//    Internet email: Matthew.Whiting [at] atnf.csiro.au
//    Postal address: Dr. Matthew Whiting
//                    Australia Telescope National Facility, CSIRO
//                    PO Box 76
//                    Epping NSW 1710
//                    AUSTRALIA
// -----------------------------------------------------------------------
#include <fstream>
#include <iostream>
#include <iomanip>
#include <vector>
#include <duchamp/duchamp.hh>
#include <duchamp/PixelMap/Object3D.hh>
#include <duchamp/Cubes/cubes.hh>
#include <duchamp/Detection/detection.hh>
#include <duchamp/ATrous/atrous.hh>
#include <duchamp/Utils/utils.hh>
#include <duchamp/Utils/feedback.hh>
#include <duchamp/Utils/Statistics.hh>

using namespace PixelInfo;
using namespace Statistics;

namespace duchamp
{

  void Cube::ReconSearch()
  {
    /// The Cube is first reconstructed, using Cube::ReconCube().
    /// The statistics of the cube are calculated next.
    /// It is then searched, using searchReconArray.
    /// The resulting object list is stored in the Cube, and outputted
    ///  to the log file if the user so requests.

    if(!this->par.getFlagATrous()){
      duchampWarning("ReconSearch",
                     "You've requested a reconSearch, but not allocated space for the reconstructed array.\nDoing the basic CubicSearch.\n");
      this->CubicSearch();
    }
    else {
  
      this->ReconCube();

      if(this->par.isVerbose()) std::cout << "  ";

      this->setCubeStats();
    
      if(this->par.isVerbose()) std::cout << "  Searching... " << std::flush;
  
      *this->objectList = searchReconArray(this->axisDim,this->array,
                                           this->recon,this->par,this->Stats);

      if(this->par.isVerbose()) std::cout << "  Updating detection map... " 
                                          << std::flush;
      this->updateDetectMap();
      if(this->par.isVerbose()) std::cout << "Done.\n";

      if(this->par.getFlagLog()){
        if(this->par.isVerbose()) 
          std::cout << "  Logging intermediate detections... " << std::flush;
        this->logDetectionList();
        if(this->par.isVerbose()) std::cout << "Done.\n";
      }

    }
  }

  /////////////////////////////////////////////////////////////////////////////
  void Cube::ReconCube()
  {
    /// A front-end to the various reconstruction functions, the choice of 
    ///  which is determined by the use of the reconDim parameter.
    /// Differs from ReconSearch only in that no searching is done.

    int dimRecon = this->par.getReconDim();
    // Test whether we have eg. an image, but have requested a 3-d 
    //  reconstruction.
    // If dimension of data array is less than dimRecon, change dimRecon 
    //  to the dimension of the array.
    int numGoodDim = this->head.getNumAxes();
    if(numGoodDim<dimRecon){
      dimRecon = numGoodDim;
      this->par.setReconDim(dimRecon);
      std::stringstream errmsg;
      errmsg << "You requested a " << dimRecon << "-dimensional reconstruction,"
             << " but the FITS file is only " << numGoodDim << "-dimensional.\n"
             << "Changing reconDim to " << numGoodDim << ".\n";
      duchampWarning("Reconstruction",errmsg.str());
    }

    switch(dimRecon)
      {
      case 1: this->ReconCube1D(); break;
      case 2: this->ReconCube2D(); break;
      case 3: this->ReconCube3D(); break;
      default:
        if(dimRecon<=0){
          std::stringstream errmsg;
          errmsg << "reconDim (" << dimRecon
                 << ") is less than 1. Performing 1-D reconstruction.\n";
          duchampWarning("Reconstruction", errmsg.str());
          this->par.setReconDim(1);
          this->ReconCube1D();
        }
        else if(dimRecon>3){ 
          //this probably won't happen with new code above, but just in case...
          std::stringstream errmsg;
          errmsg << "reconDim (" << dimRecon
                 << ") is more than 3. Performing 3-D reconstruction.\n";
          duchampWarning("Reconstruction", errmsg.str());
          this->par.setReconDim(3);
          this->ReconCube3D();
        }
        break;
      }
  }

  /////////////////////////////////////////////////////////////////////////////
  void Cube::ReconCube1D()
  {
    /// This reconstructs a cube by performing a 1D a trous reconstruction
    ///  in the spectrum of each spatial pixel.

    long xySize = this->axisDim[0] * this->axisDim[1];

    long zdim = this->axisDim[2];

    ProgressBar bar;
    if(!this->reconExists){
      if(this->par.isVerbose()){
        std::cout<<"  Reconstructing... ";
        bar.init(xySize);
      }
      for(int npix=0; npix<xySize; npix++){

        if( this->par.isVerbose() ) bar.update(npix+1);

        float *spec = new float[zdim];
        float *newSpec = new float[zdim];
        for(int z=0;z<zdim;z++) spec[z] = this->array[z*xySize + npix];
        bool verboseFlag = this->par.isVerbose();
        this->par.setVerbosity(false);
        atrous1DReconstruct(this->axisDim[2],spec,newSpec,this->par);
        this->par.setVerbosity(verboseFlag);
        for(int z=0;z<zdim;z++) this->recon[z*xySize+npix] = newSpec[z];
        delete [] spec;
        delete [] newSpec;
      }
      this->reconExists = true;
      if(this->par.isVerbose()){
        bar.fillSpace(" All Done.");
        printSpace(22);
        std::cout << "\n";
      }
    }

  }

  /////////////////////////////////////////////////////////////////////////////
  void Cube::ReconCube2D()
  {
    /// This reconstructs a cube by performing a 2D a trous reconstruction
    ///  in each spatial image (ie. each channel map) of the cube.

    long xySize = this->axisDim[0] * this->axisDim[1];
    ProgressBar bar;
    bool useBar = (this->axisDim[2]>1);

    if(!this->reconExists){
      if(this->par.isVerbose()) std::cout<<"  Reconstructing... ";
      if(useBar&&this->par.isVerbose()) bar.init(this->axisDim[2]);
      for(int z=0;z<this->axisDim[2];z++){

        if( this->par.isVerbose() && useBar ) bar.update((z+1));

        if(!this->par.isInMW(z)){
          float *im = new float[xySize];
          float *newIm = new float[xySize];
          for(int npix=0; npix<xySize; npix++) 
            im[npix] = this->array[z*xySize+npix];
          bool verboseFlag = this->par.isVerbose();
          this->par.setVerbosity(false);
          atrous2DReconstruct(this->axisDim[0],this->axisDim[1],
                              im,newIm,this->par);
          this->par.setVerbosity(verboseFlag);
          for(int npix=0; npix<xySize; npix++) 
            this->recon[z*xySize+npix] = newIm[npix];
          delete [] im;
          delete [] newIm;
        }
        else {
          for(int i=z*xySize; i<(z+1)*xySize; i++) 
            this->recon[i] = this->array[i];
        }
      }
      this->reconExists = true;
      if(this->par.isVerbose()) {
        if(useBar) bar.fillSpace(" All Done.");
        printSpace(22);
        std::cout << "\n";
      }
    }
  }

  /////////////////////////////////////////////////////////////////////////////
  void Cube::ReconCube3D()
  {
    /// This performs a full 3D a trous reconstruction of the cube

    if(this->axisDim[2]==1) this->ReconCube2D();
    else {

      if(!this->reconExists){
        if(this->par.isVerbose()) std::cout<<"  Reconstructing... "<<std::flush;
        atrous3DReconstruct(this->axisDim[0],this->axisDim[1],this->axisDim[2],
                            this->array,this->recon,this->par);
        this->reconExists = true;
        if(this->par.isVerbose()) {
          std::cout << "  All Done.";
          printSpace(22);
          std::cout << "\n";
        }
      }

    }
  }

  /////////////////////////////////////////////////////////////////////////////
  std::vector <Detection> searchReconArray(long *dim, float *originalArray, 
                                           float *reconArray, Param &par,
                                           StatsContainer<float> &stats)
  {

  if(par.getSearchType()=="spectral")
    return searchReconArraySpectral(dim,originalArray,reconArray,par,stats);
  else if(par.getSearchType()=="spatial")
    return searchReconArraySpatial(dim,originalArray,reconArray,par,stats);
  else{
    std::stringstream ss;
    ss << "Unknown search type : " << par.getSearchType();
    duchampError("searchReconArray",ss.str());
    return std::vector<Detection>(0);
  }
  }
  /////////////////////////////////////////////////////////////////////////////
  std::vector <Detection> searchReconArraySpectral(long *dim, float *originalArray, 
                                                   float *reconArray, Param &par,
                                                   StatsContainer<float> &stats)
  {
    ///  This searches for objects in a cube that has been reconstructed.
    /// 
    ///  The search is conducted just in each spatial pixel (xdim*ydim 1D 
    ///   searches).
    ///  The searches are done on the reconstructed array, although the detected
    ///   objects have fluxes drawn from the corresponding pixels of the original
    ///   array.
    /// 
    ///  \param dim Array of dimension sizes
    ///  \param originalArray Original, un-reconstructed image array.
    ///  \param reconArray Reconstructed image array
    ///  \param par The Param set.
    ///  \param stats The StatsContainer that defines what a detection is.
    /// 
    ///  \return A vector of Detections resulting from the search.

    std::vector <Detection> outputList;
    long zdim = dim[2];
    long xySize = dim[0] * dim[1];
    int num=0;

    // First search --  in each spectrum.
    if(zdim > 1){

      ProgressBar bar;
      if(par.isVerbose()) bar.init(xySize);

      bool *doPixel = new bool[xySize];
      // doPixel is a bool array to say whether to look in a given spectrum
      for(int npix=0; npix<xySize; npix++){
        doPixel[npix] = false;
        for(int z=0;z<zdim;z++) {
          doPixel[npix] = doPixel[npix] || 
            (!par.isBlank(originalArray[npix]) && !par.isInMW(z));
        }
        // doPixel[i] is false only when there are no good pixels in spectrum
        //  of pixel #i.
      }

      long *specdim = new long[2];
      specdim[0] = zdim; specdim[1]=1;
      Image *spectrum = new Image(specdim);
      delete [] specdim;
      spectrum->saveParam(par);
      spectrum->saveStats(stats);
      spectrum->setMinSize(par.getMinChannels());
      // NB the beam is not used after this point
      // spectrum->pars().setBeamSize(2.); 
      // // beam size: for spectrum, only neighbouring channels correlated

      for(int y=0; y<dim[1]; y++){
        for(int x=0; x<dim[0]; x++){

          int npix = y*dim[0] + x;
          if( par.isVerbose() ) bar.update(npix+1);

          if(doPixel[npix]){

            spectrum->extractSpectrum(reconArray,dim,npix);
            spectrum->removeMW(); // only works if flagMW is true
            std::vector<Scan> objlist = spectrum->findSources1D();
            std::vector<Scan>::iterator obj;
            num += objlist.size();
            for(obj=objlist.begin();obj!=objlist.end();obj++){
              Detection newObject;
              // Fix up coordinates of each pixel to match original array
              for(int z=obj->getX();z<=obj->getXmax();z++) {
                newObject.addPixel(x,y,z);
              }
              newObject.setOffsets(par);
              if(par.getFlagTwoStageMerging()) mergeIntoList(newObject,outputList,par);
              else outputList.push_back(newObject);
            }
          }

        }
      }

      delete spectrum;
      delete [] doPixel;
      
      if(par.isVerbose()){
        bar.remove();
        std::cout << "Found " << num << ".\n";
      }

    }

    return outputList;
  }

  /////////////////////////////////////////////////////////////////////////////
  std::vector <Detection> searchReconArraySpatial(long *dim, float *originalArray,
                                                  float *reconArray, Param &par,
                                                  StatsContainer<float> &stats)
  {
    ///  This searches for objects in a cube that has been reconstructed.
    /// 
    ///  The search is conducted only in each channel image (zdim 2D
    ///  searches).
    /// 
    ///  The searches are done on the reconstructed array, although the detected
    ///   objects have fluxes drawn from the corresponding pixels of the original
    ///   array.
    /// 
    ///  \param dim Array of dimension sizes
    ///  \param originalArray Original, un-reconstructed image array.
    ///  \param reconArray Reconstructed image array
    ///  \param par The Param set.
    ///  \param stats The StatsContainer that defines what a detection is.
    /// 
    ///  \return A vector of Detections resulting from the search.

    std::vector <Detection> outputList;
    long zdim = dim[2];
    int num=0;
    ProgressBar bar;
    bool useBar = (zdim>1);
    if(useBar&&par.isVerbose()) bar.init(zdim);
  
    long *imdim = new long[2];
    imdim[0] = dim[0]; imdim[1] = dim[1];
    Image *channelImage = new Image(imdim);
    delete [] imdim;
    channelImage->saveParam(par);
    channelImage->saveStats(stats);
    channelImage->setMinSize(1);

    for(int z=0; z<zdim; z++){  // loop over all channels

      if( par.isVerbose() && useBar ) bar.update(z+1);

      if(!par.isInMW(z)){ 
        // purpose of this is to ignore the Milky Way channels 
        //  if we are flagging them

        channelImage->extractImage(reconArray,dim,z);
        std::vector<Object2D> objlist = channelImage->findSources2D();
        std::vector<Object2D>::iterator obj;
        num += objlist.size();
        for(obj=objlist.begin();obj!=objlist.end();obj++){
          Detection newObject;
          newObject.addChannel(z,*obj);
          newObject.setOffsets(par);
          if(par.getFlagTwoStageMerging()) mergeIntoList(newObject,outputList,par);
          else outputList.push_back(newObject);
        }
      }
    
    }

    delete channelImage;

    if(par.isVerbose()){
      if(useBar) bar.remove();
      std::cout << "Found " << num << ".\n";
    }


    return outputList;
  }

}