bd-j / prospector

Python code for Stellar Population Inference from Spectra and SEDs
http://prospect.readthedocs.io
MIT License
159 stars 75 forks source link

Problem with the StochasticSFH model #342

Closed jensmelinder closed 2 months ago

jensmelinder commented 2 months ago

Hi, I'm having some problems running the StochasticSFH binned SFH treatment. My code runs fine with the parametric fitting and the binned Dirichlet mode SFH, but with StochasticSFH I get a warning (see below) and then the code fails with an error message. I'm not sure these are related (it could be me giving the wrong inputs to initialize the model).

I'm fitting a SED for a high-z galaxy with known spectroscopic redshift, and I'm thus using the priors and recommendations from Wan+24, with the exception of dust where I use the simpler 1-parameter model (dust:0 in FSPS), and keeping the gas parameters fixed.

I'm using prospector v1.4.0.

My code for building the model is:

def build_model_stochastic(object_redshift=0.0, fixed_metallicity=None, add_duste=False,
                add_neb=False, luminosity_distance=0.0, **extras):
    """Construct a nonparametric model.  This method defines a number of parameter
    specification dictionaries and uses them to initialize a
    `models.sedmodel.SedModel` object.

    :param object_redshift:
        If given, given the model redshift to this value.

    :param add_dust: (optional, default: False)
        Switch to add (fixed) parameters relevant for dust emission.

    :param add_neb: (optional, default: False)
        Switch to add (fixed) parameters relevant for nebular emission, and
        turn nebular emission on.

    :param luminosity_distance: (optional)
        If present, add a `"lumdist"` parameter to the model, and set it's
        value (in Mpc) to this.  This allows one to decouple redshift from
        distance, and fit, e.g., absolute magnitudes (by setting
        luminosity_distance to 1e-5 (10pc))
    """

    # Non-parametric template, Stochastic SFH
    model_params = TemplateLibrary["stochastic_sfh"]

    # Add lumdist parameter.  If this is not added then the distance is
    # controlled by the "zred" parameter and a WMAP9 cosmology.
    if luminosity_distance > 0:
        model_params["lumdist"] = {"N": 1, "isfree": False,
                                   "init": luminosity_distance, "units": "Mpc"}

    # set age bins for non-parametric fit
    agelims= np.r_[1e6,5e6,np.logspace(7,np.log10(300e6),6)]
    agelims= np.log10(agelims)
    agebin_lolims= agelims[:-1]
    agebin_uplims= agelims[1:]
    agebins = [[agebin_lolims[ii],agebin_uplims[ii]] for ii in range(len(agebin_lolims))] 
    model_params['agebins']['init'] = agebins
    if object_redshift != 0.0:
        # make sure zred is fixed
        model_params["zred"]['isfree'] = False
        # And set the value to the object_redshift keyword
        model_params["zred"]['init'] = object_redshift
    if add_neb:
        # Add nebular emission (with fixed parameters)
        model_params.update(TemplateLibrary["nebular"])
        model_params["gas_logz"]["init"] = -1.
        model_params["gas_logu"]["init"] = -2.
    if add_duste:
        # Add dust emission (with fixed dust SED parameters)
        model_params.update(TemplateLibrary["dust_emission"])

    model_params.update(TemplateLibrary["igm"])
    model_params["igm_factor"]["isfree"] = False

    # High-z prior on metallicity from , Wan+24, Tachella+23
    model_params["logzsol"]["prior"] = priors.ClippedNormal(mini=-2,maxi=0.0,mean=-1.5,sigma=0.5)
    # mass prior, somewhat lower limit, Wan+24
    model_params["logmass"]["prior"] = priors.Uniform(mini=6.0,maxi=12.0)

    #Priors on Stochastic SFR params, from Wan+24
    model_params["sigma_reg"]["prior"] = priors.LogUniform(mini=0.1,maxi=5)
    model_params["tau_eq"]["prior"] = priors.LogUniform(mini=0.05,maxi=0.5)
    model_params["sigma_dyn"]["prior"] = priors.LogUniform(mini=0.01,maxi=1.0)
    model_params["tau_dyn"]["prior"] = priors.LogUniform(mini=0.001,maxi=0.02)
    model_params["tau_in"]['init'] = cosm.age(object_redshift).value # Hubble time at z (Gyr)

    model_params["dust2"]["prior"] = priors.ClippedNormal(mini=0,maxi=4.0,mean=0.,sigma=0.5)

    # Now instantiate the model using this new dictionary of parameter specifications
    model_params = prospect.models.templates.adjust_stochastic_params(model_params)
    model = sedmodel.SpecModel(model_params)

    return model

Running this function throws the warning:

 RuntimeWarning: logsfr_ratios has wrong length prior, should be 6
  warnings.warn(msg.format(p, n), RuntimeWarning)

However, as far as I can see the prior seems to have the correct length. It also throws the same error with the default parameters. As far as I understand, the number of logsfr_ratios (and priors on them) should be one less than the number of agebins (so 6 ratios and 7 agebins in my case).

