diff --git a/src/syncfit/_version.py b/src/syncfit/_version.py index 260c070..f9aa3e1 100644 --- a/src/syncfit/_version.py +++ b/src/syncfit/_version.py @@ -1 +1 @@ -__version__ = "0.3.1" +__version__ = "0.3.2" diff --git a/src/syncfit/analysis/plotter.py b/src/syncfit/analysis/plotter.py index 91e62bd..ea8fa50 100644 --- a/src/syncfit/analysis/plotter.py +++ b/src/syncfit/analysis/plotter.py @@ -38,8 +38,8 @@ def plot_chains(sampler, labels, fig=None, axes=None): return fig, axes -def plot_best_fit(model, sampler, nu, F, Ferr, nkeep=1000, lum_dist=None, - nu_arr=None, method='random', fig=None, ax=None, day=None): +def plot_best_fit(model, sampler, nu, F, Ferr, nkeep=1000, + nu_arr=None, method='random', fig=None, ax=None, day=None, **kwargs): ''' Plot best fit model @@ -62,8 +62,11 @@ def plot_best_fit(model, sampler, nu, F, Ferr, nkeep=1000, lum_dist=None, matplotlib fig, ax ''' if isinstance(model, MQModel) and model.t is not None: - t = model.t - + kwargs['t'] = model.t + + if model.p is not None: + kwargs['p'] = model.p + flat_samples, log_prob = extract_output(sampler) if method == 'max': @@ -82,19 +85,10 @@ def plot_best_fit(model, sampler, nu, F, Ferr, nkeep=1000, lum_dist=None, if ax is None: fig, ax = plt.subplots(figsize=(4,4)) - - if lum_dist is not None and model.t is None: - kwargs = dict(lum_dist=lum_dist) - elif lum_dist is not None and model.t is not None: - kwargs = dict(lum_dist=lum_dist, t=t) - else: - kwargs={} for val in toplot: - if model.p is not None: - res = model.SED(nu_plot, model.p, *val, **kwargs) - else: - res = model.SED(nu_plot, *val, **kwargs) + packed_theta = model.pack_theta(val, **kwargs) + res = model.SED(nu_plot, **packed_theta) ax.plot(nu_plot, res, '-', color='cornflowerblue', lw = 0.5, alpha = 0.1) diff --git a/src/syncfit/mcmc.py b/src/syncfit/mcmc.py index 1935cb2..1f0f40f 100644 --- a/src/syncfit/mcmc.py +++ b/src/syncfit/mcmc.py @@ -8,6 +8,7 @@ import emcee import dynesty from multiprocessing import Pool +from warnings import warn from .analysis import * from .models.mq_model import MQModel from .models.syncfit_model import SyncfitModel @@ -48,11 +49,11 @@ def do_dynesty(nu:list[float], F_mJy:list[float], F_error:list[float], test_model = model() # just for now if isinstance(test_model, MQModel) and (lum_dist is None): raise ValueError('lum_dist and t reequired for MQModel!') - + if isinstance(test_model, MQModel): - model = model(p=fix_p, t=t) + model = model(prior=prior, p=fix_p, t=t) else: - model = model(p=fix_p) + model = model(prior=prior, p=fix_p) # get the extra args dynesty_args = model.get_kwargs(nu, F_mJy, F_error, lum_dist=lum_dist, t=t) diff --git a/src/syncfit/models/b1b2_b3b4_weighted_model.py b/src/syncfit/models/b1b2_b3b4_weighted_model.py index d8e0850..5c2f3d4 100644 --- a/src/syncfit/models/b1b2_b3b4_weighted_model.py +++ b/src/syncfit/models/b1b2_b3b4_weighted_model.py @@ -10,21 +10,24 @@ class B1B2_B3B4_Weighted(SyncfitModel): the B3B4 model. The idea of this model is from XXXYXYXYX et al. (YYYY). ''' - def __init__(self, p=None): + def __init__(self, prior=None, p=None): # then set the default prior for this model - if p is None: - self.prior = dict( - p=[2,4], - log_F_nu=[-4,2], - log_nu_a=[6,12], - log_nu_m=[6,12] - ) + if prior is None: + if p is None: + self.prior = dict( + p=[2,4], + log_F_nu=[-4,2], + log_nu_a=[6,12], + log_nu_m=[6,12] + ) + else: + self.prior = dict( + log_F_nu=[-4,2], + log_nu_a=[6,12], + log_nu_m=[6,12] + ) else: - self.prior = dict( - log_F_nu=[-4,2], - log_nu_a=[6,12], - log_nu_m=[6,12] - ) + self.prior = prior super().__init__(self.prior, p=p) diff --git a/src/syncfit/models/b1b2_model.py b/src/syncfit/models/b1b2_model.py index 2c3430f..5d0b6f7 100644 --- a/src/syncfit/models/b1b2_model.py +++ b/src/syncfit/models/b1b2_model.py @@ -10,23 +10,26 @@ class B1B2(SyncfitModel): (nu_m). This model uses nu_m > nu_a, the opposite of the B4B5 model. ''' - def __init__(self, p=None): + def __init__(self, prior=None, p=None): # then set the default prior for this model - if p is None: - self.prior = dict( - p=[2,4], - log_F_nu=[-4,2], - log_nu_a=[6,11], - log_nu_m=[7,15] - ) + if prior is None: + if p is None: + self.prior = dict( + p=[2,4], + log_F_nu=[-4,2], + log_nu_a=[6,11], + log_nu_m=[7,15] + ) + else: + self.prior = dict( + log_F_nu=[-4,2], + log_nu_a=[6,11], + log_nu_m=[7,15] + ) else: - self.prior = dict( - log_F_nu=[-4,2], - log_nu_a=[6,11], - log_nu_m=[7,15] - ) + self.prior = prior - super().__init__(p=p) + super().__init__(self.prior, p=p) # the model, must be named SED!!! def SED(self, nu, p, log_F_nu, log_nu_a, log_nu_m, **kwargs): @@ -51,7 +54,7 @@ def lnprior(self, theta, nu, F, upperlimit, **kwargs): Logarithmic prior function that can be changed based on the SED model. ''' uppertest = SyncfitModel._is_below_upperlimits( - nu, F, upperlimit, theta, self.SED + nu, F, upperlimit, theta, self.SED, **kwargs ) packed_theta = self.pack_theta(theta) diff --git a/src/syncfit/models/b4b5_model.py b/src/syncfit/models/b4b5_model.py index 3f1ca97..21c0296 100644 --- a/src/syncfit/models/b4b5_model.py +++ b/src/syncfit/models/b4b5_model.py @@ -11,24 +11,27 @@ class B4B5(SyncfitModel): subclass this class and overwrite lnprior to redefine this. ''' - def __init__(self, p=None): + def __init__(self, prior=None, p=None): # then set the default prior for this model - if p is None: - self.prior = dict( - p=[2,4], - log_F_nu=[-4,2], - log_nu_a=[6,11], - log_nu_m=[-np.inf,6] - ) + if prior is None: + if p is None: + self.prior = dict( + p=[2,4], + log_F_nu=[-4,2], + log_nu_a=[6,11], + log_nu_m=[-np.inf,6] + ) + else: + self.prior = dict( + log_F_nu=[-4,2], + log_nu_a=[6,11], + log_nu_m=[-np.inf,6] + ) else: - self.prior = dict( - log_F_nu=[-4,2], - log_nu_a=[6,11], - log_nu_m=[-np.inf,6] - ) + self.prior = prior - super().__init__(p=p) + super().__init__(self.prior, p=p) # the model, must be named SED!!! def SED(self, nu, p, log_F_nu, log_nu_a, log_nu_m, **kwargs): @@ -53,7 +56,7 @@ def lnprior(self, theta, nu, F, upperlimit, **kwargs): Logarithmic prior function that can be changed based on the SED model. ''' uppertest = SyncfitModel._is_below_upperlimits( - nu, F, upperlimit, theta, self.SED + nu, F, upperlimit, theta, self.SED, **kwargs ) packed_theta = self.pack_theta(theta) diff --git a/src/syncfit/models/b4b5b3_model.py b/src/syncfit/models/b4b5b3_model.py index 30f98ab..2a0153e 100644 --- a/src/syncfit/models/b4b5b3_model.py +++ b/src/syncfit/models/b4b5b3_model.py @@ -10,24 +10,27 @@ class B4B5B3(SyncfitModel): and minimum energy break (nu_m). This model always requires that nu_m < nu_a < nu_c. ''' - def __init__(self, p=None): + def __init__(self, prior=None, p=None): # then set the default prior for this model - if p is None: - self.prior = dict( - p=[2,4], - log_F_nu=[-4,2], - log_nu_a=[6,11], - log_nu_m=[0, 6], - log_nu_c=[7,15] - ) + if prior is None: + if p is None: + self.prior = dict( + p=[2,4], + log_F_nu=[-4,2], + log_nu_a=[6,11], + log_nu_m=[0, 6], + log_nu_c=[7,15] + ) + else: + self.prior = dict( + log_F_nu=[-4,2], + log_nu_a=[6,11], + log_nu_m=[0, 6], + log_nu_c=[7,15] + ) else: - self.prior = dict( - log_F_nu=[-4,2], - log_nu_a=[6,11], - log_nu_m=[0, 6], - log_nu_c=[7,15] - ) - + self.prior = prior + super().__init__(self.prior, p=p) # the model, must be named SED!!! @@ -58,7 +61,7 @@ def lnprior(self, theta, nu, F, upperlimit, **kwargs): Logarithmic prior function that can be changed based on the SED model. ''' uppertest = SyncfitModel._is_below_upperlimits( - nu, F, upperlimit, theta, self.SED + nu, F, upperlimit, theta, self.SED, **kwargs ) packed_theta = self.pack_theta(theta) diff --git a/src/syncfit/models/b5_model.py b/src/syncfit/models/b5_model.py index 8c64aec..8e14c9a 100644 --- a/src/syncfit/models/b5_model.py +++ b/src/syncfit/models/b5_model.py @@ -9,19 +9,22 @@ class B5(SyncfitModel): Single break model for just the self-absorption break. ''' - def __init__(self, p=None): + def __init__(self, prior=None, p=None): # then set the default prior for this model - if p is None: - self.prior = dict( - p=[2,4], - log_F_nu=[-4,2], - log_nu_a=[6,11] - ) + if prior is None: + if p is None: + self.prior = dict( + p=[2,4], + log_F_nu=[-4,2], + log_nu_a=[6,11] + ) + else: + self.prior = dict( + log_F_nu=[-4,2], + log_nu_a=[6,11] + ) else: - self.prior = dict( - log_F_nu=[-4,2], - log_nu_a=[6,11] - ) + self.prior = prior super().__init__(self.prior, p=p) diff --git a/src/syncfit/models/b5b3_model.py b/src/syncfit/models/b5b3_model.py index 44ddf18..bdd43db 100644 --- a/src/syncfit/models/b5b3_model.py +++ b/src/syncfit/models/b5b3_model.py @@ -11,21 +11,24 @@ class B5B3(SyncfitModel): break. ''' - def __init__(self, p=None): + def __init__(self, prior=None, p=None): # then set the default prior for this model - if p is None: - self.prior = dict( - p=[2,4], - log_F_nu=[-4,2], - log_nu_a=[6,11], - log_nu_c=[7,15] - ) + if prior is None: + if p is None: + self.prior = dict( + p=[2,4], + log_F_nu=[-4,2], + log_nu_a=[6,11], + log_nu_c=[7,15] + ) + else: + self.prior = dict( + log_F_nu=[-4,2], + log_nu_a=[6,11], + log_nu_c=[7,15] + ) else: - self.prior = dict( - log_F_nu=[-4,2], - log_nu_a=[6,11], - log_nu_c=[7,15] - ) + self.prior = prior super().__init__(self.prior, p=p) @@ -52,7 +55,7 @@ def lnprior(self, theta, nu, F, upperlimit, **kwargs): Logarithmic prior function that can be changed based on the SED model. ''' uppertest = SyncfitModel._is_below_upperlimits( - nu, F, upperlimit, theta, self.SED + nu, F, upperlimit, theta, self.SED, **kwargs ) packed_theta = self.pack_theta(theta) diff --git a/src/syncfit/models/mq_model.py b/src/syncfit/models/mq_model.py index c19f9b8..c0ede70 100644 --- a/src/syncfit/models/mq_model.py +++ b/src/syncfit/models/mq_model.py @@ -13,27 +13,30 @@ class MQModel(SyncfitModel): - def __init__(self, p=None, t=None): - # set the default prior for this model - self.prior = dict( - log_bG_sh=[-3,3], - log_Mdot=[-10,0], - log_epsilon_e=[-3,0], - log_epsilon_B=[-3,0], - log_epsilon_T=[-3,0] - ) - - if p is None: - self.prior['p'] = [2,4] - - if t is None: - self.prior['t'] = [0.1, 1_000] # super wide prior as default - + def __init__(self, prior=None, p=None, t=None): + if prior is None: + # set the default prior for this model + self.prior = dict( + log_bG_sh=[-3,3], + log_Mdot=[-10,0], + log_epsilon_e=[-3,0], + log_epsilon_B=[-3,0], + log_epsilon_T=[-3,0] + ) + + if p is None: + self.prior['p'] = [2,4] + + if t is None: + self.prior['t'] = [0.1, 1_000] # super wide prior as default + else: + self.prior = prior + # then initiate the superclass super().__init__(self.prior, p=p) self.t = t - def SED(self, nu, p, log_bG_sh, logMdot, log_epsilon_T, log_epsilon_e, log_epsilon_B, + def SED(self, nu, p, log_bG_sh, log_Mdot, log_epsilon_T, log_epsilon_e, log_epsilon_B, t, lum_dist, **kwargs): # set microphysical and geometric parameters @@ -45,7 +48,7 @@ def SED(self, nu, p, log_bG_sh, logMdot, log_epsilon_T, log_epsilon_e, log_epsil t = (t*u.day).to(u.s).value - Mdot_over_vw = (10**logMdot*(c.M_sun/u.yr/1e8)).cgs.value + Mdot_over_vw = (10**log_Mdot*(c.M_sun/u.yr/1e8)).cgs.value Lnu = Lnu_of_nu( 10**log_bG_sh, Mdot_over_vw, nu, t, p=p, @@ -63,10 +66,10 @@ def lnprior(self, theta, nu, F, upperlimit, **kwargs): Logarithmic prior function that can be changed based on the SED model. ''' uppertest = SyncfitModel._is_below_upperlimits( - nu, F, upperlimit, theta, self.SED + nu, F, upperlimit, theta, self.SED, **kwargs ) - packed_theta = self.pack_theta(theta) + packed_theta = self.pack_theta(theta, **kwargs) all_res = [] for param, val in self.prior.items(): diff --git a/src/syncfit/models/syncfit_model.py b/src/syncfit/models/syncfit_model.py index d939dea..142feff 100644 --- a/src/syncfit/models/syncfit_model.py +++ b/src/syncfit/models/syncfit_model.py @@ -135,11 +135,9 @@ def loglik(self, theta, nu, F, F_error, **kwargs): ''' with warnings.catch_warnings(): warnings.simplefilter('ignore') - if self.p is not None: - model_result = self.SED(nu, self.p, *theta, **kwargs) - else: - model_result = self.SED(nu, *theta, **kwargs) - + packed_theta = self.pack_theta(theta, **kwargs) + model_result = self.SED(nu, **packed_theta) + if not np.any(np.isfinite(model_result)): ll = -np.inf else: @@ -151,7 +149,7 @@ def loglik(self, theta, nu, F, F_error, **kwargs): return ll @staticmethod - def _is_below_upperlimits(nu, F, upperlimits, theta, model): + def _is_below_upperlimits(nu, F, upperlimits, theta, model, **kwargs): ''' Checks that the location of theta is below any upperlimits ''' @@ -162,11 +160,9 @@ def _is_below_upperlimits(nu, F, upperlimits, theta, model): where_upperlimit = np.where(upperlimits)[0] F_upperlimits = F[where_upperlimit] - if self.p is None: - test_fluxes = model(nu, *theta)[where_upperlimit] - else: - test_fluxes = model(nu, self.p, *theta)[where_upperlimit] - + packed_theta = self.pack_theta(theta, **kwargs) + test_fluxes = model(nu, **packed_theta)[where_upperlimit] + return np.all(F_upperlimits > test_fluxes) # Some *required* abstract methods @@ -184,18 +180,21 @@ def SED(self, *args, **kwargs): ''' pass - def pack_theta(self, theta): + def pack_theta(self, theta, **kwargs): ''' Pack theta into a dictionary ''' - return {param:theta[idx] for idx, param in enumerate(self.labels)} + d = {param:theta[idx] for idx, param in enumerate(self.labels)} + for k,v in kwargs.items(): + d[k] = v + return d def lnprior(self, theta, nu, F, upperlimit, **kwargs): ''' Logarithmic prior function that can be changed based on the SED model. ''' uppertest = SyncfitModel._is_below_upperlimits( - nu, F, upperlimit, theta, self.SED + nu, F, upperlimit, theta, self.SED, **kwargs ) packed_theta = self.pack_theta(theta)