Adding the option to impose positive values on fitted eclipse maps

src/eureka/S5_lightcurve_fitting/differentiable_models/StarryModel.py

@@ -14,6 +14,7 @@
 import starry
 starry.config.quiet = True
 starry.config.lazy = True
+import pymc3 as pm
 
 from . import PyMC3Model
 from ..limb_darkening_fit import ld_profile
@@ -66,6 +67,9 @@
             l_vals = [int(self.paramtitles[ind][1])
                       for ind in ylm_params]
             self.ydeg = max(l_vals)
+        elif 'pixel_ydeg' in self.paramtitles:
+            # read l order used for pixel sampling
+            self.ydeg = self.parameters.pixel_ydeg.value
         else:
             self.ydeg = 0
 
@@ -190,8 +194,58 @@
             planet.theta0 = 180.0
             planet.t0 = temp.t0
 
+            # Pixel sampling setup
+            if 'pixel_ydeg' in self.paramtitles:
+
+                # Oversample factor of 3 is a safe bet to achieve pixels
+                # ~4L^2, but not always necessary 
+                if 'pixel_oversample' in self.paramtitles:
+                    self.oversample = self.parameters.pixel_oversample.value
+                else:
+                    self.oversample = 3
+
+                # Get pixel transform matrix and number of pixels
+                A = planet.map.get_pixel_transforms(
+                    oversample=self.oversample)[3]
+                self.npix = A.shape[1]
+
+                # Set prior to either be log normal, or normal around zero
+                pixel_prior_mean = self.parameters.pixel_prior_mean.value
+                pixel_prior_width = self.parameters.pixel_prior_width.value
+                if self.force_positivity:
+                    p = pm.LogNormal("p", mu=np.log(pixel_prior_mean/np.pi),
+                                     sigma=pixel_prior_width/pixel_prior_mean,
+                                     shape=(self.npix,))
+                else:
+                    p = pm.Normal("p", mu=pixel_prior_mean/np.pi,
+                                  sd=pixel_prior_width/np.pi,
+                                  shape=(self.npix, ))
+
+                # Transform pixels to spherical harmonics
+                self.starry_x = tt.dot(A, p)
+                # Record spherical harmonics
+                pm.Deterministic("y", self.starry_x)
+
             # Instantiate the system
             system = starry.System(star, planet, light_delay=self.compute_ltt)
+
+            if 'pixel_ydeg' in self.paramtitles:
+                # Calculate light curve by multiplying spherical harmonics by
+                # design matrix, then record
+                self.starry_X = system.design_matrix(self.time)
+                lcpiece = self.starry_X[:, 0] + tt.dot(self.starry_X[:, 1:],
+                                                       self.starry_x)
+
+                # Calculate and record map
+                map_plot = starry.Map(ydeg=self.ydeg)
+                map_plot.amp = self.starry_x[0]
+                map_plot[1:, :] = self.starry_x[1:]/self.starry_x[0]
+
+                pm.Deterministic("flux_model", lcpiece)
+                pm.Deterministic("map_grid",
+                                 np.pi*map_plot.render(projection="rect",
+                                                       res=100))
+
             self.systems.append(system)
 
     def eval(self, eval=True, channel=None, **kwargs):
@@ -241,13 +295,17 @@
                 time = split([time, ], self.nints, chan)[0]
 
             # Combine the planet and stellar flux (allowing negative rp)
-            fstar, fp = systems[chan].flux(time, total=False)
-            # Do some annoying math to allow theano functions to compile
-            # (correctly defined for -1 < rp < 1)
-            sign = (lib.ceil(rps[chan])+lib.floor(rps[chan]))
-            fstar = (fstar-1)*sign + 1
-            fp = fp*sign
-            lcpiece = fstar+fp
+            if 'pixel_ydeg' in self.paramtitles:
+                lcpiece = self.starry_X[:, 0] + tt.dot(self.starry_X[:, 1:],
+                                                       self.starry_x)
+            else:
+                fstar, fp = systems[chan].flux(time, total=False)
+                # Do some annoying math to allow theano functions to compile
+                # (correctly defined for -1 < rp < 1)
+                sign = (lib.ceil(rps[chan])+lib.floor(rps[chan]))
+                fstar = (fstar-1)*sign + 1
+                fp = fp*sign
+                lcpiece = fstar+fp
 
             if eval:
                 lcpiece = lcpiece.eval()
