diff --git a/.gitignore b/.gitignore index 22ce1b4..7f8f5a0 100644 --- a/.gitignore +++ b/.gitignore @@ -55,6 +55,7 @@ distribute-*.tar.gz .project .pydevproject .settings +settings.json docs/.DS_Store .DS_Store diff --git a/docs/fleck/jax.rst b/docs/fleck/jax.rst index 85e995c..76b9e3c 100644 --- a/docs/fleck/jax.rst +++ b/docs/fleck/jax.rst @@ -201,12 +201,22 @@ that we will observe, the stellar inclination, and spectrum of the stellar photo We can add active regions to the star with `~fleck.jax.ActiveStar.add_spot` +Spot latitudes are signed angles in radians, from ``-pi/2`` to ``pi/2``; +zero is the stellar equator. To convert a colatitude, use +``latitude = pi/2 - colatitude``. + +Instead of a spectrum, you can provide a scalar or wavelength-dependent +``contrast`` to ``add_spot``. A contrast below one describes a dark spot; +a contrast above one describes a bright region. Supplying only ``temperature`` +uses a blackbody spot spectrum. If no photospheric spectrum is provided, +``T_eff`` and the wavelength grid define a blackbody photosphere. + .. code-block:: python # add a cool spot: active_star.add_spot( lon=-0.2, # [rad] - lat=1.65, # [rad] + lat=np.pi/2 - 1.65, # signed latitude [rad] rad=0.15, # [R_star] spectrum=cool.flux.value, temperature=cool.meta['PHXTEFF'] @@ -215,7 +225,7 @@ We can add active regions to the star with `~fleck.jax.ActiveStar.add_spot` # add a hot spot: active_star.add_spot( lon=0.95, - lat=1.75, + lat=np.pi/2 - 1.75, rad=0.08, spectrum=hot.flux.value, temperature=hot.meta['PHXTEFF'] @@ -376,7 +386,7 @@ code cell above, makes a lot of plotting calls to visualize the results of # add a cool spot: active_star.add_spot( lon=-0.2, # [rad] - lat=1.65, # [rad] + lat=np.pi/2 - 1.65, # signed latitude [rad] rad=0.15, # [R_star] spectrum=cool.flux.value, temperature=cool.meta['PHXTEFF'] @@ -385,7 +395,7 @@ code cell above, makes a lot of plotting calls to visualize the results of # add a hot spot: active_star.add_spot( lon=0.95, - lat=1.75, + lat=np.pi/2 - 1.75, rad=0.08, spectrum=hot.flux.value, temperature=hot.meta['PHXTEFF'] diff --git a/fleck/core.py b/fleck/core.py index 1482229..eff5dd1 100644 --- a/fleck/core.py +++ b/fleck/core.py @@ -278,6 +278,19 @@ def light_curve(self, spot_lons, spot_lats, spot_radii, inc_stellar, np.cos(Omega) * np.sin(omega + f) * np.cos(I)) Z = planet.a * np.sin(omega + f) * np.sin(I) + # Enforce that the transit chord lies below the stellar equator + # in the same "observer oriented" frame used for the spots/plot. + in_front = (X < 0) & (np.abs(Y) < 1 + planet.rp) + if np.any(in_front): + idx = np.where(in_front)[0] + # approximate mid-transit as the in-front point with minimum |Y| + mid_idx = idx[np.argmin(np.abs(Y[in_front]))] + z_star_mid = -Z[mid_idx] + if z_star_mid > 0: + # Planet is crossing the lat > 0 hemisphere; reflect it + # across the equator so it crosses lat < 0 instead. + Z = -Z + # Create a shapely circle object for the planet's silhouette only # when the planet is in front of the star, otherwise append `None` planet_disk = [circle([-Y[i], -Z[i]], planet.rp) diff --git a/fleck/jax.py b/fleck/jax.py index aa619d9..6c9dc63 100644 --- a/fleck/jax.py +++ b/fleck/jax.py @@ -30,7 +30,8 @@ class ActiveStar: optional planetary transit models and spot occultations. """ - n_mc = 10_000 # Number of Monte Carlo samples to use when computing planet+spot overlap + # Number of Monte Carlo samples to use when computing planet+spot overlap + n_mc = 10_000 key = random.PRNGKey(0) # random key seed def __init__( @@ -45,7 +46,8 @@ def __init__( inclination=empty, wavelength=None, phot=None, - P_rot=3.3 + P_rot=3.3, + contrast=empty, ): """ Parameters @@ -55,7 +57,7 @@ def __init__( lon : array Active region longitudes in radians on (0, 2pi) lat : array - Active region latitudes in radians on (0, pi) + Active region latitudes in radians on (-pi/2, pi/2) rad : array Active region radii in units of stellar radii spectrum : array @@ -69,22 +71,136 @@ def __init__( wavelength : array Wavelength for each flux observation in ``phot`` [meters] phot : array - Photospheric flux at each ``wavelength``. + Photospheric flux at each ``wavelength``. If omitted, use a + blackbody when ``T_eff`` and ``wavelength`` are provided, + otherwise use unity. P_rot : float Stellar rotation period + contrast : float or array + Ratio of the active region flux to the photospheric + flux at each ``wavelength``. If provided, this will be + used directly; otherwise it is computed from ``spectrum`` + and ``phot`` when possible. """ self.times = jnp.array(times) self.lon = jnp.array(lon) self.lat = jnp.array(lat) self.rad = jnp.array(rad) - self.spectrum = jnp.array(spectrum) + self.spectrum = jnp.array([]) if spectrum is None else jnp.array(spectrum) self.T_eff = T_eff - self.temperature = jnp.array(temperature) + self.temperature = (jnp.array([]) if temperature is None + else jnp.atleast_1d(jnp.asarray(temperature))) self.inclination = inclination - self.wavelength = wavelength - self.phot = phot + # Allow wavelength/phot to remain optional, but store them as + # JAX arrays when provided. + self.wavelength = None if wavelength is None else jnp.array(wavelength) + self.phot = None if phot is None else jnp.array(phot) self.P_rot = P_rot + # Raw contrast as passed by the user (may be scalar, 1D, 2D, or + # empty). We normalise shapes after dimension inference. + self.contrast = jnp.array([]) if contrast is None else jnp.array(contrast) + + # ------------------------------------------------------------------ + # Dimension bookkeeping: infer n_times, n_spots, n_lambda + # ------------------------------------------------------------------ + self.n_times = int(self.times.size) + self.n_spots = int(self.rad.size) + self.n_lambda = None + + if ( + self.wavelength is not None and + getattr(self.wavelength, "size", 0) != 0 + ): + # 1) Prefer the explicit wavelength grid, if provided + self.n_lambda = int(self.wavelength.size) + + elif getattr(self.spectrum, "size", 0) != 0: + # 2) Otherwise, use the spectrum shape + + if self.spectrum.ndim == 1: + # Treat a 1D spectrum as (n_lambda,) + self.n_lambda = int(self.spectrum.shape[0]) + else: + # Assume the last axis is wavelength + self.n_lambda = int(self.spectrum.shape[-1]) + + elif getattr(self.phot, "size", 0) != 0: + # 3) Otherwise, use the photosphere + + self.n_lambda = self.phot.size + + elif getattr(self.contrast, "size", 0) != 0: + # 4) Finally, fall back to the contrast shape, using n_spots to + # break degeneracies between (n_spots,) and (n_lambda,) + + c = self.contrast + if c.ndim == 0: + self.n_lambda = 1 + elif c.ndim == 1: + if self.n_spots > 0 and c.shape[0] == self.n_spots: + # One contrast per spot, same for all wavelengths + self.n_lambda = 1 + else: + # One contrast per wavelength bin, same for all spots + self.n_lambda = int(c.shape[0]) + elif c.ndim == 2: + # Last axis is wavelength + self.n_lambda = int(c.shape[-1]) + + else: + # Absolute last resort: a single "white-light" wavelength bin + self.n_lambda = 1 + + has_temperatures = self.temperature.size != 0 + self.temperature = (jnp.broadcast_to(self.temperature, (self.n_spots,)) + if has_temperatures + else jnp.full((self.n_spots,), jnp.nan)) + + # Normalise the supplied contrast argument into something at least + # 2D where possible. + if self.contrast.size != 0: + if self.contrast.ndim == 0: + # Single scalar contrast: broadcast to all spots in a single + # wavelength bin. + self.contrast = self.contrast * jnp.ones((self.n_spots, 1)) + elif self.contrast.ndim == 1: + if self.n_spots > 0 and self.contrast.shape[0] == self.n_spots: + # Per-spot contrast for a single wavelength bin. + self.contrast = self.contrast[:, None] + else: + # Per-wavelength contrast for a single spot. + self.contrast = self.contrast[None, :] + # If contrast is already 2D+ we assume the user passed + # (n_spots, n_lambda) and leave it as-is. + + if self.phot is None: + if self.T_eff is not None and self.wavelength is not None: + self.phot = self._blackbody(self.wavelength, self.T_eff) + else: + self.phot = jnp.ones((self.n_lambda,)) + self.phot = jnp.broadcast_to(self.phot, (self.n_lambda,)) + + # Keep spectral arrays aligned along spot and wavelength axes. + shape = (self.n_spots, self.n_lambda) + if self.contrast.size != 0: + self.contrast = jnp.broadcast_to(self.contrast, shape) + self.spectrum = self.contrast * self.phot[None, :] + else: + if self.spectrum.size == 0 and has_temperatures: + if self.wavelength is None: + raise ValueError("Temperature-based spots require wavelengths.") + self.spectrum = self._blackbody( + self.wavelength[None, :], self.temperature[:, None]) + if self.spectrum.size != 0: + if self.spectrum.ndim == 1 and self.n_lambda == 1: + self.spectrum = self.spectrum[:, None] + self.spectrum = jnp.broadcast_to(self.spectrum, shape) + self.contrast = self.spectrum / self.phot[None, :] + else: + self.contrast = jnp.zeros(shape) + self.spectrum = self.contrast * self.phot[None, :] + def tree_flatten(self): children = ( self.times, @@ -98,6 +214,7 @@ def tree_flatten(self): self.wavelength, self.phot, self.P_rot, + self.contrast, ) aux_data = None return children, aux_data @@ -137,7 +254,7 @@ def rotation_model(self, f0=0, t0_rot=0, u1=0, u2=0): radial_coord = 1 - jnp.geomspace(1e-5, 1, 100)[::-1] unspotted_total_flux = trapezoid( y=( - 2 * np.pi * radial_coord * + 2 * jnp.pi * radial_coord * self.limb_darkening(radial_coord, u1, u2) ), x=radial_coord @@ -145,7 +262,7 @@ def rotation_model(self, f0=0, t0_rot=0, u1=0, u2=0): # Morris 2020 Eqn 6-7 spot_model = f0 - jnp.sum( - np.pi * rad ** 2 * + jnp.pi * rad ** 2 * (1 - contrast) * self.limb_darkening(mu, u1, u2) * mask_behind_star, @@ -171,19 +288,23 @@ def spot_coords(self, times=None, t0_rot=0): Returns ------- spot_position_x : array - x-position of the active region in the observer oriented coordinate system [1]_. + x-position of the active region in the observer oriented + coordinate system [1]_. spot_position_y : array - y-position of the active region in the observer oriented coordinate system [1]_. + y-position of the active region in the observer oriented + coordinate system [1]_. spot_position_z : array - y-position of the active region in the observer oriented coordinate system [1]_. + y-position of the active region in the observer oriented + coordinate system [1]_. major_axis : array - Apparent semimajor axis of the circular active region, which is elliptical when - projected active onto the sky plane (in general) + Apparent semimajor axis of the circular active region, which is + elliptical when projected active onto the sky plane (in general) minor_axis : array - Apparent semiminor axis of the circular active region, which is elliptical when - projected active onto the sky plane (in general) + Apparent semiminor axis of the circular active region, which is + elliptical when projected active onto the sky plane (in general) angle : array - Angle between the +x-axis and the projected active region's semimajor axis + Angle between the +x-axis and the projected active region's + semimajor axis rad : array Active region radius [stellar radii] contrast: array @@ -193,17 +314,53 @@ def spot_coords(self, times=None, t0_rot=0): ---------- .. [1] Fabrycky & Winn (2009) https://arxiv.org/abs/0902.0737 """ - contrast = self.spectrum / self.phot[None, :] - - if contrast.ndim == 1: - contrast = contrast[None, :] - if times is None: times = self.times + # Use the dimensions inferred in __init__ + n_spots = int(self.n_spots) + n_lambda = int(self.n_lambda) + + # ------------------------------------------------------------------ + # Build a per-spot, per-wavelength contrast array of shape + # (n_spots, n_lambda), then broadcast to (1, n_spots, n_lambda, 1) + # for use with the f_S tensor in transit_model. + # ------------------------------------------------------------------ + if self.contrast.size != 0: + base_contrast = jnp.asarray(self.contrast) + elif self.spectrum.size != 0: + base_contrast = jnp.asarray(self.spectrum) / self.phot[None, :] + else: + # No contrast information available. + # Default to maximum contrast: spot == 0 + base_contrast = jnp.zeros((n_spots, n_lambda)) + + # Normalize to exactly (n_spots, n_lambda) + if n_spots == 0: + # No spots: keep a well-defined empty (0, n_lambda) array + # to avoid ambiguous reshapes. + base_contrast = jnp.zeros((0, n_lambda)) + elif base_contrast.ndim == 0: + base_contrast = jnp.full((n_spots, n_lambda), base_contrast) + elif base_contrast.ndim == 1: + if base_contrast.shape[0] == n_spots: + # One contrast per spot, broadcast over wavelength. + base_contrast = base_contrast[:, None]*jnp.ones((1, n_lambda)) + else: + # One contrast per wavelength, broadcast over spots. + base_contrast = base_contrast[None, :]*jnp.ones((n_spots, 1)) + elif base_contrast.shape[0] == 1 and n_spots > 1: + # Single-spot input, broadcast along spot axis first. + base_contrast = jnp.broadcast_to( + base_contrast, (n_spots,)+base_contrast.shape[1:]) + + # Final contrast tensor used downstream: + # shape: (1, n_spots, n_lambda, 1) + contrast = base_contrast[None, :, :, None] + """ Limits: - lat: (0, pi) + lat: (-pi/2, pi/2) lon: (0, 2pi) rad: (0, None) contrast: (0, inf) @@ -216,15 +373,15 @@ def spot_coords(self, times=None, t0_rot=0): 3. inclination """ - phase = jnp.expand_dims(2 * np.pi * (times - t0_rot) / self.P_rot, [1, 2, 3]) + phase = jnp.expand_dims(2 * jnp.pi * (times - t0_rot) / self.P_rot, + [1, 2, 3]) lon = jnp.expand_dims(self.lon, [0, 2, 3]) lat = jnp.expand_dims(self.lat, [0, 2, 3]) rad = jnp.expand_dims(self.rad, [0, 2, 3]) - contrast = jnp.expand_dims(contrast, [0, 3]) inclination = jnp.expand_dims(jnp.asarray(self.inclination), [0, 1, 2]) - comp_inclination = np.pi / 2 - inclination - phi = np.pi / 2 - phase - lon + comp_inclination = jnp.pi / 2 - inclination + phi = jnp.pi / 2 - phase - lon sin_lat = jnp.sin(lat) cos_lat = jnp.cos(lat) @@ -232,13 +389,13 @@ def spot_coords(self, times=None, t0_rot=0): cos_c_inc = jnp.cos(comp_inclination) spot_position_x = ( - jnp.cos(phi - np.pi / 2) * sin_c_inc * sin_lat + - cos_c_inc * cos_lat + jnp.cos(phi - jnp.pi / 2) * sin_c_inc * cos_lat + + cos_c_inc * sin_lat ) - spot_position_y = -jnp.sin(phi - np.pi / 2) * sin_lat + spot_position_y = -jnp.sin(phi - jnp.pi / 2) * cos_lat spot_position_z = ( - cos_lat * sin_c_inc - - jnp.sin(phi) * cos_c_inc * sin_lat + sin_lat * sin_c_inc - + jnp.sin(phi) * cos_c_inc * cos_lat ) rsq = spot_position_x ** 2 + spot_position_y ** 2 @@ -252,7 +409,8 @@ def spot_coords(self, times=None, t0_rot=0): major_axis, minor_axis, angle, rad, contrast ) - def add_spot(self, lon, lat, rad, contrast=None, temperature=None, spectrum=None): + def add_spot(self, lon, lat, rad, contrast=None, temperature=None, + spectrum=None): """ Add an active region to the stellar model. @@ -261,38 +419,54 @@ def add_spot(self, lon, lat, rad, contrast=None, temperature=None, spectrum=None lon : float Active region longitudes in radians on (0, 2pi) lat : float - Active region latitudes in radians on (0, pi) + Active region latitudes in radians on (-pi/2, pi/2) rad : float Active region radii in units of stellar radii - contrast : float + contrast : float or array Ratio of the active region's flux to the photospheric - flux at each ``ActiveStar.wavelength`` - spectrum : float + flux at each ``ActiveStar.wavelength``. + temperature : float + Effective temperature of the active region. If provided + (and neither ``contrast`` nor ``spectrum`` are given), + a blackbody spectrum will be computed. + spectrum : float or array The spectrum of the active region on the same wavelength - grid is ``ActiveStar.phot`` + grid as ``ActiveStar.phot``. """ - if contrast is None and spectrum is None and temperature is not None: - self.phot = self._blackbody(self.wavelength, self.T_eff) + if contrast is not None: + # User-supplied contrast: convert to spectrum + contrast = jnp.broadcast_to(jnp.asarray(contrast).squeeze(), + (self.n_lambda,)) + # Derive the spectrum from the contrast + spectrum = contrast * self.phot + elif spectrum is not None: + # User-supplied spectrum: ensure correct shape + spectrum = jnp.broadcast_to(jnp.asarray(spectrum).squeeze(), + (self.n_lambda,)) + # Derive the contrast from the spectrum + contrast = spectrum / self.phot + elif temperature is not None: + # No contrast or spectrum, but a temperature: compute a blackbody + if self.wavelength is None: + raise ValueError("Temperature-based spots require wavelengths.") spectrum = self._blackbody(self.wavelength, temperature) - - for attr, new_value in zip("lon, lat, rad, spectrum, temperature".split(', '), - [lon, lat, rad, spectrum, temperature]): - - prop = getattr(self, attr) - - if not hasattr(new_value, 'ndim'): - new_value = jnp.array([new_value]) - - if prop is not None: - if prop.ndim > 1 or (len(prop) > 1 and len(prop) == len(new_value)): - new_value = jnp.vstack([prop, new_value]) - else: - new_value = jnp.concatenate([prop, new_value]) - - setattr(self, attr, new_value) - + contrast = spectrum / self.phot + else: + contrast = jnp.zeros((self.n_lambda,)) + spectrum = contrast * self.phot + + for attr, value in (("lon", lon), ("lat", lat), ("rad", rad), + ("temperature", jnp.nan if temperature is None + else temperature)): + setattr(self, attr, jnp.concatenate( + [jnp.ravel(getattr(self, attr)), jnp.asarray(value).reshape(1)])) + self.spectrum = jnp.concatenate([self.spectrum, spectrum[None, :]], axis=0) + self.contrast = jnp.concatenate([self.contrast, contrast[None, :]], axis=0) + self.n_spots = int(self.rad.size) + + @staticmethod @jit - def _blackbody(self, wavelength_meters, temperature): + def _blackbody(wavelength_meters, temperature): """ Compute a blackbody spectrum. """ @@ -311,15 +485,15 @@ def limb_darkening(self, mu, u1, u2): Compute quadratic limb darkening as a function of :math:`\\mu`. """ return ( - 1 / np.pi * + 1 / jnp.pi * (1 - u1 * (1 - mu) - u2 * (1 - mu) ** 2) / (1 - u1 / 3 - u2 / 6) ) @jit def transit_model(self, t0, period, rp, a, inclination, - omega=np.pi / 2, ecc=0, f0=1, t0_rot=0, - u1=0, u2=0): + omega=jnp.pi/2, ecc=0., f0=1., t0_rot=0., + u1=0., u2=0.): """ Compute spectrophotometry with rotation and a planetary transit. @@ -332,8 +506,8 @@ def transit_model(self, t0, period, rp, a, inclination, Mid-transit time period : float Orbital period of the transiting planet - rp : float - Exoplanet radius in units of stellar radii + rp : float or array + Exoplanet radius in units of stellar radii for each wavelength a : float Planetary semi-major axis in units of stellar radii inclination : float @@ -346,10 +520,10 @@ def transit_model(self, t0, period, rp, a, inclination, Out-of-transit flux for an unspotted star, default is one. t0_rot : float Zero-point in time for stellar rotation, default is zero - u1 : float - Limb-darkening parameter :math:`u_1` - u2 : float - Limb-darkening parameter :math:`u_2` + u1 : float or array + Limb-darkening parameter :math:`u_1` for each wavelength + u2 : float or array + Limb-darkening parameter :math:`u_2` for each wavelength Returns ------- @@ -359,14 +533,25 @@ def transit_model(self, t0, period, rp, a, inclination, The apparent squared ratio of planet-to-star radius with stellar spectral contamination by active regions X : array - x-position of the planet in the observer oriented coordinate system [1]_. + x-position of the planet in the observer oriented coordinate + system [1]_. Y : array - y-position of the planet in the observer oriented coordinate system [1]_. + y-position of the planet in the observer oriented coordinate + system [1]_. References ---------- .. [1] Fabrycky & Winn (2009) https://arxiv.org/abs/0902.0737 """ + # Photosphere: ensure a 1D array of length n_lambda + n_lambda = int(self.n_lambda) + phot = ( + jnp.asarray(self.phot) + if self.phot is not None + else jnp.ones((n_lambda,)) + ) + + # Limb-darkening coefficients u1 = jnp.atleast_1d(u1) u2 = jnp.atleast_1d(u2) u_ld = jnp.column_stack([u1, u2]) @@ -378,16 +563,16 @@ def transit_model(self, t0, period, rp, a, inclination, ) = self.spot_coords(t0_rot=t0_rot) rsq = spot_position_x ** 2 + spot_position_y ** 2 - mu = jnp.sqrt(1 - rsq) + mu = jnp.sqrt(1. - rsq) mask_behind_star = jnp.where( - spot_position_z < 0, mu, 0 + spot_position_z < 0., mu, 0. ) - radial_coord = 1 - jnp.geomspace(1e-5, 1, 100)[::-1] + radial_coord = 1. - jnp.geomspace(1e-5, 1., 100)[::-1] unspotted_total_flux = trapezoid( y=( - 2 * np.pi * radial_coord[:, None] * + 2. * jnp.pi * radial_coord[:, None] * self.limb_darkening( radial_coord[:, None], *u_ld.T ) @@ -403,48 +588,84 @@ def transit_model(self, t0, period, rp, a, inclination, # Morris 2020 Eqn 6-7 out_of_transit = f0 - jnp.sum( - np.pi * rad ** 2 * - (1 - contrast) * + jnp.pi * rad ** 2 * + (1. - contrast) * limb_dark * mask_behind_star / unspotted_total_flux[None, None, :, None], axis=1 ) - f_S = rad ** 2 * mu * (spot_position_z < 0).astype(int) + f_S = rad ** 2 * mu * (spot_position_z < 0.).astype(int) - # compute the transit model - mean_anomaly = 2 * np.pi * (self.times - t0) / period - true_anomaly = jnp.arctan2( - *jaxoplanet.core.kepler(M=mean_anomaly, ecc=ecc) - ) + # ------------------------------------------------------------------ + # Orbital geometry (independent of spots) + # ------------------------------------------------------------------ + ecc = jnp.asarray(ecc) + omega = jnp.asarray(omega) - # Winn 2011 Eqn 1 - r = a * (1 - ecc ** 2) / (1 + ecc * jnp.cos(true_anomaly)) + # Index of time sample closest to mid-transit + t_ind = jnp.argmin(jnp.abs(self.times - t0)) - # Winn 2011 Eqn 3-4 - X = -r * jnp.cos(omega + true_anomaly) - Y = -r * jnp.sin(omega + true_anomaly) * jnp.cos(inclination) + # True anomaly at (inferior) conjunction in this geometry + f_tr = 0.5 * jnp.pi - omega - photosphere = (1 - f_S[..., 0].sum(axis=1)) * self.phot[None, :] + # Convert f_tr -> E_tr -> M_tr: + # tan(f/2) = sqrt((1+e)/(1-e)) * tan(E/2) + # M = E - e sin(E) + sqrt_factor = jnp.sqrt((1.0 - ecc) / (1.0 + ecc + 1e-15)) + E_tr = 2.0 * jnp.arctan(jnp.tan(0.5 * f_tr) * sqrt_factor) + M_tr = E_tr - ecc * jnp.sin(E_tr) - spot_coverages, spot_spectra = jnp.broadcast_arrays( - f_S[..., 0], self.spectrum[None, ...] + # Mean anomaly as a function of time with t0 at mid-transit + mean_anomaly = ( + 2.0 * jnp.pi * (self.times - t0) / period + M_tr ) - - time_series_spectrum = jnp.squeeze( - # photospheric component: - photosphere + - - # sum of the active region components: - jnp.sum(spot_coverages * spot_spectra, axis=1) + true_anomaly = jnp.arctan2( + *jaxoplanet.core.kepler(M=mean_anomaly, ecc=ecc) ) - rp = jnp.broadcast_to(rp, self.wavelength.shape) + # Winn 2011 Eqn 1: sky-projected star-planet distance + r = a * (1. - ecc ** 2) / (1. + ecc * jnp.cos(true_anomaly)) - # if one pair of limb-darkening coefficients are given, - # broadcast up to the shape of `rp` + # Winn 2011 Eqn 3-4: sky-projected coordinates + X = -r * jnp.cos(omega + true_anomaly) + Y = -r * jnp.sin(omega + true_anomaly) * jnp.cos(inclination) + + # -------------------------------------------------------------- + # Enforce that the transit chord lies below the stellar equator + # in the (X, Y) plane used for spot occultations: + # + # Choose the branch with Y(t0) <= 0 by reflecting across the + # equator if necessary. This does NOT change b(t) or the + # light curve because b = hypot(X, Y) is sign-invariant. + # -------------------------------------------------------------- + Y_mid = Y[t_ind] + sign_Y = jnp.where(Y_mid > 0., -1., 1.) + Y = sign_Y * Y + + # Total fractional coverage by *any* spot + coverage_sum = f_S[..., 0].sum(axis=1) + # (n_times, 1) + + # Spot-weighted contrast + spot_weighted_contrast = jnp.sum( + f_S * contrast, axis=1 + )[..., 0] + # (n_times, n_lambda) + + flux_factor = (1.0 - coverage_sum) + spot_weighted_contrast + time_series_spectrum = phot[None, :] * flux_factor + # (n_times, n_lambda) + + rp = jnp.broadcast_to(rp, (n_lambda,)) + # (n_lambda,) + + # Broadcast limb-darkening if needed u_ld = u_ld * jnp.ones((rp.shape[0], 1)) + # (n_lambda, 2) + + # Uncontaminated transit model transit = vmap( lambda rp, u: jaxoplanet.core.light_curve( u=u, b=jnp.hypot(X, Y), r=rp @@ -452,10 +673,9 @@ def transit_model(self, t0, period, rp, a, inclination, )(rp, u_ld) contaminated_transit = ( - time_series_spectrum - jnp.abs(transit) * self.phot[None, :] + time_series_spectrum - jnp.abs(transit) * phot[None, :] ) / time_series_spectrum - t_ind = jnp.argmin(jnp.abs(self.times - t0)) uncontaminated_max_depth = - transit[t_ind] contaminated_max_depth = ( contaminated_transit.max(0) - contaminated_transit[t_ind] @@ -464,13 +684,14 @@ def transit_model(self, t0, period, rp, a, inclination, depth_ratio = contaminated_max_depth / uncontaminated_max_depth apparent_rprs2 = rp ** 2 * depth_ratio + # Planet–spot distance in the (X, Y) plane planet_spot_distance = jnp.hypot( spot_position_y - X[:, None, None, None], - spot_position_x - Y[:, None, None, None] + spot_position_x - Y[:, None, None, None], ) occultation_possible = jnp.squeeze( (planet_spot_distance < (major_axis + rp.mean())) & - (spot_position_z < 0) + (spot_position_z < 0.) ) @jit @@ -493,27 +714,30 @@ def time_step( radius=rp, occultation_possible=occultation_possible[j], ), - lambda *args: jnp.zeros((spot_position_x.shape[1], self.n_mc), dtype=bool), + lambda *args: jnp.zeros((spot_position_x.shape[1], self.n_mc), + dtype=bool), ) occultation_per_time_per_spot_per_mc_sample = lax.scan( time_step, 0.0, jnp.arange(self.times.shape[0]) - )[1] # shape: (n_times, n_spots, n_mc_samples) + )[1] + # (n_times, n_spots, n_mc_samples) frac_occulted_per_time_per_spot = jnp.count_nonzero( occultation_per_time_per_spot_per_mc_sample, axis=2 ) / self.n_mc occultation = ( - (1 - contrast) * + (1. - contrast) * jnp.expand_dims(frac_occulted_per_time_per_spot, axis=(2, 3)) ) - scaled_occultation = (1 - contaminated_transit) * jnp.sum(occultation, axis=1)[..., 0] + scaled_occultation = ((1. - contaminated_transit) * + jnp.sum(occultation, axis=1)[..., 0]) spectrum_at_transit = time_series_spectrum[t_ind] return ( - out_of_transit[..., 0] * (contaminated_transit + scaled_occultation), + out_of_transit[..., 0]*(contaminated_transit+scaled_occultation), apparent_rprs2, X, Y, spectrum_at_transit ) @@ -525,9 +749,11 @@ def _area_union_per_time( ): # Monte Carlo sampling for points inside the planet's disk: key, subkey = random.split(self.key) - theta_p = random.uniform(key, minval=0, maxval=2 * np.pi, shape=(self.n_mc,)) + theta_p = random.uniform(key, minval=0, maxval=2 * jnp.pi, + shape=(self.n_mc,)) key, subkey = random.split(key) - rad_p = radius * random.uniform(subkey, minval=0, maxval=1, shape=(self.n_mc,)) ** 0.5 + rad_p = radius * random.uniform(subkey, minval=0, maxval=1, + shape=(self.n_mc,)) ** 0.5 xp = rad_p * jnp.cos(theta_p) + x0_circle yp = rad_p * jnp.sin(theta_p) + y0_circle @@ -541,14 +767,26 @@ def _area_union_per_time( beta = jnp.atleast_1d(beta) angle = jnp.atleast_1d(angle) + # If there are no active regions (no spots), return an empty + # (0, n_mc) array and skip the per-spot scan entirely. This + # prevents index errors when `occultation_possible` is empty. + if ( + x0_ellipse.size == 0 or + getattr(occultation_possible, "size", 0) == 0 + ): + return jnp.zeros((0, self.n_mc), dtype=bool) + @jit def find_overlap(k): - # find overlap between the planet and the elliptical region (projected circular spot) + # find overlap between the planet and the elliptical region + # (projected circular spot) in_ellipse = jnp.hypot( ((xp - x0_ellipse[k]) * jnp.cos(jnp.radians(angle[k])) + - (yp - y0_ellipse[k]) * jnp.sin(jnp.radians(angle[k]))) / alpha[k], + (yp - y0_ellipse[k]) * jnp.sin(jnp.radians(angle[k])) + ) / alpha[k], ((xp - x0_ellipse[k]) * jnp.sin(jnp.radians(angle[k])) - - (yp - y0_ellipse[k]) * jnp.cos(jnp.radians(angle[k]))) / beta[k] + (yp - y0_ellipse[k]) * jnp.cos(jnp.radians(angle[k])) + ) / beta[k] ) < 1 return in_ellipse & on_star @@ -564,12 +802,13 @@ def spot_step(carry, k, occultation_possible=occultation_possible): lambda *args: jnp.zeros(self.n_mc, dtype=bool) ) - monte_carlo_occulted_area = lax.scan(spot_step, 0, jnp.arange(x0_ellipse.shape[0]))[1] + monte_carlo_occulted_area = lax.scan( + spot_step, 0, jnp.arange(x0_ellipse.shape[0]))[1] return monte_carlo_occulted_area def plot_star(self, t0, rp, a, inclination, - ecc=0, t0_rot=0, multiply_radii=1, + ecc=0, omega=jnp.pi/2, t0_rot=0, multiply_radii=1, ax=None, annotate=False): """ Plot a 2D representation of the star and transit chord. @@ -586,11 +825,13 @@ def plot_star(self, t0, rp, a, inclination, Planetary orbital inclination [radians] ecc : float Orbital eccentricity, default is zero. + omega : float + Argument of periapse [radians], default is :math:`\\pi/2`. t0_rot : float Zero-point in time for stellar rotation, default is zero multiply_radii : float - Visually represent scaled-up active regions where the radii are increased - by factor ``multiply_radii``, default is one. + Visually represent scaled-up active regions where the radii are + increased by factor ``multiply_radii``, default is one. ax : matplotlib.axes.Axes Add the visualization to this matplotlib axis annotate : bool @@ -601,24 +842,68 @@ def plot_star(self, t0, rp, a, inclination, if ax is None: ax = plt.gca() - log_temps = np.log10(self.temperature) + temperatures = np.asarray(self.temperature) + known_temperatures = np.isfinite(temperatures) & (temperatures > 0) + color_temperatures = temperatures[known_temperatures] + if self.T_eff is not None: + color_temperatures = np.append(color_temperatures, self.T_eff) + log_temps = np.log10(color_temperatures) + log_min = log_temps.min() if log_temps.size else 0. + log_max = log_temps.max() if log_temps.size else 0. def temp_cmap(x): - return to_hex( - plt.cm.YlOrRd_r( - (np.log10(x) - min(log_temps)) / - (max(log_temps) - min(log_temps)) * 0.6 + 0.4 - ) - ) + fraction = ((np.log10(x) - log_min) / (log_max - log_min) + if log_max > log_min else 1.) + return to_hex(plt.cm.YlOrRd_r(fraction * 0.6 + 0.4)) - star = plt.Circle((0, 0), 1, color=to_hex(temp_cmap(self.T_eff))) + if self.T_eff is None: + color = 'white' + else: + color = to_hex(temp_cmap(self.T_eff)) + + star = plt.Circle((0, 0), 1, facecolor=color, edgecolor='k') ax.add_patch(star) ax.set(xlim=[-1.05, 1.05], ylim=[-1.05, 1.05]) - squeezed_coords = list(map( - jnp.squeeze, self.spot_coords(times=jnp.array([t0]), t0_rot=t0_rot) - )) - for i, (x, y, z, _, _, _, _, angle) in enumerate(zip(*squeezed_coords)): + ( + spot_position_x, + spot_position_y, + spot_position_z, + _, + _, + angle, + _, + contrast, + ) = self.spot_coords(times=jnp.array([t0, ]), t0_rot=t0_rot) + + # Slice out that time and drop the singleton trailing dims. + # After this, each has shape (n_spots,). + spot_x = spot_position_x[0, :, 0, 0] + spot_y = spot_position_y[0, :, 0, 0] + spot_z = spot_position_z[0, :, 0, 0] + # contrast is usually (1, n_spots, n_lambda, 1) + # for plotting, pick λ = 0 + contr = contrast[0, :, 0, 0] + + # Convert to NumPy for plotting (avoids JAX tracing issues) + spot_x_np = np.asarray(spot_x) + spot_y_np = np.asarray(spot_y) + spot_z_np = np.asarray(spot_z) + contr_np = np.asarray(contr) + + for i in range(self.n_spots): + x = spot_x_np[i] + y = spot_y_np[i] + z = spot_z_np[i] + c = contr_np[i] + + if not known_temperatures[i]: + color = 'k' if c <= 1 else 'white' + alpha = np.clip(abs(1-c), 0, 1) + else: + color = temp_cmap(temperatures[i]) + alpha = 1 + if z < 0: rsq = x ** 2 + y ** 2 @@ -626,12 +911,17 @@ def temp_cmap(x): angle = -np.degrees(np.arctan2(y, x)) ell = Ellipse( (y, x), width=multiply_radii * 2 * self.rad[i], - height=multiply_radii * 2 * self.rad[i] * short, angle=angle, - facecolor=temp_cmap(self.temperature[i]), edgecolor='k' + height=multiply_radii * 2 * self.rad[i] * short, + angle=angle, facecolor=color, alpha=alpha, edgecolor='k' ) ax.add_patch(ell) - if annotate: + if annotate and not known_temperatures[i]: + ax.annotate( + f"{i+1}: {c:.2f}", (y, x), + va='center', ha='center', fontsize=6 + ) + elif annotate: ax.annotate( f"{i+1}: {int(self.temperature[i])} K", (y, x), va='center', ha='center', fontsize=6 @@ -640,7 +930,7 @@ def temp_cmap(x): ax.set_aspect('equal') b = (a * np.cos(inclination) * (1 - ecc ** 2) / - (1 + ecc * np.sin(np.pi / 2))) + (1 + ecc * np.sin(omega))) if hasattr(rp, '__len__'): rp = rp.mean() @@ -702,20 +992,27 @@ def bin_spectrum(spectrum, bins=None, log=True, min=None, max=None, **kwargs): Parameters ---------- spectrum : `specutils.Spectrum1D` + Spectrum to be binned + bins : int or ~numpy.ndarray + Number of bins, or the bin edges log : bool If true, compute bin edges based on the log base 10 of the frequency. - bins : int or ~numpy.ndarray - Number of bins, or the bin edges + min : ~astropy.units.Quantity + Minimum wavelength to include in the binned spectrum. + max : ~astropy.units.Quantity + Maximum wavelength to include in the binned spectrum. + **kwargs : dict + Additional keyword arguments passed to `scipy.stats.binned_statistic` Returns ------- - new_spectrum : + new_spectrum : `specutils.Spectrum1D` + Binned spectrum """ - nirspec_wl_range = (spectrum.wavelength > min) & (spectrum.wavelength < max) - - wavelength = spectrum.wavelength[nirspec_wl_range] - flux = spectrum.flux[nirspec_wl_range] + in_range = (spectrum.wavelength > min) & (spectrum.wavelength < max) + wavelength = spectrum.wavelength[in_range] + flux = spectrum.flux[in_range] if log: wl_axis = np.log10(wavelength.to(u.um).value) @@ -741,20 +1038,33 @@ def bin_spectrum(spectrum, bins=None, log=True, min=None, max=None, **kwargs): nans = np.isnan(bs.statistic) interp_fluxes = bs.statistic.copy() if np.any(nans) and all( - map(lambda x: len(x) > 0, [wl_bins[nans], wl_bins[~nans], bs.statistic[~nans]]) + map(lambda x: len(x) > 0, [wl_bins[nans], wl_bins[~nans], + bs.statistic[~nans]]) ): - interp_fluxes[nans] = np.interp(wl_bins[nans], wl_bins[~nans], bs.statistic[~nans]) + interp_fluxes[nans] = np.interp(wl_bins[nans], wl_bins[~nans], + bs.statistic[~nans]) return Spectrum1D( - flux=interp_fluxes * flux.unit, spectral_axis=wl_bins, meta=spectrum.meta + flux=interp_fluxes * flux.unit, spectral_axis=wl_bins, + meta=spectrum.meta ) def spectral_binning(y, all_x, all_y): """ Spectral binning via trapezoidal approximation. + + Parameters + ---------- + y : array + Values to be binned + all_x : array + Full x-axis of the spectrum + all_y : array + Full y-axis of the spectrum """ min_ind = np.argwhere(all_y == y[0])[0, 0] max_ind = np.argwhere(all_y == y[-1])[0, 0] if max_ind > min_ind and y.shape == all_x[min_ind:max_ind + 1].shape: - return np.trapz(y, all_x[min_ind:max_ind + 1]) / (all_x[max_ind] - all_x[min_ind]) + return (np.trapz(y, all_x[min_ind:max_ind + 1]) / + (all_x[max_ind] - all_x[min_ind])) return np.nan diff --git a/fleck/tests/test_jax.py b/fleck/tests/test_jax.py index ded831c..d6629ea 100644 --- a/fleck/tests/test_jax.py +++ b/fleck/tests/test_jax.py @@ -37,12 +37,12 @@ def test_rotation(): ) active_star.lon = jnp.radians(jnp.array([lon1, lon2])) - active_star.lat = jnp.radians(jnp.array([lat1, lat2])) - np.pi / 2 + active_star.lat = jnp.radians(jnp.array([lat1, lat2])) active_star.rad = jnp.array([rad1, rad2]) active_star.spectrum = jnp.array([0.7, 0.7]) fleck_lc = active_star.rotation_model( - f0=1, t0_rot=0.5, u1=u_ld[0], u2=u_ld[1] + f0=1, t0_rot=0, u1=u_ld[0], u2=u_ld[1] )[0][:, 0, 0] # Assert matches STSP results to within 50 ppm: @@ -80,7 +80,7 @@ def test_jax_transit(): active_star = ActiveStar( times=times, lon=np.ravel(spot_lons.to_value(u.rad)), - lat=np.ravel(np.pi/2 + spot_lats.to_value(u.rad)), + lat=np.ravel(spot_lats.to_value(u.rad)), rad=np.ravel(spot_radii), inclination=inc_stellar.to_value(u.rad), phot=jnp.ones(10), @@ -99,3 +99,133 @@ def test_jax_transit(): assert np.max(np.abs(stsp_lc - jax_lc[:, 0]) / jax_lc[:, 0]) < 5.0e-4 assert np.std(stsp_lc / jax_lc[:, 0] - 1) < 1.5e-4 + + +@pytest.fixture +def star_inputs(): + return dict(times=jnp.linspace(-0.05, 0.05, 11), lon=[0.], lat=[0.], + rad=[0.1], inclination=np.pi/2, wavelength=[1e-6], + phot=[1.], P_rot=10.) + + +@pytest.mark.parametrize('contrast', [0.5, [0.5], np.array([0.5]), + jnp.array(0.5), jnp.array([0.5])]) +def test_contrast_inputs_produce_half_dark_spot(star_inputs, contrast): + star = ActiveStar(**star_inputs, contrast=contrast) + np.testing.assert_allclose(star.contrast, [[0.5]]) + flux = np.asarray(star.rotation_model(f0=1)[0]) + # At disk center a spot removes its area times its missing flux fraction. + assert flux[len(flux)//2, 0, 0] == pytest.approx(1 - 0.1**2*0.5, + abs=1e-6) + + +@pytest.mark.parametrize('contrast, expected', [ + (0.5, [[0.5, 0.5, 0.5], [0.5, 0.5, 0.5]]), + ([0.4, 0.7], [[0.4, 0.4, 0.4], [0.7, 0.7, 0.7]]), + ([0.4, 0.5, 0.7], [[0.4, 0.5, 0.7], [0.4, 0.5, 0.7]]), +]) +def test_contrast_broadcasts_over_spots_and_wavelengths(contrast, expected): + star = ActiveStar(times=[0.], lon=[0., 0.2], lat=[0., 0.1], + rad=[0.1, 0.05], inclination=np.pi/2, + wavelength=[1e-6, 2e-6, 3e-6], phot=[2., 3., 4.], + contrast=contrast) + np.testing.assert_allclose(star.contrast, expected) + np.testing.assert_allclose(star.spectrum, np.asarray(expected)*[2., 3., 4.]) + assert star.rotation_model(f0=1)[0].shape == (1, 3, 1) + + +def test_add_first_spot_and_mix_temperature_and_contrast(): + star = ActiveStar(times=jnp.linspace(-0.05, 0.05, 11), + inclination=np.pi/2, wavelength=[1e-6, 2e-6], + T_eff=5000.) + original_phot = np.asarray(star.phot).copy() + star.add_spot(lon=0., lat=0., rad=0.1, contrast=0.5) + star.add_spot(lon=0.2, lat=-0.1, rad=0.05, temperature=4000.) + star.add_spot(lon=-0.2, lat=0.1, rad=0.05, spectrum=0.8*original_phot) + + assert star.n_spots == 3 + np.testing.assert_allclose(star.lon, [0., 0.2, -0.2]) + np.testing.assert_allclose(star.lat, [0., -0.1, 0.1]) + np.testing.assert_allclose(star.rad, [0.1, 0.05, 0.05]) + np.testing.assert_allclose(star.contrast[[0, 2], :], [[0.5, 0.5], [0.8, 0.8]]) + np.testing.assert_allclose(star.phot, original_phot) + assert star.spectrum.shape == (3, 2) + assert star.contrast.shape == (3, 2) + assert np.isnan(np.asarray(star.temperature)[[0, 2]]).all() + assert star.temperature[1] == 4000. + assert np.isfinite(star.rotation_model(f0=1)[0]).all() + + +@pytest.mark.parametrize('construct_spot', [True, False]) +def test_temperature_spot_uses_blackbody_photosphere(construct_spot): + wavelength = np.array([1e-6, 2e-6]) + kwargs = dict(times=[0.], inclination=np.pi/2, wavelength=wavelength, + T_eff=5000.) + if construct_spot: + star = ActiveStar(**kwargs, lon=[0.], lat=[0.], rad=[0.1], + temperature=[4000.]) + else: + star = ActiveStar(**kwargs) + star.add_spot(lon=0., lat=0., rad=0.1, temperature=4000.) + hc_over_k = 6.62607015e-34*299792458./1.380649e-23 + expected = (np.expm1(hc_over_k/(wavelength*5000.)) / + np.expm1(hc_over_k/(wavelength*4000.))) + np.testing.assert_allclose(star.contrast[0], expected, rtol=1e-5) + assert np.all((star.contrast > 0) & (star.contrast < 1)) + + +def test_explicit_photosphere_is_preserved(star_inputs): + star = ActiveStar(**star_inputs, T_eff=5000., contrast=0.5) + np.testing.assert_array_equal(star.phot, [1.]) + np.testing.assert_allclose(star.spectrum, [[0.5]]) + + +@pytest.mark.parametrize('temperature', [None, 5000.]) +@pytest.mark.parametrize('contrast', [0.5, 1.5]) +def test_plot_contrast_only_spot_without_temperature(star_inputs, temperature, + contrast): + import matplotlib.pyplot as plt + + star = ActiveStar(**star_inputs, contrast=contrast, T_eff=temperature) + fig, ax = plt.subplots() + try: + star.plot_star(0., 0.1, 10., np.deg2rad(89), ax=ax, annotate=True) + fig.canvas.draw() + assert len(ax.patches) == 2 + assert ax.patches[1].get_alpha() == pytest.approx(0.5) + assert ax.texts[0].get_text() == f'1: {contrast:.2f}' + finally: + plt.close(fig) + + +@pytest.mark.parametrize('temperatures', [[4000., 5000.], [5000., 5000.], + [4000., 6000.]]) +def test_plot_temperature_colors_follow_temperature(temperatures): + import matplotlib.pyplot as plt + + star = ActiveStar(times=[0.], lon=[-0.2, 0.2], lat=[0., 0.], + rad=[0.1, 0.1], inclination=np.pi/2, + wavelength=[1e-6], phot=[1.], contrast=0.5, + T_eff=5000., temperature=temperatures) + fig, ax = plt.subplots() + try: + star.plot_star(0., 0.1, 10., np.deg2rad(89), ax=ax, annotate=True) + fig.canvas.draw() + colors = np.array([patch.get_facecolor() for patch in ax.patches]) + assert np.isfinite(colors).all() + if temperatures[0] == temperatures[1]: + np.testing.assert_array_equal(colors[0], colors[1]) + else: + assert not np.array_equal(colors[0], colors[1]) + assert not np.array_equal(colors[1], colors[2]) + assert len(ax.texts) == 2 + finally: + plt.close(fig) + + +def test_contrast_is_differentiable(star_inputs): + def center_flux(contrast): + star = ActiveStar(**star_inputs, contrast=contrast) + return star.rotation_model(f0=1)[0][5, 0, 0] + + assert float(jax.grad(center_flux)(0.5)) == pytest.approx(0.1**2, abs=1e-6)