<commit_before>/***************************************************************************
              fbellini@cern.ch - last modified on 28/11/2013
//
//Lauches KStar analysis with rsn mini package
//Allows basic configuration of pile-up check and event cuts
//
****************************************************************************/
enum pairYCutSet { kPairDefault,
		   kCentralTight,
		   kpADefault,
		   kpACentral
		 };

enum eventCutSet { kEvtDefault = 0,  //data
		   kEvtDefaultVtx8,
		   kEvtDefaultVtx5,
		   kNoPileUpCut,
		   kMCEvt,         //MC events cuts
		   kMCEvtDefault};

AliRsnMiniAnalysisTask * AddTaskF0
(
 TString     outNameSuffix = "f0",                         //suffix for output container
 Bool_t      isMC          = 0,                            //MC flag
 AliPIDResponse::EBeamType collSys = AliPIDResponse::kPP, //=0, kPPB=1, kPBPB=2 outNameSuffix
 UInt_t      triggerMask   = AliVEvent::kINT7,             //trigger selection
 Bool_t      enaMultSel    = kTRUE,                        //enable multiplicity axis
 Int_t       evtCutSetID   = eventCutSet::kEvtDefault,     //event selection
 Int_t       pairCutSetID  = pairYCutSet::kPairDefault,    //pair cuts
 Int_t       aodFilterBit  = 5,                            //AOD filter bit for AOD only
 AliRsnCutSetDaughterParticle::ERsnDaughterCutSet cutPiPid = AliRsnCutSetDaughterParticle::kTPCpidTOFveto3s, //PID cut
 Float_t     nsigma        = 3.0,   //PID cut
 Float_t     masslow       = 0.3,   //inv mass range lower boundary
 Float_t     massup        = 1.3,   //inv mass range upper boundary
 Int_t       nbins         = 1000,  //inv mass: N bins
 Bool_t      enableMonitor = kTRUE) //enable single track QA
{

  //-------------------------------------------
  // event cuts
  //-------------------------------------------
  /*Trigger selection
  Std minimum bias trigger AliVEvent::kMB, corresponding to (V0A | V0C | SPD) to be used with
  - pp 7 TeV (2010 data)
  - PbPb 2.76 TeV (2010 data and 2011 data)
  - pp 2.76 TeV (2011 data)
  Centrality triggers AliVEvent::kCentral, AliVEvent::kSemicentral to be used with
  - PbPb 2.76 TeV (2011 data)
  Std minimum bias trigger AliVEvent::kINT7, corrsesponding to (V0A & V0C) to be used with
  - pA 5.02 TeV (2013 data)
  - pp 13 TeV (2015 data)
  */
  Double_t vtxZcut            = 10.0; //cm, default cut on vtx z
  Bool_t rejectPileUp         = kTRUE; //rejects pileup from SPD
  Bool_t useMVPileUpSelection = kFALSE; //alternative pile-up rejection, default is SPD
  Int_t  MinPlpContribSPD     = 5; //default value if used
  Int_t  MinPlpContribMV      = 5; //default value if used
  // Bool_t selectDPMJETevtNSDpA = kFALSE; //cut to select true NSD events in DPMJET

  if (evtCutSetID==eventCutSet::kEvtDefaultVtx8) vtxZcut = 8.0; //cm
  if (evtCutSetID==eventCutSet::kEvtDefaultVtx8) vtxZcut = 5.0; //cm
  if (evtCutSetID==eventCutSet::kNoPileUpCut) rejectPileUp = kFALSE;

  if (evtCutSetID>=eventCutSet::kMCEvt) {
    vtxZcut = 1.0e3; //cm
  }

  if (evtCutSetID==eventCutSet::kMCEvtDefault) {
    vtxZcut = 10.0; //cm
  }

  //-------------------------------------------
  //pair rapidity cut
  //-------------------------------------------
  Double_t minYlab = -0.5;
  Double_t maxYlab =  0.5;

  if (pairCutSetID==pairYCutSet::kCentralTight) { //|y_cm|<0.3
    minYlab = -0.3;    maxYlab = 0.3;
  }

  if (pairCutSetID==pairYCutSet::kpADefault) { //-0.5<y_cm<0.0
    minYlab = -0.465;    maxYlab = 0.035;
  }

  if (pairCutSetID==pairYCutSet::kpACentral) { //|y_cm|<0.3
    minYlab = -0.765;    maxYlab = -0.165;
  }

  //-------------------------------------------
  //mixing settings
  //-------------------------------------------
  Int_t       nmix = 5;
  Float_t     maxDiffVzMix = 1.0;
  Float_t     maxDiffMultMix = 5.0;

  //
  // -- INITIALIZATION ----------------------------------------------------------------------------
  //
  // retrieve analysis manager
  AliAnalysisManager *mgr = AliAnalysisManager::GetAnalysisManager();
  if (!mgr) {
    ::Error("AddTaskF0", "No analysis manager to connect to.");
    return NULL;
  }

  // create the task and configure
  TString taskName = Form("RsnTaskF0");
  AliRsnMiniAnalysisTask *task = new AliRsnMiniAnalysisTask(taskName.Data(), isMC);

  //trigger
  task->UseESDTriggerMask(triggerMask); //ESD
  //task->SelectCollisionCandidates(triggerMask); //AOD

  //-----------------------------------------------------------------------------------------------
  // -- MULTIPLICITY/CENTRALITY -------------------------------------------------------------------
  //-----------------------------------------------------------------------------------------------
  if (collSys==AliPIDResponse::kPP) task->UseMultiplicity("AliMultSelection_V0M");
  if (collSys==AliPIDResponse::kPPB) task->UseCentrality("AliMultSelection_V0A");
  if (collSys==AliPIDResponse::kPBPB) task->UseCentrality("AliMultSelection_V0M");

  //-----------------------------------------------------------------------------------------------
  // -- EVENT MIXING CONFIG -----------------------------------------------------------------------
  //-----------------------------------------------------------------------------------------------
  task->UseContinuousMix();
  task->SetNMix(nmix);
  task->SetMaxDiffVz(maxDiffVzMix);
  task->SetMaxDiffMult(maxDiffMultMix);
  ::Info("AddTaskF0", Form("Event mixing configuration: \n events to mix = %i \n max diff. vtxZ = cm %5.3f \n max diff multi = %5.3f", nmix, maxDiffVzMix, maxDiffMultMix));

  //-----------------------------------------------------------------------------------------------
  // -- EVENT SELECTION ---------------------------------------------------------------------------
  //-----------------------------------------------------------------------------------------------
   AliRsnCutPrimaryVertex *cutVertex = new AliRsnCutPrimaryVertex("cutVertex", 10.0, 0, kFALSE);
  if (isPP && (!isMC)) cutVertex->SetCheckPileUp(rejectPileUp);   // set the check for pileup
  cutVertex->SetCheckZResolutionSPD(); ::Info("AddTaskF0", Form("CheckZResolutionSPD:              ON"));
  cutVertex->SetCheckDispersionSPD(); ::Info("AddTaskF0", Form("CheckDispersionSPD:               ON"));
  cutVertex->SetCheckZDifferenceSPDTrack(); ::Info("AddTaskF0", Form("CheckZDifferenceSPDTrack:         ON"));

  //set check for pileup in 2013
  AliRsnCutEventUtils *cutEventUtils = new AliRsnCutEventUtils("cutEventUtils", kFALSE, rejectPileUp);
  cutEventUtils->SetCheckIncompleteDAQ(kTRUE);
  cutEventUtils->SetCheckSPDClusterVsTrackletBG();
  ::Info("AddTaskF0", Form("CheckIncompleteDAQ:                  ON"));
  ::Info("AddTaskF0", Form("SetCheckSPDClusterVsTrackletBG:      ON"));

  if (useMVPileUpSelection){
    cutEventUtils->SetUseMVPlpSelection(useMVPileUpSelection);
    cutEventUtils->SetMinPlpContribMV(MinPlpContribMV);
    cutEventUtils->SetMinPlpContribSPD(MinPlpContribSPD);
    ::Info("AddTaskF0", Form("Multiple-vtx Pile-up rejection:      ON \nSettings: MinPlpContribMV = %i, MinPlpContribSPD = %i", MinPlpContribMV, MinPlpContribSPD));
  } else {
    cutEventUtils->SetMinPlpContribSPD(MinPlpContribSPD);
    ::Info("AddTaskF0", Form("SPD Pile-up rejection:                      ON \nSettings: MinPlpContribSPD = %i", MinPlpContribSPD));
  }
  ::Info("AddTaskF0", Form("Pile-up rejection mode:                %s", (rejectPileUp?"ON":"OFF")));

  AliRsnCutSet *eventCuts = new AliRsnCutSet("eventCuts", AliRsnTarget::kEvent);
  eventCuts->AddCut(cutEventUtils);
  eventCuts->AddCut(cutVertex);
  eventCuts->SetCutScheme(Form("%s&%s", cutEventUtils->GetName(), cutVertex->GetName()));
  task->SetEventCuts(eventCuts);

  //connect task
  mgr->AddTask(task);

  //-----------------------------------------------------------------------------------------------
  // -- EVENT SELECTION MONITOR -------------------------------------------------------------------
  //-----------------------------------------------------------------------------------------------
  //vertex position monitoring
  Int_t vtxID = task->CreateValue(AliRsnMiniValue::kVz, kFALSE);
  AliRsnMiniOutput *outVtx = task->CreateOutput("eventVtx", "HIST", "EVENT");
  outVtx->AddAxis(vtxID, 500, -50.0, 50.0);

  //multiplicity or centrality monitoring -- if enabled
  if (enaMultSel){

    //multiplicity or centrality with forward estimator from AliMulSelectionTask
    Int_t multID = task->CreateValue(AliRsnMiniValue::kMult, kFALSE);
    //reference multiplicity (default with global tracks with good quality, if not available uses tracklets)
    Int_t multRefID = task->CreateValue(AliRsnMiniValue::kRefMult, kFALSE);

    AliRsnMiniOutput *outMult = task->CreateOutput("eventMult", "HIST", "EVENT");
    outMult->AddAxis(multID, 100, 0.0, 100.0);

    AliRsnMiniOutput *outRefMult = task->CreateOutput("eventRefMult", "HIST", "EVENT");
    outRefMult->AddAxis(multRefID, 400, 0.0, 400.0);

    TH2F* hvz = new TH2F("hVzVsCent",Form("Vertex position vs centrality"), 101, 0., 101., 500, -50.0, 50.0);
    if (collSys==AliPIDResponse::kPPB)
      hvz->GetXaxis()->SetTitle("V0A");
    else
      hvz->GetXaxis()->SetTitle("V0M");
    hvz->GetYaxis()->SetTitle("z_{vtx} (cm)");
    task->SetEventQAHist("vz", hvz);

    TH2F* hRefMultiVsCent = new TH2F("hRefMultiVsCent",Form("Reference multiplicity vs centrality"), 101, 0., 101., 400, 0., 400.);
    if (collSys==AliPIDResponse::kPPB)
      hRefMultiVsCent->GetXaxis()->SetTitle("V0A");
    else
      hRefMultiVsCent->GetXaxis()->SetTitle("V0M");
    hRefMultiVsCent->GetYaxis()->SetTitle("GLOBAL");
    task->SetEventQAHist("refmulti",hRefMultiVsCent);

    TH2F* hMultiVsCent = new TH2F("hMultiVsCent",Form("Multiplicity vs centrality"), 101, 0., 101., 400, 0., 400.);
    if (collSys==AliPIDResponse::kPPB)
      hMultiVsCent->GetXaxis()->SetTitle("V0A");
    else
      hMultiVsCent->GetXaxis()->SetTitle("V0M");
    hMultiVsCent->GetYaxis()->SetTitle("QUALITY");
    task->SetEventQAHist("multicent",hMultiVsCent);
  }

  //-----------------------------------------------------------------------------------------------
  // -- PAIR CUTS (common to all resonances) ------------------------------------------------------
  //-----------------------------------------------------------------------------------------------
  AliRsnCutMiniPair *cutY = new AliRsnCutMiniPair("cutRapidity", AliRsnCutMiniPair::kRapidityRange);
  cutY->SetRangeD(minYlab, maxYlab);

  AliRsnCutSet *cutsPair = new AliRsnCutSet("pairCuts", AliRsnTarget::kMother);
  cutsPair->AddCut(cutY);
  cutsPair->SetCutScheme(cutY->GetName());

  //-----------------------------------------------------------------------------------------------
  // -- CONFIG ANALYSIS --------------------------------------------------------------------------
  //-----------------------------------------------------------------------------------------------
  gROOT->LoadMacro("$ALICE_PHYSICS/PWGLF/RESONANCES/macros/mini/ConfigF0.C");
  if (!ConfigF0(task, isMC, collSys, cutsPair, enaMultSel, masslow, massup, nbins, aodFilterBit, cutPiPid, nsigma, enableMonitor) ) return 0x0;

  //-----------------------------------------------------------------------------------------------
  // -- CONTAINERS --------------------------------------------------------------------------------
  //-----------------------------------------------------------------------------------------------
  TString outputFileName = AliAnalysisManager::GetCommonFileName();
  Printf("AddTaskF0 - Set OutputFileName : \n %s\n", outputFileName.Data() );
  AliAnalysisDataContainer *output = mgr->CreateContainer(Form("RsnOut_%s",outNameSuffix.Data()),
							  TList::Class(),
							  AliAnalysisManager::kOutputContainer,
							  outputFileName);
  mgr->ConnectInput(task, 0, mgr->GetCommonInputContainer());
  mgr->ConnectOutput(task, 1, output);

  return task;
}
<commit_msg>Bug fix: remove isPP flag and change with coll. system option.<commit_after>/***************************************************************************
              fbellini@cern.ch - last modified on 28/11/2013
//
//Lauches KStar analysis with rsn mini package
//Allows basic configuration of pile-up check and event cuts
//
****************************************************************************/
enum pairYCutSet { kPairDefault,
		   kCentralTight,
		   kpADefault,
		   kpACentral
		 };

enum eventCutSet { kEvtDefault = 0,  //data
		   kEvtDefaultVtx8,
		   kEvtDefaultVtx5,
		   kNoPileUpCut,
		   kMCEvt,         //MC events cuts
		   kMCEvtDefault};