@@ -373,4 +431,13 @@
 
             # Instantiate the system
             sys = starry.System(star, planet, light_delay=self.compute_ltt)
+
+            if 'pixel_ydeg' in self.paramtitles:
+                # import pixel values and convert to spherical harmonics
+                A = planet.map.get_pixel_transforms(
+                    oversample=self.oversample)[3]
+                p_fit = newparams[-self.npix:]
+                self.starry_x = tt.dot(A, p_fit)
+                self.starry_X = sys.design_matrix(self.time)
+
             self.fit.systems.append(sys)

src/eureka/S5_lightcurve_fitting/fitters.py

@@ -333,7 +333,7 @@ def emceefitter(lc, model, meta, log, **kwargs):
                                     args=(lc, model, prior1, prior2,
                                           priortype, freenames),
                                     pool=pool)
-    log.writelog('Running emcee burn-in...')
+    log.writelog('Running emcee...')
     sampler.run_mcmc(pos, meta.run_nsteps, progress=True)
     # state = sampler.run_mcmc(pos, meta.run_nsteps, progress=True)
     # # Log some details about the burn-in phase

src/eureka/S5_lightcurve_fitting/gradient_fitters.py

@@ -12,6 +12,7 @@
 from . import plots_s5 as plots
 from .fitters import group_variables, load_old_fitparams, save_fit
 from ..lib.split_channels import get_trim
+import arviz as az
 
 
 def exoplanetfitter(lc, model, meta, log, calling_function='exoplanet',
@@ -185,6 +186,17 @@
 
     model.setup(lc.time, lc.flux, lc.unc)
 
+    if hasattr(meta, 'exoplanet_first') and meta.exoplanet_first:
+        log.writelog('\nCalling exoplanet first...')
+        # RUN LEAST SQUARES
+        opt_sol = exoplanetfitter(lc, model, meta, log,
+                                  calling_function='nuts', **kwargs)
+        freepars = opt_sol.fit_params
+
+        if meta.rescale_err:
+            # Scale uncertainties with reduced chi-squared
+            lc.unc *= np.sqrt(opt_sol.chi2red)
+
     start = {}
     for name, val in zip(freenames, freepars):
         start[name] = val
@@ -232,6 +244,16 @@
     upper_errs = q[2]-q[1]
     lower_errs = q[1]-q[0]
 
+    # If pixel sampling, append best fit pixels to fit_params
+    if "pixel_ydeg" in indep_vars:
+        trace_az = az.from_pymc3(trace, model=model.model)
+        trace_fname = meta.outputdir+"S5_trace_map.nc"
+        log.writelog(f'Saving map trace to {trace_fname}...')
+        trace_az.to_netcdf(trace_fname)
+
+        fit_params = np.append(fit_params,
+                               trace_az['posterior']['p'][0, 0])
+
     model.update(fit_params)
     model.errs = dict(zip(freenames, errs))
     if "scatter_ppm" in freenames:
@@ -291,6 +313,12 @@
                                   'sinusoid_pc' in meta.run_myfuncs)):
         plots.plot_phase_variations(lc, model, meta, fitter='nuts')
 
+    # Show the inferred planetary brightness map
+    if meta.isplots_S5 >= 1 and "pixel_ydeg" in indep_vars:
+        eclipse_maps = np.transpose(trace_az.posterior.stack(
+            sample=("chain", "draw"))['map_grid'][:], [2, 0, 1])
+        plots.plot_eclipse_map(lc, eclipse_maps, meta)
+
     # Plot Allan plot
     if meta.isplots_S5 >= 3:
         plots.plot_rms(lc, model, meta, fitter='nuts')

