<commit_before>/*=========================================================================
 *
 *  Copyright Insight Software Consortium
 *
 *  Licensed under the Apache License, Version 2.0 (the "License");
 *  you may not use this file except in compliance with the License.
 *  You may obtain a copy of the License at
 *
 *         http://www.apache.org/licenses/LICENSE-2.0.txt
 *
 *  Unless required by applicable law or agreed to in writing, software
 *  distributed under the License is distributed on an "AS IS" BASIS,
 *  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 *  See the License for the specific language governing permissions and
 *  limitations under the License.
 *
 *=========================================================================*/
#ifndef __itkTransformFileWriter_hxx
#define __itkTransformFileWriter_hxx

#include "itkTransformFileWriter.h"
#include "itkTransformFactoryBase.h"
#include "itkTransformIOFactory.h"
#include "itkCompositeTransformIOHelper.h"

namespace itk
{

namespace
{

/*
 This helper is used to:
 Create and set a new type transform that have requested output precision type.
 */
template<typename TOutputScalar, typename TInputScalar>
struct TransformIOHelper
{
  typedef TransformBaseTemplate<TInputScalar>     InputTransformType;
  typedef TransformBaseTemplate<TOutputScalar>    OutputTransformType;
  typedef typename OutputTransformType::Pointer   OutputTransformPointer;

  /*
   This function gets the type name of the input transform and creates
   a new transform object based on the requested precision type.
   */
  static OutputTransformPointer
  CreateNewTypeTransform(std::string transformName)
  {
    // Transform name is modified to have the output precision type.
    TransformIOBaseTemplate<TOutputScalar>::CorrectTransformPrecisionType( transformName );

    OutputTransformPointer convertedTransform;
    // Instantiate the transform
    LightObject::Pointer i = ObjectFactoryBase::CreateInstance ( transformName.c_str() );
    convertedTransform = dynamic_cast< OutputTransformType * >( i.GetPointer() );
    if( convertedTransform.IsNull() )
      {
      itkGenericExceptionMacro ( << "Could not create an instance of " << transformName );
      }
    convertedTransform->UnRegister();
    return convertedTransform;
  }

  /* Converts the value type of transform parameters to the output precision type */
  static OptimizerParameters< TOutputScalar >
  ConvertParametersType(const OptimizerParameters< TInputScalar >  &sourceParams)
  {
    OptimizerParameters< TOutputScalar > outputParams;
    outputParams.SetSize( sourceParams.GetSize() );
    for( SizeValueType i = 0; i < sourceParams.GetSize(); ++i )
      {
      outputParams[i] = static_cast<TOutputScalar>( sourceParams[i] );
      }
    return outputParams;
  }

  /* Set fixed parameters and parameters of the new type created transform */
  static void SetAllParameters(const InputTransformType *transform, OutputTransformPointer& convertedTransform)
  {
    // The precision type of the input transform parameters should be converted to the requested output precision
    convertedTransform->SetFixedParameters( ConvertParametersType( transform->GetFixedParameters() ) );
    convertedTransform->SetParameters( ConvertParametersType( transform->GetParameters() ) );
  }
};

/* Changes the precision type of input transform to the requested precision type */
template<typename TOutputScalar, typename TInputScalar>
inline void AddToTransformList(typename TransformBaseTemplate<TInputScalar>::ConstPointer &transform, std::list< typename TransformBaseTemplate<TOutputScalar>::ConstPointer > & transformList)
  {
  /* Pushes the converted transform to the input transform list */
  typedef TransformBaseTemplate<TInputScalar>        InputTransformType;
  typedef typename InputTransformType::ConstPointer  InputTransformConstPointer;
  typedef std::list< InputTransformConstPointer >    InputConstTransformListType;
  typedef TransformBaseTemplate<TOutputScalar>       OutputTransformType;
  typedef typename OutputTransformType::Pointer      OutputTransformPointer;
  typedef typename OutputTransformType::ConstPointer OutputTransformConstPointer;
  typedef std::list< OutputTransformPointer >        OutputTransformListType;
  typedef std::list< OutputTransformConstPointer >   OutputConstTransformListType;

  const std::string transformName = transform->GetTransformTypeAsString();
  OutputTransformPointer convertedTransform;

  typedef TransformIOHelper<TOutputScalar, TInputScalar> IOhelper;

  // Composite and DisplacementFieldTransform transforms should be treated differently.
  if( transformName.find("CompositeTransform") == std::string::npos )
    {
    convertedTransform = IOhelper::CreateNewTypeTransform( transformName );
    IOhelper::SetAllParameters(transform, convertedTransform);
    }
  else
    {
    /*
     Following steps are needed to process a composite transform:
     1) Use the compositeTransformIOHelper to get the input transforms list.
     2) Iterate through the input transform list, convert each sub transform and put them in the output transform list.
     3) Use a composite IO Helper agian to set the output transform list into the converted composite transform.
    */
    CompositeTransformIOHelperTemplate<TInputScalar> inputHelper;
    InputConstTransformListType inputTransformList = inputHelper.GetTransformList( transform );

    // create output transform list
    OutputTransformListType compositeTransformList;
    /*
     The first transform of the output transform list should be a composite transform
     we push back just an empty composite transform as it will be skipped by the outputHelper.
    */
    OutputTransformPointer outputComposite = IOhelper::CreateNewTypeTransform( transformName );
    compositeTransformList.push_back( outputComposite.GetPointer() );

    // Now we iterate through input list and convert each sub transform to a new transform with requested precision type.
    typename InputConstTransformListType::iterator it = inputTransformList.begin();
    // composite transform is the first transform of the input transform list
    ++it; // skip the composite transform
    for(; it != inputTransformList.end(); ++it)
       {
       // get the input sub transform
       const InputTransformType *inSub = dynamic_cast< const InputTransformType *>( (*it).GetPointer() );
       // convert each sub transform and push them to the output transform list
       std::string inSubName = inSub->GetTransformTypeAsString();
       OutputTransformPointer convertedSub = IOhelper::CreateNewTypeTransform( inSubName );
       IOhelper::SetAllParameters( inSub, convertedSub );
       // push back the converted sub transform to the composite transform list
       compositeTransformList.push_back( convertedSub.GetPointer() );
       }

    convertedTransform = IOhelper::CreateNewTypeTransform( transformName ); // new composite transform
    CompositeTransformIOHelperTemplate<TOutputScalar> outputHelper;
    // set the output transform list into the new composite transform
    outputHelper.SetTransformList(convertedTransform.GetPointer(), compositeTransformList);
    }

  transformList.push_back( OutputTransformConstPointer(convertedTransform.GetPointer()) );
}

/* Precision type conversion is not needed when the input transform already has the requested precision type */
template<>
inline void AddToTransformList<double,double>(TransformBaseTemplate<double>::ConstPointer &transform, std::list< TransformBaseTemplate<double>::ConstPointer > & transformList)
{
  transformList.push_back( TransformBaseTemplate<double>::ConstPointer(transform) );
}

template<>
inline void AddToTransformList<float,float>(TransformBaseTemplate<float>::ConstPointer &transform, std::list< TransformBaseTemplate<float>::ConstPointer > & transformList)
{
  transformList.push_back( TransformBaseTemplate<float>::ConstPointer(transform) );
}

} // end anonymous namespace

template<typename ScalarType>
TransformFileWriterTemplate<ScalarType>
::TransformFileWriterTemplate() :
  m_FileName(""),
  m_AppendMode(false)
{
  TransformFactoryBase::RegisterDefaultTransforms();
}

template<typename ScalarType>
TransformFileWriterTemplate<ScalarType>
::~TransformFileWriterTemplate()
{
}

/** Set the writer to append to the specified file */
template<typename ScalarType>
void TransformFileWriterTemplate<ScalarType>
::SetAppendOn()
{
  this->SetAppendMode(true);
}

/** Set the writer to overwrite the specified file - This is the
* default mode. */
template<typename ScalarType>
void TransformFileWriterTemplate<ScalarType>
::SetAppendOff()
{
  this->SetAppendMode(false);
}

/** Set the writer mode (append/overwrite). */
template<typename ScalarType>
void TransformFileWriterTemplate<ScalarType>
::SetAppendMode(bool mode)
{
  this->m_AppendMode = mode;
}

/** Get the writer mode. */
template<typename ScalarType>
bool TransformFileWriterTemplate<ScalarType>
::GetAppendMode()
{
  return ( this->m_AppendMode );
}

template<typename ScalarType>
void TransformFileWriterTemplate<ScalarType>
::PushBackTransformList(const Object *transObj)
{
  TransformBaseTemplate<double>::ConstPointer dblptr = dynamic_cast<const TransformBaseTemplate<double> *>( transObj );
  if( dblptr.IsNotNull() )
    {
    AddToTransformList<ScalarType,double>( dblptr, m_TransformList );
    }
  else
    {
    TransformBaseTemplate<float>::ConstPointer fltptr = dynamic_cast<const TransformBaseTemplate<float> *>( transObj );
    if( fltptr.IsNotNull() )
      {
      AddToTransformList<ScalarType,float>( fltptr, m_TransformList );
      }
    else
      {
      itkExceptionMacro("The input of writer should be whether a double precision"
                        "or a single precision transform type.");
      }
    }
}

template<>
void TransformFileWriterTemplate<double>
::PushBackTransformList(const Object *transObj)
{
  TransformBaseTemplate<double>::ConstPointer dblptr = dynamic_cast<const TransformBaseTemplate<double> *>( transObj );
  if( dblptr.IsNotNull() )
    {
    AddToTransformList<double, double>( dblptr, m_TransformList );
    }
  else
    {
    TransformBaseTemplate<float>::ConstPointer fltptr = dynamic_cast<const TransformBaseTemplate<float> *>( transObj );
    if( fltptr.IsNotNull() )
      {
      AddToTransformList<double, float>( fltptr, m_TransformList );
      }
    else
      {
      itkExceptionMacro("The input of writer should be whether a double precision"
                        "or a single precision transform type.");
      }
    }
}

template<>
void TransformFileWriterTemplate<float>
::PushBackTransformList(const Object *transObj)
{
  TransformBaseTemplate<double>::ConstPointer dblptr = dynamic_cast<const TransformBaseTemplate<double> *>( transObj );
  if( dblptr.IsNotNull() )
    {
    AddToTransformList<float, double>( dblptr, m_TransformList );
    }
  else
    {
    TransformBaseTemplate<float>::ConstPointer fltptr = dynamic_cast<const TransformBaseTemplate<float> *>( transObj );
    if( fltptr.IsNotNull() )
      {
      AddToTransformList<float, float>( fltptr, m_TransformList );
      }
    else
      {
      itkExceptionMacro("The input of writer should be whether a double precision"
                        "or a single precision transform type.");
      }
    }
}

/** Set the input transform and reinitialize the list of transforms */
template<typename ScalarType>
void TransformFileWriterTemplate<ScalarType>
::SetInput(const Object *transform)
{
  m_TransformList.clear();
  this->PushBackTransformList(transform);
}

template<typename ScalarType>
const typename TransformFileWriterTemplate<ScalarType>::TransformType *
TransformFileWriterTemplate<ScalarType>
::GetInput()
{
  ConstTransformPointer res = *(m_TransformList.begin());
  return res.GetPointer();
}

/** Add a transform to be written */
template<typename ScalarType>
void TransformFileWriterTemplate<ScalarType>
::AddTransform(const Object *transform)
{
  /* Check for a CompositeTransform.
   * The convention is that there should be one, and it should
   * be the first transform in the file
   */
  const std::string transformName = transform->GetNameOfClass();
  if( transformName.find("CompositeTransform") != std::string::npos )
    {
    if(this->m_TransformList.size() > 0)
      {
      itkExceptionMacro("Can only write a transform of type CompositeTransform "
                        "as the first transform in the file.");
      }
    }

  this->PushBackTransformList(transform);
}

template<typename ScalarType>
void TransformFileWriterTemplate<ScalarType>
::Update()
{
  if ( m_FileName == "" )
    {
    itkExceptionMacro ("No file name given");
    }
  typename TransformIOBaseTemplate<ScalarType>::Pointer transformIO =
    TransformIOFactoryTemplate<ScalarType>::CreateTransformIO( m_FileName.c_str(), WriteMode );
  if ( transformIO.IsNull() )
    {
    itkExceptionMacro("Can't Create IO object for file "
                      << m_FileName);
    }
  transformIO->SetAppendMode(this->m_AppendMode);
  transformIO->SetFileName(m_FileName);
  transformIO->SetTransformList(this->m_TransformList);
  transformIO->Write();
}

template<typename ScalarType>
void TransformFileWriterTemplate<ScalarType>
::PrintSelf(std::ostream & os, Indent indent) const
{
  Superclass::PrintSelf(os, indent);

  os << indent << "FileName: " << m_FileName << std::endl;
}

} // namespace itk