AliRsnMiniAnalysisTask * AddTaskF0
(
 TString     outNameSuffix = "f0",                         //suffix for output container
 Bool_t      isMC          = 0,                            //MC flag
 AliPIDResponse::EBeamType collSys = AliPIDResponse::kPP, //=0, kPPB=1, kPBPB=2 outNameSuffix
 UInt_t      triggerMask   = AliVEvent::kINT7,             //trigger selection
 Bool_t      enaMultSel    = kTRUE,                        //enable multiplicity axis
 Int_t       evtCutSetID   = eventCutSet::kEvtDefault,     //event selection
 Int_t       pairCutSetID  = pairYCutSet::kPairDefault,    //pair cuts
 Int_t       aodFilterBit  = 5,                            //AOD filter bit for AOD only
 AliRsnCutSetDaughterParticle::ERsnDaughterCutSet cutPiPid = AliRsnCutSetDaughterParticle::kTPCpidTOFveto3s, //PID cut
 Float_t     nsigma        = 3.0,   //PID cut
 Float_t     masslow       = 0.3,   //inv mass range lower boundary
 Float_t     massup        = 1.3,   //inv mass range upper boundary
 Int_t       nbins         = 1000,  //inv mass: N bins
 Bool_t      enableMonitor = kTRUE) //enable single track QA
{

  //-------------------------------------------
  // event cuts
  //-------------------------------------------
  /*Trigger selection
  Std minimum bias trigger AliVEvent::kMB, corresponding to (V0A | V0C | SPD) to be used with
  - pp 7 TeV (2010 data)
  - PbPb 2.76 TeV (2010 data and 2011 data)
  - pp 2.76 TeV (2011 data)
  Centrality triggers AliVEvent::kCentral, AliVEvent::kSemicentral to be used with
  - PbPb 2.76 TeV (2011 data)
  Std minimum bias trigger AliVEvent::kINT7, corrsesponding to (V0A & V0C) to be used with
  - pA 5.02 TeV (2013 data)
  - pp 13 TeV (2015 data)
  */
  Double_t vtxZcut            = 10.0; //cm, default cut on vtx z
  Bool_t rejectPileUp         = kTRUE; //rejects pileup from SPD
  Bool_t useMVPileUpSelection = kFALSE; //alternative pile-up rejection, default is SPD
  Int_t  MinPlpContribSPD     = 5; //default value if used
  Int_t  MinPlpContribMV      = 5; //default value if used
  // Bool_t selectDPMJETevtNSDpA = kFALSE; //cut to select true NSD events in DPMJET

  if (evtCutSetID==eventCutSet::kEvtDefaultVtx8) vtxZcut = 8.0; //cm
  if (evtCutSetID==eventCutSet::kEvtDefaultVtx8) vtxZcut = 5.0; //cm
  if (evtCutSetID==eventCutSet::kNoPileUpCut) rejectPileUp = kFALSE;

  if (evtCutSetID>=eventCutSet::kMCEvt) {
    vtxZcut = 1.0e3; //cm
  }

  if (evtCutSetID==eventCutSet::kMCEvtDefault) {
    vtxZcut = 10.0; //cm
  }

  //-------------------------------------------
  //pair rapidity cut
  //-------------------------------------------
  Double_t minYlab = -0.5;
  Double_t maxYlab =  0.5;

  if (pairCutSetID==pairYCutSet::kCentralTight) { //|y_cm|<0.3
    minYlab = -0.3;    maxYlab = 0.3;
  }

  if (pairCutSetID==pairYCutSet::kpADefault) { //-0.5<y_cm<0.0
    minYlab = -0.465;    maxYlab = 0.035;
  }

  if (pairCutSetID==pairYCutSet::kpACentral) { //|y_cm|<0.3
    minYlab = -0.765;    maxYlab = -0.165;
  }

  //-------------------------------------------
  //mixing settings
  //-------------------------------------------
  Int_t       nmix = 5;
  Float_t     maxDiffVzMix = 1.0;
  Float_t     maxDiffMultMix = 5.0;

  //
  // -- INITIALIZATION ----------------------------------------------------------------------------
  //
  // retrieve analysis manager
  AliAnalysisManager *mgr = AliAnalysisManager::GetAnalysisManager();
  if (!mgr) {
    ::Error("AddTaskF0", "No analysis manager to connect to.");
    return NULL;
  }

  // create the task and configure
  TString taskName = Form("RsnTaskF0");
  AliRsnMiniAnalysisTask *task = new AliRsnMiniAnalysisTask(taskName.Data(), isMC);

  //trigger
  task->UseESDTriggerMask(triggerMask); //ESD
  //task->SelectCollisionCandidates(triggerMask); //AOD

  //-----------------------------------------------------------------------------------------------
  // -- MULTIPLICITY/CENTRALITY -------------------------------------------------------------------
  //-----------------------------------------------------------------------------------------------
  if (collSys==AliPIDResponse::kPP) task->UseMultiplicity("AliMultSelection_V0M");
  if (collSys==AliPIDResponse::kPPB) task->UseCentrality("AliMultSelection_V0A");
  if (collSys==AliPIDResponse::kPBPB) task->UseCentrality("AliMultSelection_V0M");

  //-----------------------------------------------------------------------------------------------
  // -- EVENT MIXING CONFIG -----------------------------------------------------------------------
  //-----------------------------------------------------------------------------------------------
  task->UseContinuousMix();
  task->SetNMix(nmix);
  task->SetMaxDiffVz(maxDiffVzMix);
  task->SetMaxDiffMult(maxDiffMultMix);
  ::Info("AddTaskF0", Form("Event mixing configuration: \n events to mix = %i \n max diff. vtxZ = cm %5.3f \n max diff multi = %5.3f", nmix, maxDiffVzMix, maxDiffMultMix));

  //-----------------------------------------------------------------------------------------------
  // -- EVENT SELECTION ---------------------------------------------------------------------------
  //-----------------------------------------------------------------------------------------------
   AliRsnCutPrimaryVertex *cutVertex = new AliRsnCutPrimaryVertex("cutVertex", 10.0, 0, kFALSE);
   if ((collSys==AliPIDResponse::kPP) && (!isMC)) cutVertex->SetCheckPileUp(rejectPileUp);   // set the check for pileup
   cutVertex->SetCheckZResolutionSPD(); ::Info("AddTaskF0", Form("CheckZResolutionSPD:              ON"));
  cutVertex->SetCheckDispersionSPD(); ::Info("AddTaskF0", Form("CheckDispersionSPD:               ON"));
  cutVertex->SetCheckZDifferenceSPDTrack(); ::Info("AddTaskF0", Form("CheckZDifferenceSPDTrack:         ON"));

  //set check for pileup in 2013
  AliRsnCutEventUtils *cutEventUtils = new AliRsnCutEventUtils("cutEventUtils", kFALSE, rejectPileUp);
  cutEventUtils->SetCheckIncompleteDAQ(kTRUE);
  cutEventUtils->SetCheckSPDClusterVsTrackletBG();
  ::Info("AddTaskF0", Form("CheckIncompleteDAQ:                  ON"));
  ::Info("AddTaskF0", Form("SetCheckSPDClusterVsTrackletBG:      ON"));

  if (useMVPileUpSelection){
    cutEventUtils->SetUseMVPlpSelection(useMVPileUpSelection);
    cutEventUtils->SetMinPlpContribMV(MinPlpContribMV);
    cutEventUtils->SetMinPlpContribSPD(MinPlpContribSPD);
    ::Info("AddTaskF0", Form("Multiple-vtx Pile-up rejection:      ON \nSettings: MinPlpContribMV = %i, MinPlpContribSPD = %i", MinPlpContribMV, MinPlpContribSPD));
  } else {
    cutEventUtils->SetMinPlpContribSPD(MinPlpContribSPD);
    ::Info("AddTaskF0", Form("SPD Pile-up rejection:                      ON \nSettings: MinPlpContribSPD = %i", MinPlpContribSPD));
  }
  ::Info("AddTaskF0", Form("Pile-up rejection mode:                %s", (rejectPileUp?"ON":"OFF")));

  AliRsnCutSet *eventCuts = new AliRsnCutSet("eventCuts", AliRsnTarget::kEvent);
  eventCuts->AddCut(cutEventUtils);
  eventCuts->AddCut(cutVertex);
  eventCuts->SetCutScheme(Form("%s&%s", cutEventUtils->GetName(), cutVertex->GetName()));
  task->SetEventCuts(eventCuts);

  //connect task
  mgr->AddTask(task);

  //-----------------------------------------------------------------------------------------------
  // -- EVENT SELECTION MONITOR -------------------------------------------------------------------
  //-----------------------------------------------------------------------------------------------
  //vertex position monitoring
  Int_t vtxID = task->CreateValue(AliRsnMiniValue::kVz, kFALSE);
  AliRsnMiniOutput *outVtx = task->CreateOutput("eventVtx", "HIST", "EVENT");
  outVtx->AddAxis(vtxID, 500, -50.0, 50.0);

  //multiplicity or centrality monitoring -- if enabled
  if (enaMultSel){

    //multiplicity or centrality with forward estimator from AliMulSelectionTask
    Int_t multID = task->CreateValue(AliRsnMiniValue::kMult, kFALSE);
    //reference multiplicity (default with global tracks with good quality, if not available uses tracklets)
    Int_t multRefID = task->CreateValue(AliRsnMiniValue::kRefMult, kFALSE);

    AliRsnMiniOutput *outMult = task->CreateOutput("eventMult", "HIST", "EVENT");
    outMult->AddAxis(multID, 100, 0.0, 100.0);

    AliRsnMiniOutput *outRefMult = task->CreateOutput("eventRefMult", "HIST", "EVENT");
    outRefMult->AddAxis(multRefID, 400, 0.0, 400.0);

    TH2F* hvz = new TH2F("hVzVsCent",Form("Vertex position vs centrality"), 101, 0., 101., 500, -50.0, 50.0);
    if (collSys==AliPIDResponse::kPPB)
      hvz->GetXaxis()->SetTitle("V0A");
    else
      hvz->GetXaxis()->SetTitle("V0M");
    hvz->GetYaxis()->SetTitle("z_{vtx} (cm)");
    task->SetEventQAHist("vz", hvz);

    TH2F* hRefMultiVsCent = new TH2F("hRefMultiVsCent",Form("Reference multiplicity vs centrality"), 101, 0., 101., 400, 0., 400.);
    if (collSys==AliPIDResponse::kPPB)
      hRefMultiVsCent->GetXaxis()->SetTitle("V0A");
    else
      hRefMultiVsCent->GetXaxis()->SetTitle("V0M");
    hRefMultiVsCent->GetYaxis()->SetTitle("GLOBAL");
    task->SetEventQAHist("refmulti",hRefMultiVsCent);

    TH2F* hMultiVsCent = new TH2F("hMultiVsCent",Form("Multiplicity vs centrality"), 101, 0., 101., 400, 0., 400.);
    if (collSys==AliPIDResponse::kPPB)
      hMultiVsCent->GetXaxis()->SetTitle("V0A");
    else
      hMultiVsCent->GetXaxis()->SetTitle("V0M");
    hMultiVsCent->GetYaxis()->SetTitle("QUALITY");
    task->SetEventQAHist("multicent",hMultiVsCent);
  }

  //-----------------------------------------------------------------------------------------------
  // -- PAIR CUTS (common to all resonances) ------------------------------------------------------
  //-----------------------------------------------------------------------------------------------
  AliRsnCutMiniPair *cutY = new AliRsnCutMiniPair("cutRapidity", AliRsnCutMiniPair::kRapidityRange);
  cutY->SetRangeD(minYlab, maxYlab);

  AliRsnCutSet *cutsPair = new AliRsnCutSet("pairCuts", AliRsnTarget::kMother);
  cutsPair->AddCut(cutY);
  cutsPair->SetCutScheme(cutY->GetName());

  //-----------------------------------------------------------------------------------------------
  // -- CONFIG ANALYSIS --------------------------------------------------------------------------
  //-----------------------------------------------------------------------------------------------
  gROOT->LoadMacro("$ALICE_PHYSICS/PWGLF/RESONANCES/macros/mini/ConfigF0.C");
  if (!ConfigF0(task, isMC, collSys, cutsPair, enaMultSel, masslow, massup, nbins, aodFilterBit, cutPiPid, nsigma, enableMonitor) ) return 0x0;

  //-----------------------------------------------------------------------------------------------
  // -- CONTAINERS --------------------------------------------------------------------------------
  //-----------------------------------------------------------------------------------------------
  TString outputFileName = AliAnalysisManager::GetCommonFileName();
  Printf("AddTaskF0 - Set OutputFileName : \n %s\n", outputFileName.Data() );
  AliAnalysisDataContainer *output = mgr->CreateContainer(Form("RsnOut_%s",outNameSuffix.Data()),
							  TList::Class(),
							  AliAnalysisManager::kOutputContainer,
							  outputFileName);
  mgr->ConnectInput(task, 0, mgr->GetCommonInputContainer());
  mgr->ConnectOutput(task, 1, output);

  return task;
}
<|endoftext|>


<commit_before>#pragma once

#include <cmath>
#include <vector>