src/eureka/S5_lightcurve_fitting/plots_s5.py

@@ -97,11 +97,42 @@
             nbin_plot = len(time)
         else:
             nbin_plot = meta.nbin_plot
+
         binned_time = util.binData_time(time, time, nbin_plot)
         binned_flux = util.binData_time(flux, time, nbin_plot)
         binned_unc = util.binData_time(unc, time, nbin_plot, err=True)
         binned_normflux = util.binData_time(flux/model_sys-gp, time, nbin_plot)
         binned_res = util.binData_time(residuals, time, nbin_plot)
+        binned_model_phys = util.binData_time(model_phys, time, nbin_plot)
+
+        if hasattr(meta, 'record_plot_data'):
+            if meta.record_plot_data:
+                np.savez_compressed(meta.outputdir+'time.npz',
+                                    data=time.data, mask=time.mask)
+                np.savez_compressed(meta.outputdir+'flux.npz',
+                                    data=flux.data, mask=flux.mask)
+                np.savez_compressed(meta.outputdir+'unc.npz',
+                                    data=unc.data, mask=unc.mask)
+                np.savez_compressed(meta.outputdir+'model_phys.npz',
+                                    data=model_phys.data)
+                np.savez_compressed(meta.outputdir+'model_sys.npz',
+                                    data=model_sys.data)
+
+                np.savez_compressed(meta.outputdir+'binned_time.npz',
+                                    data=binned_time.data,
+                                    mask=binned_time.mask)
+                np.savez_compressed(meta.outputdir+'binned_flux.npz',
+                                    data=binned_flux.data,
+                                    mask=binned_flux.mask)
+                np.savez_compressed(meta.outputdir+'binned_unc.npz',
+                                    data=binned_unc.data,
+                                    mask=binned_unc.mask)
+                np.savez_compressed(meta.outputdir+'binned_normflux.npz',
+                                    data=binned_normflux.data)
+                np.savez_compressed(meta.outputdir+'binned_res.npz',
+                                    data=binned_res.data)
+                np.savez_compressed(meta.outputdir+'binned_model_phys.npz',
+                                    data=binned_model_phys.data)
 
         fig = plt.figure(5101, figsize=(8, 6))
         plt.clf()
@@ -781,3 +812,116 @@
         fig.savefig(meta.outputdir+fname, bbox_inches='tight', dpi=300)
         if not meta.hide_plots:
             plt.pause(0.2)
