"""
Module for simulating the Rossiter-McLaughlin effect. This is experimental for now.
"""
import matplotlib.pyplot as plt
import numpy as np
from astropy import constants as const
from numba import njit
from tqdm import tqdm
const_c = const.c.to("km/s").value
# @njit
[docs]
def make_grid(n_r, n_theta):
# well we don't want the edges of the grid, we want the midpoint.
dr = 1 / n_r
r_range = np.linspace(dr / 2, 1, n_r) # in stellar radius units
theta_range = np.linspace(0, 2 * np.pi, n_theta) # in radians
# fine just make it uniform theta grid and weight later on by dA.
r_vals = np.tile(r_range, n_theta)
theta_vals = np.repeat(theta_range, n_r)
# define a new theta at each radius
grid = np.zeros((n_theta * n_r, 2)) # not quite sure this is correct. return!
grid[:, 0] = r_vals
grid[:, 1] = theta_vals
return grid
# @njit
[docs]
def doppler_shift_grid(grid, flux, wavelengths, v_rot):
"""
the grid is the (r, theta) values. takes in a 1D spectrum.
v_rot is in meters per second.
outputs the spectrum grid. each row is a different point on the grid. each column is a different wavelength.
"""
# yeah this is where things get funky lol
spectrum_grid = np.zeros((grid.shape[0], len(wavelengths)))
# for i, point in tqdm(enumerate(grid), desc='doppler shifting', total=grid.shape[0]):
for i, point in enumerate(grid):
theta, r = point
x = r * np.cos(theta)
# calculate the v_dopp at that point assuming rotational broadening
v_dopp = (
x * v_rot
) # x is normalized from 0 to 1. this is now in m/s. TODO: cjheck this
delta_lambda = (
wavelengths * v_dopp / const_c
) # this only works if wavelengths in meters. TODO: check this
# interpolate the spectrum onto the new wavelength grid
spectrum_grid[i] = np.interp(wavelengths + delta_lambda, wavelengths, flux)
# todo: make the array math line up. and be less confusing.
# doppler shift based on the wavel
return spectrum_grid
# @njit
[docs]
def calc_planet_locations(phases, r_star, inc, lambda_misalign, a):
"""
calculate the x and y positions of the planet on the stellar disk.
the phase is at 0.0 when the planet is in front of the star. 0.5 when it's behind the star.
phases are in units of orbital phase. not radians.
a and rstar just need to be in the same units.
inclination is in radians
lambda_misalign is in radians.
we need
note: theta is a bit all over the place for the center of the star. but that's more or less physical : )
Assumes a circular orbit, for now...
:param phases:
:param r_p:
:param r_star:
:param inc:
:param lambda_misalign:
:param kp:
:return:
"""
# todo: add the spatial extent of the planet. for now, just assume it's a point source.
positions = np.zeros(
(len(phases), 2)
) # r and theta positions! need to know the lambda misalignment though lol
for i, phase in enumerate(phases):
# calculate the x and y positions of the planet on the stellar disk
x = a * np.sin(2 * np.pi * phase) / r_star
y = np.cos(inc) * a * np.cos(2 * np.pi * phase) / r_star
# calculate the r and theta positions
r = np.sqrt(x**2 + y**2)
theta = np.arctan2(y, x)
# now rotate the theta position by the misalignment angle
theta += lambda_misalign
positions[i] = r, theta
return positions
# @njit
[docs]
def occult_grid(occulted_grid, grid, planet_location, r_p):
"""
take the grid and the planet locations and occult the planet. grid is SPATIAL grid.
r_p is given in stellar radii — so, it's in grid units.
:param grid:
:param planet_locations:
:param r_p:
:return:
"""
# calculate the radial distance from the planet
r1 = grid[:, 0]
r2 = planet_location[0]
theta1 = grid[:, 1]
theta2 = planet_location[1]
# calculate the distance between the planet and the grid points
r = np.sqrt(
np.clip(r1**2 + r2**2 - 2 * r1 * r2 * np.cos(theta1 - theta2), 0, None)
)
# find the points where the planet is
planet_points = np.where(r <= r_p)
# set those points to zero
occulted_grid[planet_points] = 0.0
return occulted_grid
# @njit
[docs]
def sum_grid(occulted_grid, areas):
# todo: calculate the areas ahead of time. then just multiply by the spectrum.
"""
sum up the grid. normalize by area.
grid is the r theta situation
:param grid:
:return:
"""
# pdb.set_trace()
# what if i summed a slice
summed_spectrum = np.dot((occulted_grid).T, areas) / np.sum(areas)
# summed_spectrum = np.zeros(spectrum_grid.shape[1])
# # can probably just multiply and broadcaset
# for i, area in enumerate(areas):
# spectrum = spectrum_grid[i]
# summed_spectrum += area * spectrum
# now it's in flux units that have area. it's ok if it's off by a multiplicative factor, I thnk. todo: check this.
return summed_spectrum
[docs]
def calc_areas(grid):
"""
calculate the area of each cell in the grid.
if it's an r, theta grid, then the area is r^2 * pi / (n_r * n_theta)
no, that's not true. it involves dr. and dtheta. but we can just do it in r, theta space.
so, the
:param grid:
:return:
"""
areas = np.zeros(grid.shape[0])
dr = grid[1, 0] - grid[0, 0]
dtheta = np.max(np.diff(grid[:, 1]))
# pdb.set_trace()
for i, point in enumerate(grid):
r, theta = point
r_outer = r + dr / 2
r_inner = r - dr / 2
area = 0.5 * (r_outer**2 - r_inner**2) * dtheta
areas[i] = area
return areas
[docs]
def make_stellar_disk(
flux,
wavelengths,
v_rot,
phases,
r_star,
inc,
lambda_misalign,
a,
r_p,
n_theta=10,
n_r=10,
):
"""
based on the x and y position of the planet on the disk, make the stellar spectrum.
note: doesn't incorporate CLV variations if I just use a phoenix model.
and if i use an empirical spectrum, I'm not blocking out CLV variations.
take the average of all these approaches? lol
v_rot is in km/s
r_p in meters, and same with rstar and a
"""
# step 1: make the stellar grid. it's in polar so we don't have to worry about edge clipping. careful about summing up the area, though.
grid = make_grid(n_r, n_theta)
# step 2: doppler shift onto a bunch of different mu values. ignores CLV variations.
spectrum_grid = doppler_shift_grid(
grid, flux, wavelengths, v_rot
) # spectrum as a function of r and theta
# step 3: where is the planet on the grid?
planet_locations = calc_planet_locations(
phases, r_star, inc, lambda_misalign, a
) # r and theta positions of the planet as a function of exposure
# loop over the next bit: occurs at every phase.
summed_grids = np.zeros((len(phases), len(wavelengths)))
areas = calc_areas(grid)
for i, phase in enumerate(phases):
occulted_grid = np.copy(spectrum_grid)
planet_location = planet_locations[i]
# step 4: occult those spots on the grid. oh. yeah it's a mask. so just multiply by spectrum grid!
occulted_grid = occult_grid(occulted_grid, grid, planet_location, r_p / r_star)
# step 5: add up the rest of the spectra, normalized by area!
summed_grids[i] = sum_grid(occulted_grid, areas)
# clean up memory. these should be pretty big arrays...
# step 6: what's that overcorrection factor?
correction_factor = (
summed_grids / flux
) # this is what you get when you divide flux(in_transit) / flux(out).
# you'll want to divide your correction factor by the data.
# because wait, shouldn't the summed grids basically be on the order of the spectrum?
return summed_grids, correction_factor, areas, occulted_grid
# can do this convolution style assuming the same stellar line profile...
# frame it as an area calculation. total minus the little bit you're
# but are phoenix models already broadened? nope, they're not!