argiopetech/base

Add adaptive burnin

Closed this issue · 2 comments

This existed previously in Base8.

 ///////////////////////////////////
// Adjust step sizes for a chain //
///////////////////////////////////
void adjustChainStepSizes(struct chain *mc, const struct mcmcControl *ctrl){
  mc->clust.stepSize[AGE] *= adjustStepSizes(&mc->acceptClust[AGE],&mc->rejectClust[AGE]);

  if(ctrl->currentBlock != MASS_RATIO_BLOCK) {
    if(mc->clust.priorVar[MOD] > EPSILON)
      mc->clust.stepSize[MOD] *= adjustStepSizes(&mc->acceptClust[MOD],&mc->rejectClust[MOD]);
    if(mc->clust.priorVar[FEH] > EPSILON)
      mc->clust.stepSize[FEH] *= adjustStepSizes(&mc->acceptClust[FEH],&mc->rejectClust[FEH]);
    if(mc->clust.priorVar[ABS] > EPSILON)
      mc->clust.stepSize[ABS] *= adjustStepSizes(&mc->acceptClust[ABS],&mc->rejectClust[ABS]);
    if(mc->clust.evoModels.mainSequenceEvol == CHABHELIUM)
      mc->clust.stepSize[YYY] *= adjustStepSizes(&mc->acceptClust[YYY],&mc->rejectClust[YYY]);
  }

  int j;
  for(j = 0; j < mc->clust.nStars; j++) {
    mc->stars[j].UStepSize         *= adjustStepSizes(&mc->acceptMass[j], &mc->rejectMass[j]);
    mc->stars[j].massRatioStepSize *= adjustStepSizes(&mc->acceptMassRatio[j], &mc->rejectMassRatio[j]);
    if(mc->stars[j].massRatioStepSize > 1.0) mc->stars[j].massRatioStepSize = 1.0;
  }
} // adjustChainStepSizes

///////////////////////////
//// Adjust Step Sizes ////
///////////////////////////
double adjustStepSizes(int *accept, int *reject){

  int i, iupper = 4, ilower = 7;
  double fracAccept;
  double fa[7] = {0.9, 0.7, 0.5, 0.3, 0.05, 0.15, 0.2};
  double w[7] = {2.0, 1.8, 1.5, 1.2, 0.5, 1/1.8, 1/1.5};

  if(*accept + *reject < 200) return 1.0;

  fracAccept = *accept / (double) (*accept + *reject);
  *accept = 0;
  *reject = 0;

  for(i=0;i<iupper;i++){
    if(fracAccept > fa[i]) return w[i];
  }
  for(i = 4; i < ilower; i++){
    if(fracAccept < fa[i]) return w[i];
  }
  return 1.0;
} // adjustStepSizes

Add code to perform a basic Monte-Carlo process on parameter step sizes by scaling parameters (either individually or all at once) with the goal of fitting acceptance ratio to a YAML-configurable range (default 20-40%). This process should run for a maximum of n iterations and should fall back to one of the older burnin methods if an acceptable acceptance ratio has not been reached.

-- | A general representation of a Markov Chain with acceptance and rejectance tracking
data Chain = Chain { accepted :: Int -- ^ Number of accepted proposals
                   , rejected :: Int -- ^ Number of rejected proposals
                   }

-- | Calculates the acceptance ratio of a `Chain`
acceptanceRatio :: Chain -> Double
acceptanceRatio c = let a = fromIntegral $ accepted c
                        r = fromIntegral $ rejected c
                    in a / r

-- | A data structure which keeps track of step sizes for the proposal process
data StepSizes = StepSizes { age :: Double
                           , feh :: Double
                           , mod :: Double
                           , abs :: Double
                           , yyy :: Double
                           , ifmrIntercept :: Double
                           , ifmrSlope :: Double
                           , ifmrQuadCoef :: Double
                           }

-- | Some pre-set step sizes used as a starting point for the burnin process
defaultStepSizes = undefined

-- | Scales all parameter step sizes in a `StepSizes` by some scaling factor
scale :: StepSizes -> Double -> StepSizes
scale (StepSizes ag fe m ab y i1 i2 i3) f =
    StepSizes (f * ag) (f * fe) (f * m) (f * ab) (f * y) (f * i1) (f * i2) (f * i3)

-- | Runs an adaptive burnin for a certain number of iterations (hard-coded to 3 here), falling back to `fallback` if an acceptable ratio is not reached
runAdaptiveBurnin :: StepSizes
runAdaptiveBurnin = go defaultStepSizes 1
    where go _ 3 = fallback -- Switch to a fallback burnin if we've tried too many times
          go s n = let r = acceptanceRatio $ burnin s -- Run the burnin
                   in if inBounds r
                        then s -- Return the current step sizes if we're in bounds
                        else go (scale s $ scaleFactor r) (n + 1) -- Otherwise, scale the StepSizes by a pre-set scaling factor and try again
              where inBounds r | r > 0.2 && r < 0.4 = True -- Shooting for between 20% and 40%
                               | otherwise          = False
                    scaleFactor r | r > 0.9 = 2.0 -- This is from BASE-8
                                  | r > 0.7 = 1.8
                                  | r > 0.5 = 1.5
                                  | r > 0.4 = 1.2 -- Changed from `r > 0.3` to reflect the current desire for wider bounds
                                  | r < 0.2  = 1 / 1.5
                                  | r < 0.15 = 1 / 1.8
                                  | r < 0.05 = 0.5
                                  | otherwise = 1.0

-- | Fallback is one of the already-defined burnin methods
fallback :: StepSizes
fallback = undefined

-- | The actual MCMC process
burnin :: StepSizes -> Chain
burnin = undefined