diff --git a/src/syncfit/_version.py b/src/syncfit/_version.py index d31c31e..493f741 100644 --- a/src/syncfit/_version.py +++ b/src/syncfit/_version.py @@ -1 +1 @@ -__version__ = "0.2.3" +__version__ = "0.3.0" diff --git a/src/syncfit/analysis/plotter.py b/src/syncfit/analysis/plotter.py index 11f456b..91e62bd 100644 --- a/src/syncfit/analysis/plotter.py +++ b/src/syncfit/analysis/plotter.py @@ -4,6 +4,7 @@ import matplotlib.pyplot as plt from .util import * +from ..models import MQModel def plot_chains(sampler, labels, fig=None, axes=None): ''' @@ -37,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, t=None, - p=None, method='random', fig=None, ax=None, day=None): +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): ''' Plot best fit model @@ -49,7 +50,7 @@ def plot_best_fit(model, sampler, nu, F, Ferr, nkeep=1000, lum_dist=None, t=None F [list]: The observed flux densities Ferr [list]: The observed flux error nkeep [int]: Number of values to keep - p [float]: p-value used, if not None + nu_arr [list]: List of nus for the best fit lines to be plot with method [str]: Either 'max' or 'last' or 'random', default is max. - max: takes the nkeep maximum probability values - last: takes the last nkeep values from the chain @@ -60,7 +61,9 @@ def plot_best_fit(model, sampler, nu, F, Ferr, nkeep=1000, lum_dist=None, t=None Returns: matplotlib fig, ax ''' - + if isinstance(model, MQModel) and model.t is not None: + t = model.t + flat_samples, log_prob = extract_output(sampler) if method == 'max': @@ -71,20 +74,25 @@ def plot_best_fit(model, sampler, nu, F, Ferr, nkeep=1000, lum_dist=None, t=None toplot = flat_samples[-nkeep*10:][np.random.randint(0, nkeep*10, nkeep)] else: raise ValueError('method must be either last or max!') - - nu_plot = np.arange(1e8,3e11,1e7) + if nu_arr is None: + nu_plot = np.arange(1e8,3e11,1e7) + else: + nu_plot = nu_arr + if ax is None: fig, ax = plt.subplots(figsize=(4,4)) - if t is not None or lum_dist is not None: - kwargs = dict(t=t,lum_dist=lum_dist) + 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 p is not None: - res = model.SED(nu_plot, p, *val, **kwargs) + if model.p is not None: + res = model.SED(nu_plot, model.p, *val, **kwargs) else: res = model.SED(nu_plot, *val, **kwargs) diff --git a/src/syncfit/mcmc.py b/src/syncfit/mcmc.py index c69e002..1935cb2 100644 --- a/src/syncfit/mcmc.py +++ b/src/syncfit/mcmc.py @@ -15,7 +15,7 @@ def do_dynesty(nu:list[float], F_mJy:list[float], F_error:list[float], lum_dist:float=None, t:float=None, model:SyncfitModel=MQModel, fix_p:float=None, - upperlimits:list[bool]=None, ncores:int=1, seed:int=None, + upperlimits:list[bool]=None, ncores:int=1, seed:int=None, prior=None, run_kwargs={}, dynesty_kwargs={}, logprob_kwargs={} ) -> tuple[list[float],list[float]]: """ @@ -28,12 +28,14 @@ def do_dynesty(nu:list[float], F_mJy:list[float], F_error:list[float], model (SyncfitModel): Model class to use from syncfit.fitter.models. Can also be a custom model but it must be a subclass of SyncfitModel! lum_dist (float): luminosity distance in cgs units. Only needed for MQModel. Default is None. - t (flost): observation time in seconds. Only needed for MQModel. Default is None. + t (flost): observation time in days. Only needed for MQModel. Default is None. fix_p (float): Will fix the p value to whatever you give, do not provide p in theta_init if this is the case! upperlimits (list[bool]): True if the point is an upperlimit, False otherwise. ncores (int) : The number of cores to run on, default is 1 and won't multiprocess seed (int): The seed for the random number generator passed to dynesty, + prior (dict) : dictionary defining the prior ranges. Keys must be same as model.get_labels(). + Value should be a list of length 2 like [min, max], both exclusive. run_kwargs (dict) : kwargs to pass to dynesty.run_sampler dynesty_kwargs (dict) : kwargs to pass to dynesty.DynamicNestedSampler logprob_kwargs (dict) : kwargs to pass to the logprob. For the most part this is @@ -42,19 +44,35 @@ def do_dynesty(nu:list[float], F_mJy:list[float], F_error:list[float], Returns: flat_samples, log_prob """ - if isinstance(model(), MQModel) and (lum_dist is None or t is None): + # instantiate a new model object + 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) + else: + model = model(p=fix_p) # get the extra args - dynesty_args = model.get_kwargs(nu, F_mJy, F_error, lum_dist=lum_dist, t=t, p=fix_p) + dynesty_args = model.get_kwargs(nu, F_mJy, F_error, lum_dist=lum_dist, t=t) # combine these with the logprob_kwargs # make the logprob_kwargs second so it overwrites anything we set here dynesty_args = dynesty_args | logprob_kwargs - ndim = len(model.get_labels(p=fix_p)) + ndim = model.ndim rstate = np.random.default_rng(seed) - + + # set the model prior instance variable + if prior is not None: + model.prior = prior + + if set(model.prior.keys()) != set(model.labels): + raise ValueError( + f'Prior dictionary keys ({model.prior.keys()}) do not match the labels ({model.labels})!' + ) + # construct the sampler and run it # NOTE: I give it the lnprob instead of loglik because there can be some other # priors that are built into the lnprob that can not be in the dynesty prior @@ -75,13 +93,13 @@ def do_dynesty(nu:list[float], F_mJy:list[float], F_error:list[float], 'The override decorator syntax is not currently supported for dynesty!' ) - return dsampler + return model, dsampler def do_emcee(theta_init:list[float], nu:list[float], F_mJy:list[float], F_error:list[float], lum_dist:float=None, t:float=None, model:SyncfitModel=SyncfitModel, niter:int=2000, nwalkers:int=100, fix_p:float=None, upperlimits:list[bool]=None, - day:str=None, plot:bool=False, ncores:int=1 + day:str=None, plot:bool=False, ncores:int=1, prior=None ) -> tuple[list[float],list[float]]: """ Fit the data with the given model using the emcee package. @@ -106,7 +124,10 @@ def do_emcee(theta_init:list[float], nu:list[float], F_mJy:list[float], Returns: flat_samples, log_prob """ - if isinstance(model(), MQModel) and (lum_dist is None or t is None): + # instantiate a new model object + model = model(p=fix_p) + + if isinstance(model, MQModel) and (lum_dist is None or t is None): raise ValueError('lum_dist and t reequired for MQModel!') ### Fill in initial guesses and number of parameters @@ -121,12 +142,15 @@ def do_emcee(theta_init:list[float], nu:list[float], F_mJy:list[float], upperlimits = np.array(upperlimits) pos, labels, emcee_args = model.unpack_util(theta_init, nu, F_mJy, F_error, - nwalkers, p=fix_p, lum_dist=lum_dist, + nwalkers, lum_dist=lum_dist, t=t, upperlimit=upperlimits) # setup and run the MCMC nwalkers, ndim = pos.shape + if set(model.prior.keys()) != set(model.labels): + raise ValueError('Prior dictionary keys do not match the labels!') + with Pool(ncores) as pool: sampler = emcee.EnsembleSampler( nwalkers, @@ -156,4 +180,4 @@ def do_emcee(theta_init:list[float], nu:list[float], F_mJy:list[float], else: fig, ax = plot_best_fit(model, sampler, emcee_args['nu'], emcee_args['F']) - return sampler + return model, sampler diff --git a/src/syncfit/models/b1b2_b3b4_weighted_model.py b/src/syncfit/models/b1b2_b3b4_weighted_model.py index fd354bc..fe5d309 100644 --- a/src/syncfit/models/b1b2_b3b4_weighted_model.py +++ b/src/syncfit/models/b1b2_b3b4_weighted_model.py @@ -9,15 +9,34 @@ class B1B2_B3B4_Weighted(SyncfitModel): This is a specialized model that uses a weighted combination of the B1B2 model and the B3B4 model. The idea of this model is from XXXYXYXYX et al. (YYYY). ''' + + def __init__(self, p=None): + super().__init__(p=p) + + # 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] + ) + else: + self.prior = dict( + log_F_nu=[-4,2], + log_nu_a=[6,12], + log_nu_m=[6,12] + ) + - def get_labels(p=None): + def get_labels(self, p=None): if p is None: - return ['p','log F_v', 'log nu_a','log nu_m'] + return ['p','log_F_nu', 'log_nu_a','log_nu_m'] else: - return ['log F_v', 'log nu_a','log nu_m'] + return ['log_F_nu', 'log_nu_a','log_nu_m'] # the model, must be named SED!!! - def SED(nu, p, log_F_nu, log_nu_a, log_nu_m, **kwargs): + def SED(self, nu, p, log_F_nu, log_nu_a, log_nu_m, **kwargs): ### Spectrum 1 b1 = 2 b2 = 1/3 @@ -56,47 +75,3 @@ def SED(nu, p, log_F_nu, log_nu_a, log_nu_m, **kwargs): F = (w1*F1+w2*F2) / (w1+w2) return F - - def lnprior(theta, nu, F, upperlimit, p=None, **kwargs): - ''' Priors: ''' - uppertest = SyncfitModel._is_below_upperlimits( - nu, F, upperlimit, theta, B1B2_B3B4_Weighted.SED, p=p - ) - - if p is None: - p, log_F_nu, log_nu_a, log_nu_m= theta - else: - log_F_nu, log_nu_a, log_nu_m= theta - - if 2< p < 4 and -4 < log_F_nu < 2 and 6 < log_nu_a < 12 and 6 < log_nu_m < 12 and uppertest: - return 0.0 - - else: - return -np.inf - - def dynesty_prior(theta, nu, F, upperlimit, p=None, **kwargs): - ''' - Prior transform for dynesty - ''' - if p is None: - p, log_F_nu, log_nu_a, log_nu_m= theta - fixed_p = False, - else: - fixed_p = True - log_F_nu, log_nu_a, log_nu_m= theta - - # log_F_nu between -4 and 2 - log_F_nu = log_F_nu*6 - 4 - - # log_nu_a between 6 and 11 - log_nu_a = log_nu_a*5 + 6 - - # same transform to log_nu_c - log_nu_m = log_nu_m*5 + 6 - - if not fixed_p: - # p should be between 2 and 4 - p = 2*p + 2 - - return p,log_F_nu,log_nu_a,log_nu_m - return log_F_nu,log_nu_a,log_nu_m diff --git a/src/syncfit/models/b1b2_model.py b/src/syncfit/models/b1b2_model.py index 4db3b6c..55a32bd 100644 --- a/src/syncfit/models/b1b2_model.py +++ b/src/syncfit/models/b1b2_model.py @@ -10,14 +10,33 @@ class B1B2(SyncfitModel): (nu_m). This model uses nu_m > nu_a, the opposite of the B4B5 model. ''' - def get_labels(p=None): + def __init__(self, p=None): + super().__init__(p=p) + + # 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] + ) + else: + self.prior = dict( + log_F_nu=[-4,2], + log_nu_a=[6,11], + log_nu_m=[7,15] + ) + + + def get_labels(self, p=None): if p is None: - return ['p','log F_v', 'log nu_a','log nu_m'] + return ['p','log_F_nu', 'log_nu_a','log_nu_m'] else: - return ['log F_v', 'log nu_a','log nu_m'] + return ['log_F_nu', 'log_nu_a','log_nu_m'] # the model, must be named SED!!! - def SED(nu, p, log_F_nu, log_nu_a, log_nu_m, **kwargs): + def SED(self, nu, p, log_F_nu, log_nu_a, log_nu_m, **kwargs): b1 = 2 b2 = 1/3 b3 = (1-p)/2 @@ -34,49 +53,25 @@ def SED(nu, p, log_F_nu, log_nu_a, log_nu_m, **kwargs): return F_nu * term1 * term2 - def lnprior(theta, nu, F, upperlimit, p=None, **kwargs): - ''' Priors: ''' + 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, B1B2.SED, p=p + nu, F, upperlimit, theta, self.SED ) - - if p is None: - p, log_F_nu, log_nu_a, log_nu_m = theta - else: - log_F_nu, log_nu_a, log_nu_m = theta - if 2< p < 4 and -4 < log_F_nu < 2 and 6 < log_nu_a < 12 and log_nu_m > log_nu_a and uppertest: + packed_theta = self.pack_theta(theta) + + all_res = [] + for param, val in self.prior.items(): + res = val[0] < packed_theta[param] < val[1] + all_res.append(res) + + if (all(all_res) and + uppertest and + packed_theta['log_nu_m'] > packed_theta['log_nu_a'] + ): return 0.0 - else: return -np.inf - - - def dynesty_transform(theta, nu, F, upperlimit, p=None, **kwargs): - ''' - Prior transform for dynesty - ''' - - if p is None: - p, log_F_nu, log_nu_a, log_nu_m = theta - fixed_p = False, - else: - fixed_p = True - log_F_nu, log_nu_a, log_nu_m = theta - - - # log_F_nu between -4 and 2 - log_F_nu = log_F_nu*6 - 4 - - # log_nu_a between 6 and 11 - log_nu_a = log_nu_a*5 + 6 - - # same transform to log_nu_m - log_nu_m = log_nu_m*5 + 6 - - if not fixed_p: - # p should be between 2 and 4 - p = 2*p + 2 - - return p,log_F_nu,log_nu_a,log_nu_m - return log_F_nu,log_nu_a,log_nu_m diff --git a/src/syncfit/models/b4b5_model.py b/src/syncfit/models/b4b5_model.py index 2acd312..626fd9d 100644 --- a/src/syncfit/models/b4b5_model.py +++ b/src/syncfit/models/b4b5_model.py @@ -11,14 +11,33 @@ class B4B5(SyncfitModel): subclass this class and overwrite lnprior to redefine this. ''' - def get_labels(p=None): + def __init__(self, p=None): + super().__init__(p=p) + + # 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] + ) + else: + self.prior = dict( + log_F_nu=[-4,2], + log_nu_a=[6,11], + log_nu_m=[-np.inf,6] + ) + + + def get_labels(self, p=None): if p is None: - return ['p','log F_v', 'log nu_a','log nu_m'] + return ['p','log_F_nu', 'log_nu_a','log_nu_m'] else: - return ['log F_v', 'log nu_a','log nu_m'] + return ['log_F_nu', 'log_nu_a','log_nu_m'] # the model, must be named SED!!! - def SED(nu, p, log_F_nu, log_nu_a, log_nu_m, **kwargs): + def SED(self, nu, p, log_F_nu, log_nu_a, log_nu_m, **kwargs): b1 = 2 b2 = 5/2 b3 = (1-p)/2 @@ -35,48 +54,25 @@ def SED(nu, p, log_F_nu, log_nu_a, log_nu_m, **kwargs): return F_nu * term1 * term2 - def lnprior(theta, nu, F, upperlimit, p=None, **kwargs): - ''' Priors: ''' + 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, B4B5.SED, p=p + nu, F, upperlimit, theta, self.SED ) - if p is None: - p, log_F_nu, log_nu_a, log_nu_m = theta - else: - log_F_nu, log_nu_a, log_nu_m = theta - - if 2< p < 4 and -4 < log_F_nu < 2 and 6 < log_nu_a < 11 and log_nu_m < log_nu_a and uppertest: + packed_theta = self.pack_theta(theta) + + all_res = [] + for param, val in self.prior.items(): + res = val[0] < packed_theta[param] < val[1] + all_res.append(res) + + if (all(all_res) and + uppertest and + packed_theta['log_nu_a'] > packed_theta['log_nu_m'] + ): return 0.0 - else: return -np.inf - - def dynesty_transform(theta, nu, F, upperlimit, p=None, **kwargs): - ''' - Prior transform for dynesty - ''' - - if p is None: - p, log_F_nu, log_nu_a, log_nu_m = theta - fixed_p = False - else: - fixed_p = True - log_F_nu, log_nu_a, log_nu_m = theta - - - # log_F_nu between -4 and 2 - log_F_nu = log_F_nu*6 - 4 - - # log_nu_a between 6 and 11 - log_nu_a = log_nu_a*5 + 6 - - # same transform to log_nu_m - log_nu_m = log_nu_m*5 + 6 - - if not fixed_p: - # p should be between 2 and 4 - p = 2*p + 2 - - return p,log_F_nu,log_nu_a,log_nu_m - return log_F_nu,log_nu_a,log_nu_m diff --git a/src/syncfit/models/b4b5b3_model.py b/src/syncfit/models/b4b5b3_model.py index 14d51e7..62846a4 100644 --- a/src/syncfit/models/b4b5b3_model.py +++ b/src/syncfit/models/b4b5b3_model.py @@ -10,14 +10,34 @@ class B4B5B3(SyncfitModel): and minimum energy break (nu_m). This model always requires that nu_m < nu_a < nu_c. ''' - def get_labels(p=None): + def __init__(self, p=None): + super().__init__(p=p) + + # 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] + ) + else: + self.prior = dict( + log_F_nu=[-4,2], + log_nu_a=[6,11], + log_nu_m=[0, 6], + log_nu_c=[7,15] + ) + + def get_labels(self, p=None): if p is None: - return ['p','log F_v', 'log nu_a','log nu_m', 'log nu_c'] + return ['p','log_F_nu', 'log_nu_a','log_nu_m', 'log_nu_c'] else: - return ['log F_v', 'log nu_a','log nu_m', 'log nu_c'] + return ['log_F_nu', 'log_nu_a','log_nu_m', 'log_nu_c'] # the model, must be named SED!!! - def SED(nu, p, log_F_nu, log_nu_a, log_nu_m, log_nu_c, **kwargs): + def SED(self, nu, p, log_F_nu, log_nu_a, log_nu_m, log_nu_c, **kwargs): b1 = 2 b2 = 5/2 b3 = (1-p)/2 @@ -39,50 +59,26 @@ def SED(nu, p, log_F_nu, log_nu_a, log_nu_m, log_nu_c, **kwargs): return F_nu * term1 * term2 * term3 - def lnprior(theta, nu, F, upperlimit, p=None, **kwargs): - ''' Priors: ''' + 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, B4B5B3.SED, p=p + nu, F, upperlimit, theta, self.SED ) - if p is None: - p, log_F_nu, log_nu_a, log_nu_m, log_nu_c= theta - else: - log_F_nu, log_nu_a, log_nu_m, log_nu_c= theta - if 2< p < 4 and -4 < log_F_nu < 2 and 6 < log_nu_a < 11 and log_nu_m < log_nu_a and log_nu_a < log_nu_c and uppertest: + packed_theta = self.pack_theta(theta) + + all_res = [] + for param, val in self.prior.items(): + res = val[0] < packed_theta[param] < val[1] + all_res.append(res) + + if (all(all_res) and + uppertest and + packed_theta['log_nu_c'] > packed_theta['log_nu_a'] and + packed_theta['log_nu_m'] < packed_theta['log_nu_a'] + ): return 0.0 - else: return -np.inf - - def dynesty_transform(theta, nu, F, upperlimit, p=None, **kwargs): - ''' - Prior transform for dynesty - ''' - - if p is None: - p, log_F_nu, log_nu_a, log_nu_m, log_nu_c = theta - fixed_p = False, - else: - fixed_p = True - log_F_nu, log_nu_a, log_nu_m, log_nu_c = theta - - - # log_F_nu between -4 and 2 - log_F_nu = log_F_nu*6 - 4 - - # log_nu_a between 6 and 11 - log_nu_a = log_nu_a*5 + 6 - - # same transform to log_nu_c - log_nu_c = log_nu_c*5 + 6 - - # same for log_nu_m - log_nu_m = log_nu_m*5 + 6 - - if not fixed_p: - # p should be between 2 and 4 - p = 2*p + 2 - - return p,log_F_nu,log_nu_a,log_nu_m,log_nu_c - return log_F_nu,log_nu_a,log_nu_m,log_nu_c diff --git a/src/syncfit/models/b5_model.py b/src/syncfit/models/b5_model.py index 7c4de1b..3da8577 100644 --- a/src/syncfit/models/b5_model.py +++ b/src/syncfit/models/b5_model.py @@ -9,14 +9,30 @@ class B5(SyncfitModel): Single break model for just the self-absorption break. ''' - def get_labels(p=None): + def __init__(self, p=None): + super().__init__(p=p) + + # 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] + ) + else: + self.prior = dict( + log_F_nu=[-4,2], + log_nu_a=[6,11] + ) + + def get_labels(self, p=None): if p is None: - return ['p','log F_v', 'log nu_a'] + return ['p','log_F_nu', 'log_nu_a'] else: - return ['log F_v', 'log nu_a'] + return ['log_F_nu', 'log_nu_a'] # the model, must be named SED!!! - def SED(nu, p, log_F_nu, log_nu_a, **kwargs): + def SED(self, nu, p, log_F_nu, log_nu_a, **kwargs): b1 = 5/2 b2 = (1-p)/2 s = 1.25-0.18*p @@ -27,49 +43,3 @@ def SED(nu, p, log_F_nu, log_nu_a, **kwargs): term = ((nu/nu_a)**(-s*b1)+(nu/nu_a)**(-s*b2)) return F_nu*term**(-1/s) - - def lnprior(theta, nu, F, upperlimit, p=None, **kwargs): - ''' Priors: ''' - uppertest = SyncfitModel._is_below_upperlimits( - nu, F, upperlimit, theta, B5.SED, p=p - ) - - if p is None: - p, log_F_nu, log_nu_a = theta - else: - log_F_nu, log_nu_a = theta - - if 2< p < 4 and -4 < log_F_nu < 2 and 6 < log_nu_a < 11 and uppertest: - return 0.0 - - else: - return -np.inf - - - def dynesty_transform(theta, nu, F, upperlimit, p=None, **kwargs): - ''' - Prior transform for dynesty - ''' - - if p is None: - p, log_F_nu, log_nu_a = theta - fixed_p = False - else: - fixed_p = True - log_F_nu, log_nu_a = theta - - - # log_F_nu between -4 and 2 - log_F_nu = log_F_nu*6 - 4 - - # log_nu_a between 6 and 11 - log_nu_a = log_nu_a*5 + 6 - - - if not fixed_p: - # p should be between 2 and 4 - p = 2*p + 2 - - return p,log_F_nu,log_nu_a - return log_F_nu,log_nu_a - diff --git a/src/syncfit/models/b5b3_model.py b/src/syncfit/models/b5b3_model.py index cb85d86..68a2c15 100644 --- a/src/syncfit/models/b5b3_model.py +++ b/src/syncfit/models/b5b3_model.py @@ -11,15 +11,33 @@ class B5B3(SyncfitModel): break. ''' + def __init__(self, p=None): + super().__init__(p=p) + + # 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] + ) + else: + self.prior = dict( + log_F_nu=[-4,2], + log_nu_a=[6,11], + log_nu_c=[7,15] + ) + # Write some getters for things that are model specific - def get_labels(p=None): + def get_labels(self, p=None): if p is None: - return ['p','log F_v', 'log nu_a','log nu_c'] + return ['p','log_F_nu', 'log_nu_a','log_nu_c'] else: - return ['log F_v', 'log nu_a','log nu_c'] + return ['log_F_nu', 'log_nu_a','log_nu_c'] # the model, must be named SED!!! - def SED(nu, p, log_F_nu, log_nu_a, log_nu_c, **kwargs): + def SED(self, nu, p, log_F_nu, log_nu_a, log_nu_c, **kwargs): b1 = 5/2 b2 = (1-p)/2 b3 = -p/2 @@ -36,48 +54,25 @@ def SED(nu, p, log_F_nu, log_nu_a, log_nu_c, **kwargs): return F_nu * term1 * term2 - def lnprior(theta, nu, F, upperlimit, p=None, **kwargs): - ''' Priors: ''' + 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, B5B3.SED, p=p + nu, F, upperlimit, theta, self.SED ) - - if p is None: - p, log_F_nu, log_nu_a, log_nu_c = theta - else: - log_F_nu, log_nu_a, log_nu_c = theta - if 2< p < 4 and -4 < log_F_nu < 2 and 6 < log_nu_a < 11 and log_nu_c > log_nu_a and uppertest: + packed_theta = self.pack_theta(theta) + + all_res = [] + for param, val in self.prior.items(): + res = val[0] < packed_theta[param] < val[1] + all_res.append(res) + + if (all(all_res) and + uppertest and + packed_theta['log_nu_c'] > packed_theta['log_nu_a'] + ): return 0.0 - else: return -np.inf - - def dynesty_transform(theta, nu, F, upperlimit, p=None, **kwargs): - ''' - Prior transform for dynesty - ''' - - if p is None: - p, log_F_nu, log_nu_a, log_nu_c = theta - fixed_p = False - else: - fixed_p = True - log_F_nu, log_nu_a, log_nu_c = theta - - - # log_F_nu between -4 and 2 - log_F_nu = log_F_nu*6 - 4 - - # log_nu_a between 6 and 11 - log_nu_a = log_nu_a*5 + 6 - - # same transform to log_nu_c - log_nu_c = log_nu_c*5 + 6 - - if not fixed_p: - # p should be between 2 and 4 - p = 2*p + 2 - - return p,log_F_nu,log_nu_a,log_nu_c - return log_F_nu,log_nu_a,log_nu_c diff --git a/src/syncfit/models/mq_model.py b/src/syncfit/models/mq_model.py index 478f546..4a8986d 100644 --- a/src/syncfit/models/mq_model.py +++ b/src/syncfit/models/mq_model.py @@ -12,15 +12,38 @@ from astropy import constants as c class MQModel(SyncfitModel): - - def get_labels(p=None): + + def __init__(self, p=None, t=None): + super().__init__(p=p) + self.t = t + self.labels = self.get_labels(p=p, t=t) + self.ndim = len(self.labels) + + # then 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: - return ['p', 'log_bG_sh', 'log_Mdot', 'log_epsilon_T', 'log_epsilon_e', 'log_epsilon_B'] - else: - return ['log_bG_sh', 'log_Mdot', 'log_epsilon_T', 'log_epsilon_e', 'log_epsilon_B'] + self.prior['p'] = [2,4] - def SED(nu, p, log_bG_sh, logMdot, log_epsilon_T, log_epsilon_e, log_epsilon_B, - lum_dist, t, **kwargs): + if t is None: + self.prior['t'] = [0.1, 1_000] # super wide prior as default + + def get_labels(self, p=None, t=None): + base_labels = ['log_bG_sh', 'log_Mdot', 'log_epsilon_T', 'log_epsilon_e', 'log_epsilon_B'] + if p is None: + base_labels = ['p'] + base_labels + if t is None: + base_labels.append('t') + return base_labels + + def SED(self, nu, p, log_bG_sh, logMdot, log_epsilon_T, log_epsilon_e, log_epsilon_B, + t, lum_dist, **kwargs): # set microphysical and geometric parameters # log_epsilon_e = -1 @@ -29,6 +52,8 @@ def SED(nu, p, log_bG_sh, logMdot, log_epsilon_T, log_epsilon_e, log_epsilon_B, f = 3.0/16.0 ell_dec = 1.0 + t = (t*u.day).to(u.s).value + Mdot_over_vw = (10**logMdot*(c.M_sun/u.yr/1e8)).cgs.value Lnu = Lnu_of_nu( @@ -41,67 +66,26 @@ def SED(nu, p, log_bG_sh, logMdot, log_epsilon_T, log_epsilon_e, log_epsilon_B, Fnu = (Lnu / (4*np.pi*(lum_dist_cm)**2)).to(u.mJy) # mJy return Fnu.value - - def lnprior(theta, nu, F, upperlimit, p=None, **kwargs): + + def lnprior(self, theta, nu, F, upperlimit, **kwargs): ''' - The prior + Logarithmic prior function that can be changed based on the SED model. ''' uppertest = SyncfitModel._is_below_upperlimits( - nu, F, upperlimit, theta, MQModel.SED, p=p + nu, F, upperlimit, theta, self.SED ) - - if p is None: - p, log_bG_sh, logMdot, epsilon_T, epsilon_e, epsilon_B = theta - else: - log_bG_sh, logMdot, epsilon_T, epsilon_e, epsilon_B = theta - - if (uppertest and - 2 < p < 4 and - -3 < log_bG_sh < 3 and - -10 < logMdot < 0 and - -6 < epsilon_e < 0 and - -6 < epsilon_T < 0 and - -6 < epsilon_B < 0 and - 0 <= 10**epsilon_e + 10**epsilon_B + 10**epsilon_T <= 1): + packed_theta = self.pack_theta(theta) + + all_res = [] + for param, val in self.prior.items(): + res = val[0] < packed_theta[param] < val[1] + all_res.append(res) + + if (all(all_res) and + uppertest and + 0 <= 10**packed_theta['log_epsilon_e'] + 10**packed_theta['log_epsilon_B'] + 10**packed_theta['log_epsilon_T'] <= 1 + ): return 0.0 else: return -np.inf - - def dynesty_transform(theta, nu, F, upperlimit, p=None, **kwargs): - ''' - Prior transform function for dynesty - - theta is expected to be a 1D array with each value in the range [0,1) - so we need to transform each to parameter space - ''' - - if p is None: - fixed_p = False - p, log_bG_sh, logMdot, epsilon_T, epsilon_e, epsilon_B = theta - else: - fixed_p = True - log_bG_sh, logMdot, epsilon_T, epsilon_e, epsilon_B = theta - - # log_bG_sh should be between -2 and 2 - log_bG_sh = log_bG_sh*6 - 3 - - # -10 < logMdot < 0 - logMdot*=-10 - - # -3 < epsilon_e < 0 - epsilon_e*=-6 - epsilon_B*=-6 - - # -3 < epsilon_T < 0 - epsilon_T*=-6 - - if not fixed_p: - # p should be between 2 and 4 - p = 2*p + 2 - - # p between 2.5 and 3.5, let's be a little more restrictive - #p += 2.5 - - return p,log_bG_sh,logMdot,epsilon_T,epsilon_e, epsilon_B - return log_bG_sh,logMdot,epsilon_T,epsilon_e, epsilon_B diff --git a/src/syncfit/models/syncfit_model.py b/src/syncfit/models/syncfit_model.py index 3e2ddaa..14a6cac 100644 --- a/src/syncfit/models/syncfit_model.py +++ b/src/syncfit/models/syncfit_model.py @@ -26,6 +26,12 @@ class SyncfitModel(object, metaclass=_SyncfitModelMeta): while also allowing users to customize their own. ''' + def __init__(self, p=None): + self.labels = self.get_labels(p=p) + self.ndim = len(self.labels) + self.prior = None + self.p = p + # Write some getters for things that are model specific # THESE WILL BE THE SAME ACROSS ALL MODELS! @staticmethod @@ -53,7 +59,7 @@ def get_pos(theta_init:list, nwalkers:int) -> list[float]: @staticmethod def get_kwargs(nu:list, F_mJy:list, F_error:list, lum_dist:float=None, - t:float=None, p:float=None, upperlimit:list=None) -> dict: + t:float=None, upperlimit:list=None) -> dict: ''' Packages up the args to be passed into the model based on the user input. @@ -72,9 +78,6 @@ def get_kwargs(nu:list, F_mJy:list, F_error:list, lum_dist:float=None, base_args = {'nu':nu, 'F':F, 'F_error':F_error, 'upperlimit':upperlimit} - if p is not None: - base_args['p'] = p - if lum_dist is not None: base_args['lum_dist'] = lum_dist @@ -84,9 +87,8 @@ def get_kwargs(nu:list, F_mJy:list, F_error:list, lum_dist:float=None, return base_args # package those up for easy getting in do_emcee - @classmethod - def unpack_util(cls, theta_init, nu, F_mJy, F_error, nwalkers, lum_dist=None, - t=None, p=None, upperlimit=None): + def unpack_util(self, theta_init, nu, F_mJy, F_error, nwalkers, lum_dist=None, + t=None, upperlimit=None): ''' A wrapper on the utility functions. @@ -98,12 +100,11 @@ def unpack_util(cls, theta_init, nu, F_mJy, F_error, nwalkers, lum_dist=None, p (float): A p-value to pass to the model, only used if p-value is fixed nwalkers (int): THe number of walkers to use ''' - return (cls.get_pos(theta_init,nwalkers), - cls.get_labels(p=p), - cls.get_kwargs(nu, F_mJy, F_error, p, upperlimit)) + return (self.get_pos(theta_init,nwalkers), + self.get_labels(p=self.p), + self.get_kwargs(nu, F_mJy, F_error, upperlimit)) - @classmethod - def lnprob(cls, theta:list, **kwargs): + def lnprob(self, theta:list, **kwargs): '''Keep or throw away step likelihood and priors Args: @@ -113,14 +114,13 @@ def lnprob(cls, theta:list, **kwargs): Returns: The likelihood of the data at that location ''' - lp = cls.lnprior(theta, **kwargs) + lp = self.lnprior(theta, **kwargs) if not np.isfinite(lp): return -np.inf else: - return lp + cls.loglik(theta, **kwargs) + return lp + self.loglik(theta, **kwargs) - @classmethod - def loglik(cls, theta, nu, F, F_error, p=None, **kwargs): + def loglik(self, theta, nu, F, F_error, **kwargs): '''Log Likelihood function Args: @@ -135,10 +135,10 @@ def loglik(cls, theta, nu, F, F_error, p=None, **kwargs): ''' with warnings.catch_warnings(): warnings.simplefilter('ignore') - if p is not None: - model_result = cls.SED(nu, p, *theta, **kwargs) + if self.p is not None: + model_result = self.SED(nu, self.p, *theta, **kwargs) else: - model_result = cls.SED(nu, *theta, **kwargs) + model_result = self.SED(nu, *theta, **kwargs) if not np.any(np.isfinite(model_result)): ll = -np.inf @@ -151,7 +151,7 @@ def loglik(cls, theta, nu, F, F_error, p=None, **kwargs): return ll @staticmethod - def _is_below_upperlimits(nu, F, upperlimits, theta, model, p=None): + def _is_below_upperlimits(nu, F, upperlimits, theta, model): ''' Checks that the location of theta is below any upperlimits ''' @@ -162,47 +162,94 @@ def _is_below_upperlimits(nu, F, upperlimits, theta, model, p=None): where_upperlimit = np.where(upperlimits)[0] F_upperlimits = F[where_upperlimit] - if p is None: + if self.p is None: test_fluxes = model(nu, *theta)[where_upperlimit] else: - test_fluxes = model(nu, p, *theta)[where_upperlimit] + test_fluxes = model(nu, self.p, *theta)[where_upperlimit] return np.all(F_upperlimits > test_fluxes) # Some *required* abstract methods - @staticmethod - @abstractmethod - def get_labels(*args, **kwargs): + def get_labels(self, *args, **kwargs): ''' Describes a list of labels used in the return values of the mcmc chain. This varies depending on the inputs to the MCMC. ''' pass - @staticmethod @abstractmethod - def SED(*args, **kwargs): + def SED(self, *args, **kwargs): ''' Describes the SED model for the model that subclasses this BaseModel ''' pass + + def pack_theta(self, theta): + ''' + Pack theta into a dictionary + ''' + return {param:theta[idx] for idx, param in enumerate(self.labels)} - @staticmethod - @abstractmethod - def lnprior(*args, **kwargs): + def lnprior(self, theta, nu, F, upperlimit, **kwargs): ''' Logarithmic prior function that can be changed based on the SED model. ''' - pass + uppertest = SyncfitModel._is_below_upperlimits( + nu, F, upperlimit, theta, self.SED + ) - @staticmethod - @abstractmethod - def dynesty_transform(*args, **kwargs): + packed_theta = self.pack_theta(theta) + + all_res = [] + for param, val in self.prior.items(): + res = val[0] < packed_theta[param] < val[1] + all_res.append(res) + + if all(all_res) and uppertest: + return 0.0 + else: + return -np.inf + + def _transform_dynesty_transform(self, param, val): ''' - Prior transformation function passed to dynesty + Subfunction for transforming the dynesty inputs to the correct prior range ''' - pass - + diff = abs(val[0]-val[1]) + if val[0] == val[1] or (val[0] == 0 and val[1] == 0) or val[0] > val[1]: + raise ValueError(f'Invalid prior range for {param}!') + + elif val[0] == 0 and val[1] != 0: + return param*val[1] + + elif val[0] != 0 and val[1] == 0: + return param*val[0] + + elif val[0] > 0 and val[1] > 0: + linear_shift = min(val) + return param*diff + linear_shift + + elif val[0] < 0 and val[1] > 0: + linear_shift = val[1] + return param*diff - linear_shift + + elif val[0] < 0 and val[0] < 0: + linear_shift = abs(max(val)) + return -param*diff - linear_shift + + else: + raise ValueError('This is a prior range we had not considered!') + + def dynesty_transform(self, theta, **kwargs): + ''' + Transform the input values to the correct prior ranges for dynesty + ''' + packed_theta = self.pack_theta(theta) + return tuple( + self._transform_dynesty_transform( + packed_theta[p],v + ) for p,v in self.prior.items() + ) + # override __subclasshook__ @classmethod def __subclasshook__(cls, C):