#endif
<commit_msg>COMP: Remove multiple definition of PushBackTransformList.<commit_after>/*=========================================================================
 *
 *  Copyright Insight Software Consortium
 *
 *  Licensed under the Apache License, Version 2.0 (the "License");
 *  you may not use this file except in compliance with the License.
 *  You may obtain a copy of the License at
 *
 *         http://www.apache.org/licenses/LICENSE-2.0.txt
 *
 *  Unless required by applicable law or agreed to in writing, software
 *  distributed under the License is distributed on an "AS IS" BASIS,
 *  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 *  See the License for the specific language governing permissions and
 *  limitations under the License.
 *
 *=========================================================================*/
#ifndef __itkTransformFileWriter_hxx
#define __itkTransformFileWriter_hxx

#include "itkTransformFileWriter.h"
#include "itkTransformFactoryBase.h"
#include "itkTransformIOFactory.h"
#include "itkCompositeTransformIOHelper.h"

namespace itk
{

namespace
{

/*
 This helper is used to:
 Create and set a new type transform that have requested output precision type.
 */
template<typename TOutputScalar, typename TInputScalar>
struct TransformIOHelper
{
  typedef TransformBaseTemplate<TInputScalar>     InputTransformType;
  typedef TransformBaseTemplate<TOutputScalar>    OutputTransformType;
  typedef typename OutputTransformType::Pointer   OutputTransformPointer;

  /*
   This function gets the type name of the input transform and creates
   a new transform object based on the requested precision type.
   */
  static OutputTransformPointer
  CreateNewTypeTransform(std::string transformName)
  {
    // Transform name is modified to have the output precision type.
    TransformIOBaseTemplate<TOutputScalar>::CorrectTransformPrecisionType( transformName );

    OutputTransformPointer convertedTransform;
    // Instantiate the transform
    LightObject::Pointer i = ObjectFactoryBase::CreateInstance ( transformName.c_str() );
    convertedTransform = dynamic_cast< OutputTransformType * >( i.GetPointer() );
    if( convertedTransform.IsNull() )
      {
      itkGenericExceptionMacro ( << "Could not create an instance of " << transformName );
      }
    convertedTransform->UnRegister();
    return convertedTransform;
  }

  /* Converts the value type of transform parameters to the output precision type */
  static OptimizerParameters< TOutputScalar >
  ConvertParametersType(const OptimizerParameters< TInputScalar >  &sourceParams)
  {
    OptimizerParameters< TOutputScalar > outputParams;
    outputParams.SetSize( sourceParams.GetSize() );
    for( SizeValueType i = 0; i < sourceParams.GetSize(); ++i )
      {
      outputParams[i] = static_cast<TOutputScalar>( sourceParams[i] );
      }
    return outputParams;
  }

  /* Set fixed parameters and parameters of the new type created transform */
  static void SetAllParameters(const InputTransformType *transform, OutputTransformPointer& convertedTransform)
  {
    // The precision type of the input transform parameters should be converted to the requested output precision
    convertedTransform->SetFixedParameters( ConvertParametersType( transform->GetFixedParameters() ) );
    convertedTransform->SetParameters( ConvertParametersType( transform->GetParameters() ) );
  }
};

/* Changes the precision type of input transform to the requested precision type */
template<typename TOutputScalar, typename TInputScalar>
inline void AddToTransformList(typename TransformBaseTemplate<TInputScalar>::ConstPointer &transform, std::list< typename TransformBaseTemplate<TOutputScalar>::ConstPointer > & transformList)
  {
  /* Pushes the converted transform to the input transform list */
  typedef TransformBaseTemplate<TInputScalar>        InputTransformType;
  typedef typename InputTransformType::ConstPointer  InputTransformConstPointer;
  typedef std::list< InputTransformConstPointer >    InputConstTransformListType;
  typedef TransformBaseTemplate<TOutputScalar>       OutputTransformType;
  typedef typename OutputTransformType::Pointer      OutputTransformPointer;
  typedef typename OutputTransformType::ConstPointer OutputTransformConstPointer;
  typedef std::list< OutputTransformPointer >        OutputTransformListType;
  typedef std::list< OutputTransformConstPointer >   OutputConstTransformListType;

  const std::string transformName = transform->GetTransformTypeAsString();
  OutputTransformPointer convertedTransform;

  typedef TransformIOHelper<TOutputScalar, TInputScalar> IOhelper;

  // Composite and DisplacementFieldTransform transforms should be treated differently.
  if( transformName.find("CompositeTransform") == std::string::npos )
    {
    convertedTransform = IOhelper::CreateNewTypeTransform( transformName );
    IOhelper::SetAllParameters(transform, convertedTransform);
    }
  else
    {
    /*
     Following steps are needed to process a composite transform:
     1) Use the compositeTransformIOHelper to get the input transforms list.
     2) Iterate through the input transform list, convert each sub transform and put them in the output transform list.
     3) Use a composite IO Helper agian to set the output transform list into the converted composite transform.
    */
    CompositeTransformIOHelperTemplate<TInputScalar> inputHelper;
    InputConstTransformListType inputTransformList = inputHelper.GetTransformList( transform );

    // create output transform list
    OutputTransformListType compositeTransformList;
    /*
     The first transform of the output transform list should be a composite transform
     we push back just an empty composite transform as it will be skipped by the outputHelper.
    */
    OutputTransformPointer outputComposite = IOhelper::CreateNewTypeTransform( transformName );
    compositeTransformList.push_back( outputComposite.GetPointer() );

    // Now we iterate through input list and convert each sub transform to a new transform with requested precision type.
    typename InputConstTransformListType::iterator it = inputTransformList.begin();
    // composite transform is the first transform of the input transform list
    ++it; // skip the composite transform
    for(; it != inputTransformList.end(); ++it)
       {
       // get the input sub transform
       const InputTransformType *inSub = dynamic_cast< const InputTransformType *>( (*it).GetPointer() );
       // convert each sub transform and push them to the output transform list
       std::string inSubName = inSub->GetTransformTypeAsString();
       OutputTransformPointer convertedSub = IOhelper::CreateNewTypeTransform( inSubName );
       IOhelper::SetAllParameters( inSub, convertedSub );
       // push back the converted sub transform to the composite transform list
       compositeTransformList.push_back( convertedSub.GetPointer() );
       }

    convertedTransform = IOhelper::CreateNewTypeTransform( transformName ); // new composite transform
    CompositeTransformIOHelperTemplate<TOutputScalar> outputHelper;
    // set the output transform list into the new composite transform
    outputHelper.SetTransformList(convertedTransform.GetPointer(), compositeTransformList);
    }

  transformList.push_back( OutputTransformConstPointer(convertedTransform.GetPointer()) );
}

/* Precision type conversion is not needed when the input transform already has the requested precision type */
template<>
inline void AddToTransformList<double,double>(TransformBaseTemplate<double>::ConstPointer &transform, std::list< TransformBaseTemplate<double>::ConstPointer > & transformList)
{
  transformList.push_back( TransformBaseTemplate<double>::ConstPointer(transform) );
}

template<>
inline void AddToTransformList<float,float>(TransformBaseTemplate<float>::ConstPointer &transform, std::list< TransformBaseTemplate<float>::ConstPointer > & transformList)
{
  transformList.push_back( TransformBaseTemplate<float>::ConstPointer(transform) );
}

} // end anonymous namespace

template<typename ScalarType>
TransformFileWriterTemplate<ScalarType>
::TransformFileWriterTemplate() :
  m_FileName(""),
  m_AppendMode(false)
{
  TransformFactoryBase::RegisterDefaultTransforms();
}

template<typename ScalarType>
TransformFileWriterTemplate<ScalarType>
::~TransformFileWriterTemplate()
{
}

/** Set the writer to append to the specified file */
template<typename ScalarType>
void TransformFileWriterTemplate<ScalarType>
::SetAppendOn()
{
  this->SetAppendMode(true);
}

/** Set the writer to overwrite the specified file - This is the
* default mode. */
template<typename ScalarType>
void TransformFileWriterTemplate<ScalarType>
::SetAppendOff()
{
  this->SetAppendMode(false);
}

/** Set the writer mode (append/overwrite). */
template<typename ScalarType>
void TransformFileWriterTemplate<ScalarType>
::SetAppendMode(bool mode)
{
  this->m_AppendMode = mode;
}

/** Get the writer mode. */
template<typename ScalarType>
bool TransformFileWriterTemplate<ScalarType>
::GetAppendMode()
{
  return ( this->m_AppendMode );
}

template<>
void TransformFileWriterTemplate<double>
::PushBackTransformList(const Object *transObj)
{
  TransformBaseTemplate<double>::ConstPointer dblptr = dynamic_cast<const TransformBaseTemplate<double> *>( transObj );
  if( dblptr.IsNotNull() )
    {
    AddToTransformList<double, double>( dblptr, m_TransformList );
    }
  else
    {
    TransformBaseTemplate<float>::ConstPointer fltptr = dynamic_cast<const TransformBaseTemplate<float> *>( transObj );
    if( fltptr.IsNotNull() )
      {
      AddToTransformList<double, float>( fltptr, m_TransformList );
      }
    else
      {
      itkExceptionMacro("The input of writer should be whether a double precision"
                        "or a single precision transform type.");
      }
    }
}

template<>
void TransformFileWriterTemplate<float>
::PushBackTransformList(const Object *transObj)
{
  TransformBaseTemplate<double>::ConstPointer dblptr = dynamic_cast<const TransformBaseTemplate<double> *>( transObj );
  if( dblptr.IsNotNull() )
    {
    AddToTransformList<float, double>( dblptr, m_TransformList );
    }
  else
    {
    TransformBaseTemplate<float>::ConstPointer fltptr = dynamic_cast<const TransformBaseTemplate<float> *>( transObj );
    if( fltptr.IsNotNull() )
      {
      AddToTransformList<float, float>( fltptr, m_TransformList );
      }
    else
      {
      itkExceptionMacro("The input of writer should be whether a double precision"
                        "or a single precision transform type.");
      }
    }
}

/** Set the input transform and reinitialize the list of transforms */
template<typename ScalarType>
void TransformFileWriterTemplate<ScalarType>
::SetInput(const Object *transform)
{
  m_TransformList.clear();
  this->PushBackTransformList(transform);
}

template<typename ScalarType>
const typename TransformFileWriterTemplate<ScalarType>::TransformType *
TransformFileWriterTemplate<ScalarType>
::GetInput()
{
  ConstTransformPointer res = *(m_TransformList.begin());
  return res.GetPointer();
}

/** Add a transform to be written */
template<typename ScalarType>
void TransformFileWriterTemplate<ScalarType>
::AddTransform(const Object *transform)
{
  /* Check for a CompositeTransform.
   * The convention is that there should be one, and it should
   * be the first transform in the file
   */
  const std::string transformName = transform->GetNameOfClass();
  if( transformName.find("CompositeTransform") != std::string::npos )
    {
    if(this->m_TransformList.size() > 0)
      {
      itkExceptionMacro("Can only write a transform of type CompositeTransform "
                        "as the first transform in the file.");
      }
    }

  this->PushBackTransformList(transform);
}

template<typename ScalarType>
void TransformFileWriterTemplate<ScalarType>
::Update()
{
  if ( m_FileName == "" )
    {
    itkExceptionMacro ("No file name given");
    }
  typename TransformIOBaseTemplate<ScalarType>::Pointer transformIO =
    TransformIOFactoryTemplate<ScalarType>::CreateTransformIO( m_FileName.c_str(), WriteMode );
  if ( transformIO.IsNull() )
    {
    itkExceptionMacro("Can't Create IO object for file "
                      << m_FileName);
    }
  transformIO->SetAppendMode(this->m_AppendMode);
  transformIO->SetFileName(m_FileName);
  transformIO->SetTransformList(this->m_TransformList);
  transformIO->Write();
}

template<typename ScalarType>
void TransformFileWriterTemplate<ScalarType>
::PrintSelf(std::ostream & os, Indent indent) const
{
  Superclass::PrintSelf(os, indent);

  os << indent << "FileName: " << m_FileName << std::endl;
}

} // namespace itk

