Пример #1
0
def main(comm):
  """
  This is a hacked version of MC3's func.py.
  This function directly call's the modeling function for the BART project.

  Modification History:
  ---------------------
  2014-04-19  patricio  Initial implementation.  [email protected]
  2014-06-25  patricio  Added support for inner-MPI loop.
  """
  # Parse arguments:
  cparser = argparse.ArgumentParser(description=__doc__, add_help=False,
                         formatter_class=argparse.RawDescriptionHelpFormatter)
  # Add config file option:
  cparser.add_argument("-c", "--config_file", 
                       help="Configuration file", metavar="FILE")
  # Remaining_argv contains all other command-line-arguments:
  args, remaining_argv = cparser.parse_known_args()

  # Get parameters from configuration file:
  cfile = args.config_file
  if cfile:
    config = ConfigParser.SafeConfigParser()
    config.optionxform = str
    config.read([cfile])
    defaults = dict(config.items("MCMC"))
  else:
    defaults = {}
  parser = argparse.ArgumentParser(parents=[cparser])
  parser.add_argument("--func",      dest="func",      type=mu.parray, 
                                     action="store",  default=None)
  parser.add_argument("--indparams", dest="indparams", type=mu.parray, 
                                     action="store",   default=[])
  parser.add_argument("--params",    dest="params",    type=mu.parray,
                                     action="store",   default=None,
                      help="Model-fitting parameters [default: %(default)s]")
  parser.add_argument("--molfit",    dest="molfit",    type=mu.parray,
                                     action="store",   default=None,
                      help="Molecules fit [default: %(default)s]")
  parser.add_argument("--Tmin",      dest="Tmin",      type=float,
                      action="store",  default=400.0,
                      help="Lower Temperature boundary [default: %(default)s]")
  parser.add_argument("--Tmax",      dest="Tmax",      type=float,
                      action="store",  default=3000.0,
                      help="Higher Temperature boundary [default: %(default)s]")
  parser.add_argument("--quiet",             action="store_true",
                      help="Set verbosity level to minimum",
                      dest="quiet")
  # Input-Converter Options:
  group = parser.add_argument_group("Input Converter Options")
  group.add_argument("--atmospheric_file",  action="store",
                     help="Atmospheric file [default: %(default)s]",
                     dest="atmfile", type=str,    default=None)
  group.add_argument("--PTtype",            action="store",
                     help="PT profile type.",
                     dest="PTtype",  type=str,    default="none")
                     #choices=('line', 'madhu'))
  group.add_argument("--tint",              action="store",
                     help="Internal temperature of the planet [default: "
                     "%(default)s].",
                     dest="tint",    type=float,  default=100.0)
  # transit Options:
  group = parser.add_argument_group("transit Options")
  group.add_argument("--config",  action="store",
                     help="transit configuration file [default: %(default)s]",
                     dest="config", type=str,    default=None)
  # Output-Converter Options:
  group = parser.add_argument_group("Output Converter Options")
  group.add_argument("--filter",                 action="store",
                     help="Waveband filter name [default: %(default)s]",
                     dest="filter",   type=mu.parray, default=None)
  group.add_argument("--tep_name",          action="store",
                     help="A TEP file [default: %(default)s]",
                     dest="tep_name", type=str,    default=None)
  group.add_argument("--kurucz_file",           action="store",
                     help="Stellar Kurucz file [default: %(default)s]",
                     dest="kurucz",   type=str,       default=None)
  group.add_argument("--solution",                    action="store",
                     help="Solution geometry [default: %(default)s]",
                     dest="solution", type=str,       default="None",
                     choices=('transit', 'eclipse'))

  parser.set_defaults(**defaults)
  args2, unknown = parser.parse_known_args(remaining_argv)

  # Quiet all threads except rank 0:
  rank = comm.Get_rank()
  verb = rank == 0

  # Get (Broadcast) the number of parameters and iterations from MPI:
  array1 = np.zeros(2, np.int)
  mu.comm_bcast(comm, array1)
  npars, niter = array1

  # :::::::  Initialize the Input converter ::::::::::::::::::::::::::
  atmfile  = args2.atmfile
  molfit   = args2.molfit
  PTtype   = args2.PTtype
  params   = args2.params
  tepfile  = args2.tep_name
  tint     = args2.tint
  Tmin     = args2.Tmin
  Tmax     = args2.Tmax
  solution = args2.solution  # Solution type

  # Extract necessary values from the TEP file:
  tep = rd.File(tepfile)
  # Stellar temperature in K:
  tstar = float(tep.getvalue('Ts')[0])
  # Stellar radius (in meters):
  rstar = float(tep.getvalue('Rs')[0]) * c.Rsun
  # Semi-major axis (in meters):
  sma   = float(tep.getvalue( 'a')[0]) * sc.au
  # Planetary radius (in meters):
  rplanet = float(tep.getvalue('Rp')[0]) * c.Rjup
  # Planetary mass (in kg):
  mplanet = float(tep.getvalue('Mp')[0]) * c.Mjup

  # Number of fitting parameters:
  nfree   = len(params)                 # Total number of free parameters
  nmolfit = len(molfit)                 # Number of molecular free parameters
  nradfit = int(solution == 'transit')  # 1 for transit, 0 for eclipse
  nPT     = nfree - nmolfit - nradfit   # Number of PT free parameters

  # Read atmospheric file to get data arrays:
  species, pressure, temp, abundances = mat.readatm(atmfile)
  # Reverse pressure order (for PT to work):
  pressure = pressure[::-1]
  nlayers  = len(pressure)   # Number of atmospheric layers
  nspecies = len(species)    # Number of species in the atmosphere
  mu.msg(verb, "There are {:d} layers and {:d} species.".format(nlayers,
                                                                nspecies))
  # Find index for Hydrogen and Helium:
  species = np.asarray(species)
  iH2     = np.where(species=="H2")[0]
  iHe     = np.where(species=="He")[0]
  # Get H2/He abundance ratio:
  ratio = (abundances[:,iH2] / abundances[:,iHe]).squeeze()
  # Find indices for the metals:
  imetals = np.where((species != "He") & (species != "H2"))[0]
  # Index of molecular abundances being modified:
  imol = np.zeros(nmolfit, dtype='i')
  for i in np.arange(nmolfit):
    imol[i] = np.where(np.asarray(species) == molfit[i])[0]

  # Pressure-Temperature profile:
  PTargs = [PTtype]
  if PTtype == "line":
    # Planetary surface gravity (in cm s-2):
    gplanet = 100.0 * sc.G * mplanet / rplanet**2
    # Additional PT arguments:
    PTargs += [rstar, tstar, tint, sma, gplanet]

  # Allocate arrays for receiving and sending data to master:
  freepars = np.zeros(nfree,                 dtype='d')
  profiles = np.zeros((nspecies+1, nlayers), dtype='d')
  # This are sub-sections of profiles, containing just the temperature and
  # the abundance profiles, respectively:
  tprofile  = profiles[0, :]
  aprofiles = profiles[1:,:]

  # Store abundance profiles:
  for i in np.arange(nspecies):
    aprofiles[i] = abundances[:, i]

  # :::::::  Spawn transit code  :::::::::::::::::::::::::::::::::::::
  # # transit configuration file:
  transitcfile = args2.tconfig
 
  # FINDME: Find a way to set verb to the transit subprocesses.
  # Silence all threads except rank 0:
  # if verb == 0:
  #   rargs = ["--quiet"]
  # else:
  #   rargs = []

  # Initialize the transit python module:
  transit_args = ["transit", "-c", transitcfile]
  trm.transit_init(len(transit_args), transit_args)

  # Get wavenumber array from transit:
  nwave  = trm.get_no_samples()
  specwn = trm.get_waveno_arr(nwave)

  # :::::::  Output Converter  :::::::::::::::::::::::::::::::::::::::
  ffile    = args2.filter    # Filter files
  kurucz   = args2.kurucz    # Kurucz file

  # Log10(stellar gravity)
  gstar = float(tep.getvalue('loggstar')[0])
  # Planet-to-star radius ratio:
  rprs  = rplanet / rstar
  mu.msg(verb, "OCON FLAG 10: {}, {}, {}".format(tstar, gstar, rprs))

  nfilters = len(ffile)  # Number of filters:

  # FINDME: Separate filter/stellar interpolation?
  # Get stellar model:
  starfl, starwn, tmodel, gmodel = w.readkurucz(kurucz, tstar, gstar)
  # Read and resample the filters:
  nifilter  = [] # Normalized interpolated filter
  istarfl   = [] # interpolated stellar flux
  wnindices = [] # wavenumber indices used in interpolation
  for i in np.arange(nfilters):
    # Read filter:
    filtwaven, filttransm = w.readfilter(ffile[i])
    # Check that filter boundaries lie within the spectrum wn range:
    if filtwaven[0] < specwn[0] or filtwaven[-1] > specwn[-1]:
      mu.exit(message="Wavenumber array ({:.2f} - {:.2f} cm-1) does not "
              "cover the filter[{:d}] wavenumber range ({:.2f} - {:.2f} "
              "cm-1).".format(specwn[0], specwn[-1], i, filtwaven[0],
                                                        filtwaven[-1]))

    # Resample filter and stellar spectrum:
    nifilt, strfl, wnind = w.resample(specwn, filtwaven, filttransm,
                                              starwn,    starfl)
    mu.msg(verb, "OCON FLAG 67: mean star flux: %.3e"%np.mean(strfl))
    nifilter.append(nifilt)
    istarfl.append(strfl)
    wnindices.append(wnind)

  # Allocate arrays for receiving and sending data to master:
  spectrum = np.zeros(nwave,    dtype='d')
  bandflux = np.zeros(nfilters, dtype='d')

  # Allocate array to receive parameters from MPI:
  params = np.zeros(npars, np.double)

  # ::::::  Main MCMC Loop  ::::::::::::::::::::::::::::::::::::::::::
  # ::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::

  while niter >= 0:
    niter -= 1
    # Receive parameters from MCMC:
    mu.comm_scatter(comm, params)
    #mu.msg(verb, "ICON FLAG 71: incon pars: {:s}".
    #             format(str(params).replace("\n", "")))

    # Input converter calculate the profiles:
    try:
      tprofile[:] = pt.PT_generator(pressure, params[0:nPT], PTargs)[::-1]
    except ValueError:
      mu.msg(verb, 'Input parameters give non-physical profile.')
      # FINDME: what to do here?

    # If the temperature goes out of bounds:
    if np.any(tprofile < Tmin) or np.any(tprofile > Tmax):
      print("Out of bounds")
      mu.comm_gather(comm, -np.ones(nfilters), MPI.DOUBLE)
      continue

    #mu.msg(verb, "T pars: \n{}\n".format(PTargs))
    mu.msg(verb-20, "Temperature profile: {}".format(tprofile))
    # Scale abundance profiles:
    for i in np.arange(nmolfit):
      m = imol[i]
      # Use variable as the log10:
      aprofiles[m] = abundances[:, m] * 10.0**params[nPT+nradfit+i]
    # Update H2, He abundances so sum(abundances) = 1.0 in each layer:
    q = 1.0 - np.sum(aprofiles[imetals], axis=0)
    aprofiles[iH2] = ratio * q / (1.0 + ratio)
    aprofiles[iHe] =         q / (1.0 + ratio)
    # print("qH2O: {}, Qmetals: {}, QH2: {}  p: {}".format(params[nPT],
    #                               q[50], profiles[iH2+1,50], profiles[:,50]))

    # Set the 'surface' level:
    if solution == "transit":
      trm.set_radius(params[nPT])

    if rank == 1:
      print("Iteration: {:05}".format(niter))
    # Let transit calculate the model spectrum:
    spectrum = trm.run_transit(profiles.flatten(), nwave)

    # Output converter band-integrate the spectrum:
    # Calculate the band-integrated intensity per filter:
    for i in np.arange(nfilters):
      if   solution == "eclipse":
        fluxrat = (spectrum[wnindices[i]]/istarfl[i]) * rprs*rprs
        bandflux[i] = w.bandintegrate(fluxrat, specwn,
                                      nifilter[i], wnindices[i])
      elif solution == "transit":
        bandflux[i] = w.bandintegrate(spectrum[wnindices[i]], specwn,
                                      nifilter[i], wnindices[i])

    # Send resutls back to MCMC:
    #mu.msg(verb, "OCON FLAG 95: Flux band integrated ({})".format(bandflux))
    #mu.msg(verb, "{}".format(params[nPT:]))
    mu.comm_gather(comm, bandflux, MPI.DOUBLE)
    #mu.msg(verb, "OCON FLAG 97: Sent results back to MCMC")

  # ::::::  End main Loop  :::::::::::::::::::::::::::::::::::::::::::
  # ::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::

  # Close communications and disconnect:
  mu.comm_disconnect(comm)
  mu.msg(verb, "FUNC FLAG 99: func out")

  # Close the transit communicators:
  trm.free_memory()
  mu.msg(verb, "FUNC FLAG OUT ~~ 100 ~~")