+
+
+def plot_eclipse_map(lc, flux_maps, meta):
+    """Plot fitted eclipse map and lat-lon slices (Figs 5105)
+
+    Parameters
+    ----------
+    lc : eureka.S5_lightcurve_fitting.lightcurve.LightCurve
+        The lightcurve data object.
+    flux_maps : array
+        The posterior distribution of the fitted maps
+    meta : eureka.lib.readECF.MetaClass
+        The metadata object.
+
+    Notes
+    -----
+    History:
+
+    - August 14, 2023 Mark Hammond
+        Initial version.
+    """
+    for i, channel in enumerate(lc.fitted_channels):
+        fig = plt.figure(5105, figsize=(12, 3))
+        fig.clf()
+        axs = fig.subplots(1, 3, width_ratios=[1.4, 1, 1])
+
+        lons = np.linspace(-180, 180, np.shape(flux_maps)[2])
+        lats = np.linspace(-90, 90, np.shape(flux_maps)[1])
+
+        # Quantiles
+        p1 = 0.841
+        p2 = 0.977
+        p3 = 0.998
+
+        # Plot median map
+        ca = axs[0].contourf(lons, lats,
+                             np.quantile(1e6*flux_maps[:], 0.5, axis=0),
+                             cmap='RdBu_r')
+        axs[0].axhline(0, color='C0', ls='--')
+        axs[0].axvline(0, color='C3', ls='--')
+        axs[0].set_xticks([-180, -90, 0, 90, 180])
+        axs[0].set_yticks([-90, -45, 0, 45, 90])
+        axs[0].set_xlim([-180, 180])
+        axs[0].set_ylim([-90, 90])
+        plt.colorbar(ca, ax=axs[0], pad=0.02)
+
+        # Plot slice along equator
+        lat0 = int(np.shape(flux_maps)[2]/2)
+        axs[1].fill_between(lons,
+                            1e6*np.quantile(flux_maps[:, lat0], p1, axis=0),
+                            1e6*np.quantile(flux_maps[:, lat0], 1-p1, axis=0),
+                            color='C0', alpha=0.3, ls='None')
+        axs[1].fill_between(lons,
+                            1e6*np.quantile(flux_maps[:, lat0], p2, axis=0),
+                            1e6*np.quantile(flux_maps[:, lat0], 1-p2, axis=0),
+                            color='C0', alpha=0.3, ls='None')
+        axs[1].fill_between(lons,
+                            1e6*np.quantile(flux_maps[:, lat0], p3, axis=0),
+                            1e6*np.quantile(flux_maps[:, lat0], 1-p3, axis=0),
+                            color='C0', alpha=0.3, ls='None')
+        axs[1].set_xlim([-180, 180])
+        axs[1].set_xticks([-180, -90, 0, 90, 180])
+
+        # Plot slice along equator
+        lon0 = int(np.shape(flux_maps)[1]/2)
+        axs[2].fill_between(
+            lats,
+            1e6*np.quantile(flux_maps[:, :, lon0], p1, axis=0),
+            1e6*np.quantile(flux_maps[:, :, lon0], 1-p1, axis=0),
+            color='C3', alpha=0.3, ls='None')
+        axs[2].fill_between(
+            lats,
+            1e6*np.quantile(flux_maps[:, :, lon0], p2, axis=0),
+            1e6*np.quantile(flux_maps[:, :, lon0], 1-p2, axis=0),
+            color='C3', alpha=0.3, ls='None')
+        axs[2].fill_between(
+            lats,
+            1e6*np.quantile(flux_maps[:, :, lon0], p3, axis=0),
+            1e6*np.quantile(flux_maps[:, :, lon0], 1-p3, axis=0),
+            color='C3', alpha=0.3, ls='None')
+        axs[2].set_xlim([-90, 90])
+        axs[2].set_xticks([-90, -45, 0, 45, 90])
+
+        axs[0].set_title('Median Map', fontsize=12)
+        axs[1].set_title('Fit Along Equator', fontsize=12)
+        axs[2].set_title(r'Fit Along 0$^{\circ}$ Longitude', fontsize=12)
+        axs[0].set_ylabel('Latitude', labelpad=-10)
+        axs[0].set_xlabel('Longitude')
+        axs[1].set_xlabel('Longitude')
+        axs[2].set_xlabel('Latitude')
+        axs[1].set_ylabel(r'$F_{\rm p}/F_{\rm s}$ (ppm)', fontsize=11,
+                          labelpad=-4)
+        axs[2].set_ylabel(r'$F_{\rm p}/F_{\rm s}$ (ppm)', fontsize=11,
+                          labelpad=-4)
+
+        plt.subplots_adjust(left=0.07,
+                            bottom=0.17,
+                            right=0.98,
+                            top=0.81,
+                            wspace=0.38,
+                            hspace=0.22)
+
+        if lc.white:
+            fname_tag = 'white'
+        else:
+            ch_number = str(channel).zfill(len(str(lc.nchannel)))
+            fname_tag = f'ch{ch_number}'
+
+        fname = (f'figs{os.sep}fig5105_{fname_tag}_eclipse_map' +
+                 plots.figure_filetype)
+        fig.savefig(meta.outputdir+fname, bbox_inches='tight', dpi=300)
+        if not meta.hide_plots:
+            plt.pause(0.2)