namespace bayesian_ab {

double logbeta(double a, double b) {
  return std::lgamma(a) + std::lgamma(b) - std::lgamma(a + b);
}

double prob_b_beats_a(int alpha_a, int beta_a, int alpha_b, int beta_b) {
  double total = 0.0;
  double logbeta_aa_ba = logbeta(alpha_a, beta_a);
  double beta_ba = beta_b + beta_a;

  for (auto i = 0; i < alpha_b; i++) {
    total += exp(logbeta(alpha_a + i, beta_ba) - log(beta_b + i) - logbeta(1 + i, beta_b) - logbeta_aa_ba);
  }

  return total;
}

double prob_c_beats_a_and_b(int alpha_a, int beta_a, int alpha_b, int beta_b, int alpha_c, int beta_c) {
  double total = 0.0;

  double logbeta_ac_bc = logbeta(alpha_c, beta_c);

  std::vector<double> log_bb_j;
  log_bb_j.reserve(alpha_b);
  std::vector<double> logbeta_j_bb;
  logbeta_j_bb.reserve(alpha_b);

  for (auto j = 0; j < alpha_b; j++) {
    log_bb_j.push_back(log(beta_b + j));
    logbeta_j_bb.push_back(logbeta(1 + j, beta_b));
  }

  double abc = beta_a + beta_b + beta_c;
  std::vector<double> logbeta_ac_i_j;
  logbeta_ac_i_j.reserve(alpha_a + alpha_b);

  for (auto i = 0; i < alpha_a + alpha_b; i++) {
    logbeta_ac_i_j.push_back(logbeta(alpha_c + i, abc));
  }

  for (auto i = 0; i < alpha_a; i++) {
    double log_ba_i = log(beta_a + i);
    double logbeta_i_ba = logbeta(1 + i, beta_a);

    for (auto j = 0; j < alpha_b; j++) {
      total += exp(logbeta_ac_i_j[i + j] -
        log_ba_i - log_bb_j[j] -
        logbeta_i_ba - logbeta_j_bb[j] -
        logbeta_ac_bc);
    }
  }

  return 1 - prob_b_beats_a(alpha_c, beta_c, alpha_a, beta_a) -
    prob_b_beats_a(alpha_c, beta_c, alpha_b, beta_b) + total;
}

double prob_d_beats_a_and_b_and_c(int alpha_a, int beta_a, int alpha_b, int beta_b, int alpha_c, int beta_c, int alpha_d, int beta_d) {
  double total = 0.0;

  double logbeta_ad_bd = logbeta(alpha_d, beta_d);

  std::vector<double> log_bb_j;
  log_bb_j.reserve(alpha_b);
  std::vector<double> logbeta_j_bb;
  logbeta_j_bb.reserve(alpha_b);

  for (auto j = 0; j < alpha_b; j++) {
    log_bb_j.push_back(log(beta_b + j));
    logbeta_j_bb.push_back(logbeta(1 + j, beta_b));
  }

  std::vector<double> log_bc_k;
  log_bc_k.reserve(alpha_c);
  std::vector<double> logbeta_k_bc;
  logbeta_k_bc.reserve(alpha_c);

  for (auto k = 0; k < alpha_c; k++) {
    log_bc_k.push_back(log(beta_c + k));
    logbeta_k_bc.push_back(logbeta(1 + k, beta_c));
  }

  double abcd = beta_a + beta_b + beta_c + beta_d;
  std::vector<double> logbeta_bd_i_j_k;
  logbeta_bd_i_j_k.reserve(alpha_a + alpha_b + alpha_c);

  for (auto i = 0; i < alpha_a + alpha_b + alpha_c; i++) {
    logbeta_bd_i_j_k.push_back(logbeta(alpha_d + i, abcd));
  }

  for (auto i = 0; i < alpha_a; i++) {
    double sum_i = -log(beta_a + i) - logbeta(1 + i, beta_a) - logbeta_ad_bd;

    for (auto j = 0; j < alpha_b; j++) {
      // TODO cache
      double sum_j = sum_i - log_bb_j[j] - logbeta_j_bb[j];

      for (auto k = 0; k < alpha_c; k++) {
        total += exp(sum_j + logbeta_bd_i_j_k[i + j + k] - log_bc_k[k] - logbeta_k_bc[k]);
      }
    }
  }

  return 1 - prob_b_beats_a(alpha_a, beta_a, alpha_d, beta_d) -
    prob_b_beats_a(alpha_b, beta_b, alpha_d, beta_d) -
    prob_b_beats_a(alpha_c, beta_c, alpha_d, beta_d) +
    prob_c_beats_a_and_b(alpha_a, beta_a, alpha_b, beta_b, alpha_d, beta_d) +
    prob_c_beats_a_and_b(alpha_a, beta_a, alpha_c, beta_c, alpha_d, beta_d) +
    prob_c_beats_a_and_b(alpha_b, beta_b, alpha_c, beta_c, alpha_d, beta_d) - total;
}

}
<commit_msg>Reduce calculations for prob_c_beats_a_and_b<commit_after>#pragma once

#include <cmath>
#include <vector>

namespace bayesian_ab {

double logbeta(double a, double b) {
  return std::lgamma(a) + std::lgamma(b) - std::lgamma(a + b);
}

double prob_b_beats_a(int alpha_a, int beta_a, int alpha_b, int beta_b) {
  double total = 0.0;
  double logbeta_aa_ba = logbeta(alpha_a, beta_a);
  double beta_ba = beta_b + beta_a;

  for (auto i = 0; i < alpha_b; i++) {
    total += exp(logbeta(alpha_a + i, beta_ba) - log(beta_b + i) - logbeta(1 + i, beta_b) - logbeta_aa_ba);
  }

  return total;
}

double prob_c_beats_a_and_b(int alpha_a, int beta_a, int alpha_b, int beta_b, int alpha_c, int beta_c) {
  double total = 0.0;

  double logbeta_ac_bc = logbeta(alpha_c, beta_c);

  std::vector<double> log_bb_j;
  log_bb_j.reserve(alpha_b);
  std::vector<double> logbeta_j_bb;
  logbeta_j_bb.reserve(alpha_b);

  for (auto j = 0; j < alpha_b; j++) {
    log_bb_j.push_back(log(beta_b + j));
    logbeta_j_bb.push_back(logbeta(1 + j, beta_b));
  }

  double abc = beta_a + beta_b + beta_c;
  std::vector<double> logbeta_ac_i_j;
  logbeta_ac_i_j.reserve(alpha_a + alpha_b);

  for (auto i = 0; i < alpha_a + alpha_b; i++) {
    logbeta_ac_i_j.push_back(logbeta(alpha_c + i, abc));
  }

  for (auto i = 0; i < alpha_a; i++) {
    double sum_i = -log(beta_a + i) - logbeta(1 + i, beta_a) - logbeta_ac_bc;

    for (auto j = 0; j < alpha_b; j++) {
      total += exp(sum_i + logbeta_ac_i_j[i + j] - log_bb_j[j] - logbeta_j_bb[j]);
    }
  }

  return 1 - prob_b_beats_a(alpha_c, beta_c, alpha_a, beta_a) -
    prob_b_beats_a(alpha_c, beta_c, alpha_b, beta_b) + total;
}

double prob_d_beats_a_and_b_and_c(int alpha_a, int beta_a, int alpha_b, int beta_b, int alpha_c, int beta_c, int alpha_d, int beta_d) {
  double total = 0.0;

  double logbeta_ad_bd = logbeta(alpha_d, beta_d);

  std::vector<double> log_bb_j;
  log_bb_j.reserve(alpha_b);
  std::vector<double> logbeta_j_bb;
  logbeta_j_bb.reserve(alpha_b);

  for (auto j = 0; j < alpha_b; j++) {
    log_bb_j.push_back(log(beta_b + j));
    logbeta_j_bb.push_back(logbeta(1 + j, beta_b));
  }

  std::vector<double> log_bc_k;
  log_bc_k.reserve(alpha_c);
  std::vector<double> logbeta_k_bc;
  logbeta_k_bc.reserve(alpha_c);

  for (auto k = 0; k < alpha_c; k++) {
    log_bc_k.push_back(log(beta_c + k));
    logbeta_k_bc.push_back(logbeta(1 + k, beta_c));
  }

  double abcd = beta_a + beta_b + beta_c + beta_d;
  std::vector<double> logbeta_bd_i_j_k;
  logbeta_bd_i_j_k.reserve(alpha_a + alpha_b + alpha_c);

  for (auto i = 0; i < alpha_a + alpha_b + alpha_c; i++) {
    logbeta_bd_i_j_k.push_back(logbeta(alpha_d + i, abcd));
  }

  for (auto i = 0; i < alpha_a; i++) {
    double sum_i = -log(beta_a + i) - logbeta(1 + i, beta_a) - logbeta_ad_bd;

    for (auto j = 0; j < alpha_b; j++) {
      // TODO cache
      double sum_j = sum_i - log_bb_j[j] - logbeta_j_bb[j];

      for (auto k = 0; k < alpha_c; k++) {
        total += exp(sum_j + logbeta_bd_i_j_k[i + j + k] - log_bc_k[k] - logbeta_k_bc[k]);
      }
    }
  }

  return 1 - prob_b_beats_a(alpha_a, beta_a, alpha_d, beta_d) -
    prob_b_beats_a(alpha_b, beta_b, alpha_d, beta_d) -
    prob_b_beats_a(alpha_c, beta_c, alpha_d, beta_d) +
    prob_c_beats_a_and_b(alpha_a, beta_a, alpha_b, beta_b, alpha_d, beta_d) +
    prob_c_beats_a_and_b(alpha_a, beta_a, alpha_c, beta_c, alpha_d, beta_d) +
    prob_c_beats_a_and_b(alpha_b, beta_b, alpha_c, beta_c, alpha_d, beta_d) - total;
}

}
<|endoftext|>


<commit_before>/****************************************************************/
/* MOOSE - Multiphysics Object Oriented Simulation Environment  */
/*                                                              */
/*          All contents are licensed under LGPL V2.1           */
/*             See LICENSE for full restrictions                */
/****************************************************************/
#include "LinearIsoElasticPFDamage.h"
#include "libmesh/utility.h"

template<>
InputParameters validParams<LinearIsoElasticPFDamage>()
{
  InputParameters params = validParams<ComputeStressBase>();
  params.addClassDescription("Phase-field fracture model energy contribution to damage growth-isotropic elasticity and undamaged stress under compressive strain");
  params.addRequiredCoupledVar("c","Order parameter for damage");
  params.addParam<Real>("kdamage",1e-6,"Stiffness of damaged matrix");

  return params;
}

LinearIsoElasticPFDamage::LinearIsoElasticPFDamage(const InputParameters & parameters) :
    ComputeStressBase(parameters),
    _c(coupledValue("c")),
    _kdamage(getParam<Real>("kdamage")),
    _G0_pos(declareProperty<Real>("G0_pos")),
    _dstress_dc(declarePropertyDerivative<RankTwoTensor>(_base_name + "stress", getVar("c", 0)->name())),
    _dG0_pos_dstrain(declareProperty<RankTwoTensor>("dG0_pos_dstrain")),
    _etens(LIBMESH_DIM),
    _epos(LIBMESH_DIM),
    _eigval(LIBMESH_DIM)
{
}

void LinearIsoElasticPFDamage::computeQpStress()
{
  updateVar();
  updateJacobian();
}

void
LinearIsoElasticPFDamage::updateVar()
{
  RankTwoTensor stress0pos, stress0neg, stress0;
  //Isotropic elasticity is assumed
  Real lambda = _elasticity_tensor[_qp](0,0,1,1);
  Real mu = _elasticity_tensor[_qp](0,1,0,1);
  Real c = _c[_qp];
  Real xfac = Utility::pow<2>(1.0-c) + _kdamage;

  _mechanical_strain[_qp].symmetricEigenvaluesEigenvectors(_eigval, _eigvec);

  //Tensors of outerproduct of eigen vectors
  for (unsigned int i = 0; i < LIBMESH_DIM; ++i)
    for (unsigned int j = 0; j < LIBMESH_DIM; ++j)
      for (unsigned int k = 0; k < LIBMESH_DIM; ++k)
        _etens[i](j,k) = _eigvec(j,i) * _eigvec(k,i);

  Real etr = 0.0;
  for (unsigned int i = 0; i < LIBMESH_DIM; ++i)
    etr += _eigval[i];

  Real etrpos = (std::abs(etr) + etr) / 2.0;
  Real etrneg = (std::abs(etr) - etr) / 2.0;

  for (unsigned int i = 0; i < LIBMESH_DIM; ++i)
  {
    stress0pos += _etens[i] * (lambda * etrpos + 2.0 * mu * (std::abs(_eigval[i]) + _eigval[i]) / 2.0);
    stress0neg += _etens[i] * (lambda * etrneg + 2.0 * mu * (std::abs(_eigval[i]) - _eigval[i]) / 2.0);
  }

  //Damage associated with positive component of stress
  _stress[_qp] = stress0pos * xfac - stress0neg;

  for (unsigned int i = 0; i < LIBMESH_DIM; ++i)
    _epos[i] = (std::abs(_eigval[i]) + _eigval[i]) / 2.0;

  Real val = 0.0;
  for (unsigned int i = 0; i < LIBMESH_DIM; ++i)
    val += Utility::pow<2>(_epos[i]);
  val *= mu;

  //Energy with positive principal strains
  _G0_pos[_qp] = lambda * Utility::pow<2>(etrpos) / 2.0 + val;
  //Used in PFFracBulkRate Jacobian
  _dG0_pos_dstrain[_qp] = stress0pos;
  //Used in StressDivergencePFFracTensors Jacobian
  _dstress_dc[_qp] = -stress0pos * (2.0 * (1.0 - c));
}