#endif
<|endoftext|>


<commit_before>///*===================================================================

//The Medical Imaging Interaction Toolkit (MITK)

//Copyright (c) German Cancer Research Center,
//Division of Medical and Biological Informatics.
//All rights reserved.

//This software is distributed WITHOUT ANY WARRANTY; without
//even the implied warranty of MERCHANTABILITY or FITNESS FOR
//A PARTICULAR PURPOSE.

//See LICENSE.txt or http://www.mitk.org for details.

//===================================================================*/
#include <mitkTestFixture.h>
#include <mitkTestingMacros.h>

#include <mitkPASpectralUnmixingFilterBase.h>
#include <mitkPALinearSpectralUnmixingFilter.h>
#include <mitkImageReadAccessor.h>

class mitkSpectralUnmixingTestSuite : public mitk::TestFixture
{
  CPPUNIT_TEST_SUITE(mitkSpectralUnmixingTestSuite);
  MITK_TEST(testSUAlgorithm);
  CPPUNIT_TEST_SUITE_END();

private:
  mitk::pa::SpectralUnmixingFilterBase::Pointer m_SpectralUnmixingFilter;
  mitk::Image::Pointer inputImage;
  std::vector<float> m_inputWavelengths;
  std::vector<float> m_CorrectResult;

public:

  void setUp() override
  {
    MITK_INFO << "setUp ... ";
    //Set empty input image:
    inputImage = mitk::Image::New();
    mitk::PixelType pixelType = mitk::MakeScalarPixelType<float>();
    const int NUMBER_OF_SPATIAL_DIMENSIONS = 3;
    auto* dimensions = new unsigned int[NUMBER_OF_SPATIAL_DIMENSIONS];

    dimensions[0] = 1;
    dimensions[1] = 1;
    dimensions[2] = 3;

    //Initialie empty input image:
    inputImage->Initialize(pixelType, NUMBER_OF_SPATIAL_DIMENSIONS, dimensions);

    //Set wavelengths for unmixing:
    m_inputWavelengths.push_back(750);
    m_inputWavelengths.push_back(800);
    m_inputWavelengths.push_back(850);

    //Set fraction of Hb and HbO2 to unmix:
    float fracHb = 0.3;
    float fracHbO2 = 0.67;

    //Fractions are also correct unmixing result:
    m_CorrectResult.push_back(fracHbO2);
    m_CorrectResult.push_back(fracHb);

    //Calculate values of wavelengths (750,800,850 nm) and multiply with fractions to get pixel values:
    float px1 = fracHb * 7.52 + fracHbO2 * 2.77;
    float px2 = fracHb * 4.08 + fracHbO2 * 4.37;
    float px3 = fracHb * 3.7 + fracHbO2 * 5.67;
    std::vector<float> m_Value{px1,px2,px3};

    float* data = new float[3];
    data[0] = m_Value[0];
    data[1] = m_Value[1];
    data[2] = m_Value[2];

    MITK_INFO << "data0 " << data[0];
    MITK_INFO << "data1 " << data[1];
    MITK_INFO << "data2 " << data[2];

    inputImage->SetImportVolume(data, mitk::Image::ImportMemoryManagementType::CopyMemory);
    delete[] data;

    MITK_INFO << "[DONE]";
  }

  void testSUAlgorithm()
  {
    MITK_INFO << "START FILTER TEST ... ";
    // Set input image
    auto m_SpectralUnmixingFilter = mitk::pa::LinearSpectralUnmixingFilter::New();
    m_SpectralUnmixingFilter->SetInput(inputImage);

    // Set Algortihm to filter
    m_SpectralUnmixingFilter->SetAlgorithm(mitk::pa::LinearSpectralUnmixingFilter::AlgortihmType::householderQr);

    //Set wavelengths to filter
    for (unsigned int imageIndex = 0; imageIndex < m_inputWavelengths.size(); imageIndex++)
    {
      unsigned int wavelength = m_inputWavelengths[imageIndex];
      m_SpectralUnmixingFilter->AddWavelength(wavelength);
    }

    //Set Chromophores to filter
    m_SpectralUnmixingFilter->AddChromophore(
      mitk::pa::SpectralUnmixingFilterBase::ChromophoreType::OXYGENATED_HEMOGLOBIN);

    m_SpectralUnmixingFilter->AddChromophore(
      mitk::pa::SpectralUnmixingFilterBase::ChromophoreType::DEOXYGENATED_HEMOGLOBIN);

    m_SpectralUnmixingFilter->Update();
    float threshold = 1e-5;

    for (int i = 0; i < 2; ++i)
    {
      mitk::Image::Pointer output = m_SpectralUnmixingFilter->GetOutput(i);
      mitk::ImageReadAccessor readAccess(output);
      const float* inputDataArray = ((const float*)readAccess.GetData());
      auto pixel = inputDataArray[0];

      MITK_INFO << "CorrectResult: " << m_CorrectResult[i];
      MITK_INFO << "FilterResult: " << pixel;

      CPPUNIT_ASSERT(std::abs(pixel - m_CorrectResult[i])<threshold);
    }
    MITK_INFO << "FILTER TEST SUCCESFULL :)";
  }

  void tearDown() override
  {
    m_SpectralUnmixingFilter = nullptr;
    inputImage = nullptr;
    m_inputWavelengths.clear();
    m_CorrectResult.clear();
    MITK_INFO << "tearDown ... [DONE]";
  }
};

MITK_TEST_SUITE_REGISTRATION(mitkSpectralUnmixing)
<commit_msg>add all solvers to test.<commit_after>///*===================================================================

//The Medical Imaging Interaction Toolkit (MITK)

//Copyright (c) German Cancer Research Center,
//Division of Medical and Biological Informatics.
//All rights reserved.

//This software is distributed WITHOUT ANY WARRANTY; without
//even the implied warranty of MERCHANTABILITY or FITNESS FOR
//A PARTICULAR PURPOSE.

//See LICENSE.txt or http://www.mitk.org for details.

//===================================================================*/
#include <mitkTestFixture.h>
#include <mitkTestingMacros.h>

#include <mitkPASpectralUnmixingFilterBase.h>
#include <mitkPALinearSpectralUnmixingFilter.h>
#include <mitkImageReadAccessor.h>

class mitkSpectralUnmixingTestSuite : public mitk::TestFixture
{
  CPPUNIT_TEST_SUITE(mitkSpectralUnmixingTestSuite);
  MITK_TEST(testSUAlgorithm);
  CPPUNIT_TEST_SUITE_END();

private:
  mitk::pa::SpectralUnmixingFilterBase::Pointer m_SpectralUnmixingFilter;
  mitk::Image::Pointer inputImage;
  std::vector<float> m_inputWavelengths;
  std::vector<float> m_CorrectResult;
  float threshold;

public:

