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/fitting.py b/prospect/fitting/fitting.py index 27c107d7..70e8cf34 100755 --- a/prospect/fitting/fitting.py +++ b/prospect/fitting/fitting.py @@ -185,13 +185,18 @@ def fit_model(observations, model, sps, lnprobfn=lnprobfn, # Make sure obs has required keys [obs.rectify() for obs in observations] + 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 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,7 +209,8 @@ def fit_model(observations, model, sps, lnprobfn=lnprobfn, run_sampler = run_ensemble elif nested_sampler: run_sampler = run_nested - kwargs["fitter"] = nested_sampler + # put nested_sampler back into kwargs for lower level functions + kwargs["nested_sampler"] = nested_sampler else: return output @@ -412,14 +418,20 @@ def run_ensemble(observations, model, sps, lnprobfn=lnprobfn, sampler, burn_loc0, burn_prob0 = out ts = time.time() - go - return sampler, ts + try: + sampler.duration = ts + except: + pass + + return sampler def run_nested(observations, model, sps, lnprobfn=lnprobfn, - fitter="dynesty", + nested_sampler="dynesty", nested_nlive=1000, nested_neff=1000, + verbose=False, **kwargs): """Thin wrapper on :py:class:`prospect.fitting.nested.run_nested_sampler` @@ -457,18 +469,16 @@ def run_nested(observations, model, sps, likelihood = wrap_lnp(lnprobfn, observations, model, sps, nested=True) # which keywords do we have for this fitter? - ns_kwargs, nr_kwargs = parse_nested_kwargs(fitter=fitter, - **kwargs) + ns_kwargs, nr_kwargs = parse_nested_kwargs(nested_sampler=nested_sampler,**kwargs) - go = time.time() output = run_nested_sampler(model, likelihood, - fitter=fitter, - verbose=False, + nested_sampler=nested_sampler, + verbose=verbose, nested_nlive=nested_nlive, nested_neff=nested_neff, nested_sampler_kwargs=ns_kwargs, nested_run_kwargs=nr_kwargs) - ts = time.time() - go + info, result_obj = output - return output, ts + return info diff --git a/prospect/fitting/nested.py b/prospect/fitting/nested.py index 80083e34..3b467bf8 100644 --- a/prospect/fitting/nested.py +++ b/prospect/fitting/nested.py @@ -1,35 +1,34 @@ import time import numpy as np - __all__ = ["run_nested_sampler", "parse_nested_kwargs"] -def parse_nested_kwargs(fitter=None, **kwargs): +def parse_nested_kwargs(nested_sampler=None, **kwargs): - # todo: + # TODO: # something like 'enlarge' # something like 'bootstrap' or N_networks or? sampler_kwargs = {} run_kwargs = {} - if fitter == "dynesty": + if nested_sampler == "dynesty": sampler_kwargs["bound"] = kwargs["nested_bound"] - sampler_kwargs["method"] = kwargs["nested_sample"] + sampler_kwargs["sample"] = kwargs["nested_sample"] sampler_kwargs["walks"] = kwargs["nested_walks"] run_kwargs["dlogz_init"] = kwargs["nested_dlogz"] - elif fitter == "ultranest": + elif nested_sampler == "ultranest": #run_kwargs["dlogz"] = kwargs["nested_dlogz"] pass - elif fitter == "nautilus": + elif nested_sampler == "nautilus": pass else: # say what? - raise ValueError(f"{fitter} not a valid fitter") + raise ValueError(f"{nested_sampler} not a valid fitter") return sampler_kwargs, run_kwargs @@ -37,7 +36,7 @@ def parse_nested_kwargs(fitter=None, **kwargs): def run_nested_sampler(model, likelihood_function, - fitter="dynesty", + nested_sampler="dynesty", nested_nlive=1000, nested_neff=1000, verbose=False, @@ -45,29 +44,41 @@ def run_nested_sampler(model, 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. """ go = time.time() # --- Nautilus --- - if fitter == "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 fitter == "ultranest": + if nested_sampler == "ultranest": from ultranest import ReactiveNestedSampler - sampler = ReactiveNestedSampler(model.theta_labels, + parameter_names = model.theta_labels() + sampler = ReactiveNestedSampler(parameter_names, likelihood_function, model.prior_transform, **nested_sampler_kwargs) @@ -75,13 +86,14 @@ def run_nested_sampler(model, 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 fitter == "dynesty": + if nested_sampler == "dynesty": from dynesty import DynamicNestedSampler sampler = DynamicNestedSampler(likelihood_function, @@ -92,18 +104,20 @@ def run_nested_sampler(model, 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 fitter == "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"] @@ -111,12 +125,11 @@ def run_nested_sampler(model, dur = time.time() - go - return dict(points=points, log_weight=log_w, log_like=log_like) + return dict(points=points, log_weight=log_w, log_like=log_like, duration=dur), obj # OMG -SAMPLER_KWARGS = { - +NESTED_KWARGS = { "dynesty_sampler_kwargs" : dict(nlive=None, bound='multi', sample='auto', diff --git a/prospect/io/write_results.py b/prospect/io/write_results.py index 79743d10..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("nested", False): - chain, extras = nested_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,14 +196,16 @@ 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 @@ -208,7 +218,8 @@ def nested_to_struct(nested_out, model): 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'] + lnlike=nested_out['log_like'], + duration=nested_out.get("duration", 0.0) ) return chaincat, extras diff --git a/tests/tests_samplers.py b/tests/tests_samplers.py new file mode 100644 index 00000000..b139bfa3 --- /dev/null +++ b/tests/tests_samplers.py @@ -0,0 +1,118 @@ +#!/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 + + +#@pytest.fixture +def get_sps(): + sps = CSPSpecBasis(zcontinuous=1) + return sps + + +def build_model(add_neb=False, add_outlier=False): + 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) + args = parser.parse_args() + run_params = vars(args) + run_params["parameter"] + + # 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) + + for sampler in samplers: + print(sampler, results[sampler]["duration"]) + + + 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