void
LinearIsoElasticPFDamage::updateJacobian()
{
  _Jacobian_mult[_qp] = _elasticity_tensor[_qp];
}
<commit_msg>Overcome the issue when c goes above 1.0 (#8197)<commit_after>/****************************************************************/
/* MOOSE - Multiphysics Object Oriented Simulation Environment  */
/*                                                              */
/*          All contents are licensed under LGPL V2.1           */
/*             See LICENSE for full restrictions                */
/****************************************************************/
#include "LinearIsoElasticPFDamage.h"
#include "libmesh/utility.h"

template<>
InputParameters validParams<LinearIsoElasticPFDamage>()
{
  InputParameters params = validParams<ComputeStressBase>();
  params.addClassDescription("Phase-field fracture model energy contribution to damage growth-isotropic elasticity and undamaged stress under compressive strain");
  params.addRequiredCoupledVar("c","Order parameter for damage");
  params.addParam<Real>("kdamage",1e-6,"Stiffness of damaged matrix");

  return params;
}

LinearIsoElasticPFDamage::LinearIsoElasticPFDamage(const InputParameters & parameters) :
    ComputeStressBase(parameters),
    _c(coupledValue("c")),
    _kdamage(getParam<Real>("kdamage")),
    _G0_pos(declareProperty<Real>("G0_pos")),
    _dstress_dc(declarePropertyDerivative<RankTwoTensor>(_base_name + "stress", getVar("c", 0)->name())),
    _dG0_pos_dstrain(declareProperty<RankTwoTensor>("dG0_pos_dstrain")),
    _etens(LIBMESH_DIM),
    _epos(LIBMESH_DIM),
    _eigval(LIBMESH_DIM)
{
}

void LinearIsoElasticPFDamage::computeQpStress()
{
  updateVar();
  updateJacobian();
}

void
LinearIsoElasticPFDamage::updateVar()
{
  RankTwoTensor stress0pos, stress0neg, stress0;
  //Isotropic elasticity is assumed
  Real lambda = _elasticity_tensor[_qp](0,0,1,1);
  Real mu = _elasticity_tensor[_qp](0,1,0,1);
  Real c = _c[_qp];
  Real xfac = _kdamage;
  if (c < 1.0)
    xfac = Utility::pow<2>(1 - c) + _kdamage;

  _mechanical_strain[_qp].symmetricEigenvaluesEigenvectors(_eigval, _eigvec);

  //Tensors of outerproduct of eigen vectors
  for (unsigned int i = 0; i < LIBMESH_DIM; ++i)
    for (unsigned int j = 0; j < LIBMESH_DIM; ++j)
      for (unsigned int k = 0; k < LIBMESH_DIM; ++k)
        _etens[i](j,k) = _eigvec(j,i) * _eigvec(k,i);

  Real etr = 0.0;
  for (unsigned int i = 0; i < LIBMESH_DIM; ++i)
    etr += _eigval[i];

  Real etrpos = (std::abs(etr) + etr) / 2.0;
  Real etrneg = (std::abs(etr) - etr) / 2.0;

  for (unsigned int i = 0; i < LIBMESH_DIM; ++i)
  {
    stress0pos += _etens[i] * (lambda * etrpos + 2.0 * mu * (std::abs(_eigval[i]) + _eigval[i]) / 2.0);
    stress0neg += _etens[i] * (lambda * etrneg + 2.0 * mu * (std::abs(_eigval[i]) - _eigval[i]) / 2.0);
  }

  //Damage associated with positive component of stress
  _stress[_qp] = stress0pos * xfac - stress0neg;

  for (unsigned int i = 0; i < LIBMESH_DIM; ++i)
    _epos[i] = (std::abs(_eigval[i]) + _eigval[i]) / 2.0;

  Real val = 0.0;
  for (unsigned int i = 0; i < LIBMESH_DIM; ++i)
    val += Utility::pow<2>(_epos[i]);
  val *= mu;

  //Energy with positive principal strains
  _G0_pos[_qp] = lambda * Utility::pow<2>(etrpos) / 2.0 + val;
  //Used in PFFracBulkRate Jacobian
  _dG0_pos_dstrain[_qp] = stress0pos;
  //Used in StressDivergencePFFracTensors Jacobian
  if (c < 1.0)
    _dstress_dc[_qp] = -stress0pos * (2.0 * (1.0 - c));
  else
    _dstress_dc[_qp] = 0.0;
}

void
LinearIsoElasticPFDamage::updateJacobian()
{
  _Jacobian_mult[_qp] = _elasticity_tensor[_qp];
}
<|endoftext|>


<commit_before>//=============================================================================================================
/**
* @file     mnelaunchcontrol.cpp
* @author   Christoph Dinh <chdinh@nmr.mgh.harvard.edu>;
*           Matti Hamalainen <msh@nmr.mgh.harvard.edu>
* @version  1.0
* @date     November, 2016
*
* @section  LICENSE
*
* Copyright (C) 2016, Christoph Dinh. All rights reserved.
*
* Redistribution and use in source and binary forms, with or without modification, are permitted provided that
* the following conditions are met:
*     * Redistributions of source code must retain the above copyright notice, this list of conditions and the
*       following disclaimer.
*     * Redistributions in binary form must reproduce the above copyright notice, this list of conditions and
*       the following disclaimer in the documentation and/or other materials provided with the distribution.
*     * Neither the name of MNE-CPP authors nor the names of its contributors may be used
*       to endorse or promote products derived from this software without specific prior written permission.
*
* THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND ANY EXPRESS OR IMPLIED
* WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A
* PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT,
* INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,
* PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION)
* HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING
* NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
* POSSIBILITY OF SUCH DAMAGE.
*
*
* @brief    MNELaunchControl class implementation
*
*/

//*************************************************************************************************************
//=============================================================================================================
// INCLUDES
//=============================================================================================================

#include "mnelaunchcontrol.h"

//*************************************************************************************************************
//=============================================================================================================
// Qt INCLUDES
//=============================================================================================================

#include <QDebug>
#include <QCoreApplication>
#include <QFile>

//*************************************************************************************************************
//=============================================================================================================
// DEFINE MEMBER METHODS
//=============================================================================================================

MNELaunchControl::MNELaunchControl(QObject *parent)
: QObject(parent)
{
    m_requiredSampleFiles   <<  "/MNE-sample-data/MEG/sample/sample_audvis_raw.fif"
                            <<  "/MNE-sample-data/MEG/sample/sample_audvis-meg-eeg-oct-6-fwd.fif"
                            <<  "/MNE-sample-data/subjects/sample/label/lh.aparc.a2009s.annot"
                            <<  "/MNE-sample-data/subjects/sample/label/rh.aparc.a2009s.annot";
}

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

void MNELaunchControl::invokeScan()
{
    invokeApplication("mne_scan",QStringList());
}

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

void MNELaunchControl::invokeBrowse()
{
    invokeApplication("mne_browse",QStringList());
}

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

void MNELaunchControl::invokeAnalyze()
{
    invokeApplication("mne_analyze",QStringList());
}

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

void MNELaunchControl::invokeApplication(const QString &application, const QStringList &arguments)
{
    QFile file(QCoreApplication::applicationDirPath()+ "/" + application);
    #if defined(Q_OS_WIN)
        file.setFileName(file.fileName() + ".exe");
    #elif defined(Q_OS_MACOS)
        //On MacOS we use .app bundle structures. Executable path is: .app/Contents/MacOS/
        file.setFileName("../../../" + application + ".app/Contents/MacOS/" + application);
    #endif

    qDebug()<<"file.fileName()"<<file.fileName();
    qDebug()<<"QCoreApplication::applicationDirPath()"<<QCoreApplication::applicationDirPath();

    if(file.exists()) {
        try {
            QPointer<QProcess> process( new QProcess );

            #if defined(Q_OS_MACOS)
                process->start("open", {file.fileName()});
            #else
                process->start(file.fileName(), arguments);
            #endif

            m_ListProcesses.append(process);
        } catch (int e) {
            qWarning() << "Not able to start" << file.fileName() << ". Error:" << e;
        }
    }
    else {
        qWarning() << "Application" << file.fileName() << "does not exists.";
    }
}

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

bool MNELaunchControl::getSampleDataAvailable() const
{
    bool sampleFilesExist = true;

    for(QString fileName : m_requiredSampleFiles)
    {
        QFile file (QCoreApplication::applicationDirPath() + fileName);
        if ( !file.exists() ) {
            qWarning() << "Missing sample file" << file.fileName();
            sampleFilesExist = false;
        }
    }
    return sampleFilesExist;
}
<commit_msg>[LIB-216] Added absolute path to mne_launch<commit_after>//=============================================================================================================
/**
* @file     mnelaunchcontrol.cpp
* @author   Christoph Dinh <chdinh@nmr.mgh.harvard.edu>;
*           Matti Hamalainen <msh@nmr.mgh.harvard.edu>
* @version  1.0
* @date     November, 2016
*
* @section  LICENSE
*
* Copyright (C) 2016, Christoph Dinh. All rights reserved.
*
* Redistribution and use in source and binary forms, with or without modification, are permitted provided that
* the following conditions are met:
*     * Redistributions of source code must retain the above copyright notice, this list of conditions and the
*       following disclaimer.
*     * Redistributions in binary form must reproduce the above copyright notice, this list of conditions and
*       the following disclaimer in the documentation and/or other materials provided with the distribution.
*     * Neither the name of MNE-CPP authors nor the names of its contributors may be used
*       to endorse or promote products derived from this software without specific prior written permission.
*
* THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND ANY EXPRESS OR IMPLIED
* WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A
* PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT,
* INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,
* PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION)
* HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING
* NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
* POSSIBILITY OF SUCH DAMAGE.
*
*
* @brief    MNELaunchControl class implementation
*
*/

//*************************************************************************************************************
//=============================================================================================================
// INCLUDES
//=============================================================================================================

#include "mnelaunchcontrol.h"

//*************************************************************************************************************
//=============================================================================================================
// Qt INCLUDES
//=============================================================================================================

#include <QDebug>
#include <QCoreApplication>
#include <QFile>

//*************************************************************************************************************
//=============================================================================================================
// DEFINE MEMBER METHODS
//=============================================================================================================

MNELaunchControl::MNELaunchControl(QObject *parent)
: QObject(parent)
{
    m_requiredSampleFiles   <<  "/MNE-sample-data/MEG/sample/sample_audvis_raw.fif"
                            <<  "/MNE-sample-data/MEG/sample/sample_audvis-meg-eeg-oct-6-fwd.fif"
                            <<  "/MNE-sample-data/subjects/sample/label/lh.aparc.a2009s.annot"
                            <<  "/MNE-sample-data/subjects/sample/label/rh.aparc.a2009s.annot";
}

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

void MNELaunchControl::invokeScan()
{
    invokeApplication("mne_scan",QStringList());
}

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

void MNELaunchControl::invokeBrowse()
{
    invokeApplication("mne_browse",QStringList());
}

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

void MNELaunchControl::invokeAnalyze()
{
    invokeApplication("mne_analyze",QStringList());
}

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

void MNELaunchControl::invokeApplication(const QString &application, const QStringList &arguments)
{
    QFile file(QCoreApplication::applicationDirPath()+ "/" + application);
    #if defined(Q_OS_WIN)
        file.setFileName(file.fileName() + ".exe");
    #elif defined(Q_OS_MACOS)
        //On MacOS we use .app bundle structures. MNE Launch executable path is: mne_launch.app/Contents/MacOS/
        file.setFileName("QCoreApplication::applicationDirPath() + "/../../../" + application + ".app");
    #endif

    if(file.exists()) {
        try {
            QPointer<QProcess> process( new QProcess );

            #if defined(Q_OS_MACOS)
                process->start("open", {file.fileName()});
            #else
                process->start(file.fileName(), arguments);
            #endif

            m_ListProcesses.append(process);
        } catch (int e) {
            qWarning() << "Not able to start" << file.fileName() << ". Error:" << e;
        }
    }
    else {
        qWarning() << "Application" << file.fileName() << "does not exists.";
    }
}

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

bool MNELaunchControl::getSampleDataAvailable() const
{
    bool sampleFilesExist = true;

    for(QString fileName : m_requiredSampleFiles)
    {
        QFile file (QCoreApplication::applicationDirPath() + fileName);
        if ( !file.exists() ) {
            qWarning() << "Missing sample file" << file.fileName();
            sampleFilesExist = false;
        }
    }
    return sampleFilesExist;
}
<|endoftext|>


<commit_before>

// AliTRDonlineTrackletFilter implements an analysis task which reads
// the tracklets from the TRD.Tracklets.root file and makes them
// available to other analysis tasks. This allows to access the full
// tracklet information without repeating the techncicalities in each
// analysis task. The tracklets are made available in an exchange
// container on the output slot 1.

#include "TFile.h"
#include "TTree.h"
#include "TChain.h"
#include "TH1F.h"
#include "TH2F.h"
#include "TProfile.h"
#include "TCanvas.h"

#include "AliAnalysisManager.h"
#include "AliESDEvent.h"
#include "AliAODEvent.h"
#include "AliMCEvent.h"
#include "AliESDInputHandler.h"
#include "AliAODHandler.h"
#include "AliMCEventHandler.h"
#include "AliLog.h"
#include "AliESDTrdTrack.h"

#include "AliTRDtrackletMCM.h"
#include "AliTRDtrackletWord.h"
#include "AliVParticle.h"
#include "AliMCParticle.h"

#include "AliTRDonlineTrackletFilter.h"

ClassImp(AliTRDonlineTrackletFilter)

AliTRDonlineTrackletFilter::AliTRDonlineTrackletFilter(const char *name) :
  AliAnalysisTask(name, ""),
  fESD(0x0),
  fInputHandler(0x0),
  fInputEvent(0x0),
  fOutputAOD(0x0),
  fMCEvent(0x0),
  fTrackletsRaw(new TClonesArray("AliTRDtrackletWord")),
  fTrackletsSim(new TClonesArray("AliTRDtrackletMCM")),
  fTrackletTree(0x0),
  fGeo(new AliTRDgeometry),
  fNevent(0),
  fPath(""),
  fTrackletFile(0x0),
  fNEventsPerFile(0),
  fEvent(0),
  fFileNumber(0),
  fTrackletTreeSim(0x0),
  fTrackletTreeRaw(0x0)
{
  // ctor

  DefineInput(0, TChain::Class());

  DefineOutput(0, TTree::Class());
  DefineOutput(1, TTree::Class());
}

AliTRDonlineTrackletFilter::~AliTRDonlineTrackletFilter()
{
  // dtor

  delete fTrackletsRaw;
  delete fTrackletsSim;
  delete fGeo;
}

void AliTRDonlineTrackletFilter::ConnectInputData(const Option_t */* option */)
{
  // connect all the input handlers to access the data

  fInputHandler = (AliInputEventHandler*) AliAnalysisManager::GetAnalysisManager()->GetInputEventHandler();
  if (fInputHandler)
    fInputEvent = fInputHandler->GetEvent();

  AliMCEventHandler *mcH = (AliMCEventHandler*) AliAnalysisManager::GetAnalysisManager()->GetMCtruthEventHandler();
  if (mcH)
    fMCEvent = mcH->MCEvent();
}

void AliTRDonlineTrackletFilter::CreateOutputObjects()
{
  // create the output objects used to make the tracklets available to
  // other tasks

  OpenFile(1);

  fTrackletTree = new TTree("TRDtrackletfilter", "on-line tracklets");
  fTrackletTree->Branch("tracklets_sim", fTrackletsSim);
  fTrackletTree->Branch("tracklets_raw", fTrackletsRaw);

  PostData(1, fTrackletTree);
}

Bool_t AliTRDonlineTrackletFilter::Notify()
{
  // we use notify to figure out which TRD.Tracklets.root file we have
  // to use for accessing the tracklets

  TString filename(AliAnalysisManager::GetAnalysisManager()->GetTree()->GetCurrentFile()->GetName());

  AliDebug(1, Form("Now reading from %s", filename.Data()));

  if (filename.Contains("AliAOD.root"))
    filename.ReplaceAll("AliAOD.root", "");
  else if (filename.Contains("NewAliESDs.root"))
    filename.ReplaceAll("NewAliESDs.root", "");
  else if (filename.Contains("AliESDs.root"))
    filename.ReplaceAll("AliESDs.root", "");
  else if (filename.Contains("galice.root"))
    filename.ReplaceAll("galice.root", "");
  else if (filename.BeginsWith("root:"))
    filename.Append("?ZIP=");

  fPath = filename;

  fTrackletFile = TFile::Open(Form("%sTRD.Tracklets.root", fPath.Data()));

  if (!fTrackletFile) {
    AliError(Form("No tracklet file in %s", fPath.Data()));
    return kFALSE;
  }

  fNEventsPerFile = fTrackletFile->GetNkeys() - fTrackletFile->GetNProcessIDs(); //???

  fEvent = -1;
  fFileNumber = 0;

  return kTRUE;
}

void AliTRDonlineTrackletFilter::Exec(const Option_t * /* option */)
{
  // execute this for each event

  if (!LoadEvent())
    return;

  // ----- simulated tracklets -----
  if (fTrackletTreeSim) {
    AliTRDtrackletMCM *trkl = 0x0;

    TBranch *br = fTrackletTreeSim->GetBranch("mcmtrklbranch");
    br->SetAddress(&trkl);

    for (Int_t iTracklet = 0; iTracklet < br->GetEntries(); iTracklet++) {
      br->GetEntry(iTracklet);
      new ((*fTrackletsSim)[fTrackletsSim->GetEntriesFast()]) AliTRDtrackletMCM(*trkl);
    }
  }

  // ----- raw tracklets -----
  if (fTrackletTreeRaw) {
    Int_t hc;
    TClonesArray *trklArray = 0x0;
    fTrackletTreeRaw->SetBranchAddress("hc", &hc);
    fTrackletTreeRaw->SetBranchAddress("trkl", &trklArray);
    for (Int_t iHCidx = 0; iHCidx < fTrackletTreeRaw->GetEntries(); iHCidx++) {
      fTrackletTreeRaw->GetEntry(iHCidx);
      if (trklArray) {
	for (Int_t iTracklet = 0; iTracklet < trklArray->GetEntries(); iTracklet++) {
	  AliTRDtrackletWord *trklWord = (AliTRDtrackletWord*) ((*trklArray)[iTracklet]);
	  trklWord->SetDetector(hc/2);
	  new ((*fTrackletsRaw)[fTrackletsRaw->GetEntriesFast()]) AliTRDtrackletWord(*trklWord);
	}
      }
    }
  }

  AliDebug(1, Form("%i tracklets", fTrackletsSim->GetEntriesFast()));
  fTrackletTree->SetBranchAddress("tracklets_sim", &fTrackletsSim);
  fTrackletTree->SetBranchAddress("tracklets_raw", &fTrackletsRaw);
  fTrackletTree->Fill();
  PostData(1, fTrackletTree);
}

void AliTRDonlineTrackletFilter::LocalInit()
{

}

void AliTRDonlineTrackletFilter::Terminate(const Option_t * /* option */)
{

}

Bool_t AliTRDonlineTrackletFilter::LoadEvent()
{
  // load tracklets for the current event

  // ----- cleaning -----
  fTrackletsSim->Delete();
  fTrackletsRaw->Delete();

  // ----- initialization -----
  if (!fInputEvent) {
    AliError("No event found!");
    return kFALSE;
  }
  fESD = dynamic_cast<AliESDEvent*> (fInputEvent);

  fEvent++;
  Int_t inew = fNEventsPerFile > 0 ? fEvent / fNEventsPerFile : 0;
  if ( inew != fFileNumber) {
    fFileNumber++;

    delete fTrackletFile;
    fTrackletFile = TFile::Open(Form("%sTRD.Tracklets%d.root", fPath.Data(), fFileNumber));

    if (!fTrackletFile) {
      AliError("No tracklet file");
      return kFALSE;
    }
  }

  if (!fTrackletFile) {
    AliError("no tracklet file");
    return kFALSE;
  }

  // tracklets from simulation
  char treename[30];
  snprintf(treename, 30, "Event%d/tracklets", fEvent);

  fTrackletTreeSim = (TTree*) fTrackletFile->Get(treename);

  // tracklets from raw
  char treenameRaw[30];
  snprintf(treenameRaw, 30, "Event%d/tracklets-raw", fEvent);
  fTrackletTreeRaw = (TTree*) fTrackletFile->Get(treenameRaw);

  return kTRUE;
}

<commit_msg>Revert "- test commit"<commit_after>// AliTRDonlineTrackletFilter implements an analysis task which reads
// the tracklets from the TRD.Tracklets.root file and makes them
// available to other analysis tasks. This allows to access the full
// tracklet information without repeating the techncicalities in each
// analysis task. The tracklets are made available in an exchange
// container on the output slot 1.

#include "TFile.h"
#include "TTree.h"
#include "TChain.h"
#include "TH1F.h"
#include "TH2F.h"
#include "TProfile.h"
#include "TCanvas.h"

#include "AliAnalysisManager.h"
#include "AliESDEvent.h"
#include "AliAODEvent.h"
#include "AliMCEvent.h"
#include "AliESDInputHandler.h"
#include "AliAODHandler.h"
#include "AliMCEventHandler.h"
#include "AliLog.h"
#include "AliESDTrdTrack.h"

#include "AliTRDtrackletMCM.h"
#include "AliTRDtrackletWord.h"
#include "AliVParticle.h"
#include "AliMCParticle.h"

#include "AliTRDonlineTrackletFilter.h"

ClassImp(AliTRDonlineTrackletFilter)

AliTRDonlineTrackletFilter::AliTRDonlineTrackletFilter(const char *name) :
  AliAnalysisTask(name, ""),
  fESD(0x0),
  fInputHandler(0x0),
  fInputEvent(0x0),
  fOutputAOD(0x0),
  fMCEvent(0x0),
  fTrackletsRaw(new TClonesArray("AliTRDtrackletWord")),
  fTrackletsSim(new TClonesArray("AliTRDtrackletMCM")),
  fTrackletTree(0x0),
  fGeo(new AliTRDgeometry),
  fNevent(0),
  fPath(""),
  fTrackletFile(0x0),
  fNEventsPerFile(0),
  fEvent(0),
  fFileNumber(0),
  fTrackletTreeSim(0x0),
  fTrackletTreeRaw(0x0)
{
  // ctor

  DefineInput(0, TChain::Class());

  DefineOutput(0, TTree::Class());
  DefineOutput(1, TTree::Class());
}

AliTRDonlineTrackletFilter::~AliTRDonlineTrackletFilter()
{
  // dtor

  delete fTrackletsRaw;
  delete fTrackletsSim;
  delete fGeo;
}

void AliTRDonlineTrackletFilter::ConnectInputData(const Option_t */* option */)
{
  // connect all the input handlers to access the data

  fInputHandler = (AliInputEventHandler*) AliAnalysisManager::GetAnalysisManager()->GetInputEventHandler();
  if (fInputHandler)
    fInputEvent = fInputHandler->GetEvent();

  AliMCEventHandler *mcH = (AliMCEventHandler*) AliAnalysisManager::GetAnalysisManager()->GetMCtruthEventHandler();
  if (mcH)
    fMCEvent = mcH->MCEvent();
}

void AliTRDonlineTrackletFilter::CreateOutputObjects()
{
  // create the output objects used to make the tracklets available to
  // other tasks

  OpenFile(1);

  fTrackletTree = new TTree("TRDtrackletfilter", "on-line tracklets");
  fTrackletTree->Branch("tracklets_sim", fTrackletsSim);
  fTrackletTree->Branch("tracklets_raw", fTrackletsRaw);

  PostData(1, fTrackletTree);
}

Bool_t AliTRDonlineTrackletFilter::Notify()
{
  // we use notify to figure out which TRD.Tracklets.root file we have
  // to use for accessing the tracklets

  TString filename(AliAnalysisManager::GetAnalysisManager()->GetTree()->GetCurrentFile()->GetName());

  AliDebug(1, Form("Now reading from %s", filename.Data()));

  if (filename.Contains("AliAOD.root"))
    filename.ReplaceAll("AliAOD.root", "");
  else if (filename.Contains("NewAliESDs.root"))
    filename.ReplaceAll("NewAliESDs.root", "");
  else if (filename.Contains("AliESDs.root"))
    filename.ReplaceAll("AliESDs.root", "");
  else if (filename.Contains("galice.root"))
    filename.ReplaceAll("galice.root", "");
  else if (filename.BeginsWith("root:"))
    filename.Append("?ZIP=");

  fPath = filename;

  fTrackletFile = TFile::Open(Form("%sTRD.Tracklets.root", fPath.Data()));

  if (!fTrackletFile) {
    AliError(Form("No tracklet file in %s", fPath.Data()));
    return kFALSE;
  }

  fNEventsPerFile = fTrackletFile->GetNkeys() - fTrackletFile->GetNProcessIDs(); //???

  fEvent = -1;
  fFileNumber = 0;

  return kTRUE;
}

void AliTRDonlineTrackletFilter::Exec(const Option_t * /* option */)
{
  // execute this for each event

  if (!LoadEvent())
    return;

  // ----- simulated tracklets -----
  if (fTrackletTreeSim) {
    AliTRDtrackletMCM *trkl = 0x0;

    TBranch *br = fTrackletTreeSim->GetBranch("mcmtrklbranch");
    br->SetAddress(&trkl);

    for (Int_t iTracklet = 0; iTracklet < br->GetEntries(); iTracklet++) {
      br->GetEntry(iTracklet);
      new ((*fTrackletsSim)[fTrackletsSim->GetEntriesFast()]) AliTRDtrackletMCM(*trkl);
    }
  }

  // ----- raw tracklets -----
  if (fTrackletTreeRaw) {
    Int_t hc;
    TClonesArray *trklArray = 0x0;
    fTrackletTreeRaw->SetBranchAddress("hc", &hc);
    fTrackletTreeRaw->SetBranchAddress("trkl", &trklArray);
    for (Int_t iHCidx = 0; iHCidx < fTrackletTreeRaw->GetEntries(); iHCidx++) {
      fTrackletTreeRaw->GetEntry(iHCidx);
      if (trklArray) {
	for (Int_t iTracklet = 0; iTracklet < trklArray->GetEntries(); iTracklet++) {
	  AliTRDtrackletWord *trklWord = (AliTRDtrackletWord*) ((*trklArray)[iTracklet]);
	  trklWord->SetDetector(hc/2);
	  new ((*fTrackletsRaw)[fTrackletsRaw->GetEntriesFast()]) AliTRDtrackletWord(*trklWord);
	}
      }
    }
  }

  AliDebug(1, Form("%i tracklets", fTrackletsSim->GetEntriesFast()));
  fTrackletTree->SetBranchAddress("tracklets_sim", &fTrackletsSim);
  fTrackletTree->SetBranchAddress("tracklets_raw", &fTrackletsRaw);
  fTrackletTree->Fill();
  PostData(1, fTrackletTree);
}

void AliTRDonlineTrackletFilter::LocalInit()
{

}

void AliTRDonlineTrackletFilter::Terminate(const Option_t * /* option */)
{

}

Bool_t AliTRDonlineTrackletFilter::LoadEvent()
{
  // load tracklets for the current event

  // ----- cleaning -----
  fTrackletsSim->Delete();
  fTrackletsRaw->Delete();

  // ----- initialization -----
  if (!fInputEvent) {
    AliError("No event found!");
    return kFALSE;
  }
  fESD = dynamic_cast<AliESDEvent*> (fInputEvent);

  fEvent++;
  Int_t inew = fNEventsPerFile > 0 ? fEvent / fNEventsPerFile : 0;
  if ( inew != fFileNumber) {
    fFileNumber++;

    delete fTrackletFile;
    fTrackletFile = TFile::Open(Form("%sTRD.Tracklets%d.root", fPath.Data(), fFileNumber));

    if (!fTrackletFile) {
      AliError("No tracklet file");
      return kFALSE;
    }
  }

  if (!fTrackletFile) {
    AliError("no tracklet file");
    return kFALSE;
  }

  // tracklets from simulation
  char treename[30];
  snprintf(treename, 30, "Event%d/tracklets", fEvent);

  fTrackletTreeSim = (TTree*) fTrackletFile->Get(treename);

  // tracklets from raw
  char treenameRaw[30];
  snprintf(treenameRaw, 30, "Event%d/tracklets-raw", fEvent);
  fTrackletTreeRaw = (TTree*) fTrackletFile->Get(treenameRaw);

  return kTRUE;
}

<|endoftext|>


<commit_before>//(c) 2016 by Authors
//This file is a part of ABruijn program.
//Released under the BSD license (see LICENSE file)

#include <limits>
#include <algorithm>
#include <iomanip>
#include <stack>

#include "../common/config.h"
#include "../common/logger.h"
#include "../common/parallel.h"
#include "extender.h"

Extender::ExtensionInfo Extender::extendContig(FastaRecord::Id startRead)
{

	std::unordered_set<FastaRecord::Id> currentReads;
	currentReads.insert(startRead);
	currentReads.insert(startRead.rc());

	bool rightExtension = true;
	FastaRecord::Id currentRead = startRead;
	std::vector<int> numExtensions;
	std::vector<int> overlapSizes;
	ExtensionInfo exInfo;
	exInfo.reads.push_back(startRead);
	exInfo.assembledLength = _readsContainer.seqLen(startRead);

	auto leftExtendsStart = [startRead, this](const FastaRecord::Id readId)
	{
		for (auto& ovlp : _ovlpContainer.lazySeqOverlaps(readId))
		{
			if (ovlp.extId == startRead &&
				ovlp.leftShift > (int)Config::get("maximum_jump"))
			{
				return true;
			}
		}
		return false;
	};

	while(true)
	{
		auto& overlaps = _ovlpContainer.lazySeqOverlaps(currentRead);
		std::vector<OverlapRange> extensions;
		//int innerOverlaps = 0;
		for (auto& ovlp : overlaps)
		{
			//if (_innerReads.contains(ovlp.extId)) ++innerOverlaps;
			if (this->extendsRight(ovlp)) extensions.push_back(ovlp);
		}
		numExtensions.push_back(extensions.size());

		//sort from longes to shortest overlap
		std::sort(extensions.begin(), extensions.end(),
				  [](const OverlapRange& a, const OverlapRange& b)
					 {return a.curRange() > b.curRange();});

		//checking if read overlaps with one of the used reads
		bool foundExtension = false;

		//getting extension
		int minExtensions = std::max(1, (int)extensions.size() /
										(int)Config::get("max_coverage_drop_rate"));

		/*Logger::get().debug() << _readsContainer.seqName(currentRead)
			<< "\t" << overlaps.size();
		for (auto& ovlp : overlaps)
		{
			if (this->extendsRight(ovlp))
			{
				Logger::get().debug() << "\t" <<
					_readsContainer.seqName(ovlp.extId) << "\t"
					<< "o:" << ovlp.curRange() << "\tchim:"
					<< _chimDetector.isChimeric(ovlp.extId)
					<< "\trep:" << this->isRightRepeat(ovlp.extId)
					<< "\text:" << this->countRightExtensions(ovlp.extId);
			}
		}*/

		OverlapRange* maxExtension = nullptr;
		for (auto& ovlp : extensions)
		{
			if(leftExtendsStart(ovlp.extId)) continue;

			//try to find a good one
			if (!_chimDetector.isChimeric(ovlp.extId) &&
				this->countRightExtensions(ovlp.extId) > minExtensions)
			{
				foundExtension = true;
				maxExtension = &ovlp;
				exInfo.assembledLength += ovlp.rightShift;
				currentRead = ovlp.extId;
				break;
			}

			if (!maxExtension || maxExtension->rightShift < ovlp.rightShift)
			{
				maxExtension = &ovlp;
			}
		}
		//in case of suspicious extension make the farthest jump possible
		if (!foundExtension && maxExtension)
		{
			++exInfo.numSuspicious;
			exInfo.assembledLength += maxExtension->rightShift;
			foundExtension = true;
			currentRead = maxExtension->extId;
		}

		bool overlapsVisited = _innerReads.contains(currentRead);
		if (foundExtension)
		{
			exInfo.reads.push_back(currentRead);
			overlapSizes.push_back(maxExtension->curRange());
			overlapsVisited |= currentReads.count(currentRead);
		}
		else
		{
			rightExtension ? exInfo.leftTip = true : exInfo.rightTip = true;
		}

		if (!foundExtension || overlapsVisited)
		{
			//Logger::get().debug() << "Not found: " << !foundExtension <<
			//	" overlaps visited: " << overlapsVisited;

			//right extension done, try to extend left from start read
			if (rightExtension && !exInfo.reads.empty())
			{
				exInfo.stepsToTurn = exInfo.reads.size();
				rightExtension = false;
				currentRead = exInfo.reads.front().rc();
				std::reverse(exInfo.reads.begin(), exInfo.reads.end());
				for (size_t i = 0; i < exInfo.reads.size(); ++i)
				{
					exInfo.reads[i] = exInfo.reads[i].rc();
				}
			}
			//done with extension
			else
			{
				break;
			}
		}

		currentReads.insert(currentRead);
		currentReads.insert(currentRead.rc());
	}

	if (!numExtensions.empty())
	{
		exInfo.meanOverlaps = median(numExtensions);
	}
	if (!overlapSizes.empty())
	{
		exInfo.avgOverlapSize = median(overlapSizes);
		exInfo.minOverlapSize = *std::min_element(overlapSizes.begin(),
												  overlapSizes.end());
	}

	return exInfo;
}

void Extender::assembleContigs()
{
	static const int MAX_JUMP = Config::get("maximum_jump");
	Logger::get().info() << "Extending reads";
	_chimDetector.estimateGlobalCoverage();
	_ovlpContainer.overlapDivergenceStats();
	_innerReads.clear();
	cuckoohash_map<FastaRecord::Id, size_t> coveredReads;

	std::mutex indexMutex;
	auto processRead = [this, &indexMutex, &coveredReads]
		(FastaRecord::Id startRead)
	{
		if (coveredReads.contains(startRead)) return true;

		int numInnerOvlp = 0;
		for (auto& ovlp : _ovlpContainer.lazySeqOverlaps(startRead))
		{
			if (_innerReads.contains(ovlp.extId)) ++numInnerOvlp;
			if (ovlp.extId == startRead ||
				ovlp.extId.rc() == startRead) return true;
		}
		if (numInnerOvlp > 0) return true;

		int maxExtensions = _chimDetector.getOverlapCoverage() * 10;
		if (_chimDetector.isChimeric(startRead) ||
			this->countRightExtensions(startRead) > maxExtensions ||
			this->countRightExtensions(startRead.rc()) > maxExtensions) return true;

		//Good to go!
		ExtensionInfo exInfo = this->extendContig(startRead);
		if (exInfo.reads.size() <
			(size_t)Config::get("min_reads_in_contig")) return true;

		//Exclusive part - updating the overall assembly
		std::lock_guard<std::mutex> guard(indexMutex);

		int innerCount = 0;
		for (auto& readId : exInfo.reads)
		{
			if (_innerReads.contains(readId)) ++innerCount;
		}
		int innerThreshold = std::min((int)Config::get("max_inner_reads"),
									  int((float)Config::get("max_inner_fraction") *
										  exInfo.reads.size()));
		if (innerCount > innerThreshold)
		{
			Logger::get().debug() << "Discarded contig with "
				<< exInfo.reads.size() << " reads and "
				<< innerCount << " inner overlaps";
			return false;
		}

		Logger::get().debug() << "Assembled contig "
			<< std::to_string(_readLists.size() + 1)
			<< "\n\tWith " << exInfo.reads.size() << " reads"
			<< "\n\tStart read: " << _readsContainer.seqName(startRead)
			<< "\n\tAt position: " << exInfo.stepsToTurn
			<< "\n\tleftTip: " << exInfo.leftTip
			<< " rightTip: " << exInfo.rightTip
			<< "\n\tSuspicios: " << exInfo.numSuspicious
			<< "\n\tMean extensions: " << exInfo.meanOverlaps
			<< "\n\tAvg overlap len: " << exInfo.avgOverlapSize
			<< "\n\tMin overlap len: " << exInfo.minOverlapSize
			<< "\n\tInner reads: " << innerCount
			<< "\n\tLength: " << exInfo.assembledLength;

		//update inner read index
		std::unordered_set<FastaRecord::Id> rightExtended;
		std::unordered_set<FastaRecord::Id> leftExtended;
		for (auto& readId : exInfo.reads)
		{
			coveredReads.insert(readId, true);
			coveredReads.insert(readId.rc(), true);
			_innerReads.insert(readId, true);
			_innerReads.insert(readId.rc(), true);

			//repetitive read - will bring to many "off-target" reads
			int maxExtensions = exInfo.meanOverlaps * 2;
			if (this->countRightExtensions(readId) > maxExtensions||
				this->countRightExtensions(readId.rc()) > maxExtensions) continue;

			//so each read is covered from the left and right
			for (auto& ovlp : _ovlpContainer.lazySeqOverlaps(readId))
			{
				coveredReads.insert(ovlp.extId, true);
				coveredReads.insert(ovlp.extId.rc(), true);

				//contained inside the other read - probably repetitive
				//if (ovlp.leftShift > MAX_JUMP &&
				//	ovlp.rightShift < -MAX_JUMP) continue;

				if (ovlp.leftShift > MAX_JUMP)
				{
					leftExtended.insert(ovlp.extId);
					rightExtended.insert(ovlp.extId.rc());
				}
				if (ovlp.rightShift < -MAX_JUMP)
				{
					rightExtended.insert(ovlp.extId);
					leftExtended.insert(ovlp.extId.rc());
				}
			}
		}
		for (auto& read : rightExtended)
		{
			if (leftExtended.count(read))
			{
				_innerReads.insert(read, true);
				_innerReads.insert(read.rc(), true);
			}
		}

		Logger::get().debug() << "Inner: " <<
			_innerReads.size() << " covered: " << coveredReads.size()
			<< " total: "<< _readsContainer.iterSeqs().size();

		_readLists.push_back(std::move(exInfo));
		return true;
	};

	std::function<void(const FastaRecord::Id&)> threadWorker =
		[processRead] (const FastaRecord::Id& readId)
	{
		processRead(readId);
	};
	std::vector<FastaRecord::Id> allReads;
	for (auto& seq : _readsContainer.iterSeqs())
	{
		allReads.push_back(seq.id);
	}
	std::random_shuffle(allReads.begin(), allReads.end());
	processInParallel(allReads, threadWorker,
					  Parameters::get().numThreads, true);
	_ovlpContainer.ensureTransitivity(/*only max*/ true);

	bool addSingletons = (bool)Config::get("add_unassembled_reads");
	if (addSingletons)
	{
		std::vector<FastaRecord::Id> sortedByLength;
		for (auto& seq : _readsContainer.iterSeqs())
		{
			if (seq.id.strand() && !_innerReads.contains(seq.id) &&
				_readsContainer.seqLen(seq.id) > Parameters::get().minimumOverlap)
			{
				sortedByLength.push_back(seq.id);
			}
		}
		std::sort(sortedByLength.begin(), sortedByLength.end(),
				  [this](FastaRecord::Id idOne, FastaRecord::Id idTwo)
				  	{return _readsContainer.seqLen(idOne) >
							_readsContainer.seqLen(idTwo);});

		int singletonsAdded = 0;
		std::unordered_set<FastaRecord::Id> coveredLocal;
		for (auto readId : sortedByLength)
		{
			if (!coveredLocal.count(readId))
			{
				for (auto& ovlp : _ovlpContainer.lazySeqOverlaps(readId))
				{
					if (ovlp.leftShift >= 0 && ovlp.rightShift <= 0)
					{
						coveredLocal.insert(ovlp.extId);
						coveredLocal.insert(ovlp.extId.rc());
					}
				}
				ExtensionInfo path;
				path.singleton = true;
				path.reads.push_back(readId);
				_readLists.push_back(path);
				++singletonsAdded;
			}
		}
		Logger::get().info() << "Added " << singletonsAdded << " singleton reads";
	}

	this->convertToContigs();
	Logger::get().info() << "Assembled " << _contigPaths.size() << " draft contigs";
}

void Extender::convertToContigs()
{
	for (auto& exInfo : _readLists)
	{
		ContigPath path;
		if (!exInfo.singleton)
		{
			path.name = "contig_" + std::to_string(_contigPaths.size() + 1);
		}
		else
		{
			path.name = "read_" + std::to_string(_contigPaths.size() + 1);
		}

		for (size_t i = 0; i < exInfo.reads.size() - 1; ++i)
		{
			bool found = false;
			OverlapRange readsOvlp;

			for (auto& ovlp : _ovlpContainer.lazySeqOverlaps(exInfo.reads[i]))
			{
				if (ovlp.extId == exInfo.reads[i + 1])
				{
					readsOvlp = ovlp;
					found = true;
					break;
				}
			}
			if (!found) throw std::runtime_error("Ovlp not found!");

			path.sequences.push_back(_readsContainer.getSeq(exInfo.reads[i]));
			path.overlaps.push_back(readsOvlp);
		}
		path.sequences.push_back(_readsContainer.getSeq(exInfo.reads.back()));
		_contigPaths.push_back(std::move(path));
	}
}

int Extender::countRightExtensions(FastaRecord::Id readId) const
{
	int count = 0;
	for (auto& ovlp : _ovlpContainer.lazySeqOverlaps(readId))
	{
		if (this->extendsRight(ovlp)) ++count;
	}
	return count;
}

/*bool Extender::isRightRepeat(FastaRecord::Id readId) const
{
	int maxExtensions = _chimDetector.getOverlapCoverage() * 10;
	return this->countRightExtensions(readId) >= maxExtensions;
}*/

bool Extender::extendsRight(const OverlapRange& ovlp) const
{
	static const int MAX_JUMP = Config::get("maximum_jump");
	return ovlp.rightShift > MAX_JUMP;
}
<commit_msg>more accurate disjointig generation<commit_after>//(c) 2016 by Authors
//This file is a part of ABruijn program.
//Released under the BSD license (see LICENSE file)

#include <limits>
#include <algorithm>
#include <iomanip>
#include <stack>

#include "../common/config.h"
#include "../common/logger.h"
#include "../common/parallel.h"
#include "extender.h"

Extender::ExtensionInfo Extender::extendContig(FastaRecord::Id startRead)
{

	std::unordered_set<FastaRecord::Id> currentReads;
	currentReads.insert(startRead);
	currentReads.insert(startRead.rc());

	bool rightExtension = true;
	FastaRecord::Id currentRead = startRead;
	std::vector<int> numExtensions;
	std::vector<int> overlapSizes;
	ExtensionInfo exInfo;
	exInfo.reads.push_back(startRead);
	exInfo.assembledLength = _readsContainer.seqLen(startRead);

	auto leftExtendsStart = [startRead, this](const FastaRecord::Id readId)
	{
		for (auto& ovlp : _ovlpContainer.lazySeqOverlaps(readId))
		{
			if (ovlp.extId == startRead &&
				ovlp.leftShift > (int)Config::get("maximum_jump"))
			{
				return true;
			}
		}
		return false;
	};

	while(true)
	{
		auto& overlaps = _ovlpContainer.lazySeqOverlaps(currentRead);
		std::vector<OverlapRange> extensions;
		//int innerOverlaps = 0;
		for (auto& ovlp : overlaps)
		{
			//if (_innerReads.contains(ovlp.extId)) ++innerOverlaps;
			if (this->extendsRight(ovlp)) extensions.push_back(ovlp);
		}
		numExtensions.push_back(extensions.size());

		//sort from longes to shortest overlap
		std::sort(extensions.begin(), extensions.end(),
				  [](const OverlapRange& a, const OverlapRange& b)
					 {return a.curRange() > b.curRange();});

		//checking if read overlaps with one of the used reads
		bool foundExtension = false;

		//getting extension
		int minExtensions = std::max(1, (int)extensions.size() /
										(int)Config::get("max_coverage_drop_rate"));

		/*Logger::get().debug() << _readsContainer.seqName(currentRead)
			<< "\t" << overlaps.size();
		for (auto& ovlp : overlaps)
		{
			if (this->extendsRight(ovlp))
			{
				Logger::get().debug() << "\t" <<
					_readsContainer.seqName(ovlp.extId) << "\t"
					<< "o:" << ovlp.curRange() << "\tchim:"
					<< _chimDetector.isChimeric(ovlp.extId)
					<< "\trep:" << this->isRightRepeat(ovlp.extId)
					<< "\text:" << this->countRightExtensions(ovlp.extId);
			}
		}*/

		OverlapRange* maxExtension = nullptr;
		for (auto& ovlp : extensions)
		{
			if(leftExtendsStart(ovlp.extId)) continue;

			//try to find a good one
			if (!_chimDetector.isChimeric(ovlp.extId) &&
				this->countRightExtensions(ovlp.extId) > minExtensions)
			{
				foundExtension = true;
				maxExtension = &ovlp;
				exInfo.assembledLength += ovlp.rightShift;
				currentRead = ovlp.extId;
				break;
			}

			if (!maxExtension || maxExtension->rightShift < ovlp.rightShift)
			{
				maxExtension = &ovlp;
			}
		}
		//in case of suspicious extension make the farthest jump possible
		if (!foundExtension && maxExtension)
		{
			++exInfo.numSuspicious;
			exInfo.assembledLength += maxExtension->rightShift;
			foundExtension = true;
			currentRead = maxExtension->extId;
		}

		bool overlapsVisited = _innerReads.contains(currentRead);
		if (foundExtension)
		{
			exInfo.reads.push_back(currentRead);
			overlapSizes.push_back(maxExtension->curRange());
			overlapsVisited |= currentReads.count(currentRead);
		}
		else
		{
			rightExtension ? exInfo.leftTip = true : exInfo.rightTip = true;
		}

		if (!foundExtension || overlapsVisited)
		{
			//Logger::get().debug() << "Not found: " << !foundExtension <<
			//	" overlaps visited: " << overlapsVisited;

			//right extension done, try to extend left from start read
			if (rightExtension && !exInfo.reads.empty())
			{
				exInfo.stepsToTurn = exInfo.reads.size();
				rightExtension = false;
				currentRead = exInfo.reads.front().rc();
				std::reverse(exInfo.reads.begin(), exInfo.reads.end());
				for (size_t i = 0; i < exInfo.reads.size(); ++i)
				{
					exInfo.reads[i] = exInfo.reads[i].rc();
				}
			}
			//done with extension
			else
			{
				break;
			}
		}

		currentReads.insert(currentRead);
		currentReads.insert(currentRead.rc());
	}

	if (!numExtensions.empty())
	{
		exInfo.meanOverlaps = median(numExtensions);
	}
	if (!overlapSizes.empty())
	{
		exInfo.avgOverlapSize = median(overlapSizes);
		exInfo.minOverlapSize = *std::min_element(overlapSizes.begin(),
												  overlapSizes.end());
	}

	return exInfo;
}

void Extender::assembleContigs()
{
	static const int MAX_JUMP = Config::get("maximum_jump");
	Logger::get().info() << "Extending reads";
	_chimDetector.estimateGlobalCoverage();
	_ovlpContainer.overlapDivergenceStats();
	_innerReads.clear();
	cuckoohash_map<FastaRecord::Id, size_t> coveredReads;

	int totalReads = 0;
	for (auto& read : _readsContainer.iterSeqs())
	{
		if ((int)read.sequence.length() >
			Parameters::get().minimumOverlap) ++totalReads;
	}

	std::mutex indexMutex;
	auto processRead = [this, &indexMutex, &coveredReads, totalReads]
		(FastaRecord::Id startRead)
	{
		//if (coveredReads.contains(startRead)) return true;
		if (_innerReads.contains(startRead)) return true;

		int numInnerOvlp = 0;
		for (auto& ovlp : _ovlpContainer.lazySeqOverlaps(startRead))
		{
			if (_innerReads.contains(ovlp.extId)) ++numInnerOvlp;
			if (ovlp.extId == startRead ||
				ovlp.extId.rc() == startRead) return true;
		}
		if (numInnerOvlp > 0) return true;

		int maxExtensions = _chimDetector.getOverlapCoverage() * 10;
		if (_chimDetector.isChimeric(startRead) ||
			this->countRightExtensions(startRead) > maxExtensions ||
			this->countRightExtensions(startRead.rc()) > maxExtensions) return true;

		//Good to go!
		ExtensionInfo exInfo = this->extendContig(startRead);
		if (exInfo.reads.size() <
			(size_t)Config::get("min_reads_in_contig")) return true;

		//Exclusive part - updating the overall assembly
		std::lock_guard<std::mutex> guard(indexMutex);

		int innerCount = 0;
		//do not count dirst and last reads - they are inner by defalut
		for (size_t i = 1; i < exInfo.reads.size() - 1; ++i)
		{
			if (_innerReads.contains(exInfo.reads[i])) ++innerCount;
		}
		int innerThreshold = std::min((int)Config::get("max_inner_reads"),
									  int((float)Config::get("max_inner_fraction") *
										  exInfo.reads.size()));
		if (innerCount > innerThreshold)
		{
			Logger::get().debug() << "Discarded contig with "
				<< exInfo.reads.size() << " reads and "
				<< innerCount << " inner overlaps";
			return false;
		}

		Logger::get().debug() << "Assembled contig "
			<< std::to_string(_readLists.size() + 1)
			<< "\n\tWith " << exInfo.reads.size() << " reads"
			<< "\n\tStart read: " << _readsContainer.seqName(startRead)
			<< "\n\tAt position: " << exInfo.stepsToTurn
			<< "\n\tleftTip: " << exInfo.leftTip
			<< " rightTip: " << exInfo.rightTip
			<< "\n\tSuspicios: " << exInfo.numSuspicious
			<< "\n\tMean extensions: " << exInfo.meanOverlaps
			<< "\n\tAvg overlap len: " << exInfo.avgOverlapSize
			<< "\n\tMin overlap len: " << exInfo.minOverlapSize
			<< "\n\tInner reads: " << innerCount
			<< "\n\tLength: " << exInfo.assembledLength;

		//update inner read index
		std::unordered_set<FastaRecord::Id> rightExtended;
		std::unordered_set<FastaRecord::Id> leftExtended;
		for (auto& readId : exInfo.reads)
		{
			coveredReads.insert(readId, true);
			coveredReads.insert(readId.rc(), true);
			_innerReads.insert(readId, true);
			_innerReads.insert(readId.rc(), true);

			//repetitive read - will bring to many "off-target" reads
			int maxExtensions = exInfo.meanOverlaps * 2;
			if (this->countRightExtensions(readId) > maxExtensions||
				this->countRightExtensions(readId.rc()) > maxExtensions) continue;

			//so each read is covered from the left and right
			for (auto& ovlp : _ovlpContainer.lazySeqOverlaps(readId))
			{
				coveredReads.insert(ovlp.extId, true);
				coveredReads.insert(ovlp.extId.rc(), true);

				//contained inside the other read - probably repetitive
				//if (ovlp.leftShift > MAX_JUMP &&
				//	ovlp.rightShift < -MAX_JUMP) continue;

				if (ovlp.leftShift > MAX_JUMP)
				{
					leftExtended.insert(ovlp.extId);
					rightExtended.insert(ovlp.extId.rc());
				}
				if (ovlp.rightShift < -MAX_JUMP)
				{
					rightExtended.insert(ovlp.extId);
					leftExtended.insert(ovlp.extId.rc());
				}
			}
		}
		for (auto& read : rightExtended)
		{
			if (leftExtended.count(read))
			{
				_innerReads.insert(read, true);
				_innerReads.insert(read.rc(), true);
			}
		}

		Logger::get().debug() << "Inner: " <<
			_innerReads.size() << " covered: " << coveredReads.size()
			<< " total: "<< totalReads;

		_readLists.push_back(std::move(exInfo));
		return true;
	};

	std::function<void(const FastaRecord::Id&)> threadWorker =
		[processRead] (const FastaRecord::Id& readId)
	{
		processRead(readId);
	};
	std::vector<FastaRecord::Id> allReads;
	for (auto& seq : _readsContainer.iterSeqs())
	{
		allReads.push_back(seq.id);
	}
	std::random_shuffle(allReads.begin(), allReads.end());
	processInParallel(allReads, threadWorker,
					  Parameters::get().numThreads, true);
	_ovlpContainer.ensureTransitivity(/*only max*/ true);

	bool addSingletons = (bool)Config::get("add_unassembled_reads");
	if (addSingletons)
	{
		std::vector<FastaRecord::Id> sortedByLength;
		for (auto& seq : _readsContainer.iterSeqs())
		{
			if (seq.id.strand() && !_innerReads.contains(seq.id) &&
				_readsContainer.seqLen(seq.id) > Parameters::get().minimumOverlap)
			{
				sortedByLength.push_back(seq.id);
			}
		}
		std::sort(sortedByLength.begin(), sortedByLength.end(),
				  [this](FastaRecord::Id idOne, FastaRecord::Id idTwo)
				  	{return _readsContainer.seqLen(idOne) >
							_readsContainer.seqLen(idTwo);});

		int singletonsAdded = 0;
		std::unordered_set<FastaRecord::Id> coveredLocal;
		for (auto readId : sortedByLength)
		{
			if (!coveredLocal.count(readId))
			{
				for (auto& ovlp : _ovlpContainer.lazySeqOverlaps(readId))
				{
					if (ovlp.leftShift >= 0 && ovlp.rightShift <= 0)
					{
						coveredLocal.insert(ovlp.extId);
						coveredLocal.insert(ovlp.extId.rc());
					}
				}
				ExtensionInfo path;
				path.singleton = true;
				path.reads.push_back(readId);
				_readLists.push_back(path);
				++singletonsAdded;
			}
		}
		Logger::get().info() << "Added " << singletonsAdded << " singleton reads";
	}

	this->convertToContigs();
	Logger::get().info() << "Assembled " << _contigPaths.size() << " draft contigs";
}

void Extender::convertToContigs()
{
	for (auto& exInfo : _readLists)
	{
		ContigPath path;
		if (!exInfo.singleton)
		{
			path.name = "contig_" + std::to_string(_contigPaths.size() + 1);
		}
		else
		{
			path.name = "read_" + std::to_string(_contigPaths.size() + 1);
		}

		for (size_t i = 0; i < exInfo.reads.size() - 1; ++i)
		{
			bool found = false;
			OverlapRange readsOvlp;

			for (auto& ovlp : _ovlpContainer.lazySeqOverlaps(exInfo.reads[i]))
			{
				if (ovlp.extId == exInfo.reads[i + 1])
				{
					readsOvlp = ovlp;
					found = true;
					break;
				}
			}
			if (!found) throw std::runtime_error("Ovlp not found!");

			path.sequences.push_back(_readsContainer.getSeq(exInfo.reads[i]));
			path.overlaps.push_back(readsOvlp);
		}
		path.sequences.push_back(_readsContainer.getSeq(exInfo.reads.back()));
		_contigPaths.push_back(std::move(path));
	}
}

int Extender::countRightExtensions(FastaRecord::Id readId) const
{
	int count = 0;
	for (auto& ovlp : _ovlpContainer.lazySeqOverlaps(readId))
	{
		if (this->extendsRight(ovlp)) ++count;
	}
	return count;
}

/*bool Extender::isRightRepeat(FastaRecord::Id readId) const
{
	int maxExtensions = _chimDetector.getOverlapCoverage() * 10;
	return this->countRightExtensions(readId) >= maxExtensions;
}*/

bool Extender::extendsRight(const OverlapRange& ovlp) const
{
	static const int MAX_JUMP = Config::get("maximum_jump");
	return ovlp.rightShift > MAX_JUMP;
}
<|endoftext|>


<commit_before>/*
 * $Id$
 *
 * Copyright(C) 1998-2004 Satoshi Nakamura
 * All rights reserved.
 *
 */

#include <qmaccount.h>
#include <qmapplication.h>
#include <qmfilenames.h>
#include <qmfolder.h>
#include <qmmacro.h>

#include <qsconv.h>
#include <qsutil.h>

#include <algorithm>

#include "color.h"

using namespace qm;
using namespace qs;

/****************************************************************************
 *
 * ColorManager
 *
 */

qm::ColorManager::ColorManager()
{
	load();
}

qm::ColorManager::~ColorManager()
{
	std::for_each(listColorSet_.begin(),
		listColorSet_.end(), deleter<ColorSet>());
}

const ColorSet* qm::ColorManager::getColorSet(Folder* pFolder) const
{
	assert(pFolder);

	for (ColorSetList::const_iterator it = listColorSet_.begin(); it != listColorSet_.end(); ++it) {
		const ColorSet* pColor = *it;
		if (pColor->match(pFolder))
			return pColor;
	}

	return 0;
}

void qm::ColorManager::addColorSet(std::auto_ptr<ColorSet> pSet)
{
	listColorSet_.push_back(pSet.get());
	pSet.release();
}

bool qm::ColorManager::load()
{
	wstring_ptr wstrPath(
		Application::getApplication().getProfilePath(FileNames::COLORS_XML));

	W2T(wstrPath.get(), ptszPath);
	if (::GetFileAttributes(ptszPath) != 0xffffffff) {
		XMLReader reader;
		ColorContentHandler handler(this);
		reader.setContentHandler(&handler);
		if (!reader.parse(wstrPath.get()))
			return false;
	}

	return true;
}

/****************************************************************************
 *
 * ColorSet
 *
 */

qm::ColorSet::ColorSet(std::auto_ptr<RegexPattern> pAccountName,
					   std::auto_ptr<RegexPattern> pFolderName) :
	pAccountName_(pAccountName),
	pFolderName_(pFolderName)
{
}

qm::ColorSet::~ColorSet()
{
	std::for_each(listEntry_.begin(),
		listEntry_.end(), deleter<ColorEntry>());
}

bool qm::ColorSet::match(Folder* pFolder) const
{
	assert(pFolder);

	if (pAccountName_.get() &&
		!pAccountName_->match(pFolder->getAccount()->getName()))
		return false;

	if (pFolderName_.get()) {
		wstring_ptr wstrFullName(pFolder->getFullName());
		if (!pFolderName_->match(wstrFullName.get()))
			return false;
	}

	return true;
}

COLORREF qm::ColorSet::getColor(MacroContext* pContext) const
{
	assert(pContext);

	for (EntryList::const_iterator it = listEntry_.begin(); it != listEntry_.end(); ++it) {
		const ColorEntry* pEntry = *it;
		if (pEntry->match(pContext))
			return pEntry->getColor();
	}

	return 0xff000000;
}

void qm::ColorSet::addEntry(std::auto_ptr<ColorEntry> pEntry)
{
	listEntry_.push_back(pEntry.get());
	pEntry.release();
}

/****************************************************************************
 *
 * ColorEntry
 *
 */

qm::ColorEntry::ColorEntry(std::auto_ptr<Macro> pMacro,
						   COLORREF cr) :
	pMacro_(pMacro),
	cr_(cr)
{
}

qm::ColorEntry::~ColorEntry()
{
}

bool qm::ColorEntry::match(MacroContext* pContext) const
{
	assert(pContext);
	assert(pContext->getMessageHolder());

	MacroValuePtr pValue(pMacro_->value(pContext));
	return pValue->boolean();
}

COLORREF qm::ColorEntry::getColor() const
{
	return cr_;
}

/****************************************************************************
 *
 * ColorContentHandler
 *
 */

qm::ColorContentHandler::ColorContentHandler(ColorManager* pManager) :
	pManager_(pManager),
	state_(STATE_ROOT),
	pColorSet_(0),
	parser_(MacroParser::TYPE_COLOR)
{
}

qm::ColorContentHandler::~ColorContentHandler()
{
}

bool qm::ColorContentHandler::startElement(const WCHAR* pwszNamespaceURI,
										   const WCHAR* pwszLocalName,
										   const WCHAR* pwszQName,
										   const Attributes& attributes)
{
	if (wcscmp(pwszLocalName, L"colors") == 0) {
		if (state_ != STATE_ROOT)
			return false;
		if (attributes.getLength() != 0)
			return false;
		state_ = STATE_COLORS;
	}
	else if (wcscmp(pwszLocalName, L"colorSet") == 0) {
		if (state_ != STATE_COLORS)
			return false;

		const WCHAR* pwszAccount = 0;
		const WCHAR* pwszFolder = 0;
		for (int n = 0; n < attributes.getLength(); ++n) {
			const WCHAR* pwszAttrName = attributes.getLocalName(n);
			if (wcscmp(pwszAttrName, L"account") == 0)
				pwszAccount = attributes.getValue(n);
			else if (wcscmp(pwszAttrName, L"folder") == 0)
				pwszFolder = attributes.getValue(n);
			else
				return false;
		}

		assert(!pColorSet_);

		RegexCompiler compiler;

		std::auto_ptr<RegexPattern> pAccountName;
		if (pwszAccount) {
			pAccountName = compiler.compile(pwszAccount);
			if (!pAccountName.get())
				return false;
		}

		std::auto_ptr<RegexPattern> pFolderName;
		if (pwszFolder) {
			pFolderName = compiler.compile(pwszFolder);
			if (!pFolderName.get())
				return false;
		}

		std::auto_ptr<ColorSet> pColorSet(new ColorSet(pAccountName, pFolderName));
		pColorSet_ = pColorSet.get();
		pManager_->addColorSet(pColorSet);

		state_ = STATE_COLORSET;
	}
	else if (wcscmp(pwszLocalName, L"color") == 0) {
		if (state_ != STATE_COLORSET)
			return false;

		const WCHAR* pwszMatch = 0;
		for (int n = 0; n < attributes.getLength(); ++n) {
			const WCHAR* pwszAttrName = attributes.getLocalName(n);
			if (wcscmp(pwszAttrName, L"match") == 0)
				pwszMatch = attributes.getValue(n);
			else
				return false;
		}
		if (!pwszMatch)
			return false;

		assert(!pMacro_.get());
		pMacro_ = parser_.parse(pwszMatch);
		if (!pMacro_.get())
			return false;

		state_ = STATE_COLOR;
	}
	else {
		return false;
	}

	return true;
}

bool qm::ColorContentHandler::endElement(const WCHAR* pwszNamespaceURI,
										 const WCHAR* pwszLocalName,
										 const WCHAR* pwszQName)
{
	if (wcscmp(pwszLocalName, L"colors") == 0) {
		assert(state_ == STATE_COLORS);
		state_ = STATE_ROOT;
	}
	else if (wcscmp(pwszLocalName, L"colorSet") == 0) {
		assert(state_ = STATE_COLORSET);
		pColorSet_ = 0;
		state_ = STATE_COLORS;
	}
	else if (wcscmp(pwszLocalName, L"color") == 0) {
		assert(state_ == STATE_COLOR);

		Color color(buffer_.getCharArray());
		if (color.getColor() == 0xffffffff)
			return false;
		std::auto_ptr<ColorEntry> pEntry(new ColorEntry(pMacro_, color.getColor()));
		pMacro_.reset(0);
		buffer_.remove();

		assert(pColorSet_);
		pColorSet_->addEntry(pEntry);

		state_ = STATE_COLORSET;
	}
	else {
		return false;
	}

	return true;
}

bool qm::ColorContentHandler::characters(const WCHAR* pwsz,
										 size_t nStart,
										 size_t nLength)
{
	if (state_ == STATE_COLOR) {
		buffer_.append(pwsz + nStart, nLength);
	}
	else {
		const WCHAR* p = pwsz + nStart;
		for (size_t n = 0; n < nLength; ++n, ++p) {
			if (*p != L' ' && *p != L'\t' && *p != '\n')
				return false;
		}
	}

	return true;
}
<commit_msg>Fix qm::ColorEntry::match causes crush when evaluating macro failed.<commit_after>/*
 * $Id$
 *
 * Copyright(C) 1998-2004 Satoshi Nakamura
 * All rights reserved.
 *
 */

#include <qmaccount.h>
#include <qmapplication.h>
#include <qmfilenames.h>
#include <qmfolder.h>
#include <qmmacro.h>

#include <qsconv.h>
#include <qsutil.h>

#include <algorithm>

#include "color.h"

using namespace qm;
using namespace qs;

/****************************************************************************
 *
 * ColorManager
 *
 */

qm::ColorManager::ColorManager()
{
	load();
}

qm::ColorManager::~ColorManager()
{
	std::for_each(listColorSet_.begin(),
		listColorSet_.end(), deleter<ColorSet>());
}

const ColorSet* qm::ColorManager::getColorSet(Folder* pFolder) const
{
	assert(pFolder);

	for (ColorSetList::const_iterator it = listColorSet_.begin(); it != listColorSet_.end(); ++it) {
		const ColorSet* pColor = *it;
		if (pColor->match(pFolder))
			return pColor;
	}

	return 0;
}

void qm::ColorManager::addColorSet(std::auto_ptr<ColorSet> pSet)
{
	listColorSet_.push_back(pSet.get());
	pSet.release();
}

bool qm::ColorManager::load()
{
	wstring_ptr wstrPath(
		Application::getApplication().getProfilePath(FileNames::COLORS_XML));

	W2T(wstrPath.get(), ptszPath);
	if (::GetFileAttributes(ptszPath) != 0xffffffff) {
		XMLReader reader;
		ColorContentHandler handler(this);
		reader.setContentHandler(&handler);
		if (!reader.parse(wstrPath.get()))
			return false;
	}

	return true;
}

/****************************************************************************
 *
 * ColorSet
 *
 */

qm::ColorSet::ColorSet(std::auto_ptr<RegexPattern> pAccountName,
					   std::auto_ptr<RegexPattern> pFolderName) :
	pAccountName_(pAccountName),
	pFolderName_(pFolderName)
{
}

qm::ColorSet::~ColorSet()
{
	std::for_each(listEntry_.begin(),
		listEntry_.end(), deleter<ColorEntry>());
}

bool qm::ColorSet::match(Folder* pFolder) const
{
	assert(pFolder);

	if (pAccountName_.get() &&
		!pAccountName_->match(pFolder->getAccount()->getName()))
		return false;

	if (pFolderName_.get()) {
		wstring_ptr wstrFullName(pFolder->getFullName());
		if (!pFolderName_->match(wstrFullName.get()))
			return false;
	}

	return true;
}

COLORREF qm::ColorSet::getColor(MacroContext* pContext) const
{
	assert(pContext);

	for (EntryList::const_iterator it = listEntry_.begin(); it != listEntry_.end(); ++it) {
		const ColorEntry* pEntry = *it;
		if (pEntry->match(pContext))
			return pEntry->getColor();
	}

	return 0xff000000;
}

void qm::ColorSet::addEntry(std::auto_ptr<ColorEntry> pEntry)
{
	listEntry_.push_back(pEntry.get());
	pEntry.release();
}

/****************************************************************************
 *
 * ColorEntry
 *
 */

qm::ColorEntry::ColorEntry(std::auto_ptr<Macro> pMacro,
						   COLORREF cr) :
	pMacro_(pMacro),
	cr_(cr)
{
}

qm::ColorEntry::~ColorEntry()
{
}

bool qm::ColorEntry::match(MacroContext* pContext) const
{
	assert(pContext);
	assert(pContext->getMessageHolder());

	MacroValuePtr pValue(pMacro_->value(pContext));
	return pValue.get() && pValue->boolean();
}

COLORREF qm::ColorEntry::getColor() const
{
	return cr_;
}

/****************************************************************************
 *
 * ColorContentHandler
 *
 */

qm::ColorContentHandler::ColorContentHandler(ColorManager* pManager) :
	pManager_(pManager),
	state_(STATE_ROOT),
	pColorSet_(0),
	parser_(MacroParser::TYPE_COLOR)
{
}

qm::ColorContentHandler::~ColorContentHandler()
{
}

bool qm::ColorContentHandler::startElement(const WCHAR* pwszNamespaceURI,
										   const WCHAR* pwszLocalName,
										   const WCHAR* pwszQName,
										   const Attributes& attributes)
{
	if (wcscmp(pwszLocalName, L"colors") == 0) {
		if (state_ != STATE_ROOT)
			return false;
		if (attributes.getLength() != 0)
			return false;
		state_ = STATE_COLORS;
	}
	else if (wcscmp(pwszLocalName, L"colorSet") == 0) {
		if (state_ != STATE_COLORS)
			return false;

		const WCHAR* pwszAccount = 0;
		const WCHAR* pwszFolder = 0;
		for (int n = 0; n < attributes.getLength(); ++n) {
			const WCHAR* pwszAttrName = attributes.getLocalName(n);
			if (wcscmp(pwszAttrName, L"account") == 0)
				pwszAccount = attributes.getValue(n);
			else if (wcscmp(pwszAttrName, L"folder") == 0)
				pwszFolder = attributes.getValue(n);
			else
				return false;
		}

		assert(!pColorSet_);

		RegexCompiler compiler;

		std::auto_ptr<RegexPattern> pAccountName;
		if (pwszAccount) {
			pAccountName = compiler.compile(pwszAccount);
			if (!pAccountName.get())
				return false;
		}

		std::auto_ptr<RegexPattern> pFolderName;
		if (pwszFolder) {
			pFolderName = compiler.compile(pwszFolder);
			if (!pFolderName.get())
				return false;
		}

		std::auto_ptr<ColorSet> pColorSet(new ColorSet(pAccountName, pFolderName));
		pColorSet_ = pColorSet.get();
		pManager_->addColorSet(pColorSet);

		state_ = STATE_COLORSET;
	}
	else if (wcscmp(pwszLocalName, L"color") == 0) {
		if (state_ != STATE_COLORSET)
			return false;

		const WCHAR* pwszMatch = 0;
		for (int n = 0; n < attributes.getLength(); ++n) {
			const WCHAR* pwszAttrName = attributes.getLocalName(n);
			if (wcscmp(pwszAttrName, L"match") == 0)
				pwszMatch = attributes.getValue(n);
			else
				return false;
		}
		if (!pwszMatch)
			return false;

		assert(!pMacro_.get());
		pMacro_ = parser_.parse(pwszMatch);
		if (!pMacro_.get())
			return false;

		state_ = STATE_COLOR;
	}
	else {
		return false;
	}

	return true;
}

bool qm::ColorContentHandler::endElement(const WCHAR* pwszNamespaceURI,
										 const WCHAR* pwszLocalName,
										 const WCHAR* pwszQName)
{
	if (wcscmp(pwszLocalName, L"colors") == 0) {
		assert(state_ == STATE_COLORS);
		state_ = STATE_ROOT;
	}
	else if (wcscmp(pwszLocalName, L"colorSet") == 0) {
		assert(state_ = STATE_COLORSET);
		pColorSet_ = 0;
		state_ = STATE_COLORS;
	}
	else if (wcscmp(pwszLocalName, L"color") == 0) {
		assert(state_ == STATE_COLOR);

		Color color(buffer_.getCharArray());
		if (color.getColor() == 0xffffffff)
			return false;
		std::auto_ptr<ColorEntry> pEntry(new ColorEntry(pMacro_, color.getColor()));
		pMacro_.reset(0);
		buffer_.remove();

		assert(pColorSet_);
		pColorSet_->addEntry(pEntry);

		state_ = STATE_COLORSET;
	}
	else {
		return false;
	}

	return true;
}

bool qm::ColorContentHandler::characters(const WCHAR* pwsz,
										 size_t nStart,
										 size_t nLength)
{
	if (state_ == STATE_COLOR) {
		buffer_.append(pwsz + nStart, nLength);
	}
	else {
		const WCHAR* p = pwsz + nStart;
		for (size_t n = 0; n < nLength; ++n, ++p) {
			if (*p != L' ' && *p != L'\t' && *p != '\n')
				return false;
		}
	}

	return true;
}
<|endoftext|>