From bdd4b2f80375d72f8b7739d9073391e9f5ca8ab1 Mon Sep 17 00:00:00 2001 From: Taylor Bell Date: Thu, 20 Jun 2024 17:30:08 -0700 Subject: [PATCH] Fully implementing multi-planet models for starry --- docs/source/ecf.rst | 4 + .../differentiable_models/StarryModel.py | 281 ++++++++---------- 2 files changed, 132 insertions(+), 153 deletions(-) diff --git a/docs/source/ecf.rst b/docs/source/ecf.rst index 32c0f524..5d83ccc8 100644 --- a/docs/source/ecf.rst +++ b/docs/source/ecf.rst @@ -892,6 +892,10 @@ force_positivity '''''''''''''''' Optional boolean. Used by the sinusoid_pc and poet_pc models. If True, force positive phase variations (phase variations that never go below the bottom of the eclipse). Physically speaking, a negative phase curve is impossible, but strictly enforcing this can hide issues with the decorrelation or potentially bias your measured minimum flux level. Either way, use caution when choosing the value of this parameter. +mutualOccultations +'''''''''''''''''' +Optional boolean, only relevant for starry astrophysical models. If True (default), then the model will account for planet-planet occultations; if False, then the model will not include planet-planet occultations (and will likely take longer since each planet needs to be modelled separately). + manual_clip ''''''''''' Optional. A list of lists specifying the start and end integration numbers for manual removal. E.g., to remove the first 20 data points specify [[0,20]], and to also remove the last 20 data points specify [[0,20],[-20,None]]. If you want to clip the 10th integration, this would be index 9 since python uses zero-indexing. And the manual_clip start and end values are used to slice a numpy array, so they follow the same convention of *inclusive* start index and *exclusive* end index. In other words, to trim the 10th integrations, you would set manual_clip to [[9,10]]. diff --git a/src/eureka/S5_lightcurve_fitting/differentiable_models/StarryModel.py b/src/eureka/S5_lightcurve_fitting/differentiable_models/StarryModel.py index f165643a..c8ea442d 100644 --- a/src/eureka/S5_lightcurve_fitting/differentiable_models/StarryModel.py +++ b/src/eureka/S5_lightcurve_fitting/differentiable_models/StarryModel.py @@ -16,18 +16,11 @@ starry.config.lazy = True from . import PyMC3Model -# from .AstroModel import PlanetParams +from .AstroModel import PlanetParams from ..limb_darkening_fit import ld_profile from ...lib.split_channels import split -class temp_class: - # FINDME: I shouldn't need this anymore once I've - # fully setup PlanetParams - def __init__(self): - pass - - class StarryModel(PyMC3Model): def __init__(self, **kwargs): """Initialize the model. @@ -117,93 +110,83 @@ def setup(self, newparams): """ self.systems = [] self.rps = [] - for c in range(self.nchannel_fitted): - # To save ourselves from tonnes of getattr lines, let's make a - # new object without the _ch# parts of the parnames - # For example, this way we can do `temp.u1` rather than - # `getattr(self.model, 'u1_ch'+c)`. - # - # FINDME: I shouldn't need this anymore once I've - # fully setup PlanetParams - temp = temp_class() - for key in self.model.named_vars.keys(): - channum = key.split('_ch')[-1].split('_')[0] - if ((channum.isnumeric() and int(channum) == c) - or not channum.isnumeric()): - # Remove the _ch part of the parname but leave any - # other underscores intact - name = key.split('_ch')[0] - setattr(temp, name, getattr(self.model, key)) + for chan in range(self.nchannel_fitted): + # Initialize PlanetParams object + pl_params = PlanetParams(self, 0, chan, eval=False) # Initialize star object star = starry.Primary(starry.Map(udeg=self.udeg), - m=0, r=temp.Rs) - + m=0, r=pl_params.Rs) if hasattr(self.parameters, 'limb_dark'): if self.parameters.limb_dark.value == 'kipping2013': # Transform stellar variables to uniform used by starry - star.map[1] = 2*tt.sqrt(temp.u1)*temp.u2 - star.map[2] = tt.sqrt(temp.u1)*(1-2*temp.u2) + star.map[1] = 2*tt.sqrt(pl_params.u1)*pl_params.u2 + star.map[2] = tt.sqrt(pl_params.u1)*(1-2*pl_params.u2) elif self.parameters.limb_dark.value == 'quadratic': - star.map[1] = temp.u1 - star.map[2] = temp.u2 + star.map[1] = pl_params.u1 + star.map[2] = pl_params.u2 elif self.parameters.limb_dark.value == 'linear': - star.map[1] = temp.u1 + star.map[1] = pl_params.u1 elif self.parameters.limb_dark.value != 'uniform': - message = (f'ERROR: starryModel is not yet able to ' + message = (f'ERROR: Our StarryModel is not yet able to ' f'handle {self.parameters.limb_dark.value} ' f'limb darkening.\n' f' limb_dark must be one of uniform, ' f'linear, quadratic, or kipping2013.') raise ValueError(message) - if not hasattr(temp, 'fp'): - planet_map = starry.Map(ydeg=self.ydeg, amp=0) - else: - planet_map = starry.Map(ydeg=self.ydeg) - planet_map_temp = starry.Map(ydeg=self.ydeg) - for ell in range(1, self.ydeg+1): - for m in range(-ell, ell+1): - if hasattr(temp, f'Y{ell}{m}'): - planet_map[ell, m] = getattr(temp, - f'Y{ell}{m}') - planet_map_temp[ell, m] = getattr(temp, - f'Y{ell}{m}') - amp = temp.fp/tt.abs_(planet_map_temp.flux(theta=0)[0]) - planet_map.amp = amp - self.rps.append(temp.rp) - - # Solve Keplerian orbital period equation for system mass - # (otherwise starry is going to mess with P or a...) - a = temp.a*temp.Rs*const.R_sun.value - p = temp.per*(24.*3600.) - Mp = ((2.*np.pi*a**(3./2.))/p)**2/const.G.value/const.M_sun.value - - # The following code should work but doesn't see to work well - # self.model.ecc = tt.sqrt(temp.ecosw**2 + temp.esinw**2) - # longitude of periastron needs to be in degrees for batman! - # self.model.w = tt.arctan2(temp.esinw, temp.ecosw)*180./np.pi - - # Initialize planet object - planet = starry.Secondary( - planet_map, - m=Mp, - # Convert radius ratio to R_star units - r=tt.abs_(temp.rp)*temp.Rs, - a=temp.a, - # Another option to set inclination using impact parameter - # inc=tt.arccos(b/a)*180/np.pi - inc=temp.inc, - ecc=temp.ecc, - w=temp.w - ) - planet.porb = temp.per - planet.prot = temp.per - planet.theta0 = 180.0 - planet.t0 = temp.t0 + # Setup each planet + planets = [] + for pid in range(self.num_planets): + # Initialize PlanetParams object for this planet + pl_params = PlanetParams(self, pid, chan) + + if not hasattr(pl_params, 'fp'): + planet_map = starry.Map(ydeg=self.ydeg, amp=0) + else: + planet_map = starry.Map(ydeg=self.ydeg) + planet_map_temp = starry.Map(ydeg=self.ydeg) + for ell in range(1, self.ydeg+1): + for m in range(-ell, ell+1): + if hasattr(pl_params, f'Y{ell}{m}'): + planet_map[ell, m] = getattr(pl_params, + f'Y{ell}{m}') + planet_map_temp[ell, m] = getattr(pl_params, + f'Y{ell}{m}') + amp = pl_params.fp/tt.abs_( + planet_map_temp.flux(theta=0)[0]) + planet_map.amp = amp + self.rps.append(pl_params.rp) + + # Solve Keplerian orbital period equation for system mass + # (otherwise starry is going to mess with P or a...) + a = pl_params.a*pl_params.Rs*const.R_sun.value + p = pl_params.per*(24*3600) + Mp = ((2*np.pi*a**(3/2))/p)**2/const.G.value/const.M_sun.value + + # Initialize planet object + planet = starry.Secondary( + planet_map, + m=Mp, + # Convert radius ratio to R_star units + r=tt.abs_(pl_params.rp)*pl_params.Rs, + a=pl_params.a, + # Another option to set inclination using impact parameter + # inc=tt.arccos(b/a)*180/np.pi + inc=pl_params.inc, + ecc=pl_params.ecc, + w=pl_params.w + ) + planet.porb = pl_params.per + planet.prot = pl_params.per + planet.theta0 = 180.0 + planet.t0 = pl_params.t0 + + planets.append(planet) # Instantiate the system - system = starry.System(star, planet, light_delay=self.compute_ltt) + system = starry.System(star, *planets, + light_delay=self.compute_ltt) self.systems.append(system) self.update(newparams) @@ -236,7 +219,8 @@ def eval(self, eval=True, channel=None, piecewise=False, **kwargs): nchan = 1 channels = [channel, ] - # Currently can't separate starry models (given mutual occultations) + # Currently can't separately evaluate starry models + # (given mutual occultations) pid_iter = range(self.num_planets) if eval: @@ -304,11 +288,13 @@ def eval(self, eval=True, channel=None, piecewise=False, **kwargs): return returnVal - def compute_fp(self, theta=0): + def compute_fp(self, pid=0, theta=0): """Compute the planetary flux at an arbitrary orbital position. Parameters ---------- + pid : int; optional + Planet ID, default is 0. theta : int, ndarray; optional The orbital angle(s) in degrees with respect to mid-eclipse. Defaults to 0. @@ -321,7 +307,7 @@ def compute_fp(self, theta=0): with self.model: fps = [] for system in self.fit.systems: - planet_map = system.secondaries[0].map + planet_map = system.secondaries[pid].map fps.append(planet_map.flux(theta=theta).eval()) return np.array(fps) @@ -340,91 +326,80 @@ def update(self, newparams, **kwargs): self.fit.systems = [] self.fit_rps = [] - for c in range(self.nchannel_fitted): - # To save ourselves from tonnes of getattr lines, let's make a - # new object without the _ch# parts of the parnames - # For example, this way we can do `temp.u1` rather than - # `getattr(self.fit, 'u1_ch'+c)`. - # - # FINDME: I shouldn't need this anymore once I've - # fully setup PlanetParams - temp = temp_class() - for key in self.fit.__dict__.keys(): - channum = key.split('_ch')[-1].split('_')[0] - if ((channum.isnumeric() and int(channum) == c) - or not channum.isnumeric()): - # Remove the _ch part of the parname but leave any - # other underscores intact - name = key.split('_ch')[0] - setattr(temp, name, getattr(self.fit, key)) + for chan in range(self.nchannel_fitted): + # Initialize PlanetParams object + pl_params = PlanetParams(self, 0, chan, eval=True) # Initialize star object star = starry.Primary(starry.Map(udeg=self.udeg), - m=0, r=temp.Rs) - + m=0, r=pl_params.Rs) if hasattr(self.parameters, 'limb_dark'): if self.parameters.limb_dark.value == 'kipping2013': # Transform stellar variables to uniform used by starry - star.map[1] = 2*np.sqrt(temp.u1)*temp.u2 - star.map[2] = np.sqrt(temp.u1)*(1-2*temp.u2) + star.map[1] = 2*np.sqrt(pl_params.u1)*pl_params.u2 + star.map[2] = np.sqrt(pl_params.u1)*(1-2*pl_params.u2) elif self.parameters.limb_dark.value == 'quadratic': - star.map[1] = temp.u1 - star.map[2] = temp.u2 + star.map[1] = pl_params.u1 + star.map[2] = pl_params.u2 elif self.parameters.limb_dark.value == 'linear': - star.map[1] = temp.u1 + star.map[1] = pl_params.u1 elif self.parameters.limb_dark.value != 'uniform': - message = (f'ERROR: starryModel is not yet able to handle ' - f'{self.parameters.limb_dark.value} ' - f'limb_dark.\n' + message = (f'ERROR: Our StarryModel is not yet able to ' + f'handle {self.parameters.limb_dark.value} ' + f'limb darkening.\n' f' limb_dark must be one of uniform, ' f'linear, quadratic, or kipping2013.') raise ValueError(message) - if not hasattr(temp, 'fp'): - planet_map = starry.Map(ydeg=self.ydeg, amp=0) - else: - planet_map = starry.Map(ydeg=self.ydeg) - planet_map_temp = starry.Map(ydeg=self.ydeg) - for ell in range(1, self.ydeg+1): - for m in range(-ell, ell+1): - if hasattr(temp, f'Y{ell}{m}'): - planet_map[ell, m] = getattr(temp, - f'Y{ell}{m}') - planet_map_temp[ell, m] = getattr(temp, - f'Y{ell}{m}') - amp = temp.fp/np.abs(planet_map_temp.flux(theta=0)[0]) - planet_map.amp = amp - self.fit_rps.append(temp.rp) - - # Solve Keplerian orbital period equation for system mass - # (otherwise starry is going to mess with P or a...) - a = temp.a*temp.Rs*const.R_sun.value - p = temp.per*(24.*3600.) - Mp = ((2.*np.pi*a**(3./2.))/p)**2/const.G.value/const.M_sun.value - - # The following code should work but doesn't see to work well - # ecc = np.sqrt(temp.ecosw**2 + temp.esinw**2) - # longitude of periastron needs to be in degrees for batman! - # w = np.arctan2(temp.esinw, temp.ecosw)*180./np.pi - - # Initialize planet object - planet = starry.Secondary( - planet_map, - m=Mp, - # Convert radius to R_star units - r=np.abs(temp.rp)*temp.Rs, - a=temp.a, - # Another option to set inclination using impact parameter - # inc=tt.arccos(b/a)*180/np.pi - inc=temp.inc, - ecc=temp.ecc, - w=temp.w - ) - planet.porb = temp.per - planet.prot = temp.per - planet.theta0 = 180.0 - planet.t0 = temp.t0 + # Setup each planet + planets = [] + for pid in range(self.num_planets): + # Initialize PlanetParams object for this planet + pl_params = PlanetParams(self, pid, chan, eval=True) + + if not hasattr(pl_params, 'fp'): + planet_map = starry.Map(ydeg=self.ydeg, amp=0) + else: + planet_map = starry.Map(ydeg=self.ydeg) + planet_map_temp = starry.Map(ydeg=self.ydeg) + for ell in range(1, self.ydeg+1): + for m in range(-ell, ell+1): + if hasattr(pl_params, f'Y{ell}{m}'): + planet_map[ell, m] = getattr(pl_params, + f'Y{ell}{m}') + planet_map_temp[ell, m] = getattr(pl_params, + f'Y{ell}{m}') + amp = pl_params.fp/np.abs(planet_map_temp.flux(theta=0)[0]) + planet_map.amp = amp + self.fit_rps.append(pl_params.rp) + + # Solve Keplerian orbital period equation for system mass + # (otherwise starry is going to mess with P or a...) + a = pl_params.a*pl_params.Rs*const.R_sun.value + p = pl_params.per*(24*3600) + Mp = ((2*np.pi*a**(3/2))/p)**2/const.G.value/const.M_sun.value + + # Initialize planet object + planet = starry.Secondary( + planet_map, + m=Mp, + # Convert radius ratio to R_star units + r=np.abs(pl_params.rp)*pl_params.Rs, + a=pl_params.a, + # Another option to set inclination using impact parameter + # inc=tt.arccos(b/a)*180/np.pi + inc=pl_params.inc, + ecc=pl_params.ecc, + w=pl_params.w + ) + planet.porb = pl_params.per + planet.prot = pl_params.per + planet.theta0 = 180.0 + planet.t0 = pl_params.t0 + + planets.append(planet) # Instantiate the system - system = starry.System(star, planet, light_delay=self.compute_ltt) + system = starry.System(star, *planets, + light_delay=self.compute_ltt) self.fit.systems.append(system)