Пример #2
0
def main(comm):
    """
  Wrapper of modeling function for MCMC under MPI protocol.

  Modification History:
  ---------------------
  2014-04-19  patricio  Initial implementation.  [email protected]
  """
    # Parse arguments:
    cparser = argparse.ArgumentParser(
        description=__doc__,
        add_help=False,
        formatter_class=argparse.RawDescriptionHelpFormatter)
    # Add config file option:
    cparser.add_argument("-c",
                         "--config_file",
                         help="Configuration file",
                         metavar="FILE")
    # Remaining_argv contains all other command-line-arguments:
    args, remaining_argv = cparser.parse_known_args()

    # Get parameters from configuration file (if exists):
    cfile = args.config_file  # The configuration file
    if cfile:
        config = ConfigParser.SafeConfigParser()
        config.read([cfile])
        defaults = dict(config.items("MCMC"))
    else:
        defaults = {}
    parser = argparse.ArgumentParser(parents=[cparser])
    parser.add_argument("-f",
                        "--func",
                        dest="func",
                        type=mu.parray,
                        action="store",
                        default=None)
    parser.add_argument("-i",
                        "--indparams",
                        dest="indparams",
                        type=mu.parray,
                        action="store",
                        default=[])
    parser.set_defaults(**defaults)
    args2, unknown = parser.parse_known_args(remaining_argv)

    # Get indparams from configuration file:
    if args2.indparams != [] and os.path.isfile(args2.indparams[0]):
        indparams = mu.read2array(args2.indparams[0], square=False)

    # Get func from configuration file:
    if len(args2.func) == 3:
        sys.path.append(args2.func[2])
        exec('from %s import %s as func' % (args2.func[1], args2.func[0]))

    # Get the number of parameters and iterations from MPI:
    array1 = np.zeros(2, np.int)
    mu.comm_bcast(comm, array1)
    npars, niter = array1

    # Allocate array to receive parameters from MPI:
    params = np.zeros(npars, np.double)
    while niter >= 0:
        # Receive parameters from master:
        mu.comm_scatter(comm, params)
        # Evaluate model:
        fargs = [params] + indparams  # List of function's arguments
        model = func(*fargs)
        # Send resutls:
        mu.comm_gather(comm, model, MPI.DOUBLE)
        niter -= 1

    # Close communications and disconnect:
    mu.exit(comm)
Пример #3
0
def main(comm):
  """
  Wrapper of modeling function for MCMC under MPI protocol.

  Modification History:
  ---------------------
  2014-04-19  patricio  Initial implementation.  [email protected]
  2014-06-25  patricio  Added support for inner-MPI loop.
  2014-10-23  patricio  Removed inner-MPI loop.
  """
  # Parse arguments:
  cparser = argparse.ArgumentParser(description=__doc__, add_help=False,
                         formatter_class=argparse.RawDescriptionHelpFormatter)
  # Add config file option:
  cparser.add_argument("-c", "--config_file", 
                       help="Configuration file", metavar="FILE")
  # Remaining_argv contains all other command-line-arguments:
  args, remaining_argv = cparser.parse_known_args()

  # Get parameters from configuration file:
  cfile = args.config_file
  if cfile:
    config = ConfigParser.SafeConfigParser()
    config.read([cfile])
    defaults = dict(config.items("MCMC"))
  else:
    defaults = {}
  parser = argparse.ArgumentParser(parents=[cparser])
  parser.add_argument("-f", "--func",      dest="func",     type=mu.parray, 
                                           action="store",  default=None)
  parser.add_argument("-i", "--indparams", dest="indparams", type=mu.parray, 
                                           action="store",   default=[])
  parser.set_defaults(**defaults)
  args2, unknown = parser.parse_known_args(remaining_argv)

  # Add path to func:
  if len(args2.func) == 3:
    sys.path.append(args2.func[2])

  exec('from {:s} import {:s} as func'.format(args2.func[1], args2.func[0]))
  # Get indparams from configuration file:
  if args2.indparams != [] and os.path.isfile(args2.indparams[0]):
    indparams = mu.readbin(args2.indparams[0])


  # Get the number of parameters and iterations from MPI:
  array1 = np.zeros(2, np.int)
  mu.comm_bcast(comm, array1)
  npars, niter = array1

  # Allocate array to receive parameters from MPI:
  params = np.zeros(npars, np.double)

  # Main MCMC Loop:
  while niter >= 0:
    # Receive parameters from MCMC:
    mu.comm_scatter(comm, params)

    # Evaluate model:
    fargs = [params] + indparams  # List of function's arguments
    model = func(*fargs)

    # Send resutls:
    mu.comm_gather(comm, model, MPI.DOUBLE)
    niter -= 1

  # Close communications and disconnect:
  mu.comm_disconnect(comm)