I include nebular emission and IGM transmission in the model. I also have a build_obs function that reads the data (this seems to be ok since the parametric and Dirichlet SFH runs work fine), and I initialize FSPS with:

    def build_sps_nonpar(zcontinuous=1, **extras):
        sps = FastStepBasis(zcontinuous=zcontinuous)
    return sps

I then run the actual fitting with:

    obs, model, sps, noise = build_all(**config) # runs build_model, build_obs, and build_sps_nonpar
    config["sps_libraries"] = sps.ssp.libraries
    print(model)

    # set dynesty pars
    config['nlive_init']=400
    config['nested_method']='rwalk'
    config['nested_target_n_effective']=1000
    config['nested_dlogz_init']=0.05

    if args.debug:
        sys.exit()

    # --- Set up output ---
    ts = time.strftime("%y%b%d-%H.%M", time.localtime())
    hfile = f"{args.outfile}_{ts}_result.h5"

    #  --- Run the actual fit ---
    output = fit_model(obs, model, sps, **config)

The run then fails with this error message (also including the model params in the beginning):

<class 'prospect.models.sedmodel.SpecModel'>

Free Parameters: (name: prior) 
-----------
  logzsol: <class 'prospect.models.priors.ClippedNormal'>(mean=-1.5,sigma=0.5,mini=-2,maxi=0.0)
  dust2: <class 'prospect.models.priors.ClippedNormal'>(mean=0.0,sigma=0.5,mini=0,maxi=4.0)
  logmass: <class 'prospect.models.priors.Uniform'>(mini=6.0,maxi=12.0)
  sigma_reg: <class 'prospect.models.priors.LogUniform'>(mini=0.1,maxi=5)
  tau_eq: <class 'prospect.models.priors.LogUniform'>(mini=0.05,maxi=0.5)
  sigma_dyn: <class 'prospect.models.priors.LogUniform'>(mini=0.01,maxi=1.0)
  tau_dyn: <class 'prospect.models.priors.LogUniform'>(mini=0.001,maxi=0.02)
  logsfr_ratios: <class 'prospect.models.priors.MultiVariateNormal'>(mean=[0. 0. 0. 0. 0. 0.],Sigma=[[ 3.71928895e-05 -1.39799705e-06 -1.11035102e-07  5.16272799e-06
   1.57321908e-05  2.79403436e-05]
 [-1.39799705e-06  5.29313065e-05 -2.20217602e-06  5.73263833e-06
   2.22686884e-05  4.09477963e-05]
 [-1.11035102e-07 -2.20217602e-06  1.00425345e-04  2.51959040e-06
   4.09051336e-05  8.22972246e-05]
 [ 5.16272799e-06  5.73263833e-06  2.51959040e-06  1.94104049e-04
   5.37636955e-05  1.68421019e-04]
 [ 1.57321908e-05  2.22686884e-05  4.09051336e-05  5.37636955e-05
   4.11967871e-04  3.04491000e-04]
 [ 2.79403436e-05  4.09477963e-05  8.22972246e-05  1.68421019e-04
   3.04491000e-04  1.07277772e-03]])

Fixed Parameters: (name: value [, depends_on]) 
-----------
  zred: [10] 
  mass: [1000000.] <function logsfr_ratios_to_masses at 0x7f5bcfbd5a20>
  sfh: [3] 
  imf_type: [2] 
  dust_type: [0] 
  agebins: [[6.         6.69897   ]
 [6.69897    7.        ]
 [7.         7.29542425]
 [7.29542425 7.5908485 ]
 [7.5908485  7.88627275]
 [7.88627275 8.181697  ]
 [8.181697   8.47712125]] 
  tau_in: [0.4792] 
  add_neb_emission: [ True] 
  add_neb_continuum: [ True] 
  nebemlineinspec: [ True] 
  gas_logz: [-1.] <function stellar_logzsol at 0x7f5c07c1f400>
  gas_logu: [-2.] 
  add_igm_absorption: [ True] 
  igm_factor: [1.] 
/disk/miri-b1/jeand/tools/miniconda3/envs/jwst_39/lib/python3.10/site-packages/prospect/models/priors.py:125: RuntimeWarning: divide by zero encountered in log
  lnp = np.log(p)
Exception while calling loglikelihood function:
  params: [-1.54365273e+00  8.04830203e-01  1.02824321e+01  1.65776724e-01
  4.46876046e-01  7.48602220e-02  1.79261129e-03  8.51188379e-03
 -7.60758646e-04 -6.03178340e-03 -1.15966502e-02  2.61531753e-02
  1.34671979e-02]
  args: []
  kwargs: {}
  exception:
Traceback (most recent call last):
  File "/disk/miri-b1/jeand/tools/miniconda3/envs/jwst_39/lib/python3.10/site-packages/dynesty/dynesty.py", line 913, in __call__
    return self.func(np.asarray(x).copy(), *self.args, **self.kwargs)
  File "/disk/miri-b1/jeand/tools/miniconda3/envs/jwst_39/lib/python3.10/site-packages/prospect/fitting/fitting.py", line 94, in lnprobfn
    if not np.isfinite(lnp_prior):
