text: #!/usr/bin/env python
  # -*- coding: utf-8 -*-
  
  """fitting.py -- Default posterior probability function and high-level fitting
  methods for prospector
  """
  
  
  import time
  from functools import partial as argfix
  
  import numpy as np
  from scipy.optimize import minimize, least_squares
  
  
  from .minimizer import minimize_wrapper, minimizer_ball
  from .ensemble import run_emcee_sampler
  from .nested import run_dynesty_sampler
  from ..likelihood import lnlike_spec, lnlike_phot, chi_spec, chi_phot, write_log
  from ..utils.obsutils import fix_obs
  
  
  __all__ = ["lnprobfn", "fit_model",
             "run_minimize", "run_emcee", "run_dynesty"
             ]
  
  
  def lnprobfn(theta, model=None, obs=None, sps=None, noise=(None, None),
               residuals=False, nested=False, verbose=False):
      """Given a parameter vector and optionally a dictionary of observational
      ata and a model object, return the ln of the posterior. This requires that
      an sps object (and if using spectra and gaussian processes, a GP object) be
      instantiated.
  
      :param theta:
          Input parameter vector, ndarray of shape (ndim,)
  
      :param model:
          SedModel model object, with attributes including ``params``, a
          dictionary of model parameter state.  It must also have
          :py:method:`prior_product`, and :py:method:`mean_model` methods
          defined.
  
      :param obs:
          A dictionary of observational data.  The keys should be
            *``wavelength``  (angstroms)
            *``spectrum``    (maggies)
            *``unc``         (maggies)
            *``maggies``     (photometry in maggies)
            *``maggies_unc`` (photometry uncertainty in maggies)
            *``filters``     (iterable of :py:class:`sedpy.observate.Filter`)
            * and optional spectroscopic ``mask`` and ``phot_mask``
              (same length as `spectrum` and `maggies` respectively,
               True means use the data points)
  
      :param sps:
          A :py:class:`prospect.sources.SSPBasis` object or subclass thereof, or
          any object with a ``get_spectrum`` method that will take a dictionary
          of model parameters and return a spectrum, photometry, and ancillary
          information.
  
      :param noise: (optional, default: (None, None))
          A 2-element tuple of :py:class:`prospect.likelihood.NoiseModel` objects.
  
      :param residuals: (optional, default: False)
          A switch to allow vectors of :math:`\chi` values to be returned instead
          of a scalar posterior probability.  This can be useful for
          least-squares optimization methods. Note that prior probabilities are
          not included in this calculation.
  
      :param nested: (optional, default: False)
          If ``True``, do not add the ln-prior probability to the ln-likelihood
          when computing the ln-posterior.  For nested sampling algorithms the
          prior probability is incorporated in the way samples are drawn, so
          should not be included here.
  
      :returns lnp:
          Ln posterior probability, unless `residuals=True` in which case a
          vector of :math:`\chi` values is returned.
      """
      if residuals:
          lnnull = np.zeros(obs["ndof"]) - 1e18 # np.infty
          #lnnull = -np.infty
      else:
          lnnull = -np.infty
  
      # --- Calculate prior probability and exit if not within prior ---
      lnp_prior = model.prior_product(theta, nested=nested)
      if not np.isfinite(lnp_prior):
          return lnnull
  
      # --- Generate mean model ---
      try:
          t1 = time.time()
          spec, phot, x = model.mean_model(theta, obs, sps=sps)
          d1 = time.time() - t1
      except(ValueError):
          return lnnull
  
      # --- Optionally return chi vectors for least-squares ---
      # note this does not include priors!
      if residuals:
          chispec = chi_spec(spec, obs)
          chiphot = chi_phot(phot, obs)
          return np.concatenate([chispec, chiphot])
  
      #  --- Update Noise Model ---
      spec_noise, phot_noise = noise
      vectors = {}  # These should probably be copies....
      if spec_noise is not None:
          spec_noise.update(**model.params)
          vectors.update({'spec': spec, "unc": obs['unc']})
          vectors.update({'sed': model._spec, 'cal': model._speccal})
      if phot_noise is not None:
          phot_noise.update(**model.params)
          vectors.update({'phot': phot, 'phot_unc': obs['maggies_unc']})
  
      #  --- Mixture Model ---
      f_outlier_spec = model.params.get('f_outlier_spec',0.0)
      if (f_outlier_spec != 0.0):
          sigma_outlier_spec = model.params.get('nsigma_outlier_spec',10)
          vectors.update({'nsigma_outlier_spec': sigma_outlier_spec})
      f_outlier_phot = model.params.get('f_outlier_phot',0.0)
      if (f_outlier_phot != 0.0):
          sigma_outlier_phot = model.params.get('nsigma_outlier_phot',10)
          vectors.update({'nsigma_outlier_phot': sigma_outlier_phot})
  
      # --- Calculate likelihoods ---
      t1 = time.time()
      lnp_spec = lnlike_spec(spec, obs=obs,
                             f_outlier_spec=f_outlier_spec,
                             spec_noise=spec_noise, **vectors)
      lnp_phot = lnlike_phot(phot, obs=obs,
                             f_outlier_phot=f_outlier_phot,
                             phot_noise=phot_noise, **vectors)
      d2 = time.time() - t1
      if verbose:
          write_log(theta, lnp_prior, lnp_spec, lnp_phot, d1, d2)
  
      return lnp_prior + lnp_phot + lnp_spec
  
  
  def wrap_lnp(lnpfn, obs, model, sps, **lnp_kwargs):
      return argfix(lnpfn, obs=obs, model=model, sps=sps,
                    **lnp_kwargs)
  
  
  def fit_model(obs, model, sps, noise=(None, None), lnprobfn=lnprobfn,
                optimize=False, emcee=False, dynesty=True, **kwargs):
      """Fit a model to observations using a number of different methods
  
      :param obs:
          The ``obs`` dictionary containing the data to fit to, which will be
          passed to ``lnprobfn``.
  
      :param model:
          An instance of the :py:class:`prospect.models.SedModel` class
          containing the model parameterization and parameter state.  It will be
          passed to ``lnprobfn``.
  
      :param sps:
          An instance of a :py:class:`prospect.sources.SSPBasis` (sub-)class.
          Alternatively, anything with a compatible :py:method:`get_spectrum` can
          be used here. It will be passed to ``lnprobfn``
  
      :param noise: (optional, default: (None, None))
          A tuple of NoiseModel objects for the spectroscopy and photometry
          respectively.  Can also be (None, None) in which case simple chi-square
          will be used.
  
      :param lnprobfn: (optional, default: lnprobfn)
          A posterior probability function that can take ``obs``, ``model``,
          ``sps``, and ``noise`` as keywords. By default use the
          :py:method:`lnprobfn` defined above.
  
      :param optimize: (optional, default: False)
          If ``True``, conduct a round of optimization before sampling from the
          posterior.  The model state will be set to the best value at the end of
          optimization before continuing on to sampling or returning.  Parameters
          controlling the optimization can be passed via ``kwargs``, including
  
          * min_method: 'lm' | 'powell'
          * nmin: number of minimizations to do.  Beyond the first, minimizations
            will be started from draws from the prior. 
          * min_opts: dictionary of minimization options passed to the
            scipy.optimize.minimize method.
          
          See :py:method:`run_minimize` for details.
  
      :param emcee:  (optional, default: False)
          If ``True``, sample from the posterior using emcee.  Additonal
          parameters controlling emcee can be passed via **kwargs.  These include
  
          * initial_positions: A set of initial positions for the walkers
          * hfile: an open h5py.File file handle for writing result incrementally
          
          Many additional emcee parameters can be provided here, see
          :py:method:`run_emcee` for details.
  
      :param dynesty:
          If ``True``, sample from the posterior using dynesty.  Additonal
          parameters controlling dynesty can be passed via **kwargs. See
          :py:method:`run_dynesty` for details.
  
      :returns output:
          A dictionary with two keys, 'optimization' and 'sampling'.  The value
          of each of these is a 2-tuple with results in the first element and
          durations (in seconds) in the second element.
      """
      # Make sure obs has required keys
      obs = fix_obs(obs)
  
      if emcee & dynesty:
          msg = ("Cannot run both emcee and dynesty fits "
                 "in a single call to fit_model")
          raise(ValueError, msg)
  
      output = {"optimization": (None, 0.),
                "sampling": (None, 0.)}
  
      if optimize:
          optres, topt, best = run_minimize(obs, model, sps, noise,
                                            lnprobfn=lnprobfn, **kwargs)
          # set to the best
          model.set_parameters(optres[best].x)
          output["optimization"] = (optres, topt)
  
      if emcee:
          run_sampler = run_emcee
      elif dynesty:
          run_sampler = run_dynesty
      else:
          return output
  
      output["sampling"] = run_sampler(obs, model, sps, noise,
                                       lnprobfn=lnprobfn, **kwargs)
      return output
  
  
  def run_minimize(obs=None, model=None, sps=None, noise=None, lnprobfn=lnprobfn,
                   min_method='lm', min_opts={}, nmin=1, pool=None, **extras):
      """Run a minimization.  This wraps the lnprobfn fixing the ``obs``,
      ``model``, ``noise``, and ``sps`` objects, and then runs a minimization of
      -lnP using scipy.optimize methods.
      
      :param obs:
          The ``obs`` dictionary containing the data to fit to, which will be
          passed to ``lnprobfn``.
  
      :param model:
          An instance of the :py:class:`prospect.models.SedModel` class
          containing the model parameterization and parameter state.  It will be
          passed to ``lnprobfn``.
  
      :param sps:
          An instance of a :py:class:`prospect.sources.SSPBasis` (sub-)class.
          Alternatively, anything with a compatible :py:method:`get_spectrum` can
          be used here. It will be passed to ``lnprobfn``
  
      :param noise: (optional) 
          If given, a tuple of :py:class:`NoiseModel` objects passed to
          ``lnprobfn``.
  
      :param lnprobfn: (optional, default: lnprobfn)
          A posterior probability function that can take ``obs``, ``model``,
          ``sps``, and ``noise`` as keywords. By default use the
          :py:method:`lnprobfn` defined above.
  
      :param min_method: (optional, default: 'lm')
          Method to use for minimization
          * 'lm': Levenberg-Marquardt
          * 'powell': Powell line search method
      
      :param nmin: (optional, default: 1)
          Number of minimizations to do.  Beyond the first, minimizations will be
          started from draws from the prior.
  
      :param min_opts: (optional, default: {})
          Dictionary of minimization options passed to the scipy.optimize method.
          These include things like 'xtol', 'ftol', etc..
  
      :param pool: (optional, default: None)
          A pool to use for parallel optimization from multiple initial positions.
  
      :returns results:
          A list of `scipy.optimize.OptimizeResult` objects.
  
      :returns tm:
          Wall time used for the minimization, in seconds.
  
      :returns best:
          The index of the results list containing the lowest chi-square result.
      """
      initial = model.theta.copy()
      
      lsq = ["lm"]
      scalar = ["powell"]
  
      # --- Set some options based on minimization method ---
      if min_method in lsq:
          algorithm = least_squares
          residuals = True
          min_opts["x_scale"] = "jac"
      elif min_method in scalar:
          algorithm = minimize
          residuals = False
  
      args = []
      loss = argfix(lnprobfn, obs=obs, model=model, sps=sps,
                    noise=noise, residuals=residuals)
      minimizer = minimize_wrapper(algorithm, loss, [], min_method, min_opts)
      qinit = minimizer_ball(initial, nmin, model)
  
      if pool is not None:
              M = pool.map
      else:
              M = map
  
      t = time.time()
      results = list(M(minimizer, [np.array(q) for q in qinit]))
      tm = time.time() - t
  
      if min_method in lsq:
          chisq = [np.sum(r.fun**2) for r in results]
          best = np.argmin(chisq)
      elif min_method in scalar:
          best = np.argmin([p.fun for p in results])
  
      return results, tm, best
  
  
  def run_emcee(obs, model, sps, noise, lnprobfn=lnprobfn,
                hfile=None, initial_positions=None,
                **kwargs):
      """Run emcee, optionally including burn-in and convergence checking.  Thin
      wrapper on :py:class:`prospect.fitting.ensemble.run_emcee_sampler`
  
      :param obs:
          The ``obs`` dictionary containing the data to fit to, which will be
          passed to ``lnprobfn``.
  
      :param model:
          An instance of the :py:class:`prospect.models.SedModel` class
          containing the model parameterization and parameter state.  It will be
          passed to ``lnprobfn``.
  
      :param sps:
          An instance of a :py:class:`prospect.sources.SSPBasis` (sub-)class.
          Alternatively, anything with a compatible :py:method:`get_spectrum` can
          be used here. It will be passed to ``lnprobfn``
  
      :param noise:
          A tuple of :py:class:`NoiseModel` objects passed to ``lnprobfn``.
  
      :param lnprobfn: (optional, default: lnprobfn)
          A posterior probability function that can take ``obs``, ``model``,
          ``sps``, and ``noise`` as keywords. By default use the
          :py:method:`lnprobfn` defined above.
  
      :param hfile: (optional, default: None)
          A file handle for a :py:class:`h5py.File` object that will be written
          to incremantally during sampling.
  
      :param initial_positions: (optional, default: None)
          If given, a set of initial positions for the emcee walkers.  Must have
          shape (nwalkers, ndim).  Rounds of burn-in will be skipped if this
          parameter is present.
  
      Extra Parameters
      --------
  
      :param nwalkers:
          The number of walkers to use.  If None, use the nearest power of two to
          ``ndim * walker_factor``.
  
      :param niter:
          Number of iterations for the production run
  
      :param nburn:
          List of the number of iterations to run in each round of burn-in (for
          removing stuck walkers.) E.g. `nburn=[32, 64]` will run the sampler for
          32 iterations before reinitializing and then run the sampler for
          another 64 iterations before starting the production run.
  
      :param storechain: (default: True)
          If using HDF5 output, setting this to False will keep the chain from
          being held in memory by the sampler object.
  
      :param pool: (optional)
          A ``Pool`` object, either from ``multiprocessing`` or from
          ``emcee.mpi_pool``.
  
      :param interval:
          Fraction of the full run at which to flush to disk, if using hdf5 for
          output.
  
      :param convergence_check_interval:
          How often to assess convergence, in number of iterations. If this is
          not `None`, then the KL convergence test is run.
  
      :param convergence_chunks:
          The number of iterations to combine when creating the marginalized
          parameter probability functions.
  
      :param convergence_stable_points_criteria:
          The number of stable convergence checks that the chain must pass before
          being declared stable.
  
      Returns
      --------
  
      :returns sampler:
          An instance of :py:class:`emcee.EnsembleSampler`.
  
      :returns ts:
          Duration of sampling (including burn-in) in seconds of wall time.
      """
      q = model.theta.copy()
  
      postkwargs = {"obs": obs,
                    "model": model,
                    "sps": sps,
                    "noise": noise,
                    "nested": False,
                    }
  
      # Could try to make signatures for these two methods the same....
      if initial_positions is not None:
          meth = restart_emcee_sampler
          t = time.time()
          out = meth(lnprobfn, initial_positions, hdf5=hfile,
                     postkwargs=postkwargs, **kwargs)
          sampler = out
          ts = time.time() - t
      else:
          meth = run_emcee_sampler
          t = time.time()
          out = meth(lnprobfn, q, model, hdf5=hfile,
                     postkwargs=postkwargs, **kwargs)
          sampler, burn_p0, burn_prob0 = out
          ts = time.time() - t
  
      return sampler, ts
  
  
  def run_dynesty(obs, model, sps, noise, lnprobfn=lnprobfn,
                  pool=None, nested_posterior_thresh=0.05, **kwargs):
      """Thin wrapper on :py:class:`prospect.fitting.nested.run_dynesty_sampler`
  
      :param obs:
          The ``obs`` dictionary containing the data to fit to, which will be
          passed to ``lnprobfn``.
  
      :param model:
          An instance of the :py:class:`prospect.models.SedModel` class
          containing the model parameterization and parameter state.  It will be
          passed to ``lnprobfn``.
  
      :param sps:
          An instance of a :py:class:`prospect.sources.SSPBasis` (sub-)class.
          Alternatively, anything with a compatible :py:method:`get_spectrum` can
          be used here. It will be passed to ``lnprobfn``
  
      :param noise:
          A tuple of :py:class:`prospect.likelihood.NoiseModel` objects passed to
          ``lnprobfn``.
  
      :param lnprobfn: (optional, default: :py:method:`lnprobfn`)
          A posterior probability function that can take ``obs``, ``model``,
          ``sps``, and ``noise`` as keywords. This function must also take a
          ``nested`` keyword.
  
      Extra Parameters
      --------
      :param nested_bound: (optional, default: 'multi')
  
      :param nested_sample: (optional, default: 'unif')
  
      :param nested_nlive_init: (optional, default: 100)
  
      :param nested_nlive_batch: (optional, default: 100)
  
      :param nested_dlogz_init: (optional, default: 0.02)
  
      :param nested_maxcall: (optional, default: None)
  
      :param nested_walks: (optional, default: 25)
  
      Returns
      --------
  
      :returns result:
          An instance of :py:class:`dynesty.results.Results`.
  
      :returns ts:
          Duration of sampling in seconds of wall time.
      """
      from dynesty.dynamicsampler import stopping_function, weight_function
      nested_stop_kwargs = {"post_thresh": nested_posterior_thresh}
  
      lnp = wrap_lnp(lnprobfn, obs, model, sps, noise=noise,
                     nested=True)
  
      # Need to deal with postkwargs...
  
      t = time.time()
      dynestyout = run_dynesty_sampler(lnp, model.prior_transform, model.ndim,
                                       stop_function=stopping_function,
                                       wt_function=weight_function,
                                       nested_stop_kwargs=nested_stop_kwargs,
                                       pool=pool, **kwargs)
      ts = time.time() - t
  
      return dynestyout, ts
  