Пример #4
0
def mcmc(data,             uncert=None,   func=None,     indparams=[],
         params=None,      pmin=None,     pmax=None,     stepsize=None,
         prior=None,       priorlow=None, priorup=None,  numit=10,
         nchains=10,       walk='demc',   wlike=False,   leastsq=True,
         chisqscale=False, grtest=True,   grexit=False,  burnin=0,
         thinning=1,       plots=False,   savefile=None, savemodel=None,
         comm=None,        resume=False,  log=None,      rms=False):
  """
  This beautiful piece of code runs a Markov-chain Monte Carlo algoritm.

  Parameters
  ----------
  data: 1D ndarray
     Dependent data fitted by func.
  uncert: 1D ndarray
     Uncertainty of data.
  func: callable or string-iterable
     The callable function that models data as:
        model = func(params, *indparams)
     Or an iterable (list, tuple, or ndarray) of 3 strings:
        (funcname, modulename, path)
     that specify the function name, function module, and module path.
     If the module is already in the python-path scope, path can be omitted.
  indparams: tuple
     Additional arguments required by func.
  params: 1D or 2D ndarray
     Set of initial fitting parameters for func.  If 2D, of shape
     (nparams, nchains), it is assumed that it is one set for each chain.
  pmin: 1D ndarray
     Lower boundaries of the posteriors.
  pmax: 1D ndarray
     Upper boundaries of the posteriors.
  stepsize: 1D ndarray
     Proposal jump scale.  If a values is 0, keep the parameter fixed.
     Negative values indicate a shared parameter (See Note 1).
  prior: 1D ndarray
     Parameter prior distribution means (See Note 2).
  priorlow: 1D ndarray
     Lower prior uncertainty values (See Note 2).
  priorup: 1D ndarray
     Upper prior uncertainty values (See Note 2).
  numit: Scalar
     Total number of iterations.
  nchains: Scalar
     Number of simultaneous chains to run.
  walk: String
     Random walk algorithm:
     - 'mrw':  Metropolis random walk.
     - 'demc': Differential Evolution Markov chain.
  wlike: Boolean
     If True, calculate the likelihood in a wavelet-base.  This requires
     three additional parameters (See Note 3).
  leastsq: Boolean
     Perform a least-square minimization before the MCMC run.
  chisqscale: Boolean
     Scale the data uncertainties such that the reduced chi-squared = 1.
  grtest: Boolean
     Run Gelman & Rubin test.
  grexit: Boolean
     Exit the MCMC loop if the MCMC satisfies GR two consecutive times.
  burnin: Scalar
     Burned-in (discarded) number of iterations at the beginning
     of the chains.
  thinning: Integer
     Thinning factor of the chains (use every thinning-th iteration) used
     in the GR test and plots.
  plots: Boolean
     If True plot parameter traces, pairwise-posteriors, and posterior
     histograms.
  savefile: String
     If not None, filename to store allparams (with np.save).
  savemodel: String
     If not None, filename to store the values of the evaluated function
     (with np.save).
  comm: MPI Communicator
     A communicator object to transfer data through MPI.
  resume: Boolean
     If True resume a previous run.
  log: FILE pointer
     File object to write log into.

  Returns
  -------
  allparams: 2D ndarray
     An array of shape (nfree, numit-nchains*burnin) with the MCMC
     posterior distribution of the fitting parameters.
  bestp: 1D ndarray
     Array of the best fitting parameters.

  Notes
  -----
  1.- To set one parameter equal to another, set its stepsize to the
      negative index in params (Starting the count from 1); e.g.: to set
      the second parameter equal to the first one, do: stepsize[1] = -1.
  2.- If any of the fitting parameters has a prior estimate, e.g.,
        param[i] = p0 +up/-low,
      with up and low the 1sigma uncertainties.  This information can be
      considered in the MCMC run by setting:
      prior[i]    = p0
      priorup[i]  = up
      priorlow[i] = low
      All three: prior, priorup, and priorlow must be set and, furthermore,
      priorup and priorlow must be > 0 to be considered as prior.
  3.- FINDME WAVELET LIKELIHOOD

  Examples
  --------
  >>> # See examples: https://github.com/pcubillos/MCcubed/tree/master/examples

  Previous (uncredited) developers
  --------------------------------
  Kevin Stevenson    UCF  [email protected]
  """

  mu.msg(1, "{:s}\n  Multi-Core Markov-Chain Monte Carlo (MC3).\n"
            "  Version {:d}.{:d}.{:d}.\n"
            "  Copyright (c) 2015-2016 Patricio Cubillos and collaborators.\n"
            "  MC3 is open-source software under the MIT license "
            "(see LICENSE).\n{:s}\n\n".
            format(mu.sep, ver.MC3_VER, ver.MC3_MIN, ver.MC3_REV, mu.sep), log)

  # Import the model function:
  if type(func) in [list, tuple, np.ndarray]:
    if func[0] != 'hack':
      if len(func) == 3:
        sys.path.append(func[2])
      exec('from %s import %s as func'%(func[1], func[0]))
  elif not callable(func):
    mu.error("'func' must be either, a callable, or an iterable (list, "
             "tuple, or ndarray) of strings with the model function, file, "
             "and path names.", log)

  if np.ndim(params) == 1:  # Force it to be 2D (one for each chain)
    params  = np.atleast_2d(params)
  nparams = len(params[0])  # Number of model params
  ndata   = len(data)       # Number of data values
  # Set default uncertainties:
  if uncert is None:
    uncert = np.ones(ndata)
  # Set default boundaries:
  if pmin is None:
    pmin = np.zeros(nparams) - np.inf
  if pmax is None:
    pmax = np.zeros(nparams) + np.inf
  # Set default stepsize:
  if stepsize is None:
    stepsize = 0.1 * np.abs(params[0])
  # Set prior parameter indices:
  if (prior is None) or (priorup is None) or (priorlow is None):
    prior   = priorup = priorlow = np.zeros(nparams)  # Zero arrays
  iprior = np.where(priorlow != 0)[0]
  ilog   = np.where(priorlow <  0)[0]

  # Check that initial values lie within the boundaries:
  if np.any(np.asarray(params) < pmin):
    mu.error("One or more of the initial-guess values:\n{:s}\n are smaller "
      "than their lower boundaries:\n{:s}".format(str(params), str(pmin)), log)
  if np.any(np.asarray(params) > pmax):
    mu.error("One or more of the initial-guess values:\n{:s}\n are greater "
      "than their higher boundaries:\n{:s}".format(str(params), str(pmax)), log)

  nfree     = np.sum(stepsize > 0)        # Number of free parameters
  chainsize = int(np.ceil(numit/nchains)) # Number of iterations per chain
  ifree     = np.where(stepsize > 0)[0]   # Free   parameter indices
  ishare    = np.where(stepsize < 0)[0]   # Shared parameter indices
  # Number of model parameters (excluding wavelet parameters):
  if wlike:
    mpars  = nparams - 3
  else:
    mpars  = nparams

  if chainsize < burnin:
    mu.error("The number of burned-in samples ({:d}) is greater than "
             "the number of iterations per chain ({:d}).".
             format(burnin, chainsize), log)

  # Intermediate steps to run GR test and print progress report:
  intsteps   = chainsize / 10

  # Allocate arrays with variables:
  numaccept  = np.zeros(nchains)          # Number of accepted proposal jumps
  outbounds  = np.zeros((nchains, nfree), np.int)    # Out of bounds proposals
  allparams  = np.zeros((nchains, nfree, chainsize)) # Parameter's record
  if savemodel is not None:
    allmodel = np.zeros((nchains, ndata, chainsize)) # Fit model

  if resume:
    oldparams = np.load(savefile)
    nold = np.shape(oldparams)[2] # Number of old-run iterations
    allparams = np.dstack((oldparams, allparams))
    if savemodel is not None:
      allmodel  = np.dstack((np.load(savemodel), allmodel))
    # Set params to the last-iteration state of the previous run:
    params = np.repeat(params, nchains, 0)
    params[:,ifree] = oldparams[:,:,-1]
  else:
    nold = 0

  # Set MPI flag:
  mpi = comm is not None

  if mpi:
    from mpi4py import MPI
    # Send sizes info to other processes:
    array1 = np.asarray([mpars, chainsize], np.int)
    mu.comm_bcast(comm, array1, MPI.INT)

  # DEMC parameters:
  gamma  = 2.4 / np.sqrt(2*nfree)
  gamma2 = 0.001  # Jump scale factor of support distribution

  # Least-squares minimization:
  if leastsq:
    fitargs = (params[0], func, data, uncert, indparams, stepsize, pmin, pmax,
               prior, priorlow, priorup)
    fitchisq, dummy = mf.modelfit(params[0,ifree], args=fitargs)
    fitbestp = np.copy(params[0, ifree])
    mu.msg(1, "Least-squares best-fitting parameters: \n{:s}\n\n".
              format(str(fitbestp)), log)

  # Replicate to make one set for each chain: (nchains, nparams):
  if np.shape(params)[0] != nchains:
    params = np.repeat(params, nchains, 0)
    # Start chains with an initial jump:
    for p in ifree:
      # For each free param, use a normal distribution: 
      params[1:, p] = np.random.normal(params[0, p], stepsize[p], nchains-1)
      # Stay within pmin and pmax boundaries:
      params[np.where(params[:, p] < pmin[p]), p] = pmin[p]
      params[np.where(params[:, p] > pmax[p]), p] = pmax[p]
  
  # Update shared parameters:
  for s in ishare:
    params[:, s] = params[:, -int(stepsize[s])-1]

  # Calculate chi-squared for model using current params:
  models = np.zeros((nchains, ndata))
  if mpi:
    # Scatter (send) parameters to func:
    mu.comm_scatter(comm, params[:,0:mpars].flatten(), MPI.DOUBLE)
    # Gather (receive) evaluated models:
    mpimodels = np.zeros(nchains*ndata, np.double)
    mu.comm_gather(comm, mpimodels)
    # Store them in models variable:
    models = np.reshape(mpimodels, (nchains, ndata))
  else:
    for c in np.arange(nchains):
      fargs = [params[c, 0:mpars]] + indparams  # List of function's arguments
      models[c] = func(*fargs)

  # Calculate chi-squared for each chain:
  currchisq = np.zeros(nchains)
  c2        = np.zeros(nchains)  # No-Jeffrey's chisq
  for c in np.arange(nchains):
    if wlike: # Wavelet-based likelihood (chi-squared, actually)
      currchisq[c], c2[c] = dwt.wlikelihood(params[c, mpars:], models[c]-data,
                 (params[c]-prior)[iprior], priorlow[iprior], priorlow[iprior])
    else:
      currchisq[c], c2[c] = cs.chisq(models[c], data, uncert,
                 (params[c]-prior)[iprior], priorlow[iprior], priorlow[iprior])

  # Scale data-uncertainties such that reduced chisq = 1:
  if chisqscale:
    chifactor = np.sqrt(np.amin(currchisq)/(ndata-nfree))
    uncert *= chifactor
    # Re-calculate chisq with the new uncertainties:
    for c in np.arange(nchains):
      if wlike: # Wavelet-based likelihood (chi-squared, actually)
        currchisq[c], c2[c] = dwt.wlikelihood(params[c,mpars:], models[c]-data,
                 (params[c]-prior)[iprior], priorlow[iprior], priorlow[iprior])
      else:
        currchisq[c], c2[c] = cs.chisq(models[c], data, uncert,
                 (params[c]-prior)[iprior], priorlow[iprior], priorlow[iprior])
    if leastsq:
      fitchisq = currchisq[0]

  # Get lowest chi-square and best fitting parameters:
  bestchisq = np.amin(c2)
  bestp     = np.copy(params[np.argmin(c2)])
  bestmodel = np.copy(models[np.argmin(c2)])

  if savemodel is not None:
    allmodel[:,:,0] = models

  # Set up the random walks:
  if   walk == "mrw":
    # Generate proposal jumps from Normal Distribution for MRW:
    mstep   = np.random.normal(0, stepsize[ifree], (chainsize, nchains, nfree))
  elif walk == "demc":
    # Support random distribution:
    support = np.random.normal(0, stepsize[ifree], (chainsize, nchains, nfree))
    # Generate indices for the chains such r[c] != c:
    r1 = np.random.randint(0, nchains-1, (nchains, chainsize))
    r2 = np.random.randint(0, nchains-1, (nchains, chainsize))
    for c in np.arange(nchains):
      r1[c][np.where(r1[c]==c)] = nchains-1
      r2[c][np.where(r2[c]==c)] = nchains-1

  # Uniform random distribution for the Metropolis acceptance rule:
  unif = np.random.uniform(0, 1, (chainsize, nchains))

  # Proposed iteration parameters and chi-square (per chain):
  nextp     = np.copy(params)    # Proposed parameters
  nextchisq = np.zeros(nchains)  # Chi square of nextp 

  # Gelman-Rubin exit flag:
  grflag = False

  # ::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::
  # Start loop:
  mu.msg(1, "Start MCMC chains  ({:s})".format(time.ctime()), log)
  for i in np.arange(chainsize):
    # Proposal jump:
    if   walk == "mrw":
      jump = mstep[i]
    elif walk == "demc":
      jump = (gamma  * (params[r1[:,i]]-params[r2[:,i]])[:,ifree] +
              gamma2 * support[i]                                 )
    # Propose next point:
    nextp[:,ifree] = params[:,ifree] + jump

    # Check it's within boundaries: 
    outpars = np.asarray(((nextp < pmin) | (nextp > pmax))[:,ifree])
    outflag  = np.any(outpars, axis=1)
    outbounds += ((nextp < pmin) | (nextp > pmax))[:,ifree]
    for p in ifree:
      nextp[np.where(nextp[:, p] < pmin[p]), p] = pmin[p]
      nextp[np.where(nextp[:, p] > pmax[p]), p] = pmax[p]

    # Update shared parameters:
    for s in ishare:
      nextp[:, s] = nextp[:, -int(stepsize[s])-1]

    # Evaluate the models for the proposed parameters:
    if mpi:
      mu.comm_scatter(comm, nextp[:,0:mpars].flatten(), MPI.DOUBLE)
      mu.comm_gather(comm, mpimodels)
      models = np.reshape(mpimodels, (nchains, ndata))
    else:
      for c in np.where(~outflag)[0]:
        fargs = [nextp[c, 0:mpars]] + indparams  # List of function's arguments
        models[c] = func(*fargs)

    # Calculate chisq:
    for c in np.where(~outflag)[0]:
      if wlike: # Wavelet-based likelihood (chi-squared, actually)
        nextchisq[c], c2[c] = dwt.wlikelihood(nextp[c,mpars:], models[c]-data,
                 (nextp[c]-prior)[iprior], priorlow[iprior], priorlow[iprior])
      else:
        nextchisq[c], c2[c] = cs.chisq(models[c], data, uncert,
                 (nextp[c]-prior)[iprior], priorlow[iprior], priorlow[iprior])

    # Reject out-of-bound jumps:
    nextchisq[np.where(outflag)] = np.inf
    # Evaluate which steps are accepted and update values:
    accept = np.exp(0.5 * (currchisq - nextchisq))
    accepted = accept >= unif[i]
    if i >= burnin:
      numaccept += accepted
    # Update params and chi square:
    params   [accepted] = nextp    [accepted]
    currchisq[accepted] = nextchisq[accepted]

    # Check lowest chi-square:
    if np.amin(c2) < bestchisq:
      bestp     = np.copy(params[np.argmin(c2)])
      bestmodel = np.copy(models[np.argmin(c2)])
      bestchisq = np.amin(c2)

    # Store current iteration values:
    allparams[:,:,i+nold] = params[:, ifree]
    if savemodel is not None:
      models[~accepted] = allmodel[~accepted,:,i+nold-1]
      allmodel[:,:,i+nold] = models
  
    # Print intermediate info:
    if ((i+1) % intsteps == 0) and (i > 0):
      mu.progressbar((i+1.0)/chainsize, log)
      mu.msg(1, "Out-of-bound Trials:\n {:s}".
                 format(np.sum(outbounds, axis=0)), log)
      mu.msg(1, "Best Parameters:   (chisq={:.4f})\n{:s}".
                 format(bestchisq, str(bestp)), log)

      # Gelman-Rubin statistic:
      if grtest and (i+nold) > burnin:
        psrf = gr.convergetest(allparams[:, :, burnin:i+nold+1:thinning])
        mu.msg(1, "Gelman-Rubin statistic for free parameters:\n{:s}".
                  format(psrf), log)
        if np.all(psrf < 1.01):
          mu.msg(1, "All parameters have converged to within 1% of unity.", log)
          # End the MCMC if all parameters satisfy GR two consecutive times:
          if grexit and grflag:
            # Let the workers know that the MCMC is stopping:
            if mpi:
              endflag = np.tile(np.inf, nchains*mpars)
              mu.comm_scatter(comm, endflag, MPI.DOUBLE)
            break
          grflag = True
        else:
          grflag = False
      # Save current results:
      if savefile is not None:
        np.save(savefile, allparams[:,:,0:i+nold])
      if savemodel is not None:
        np.save(savemodel, allmodel[:,:,0:i+nold])

  # Stack together the chains:
  chainlen = nold + i+1
  allstack = allparams[0, :, burnin:chainlen]
  for c in np.arange(1, nchains):
    allstack = np.hstack((allstack, allparams[c, :, burnin:chainlen]))
  # And the models:
  if savemodel is not None:
    modelstack = allmodel[0,:,burnin:chainlen]
    for c in np.arange(1, nchains):
      modelstack = np.hstack((modelstack, allmodel[c, :, burnin:chainlen]))

  # Print out Summary:
  mu.msg(1, "\nFin, MCMC Summary:\n------------------", log)

  nsample   = (i+1-burnin)*nchains
  ntotal    = np.size(allstack[0])
  BIC       = bestchisq + nfree*np.log(ndata)
  redchisq  = bestchisq/(ndata-nfree)
  sdr       = np.std(bestmodel-data)

  fmtlen = len(str(ntotal))
  mu.msg(1, "Burned in iterations per chain: {:{}d}".
             format(burnin,   fmtlen), log, 1)
  mu.msg(1, "Number of iterations per chain: {:{}d}".
             format(i+1, fmtlen), log, 1)
  mu.msg(1, "MCMC sample size:               {:{}d}".
             format(nsample,  fmtlen), log, 1)
  mu.msg(resume, "Total MCMC sample size:         {:{}d}".
             format(ntotal, fmtlen), log, 1)
  mu.msg(1, "Acceptance rate:   {:.2f}%\n ".
             format(np.sum(numaccept)*100.0/nsample), log, 1)

  meanp   = np.mean(allstack, axis=1) # Parameters mean
  uncertp = np.std(allstack,  axis=1) # Parameter standard deviation
  mu.msg(1, "Best-fit params    Uncertainties   Signal/Noise       Sample "
            "Mean", log, 1)
  for i in np.arange(nfree):
    mu.msg(1, "{: 15.7e}  {: 15.7e}   {:12.2f}   {: 15.7e}".
               format(bestp[ifree][i], uncertp[i],
                      np.abs(bestp[ifree][i])/uncertp[i], meanp[i]), log, 1)

  if leastsq and np.any(np.abs((bestp[ifree]-fitbestp)/fitbestp) > 1e-08):
    np.set_printoptions(precision=8)
    mu.warning("MCMC found a better fit than the minimizer:\n"
               " MCMC best-fitting parameters:       (chisq={:.8g})\n  {:s}\n"
               " Minimizer best-fitting parameters:  (chisq={:.8g})\n"
               "  {:s}".format(bestchisq, str(bestp[ifree]), 
                               fitchisq,  str(fitbestp)), log)

  fmtl = len("%.4f"%BIC)  # Length of string formatting
  mu.msg(1, " ", log)
  if chisqscale:
    mu.msg(1, "sqrt(reduced chi-squared) factor: {:{}.4f}".
               format(chifactor, fmtl), log, 1)
  mu.msg(1,   "Best-parameter's chi-squared:     {:{}.4f}".
               format(bestchisq, fmtl), log, 1)
  mu.msg(1,   "Bayesian Information Criterion:   {:{}.4f}".
               format(BIC,       fmtl), log, 1)
  mu.msg(1,   "Reduced chi-squared:              {:{}.4f}".
               format(redchisq,  fmtl), log, 1)
  mu.msg(1,   "Standard deviation of residuals:  {:.6g}\n".format(sdr), log, 1)


  if rms:
    rms, rmse, stderr, bs = ta.binrms(bestmodel-data)

  if plots:
    print("Plotting figures ...")
    # Extract filename from savefile:
    if savefile is not None:
      if savefile.rfind(".") == -1:
        fname = savefile[savefile.rfind("/")+1:] # Cut out file extention.
      else:
        fname = savefile[savefile.rfind("/")+1:savefile.rfind(".")]
    else:
      fname = "MCMC"
    # Trace plot:
    mp.trace(allstack,     thinning=thinning, savefile=fname+"_trace.png",
             sep=np.size(allstack[0])/nchains)
    # Pairwise posteriors:
    mp.pairwise(allstack,  thinning=thinning, savefile=fname+"_pairwise.png")
    # Histograms:
    mp.histogram(allstack, thinning=thinning, savefile=fname+"_posterior.png")
    # RMS vs bin size:
    if rms:
      mp.RMS(bs, rms, stderr, rmse, binstep=len(bs)/500+1,
                                              savefile=fname+"_RMS.png")
    if indparams != [] and np.size(indparams[0]) == ndata:
      mp.modelfit(data, uncert, indparams[0], bestmodel,
                                              savefile=fname+"_model.png")

  # Save definitive results:
  if savefile is not None:
    np.save(savefile,  allparams[:,:,:chainlen])
  if savemodel is not None:
    np.save(savemodel, allmodel [:,:,:chainlen])

  return allstack, bestp
