diff --git a/MonoTools/fit.py b/MonoTools/fit.py index 62cb6ef..44e0943 100755 --- a/MonoTools/fit.py +++ b/MonoTools/fit.py @@ -207,7 +207,14 @@ def save_model_to_file(self, savefile=None, limit_size=False): self.get_savename(how='save') savefile=self.savenames[0]+'_model.pickle' if hasattr(self,'trace'): - self.trace.to_netcdf(self.savenames[0]+'_trace.nc') + try: + self.trace.to_netcdf(self.savenames[0]+'_trace.nc') + except: + try: + #Stacking/unstacking removes Multitrace objects: + self.trace.unstack().to_netcdf(self.savenames[0]+'_trace.nc') + except: + print("Still a save error after unstacking") excl_types=[az.InferenceData] cloudpickle.dumps({attr:getattr(self,attr) for attr in self.__dict__},open(savefile,'wb')) @@ -4643,7 +4650,7 @@ def plot_corner(self,corner_vars=None,use_marg=True,truths=None): if not self.assume_circ: corner_vars+=['multi_ecc','multi_omega'] ''' - samples = pm.trace_to_dataframe(self.trace, varnames=corner_vars) + samples =self.make_table(cols=corner_vars) #print(samples.shape,samples.columns) assert samples.shape[1]<50 @@ -4715,13 +4722,13 @@ def make_table(self,short=True,save=True,cols=['all']): cols_to_remove+=['mono_'+col+'s','duo_'+col+'s'] medvars=[var for var in self.trace if not np.any([icol in var for icol in cols_to_remove])] #print(cols_to_remove, medvars) - df = pm.summary(self.trace,var_names=medvars,stat_funcs={"5%": lambda x: np.percentile(x, 5), + df = pm.summary(self.trace.unstack(),var_names=medvars,stat_funcs={"5%": lambda x: np.percentile(x, 5), "-$1\sigma$": lambda x: np.percentile(x, 15.87), "median": lambda x: np.percentile(x, 50), "+$1\sigma$": lambda x: np.percentile(x, 84.13), "95%": lambda x: np.percentile(x, 95)},round_to=5) else: - df = pm.summary(self.trace,var_names=cols,stat_funcs={"5%": lambda x: np.percentile(x, 5), + df = pm.summary(self.trace.unstack(),var_names=cols,stat_funcs={"5%": lambda x: np.percentile(x, 5), "-$1\sigma$": lambda x: np.percentile(x, 15.87), "median": lambda x: np.percentile(x, 50), "+$1\sigma$": lambda x: np.percentile(x, 84.13), @@ -4880,23 +4887,22 @@ def predict_future_transits(self, time_start=None, time_end=None, time_dur=180, print("time range",Time(time_start+self.lc.jd_base,format='jd').isot, "->",Time(time_end+self.lc.jd_base,format='jd').isot) - if check_TESS: sect_start_ends=self.check_TESS() all_trans_fin=pd.DataFrame() loopplanets = self.duos+self.trios+self.multis if include_multis else self.duos+self.trios - + all_unq_trans=[] for pl in loopplanets: all_trans=pd.DataFrame() if pl in self.duos+self.trios: sum_all_probs=np.logaddexp.reduce(np.nanmedian(self.trace['logprob_marg_'+pl],axis=0)) - trans_p0=np.floor(np.nanmedian(time_start - self.trace['t0_2_'+pl])/np.nanmedian(self.trace['per_'+pl],axis=0)) - trans_p1=np.ceil(np.nanmedian(time_end - self.trace['t0_2_'+pl])/np.nanmedian(self.trace['per_'+pl],axis=0)) + trans_p0=np.floor(np.nanmedian(time_start - self.trace['t0_2_'+pl].values)/np.nanmedian(self.trace['per_'+pl].values,axis=0)) + trans_p1=np.ceil(np.nanmedian(time_end - self.trace['t0_2_'+pl].values)/np.nanmedian(self.trace['per_'+pl].values,axis=0)) n_trans=trans_p1-trans_p0 elif pl in self.multis: - trans_p0=[np.floor(np.nanmedian(time_start - self.trace['t0_'+pl])/np.nanmedian(self.trace['per_'+pl],axis=0))] - trans_p1=[np.ceil(np.nanmedian(time_end - self.trace['t0_'+pl])/np.nanmedian(self.trace['per_'+pl],axis=0))] + trans_p0=[np.floor(np.nanmedian(time_start - self.trace['t0_'+pl].values)/np.nanmedian(self.trace['per_'+pl].values))] + trans_p1=[np.ceil(np.nanmedian(time_end - self.trace['t0_'+pl].values)/np.nanmedian(self.trace['per_'+pl].values))] n_trans=[trans_p1[0]-trans_p0[0]] #print(pl,trans_p0,trans_p1,n_trans) #print(np.nanmedian(self.trace['t0_2_'+pl])+np.nanmedian(self.trace['per_'+pl],axis=0)*trans_p0) @@ -4909,21 +4915,21 @@ def predict_future_transits(self, time_start=None, time_end=None, time_dur=180, if 'tdur' in self.fit_params or pl in self.multis: dur=np.nanpercentile(self.trace['tdur_'+pl],percentiles) naliases=[0] if pl in self.multis else np.arange(self.planets[pl]['npers']) + idfs=[] for nd in naliases: if n_trans[nd]>0: if pl in self.duos+self.trios: int_alias=int(self.planets[pl]['period_int_aliases'][nd]) - transits=np.nanpercentile(np.vstack([self.trace['t0_2_'+pl]+ntr*self.trace['per_'+pl][:,nd] for ntr in np.arange(trans_p0[nd],trans_p1[nd])]),percentiles,axis=1) + transits=np.nanpercentile(np.vstack([self.trace['t0_2_'+pl].values+ntr*self.trace['per_'+pl].values[:,nd] for ntr in np.arange(trans_p0[nd],trans_p1[nd])]),percentiles,axis=1) if 'tdur' in self.marginal_params: dur=np.nanpercentile(self.trace['tdur_'+pl][:,nd],percentiles) logprobs=np.nanmedian(self.trace['logprob_marg_'+pl][:,nd])-sum_all_probs else: - transits=np.nanpercentile(np.vstack([self.trace['t0_'+pl]+ntr*self.trace['per_'+pl] for ntr in np.arange(trans_p0[nd],trans_p1[nd])]),percentiles,axis=1) + transits=np.nanpercentile(np.column_stack([self.trace['t0_'+pl].values+ntr*self.trace['per_'+pl].values for ntr in np.arange(trans_p0[nd],trans_p1[nd],1.0)]),percentiles,axis=0) int_alias=1 logprobs=np.array([0.0]) #Getting the aliases for this: - - idf=pd.DataFrame({'transit_mid_date':Time(transits[2]+self.lc.jd_base,format='jd').isot, + idfs+=[pd.DataFrame({'transit_mid_date':Time(transits[2]+self.lc.jd_base,format='jd').isot, 'transit_mid_med':transits[2], 'transit_dur_med':np.tile(dur[2],len(transits[2])), 'transit_dur_-1sig':np.tile(dur[1],len(transits[2])), @@ -4945,8 +4951,8 @@ def predict_future_transits(self, time_start=None, time_end=None, time_dur=180, 'prob':np.tile(np.exp(logprobs),len(transits[2])), 'planet_name':np.tile('multi_'+pl,len(transits[2])) if pl in self.multis else np.tile('duo_'+pl,len(transits[2])), 'alias_n':np.tile(nd,len(transits[2])), - 'alias_p':np.tile(np.nanmedian(self.trace['per_'+pl][:,nd]),len(transits[2])) if pl in self.duos+self.trios else np.nanmedian(self.trace['per_'+pl])}) - all_trans=all_trans.append(idf) + 'alias_p':np.tile(np.nanmedian(self.trace['per_'+pl].values[:,nd]),len(transits[2])) if pl in self.duos+self.trios else np.tile(np.nanmedian(self.trace['per_'+pl].values),len(transits[2]))})] + all_trans=pd.concat(idfs) unq_trans = all_trans.sort_values('log_prob').copy().drop_duplicates('transit_fractions') unq_trans = unq_trans.set_index(np.arange(len(unq_trans))) unq_trans['aliases_ns']=unq_trans['alias_n'].values.astype(str) @@ -4960,7 +4966,8 @@ def predict_future_transits(self, time_start=None, time_end=None, time_dur=180, unq_trans.loc[i,'aliases_ps']=','.join(list(np.round(oths['alias_p'].values,4).astype(str))) unq_trans.loc[i,'num_aliases']=len(oths) unq_trans.loc[i,'total_prob']=np.sum(oths['prob'].values) - all_trans_fin=all_trans_fin.append(unq_trans) + all_unq_trans+=[unq_trans] + all_trans_fin=pd.concat(all_unq_trans) all_trans_fin = all_trans_fin.loc[(all_trans_fin['transit_end_+2sig']>time_start)*(all_trans_fin['transit_start_-2sig']