  void setUp() override
  {
    MITK_INFO << "setUp ... ";
    //Set empty input image:
    inputImage = mitk::Image::New();
    mitk::PixelType pixelType = mitk::MakeScalarPixelType<float>();
    const int NUMBER_OF_SPATIAL_DIMENSIONS = 3;
    auto* dimensions = new unsigned int[NUMBER_OF_SPATIAL_DIMENSIONS];

    dimensions[0] = 1;
    dimensions[1] = 1;
    dimensions[2] = 3;

    inputImage->Initialize(pixelType, NUMBER_OF_SPATIAL_DIMENSIONS, dimensions);

    //Set wavelengths for unmixing:
    m_inputWavelengths.push_back(750);
    m_inputWavelengths.push_back(800);
    m_inputWavelengths.push_back(850);

    //Set fraction of Hb and HbO2 to unmix:
    float fracHb = 0.396579;
    float fracHbO2 = 0.594845;
    m_CorrectResult.push_back(fracHbO2);
    m_CorrectResult.push_back(fracHb);

    //Set threshold 1 digit
    threshold = 0.000001;

    //Multiply values of wavelengths (750,800,850 nm) with fractions to get pixel values:
    float px1 = fracHb * 7.52 + fracHbO2 * 2.77;
    float px2 = fracHb * 4.08 + fracHbO2 * 4.37;
    float px3 = fracHb * 3.7 + fracHbO2 * 5.67;
    std::vector<float> m_Value{px1,px2,px3};

    float* data = new float[3];
    data[0] = m_Value[0];
    data[1] = m_Value[1];
    data[2] = m_Value[2];

    inputImage->SetImportVolume(data, mitk::Image::ImportMemoryManagementType::CopyMemory);
    delete[] data;
    MITK_INFO << "[DONE]";
  }

  void testSUAlgorithm()
  {
    MITK_INFO << "START FILTER TEST ... ";
    // Set input image
    auto m_SpectralUnmixingFilter = mitk::pa::LinearSpectralUnmixingFilter::New();
    m_SpectralUnmixingFilter->SetInput(inputImage);

    // Set Algortihm to filter
    //m_SpectralUnmixingFilter->SetAlgorithm(mitk::pa::LinearSpectralUnmixingFilter::AlgortihmType::householderQr);
    //m_SpectralUnmixingFilter->SetAlgorithm(mitk::pa::LinearSpectralUnmixingFilter::AlgortihmType::ldlt);
    //m_SpectralUnmixingFilter->SetAlgorithm(mitk::pa::LinearSpectralUnmixingFilter::AlgortihmType::llt);
    m_SpectralUnmixingFilter->SetAlgorithm(mitk::pa::LinearSpectralUnmixingFilter::AlgortihmType::colPivHouseholderQr);
    //m_SpectralUnmixingFilter->SetAlgorithm(mitk::pa::LinearSpectralUnmixingFilter::AlgortihmType::jacobiSvd);
    //m_SpectralUnmixingFilter->SetAlgorithm(mitk::pa::LinearSpectralUnmixingFilter::AlgortihmType::fullPivLu);
    //m_SpectralUnmixingFilter->SetAlgorithm(mitk::pa::LinearSpectralUnmixingFilter::AlgortihmType::fullPivHouseholderQr);
    //m_SpectralUnmixingFilter->SetAlgorithm(mitk::pa::LinearSpectralUnmixingFilter::AlgortihmType::test);

    //Set wavelengths to filter
    for (unsigned int imageIndex = 0; imageIndex < m_inputWavelengths.size(); imageIndex++)
    {
      unsigned int wavelength = m_inputWavelengths[imageIndex];
      m_SpectralUnmixingFilter->AddWavelength(wavelength);
    }

    //Set Chromophores to filter
    m_SpectralUnmixingFilter->AddChromophore(
      mitk::pa::SpectralUnmixingFilterBase::ChromophoreType::OXYGENATED_HEMOGLOBIN);
    m_SpectralUnmixingFilter->AddChromophore(
      mitk::pa::SpectralUnmixingFilterBase::ChromophoreType::DEOXYGENATED_HEMOGLOBIN);

    m_SpectralUnmixingFilter->Update();

    for (int i = 0; i < 2; ++i)
    {
      mitk::Image::Pointer output = m_SpectralUnmixingFilter->GetOutput(i);
      mitk::ImageReadAccessor readAccess(output);
      const float* inputDataArray = ((const float*)readAccess.GetData());
      auto pixel = inputDataArray[0];
      /*For pixel values and results look at: [...]mitk-superbuild\MITK-build\Modules\PhotoacousticsLib\test\*/
      CPPUNIT_ASSERT(std::abs(pixel - m_CorrectResult[i])<threshold);
    }
    MITK_INFO << "FILTER TEST SUCCESFULL :)";
  }

  void tearDown() override
  {
    m_SpectralUnmixingFilter = nullptr;
    inputImage = nullptr;
    m_inputWavelengths.clear();
    m_CorrectResult.clear();
    MITK_INFO << "tearDown ... [DONE]";
  }
};

MITK_TEST_SUITE_REGISTRATION(mitkSpectralUnmixing)
<|endoftext|>


<commit_before>/// \file AddTaskNucleiYield.C
/// \brief Simple macro to add the task to a grid job
///
/// The task is here added several times to analyse different particle species and to investigate
/// different set of cuts in only one run.
///
/// \author Maximiliano Puccio <maximiliano.puccio@cern.ch>, University and INFN Torino
/// \date Feb 19th, 2015

#if !defined(__CINT__) || defined(__MAKECINT__)
#include <Rtypes.h>
#include <TString.h>
#include "AliAnalysisTaskNucleiYield.h"
#include "AliAnalysisManager.h"
#include "AliAnalysisDataContainer.h"
#include "AliPID.h"
#endif

AliAnalysisTaskNucleiYield* AddTaskNucleiYield__pp2016(Bool_t isMC = kFALSE,
    AliPID::EParticleType part = AliPID::kProton,
    Int_t pdgCode = 2212,
    TString tskname = "proton",
    TString suffix = "") {

  // Get the current analysis manager
  AliAnalysisManager *mgr = AliAnalysisManager::GetAnalysisManager();
  if (!mgr) {
    Error("AddTaskNucleiYield", "No analysis manager found.");
    return 0x0;
  }

  // Check the analysis type using the event handlers connected to the analysis manager.
  if (!mgr->GetInputEventHandler()) {
    ::Error("AddTaskProtonYield__pp5TeV", "This task requires an input event handler");
    return 0x0;
  }

  tskname.Append(Form("%s",suffix.Data()));

  AliAnalysisTaskNucleiYield *task = new AliAnalysisTaskNucleiYield(tskname);

  task->SetParticleType(part);
  task->SetPDG(pdgCode);
  task->SetIsMC(isMC);
  task->SetDCABins(80,-0.5,0.5);
	task->SaveTrees();

  task->SetRequireTPCpidSigmas(3.f);
  float cent[14] = {-5.f,0.f,1.f,5.f,10.f,20.f,30.f,40.f,50.f,60.f,70.f,80.f,90.f,100.f};
  task->SetCentBins(13, cent);
  task->SetUseFlattening(false);
  float pt[20] = {
    0.5f,0.6f,0.7f,0.8f,0.9f,1.0f,1.1f,1.2f,1.4f,1.6f,
    1.8f,2.0f,2.2f,2.6f,3.0f,3.4f,3.8f,4.4f,5.0f,6.0f
  };
  task->SetPtBins(19,pt);

  float dcabins[53] = {
    -1.30,-1.20,-1.10,-1.00,-0.90,-0.80,-0.70,-0.60,-0.50,-0.40,
    -0.35,-0.30,-0.25,-0.20,-0.15,-0.12,-0.10,-0.09,-0.08,-0.07,
    -0.06,-0.05,-0.04,-0.03,-0.02,-0.01, 0.00, 0.01, 0.02, 0.03,
     0.04, 0.05, 0.06, 0.07, 0.08, 0.09, 0.10, 0.12, 0.15, 0.20,
     0.25, 0.30, 0.35, 0.40, 0.50, 0.60, 0.70, 0.80, 0.90, 1.00,
     1.10, 1.20, 1.30
  };
  task->SetDCABins(52,dcabins);

  mgr->AddTask(task);

  TString output = "AnalysisResults.root";
  AliAnalysisDataContainer *cont = mgr->CreateContainer(Form("nuclei_%s",tskname.Data()),
      TList::Class(),
      AliAnalysisManager::kOutputContainer,
      output.Data());
  mgr->ConnectInput  (task,  0, mgr->GetCommonInputContainer());
  mgr->ConnectOutput (task,  1, cont);
  return task;
}
<commit_msg>Fix addrtask for proton tree<commit_after>#if !defined(__CINT__) || defined(__MAKECINT__)
#include <Rtypes.h>
#include <TString.h>
#include "AliAnalysisTaskNucleiYield.h"
#include "AliAnalysisManager.h"
#include "AliAnalysisDataContainer.h"
#include "AliPID.h"
#endif

AliAnalysisTaskNucleiYield* AddTaskNucleiYield_pp5TeV(Bool_t isMC = kFALSE,
    AliPID::EParticleType part = AliPID::kProton,
    Int_t pdgCode = 2212,
    TString tskname = "proton",
    TString suffix = "") {

  // Get the current analysis manager
  AliAnalysisManager *mgr = AliAnalysisManager::GetAnalysisManager();
  if (!mgr) {
    Error("AddTaskProtonYield_pp5TeV", "No analysis manager found.");
    return 0x0;
  }

  // Check the analysis type using the event handlers connected to the analysis manager.
  if (!mgr->GetInputEventHandler()) {
    ::Error("AddTaskProtonYield_pp5TeV", "This task requires an input event handler");
    return 0x0;
  }

  tskname.Append(Form("%s",suffix.Data()));

  AliAnalysisTaskNucleiYield *task = new AliAnalysisTaskNucleiYield(tskname);

  task->SetParticleType(part);
  task->SetPDG(pdgCode);
  task->SetIsMC(isMC);
  task->SetDCABins(80,-0.5,0.5);
	task->SaveTrees();

  task->SetRequireTPCpidSigmas(3.f);
  float cent[14] = {-5.f,0.f,1.f,5.f,10.f,20.f,30.f,40.f,50.f,60.f,70.f,80.f,90.f,100.f};
  task->SetCentBins(13, cent);
  task->SetUseFlattening(false);
  float pt[20] = {
    0.5f,0.6f,0.7f,0.8f,0.9f,1.0f,1.1f,1.2f,1.4f,1.6f,
    1.8f,2.0f,2.2f,2.6f,3.0f,3.4f,3.8f,4.4f,5.0f,6.0f
  };
  task->SetPtBins(19,pt);

  float dcabins[53] = {
    -1.30,-1.20,-1.10,-1.00,-0.90,-0.80,-0.70,-0.60,-0.50,-0.40,
    -0.35,-0.30,-0.25,-0.20,-0.15,-0.12,-0.10,-0.09,-0.08,-0.07,
    -0.06,-0.05,-0.04,-0.03,-0.02,-0.01, 0.00, 0.01, 0.02, 0.03,
     0.04, 0.05, 0.06, 0.07, 0.08, 0.09, 0.10, 0.12, 0.15, 0.20,
     0.25, 0.30, 0.35, 0.40, 0.50, 0.60, 0.70, 0.80, 0.90, 1.00,
     1.10, 1.20, 1.30
  };
  task->SetDCABins(52,dcabins);

  mgr->AddTask(task);

  TString output = "AnalysisResults.root";
  AliAnalysisDataContainer *cont = mgr->CreateContainer(Form("nuclei_%s",tskname.Data()),
      TList::Class(),
      AliAnalysisManager::kOutputContainer,
      output.Data());
  mgr->ConnectInput  (task,  0, mgr->GetCommonInputContainer());
  mgr->ConnectOutput (task,  1, cont);
  return task;
}
<|endoftext|>