ValueError: The truth value of an array with more than one element is ambiguous. Use a.any() or a.all()
Traceback (most recent call last):
  File "/disk/miri-b1/jeand/sedfitting/prospector/run_prosp.py", line 359, in <module>
    output = fit_model(obs, model, sps, **config)
  File "/disk/miri-b1/jeand/tools/miniconda3/envs/jwst_39/lib/python3.10/site-packages/prospect/fitting/fitting.py", line 259, in fit_model
    output["sampling"] = run_sampler(obs, model, sps, noise,
  File "/disk/miri-b1/jeand/tools/miniconda3/envs/jwst_39/lib/python3.10/site-packages/prospect/fitting/fitting.py", line 532, in run_dynesty
    dynestyout = run_dynesty_sampler(lnp, model.prior_transform, model.ndim,
  File "/disk/miri-b1/jeand/tools/miniconda3/envs/jwst_39/lib/python3.10/site-packages/prospect/fitting/nested.py", line 84, in run_dynesty_sampler
    for results in dsampler.sample_initial(nlive=nested_nlive_init,
  File "/disk/miri-b1/jeand/tools/miniconda3/envs/jwst_39/lib/python3.10/site-packages/dynesty/dynamicsampler.py", line 1289, in sample_initial
    blobs), logvol_init, init_ncalls = _initialize_live_points(
  File "/disk/miri-b1/jeand/tools/miniconda3/envs/jwst_39/lib/python3.10/site-packages/dynesty/dynamicsampler.py", line 438, in _initialize_live_points
    cur_live_logl = loglikelihood.map(np.asarray(cur_live_v))
  File "/disk/miri-b1/jeand/tools/miniconda3/envs/jwst_39/lib/python3.10/site-packages/dynesty/utils.py", line 171, in map
    ret = list([
  File "/disk/miri-b1/jeand/tools/miniconda3/envs/jwst_39/lib/python3.10/site-packages/dynesty/utils.py", line 171, in <listcomp>
    ret = list([
  File "/disk/miri-b1/jeand/tools/miniconda3/envs/jwst_39/lib/python3.10/site-packages/dynesty/dynesty.py", line 913, in __call__
    return self.func(np.asarray(x).copy(), *self.args, **self.kwargs)
  File "/disk/miri-b1/jeand/tools/miniconda3/envs/jwst_39/lib/python3.10/site-packages/prospect/fitting/fitting.py", line 94, in lnprobfn
    if not np.isfinite(lnp_prior):
ValueError: The truth value of an array with more than one element is ambiguous. Use a.any() or a.all()

I'm guessing this indicates that there are problems with the calculated values in lnp_prior, but I'm at a loss trying to figure out what went wrong, and if it is related to the earlier runtime warning. Would be grateful for any help with this.

bd-j commented 2 months ago

Hi @jensmelinder (cc @jwan)

For the stochastic SFH prior you need to use a special SpecModel subclass that overrides the prior_product method:

from prospect.models import sedmodel, templates
model_params = template.TemplateLibrary["stochastic_sfh"]
model = sedmodel.HyperSpecModel(model_params)
print(model.prior_product(model.theta))

In v1.4.0 this yields a scalar for me, which should fix the error during fitting. We should improve the documentation to mention this. Note that HyperSpecModel is just

from prospect.models.hyperparameters import ProspectorHyperParams
from prospect.models.sedmodel import SpecModel
class HyperSpecModel(ProspectorHyperParams, SpecModel):
    pass

The warning has to do with a bug in how prior lengths are computed, which because of the covariance matrix for logsfr_ratios returns (N_bin -1)**2 instead of just N_bin-1 for that prior. I think it can be safely ignored.

Also, just want to make sure that you really want to ignore stars younger than 1 Myr, which can be important for the SED, with the agebins that you are using.

Hope this helps, Ben

jensmelinder commented 2 months ago

Thanks a lot! I figured that there was some documentation missing there. Will try with HyperSpecModel.

Yeah, I was just testing with the agebins and wanted something non-zero to not get log errors (but perhaps the code deals with that anyway).

bd-j commented 2 months ago

Yes, 0 corresponds to 1 year - the code extrapolates the youngest available SSP to cover these very recent ages.

jensmelinder commented 2 months ago

The fitting worked just fine after making the change. Feel free to close this (I'll check and close it within a few days otherwise). Thanks again for the prompt reply.

Just one more question about the Stochastic SFH if I may. To transform the logsfr_ratios into SFRs can I use the Continuity SFH transform function (transforms.logsfr_ratios_to_sfrs)? I couldn't find a convenience function for this for the Stochastic SFH, but I don't think that conversion depends on the prior so it should be ok?

bd-j commented 2 months ago

Great, glad it worked. And yes, it should be ok to use the regular transform for SFRs, as you say the conversion itself does not depend on the prior.