Пример #5
0
def mcmc(data,         uncert=None,      func=None,     indparams=[],
         params=None,  pmin=None,        pmax=None,     stepsize=None,
         prior=None,   priorlow=None,    priorup=None,
         numit=10,     nchains=10,       walk='demc',   wlike=False,
         leastsq=True, chisqscale=False, grtest=True,   burnin=0,
         thinning=1,   plots=False,      savefile=None, savemodel=None,
         comm=None,    resume=False,     log=None,      rms=False):
  """
  This beautiful piece of code runs a Markov-chain Monte Carlo algoritm.

  Parameters:
  -----------
  data: 1D ndarray
     Dependent data fitted by func.
  uncert: 1D ndarray
     Uncertainty of data.
  func: callable or string-iterable
     The callable function that models data as:
        model = func(params, *indparams)
     Or an iterable (list, tuple, or ndarray) of 3 strings:
        (funcname, modulename, path)
     that specify the function name, function module, and module path.
     If the module is already in the python-path scope, path can be omitted.
  indparams: tuple
     Additional arguments required by func.
  params: 1D or 2D ndarray
     Set of initial fitting parameters for func.  If 2D, of shape
     (nparams, nchains), it is assumed that it is one set for each chain.
  pmin: 1D ndarray
     Lower boundaries of the posteriors.
  pmax: 1D ndarray
     Upper boundaries of the posteriors.
  stepsize: 1D ndarray
     Proposal jump scale.  If a values is 0, keep the parameter fixed.
     Negative values indicate a shared parameter (See Note 1).
  prior: 1D ndarray
     Parameter prior distribution means (See Note 2).
  priorlow: 1D ndarray
     Lower prior uncertainty values (See Note 2).
  priorup: 1D ndarray
     Upper prior uncertainty values (See Note 2).
  numit: Scalar
     Total number of iterations.
  nchains: Scalar
     Number of simultaneous chains to run.
  walk: String
     Random walk algorithm:
     - 'mrw':  Metropolis random walk.
     - 'demc': Differential Evolution Markov chain.
  wlike: Boolean
     If True, calculate the likelihood in a wavelet-base.  This requires
     three additional parameters (See Note 3).
  leastsq: Boolean
     Perform a least-square minimization before the MCMC run.
  chisqscale: Boolean
     Scale the data uncertainties such that the reduced chi-squared = 1.
  grtest: Boolean
     Run Gelman & Rubin test.
  burnin: Scalar
     Burned-in (discarded) number of iterations at the beginning
     of the chains.
  thinning: Integer
     Thinning factor of the chains (use every thinning-th iteration) used
     in the GR test and plots.
  plots: Boolean
     If True plot parameter traces, pairwise-posteriors, and posterior
     histograms.
  savefile: String
     If not None, filename to store allparams (with np.save).
  savemodel: String
     If not None, filename to store the values of the evaluated function
     (with np.save).
  comm: MPI Communicator
     A communicator object to transfer data through MPI.
  resume: Boolean
     If True resume a previous run.
  log: FILE pointer
     File object to write log into.

  Returns:
  --------
  allparams: 2D ndarray
     An array of shape (nfree, numit-nchains*burnin) with the MCMC
     posterior distribution of the fitting parameters.
  bestp: 1D ndarray
     Array of the best fitting parameters.

  Notes:
  ------
  1.- To set one parameter equal to another, set its stepsize to the
      negative index in params (Starting the count from 1); e.g.: to set
      the second parameter equal to the first one, do: stepsize[1] = -1.
  2.- If any of the fitting parameters has a prior estimate, e.g.,
        param[i] = p0 +up/-low,
      with up and low the 1sigma uncertainties.  This information can be
      considered in the MCMC run by setting:
      prior[i]    = p0
      priorup[i]  = up
      priorlow[i] = low
      All three: prior, priorup, and priorlow must be set and, furthermore,
      priorup and priorlow must be > 0 to be considered as prior.
  3.- FINDME WAVELET LIKELIHOOD

  Examples:
  ---------
  >>> # See examples: https://github.com/pcubillos/MCcubed/tree/master/examples

  Developers:
  -----------
  Kevin Stevenson    UCF  [email protected]
  Patricio Cubillos  UCF  [email protected]

  Modification History:
  ---------------------
    2008-05-02  kevin     Initial implementation
    2008-06-21  kevin     Finished updating
    2009-11-01  kevin     Updated for multi events:
    2010-06-09  kevin     Updated for ipspline, nnint & bilinint
    2011-07-06  kevin     Updated for Gelman-Rubin statistic
    2011-07-22  kevin     Added principal component analysis
    2011-10-11  kevin     Added priors
    2012-09-03  patricio  Added Differential Evolution MC. Documented.
    2013-01-31  patricio  Modified for general purposes.
    2013-02-21  patricio  Added support distribution for DEMC.
    2014-03-31  patricio  Modified to be completely agnostic of the
                          fitting function, updated documentation.
    2014-04-17  patricio  Revamped use of 'func': no longer requires a
                          wrapper.  Alternatively, can take a string list with
                          the function, module, and path names.
    2014-04-19  patricio  Added savefile, thinning, plots, and mpi arguments.
    2014-05-04  patricio  Added Summary print out.
    2014-05-09  patricio  Added Wavelet-likelihood calculation.
    2014-05-09  patricio  Changed figure types from pdf to png, because it's
                          much faster.
    2014-05-26  patricio  Changed mpi bool argument by comm.  Re-engineered
                          MPI communications to make direct calls to func.
    2014-06-09  patricio  Fixed glitch with leastsq+informative priors.
    2014-10-17  patricio  Added savemodel argument.
    2014-10-23  patricio  Added support for func hack.
    2015-02-04  patricio  Added resume argument.
    2015-05-15  patricio  Added log argument.
  """

  # Import the model function:
  if type(func) in [list, tuple, np.ndarray]:
    if func[0] != 'hack':
      if len(func) == 3:
        sys.path.append(func[2])
      exec('from %s import %s as func'%(func[1], func[0]))
  elif not callable(func):
    mu.error("'func' must be either, a callable, or an iterable (list, "
             "tuple, or ndarray) of strings with the model function, file, "
             "and path names.", log)

  if np.ndim(params) == 1:  # Force it to be 2D (one for each chain)
    params  = np.atleast_2d(params)
  nparams = len(params[0])  # Number of model params
  ndata   = len(data)       # Number of data values
  # Set default uncertainties:
  if uncert is None:
    uncert = np.ones(ndata)
  # Set default boundaries:
  if pmin is None:
    pmin = np.zeros(nparams) - np.inf
  if pmax is None:
    pmax = np.zeros(nparams) + np.inf
  # Set default stepsize:
  if stepsize is None:
    stepsize = 0.1 * np.abs(params[0])
  # Set prior parameter indices:
  if (prior is None) or (priorup is None) or (priorlow is None):
    prior   = priorup = priorlow = np.zeros(nparams)  # Zero arrays
  iprior = np.where(priorlow != 0)[0]
  ilog   = np.where(priorlow <  0)[0]

  nfree    = np.sum(stepsize > 0)        # Number of free parameters
  chainlen = int(np.ceil(numit/nchains)) # Number of iterations per chain
  ifree    = np.where(stepsize > 0)[0]   # Free   parameter indices
  ishare   = np.where(stepsize < 0)[0]   # Shared parameter indices
  # Number of model parameters (excluding wavelet parameters):
  if wlike:
    mpars  = nparams - 3
  else:
    mpars  = nparams

  # Intermediate steps to run GR test and print progress report:
  intsteps   = chainlen / 10

  # Allocate arrays with variables:
  numaccept  = np.zeros(nchains)          # Number of accepted proposal jumps
  outbounds  = np.zeros((nchains, nfree), np.int)   # Out of bounds proposals
  allparams  = np.zeros((nchains, nfree, chainlen)) # Parameter's record
  if savemodel is not None:
    allmodel = np.zeros((nchains, ndata, chainlen)) # Fit model

  if resume:
    oldparams = np.load(savefile)
    nold = np.shape(oldparams)[2] # Number of old-run iterations
    allparams = np.dstack((oldparams, allparams))
    if savemodel is not None:
      allmodel  = np.dstack((np.load(savemodel), allmodel))
    # Set params to the last-iteration state of the previous run:
    params = np.repeat(params, nchains, 0)
    params[:,ifree] = oldparams[:,:,-1]
  else:
    nold = 0

  # Set MPI flag:
  mpi = comm is not None

  if mpi:
    from mpi4py import MPI
    # Send sizes info to other processes:
    array1 = np.asarray([mpars, chainlen], np.int)
    mu.comm_bcast(comm, array1, MPI.INT)

  # DEMC parameters:
  gamma  = 2.4 / np.sqrt(2*nfree)
  gamma2 = 0.001  # Jump scale factor of support distribution

  # Least-squares minimization:
  if leastsq:
    fitargs = (params[0], func, data, uncert, indparams, stepsize, pmin, pmax,
               prior, priorlow, priorup)
    fitchisq, dummy = mf.modelfit(params[0,ifree], args=fitargs)
    fitbestp = np.copy(params[0, ifree])
    mu.msg(1, "Least-squares best-fitting parameters: \n{:s}\n\n".
              format(str(fitbestp)), log)

  # Replicate to make one set for each chain: (nchains, nparams):
  if np.shape(params)[0] != nchains:
    params = np.repeat(params, nchains, 0)
    # Start chains with an initial jump:
    for p in ifree:
      # For each free param, use a normal distribution: 
      params[1:, p] = np.random.normal(params[0, p], stepsize[p], nchains-1)
      # Stay within pmin and pmax boundaries:
      params[np.where(params[:, p] < pmin[p]), p] = pmin[p]
      params[np.where(params[:, p] > pmax[p]), p] = pmax[p]
  
  # Update shared parameters:
  for s in ishare:
    params[:, s] = params[:, -int(stepsize[s])-1]

  # Calculate chi-squared for model using current params:
  models = np.zeros((nchains, ndata))
  if mpi:
    # Scatter (send) parameters to func:
    mu.comm_scatter(comm, params[:,0:mpars].flatten(), MPI.DOUBLE)
    # Gather (receive) evaluated models:
    mpimodels = np.zeros(nchains*ndata, np.double)
    mu.comm_gather(comm, mpimodels)
    # Store them in models variable:
    models = np.reshape(mpimodels, (nchains, ndata))
  else:
    for c in np.arange(nchains):
      fargs = [params[c, 0:mpars]] + indparams  # List of function's arguments
      models[c] = func(*fargs)

  # Calculate chi-squared for each chain:
  currchisq = np.zeros(nchains)
  c2        = np.zeros(nchains)  # No-Jeffrey's chisq
  for c in np.arange(nchains):
    if wlike: # Wavelet-based likelihood (chi-squared, actually)
      currchisq[c], c2[c] = dwt.wlikelihood(params[c, mpars:], models[c]-data,
                 (params[c]-prior)[iprior], priorlow[iprior], priorlow[iprior])
    else:
      currchisq[c], c2[c] = cs.chisq(models[c], data, uncert,
                 (params[c]-prior)[iprior], priorlow[iprior], priorlow[iprior])

  # Scale data-uncertainties such that reduced chisq = 1:
  if chisqscale:
    chifactor = np.sqrt(np.amin(currchisq)/(ndata-nfree))
    uncert *= chifactor
    # Re-calculate chisq with the new uncertainties:
    for c in np.arange(nchains):
      if wlike: # Wavelet-based likelihood (chi-squared, actually)
        currchisq[c], c2[c] = dwt.wlikelihood(params[c,mpars:], models[c]-data,
                 (params[c]-prior)[iprior], priorlow[iprior], priorlow[iprior])
      else:
        currchisq[c], c2[c] = cs.chisq(models[c], data, uncert,
                 (params[c]-prior)[iprior], priorlow[iprior], priorlow[iprior])
    if leastsq:
      fitchisq = currchisq[0]

  # Get lowest chi-square and best fitting parameters:
  bestchisq = np.amin(c2)
  bestp     = np.copy(params[np.argmin(c2)])
  bestmodel = np.copy(models[np.argmin(c2)])

  if savemodel is not None:
    allmodel[:,:,0] = models

  # Set up the random walks:
  if   walk == "mrw":
    # Generate proposal jumps from Normal Distribution for MRW:
    mstep   = np.random.normal(0, stepsize[ifree], (chainlen, nchains, nfree))
  elif walk == "demc":
    # Support random distribution:
    support = np.random.normal(0, stepsize[ifree], (chainlen, nchains, nfree))
    # Generate indices for the chains such r[c] != c:
    r1 = np.random.randint(0, nchains-1, (nchains, chainlen))
    r2 = np.random.randint(0, nchains-1, (nchains, chainlen))
    for c in np.arange(nchains):
      r1[c][np.where(r1[c]==c)] = nchains-1
      r2[c][np.where(r2[c]==c)] = nchains-1

  # Uniform random distribution for the Metropolis acceptance rule:
  unif = np.random.uniform(0, 1, (chainlen, nchains))

  # Proposed iteration parameters and chi-square (per chain):
  nextp     = np.copy(params)    # Proposed parameters
  nextchisq = np.zeros(nchains)  # Chi square of nextp 

  # Start loop:
  mu.msg(1, "Start MCMC chains  ({:s})".format(time.ctime()), log)
  for i in np.arange(chainlen):
    # Proposal jump:
    if   walk == "mrw":
      jump = mstep[i]
    elif walk == "demc":
      jump = (gamma  * (params[r1[:,i]]-params[r2[:,i]])[:,ifree] +
              gamma2 * support[i]                                 )
    # Propose next point:
    nextp[:,ifree] = params[:,ifree] + jump

    # Check it's within boundaries: 
    outpars = np.asarray(((nextp < pmin) | (nextp > pmax))[:,ifree])
    outflag  = np.any(outpars, axis=1)
    outbounds += ((nextp < pmin) | (nextp > pmax))[:,ifree]
    for p in ifree:
      nextp[np.where(nextp[:, p] < pmin[p]), p] = pmin[p]
      nextp[np.where(nextp[:, p] > pmax[p]), p] = pmax[p]

    # Update shared parameters:
    for s in ishare:
      nextp[:, s] = nextp[:, -int(stepsize[s])-1]

    # Evaluate the models for the proposed parameters:
    if mpi:
      mu.comm_scatter(comm, nextp[:,0:mpars].flatten(), MPI.DOUBLE)
      mu.comm_gather(comm, mpimodels)
      models = np.reshape(mpimodels, (nchains, ndata))
    else:
      for c in np.where(~outflag)[0]:
        fargs = [nextp[c, 0:mpars]] + indparams  # List of function's arguments
        models[c] = func(*fargs)

    # Calculate chisq:
    for c in np.where(~outflag)[0]:
      if wlike: # Wavelet-based likelihood (chi-squared, actually)
        nextchisq[c], c2[c] = dwt.wlikelihood(nextp[c,mpars:], models[c]-data,
                 (nextp[c]-prior)[iprior], priorlow[iprior], priorlow[iprior])
      else:
        nextchisq[c], c2[c] = cs.chisq(models[c], data, uncert,
                 (nextp[c]-prior)[iprior], priorlow[iprior], priorlow[iprior])

    # Reject out-of-bound jumps:
    nextchisq[np.where(outflag)] = np.inf
    # Evaluate which steps are accepted and update values:
    accept = np.exp(0.5 * (currchisq - nextchisq))
    accepted = accept >= unif[i]
    if i >= burnin:
      numaccept += accepted
    # Update params and chi square:
    params   [accepted] = nextp    [accepted]
    currchisq[accepted] = nextchisq[accepted]

    # Check lowest chi-square:
    if np.amin(c2) < bestchisq:
      bestp     = np.copy(params[np.argmin(c2)])
      bestmodel = np.copy(models[np.argmin(c2)])
      bestchisq = np.amin(c2)

    # Store current iteration values:
    allparams[:,:,i+nold] = params[:, ifree]
    if savemodel is not None:
      models[~accepted] = allmodel[~accepted,:,i+nold-1]
      allmodel[:,:,i+nold] = models
  
    # Print intermediate info:
    if ((i+1) % intsteps == 0) and (i > 0):
      mu.progressbar((i+1.0)/chainlen, log)
      mu.msg(1, "Out-of-bound Trials:\n {:s}".
                 format(np.sum(outbounds, axis=0)), log)
      mu.msg(1, "Best Parameters:   (chisq={:.4f})\n{:s}".
                 format(bestchisq, str(bestp)), log)

      # Gelman-Rubin statistic:
      if grtest and (i+nold) > burnin:
        psrf = gr.convergetest(allparams[:, :, burnin:i+nold+1:thinning])
        mu.msg(1, "Gelman-Rubin statistic for free parameters:\n{:s}".
                  format(psrf), log)
        if np.all(psrf < 1.01):
          mu.msg(1, "All parameters have converged to within 1% of unity.", log)
      # Save current results:
      if savefile is not None:
        np.save(savefile, allparams[:,:,0:i+nold])
      if savemodel is not None:
        np.save(savemodel, allmodel[:,:,0:i+nold])

  # Stack together the chains:
  allstack = allparams[0, :, burnin:]
  for c in np.arange(1, nchains):
    allstack = np.hstack((allstack, allparams[c, :, burnin:]))
  # And the models:
  if savemodel is not None:
    modelstack = allmodel[0,:,burnin:]
    for c in np.arange(1, nchains):
      modelstack = np.hstack((modelstack, allmodel[c, :, burnin:]))

  # Print out Summary:
  mu.msg(1, "\nFin, MCMC Summary:\n------------------", log)

  nsample   = (chainlen-burnin)*nchains # This sample
  ntotal    = (nold+chainlen-burnin)*nchains
  BIC       = bestchisq + nfree*np.log(ndata)
  redchisq  = bestchisq/(ndata-nfree)
  sdr       = np.std(bestmodel-data)

  fmtlen = len(str(ntotal))
  mu.msg(1, "Burned in iterations per chain: {:{}d}".
             format(burnin,   fmtlen), log, 1)
  mu.msg(1, "Number of iterations per chain: {:{}d}".
             format(chainlen, fmtlen), log, 1)
  mu.msg(1, "MCMC sample size:               {:{}d}".
             format(nsample,  fmtlen), log, 1)
  mu.msg(resume, "Total MCMC sample size:         {:{}d}".
             format(ntotal, fmtlen), log, 1)
  mu.msg(1, "Acceptance rate:   {:.2f}%\n ".
             format(np.sum(numaccept)*100.0/nsample), log, 1)

  meanp   = np.mean(allstack, axis=1) # Parameters mean
  uncertp = np.std(allstack,  axis=1) # Parameter standard deviation
  mu.msg(1, "Best-fit params    Uncertainties   Signal/Noise       Sample "
            "Mean", log, 1)
  for i in np.arange(nfree):
    mu.msg(1, "{: 15.7e}  {: 15.7e}   {:12.2f}   {: 15.7e}".
               format(bestp[ifree][i], uncertp[i],
                      np.abs(bestp[ifree][i])/uncertp[i], meanp[i]), log, 1)

  if leastsq and np.any(np.abs((bestp[ifree]-fitbestp)/fitbestp) > 1e-08):
    np.set_printoptions(precision=8)
    mu.warning("MCMC found a better fit than the minimizer:\n"
               " MCMC best-fitting parameters:       (chisq={:.8g})\n  {:s}\n"
               " Minimizer best-fitting parameters:  (chisq={:.8g})\n"
               "  {:s}".format(bestchisq, str(bestp[ifree]), 
                               fitchisq,  str(fitbestp)), log)

  fmtl = len("%.4f"%BIC)  # Length of string formatting
  mu.msg(1, " ", log)
  if chisqscale:
    mu.msg(1, "sqrt(reduced chi-squared) factor: {:{}.4f}".
               format(chifactor, fmtl), log, 1)
  mu.msg(1,   "Best-parameter's chi-squared:     {:{}.4f}".
               format(bestchisq, fmtl), log, 1)
  mu.msg(1,   "Bayesian Information Criterion:   {:{}.4f}".
               format(BIC,       fmtl), log, 1)
  mu.msg(1,   "Reduced chi-squared:              {:{}.4f}".
               format(redchisq,  fmtl), log, 1)
  mu.msg(1,   "Standard deviation of residuals:  {:.6g}\n".format(sdr), log, 1)


  if rms:
    rms, rmse, stderr, bs = ta.binrms(bestmodel-data)

  if plots:
    print("Plotting figures ...")
    # Extract filename from savefile:
    if savefile is not None:
      if savefile.rfind(".") == -1:
        fname = savefile[savefile.rfind("/")+1:] # Cut out file extention.
      else:
        fname = savefile[savefile.rfind("/")+1:savefile.rfind(".")]
    else:
      fname = "MCMC"
    # Trace plot:
    mp.trace(allstack,     thinning=thinning, savefile=fname+"_trace.png",
             sep=np.size(allstack[0])/nchains)
    # Pairwise posteriors:
    mp.pairwise(allstack,  thinning=thinning, savefile=fname+"_pairwise.png")
    # Histograms:
    mp.histogram(allstack, thinning=thinning, savefile=fname+"_posterior.png")
    # RMS vs bin size:
    if rms:
      mp.RMS(bs, rms, stderr, rmse, binstep=len(bs)/500+1,
                                              savefile=fname+"_RMS.png")
    if indparams != [] and np.size(indparams[0]) == ndata:
      mp.modelfit(data, uncert, indparams[0], bestmodel,
                                              savefile=fname+"_model.png")

  # Save definitive results:
  if savefile is not None:
    np.save(savefile,  allparams)
  if savemodel is not None:
    np.save(savemodel, allmodel)

  return allstack, bestp