<commit_before>/**************************************************************************
 * Copyright(c) 1998-2015, ALICE Experiment at CERN, All rights reserved. *
 *                                                                        *
 * Author: The ALICE Off-line Project.                                    *
 * Contributors are mentioned in the code where appropriate.              *
 *                                                                        *
 * Permission to use, copy, modify and distribute this software and its   *
 * documentation strictly for non-commercial purposes is hereby granted   *
 * without fee, proviyaded that the above copyright notice appears in all *
 * copies and that both the copyright notice and this permission notice   *
 * appear in the supporting documentation. The authors make no claims     *
 * about the suitability of this software for any purapose. It is         *
 * provided "as is" without express or implied warranty.                  *
 **************************************************************************/

///////////////////////////////////////////////////////////////////////////////
//                                                                           //
// This analysis extracts information into TTree for pT-spectra of deuterons //
// Based on AliAnalysisDeuteronpA task of J. Anielski for deuteron analysis  //
// and AliAnalysisTaskExtractV0 by D. Chinellato for TTree interface         //
// L.Barnby October 2015                                                     //
///////////////////////////////////////////////////////////////////////////////

#include "TTree.h"
#include "TH1F.h"
#include "TList.h"
#include "TChain.h"

#include "AliAnalysisTaskSE.h"
#include "AliAnalysisManager.h"

#include "AliPID.h"
#include "AliESDtrackCuts.h"
#include "AliESDVertex.h"
#include "AliESDEvent.h"
#include "AliESDInputHandler.h"
#include "AliESDtrack.h"
#include "AliESDpid.h"
#include "AliCentrality.h"
#include "AliESDUtils.h"
#include "AliMultiplicity.h"

#include "AliMCEventHandler.h"
#include "AliMCEvent.h"
#include "AliStack.h"
#include "AliLog.h"
#include "AliAnalysisUtils.h"

#include "AliAnalysisDeuteronTree.h"

ClassImp(AliAnalysisDeuteronTree)

//________________________________________________________________________
AliAnalysisDeuteronTree::AliAnalysisDeuteronTree()
: AliAnalysisTaskSE(), fListHist(0), fTree(0),fESDtrackCuts(0),fPIDResponse(0),
fMCtrue(0),
fRapCMSpA(0),
fUtils(0),
fTimeStamp(0),
fCentrality(0),
fPtCor(0),
fPt(0),
fMom(0),
fRapd(0),
fNsigmaTPCd(0),
fNsigmaTOFd(0),
fDcaXYd(0),
fMcCode(0),
fNpion(0),
fhZVertex(0),
fhCentrality(0),
fhNsigmaTPCvsMom(0),
fhNsigmaTOFvsMom(0),
fMomCorrConstA(0),
fMomCorrConstB(0),
fMomCorrPower(0),
fMinPtCut(0),
fMinPtTOFCut(0),
fNsigmaTOFdMax(0),
fNsigmaTOFdMin(0),
fNsigmaTPCdMax(0),
fNsigmaTPCdMin(0)
{
    //Default constructor
}

//________________________________________________________________________
AliAnalysisDeuteronTree::AliAnalysisDeuteronTree(const char *name)
: AliAnalysisTaskSE(name), fListHist(0), fTree(0), fESDtrackCuts(0),fPIDResponse(0),
fMCtrue(0),
fRapCMSpA(0),
fUtils(0),
fTimeStamp(0),
fCentrality(0),
fPtCor(0),
fPt(0),
fMom(0),
fRapd(0),
fNsigmaTPCd(0),
fNsigmaTOFd(0),
fDcaXYd(0),
fMcCode(0),
fNpion(0),
fhZVertex(0),
fhCentrality(0),
fhNsigmaTPCvsMom(0),
fhNsigmaTOFvsMom(0),
fMomCorrConstA(0),
fMomCorrConstB(0),
fMomCorrPower(0),
fMinPtCut(0),
fMinPtTOFCut(0),
fNsigmaTOFdMax(0),
fNsigmaTOFdMin(0),
fNsigmaTPCdMax(0),
fNsigmaTPCdMin(0)
{
    //Standard constructor
    fMCtrue = kFALSE;
    fRapCMSpA = kTRUE;

    //fESDtrackCuts = new AliESDtrackCuts("AliESDtrackCuts","AliESDtrackCuts");
    Initialize();

    //AliInfo("About to define TChain input");
    //DefineInput(0, TChain::Class());

    AliInfo("About to define TList output");
    DefineOutput(1, TList::Class());

    AliInfo("About to define TTree output");
    DefineOutput(2, TTree::Class());

    AliInfo("Constructor Finished");
}

//________________________________________________________________________
void AliAnalysisDeuteronTree::Initialize()
{
    //
    // updating parameters in case of changes
    //
    AliInfo("Initialization started");

    fESDtrackCuts = AliESDtrackCuts::GetStandardITSTPCTrackCuts2011(kFALSE,kTRUE);
    if (fESDtrackCuts == 0x0) {
        AliWarning("No ESDtrackCuts");
    }
    fESDtrackCuts->SetMaxDCAToVertexXY(3);
    fESDtrackCuts->SetMaxDCAToVertexZ(2);
    fESDtrackCuts->SetEtaRange(-0.9,0.9);
    //

    // Momentum correction
    fMomCorrConstA = 0.333303;
    fMomCorrConstB = 0.651111;
    fMomCorrPower = 5.27268;

    fMinPtCut = 0.5;
    fMinPtTOFCut = 1.0;
    fNsigmaTOFdMax = 5.0;
    fNsigmaTOFdMin = -5.0;
    fNsigmaTPCdMax = 5.0;
    fNsigmaTPCdMin = -5.0;

    AliInfo("Initialization complete");
}

//________________________________________________________________________
AliAnalysisDeuteronTree::~AliAnalysisDeuteronTree()
{
    if (fListHist) {
        delete fListHist;
        fListHist = 0x0;
    }
    if (fTree) {
        delete fTree;
        fTree = 0x0;
    }
    //cleanup esd track cuts object too...
    if (fESDtrackCuts) {
        delete fESDtrackCuts;
        fESDtrackCuts = 0x0;
    }
    if (fPIDResponse) {
        delete fPIDResponse;
        fPIDResponse = 0x0;
    }
    if (fUtils) {
        delete fUtils;
        fUtils = 0x0;
    }

}

//________________________________________________________________________
void AliAnalysisDeuteronTree::UserCreateOutputObjects()
{
    AliInfo("Creation started");

    // Proper handling of output objects ("magic" lines)
    // For file-resident Tree
    OpenFile(2);
    fTree = new TTree("fTree","dCandidates");

    fListHist = new TList();
    fListHist->SetOwner();

    // Get PID response object
    AliAnalysisManager *man=AliAnalysisManager::GetAnalysisManager();
    if(!man)
        AliFatal("Could not find manager");
    AliInputEventHandler* inputHandler = dynamic_cast<AliInputEventHandler*> (man->GetInputEventHandler());
    if(!inputHandler)
        AliFatal("No input event handler");
        fPIDResponse = dynamic_cast<AliPIDResponse *>(inputHandler->GetPIDResponse());
    if (!fPIDResponse)
        AliFatal("PIDResponse object was not created"); // Escalated to fatal. Task is unusable without PID response.

    //Create analysis utils for event selection and pileup rejection
    fUtils = new AliAnalysisUtils();
    fUtils->SetCutOnZVertexSPD(kFALSE);

    // histograms
    fhZVertex = new TH1F("hZVertex","Vertex Z position;Z_{VTX} (cm);Counts/0.25 cm",100,-25,25);
    fListHist->Add(fhZVertex);
    fhCentrality = new TH1F("hCentrality",";Centrality(V0A) %;Counts/1%",100,0,100);
    fListHist->Add(fhCentrality);

    fhNsigmaTPCvsMom = new TH2F("hNsigmaTPCvsMom","n_{#sigma} TPC distribution;p (GeV/#it{c});n_{#sigma}^deuteron",300,0.0,6.0,200,-15.0,5);
    fListHist->Add(fhNsigmaTPCvsMom);

    fhNsigmaTOFvsMom = new TH2F("hNsigmaTOFvsMom","n_{#sigma} TOF distribution;p (GeV/#it{c});n_{#sigma}^deuteron",300,0.0,6.0,200,-15.0,5);
    fListHist->Add(fhNsigmaTOFvsMom);

    // fTree Branch definitions
    fTree->Branch("fCentrality",&fCentrality,"fCentrality/F");
    fTree->Branch("fPtCor",&fPtCor,"fPtCor/F");
    fTree->Branch("fPt",&fPt,"fPt/F");
    fTree->Branch("fMom",&fMom,"fMom/F");
    fTree->Branch("fRapd",&fRapd,"fRapd/F");
    fTree->Branch("fNsigmaTPCd",&fNsigmaTPCd,"fNsigmaTPCd/F");
    fTree->Branch("fNsigmaTOFd",&fNsigmaTOFd,"fNsigmaTOFd/F");
    fTree->Branch("fDcaXYd",&fDcaXYd,"fDcaXYd/F");
    fTree->Branch("fMcCode",&fMcCode,"fMcCode/I");
    fTree->Branch("fNpion",&fNpion,"fNpion/I");
    PostData(1, fListHist);
    PostData(2, fTree);

    AliInfo("Creation finished");
}

