diff --git a/.github/workflows/tests.yml b/.github/workflows/tests.yml index f7b7c073..0b720d3c 100644 --- a/.github/workflows/tests.yml +++ b/.github/workflows/tests.yml @@ -12,7 +12,7 @@ jobs: strategy: fail-fast: false matrix: - python-version: ["3.9", "3.10"] + python-version: ["3.9", "3.11"] os: [ubuntu-latest, macos-latest] steps: - name: Clone the repo diff --git a/demo/demo_params.py b/demo/demo_params.py index b4d8a769..36fa1ed2 100644 --- a/demo/demo_params.py +++ b/demo/demo_params.py @@ -148,7 +148,7 @@ def build_obs(objid=0, phottable='demo_photometry.dat', # import astropy.io.fits as pyfits # catalog = pyfits.getdata(phottable) - from prospect.data.observation import Photometry, Spectrum + from prospect.observation import Photometry, Spectrum # Here we will read in an ascii catalog of magnitudes as a numpy structured # array @@ -262,10 +262,12 @@ def build_all(**kwargs): output = fit_model(obs, model, sps, **config) print("writing to {}".format(hfile)) - writer.write_hdf5(hfile, run_params, model, obs, - output["sampling"][0], output["optimization"][0], - tsample=output["sampling"][1], - toptimize=output["optimization"][1], + writer.write_hdf5(hfile, + config, + model, + obs, + output["sampling"], + output["optimization"], sps=sps ) diff --git a/demo/tutorial.rst b/demo/tutorial.rst index dff8a13b..d2396367 100644 --- a/demo/tutorial.rst +++ b/demo/tutorial.rst @@ -164,25 +164,26 @@ that for now. Running a fit ---------------------- -There are two kinds of fitting packages that can be used with |Codename|. The -first is ``emcee`` which implements ensemble MCMC sampling, and the second is -``dynesty``, which implements dynamic nested sampling. It is also possible to -perform optimization. If ``emcee`` is used, the result of the optimization will -be used to initalize the ensemble of walkers. +There are a few kinds of fitting packages that can be used with |Codename|. The +first is ``emcee`` which implements ensemble MCMC sampling. A number of nested +samplers are also available (``dynesty``, ``nautilus``, and ``ultranest``). It +is also possible to perform optimization. If ``emcee`` is used, the result of +the optimization will be used to initalize the ensemble of walkers. The choice of which fitting algorithms to use is based on command line flags -(``--optimization``, ``--emcee``, and ``--dynesty``.) If no flags are set the -model and data objects will be generated and stored in the output file, but no -fitting will take place. To run the fit on object number 0 using ``emcee`` after -an initial optimization, we would do the following at the command line +(``--optimization``, ``--emcee``, and ``--nested_sampler ``.) +If no flags are set the model and data objects will be generated and stored in +the output file, but no fitting will take place. To run the fit on object number +0 using ``emcee`` after an initial optimization, we would do the following at +the command line .. code-block:: shell python demo_params.py --objid=0 --emcee --optimize \ --outfile=demo_obj0_emcee -If we wanted to change something about the MCMC parameters, or fit a different object, -we could also do that at the command line +If we wanted to change something about the MCMC parameters, or fit a different +object, we could also do that at the command line .. code-block:: shell @@ -193,7 +194,7 @@ And if we want to use nested sampling with ``dynesty`` we would do the following .. code-block:: shell - python demo_params.py --objid=0 --dynesty \ + python demo_params.py --objid=0 --nested_sampler dynesty \ --outfile=demo_obj0_dynesty Finally, it is sometimes useful to run the script from the interpreter to do diff --git a/doc/quickstart.rst b/doc/quickstart.rst index eabc38bf..376ef785 100644 --- a/doc/quickstart.rst +++ b/doc/quickstart.rst @@ -154,12 +154,12 @@ it should still take of order tens of minutes. output = fit_model(obs, model, sps, lnprobfn=lnprobfn, optimize=False, dynesty=True, **fitting_kwargs) - result, duration = output["sampling"] + result = output["sampling"] The ``result`` is a dictionary with keys giving the Monte Carlo samples of parameter values and the corresponding posterior probabilities. Because we are -using dynesty, we also get weights associated with each parameter sample in the -chain. +using nested sampling, we also get weights associated with each parameter sample +in the chain. Typically we'll want to save the fit information. We can save the output of the sampling along with other information about the model and the data that was fit @@ -168,12 +168,16 @@ as follows: .. code:: python from prospect.io import write_results as writer - hfile = "./quickstart_dynesty_mcmc.h5" - writer.write_hdf5(hfile, {}, model, observations, - output["sampling"][0], None, - sps=sps, - tsample=output["sampling"][1], - toptimize=0.0) + writer.write_hdf5("./quickstart_dynesty_mcmc.h5", + config=fitting_kwargs, + model=model, + obs=observations, + output["sampling"], + None, + sps=sps) + +Note that this doesn't include all the config information that would normally be stored (see :ref:`usage`) + Make plots ---------- @@ -187,14 +191,14 @@ information we can use the built-in reader. hfile = "./quickstart_dynesty_mcmc.h5" out, out_obs, out_model = reader.results_from(hfile) -This gives a dictionary of useful information (``out``), as well as the obs -dictionary that we were using and, in some cases, a reconsitituted model object. -However, that is only possible if the model generation code is saved to the -results file, in the form of the text for a `build_model()` function. Here we -will use just use the model object that we've already generated. +This gives a dictionary of useful information (``out``), as well as the obs data +that we were using and, in some cases, a reconsitituted model object. However, +that is *only* possible if the model generation code is saved to the results file, +in the form of the text for a `build_model()` function. Here we will use just +use the model object that we've already generated. -Now we will do some plotting. First, lets make a corner plot of the posterior. -We'll mark the highest probablity posterior sample as well. +First, lets make a corner plot of the posterior. We'll mark the highest +probablity posterior sample as well. .. code:: python diff --git a/doc/usage.rst b/doc/usage.rst index 950f0931..29840131 100644 --- a/doc/usage.rst +++ b/doc/usage.rst @@ -13,7 +13,7 @@ execution: .. code-block:: shell - python parameter_file.py --dynesty + python parameter_file.py --nested_sampler dynesty --nested_target_n_effective 512 Additional command line options can be given (see below) e.g. @@ -50,6 +50,7 @@ writes output. # --- Configure --- args = parser.parse_args() config = vars(args) + # allows parameter file text to be stored for model regeneration config["param_file"] = __file__ # --- Get fitting ingredients --- @@ -68,12 +69,14 @@ writes output. output = fit_model(obs, model, sps, **config) print("writing to {}".format(hfile)) - writer.write_hdf5(hfile, run_params, model, obs, - output["sampling"][0], output["optimization"][0], - tsample=output["sampling"][1], - toptimize=output["optimization"][1], - sps=sps - ) + writer.write_hdf5(hfile, + config=config, + model=model, + obs=obs, + output["sampling"], + output["optimization"], + sps=sps + ) try: hfile.close() @@ -147,7 +150,8 @@ might actually increase the total CPU usage). To use MPI a "pool" of cores must be made available; each core will instantiate the fitting ingredients separately, and a single core in the pool will then conduct the fit, distributing likelihood requests to the other cores in the -pool. This requires changes to the final code block that instantiates and runs the fit: +pool. This requires changes to the final code block that instantiates and runs +the fit: .. code-block:: python @@ -189,7 +193,7 @@ pool. This requires changes to the final code block that instantiates and runs # Evaluate SPS over logzsol grid in order to get necessary data in cache/memory # for each MPI process. Otherwise, you risk creating a lag between the MPI tasks - # caching data depending which can slow down the parallelization + # caching SSPs which can slow down the parallelization if (withmpi) & ('logzsol' in model.free_params): dummy_obs = dict(filters=None, wavelength=None) @@ -225,13 +229,15 @@ pool. This requires changes to the final code block that instantiates and runs # Set up an output file and write ts = time.strftime("%y%b%d-%H.%M", time.localtime()) - hfile = "{0}_{1}_mcmc.h5".format(args.outfile, ts) - writer.write_hdf5(hfile, run_params, model, obs, - output["sampling"][0], output["optimization"][0], - tsample=output["sampling"][1], - toptimize=output["optimization"][1], - sps=sps) - + hfile = f"{args.outfile}_worker{comm.rank()}_{ts}_mcmc.h5" + writer.write_hdf5(hfile, + run_params=run_params, + model=model, + obs=obs, + output["sampling"], + output["optimization"], + sps=sps + ) try: hfile.close() except(AttributeError): @@ -243,7 +249,7 @@ Then, to run this file using mpi it can be called from the command line with som mpirun -np python parameter_file.py --emcee # or - mpirun -np python parameter_file.py --dynesty + mpirun -np python parameter_file.py --nested_sampler dynesty Note that only model evaluation is parallelizable with `dynesty`, and many operations (e.g. new point proposal) are still done in serial. This means that diff --git a/prospect/fitting/__init__.py b/prospect/fitting/__init__.py index 80cbb8f7..0a6bbc6d 100644 --- a/prospect/fitting/__init__.py +++ b/prospect/fitting/__init__.py @@ -1,10 +1,10 @@ from .ensemble import run_emcee_sampler, restart_emcee_sampler from .minimizer import reinitialize -from .nested import run_dynesty_sampler +from .nested import run_nested_sampler from .fitting import fit_model, lnprobfn, run_minimize __all__ = ["fit_model", "lnprobfn", # below should all be removed/deprecated "run_emcee_sampler", "restart_emcee_sampler", - "run_dynesty_sampler", + "run_nested_sampler", "run_minimize", "reinitialize"] diff --git a/prospect/fitting/fitting.py b/prospect/fitting/fitting.py index d222d07c..aa4c0a22 100755 --- a/prospect/fitting/fitting.py +++ b/prospect/fitting/fitting.py @@ -15,12 +15,12 @@ from .minimizer import minimize_wrapper, minimizer_ball from .ensemble import run_emcee_sampler -from .nested import run_dynesty_sampler +from .nested import run_nested_sampler, parse_nested_kwargs from ..likelihood.likelihood import compute_chi, compute_lnlike __all__ = ["lnprobfn", "fit_model", - "run_minimize", "run_emcee", "run_dynesty" + "run_minimize", "run_ensemble", "run_nested" ] @@ -72,7 +72,7 @@ def lnprobfn(theta, model=None, observations=None, sps=None, """ if residuals: ndof = np.sum([obs["ndof"] for obs in observations]) - lnnull = np.zeros(ndof) - 1e18 # -np.infty + lnnull = np.zeros(ndof) - 1e18 # -np.inf else: lnnull = -np.inf @@ -123,7 +123,8 @@ def wrap_lnp(lnpfn, observations, model, sps, **lnp_kwargs): def fit_model(observations, model, sps, lnprobfn=lnprobfn, - optimize=False, emcee=False, dynesty=True, **kwargs): + optimize=False, emcee=False, nested_sampler="", + **kwargs): """Fit a model to observations using a number of different methods Parameters @@ -167,7 +168,7 @@ def fit_model(observations, model, sps, lnprobfn=lnprobfn, + ``hfile``: an open h5py.File file handle for writing result incrementally Many additional emcee parameters can be provided here, see - :py:func:`run_emcee` for details. + :py:func:`run_ensemble` for details. dynesty : bool (optional, default: True) If ``True``, sample from the posterior using dynesty. Additonal @@ -184,17 +185,18 @@ def fit_model(observations, model, sps, lnprobfn=lnprobfn, # Make sure obs has required keys [obs.rectify() for obs in observations] - if emcee & dynesty: - msg = ("Cannot run both emcee and dynesty fits " + if (emcee) & (bool(nested_sampler)): + msg = ("Cannot run both emcee and nested fits " "in a single call to fit_model") raise(ValueError, msg) - if (not emcee) & (not dynesty) & (not optimize): + + if (not bool(emcee)) & (not bool(nested_sampler)) & (not optimize): msg = ("No sampling or optimization routine " "specified by user; returning empty results") warnings.warn(msg) - output = {"optimization": (None, 0.), - "sampling": (None, 0.)} + output = {"optimization": None, + "sampling": None} if optimize: optres, topt, best = run_minimize(observations, model, sps, @@ -204,14 +206,17 @@ def fit_model(observations, model, sps, lnprobfn=lnprobfn, output["optimization"] = (optres, topt) if emcee: - run_sampler = run_emcee - elif dynesty: - run_sampler = run_dynesty + run_sampler = run_ensemble + elif nested_sampler: + run_sampler = run_nested + # put nested_sampler back into kwargs for lower level functions + kwargs["nested_sampler"] = nested_sampler else: return output output["sampling"] = run_sampler(observations, model, sps, - lnprobfn=lnprobfn, **kwargs) + lnprobfn=lnprobfn, + **kwargs) return output @@ -305,8 +310,8 @@ def run_minimize(observations=None, model=None, sps=None, lnprobfn=lnprobfn, return results, tm, best -def run_emcee(observations, model, sps, lnprobfn=lnprobfn, - hfile=None, initial_positions=None, **kwargs): +def run_ensemble(observations, model, sps, 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` @@ -387,6 +392,7 @@ def run_emcee(observations, model, sps, lnprobfn=lnprobfn, q = model.theta.copy() postkwargs = {} + # Hack for MPI pools to access the global namespace for item in ['observations', 'model', 'sps']: val = eval(item) if val is not None: @@ -396,26 +402,38 @@ def run_emcee(observations, model, sps, lnprobfn=lnprobfn, # Could try to make signatures for these two methods the same.... if initial_positions is not None: raise NotImplementedError - meth = restart_emcee_sampler - t = time.time() - out = meth(lnprobfn, initial_positions, hdf5=hfile, - postkwargs=postkwargs, **kwargs) + go = time.time() + out = restart_emcee_sampler(lnprobfn, initial_positions, + hdf5=hfile, + postkwargs=postkwargs, + **kwargs) sampler = out - ts = time.time() - t + ts = time.time() - go 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 + go = time.time() + out = run_emcee_sampler(lnprobfn, q, model, + hdf5=hfile, + postkwargs=postkwargs, + **kwargs) + sampler, burn_loc0, burn_prob0 = out + ts = time.time() - go + + try: + sampler.duration = ts + except: + pass - return sampler, ts + return sampler -def run_dynesty(observations, model, sps, lnprobfn=lnprobfn, - pool=None, nested_target_n_effective=10000, **kwargs): - """Thin wrapper on :py:class:`prospect.fitting.nested.run_dynesty_sampler` +def run_nested(observations, model, sps, + lnprobfn=lnprobfn, + nested_sampler="dynesty", + nested_nlive=1000, + nested_target_n_effective=1000, + verbose=False, + **kwargs): + """Thin wrapper on :py:class:`prospect.fitting.nested.run_nested_sampler` Parameters ---------- @@ -436,43 +454,38 @@ def run_dynesty(observations, model, sps, lnprobfn=lnprobfn, ``model``, and ``sps`` as keywords. By default use the :py:func:`lnprobfn` defined above. - Extra Parameters - -------- - nested_bound: (optional, default: 'multi') - - nested_sample: (optional, default: 'unif') - - nested_nlive_init: (optional, default: 100) + nested_target_n_effective : int + Target number of effective samples - nested_nlive_batch: (optional, default: 100) - - nested_dlogz_init: (optional, default: 0.02) - - nested_maxcall: (optional, default: None) - - nested_walks: (optional, default: 25) + nested_nlive : int + Number of live points for the nested sampler. Meaning somewhat + dependent on the chosen sampler Returns -------- - result: - An instance of :py:class:`dynesty.results.Results`. + result: Dictionary + Will have keys: + * points : parameter location of the samples + * log_weight : ln of the weights of each sample + * log_like : ln of the likelihoods of each sample t_wall : float Duration of sampling in seconds of wall time. """ - from dynesty.dynamicsampler import stopping_function, weight_function - nested_stop_kwargs = {"target_n_effective": nested_target_n_effective} - - lnp = wrap_lnp(lnprobfn, observations, model, sps, 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 + # wrap the probability fiunction, making sure it's a likelihood + likelihood = wrap_lnp(lnprobfn, observations, model, sps, nested=True) + + # which keywords do we have for this fitter? + ns_kwargs, nr_kwargs = parse_nested_kwargs(nested_sampler=nested_sampler,**kwargs) + + output = run_nested_sampler(model, + likelihood, + nested_sampler=nested_sampler, + verbose=verbose, + nested_nlive=nested_nlive, + nested_neff=nested_target_n_effective, + nested_sampler_kwargs=ns_kwargs, + nested_run_kwargs=nr_kwargs) + info, result_obj = output + + return info diff --git a/prospect/fitting/nested.py b/prospect/fitting/nested.py index 73d7ed1e..f0f25185 100644 --- a/prospect/fitting/nested.py +++ b/prospect/fitting/nested.py @@ -1,199 +1,259 @@ -import sys, time +import time import numpy as np -from numpy.random import normal, multivariate_normal - -try: - import nestle -except(ImportError): - pass - - -try: - import dynesty - from dynesty.utils import * - from dynesty.dynamicsampler import _kld_error -except(ImportError): - pass - - -__all__ = ["run_nestle_sampler", "run_dynesty_sampler"] - - -def run_nestle_sampler(lnprobfn, model, verbose=True, - callback=None, - nestle_method='multi', nestle_npoints=200, - nestle_maxcall=int(1e6), nestle_update_interval=None, - **kwargs): - - result = nestle.sample(lnprobfn, model.prior_transform, model.ndim, - method=nestle_method, npoints=nestle_npoints, - callback=callback, maxcall=nestle_maxcall, - update_interval=nestle_update_interval) - return result - - -def run_dynesty_sampler(lnprobfn, prior_transform, ndim, - verbose=True, - # sampler kwargs - nested_bound='multi', - nested_sample='unif', - nested_walks=25, - nested_update_interval=0.6, - nested_bootstrap=0, - pool=None, - use_pool={}, - queue_size=1, - # init sampling kwargs - nested_nlive_init=100, - nested_dlogz_init=0.02, - nested_maxiter_init=None, - nested_maxcall_init=None, - nested_live_points=None, - # batch sampling kwargs - nested_maxbatch=None, - nested_nlive_batch=100, - nested_maxiter_batch=None, - nested_maxcall_batch=None, - nested_use_stop=True, - # overall kwargs - nested_maxcall=None, - nested_maxiter=None, - nested_first_update={}, - stop_function=None, - wt_function=None, - nested_weight_kwargs={'pfrac': 1.0}, - nested_stop_kwargs={}, - nested_save_bounds=False, - print_progress=True, - **extras): - - # instantiate sampler - dsampler = dynesty.DynamicNestedSampler(lnprobfn, prior_transform, ndim, - bound=nested_bound, - sample=nested_sample, - walks=nested_walks, - bootstrap=nested_bootstrap, - update_interval=nested_update_interval, - pool=pool, queue_size=queue_size, use_pool=use_pool - ) - - # generator for initial nested sampling - ncall = dsampler.ncall - niter = dsampler.it - 1 - tstart = time.time() - for results in dsampler.sample_initial(nlive=nested_nlive_init, - dlogz=nested_dlogz_init, - maxcall=nested_maxcall_init, - maxiter=nested_maxiter_init, - live_points=nested_live_points): - - try: - # dynesty >= 2.0 - (worst, ustar, vstar, loglstar, logvol, - logwt, logz, logzvar, h, nc, worst_it, - propidx, propiter, eff, delta_logz, blob) = results - except(ValueError): - # dynsety < 2.0 - (worst, ustar, vstar, loglstar, logvol, - logwt, logz, logzvar, h, nc, worst_it, - propidx, propiter, eff, delta_logz) = results - - if delta_logz > 1e6: - delta_logz = np.inf - ncall += nc - niter += 1 - - if print_progress: - with np.errstate(invalid='ignore'): - logzerr = np.sqrt(logzvar) - sys.stderr.write("\riter: {:d} | batch: {:d} | nc: {:d} | " - "ncall: {:d} | eff(%): {:6.3f} | " - "logz: {:6.3f} +/- {:6.3f} | " - "dlogz: {:6.3f} > {:6.3f} " - .format(niter, 0, nc, ncall, eff, logz, - logzerr, delta_logz, nested_dlogz_init)) - sys.stderr.flush() - - ndur = time.time() - tstart - if verbose: - print('\ndone dynesty (initial) in {0}s'.format(ndur)) - - if nested_maxcall is None: - nested_maxcall = sys.maxsize - if nested_maxbatch is None: - nested_maxbatch = sys.maxsize - if nested_maxcall_batch is None: - nested_maxcall_batch = sys.maxsize - if nested_maxiter is None: - nested_maxiter = sys.maxsize - if nested_maxiter_batch is None: - nested_maxiter_batch = sys.maxsize - - # generator for dynamic sampling - tstart = time.time() - for n in range(dsampler.batch, nested_maxbatch): - # Update stopping criteria. - dsampler.sampler.save_bounds = False - res = dsampler.results - mcall = min(nested_maxcall - ncall, nested_maxcall_batch) - miter = min(nested_maxiter - niter, nested_maxiter_batch) - if nested_use_stop: - if dsampler.use_pool_stopfn: - M = dsampler.M - else: - M = map - stop, stop_vals = stop_function(res, nested_stop_kwargs, - rstate=dsampler.rstate, M=M, - return_vals=True) - stop_val = stop_vals[2] - else: - stop = False - stop_val = np.NaN - - # If we have either likelihood calls or iterations remaining, - # run our batch. - if mcall > 0 and miter > 0 and not stop: - # Compute our sampling bounds using the provided - # weight function. - logl_bounds = wt_function(res, nested_weight_kwargs) - lnz, lnzerr = res.logz[-1], res.logzerr[-1] - for results in dsampler.sample_batch(nlive_new=nested_nlive_batch, - logl_bounds=logl_bounds, - maxiter=miter, - maxcall=mcall, - save_bounds=nested_save_bounds): - - try: - # dynesty >= 2.0 - (worst, ustar, vstar, loglstar, nc, - worst_it, propidx, propiter, eff, blob) = results - except(ValueError): - # dynesty < 2.0 - (worst, ustar, vstar, loglstar, nc, - worst_it, propidx, propiter, eff) = results - ncall += nc - niter += 1 - if print_progress: - sys.stderr.write("\riter: {:d} | batch: {:d} | " - "nc: {:d} | ncall: {:d} | " - "eff(%): {:6.3f} | " - "loglstar: {:6.3f} < {:6.3f} " - "< {:6.3f} | " - "logz: {:6.3f} +/- {:6.3f} | " - "stop: {:6.3f} " - .format(niter, n+1, nc, ncall, - eff, logl_bounds[0], loglstar, - logl_bounds[1], lnz, lnzerr, - stop_val)) - sys.stderr.flush() - dsampler.combine_runs() - else: - # We're done! - break - - ndur = time.time() - tstart + +__all__ = ["run_nested_sampler", "parse_nested_kwargs"] + + +def parse_nested_kwargs(nested_sampler=None, **kwargs): + + # TODO: + # something like 'enlarge' + # something like 'bootstrap' or N_networks or? + + sampler_kwargs = {} + run_kwargs = {} + + if nested_sampler == "dynesty": + sampler_kwargs["bound"] = kwargs["nested_bound"] + sampler_kwargs["sample"] = kwargs["nested_sample"] + sampler_kwargs["walks"] = kwargs["nested_walks"] + run_kwargs["dlogz_init"] = kwargs["nested_dlogz"] + + elif nested_sampler == "ultranest": + #run_kwargs["dlogz"] = kwargs["nested_dlogz"] + pass + + elif nested_sampler == "nautilus": + pass + + else: + # say what? + raise ValueError(f"{nested_sampler} not a valid fitter") + + return sampler_kwargs, run_kwargs + + + +def run_nested_sampler(model, + likelihood_function, + nested_sampler="dynesty", + nested_nlive=1000, + nested_neff=1000, + verbose=False, + nested_run_kwargs={}, + nested_sampler_kwargs={}): + """We give a model -- parameter discription and prior transform -- and a + likelihood function. We get back samples, weights, and likelihood values. + + Returns + ------- + samples : 3-tuple of ndarrays (loc, logwt, loglike) + Loctions, log-weights, and log-likelihoods for the samples + + obj : Object + The sampling object. This will depend on the nested sampler being used. + """ if verbose: - print('done dynesty (dynamic) in {0}s'.format(ndur)) + print(f"running {nested_sampler} for {nested_neff} effective samples") + + go = time.time() + + # --- Nautilus --- + if nested_sampler == "nautilus": + from nautilus import Sampler + + sampler = Sampler(model.prior_transform, + likelihood_function, + n_dim=model.ndim, + pass_dict=False, # likelihood expects array, not dict + n_live=nested_nlive, + **nested_sampler_kwargs) + sampler.run(n_eff=nested_neff, + verbose=verbose, + **nested_run_kwargs) + obj = sampler + + points, log_w, log_like = sampler.posterior() + + # --- Ultranest --- + if nested_sampler == "ultranest": + + from ultranest import ReactiveNestedSampler + parameter_names = model.theta_labels() + sampler = ReactiveNestedSampler(parameter_names, + likelihood_function, + model.prior_transform, + **nested_sampler_kwargs) + result = sampler.run(min_ess=nested_neff, + min_num_live_points=nested_nlive, + show_status=verbose, + **nested_run_kwargs) + obj = result + + points = np.array(result['weighted_samples']['points']) + log_w = np.log(np.array(result['weighted_samples']['weights'])) + log_like = np.array(result['weighted_samples']['logl']) + + # --- Dynesty --- + if nested_sampler == "dynesty": + from dynesty import DynamicNestedSampler + + sampler = DynamicNestedSampler(likelihood_function, + model.prior_transform, + model.ndim, + nlive=nested_nlive, + **nested_sampler_kwargs) + sampler.run_nested(n_effective=nested_neff, + print_progress=verbose, + **nested_run_kwargs) + obj = sampler + + points = sampler.results["samples"] + log_w = sampler.results["logwt"] + log_like = sampler.results["logl"] + + # --- Nestle --- + if nested_sampler == "nestle": + import nestle + result = nestle.sample(likelihood_function, + model.prior_transform, + model.ndim, + **nested_sampler_kwargs) + obj = result + + points = result["samples"] + log_w = result["logwt"] + log_like = result["logl"] + + dur = time.time() - go + + return dict(points=points, log_weight=log_w, log_like=log_like, duration=dur), obj - return dsampler.results +# OMG +NESTED_KWARGS = { +"dynesty_sampler_kwargs" : dict(nlive=None, + bound='multi', + sample='auto', + #periodic=None, + #reflective=None, + update_interval=None, + first_update=None, + npdim=None, + #rstate=None, + queue_size=None, + pool=None, + use_pool=None, + #logl_args=None, + #logl_kwargs=None, + #ptform_args=None, + #ptform_kwargs=None, + #gradient=None, + #grad_args=None, + #grad_kwargs=None, + #compute_jac=False, + enlarge=None, + bootstrap=None, + walks=None, + facc=0.5, + slices=None, + fmove=0.9, + max_move=100, + #update_func=None, + ncdim=None, + blob=False, + #save_history=False, + #history_filename=None) + ), +"dynesty_run_kwargs" : dict(nlive_init=None, # nlive0 + maxiter_init=None, + maxcall_init=None, + dlogz_init=0.01, + logl_max_init=np.inf, + n_effective_init=np.inf, # deprecated + nlive_batch=None, #nlive0 + wt_function=None, + wt_kwargs=None, + maxiter_batch=None, + maxcall_batch=None, + maxiter=None, + maxcall=None, + maxbatch=None, + n_effective=None, + stop_function=None, + stop_kwargs=None, + #use_stop=True, + #save_bounds=True, # doesn't hurt...? + print_progress=True, + #print_func=None, + live_points=None, + #resume=False, + #checkpoint_file=None, + #checkpoint_every=60) + ), +"ultranest_sampler_kwargs":dict(transform=None, + #derived_param_names=[], + #wrapped_params=None, + #resume='subfolder', + #run_num=None, + #log_dir=None, + #num_test_samples=2, + draw_multiple=True, + num_bootstraps=30, + #vectorized=False, + ndraw_min=128, + ndraw_max=65536, + storage_backend='hdf5', + warmstart_max_tau=-1, + ), +"ultranest_run_kwargs":dict(update_interval_volume_fraction=0.8, + update_interval_ncall=None, + #log_interval=None, + show_status=True, + #viz_callback='auto', + dlogz=0.5, + dKL=0.5, + frac_remain=0.01, + Lepsilon=0.001, + min_ess=400, + max_iters=None, + max_ncalls=None, + max_num_improvement_loops=-1, + min_num_live_points=400, + cluster_num_live_points=40, + insertion_test_window=10, + insertion_test_zscore_threshold=4, + region_class="MLFriends", + #widen_before_initial_plateau_num_warn=10000, + #widen_before_initial_plateau_num_max=50000, + ), +"nautilus_sampler_kwargs":dict(n_live=2000, + n_update=None, + enlarge_per_dim=1.1, + n_points_min=None, + split_threshold=100, + n_networks=4, + neural_network_kwargs={}, + #prior_args=[], + #prior_kwargs={}, + #likelihood_args=[], + #likelihood_kwargs={}, + n_batch=None, + n_like_new_bound=None, + #vectorized=False, + #pass_dict=None, + pool=None, + seed=None, + blobs_dtype=None, + #filepath=None, + #resume=True + ), +"nautilus_run_kwargs":dict(f_live=0.01, + n_shell=1, + n_eff=10000, + n_like_max=np.inf, + discard_exploration=False, + timeout=np.inf, + verbose=False + ), +} diff --git a/prospect/io/write_results.py b/prospect/io/write_results.py index 2d27c67a..ad08cfb1 100644 --- a/prospect/io/write_results.py +++ b/prospect/io/write_results.py @@ -70,41 +70,43 @@ def paramfile_string(param_file=None, **extras): return pstr - -def write_hdf5(hfile, run_params, model, obs, - sampler=None, - optimize_result_list=None, - tsample=0.0, toptimize=0.0, - sampling_initial_center=[], +def write_hdf5(hfile, + config={}, + model=None, + obs=None, + sampling_result=None, + optimize_result_tuple=None, write_model_params=True, - sps=None, **extras): + sps=None, + **extras): """Write output and information to an HDF5 file object (or group). - :param hfile: + hfile : string or `h5py.File` File to which results will be written. Can be a string name or an `h5py.File` object handle. - :param run_params: + run_params : dict-like The dictionary of arguments used to build and fit a model. - :param model: - The `prospect.models.SedModel` object. + model : Instance of :py:class:`prospect.models.SpecModel` + The object. - :param obs: - The dictionary of observations that were fit. + obs : list of Observation() instances + The observations that were fit. - :param sampler: - The `emcee` or `dynesty` sampler object used to draw posterior samples. - Can be `None` if only optimization was performed. + sampling_result : EnsembleSampler() or dict + The `emcee` sampler used to draw posterior samples or nested sampler + output. Can be `None` if only optimization was performed. - :param optimize_result_list: + optimize_result_tuple : 2-tuple of (list, float) A list of `scipy.optimize.OptimizationResult` objects generated during - the optimization stage. Can be `None` if no optimization is performed + the optimization stage, and a float giving the duration of sampling. Can + be `None` if no optimization is performed - param sps: (optional, default: None) + sps : instance of :py:class:`prospect.sources.SSPBasis` (optional, default: None) If a `prospect.sources.SSPBasis` object is supplied, it will be used to - generate and store + generate and store best fit values (not implemented) """ # If ``hfile`` is not a file object, assume it is a filename and open if isinstance(hfile, str): @@ -112,12 +114,15 @@ def write_hdf5(hfile, run_params, model, obs, else: hf = hfile + assert (model is not None), "Must pass a prospector model" + run_params = config + # ---------------------- # Sampling info if run_params.get("emcee", False): - chain, extras = emcee_to_struct(sampler, model) - elif run_params.get("dynesty", False): - chain, extras = dynesty_to_struct(sampler, model) + chain, extras = emcee_to_struct(sampling_result, model) + elif bool(run_params.get("nested_sampler", False)): + chain, extras = nested_to_struct(sampling_result, model) else: chain, extras = None, None write_sampling_h5(hf, chain, extras) @@ -125,8 +130,9 @@ def write_hdf5(hfile, run_params, model, obs, # ---------------------- # Observational data - write_obs_to_h5(hf, obs) - hf.flush() + if obs is not None: + write_obs_to_h5(hf, obs) + hf.flush() # ---------------------- # High level parameter and version info @@ -134,15 +140,15 @@ def write_hdf5(hfile, run_params, model, obs, for k, v in meta.items(): hf.attrs[k] = v hf.flush() - hf.attrs['sampling_duration'] = json.dumps(tsample) # ----------------- # Optimizer info - hf.attrs['optimizer_duration'] = json.dumps(toptimize) - if optimize_result_list is not None: - out = optresultlist_to_ndarray(optimize_result_list) - mgroup = hf.create_group('optimization') - mdat = mgroup.create_dataset('optimizer_results', data=out) + if optimize_result_tuple is not None: + optimize_list, toptimize = optimize_result_tuple + optarr = optresultlist_to_ndarray(optimize_list) + opt = hf.create_group('optimization') + _ = opt.create_dataset('optimizer_results', data=optarr) + opt.attrs["optimizer_duration"] = json.dumps(toptimize) # --------------- # Best fitting model in space of data @@ -151,7 +157,7 @@ def write_hdf5(hfile, run_params, model, obs, pass #from ..plotting.utils import best_sample #pbest = best_sample(hf["sampling"]) - #spec, phot, mfrac = model.predict(pbest, obs=obs, sps=sps) + #predictions, mfrac = model.predict(pbest, obs=obs, sps=sps) #best = hf.create_group("bestfit") #best.create_dataset("spectrum", data=spec) #best.create_dataset("photometry", data=phot) @@ -168,6 +174,8 @@ def write_hdf5(hfile, run_params, model, obs, def metadata(run_params, model, write_model_params=True): + """Generate a metadata dictionary, with serialized entries. + """ meta = dict(run_params=run_params, paramfile_text=paramfile_string(**run_params)) if write_model_params: @@ -188,32 +196,31 @@ def emcee_to_struct(sampler, model): # preamble samples = sampler.get_chain(flat=True) lnprior = model.prior_product(samples) + lnpost = sampler.get_log_prob(flat=True) # chaincat & extras chaincat = chain_to_struct(samples, model=model) extras = dict(weights=None, - lnprobability=sampler.get_log_prob(flat=True), - lnlike=sampler.get_log_prob(flat=True) - lnprior, + lnprobability=lnpost, + lnlike=lnpost - lnprior, acceptance=sampler.acceptance_fraction, - rstate=sampler.random_state) + rstate=sampler.random_state, + duration=sampler.getattr("duration", 0.0)) return chaincat, extras -def dynesty_to_struct(dyout, model): +def nested_to_struct(nested_out, model): # preamble - lnprior = model.prior_product(dyout['samples']) + lnprior = model.prior_product(nested_out['points']) # chaincat & extras - chaincat = chain_to_struct(dyout["samples"], model=model) - extras = dict(weights=np.exp(dyout['logwt']-dyout['logz'][-1]), - lnprobability=dyout['logl'] + lnprior, - lnlike=dyout['logl'], - efficiency=np.atleast_1d(dyout['eff']), - logz=np.atleast_1d(dyout['logz']), - ncall=np.atleast_1d(dyout['ncall']) - #ncall=json.dumps(dyout['ncall'].tolist()) - ) + chaincat = chain_to_struct(nested_out['points'], model=model) + extras = dict(weights=np.exp(nested_out['log_weight']), + lnprobability=nested_out['log_like'] + lnprior, + lnlike=nested_out['log_like'], + duration=nested_out.get("duration", 0.0) + ) return chaincat, extras diff --git a/prospect/models/parameters.py b/prospect/models/parameters.py index 01b2509e..37624bef 100644 --- a/prospect/models/parameters.py +++ b/prospect/models/parameters.py @@ -10,13 +10,9 @@ from copy import deepcopy import warnings import numpy as np -import json, pickle from . import priors from .templates import describe -import scipy -from . import hyperparam_transforms as transforms -#from gp_sfh import * -#import gp_sfh_kernels + __all__ = ["ProspectorParams"] @@ -67,10 +63,10 @@ def __init__(self, configuration, verbose=True, param_order=None, **kwargs): """ self.init_config = deepcopy(configuration) self.parameter_order = param_order - if type(configuration) == list: + if isinstance(configuration, list): self.config_list = configuration self.config_dict = plist_to_pdict(self.config_list) - elif type(configuration) == dict: + elif isinstance(configuration, dict): self.config_dict = configuration self.config_list = pdict_to_plist(self.config_dict, order=param_order) else: @@ -188,7 +184,7 @@ def _prior_product(self, theta, **extras): parameter values. """ lnp_prior = 0 - + for k, inds in list(self.theta_index.items()): func = self.config_dict[k]['prior'] this_prior = np.sum(func(theta[..., inds]), axis=-1) @@ -209,10 +205,10 @@ def prior_transform(self, unit_coords): theta = np.zeros(len(unit_coords)) for k, inds in list(self.theta_index.items()): - + func = self.config_dict[k]['prior'].unit_transform theta[inds] = func(unit_coords[inds]) - + return theta def propagate_parameter_dependencies(self): diff --git a/prospect/utils/prospect_args.py b/prospect/utils/prospect_args.py index 15c3d31a..7e257a5e 100644 --- a/prospect/utils/prospect_args.py +++ b/prospect/utils/prospect_args.py @@ -14,7 +14,7 @@ def show_default_args(): parser.print_help() -def get_parser(fitters=["optimize", "emcee", "dynesty"]): +def get_parser(fitters=["optimize", "emcee", "nested"]): """Get a default prospector argument parser """ @@ -46,8 +46,8 @@ def get_parser(fitters=["optimize", "emcee", "dynesty"]): if "emcee" in fitters: parser = add_emcee_args(parser) - if "dynesty" in fitters: - parser = add_dynesty_args(parser) + if "nested" in fitters: + parser = add_nested_args(parser) return parser @@ -108,10 +108,22 @@ def add_emcee_args(parser): return parser -def add_dynesty_args(parser): +def add_nested_args(parser): # --- dynesty parameters --- - parser.add_argument("--dynesty", action="store_true", - help="If set, do nested sampling with dynesty.") + parser.add_argument("--nested_sampler", type=str, default="", + choices=["", "dynesty", "nautilus", "ultranest"], + help="do sampling with this sampler") + + parser.add_argument("--nested_nlive", dest="nested_nlive", type=int, default=1000, + help="Number of live points for the nested sampling run.") + + parser.add_argument("--nested_target_n_effective", type=int, default=10000, + help=("Stop when the number of *effective* posterior samples as estimated " + "by dynesty reaches the target number.")) + + parser.add_argument("--nested_dlogz", type=float, default=0.05, + help=("Stop the initial run when the remaining evidence is estimated " + "to be less than this.")) parser.add_argument("--nested_bound", type=str, default="multi", choices=["single", "multi", "balls", "cubes"], @@ -119,42 +131,17 @@ def add_dynesty_args(parser): "One of single | multi | balls | cubes")) parser.add_argument("--nested_sample", "--nested_method", type=str, dest="nested_sample", - default="slice", choices=["unif", "rwalk", "slice"], + default="auto", choices=["auto", "unif", "rwalk", "slice", "hslice"], help=("Method for drawing new points during sampling. " - "One of unif | rwalk | slice")) + "One of auto | unif | rwalk | slice | hslice")) parser.add_argument("--nested_walks", type=int, default=48, help=("Number of Metropolis steps to take when " "`nested_sample` is 'rwalk'")) - parser.add_argument("--nlive_init", dest="nested_nlive_init", type=int, default=100, - help="Number of live points for the intial nested sampling run.") - - parser.add_argument("--nlive_batch", dest="nested_nlive_batch", type=int, default=100, - help="Number of live points for the dynamic nested sampling batches") - - parser.add_argument("--nested_dlogz_init", type=float, default=0.05, - help=("Stop the initial run when the remaining evidence is estimated " - "to be less than this.")) - - parser.add_argument("--nested_maxcall", type=int, default=int(5e7), - help=("Maximum number of likelihood calls during nested sampling. " - "This will only be enforced after the initial pass")) - - parser.add_argument("--nested_maxiter", type=int, default=int(1e6), - help=("Maximum number of iterations during nested sampling. " - "This will only be enforced after the initial pass")) - - parser.add_argument("--nested_maxbatch", type=int, default=10, - help="Maximum number of dynamic batches.") - parser.add_argument("--nested_bootstrap", type=int, default=0, - help=("Number of bootstrap resamplings to use when estimating " - "ellipsoid expansion factor.")) - - parser.add_argument("--nested_target_n_effective", type=int, default=10000, - help=("Stop when the number of *effective* posterior samples as estimated " - "by dynesty reaches the target number.")) + help=("Number of bootstrap resamplings to use when estimating " + "ellipsoid expansion factor.")) return parser diff --git a/tests/tests_samplers.py b/tests/tests_samplers.py new file mode 100644 index 00000000..8c47d85f --- /dev/null +++ b/tests/tests_samplers.py @@ -0,0 +1,123 @@ +#!/usr/bin/env python +# -*- coding: utf-8 -*- + +import sys +import numpy as np + +#import pytest + +from prospect.utils.prospect_args import get_parser +from prospect.sources import CSPSpecBasis +from prospect.models import SpecModel, templates +from prospect.observation import Photometry +from prospect.fitting import fit_model +from prospect.fitting.nested import parse_nested_kwargs +from prospect.io.write_results import write_hdf5 +from prospect.io.read_results import results_from + +#@pytest.fixture +def get_sps(**kwargs): + sps = CSPSpecBasis(zcontinuous=1) + return sps + + +def build_model(add_neb=False, add_outlier=False, **kwargs): + model_params = templates.TemplateLibrary["parametric_sfh"] + model_params["logzsol"]["isfree"] = False # built for speed + if add_neb: # skip for speed + model_params.update(templates.TemplateLibrary["nebular"]) + if add_outlier: + model_params.update(templates.TemplateLibrary["outlier_model"]) + model_params["f_outlier_phot"]["isfree"] = True + model_params["f_outlier_phot"]["init"] = 0.05 + return SpecModel(model_params) + + +def build_obs(**kwargs): + + from astroquery.sdss import SDSS + from astropy.coordinates import SkyCoord + bands = "ugriz" + mcol = [f"cModelMag_{b}" for b in bands] + ecol = [f"cModelMagErr_{b}" for b in bands] + cat = SDSS.query_crossid(SkyCoord(ra=204.46376, dec=35.79883, unit="deg"), + data_release=16, + photoobj_fields=mcol + ecol + ["specObjID"]) + shdus = SDSS.get_spectra(plate=2101, mjd=53858, fiberID=220)[0] + assert int(shdus[2].data["SpecObjID"][0]) == cat[0]["specObjID"] + + fnames = [f"sdss_{b}0" for b in bands] + maggies = np.array([10**(-0.4 * cat[0][f"cModelMag_{b}"]) for b in bands]) + magerr = np.array([cat[0][f"cModelMagErr_{b}"] for b in bands]) + magerr = np.hypot(magerr, 0.05) + + phot = Photometry(filters=fnames, flux=maggies, uncertainty=magerr*maggies/1.086, + name=f'sdss_phot_specobjID{cat[0]["specObjID"]}') + + obslist = [phot] + [obs.rectify() for obs in obslist] + return obslist + + +if __name__ == "__main__": + + parser = get_parser() + parser.set_defaults(nested_target_n_effective=256, + nested_nlive=512, + verbose=0) + args = parser.parse_args() + run_params = vars(args) + run_params["param_file"] = __file__ + + # test the parsing + run_params["nested_sampler"] = "dynesty" + nr, ns = parse_nested_kwargs(**run_params) + print(nr) + print(ns) + + # build stuff + model = build_model() + obs = build_obs() + sps = get_sps() + + # do a first model caching + (phot,), mfrac = model.predict(model.theta, obs, sps=sps) + print(model) + + # loop over samplers + results = {} + samplers = ["nautilus", "ultranest", "dynesty"] + for sampler in samplers: + print(sampler) + run_params["nested_sampler"] = sampler + out = fit_model(obs, model, sps, **run_params) + results[sampler] = out["sampling"] + hfile = f"./{sampler}_test.h5" + write_hdf5(hfile, + run_params, + model, + obs, + results[sampler], + None, + sps=sps) + ires, iobs, im = results_from(hfile) + assert (im is not None) + + # compare runtime + for sampler in samplers: + print(sampler, results[sampler]["duration"]) + + # compare posteriors + colors = ["royalblue", "darkorange", "firebrick"] + import matplotlib.pyplot as pl + from prospect.plotting import corner + ndim = model.ndim + cfig, axes = pl.subplots(ndim, ndim, figsize=(10,9)) + for sampler, color in zip(samplers, colors): + out = results[sampler] + axes = corner.allcorner(out["points"].T, + model.theta_labels(), + axes, + color=color, + weights= np.exp(out["log_weight"]), + show_titles=True) \ No newline at end of file