Пример #6
0
def main(cfile=None):
  """
  Multi-Processor Markov-Chain Monte Carlo (MPMC)

  This code calls MCMC to work under an MPI multiprocessor protocol or
  single-thread mode.  When using MPI it will launch one CPU per MCMC chain
  to work in parallel.

  Parameters:
  -----------
  cfile: String
     Filename of a configuration file.

  Modification History:
  ---------------------
  2014-04-19  patricio  Initial implementation.  [email protected]
  2014-05-04  patricio  Added cfile argument for Interpreter support.
  """
  # Parse arguments:
  cparser = argparse.ArgumentParser(description=__doc__, add_help=False,
                         formatter_class=argparse.RawDescriptionHelpFormatter)
  # Add config file option:
  cparser.add_argument("-c", "--config_file",
                       help="Specify config file", metavar="FILE")
  # Remaining_argv contains all other command-line-arguments:
  args, remaining_argv = cparser.parse_known_args()

  # Take configuration file from command-line if not given as an argument:
  if cfile is None:
    cfile = args.config_file

  # Incorrect configuration file name:
  if cfile is not None and not os.path.isfile(cfile):
    mu.exit(None, message="Configuration file: '%s' not found."%cfile)
  if cfile:
    config = ConfigParser.SafeConfigParser()
    config.read([cfile])
    defaults = dict(config.items("MCMC"))
  else:
    defaults = {}
  parser = argparse.ArgumentParser(parents=[cparser], add_help=False)
  parser.add_argument("-x", "--nchains",   dest="nchains", type=int,
                                           action="store", default=10)
  parser.add_argument(       "--mpi",      dest="mpi",     type=eval,
                    help="Run under MPI multiprocessing [default %(default)s]",
                     action="store", default=False)
  parser.add_argument("-T", "--tracktime", dest="tractime", action="store_true")
  parser.add_argument("-h", "--help",      dest="help", action="store_true")

  parser.set_defaults(**defaults)
  args2, unknown = parser.parse_known_args(remaining_argv)

  # The number of processors is the number of chains:
  nprocs    = args2.nchains
  mpi       = args2.mpi
  # Hidden feature to track the execution of the time per loop for MPI:
  tracktime = args2.tractime

  # Get source dir:
  mcfile = mc.__file__
  iright = mcfile.rfind('/')
  if iright == -1:
    sdir = "."
  else:
    sdir = mcfile[:iright]

  # If asked for help:
  if args2.help:
    subprocess.call([sdir + "/mcmc.py --help"], shell=True)
    sys.exit(0)

  if not mpi:
    subprocess.call([sdir+"/mcmc.py " + " ".join(sys.argv[1:])], shell=True)
  else:
    if tracktime:
      start_mpi = timeit.default_timer()
    # Call MCMC:
    args = [sdir+"/mcmc.py", "-c"+cfile] + remaining_argv
    comm1 = MPI.COMM_SELF.Spawn(sys.executable, args=args, maxprocs=1)

    # Get OK flag from MCMC:
    abort = np.array([0])
    mu.comm_gather(comm1, abort)
    if abort[0]:
      mu.exit(comm1)

    # Call wrapper of model function:
    args = [sdir+"/func.py", "-c"+cfile] + remaining_argv
    comm2 = MPI.COMM_SELF.Spawn(sys.executable, args=args, maxprocs=nprocs)

    # MPI get sizes from MCMC:
    array1 = np.zeros(3, np.int)
    mu.comm_gather(comm1, array1)
    npars, ndata, niter = array1
    # MPI Broadcast to workers:
    mu.comm_bcast(comm2, np.asarray([npars, niter], np.int), MPI.INT)

    # get npars, ndata from MCMC:
    mpipars   = np.zeros(npars*nprocs, np.double)
    mpimodels = np.zeros(ndata*nprocs, np.double)

    if tracktime:
      start_loop = timeit.default_timer()
      loop_timer = []
      loop_timer2 = []
    while niter >= 0:
      if tracktime:
        loop_timer.append(timeit.default_timer())
      # Gather (receive) parameters from MCMC:
      mu.comm_gather(comm1, mpipars)

      # Scatter (send) parameters to funcwrapper:
      mu.comm_scatter(comm2, mpipars, MPI.DOUBLE)
      # Gather (receive) models:
      mu.comm_gather(comm2, mpimodels)

      # Scatter (send) results to MCMC:
      mu.comm_scatter(comm1, mpimodels, MPI.DOUBLE)
      niter -= 1
      if tracktime:
        loop_timer2.append(timeit.default_timer() - loop_timer[-1])

    if tracktime:
      stop = timeit.default_timer()

    # Close communications and disconnect:
    if mpi:
      mu.comm_disconnect(comm1)
      mu.comm_disconnect(comm2)

    #if bench == True:
    if tracktime:
      print("Total execution time:   %10.6f s"%(stop - start))
      print("Time to initialize MPI: %10.6f s"%(start_loop - start_mpi))
      print("Time to run first loop: %10.6f s"%(loop_timer[1] - loop_timer[0]))
      print("Time to run last loop:  %10.6f s"%(loop_timer[-1]- loop_timer[-2]))
      print("Time to run avg loop:   %10.6f s"%(np.mean(loop_timer2)))