//________________________________________________________________________
void AliAnalysisDeuteronTree::UserExec(Option_t *option){

    AliInfo("UserExec started");

    AliESDEvent *lESDevent = 0x0;
    lESDevent = dynamic_cast<AliESDEvent*>( InputEvent() );
    if (!lESDevent) {
        AliWarning("ERROR: lESDevent not available \n");
        return;
    }

    AliCentrality *centrality = lESDevent->GetCentrality();

    // Later add code for accessing MC information
    // AliMCEventHandler, AliMCEvent, AliStack

    // Event checks
    //first event in chunck --> continue
    if(fUtils->IsFirstEventInChunk(lESDevent)) {
        PostData(1, fListHist);
        PostData(2, fTree);
        return;
    }
    //check for pileup
    if (fUtils->IsPileUpEvent(lESDevent)){
        PostData(1, fListHist);
        PostData(2, fTree);
       return;
    }
    //vertex cuts
    Bool_t isVtxOk = kTRUE;
    if(!fUtils->IsVertexSelected2013pA(lESDevent)) isVtxOk = kFALSE;

    if (!fESDtrackCuts) {
        AliFatal("ERROR: fESDtrackCuts not available"); // No sense to continue without track cuts
        return;
    }

    // Physics selection should be taken care of by analysis manager
    // and task->SetCollisionsCandidates(AliVEvent::kMB) in the macro

    const AliESDVertex *vertex = lESDevent->GetPrimaryVertex();
    if (vertex && isVtxOk) fhZVertex->Fill(vertex->GetZ());

    fCentrality = centrality->GetCentralityPercentile("V0A");
    fhCentrality->Fill(fCentrality);

    fTimeStamp = lESDevent->GetTimeStamp();

    Float_t dca[2], cov[3]; // dca_xy, dca_z, sigma_xy, sigma_xy_z, sigma_z for the vertex cut
    Double_t lTOFNsigma, lTPCNsigma, lEnergyDeuteron, lPz;
    Int_t lNpion;

    for (Int_t i=0;i<lESDevent->GetNumberOfTracks();++i) {
        //
        AliESDtrack *track = lESDevent->GetTrack(i);

        if (!track->GetInnerParam()) continue;

        Double_t ptot = track->GetInnerParam()->GetP(); // momentum for dEdx determination
        fMom=ptot*track->Charge(); // Shouldn't this be corrected for the different mass assumption as below
        // i.e find uncorrected pz, correct pt and then recombine to get corrected ptot?

        //
        // momentum correction for different mass assumption in tracking
        //
        fPt = track->Pt();
        fPtCor = track->Pt()/(1 - fMomCorrConstA/TMath::Power(track->Pt() + fMomCorrConstB, fMomCorrPower));

        track->GetImpactParameters(dca, cov);
        if (!fESDtrackCuts->AcceptTrack(track)) continue;
        if (track->Pt()< fMinPtCut) continue;

        AliPIDResponse::EDetPidStatus lETOFStatus, lETPCStatus;
        lETOFStatus = fPIDResponse->NumberOfSigmas(AliPIDResponse::kTOF, track, AliPID::kDeuteron, lTOFNsigma);
        fNsigmaTOFd = -999;
        if(lETOFStatus==AliPIDResponse::kDetPidOk) fNsigmaTOFd = lTOFNsigma;

        lETPCStatus = fPIDResponse->NumberOfSigmas(AliPIDResponse::kTPC, track, AliPID::kDeuteron, lTPCNsigma);
        fNsigmaTPCd = -999;
        if(lETPCStatus==AliPIDResponse::kDetPidOk) fNsigmaTPCd = lTPCNsigma;

        fDcaXYd = dca[0];

        lPz = TMath::Sqrt(ptot*ptot - fPt*fPt);
        lEnergyDeuteron = TMath::Sqrt(fPt*fPt + lPz*lPz +
                                      AliPID::ParticleMass(AliPID::kDeuteron)*AliPID::ParticleMass(AliPID::kDeuteron));
        fRapd = 0.5*TMath::Log((lEnergyDeuteron + lPz)/(lEnergyDeuteron - lPz));
        fMcCode = 0;

        fhNsigmaTPCvsMom->Fill(ptot,fNsigmaTPCd);
        fhNsigmaTOFvsMom->Fill(track->GetP(),fNsigmaTOFd);

        // Eventually need some selections here to stop the tree getting too big (because of all the non-deuterons)
        if (fNsigmaTPCd < fNsigmaTPCdMin || fNsigmaTPCd > fNsigmaTPCdMax) continue; // within TPC dEd/dx
        if (fPt > fMinPtTOFCut && (fNsigmaTOFd > fNsigmaTOFdMin && fNsigmaTOFd < fNsigmaTOFdMax)) {
            // Greater than a certain pt, must also satisfy TOF Nsigma
            lNpion = 0;
            // Only in this case loop to get pions
            for (Int_t j=0; j<lESDevent->GetNumberOfTracks();++j) {
                if (i!=j) { // The deuteron candidate should not be used as the pion
                    AliESDtrack *pion = lESDevent->GetTrack(j);
                    if (!fESDtrackCuts->AcceptTrack(pion)) continue;
                    if (!pion->GetInnerParam()) continue;
                    lETPCStatus = fPIDResponse->NumberOfSigmas(AliPIDResponse::kTPC, pion, AliPID::kPion, lTPCNsigma);
                    if (TMath::Abs(lTPCNsigma)<3.5) {
                        lNpion++;
                    }
                }

            }
            fNpion = lNpion;
            fTree->Fill();
        }

    } // End track loop

    // Output data
    PostData(1, fListHist);
    PostData(2, fTree);
    return;

}

void AliAnalysisDeuteronTree::Terminate(Option_t *)
{
    // Draw control histograms (if any)
}<commit_msg>Fix logic for selecting d candidates<commit_after>/**************************************************************************
 * Copyright(c) 1998-2015, ALICE Experiment at CERN, All rights reserved. *
 *                                                                        *
 * Author: The ALICE Off-line Project.                                    *
 * Contributors are mentioned in the code where appropriate.              *
 *                                                                        *
 * Permission to use, copy, modify and distribute this software and its   *
 * documentation strictly for non-commercial purposes is hereby granted   *
 * without fee, proviyaded that the above copyright notice appears in all *
 * copies and that both the copyright notice and this permission notice   *
 * appear in the supporting documentation. The authors make no claims     *
 * about the suitability of this software for any purapose. It is         *
 * provided "as is" without express or implied warranty.                  *
 **************************************************************************/

///////////////////////////////////////////////////////////////////////////////
//                                                                           //
// This analysis extracts information into TTree for pT-spectra of deuterons //
// Based on AliAnalysisDeuteronpA task of J. Anielski for deuteron analysis  //
// and AliAnalysisTaskExtractV0 by D. Chinellato for TTree interface         //
// L.Barnby October 2015                                                     //
///////////////////////////////////////////////////////////////////////////////

#include "TTree.h"
#include "TH1F.h"
#include "TList.h"
#include "TChain.h"

#include "AliAnalysisTaskSE.h"
#include "AliAnalysisManager.h"

#include "AliPID.h"
#include "AliESDtrackCuts.h"
#include "AliESDVertex.h"
#include "AliESDEvent.h"
#include "AliESDInputHandler.h"
#include "AliESDtrack.h"
#include "AliESDpid.h"
#include "AliCentrality.h"
#include "AliESDUtils.h"
#include "AliMultiplicity.h"

#include "AliMCEventHandler.h"
#include "AliMCEvent.h"
#include "AliStack.h"
#include "AliLog.h"
#include "AliAnalysisUtils.h"

#include "AliAnalysisDeuteronTree.h"

ClassImp(AliAnalysisDeuteronTree)

//________________________________________________________________________
AliAnalysisDeuteronTree::AliAnalysisDeuteronTree()
: AliAnalysisTaskSE(), fListHist(0), fTree(0),fESDtrackCuts(0),fPIDResponse(0),
fMCtrue(0),
fRapCMSpA(0),
fUtils(0),
fTimeStamp(0),
fCentrality(0),
fPtCor(0),
fPt(0),
fMom(0),
fRapd(0),
fNsigmaTPCd(0),
fNsigmaTOFd(0),
fDcaXYd(0),
fMcCode(0),
fNpion(0),
fhZVertex(0),
fhCentrality(0),
fhNsigmaTPCvsMom(0),
fhNsigmaTOFvsMom(0),
fMomCorrConstA(0),
fMomCorrConstB(0),
fMomCorrPower(0),
fMinPtCut(0),
fMinPtTOFCut(0),
fNsigmaTOFdMax(0),
fNsigmaTOFdMin(0),
fNsigmaTPCdMax(0),
fNsigmaTPCdMin(0)
{
    //Default constructor
}

//________________________________________________________________________
AliAnalysisDeuteronTree::AliAnalysisDeuteronTree(const char *name)
: AliAnalysisTaskSE(name), fListHist(0), fTree(0), fESDtrackCuts(0),fPIDResponse(0),
fMCtrue(0),
fRapCMSpA(0),
fUtils(0),
fTimeStamp(0),
fCentrality(0),
fPtCor(0),
fPt(0),
fMom(0),
fRapd(0),
fNsigmaTPCd(0),
fNsigmaTOFd(0),
fDcaXYd(0),
fMcCode(0),
fNpion(0),
fhZVertex(0),
fhCentrality(0),
fhNsigmaTPCvsMom(0),
fhNsigmaTOFvsMom(0),
fMomCorrConstA(0),
fMomCorrConstB(0),
fMomCorrPower(0),
fMinPtCut(0),
fMinPtTOFCut(0),
fNsigmaTOFdMax(0),
fNsigmaTOFdMin(0),
fNsigmaTPCdMax(0),
fNsigmaTPCdMin(0)
{
    //Standard constructor
    fMCtrue = kFALSE;
    fRapCMSpA = kTRUE;

    //fESDtrackCuts = new AliESDtrackCuts("AliESDtrackCuts","AliESDtrackCuts");
    Initialize();

    //AliInfo("About to define TChain input");
    //DefineInput(0, TChain::Class());

    AliInfo("About to define TList output");
    DefineOutput(1, TList::Class());

    AliInfo("About to define TTree output");
    DefineOutput(2, TTree::Class());

    AliInfo("Constructor Finished");
}

//________________________________________________________________________
void AliAnalysisDeuteronTree::Initialize()
{
    //
    // updating parameters in case of changes
    //
    AliInfo("Initialization started");

    fESDtrackCuts = AliESDtrackCuts::GetStandardITSTPCTrackCuts2011(kFALSE,kTRUE);
    if (fESDtrackCuts == 0x0) {
        AliWarning("No ESDtrackCuts");
    }
    fESDtrackCuts->SetMaxDCAToVertexXY(3);
    fESDtrackCuts->SetMaxDCAToVertexZ(2);
    fESDtrackCuts->SetEtaRange(-0.9,0.9);
    //

    // Momentum correction
    fMomCorrConstA = 0.333303;
    fMomCorrConstB = 0.651111;
    fMomCorrPower = 5.27268;

    fMinPtCut = 0.5;
    fMinPtTOFCut = 1.0;
    fNsigmaTOFdMax = 5.0;
    fNsigmaTOFdMin = -5.0;
    fNsigmaTPCdMax = 5.0;
    fNsigmaTPCdMin = -5.0;

    AliInfo("Initialization complete");
}

//________________________________________________________________________
AliAnalysisDeuteronTree::~AliAnalysisDeuteronTree()
{
    if (fListHist) {
        delete fListHist;
        fListHist = 0x0;
    }
    if (fTree) {
        delete fTree;
        fTree = 0x0;
    }
    //cleanup esd track cuts object too...
    if (fESDtrackCuts) {
        delete fESDtrackCuts;
        fESDtrackCuts = 0x0;
    }
    if (fPIDResponse) {
        delete fPIDResponse;
        fPIDResponse = 0x0;
    }
    if (fUtils) {
        delete fUtils;
        fUtils = 0x0;
    }

}

//________________________________________________________________________
void AliAnalysisDeuteronTree::UserCreateOutputObjects()
{
    AliInfo("Creation started");

    // Proper handling of output objects ("magic" lines)
    // For file-resident Tree
    OpenFile(2);
    fTree = new TTree("fTree","dCandidates");

    fListHist = new TList();
    fListHist->SetOwner();

    // Get PID response object
    AliAnalysisManager *man=AliAnalysisManager::GetAnalysisManager();
    if(!man)
        AliFatal("Could not find manager");
    AliInputEventHandler* inputHandler = dynamic_cast<AliInputEventHandler*> (man->GetInputEventHandler());
    if(!inputHandler)
        AliFatal("No input event handler");
        fPIDResponse = dynamic_cast<AliPIDResponse *>(inputHandler->GetPIDResponse());
    if (!fPIDResponse)
        AliFatal("PIDResponse object was not created"); // Escalated to fatal. Task is unusable without PID response.

    //Create analysis utils for event selection and pileup rejection
    fUtils = new AliAnalysisUtils();
    fUtils->SetCutOnZVertexSPD(kFALSE);

    // histograms
    fhZVertex = new TH1F("hZVertex","Vertex Z position;Z_{VTX} (cm);Counts/0.25 cm",100,-25,25);
    fListHist->Add(fhZVertex);
    fhCentrality = new TH1F("hCentrality",";Centrality(V0A) %;Counts/1%",100,0,100);
    fListHist->Add(fhCentrality);

    fhNsigmaTPCvsMom = new TH2F("hNsigmaTPCvsMom","n_{#sigma} TPC distribution;p (GeV/#it{c});n_{#sigma}^deuteron",300,0.0,6.0,200,-15.0,5);
    fListHist->Add(fhNsigmaTPCvsMom);

    fhNsigmaTOFvsMom = new TH2F("hNsigmaTOFvsMom","n_{#sigma} TOF distribution;p (GeV/#it{c});n_{#sigma}^deuteron",300,0.0,6.0,200,-15.0,5);
    fListHist->Add(fhNsigmaTOFvsMom);

    // fTree Branch definitions
    fTree->Branch("fCentrality",&fCentrality,"fCentrality/F");
    fTree->Branch("fPtCor",&fPtCor,"fPtCor/F");
    fTree->Branch("fPt",&fPt,"fPt/F");
    fTree->Branch("fMom",&fMom,"fMom/F");
    fTree->Branch("fRapd",&fRapd,"fRapd/F");
    fTree->Branch("fNsigmaTPCd",&fNsigmaTPCd,"fNsigmaTPCd/F");
    fTree->Branch("fNsigmaTOFd",&fNsigmaTOFd,"fNsigmaTOFd/F");
    fTree->Branch("fDcaXYd",&fDcaXYd,"fDcaXYd/F");
    fTree->Branch("fMcCode",&fMcCode,"fMcCode/I");
    fTree->Branch("fNpion",&fNpion,"fNpion/I");
    PostData(1, fListHist);
    PostData(2, fTree);

    AliInfo("Creation finished");
}

//________________________________________________________________________
void AliAnalysisDeuteronTree::UserExec(Option_t *option){

    AliInfo("UserExec started");

    AliESDEvent *lESDevent = 0x0;
    lESDevent = dynamic_cast<AliESDEvent*>( InputEvent() );
    if (!lESDevent) {
        AliWarning("ERROR: lESDevent not available \n");
        return;
    }

    AliCentrality *centrality = lESDevent->GetCentrality();

    // Later add code for accessing MC information
    // AliMCEventHandler, AliMCEvent, AliStack

    // Event checks
    //first event in chunck --> continue
    if(fUtils->IsFirstEventInChunk(lESDevent)) {
        PostData(1, fListHist);
        PostData(2, fTree);
        return;
    }
    //check for pileup
    if (fUtils->IsPileUpEvent(lESDevent)){
        PostData(1, fListHist);
        PostData(2, fTree);
       return;
    }
    //vertex cuts
    Bool_t isVtxOk = kTRUE;
    if(!fUtils->IsVertexSelected2013pA(lESDevent)) isVtxOk = kFALSE;

    if (!fESDtrackCuts) {
        AliFatal("ERROR: fESDtrackCuts not available"); // No sense to continue without track cuts
        return;
    }

    // Physics selection should be taken care of by analysis manager
    // and task->SetCollisionsCandidates(AliVEvent::kMB) in the macro

    const AliESDVertex *vertex = lESDevent->GetPrimaryVertex();
    if (vertex && isVtxOk) fhZVertex->Fill(vertex->GetZ());

    fCentrality = centrality->GetCentralityPercentile("V0A");
    fhCentrality->Fill(fCentrality);

    fTimeStamp = lESDevent->GetTimeStamp();

    Float_t dca[2], cov[3]; // dca_xy, dca_z, sigma_xy, sigma_xy_z, sigma_z for the vertex cut
    Double_t lTOFNsigma, lTPCNsigma, lEnergyDeuteron, lPz;
    Int_t lNpion;

    for (Int_t i=0;i<lESDevent->GetNumberOfTracks();++i) {
        //
        AliESDtrack *track = lESDevent->GetTrack(i);

        if (!track->GetInnerParam()) continue;

        Double_t ptot = track->GetInnerParam()->GetP(); // momentum for dEdx determination
        fMom=ptot*track->Charge(); // Shouldn't this be corrected for the different mass assumption as below
        // i.e find uncorrected pz, correct pt and then recombine to get corrected ptot?

        //
        // momentum correction for different mass assumption in tracking
        //
        fPt = track->Pt();
        fPtCor = track->Pt()/(1 - fMomCorrConstA/TMath::Power(track->Pt() + fMomCorrConstB, fMomCorrPower));

        track->GetImpactParameters(dca, cov);
        if (!fESDtrackCuts->AcceptTrack(track)) continue;
        if (track->Pt()< fMinPtCut) continue;

        AliPIDResponse::EDetPidStatus lETOFStatus, lETPCStatus;
        lETOFStatus = fPIDResponse->NumberOfSigmas(AliPIDResponse::kTOF, track, AliPID::kDeuteron, lTOFNsigma);
        fNsigmaTOFd = -999;
        if(lETOFStatus==AliPIDResponse::kDetPidOk) fNsigmaTOFd = lTOFNsigma;

        lETPCStatus = fPIDResponse->NumberOfSigmas(AliPIDResponse::kTPC, track, AliPID::kDeuteron, lTPCNsigma);
        fNsigmaTPCd = -999;
        if(lETPCStatus==AliPIDResponse::kDetPidOk) fNsigmaTPCd = lTPCNsigma;

        fDcaXYd = dca[0];

        lPz = TMath::Sqrt(ptot*ptot - fPt*fPt);
        lEnergyDeuteron = TMath::Sqrt(fPt*fPt + lPz*lPz +
                                      AliPID::ParticleMass(AliPID::kDeuteron)*AliPID::ParticleMass(AliPID::kDeuteron));
        fRapd = 0.5*TMath::Log((lEnergyDeuteron + lPz)/(lEnergyDeuteron - lPz));
        fMcCode = 0;

        fhNsigmaTPCvsMom->Fill(ptot,fNsigmaTPCd);
        fhNsigmaTOFvsMom->Fill(track->GetP(),fNsigmaTOFd);

        // Need some selections here to stop the tree getting too big (because of all the non-deuterons)
        // outside TPC dEd/dx band
        if (fNsigmaTPCd < fNsigmaTPCdMin || fNsigmaTPCd > fNsigmaTPCdMax) continue;
        // over TOF pt threshold and outside TOF band
        if ((fPt > fMinPtTOFCut && (fNsigmaTOFd < fNsigmaTOFdMin || fNsigmaTOFd > fNsigmaTOFdMax))) continue;

        // Only in this case loop to get pions
        lNpion = 0;
        for (Int_t j=0; j<lESDevent->GetNumberOfTracks();++j) {
            if (i!=j) { // The deuteron candidate should not be used as the pion
                AliESDtrack *pion = lESDevent->GetTrack(j);
                if (!fESDtrackCuts->AcceptTrack(pion)) continue;
                if (!pion->GetInnerParam()) continue;
                lETPCStatus = fPIDResponse->NumberOfSigmas(AliPIDResponse::kTPC, pion, AliPID::kPion, lTPCNsigma);
                if (TMath::Abs(lTPCNsigma)<3.5) {
                    lNpion++;
                }
            }

        }
        fNpion = lNpion;
        fTree->Fill();

    } // End track loop

    // Output data
    PostData(1, fListHist);
    PostData(2, fTree);
    return;

}

void AliAnalysisDeuteronTree::Terminate(Option_t *)
{
    // Draw control histograms (if any)
}<|endoftext|>


<commit_before>#define NDEBUG

#include <QLinkedList>
#include <QVector>
#include <QSet>
#include <QMap>
#include <QHash>

#include <stdlib.h>
#include <stdio.h>
#include <stdbool.h>
#include <string.h>
#include <sys/types.h>
#include <sys/resource.h>
#include <sys/time.h>

using namespace std;

unsigned long g_result;

static unsigned long long
cputime (void)
{
  struct rusage rus;
  getrusage (0, &rus);
  return rus.ru_utime.tv_sec * 1000000ULL + rus.ru_utime.tv_usec;
}

static unsigned int randValue = 0;
static void rand_init(void)
{
  randValue = 0;
}
static unsigned int rand_get(void)
{
  randValue = randValue * 31421U + 6927U;
  return randValue;
}

static void test_function(const char str[], size_t n, void (*func)(size_t))
{
  unsigned long long start, end;
  //  (*func)(n);
  start = cputime();
  (*func)(n);
  end = cputime();
  end = (end - start) / 1000U;
  printf ("%s %Lu ms for n = %lu\n", str, end, n);
}

/********************************************************************************************/

static void test_list(size_t n)
{
  rand_init();
  QLinkedList<unsigned int> a1, a2;
  for(size_t i = 0; i < n; i++) {
    a1.push_back(rand_get() );
    a2.push_back(rand_get() );
  }
  unsigned int s = 0;
  for (QLinkedList<unsigned int>::const_iterator ci1 = a1.begin(), ci2 = a2.begin (); ci1 != a1.end(); ++ci1, ++ci2) {
    s += *ci1 * *ci2;
  }
  g_result = s;
}

/********************************************************************************************/

static void test_array(size_t n)
{
  rand_init();
  QVector<unsigned int> a1, a2;
  for(size_t i = 0; i < n; i++) {
    a1.push_back(rand_get() );
    a2.push_back(rand_get() );
  }
  unsigned int s = 0;
  for(unsigned long i = 0; i < n; i++) {
    s += a1[i] * a2[i];
  }
  g_result = s;
}

/********************************************************************************************/

/* NOTE: QSet doesn't perform full ordering. I can't find any container
   which is a set with total ordering within Qt. */
#if 0
static void test_rbtree(size_t n)
{
  rand_init();
  QSet<unsigned long> tree;

  for (size_t i = 0; i < n; i++) {
    tree.insert(rand_get());
  }

  unsigned int s = 0;
  for (size_t i = 0; i < n; i++) {
    QSet<unsigned long>::iterator it = tree.find(rand_get());
    if (it != tree.end())
      s += *it;
  }
  g_result = s;
}
#endif

/********************************************************************************************/

static void
test_dict1(unsigned long  n)
{
  rand_init();
  QMap<unsigned long, unsigned long> dict;

  for (size_t i = 0; i < n; i++) {
    dict[rand_get()] = rand_get();
  }

  unsigned int s = 0;
  for (size_t i = 0; i < n; i++) {
    QMap<unsigned long, unsigned long>::iterator it = dict.find(rand_get());
    if (it != dict.end())
      s += it.value();
  }
  g_result = s;
}

/********************************************************************************************/

static void
test_dict2(unsigned long  n)
{
  rand_init();
  QHash<unsigned long, unsigned long> dict;

  for (size_t i = 0; i < n; i++) {
    dict[rand_get()] = rand_get();
  }

  unsigned int s = 0;
  for (size_t i = 0; i < n; i++) {
    QHash<unsigned long, unsigned long>::iterator it = dict.find(rand_get());
    if (it != dict.end())
      s += it.value();
  }
  g_result = s;
}

/********************************************************************************************/
struct char_array_s {
  char a[256];
  char_array_s () { a[0] = 0 ; }
  char_array_s ( const char_array_s & other) { strcpy(a, other.a); }
  bool operator==(const char_array_s &other) const { return strcmp(a, other.a) == 0; }
};

uint qHash(const char_array_s &k)
{
  size_t hash = 0;
  const char *s = k.a;
  while (*s) hash = hash * 31421 + (*s++) + 6927;
  return hash;
};

static void
test_dict_big(unsigned long  n)
{
  rand_init();
  QHash<char_array_s, char_array_s> dict;

  for (size_t i = 0; i < n; i++) {
    char_array_s s1, s2;
    sprintf(s1.a, "%u", rand_get());
    sprintf(s2.a, "%u", rand_get());
    dict[s1] = s2;
  }

  unsigned int s = 0;
  for (size_t i = 0; i < n; i++) {
    char_array_s s1;
    sprintf(s1.a, "%u", rand_get());
    QHash<char_array_s, char_array_s>::iterator it = dict.find(s1);
    if (it != dict.end())
      s ++;
  }
  g_result = s;
}

/********************************************************************************************/

int main(int argc, const char *argv[])
{
  int n = (argc > 1) ? atoi(argv[1]) : 0;
  if (n == 1)
    test_function("List   time",10000000, test_list);
  if (n == 2)
    test_function("Array  time", 100000000, test_array);
  /* if (n == 3)
     test_function("Rbtree time", 1000000, test_rbtree); */
  if (n == 4)
    test_function("Dict(m)time", 1000000, test_dict1);
  if (n == 5)
    test_function("Dict(u)time", 1000000, test_dict2);
  if (n == 6)
    test_function("DictB  time", 1000000, test_dict_big);
}
<commit_msg>Implement test_rbtree in bench-qt.<commit_after>#define NDEBUG

#include <QLinkedList>
#include <QVector>
#include <QSet>
#include <QMap>
#include <QHash>

#include <stdlib.h>
#include <stdio.h>
#include <stdbool.h>
#include <string.h>
#include <sys/types.h>
#include <sys/resource.h>
#include <sys/time.h>

using namespace std;

unsigned long g_result;

static unsigned long long
cputime (void)
{
  struct rusage rus;
  getrusage (0, &rus);
  return rus.ru_utime.tv_sec * 1000000ULL + rus.ru_utime.tv_usec;
}

static unsigned int randValue = 0;
static void rand_init(void)
{
  randValue = 0;
}
static unsigned int rand_get(void)
{
  randValue = randValue * 31421U + 6927U;
  return randValue;
}

static void test_function(const char str[], size_t n, void (*func)(size_t))
{
  unsigned long long start, end;
  //  (*func)(n);
  start = cputime();
  (*func)(n);
  end = cputime();
  end = (end - start) / 1000U;
  printf ("%s %Lu ms for n = %lu\n", str, end, n);
}

/********************************************************************************************/

static void test_list(size_t n)
{
  rand_init();
  QLinkedList<unsigned int> a1, a2;
  for(size_t i = 0; i < n; i++) {
    a1.push_back(rand_get() );
    a2.push_back(rand_get() );
  }
  unsigned int s = 0;
  for (QLinkedList<unsigned int>::const_iterator ci1 = a1.begin(), ci2 = a2.begin (); ci1 != a1.end(); ++ci1, ++ci2) {
    s += *ci1 * *ci2;
  }
  g_result = s;
}

/********************************************************************************************/

static void test_array(size_t n)
{
  rand_init();
  QVector<unsigned int> a1, a2;
  for(size_t i = 0; i < n; i++) {
    a1.push_back(rand_get() );
    a2.push_back(rand_get() );
  }
  unsigned int s = 0;
  for(unsigned long i = 0; i < n; i++) {
    s += a1[i] * a2[i];
  }
  g_result = s;
}

/********************************************************************************************/

/* NOTE: QSet doesn't perform full ordering. I can't find any container
   which is a set with total ordering within Qt. So we have to fake it. */
static void test_rbtree(size_t n)
{
  rand_init();
  QMap<unsigned long, QHashDummyValue> tree;

  for (size_t i = 0; i < n; i++) {
    tree[rand_get()] = QHashDummyValue();
  }

  unsigned int s = 0;
  for (size_t i = 0; i < n; i++) {
    QMap<unsigned long, QHashDummyValue>::iterator it = tree.find(rand_get());
    if (it != tree.end())
      s += it.key();
  }
  g_result = s;
}

/********************************************************************************************/

static void
test_dict1(unsigned long  n)
{
  rand_init();
  QMap<unsigned long, unsigned long> dict;

  for (size_t i = 0; i < n; i++) {
    dict[rand_get()] = rand_get();
  }

  unsigned int s = 0;
  for (size_t i = 0; i < n; i++) {
    QMap<unsigned long, unsigned long>::iterator it = dict.find(rand_get());
    if (it != dict.end())
      s += it.value();
  }
  g_result = s;
}

/********************************************************************************************/

static void
test_dict2(unsigned long  n)
{
  rand_init();
  QHash<unsigned long, unsigned long> dict;

  for (size_t i = 0; i < n; i++) {
    dict[rand_get()] = rand_get();
  }

  unsigned int s = 0;
  for (size_t i = 0; i < n; i++) {
    QHash<unsigned long, unsigned long>::iterator it = dict.find(rand_get());
    if (it != dict.end())
      s += it.value();
  }
  g_result = s;
}

/********************************************************************************************/
struct char_array_s {
  char a[256];
  char_array_s () { a[0] = 0 ; }
  char_array_s ( const char_array_s & other) { strcpy(a, other.a); }
  bool operator==(const char_array_s &other) const { return strcmp(a, other.a) == 0; }
};

uint qHash(const char_array_s &k)
{
  size_t hash = 0;
  const char *s = k.a;
  while (*s) hash = hash * 31421 + (*s++) + 6927;
  return hash;
};

static void
test_dict_big(unsigned long  n)
{
  rand_init();
  QHash<char_array_s, char_array_s> dict;

  for (size_t i = 0; i < n; i++) {
    char_array_s s1, s2;
    sprintf(s1.a, "%u", rand_get());
    sprintf(s2.a, "%u", rand_get());
    dict[s1] = s2;
  }

  unsigned int s = 0;
  for (size_t i = 0; i < n; i++) {
    char_array_s s1;
    sprintf(s1.a, "%u", rand_get());
    QHash<char_array_s, char_array_s>::iterator it = dict.find(s1);
    if (it != dict.end())
      s ++;
  }
  g_result = s;
}

/********************************************************************************************/

int main(int argc, const char *argv[])
{
  int n = (argc > 1) ? atoi(argv[1]) : 0;
  if (n == 1)
    test_function("List   time",10000000, test_list);
  if (n == 2)
    test_function("Array  time", 100000000, test_array);
  if (n == 3)
     test_function("Rbtree time", 1000000, test_rbtree);
  if (n == 4)
    test_function("Dict(m)time", 1000000, test_dict1);
  if (n == 5)
    test_function("Dict(u)time", 1000000, test_dict2);
  if (n == 6)
    test_function("DictB  time", 1000000, test_dict_big);
}
<|endoftext|>